From 1ac32df817c5104bf46310a340010756cc91893a Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sat, 25 Apr 2026 00:36:00 -0700 Subject: [PATCH] =?UTF-8?q?feat(migration):=20phase=203=20=E2=80=94=20mult?= =?UTF-8?q?i-agent=20graph=20tools=20+=20Runner=20bridge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Six SDK function tools that drive the AgentMessageBus from Phase 0, replacing the legacy _agent_graph / _agent_messages / _agent_instances globals: - view_agent_graph: render parent/child tree from bus.parent_of with a per-status summary (running / waiting / completed / crashed / stopped). - agent_status: per-agent lifecycle + pending-message count snapshot. - send_message_to_agent: queue into bus.inboxes; rejects sends to finalized targets so the model gets feedback rather than a silent drop (the bus's own send method drops to support the C13 cleanup, but the tool surfaces it as a structured error). - wait_for_message: poll inbox once per second up to timeout. Polling rather than asyncio.Event because a missed wakeup on Event would be hard to debug; the bus already serializes through its own lock. - create_agent: spawn a child via asyncio.create_task(Runner.run(...)). Pulls an agent_factory callable from ctx.context (the Phase 5 root assembly is the one that wires it in). Registers the child with the bus before the task starts, stores the task handle in bus.tasks so cancel_descendants can cascade (C9), builds the child's identity block + optional inherited parent context, and runs the child with StrixOrchestrationHooks. - agent_finish: subagent-only termination. Flips agent_finish_called so the on_agent_end hook records "completed" instead of "crashed" (C8), and posts a structured XML envelope to the parent's inbox. run_config_factory.make_agent_context grows two fields: sandbox_client (reused across child runs) and agent_factory (Phase 3 needs it; Phase 5 fills it in). PLC0415 fixed by hoisting the openai.types.shared.Reasoning import to module-level. Tests: 17 new tests in test_sdk_graph_tools.py — registration, all six tools' happy and error paths, real AgentMessageBus integration so the tools exercise production code paths, create_agent verified for spawn shape (task created, bus registered, identity block in input) plus a bus.cancel_descendants integration check. Refs: PLAYBOOK.md §4.3, AUDIT_R2 §1.4 (cancel_descendants), AUDIT_R3 C8. Co-Authored-By: Claude Opus 4.7 (1M context) --- pyproject.toml | 1 + strix/run_config_factory.py | 13 +- .../agents_graph/agents_graph_sdk_tools.py | 483 +++++++++++++++++ tests/tools/test_sdk_graph_tools.py | 512 ++++++++++++++++++ 4 files changed, 1007 insertions(+), 2 deletions(-) create mode 100644 strix/tools/agents_graph/agents_graph_sdk_tools.py create mode 100644 tests/tools/test_sdk_graph_tools.py diff --git a/pyproject.toml b/pyproject.toml index cc961a2..6dd6251 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -298,6 +298,7 @@ ignore = [ "strix/tools/terminal/terminal_sdk_tool.py" = ["TC002"] "strix/tools/python/python_sdk_tool.py" = ["TC002"] "strix/tools/proxy/proxy_sdk_tools.py" = ["TC002"] +"strix/tools/agents_graph/agents_graph_sdk_tools.py" = ["TC002"] # Sandbox dispatch helper has many short-circuit error returns (auth fail, # size cap, decode fail, etc). Each is a distinct, documented failure mode # the model needs to see verbatim — collapsing them harms readability. diff --git a/strix/run_config_factory.py b/strix/run_config_factory.py index fec7562..06ab512 100644 --- a/strix/run_config_factory.py +++ b/strix/run_config_factory.py @@ -27,6 +27,7 @@ from agents.retry import ( retry_policies, ) from agents.sandbox import SandboxRunConfig +from openai.types.shared import Reasoning from strix.llm.multi_provider_setup import build_multi_provider from strix.orchestration.filter import inject_messages_filter @@ -130,8 +131,6 @@ def make_run_config( ), ) if reasoning_effort is not None: - from openai.types.shared import Reasoning - base_settings = base_settings.resolve( ModelSettings(reasoning=Reasoning(effort=reasoning_effort)), ) @@ -174,6 +173,8 @@ def make_agent_context( is_whitebox: bool = False, diff_scope: dict[str, Any] | None = None, run_id: str | None = None, + sandbox_client: Any | None = None, + agent_factory: Any | None = None, ) -> dict[str, Any]: """Build the per-agent ``context`` dict passed to ``Runner.run(context=...)``. @@ -184,10 +185,17 @@ def make_agent_context( C21 (AUDIT_R3): includes ``is_whitebox``, ``diff_scope`` and ``run_id`` fields that the legacy code relied on but the original PLAYBOOK §2.10 skeleton omitted. + + ``agent_factory`` is a callable ``(name, skills) -> agents.Agent`` used by + the ``create_agent`` graph tool to spin up children. The actual factory + lives in the Phase 4/5 root-assembly module; Phase 3 only requires that + it be present in context. ``sandbox_client`` is the host-side Docker + subclass; ``create_agent`` reuses it across child runs. """ return { "bus": bus, "sandbox_session": sandbox_session, + "sandbox_client": sandbox_client, "sandbox_token": sandbox_token, "tool_server_host_port": tool_server_host_port, "caido_host_port": caido_host_port, @@ -203,4 +211,5 @@ def make_agent_context( "is_whitebox": is_whitebox, "diff_scope": diff_scope, "run_id": run_id, + "agent_factory": agent_factory, } diff --git a/strix/tools/agents_graph/agents_graph_sdk_tools.py b/strix/tools/agents_graph/agents_graph_sdk_tools.py new file mode 100644 index 0000000..318b636 --- /dev/null +++ b/strix/tools/agents_graph/agents_graph_sdk_tools.py @@ -0,0 +1,483 @@ +"""SDK function-tool wrappers for the multi-agent graph tools. + +Six tools that read/write the :class:`AgentMessageBus` (built in Phase 0, +``strix.orchestration.bus``): + +- ``view_agent_graph``: render the parent/child tree from ``bus.parent_of``. +- ``agent_status``: per-agent status + pending message count. +- ``send_message_to_agent``: peer-to-peer message into a child/sibling inbox. +- ``wait_for_message``: poll our own inbox until a message arrives or the + timeout expires (the legacy harness's "I'm idle, wake me on inbox"). +- ``create_agent``: spawn a child via ``asyncio.create_task(Runner.run(...))``; + registers the child with the bus and stores its task handle so root cancels + cascade (C9, ``bus.cancel_descendants``). +- ``agent_finish``: subagents only — flips ``agent_finish_called`` so the + on_agent_end hook records "completed" rather than "crashed" (C8), and + posts a structured completion report to the parent's inbox. + +The legacy ``strix.tools.agents_graph.agents_graph_actions`` is left +untouched — it still drives the legacy harness. These wrappers only +target the bus and don't touch the legacy ``_agent_graph`` dict. + +References: + - PLAYBOOK.md §4.3 + - AUDIT_R2.md §1.4 (cancel_descendants — Runner.run task handle stored + in bus.tasks so a root cancel walks the tree) + - AUDIT_R3.md C8 (crash detection via on_agent_end + agent_finish_called) +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import uuid +from typing import TYPE_CHECKING, Any, Literal + +from agents import RunContextWrapper, Runner +from agents.items import TResponseInputItem + +from strix.orchestration.hooks import StrixOrchestrationHooks +from strix.run_config_factory import make_agent_context, make_run_config +from strix.tools._decorator import strix_tool + + +if TYPE_CHECKING: + from collections.abc import Callable + + from agents import Agent as SDKAgent + + +logger = logging.getLogger(__name__) + + +def _dump(result: dict[str, Any]) -> str: + return json.dumps(result, ensure_ascii=False, default=str) + + +@strix_tool(timeout=30) +async def view_agent_graph(ctx: RunContextWrapper) -> str: + """Render the multi-agent tree starting from each root. + + Output is a single string the model can parse: indented bullet list, + one line per agent, status in brackets. Roots are agents whose + ``parent_of[id]`` is ``None``. + """ + inner = ctx.context if isinstance(ctx.context, dict) else {} + bus = inner.get("bus") + me = inner.get("agent_id") + if bus is None: + return _dump({"success": False, "error": "Bus not initialized in context."}) + + async with bus._lock: + parent_of = dict(bus.parent_of) + statuses = dict(bus.statuses) + names = dict(bus.names) + + lines: list[str] = [] + + def render(aid: str, depth: int) -> None: + status = statuses.get(aid, "?") + marker = " ← you" if aid == me else "" + lines.append(f"{' ' * depth}- {names.get(aid, aid)} ({aid}) [{status}]{marker}") + for child, p in parent_of.items(): + if p == aid: + render(child, depth + 1) + + roots = [aid for aid, parent in parent_of.items() if parent is None] + for root in roots: + render(root, 0) + + summary = { + "total": len(parent_of), + "running": sum(1 for s in statuses.values() if s == "running"), + "waiting": sum(1 for s in statuses.values() if s == "waiting"), + "completed": sum(1 for s in statuses.values() if s == "completed"), + "crashed": sum(1 for s in statuses.values() if s == "crashed"), + "stopped": sum(1 for s in statuses.values() if s == "stopped"), + } + return _dump( + { + "success": True, + "graph_structure": "\n".join(lines) or "(no agents)", + "summary": summary, + } + ) + + +@strix_tool(timeout=30) +async def agent_status(ctx: RunContextWrapper, agent_id: str) -> str: + """Inspect one agent's lifecycle state and pending message count.""" + inner = ctx.context if isinstance(ctx.context, dict) else {} + bus = inner.get("bus") + if bus is None: + return _dump({"success": False, "error": "Bus not initialized in context."}) + + async with bus._lock: + if agent_id not in bus.statuses: + return _dump( + { + "success": False, + "error": f"Unknown agent_id: {agent_id}", + } + ) + return _dump( + { + "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, [])), + } + ) + + +@strix_tool(timeout=30) +async def send_message_to_agent( + ctx: RunContextWrapper, + target_agent_id: str, + message: str, + message_type: Literal["query", "instruction", "information"] = "information", + priority: Literal["low", "normal", "high", "urgent"] = "normal", +) -> str: + """Queue a message for another agent's inbox. + + The target's next ``inject_messages_filter`` pass (top of its next LLM + turn) drains the inbox and surfaces the message wrapped in + ````. Messages to a finalized agent are dropped + silently by the bus (C13). + """ + inner = ctx.context if isinstance(ctx.context, dict) else {} + bus = inner.get("bus") + me = inner.get("agent_id") + if bus is None or me is None: + return _dump({"success": False, "error": "Bus or agent_id missing in context."}) + + async with bus._lock: + if target_agent_id not in bus.statuses: + return _dump( + { + "success": False, + "error": f"Target agent '{target_agent_id}' not found.", + } + ) + target_status = bus.statuses.get(target_agent_id) + + if target_status in ("completed", "crashed", "stopped"): + return _dump( + { + "success": False, + "error": f"Target agent '{target_agent_id}' is {target_status}; message dropped.", + } + ) + + msg_id = f"msg_{uuid.uuid4().hex[:8]}" + await bus.send( + target_agent_id, + { + "id": msg_id, + "from": me, + "content": message, + "type": message_type, + "priority": priority, + }, + ) + return _dump( + { + "success": True, + "message_id": msg_id, + "target_agent_id": target_agent_id, + "delivery_status": "queued", + } + ) + + +# Polling cadence for ``wait_for_message``. 1s matches the PLAYBOOK +# skeleton; tighter would burn CPU, slacker would feel laggy when a sibling +# delivers a message right after the wait starts. +_WAIT_POLL_SECONDS = 1.0 + + +@strix_tool(timeout=601) +async def wait_for_message( + ctx: RunContextWrapper, + reason: str = "Waiting for messages from other agents", + timeout_seconds: int = 600, +) -> str: + """Block this agent's turn until a message arrives or ``timeout_seconds``. + + Implementation polls ``bus.inboxes`` once per second. Cheaper than an + asyncio.Event because the message bus already serializes through its + own lock — a missed wakeup on Event would be subtle to debug, while + polling is trivially correct. + + Args: + reason: Human-readable note shown in graph snapshots while waiting. + timeout_seconds: Cap on the wait. 600s matches the legacy default. + """ + inner = ctx.context if isinstance(ctx.context, dict) else {} + bus = inner.get("bus") + me = inner.get("agent_id") + if bus is None or me is None: + return _dump({"success": False, "error": "Bus or agent_id missing in context."}) + + async with bus._lock: + bus.statuses[me] = "waiting" + + deadline = asyncio.get_event_loop().time() + timeout_seconds + 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 _dump( + { + "success": True, + "status": "message_arrived", + "pending_messages": pending, + "reason": reason, + } + ) + 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" + + return _dump( + { + "success": True, + "status": "timeout", + "timeout_seconds": timeout_seconds, + "reason": reason, + "note": "No messages within timeout — continue work or call agent_finish.", + } + ) + + +@strix_tool(timeout=120) +async def create_agent( + ctx: RunContextWrapper, + name: str, + task: str, + inherit_context: bool = True, + skills: list[str] | None = None, +) -> str: + """Spawn a child agent that runs in parallel via ``asyncio.create_task``. + + The child's ``Runner.run`` task handle is stored in ``bus.tasks[child_id]`` + so a root-level cancel can cascade to descendants (C9). The child is + registered with the bus before the task starts so messages aimed at it + don't get dropped during the brief register→start window. + + Args: + name: Human-readable child name (also stored in ``bus.names``). + task: The task description handed to the child agent. + inherit_context: When True, the child receives a copy of the parent's + input items as background context, wrapped in + ````. Default True. + skills: Optional list of skill names the child should preload. + + Returns a JSON-encoded ``{"success": ..., "agent_id": ...}``. + """ + inner = ctx.context if isinstance(ctx.context, dict) else {} + bus = inner.get("bus") + parent_id = inner.get("agent_id") + factory: Callable[..., SDKAgent] | None = inner.get("agent_factory") + + if bus is None or parent_id is None: + return _dump({"success": False, "error": "Bus or agent_id missing in context."}) + if factory is None: + return _dump( + { + "success": False, + "error": ( + "No agent_factory in context. " + "The root assembly must inject one via make_agent_context." + ), + } + ) + + child_id = uuid.uuid4().hex[:8] + + try: + child_agent = factory(name=name, skills=skills or []) + except Exception as e: + logger.exception("agent_factory raised while building child '%s'", name) + return _dump( + { + "success": False, + "error": f"agent_factory failed: {e!s}", + } + ) + + await bus.register(child_id, name, parent_id) + + # Build the child's input. Identity injection mirrors the legacy + # envelope so the child's system prompt's existing + # rules around self-identity still apply. + parent_history = inner.get("parent_input_items") if inherit_context else None + initial_input: list[TResponseInputItem] = [] + if parent_history: + initial_input.append( + { + "role": "user", + "content": "", + } + ) + initial_input.extend(parent_history) + initial_input.append( + { + "role": "user", + "content": "", + } + ) + initial_input.append( + { + "role": "user", + "content": ( + f"\n" + f"You are agent {name} ({child_id}). Parent is {parent_id}.\n" + f"Maintain self-identity. Use agent_finish when complete.\n" + f"" + ), + } + ) + initial_input.append({"role": "user", "content": task}) + + child_ctx = make_agent_context( + bus=bus, + sandbox_session=inner.get("sandbox_session"), + sandbox_client=inner.get("sandbox_client"), + sandbox_token=inner.get("sandbox_token"), + tool_server_host_port=inner.get("tool_server_host_port"), + caido_host_port=inner.get("caido_host_port"), + agent_id=child_id, + agent_name=name, + parent_id=parent_id, + tracer=inner.get("tracer"), + model=inner.get("model", "strix/claude-sonnet-4.6"), + model_settings=inner.get("model_settings"), + max_turns=int(inner.get("max_turns", 300)), + is_whitebox=bool(inner.get("is_whitebox", False)), + diff_scope=inner.get("diff_scope"), + run_id=inner.get("run_id"), + agent_factory=factory, + ) + + child_run_config = make_run_config( + sandbox_session=inner.get("sandbox_session"), + sandbox_client=inner.get("sandbox_client"), + model=inner.get("model", "strix/claude-sonnet-4.6"), + model_settings_override=inner.get("model_settings"), + ) + + task_handle = asyncio.create_task( + Runner.run( + child_agent, + input=initial_input, + run_config=child_run_config, + context=child_ctx, + hooks=StrixOrchestrationHooks(), + max_turns=int(inner.get("max_turns", 300)), + ), + name=f"agent-{name}-{child_id}", + ) + async with bus._lock: + bus.tasks[child_id] = task_handle + + return _dump( + { + "success": True, + "agent_id": child_id, + "name": name, + "parent_id": parent_id, + "message": f"Spawned '{name}' ({child_id}) running in parallel.", + } + ) + + +@strix_tool(timeout=30) +async def agent_finish( + ctx: RunContextWrapper, + result_summary: str, + findings: list[str] | None = None, + success: bool = True, + report_to_parent: bool = True, + final_recommendations: list[str] | None = None, +) -> str: + """Subagent-only termination: post a completion report and signal the SDK. + + Sets ``ctx.context['agent_finish_called'] = True`` so the on_agent_end + hook records "completed" rather than "crashed". The SDK terminates the + child's loop because every child is built with + ``tool_use_behavior={"stop_at_tool_names": ["agent_finish"]}`` (C4). + + Root agents must call ``finish_scan`` instead. This tool refuses to run + when ``parent_id`` is None. + """ + inner = ctx.context if isinstance(ctx.context, dict) else {} + bus = inner.get("bus") + me = inner.get("agent_id") + if bus is None or me is None: + return _dump({"success": False, "error": "Bus or agent_id missing in context."}) + + parent_id = inner.get("parent_id") + if parent_id is None: + return _dump( + { + "success": False, + "agent_completed": False, + "error": ( + "agent_finish is for subagents. Root/main agents must call finish_scan instead." + ), + "parent_notified": False, + } + ) + + inner["agent_finish_called"] = True + + parent_notified = False + if report_to_parent: + findings_xml = "\n".join(f" {f}" for f in (findings or [])) + rec_xml = "\n".join( + f" {r}" for r in (final_recommendations or []) + ) + async with bus._lock: + agent_name = bus.names.get(me, me) + report = ( + f"\n" + f" {result_summary}\n" + f" \n{findings_xml}\n \n" + f" \n{rec_xml}\n \n" + f"" + ) + await bus.send( + parent_id, + { + "id": f"report_{uuid.uuid4().hex[:8]}", + "from": me, + "content": report, + "type": "completion", + "priority": "high", + }, + ) + parent_notified = True + + return _dump( + { + "success": True, + "agent_completed": True, + "parent_notified": parent_notified, + "agent_id": me, + "summary": result_summary, + "findings_count": len(findings or []), + "has_recommendations": bool(final_recommendations), + } + ) diff --git a/tests/tools/test_sdk_graph_tools.py b/tests/tools/test_sdk_graph_tools.py new file mode 100644 index 0000000..0eb0daf --- /dev/null +++ b/tests/tools/test_sdk_graph_tools.py @@ -0,0 +1,512 @@ +"""Phase 3 tests for the multi-agent graph SDK tools. + +Six tools: view_agent_graph, agent_status, send_message_to_agent, +wait_for_message, create_agent, agent_finish. + +Strategy: + +- Build a real ``AgentMessageBus`` in each test and put it under + ``ctx.context['bus']`` so the tools exercise the same code path + production runs do. +- ``create_agent`` is the only tool that touches ``Runner.run`` and + ``asyncio.create_task``. Its test injects a stub agent factory and + patches the SDK's ``Runner.run`` to a sentinel coroutine — we verify + the spawn shape (bus.register called, task handle stored, identity + block in initial input) without spinning up a real LLM. +- ``wait_for_message`` is exercised on both branches: a message arrives + mid-poll, and the timeout path. +""" + +from __future__ import annotations + +import asyncio +import json +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import AsyncMock, patch + +import pytest +from agents.tool import FunctionTool + +from strix.orchestration.bus import AgentMessageBus +from strix.tools.agents_graph.agents_graph_sdk_tools import ( + agent_finish, + agent_status, + create_agent, + send_message_to_agent, + view_agent_graph, + wait_for_message, +) + + +@dataclass +class _Ctx: + context: dict[str, Any] = field(default_factory=dict) + + +async def _invoke(tool: FunctionTool, ctx: _Ctx, **kwargs: Any) -> dict[str, Any]: + from agents.tool_context import ToolContext + + tool_ctx = ToolContext( + context=ctx.context, + usage=None, + tool_name=tool.name, + tool_call_id="test-call-id", + tool_arguments=json.dumps(kwargs), + ) + result = await tool.on_invoke_tool(tool_ctx, json.dumps(kwargs)) + assert isinstance(result, str) + decoded = json.loads(result) + assert isinstance(decoded, dict) + return decoded + + +async def _make_bus_with_agents() -> AgentMessageBus: + """Bus prepopulated with a root + two registered children.""" + bus = AgentMessageBus() + await bus.register("root-1", "root", parent_id=None) + await bus.register("child-A", "scanner", parent_id="root-1") + await bus.register("child-B", "exploiter", parent_id="root-1") + return bus + + +def _ctx_for(bus: AgentMessageBus, agent_id: str = "root-1") -> _Ctx: + return _Ctx(context={"bus": bus, "agent_id": agent_id, "parent_id": None}) + + +# --- registration --------------------------------------------------------- + + +def test_all_graph_tools_are_function_tools() -> None: + for tool in ( + view_agent_graph, + agent_status, + send_message_to_agent, + wait_for_message, + create_agent, + agent_finish, + ): + assert isinstance(tool, FunctionTool) + + +# --- view_agent_graph ----------------------------------------------------- + + +@pytest.mark.asyncio +async def test_view_agent_graph_renders_tree() -> None: + bus = await _make_bus_with_agents() + out = await _invoke(view_agent_graph, _ctx_for(bus)) + assert out["success"] is True + assert "root (root-1)" in out["graph_structure"] + assert "scanner (child-A)" in out["graph_structure"] + assert "exploiter (child-B)" in out["graph_structure"] + # The "you" marker should be on the calling agent's line. + you_lines = [line for line in out["graph_structure"].splitlines() if "← you" in line] + assert len(you_lines) == 1 + assert "root-1" in you_lines[0] + assert out["summary"]["total"] == 3 + assert out["summary"]["running"] == 3 + + +@pytest.mark.asyncio +async def test_view_agent_graph_handles_missing_bus() -> None: + out = await _invoke(view_agent_graph, _Ctx(context={})) + assert out["success"] is False + assert "Bus" in out["error"] + + +# --- agent_status --------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_agent_status_returns_state() -> None: + bus = await _make_bus_with_agents() + out = await _invoke(agent_status, _ctx_for(bus), agent_id="child-A") + assert out["success"] is True + assert out["agent_id"] == "child-A" + assert out["name"] == "scanner" + assert out["status"] == "running" + assert out["parent_id"] == "root-1" + assert out["pending_messages"] == 0 + + +@pytest.mark.asyncio +async def test_agent_status_unknown_id() -> None: + bus = await _make_bus_with_agents() + out = await _invoke(agent_status, _ctx_for(bus), agent_id="nope") + assert out["success"] is False + assert "Unknown" in out["error"] + + +# --- send_message_to_agent ----------------------------------------------- + + +@pytest.mark.asyncio +async def test_send_message_queues_into_target_inbox() -> None: + bus = await _make_bus_with_agents() + out = await _invoke( + send_message_to_agent, + _ctx_for(bus, agent_id="child-A"), + target_agent_id="child-B", + message="hello sibling", + priority="high", + ) + assert out["success"] is True + assert out["delivery_status"] == "queued" + # Drain via the same API the filter uses; confirm the message landed. + msgs = await bus.drain("child-B") + assert len(msgs) == 1 + assert msgs[0]["from"] == "child-A" + assert msgs[0]["content"] == "hello sibling" + assert msgs[0]["priority"] == "high" + + +@pytest.mark.asyncio +async def test_send_message_unknown_target() -> None: + bus = await _make_bus_with_agents() + out = await _invoke( + send_message_to_agent, + _ctx_for(bus), + target_agent_id="ghost", + message="hi", + ) + assert out["success"] is False + assert "not found" in out["error"] + + +@pytest.mark.asyncio +async def test_send_message_to_finalized_agent_is_rejected() -> None: + """A finalized target should not silently swallow messages.""" + bus = await _make_bus_with_agents() + await bus.finalize("child-A", "completed") + out = await _invoke( + send_message_to_agent, + _ctx_for(bus), + target_agent_id="child-A", + message="too late", + ) + assert out["success"] is False + # finalize() also clears parent_of/names, so the user-visible state is + # "completed" — confirm the wrapper treats finalized agents as + # undeliverable rather than queuing into a dropped inbox. + assert "completed" in out["error"] or "not found" in out["error"] + + +# --- wait_for_message ---------------------------------------------------- + + +@pytest.mark.asyncio +async def test_wait_for_message_returns_when_message_arrives() -> None: + bus = await _make_bus_with_agents() + + async def deliver_after_short_pause() -> None: + await asyncio.sleep(0.1) + await bus.send("child-A", {"from": "child-B", "content": "ping", "type": "info"}) + + sender = asyncio.create_task(deliver_after_short_pause()) + out = await _invoke( + wait_for_message, + _ctx_for(bus, agent_id="child-A"), + timeout_seconds=3, + ) + await sender + assert out["success"] is True + assert out["status"] == "message_arrived" + assert out["pending_messages"] >= 1 + # Status must be returned to "running" after the wait completes. + assert bus.statuses["child-A"] == "running" + + +@pytest.mark.asyncio +async def test_wait_for_message_times_out() -> None: + bus = await _make_bus_with_agents() + out = await _invoke( + wait_for_message, + _ctx_for(bus, agent_id="child-A"), + timeout_seconds=1, + ) + assert out["success"] is True + assert out["status"] == "timeout" + assert out["timeout_seconds"] == 1 + assert bus.statuses["child-A"] == "running" + + +# --- create_agent -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_agent_requires_factory_in_context() -> None: + bus = await _make_bus_with_agents() + out = await _invoke( + create_agent, + _ctx_for(bus), + name="recon-bot", + task="enumerate hosts", + ) + assert out["success"] is False + assert "agent_factory" in out["error"] + + +@pytest.mark.asyncio +async def test_create_agent_spawns_and_registers_child() -> None: + """Verify the spawn shape without running a real LLM.""" + bus = await _make_bus_with_agents() + + factory_calls: list[dict[str, Any]] = [] + + def fake_factory(*, name: str, skills: list[str]) -> Any: + factory_calls.append({"name": name, "skills": list(skills)}) + # The Runner.run patch below ignores this object; any sentinel + # works. + return object() + + runner_calls: list[dict[str, Any]] = [] + + async def fake_runner_run(*args: Any, **kwargs: Any) -> Any: + # Capture the input + max_turns so the test can assert the + # identity block + delegation envelope are present. + runner_calls.append({"args": args, "kwargs": kwargs}) + await asyncio.sleep(0) # cooperate so create_task can return + return None + + ctx = _Ctx( + context={ + "bus": bus, + "agent_id": "root-1", + "parent_id": None, + "agent_factory": fake_factory, + "sandbox_session": None, + "sandbox_client": None, + "sandbox_token": "token", + "tool_server_host_port": 12345, + "caido_host_port": None, + "tracer": None, + "model": "strix/claude-sonnet-4.6", + "model_settings": None, + "max_turns": 300, + "is_whitebox": False, + } + ) + + with patch( + "strix.tools.agents_graph.agents_graph_sdk_tools.Runner.run", + side_effect=fake_runner_run, + ): + out = await _invoke( + create_agent, + ctx, + name="recon-bot", + task="enumerate hosts", + inherit_context=False, + skills=["recon"], + ) + # The spawned task must be allowed to run so Runner.run side-effect + # records the call. + new_id = out["agent_id"] + await asyncio.gather(*(t for t in bus.tasks.values()), return_exceptions=True) + + assert out["success"] is True + assert factory_calls == [{"name": "recon-bot", "skills": ["recon"]}] + assert len(runner_calls) == 1 + + # Bus state: child registered, task stored. + assert new_id in bus.statuses + assert bus.parent_of[new_id] == "root-1" + assert bus.names[new_id] == "recon-bot" + assert new_id in bus.tasks + + # Initial input shape: identity block + task message at the end. + initial_input = runner_calls[0]["kwargs"]["input"] + assert any( + isinstance(item, dict) and "agent_delegation" in item.get("content", "") + for item in initial_input + ) + assert initial_input[-1]["content"] == "enumerate hosts" + + +@pytest.mark.asyncio +async def test_create_agent_inherits_parent_history() -> None: + bus = await _make_bus_with_agents() + + def fake_factory(*, name: str, skills: list[str]) -> Any: + return object() + + runner_calls: list[dict[str, Any]] = [] + + async def fake_runner_run(*args: Any, **kwargs: Any) -> Any: + runner_calls.append(kwargs) + return None + + parent_history = [ + {"role": "user", "content": "scope: example.com"}, + {"role": "assistant", "content": "I'll start with subdomain enum."}, + ] + ctx = _Ctx( + context={ + "bus": bus, + "agent_id": "root-1", + "parent_id": None, + "agent_factory": fake_factory, + "parent_input_items": parent_history, + "sandbox_session": None, + "sandbox_client": None, + "sandbox_token": "token", + "tool_server_host_port": 12345, + "caido_host_port": None, + "tracer": None, + "model": "strix/claude-sonnet-4.6", + "model_settings": None, + "max_turns": 300, + } + ) + + with patch( + "strix.tools.agents_graph.agents_graph_sdk_tools.Runner.run", + side_effect=fake_runner_run, + ): + await _invoke( + create_agent, + ctx, + name="child", + task="do thing", + inherit_context=True, + ) + await asyncio.gather(*(t for t in bus.tasks.values()), return_exceptions=True) + + initial_input = runner_calls[0]["input"] + contents = [item.get("content", "") for item in initial_input] + assert "" in contents + assert "" in contents + # Parent's exact items should be in between. + assert any(c == "scope: example.com" for c in contents) + + +# --- agent_finish -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_agent_finish_rejects_root() -> None: + bus = await _make_bus_with_agents() + out = await _invoke( + agent_finish, + _ctx_for(bus, agent_id="root-1"), + result_summary="done", + ) + assert out["success"] is False + assert "agent_finish is for subagents" in out["error"] + + +@pytest.mark.asyncio +async def test_agent_finish_posts_report_to_parent_inbox() -> None: + bus = await _make_bus_with_agents() + ctx = _Ctx( + context={ + "bus": bus, + "agent_id": "child-A", + "parent_id": "root-1", + "agent_finish_called": False, + } + ) + out = await _invoke( + agent_finish, + ctx, + result_summary="found 3 issues", + findings=["xss in /search", "open redirect", "stored xss"], + final_recommendations=["sanitize search input"], + success=True, + ) + assert out["success"] is True + assert out["agent_completed"] is True + assert out["parent_notified"] is True + + # Side effects: agent_finish_called flipped (so on_agent_end records + # "completed", not "crashed"), and the parent's inbox got the report. + assert ctx.context["agent_finish_called"] is True + parent_msgs = await bus.drain("root-1") + assert len(parent_msgs) == 1 + msg = parent_msgs[0] + assert msg["type"] == "completion" + assert msg["from"] == "child-A" + assert "found 3 issues" in msg["content"] + assert "xss in /search" in msg["content"] + assert "sanitize search input" in msg["content"] + + +@pytest.mark.asyncio +async def test_agent_finish_skips_parent_when_report_to_parent_false() -> None: + bus = await _make_bus_with_agents() + ctx = _Ctx( + context={ + "bus": bus, + "agent_id": "child-A", + "parent_id": "root-1", + "agent_finish_called": False, + } + ) + out = await _invoke( + agent_finish, + ctx, + result_summary="silent done", + report_to_parent=False, + ) + assert out["success"] is True + assert out["parent_notified"] is False + assert ctx.context["agent_finish_called"] is True + parent_msgs = await bus.drain("root-1") + assert parent_msgs == [] + + +# --- bus integration sanity --------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_agent_spawn_is_cancelable_via_bus() -> None: + """Verify bus.cancel_descendants reaches a child task we just spawned.""" + bus = await _make_bus_with_agents() + + def fake_factory(*, name: str, skills: list[str]) -> Any: + return object() + + # Long-lived child that yields control so we can cancel it before it + # finishes naturally. + async def slow_runner_run(*args: Any, **kwargs: Any) -> Any: + await asyncio.sleep(60) + return None + + ctx = _Ctx( + context={ + "bus": bus, + "agent_id": "root-1", + "parent_id": None, + "agent_factory": fake_factory, + "sandbox_session": None, + "sandbox_client": None, + "sandbox_token": "t", + "tool_server_host_port": 12345, + "caido_host_port": None, + "tracer": None, + "model": "strix/claude-sonnet-4.6", + "model_settings": None, + "max_turns": 300, + } + ) + + runner_mock = AsyncMock(side_effect=slow_runner_run) + with patch( + "strix.tools.agents_graph.agents_graph_sdk_tools.Runner.run", + new=runner_mock, + ): + out = await _invoke( + create_agent, + ctx, + name="long-running", + task="do thing", + inherit_context=False, + ) + child_id = out["agent_id"] + # Let the task actually start. + await asyncio.sleep(0.05) + await bus.cancel_descendants("root-1") + + # The cancel should have propagated; the task is done (cancelled). + assert bus.tasks[child_id].done()