0xallam
5fd2a64562
Three critical correctness fixes + six TUI/audit/UX fixes from the
parallel-agent audit. All changes verified by an end-to-end smoke
that builds, persists, and re-hydrates state across two simulated
process boundaries.
Critical (resume integrity):
1. ``bus.cancel_descendants_graceful`` now calls ``_maybe_snapshot``
after mutating the ``stopping`` set. Previously, a process crash
between user-initiated graceful-stop and the next finalize lost
the stop signal — respawned agents would run forever instead of
exiting. ``_respawn_subagents`` also gains a guard that skips
agents in ``stopping`` so a previously-cancelled agent is not
resurrected on resume.
2. ``Tracer.hydrate_from_run_dir`` now **raises** on corrupt
``vulnerabilities.json`` instead of swallowing the exception. The
prior behaviour silently reset ``vulnerability_reports`` to empty,
so the next ``add_vulnerability_report`` would allocate ``vuln-0001``
and overwrite the prior MD on disk — silent data loss.
3. ``--instruction`` passed on resume now reaches the model. The CLI
captures whether the user explicitly passed an instruction
(``args.user_explicit_instruction``) before ``_load_resume_state``
loads the persisted one. ``run_strix_scan`` reads
``scan_config["resume_instruction"]`` and, on resume, sends the
new instruction to root's bus inbox before calling
``run_with_continuation`` (which uses ``initial_input=[]`` for SDK
replay). The inject filter surfaces it on the next turn.
4. ``--resume X`` errors loudly when ``scan_state.json`` exists but
``bus.json`` doesn't. Previously this silently fresh-started in
the same dir, confusing the user who explicitly asked to resume.
TUI / audit / UX:
5. ``Tracer.hydrate_from_run_dir`` now reads ``bus.json`` too and
pre-populates ``tracer.agents`` from the snapshot's ``statuses`` /
``names`` / ``parent_of``. Before this, the TUI tree on resume
showed only currently-running agents; completed/crashed children
from the prior run were invisible.
6. ``Tracer.hydrate_from_run_dir`` also seeds ``self._llm_stats`` from
``bus.stats_live + bus.stats_completed`` so the resume's footer
shows cumulative tokens / requests across the prior run plus the
resume segment, instead of resetting to zero.
7. ``Tracer.save_run_data`` now also writes ``run_metadata.json``
(start_time, run_id, run_name, targets, status), and
``hydrate_from_run_dir`` restores ``start_time`` from it. Prior
behaviour reset start_time to ``now()`` on every Tracer init,
breaking the final report's duration calc on resumed scans.
8. Per-agent todos persist to ``{run_dir}/todos.json`` (atomic write
on every CRUD). ``hydrate_todos_from_disk`` (called from
``run_strix_scan``) reloads them so respawned subagents find
their lists intact. Previously, the module-level
``_todos_storage`` was lost on every process restart.
9. ``_load_resume_state`` validates each ``cloned_repo_path`` from
the persisted ``scan_state.json`` still exists on disk. Previously
a deleted clone dir would let the resume proceed with an empty
source tree, with agents silently scanning nothing.
Bonus: ``bus.finalize`` no longer pops ``parent_of`` and ``names``
for finalized agents. Routing protection (don't accept ``send`` to
finalized agents) comes from the ``statuses[id]`` terminal-state
check in ``send`` itself, so dropping those keys was overzealous and
made completed children invisible in ``view_agent_graph`` and the
TUI tree.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
462 lines
19 KiB
Python
462 lines
19 KiB
Python
"""``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)
|