diff --git a/reai_toolkit/app/components/tabs/chat_tab.py b/reai_toolkit/app/components/tabs/chat_tab.py index e811065..7b0cc34 100644 --- a/reai_toolkit/app/components/tabs/chat_tab.py +++ b/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: diff --git a/reai_toolkit/app/coordinators/chat_coordinator.py b/reai_toolkit/app/coordinators/chat_coordinator.py index c78aea5..59e510c 100644 --- a/reai_toolkit/app/coordinators/chat_coordinator.py +++ b/reai_toolkit/app/coordinators/chat_coordinator.py @@ -29,7 +29,6 @@ ConfirmTool, EventAction, SendMessage, - SetToolFunctions, build_initial_state, chat_reducer, initial_state, @@ -37,7 +36,6 @@ 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 "" diff --git a/reai_toolkit/app/services/chat/reducer.py b/reai_toolkit/app/services/chat/reducer.py index 9c2f736..f1bf111 100644 --- a/reai_toolkit/app/services/chat/reducer.py +++ b/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, diff --git a/reai_toolkit/app/services/chat/schema.py b/reai_toolkit/app/services/chat/schema.py index 4936c89..ae0ae84 100644 --- a/reai_toolkit/app/services/chat/schema.py +++ b/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 diff --git a/tests/unit/chat/test_reducer.py b/tests/unit/chat/test_reducer.py index 09fe16e..5b037fe 100644 --- a/tests/unit/chat/test_reducer.py +++ b/tests/unit/chat/test_reducer.py @@ -6,7 +6,6 @@ ConfirmTool, EventAction, SendMessage, - SetToolFunctions, build_initial_state, chat_reducer, initial_state, @@ -15,6 +14,8 @@ AssistantMessage, ChatEvent, ContextCompacted, + EntityRef, + EntityUpdate, FunctionRef, Step, ToolCall, @@ -210,17 +211,47 @@ def test_api_error_action_finalizes(): assert state.items[-1].is_streaming is False -def test_set_tool_functions_attaches_links(): +def _rename_result_events(): + return [ + _ev("TOOL_CALL_START", tool_call_id="t1", tool_name="rename_functions"), + _ev( + "TOOL_CALL_RESULT", + tool_call_id="t1", + tool_name="rename_functions", + updated=[ + EntityUpdate( + type="function", + ids=[2015699787], + refs=[ + EntityRef(id=2015699787, name="region_position", vaddr=4198416) + ], + ) + ], + ), + ] + + +def test_tool_call_result_attaches_function_links_from_refs(): + state = _fold(_rename_result_events()) + tool = [i for i in state.items if isinstance(i, ToolCall)][0] + assert tool.functions == [FunctionRef(ea=4198416, name="region_position")] + + +def test_tool_call_result_links_persist_on_replay(): + state = build_initial_state(_rename_result_events()) + tool = [i for i in state.items if isinstance(i, ToolCall)][0] + assert tool.functions == [FunctionRef(ea=4198416, name="region_position")] + + +def test_tool_call_result_without_refs_leaves_no_links(): state = _fold( [ - _ev("TOOL_CALL_START", tool_call_id="t1", tool_name="rename_functions"), - _ev("TOOL_CALL_RESULT", tool_call_id="t1", tool_name="rename_functions"), + _ev("TOOL_CALL_START", tool_call_id="t1", tool_name="do"), + _ev("TOOL_CALL_RESULT", tool_call_id="t1", tool_name="do"), ] ) - refs = [FunctionRef(ea=0x408140, name="chat_agent_renamed")] - state = chat_reducer(state, SetToolFunctions(tool_call_id="t1", functions=refs)) tool = [i for i in state.items if isinstance(i, ToolCall)][0] - assert tool.functions == refs + assert tool.functions is None def test_build_initial_state_replays_user_and_events(): diff --git a/tests/unit/chat/test_render.py b/tests/unit/chat/test_render.py index 00197d9..d82dc83 100644 --- a/tests/unit/chat/test_render.py +++ b/tests/unit/chat/test_render.py @@ -14,10 +14,14 @@ render_transcript_markdown, title_case, ) +from reai_toolkit.app.services.chat.reducer import build_initial_state from reai_toolkit.app.services.chat.schema import ( AssistantMessage, + ChatEvent, ChatState, ContextCompacted, + EntityRef, + EntityUpdate, FunctionRef, RunError, Step, @@ -85,6 +89,30 @@ def test_render_function_jump_links(): assert "[chat_agent_renamed](ida://jump/4227392)" in md +def test_render_jump_links_persist_from_replayed_refs(): + events = [ + ChatEvent( + type="TOOL_CALL_START", tool_call_id="t1", tool_name="rename_functions" + ), + ChatEvent( + type="TOOL_CALL_RESULT", + tool_call_id="t1", + tool_name="rename_functions", + updated=[ + EntityUpdate( + type="function", + ids=[2015699787], + refs=[ + EntityRef(id=2015699787, name="region_position", vaddr=4198416) + ], + ) + ], + ), + ] + md = render_transcript_markdown(build_initial_state(events)) + assert "[region_position](ida://jump/4198416)" in md + + def test_find_pending_confirmation(): state = ChatState( items=[ diff --git a/tests/unit/chat/test_sse_parser.py b/tests/unit/chat/test_sse_parser.py index e1ec4eb..6a4815f 100644 --- a/tests/unit/chat/test_sse_parser.py +++ b/tests/unit/chat/test_sse_parser.py @@ -1,6 +1,6 @@ """SSE frame parser tests — feed fake byte chunks, no network.""" -from reai_toolkit.app.services.chat.schema import normalize_event +from reai_toolkit.app.services.chat.schema import EntityRef, normalize_event from reai_toolkit.app.services.chat.sse import ( event_from_frame, iter_sse_events, @@ -143,3 +143,17 @@ def test_tool_result_entity_updates_parsed(): assert ev.updated is not None assert ev.updated[0].type == "function" assert ev.updated[0].ids == [1, 2] + assert ev.updated[0].refs == [] + + +def test_tool_result_entity_refs_parsed(): + frame = _frame( + '{"type": 12, "data": {"tool_call_id": "t", "tool_name": "rename_functions", ' + '"updated": [{"type": "function", "ids": [2015699787], ' + '"refs": [{"id": 2015699787, "name": "region_position", "vaddr": 4198416}]}]}}' + ) + ev = list(iter_sse_events([frame]))[0] + assert ev.updated is not None + assert ev.updated[0].refs == [ + EntityRef(id=2015699787, name="region_position", vaddr=4198416) + ]