feat(PLU-311): persist function jumps in chat history, render bug fix

reai_toolkit/app/components/tabs/chat_tab.py

@@ -10,6 +10,8 @@
 
 from __future__ import annotations
 
+import queue
+import sys
 import threading
 from typing import Callable, Optional
 
@@ -26,10 +28,37 @@
 from reai_toolkit.app.services.chat.schema import ChatState
 
 
-class ChatStreamWorker(QtCore.QObject):
-    """Runs one chat turn (optional create → optional send → stream) off-thread."""
+class _StderrSlotWarningFilter:
+    def __init__(self, wrapped) -> None:
+        self._wrapped = wrapped
+        self._local = threading.local()
+
+    def write(self, text):
+        buf = getattr(self._local, "buf", "") + text
+        while "\n" in buf:
+            line, buf = buf.split("\n", 1)
+            if "_StreamRelay::handle_event" not in line:
+                self._wrapped.write(line + "\n")
+        self._local.buf = buf
+        return len(text)
+
+    def flush(self):
+        buf = getattr(self._local, "buf", "")
+        self._local.buf = ""
+        if buf and "_StreamRelay::handle_event" not in buf:
+            self._wrapped.write(buf)
+        self._wrapped.flush()
+
+    def __getattr__(self, name):
+        return getattr(self._wrapped, name)
+
 
-    event_ready = Signal(object)
+if not isinstance(sys.stderr, _StderrSlotWarningFilter):
+    sys.stderr = _StderrSlotWarningFilter(sys.stderr)
+
+
+class ChatStreamWorker(QtCore.QObject):
+    event_ready = Signal()
     conversation_created = Signal(str)
     errored = Signal(str)
     finished = Signal()
@@ -50,6 +79,7 @@ def __init__(
         self._last_event_id = last_event_id
         self._stopped = False
         self._stop_event = threading.Event()
+        self._event_queue: "queue.Queue" = queue.Queue()
 
     @Slot()
     def run(self) -> None:
@@ -78,7 +108,8 @@ def run(self) -> None:
             for ev in self._svc.stream(conv_id, self._stop_event, self._last_event_id):
                 if self._stopped:
                     break
-                self.event_ready.emit(ev)
+                self._event_queue.put(ev)
+                self.event_ready.emit()
         except Exception as e:
             self.errored.emit(str(e))
         finally:
@@ -94,19 +125,29 @@ def stop(self) -> None:
 
 
 class _StreamRelay(QtCore.QObject):
-    """UI-thread QObject that forwards worker signals to plain callbacks."""
-
     def __init__(self) -> None:
         super().__init__()
+        self._event_queue: Optional["queue.Queue"] = None
         self.on_event: Optional[Callable] = None
         self.on_conversation_created: Optional[Callable] = None
         self.on_error: Optional[Callable] = None
         self.on_finished: Optional[Callable] = None
 
-    @Slot(object)
-    def handle_event(self, ev) -> None:
-        if self.on_event:
-            self.on_event(ev)
+    def bind_queue(self, q: "queue.Queue") -> None:
+        self._event_queue = q
+
+    @Slot()
+    def handle_event(self) -> None:
+        q = self._event_queue
+        if q is None:
+            return
+        while True:
+            try:
+                ev = q.get_nowait()
+            except queue.Empty:
+                break
+            if self.on_event:
+                self.on_event(ev)
 
     @Slot(str)
     def handle_conversation_created(self, uuid: str) -> None:
@@ -204,17 +245,19 @@ def OnCreate(self, form) -> None:
         f = self._title_label.font()
         f.setBold(True)
         self._title_label.setFont(f)
-        self._context_label = QtWidgets.QLabel("")
-        self._context_label.setStyleSheet("color: gray;")
         history_btn = QtWidgets.QPushButton("History")
         new_btn = QtWidgets.QPushButton("New chat")
         header.addWidget(self._title_label)
         header.addStretch(1)
-        header.addWidget(self._context_label)
         header.addWidget(history_btn)
         header.addWidget(new_btn)
         root.addLayout(header)
 
+        self._context_label = QtWidgets.QLabel("")
+        self._context_label.setStyleSheet("color: gray;")
+        self._context_label.setWordWrap(True)
+        root.addWidget(self._context_label)
+
         self._history_list = QtWidgets.QListWidget()
         self._history_list.setVisible(False)
         self._history_list.setMaximumHeight(140)
@@ -265,7 +308,7 @@ def OnCreate(self, form) -> None:
 
         self._render_timer = QtCore.QTimer(self._parent_window)
         self._render_timer.setSingleShot(True)
-        self._render_timer.setInterval(40)
+        self._render_timer.setInterval(100)
         self._render_timer.timeout.connect(self._flush_render)
 
         self._relay = _StreamRelay()
@@ -376,6 +419,7 @@ def start_stream_worker(self, worker: ChatStreamWorker) -> None:
         worker.moveToThread(thread)
         thread.started.connect(worker.run)
         if self._relay is not None:
+            self._relay.bind_queue(worker._event_queue)
             worker.event_ready.connect(self._relay.handle_event)
             worker.conversation_created.connect(self._relay.handle_conversation_created)
             worker.errored.connect(self._relay.handle_error)
@@ -429,6 +473,9 @@ def _handle_stream_finished(self) -> None:
         self._streaming = False
         if self.on_stream_finished:
             self.on_stream_finished()
+        if self._render_timer is not None and self._render_timer.isActive():
+            self._render_timer.stop()
+        self._flush_render()
 
     def _on_send_clicked(self) -> None:
         if self._input is None:

reai_toolkit/app/coordinators/chat_coordinator.py

@@ -29,15 +29,13 @@
     ConfirmTool,
     EventAction,
     SendMessage,
-    SetToolFunctions,
     build_initial_state,
     chat_reducer,
     initial_state,
 )
 from reai_toolkit.app.services.chat.schema import (
     ChatState,
     ConversationContextDTO,
-    FunctionRef,
 )
 from reai_toolkit.hooks.reactive import ChatContextHooks
 
