From c1108d1081d4d19038b9762d25b583f6b8244832 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Fri, 3 Jul 2026 16:55:48 +0100 Subject: [PATCH] feat(PRO-298): chat agent base impl --- reai_toolkit/app/app.py | 4 + .../app/components/tabs/chat_render.py | 93 ++++ reai_toolkit/app/components/tabs/chat_tab.py | 465 ++++++++++++++++++ reai_toolkit/app/coordinator.py | 9 + .../app/coordinators/chat_coordinator.py | 372 ++++++++++++++ reai_toolkit/app/factory.py | 5 + reai_toolkit/app/services/chat/__init__.py | 5 + .../app/services/chat/chat_service.py | 274 +++++++++++ reai_toolkit/app/services/chat/reducer.py | 328 ++++++++++++ reai_toolkit/app/services/chat/schema.py | 253 ++++++++++ reai_toolkit/app/services/chat/sse.py | 79 +++ reai_toolkit/hooks/menu.py | 28 ++ reai_toolkit/hooks/popup.py | 38 ++ tests/unit/chat/test_reducer.py | 239 +++++++++ tests/unit/chat/test_render.py | 97 ++++ tests/unit/chat/test_sdk_schemas.py | 53 ++ tests/unit/chat/test_service.py | 161 ++++++ tests/unit/chat/test_sse_parser.py | 145 ++++++ 18 files changed, 2648 insertions(+) create mode 100644 reai_toolkit/app/components/tabs/chat_render.py create mode 100644 reai_toolkit/app/components/tabs/chat_tab.py create mode 100644 reai_toolkit/app/coordinators/chat_coordinator.py create mode 100644 reai_toolkit/app/services/chat/__init__.py create mode 100644 reai_toolkit/app/services/chat/chat_service.py create mode 100644 reai_toolkit/app/services/chat/reducer.py create mode 100644 reai_toolkit/app/services/chat/schema.py create mode 100644 reai_toolkit/app/services/chat/sse.py create mode 100644 tests/unit/chat/test_reducer.py create mode 100644 tests/unit/chat/test_render.py create mode 100644 tests/unit/chat/test_sdk_schemas.py create mode 100644 tests/unit/chat/test_service.py create mode 100644 tests/unit/chat/test_sse_parser.py diff --git a/reai_toolkit/app/app.py b/reai_toolkit/app/app.py index 80a494a..a9b8dd4 100644 --- a/reai_toolkit/app/app.py +++ b/reai_toolkit/app/app.py @@ -5,6 +5,7 @@ ) from reai_toolkit.app.services.analysis_sync.analysis_sync import AnalysisSyncService from reai_toolkit.app.services.auth.auth_service import AuthService +from reai_toolkit.app.services.chat.chat_service import ChatService from reai_toolkit.app.services.existing_analyses.existing_analyses_service import ( ExistingAnalysesService, ) @@ -64,4 +65,7 @@ def __init__(self, ida_version: str = "UNKNOWN", plugin_version: str = "UNKNOWN" self.similarity_service = SimilarityService( netstore_service=self.netstore_service, sdk_config=sdk_config ) + self.chat_service = ChatService( + netstore_service=self.netstore_service, sdk_config=sdk_config + ) diff --git a/reai_toolkit/app/components/tabs/chat_render.py b/reai_toolkit/app/components/tabs/chat_render.py new file mode 100644 index 0000000..0882655 --- /dev/null +++ b/reai_toolkit/app/components/tabs/chat_render.py @@ -0,0 +1,93 @@ +"""Pure presentation helpers for the Agent Chat panel. + +Turns a :class:`ChatState` into a markdown transcript. No Qt / IDA imports, so it +is unit-testable without a GUI (Qt only imports in the GUI version of IDA). +""" + +from __future__ import annotations + +from typing import Optional + +from reai_toolkit.app.services.chat.schema import ChatState, ToolConfirmation + + +def title_case(name: str) -> str: + pretty = (name or "").replace("_", " ").strip() + return pretty.title() if pretty else (name or "tool") + + +def tool_marker(is_error: bool, status: str) -> str: + if is_error: + return "✗" + return "✓" if status == "finished" else "…" + + +JUMP_SCHEME = "ida://jump/" + + +def jump_href(ea: int) -> str: + return f"{JUMP_SCHEME}{ea}" + + +def parse_jump_href(url: str) -> Optional[int]: + if not url.startswith(JUMP_SCHEME): + return None + try: + return int(url[len(JUMP_SCHEME):]) + except (ValueError, TypeError): + return None + + +def _function_links(functions) -> str: + parts = [f"[{f.name}]({jump_href(f.ea)})" for f in functions if f.name] + return "↪ " + " · ".join(parts) if parts else "" + + +def render_transcript_markdown(state: ChatState) -> str: + """Build a single markdown document from the chat items.""" + parts: list[str] = [] + for item in state.items: + kind = item.kind + if kind == "user-message": + parts.append(f"**You:** {item.content}") + elif kind == "assistant-message": + text = item.content or "" + if item.is_streaming: + text = f"{text} ▍" + parts.append(text if text.strip() else "_…_") + elif kind == "tool-call": + parts.append(f"`{tool_marker(item.is_error, item.status)} {title_case(item.name)}`") + if item.functions: + links = _function_links(item.functions) + if links: + parts.append(links) + elif kind == "step": + if item.status == "running": + parts.append(f"_{item.step_name}…_") + elif kind == "tool-confirmation": + tool = title_case(item.tool_name) + if item.status == "pending": + parts.append(f"> ⚠ **Approval needed** — `{tool}`") + elif item.status == "approved": + parts.append(f"> ✓ Approved — `{tool}`") + else: + parts.append(f"> ✗ Rejected — `{tool}`") + elif kind == "context-compacted": + parts.append("_— context compacted —_") + + if state.run_status == "running": + last = state.items[-1] if state.items else None + if last is None or last.kind == "user-message": + parts.append("_Thinking…_") + + if state.run_status == "error" and state.run_error is not None: + parts.append(f"> ⚠ **Error:** {state.run_error.message}") + + return "\n\n".join(parts) + + +def find_pending_confirmation(state: ChatState) -> Optional[ToolConfirmation]: + for item in reversed(state.items): + if isinstance(item, ToolConfirmation) and item.status == "pending": + return item + return None diff --git a/reai_toolkit/app/components/tabs/chat_tab.py b/reai_toolkit/app/components/tabs/chat_tab.py new file mode 100644 index 0000000..86b71eb --- /dev/null +++ b/reai_toolkit/app/components/tabs/chat_tab.py @@ -0,0 +1,465 @@ +"""Dockable Agent Chat panel + background streaming worker. + +`ChatPanel` is a dumb renderer over a :class:`ChatState`; all state/business +logic lives in :class:`ChatCoordinator`. `ChatStreamWorker` runs the network +turn (create → send → stream) on a `QThread` and emits Qt signals; a QObject +relay with UI-thread affinity guarantees those signals are delivered on the UI +thread (a bound method of the non-QObject `PluginForm` would otherwise run on the +worker thread under an auto/direct connection). +""" + +from __future__ import annotations + +import threading +from typing import Callable, Optional + +import ida_kernwin as kw +from loguru import logger + +from reai_toolkit.app.components.tabs.chat_render import ( + find_pending_confirmation, + parse_jump_href, + render_transcript_markdown, + title_case, +) +from reai_toolkit.app.core.qt_compat import QtCore, QtWidgets, Signal, Slot +from reai_toolkit.app.services.chat.schema import ChatState + + +class ChatStreamWorker(QtCore.QObject): + """Runs one chat turn (optional create → optional send → stream) off-thread.""" + + event_ready = Signal(object) + conversation_created = Signal(str) + errored = Signal(str) + finished = Signal() + + def __init__( + self, + chat_service, + conversation_id: Optional[str], + content: Optional[str], + context, + last_event_id: Optional[int] = None, + ) -> None: + super().__init__() + self._svc = chat_service + self._conversation_id = conversation_id + self._content = content + self._context = context + self._last_event_id = last_event_id + self._stopped = False + self._stop_event = threading.Event() + + @Slot() + def run(self) -> None: + try: + conv_id = self._conversation_id + if conv_id is None: + res = self._svc.create_conversation(self._context) + if not res.success or not res.data: + self.errored.emit(res.error_message or "Failed to create conversation") + return + conv_id = res.data + self.conversation_created.emit(conv_id) + + if self._stopped: + return + + if self._content is not None: + res = self._svc.send_message(conv_id, self._content, self._context) + if not res.success: + self.errored.emit(res.error_message or "Failed to send message") + return + + if self._stopped: + return + + for ev in self._svc.stream(conv_id, self._stop_event, self._last_event_id): + if self._stopped: + break + self.event_ready.emit(ev) + except Exception as e: + self.errored.emit(str(e)) + finally: + self.finished.emit() + + def stop(self) -> None: + self._stopped = True + self._stop_event.set() + try: + self._svc.close_active_stream() + except Exception: + pass + + +class _StreamRelay(QtCore.QObject): + """UI-thread QObject that forwards worker signals to plain callbacks.""" + + def __init__(self) -> None: + super().__init__() + 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) + + @Slot(str) + def handle_conversation_created(self, uuid: str) -> None: + if self.on_conversation_created: + self.on_conversation_created(uuid) + + @Slot(str) + def handle_error(self, msg: str) -> None: + if self.on_error: + self.on_error(msg) + + @Slot() + def handle_finished(self) -> None: + if self.on_finished: + self.on_finished() + + +class _ChatInput(QtWidgets.QPlainTextEdit): + submit = Signal() + + def keyPressEvent(self, event) -> None: + key = event.key() + is_enter = key in (QtCore.Qt.Key_Return, QtCore.Qt.Key_Enter) + shift = bool(event.modifiers() & QtCore.Qt.ShiftModifier) + if is_enter and not shift: + self.submit.emit() + return + super().keyPressEvent(event) + + +class ChatPanel(kw.PluginForm): + TITLE = "RevEng.AI — Agent Chat" + + def __init__(self, on_closed: Optional[Callable[[], None]] = None) -> None: + super().__init__() + self._on_closed = on_closed + + self.on_send: Optional[Callable[[str], None]] = None + self.on_stop: Optional[Callable[[], None]] = None + self.on_confirm: Optional[Callable[[str, bool], None]] = None + self.on_new_chat: Optional[Callable[[], None]] = None + self.on_select_conversation: Optional[Callable[[str], None]] = None + self.on_request_history: Optional[Callable[[], None]] = None + self.on_jump: Optional[Callable[[int], None]] = None + self.on_stream_event: Optional[Callable] = None + self.on_stream_conversation_created: Optional[Callable[[str], None]] = None + self.on_stream_error: Optional[Callable[[str], None]] = None + self.on_stream_finished: Optional[Callable[[], None]] = None + + self._parent_window: Optional[QtWidgets.QWidget] = None + self._title_label: Optional[QtWidgets.QLabel] = None + self._context_label: Optional[QtWidgets.QLabel] = None + self._history_list: Optional[QtWidgets.QListWidget] = None + self._transcript: Optional[QtWidgets.QTextBrowser] = None + self._confirm_bar: Optional[QtWidgets.QWidget] = None + self._confirm_msg: Optional[QtWidgets.QLabel] = None + self._input: Optional[_ChatInput] = None + self._send_btn: Optional[QtWidgets.QPushButton] = None + self._stop_btn: Optional[QtWidgets.QPushButton] = None + + self._relay: Optional[_StreamRelay] = None + self._thread: Optional[QtCore.QThread] = None + self._worker: Optional[ChatStreamWorker] = None + self._streaming = False + self._zombies: list = [] + + self._render_timer: Optional[QtCore.QTimer] = None + self._pending_state: Optional[ChatState] = None + self._pending_confirm_id: Optional[str] = None + + def Create(self, title) -> bool: + flags = getattr(kw.PluginForm, "WOPN_DP_TAB", 0) | getattr( + kw.PluginForm, "WOPN_RESTORE", 0 + ) + name = str(title) if title else self.TITLE + ok = self.Show(name, flags) + if not ok: + logger.error("Failed to show Agent Chat tab") + else: + try: + kw.set_dock_pos(name, "Pseudocode-A", kw.DP_RIGHT) + except Exception: + pass + return ok + + def OnCreate(self, form) -> None: + self._parent_window = self.FormToPyQtWidget(form) + root = QtWidgets.QVBoxLayout(self._parent_window) + root.setContentsMargins(6, 6, 6, 6) + root.setSpacing(6) + + header = QtWidgets.QHBoxLayout() + self._title_label = QtWidgets.QLabel("AI Agent") + 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._history_list = QtWidgets.QListWidget() + self._history_list.setVisible(False) + self._history_list.setMaximumHeight(140) + root.addWidget(self._history_list) + + self._transcript = QtWidgets.QTextBrowser() + self._transcript.setOpenLinks(False) + self._transcript.setOpenExternalLinks(False) + self._transcript.anchorClicked.connect(self._on_anchor_clicked) + root.addWidget(self._transcript, 1) + + self._confirm_bar = QtWidgets.QWidget() + cbar = QtWidgets.QHBoxLayout(self._confirm_bar) + cbar.setContentsMargins(0, 0, 0, 0) + self._confirm_msg = QtWidgets.QLabel("") + self._confirm_msg.setWordWrap(True) + approve_btn = QtWidgets.QPushButton("Approve") + reject_btn = QtWidgets.QPushButton("Reject") + cbar.addWidget(self._confirm_msg, 1) + cbar.addWidget(approve_btn) + cbar.addWidget(reject_btn) + self._confirm_bar.setVisible(False) + root.addWidget(self._confirm_bar) + + input_row = QtWidgets.QHBoxLayout() + self._input = _ChatInput() + self._input.setPlaceholderText("Ask a question… (Enter to send, Shift+Enter for newline)") + self._input.setMaximumHeight(80) + self._send_btn = QtWidgets.QPushButton("Send") + self._stop_btn = QtWidgets.QPushButton("Stop") + self._stop_btn.setEnabled(False) + btn_col = QtWidgets.QVBoxLayout() + btn_col.addWidget(self._send_btn) + btn_col.addWidget(self._stop_btn) + input_row.addWidget(self._input, 1) + input_row.addLayout(btn_col) + root.addLayout(input_row) + + history_btn.clicked.connect(self._on_history_clicked) + new_btn.clicked.connect(self._on_new_chat_clicked) + approve_btn.clicked.connect(lambda: self._on_confirm_clicked(True)) + reject_btn.clicked.connect(lambda: self._on_confirm_clicked(False)) + self._send_btn.clicked.connect(self._on_send_clicked) + self._stop_btn.clicked.connect(self._on_stop_clicked) + self._input.submit.connect(self._on_send_clicked) + self._history_list.itemActivated.connect(self._on_history_item) + self._history_list.itemClicked.connect(self._on_history_item) + + self._render_timer = QtCore.QTimer(self._parent_window) + self._render_timer.setSingleShot(True) + self._render_timer.setInterval(40) + self._render_timer.timeout.connect(self._flush_render) + + self._relay = _StreamRelay() + self._relay.on_event = self._handle_stream_event + self._relay.on_conversation_created = self._handle_conversation_created + self._relay.on_error = self._handle_stream_error + self._relay.on_finished = self._handle_stream_finished + + def OnClose(self, form) -> None: + self.stop_stream_worker() + if self._render_timer is not None: + try: + self._render_timer.stop() + except Exception: + pass + if callable(self._on_closed): + try: + self._on_closed() + except Exception as e: + logger.warning(f"Agent Chat on_closed failed: {e}") + self._transcript = None + self._input = None + self._parent_window = None + self._render_timer = None + + def focus(self) -> None: + if self._parent_window: + try: + kw.activate_widget(self._parent_window, True) + except Exception: + pass + + def request_render(self, state: ChatState) -> None: + self._pending_state = state + if self._render_timer is None: + self._flush_render() + elif not self._render_timer.isActive(): + self._render_timer.start() + + def set_context_chip(self, text: str) -> None: + if self._context_label is not None: + self._context_label.setText(text or "") + + def set_history(self, summaries) -> None: + if self._history_list is None: + return + self._history_list.clear() + for s in summaries: + label = s.title or "Untitled" + item = QtWidgets.QListWidgetItem(label) + item.setData(QtCore.Qt.UserRole, s.conversation_uuid) + self._history_list.addItem(item) + + def clear_input(self) -> None: + if self._input is not None: + self._input.clear() + + def focus_input(self) -> None: + if self._input is not None: + self._input.setFocus() + + def _flush_render(self) -> None: + state = self._pending_state + if state is None or self._transcript is None: + return + + sb = self._transcript.verticalScrollBar() + at_bottom = sb is None or sb.value() >= sb.maximum() - 4 + md = render_transcript_markdown(state) + if hasattr(self._transcript, "setMarkdown"): + self._transcript.setMarkdown(md) + else: + self._transcript.setPlainText(md) + if at_bottom and sb is not None: + sb.setValue(sb.maximum()) + + running = state.run_status == "running" + if self._send_btn is not None: + self._send_btn.setEnabled(not running) + if self._stop_btn is not None: + self._stop_btn.setEnabled(running) + if self._title_label is not None: + self._title_label.setText(state.title or "AI Agent") + + pending = find_pending_confirmation(state) + if self._confirm_bar is not None: + if pending is not None: + self._pending_confirm_id = pending.id + if self._confirm_msg is not None: + self._confirm_msg.setText( + pending.message or f"Approve tool '{title_case(pending.tool_name)}'?" + ) + self._confirm_bar.setVisible(True) + else: + self._pending_confirm_id = None + self._confirm_bar.setVisible(False) + + def start_stream_worker(self, worker: ChatStreamWorker) -> None: + self.stop_stream_worker() + thread = QtCore.QThread() + worker.moveToThread(thread) + thread.started.connect(worker.run) + if self._relay is not None: + worker.event_ready.connect(self._relay.handle_event) + worker.conversation_created.connect(self._relay.handle_conversation_created) + worker.errored.connect(self._relay.handle_error) + worker.finished.connect(self._relay.handle_finished) + worker.finished.connect(thread.quit) + self._thread = thread + self._worker = worker + self._streaming = True + thread.start() + + def stop_stream_worker(self) -> None: + self._zombies = [ + (t, w) for (t, w) in self._zombies if t is not None and not t.isFinished() + ] + + worker, thread = self._worker, self._thread + self._worker = None + self._thread = None + self._streaming = False + + if worker is not None: + try: + worker.stop() + except Exception: + pass + if thread is None: + return + try: + thread.quit() + if not thread.wait(3000): + self._zombies.append((thread, worker)) + except Exception: + pass + + def is_streaming(self) -> bool: + return self._streaming + + def _handle_stream_event(self, ev) -> None: + if self.on_stream_event: + self.on_stream_event(ev) + + def _handle_conversation_created(self, uuid: str) -> None: + if self.on_stream_conversation_created: + self.on_stream_conversation_created(uuid) + + def _handle_stream_error(self, msg: str) -> None: + if self.on_stream_error: + self.on_stream_error(msg) + + def _handle_stream_finished(self) -> None: + self._streaming = False + if self.on_stream_finished: + self.on_stream_finished() + + def _on_send_clicked(self) -> None: + if self._input is None: + return + text = self._input.toPlainText() + if not text.strip(): + return + self._input.clear() + if self.on_send: + self.on_send(text) + + def _on_stop_clicked(self) -> None: + if self.on_stop: + self.on_stop() + + def _on_confirm_clicked(self, approved: bool) -> None: + if self._pending_confirm_id and self.on_confirm: + self.on_confirm(self._pending_confirm_id, approved) + + def _on_new_chat_clicked(self) -> None: + if self.on_new_chat: + self.on_new_chat() + + def _on_history_clicked(self) -> None: + if self._history_list is None: + return + show = not self._history_list.isVisible() + self._history_list.setVisible(show) + if show and self.on_request_history: + self.on_request_history() + + def _on_history_item(self, item) -> None: + if item is None: + return + uuid = item.data(QtCore.Qt.UserRole) + if uuid and self.on_select_conversation: + self.on_select_conversation(uuid) + + def _on_anchor_clicked(self, url) -> None: + ea = parse_jump_href(url.toString()) + if ea is not None and self.on_jump: + self.on_jump(ea) diff --git a/reai_toolkit/app/coordinator.py b/reai_toolkit/app/coordinator.py index d3fbb8e..b92ce29 100644 --- a/reai_toolkit/app/coordinator.py +++ b/reai_toolkit/app/coordinator.py @@ -6,6 +6,7 @@ from reai_toolkit.app.coordinators.ai_decomp_coordinator import AiDecompCoordinator from reai_toolkit.app.coordinators.auth_coordinator import AuthCoordinator from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator +from reai_toolkit.app.coordinators.chat_coordinator import ChatCoordinator from reai_toolkit.app.coordinators.create_analysis_coordinator import ( CreateAnalysisCoordinator, ) @@ -90,6 +91,14 @@ def __init__(self, app, factory, log): similarity_service=app.similarity_service ) + self.chatc: ChatCoordinator = ChatCoordinator( + app=app, + factory=factory, + log=log, + chat_service=app.chat_service, + ai_decomp_coord=self.ai_decompc, + ) + def run_dialog(self): """Run all necessary dialogs at startup.""" pass diff --git a/reai_toolkit/app/coordinators/chat_coordinator.py b/reai_toolkit/app/coordinators/chat_coordinator.py new file mode 100644 index 0000000..f00fea6 --- /dev/null +++ b/reai_toolkit/app/coordinators/chat_coordinator.py @@ -0,0 +1,372 @@ +"""Agent Chat coordinator. + +Owns the :class:`ChatState`, the active conversation id, and the streaming worker +lifecycle; wires the dockable :class:`ChatPanel` to the :class:`ChatService`. The +panel is a dumb renderer — every state transition goes through the pure +``chat_reducer`` here, mirroring the Dashboard's ``useChat`` hook. +""" + +from __future__ import annotations + +import threading +import uuid +from logging import Logger +from typing import TYPE_CHECKING, Optional + +import ida_funcs +import ida_kernwin +import ida_name +from libbs.decompilers.ida.compat import execute_read, execute_ui + +from reai_toolkit.app.app import App +from reai_toolkit.app.components.tabs.chat_tab import ChatPanel, ChatStreamWorker +from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator +from reai_toolkit.app.factory import DialogFactory +from reai_toolkit.app.services.chat.chat_service import ChatService +from reai_toolkit.app.services.chat.reducer import ( + ApiError, + Cancel, + ConfirmTool, + EventAction, + SendMessage, + SetToolFunctions, + build_initial_state, + chat_reducer, + initial_state, +) +from reai_toolkit.app.services.chat.schema import ( + ChatState, + ConversationContextDTO, + FunctionRef, +) + +if TYPE_CHECKING: + from reai_toolkit.app.coordinators.ai_decomp_coordinator import AiDecompCoordinator + +AI_DECOMP_TOOL_HINTS = ("decomp",) + + +class ChatCoordinator(BaseCoordinator): + def __init__( + self, + *, + app: "App", + factory: "DialogFactory", + log: Logger, + chat_service: ChatService, + ai_decomp_coord: "Optional[AiDecompCoordinator]" = None, + ) -> None: + super().__init__(app=app, factory=factory, log=log) + self.chat_service: ChatService = chat_service + self._ai_decomp_coord: "Optional[AiDecompCoordinator]" = ai_decomp_coord + self._panel: Optional[ChatPanel] = None + self._state: ChatState = initial_state() + self._conversation_id: Optional[str] = None + self._last_event_id: Optional[int] = None + self._last_context: Optional[ConversationContextDTO] = 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._panel.focus() + self._update_context_chip() + if prefill_context: + self._panel.focus_input() + + if self._conversation_id is None and not self._state.items: + last = self._get_persisted_conversation() + if last: + self.load_conversation(last) + return + self._render() + + def _wire_panel(self, panel: ChatPanel) -> None: + panel.on_send = self.send + panel.on_stop = self.stop + panel.on_confirm = self.confirm_tool + panel.on_new_chat = self.new_conversation + panel.on_select_conversation = self.load_conversation + panel.on_request_history = self.request_history + panel.on_jump = self.jump_to + panel.on_stream_event = self.on_stream_event + panel.on_stream_conversation_created = self.on_stream_conversation_created + panel.on_stream_error = self.on_stream_error + panel.on_stream_finished = self.on_stream_finished + + def _on_pane_closed(self) -> None: + if self._panel is not None: + self._panel.stop_stream_worker() + self._panel = None + self.log.info("Agent Chat panel closed.") + + def send(self, text: str) -> None: + content = (text or "").strip() + if not content or self._state.run_status == "running": + return + self._state = chat_reducer( + self._state, SendMessage(id=uuid.uuid4().hex, content=content) + ) + self._render() + + context = self._resolve_context() + self._last_context = context + self._update_context_chip() + worker = ChatStreamWorker( + chat_service=self.chat_service, + conversation_id=self._conversation_id, + content=content, + context=context, + last_event_id=None, + ) + if self._panel is not None: + self._panel.start_stream_worker(worker) + + def stop(self) -> None: + if self._panel is not None: + self._panel.stop_stream_worker() + self._state = chat_reducer(self._state, Cancel()) + self._render() + if self._conversation_id: + self._post_async(self.chat_service.cancel_run, self._conversation_id) + + def confirm_tool(self, tool_call_id: str, approved: bool) -> None: + self._state = chat_reducer( + self._state, ConfirmTool(id=tool_call_id, approved=approved) + ) + self._render() + if not self._conversation_id: + return + self._post_async(self.chat_service.confirm_tool, self._conversation_id, approved) + if ( + approved + and self._panel is not None + and not self._panel.is_streaming() + ): + worker = ChatStreamWorker( + chat_service=self.chat_service, + conversation_id=self._conversation_id, + content=None, + context=self._resolve_context(), + last_event_id=self._last_event_id, + ) + self._panel.start_stream_worker(worker) + + def new_conversation(self) -> None: + if self._panel is not None: + self._panel.stop_stream_worker() + self._conversation_id = None + self._last_event_id = None + self._state = initial_state() + self._render() + + def load_conversation(self, conversation_uuid: str) -> None: + if self._panel is not None: + self._panel.stop_stream_worker() + + def _work() -> None: + res = self.chat_service.get_conversation(conversation_uuid) + if not res.success or res.data is None: + if res.error_message: + self.show_error_dialog(message=res.error_message) + return + replay = res.data + state = build_initial_state(replay.events) + state.title = replay.title + self._apply_loaded_state(conversation_uuid, state) + + threading.Thread(target=_work, daemon=True).start() + + def request_history(self) -> None: + def _work() -> None: + res = self.chat_service.list_conversations() + if res.success and res.data is not None: + self._set_history(res.data) + elif res.error_message: + self.log.warning(f"Failed to list conversations: {res.error_message}") + + threading.Thread(target=_work, daemon=True).start() + + def on_stream_event(self, ev) -> None: + self._state = chat_reducer(self._state, EventAction(ev)) + if ev.event_id is not None: + self._last_event_id = ev.event_id + if ev.type == "TOOL_CALL_RESULT" and not ev.is_error: + self._handle_tool_result(ev) + self._render() + + def on_stream_conversation_created(self, conversation_uuid: str) -> None: + self._conversation_id = conversation_uuid + self._persist_conversation(conversation_uuid) + + def on_stream_error(self, msg: str) -> None: + self._state = chat_reducer(self._state, ApiError(message=msg)) + self._render() + + def on_stream_finished(self) -> None: + return + + 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) + 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 _work() -> None: + func_map = self.app.netstore_service.get_function_mapping() + if func_map is None: + return + refs: list[FunctionRef] = [] + applied = False + for fid in func_ids: + ea = func_map.function_map.get(str(fid)) + if ea is None: + continue + res = self.chat_service.get_function_name(fid) + if not res.success or not res.data: + continue + name = res.data + 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) + if refs: + self._attach_function_links(tool_call_id, refs) + if applied: + self.refresh_disassembly_view() + + 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: + ida_kernwin.jumpto(ea) + except Exception as e: + self.log.debug(f"jumpto({ea}) failed: {e}") + + def _maybe_open_viewer(self, ev) -> None: + if self._ai_decomp_coord is None: + return + 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]: + 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)) + + def _render(self) -> None: + if self._panel is not None: + self._panel.request_render(self._state) + + @execute_ui + def _apply_loaded_state(self, conversation_uuid: str, state: ChatState) -> None: + self._conversation_id = conversation_uuid + self._state = state + self._last_event_id = None + self._persist_conversation(conversation_uuid) + self._render() + + @execute_ui + def _set_history(self, summaries) -> None: + if self._panel is not None: + self._panel.set_history(summaries) + + @execute_read + def _resolve_context(self) -> ConversationContextDTO: + analysis_id = None + function_id = None + try: + analysis_id = self.app.netstore_service.get_analysis_id() + func_map = self.app.netstore_service.get_function_mapping() + func = ida_funcs.get_func(ida_kernwin.get_screen_ea()) + if func_map is not None and func is not None: + function_id = func_map.inverse_function_map.get(str(func.start_ea)) + except Exception as e: + 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: + if self._panel is not None: + self._panel.set_context_chip(self._resolve_context_chip_text()) + + @execute_read + def _resolve_context_chip_text(self) -> 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 + bits = [] + if name: + 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" + except Exception: + return "" + + def _post_async(self, fn, *args) -> None: + def _work() -> None: + try: + res = fn(*args) + if res is not None and not res.success and res.error_message: + self.show_error_dialog(message=res.error_message) + except Exception as e: + self.log.warning(f"Async chat request failed: {e}") + + threading.Thread(target=_work, daemon=True).start() + + def _conversation_key(self, analysis_id: int) -> str: + return f"chat:last_conversation:{analysis_id}" + + def _persist_conversation(self, conversation_uuid: str) -> None: + try: + analysis_id = self.app.netstore_service.get_analysis_id() + if analysis_id is not None: + self.app.netstore_service.put_global( + self._conversation_key(analysis_id), conversation_uuid + ) + except Exception as e: + self.log.debug(f"Failed to persist conversation id: {e}") + + def _get_persisted_conversation(self) -> Optional[str]: + try: + analysis_id = self.app.netstore_service.get_analysis_id() + if analysis_id is None: + return None + return self.app.netstore_service.get_global( + self._conversation_key(analysis_id) + ) + except Exception: + return None diff --git a/reai_toolkit/app/factory.py b/reai_toolkit/app/factory.py index c2d5fa4..bc7e5d6 100644 --- a/reai_toolkit/app/factory.py +++ b/reai_toolkit/app/factory.py @@ -51,6 +51,11 @@ def ai_decomp(self, on_closed: Callable[..., Any]) -> AIDecompView: return AIDecompView(on_closed=on_closed) + def chat(self, on_closed: Callable[..., Any]): + from reai_toolkit.app.components.tabs.chat_tab import ChatPanel + + return ChatPanel(on_closed=on_closed) + def function_matching( self, valid_functions: list[ValidFunction], func_map: dict[str, int] ): diff --git a/reai_toolkit/app/services/chat/__init__.py b/reai_toolkit/app/services/chat/__init__.py new file mode 100644 index 0000000..6025f5e --- /dev/null +++ b/reai_toolkit/app/services/chat/__init__.py @@ -0,0 +1,5 @@ +"""Agent Chat feature — streaming conversational AI agent. + +Ports the Dashboard's Agent Chat subsystem (event schema, reducer state machine +and SSE streaming client) onto the vendored ``revengai.ConversationsApi``. +""" diff --git a/reai_toolkit/app/services/chat/chat_service.py b/reai_toolkit/app/services/chat/chat_service.py new file mode 100644 index 0000000..c8f64e5 --- /dev/null +++ b/reai_toolkit/app/services/chat/chat_service.py @@ -0,0 +1,274 @@ +"""Streaming chat client for the Agent Chat feature. + +Wraps the vendored ``revengai.ConversationsApi`` (create / send / stream / confirm +/ cancel / list / get) behind the plugin's :class:`BaseService` conventions +(``yield_api_client`` + :class:`GenericApiReturn`). The SSE frame parsing lives in +the pure :mod:`sse` module; this class only owns the network lifecycle. +""" + +from __future__ import annotations + +import threading +from typing import Generator, Optional +from uuid import UUID + +import urllib3 +from loguru import logger +from revengai import ( + ApiException, + ConfirmToolInputBody, + Configuration, + ConversationContext, + ConversationsApi, + CreateConversationRequest, + FunctionsCoreApi, + SendMessageRequest, +) + +from reai_toolkit.app.core.netstore_service import SimpleNetStore +from reai_toolkit.app.core.shared_schema import GenericApiReturn +from reai_toolkit.app.core.utils import parse_exception +from reai_toolkit.app.interfaces.base_service import BaseService +from reai_toolkit.app.services.chat.schema import ( + ROLE_USER, + ChatEvent, + ConversationContextDTO, + ConversationReplay, + ConversationSummary, + StoredEvent, + UserMessageReplay, + normalize_event, + resolve_type, +) +from reai_toolkit.app.services.chat.sse import iter_sse_events + +CONNECT_TIMEOUT: float = 10.0 +READ_TIMEOUT: float = 300.0 + + +class ChatService(BaseService): + def __init__( + self, netstore_service: SimpleNetStore, sdk_config: Configuration + ) -> None: + super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) + self._active_response: Optional[urllib3.HTTPResponse] = None + + def create_conversation( + self, context: Optional[ConversationContextDTO], title: Optional[str] = None + ) -> GenericApiReturn[str]: + req = CreateConversationRequest(context=self._to_sdk_context(context), title=title) + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + conv = ConversationsApi(api_client).create_conversation(req) + return GenericApiReturn(success=True, data=conv.conversation_uuid) + except ApiException as e: + return self._api_error(e) + except Exception as e: + return GenericApiReturn(success=False, error_message=f"Unexpected error: {e}") + + def send_message( + self, + conversation_id: str, + content: str, + context: Optional[ConversationContextDTO], + ) -> GenericApiReturn[None]: + req = SendMessageRequest(content=content, context=self._to_sdk_context(context)) + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + ConversationsApi(api_client).send_message(UUID(conversation_id), req) + return GenericApiReturn(success=True) + except ApiException as e: + if e.status == 409: + return GenericApiReturn(success=True) + return self._api_error(e) + except Exception as e: + return GenericApiReturn(success=False, error_message=f"Unexpected error: {e}") + + def confirm_tool( + self, conversation_id: str, approved: bool + ) -> GenericApiReturn[None]: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + ConversationsApi(api_client).confirm_tool( + UUID(conversation_id), ConfirmToolInputBody(approved=approved) + ) + return GenericApiReturn(success=True) + except ApiException as e: + if e.status == 404: + return GenericApiReturn(success=True) + return self._api_error(e) + except Exception as e: + return GenericApiReturn(success=False, error_message=f"Unexpected error: {e}") + + def cancel_run(self, conversation_id: str) -> GenericApiReturn[None]: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + ConversationsApi(api_client).cancel_run(UUID(conversation_id)) + return GenericApiReturn(success=True) + except ApiException as e: + if e.status == 404: + return GenericApiReturn(success=True) + return self._api_error(e) + except Exception as e: + return GenericApiReturn(success=False, error_message=f"Unexpected error: {e}") + + def list_conversations(self) -> GenericApiReturn[list[ConversationSummary]]: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + convs = ConversationsApi(api_client).list_conversations() + summaries = [ + ConversationSummary( + conversation_uuid=c.conversation_uuid, + title=c.title, + updated_at=str(c.updated_at) if c.updated_at is not None else None, + ) + for c in (convs or []) + ] + return GenericApiReturn(success=True, data=summaries) + except ApiException as e: + return self._api_error(e) + except Exception as e: + return GenericApiReturn(success=False, error_message=f"Unexpected error: {e}") + + def get_conversation( + self, conversation_id: str + ) -> GenericApiReturn[ConversationReplay]: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + conv = ConversationsApi(api_client).get_conversation(UUID(conversation_id)) + replay = ConversationReplay( + conversation_uuid=conv.conversation_uuid, + title=conv.title, + events=self._replay_events(conv.events or []), + ) + return GenericApiReturn(success=True, data=replay) + except ApiException as e: + return self._api_error(e) + except Exception as e: + return GenericApiReturn(success=False, error_message=f"Unexpected error: {e}") + + def get_function_name(self, function_id: int) -> GenericApiReturn[str]: + """Fetch a single function's current name from the backend (the same + info the agent's 'Get Function Info' tool reads).""" + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + details = FunctionsCoreApi(api_client).get_function_details_0( + function_id=function_id + ) + return GenericApiReturn(success=True, data=details.function_name) + except ApiException as e: + return self._api_error(e) + except Exception as e: + return GenericApiReturn(success=False, error_message=f"Unexpected error: {e}") + + def stream( + self, + conversation_id: str, + stop_event: threading.Event, + last_event_id: Optional[int] = None, + ) -> Generator[ChatEvent, None, None]: + """Yield normalized :class:`ChatEvent`\\ s from the live SSE stream. + + Uses ``stream_events_without_preload_content`` to get the raw urllib3 + response and iterates it incrementally. Ends gracefully on a terminal + event, a set ``stop_event``, or a dropped/closed connection. + """ + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + resp = ConversationsApi(api_client).stream_events_without_preload_content( + id=UUID(conversation_id), + last_event_id=last_event_id, + _request_timeout=(CONNECT_TIMEOUT, READ_TIMEOUT), + ) + self._active_response = resp + try: + yield from iter_sse_events( + resp.stream(1024, decode_content=True), + stop=stop_event.is_set, + ) + except (urllib3.exceptions.HTTPError, OSError) as e: + logger.debug(f"[Chat] SSE stream ended: {e}") + return + finally: + self._active_response = None + try: + resp.release_conn() + except Exception: + pass + + def close_active_stream(self) -> None: + """Interrupt a blocked stream read from another thread (cancellation).""" + resp = self._active_response + if resp is not None: + try: + resp.close() + except Exception: + pass + + @staticmethod + def _to_sdk_context( + dto: Optional[ConversationContextDTO], + ) -> Optional[ConversationContext]: + if dto is None or dto.is_empty(): + return None + return ConversationContext( + analysis_id=dto.analysis_id, function_id=dto.function_id + ) + + @staticmethod + def _replay_events(sdk_events: list) -> list[StoredEvent]: + """Reconstruct stored events for history replay. + + Mirrors ``getConversation`` in agentApi.ts: role-USER TEXT_MESSAGE_* + events collapse into a single :class:`UserMessageReplay`; everything else + is normalized like a live frame. + """ + out: list[StoredEvent] = [] + emitted_user_ids: set[str] = set() + current_user_id: Optional[str] = None + current_user_content = "" + + for ev in sdk_events: + etype = resolve_type(getattr(ev, "type", None)) + role = getattr(ev, "role", None) + data = getattr(ev, "data", None) + data = data if isinstance(data, dict) else {} + + if role == ROLE_USER: + if etype == "TEXT_MESSAGE_START": + msg_id = str( + data.get("message_id") or getattr(ev, "event_id", "") or "" + ) + if msg_id not in emitted_user_ids: + current_user_id = msg_id + current_user_content = "" + else: + current_user_id = None + elif etype == "TEXT_MESSAGE_CONTENT": + if current_user_id is not None: + current_user_content += str(data.get("delta") or "") + elif etype == "TEXT_MESSAGE_END": + if current_user_id is not None: + out.append( + UserMessageReplay( + id=current_user_id, content=current_user_content + ) + ) + emitted_user_ids.add(current_user_id) + current_user_id = None + current_user_content = "" + continue + + norm = normalize_event(getattr(ev, "type", None), data) + if norm is not None: + out.append(norm) + return out + + @staticmethod + def _api_error(e: ApiException) -> GenericApiReturn: + error_response = parse_exception(e) + if error_response and error_response.errors and len(error_response.errors) > 0: + return GenericApiReturn( + success=False, + error_message=f"{error_response.errors[0].code}: {error_response.errors[0].message}", + ) + return GenericApiReturn(success=False, error_message=f"API Exception: {e}") diff --git a/reai_toolkit/app/services/chat/reducer.py b/reai_toolkit/app/services/chat/reducer.py new file mode 100644 index 0000000..9c2f736 --- /dev/null +++ b/reai_toolkit/app/services/chat/reducer.py @@ -0,0 +1,328 @@ +"""Pure reducer state machine for the Agent Chat feature. + +Direct port of ``Dashboard/components/features/AgentChat/reducer.ts``. Operates +only on :mod:`schema` dataclasses — no Qt / IDA / SDK imports — so it runs under +plain ``pytest``. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass, replace +from typing import Callable, Optional, Union + +from reai_toolkit.app.services.chat.schema import ( + AssistantMessage, + ChatEvent, + ChatItem, + ChatState, + ContextCompacted, + RunError, + Step, + StoredEvent, + ToolCall, + ToolConfirmation, + UserMessage, + UserMessageReplay, +) + + +@dataclass +class SendMessage: + id: str + content: str + + +@dataclass +class Cancel: + pass + + +@dataclass +class ConfirmTool: + id: str + approved: bool + + +@dataclass +class EventAction: + event: ChatEvent + + +@dataclass +class ApiError: + message: str + code: Optional[str] = None + doc_url: Optional[str] = None + + +@dataclass +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 +] + + +def _new_id() -> str: + return uuid.uuid4().hex + + +def _find_last_and_update( + items: list[ChatItem], + guard: Callable[[ChatItem], bool], + update: Callable[[ChatItem], ChatItem], +) -> list[ChatItem]: + """Update the last item matching ``guard``. Returns the same list unchanged + if no item matches (mirrors ``findAndUpdate`` in reducer.ts).""" + for idx in range(len(items) - 1, -1, -1): + if guard(items[idx]): + next_items = list(items) + next_items[idx] = update(items[idx]) + return next_items + return items + + +def _finalize_running_items(items: list[ChatItem]) -> list[ChatItem]: + """Flip every in-progress item to a terminal state (CANCEL / *_ERROR).""" + out: list[ChatItem] = [] + for item in items: + if isinstance(item, AssistantMessage) and item.is_streaming: + out.append(replace(item, is_streaming=False)) + elif isinstance(item, (ToolCall, Step)) and item.status == "running": + out.append(replace(item, status="finished")) + elif isinstance(item, ToolConfirmation) and item.status == "pending": + out.append(replace(item, status="rejected")) + else: + out.append(item) + return out + + +def initial_state() -> ChatState: + return ChatState() + + +def chat_reducer(state: ChatState, action: ChatAction) -> ChatState: + if isinstance(action, SendMessage): + return replace( + state, + items=[*state.items, UserMessage(id=action.id, content=action.content)], + run_status="running", + run_error=None, + ) + + if isinstance(action, Reset): + return build_initial_state(action.stored_events) + + if isinstance(action, Cancel): + return replace( + state, + run_status="idle", + items=_finalize_running_items(state.items), + ) + + if isinstance(action, ConfirmTool): + return replace( + state, + run_status="running" if action.approved else state.run_status, + items=_find_last_and_update( + state.items, + lambda it: isinstance(it, ToolConfirmation) and it.id == action.id, + lambda it: replace( + it, status="approved" if action.approved else "rejected" + ), + ), + ) + + if isinstance(action, ApiError): + return replace( + state, + run_status="error", + run_error=RunError( + message=action.message, code=action.code, doc_url=action.doc_url + ), + 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 _reduce_event(state: ChatState, event: ChatEvent) -> ChatState: + t = event.type + + if t == "RUN_STARTED": + return replace(state, run_status="running", run_error=None) + + if t == "RUN_FINISHED": + return replace(state, run_status="idle") + + if t == "RUN_ERROR": + return replace( + state, + run_status="error", + run_error=RunError(message=event.message or "Unknown error"), + items=_finalize_running_items(state.items), + ) + + if t == "STEP_STARTED": + return replace( + state, + items=[ + *state.items, + Step(id=_new_id(), step_name="Processing", status="running"), + ], + ) + + if t == "STEP_FINISHED": + return replace( + state, + items=_find_last_and_update( + state.items, + lambda it: isinstance(it, Step) and it.status == "running", + lambda it: replace(it, status="finished"), + ), + ) + + if t == "TEXT_MESSAGE_START": + if any( + isinstance(it, AssistantMessage) and it.id == event.message_id + for it in state.items + ): + return state + return replace( + state, + items=[ + *state.items, + AssistantMessage( + id=event.message_id or _new_id(), content="", is_streaming=True + ), + ], + ) + + if t == "TEXT_MESSAGE_CONTENT": + return replace( + state, + items=_find_last_and_update( + state.items, + lambda it: isinstance(it, AssistantMessage) + and it.id == event.message_id, + lambda it: replace(it, content=it.content + (event.delta or "")), + ), + ) + + if t == "TEXT_MESSAGE_END": + return replace( + state, + items=_find_last_and_update( + state.items, + lambda it: isinstance(it, AssistantMessage) + and it.id == event.message_id, + lambda it: replace(it, is_streaming=False), + ), + ) + + if t == "TOOL_CALL_START": + return replace( + state, + items=[ + *state.items, + ToolCall( + id=event.tool_call_id or _new_id(), + name=event.tool_name or "", + status="running", + is_error=False, + ), + ], + ) + + if t == "TOOL_CALL_END": + return replace( + state, + 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"), + ), + ) + + if t == "TOOL_CALL_RESULT": + 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), + ) + items = _find_last_and_update( + items, + lambda it: isinstance(it, ToolConfirmation) + and it.id == event.tool_call_id + and it.status == "pending", + lambda it: replace(it, status="rejected" if event.is_error else "approved"), + ) + return replace(state, items=items) + + if t == "TITLE_UPDATED": + return replace(state, title=event.title) + + if t == "RUN_CANCELLED": + return replace( + state, + run_status="idle", + items=_finalize_running_items(state.items), + ) + + if t == "CONTEXT_COMPACTED": + return replace(state, items=[*state.items, ContextCompacted(id=_new_id())]) + + if t == "TOOL_CONFIRMATION_REQUIRED": + return replace( + state, + items=[ + *state.items, + ToolConfirmation( + id=event.tool_call_id or _new_id(), + tool_name=event.tool_name or "", + message=event.message or "", + status="pending", + ), + ], + ) + + return state + + +def build_initial_state(stored_events: Optional[list[StoredEvent]]) -> ChatState: + """Replay stored/normalized events into a :class:`ChatState` (history reload). + + Mirrors ``buildInitialState`` in reducer.ts: a ``UserMessageReplay`` drives a + ``SEND_MESSAGE`` action, everything else an ``EVENT`` action. + """ + state = initial_state() + if not stored_events: + return state + for ev in stored_events: + if isinstance(ev, UserMessageReplay): + state = chat_reducer(state, SendMessage(id=ev.id, content=ev.content)) + elif isinstance(ev, ChatEvent): + state = chat_reducer(state, EventAction(ev)) + return state diff --git a/reai_toolkit/app/services/chat/schema.py b/reai_toolkit/app/services/chat/schema.py new file mode 100644 index 0000000..4936c89 --- /dev/null +++ b/reai_toolkit/app/services/chat/schema.py @@ -0,0 +1,253 @@ +"""Pure data model + event normalization for the Agent Chat feature. + +Port of the Dashboard's ``utils/v2/agent/events.ts`` (+ ``agentApi.ts`` event +decoding) and ``components/features/AgentChat/types.ts``. This module has NO Qt / +IDA / SDK imports, so it is unit-testable in isolation. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, ClassVar, Optional, Union + +EVENT_TYPE_NAMES: dict[int, str] = { + 1: "RUN_STARTED", + 2: "RUN_FINISHED", + 3: "RUN_ERROR", + 4: "STEP_STARTED", + 5: "STEP_FINISHED", + 6: "TEXT_MESSAGE_START", + 7: "TEXT_MESSAGE_CONTENT", + 8: "TEXT_MESSAGE_END", + 9: "TOOL_CALL_START", + 10: "TOOL_CALL_ARGS_DELTA", + 11: "TOOL_CALL_END", + 12: "TOOL_CALL_RESULT", + 13: "TITLE_UPDATED", + 14: "RUN_CANCELLED", + 15: "CONTEXT_COMPACTED", + 16: "TOOL_CONFIRMATION_REQUIRED", +} + +TERMINAL_EVENTS: frozenset[str] = frozenset( + {"RUN_FINISHED", "RUN_ERROR", "RUN_CANCELLED"} +) + +ROLE_USER = 2 +ROLE_SYSTEM = 3 +ROLE_TOOL = 4 + + +@dataclass +class EntityUpdate: + """A backend entity a tool result reports as changed. Drives view refresh.""" + + type: str + ids: list[int] + + +@dataclass +class ChatEvent: + """A normalized agent event. Fields are a flattened superset of every event + type; only the fields relevant to ``type`` are populated (mirrors the FE's + discriminated union after ``parseApiEvent``).""" + + type: str + message_id: Optional[str] = None + delta: Optional[str] = None + tool_call_id: Optional[str] = None + tool_name: Optional[str] = None + result: Optional[str] = None + is_error: bool = False + updated: Optional[list[EntityUpdate]] = None + message: Optional[str] = None + title: Optional[str] = None + tool_args: str = "" + role: Optional[str] = None + event_id: Optional[int] = None + + +def resolve_type(type_field: Any) -> Optional[str]: + """Resolve a wire ``type`` (string name or integer 1..16) to its name.""" + if isinstance(type_field, bool): + return None + if isinstance(type_field, str): + return type_field or None + if isinstance(type_field, int): + return EVENT_TYPE_NAMES.get(type_field) + return None + + +def _parse_entity_updates(raw: Any) -> Optional[list[EntityUpdate]]: + if not isinstance(raw, list): + return None + out: list[EntityUpdate] = [] + for item in raw: + if not isinstance(item, dict): + continue + etype = item.get("type") + ids = item.get("ids") + if etype not in ("function", "analysis", "capabilities"): + continue + if not isinstance(ids, list): + continue + try: + out.append(EntityUpdate(type=etype, ids=[int(i) for i in ids])) + except (TypeError, ValueError): + continue + return out or None + + +def normalize_event(type_field: Any, leaf: Optional[dict]) -> Optional[ChatEvent]: + """Normalize a raw SSE ``{type, data}`` frame into a :class:`ChatEvent`. + + ``leaf`` is the nested ``data`` object holding snake_case leaf fields. + Returns ``None`` for unknown/undecodable event types (mirrors the FE's + ``parseApiEvent`` returning null on parse failure). + """ + etype = resolve_type(type_field) + if etype is None: + return None + data: dict = leaf if isinstance(leaf, dict) else {} + + ev = ChatEvent(type=etype) + ev.message_id = data.get("message_id") + ev.delta = data.get("delta") + ev.tool_call_id = data.get("tool_call_id") + ev.tool_name = data.get("tool_name") + ev.result = data.get("result") + ev.is_error = bool(data.get("is_error", False)) + ev.updated = _parse_entity_updates(data.get("updated")) + ev.title = data.get("title") + ev.tool_args = data.get("tool_args") or "" + ev.message = data.get("message") or data.get("error") + if etype == "RUN_ERROR" and not ev.message: + ev.message = "Unknown error" + + if etype == "TEXT_MESSAGE_START": + ev.role = "assistant" + else: + role = data.get("role") + ev.role = role if isinstance(role, str) else None + return ev + + +@dataclass +class UserMessage: + kind: ClassVar[str] = "user-message" + id: str + content: str + + +@dataclass +class AssistantMessage: + kind: ClassVar[str] = "assistant-message" + id: str + content: str + is_streaming: bool + + +@dataclass +class FunctionRef: + """A function the agent touched, resolved to a local address for navigation.""" + + ea: int + name: str + + +@dataclass +class ToolCall: + kind: ClassVar[str] = "tool-call" + id: str + name: str + status: str + is_error: bool + functions: Optional[list[FunctionRef]] = None + + +@dataclass +class Step: + kind: ClassVar[str] = "step" + id: str + step_name: str + status: str + + +@dataclass +class ToolConfirmation: + kind: ClassVar[str] = "tool-confirmation" + id: str + tool_name: str + message: str + status: str + + +@dataclass +class ContextCompacted: + kind: ClassVar[str] = "context-compacted" + id: str + + +ChatItem = Union[ + UserMessage, + AssistantMessage, + ToolCall, + Step, + ToolConfirmation, + ContextCompacted, +] + + +@dataclass +class RunError: + message: str + code: Optional[str] = None + doc_url: Optional[str] = None + + +@dataclass +class ChatState: + items: list[ChatItem] = field(default_factory=list) + title: Optional[str] = None + run_status: str = "idle" + run_error: Optional[RunError] = None + + +@dataclass +class ConversationContextDTO: + """Plain context passed between coordinator and service (kept SDK-free).""" + + analysis_id: Optional[int] = None + function_id: Optional[int] = None + + def is_empty(self) -> bool: + return self.analysis_id is None and self.function_id is None + + +@dataclass +class UserMessageReplay: + """Reconstructed user message emitted only during history replay.""" + + id: str + content: str + + +StoredEvent = Union[ChatEvent, UserMessageReplay] + + +@dataclass +class ConversationSummary: + """Lightweight conversation row for the history browser.""" + + conversation_uuid: str + title: Optional[str] = None + updated_at: Optional[str] = None + + +@dataclass +class ConversationReplay: + """A full conversation reloaded from history, normalized for the reducer.""" + + conversation_uuid: str + title: Optional[str] + events: list[StoredEvent] diff --git a/reai_toolkit/app/services/chat/sse.py b/reai_toolkit/app/services/chat/sse.py new file mode 100644 index 0000000..19644d2 --- /dev/null +++ b/reai_toolkit/app/services/chat/sse.py @@ -0,0 +1,79 @@ +"""Pure Server-Sent-Events frame parsing for the Agent Chat stream. + +Ports the SSE read loop of ``Dashboard/utils/v2/agent/agentApi.ts::streamEvents``: +line-based framing, only ``data:`` frames, buffering across chunk boundaries, +``[DONE]`` handling and stopping on terminal events. No Qt / IDA / SDK imports, +so it is unit-testable against a fake byte-chunk iterator. +""" + +from __future__ import annotations + +import json +from typing import Callable, Iterable, Iterator, Optional + +from reai_toolkit.app.services.chat.schema import ( + TERMINAL_EVENTS, + ChatEvent, + normalize_event, +) + + +def parse_sse_data_line(line: str) -> Optional[dict]: + """Decode a single SSE line into its JSON object, or ``None``. + + Only ``data:`` lines carry payloads; ``event:`` / ``id:`` / comment (``:``) + lines and the ``[DONE]`` sentinel are ignored. + """ + line = line.strip() + if not line.startswith("data:"): + return None + payload = line[len("data:"):].strip() + if not payload or payload == "[DONE]": + return None + try: + obj = json.loads(payload) + except (ValueError, TypeError): + return None + return obj if isinstance(obj, dict) else None + + +def event_from_frame(obj: dict) -> Optional[ChatEvent]: + """Turn a decoded ``{type, event_id, data}`` envelope into a ChatEvent.""" + ev = normalize_event(obj.get("type"), obj.get("data")) + if ev is None: + return None + eid = obj.get("event_id") + if isinstance(eid, int) and not isinstance(eid, bool): + ev.event_id = eid + return ev + + +def iter_sse_events( + chunks: Iterable[bytes], + stop: Optional[Callable[[], bool]] = None, +) -> Iterator[ChatEvent]: + """Yield :class:`ChatEvent`\\ s parsed from an iterable of raw byte chunks. + + ``stop`` is polled between chunks for cooperative cancellation. Iteration + stops after a terminal event (RUN_FINISHED / RUN_ERROR / RUN_CANCELLED). + """ + buf = b"" + for chunk in chunks: + if stop is not None and stop(): + return + if not chunk: + continue + buf += chunk + parts = buf.split(b"\n") + buf = parts.pop() + for raw in parts: + line = raw.rstrip(b"\r").decode("utf-8", "replace") + obj = parse_sse_data_line(line) + if obj is None: + continue + ev = event_from_frame(obj) + if ev is None: + continue + yield ev + if ev.type in TERMINAL_EVENTS: + return diff --git a/reai_toolkit/hooks/menu.py b/reai_toolkit/hooks/menu.py index 83d9ba7..19036d5 100644 --- a/reai_toolkit/hooks/menu.py +++ b/reai_toolkit/hooks/menu.py @@ -156,6 +156,22 @@ def update(self, ctx) -> int: return ida_kernwin.AST_ENABLE +class AgentChatH(ida_kernwin.action_handler_t): + def __init__(self, coordinator: Coordinator) -> None: + super().__init__() + self.coordinator: Coordinator = coordinator + + def activate(self, ctx) -> int: + self.coordinator.chatc.run_dialog() + return 1 + + def update(self, ctx) -> int: + is_authed: bool = self.coordinator.app.auth_service.is_authenticated() + if is_authed is False or menu_hook_globals.ANALYSIS_ID is None: + return ida_kernwin.AST_DISABLE + return ida_kernwin.AST_ENABLE + + def _safe_detach(aid: str, path: str) -> None: try: ida_kernwin.detach_action_from_menu(path, aid) @@ -182,6 +198,7 @@ def register_menu_hooks(coordinator: Coordinator, plugin_version: str) -> dict: "existing_analysis": ExistingAnalysesH(coordinator), "detach_analysis": DetachAnalysisH(coordinator), "view_analysis": ViewAnalysisH(coordinator), + "agent_chat": AgentChatH(coordinator), } all_aids: list[str] = [ @@ -191,6 +208,7 @@ def register_menu_hooks(coordinator: Coordinator, plugin_version: str) -> dict: "reai:sync_and_poll", "reai:view_analysis", "reai:function_match", + "reai:agent_chat", "reai:auth", "reai:help", "reai:about", @@ -254,6 +272,12 @@ def register_menu_hooks(coordinator: Coordinator, plugin_version: str) -> dict: ) ) + ida_kernwin.register_action( + ida_kernwin.action_desc_t( + "reai:agent_chat", "Agent Chat", _handlers["agent_chat"] + ) + ) + ida_kernwin.create_menu("reai:menubar", MENU_TITLE, "View") ida_kernwin.attach_action_to_menu( MENU_ROOT + "Analysis/", "reai:analyse", ida_kernwin.SETMENU_APP @@ -278,6 +302,10 @@ def register_menu_hooks(coordinator: Coordinator, plugin_version: str) -> dict: MENU_ROOT + "Function Matching", "reai:function_match", ida_kernwin.SETMENU_APP ) + ida_kernwin.attach_action_to_menu( + MENU_ROOT + "Agent Chat", "reai:agent_chat", ida_kernwin.SETMENU_APP + ) + ida_kernwin.attach_action_to_menu( MENU_ROOT + "Configure", "reai:auth", ida_kernwin.SETMENU_APP ) diff --git a/reai_toolkit/hooks/popup.py b/reai_toolkit/hooks/popup.py index fae2eca..f4ad79b 100644 --- a/reai_toolkit/hooks/popup.py +++ b/reai_toolkit/hooks/popup.py @@ -9,6 +9,7 @@ VIEW_ACTION_NAME = "revengai:function_portal_view" MATCHING_ACTION_NAME = "revengai:function_matching_view" SIMILARITY_ACTION_NAME = "revengai:function_similarity_view" +CHAT_ACTION_NAME = "revengai:ask_about_function" _HANDLERS = {} @@ -140,12 +141,46 @@ def _register_function_similarity_action(coordinator: Coordinator): return SIMILARITY_ACTION_NAME +class AskAboutFunctionH(kw.action_handler_t): + """Opens the Agent Chat panel prefilled with the current function context.""" + + def __init__(self, coordinator: Coordinator): + super().__init__() + self.coordinator: Coordinator = coordinator + + def activate(self, ctx): + self.coordinator.chatc.run_dialog(prefill_context=True) + return 1 + + def update(self, ctx): + return kw.AST_ENABLE_FOR_WIDGET + + +def _register_chat_action(coordinator: Coordinator): + try: + kw.unregister_action(CHAT_ACTION_NAME) + except Exception: + pass + h = AskAboutFunctionH(coordinator) + _HANDLERS[CHAT_ACTION_NAME] = h + desc = kw.action_desc_t( + CHAT_ACTION_NAME, + "Ask RevEng.AI about this function", + h, + None, + "Open the RevEng.AI Agent Chat for this function", + ) + kw.register_action(desc) + return CHAT_ACTION_NAME + + def build_hooks(coordinator: Coordinator): """Attach our action directly to all popups (IDA 9.1 compatible).""" decomp_action = _register_decomp_action(coordinator) func_view_action = _register_function_view_action(coordinator) func_matching_action = _register_function_matching_action(coordinator) func_similarity_action = _register_function_similarity_action(coordinator) + chat_action = _register_chat_action(coordinator) def _on_popup(widget, popup_handle) -> None: try: @@ -169,6 +204,9 @@ def _on_popup(widget, popup_handle) -> None: kw.attach_action_to_popup( widget, popup_handle, func_similarity_action, "RevEng.AI/", 0 ) + kw.attach_action_to_popup( + widget, popup_handle, chat_action, "RevEng.AI/", 0 + ) except Exception as e: logger.error(f"[RevEng.AI] attach_action_to_popup error: {e}") diff --git a/tests/unit/chat/test_reducer.py b/tests/unit/chat/test_reducer.py new file mode 100644 index 0000000..09fe16e --- /dev/null +++ b/tests/unit/chat/test_reducer.py @@ -0,0 +1,239 @@ +"""Reducer state-machine tests — mirror every case of the Dashboard reducer.ts.""" + +from reai_toolkit.app.services.chat.reducer import ( + ApiError, + Cancel, + ConfirmTool, + EventAction, + SendMessage, + SetToolFunctions, + build_initial_state, + chat_reducer, + initial_state, +) +from reai_toolkit.app.services.chat.schema import ( + AssistantMessage, + ChatEvent, + ContextCompacted, + FunctionRef, + Step, + ToolCall, + ToolConfirmation, + UserMessage, + UserMessageReplay, +) + + +def _ev(type_, **kw): + return ChatEvent(type=type_, **kw) + + +def _fold(events, state=None): + state = state or initial_state() + for ev in events: + state = chat_reducer(state, EventAction(ev)) + return state + + +def test_send_message_appends_user_and_runs(): + state = chat_reducer(initial_state(), SendMessage(id="u1", content="hi")) + assert state.run_status == "running" + assert len(state.items) == 1 + assert isinstance(state.items[0], UserMessage) + assert state.items[0].content == "hi" + + +def test_text_message_flow_accumulates_by_id(): + state = _fold( + [ + _ev("TEXT_MESSAGE_START", message_id="a"), + _ev("TEXT_MESSAGE_CONTENT", message_id="a", delta="Hel"), + _ev("TEXT_MESSAGE_CONTENT", message_id="a", delta="lo"), + _ev("TEXT_MESSAGE_END", message_id="a"), + ] + ) + msg = state.items[-1] + assert isinstance(msg, AssistantMessage) + assert msg.content == "Hello" + assert msg.is_streaming is False + + +def test_text_message_start_is_deduped(): + state = _fold( + [ + _ev("TEXT_MESSAGE_START", message_id="a"), + _ev("TEXT_MESSAGE_START", message_id="a"), + ] + ) + assistants = [i for i in state.items if isinstance(i, AssistantMessage)] + assert len(assistants) == 1 + + +def test_tool_call_lifecycle(): + state = _fold( + [ + _ev("TOOL_CALL_START", tool_call_id="t1", tool_name="read_function"), + _ev("TOOL_CALL_END", tool_call_id="t1"), + ] + ) + tool = state.items[-1] + assert isinstance(tool, ToolCall) + assert tool.name == "read_function" + assert tool.status == "finished" + assert tool.is_error is False + + +def test_tool_call_result_sets_error_and_finished(): + state = _fold( + [ + _ev("TOOL_CALL_START", tool_call_id="t1", tool_name="do"), + _ev("TOOL_CALL_RESULT", tool_call_id="t1", tool_name="do", is_error=True), + ] + ) + tool = state.items[-1] + assert tool.status == "finished" + assert tool.is_error is True + + +def test_confirmation_then_confirm_tool_approved(): + state = _fold( + [_ev("TOOL_CONFIRMATION_REQUIRED", tool_call_id="c1", tool_name="rm", message="ok?")] + ) + state = chat_reducer(state, ConfirmTool(id="c1", approved=True)) + conf = state.items[-1] + assert isinstance(conf, ToolConfirmation) + assert conf.status == "approved" + assert state.run_status == "running" + + +def test_confirm_tool_rejected_keeps_status(): + state = _fold( + [_ev("TOOL_CONFIRMATION_REQUIRED", tool_call_id="c1", tool_name="rm", message="ok?")] + ) + state = chat_reducer(state, ConfirmTool(id="c1", approved=False)) + assert state.items[-1].status == "rejected" + assert state.run_status == "idle" + + +def test_tool_call_result_resolves_pending_confirmation(): + state = _fold( + [ + _ev("TOOL_CONFIRMATION_REQUIRED", tool_call_id="c1", tool_name="rm", message="?"), + _ev("TOOL_CALL_RESULT", tool_call_id="c1", tool_name="rm", is_error=False), + ] + ) + conf = [i for i in state.items if isinstance(i, ToolConfirmation)][0] + assert conf.status == "approved" + + +def test_tool_call_result_error_rejects_pending_confirmation(): + state = _fold( + [ + _ev("TOOL_CONFIRMATION_REQUIRED", tool_call_id="c1", tool_name="rm", message="?"), + _ev("TOOL_CALL_RESULT", tool_call_id="c1", tool_name="rm", is_error=True), + ] + ) + conf = [i for i in state.items if isinstance(i, ToolConfirmation)][0] + assert conf.status == "rejected" + + +def test_cancel_finalizes_running_items(): + state = _fold( + [ + _ev("TEXT_MESSAGE_START", message_id="a"), + _ev("TOOL_CALL_START", tool_call_id="t1", tool_name="x"), + _ev("STEP_STARTED"), + _ev("TOOL_CONFIRMATION_REQUIRED", tool_call_id="c1", tool_name="y", message="?"), + ] + ) + state = chat_reducer(state, Cancel()) + assert state.run_status == "idle" + kinds = {} + for i in state.items: + kinds[i.kind] = i + assert kinds["assistant-message"].is_streaming is False + assert kinds["tool-call"].status == "finished" + assert kinds["step"].status == "finished" + assert kinds["tool-confirmation"].status == "rejected" + + +def test_run_error_sets_error_and_finalizes(): + state = _fold([_ev("TOOL_CALL_START", tool_call_id="t1", tool_name="x")]) + state = _fold([_ev("RUN_ERROR", message="boom")], state) + assert state.run_status == "error" + assert state.run_error.message == "boom" + assert state.items[-1].status == "finished" + + +def test_run_lifecycle_status(): + state = _fold([_ev("RUN_STARTED")]) + assert state.run_status == "running" + state = _fold([_ev("RUN_FINISHED")], state) + assert state.run_status == "idle" + + +def test_run_cancelled_goes_idle(): + state = _fold([_ev("TOOL_CALL_START", tool_call_id="t", tool_name="x"), _ev("RUN_CANCELLED")]) + assert state.run_status == "idle" + assert state.items[-1].status == "finished" + + +def test_title_updated(): + state = _fold([_ev("TITLE_UPDATED", title="My chat")]) + assert state.title == "My chat" + + +def test_step_started_and_finished(): + state = _fold([_ev("STEP_STARTED")]) + assert isinstance(state.items[-1], Step) + assert state.items[-1].status == "running" + state = _fold([_ev("STEP_FINISHED")], state) + assert state.items[-1].status == "finished" + + +def test_context_compacted_appends(): + state = _fold([_ev("CONTEXT_COMPACTED")]) + assert isinstance(state.items[-1], ContextCompacted) + + +def test_tool_call_args_delta_is_noop(): + before = _fold([_ev("TOOL_CALL_START", tool_call_id="t", tool_name="x")]) + after = chat_reducer(before, EventAction(_ev("TOOL_CALL_ARGS_DELTA", tool_call_id="t", delta="{"))) + assert len(after.items) == len(before.items) + + +def test_api_error_action_finalizes(): + state = _fold([_ev("TEXT_MESSAGE_START", message_id="a")]) + state = chat_reducer(state, ApiError(message="nope", code="X", doc_url="u")) + assert state.run_status == "error" + assert state.run_error.code == "X" + assert state.items[-1].is_streaming is False + + +def test_set_tool_functions_attaches_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"), + ] + ) + 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 + + +def test_build_initial_state_replays_user_and_events(): + events = [ + UserMessageReplay(id="u1", content="what is this?"), + _ev("TEXT_MESSAGE_START", message_id="a"), + _ev("TEXT_MESSAGE_CONTENT", message_id="a", delta="an answer"), + _ev("TEXT_MESSAGE_END", message_id="a"), + _ev("TITLE_UPDATED", title="T"), + ] + state = build_initial_state(events) + assert isinstance(state.items[0], UserMessage) + assert state.items[0].content == "what is this?" + assert isinstance(state.items[1], AssistantMessage) + assert state.items[1].content == "an answer" + assert state.title == "T" diff --git a/tests/unit/chat/test_render.py b/tests/unit/chat/test_render.py new file mode 100644 index 0000000..00197d9 --- /dev/null +++ b/tests/unit/chat/test_render.py @@ -0,0 +1,97 @@ +"""Panel markdown renderer tests. + +The renderer is deliberately kept in a pure module (`chat_render`) so it is +testable headlessly. The Qt panel itself (`chat_tab`) cannot be imported under +idalib — Qt only loads in the GUI version of IDA — so, like every other Qt view +in this repo, it is covered by byte-compile + manual verification in IDA, not by +unit tests. +""" + +from reai_toolkit.app.components.tabs.chat_render import ( + find_pending_confirmation, + jump_href, + parse_jump_href, + render_transcript_markdown, + title_case, +) +from reai_toolkit.app.services.chat.schema import ( + AssistantMessage, + ChatState, + ContextCompacted, + FunctionRef, + RunError, + Step, + ToolCall, + ToolConfirmation, + UserMessage, +) + + +def test_title_case(): + assert title_case("read_function") == "Read Function" + assert title_case("") == "tool" + + +def test_render_transcript_markdown(): + state = ChatState( + items=[ + UserMessage(id="u", content="what is this?"), + AssistantMessage(id="a", content="a func", is_streaming=True), + ToolCall(id="t", name="read_function", status="running", is_error=False), + ToolConfirmation(id="c", tool_name="rename_fn", message="Rename?", status="pending"), + Step(id="s", step_name="Processing", status="running"), + ContextCompacted(id="x"), + ], + title="Demo", + run_status="running", + ) + md = render_transcript_markdown(state) + assert "**You:** what is this?" in md + assert "▍" in md + assert "Read Function" in md + assert "Approval needed" in md + assert "context compacted" in md + + +def test_render_error_state(): + state = ChatState(run_status="error", run_error=RunError(message="access denied")) + assert "access denied" in render_transcript_markdown(state) + + +def test_render_thinking_indicator(): + state = ChatState(items=[UserMessage(id="u", content="hi")], run_status="running") + assert "Thinking" in render_transcript_markdown(state) + + +def test_jump_href_roundtrip(): + assert parse_jump_href(jump_href(0x407F30)) == 0x407F30 + assert parse_jump_href("https://example.com") is None + assert parse_jump_href("ida://jump/notanumber") is None + + +def test_render_function_jump_links(): + state = ChatState( + items=[ + ToolCall( + id="t", + name="rename_functions", + status="finished", + is_error=False, + functions=[FunctionRef(ea=0x408140, name="chat_agent_renamed")], + ) + ] + ) + md = render_transcript_markdown(state) + assert "[chat_agent_renamed](ida://jump/4227392)" in md + + +def test_find_pending_confirmation(): + state = ChatState( + items=[ + ToolConfirmation(id="c1", tool_name="a", message="", status="approved"), + ToolConfirmation(id="c2", tool_name="b", message="", status="pending"), + ] + ) + pending = find_pending_confirmation(state) + assert pending is not None and pending.id == "c2" + assert find_pending_confirmation(ChatState()) is None diff --git a/tests/unit/chat/test_sdk_schemas.py b/tests/unit/chat/test_sdk_schemas.py new file mode 100644 index 0000000..3c107a0 --- /dev/null +++ b/tests/unit/chat/test_sdk_schemas.py @@ -0,0 +1,53 @@ +"""Schema-shape assertions: fail loudly if the revengai conversation API drifts.""" + +from revengai import ( + ConfirmToolInputBody, + Conversation, + ConversationContext, + ConversationsApi, + ConversationWithEvents, + CreateConversationRequest, + SendMessageRequest, +) +from revengai.models.event import Event + + +def test_conversations_api_has_expected_methods(): + for name in ( + "create_conversation", + "send_message", + "confirm_tool", + "cancel_run", + "get_conversation", + "list_conversations", + "stream_events_without_preload_content", + ): + assert hasattr(ConversationsApi, name), name + + +def test_create_conversation_request_fields(): + assert {"context", "title"} <= set(CreateConversationRequest.model_fields) + + +def test_send_message_request_fields(): + assert {"content", "context"} <= set(SendMessageRequest.model_fields) + + +def test_conversation_context_fields(): + assert {"analysis_id", "function_id"} <= set(ConversationContext.model_fields) + + +def test_confirm_tool_input_body_field(): + assert "approved" in ConfirmToolInputBody.model_fields + + +def test_conversation_fields(): + assert "conversation_uuid" in Conversation.model_fields + + +def test_conversation_with_events_fields(): + assert {"conversation_uuid", "events"} <= set(ConversationWithEvents.model_fields) + + +def test_event_fields(): + assert {"type", "role", "data", "event_id"} <= set(Event.model_fields) diff --git a/tests/unit/chat/test_service.py b/tests/unit/chat/test_service.py new file mode 100644 index 0000000..2988bf2 --- /dev/null +++ b/tests/unit/chat/test_service.py @@ -0,0 +1,161 @@ +"""ChatService tests — mock ConversationsApi, assert requests + error mapping.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from revengai import ( + ApiException, + ConversationContext, + CreateConversationRequest, + SendMessageRequest, +) + +from reai_toolkit.app.services.chat import chat_service as svc_mod +from reai_toolkit.app.services.chat.chat_service import ChatService +from reai_toolkit.app.services.chat.schema import ( + ConversationContextDTO, + UserMessageReplay, +) + + +@pytest.fixture +def service(): + return ChatService(netstore_service=MagicMock(), sdk_config=MagicMock()) + + +@pytest.fixture +def api(mocker): + mocker.patch.object(svc_mod.ChatService, "yield_api_client") + api_class = mocker.patch.object(svc_mod, "ConversationsApi") + inst = MagicMock() + api_class.return_value = inst + return inst + + +def test_create_conversation_returns_uuid_and_sends_context(service, api): + api.create_conversation.return_value = SimpleNamespace(conversation_uuid="uuid-1") + ctx = ConversationContextDTO(analysis_id=7, function_id=42) + + res = service.create_conversation(ctx, title="hi") + + assert res.success is True + assert res.data == "uuid-1" + req = api.create_conversation.call_args[0][0] + assert isinstance(req, CreateConversationRequest) + assert req.title == "hi" + assert isinstance(req.context, ConversationContext) + assert req.context.analysis_id == 7 + assert req.context.function_id == 42 + + +def test_create_conversation_error_maps(service, api): + api.create_conversation.side_effect = ApiException(status=500, reason="boom") + res = service.create_conversation(None) + assert res.success is False + assert res.error_message + + +def test_send_message_builds_request(service, api): + cid = "11111111-1111-1111-1111-111111111111" + res = service.send_message(cid, "hello", ConversationContextDTO(analysis_id=1)) + assert res.success is True + conv_arg = api.send_message.call_args[0][0] + req = api.send_message.call_args[0][1] + assert str(conv_arg) == cid + assert isinstance(req, SendMessageRequest) + assert req.content == "hello" + assert req.context.analysis_id == 1 + + +def test_send_message_409_is_tolerated(service, api): + api.send_message.side_effect = ApiException(status=409) + res = service.send_message("11111111-1111-1111-1111-111111111111", "x", None) + assert res.success is True + + +def test_send_message_other_error_fails(service, api): + api.send_message.side_effect = ApiException(status=402, reason="pay up") + res = service.send_message("11111111-1111-1111-1111-111111111111", "x", None) + assert res.success is False + + +def test_confirm_tool_404_is_tolerated(service, api): + api.confirm_tool.side_effect = ApiException(status=404) + res = service.confirm_tool("11111111-1111-1111-1111-111111111111", True) + assert res.success is True + + +def test_cancel_run_404_is_tolerated(service, api): + api.cancel_run.side_effect = ApiException(status=404) + res = service.cancel_run("11111111-1111-1111-1111-111111111111") + assert res.success is True + + +def test_list_conversations_maps_summaries(service, api): + api.list_conversations.return_value = [ + SimpleNamespace(conversation_uuid="a", title="A", updated_at="2026-01-01"), + SimpleNamespace(conversation_uuid="b", title=None, updated_at=None), + ] + res = service.list_conversations() + assert res.success is True + assert [s.conversation_uuid for s in res.data] == ["a", "b"] + assert res.data[0].title == "A" + + +def test_get_conversation_replays_events(service, api): + events = [ + SimpleNamespace(type=6, role=2, data={"message_id": "u1"}, event_id=1), + SimpleNamespace(type=7, role=2, data={"delta": "why?"}, event_id=2), + SimpleNamespace(type=8, role=2, data={}, event_id=3), + SimpleNamespace(type=6, role=1, data={"message_id": "a1"}, event_id=4), + SimpleNamespace(type=7, role=1, data={"delta": "because"}, event_id=5), + SimpleNamespace(type=8, role=1, data={}, event_id=6), + ] + api.get_conversation.return_value = SimpleNamespace( + conversation_uuid="cid", title="T", events=events + ) + res = service.get_conversation("11111111-1111-1111-1111-111111111111") + assert res.success is True + replay = res.data + assert replay.title == "T" + assert isinstance(replay.events[0], UserMessageReplay) + assert replay.events[0].content == "why?" + assert any(getattr(e, "type", None) == "TEXT_MESSAGE_CONTENT" for e in replay.events) + + +def test_get_function_name(service, mocker): + mocker.patch.object(svc_mod.ChatService, "yield_api_client") + fapi_class = mocker.patch.object(svc_mod, "FunctionsCoreApi") + inst = MagicMock() + fapi_class.return_value = inst + inst.get_function_details_0.return_value = SimpleNamespace( + function_name="chat_agent_renamed" + ) + + res = service.get_function_name(408140) + + assert res.success is True + assert res.data == "chat_agent_renamed" + inst.get_function_details_0.assert_called_once_with(function_id=408140) + + +def test_get_function_name_error_maps(service, mocker): + mocker.patch.object(svc_mod.ChatService, "yield_api_client") + fapi_class = mocker.patch.object(svc_mod, "FunctionsCoreApi") + inst = MagicMock() + fapi_class.return_value = inst + inst.get_function_details_0.side_effect = ApiException(status=404, reason="gone") + + res = service.get_function_name(1) + + assert res.success is False + assert res.error_message + + +def test_to_sdk_context_none_for_empty(): + assert ChatService._to_sdk_context(None) is None + assert ChatService._to_sdk_context(ConversationContextDTO()) is None + ctx = ChatService._to_sdk_context(ConversationContextDTO(analysis_id=3)) + assert isinstance(ctx, ConversationContext) + assert ctx.analysis_id == 3 diff --git a/tests/unit/chat/test_sse_parser.py b/tests/unit/chat/test_sse_parser.py new file mode 100644 index 0000000..e1ec4eb --- /dev/null +++ b/tests/unit/chat/test_sse_parser.py @@ -0,0 +1,145 @@ +"""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.sse import ( + event_from_frame, + iter_sse_events, + parse_sse_data_line, +) + + +def _frame(obj_json: str) -> bytes: + return f"data: {obj_json}\n".encode() + + +def test_parse_data_line_ignores_non_data(): + assert parse_sse_data_line("event: message") is None + assert parse_sse_data_line(": comment") is None + assert parse_sse_data_line("id: 5") is None + assert parse_sse_data_line("") is None + + +def test_parse_data_line_done_and_empty(): + assert parse_sse_data_line("data: [DONE]") is None + assert parse_sse_data_line("data:") is None + + +def test_parse_data_line_decodes_json(): + assert parse_sse_data_line('data: {"type": 1}') == {"type": 1} + + +def test_event_from_frame_tracks_event_id(): + ev = event_from_frame({"type": 7, "event_id": 9, "data": {"message_id": "m", "delta": "x"}}) + assert ev is not None + assert ev.type == "TEXT_MESSAGE_CONTENT" + assert ev.event_id == 9 + assert ev.delta == "x" + + +def test_integer_type_is_decoded(): + events = list(iter_sse_events([_frame('{"type": 6, "data": {"message_id": "m1"}}')])) + assert [e.type for e in events] == ["TEXT_MESSAGE_START"] + assert events[0].role == "assistant" + + +def test_string_type_passthrough(): + events = list(iter_sse_events([_frame('{"type": "TITLE_UPDATED", "data": {"title": "T"}}')])) + assert events[0].type == "TITLE_UPDATED" + assert events[0].title == "T" + + +def test_snake_case_leaf_normalization(): + frame = _frame('{"type": 9, "data": {"tool_call_id": "t1", "tool_name": "read_fn"}}') + ev = list(iter_sse_events([frame]))[0] + assert ev.tool_call_id == "t1" + assert ev.tool_name == "read_fn" + + +def test_frames_split_across_chunk_boundaries(): + whole = _frame('{"type": 7, "data": {"message_id": "m", "delta": "hello"}}') + mid = len(whole) // 2 + chunks = [whole[:mid], whole[mid:]] + events = list(iter_sse_events(chunks)) + assert len(events) == 1 + assert events[0].delta == "hello" + + +def test_crlf_line_endings(): + frame = b'data: {"type": "TITLE_UPDATED", "data": {"title": "T"}}\r\n' + events = list(iter_sse_events([frame])) + assert events[0].title == "T" + + +def test_multiple_frames_in_one_chunk(): + chunk = ( + _frame('{"type": 1}') + + _frame('{"type": 7, "data": {"message_id": "m", "delta": "a"}}') + ) + events = list(iter_sse_events([chunk])) + assert [e.type for e in events] == ["RUN_STARTED", "TEXT_MESSAGE_CONTENT"] + + +def test_event_and_id_lines_ignored(): + chunk = b"event: message\nid: 3\n" + _frame('{"type": 1}') + events = list(iter_sse_events([chunk])) + assert [e.type for e in events] == ["RUN_STARTED"] + + +def test_done_sentinel_and_blank_lines_skipped(): + chunk = b"\n" + _frame('{"type": 1}') + b"data: [DONE]\n" + events = list(iter_sse_events([chunk])) + assert [e.type for e in events] == ["RUN_STARTED"] + + +def test_terminal_event_stops_iteration(): + chunk = ( + _frame('{"type": 2}') + + _frame('{"type": 1}') + ) + events = list(iter_sse_events([chunk])) + assert [e.type for e in events] == ["RUN_FINISHED"] + + +def test_stop_callback_short_circuits(): + calls = {"n": 0} + + def stop(): + calls["n"] += 1 + return True + + events = list(iter_sse_events([_frame('{"type": 1}')], stop=stop)) + assert events == [] + + +def test_bad_json_frame_is_skipped(): + chunk = b"data: {not json}\n" + _frame('{"type": 1}') + events = list(iter_sse_events([chunk])) + assert [e.type for e in events] == ["RUN_STARTED"] + + +def test_unknown_int_type_is_dropped(): + assert normalize_event(999, {}) is None + events = list(iter_sse_events([_frame('{"type": 999, "data": {}}')])) + assert events == [] + + +def test_run_error_defaults_message(): + ev = list(iter_sse_events([_frame('{"type": 3, "data": {}}')]))[0] + assert ev.type == "RUN_ERROR" + assert ev.message == "Unknown error" + + +def test_run_error_uses_error_key_as_message(): + ev = list(iter_sse_events([_frame('{"type": 3, "data": {"error": "kaboom"}}')]))[0] + assert ev.message == "kaboom" + + +def test_tool_result_entity_updates_parsed(): + frame = _frame( + '{"type": 12, "data": {"tool_call_id": "t", "tool_name": "d", ' + '"updated": [{"type": "function", "ids": [1, 2]}]}}' + ) + ev = list(iter_sse_events([frame]))[0] + assert ev.updated is not None + assert ev.updated[0].type == "function" + assert ev.updated[0].ids == [1, 2]