feat(PRO-309): parallelise ai-decomp requests + chat sync with decomp

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,32 +45,70 @@ 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
         self._current_summary = 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

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,13 +65,15 @@ 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:
         if self._panel is 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}")