From f4834cd6f7a4162b04037d2b644cbd232d103aa6 Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sat, 25 Apr 2026 17:30:29 -0700 Subject: [PATCH] =?UTF-8?q?feat(orchestration):=20full=20parity=20with=20l?= =?UTF-8?q?egacy=20harness=20=E2=80=94=208=20gaps=20closed=20via=20SDK=20n?= =?UTF-8?q?atives?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 ```` XML envelope with the ``DO NOT echo back`` 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 ```` block. - G7: ``agent_finish`` emits the legacy ```` 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) --- strix/agents/factory.py | 2 + strix/entry.py | 4 +- strix/interface/tui.py | 13 ++- strix/orchestration/bus.py | 143 +++++++++++++++++++++++-- strix/orchestration/filter.py | 49 +++++++-- strix/orchestration/hooks.py | 36 ++++--- strix/orchestration/run_loop.py | 126 +++++++++++++++++----- strix/run_config_factory.py | 1 - strix/tools/agents_graph/tools.py | 170 ++++++++++++++++++++++++++---- 9 files changed, 459 insertions(+), 85 deletions(-) diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 3f37483..12139fa 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -32,6 +32,7 @@ from strix.tools.agents_graph.tools import ( agent_status, create_agent, send_message_to_agent, + stop_agent, view_agent_graph, wait_for_message, ) @@ -104,6 +105,7 @@ _BASE_TOOLS: tuple[Tool, ...] = ( send_message_to_agent, wait_for_message, create_agent, + stop_agent, ) diff --git a/strix/entry.py b/strix/entry.py index 1bfe8c8..4f336ad 100644 --- a/strix/entry.py +++ b/strix/entry.py @@ -36,7 +36,7 @@ from strix.runtime import session_manager if TYPE_CHECKING: - from agents.result import RunResult + from agents.result import RunResultBase logger = logging.getLogger(__name__) @@ -163,7 +163,7 @@ async def run_strix_scan( max_turns: int = STRIX_DEFAULT_MAX_TURNS, model: str | None = None, cleanup_on_exit: bool = True, -) -> RunResult: +) -> RunResultBase: """Run one Strix scan end-to-end against a freshly-prepared sandbox. Args: diff --git a/strix/interface/tui.py b/strix/interface/tui.py index d0d22e5..8c5882b 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -1727,15 +1727,24 @@ class StrixTUIApp(App): # type: ignore[misc] # Route to the agent's bus inbox. 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. + # stays responsive. After enqueuing the message, request a + # graceful interrupt of the agent's current turn so the user + # input is processed without waiting for the active LLM/tool + # call to finish — the SDK saves the in-flight turn cleanly + # before honoring ``cancel(mode="after_turn")``. 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.selected_agent_id, + 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._scan_loop, + ) self._displayed_events.clear() self._update_chat_view() diff --git a/strix/orchestration/bus.py b/strix/orchestration/bus.py index 135fc55..6a27dfb 100644 --- a/strix/orchestration/bus.py +++ b/strix/orchestration/bus.py @@ -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") diff --git a/strix/orchestration/filter.py b/strix/orchestration/filter.py index c7ddde4..09ee8a0 100644 --- a/strix/orchestration/filter.py +++ b/strix/orchestration/filter.py @@ -8,12 +8,15 @@ retries — so a single drain per turn doesn't lose messages on retry. from __future__ import annotations import logging +from datetime import UTC, datetime 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 @@ -23,9 +26,11 @@ 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. - Messages from peer agents are formatted with a labeled header so the - receiving model can attribute them. Messages from the literal sender - ``"user"`` (a real human via TUI) are added as plain user messages. + Peer-agent messages are wrapped in the legacy ```` + XML envelope (sender / metadata / content / delivery_info) so the + receiving model gets the same prompt-shape as pre-migration — including + the explicit "DO NOT echo back" instruction. 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`` @@ -49,14 +54,11 @@ async def inject_messages_filter(data: CallModelData) -> ModelInputData: if sender == "user": new_input.append({"role": "user", "content": content}) else: - msg_type = msg.get("type", "info") - priority = msg.get("priority", "normal") - header = f"[Message from agent {sender} | type={msg_type} | priority={priority}]" new_input.append( { "role": "user", - "content": f"{header}\n{content}", - } + "content": _format_inter_agent_message(bus, msg), + }, ) return ModelInputData( input=new_input, @@ -67,3 +69,34 @@ async def inject_messages_filter(data: CallModelData) -> ModelInputData: "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 in the legacy XML envelope. + + The wrapper carries an explicit "do not echo back" instruction that + the legacy harness used to keep models from quoting the entire + received message dict in their own next turn. + """ + 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") + timestamp = msg.get("timestamp") or datetime.now(UTC).isoformat() + content = str(msg.get("content", "")) + return ( + "\n" + " You have received a message from another " + "agent. Acknowledge and respond to the sender if needed; DO NOT echo " + "back this entire message block.\n" + f" {sender_name}" + f"{sender_id}\n" + f" {msg_type}" + f"{priority}" + f"{timestamp}\n" + f" {content}\n" + " This message was delivered during your task " + "execution. Please acknowledge and respond if needed." + "\n" + "" + ) diff --git a/strix/orchestration/hooks.py b/strix/orchestration/hooks.py index fb08f81..c785a6c 100644 --- a/strix/orchestration/hooks.py +++ b/strix/orchestration/hooks.py @@ -36,15 +36,28 @@ class StrixOrchestrationHooks(RunHooks[Any]): system_prompt: str | None, input_items: list[Any], ) -> None: + del agent, system_prompt try: - # Type contract guarantees ``input_items`` is list[TResponseInputItem]; - # we trust SDK here. The try/except below catches any surprise. 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(ctx.get("turn_count", 0)) - if max_turns >= 4 and cur == int(max_turns * 0.85): + 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``. + if cur >= int(max_turns * 0.85) and not stats.get("warned_85"): + stats["warned_85"] = True input_items.append( { "role": "user", @@ -52,9 +65,10 @@ class StrixOrchestrationHooks(RunHooks[Any]): "[System warning] You are at 85% of your iteration " "budget. Begin consolidating findings." ), - } + }, ) - elif max_turns >= 4 and cur == max_turns - 3: + if cur >= max_turns - 3 and not stats.get("warned_final"): + stats["warned_final"] = True input_items.append( { "role": "user", @@ -62,7 +76,7 @@ class StrixOrchestrationHooks(RunHooks[Any]): "[System warning] You have 3 iterations left. Your " "next tool call MUST be the finish tool." ), - } + }, ) except Exception: logger.exception("on_llm_start failed") @@ -94,7 +108,6 @@ class StrixOrchestrationHooks(RunHooks[Any]): output_tokens=int(getattr(usage, "output_tokens", 0) or 0), cached_tokens=cached, ) - ctx["turn_count"] = int(ctx.get("turn_count", 0)) + 1 except Exception: logger.exception("on_llm_end failed") @@ -178,11 +191,10 @@ class StrixOrchestrationHooks(RunHooks[Any]): if stays_alive: await bus.park(me) - # Reset per-cycle flags so the next ``Runner.run`` invocation - # can detect a fresh finish-tool call and re-trigger budget - # warnings against its own iteration count. + # 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 - ctx["turn_count"] = 0 else: await bus.finalize(me, final_status) except Exception: diff --git a/strix/orchestration/run_loop.py b/strix/orchestration/run_loop.py index 811a7bd..7c06262 100644 --- a/strix/orchestration/run_loop.py +++ b/strix/orchestration/run_loop.py @@ -1,21 +1,25 @@ -"""``run_with_continuation`` — interactive-mode demo-loop wrapper around ``Runner.run``. +"""``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 ``Runner.run`` returns on -``StopAtTools(...)`` and the agent is gone. +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. -This helper restores the legacy semantics using the SDK's canonical -demo-loop pattern (``agents/repl.py:run_demo_loop``): after each -``Runner.run`` cycle, ``await bus.wait_for_message(agent_id)``, drain -new messages, and re-invoke ``Runner.run`` with them as the next turn's -input. The session (if provided) preserves prior conversation across -cycles. +Behaviors restored from legacy: -Used by both the root scan loop in ``entry.run_strix_scan`` and the -child-agent loop in ``tools.agents_graph.tools.create_agent`` so every -interactive agent on the bus stays alive. +- **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 @@ -25,12 +29,14 @@ 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 RunResult + from agents.result import RunResultBase from agents.run_config import RunConfig from strix.orchestration.bus import AgentMessageBus @@ -39,6 +45,13 @@ if TYPE_CHECKING: 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, @@ -51,19 +64,8 @@ async def run_with_continuation( agent_id: str, interactive: bool, session: Session | None = None, -) -> RunResult: - """Run an agent once (non-interactive) or in a continuation loop (interactive). - - For non-interactive runs this is a thin wrapper around - ``Runner.run`` and returns its result. - - For interactive runs the function loops: after each ``Runner.run`` - returns, it awaits ``bus.wait_for_message(agent_id)``, drains any - accumulated messages from the inbox, formats them as the next - turn's user input, and invokes ``Runner.run`` again. The loop ends - when the wait gets cancelled (e.g. parent ``cancel_descendants`` or - user-issued KeyboardInterrupt). - """ +) -> RunResultBase: + """Run an agent once (non-interactive) or in a continuation loop (interactive).""" kwargs: dict[str, Any] = { "input": initial_input, "run_config": run_config, @@ -74,16 +76,38 @@ async def run_with_continuation( if session is not None: kwargs["session"] = session - result: RunResult = await Runner.run(agent, **kwargs) + # 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 while True: + if agent_id in bus.stopping: + return result + try: - await bus.wait_for_message(agent_id) + 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: return result + except TimeoutError: + kwargs["input"] = _TIMEOUT_RESUME_MESSAGE + result = await _run_streamed(agent, bus, agent_id, **kwargs) + continue pending = await bus.drain(agent_id) if not pending: @@ -95,4 +119,48 @@ async def run_with_continuation( continue kwargs["input"] = next_input - result = await Runner.run(agent, **kwargs) + result = await _run_streamed(agent, bus, agent_id, **kwargs) + + +async def _run_streamed( + agent: Any, + bus: AgentMessageBus, + agent_id: str, + **kwargs: Any, +) -> RunResultBase: + """Drive one ``Runner.run_streamed`` cycle to completion. + + Catches hard model failures (after SDK retries are exhausted) and + parks the agent in ``llm_failed`` until a user message arrives, + matching legacy ``state.llm_failed`` semantics. Programmer errors + (``UserError``), max-turn breaches, and explicit cancellation + propagate to the caller. + """ + interactive = bool(kwargs.get("context", {}).get("interactive", False)) + while True: + streamed = Runner.run_streamed(agent, **kwargs) + try: + async with bus.attach_stream(agent_id, streamed): + async for _event in streamed.stream_events(): + pass + except (UserError, MaxTurnsExceeded, asyncio.CancelledError): + raise + except (AgentsException, APIError): + if not interactive: + raise + logger.exception( + "LLM hard failure for agent %s; awaiting user resume", + agent_id, + ) + await bus.mark_llm_failed(agent_id) + await bus.wait_for_user_message(agent_id) + pending = await bus.drain(agent_id) + next_input = "\n\n".join( + str(msg.get("content", "")).strip() for msg in pending if msg.get("content") + ) + if not next_input: + continue + kwargs["input"] = next_input + continue + else: + return streamed diff --git a/strix/run_config_factory.py b/strix/run_config_factory.py index 1146330..88af371 100644 --- a/strix/run_config_factory.py +++ b/strix/run_config_factory.py @@ -150,7 +150,6 @@ def make_agent_context( "model": model, "model_settings": model_settings, "max_turns": max_turns, - "turn_count": 0, "agent_finish_called": False, "is_whitebox": is_whitebox, "interactive": interactive, diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index 8b6ee42..9642dd2 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -42,6 +42,38 @@ def _ctx(ctx: RunContextWrapper) -> dict[str, Any]: return ctx.context if isinstance(ctx.context, dict) else {} +def _render_completion_report( + *, + agent_name: str, + agent_id: str, + task: str, + success: bool, + result_summary: str, + findings: list[str], + recommendations: list[str], +) -> str: + """Render an ```` XML payload (legacy parity).""" + from datetime import UTC, datetime + from html import escape + + status = "SUCCESS" if success else "FAILED" + completion_time = datetime.now(UTC).isoformat() + findings_xml = "".join(f"{escape(f)}" for f in findings) + recs_xml = "".join(f"{escape(r)}" for r in recommendations) + return ( + "\n" + f" {escape(agent_name)}" + f"{escape(agent_id)}" + f"{escape(task)}" + f"{status}" + f"{completion_time}\n" + f" {escape(result_summary)}" + f"{findings_xml}" + f"{recs_xml}\n" + "" + ) + + @function_tool(timeout=30) async def view_agent_graph(ctx: RunContextWrapper) -> str: """Print the multi-agent tree — every agent, its parent, its status. @@ -411,21 +443,26 @@ async def create_agent( await bus.register(child_id, name, parent_id) - parent_history = inner.get("parent_input_items") if inherit_context else None + # ``ctx.turn_input`` carries the parent's full conversation up to and + # including the call that's currently invoking ``create_agent`` + # (populated by SDK at ``run_internal/turn_resolution.py:806``). + # Wrap as a single read-only block so the child sees the parent's + # reasoning as background but doesn't try to continue parent's turns. + parent_history = list(ctx.turn_input) if inherit_context and ctx.turn_input else [] initial_input: list[TResponseInputItem] = [] if parent_history: + rendered = json.dumps(parent_history, ensure_ascii=False, default=str) initial_input.append( { "role": "user", - "content": "[Inherited context from parent — read-only history]", - } - ) - initial_input.extend(parent_history) - initial_input.append( - { - "role": "user", - "content": "[End of inherited context]", - } + "content": ( + "\n" + f"{rendered}\n" + "\n" + "Use the above as background only; do not continue the " + "parent's work. Your task follows." + ), + }, ) initial_input.append( { @@ -456,6 +493,9 @@ async def create_agent( run_id=inner.get("run_id"), agent_factory=factory, ) + # Stash the task string for ``agent_finish`` to echo back in its + # XML completion report. + child_ctx["task"] = task child_run_config = make_run_config( sandbox_session=inner.get("sandbox_session"), @@ -569,17 +609,14 @@ async def agent_finish( if report_to_parent: async with bus._lock: agent_name = bus.names.get(me, me) - report = json.dumps( - { - "kind": "agent_completion_report", - "from": agent_name, - "agent_id": me, - "success": success, - "summary": result_summary, - "findings": list(findings or []), - "recommendations": list(final_recommendations or []), - }, - ensure_ascii=False, + report = _render_completion_report( + agent_name=agent_name, + agent_id=me, + task=str(inner.get("task", "")), + success=success, + result_summary=result_summary, + findings=list(findings or []), + recommendations=list(final_recommendations or []), ) await bus.send( parent_id, @@ -606,3 +643,94 @@ async def agent_finish( ensure_ascii=False, default=str, ) + + +@function_tool(timeout=30) +async def stop_agent( + ctx: RunContextWrapper, + target_agent_id: str, + cascade: bool = True, + reason: str = "", +) -> str: + """Gracefully stop a running agent (and optionally its descendants). + + 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". + + Use sparingly. Prefer ``send_message_to_agent`` (asking the agent + to wrap up) for soft-stop scenarios. Reach for ``stop_agent`` when + a child has gone off-track and won't self-correct. + + Args: + target_agent_id: The 8-char id from ``view_agent_graph`` / + ``create_agent``. Cannot stop yourself. + cascade: If ``True`` (default), also stop every descendant of + ``target_agent_id`` leaves-first. ``False`` stops only the + target. + reason: Optional human-readable reason for the stop, surfaced + in logs and telemetry. + """ + inner = _ctx(ctx) + bus = inner.get("bus") + me = inner.get("agent_id") + if bus is None or me is None: + return json.dumps( + {"success": False, "error": "Bus or agent_id missing in context."}, + ensure_ascii=False, + default=str, + ) + if target_agent_id == me: + return json.dumps( + { + "success": False, + "error": "Cannot stop yourself; call agent_finish or finish_scan instead.", + }, + 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"): + return json.dumps( + { + "success": False, + "error": f"Target agent '{target_agent_id}' is already {target_status}.", + }, + ensure_ascii=False, + default=str, + ) + + if cascade: + await bus.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") + + logger.info( + "stop_agent: target=%s cascade=%s reason=%r", + target_agent_id, + cascade, + reason, + ) + return json.dumps( + { + "success": True, + "target_agent_id": target_agent_id, + "cascade": cascade, + "reason": reason, + "note": "Cancellation is graceful — current turn completes first.", + }, + ensure_ascii=False, + default=str, + )