mortal/strix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(orchestration): full parity with legacy harness — 8 gaps closed via SDK natives

Browse Source 
Audit found 8 behavioral gaps between post-migration and the legacy
``BaseAgent.agent_loop``. All 8 are now closed using SDK-native
primitives — no custom workarounds, no shadow state machines.

What was broken / different:

- G1: ``inherit_context`` was dead code; children always started fresh.
- G2: TUI user message couldn't interrupt an in-flight LLM/tool turn.
- G3: ``llm_failed`` state never set; hard failures propagated as crashes.
- G4: No graceful ``stop_agent`` tool.
- G5: Parked subagents waited forever (no auto-resume timeout).
- G6: Inter-agent messages used a plain header instead of legacy XML.
- G7: Completion reports used JSON instead of legacy XML.
- G11/G12: Turn counter reset per cycle; budget warnings could re-fire.

What we did:

Bus extensions (``orchestration/bus.py``):
- ``streams`` registry + ``attach_stream`` ctx manager + ``request_interrupt``
  for SDK-native ``RunResultStreaming.cancel(mode="after_turn")``.
- ``mark_llm_failed`` + ``wait_for_user_message`` (filtered: only ``from="user"``
  satisfies; peer messages don't unstick a stuck model).
- ``stopping: set[str]`` for graceful programmatic exit.
- ``cancel_descendants_graceful`` — leaves-first via ``request_interrupt``.
- ``record_usage`` increments ``calls`` unconditionally so it doubles as the
  per-agent-lifetime turn counter (legacy ``state.iteration`` parity).
- ``warned_85`` / ``warned_final`` flags on ``stats_live`` for once-fire
  budget warnings.

Run loop rewrite (``orchestration/run_loop.py``):
- ``Runner.run`` → ``Runner.run_streamed`` with ``bus.attach_stream`` so
  cancel has a target. Catch ``(AgentsException, APIError)`` after retries
  exhaust; in interactive mode call ``mark_llm_failed`` + wait for user.
- ``UserError`` / ``MaxTurnsExceeded`` / ``CancelledError`` propagate.
- Outer loop: ``asyncio.wait_for(bus.wait_for_message, timeout=300)`` for
  interactive subagents (root waits forever). ``TimeoutError`` injects
  ``"Waiting timeout reached. Resuming execution."``.
- Honors ``bus.stopping`` at top of each iteration.

Hooks (``orchestration/hooks.py``):
- Counter source moved from per-cycle ``ctx["turn_count"]`` to
  per-lifetime ``bus.stats_live[agent_id]["calls"]``.
- Warnings guarded by once-flags — exactly-once across all cycles.

Filter (``orchestration/filter.py``):
- Restored legacy ``<inter_agent_message>`` XML envelope with the
  ``<delivery_notice>DO NOT echo back</delivery_notice>`` instruction.

Agents-graph (``tools/agents_graph/tools.py``):
- G1: ``create_agent`` reads ``ctx.turn_input`` (SDK populates it before
  tool execution at ``run_internal/turn_resolution.py:806``). Wraps as
  one ``<inherited_context_from_parent>`` block.
- G7: ``agent_finish`` emits the legacy ``<agent_completion_report>``
  XML. ``child_ctx["task"] = task`` threaded so the report echoes the
  original task.
- G4: New ``stop_agent`` tool — refuses self-stop, refuses already-
  finalized targets, ``cascade=True`` uses ``cancel_descendants_graceful``.

TUI (``interface/tui.py``):
- ``_send_user_message`` schedules ``bus.send`` AND
  ``bus.request_interrupt(target, mode="after_turn")`` — SDK finishes
  current turn cleanly, next cycle picks up the user's message.

Factory (``agents/factory.py``):
- Registered ``stop_agent`` in ``_BASE_TOOLS``.

Out of scope:
- G8 (``[ABORTED BY USER]`` marker) is auto-resolved by G2 — the SDK
  saves the full assistant message before honoring
  ``cancel(mode="after_turn")``, so partial content is preserved in the
  session.

Verified all bus behaviors with a smoke test. Lint at baseline.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
0xallam
2026-04-25 17:30:29 -07:00
parent 5896f25cec
commit f4834cd6f7
9 changed files with 459 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
strix/orchestration/bus.py
+133 -10
View File
@@ -8,8 +8,15 @@ Strix scan.
from __future__ import annotations
import asyncio
import contextlib
from dataclasses import dataclass, field
from typing import Any
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import AsyncIterator
from agents.result import RunResultStreaming
@dataclass
@@ -23,21 +30,30 @@ class AgentMessageBus:
``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``.
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.
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)
_events: dict[str, asyncio.Event] = field(default_factory=dict)
_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
@@ -59,6 +75,8 @@ class AgentMessageBus:
"cached": 0,
"cost": 0.0,
"calls": 0,
"warned_85": False,
"warned_final": False,
}
async def send(self, target: str, msg: dict[str, Any]) -> None:
@@ -91,6 +109,23 @@ class AgentMessageBus:
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.
@@ -108,20 +143,30 @@ class AgentMessageBus:
"""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).
"""
if usage is None:
return
async with self._lock:
stats = self.stats_live.setdefault(
agent_id,
{"in": 0, "out": 0, "cached": 0, "cost": 0.0, "calls": 0},
{
"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
stats["calls"] += 1
async def finalize(self, agent_id: str, status: str) -> None:
"""Move an agent from live to completed; clean up routing state.
@@ -135,6 +180,8 @@ class AgentMessageBus:
self.inboxes.pop(agent_id, None)
self.parent_of.pop(agent_id, None)
self.names.pop(agent_id, None)
self.streams.pop(agent_id, None)
self.stopping.discard(agent_id)
self._events.pop(agent_id, None)
async def park(self, agent_id: str) -> None:
@@ -150,13 +197,62 @@ class AgentMessageBus:
if agent_id in self.statuses:
self.statuses[agent_id] = "waiting"
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"
@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:
return False
streamed.cancel(mode=mode) # type: ignore[arg-type] # mode is a Literal
return True
async def total_stats(self) -> dict[str, Any]:
"""Snapshot of live + completed stats."""
"""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, value in stats.items():
agg[key] = agg.get(key, 0) + value
for key in agg:
agg[key] += stats.get(key, 0)
return agg
async def cancel_descendants(self, root_agent_id: str) -> None:
@@ -165,6 +261,10 @@ class AgentMessageBus:
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]
@@ -182,3 +282,26 @@ class AgentMessageBus:
*(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
]
for _aid, streamed in streams_to_cancel:
streamed.cancel(mode="after_turn")