diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index 75967f4..a04fbf8 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -1,6 +1,8 @@ from logging import Logger -from libbs.decompilers.ida.compat import execute_ui +import ida_funcs +import ida_kernwin +from libbs.decompilers.ida.compat import execute_read, execute_ui from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData from revengai.models.inline_comment import InlineComment @@ -43,10 +45,18 @@ def disable_function_tracking(self) -> None: self._decomp_hooks.unhook() self._decomp_hooks = None + def is_tracking(self) -> bool: + return self._decomp_hooks is not None + + @execute_ui + def ensure_tracking(self) -> None: + self.enable_function_tracking() + @execute_ui def run_dialog(self) -> None: - self._decomp_view = self.factory.ai_decomp(on_closed=self._on_pane_closed) - self._decomp_view.Create(self._decomp_view.TITLE) + if self._decomp_view is None: + self._decomp_view = self.factory.ai_decomp(on_closed=self._on_pane_closed) + self._decomp_view.Create(self._decomp_view.TITLE) def start_decompilation(self, ea: int) -> None: self._current_decomp = None @@ -54,21 +64,51 @@ def start_decompilation(self, ea: int) -> None: self._current_comments = None self._current_func_vaddr = ea - if self._decomp_view: - self._decomp_view.update_view_content( - "Please wait, decompilation in progress..." - ) + self.run_dialog() + + cached = self.ai_decomp_service.peek_decomp(ea) + if self._decomp_view is not None: + if cached is not None and cached.decompilation: + self._current_decomp = cached + self._rerender() + else: + self._decomp_view.update_view_content( + "Please wait, decompilation in progress..." + ) + + self._dispatch_task(ea) + + def prefetch_decompilation(self, ea: int) -> None: + self._dispatch_task(ea) + def follow_function(self, ea: int, prefetch_if_closed: bool = False) -> None: + if self._decomp_view is not None: + if self._screen_function_start() == ea: + self.start_decompilation(ea) + else: + self.prefetch_decompilation(ea) + elif prefetch_if_closed: + self.prefetch_decompilation(ea) + + @execute_read + def _screen_function_start(self) -> int | None: + func = ida_funcs.get_func(ida_kernwin.get_screen_ea()) + return func.start_ea if func else None + + def _dispatch_task(self, ea: int) -> None: self.ai_decomp_service.start_ai_decomp_task( ea=ea, - on_decomp=self._on_decomp_complete, - on_summary=self._on_summary_complete, - on_comments=self._on_comments_complete, + on_decomp=lambda response: self._on_decomp_complete(ea, response), + on_summary=lambda response: self._on_summary_complete(ea, response), + on_comments=lambda response: self._on_comments_complete(ea, response), ) def _on_decomp_complete( - self, response: GenericApiReturn[DecompilationData] + self, ea: int, response: GenericApiReturn[DecompilationData] ) -> None: + if ea != self._current_func_vaddr: + return + if response.success is False: if response.error_message: self.show_error_dialog(message=response.error_message) @@ -95,8 +135,10 @@ def _on_decomp_complete( self._rerender() def _on_summary_complete( - self, response: GenericApiReturn[SummaryData] + self, ea: int, response: GenericApiReturn[SummaryData] ) -> None: + if ea != self._current_func_vaddr: + return if not response.success: self.log.warning(f"AI summary fetch failed: {response.error_message}") return @@ -106,8 +148,10 @@ def _on_summary_complete( self._rerender() def _on_comments_complete( - self, response: GenericApiReturn[CommentsData] + self, ea: int, response: GenericApiReturn[CommentsData] ) -> None: + if ea != self._current_func_vaddr: + return if not response.success: self.log.warning(f"Inline comments fetch failed: {response.error_message}") return diff --git a/reai_toolkit/app/coordinators/chat_coordinator.py b/reai_toolkit/app/coordinators/chat_coordinator.py index f00fea6..c78aea5 100644 --- a/reai_toolkit/app/coordinators/chat_coordinator.py +++ b/reai_toolkit/app/coordinators/chat_coordinator.py @@ -39,6 +39,7 @@ ConversationContextDTO, FunctionRef, ) +from reai_toolkit.hooks.reactive import ChatContextHooks if TYPE_CHECKING: from reai_toolkit.app.coordinators.ai_decomp_coordinator import AiDecompCoordinator @@ -64,6 +65,7 @@ def __init__( self._conversation_id: Optional[str] = None self._last_event_id: Optional[int] = None self._last_context: Optional[ConversationContextDTO] = None + self._context_hooks: Optional[ChatContextHooks] = None @execute_ui def run_dialog(self, prefill_context: bool = False) -> None: @@ -71,6 +73,7 @@ def run_dialog(self, prefill_context: bool = False) -> None: self._panel = self.factory.chat(on_closed=self._on_pane_closed) self._wire_panel(self._panel) self._panel.Create(self._panel.TITLE) + self._enable_context_tracking() self._panel.focus() self._update_context_chip() if prefill_context: @@ -99,9 +102,23 @@ def _wire_panel(self, panel: ChatPanel) -> None: def _on_pane_closed(self) -> None: if self._panel is not None: self._panel.stop_stream_worker() + self._disable_context_tracking() self._panel = None self.log.info("Agent Chat panel closed.") + def _enable_context_tracking(self) -> None: + if self._context_hooks is None: + self._context_hooks = ChatContextHooks(self) + self._context_hooks.hook() + + def _disable_context_tracking(self) -> None: + if self._context_hooks is not None: + self._context_hooks.unhook() + self._context_hooks = None + + def on_screen_function_changed(self, ea: int) -> None: + self._update_context_chip(ea) + def send(self, text: str) -> None: content = (text or "").strip() if not content or self._state.run_status == "running": @@ -228,6 +245,7 @@ def _work() -> None: return refs: list[FunctionRef] = [] applied = False + last_ea: Optional[int] = None for fid in func_ids: ea = func_map.function_map.get(str(fid)) if ea is None: @@ -241,10 +259,13 @@ def _work() -> None: applied = True refs.append(FunctionRef(ea=ea, name=name)) self.jump_to(ea) + last_ea = ea if refs: self._attach_function_links(tool_call_id, refs) if applied: self.refresh_disassembly_view() + if last_ea is not None and self._ai_decomp_coord is not None: + self._ai_decomp_coord.follow_function(last_ea) threading.Thread(target=_work, daemon=True).start() @@ -268,23 +289,16 @@ def _maybe_open_viewer(self, ev) -> None: name = (ev.tool_name or "").lower() if not any(hint in name for hint in AI_DECOMP_TOOL_HINTS): return - ctx_fid = self._last_context.function_id if self._last_context else None - if ctx_fid is None: - return - target_ids = [i for u in (ev.updated or []) if u.type == "function" for i in u.ids] - if target_ids and ctx_fid not in target_ids: - return - ea = self._resolve_ea_for_function(ctx_fid) - if ea is None: - return - self._ai_decomp_coord.start_decompilation(ea) - - @execute_read - def _resolve_ea_for_function(self, function_id: int) -> Optional[int]: + if not self._ai_decomp_coord.is_tracking(): + self._ai_decomp_coord.ensure_tracking() func_map = self.app.netstore_service.get_function_mapping() if func_map is None: - return None - return func_map.function_map.get(str(function_id)) + return + target_ids = [i for u in (ev.updated or []) if u.type == "function" for i in u.ids] + for fid in target_ids: + ea = func_map.function_map.get(str(fid)) + if ea is not None: + self._ai_decomp_coord.prefetch_decompilation(ea) def _render(self) -> None: if self._panel is not None: @@ -317,16 +331,18 @@ def _resolve_context(self) -> ConversationContextDTO: self.log.debug(f"Failed to resolve chat context: {e}") return ConversationContextDTO(analysis_id=analysis_id, function_id=function_id) - def _update_context_chip(self) -> None: + def _update_context_chip(self, ea: Optional[int] = None) -> None: if self._panel is not None: - self._panel.set_context_chip(self._resolve_context_chip_text()) + self._panel.set_context_chip(self._resolve_context_chip_text(ea)) @execute_read - def _resolve_context_chip_text(self) -> str: + def _resolve_context_chip_text(self, ea: Optional[int] = None) -> str: try: analysis_id = self.app.netstore_service.get_analysis_id() - func = ida_funcs.get_func(ida_kernwin.get_screen_ea()) - name = ida_name.get_ea_name(func.start_ea) if func is not None else None + if ea is None: + func = ida_funcs.get_func(ida_kernwin.get_screen_ea()) + ea = func.start_ea if func is not None else None + name = ida_name.get_ea_name(ea) if ea is not None else None bits = [] if name: bits.append(f"fn: {name}") diff --git a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py index 39b0977..5b69b74 100644 --- a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py +++ b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py @@ -27,15 +27,35 @@ class AiDecompService(IThreadService): def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) -> None: super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) - self._on_decomp: Callable[[GenericApiReturn[DecompilationData]], None] | None = None - self._on_summary: Callable[[GenericApiReturn[SummaryData]], None] | None = None - self._on_comments: Callable[[GenericApiReturn[CommentsData]], None] | None = None self._decomp_cache: dict[int, DecompilationData] = {} self._summary_cache: dict[int, SummaryData] = {} + self._comments_cache: dict[int, CommentsData] = {} + self._inflight: dict[int, threading.Event] = {} + self._inflight_lock = threading.Lock() def thread_in_progress(self) -> bool: return self.is_worker_running() + def is_worker_running(self) -> bool: + with self._inflight_lock: + return bool(self._inflight) + + def stop_worker(self) -> None: + with self._inflight_lock: + events = list(self._inflight.values()) + self._inflight.clear() + for evt in events: + evt.set() + self._decomp_cache.clear() + self._summary_cache.clear() + self._comments_cache.clear() + + def peek_decomp(self, ea: int) -> DecompilationData | None: + function_id = self._get_function_id(start_ea=ea) + if function_id is None: + return None + return self._decomp_cache.get(function_id) + def start_ai_decomp_task( self, ea: int, @@ -43,11 +63,46 @@ def start_ai_decomp_task( on_summary: Callable[[GenericApiReturn[SummaryData]], None], on_comments: Callable[[GenericApiReturn[CommentsData]], None], ) -> None: - self.stop_worker() - self._on_decomp = on_decomp - self._on_summary = on_summary - self._on_comments = on_comments - self.start_worker(target=self._begin_ai_decomp_task, args=(ea,)) + function_id = self._get_function_id(start_ea=ea) + if function_id is None: + on_decomp(GenericApiReturn[DecompilationData](success=True, data=None)) + return + + with self._inflight_lock: + if function_id in self._inflight: + return + stop_event = threading.Event() + self._inflight[function_id] = stop_event + + worker = threading.Thread( + target=self._run_task, + args=(function_id, stop_event, on_decomp, on_summary, on_comments), + name=f"reai-aidecomp-{function_id}", + daemon=True, + ) + worker.start() + + def _run_task( + self, + function_id: int, + stop_event: threading.Event, + on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], + on_summary: Callable[[GenericApiReturn[SummaryData]], None], + on_comments: Callable[[GenericApiReturn[CommentsData]], None], + ) -> None: + try: + if stop_event.is_set(): + return + decomp = self._run_decomp_phase(function_id, stop_event, on_decomp) + if decomp is None: + return + self._run_summary_phase(function_id, stop_event, on_summary) + self._run_comments_phase(function_id, stop_event, on_comments) + except Exception as e: + logger.error(f"RevEng.AI: AI decompilation task crashed for {function_id}: {e}") + finally: + with self._inflight_lock: + self._inflight.pop(function_id, None) def _get_function_id(self, start_ea: int) -> int | None: function_map: FunctionMapping | None = self.netstore_service.get_function_mapping() @@ -112,13 +167,16 @@ def _fetch_decompilation( return None, f"Unexpected error fetching AI decompilation: {e}" def _run_decomp_phase( - self, function_id: int, stop_event: threading.Event + self, + function_id: int, + stop_event: threading.Event, + on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], ) -> DecompilationData | None: cached = self._decomp_cache.get(function_id) if cached is not None: self._safe_dispatch( stop_event, - self._on_decomp, + on_decomp, GenericApiReturn[DecompilationData](success=True, data=cached), ) return cached @@ -127,7 +185,7 @@ def _run_decomp_phase( if decomp is None: self._safe_dispatch( stop_event, - self._on_decomp, + on_decomp, GenericApiReturn[DecompilationData](success=False, error_message=fetch_err), ) return None @@ -138,7 +196,7 @@ def _run_decomp_phase( self._decomp_cache[function_id] = decomp self._safe_dispatch( stop_event, - self._on_decomp, + on_decomp, GenericApiReturn[DecompilationData](success=True, data=decomp), ) return decomp @@ -148,7 +206,7 @@ def _run_decomp_phase( if not queued_ok: self._safe_dispatch( stop_event, - self._on_decomp, + on_decomp, GenericApiReturn[DecompilationData]( success=False, error_message=queue_err ), @@ -166,7 +224,7 @@ def _run_decomp_phase( if poll_err is not None: self._safe_dispatch( stop_event, - self._on_decomp, + on_decomp, GenericApiReturn[DecompilationData]( success=False, error_message=poll_err ), @@ -177,7 +235,7 @@ def _run_decomp_phase( if final is None or not final.decompilation: self._safe_dispatch( stop_event, - self._on_decomp, + on_decomp, GenericApiReturn[DecompilationData]( success=False, error_message=final_err or "AI decompilation returned no content.", @@ -188,7 +246,7 @@ def _run_decomp_phase( self._decomp_cache[function_id] = final self._safe_dispatch( stop_event, - self._on_decomp, + on_decomp, GenericApiReturn[DecompilationData](success=True, data=final), ) return final @@ -236,13 +294,16 @@ def _fetch_summary( return None, f"Unexpected error fetching AI decompilation summary: {e}" def _run_summary_phase( - self, function_id: int, stop_event: threading.Event + self, + function_id: int, + stop_event: threading.Event, + on_summary: Callable[[GenericApiReturn[SummaryData]], None], ) -> None: cached = self._summary_cache.get(function_id) if cached is not None: self._safe_dispatch( stop_event, - self._on_summary, + on_summary, GenericApiReturn[SummaryData](success=True, data=cached), ) return @@ -251,7 +312,7 @@ def _run_summary_phase( if summary is None: self._safe_dispatch( stop_event, - self._on_summary, + on_summary, GenericApiReturn[SummaryData](success=False, error_message=err), ) return @@ -262,7 +323,7 @@ def _run_summary_phase( self._summary_cache[function_id] = summary self._safe_dispatch( stop_event, - self._on_summary, + on_summary, GenericApiReturn[SummaryData](success=True, data=summary), ) return @@ -272,7 +333,7 @@ def _run_summary_phase( if not queued_ok: self._safe_dispatch( stop_event, - self._on_summary, + on_summary, GenericApiReturn[SummaryData]( success=False, error_message=queue_err ), @@ -290,7 +351,7 @@ def _run_summary_phase( if poll_err is not None: self._safe_dispatch( stop_event, - self._on_summary, + on_summary, GenericApiReturn[SummaryData]( success=False, error_message=poll_err ), @@ -301,7 +362,7 @@ def _run_summary_phase( if final is None: self._safe_dispatch( stop_event, - self._on_summary, + on_summary, GenericApiReturn[SummaryData]( success=False, error_message=final_err ), @@ -311,7 +372,7 @@ def _run_summary_phase( self._summary_cache[function_id] = final self._safe_dispatch( stop_event, - self._on_summary, + on_summary, GenericApiReturn[SummaryData](success=True, data=final), ) @@ -358,13 +419,25 @@ def _fetch_comments( return None, f"Unexpected error fetching inline comments: {e}" def _run_comments_phase( - self, function_id: int, stop_event: threading.Event + self, + function_id: int, + stop_event: threading.Event, + on_comments: Callable[[GenericApiReturn[CommentsData]], None], ) -> None: + cached = self._comments_cache.get(function_id) + if cached is not None: + self._safe_dispatch( + stop_event, + on_comments, + GenericApiReturn[CommentsData](success=True, data=cached), + ) + return + comments, err = self._fetch_comments(function_id) if comments is None: self._safe_dispatch( stop_event, - self._on_comments, + on_comments, GenericApiReturn[CommentsData](success=False, error_message=err), ) return @@ -372,9 +445,10 @@ def _run_comments_phase( status = str(comments.task_status) if status == TaskStatus.COMPLETED.value: + self._comments_cache[function_id] = comments self._safe_dispatch( stop_event, - self._on_comments, + on_comments, GenericApiReturn[CommentsData](success=True, data=comments), ) return @@ -384,7 +458,7 @@ def _run_comments_phase( if not queued_ok: self._safe_dispatch( stop_event, - self._on_comments, + on_comments, GenericApiReturn[CommentsData]( success=False, error_message=queue_err ), @@ -402,7 +476,7 @@ def _run_comments_phase( if poll_err is not None: self._safe_dispatch( stop_event, - self._on_comments, + on_comments, GenericApiReturn[CommentsData]( success=False, error_message=poll_err ), @@ -413,16 +487,17 @@ def _run_comments_phase( if final is None: self._safe_dispatch( stop_event, - self._on_comments, + on_comments, GenericApiReturn[CommentsData]( success=False, error_message=final_err ), ) return + self._comments_cache[function_id] = final self._safe_dispatch( stop_event, - self._on_comments, + on_comments, GenericApiReturn[CommentsData](success=True, data=final), ) @@ -469,23 +544,6 @@ def _poll_workflow( return False, None - def _begin_ai_decomp_task(self, stop_event: threading.Event, start_ea: int) -> None: - function_id = self._get_function_id(start_ea=start_ea) - if function_id is None: - self._safe_dispatch( - stop_event, - self._on_decomp, - GenericApiReturn[DecompilationData](success=True, data=None), - ) - return - - decomp = self._run_decomp_phase(function_id, stop_event) - if decomp is None: - return - - self._run_summary_phase(function_id, stop_event) - self._run_comments_phase(function_id, stop_event) - def _format_api_error(e: ApiException) -> str: error_response = parse_exception(e) diff --git a/reai_toolkit/hooks/reactive.py b/reai_toolkit/hooks/reactive.py index 25de886..7f76656 100644 --- a/reai_toolkit/hooks/reactive.py +++ b/reai_toolkit/hooks/reactive.py @@ -32,6 +32,40 @@ def _function_rename(self, ea: int, new_name: str): ) +class ChatContextHooks(kw.UI_Hooks): + def __init__(self, coordinator): + super().__init__() + self.coordinator = coordinator + self._last_func_start = None + self._is_hooked = False + + def hook(self) -> bool: + if self._is_hooked: + return False + ok = super().hook() + if ok: + self._is_hooked = True + func = ida_funcs.get_func(kw.get_screen_ea()) + self._last_func_start = func.start_ea if func else None + return ok + + def unhook(self) -> None: + if self._is_hooked: + super().unhook() + self._is_hooked = False + self._last_func_start = None + + def screen_ea_changed(self, ea: int, prev_ea: int) -> None: + func = ida_funcs.get_func(ea) + if not func or func.start_ea == self._last_func_start: + return + self._last_func_start = func.start_ea + try: + self.coordinator.on_screen_function_changed(func.start_ea) + except Exception as e: + logger.error(f"[ChatContextHooks] context update failed: {e}") + + class AiDecompFunctionViewHooks(kw.UI_Hooks): """ Hook that tracks when the screen EA changes (user moves between functions), @@ -64,9 +98,10 @@ def hook(self) -> bool: if ok: # First hook, does not trigger event. Manual call to coordinator. self._is_hooked = True - ea = kw.get_screen_ea() - func = ida_funcs.get_func(ea) - self.coordinator.start_decompilation(ea=func.start_ea) + func = ida_funcs.get_func(kw.get_screen_ea()) + if func is not None: + self._last_func_start = func.start_ea + self.coordinator.start_decompilation(ea=func.start_ea) logger.info("[FunctionViewHooks] Hook registered.") else: logger.error("[FunctionViewHooks] Failed to register.") diff --git a/tests/idalib/ai_decomp/test_smoke.py b/tests/idalib/ai_decomp/test_smoke.py index 13e9c42..d05b7a8 100644 --- a/tests/idalib/ai_decomp/test_smoke.py +++ b/tests/idalib/ai_decomp/test_smoke.py @@ -1,3 +1,4 @@ +import time from unittest.mock import MagicMock import pytest @@ -43,8 +44,10 @@ def test_service_runs_to_completion_under_idalib(loaded_binary, mocker): service.start_ai_decomp_task( ea=ea, on_decomp=on_decomp, on_summary=on_summary, on_comments=on_comments ) - service._worker_thread.join(timeout=10.0) - assert not service._worker_thread.is_alive(), "worker deadlocked under idalib" + deadline = time.monotonic() + 10.0 + while service.is_worker_running() and time.monotonic() < deadline: + time.sleep(0.01) + assert not service.is_worker_running(), "worker deadlocked under idalib" on_decomp.assert_called_once() result = on_decomp.call_args[0][0] diff --git a/tests/unit/ai_decomp/test_service.py b/tests/unit/ai_decomp/test_service.py index 3bfa6dd..3c21276 100644 --- a/tests/unit/ai_decomp/test_service.py +++ b/tests/unit/ai_decomp/test_service.py @@ -69,11 +69,11 @@ def _comments(items=None, status=TaskStatus.COMPLETED.value) -> CommentsData: return CommentsData.model_construct(inline_comments=items or [], task_status=status) -def _wait(thread: threading.Thread | None, timeout: float = 5.0) -> None: - if thread is None: - return - thread.join(timeout=timeout) - assert not thread.is_alive(), "worker thread did not exit" +def _wait(service, timeout: float = 5.0) -> None: + deadline = time.monotonic() + timeout + while service.is_worker_running() and time.monotonic() < deadline: + time.sleep(0.005) + assert not service.is_worker_running(), "worker did not go idle" def _run(service, ea: int = 4096): @@ -84,7 +84,7 @@ def _run(service, ea: int = 4096): on_summary=on_summary, on_comments=on_comments, ) - _wait(service._worker_thread) + _wait(service) return on_decomp, on_summary, on_comments