0xallam
1ac32df817
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 <agent_completion_report> 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) <noreply@anthropic.com>
513 lines
16 KiB
Python
513 lines
16 KiB
Python
"""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 "<inherited_context_from_parent>" in contents
|
|
assert "</inherited_context_from_parent>" 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 "<finding>xss in /search</finding>" 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()
|