@@ -223,27 +221,22 @@ def on_stream_error(self, msg: str) -> None:
         self._render()
 
     def on_stream_finished(self) -> None:
-        return
+        self._render()
 
     def _handle_tool_result(self, ev) -> None:
         if ev.updated:
             func_ids = [i for u in ev.updated if u.type == "function" for i in u.ids]
             if func_ids:
-                self._sync_and_link_functions(ev.tool_call_id, func_ids)
+                self._sync_functions(func_ids)
             elif any(u.type == "analysis" for u in ev.updated):
                 self.refresh_disassembly_view()
         self._maybe_open_viewer(ev)
 
-    def _sync_and_link_functions(self, tool_call_id: str, func_ids: list) -> None:
-        """For each function the agent touched: pull its current name into the
-        IDB (targeted, per-function — no full analysis sync), auto-jump to it as
-        it applies, and attach a clickable jump link to the tool-call."""
-
+    def _sync_functions(self, func_ids: list) -> None:
         def _work() -> None:
             func_map = self.app.netstore_service.get_function_mapping()
             if func_map is None:
                 return
-            refs: list[FunctionRef] = []
             applied = False
             last_ea: Optional[int] = None
             for fid in func_ids:
@@ -257,25 +250,16 @@ def _work() -> None:
                 if self.chat_service.update_function_name(ea, name):
                     self.chat_service.tag_function_as_renamed(name)
                     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()
+                self._refresh_context_chip(last_ea)
             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()
 
-    @execute_ui
-    def _attach_function_links(self, tool_call_id: str, refs: list) -> None:
-        self._state = chat_reducer(
-            self._state, SetToolFunctions(tool_call_id=tool_call_id, functions=refs)
-        )
-        self._render()
-
     @execute_ui
     def jump_to(self, ea: int) -> None:
         try:
@@ -335,6 +319,10 @@ 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(ea))
 
+    @execute_ui
+    def _refresh_context_chip(self, ea: Optional[int] = None) -> None:
+        self._update_context_chip(ea)
+
     @execute_read
     def _resolve_context_chip_text(self, ea: Optional[int] = None) -> str:
         try:
@@ -348,7 +336,7 @@ def _resolve_context_chip_text(self, ea: Optional[int] = None) -> str:
                 bits.append(f"fn: {name}")
             if analysis_id is not None:
                 bits.append(f"analysis #{analysis_id}")
-            return "  ·  ".join(bits) if bits else "No analysis attached"
+            return "\n".join(bits) if bits else "No analysis attached"
         except Exception:
             return ""
 

