Simplify SDK agent orchestration

This commit is contained in:
0xallam
2026-04-26 09:25:47 -07:00
parent 53188a7583
commit 5ec1e0786f
12 changed files with 1156 additions and 1309 deletions
+69
View File
@@ -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.
+69 -12
View File
@@ -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,
+12 -17
View File
@@ -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,
)
+8 -8
View File
@@ -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.
"""
-461
View File
@@ -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)
+712
View File
@@ -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,
}
-96
View File
@@ -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}"
)
-309
View File
@@ -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")
-183
View File
@@ -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
+86 -73
View File
@@ -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/<child_id>.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:
+171 -149
View File
@@ -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",
+29 -1
View File
@@ -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)