diff --git a/AGENT_ORCHESTRATION_REFACTOR_PLAN.md b/AGENT_ORCHESTRATION_REFACTOR_PLAN.md new file mode 100644 index 0000000..22052b3 --- /dev/null +++ b/AGENT_ORCHESTRATION_REFACTOR_PLAN.md @@ -0,0 +1,69 @@ +# Agent Orchestration Refactor + +## Goal + +Keep every Strix agent first-class and addressable. Root and subagents are all SDK agents with +their own `SQLiteSession`; any registered agent can be messaged, stopped, resumed, and observed. +Strix should not collapse to a manager-only pattern and should not rebuild the SDK model/tool loop. + +## SDK Boundary + +- The OpenAI Agents SDK owns model/tool execution through `Runner.run_streamed`. +- The SDK session owns per-agent conversation history and message replay. +- Strix owns only product semantics the SDK does not provide: agent ids, parent/child graph, + wake/stop signals, TUI-visible status, snapshots, and process-resume metadata. +- SDK stream events are enough for tool/usage telemetry; lifecycle should not depend on custom + `RunHooks`. + +## Current File Shape + +- `strix/orchestration/coordinator.py`: graph state, runtime handles, SDK session messaging, + wake/stop signals, snapshots, resume, stream telemetry, and the interactive continuation loop. +- `strix/orchestration/scan.py`: top-level scan assembly only: sandbox setup, root construction, + resume restore, session opening, and subagent respawn. +- `strix/tools/agents_graph/tools.py`: thin tool facade over `AgentCoordinator`. +- `strix/tools/finish/tool.py`: report finalization plus active-agent guard. + +Removed custom orchestration modules: + +- `bus.py` +- `filter.py` +- `hooks.py` +- `run_loop.py` + +`bus.json` remains only as the backward-compatible snapshot filename for old runs and tracer +hydration. It is no longer backed by `AgentMessageBus`. + +## Lifecycle Semantics + +- `running`: an SDK run cycle is active. +- `waiting`: parked and addressable. +- `completed`: task completed, parked, and still addressable in interactive mode. +- `llm_failed`: parked after SDK/model failure; only user input resumes it. +- `stopped`: gracefully stopped, parked, and addressable. +- `crashed` / `failed`: parked after failure and addressable for user recovery. + +Unknown agent id is the only invalid message target. There is no routing-closed status. + +## Tool Semantics + +- `create_agent` registers a child, opens its SDK session, and starts its SDK runner task. +- `send_message_to_agent` appends a user-role item to the target SDK session and wakes the target. +- `wait_for_message` parks immediately in interactive mode; in non-interactive mode it blocks on + the coordinator wake event and returns newly appended session items to the current run. +- `agent_finish` sends the parent completion report and returns a final-output marker; the + coordinator settles status from that marker. +- `finish_scan` refuses to finalize while active agents remain and lets the coordinator settle root + status from the successful final-output marker. +- `stop_agent` uses SDK streaming cancellation when a run is active and wakes parked agents so the + continuation loop observes the stop. + +## Remaining Regression Checks To Add + +- Completed child remains messageable and resumes from its SDK session. +- Stopped/crashed/failed child remains messageable and resumes from user input. +- Interactive `wait_for_message` parks and returns control to the continuation loop. +- Non-interactive `wait_for_message` returns the newly appended session message content. +- `finish_scan` blocks while child agents are active. +- Resume rebuilds parked child runners from `bus.json` plus per-agent session DBs. +- Graceful stop works for both active-stream and parked agents. diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 57b46f6..54aeb0f 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -6,11 +6,10 @@ Jinja prompt into one ``agents.Agent`` ready for ``Runner.run``. Two flavors: - **Root** (``is_root=True``): top-level scan agent. Carries - ``finish_scan`` and stops there. + ``finish_scan`` and stops after that tool reports ``scan_completed``. - **Child** (``is_root=False``): subagents spawned by the ``create_agent`` graph tool. Carries ``agent_finish`` and stops - there — without ``stop_at_tool_names`` the SDK loop would keep - running to ``max_turns`` even after the child reported back. + after that tool reports ``agent_completed``. Skills are baked into the system prompt at scan bring-up; there's no runtime skill-loading tool. @@ -18,10 +17,11 @@ runtime skill-loading tool. from __future__ import annotations +import json import logging -from typing import Any +from typing import TYPE_CHECKING, Any -from agents.agent import StopAtTools +from agents.agent import ToolsToFinalOutputResult from agents.sandbox import SandboxAgent from agents.sandbox.capabilities import Filesystem, Shell from agents.tool import Tool @@ -65,9 +65,68 @@ from strix.tools.todo.tools import ( from strix.tools.web_search.tool import web_search +if TYPE_CHECKING: + from agents import RunContextWrapper + from agents.tool import FunctionToolResult + + logger = logging.getLogger(__name__) +def _lifecycle_tool_completed(tool_name: str, output: Any) -> bool: + if tool_name == "agent_finish": + completion_key = "agent_completed" + elif tool_name == "finish_scan": + completion_key = "scan_completed" + else: + return False + + if not isinstance(output, str): + return False + try: + parsed = json.loads(output) + except (TypeError, ValueError): + return False + return bool(isinstance(parsed, dict) and parsed.get("success") and parsed.get(completion_key)) + + +def _wait_tool_parked(tool_name: str, output: Any) -> bool: + if tool_name != "wait_for_message" or not isinstance(output, str): + return False + try: + parsed = json.loads(output) + except (TypeError, ValueError): + return False + return bool( + isinstance(parsed, dict) + and parsed.get("success") + and parsed.get("agent_waiting") + and parsed.get("status") == "waiting" + ) + + +def _finish_tool_use_behavior( + ctx: RunContextWrapper[Any], + tool_results: list[FunctionToolResult], +) -> ToolsToFinalOutputResult: + """Stop only after a lifecycle tool reports successful completion.""" + interactive = ( + bool(ctx.context.get("interactive", False)) if isinstance(ctx.context, dict) else False + ) + for tool_result in tool_results: + if _lifecycle_tool_completed(tool_result.tool.name, tool_result.output): + return ToolsToFinalOutputResult( + is_final_output=True, + final_output=tool_result.output, + ) + if interactive and _wait_tool_parked(tool_result.tool.name, tool_result.output): + return ToolsToFinalOutputResult( + is_final_output=True, + final_output=tool_result.output, + ) + return ToolsToFinalOutputResult(is_final_output=False, final_output=None) + + # Host-side Strix tools. Sandbox shell + filesystem are added per-run # by the SDK via the ``Shell`` and ``Filesystem`` capabilities below # (they bind to the live sandbox session and emit ``exec_command`` / @@ -102,7 +161,7 @@ _BASE_TOOLS: tuple[Tool, ...] = ( scope_rules, # Stateless Python execution with proxy helpers pre-bound python_action, - # Multi-agent graph tools (the bus is in ctx.context) + # Multi-agent graph tools (the coordinator is in ctx.context) view_agent_graph, agent_status, send_message_to_agent, @@ -132,13 +191,13 @@ def build_strix_agent( Responses API only). Args: - name: Agent name. Surfaces in traces and the bus's ``names`` map. + name: Agent name. Surfaces in traces and the coordinator's ``names`` map. Defaults to ``"strix"`` for the root; create_agent passes distinct names per child. skills: Skills to preload into the system prompt. is_root: Selects the tool list and ``tool_use_behavior``. - Root carries ``finish_scan`` and stops there; child carries - ``agent_finish`` and stops there. + Root carries ``finish_scan`` and child carries ``agent_finish``; + the run only stops when the lifecycle tool result succeeds. scan_mode: ``"deep"`` etc.; routes the scan-mode skill section of the prompt template. is_whitebox: Whitebox source-aware mode toggle. Adds two extra @@ -161,10 +220,8 @@ def build_strix_agent( if is_root: tools: list[Tool] = [*_BASE_TOOLS, finish_scan] - stop_at = ("finish_scan",) else: tools = [*_BASE_TOOLS, agent_finish] - stop_at = ("agent_finish",) logger.info( "Built %s agent '%s' (skills=%d, tools=%d, scan_mode=%s, whitebox=%s)", @@ -180,7 +237,7 @@ def build_strix_agent( name=name, instructions=instructions, tools=tools, - tool_use_behavior=StopAtTools(stop_at_tool_names=list(stop_at)), + tool_use_behavior=_finish_tool_use_behavior, # model=None so ``RunConfig.model`` drives provider selection # via :func:`build_multi_provider` rather than the SDK's default. model=None, diff --git a/strix/interface/tui.py b/strix/interface/tui.py index bf1e1b1..0b0ba09 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -708,12 +708,11 @@ class StrixTUIApp(App): # type: ignore[misc] self.tracer.set_scan_config(self.scan_config) set_global_tracer(self.tracer) - # Pre-create the bus here (rather than letting ``run_strix_scan`` - # build its own) so the TUI can hold a handle for stop / chat - # routing while the scan loop runs in a worker thread. - from strix.orchestration.bus import AgentMessageBus + # Pre-create the coordinator here so the TUI can route stop/chat + # commands while the scan loop runs in a worker thread. + from strix.orchestration.coordinator import AgentCoordinator - self.bus: AgentMessageBus = AgentMessageBus() + self.coordinator = AgentCoordinator() self.agent_nodes: dict[str, TreeNode] = {} @@ -938,7 +937,6 @@ class StrixTUIApp(App): # type: ignore[misc] "completed": "🟢", "failed": "🔴", "stopped": "■", - "stopping": "○", "llm_failed": "🔴", } @@ -1108,7 +1106,6 @@ class StrixTUIApp(App): # type: ignore[misc] return t simple_statuses: dict[str, tuple[str, str]] = { - "stopping": ("Agent stopping...", ""), "stopped": ("Agent stopped", ""), "completed": ("Agent completed", ""), } @@ -1402,7 +1399,7 @@ class StrixTUIApp(App): # type: ignore[misc] loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) # Stash the loop so synchronous TUI handlers (stop / - # chat) can submit bus coroutines onto it from the + # chat) can submit coordinator coroutines onto it from the # main thread. self._scan_loop = loop @@ -1416,7 +1413,7 @@ class StrixTUIApp(App): # type: ignore[misc] image=str(image), local_sources=getattr(self.args, "local_sources", None) or [], tracer=self.tracer, - bus=self.bus, + coordinator=self.coordinator, interactive=True, ), ) @@ -1473,7 +1470,6 @@ class StrixTUIApp(App): # type: ignore[misc] "completed": "🟢", "failed": "🔴", "stopped": "■", - "stopping": "○", "llm_failed": "🔴", } @@ -1546,7 +1542,6 @@ class StrixTUIApp(App): # type: ignore[misc] "completed": "🟢", "failed": "🔴", "stopped": "■", - "stopping": "○", "llm_failed": "🔴", } @@ -1730,7 +1725,7 @@ class StrixTUIApp(App): # type: ignore[misc] agent_id=self.selected_agent_id, ) - # Route to the agent's bus inbox. The scan loop runs on a + # Route to the agent's SDK session. The scan loop runs on a # worker thread; ``run_coroutine_threadsafe`` submits the # coroutine onto that loop and returns immediately so the TUI # stays responsive. After enqueuing the message, request a @@ -1741,14 +1736,14 @@ class StrixTUIApp(App): # type: ignore[misc] if self._scan_loop is not None and not self._scan_loop.is_closed(): target_agent_id = self.selected_agent_id asyncio.run_coroutine_threadsafe( - self.bus.send( + self.coordinator.send( target_agent_id, {"from": "user", "content": message, "type": "instruction"}, ), self._scan_loop, ) asyncio.run_coroutine_threadsafe( - self.bus.request_interrupt(target_agent_id, mode="after_turn"), + self.coordinator.request_interrupt(target_agent_id, mode="after_turn"), self._scan_loop, ) @@ -1832,7 +1827,7 @@ class StrixTUIApp(App): # type: ignore[misc] agent_name = agent_data.get("name", "Unknown Agent") agent_status = agent_data.get("status", "running") - if agent_status not in ["running"]: + if agent_status not in ["running", "waiting", "llm_failed"]: return agent_name, False agent_events = self._gather_agent_events(self.selected_agent_id) @@ -1849,7 +1844,7 @@ class StrixTUIApp(App): # type: ignore[misc] def action_confirm_stop_agent(self, agent_id: str) -> None: # Graceful stop: each agent's current turn finishes (and is saved to # session) before the run loop honors the cancel. The interactive - # outer loop sees ``stopping`` and exits with status="stopped". + # outer loop parks with status="stopped". # The hard ``cancel_descendants`` path remains for KeyboardInterrupt # in entry.py where graceful isn't possible. if self._scan_loop is None or self._scan_loop.is_closed(): @@ -1857,7 +1852,7 @@ class StrixTUIApp(App): # type: ignore[misc] return logger.info("TUI: graceful stop requested for %s (cascade)", agent_id) asyncio.run_coroutine_threadsafe( - self.bus.cancel_descendants_graceful(agent_id), + self.coordinator.cancel_descendants_graceful(agent_id), self._scan_loop, ) diff --git a/strix/orchestration/__init__.py b/strix/orchestration/__init__.py index 3451984..bf6a962 100644 --- a/strix/orchestration/__init__.py +++ b/strix/orchestration/__init__.py @@ -1,12 +1,12 @@ """Strix multi-agent orchestration on top of OpenAI Agents SDK. -- :class:`AgentMessageBus` — peer-to-peer agent inbox + status + stats. -- :func:`inject_messages_filter` — SDK ``call_model_input_filter`` for - inbox drain at the top of each LLM turn. -- :class:`StrixOrchestrationHooks` — SDK ``RunHooks`` subclass for - lifecycle wiring. +- :class:`AgentCoordinator` owns Strix-specific graph/status/wake state. +- SDK ``SQLiteSession`` owns per-agent conversation history and message + transport. +- :func:`run_with_continuation` wraps SDK ``Runner.run_streamed`` only + enough to keep interactive agents addressable between bounded runs. -Import deeply (``from strix.orchestration.bus import AgentMessageBus``) -so ``import strix.orchestration`` doesn't drag every submodule's deps -in eagerly. +Import deeply (for example, ``from strix.orchestration.coordinator +import AgentCoordinator``) so ``import strix.orchestration`` doesn't +drag every submodule's deps in eagerly. """ diff --git a/strix/orchestration/bus.py b/strix/orchestration/bus.py deleted file mode 100644 index 846bc14..0000000 --- a/strix/orchestration/bus.py +++ /dev/null @@ -1,461 +0,0 @@ -"""``AgentMessageBus`` — peer-to-peer multi-agent state for one scan. - -A single ``asyncio.Lock``-protected dataclass that owns inboxes, -parent edges, statuses, and per-agent stats for the lifetime of one -Strix scan. -""" - -from __future__ import annotations - -import asyncio -import contextlib -import json -import logging -import tempfile -from dataclasses import dataclass, field -from pathlib import Path -from typing import TYPE_CHECKING, Any - - -if TYPE_CHECKING: - from collections.abc import AsyncIterator - - from agents.result import RunResultStreaming - - -logger = logging.getLogger(__name__) - - -_SNAPSHOT_VERSION = 1 - - -@dataclass -class AgentMessageBus: - """Shared state for multi-agent orchestration. - - All mutations happen under ``_lock``; readers also take the lock for - consistent snapshots. The bus owns: - - - ``inboxes``: per-agent FIFO list of pending messages (drained by the - ``inject_messages_filter`` at the top of each LLM turn). - - ``tasks``: per-agent ``asyncio.Task`` handle so the parent (or signal - handler) can cancel descendants. - - ``streams``: per-agent ``RunResultStreaming`` handle so callers can - request graceful ``cancel(mode="after_turn")`` mid-stream — the SDK - saves the current turn to session before honoring the cancel. - - ``statuses``: per-agent lifecycle state — ``running | waiting | - completed | crashed | stopped | llm_failed``. - - ``parent_of``: tree edges; root agents have ``None``. - - ``names``: human-readable per-agent names. - - ``stats_live`` / ``stats_completed``: token + call counters that hooks - keep up to date for live and finalized agents respectively. Also - carries per-agent-lifetime warning flags (``warned_85``, - ``warned_final``). - - ``stopping``: agent ids whose interactive outer-loop should exit on - next iteration instead of waiting for more messages. - """ - - inboxes: dict[str, list[dict[str, Any]]] = field(default_factory=dict) - tasks: dict[str, asyncio.Task[Any]] = field(default_factory=dict) - streams: dict[str, RunResultStreaming] = field(default_factory=dict) - statuses: dict[str, str] = field(default_factory=dict) - parent_of: dict[str, str | None] = field(default_factory=dict) - names: dict[str, str] = field(default_factory=dict) - stats_live: dict[str, dict[str, Any]] = field(default_factory=dict) - stats_completed: dict[str, dict[str, Any]] = field(default_factory=dict) - stopping: set[str] = field(default_factory=set) - metadata: dict[str, dict[str, Any]] = field(default_factory=dict) - _events: dict[str, asyncio.Event] = field(default_factory=dict) - _lock: asyncio.Lock = field(default_factory=asyncio.Lock) - _snapshot_path: Path | None = None - - def set_snapshot_path(self, path: Path) -> None: - """Wire the on-disk snapshot path. Triggers persist after each lifecycle event.""" - self._snapshot_path = path - - async def register( - self, - agent_id: str, - name: str, - parent_id: str | None, - *, - task: str | None = None, - skills: list[str] | None = None, - is_whitebox: bool = False, - scan_mode: str = "deep", - diff_scope: dict[str, Any] | None = None, - ) -> None: - """Add a new agent to the bus before its Runner.run task starts. - - ``task`` / ``skills`` / ``is_whitebox`` / ``scan_mode`` / - ``diff_scope`` are persisted on the bus's ``metadata`` map so a - process restart can rebuild the same agent from the snapshot. - """ - async with self._lock: - self.inboxes[agent_id] = [] - self.statuses[agent_id] = "running" - self.parent_of[agent_id] = parent_id - self.names[agent_id] = name - self.stats_live[agent_id] = { - "in": 0, - "out": 0, - "cached": 0, - "cost": 0.0, - "calls": 0, - "warned_85": False, - "warned_final": False, - } - self.metadata[agent_id] = { - "task": task or "", - "skills": list(skills or []), - "is_whitebox": bool(is_whitebox), - "scan_mode": scan_mode, - "diff_scope": diff_scope, - } - logger.info("bus.register %s (%s) parent=%s", agent_id, name, parent_id or "-") - await self._maybe_snapshot() - - async def send(self, target: str, msg: dict[str, Any]) -> None: - """Append a message to ``target``'s inbox. - - Messages addressed to a finalized agent are dropped silently — - :meth:`finalize` clears the inbox so they can't accumulate. - """ - async with self._lock: - if target not in self.statuses: - logger.debug("bus.send dropped (unknown target=%s)", target) - return - if self.statuses[target] in ("completed", "crashed", "stopped"): - logger.debug( - "bus.send dropped (target=%s status=%s)", - target, - self.statuses[target], - ) - return - self.inboxes.setdefault(target, []).append(msg) - event = self._events.get(target) - if event is not None: - event.set() - sender = msg.get("from", "?") - msg_type = msg.get("type", "message") - content = str(msg.get("content", "")) - logger.debug( - "bus.send %s -> %s (type=%s len=%d): %s", - sender, - target, - msg_type, - len(content), - content[:200], - ) - - async def wait_for_message(self, agent_id: str) -> None: - """Block until ``agent_id``'s inbox has at least one pending message. - - Used by the interactive-mode outer loop in :func:`run_strix_scan` to - wake on the next user message between ``Runner.run`` cycles. Cheap - if the inbox already has content (returns immediately). - """ - async with self._lock: - if self.inboxes.get(agent_id): - return - event = self._events.setdefault(agent_id, asyncio.Event()) - event.clear() - await event.wait() - - async def wait_for_user_message(self, agent_id: str) -> None: - """Block until ``agent_id``'s inbox has a message with ``from='user'``. - - Used by the ``llm_failed`` recovery path: after a hard model failure, - only direct user input should resume the agent — peer messages can't - unstick a stuck model. Re-checks the inbox after each event in case - only peer messages arrived. - """ - while True: - async with self._lock: - for msg in self.inboxes.get(agent_id, []): - if msg.get("from") == "user": - return - event = self._events.setdefault(agent_id, asyncio.Event()) - event.clear() - await event.wait() - - async def drain(self, agent_id: str) -> list[dict[str, Any]]: - """Atomically read and clear ``agent_id``'s pending messages. - - Called by ``inject_messages_filter`` before every model call. - Filter output is captured by SDK in a lambda closure for retries - (verified `model_retry.py:34-35`), so a single drain per turn does - not lose messages on retry. - """ - async with self._lock: - msgs = self.inboxes.get(agent_id, []) - self.inboxes[agent_id] = [] - if msgs: - logger.debug("bus.drain %s -> %d message(s)", agent_id, len(msgs)) - return msgs - - async def record_usage(self, agent_id: str, usage: Any) -> None: - """Accumulate per-call usage from RunHooks.on_llm_end. - - Tolerates ``usage=None`` (some providers omit usage on streaming). - Increments ``calls`` unconditionally so it doubles as a per-agent - lifetime turn counter (legacy ``state.iteration`` parity). - """ - async with self._lock: - stats = self.stats_live.setdefault( - agent_id, - { - "in": 0, - "out": 0, - "cached": 0, - "cost": 0.0, - "calls": 0, - "warned_85": False, - "warned_final": False, - }, - ) - stats["calls"] += 1 - if usage is None: - return - stats["in"] += getattr(usage, "input_tokens", 0) or 0 - stats["out"] += getattr(usage, "output_tokens", 0) or 0 - details = getattr(usage, "input_tokens_details", None) - if details is not None: - stats["cached"] += getattr(details, "cached_tokens", 0) or 0 - - async def finalize(self, agent_id: str, status: str) -> None: - """Move an agent from live to completed; clean up routing state. - - Clears live routing state (``inboxes``, ``streams``, - ``_events``, ``stopping``, ``metadata`` for the spawn args) and - moves stats from ``stats_live`` to ``stats_completed``. Keeps - ``parent_of`` and ``names`` so the agent remains visible in - ``view_agent_graph`` and the TUI tree post-finalize. Routing - protection against late ``send`` calls comes from the - ``statuses[id]`` terminal-state check in :meth:`send`, not from - removing the parent/name keys. - """ - async with self._lock: - self.statuses[agent_id] = status - self.stats_completed[agent_id] = self.stats_live.pop(agent_id, {}) - self.inboxes.pop(agent_id, None) - self.streams.pop(agent_id, None) - self.stopping.discard(agent_id) - self._events.pop(agent_id, None) - self.metadata.pop(agent_id, None) - logger.info("bus.finalize %s status=%s", agent_id, status) - await self._maybe_snapshot() - - async def park(self, agent_id: str) -> None: - """Mark an agent as ``waiting`` without finalizing. - - Used in interactive mode for the root agent between ``Runner.run`` - cycles: the run completed, but the agent stays alive on the bus - so user messages still land in its inbox until the next cycle - starts. Stats stay live (will be merged on actual finalize at - scan teardown). - """ - async with self._lock: - if agent_id in self.statuses: - self.statuses[agent_id] = "waiting" - logger.debug("bus.park %s", agent_id) - await self._maybe_snapshot() - - async def mark_llm_failed(self, agent_id: str) -> None: - """Mark an agent as ``llm_failed`` after retries exhausted. - - Mirrors legacy ``state.llm_failed`` semantics: only direct user - input can resume the agent (see :meth:`wait_for_user_message`). - Status survives until the next ``Runner.run`` cycle starts and - ``on_agent_start`` mirrors it back to ``running``, or finalize - clears it. - """ - async with self._lock: - if agent_id in self.statuses: - self.statuses[agent_id] = "llm_failed" - logger.warning("bus.mark_llm_failed %s — awaiting user resume", agent_id) - await self._maybe_snapshot() - - @contextlib.asynccontextmanager - async def attach_stream( - self, - agent_id: str, - streamed: RunResultStreaming, - ) -> AsyncIterator[None]: - """Register ``streamed`` so ``request_interrupt`` can find it; clean up after.""" - async with self._lock: - self.streams[agent_id] = streamed - try: - yield - finally: - async with self._lock: - if self.streams.get(agent_id) is streamed: - self.streams.pop(agent_id, None) - - async def request_interrupt( - self, - agent_id: str, - mode: str = "after_turn", - ) -> bool: - """Ask the agent's active streaming run to cancel gracefully. - - Returns True if a streaming run was attached (so a cancel request - was issued), False otherwise. ``mode='after_turn'`` lets the SDK - finish the current turn — including saving items to session — so - cancellation never leaves orphan tool outputs or truncated - assistant messages. ``mode='immediate'`` is the hard variant. - """ - async with self._lock: - streamed = self.streams.get(agent_id) - if streamed is None: - logger.debug("bus.request_interrupt %s — no active stream", agent_id) - return False - streamed.cancel(mode=mode) # type: ignore[arg-type] # mode is a Literal - logger.info("bus.request_interrupt %s mode=%s", agent_id, mode) - return True - - async def total_stats(self) -> dict[str, Any]: - """Snapshot of live + completed stats. Excludes warning flags.""" - async with self._lock: - agg = {"in": 0, "out": 0, "cached": 0, "cost": 0.0, "calls": 0} - for stats in (*self.stats_live.values(), *self.stats_completed.values()): - for key in agg: - agg[key] += stats.get(key, 0) - return agg - - async def cancel_descendants(self, root_agent_id: str) -> None: - """Cancel ``root_agent_id`` and every transitive child, leaves first. - - Wired into the CLI Ctrl+C handler and TUI stop button — - the SDK's ``result.cancel`` doesn't cascade to children spawned - via ``asyncio.create_task``, so we walk the tree ourselves. - - This is the **hard** path: ``task.cancel()`` raises ``CancelledError`` - immediately, which may truncate a turn mid-stream. For graceful - cascading stops use :meth:`cancel_descendants_graceful`. - """ - async with self._lock: - queue = [root_agent_id] - order: list[str] = [] - while queue: - aid = queue.pop() - order.append(aid) - queue.extend(child for child, parent in self.parent_of.items() if parent == aid) - tasks_to_cancel = [self.tasks[a] for a in reversed(order) if a in self.tasks] - logger.info( - "bus.cancel_descendants %s (hard, %d task(s))", - root_agent_id, - len(tasks_to_cancel), - ) - for task in tasks_to_cancel: - if not task.done(): - task.cancel() - # Wait for cancellations to settle so on_agent_end can mark statuses. - await asyncio.gather( - *(t for t in tasks_to_cancel if not t.done()), - return_exceptions=True, - ) - - async def cancel_descendants_graceful(self, root_agent_id: str) -> None: - """Graceful cascade: ``request_interrupt`` per node, leaves-first. - - Each node's current turn finishes (and is saved to session) before - the run loop honors the cancel. The interactive outer loop sees - the agent in ``stopping`` and returns instead of awaiting more - messages, so finalize fires with status="stopped". - """ - async with self._lock: - queue = [root_agent_id] - order: list[str] = [] - while queue: - aid = queue.pop() - order.append(aid) - queue.extend(child for child, parent in self.parent_of.items() if parent == aid) - for aid in order: - self.stopping.add(aid) - streams_to_cancel = [ - (aid, self.streams[aid]) for aid in reversed(order) if aid in self.streams - ] - logger.info( - "bus.cancel_descendants_graceful %s (%d active stream(s), %d total)", - root_agent_id, - len(streams_to_cancel), - len(order), - ) - for _aid, streamed in streams_to_cancel: - streamed.cancel(mode="after_turn") - # Persist ``stopping`` so a process crash before any child finalizes - # doesn't lose the user's stop signal — otherwise resume would - # respawn the cancelled agents and they'd run forever. - await self._maybe_snapshot() - - # ------------------------------------------------------------------ - # Snapshot / restore — persist serializable state to ``bus.json`` so - # a process restart can rebuild the topology + per-agent metadata - # and respawn each non-terminal agent. - # ------------------------------------------------------------------ - async def snapshot(self) -> dict[str, Any]: - """Return a JSON-ready deep copy of every persistable bus field. - - Held under ``_lock`` for the copy so the caller can't observe a - partial mutation. The actual JSON serialisation happens outside - the lock — that's fine for ``json.dumps``: pure-Python primitives, - no I/O. - """ - async with self._lock: - return { - "version": _SNAPSHOT_VERSION, - "inboxes": {aid: list(msgs) for aid, msgs in self.inboxes.items()}, - "statuses": dict(self.statuses), - "parent_of": dict(self.parent_of), - "names": dict(self.names), - "stats_live": {aid: dict(s) for aid, s in self.stats_live.items()}, - "stats_completed": {aid: dict(s) for aid, s in self.stats_completed.items()}, - "stopping": list(self.stopping), - "metadata": {aid: dict(m) for aid, m in self.metadata.items()}, - } - - async def restore(self, snap: dict[str, Any]) -> None: - """Repopulate from a prior :meth:`snapshot`. Tasks/streams/events - stay empty — they're rebuilt by the resume path as agents respawn. - """ - async with self._lock: - self.inboxes = {aid: list(msgs) for aid, msgs in snap.get("inboxes", {}).items()} - self.statuses = dict(snap.get("statuses", {})) - self.parent_of = dict(snap.get("parent_of", {})) - self.names = dict(snap.get("names", {})) - self.stats_live = {aid: dict(s) for aid, s in snap.get("stats_live", {}).items()} - self.stats_completed = { - aid: dict(s) for aid, s in snap.get("stats_completed", {}).items() - } - self.stopping = set(snap.get("stopping", [])) - self.metadata = {aid: dict(m) for aid, m in snap.get("metadata", {}).items()} - - async def _maybe_snapshot(self) -> None: - """Persist the current state to ``_snapshot_path`` if one is set. - - No-op when ``_snapshot_path`` is unset (e.g. in tests). Atomic - ``tempfile`` + ``os.replace`` so a crash mid-write leaves the - previous valid file intact. Errors are logged + swallowed; a - snapshot failure should never tear down the run. - """ - path = self._snapshot_path - if path is None: - return - try: - data = await self.snapshot() - payload = json.dumps(data, ensure_ascii=False, default=str) - path.parent.mkdir(parents=True, exist_ok=True) - with tempfile.NamedTemporaryFile( - mode="w", - encoding="utf-8", - dir=str(path.parent), - prefix=f".{path.name}.", - suffix=".tmp", - delete=False, - ) as tmp: - tmp.write(payload) - tmp_path = Path(tmp.name) - tmp_path.replace(path) - except Exception: - logger.exception("bus snapshot to %s failed", path) diff --git a/strix/orchestration/coordinator.py b/strix/orchestration/coordinator.py new file mode 100644 index 0000000..c382d46 --- /dev/null +++ b/strix/orchestration/coordinator.py @@ -0,0 +1,712 @@ +"""SDK-native coordinator for Strix's addressable agent graph. + +The Agents SDK owns model/tool execution and per-agent conversation +history. Strix owns only product semantics the SDK does not provide: +agent ids, the parent/child graph, wake/stop signals, TUI-visible +status, and process-resume metadata. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import tempfile +from dataclasses import dataclass, field +from datetime import UTC, datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal + +from agents import Runner +from agents.exceptions import AgentsException, MaxTurnsExceeded, UserError +from agents.memory import SQLiteSession +from openai import APIError + + +if TYPE_CHECKING: + from agents.items import TResponseInputItem + from agents.memory import Session + from agents.result import RunResultBase, RunResultStreaming + from agents.run_config import RunConfig + + +logger = logging.getLogger(__name__) + +Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed", "llm_failed"] +ACTIVE_AGENT_STATUSES = {"running", "waiting", "llm_failed"} +_SNAPSHOT_VERSION = 2 +_WAITING_TIMEOUT_SUBAGENT = 300.0 +_TIMEOUT_RESUME_MESSAGE = "Waiting timeout reached. Resuming execution." + + +@dataclass(slots=True) +class AgentRuntime: + session: Session | None = None + task: asyncio.Task[Any] | None = None + stream: RunResultStreaming | None = None + wake: asyncio.Event = field(default_factory=asyncio.Event) + tool_calls: dict[str, str] = field(default_factory=dict) + + +class AgentCoordinator: + """Single owner for graph state, SDK runtimes, messages, and resume snapshots.""" + + def __init__(self) -> None: + self.statuses: dict[str, Status] = {} + self.parent_of: dict[str, str | None] = {} + self.names: dict[str, str] = {} + self.metadata: dict[str, dict[str, Any]] = {} + self.pending_counts: dict[str, int] = {} + self.pending_user_counts: dict[str, int] = {} + self.queued_messages: dict[str, list[dict[str, Any]]] = {} + self.stats_live: dict[str, dict[str, Any]] = {} + self.runtimes: dict[str, AgentRuntime] = {} + self._lock = asyncio.Lock() + self._snapshot_path: Path | None = None + + def set_snapshot_path(self, path: Path) -> None: + self._snapshot_path = path + + async def register( + self, + agent_id: str, + name: str, + parent_id: str | None, + *, + task: str | None = None, + skills: list[str] | None = None, + is_whitebox: bool = False, + scan_mode: str = "deep", + diff_scope: dict[str, Any] | None = None, + session_path: str | Path | None = None, + ) -> None: + async with self._lock: + self.statuses[agent_id] = "running" + self.parent_of[agent_id] = parent_id + self.names[agent_id] = name + self.pending_counts.setdefault(agent_id, 0) + self.pending_user_counts.setdefault(agent_id, 0) + self.stats_live.setdefault(agent_id, _empty_stats()) + self.metadata[agent_id] = { + "task": task or "", + "skills": list(skills or []), + "is_whitebox": bool(is_whitebox), + "scan_mode": scan_mode, + "diff_scope": diff_scope, + "session_path": str(session_path) if session_path is not None else "", + } + self.runtimes.setdefault(agent_id, AgentRuntime()) + logger.info("agent.register %s (%s) parent=%s", agent_id, name, parent_id or "-") + await self._maybe_snapshot() + + async def attach_runtime( + self, + agent_id: str, + *, + session: Session | None = None, + task: asyncio.Task[Any] | None = None, + ) -> None: + async with self._lock: + runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) + if session is not None: + runtime.session = session + if task is not None: + runtime.task = task + if session is not None: + await self.flush_queued_messages(agent_id) + + async def mark_running(self, agent_id: str) -> None: + async with self._lock: + if agent_id in self.statuses: + self.statuses[agent_id] = "running" + await self._maybe_snapshot() + + async def park_waiting(self, agent_id: str) -> None: + await self.set_status(agent_id, "waiting") + + async def mark_llm_failed(self, agent_id: str) -> None: + await self.set_status(agent_id, "llm_failed") + + async def set_status(self, agent_id: str, status: Status | str) -> None: + async with self._lock: + if agent_id not in self.statuses: + return + self.statuses[agent_id] = status # type: ignore[assignment] + runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) + runtime.wake.set() + logger.info("agent.status %s=%s", agent_id, status) + await self._maybe_snapshot() + + async def send(self, target_agent_id: str, message: dict[str, Any]) -> bool: + """Deliver a user/peer message by appending it to the target SDK session.""" + should_queue = False + async with self._lock: + if target_agent_id not in self.statuses: + logger.debug("agent.send dropped unknown target=%s", target_agent_id) + return False + runtime = self.runtimes.setdefault(target_agent_id, AgentRuntime()) + session = runtime.session + should_queue = session is None or runtime.stream is not None + if should_queue: + self.queued_messages.setdefault(target_agent_id, []).append(dict(message)) + if session is not None and not should_queue: + try: + await session.add_items([self._message_to_session_item(message)]) + except Exception: + logger.exception( + "agent.send failed to append to SDK session target=%s", + target_agent_id, + ) + return False + async with self._lock: + self.pending_counts[target_agent_id] = self.pending_counts.get(target_agent_id, 0) + 1 + if message.get("from") == "user": + self.pending_user_counts[target_agent_id] = ( + self.pending_user_counts.get(target_agent_id, 0) + 1 + ) + self.runtimes.setdefault(target_agent_id, AgentRuntime()).wake.set() + if should_queue: + logger.debug( + "agent.send %s queued until SDK session is safe to append", target_agent_id + ) + await self._maybe_snapshot() + return True + + async def flush_queued_messages(self, agent_id: str) -> None: + async with self._lock: + runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) + session = runtime.session + queued = self.queued_messages.pop(agent_id, []) + if not queued: + return + if session is None: + async with self._lock: + self.queued_messages.setdefault(agent_id, []).extend(queued) + return + try: + await session.add_items([self._message_to_session_item(msg) for msg in queued]) + except Exception: + async with self._lock: + self.queued_messages.setdefault(agent_id, [])[0:0] = queued + logger.exception("agent.flush_queued_messages failed for %s", agent_id) + + async def wait_for_message(self, agent_id: str, *, user_only: bool = False) -> None: + while True: + async with self._lock: + pending = ( + self.pending_user_counts.get(agent_id, 0) + if user_only + else self.pending_counts.get(agent_id, 0) + ) + if pending > 0: + return + wake = self.runtimes.setdefault(agent_id, AgentRuntime()).wake + wake.clear() + await wake.wait() + + async def wait_for_user_message(self, agent_id: str) -> None: + await self.wait_for_message(agent_id, user_only=True) + + async def consume_wake(self, agent_id: str) -> None: + async with self._lock: + self.pending_counts[agent_id] = 0 + self.pending_user_counts[agent_id] = 0 + + async def pending_count(self, agent_id: str) -> int: + async with self._lock: + return self.pending_counts.get(agent_id, 0) + + async def recent_session_items(self, agent_id: str, count: int) -> list[TResponseInputItem]: + if count <= 0: + return [] + async with self._lock: + session = self.runtimes.get(agent_id, AgentRuntime()).session + if session is None: + return [] + items = await session.get_items() + return list(items[-count:]) + + async def request_interrupt(self, agent_id: str, mode: str = "after_turn") -> bool: + async with self._lock: + stream = self.runtimes.get(agent_id, AgentRuntime()).stream + if stream is None: + return False + stream.cancel(mode=mode) # type: ignore[arg-type] + return True + + async def request_stop(self, agent_id: str, *, interrupt: bool = True) -> None: + async with self._lock: + if agent_id not in self.statuses: + return + self.statuses[agent_id] = "stopped" + runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) + runtime.wake.set() + stream = runtime.stream + if interrupt and stream is not None: + stream.cancel(mode="after_turn") + await self._maybe_snapshot() + + async def cancel_descendants(self, agent_id: str) -> None: + tasks = [] + async with self._lock: + for aid in reversed(self._subtree_order_locked(agent_id)): + task = self.runtimes.get(aid, AgentRuntime()).task + if task is not None and not task.done(): + tasks.append(task) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + async def cancel_descendants_graceful(self, agent_id: str) -> None: + async with self._lock: + order = self._subtree_order_locked(agent_id) + for aid in reversed(order): + await self.request_stop(aid) + await self._maybe_snapshot() + + async def attach_stream( + self, + agent_id: str, + stream: RunResultStreaming, + ) -> None: + async with self._lock: + self.runtimes.setdefault(agent_id, AgentRuntime()).stream = stream + + async def detach_stream( + self, + agent_id: str, + stream: RunResultStreaming, + ) -> None: + async with self._lock: + runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) + if runtime.stream is stream: + runtime.stream = None + + async def active_agents_except(self, agent_id: str) -> list[dict[str, Any]]: + async with self._lock: + return [ + self._agent_info_locked(aid) + for aid, status in self.statuses.items() + if aid != agent_id and status in ACTIVE_AGENT_STATUSES + ] + + async def agent_info(self, agent_id: str) -> dict[str, Any] | None: + async with self._lock: + if agent_id not in self.statuses: + return None + return self._agent_info_locked(agent_id) + + async def graph_snapshot( + self, + ) -> tuple[dict[str, str | None], dict[str, Status], dict[str, str]]: + async with self._lock: + return dict(self.parent_of), dict(self.statuses), dict(self.names) + + async def record_usage(self, agent_id: str, usage: Any) -> None: + if usage is None: + return + async with self._lock: + stats = self.stats_live.setdefault(agent_id, _empty_stats()) + stats["calls"] += 1 + stats["in"] += getattr(usage, "input_tokens", 0) or 0 + stats["out"] += getattr(usage, "output_tokens", 0) or 0 + details = getattr(usage, "input_tokens_details", None) + stats["cached"] += getattr(details, "cached_tokens", 0) or 0 if details else 0 + + def _message_to_session_item(self, message: dict[str, Any]) -> TResponseInputItem: + sender = str(message.get("from", "unknown")) + content = str(message.get("content", "")) + if sender == "user": + return {"role": "user", "content": content} + sender_name = self.names.get(sender, sender) + msg_type = message.get("type", "information") + priority = message.get("priority", "normal") + return { + "role": "user", + "content": ( + f"[Message from {sender_name} ({sender}) | type={msg_type} " + f"| priority={priority}]\n{content}" + ), + } + + def _subtree_order_locked(self, agent_id: str) -> list[str]: + queue = [agent_id] + order: list[str] = [] + while queue: + aid = queue.pop() + order.append(aid) + queue.extend(child for child, parent in self.parent_of.items() if parent == aid) + return order + + def _agent_info_locked(self, agent_id: str) -> dict[str, Any]: + return { + "agent_id": agent_id, + "name": self.names.get(agent_id, agent_id), + "status": self.statuses.get(agent_id), + "parent_id": self.parent_of.get(agent_id), + "pending_messages": self.pending_counts.get(agent_id, 0), + } + + async def snapshot(self) -> dict[str, Any]: + async with self._lock: + return { + "version": _SNAPSHOT_VERSION, + "statuses": dict(self.statuses), + "parent_of": dict(self.parent_of), + "names": dict(self.names), + "metadata": {aid: dict(md) for aid, md in self.metadata.items()}, + "pending_counts": dict(self.pending_counts), + "pending_user_counts": dict(self.pending_user_counts), + "queued_messages": { + aid: [dict(msg) for msg in msgs] for aid, msgs in self.queued_messages.items() + }, + "stats_live": {aid: dict(s) for aid, s in self.stats_live.items()}, + # Kept for old tracer hydration shape. + "stats_completed": {}, + } + + async def restore(self, snap: dict[str, Any]) -> None: + async with self._lock: + self.statuses = dict(snap.get("statuses", {})) + self.parent_of = dict(snap.get("parent_of", {})) + self.names = dict(snap.get("names", {})) + self.metadata = {aid: dict(md) for aid, md in snap.get("metadata", {}).items()} + self.pending_counts = dict(snap.get("pending_counts", {})) + self.pending_user_counts = dict(snap.get("pending_user_counts", {})) + self.queued_messages = { + aid: [dict(msg) for msg in msgs] + for aid, msgs in snap.get("queued_messages", {}).items() + } + for aid in snap.get("stopping", []): + if aid in self.statuses: + self.statuses[aid] = "stopped" + self.stats_live = {aid: dict(s) for aid, s in snap.get("stats_live", {}).items()} + for aid, msgs in snap.get("inboxes", {}).items(): + # Legacy bus snapshots used inboxes. Preserve them as queued + # messages so attaching the SDK session makes them visible. + self.pending_counts[aid] = max(self.pending_counts.get(aid, 0), len(msgs)) + self.queued_messages.setdefault(aid, []).extend(dict(msg) for msg in msgs) + for aid in self.statuses: + self.runtimes.setdefault(aid, AgentRuntime()) + + async def _maybe_snapshot(self) -> None: + path = self._snapshot_path + if path is None: + return + try: + data = await self.snapshot() + payload = json.dumps(data, ensure_ascii=False, default=str) + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=str(path.parent), + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as tmp: + tmp.write(payload) + tmp_path = Path(tmp.name) + tmp_path.replace(path) + except Exception: + logger.exception("coordinator snapshot to %s failed", path) + + +def coordinator_from_context(ctx: dict[str, Any]) -> AgentCoordinator | None: + coordinator = ctx.get("coordinator") + return coordinator if isinstance(coordinator, AgentCoordinator) else None + + +async def run_with_continuation( + *, + agent: Any, + initial_input: Any, + run_config: RunConfig, + context: dict[str, Any], + max_turns: int, + coordinator: AgentCoordinator, + agent_id: str, + interactive: bool, + session: Session | None = None, + start_parked: bool = False, +) -> RunResultBase | None: + await coordinator.attach_runtime(agent_id, session=session) + waiting_timeout = await _waiting_timeout(coordinator, agent_id, interactive) + result: RunResultBase | None = None + + if not (start_parked and interactive): + result = await _run_cycle( + agent, + coordinator, + agent_id, + input_data=initial_input, + run_config=run_config, + context=context, + max_turns=max_turns, + session=session, + interactive=interactive, + ) + + if not interactive: + return result + + while True: + async with coordinator._lock: + status = coordinator.statuses.get(agent_id) + + try: + if status == "llm_failed": + await coordinator.wait_for_user_message(agent_id) + elif waiting_timeout is None: + await coordinator.wait_for_message(agent_id) + else: + await asyncio.wait_for( + coordinator.wait_for_message(agent_id), + timeout=waiting_timeout, + ) + except asyncio.CancelledError: + return result + except TimeoutError: + result = await _run_cycle( + agent, + coordinator, + agent_id, + input_data=_TIMEOUT_RESUME_MESSAGE, + run_config=run_config, + context=context, + max_turns=max_turns, + session=session, + interactive=interactive, + ) + continue + + await coordinator.consume_wake(agent_id) + result = await _run_cycle( + agent, + coordinator, + agent_id, + input_data=[], + run_config=run_config, + context=context, + max_turns=max_turns, + session=session, + interactive=interactive, + ) + + +async def _waiting_timeout( + coordinator: AgentCoordinator, + agent_id: str, + interactive: bool, +) -> float | None: + if not interactive: + return None + async with coordinator._lock: + return ( + _WAITING_TIMEOUT_SUBAGENT if coordinator.parent_of.get(agent_id) is not None else None + ) + + +async def _run_cycle( + agent: Any, + coordinator: AgentCoordinator, + agent_id: str, + *, + input_data: Any, + run_config: RunConfig, + context: dict[str, Any], + max_turns: int, + session: Session | None, + interactive: bool, +) -> RunResultBase | None: + try: + await coordinator.mark_running(agent_id) + stream = Runner.run_streamed( + agent, + input=input_data, + run_config=run_config, + context=context, + max_turns=max_turns, + session=session, + ) + await coordinator.attach_stream(agent_id, stream) + try: + async for event in stream.stream_events(): + await _handle_stream_event(coordinator, agent_id, context, event) + finally: + await coordinator.detach_stream(agent_id, stream) + await coordinator.flush_queued_messages(agent_id) + await _settle_run_result(coordinator, agent_id, stream.final_output, interactive, context) + return stream + except (AgentsException, APIError): + if not interactive: + raise + logger.exception("LLM/runtime failure for %s; waiting for user resume", agent_id) + await coordinator.mark_llm_failed(agent_id) + return None + except Exception as exc: + if not interactive: + raise + status: Status = "stopped" if isinstance(exc, MaxTurnsExceeded) else "crashed" + if isinstance(exc, UserError): + status = "failed" + logger.exception("agent run failed for %s; parking as %s", agent_id, status) + await coordinator.set_status(agent_id, status) + await _notify_parent_on_crash(coordinator, agent_id, status) + _mirror_tracer_status(context, coordinator, agent_id, status) + return None + + +async def _settle_run_result( + coordinator: AgentCoordinator, + agent_id: str, + final_output: Any, + interactive: bool, + context: dict[str, Any], +) -> None: + parsed = _parse_final_output(final_output) + async with coordinator._lock: + current_status = coordinator.statuses.get(agent_id) + + if current_status == "stopped": + status: Status = "stopped" + elif parsed.get("agent_completed") or parsed.get("scan_completed"): + status = "completed" + elif parsed.get("agent_waiting") or interactive: + status = "waiting" + else: + status = "crashed" + + await coordinator.set_status(agent_id, status) + await _notify_parent_on_crash(coordinator, agent_id, status) + _mirror_tracer_status(context, coordinator, agent_id, status) + + +def _parse_final_output(output: Any) -> dict[str, Any]: + if not isinstance(output, str): + return {} + try: + parsed = json.loads(output) + except (TypeError, ValueError): + return {} + return parsed if isinstance(parsed, dict) and parsed.get("success") else {} + + +async def _notify_parent_on_crash( + coordinator: AgentCoordinator, + agent_id: str, + status: str, +) -> None: + if status != "crashed": + return + async with coordinator._lock: + parent = coordinator.parent_of.get(agent_id) + name = coordinator.names.get(agent_id, agent_id) + if parent is None: + return + await coordinator.send( + parent, + { + "from": agent_id, + "type": "crash", + "priority": "high", + "content": ( + f"[Agent crash] {name} ({agent_id}) terminated unexpectedly. " + "Stop waiting on this child unless you want to message it again." + ), + }, + ) + + +def _mirror_tracer_status( + context: dict[str, Any], + coordinator: AgentCoordinator, + agent_id: str, + status: str, +) -> None: + tracer = context.get("tracer") + if tracer is None: + return + now = datetime.now(UTC).isoformat() + tracer.agents.setdefault( + agent_id, + { + "id": agent_id, + "name": coordinator.names.get(agent_id, agent_id), + "parent_id": coordinator.parent_of.get(agent_id), + "created_at": now, + }, + ) + tracer.agents[agent_id]["status"] = status + tracer.agents[agent_id]["updated_at"] = now + + +async def _handle_stream_event( + coordinator: AgentCoordinator, + agent_id: str, + context: dict[str, Any], + event: Any, +) -> None: + tracer = context.get("tracer") + if event.type == "raw_response_event": + response = getattr(event.data, "response", None) + usage = getattr(response, "usage", None) + if usage is not None: + await coordinator.record_usage(agent_id, usage) + if usage is not None and tracer is not None and hasattr(tracer, "record_llm_usage"): + details = getattr(usage, "input_tokens_details", None) + tracer.record_llm_usage( + input_tokens=int(getattr(usage, "input_tokens", 0) or 0), + output_tokens=int(getattr(usage, "output_tokens", 0) or 0), + cached_tokens=int(getattr(details, "cached_tokens", 0) or 0) if details else 0, + ) + return + if tracer is None or event.type != "run_item_stream_event": + return + item = event.item + raw = getattr(item, "raw_item", None) + if event.name == "tool_called": + call_id = str(getattr(raw, "call_id", None) or getattr(raw, "id", "")) + tool_name = str(getattr(raw, "name", None) or getattr(raw, "type", "tool")) + args = _parse_tool_args(getattr(raw, "arguments", None)) + runtime = coordinator.runtimes.setdefault(agent_id, AgentRuntime()) + if call_id: + runtime.tool_calls[call_id] = tool_name + tracer.log_tool_start(agent_id, tool_name, args) + elif event.name == "tool_output": + call_id = str(getattr(raw, "call_id", None) or "") + runtime = coordinator.runtimes.setdefault(agent_id, AgentRuntime()) + tool_name = runtime.tool_calls.get(call_id, "tool") + tracer.log_tool_end(agent_id, tool_name, _dump_raw(raw)) + + +def _parse_tool_args(raw: Any) -> dict[str, Any]: + if not raw: + return {} + try: + parsed = json.loads(raw) + except (TypeError, ValueError): + return {} + return parsed if isinstance(parsed, dict) else {} + + +def _dump_raw(raw: Any) -> Any: + if hasattr(raw, "model_dump"): + return raw.model_dump(exclude_unset=True) + return raw + + +def open_agent_session(agent_id: str, path: Path) -> SQLiteSession: + path.parent.mkdir(parents=True, exist_ok=True) + return SQLiteSession(session_id=agent_id, db_path=path) + + +def _empty_stats() -> dict[str, Any]: + return { + "in": 0, + "out": 0, + "cached": 0, + "cost": 0.0, + "calls": 0, + } diff --git a/strix/orchestration/filter.py b/strix/orchestration/filter.py deleted file mode 100644 index 88b73d1..0000000 --- a/strix/orchestration/filter.py +++ /dev/null @@ -1,96 +0,0 @@ -"""``inject_messages_filter`` — SDK ``call_model_input_filter`` for the bus. - -The SDK runs ``call_model_input_filter`` exactly once per turn before -the LLM call and captures the output in a closure for any subsequent -retries — so a single drain per turn doesn't lose messages on retry. -""" - -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING - -from agents.run_config import CallModelData, ModelInputData - - -if TYPE_CHECKING: - from typing import Any - - from strix.orchestration.bus import AgentMessageBus - - -logger = logging.getLogger(__name__) - - -async def inject_messages_filter(data: CallModelData) -> ModelInputData: - """Drain bus inbox and append messages as user-role items before the LLM call. - - Peer-agent messages get a one-line labeled header followed by the - body. Direct user messages (``from="user"``) are passed plain. - - Any exception inside the filter — malformed message dict, bug in - ``bus.drain``, etc. — is caught and the original ``data.model_data`` - is returned unmodified. A bug here must never tear down the run. - """ - try: - if not isinstance(data.context, dict): - return data.model_data - bus: AgentMessageBus | None = data.context.get("bus") - agent_id: str | None = data.context.get("agent_id") - if bus is None or agent_id is None: - return data.model_data - pending = await bus.drain(agent_id) - if not pending: - return data.model_data - - new_input = list(data.model_data.input) - for msg in pending: - sender = msg.get("from", "unknown") - content = msg.get("content", "") - if sender == "user": - new_input.append({"role": "user", "content": content}) - else: - new_input.append( - { - "role": "user", - "content": _format_inter_agent_message(bus, msg), - }, - ) - logger.debug( - "inject_messages_filter: appended %d message(s) to input for %s", - len(pending), - agent_id, - ) - return ModelInputData( - input=new_input, - instructions=data.model_data.instructions, - ) - except Exception: - logger.exception( - "inject_messages_filter failed; proceeding with unmodified input", - ) - return data.model_data - - -def _format_inter_agent_message(bus: AgentMessageBus, msg: dict[str, Any]) -> str: - """Render a peer-agent message as a labeled header + body. - - Format: - [Message from {name} ({id}) | type={type} | priority={priority}] - {content} - - Plain text by design — no XML wrapping, no escaping concerns. The - label line tells the receiver who sent this and why so it doesn't - confuse a peer message with its own work; the rest of the body is - delivered as-is. - """ - sender_id = str(msg.get("from", "unknown")) - sender_name = bus.names.get(sender_id, sender_id) - msg_type = msg.get("type", "information") - priority = msg.get("priority", "normal") - content = str(msg.get("content", "")) - return ( - f"[Message from {sender_name} ({sender_id}) " - f"| type={msg_type} | priority={priority}]\n" - f"{content}" - ) diff --git a/strix/orchestration/hooks.py b/strix/orchestration/hooks.py deleted file mode 100644 index 6eea93b..0000000 --- a/strix/orchestration/hooks.py +++ /dev/null @@ -1,309 +0,0 @@ -"""``StrixOrchestrationHooks`` — RunHooks wiring bus + tracer + warnings.""" - -from __future__ import annotations - -import json -import logging -from datetime import UTC, datetime -from typing import Any - -from agents.items import ModelResponse -from agents.lifecycle import RunHooks -from agents.run_context import AgentHookContext, RunContextWrapper -from agents.tool_context import ToolContext - -from strix.telemetry.logging import set_agent_id - - -logger = logging.getLogger(__name__) - - -class StrixOrchestrationHooks(RunHooks[Any]): - """Lifecycle hooks for Strix multi-agent runs. - - Wires three concerns: - - 1. Turn-budget warnings injected into ``input_items`` at 85% and - ``N - 3`` of ``max_turns``. - 2. LLM usage recording into the bus + tracer, plus mirroring the - bus's agent tree into ``tracer.agents`` for the TUI. - 3. Subagent crash detection: if ``on_agent_end`` fires without - ``agent_finish_called`` being set, posts a crash message to the - parent's inbox so the parent learns on its next turn instead of - waiting forever. - """ - - async def on_llm_start( - self, - context: RunContextWrapper[Any], - agent: Any, - system_prompt: str | None, - input_items: list[Any], - ) -> None: - del agent, system_prompt - try: - ctx = context.context - if not isinstance(ctx, dict): - return - bus = ctx.get("bus") - agent_id = ctx.get("agent_id") - if bus is None or agent_id is None: - return - stats = bus.stats_live.get(agent_id) - if stats is None: - return - max_turns = int(ctx.get("max_turns", 300)) - cur = int(stats.get("calls", 0)) - if max_turns < 4: - return - # Once-flags live alongside ``calls`` on ``bus.stats_live`` so the - # warnings fire exactly once per agent lifetime — surviving - # ``run_with_continuation`` cycles, mirroring legacy - # ``state.max_iterations_warning_sent``. - # - # The flags are mutated lock-free below. Safe because the SDK - # serializes ``on_llm_start`` per agent (one in-flight LLM call - # per ``Runner.run`` instance), so this hook is the sole writer - # to ``warned_85`` / ``warned_final`` for this agent_id. - # ``record_usage`` (which acquires the lock) only writes - # ``in`` / ``out`` / ``cached`` / ``calls`` — disjoint keys. - if cur >= int(max_turns * 0.85) and not stats.get("warned_85"): - stats["warned_85"] = True - input_items.append( - { - "role": "user", - "content": ( - "[System warning] You are at 85% of your iteration " - "budget. Begin consolidating findings." - ), - }, - ) - if cur >= max_turns - 3 and not stats.get("warned_final"): - stats["warned_final"] = True - input_items.append( - { - "role": "user", - "content": ( - "[System warning] You have 3 iterations left. Your " - "next tool call MUST be the finish tool." - ), - }, - ) - except Exception: - logger.exception("on_llm_start failed") - - async def on_llm_end( - self, - context: RunContextWrapper[Any], - agent: Any, - response: ModelResponse, - ) -> None: - del agent - try: - ctx = context.context - if not isinstance(ctx, dict): - return - usage = getattr(response, "usage", None) - agent_id = ctx.get("agent_id") - bus = ctx.get("bus") - if bus is not None and agent_id is not None: - await bus.record_usage(agent_id, usage) - tracer = ctx.get("tracer") - cached_total = 0 - input_total = 0 - output_total = 0 - if usage is not None: - details = getattr(usage, "input_tokens_details", None) - if details is not None: - cached_total = int(getattr(details, "cached_tokens", 0) or 0) - input_total = int(getattr(usage, "input_tokens", 0) or 0) - output_total = int(getattr(usage, "output_tokens", 0) or 0) - if tracer is not None and hasattr(tracer, "record_llm_usage"): - tracer.record_llm_usage( - input_tokens=input_total, - output_tokens=output_total, - cached_tokens=cached_total, - ) - logger.debug( - "LLM call done: input=%d output=%d cached=%d", - input_total, - output_total, - cached_total, - ) - except Exception: - logger.exception("on_llm_end failed") - - async def on_agent_start( - self, - context: AgentHookContext[Any], - agent: Any, - ) -> None: - # The TUI reads ``tracer.agents`` to render the agent tree; - # mirror the bus state into the tracer here so the tree - # populates as agents come online. - del agent - try: - ctx = context.context - if not isinstance(ctx, dict): - return - me = ctx.get("agent_id") - if isinstance(me, str): - # Tag every log record emitted under this agent's task - # with the agent_id, automatically. Cleared on agent end. - set_agent_id(me) - tracer = ctx.get("tracer") - bus = ctx.get("bus") - if tracer is None or bus is None or me is None: - return - name = bus.names.get(me, me) - parent = bus.parent_of.get(me) - logger.info("Agent %s (%s) started, parent=%s", me, name, parent or "-") - now = datetime.now(UTC).isoformat() - tracer.agents.setdefault( - me, - { - "id": me, - "name": name, - "parent_id": parent, - "status": bus.statuses.get(me, "running"), - "created_at": now, - "updated_at": now, - }, - ) - except Exception: - logger.exception("on_agent_start failed") - - async def on_agent_end( - self, - context: AgentHookContext[Any], - agent: Any, - output: Any, - ) -> None: - try: - ctx = context.context - if not isinstance(ctx, dict): - return - bus = ctx.get("bus") - me = ctx.get("agent_id") - if bus is None or me is None: - return - crashed = (output is None) or not ctx.get("agent_finish_called", False) - - # Interactive agents (root and children) stay alive across - # ``Runner.run`` cycles — ``run_with_continuation`` re-invokes - # ``Runner.run`` whenever the agent receives a follow-up - # message, so we just park (status=waiting) instead of - # finalizing. Crashed runs always finalize so the parent - # learns to stop waiting. - stays_alive = bool(ctx.get("interactive", False)) and not crashed - - final_status = "waiting" if stays_alive else ("crashed" if crashed else "completed") - - tracer = ctx.get("tracer") - if tracer is not None and me in tracer.agents: - tracer.agents[me]["status"] = final_status - tracer.agents[me]["updated_at"] = datetime.now(UTC).isoformat() - parent = bus.parent_of.get(me) - if crashed and parent is not None: - await bus.send( - parent, - { - "from": me, - "content": ( - f"[Agent crash] {bus.names.get(me, me)} ({me}) " - f"terminated without calling agent_finish. " - f"Stop waiting on this child." - ), - "type": "crash", - }, - ) - - calls = ( - int(bus.stats_live.get(me, {}).get("calls", 0)) if hasattr(bus, "stats_live") else 0 - ) - logger.info( - "Agent %s (%s) ended status=%s calls=%d", - me, - bus.names.get(me, me), - final_status, - calls, - ) - - if stays_alive: - await bus.park(me) - # Reset the finish flag so the next cycle can detect its own - # finish-tool call. The lifetime turn counter and warning - # flags live on ``bus.stats_live`` and persist across cycles. - ctx["agent_finish_called"] = False - else: - await bus.finalize(me, final_status) - # Clear the agent_id from the log context so any post-finalize - # work (cleanup in scan.py finally) doesn't keep the stale tag. - set_agent_id(None) - except Exception: - logger.exception("on_agent_end failed") - - async def on_tool_start( - self, - context: RunContextWrapper[Any], - agent: Any, - tool: Any, - ) -> None: - del agent - try: - ctx = context.context - if not isinstance(ctx, dict): - return - tracer = ctx.get("tracer") - if tracer is None: - return - # ``context`` is a ``ToolContext`` for function-tool calls (per the - # SDK ``RunHooks.on_tool_start`` docstring) — that's where the - # per-call args live. ``tool_input`` is the parsed dict when the - # SDK has it; otherwise parse ``tool_arguments`` (raw JSON). - args: dict[str, Any] = {} - if isinstance(context, ToolContext): - tool_input = context.tool_input - if isinstance(tool_input, dict): - args = tool_input - else: - raw = context.tool_arguments - if raw: - try: - parsed = json.loads(raw) - except (ValueError, TypeError): - parsed = None - if isinstance(parsed, dict): - args = parsed - tracer.log_tool_start(ctx.get("agent_id", "?"), tool.name, args) - args_repr = json.dumps(args, ensure_ascii=False, default=str) if args else "" - logger.debug( - "Tool %s start (args_len=%d): %s", - tool.name, - len(args_repr), - args_repr[:500], - ) - except Exception: - logger.exception("on_tool_start failed") - - async def on_tool_end( - self, - context: RunContextWrapper[Any], - agent: Any, - tool: Any, - result: str, - ) -> None: - del agent - try: - ctx = context.context - if not isinstance(ctx, dict): - return - if tool.name in ("agent_finish", "finish_scan"): - ctx["agent_finish_called"] = True - tracer = ctx.get("tracer") - if tracer is not None: - tracer.log_tool_end(ctx.get("agent_id", "?"), tool.name, result) - result_str = result if isinstance(result, str) else str(result) - logger.debug("Tool %s done (result_len=%d)", tool.name, len(result_str)) - except Exception: - logger.exception("on_tool_end failed") diff --git a/strix/orchestration/run_loop.py b/strix/orchestration/run_loop.py deleted file mode 100644 index 4f817a8..0000000 --- a/strix/orchestration/run_loop.py +++ /dev/null @@ -1,183 +0,0 @@ -"""``run_with_continuation`` — interactive-mode demo-loop wrapper around ``Runner.run_streamed``. - -Pre-migration ``BaseAgent.agent_loop`` ran forever in interactive mode, -re-entering a "waiting state" after each finish-tool call so the agent -could pick up follow-up messages from its parent (or from the user, in -the root's case). Post-migration this helper restores the legacy -semantics using the SDK's streaming Runner + ``RunResultStreaming.cancel`` -so the user can interrupt mid-turn without truncating session state. - -Behaviors restored from legacy: - -- **Mid-stream interrupt** via ``streamed.cancel(mode="after_turn")``: - TUI signals through ``bus.request_interrupt``; the SDK saves the - current turn cleanly before honoring the cancel. -- **LLM failure resume** (legacy ``state.llm_failed``): hard model - failures after retries exhausted park the agent in ``llm_failed`` - status; only direct user input can resume. -- **Waiting timeout** auto-resume (legacy ``waiting_timeout``): - interactive subagents auto-resume after 300s with a "Waiting timeout - reached" message. Interactive root waits forever. -- **Graceful stop** (legacy ``stop_agent``): ``bus.stopping`` set - causes the outer loop to return instead of awaiting more messages. -""" - -from __future__ import annotations - -import asyncio -import logging -from typing import TYPE_CHECKING, Any - -from agents import Runner -from agents.exceptions import AgentsException, MaxTurnsExceeded, UserError -from openai import APIError - - -if TYPE_CHECKING: - from agents.lifecycle import RunHooks - from agents.memory import Session - from agents.result import RunResultBase - from agents.run_config import RunConfig - - from strix.orchestration.bus import AgentMessageBus - - -logger = logging.getLogger(__name__) - - -#: Auto-resume timeout for interactive *subagents* (legacy parity). -#: Interactive root agents wait forever; non-interactive runs don't loop. -_WAITING_TIMEOUT_SUBAGENT = 300.0 - -_TIMEOUT_RESUME_MESSAGE = "Waiting timeout reached. Resuming execution." - - -async def run_with_continuation( - *, - agent: Any, - initial_input: Any, - run_config: RunConfig, - context: dict[str, Any], - hooks: RunHooks[Any], - max_turns: int, - bus: AgentMessageBus, - agent_id: str, - interactive: bool, - session: Session | None = None, -) -> RunResultBase: - """Run an agent once (non-interactive) or in a continuation loop (interactive).""" - kwargs: dict[str, Any] = { - "input": initial_input, - "run_config": run_config, - "context": context, - "hooks": hooks, - "max_turns": max_turns, - } - if session is not None: - kwargs["session"] = session - - # Interactive subagents auto-resume after a timeout to mirror legacy - # ``waiting_timeout``. Roots wait forever (legacy ``waiting_timeout=0``). - waiting_timeout: float | None = None - if interactive: - async with bus._lock: - parent_id = bus.parent_of.get(agent_id) - if parent_id is not None: - waiting_timeout = _WAITING_TIMEOUT_SUBAGENT - - result = await _run_streamed(agent, bus, agent_id, **kwargs) - - if not interactive: - return result - - logger.debug( - "run_with_continuation: entering interactive outer loop for %s (timeout=%s)", - agent_id, - waiting_timeout, - ) - while True: - if agent_id in bus.stopping: - logger.info("run_with_continuation: %s in stopping set, returning", agent_id) - return result - - try: - if waiting_timeout is None: - await bus.wait_for_message(agent_id) - else: - await asyncio.wait_for( - bus.wait_for_message(agent_id), - timeout=waiting_timeout, - ) - except asyncio.CancelledError: - logger.info("run_with_continuation: %s cancelled while waiting", agent_id) - return result - except TimeoutError: - logger.info( - "run_with_continuation: %s waiting timeout, auto-resuming", - agent_id, - ) - kwargs["input"] = _TIMEOUT_RESUME_MESSAGE - result = await _run_streamed(agent, bus, agent_id, **kwargs) - continue - - pending = await bus.drain(agent_id) - if not pending: - continue - next_input = "\n\n".join( - str(msg.get("content", "")).strip() for msg in pending if msg.get("content") - ) - if not next_input: - continue - - logger.debug( - "run_with_continuation: %s resuming with %d message(s) (input_len=%d)", - agent_id, - len(pending), - len(next_input), - ) - kwargs["input"] = next_input - result = await _run_streamed(agent, bus, agent_id, **kwargs) - - -async def _run_streamed( - agent: Any, - bus: AgentMessageBus, - agent_id: str, - **kwargs: Any, -) -> RunResultBase: - """Drive one ``Runner.run_streamed`` cycle to completion. - - Catches hard model failures (after SDK retries are exhausted) and - parks the agent in ``llm_failed`` until a user message arrives, - matching legacy ``state.llm_failed`` semantics. Programmer errors - (``UserError``), max-turn breaches, and explicit cancellation - propagate to the caller. - """ - interactive = bool(kwargs.get("context", {}).get("interactive", False)) - while True: - streamed = Runner.run_streamed(agent, **kwargs) - try: - async with bus.attach_stream(agent_id, streamed): - async for _event in streamed.stream_events(): - pass - except (UserError, MaxTurnsExceeded, asyncio.CancelledError): - raise - except (AgentsException, APIError): - if not interactive: - raise - logger.exception( - "LLM hard failure for agent %s; awaiting user resume", - agent_id, - ) - await bus.mark_llm_failed(agent_id) - await bus.wait_for_user_message(agent_id) - pending = await bus.drain(agent_id) - next_input = "\n\n".join( - str(msg.get("content", "")).strip() for msg in pending if msg.get("content") - ) - if not next_input: - continue - kwargs["input"] = next_input - continue - else: - return streamed diff --git a/strix/orchestration/scan.py b/strix/orchestration/scan.py index 8006ea7..f2c261a 100644 --- a/strix/orchestration/scan.py +++ b/strix/orchestration/scan.py @@ -1,11 +1,11 @@ """Top-level scan entry point with auto-resume. -1. Build (or take from caller) the per-scan ``AgentMessageBus``. +1. Build (or take from caller) the per-scan ``AgentCoordinator``. 2. Wire a snapshot path so every lifecycle event auto-persists ``bus.json``. 3. Acquire an advisory file lock so a second ``strix`` process can't run on the same ``scan_id`` concurrently. 4. **Resume detection**: if ``{run_dir}/bus.json`` already exists, restore - the bus, hydrate the tracer, reuse the persisted ``root_id`` instead + the coordinator, hydrate the tracer, reuse the persisted ``root_id`` instead of generating a fresh one, and respawn every non-terminal subagent from its per-child ``SQLiteSession`` before starting the root. 5. Bring up (or reuse) a sandbox session for ``scan_id``. @@ -41,10 +41,11 @@ from strix.agents.factory import build_strix_agent, make_child_factory from strix.config import load_settings from strix.llm.multi_provider_setup import build_multi_provider from strix.llm.retry import DEFAULT_RETRY -from strix.orchestration.bus import AgentMessageBus -from strix.orchestration.filter import inject_messages_filter -from strix.orchestration.hooks import StrixOrchestrationHooks -from strix.orchestration.run_loop import run_with_continuation +from strix.orchestration.coordinator import ( + AgentCoordinator, + open_agent_session, + run_with_continuation, +) from strix.runtime import session_manager from strix.telemetry.logging import set_agent_id, set_scan_id, setup_scan_logging @@ -176,12 +177,12 @@ async def run_strix_scan( image: str, local_sources: list[dict[str, str]] | None = None, tracer: Any | None = None, - bus: AgentMessageBus | None = None, + coordinator: AgentCoordinator | None = None, interactive: bool = False, max_turns: int = _MAX_TURNS, model: str | None = None, cleanup_on_exit: bool = True, -) -> RunResultBase: +) -> RunResultBase | None: """Run one Strix scan end-to-end against a freshly-prepared sandbox. Args: @@ -253,12 +254,11 @@ async def run_strix_scan( ) logger.info("LLM model resolved: %s", resolved_model) - # Caller may pre-create the bus so it can hold a handle (e.g., the - # TUI uses it to route stop / chat-input commands). Otherwise we - # own the bus internally for the scan's lifetime. - if bus is None: - bus = AgentMessageBus() - bus.set_snapshot_path(bus_path) + # Caller may pre-create the coordinator so it can route stop/chat + # commands while the scan loop runs in another thread. + if coordinator is None: + coordinator = AgentCoordinator() + coordinator.set_snapshot_path(bus_path) if tracer is not None and hasattr(tracer, "hydrate_from_run_dir"): tracer.hydrate_from_run_dir() @@ -281,8 +281,8 @@ async def run_strix_scan( raise RuntimeError( f"Cannot resume scan {scan_id}: bus.json is unreadable: {exc}", ) from exc - await bus.restore(snap) - for aid, parent in bus.parent_of.items(): + await coordinator.restore(snap) + for aid, parent in coordinator.parent_of.items(): if parent is None: root_id = aid break @@ -292,11 +292,9 @@ async def run_strix_scan( f"Cannot resume scan {scan_id}: bus.json has no root agent (parent=None)", ) logger.info( - "Resume: restored bus with %d agent(s); root=%s; %d non-terminal to respawn", - len(bus.statuses), + "Resume: restored coordinator with %d agent(s); root=%s", + len(coordinator.statuses), root_id, - sum(1 for s in bus.statuses.values() if s in {"running", "waiting", "llm_failed"}) - - 1, # subtract root ) else: root_id = uuid.uuid4().hex[:8] @@ -334,7 +332,7 @@ async def run_strix_scan( ) if not is_resume: - await bus.register( + await coordinator.register( root_id, "strix", parent_id=None, @@ -353,7 +351,7 @@ async def run_strix_scan( ) context: dict[str, Any] = { - "bus": bus, + "coordinator": coordinator, "sandbox_session": bundle["session"], "sandbox_client": bundle["client"], "caido_client": bundle["caido_client"], @@ -390,14 +388,13 @@ async def run_strix_scan( model_provider=build_multi_provider(), model_settings=model_settings, sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]), - call_model_input_filter=inject_messages_filter, tracing_disabled=False, trace_include_sensitive_data=False, ) if is_resume: await _respawn_subagents( - bus=bus, + coordinator=coordinator, sessions_dir=sessions_dir, factory=agent_factory, parent_ctx=context, @@ -413,17 +410,18 @@ async def run_strix_scan( session_db = run_dir / "session.db" root_session = SQLiteSession(session_id=scan_id, db_path=session_db) sessions_to_close.append(root_session) + await coordinator.attach_runtime(root_id, session=root_session) initial_input: Any = [] if is_resume else _build_root_task(scan_config) # Resume + new ``--instruction``: SDK replay drives root from # session.db with ``initial_input=[]``, so a brand-new instruction # passed on the resume CLI would otherwise be silently ignored. - # Inject it as a fresh user message in root's inbox; the - # ``inject_messages_filter`` will surface it on the very next turn. + # Inject it as a fresh user message in root's SDK session; the + # next run cycle will replay it with the rest of the session. resume_instruction = str(scan_config.get("resume_instruction") or "").strip() if is_resume and resume_instruction: - await bus.send( + await coordinator.send( root_id, { "from": "user", @@ -432,41 +430,38 @@ async def run_strix_scan( "content": resume_instruction, }, ) - # ``bus.send`` is one of the high-frequency mutations that - # deliberately skips ``_maybe_snapshot``. The resume-instruction - # is the one specific message we can't lose: a SIGKILL between - # the send and root's first turn would silently drop the - # user's new ``--instruction``. Force a snapshot here. - await bus._maybe_snapshot() logger.info( "Resume: injected new instruction into root inbox (len=%d)", len(resume_instruction), ) + async with coordinator._lock: + root_status = coordinator.statuses.get(root_id) + return await run_with_continuation( agent=root_agent, initial_input=initial_input, run_config=run_config, context=context, - hooks=StrixOrchestrationHooks(), max_turns=max_turns, - bus=bus, + coordinator=coordinator, agent_id=root_id, interactive=interactive, session=root_session, + start_parked=bool(interactive and is_resume and root_status != "running"), ) except BaseException as exc: logger.exception("Strix scan %s failed", scan_id) # Cancel any descendant tasks the root spawned before unwinding. # cancel_descendants is idempotent and handles the empty-tree case. if root_id is not None: - await bus.cancel_descendants(root_id) + await coordinator.cancel_descendants(root_id) # The SDK's on_agent_end hook only fires after a successful # ``Runner.run_streamed`` reaches the agent's first turn. A # failure earlier (e.g., model-provider routing, sandbox # bring-up) leaves the root stuck at status="running" — the # TUI keeps animating "Initializing" forever. Finalize it - # here so the bus + tracer reflect reality, and stash the + # here so the coordinator + tracer reflect reality, and stash the # error message for the status-line display. error_message = f"{type(exc).__name__}: {exc}" if tracer is not None and root_id in getattr(tracer, "agents", {}): @@ -474,7 +469,7 @@ async def run_strix_scan( tracer.agents[root_id]["error_message"] = error_message tracer.agents[root_id]["updated_at"] = datetime.now(UTC).isoformat() with contextlib.suppress(Exception): - await bus.finalize(root_id, "failed") + await coordinator.set_status(root_id, "failed") set_agent_id(None) raise finally: @@ -482,7 +477,7 @@ async def run_strix_scan( with contextlib.suppress(Exception): s.close() with contextlib.suppress(Exception): - await bus._maybe_snapshot() + await coordinator._maybe_snapshot() if cleanup_on_exit: logger.info("Tearing down sandbox session for scan %s", scan_id) await session_manager.cleanup(scan_id) @@ -493,7 +488,7 @@ async def run_strix_scan( async def _respawn_subagents( *, - bus: AgentMessageBus, + coordinator: AgentCoordinator, sessions_dir: Path, factory: Any, parent_ctx: dict[str, Any], @@ -502,57 +497,77 @@ async def _respawn_subagents( root_id: str, sessions_to_close: list[SQLiteSession], ) -> None: - """Re-spawn every non-terminal subagent from a restored bus snapshot. + """Re-spawn subagent runners from a restored coordinator snapshot. Each child gets its own :class:`SQLiteSession` reopened at ``sessions_dir/.db`` so the SDK replays its prior conversation. Per-child failure (missing/corrupt session DB, factory raising) finalizes that child as ``crashed`` and continues. - Terminal-status agents (``completed`` / ``crashed`` / ``stopped``) - are left alone — their stats stay in ``stats_completed`` for the - TUI, but no task respawns. + Interactive mode respawns every registered child as a parked runner + unless it was actively running before the crash. That keeps + completed/stopped/crashed/failed agents addressable: a later message + wakes the SDK session instead of being dropped into a dead inbox. """ - async with bus._lock: - # Snapshot the iteration view first so we can mutate via finalize + interactive = bool(parent_ctx.get("interactive", False)) + async with coordinator._lock: + # Snapshot the iteration view first so we can mutate via coordinator # below without "dict changed during iteration" trouble. agents_snapshot = [ - (aid, status, dict(bus.metadata.get(aid, {}))) for aid, status in bus.statuses.items() + (aid, status, dict(coordinator.metadata.get(aid, {}))) + for aid, status in coordinator.statuses.items() ] candidates: list[tuple[str, str, str | None, dict[str, Any]]] = [] - to_finalize_stopped: list[str] = [] for aid, status, md in agents_snapshot: - if status not in {"running", "waiting", "llm_failed"}: + if not interactive and status not in {"running", "waiting", "llm_failed"}: continue - if bus.parent_of.get(aid) is None or aid == root_id: + if coordinator.parent_of.get(aid) is None or aid == root_id: continue - if aid in bus.stopping: - # User clicked "stop" before the crash; don't respawn, - # and reconcile the bus so its status truthfully reflects - # "stopped" instead of staying "running" forever. - to_finalize_stopped.append(aid) - continue - candidates.append((aid, bus.names.get(aid, aid), bus.parent_of.get(aid), md)) - - # Finalize outside the lock — ``bus.finalize`` acquires it itself. - for aid in to_finalize_stopped: - logger.info("respawn-skip %s: previously-cancelled, finalizing as stopped", aid) - await bus.finalize(aid, "stopped") + md["_restored_status"] = status + candidates.append( + ( + aid, + coordinator.names.get(aid, aid), + coordinator.parent_of.get(aid), + md, + ) + ) for child_id, name, parent_id, md in candidates: try: session_path = sessions_dir / f"{child_id}.db" if not session_path.exists(): + if interactive: + logger.warning( + "respawn %s (%s): session db missing at %s; starting parked with empty session", + child_id, + name, + session_path, + ) + else: + logger.warning( + "respawn %s (%s): session db missing at %s — marking crashed", + child_id, + name, + session_path, + ) + await coordinator.set_status(child_id, "crashed") + continue + session_path.parent.mkdir(parents=True, exist_ok=True) + + restored_status = str(md.get("_restored_status") or "running") + start_parked = interactive and restored_status != "running" + + if start_parked: logger.warning( - "respawn %s (%s): session db missing at %s — finalizing as crashed", + "respawn %s (%s): starting parked from status=%s", child_id, name, - session_path, + restored_status, ) - await bus.finalize(child_id, "crashed") - continue - child_session = SQLiteSession(session_id=child_id, db_path=session_path) + child_session = open_agent_session(child_id, session_path) sessions_to_close.append(child_session) + await coordinator.attach_runtime(child_id, session=child_session) child_skills = list(md.get("skills") or []) child_agent = factory(name=name, skills=child_skills) @@ -580,7 +595,6 @@ async def _respawn_subagents( client=parent_ctx["sandbox_client"], session=parent_ctx["sandbox_session"], ), - call_model_input_filter=inject_messages_filter, tracing_disabled=False, trace_include_sensitive_data=False, ) @@ -591,17 +605,16 @@ async def _respawn_subagents( initial_input=[], run_config=child_run_config, context=child_ctx, - hooks=StrixOrchestrationHooks(), max_turns=int(parent_ctx.get("max_turns", 300)), - bus=bus, + coordinator=coordinator, agent_id=child_id, interactive=bool(parent_ctx.get("interactive", False)), session=child_session, + start_parked=start_parked, ), name=f"agent-{name}-{child_id}", ) - async with bus._lock: - bus.tasks[child_id] = task_handle + await coordinator.attach_runtime(child_id, task=task_handle) logger.info( "respawned %s (%s) parent=%s task_len=%d", child_id, @@ -612,7 +625,7 @@ async def _respawn_subagents( except Exception: logger.exception("respawn %s failed; marking crashed", child_id) with contextlib.suppress(Exception): - await bus.finalize(child_id, "crashed") + await coordinator.set_status(child_id, "crashed") def _acquire_run_lock(run_dir: Path) -> Any: diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index 1194725..e17a1ae 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -1,16 +1,15 @@ -"""Multi-agent graph tools — read/write the :class:`AgentMessageBus`. +"""Multi-agent graph tools backed by the SDK-native :class:`AgentCoordinator`. - ``view_agent_graph``: render the parent/child tree. - ``agent_status``: per-agent status + pending message count. -- ``send_message_to_agent``: queue a message in another agent's inbox. +- ``send_message_to_agent``: append a message to another agent's SDK session. - ``wait_for_message``: pause this agent until a message arrives or ``timeout_seconds`` elapses. - ``create_agent``: spawn a child via ``asyncio.create_task(Runner.run(...))``; the task handle is stored so a root-level cancel cascades to descendants. -- ``agent_finish``: subagents only — flips ``agent_finish_called`` so - the ``on_agent_end`` hook records "completed" rather than "crashed", - and posts a structured completion report to the parent's inbox. +- ``agent_finish``: subagents only — posts a structured completion + report to the parent's SDK session and returns a final-output marker. """ from __future__ import annotations @@ -25,15 +24,16 @@ from typing import TYPE_CHECKING, Any, Literal from agents import RunConfig, RunContextWrapper, function_tool from agents.items import TResponseInputItem -from agents.memory import SQLiteSession from agents.model_settings import ModelSettings from agents.sandbox import SandboxRunConfig from strix.llm.multi_provider_setup import build_multi_provider from strix.llm.retry import DEFAULT_RETRY -from strix.orchestration.filter import inject_messages_filter -from strix.orchestration.hooks import StrixOrchestrationHooks -from strix.orchestration.run_loop import run_with_continuation +from strix.orchestration.coordinator import ( + coordinator_from_context, + open_agent_session, + run_with_continuation, +) if TYPE_CHECKING: @@ -61,9 +61,9 @@ def _render_completion_report( ) -> str: """Render a child's completion report as plain structured text. - Goes into the parent's bus inbox; the inject filter prepends a - ``[Message from ...]`` header on top, so this body just carries the - contents. No XML — no escaping concerns, no parser ambiguity. + Goes into the parent's SDK session with coordinator-added sender + metadata, so this body just carries the contents. No XML — no + escaping concerns, no parser ambiguity. """ status = "SUCCESS" if success else "FAILED" completion_time = datetime.now(UTC).isoformat() @@ -101,19 +101,16 @@ async def view_agent_graph(ctx: RunContextWrapper) -> str: is marked ``← you``. """ inner = _ctx(ctx) - bus = inner.get("bus") + coordinator = coordinator_from_context(inner) me = inner.get("agent_id") - if bus is None: + if coordinator is None: return json.dumps( - {"success": False, "error": "Bus not initialized in context."}, + {"success": False, "error": "Agent coordinator not initialized in context."}, ensure_ascii=False, default=str, ) - async with bus._lock: - parent_of = dict(bus.parent_of) - statuses = dict(bus.statuses) - names = dict(bus.names) + parent_of, statuses, names = await coordinator.graph_snapshot() lines: list[str] = [] @@ -162,36 +159,26 @@ async def agent_status(ctx: RunContextWrapper, agent_id: str) -> str: ``create_agent``. """ inner = _ctx(ctx) - bus = inner.get("bus") - if bus is None: + coordinator = coordinator_from_context(inner) + if coordinator is None: return json.dumps( - {"success": False, "error": "Bus not initialized in context."}, + {"success": False, "error": "Agent coordinator not initialized in context."}, ensure_ascii=False, default=str, ) - async with bus._lock: - if agent_id not in bus.statuses: - return json.dumps( - { - "success": False, - "error": f"Unknown agent_id: {agent_id}", - }, - ensure_ascii=False, - default=str, - ) + info = await coordinator.agent_info(agent_id) + if info is None: return json.dumps( { - "success": True, - "agent_id": agent_id, - "name": bus.names.get(agent_id), - "status": bus.statuses.get(agent_id), - "parent_id": bus.parent_of.get(agent_id), - "pending_messages": len(bus.inboxes.get(agent_id, [])), + "success": False, + "error": f"Unknown agent_id: {agent_id}", }, ensure_ascii=False, default=str, ) + info["success"] = True + return json.dumps(info, ensure_ascii=False, default=str) @function_tool(timeout=30) @@ -215,7 +202,8 @@ async def send_message_to_agent( **Don't** use for routine "hello/status" pings, for context the target already has (children inherit parent history), or when parent/child completion via ``agent_finish`` already covers the - flow. Messages to a finalized agent are dropped. + flow. Messages to any registered agent are queued, regardless of + status, so a follow-up can wake a completed/stopped/failed agent. Args: target_agent_id: Recipient's 8-char id. @@ -227,39 +215,17 @@ async def send_message_to_agent( priority: ``low`` / ``normal`` / ``high`` / ``urgent``. """ inner = _ctx(ctx) - bus = inner.get("bus") + coordinator = coordinator_from_context(inner) me = inner.get("agent_id") - if bus is None or me is None: + if coordinator is None or me is None: return json.dumps( - {"success": False, "error": "Bus or agent_id missing in context."}, - ensure_ascii=False, - default=str, - ) - - async with bus._lock: - if target_agent_id not in bus.statuses: - return json.dumps( - { - "success": False, - "error": f"Target agent '{target_agent_id}' not found.", - }, - ensure_ascii=False, - default=str, - ) - target_status = bus.statuses.get(target_agent_id) - - if target_status in ("completed", "crashed", "stopped"): - return json.dumps( - { - "success": False, - "error": f"Target agent '{target_agent_id}' is {target_status}; message dropped.", - }, + {"success": False, "error": "Agent coordinator or agent_id missing in context."}, ensure_ascii=False, default=str, ) msg_id = f"msg_{uuid.uuid4().hex[:8]}" - await bus.send( + delivered = await coordinator.send( target_agent_id, { "id": msg_id, @@ -269,6 +235,15 @@ async def send_message_to_agent( "priority": priority, }, ) + if not delivered: + return json.dumps( + { + "success": False, + "error": f"Target agent '{target_agent_id}' not found or message delivery failed.", + }, + ensure_ascii=False, + default=str, + ) return json.dumps( { "success": True, @@ -281,9 +256,16 @@ async def send_message_to_agent( ) -# Tighter would burn CPU; slacker would feel laggy when a sibling -# delivers a message right after the wait starts. -_WAIT_POLL_SECONDS = 1.0 +def _session_items_payload(items: list[Any]) -> list[dict[str, Any]]: + payload: list[dict[str, Any]] = [] + for item in items: + if isinstance(item, dict): + role = item.get("role") + content = item.get("content") + payload.append({"role": role, "content": content}) + else: + payload.append({"content": str(item)}) + return payload @function_tool(timeout=601) @@ -319,51 +301,105 @@ async def wait_for_message( again. """ inner = _ctx(ctx) - bus = inner.get("bus") + coordinator = coordinator_from_context(inner) me = inner.get("agent_id") - if bus is None or me is None: + interactive = bool(inner.get("interactive", False)) + if coordinator is None or me is None: return json.dumps( - {"success": False, "error": "Bus or agent_id missing in context."}, + {"success": False, "error": "Agent coordinator or agent_id missing in context."}, ensure_ascii=False, default=str, ) - async with bus._lock: - bus.statuses[me] = "waiting" + async with coordinator._lock: + stopped = coordinator.statuses.get(me) == "stopped" + if stopped: + return json.dumps( + { + "success": True, + "status": "stopped", + "reason": reason, + "note": "Wait ended because this agent is stopped.", + }, + ensure_ascii=False, + default=str, + ) - deadline = asyncio.get_event_loop().time() + timeout_seconds + pending = await coordinator.pending_count(me) + if pending > 0: + items = await coordinator.recent_session_items(me, pending) + await coordinator.consume_wake(me) + await coordinator.mark_running(me) + return json.dumps( + { + "success": True, + "status": "message_arrived", + "pending_messages": pending, + "messages": _session_items_payload(items), + "reason": reason, + }, + ensure_ascii=False, + default=str, + ) + + if interactive: + await coordinator.park_waiting(me) + inner["agent_waiting_called"] = True + return json.dumps( + { + "success": True, + "status": "waiting", + "agent_waiting": True, + "reason": reason, + "note": "Agent parked; execution will resume when a message arrives.", + }, + ensure_ascii=False, + default=str, + ) + + await coordinator.park_waiting(me) try: - while asyncio.get_event_loop().time() < deadline: - async with bus._lock: - pending = len(bus.inboxes.get(me, [])) - if pending > 0: - async with bus._lock: - bus.statuses[me] = "running" - return json.dumps( - { - "success": True, - "status": "message_arrived", - "pending_messages": pending, - "reason": reason, - }, - ensure_ascii=False, - default=str, - ) - await asyncio.sleep(_WAIT_POLL_SECONDS) - finally: - async with bus._lock: - # Don't clobber a status another writer set (e.g., on_agent_end - # finalized us as ``stopped`` mid-wait). - if bus.statuses.get(me) == "waiting": - bus.statuses[me] = "running" + await asyncio.wait_for(coordinator.wait_for_message(me), timeout_seconds) + except TimeoutError: + await coordinator.mark_running(me) + return json.dumps( + { + "success": True, + "status": "timeout", + "timeout_seconds": timeout_seconds, + "reason": reason, + "note": "No messages within timeout — continue work or call agent_finish.", + }, + ensure_ascii=False, + default=str, + ) + + async with coordinator._lock: + stopped = coordinator.statuses.get(me) == "stopped" + if stopped: + return json.dumps( + { + "success": True, + "status": "stopped", + "reason": reason, + "note": "Wait ended because this agent is stopped.", + }, + ensure_ascii=False, + default=str, + ) + + pending = await coordinator.pending_count(me) + items = await coordinator.recent_session_items(me, pending) + await coordinator.consume_wake(me) + await coordinator.mark_running(me) return json.dumps( { "success": True, - "status": "timeout", - "timeout_seconds": timeout_seconds, + "status": "message_arrived", + "pending_messages": pending, + "messages": _session_items_payload(items), "reason": reason, - "note": "No messages within timeout — continue work or call agent_finish.", }, ensure_ascii=False, default=str, @@ -418,13 +454,13 @@ async def create_agent( skills: Comma-separated skill names. Max 5; prefer 1-3. """ inner = _ctx(ctx) - bus = inner.get("bus") + coordinator = coordinator_from_context(inner) parent_id = inner.get("agent_id") factory: Callable[..., SDKAgent] | None = inner.get("agent_factory") - if bus is None or parent_id is None: + if coordinator is None or parent_id is None: return json.dumps( - {"success": False, "error": "Bus or agent_id missing in context."}, + {"success": False, "error": "Agent coordinator or agent_id missing in context."}, ensure_ascii=False, default=str, ) @@ -456,7 +492,14 @@ async def create_agent( default=str, ) - await bus.register( + tracer = inner.get("tracer") + if tracer is not None and hasattr(tracer, "get_run_dir"): + child_session_path = tracer.get_run_dir() / "sessions" / f"{child_id}.db" + else: + run_id = inner.get("run_id") or "default" + child_session_path = Path.cwd() / "strix_runs" / str(run_id) / "sessions" / f"{child_id}.db" + + await coordinator.register( child_id, name, parent_id, @@ -465,6 +508,7 @@ async def create_agent( is_whitebox=bool(inner.get("is_whitebox", False)), scan_mode=str(inner.get("scan_mode", "deep")), diff_scope=inner.get("diff_scope"), + session_path=child_session_path, ) # ``ctx.turn_input`` carries the parent's full conversation up to and @@ -501,7 +545,7 @@ async def create_agent( initial_input.append({"role": "user", "content": task}) child_ctx: dict[str, Any] = { - "bus": bus, + "coordinator": coordinator, "sandbox_session": inner.get("sandbox_session"), "sandbox_client": inner.get("sandbox_client"), "caido_client": inner.get("caido_client"), @@ -531,17 +575,11 @@ async def create_agent( # tracer's run_dir (root-side construction in # :func:`run_strix_scan`); fall back to ``./strix_runs/{run_id}/`` if # the tracer is absent (unit-test path). - tracer = inner.get("tracer") - if tracer is not None and hasattr(tracer, "get_run_dir"): - child_session_path = tracer.get_run_dir() / "sessions" / f"{child_id}.db" - else: - run_id = inner.get("run_id") or "default" - child_session_path = Path.cwd() / "strix_runs" / str(run_id) / "sessions" / f"{child_id}.db" - child_session_path.parent.mkdir(parents=True, exist_ok=True) - child_session = SQLiteSession(session_id=child_id, db_path=child_session_path) + child_session = open_agent_session(child_id, child_session_path) sessions_list = child_ctx.get("_sessions_to_close") if isinstance(sessions_list, list): sessions_list.append(child_session) + await coordinator.attach_runtime(child_id, session=child_session) child_model_settings = ModelSettings( parallel_tool_calls=False, @@ -561,7 +599,6 @@ async def create_agent( if sandbox_session is not None else None ), - call_model_input_filter=inject_messages_filter, tracing_disabled=False, trace_include_sensitive_data=False, ) @@ -572,17 +609,15 @@ async def create_agent( initial_input=initial_input, run_config=child_run_config, context=child_ctx, - hooks=StrixOrchestrationHooks(), max_turns=int(inner.get("max_turns", 300)), - bus=bus, + coordinator=coordinator, agent_id=child_id, interactive=bool(inner.get("interactive", False)), session=child_session, ), name=f"agent-{name}-{child_id}", ) - async with bus._lock: - bus.tasks[child_id] = task_handle + await coordinator.attach_runtime(child_id, task=task_handle) logger.info( "create_agent: spawned %s (%s) parent=%s skills=%d task_len=%d", @@ -650,11 +685,11 @@ async def agent_finish( cover Y"). """ inner = _ctx(ctx) - bus = inner.get("bus") + coordinator = coordinator_from_context(inner) me = inner.get("agent_id") - if bus is None or me is None: + if coordinator is None or me is None: return json.dumps( - {"success": False, "error": "Bus or agent_id missing in context."}, + {"success": False, "error": "Agent coordinator or agent_id missing in context."}, ensure_ascii=False, default=str, ) @@ -674,13 +709,14 @@ async def agent_finish( default=str, ) - # ``agent_finish_called`` is set by ``StrixOrchestrationHooks.on_tool_end``; - # no need to set it here. + # The coordinator settles lifecycle from this JSON final output after + # the SDK run finishes; the tool only sends the parent report. + inner["agent_finish_called"] = True parent_notified = False if report_to_parent: - async with bus._lock: - agent_name = bus.names.get(me, me) + async with coordinator._lock: + agent_name = coordinator.names.get(me, me) report = _render_completion_report( agent_name=agent_name, agent_id=me, @@ -690,7 +726,7 @@ async def agent_finish( findings=list(findings or []), recommendations=list(final_recommendations or []), ) - await bus.send( + await coordinator.send( parent_id, { "id": f"report_{uuid.uuid4().hex[:8]}", @@ -737,8 +773,8 @@ async def stop_agent( Uses the SDK's ``RunResultStreaming.cancel(mode="after_turn")`` so the target's current turn finishes — including saving items to its session — before the run loop honors the cancel. The agent's - interactive outer loop sees ``stopping`` and exits without awaiting - more messages, so ``on_agent_end`` finalizes with status="stopped". + interactive outer loop parks as ``stopped``; later user/peer + messages can wake it again. Use sparingly. Prefer ``send_message_to_agent`` (asking the agent to wrap up) for soft-stop scenarios. Reach for ``stop_agent`` when @@ -754,11 +790,11 @@ async def stop_agent( in logs and telemetry. """ inner = _ctx(ctx) - bus = inner.get("bus") + coordinator = coordinator_from_context(inner) me = inner.get("agent_id") - if bus is None or me is None: + if coordinator is None or me is None: return json.dumps( - {"success": False, "error": "Bus or agent_id missing in context."}, + {"success": False, "error": "Agent coordinator or agent_id missing in context."}, ensure_ascii=False, default=str, ) @@ -771,31 +807,17 @@ async def stop_agent( ensure_ascii=False, default=str, ) - async with bus._lock: - if target_agent_id not in bus.statuses: - return json.dumps( - {"success": False, "error": f"Unknown agent_id: {target_agent_id}"}, - ensure_ascii=False, - default=str, - ) - target_status = bus.statuses.get(target_agent_id) - - if target_status in ("completed", "crashed", "stopped"): + if await coordinator.agent_info(target_agent_id) is None: return json.dumps( - { - "success": False, - "error": f"Target agent '{target_agent_id}' is already {target_status}.", - }, + {"success": False, "error": f"Unknown agent_id: {target_agent_id}"}, ensure_ascii=False, default=str, ) if cascade: - await bus.cancel_descendants_graceful(target_agent_id) + await coordinator.cancel_descendants_graceful(target_agent_id) else: - async with bus._lock: - bus.stopping.add(target_agent_id) - await bus.request_interrupt(target_agent_id, mode="after_turn") + await coordinator.request_stop(target_agent_id) logger.info( "stop_agent: target=%s cascade=%s reason=%r", diff --git a/strix/tools/finish/tool.py b/strix/tools/finish/tool.py index bbf7846..e0207c6 100644 --- a/strix/tools/finish/tool.py +++ b/strix/tools/finish/tool.py @@ -9,6 +9,8 @@ from typing import Any from agents import RunContextWrapper, function_tool +from strix.orchestration.coordinator import coordinator_from_context + logger = logging.getLogger(__name__) @@ -144,12 +146,38 @@ async def finish_scan( recommendations: Prioritized, actionable remediation. """ inner = ctx.context if isinstance(ctx.context, dict) else {} + coordinator = coordinator_from_context(inner) + me = inner.get("agent_id") + parent_id = inner.get("parent_id") + if coordinator is not None and parent_id is None and me is not None: + active_agents = await coordinator.active_agents_except(me) + else: + active_agents = [] + + if active_agents: + return json.dumps( + { + "success": False, + "scan_completed": False, + "error": "active_agents_remaining", + "message": ( + "Cannot finish scan while child agents are still active. " + "Wait for completion, send them finish instructions, or stop them first." + ), + "active_agents": active_agents, + }, + ensure_ascii=False, + default=str, + ) + result = await asyncio.to_thread( _do_finish, - parent_id=inner.get("parent_id"), + parent_id=parent_id, executive_summary=executive_summary, methodology=methodology, technical_analysis=technical_analysis, recommendations=recommendations, ) + if result.get("success") and result.get("scan_completed"): + inner["finish_scan_called"] = True return json.dumps(result, ensure_ascii=False, default=str)