reai_toolkit/app/services/chat/reducer.py

@@ -17,6 +17,7 @@
     ChatItem,
     ChatState,
     ContextCompacted,
+    FunctionRef,
     RunError,
     Step,
     StoredEvent,
@@ -61,15 +62,7 @@ class Reset:
     stored_events: Optional[list[StoredEvent]] = None
 
 
-@dataclass
-class SetToolFunctions:
-    tool_call_id: str
-    functions: list
-
-
-ChatAction = Union[
-    SendMessage, Cancel, ConfirmTool, EventAction, ApiError, Reset, SetToolFunctions
-]
+ChatAction = Union[SendMessage, Cancel, ConfirmTool, EventAction, ApiError, Reset]
 
 
 def _new_id() -> str:
@@ -152,22 +145,25 @@ def chat_reducer(state: ChatState, action: ChatAction) -> ChatState:
             items=_finalize_running_items(state.items),
         )
 
-    if isinstance(action, SetToolFunctions):
-        return replace(
-            state,
-            items=_find_last_and_update(
-                state.items,
-                lambda it: isinstance(it, ToolCall) and it.id == action.tool_call_id,
-                lambda it: replace(it, functions=action.functions),
-            ),
-        )
-
     if isinstance(action, EventAction):
         return _reduce_event(state, action.event)
 
     return state
 
 
+def _functions_from_updates(updates) -> Optional[list[FunctionRef]]:
+    if not updates:
+        return None
+    out: list[FunctionRef] = []
+    for u in updates:
+        if u.type != "function":
+            continue
+        for r in u.refs:
+            if r.name and r.vaddr:
+                out.append(FunctionRef(ea=r.vaddr, name=r.name))
+    return out or None
+
+
 def _reduce_event(state: ChatState, event: ChatEvent) -> ChatState:
     t = event.type
 
@@ -267,10 +263,16 @@ def _reduce_event(state: ChatState, event: ChatEvent) -> ChatState:
         )
 
     if t == "TOOL_CALL_RESULT":
+        functions = _functions_from_updates(event.updated)
         items = _find_last_and_update(
             state.items,
             lambda it: isinstance(it, ToolCall) and it.id == event.tool_call_id,
-            lambda it: replace(it, status="finished", is_error=event.is_error),
+            lambda it: replace(
+                it,
+                status="finished",
+                is_error=event.is_error,
+                functions=functions or it.functions,
+            ),
         )
         items = _find_last_and_update(
             items,

reai_toolkit/app/services/chat/schema.py

@@ -38,12 +38,20 @@
 ROLE_TOOL = 4
 
 
+@dataclass
+class EntityRef:
+    id: int
+    name: str = ""
+    vaddr: int = 0
+
+
 @dataclass
 class EntityUpdate:
     """A backend entity a tool result reports as changed. Drives view refresh."""
 
     type: str
     ids: list[int]
+    refs: list[EntityRef] = field(default_factory=list)
 
 
 @dataclass
@@ -78,6 +86,26 @@ def resolve_type(type_field: Any) -> Optional[str]:
     return None
 
 
+def _parse_refs(raw: Any) -> list[EntityRef]:
+    if not isinstance(raw, list):
+        return []
+    out: list[EntityRef] = []
+    for item in raw:
+        if not isinstance(item, dict) or "id" not in item:
+            continue
+        try:
+            out.append(
+                EntityRef(
+                    id=int(item["id"]),
+                    name=str(item.get("name") or ""),
+                    vaddr=int(item.get("vaddr") or 0),
+                )
+            )
+        except (TypeError, ValueError):
+            continue
+    return out
+
+
 def _parse_entity_updates(raw: Any) -> Optional[list[EntityUpdate]]:
     if not isinstance(raw, list):
         return None
@@ -92,7 +120,13 @@ def _parse_entity_updates(raw: Any) -> Optional[list[EntityUpdate]]:
         if not isinstance(ids, list):
             continue
         try:
-            out.append(EntityUpdate(type=etype, ids=[int(i) for i in ids]))
+            out.append(
+                EntityUpdate(
+                    type=etype,
+                    ids=[int(i) for i in ids],
+                    refs=_parse_refs(item.get("refs")),
+                )
+            )
         except (TypeError, ValueError):
             continue
     return out or None