Simplify SDK agent orchestration
This commit is contained in:
@@ -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
@@ -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
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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}"
|
||||
)
|
||||
@@ -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")
|
||||
@@ -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
@@ -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
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user