"""Multi-agent graph tools backed by AgentCoordinator.""" from __future__ import annotations import asyncio import json import logging import uuid from collections import Counter from datetime import UTC, datetime from typing import Any, Literal, get_args from agents import RunContextWrapper, function_tool from strix.core.agents import Status, coordinator_from_context from strix.skills import validate_requested_skills _ACTIVE_STATUSES: frozenset[str] = frozenset({"running", "waiting"}) logger = logging.getLogger(__name__) 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 a child's completion report as plain structured text. Goes into the parent's SDK session with coordinator-added sender metadata, so this body just carries the contents. No XML — no escaping concerns, no parser ambiguity. """ status = "SUCCESS" if success else "FAILED" completion_time = datetime.now(UTC).isoformat() lines: list[str] = [ f"== Completion report from {agent_name} ({agent_id}) ==", f"Status: {status}", f"Time: {completion_time}", ] if task: lines.append(f"Task: {task}") lines.append("") lines.append("Summary:") lines.append(result_summary or "(none)") if findings: lines.append("") lines.append("Findings:") lines.extend(f"- {f}" for f in findings) if recommendations: lines.append("") lines.append("Recommendations:") lines.extend(f"- {r}" for r in recommendations) return "\n".join(lines) @function_tool(timeout=30) async def view_agent_graph(ctx: RunContextWrapper) -> str: """Print the multi-agent tree — every agent, its parent, its status. Use before spawning a new agent (don't duplicate work — check whether something specialized for that task already exists) and any time you want a snapshot of who's still ``running`` / ``waiting`` / ``completed`` / ``crashed`` / ``stopped``. Output is an indented bullet list with status in brackets; the agent that called this tool is marked ``← you``. """ inner = _ctx(ctx) coordinator = coordinator_from_context(inner) me = inner.get("agent_id") if coordinator is None: return json.dumps( {"success": False, "error": "Agent coordinator not initialized in context"}, ensure_ascii=False, default=str, ) parent_of, statuses, names = await coordinator.graph_snapshot() 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) counts = Counter(statuses.values()) summary: dict[str, int] = {"total": len(parent_of)} for status_name in get_args(Status): summary[status_name] = counts.get(status_name, 0) return json.dumps( { "success": True, "graph_structure": "\n".join(lines) or "(no agents)", "summary": summary, }, ensure_ascii=False, default=str, ) @function_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: """Send a message to another agent's inbox — sparingly. Inter-agent messages are appended to the target's SDK session and interrupt any active target turn so the next run cycle sees them. Use only when essential: - Sharing a discovered finding/credential another agent needs. - Asking a specialist a focused question. - Coordinating who covers what (avoid overlap). - Telling a child to wrap up or change course. **Don't** use for routine "hello/status" pings, for context the target already has (children inherit parent history), or when parent/child completion via ``agent_finish`` already covers the flow. Messages to any registered agent wake it, regardless of status, so a follow-up can restart a completed/stopped/failed agent. Args: target_agent_id: Recipient's 8-char id. message: The full message body. Be specific — include payloads, URLs, or what you want them to do, not just headlines. message_type: ``query`` (you want a reply), ``instruction`` (you're directing them), ``information`` (FYI, no reply expected). Default ``information``. priority: ``low`` / ``normal`` / ``high`` / ``urgent``. """ inner = _ctx(ctx) coordinator = coordinator_from_context(inner) me = inner.get("agent_id") if coordinator is None or me is None: return json.dumps( {"success": False, "error": "Agent coordinator or agent_id missing in context"}, ensure_ascii=False, default=str, ) if target_agent_id == me: return json.dumps( { "success": False, "error": ( "Cannot send a message to yourself; use `think` to record a " "private note, or `agent_finish` / `finish_scan` to terminate" ), }, ensure_ascii=False, default=str, ) msg_id = f"msg_{uuid.uuid4().hex[:8]}" delivered = await coordinator.send( target_agent_id, { "id": msg_id, "from": me, "content": message, "type": message_type, "priority": priority, }, ) if not delivered: return json.dumps( { "success": False, "error": f"Target agent '{target_agent_id}' not found or message delivery failed", }, ensure_ascii=False, default=str, ) return json.dumps( { "success": True, "message_id": msg_id, "target_agent_id": target_agent_id, "delivery_status": "delivered", }, ensure_ascii=False, default=str, ) def _session_items_payload(items: list[Any]) -> list[dict[str, Any]]: payload: list[dict[str, Any]] = [] for item in items: if isinstance(item, dict): role = item.get("role") content = item.get("content") payload.append({"role": role, "content": content}) else: payload.append({"content": str(item)}) return payload @function_tool(timeout=601) async def wait_for_message( # noqa: PLR0911 ctx: RunContextWrapper, reason: str = "Waiting for messages from other agents", timeout_seconds: int = 600, ) -> str: """Pause this agent until a message lands in its inbox (or timeout). Use when you have nothing useful to do until a child/peer responds — typically after spawning subagents and you want to wait for their completion reports. The agent automatically resumes when any message arrives. **Critical caveats:** - **Never** call this if you finished your own task and have **no** child agents running — that's a permanent stall. Call ``finish_scan`` (root) or ``agent_finish`` (subagent) instead. - If you're waiting on an agent that **isn't your child**, message it first asking it to ping you when done — otherwise it has no reason to send to your inbox and you'll wait the full timeout. - Children update the parent automatically via ``agent_finish`` → no extra coordination needed. Args: reason: One-line note shown in graph snapshots while you're waiting (helps a human or sibling agent debug who's stuck on what). timeout_seconds: Hard cap (default 600s). On timeout the tool returns and you decide whether to keep working or wait again. """ inner = _ctx(ctx) coordinator = coordinator_from_context(inner) me = inner.get("agent_id") interactive = bool(inner.get("interactive", False)) if coordinator is None or me is None: return json.dumps( {"success": False, "error": "Agent coordinator or agent_id missing in context"}, ensure_ascii=False, default=str, ) async with coordinator._lock: stopped = coordinator.statuses.get(me) == "stopped" if stopped: return json.dumps( { "success": True, "wait_outcome": "stopped", "reason": reason, "note": "Wait ended because this agent is stopped.", }, ensure_ascii=False, default=str, ) pending, items = await coordinator.consume_pending(me, include_items=True) if pending > 0: await coordinator.mark_running(me) return json.dumps( { "success": True, "wait_outcome": "message_arrived", "pending_messages": pending, "messages": _session_items_payload(items), "reason": reason, }, ensure_ascii=False, default=str, ) if interactive: await coordinator.park_waiting(me) return json.dumps( { "success": True, "wait_outcome": "waiting", "reason": reason, "note": "Agent parked; execution will resume when a message arrives.", }, ensure_ascii=False, default=str, ) await coordinator.park_waiting(me) try: await asyncio.wait_for(coordinator.wait_for_message(me), timeout_seconds) except TimeoutError: await coordinator.mark_running(me) return json.dumps( { "success": True, "wait_outcome": "timeout", "timeout_seconds": timeout_seconds, "reason": reason, "note": "No messages within timeout — continue work or call agent_finish.", }, ensure_ascii=False, default=str, ) async with coordinator._lock: stopped = coordinator.statuses.get(me) == "stopped" if stopped: return json.dumps( { "success": True, "wait_outcome": "stopped", "reason": reason, "note": "Wait ended because this agent is stopped.", }, ensure_ascii=False, default=str, ) pending, items = await coordinator.consume_pending(me, include_items=True) await coordinator.mark_running(me) return json.dumps( { "success": True, "wait_outcome": "message_arrived", "pending_messages": pending, "messages": _session_items_payload(items), "reason": reason, }, ensure_ascii=False, default=str, ) @function_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 specialist child agent to run in parallel. Decompose complex pentests by handing focused subtasks to dedicated children. The child runs asynchronously — the parent continues immediately and can ``wait_for_message`` later (or just keep working in parallel). When the child calls ``agent_finish``, its completion report lands in the parent's inbox. **Before spawning, call ``view_agent_graph``** to confirm no existing agent already covers this scope — duplicate specialists waste turns and create coordination headaches. **Specialization principles:** - Most agents need at least one ``skill`` to be useful. - Aim for **1-3 related skills** per agent. Up to 5 only when the task genuinely spans them. - One skill = most focused (e.g., XSS-only). Five skills = upper bound. - Match the ``name`` to the focus (``XSS Specialist``, ``SQLi Validator``, ``Auth Specialist``). **When to spawn vs do it yourself:** - Spawn when the subtask is large, parallelizable, or needs different specialization than what you're already doing. - Don't spawn for trivial one-shot probes — just run the tool yourself. Args: name: Human-readable child name (used in graph views and ``send_message_to_agent`` flows). task: Specific objective. Be concrete — what to test, what success looks like, any constraints. inherit_context: Default ``True``. The child receives the parent's input history as background; only set ``False`` when starting a clean-slate task. skills: List of skill names (e.g. ``["xss", "sql_injection"]``). Max 5; prefer 1-3. """ inner = _ctx(ctx) coordinator = coordinator_from_context(inner) parent_id = inner.get("agent_id") spawner = inner.get("spawn_child_agent") if coordinator is None or parent_id is None: return json.dumps( {"success": False, "error": "Agent coordinator or agent_id missing in context"}, ensure_ascii=False, default=str, ) if not callable(spawner): return json.dumps( { "success": False, "error": "Scan runner did not provide a child-agent spawner in context", }, ensure_ascii=False, default=str, ) skill_list = list(skills or []) skill_error = validate_requested_skills(skill_list) if skill_error: return json.dumps( {"success": False, "error": skill_error, "agent_id": None}, ensure_ascii=False, default=str, ) parent_history = list(ctx.turn_input) if inherit_context and ctx.turn_input else [] try: result = await spawner( parent_ctx=inner, name=name, task=task, skills=skill_list, parent_history=parent_history, ) except Exception as e: logger.exception("create_agent: scan runner failed to spawn child '%s'", name) return json.dumps( {"success": False, "error": f"child spawn failed: {e!s}"}, ensure_ascii=False, default=str, ) logger.info( "create_agent: spawned %s (%s) parent=%s skills=%d task_len=%d", result.get("agent_id"), name, parent_id or "-", len(skill_list), len(task or ""), ) return json.dumps( result, ensure_ascii=False, default=str, ) @function_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 termination — post a completion report to the parent. **Subagents only.** Root agents must call ``finish_scan`` instead; this tool refuses to run for root agents. Calling this: 1. Marks the subagent as ``completed``. 2. Posts a structured completion report to the parent's inbox (when ``report_to_parent`` is true). 3. Stops this subagent's execution. **Vulnerability findings must already be filed via ``create_vulnerability_report`` before calling this.** The ``findings`` field here is for narrative summary only — it does not register vulns in the scan report. Write the summary as if the parent has no idea what you were doing: what did you test, what did you find/confirm/rule out, what's still open. Args: result_summary: What you accomplished and discovered. Concrete and specific (URLs, parameters, payloads that worked). findings: Optional bullet list of confirmed observations. For credit-bearing vulnerabilities, file ``create_vulnerability_report`` first; this is for narrative. success: Whether the assigned subtask was completed successfully. Default ``True``. report_to_parent: Whether to deliver the completion report to the parent's inbox. Default ``True``. final_recommendations: Optional next-step suggestions for the parent (e.g., "prioritize testing X", "spawn an agent to cover Y"). """ inner = _ctx(ctx) coordinator = coordinator_from_context(inner) me = inner.get("agent_id") if coordinator is None or me is None: return json.dumps( {"success": False, "error": "Agent coordinator or agent_id missing in context"}, ensure_ascii=False, default=str, ) parent_id = inner.get("parent_id") if parent_id is None: return json.dumps( { "success": False, "error": ( "agent_finish is for subagents. Root/main agents must call finish_scan instead" ), }, ensure_ascii=False, default=str, ) parent_notified = False if report_to_parent: async with coordinator._lock: agent_name = coordinator.names.get(me, me) 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 coordinator.send( parent_id, { "id": f"report_{uuid.uuid4().hex[:8]}", "from": me, "content": report, "type": "completion", "priority": "high", }, ) parent_notified = True logger.info( "agent_finish: %s success=%s findings=%d parent_notified=%s", me, success, len(findings or []), parent_notified, ) await coordinator.set_status(me, "completed") return json.dumps( { "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), }, 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 parks as ``stopped``; later user/peer messages can wake it again. 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) coordinator = coordinator_from_context(inner) me = inner.get("agent_id") if coordinator is None or me is None: return json.dumps( {"success": False, "error": "Agent coordinator 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, ) _, statuses, _ = await coordinator.graph_snapshot() if target_agent_id not in statuses: return json.dumps( {"success": False, "error": f"Unknown agent_id: {target_agent_id}"}, ensure_ascii=False, default=str, ) current_status = statuses[target_agent_id] if current_status not in _ACTIVE_STATUSES: return json.dumps( { "success": False, "error": ( f"Agent {target_agent_id} is already '{current_status}'; " "stop_agent only acts on running/waiting agents — use " "view_agent_graph to find still-active descendants and " "stop them individually, or send_message_to_agent if you " "want to wake this one with new instructions" ), "target_agent_id": target_agent_id, "current_status": current_status, }, ensure_ascii=False, default=str, ) if cascade: await coordinator.cancel_descendants_graceful(target_agent_id) else: await coordinator.request_stop(target_agent_id) 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, )