Simplify SDK-native orchestration

This commit is contained in:
0xallam
2026-04-26 11:30:00 -07:00
parent dc03f1f4ed
commit bd40884fcf
21 changed files with 1435 additions and 2133 deletions
+5 -14
View File
@@ -194,11 +194,6 @@ ignore = [
"strix/tools/**/*.py" = [ "strix/tools/**/*.py" = [
"ARG001", # Unused function argument (tools may have unused args for interface consistency) "ARG001", # Unused function argument (tools may have unused args for interface consistency)
] ]
# SDK lifecycle hooks have a fixed signature — unused args are part of the contract.
"strix/orchestration/hooks.py" = [
"ARG002", # Unused method argument (RunHooks signature is set by SDK)
"TC002", # ModelResponse, RunContextWrapper, AgentHookContext are annotation-only
]
# Anthropic cache wrapper inherits from LitellmModel; signature shadows builtin `input`. # Anthropic cache wrapper inherits from LitellmModel; signature shadows builtin `input`.
"strix/llm/anthropic_cache_wrapper.py" = [ "strix/llm/anthropic_cache_wrapper.py" = [
"A002", # Argument shadows builtin (parent signature uses `input`) "A002", # Argument shadows builtin (parent signature uses `input`)
@@ -209,9 +204,6 @@ ignore = [
# Backend factories import their backend's deps lazily so deployments # Backend factories import their backend's deps lazily so deployments
# that pick a different backend don't need every backend's libs installed. # that pick a different backend don't need every backend's libs installed.
"strix/runtime/backends.py" = ["PLC0415"] "strix/runtime/backends.py" = ["PLC0415"]
# The vulnerability MD renderer is a long monolithic format function;
# splitting per-section would obscure the structure without simplifying.
"strix/telemetry/scan_artifacts.py" = ["PLR0912", "PLR0915", "PERF401"]
"strix/runtime/docker_client.py" = [ "strix/runtime/docker_client.py" = [
"TC002", # Manifest, Container imported for annotations "TC002", # Manifest, Container imported for annotations
"TC003", # uuid imported for annotation "TC003", # uuid imported for annotation
@@ -234,13 +226,13 @@ ignore = [
"strix/agents/factory.py" = ["TC002"] "strix/agents/factory.py" = ["TC002"]
# Entry point: ``Path`` is used at runtime by the typing of the # Entry point: ``Path`` is used at runtime by the typing of the
# session_manager call; importing under TYPE_CHECKING would defer # session_manager call; importing under TYPE_CHECKING would defer
# resolution past where mypy needs it. ``_build_root_task`` legitimately # resolution past where mypy needs it. ``build_root_task`` legitimately
# walks every supported target type — splitting it into per-type # walks every supported target type — splitting it into per-type
# helpers would add indirection without simplifying anything. # helpers would add indirection without simplifying anything.
"strix/orchestration/scan.py" = ["TC003", "PLR0912", "PLR0915", "PLC0415"] "strix/orchestration/runner.py" = ["TC003", "PLR0912", "PLR0915", "PLC0415"]
# Tracer carries a long event surface and a runtime ``Callable`` # ScanStore carries scan artifact/report fields, artifact rendering, and
# annotation on ``vulnerability_found_callback``. # a runtime ``Callable`` annotation on ``vulnerability_found_callback``.
"strix/telemetry/tracer.py" = ["TC003", "PLR0912", "PLR0915", "E501"] "strix/telemetry/scan_store.py" = ["TC003", "PLR0912", "PLR0915", "E501", "PERF401"]
# Interface utility branches per scope-mode / target-type combination; # Interface utility branches per scope-mode / target-type combination;
# splitting would obscure the decision tree without simplifying it. # splitting would obscure the decision tree without simplifying it.
"strix/interface/utils.py" = ["PLR0912", "BLE001", "PLC0415"] "strix/interface/utils.py" = ["PLR0912", "BLE001", "PLC0415"]
@@ -249,7 +241,6 @@ ignore = [
"strix/interface/cli.py" = ["BLE001", "PLC0415"] "strix/interface/cli.py" = ["BLE001", "PLC0415"]
"strix/interface/tui.py" = ["BLE001", "PLC0415", "PLR0912", "PLR0915", "SIM105"] "strix/interface/tui.py" = ["BLE001", "PLC0415", "PLR0912", "PLR0915", "SIM105"]
"strix/interface/main.py" = ["BLE001", "PLC0415", "PLR0912", "PLR0915"] "strix/interface/main.py" = ["BLE001", "PLC0415", "PLR0912", "PLR0915"]
"strix/interface/streaming_parser.py" = ["PLC0415"]
"strix/interface/tool_components/agent_message_renderer.py" = ["PLC0415"] "strix/interface/tool_components/agent_message_renderer.py" = ["PLC0415"]
[tool.ruff.lint.isort] [tool.ruff.lint.isort]
+27 -13
View File
@@ -119,23 +119,37 @@ hiddenimports = [
'strix.interface.utils', 'strix.interface.utils',
'strix.interface.tool_components', 'strix.interface.tool_components',
'strix.agents', 'strix.agents',
'strix.agents.base_agent', 'strix.agents.factory',
'strix.agents.state', 'strix.agents.prompt',
'strix.agents.StrixAgent',
'strix.llm', 'strix.llm',
'strix.llm.llm', 'strix.llm.anthropic_cache_wrapper',
'strix.llm.config', 'strix.llm.dedupe',
'strix.llm.utils', 'strix.llm.multi_provider_setup',
'strix.llm.memory_compressor', 'strix.llm.retry',
'strix.orchestration',
'strix.orchestration.coordinator',
'strix.orchestration.runner',
'strix.orchestration.utils',
'strix.runtime', 'strix.runtime',
'strix.runtime.runtime', 'strix.runtime.backends',
'strix.runtime.docker_runtime', 'strix.runtime.caido_bootstrap',
'strix.runtime.docker_client',
'strix.runtime.session_manager',
'strix.telemetry', 'strix.telemetry',
'strix.telemetry.tracer', 'strix.telemetry.logging',
'strix.telemetry.posthog',
'strix.telemetry.scan_store',
'strix.tools', 'strix.tools',
'strix.tools.registry', 'strix.tools.agents_graph.tools',
'strix.tools.executor', 'strix.tools.finish.tool',
'strix.tools.argument_parser', 'strix.tools.notes.tools',
'strix.tools.proxy._calls',
'strix.tools.proxy.tools',
'strix.tools.python.tool',
'strix.tools.reporting.tool',
'strix.tools.thinking.tool',
'strix.tools.todo.tools',
'strix.tools.web_search.tool',
'strix.skills', 'strix.skills',
] ]
-2
View File
@@ -29,7 +29,6 @@ from agents.tool import Tool
from strix.agents.prompt import render_system_prompt from strix.agents.prompt import render_system_prompt
from strix.tools.agents_graph.tools import ( from strix.tools.agents_graph.tools import (
agent_finish, agent_finish,
agent_status,
create_agent, create_agent,
send_message_to_agent, send_message_to_agent,
stop_agent, stop_agent,
@@ -163,7 +162,6 @@ _BASE_TOOLS: tuple[Tool, ...] = (
python_action, python_action,
# Multi-agent graph tools (the coordinator is in ctx.context) # Multi-agent graph tools (the coordinator is in ctx.context)
view_agent_graph, view_agent_graph,
agent_status,
send_message_to_agent, send_message_to_agent,
wait_for_message, wait_for_message,
create_agent, create_agent,
+12 -12
View File
@@ -13,9 +13,9 @@ from rich.panel import Panel
from rich.text import Text from rich.text import Text
from strix.config import load_settings from strix.config import load_settings
from strix.orchestration.scan import run_strix_scan from strix.orchestration.runner import run_strix_scan
from strix.runtime import session_manager from strix.runtime import session_manager
from strix.telemetry.tracer import Tracer, set_global_tracer from strix.telemetry.scan_store import ScanStore, set_global_scan_store
from .utils import ( from .utils import (
build_live_stats_text, build_live_stats_text,
@@ -95,8 +95,9 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
"resume_instruction": getattr(args, "user_explicit_instruction", None) or "", "resume_instruction": getattr(args, "user_explicit_instruction", None) or "",
} }
tracer = Tracer(args.run_name) scan_store = ScanStore(args.run_name)
tracer.set_scan_config(scan_config) scan_store.set_scan_config(scan_config)
scan_store.hydrate_from_run_dir()
def display_vulnerability(report: dict[str, Any]) -> None: def display_vulnerability(report: dict[str, Any]) -> None:
report_id = report.get("id", "unknown") report_id = report.get("id", "unknown")
@@ -114,13 +115,13 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
console.print(vuln_panel) console.print(vuln_panel)
console.print() console.print()
tracer.vulnerability_found_callback = display_vulnerability scan_store.vulnerability_found_callback = display_vulnerability
def cleanup_on_exit() -> None: def cleanup_on_exit() -> None:
tracer.cleanup() scan_store.cleanup()
def signal_handler(_signum: int, _frame: Any) -> None: def signal_handler(_signum: int, _frame: Any) -> None:
tracer.cleanup() scan_store.cleanup()
sys.exit(1) sys.exit(1)
atexit.register(cleanup_on_exit) atexit.register(cleanup_on_exit)
@@ -129,14 +130,14 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
if hasattr(signal, "SIGHUP"): if hasattr(signal, "SIGHUP"):
signal.signal(signal.SIGHUP, signal_handler) signal.signal(signal.SIGHUP, signal_handler)
set_global_tracer(tracer) set_global_scan_store(scan_store)
def create_live_status() -> Panel: def create_live_status() -> Panel:
status_text = Text() status_text = Text()
status_text.append("Penetration test in progress", style="bold #22c55e") status_text.append("Penetration test in progress", style="bold #22c55e")
status_text.append("\n\n") status_text.append("\n\n")
stats_text = build_live_stats_text(tracer) stats_text = build_live_stats_text(scan_store)
if stats_text: if stats_text:
status_text.append(stats_text) status_text.append(stats_text)
@@ -179,7 +180,6 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
scan_id=args.run_name, scan_id=args.run_name,
image=_resolve_sandbox_image(), image=_resolve_sandbox_image(),
local_sources=getattr(args, "local_sources", None) or [], local_sources=getattr(args, "local_sources", None) or [],
tracer=tracer,
interactive=bool(getattr(args, "interactive", False)), interactive=bool(getattr(args, "interactive", False)),
) )
finally: finally:
@@ -196,7 +196,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
console.print(f"[bold red]Error during penetration test:[/] {e}") console.print(f"[bold red]Error during penetration test:[/] {e}")
raise raise
if tracer.final_scan_result: if scan_store.final_scan_result:
console.print() console.print()
final_report_text = Text() final_report_text = Text()
@@ -206,7 +206,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
Text.assemble( Text.assemble(
final_report_text, final_report_text,
"\n\n", "\n\n",
tracer.final_scan_result, scan_store.final_scan_result,
), ),
title="[bold white]STRIX", title="[bold white]STRIX",
title_align="left", title_align="left",
+11 -11
View File
@@ -37,7 +37,7 @@ from strix.interface.utils import (
validate_llm_response, validate_llm_response,
) )
from strix.telemetry import posthog from strix.telemetry import posthog
from strix.telemetry.tracer import get_global_tracer from strix.telemetry.scan_store import get_global_scan_store
HOST_GATEWAY_HOSTNAME = "host.docker.internal" HOST_GATEWAY_HOSTNAME = "host.docker.internal"
@@ -47,7 +47,7 @@ import logging # noqa: E402
# Per-scan logging is set up by ``setup_scan_logging`` from inside # Per-scan logging is set up by ``setup_scan_logging`` from inside
# ``orchestration.scan.run_strix_scan`` once the scan ``run_dir`` is # ``orchestration.runner.run_strix_scan`` once the scan ``run_dir`` is
# known — that's where ``strix.*`` levels and handlers are owned. Pre-scan # known — that's where ``strix.*`` levels and handlers are owned. Pre-scan
# work (``main()``, env validation, image pull) emits via the module # work (``main()``, env validation, image pull) emits via the module
# logger; once setup_scan_logging runs, those records start landing in # logger; once setup_scan_logging runs, those records start landing in
@@ -552,11 +552,11 @@ def _load_resume_state(args: argparse.Namespace, parser: argparse.ArgumentParser
def display_completion_message(args: argparse.Namespace, results_path: Path) -> None: def display_completion_message(args: argparse.Namespace, results_path: Path) -> None:
console = Console() console = Console()
tracer = get_global_tracer() scan_store = get_global_scan_store()
scan_completed = False scan_completed = False
if tracer and tracer.scan_results: if scan_store and scan_store.scan_results:
scan_completed = tracer.scan_results.get("scan_completed", False) scan_completed = scan_store.scan_results.get("scan_completed", False)
completion_text = Text() completion_text = Text()
if scan_completed: if scan_completed:
@@ -575,7 +575,7 @@ def display_completion_message(args: argparse.Namespace, results_path: Path) ->
target_text.append("\n ") target_text.append("\n ")
target_text.append(target_info["original"], style="white") target_text.append(target_info["original"], style="white")
stats_text = build_final_stats_text(tracer) stats_text = build_final_stats_text(scan_store)
panel_parts: list[Text | str] = [completion_text, "\n\n", target_text] panel_parts: list[Text | str] = [completion_text, "\n\n", target_text]
@@ -753,16 +753,16 @@ def main() -> None:
posthog.error("unhandled_exception", str(e)) posthog.error("unhandled_exception", str(e))
raise raise
finally: finally:
tracer = get_global_tracer() scan_store = get_global_scan_store()
if tracer: if scan_store:
posthog.end(tracer, exit_reason=exit_reason) posthog.end(scan_store, exit_reason=exit_reason)
results_path = Path("strix_runs") / args.run_name results_path = Path("strix_runs") / args.run_name
display_completion_message(args, results_path) display_completion_message(args, results_path)
if args.non_interactive: if args.non_interactive:
tracer = get_global_tracer() scan_store = get_global_scan_store()
if tracer and tracer.vulnerability_reports: if scan_store and scan_store.vulnerability_reports:
sys.exit(2) sys.exit(2)
+217 -177
View File
@@ -2,13 +2,16 @@ import argparse
import asyncio import asyncio
import atexit import atexit
import contextlib import contextlib
import json
import logging import logging
import signal import signal
import sys import sys
import threading import threading
from collections.abc import Callable from collections.abc import Callable
from datetime import UTC, datetime
from importlib.metadata import PackageNotFoundError from importlib.metadata import PackageNotFoundError
from importlib.metadata import version as pkg_version from importlib.metadata import version as pkg_version
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar from typing import TYPE_CHECKING, Any, ClassVar
@@ -31,12 +34,11 @@ from textual.widgets.tree import TreeNode
from strix.config import load_settings from strix.config import load_settings
from strix.interface.tool_components.agent_message_renderer import AgentMessageRenderer from strix.interface.tool_components.agent_message_renderer import AgentMessageRenderer
from strix.interface.tool_components.registry import get_tool_renderer
from strix.interface.tool_components.user_message_renderer import UserMessageRenderer from strix.interface.tool_components.user_message_renderer import UserMessageRenderer
from strix.interface.utils import build_tui_stats_text from strix.interface.utils import build_tui_stats_text
from strix.orchestration.scan import run_strix_scan from strix.orchestration.runner import run_strix_scan
from strix.runtime import session_manager from strix.runtime import session_manager
from strix.telemetry.tracer import Tracer, set_global_tracer from strix.telemetry.scan_store import ScanStore, set_global_scan_store
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -639,6 +641,139 @@ class VulnerabilitiesPanel(VerticalScroll): # type: ignore[misc]
self.mount(item) self.mount(item)
class TuiLiveView:
"""UI-owned projection of coordinator state and SDK session items."""
def __init__(self) -> None:
self.agents: dict[str, dict[str, Any]] = {}
self.chat_messages: list[dict[str, Any]] = []
self._next_message_id = 1
self._session_message_keys: dict[tuple[str, int], dict[str, Any]] = {}
def hydrate_from_run_dir(self, run_dir: Path) -> None:
agents_path = run_dir / "agents.json"
if not agents_path.exists():
return
try:
agents_data = json.loads(agents_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError):
return
statuses = agents_data.get("statuses") or {}
names = agents_data.get("names") or {}
parent_of = agents_data.get("parent_of") or {}
if not isinstance(statuses, dict):
return
for agent_id, status in statuses.items():
if not isinstance(agent_id, str):
continue
self.upsert_agent(
agent_id,
name=names.get(agent_id, agent_id) if isinstance(names, dict) else agent_id,
parent_id=parent_of.get(agent_id) if isinstance(parent_of, dict) else None,
status=str(status),
)
def upsert_agent(
self,
agent_id: str,
*,
name: str | None = None,
parent_id: str | None = None,
status: str | None = None,
error_message: str | None = None,
) -> None:
now = datetime.now(UTC).isoformat()
current = self.agents.setdefault(
agent_id,
{
"id": agent_id,
"name": name or agent_id,
"parent_id": parent_id,
"status": status or "running",
"created_at": now,
"updated_at": now,
},
)
if name is not None:
current["name"] = name
if parent_id is not None or "parent_id" not in current:
current["parent_id"] = parent_id
if status is not None:
current["status"] = status
if error_message:
current["error_message"] = error_message
current["updated_at"] = now
def sync_agent_messages_from_items(self, agent_id: str, items: list[Any]) -> None:
other_agents = [m for m in self.chat_messages if m.get("agent_id") != agent_id]
refreshed: list[dict[str, Any]] = []
for index, item in enumerate(items):
converted = _sdk_item_to_chat_message(item)
if converted is None:
continue
key = (agent_id, index)
existing = self._session_message_keys.get(key)
if existing is None:
existing = {
"message_id": self._next_message_id,
"timestamp": datetime.now(UTC).isoformat(),
}
self._next_message_id += 1
self._session_message_keys[key] = existing
refreshed.append(
{
"message_id": existing["message_id"],
"content": converted["content"],
"role": converted["role"],
"agent_id": agent_id,
"timestamp": existing["timestamp"],
"metadata": {"source": "sdk_session", "session_index": index},
}
)
self.chat_messages = other_agents + refreshed
def _sdk_item_to_chat_message(item: Any) -> dict[str, str] | None:
if not isinstance(item, dict):
if hasattr(item, "model_dump"):
item = item.model_dump(exclude_unset=True)
else:
return None
role = item.get("role")
if role not in {"user", "assistant"}:
return None
content = _extract_sdk_text(item.get("content"))
if not content:
return None
return {"role": str(role), "content": content}
def _extract_sdk_text(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for part in content:
if isinstance(part, str):
parts.append(part)
elif isinstance(part, dict):
text = part.get("text") or part.get("content")
if isinstance(text, str):
parts.append(text)
else:
text = getattr(part, "text", None) or getattr(part, "content", None)
if isinstance(text, str):
parts.append(text)
return "\n".join(p for p in parts if p)
return ""
class QuitScreen(ModalScreen): # type: ignore[misc] class QuitScreen(ModalScreen): # type: ignore[misc]
def compose(self) -> ComposeResult: def compose(self) -> ComposeResult:
yield Grid( yield Grid(
@@ -704,9 +839,13 @@ class StrixTUIApp(App): # type: ignore[misc]
self.args = args self.args = args
self.scan_config = self._build_scan_config(args) self.scan_config = self._build_scan_config(args)
self.tracer = Tracer(self.scan_config["run_name"]) self.scan_store = ScanStore(self.scan_config["run_name"])
self.tracer.set_scan_config(self.scan_config) self.scan_store.set_scan_config(self.scan_config)
set_global_tracer(self.tracer) self.scan_store.hydrate_from_run_dir()
set_global_scan_store(self.scan_store)
self.live_view = TuiLiveView()
self.live_view.hydrate_from_run_dir(self.scan_store.get_run_dir())
self._live_sync_future: Any | None = None
# Pre-create the coordinator here so the TUI can route stop/chat # Pre-create the coordinator here so the TUI can route stop/chat
# commands while the scan loop runs in a worker thread. # commands while the scan loop runs in a worker thread.
@@ -759,10 +898,10 @@ class StrixTUIApp(App): # type: ignore[misc]
def _setup_cleanup_handlers(self) -> None: def _setup_cleanup_handlers(self) -> None:
def cleanup_on_exit() -> None: def cleanup_on_exit() -> None:
self.tracer.cleanup() self.scan_store.cleanup()
def signal_handler(_signum: int, _frame: Any) -> None: def signal_handler(_signum: int, _frame: Any) -> None:
self.tracer.cleanup() self.scan_store.cleanup()
sys.exit(0) sys.exit(0)
atexit.register(cleanup_on_exit) atexit.register(cleanup_on_exit)
@@ -881,9 +1020,9 @@ class StrixTUIApp(App): # type: ignore[misc]
self._start_scan_thread() self._start_scan_thread()
self.set_interval(0.35, self._update_ui_from_tracer) self.set_interval(0.35, self._update_ui)
def _update_ui_from_tracer(self) -> None: def _update_ui(self) -> None:
if self.show_splash: if self.show_splash:
return return
@@ -902,8 +1041,10 @@ class StrixTUIApp(App): # type: ignore[misc]
except (ValueError, Exception): except (ValueError, Exception):
return return
self._sync_live_view()
agent_updates = False agent_updates = False
for agent_id, agent_data in list(self.tracer.agents.items()): for agent_id, agent_data in list(self.live_view.agents.items()):
if agent_id not in self._displayed_agents: if agent_id not in self._displayed_agents:
self._add_agent_node(agent_data) self._add_agent_node(agent_data)
self._displayed_agents.add(agent_id) self._displayed_agents.add(agent_id)
@@ -922,6 +1063,43 @@ class StrixTUIApp(App): # type: ignore[misc]
self._update_vulnerabilities_panel() self._update_vulnerabilities_panel()
def _sync_live_view(self) -> None:
future = self._live_sync_future
if future is not None:
if not future.done():
if self._scan_loop is not None and self._scan_loop.is_closed():
future.cancel()
self._live_sync_future = None
else:
return
else:
self._live_sync_future = None
try:
parent_of, statuses, names, session_items = future.result()
except Exception:
logger.exception("TUI live-view sync failed")
else:
for agent_id, status in statuses.items():
self.live_view.upsert_agent(
agent_id,
name=names.get(agent_id, agent_id),
parent_id=parent_of.get(agent_id),
status=status,
)
for agent_id, items in session_items.items():
self.live_view.sync_agent_messages_from_items(agent_id, items)
if self._scan_loop is None or self._scan_loop.is_closed():
return
async def collect() -> tuple[
dict[str, str | None], dict[str, Any], dict[str, str], dict[str, list[Any]]
]:
parent_of, statuses, names = await self.coordinator.graph_snapshot()
session_items = await self.coordinator.session_items_snapshot()
return parent_of, statuses, names, session_items
self._live_sync_future = asyncio.run_coroutine_threadsafe(collect(), self._scan_loop)
def _update_agent_node(self, agent_id: str, agent_data: dict[str, Any]) -> bool: def _update_agent_node(self, agent_id: str, agent_data: dict[str, Any]) -> bool:
if agent_id not in self.agent_nodes: if agent_id not in self.agent_nodes:
return False return False
@@ -937,7 +1115,6 @@ class StrixTUIApp(App): # type: ignore[misc]
"completed": "🟢", "completed": "🟢",
"failed": "🔴", "failed": "🔴",
"stopped": "", "stopped": "",
"llm_failed": "🔴",
} }
status_icon = status_indicators.get(status, "") status_icon = status_indicators.get(status, "")
@@ -1074,8 +1251,6 @@ class StrixTUIApp(App): # type: ignore[misc]
if event["type"] == "chat": if event["type"] == "chat":
content = self._render_chat_content(event["data"]) content = self._render_chat_content(event["data"])
elif event["type"] == "tool":
content = self._render_tool_content_simple(event["data"])
if content: if content:
if renderables: if renderables:
@@ -1090,7 +1265,7 @@ class StrixTUIApp(App): # type: ignore[misc]
return self._merge_renderables(renderables) return self._merge_renderables(renderables)
def _get_status_display_content( # noqa: PLR0911 def _get_status_display_content(
self, agent_id: str, agent_data: dict[str, Any] self, agent_id: str, agent_data: dict[str, Any]
) -> tuple[Text | None, Text, bool]: ) -> tuple[Text | None, Text, bool]:
status = agent_data.get("status", "running") status = agent_data.get("status", "running")
@@ -1116,18 +1291,6 @@ class StrixTUIApp(App): # type: ignore[misc]
text.append(msg) text.append(msg)
return (text, Text(), False) return (text, Text(), False)
if status == "llm_failed":
error_msg = agent_data.get("error_message", "")
text = Text()
if error_msg:
text.append(error_msg, style="red")
else:
text.append("LLM request failed", style="red")
self._stop_dot_animation()
keymap = Text()
keymap.append("Send message to retry", style="dim")
return (text, keymap, False)
if status == "failed": if status == "failed":
error_msg = agent_data.get("error_message", "") error_msg = agent_data.get("error_message", "")
text = Text() text = Text()
@@ -1173,7 +1336,7 @@ class StrixTUIApp(App): # type: ignore[misc]
return return
try: try:
agent_data = self.tracer.agents[self.selected_agent_id] agent_data = self.live_view.agents[self.selected_agent_id]
content, keymap, should_animate = self._get_status_display_content( content, keymap, should_animate = self._get_status_display_content(
self.selected_agent_id, agent_data self.selected_agent_id, agent_data
) )
@@ -1206,7 +1369,7 @@ class StrixTUIApp(App): # type: ignore[misc]
stats_content = Text() stats_content = Text()
stats_text = build_tui_stats_text(self.tracer) stats_text = build_tui_stats_text(self.scan_store)
if stats_text: if stats_text:
stats_content.append(stats_text) stats_content.append(stats_text)
@@ -1225,7 +1388,7 @@ class StrixTUIApp(App): # type: ignore[misc]
if not self._is_widget_safe(vuln_panel): if not self._is_widget_safe(vuln_panel):
return return
vulnerabilities = self.tracer.vulnerability_reports vulnerabilities = self.scan_store.vulnerability_reports
if not vulnerabilities: if not vulnerabilities:
self._safe_widget_operation(vuln_panel.add_class, "hidden") self._safe_widget_operation(vuln_panel.add_class, "hidden")
@@ -1234,8 +1397,10 @@ class StrixTUIApp(App): # type: ignore[misc]
enriched_vulns = [] enriched_vulns = []
for vuln in vulnerabilities: for vuln in vulnerabilities:
enriched = dict(vuln) enriched = dict(vuln)
report_id = vuln.get("id", "") agent_name = enriched.get("agent_name")
agent_name = self._get_agent_name_for_vulnerability(report_id) agent_id = enriched.get("agent_id")
if not agent_name and isinstance(agent_id, str):
agent_name = self._get_agent_name(agent_id)
if agent_name: if agent_name:
enriched["agent_name"] = agent_name enriched["agent_name"] = agent_name
enriched_vulns.append(enriched) enriched_vulns.append(enriched)
@@ -1243,18 +1408,6 @@ class StrixTUIApp(App): # type: ignore[misc]
self._safe_widget_operation(vuln_panel.remove_class, "hidden") self._safe_widget_operation(vuln_panel.remove_class, "hidden")
vuln_panel.update_vulnerabilities(enriched_vulns) vuln_panel.update_vulnerabilities(enriched_vulns)
def _get_agent_name_for_vulnerability(self, report_id: str) -> str | None:
"""Find the agent name that created a vulnerability report."""
for _exec_id, tool_data in list(self.tracer.tool_executions.items()):
if tool_data.get("tool_name") == "create_vulnerability_report":
result = tool_data.get("result", {})
if isinstance(result, dict) and result.get("report_id") == report_id:
agent_id = tool_data.get("agent_id")
if agent_id and agent_id in self.tracer.agents:
name: str = self.tracer.agents[agent_id].get("name", "Unknown Agent")
return name
return None
def _get_sweep_animation(self, color_palette: list[str]) -> Text: def _get_sweep_animation(self, color_palette: list[str]) -> Text:
text = Text() text = Text()
num_squares = self._sweep_num_squares num_squares = self._sweep_num_squares
@@ -1307,8 +1460,8 @@ class StrixTUIApp(App): # type: ignore[misc]
def _animate_dots(self) -> None: def _animate_dots(self) -> None:
has_active_agents = False has_active_agents = False
if self.selected_agent_id and self.selected_agent_id in self.tracer.agents: if self.selected_agent_id and self.selected_agent_id in self.live_view.agents:
agent_data = self.tracer.agents[self.selected_agent_id] agent_data = self.live_view.agents[self.selected_agent_id]
status = agent_data.get("status", "running") status = agent_data.get("status", "running")
if status in ["running", "waiting"]: if status in ["running", "waiting"]:
has_active_agents = True has_active_agents = True
@@ -1323,7 +1476,7 @@ class StrixTUIApp(App): # type: ignore[misc]
if not has_active_agents: if not has_active_agents:
has_active_agents = any( has_active_agents = any(
agent_data.get("status", "running") in ["running", "waiting"] agent_data.get("status", "running") in ["running", "waiting"]
for agent_data in self.tracer.agents.values() for agent_data in self.live_view.agents.values()
) )
if not has_active_agents: if not has_active_agents:
@@ -1331,28 +1484,12 @@ class StrixTUIApp(App): # type: ignore[misc]
self._spinner_frame_index = 0 self._spinner_frame_index = 0
def _agent_has_real_activity(self, agent_id: str) -> bool: def _agent_has_real_activity(self, agent_id: str) -> bool:
initial_tools = {"scan_start_info", "subagent_start_info"} return any(msg.get("agent_id") == agent_id for msg in self.live_view.chat_messages)
for _exec_id, tool_data in list(self.tracer.tool_executions.items()):
if tool_data.get("agent_id") == agent_id:
tool_name = tool_data.get("tool_name", "")
if tool_name not in initial_tools:
return True
return False
def _agent_vulnerability_count(self, agent_id: str) -> int: def _agent_vulnerability_count(self, agent_id: str) -> int:
count = 0 return sum(
for _exec_id, tool_data in list(self.tracer.tool_executions.items()): 1 for vuln in self.scan_store.vulnerability_reports if vuln.get("agent_id") == agent_id
if tool_data.get("agent_id") == agent_id: )
tool_name = tool_data.get("tool_name", "")
if tool_name == "create_vulnerability_report":
status = tool_data.get("status", "")
if status == "completed":
result = tool_data.get("result", {})
if isinstance(result, dict) and result.get("success"):
count += 1
return count
def _gather_agent_events(self, agent_id: str) -> list[dict[str, Any]]: def _gather_agent_events(self, agent_id: str) -> list[dict[str, Any]]:
chat_events = [ chat_events = [
@@ -1362,24 +1499,12 @@ class StrixTUIApp(App): # type: ignore[misc]
"id": f"chat_{msg['message_id']}", "id": f"chat_{msg['message_id']}",
"data": msg, "data": msg,
} }
for msg in self.tracer.chat_messages for msg in self.live_view.chat_messages
if msg.get("agent_id") == agent_id if msg.get("agent_id") == agent_id
] ]
tool_events = [ chat_events.sort(key=lambda e: (e["timestamp"], e["id"]))
{ return chat_events
"type": "tool",
"timestamp": tool_data["timestamp"],
"id": f"tool_{exec_id}",
"data": tool_data,
}
for exec_id, tool_data in list(self.tracer.tool_executions.items())
if tool_data.get("agent_id") == agent_id
]
events = chat_events + tool_events
events.sort(key=lambda e: (e["timestamp"], e["id"]))
return events
def watch_selected_agent_id(self, _agent_id: str | None) -> None: def watch_selected_agent_id(self, _agent_id: str | None) -> None:
if len(self.screen_stack) > 1 or self.show_splash: if len(self.screen_stack) > 1 or self.show_splash:
@@ -1412,7 +1537,6 @@ class StrixTUIApp(App): # type: ignore[misc]
scan_id=self.scan_config["run_name"], scan_id=self.scan_config["run_name"],
image=str(image), image=str(image),
local_sources=getattr(self.args, "local_sources", None) or [], local_sources=getattr(self.args, "local_sources", None) or [],
tracer=self.tracer,
coordinator=self.coordinator, coordinator=self.coordinator,
interactive=True, interactive=True,
), ),
@@ -1470,7 +1594,6 @@ class StrixTUIApp(App): # type: ignore[misc]
"completed": "🟢", "completed": "🟢",
"failed": "🔴", "failed": "🔴",
"stopped": "", "stopped": "",
"llm_failed": "🔴",
} }
status_icon = status_indicators.get(status, "") status_icon = status_indicators.get(status, "")
@@ -1532,7 +1655,7 @@ class StrixTUIApp(App): # type: ignore[misc]
def _copy_node_under(self, node_to_copy: TreeNode, new_parent: TreeNode) -> None: def _copy_node_under(self, node_to_copy: TreeNode, new_parent: TreeNode) -> None:
agent_id = node_to_copy.data["agent_id"] agent_id = node_to_copy.data["agent_id"]
agent_data = self.tracer.agents.get(agent_id, {}) agent_data = self.live_view.agents.get(agent_id, {})
agent_name_raw = agent_data.get("name", "Agent") agent_name_raw = agent_data.get("name", "Agent")
status = agent_data.get("status", "running") status = agent_data.get("status", "running")
@@ -1542,7 +1665,6 @@ class StrixTUIApp(App): # type: ignore[misc]
"completed": "🟢", "completed": "🟢",
"failed": "🔴", "failed": "🔴",
"stopped": "", "stopped": "",
"llm_failed": "🔴",
} }
status_icon = status_indicators.get(status, "") status_icon = status_indicators.get(status, "")
@@ -1567,7 +1689,7 @@ class StrixTUIApp(App): # type: ignore[misc]
def _reorganize_orphaned_agents(self, new_parent_id: str) -> None: def _reorganize_orphaned_agents(self, new_parent_id: str) -> None:
agents_to_move = [] agents_to_move = []
for agent_id, agent_data in list(self.tracer.agents.items()): for agent_id, agent_data in list(self.live_view.agents.items()):
if ( if (
agent_data.get("parent_id") == new_parent_id agent_data.get("parent_id") == new_parent_id
and agent_id in self.agent_nodes and agent_id in self.agent_nodes
@@ -1608,71 +1730,6 @@ class StrixTUIApp(App): # type: ignore[misc]
return AgentMessageRenderer.render_simple(content) return AgentMessageRenderer.render_simple(content)
def _render_tool_content_simple(self, tool_data: dict[str, Any]) -> Any:
tool_name = tool_data.get("tool_name", "Unknown Tool")
args = tool_data.get("args", {})
status = tool_data.get("status", "unknown")
result = tool_data.get("result")
renderer = get_tool_renderer(tool_name)
if renderer:
widget = renderer.render(tool_data)
return widget.content
text = Text()
if tool_name in ("llm_error_details", "sandbox_error_details"):
return self._render_error_details(text, tool_name, args)
text.append("→ Using tool ")
text.append(tool_name, style="bold blue")
status_styles = {
"running": ("", "yellow"),
"completed": ("", "green"),
"failed": ("", "red"),
"error": ("", "red"),
}
icon, style = status_styles.get(status, ("", "dim"))
text.append(" ")
text.append(icon, style=style)
if args:
for k, v in list(args.items())[:5]:
str_v = str(v)
if len(str_v) > 500:
str_v = str_v[:497] + "..."
text.append("\n ")
text.append(k, style="dim")
text.append(": ")
text.append(str_v)
if status in ["completed", "failed", "error"] and result:
result_str = str(result)
if len(result_str) > 1000:
result_str = result_str[:997] + "..."
text.append("\n")
text.append("Result: ", style="bold")
text.append(result_str)
return text
def _render_error_details(self, text: Any, tool_name: str, args: dict[str, Any]) -> Any:
if tool_name == "llm_error_details":
text.append("✗ LLM Request Failed", style="red")
else:
text.append("✗ Sandbox Initialization Failed", style="red")
if args.get("error"):
text.append(f"\n{args['error']}", style="bold red")
if args.get("details"):
details = str(args["details"])
if len(details) > 1000:
details = details[:997] + "..."
text.append("\nDetails: ", style="dim")
text.append(details)
return text
@on(Tree.NodeHighlighted) # type: ignore[misc] @on(Tree.NodeHighlighted) # type: ignore[misc]
def handle_tree_highlight(self, event: Tree.NodeHighlighted) -> None: def handle_tree_highlight(self, event: Tree.NodeHighlighted) -> None:
if len(self.screen_stack) > 1 or self.show_splash: if len(self.screen_stack) > 1 or self.show_splash:
@@ -1718,21 +1775,8 @@ class StrixTUIApp(App): # type: ignore[misc]
self.selected_agent_id, self.selected_agent_id,
len(message), len(message),
) )
if self.tracer: # Route to the agent's SDK session. The coordinator also interrupts
self.tracer.log_chat_message( # any active stream so the message is picked up on the next run cycle.
content=message,
role="user",
agent_id=self.selected_agent_id,
)
# Route to the agent's SDK session. 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. 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(): if self._scan_loop is not None and not self._scan_loop.is_closed():
target_agent_id = self.selected_agent_id target_agent_id = self.selected_agent_id
asyncio.run_coroutine_threadsafe( asyncio.run_coroutine_threadsafe(
@@ -1742,10 +1786,6 @@ class StrixTUIApp(App): # type: ignore[misc]
), ),
self._scan_loop, self._scan_loop,
) )
asyncio.run_coroutine_threadsafe(
self.coordinator.request_interrupt(target_agent_id, mode="after_turn"),
self._scan_loop,
)
self._displayed_events.clear() self._displayed_events.clear()
self._update_chat_view() self._update_chat_view()
@@ -1754,8 +1794,8 @@ class StrixTUIApp(App): # type: ignore[misc]
def _get_agent_name(self, agent_id: str) -> str: def _get_agent_name(self, agent_id: str) -> str:
try: try:
if self.tracer and agent_id in self.tracer.agents: if agent_id in self.live_view.agents:
agent_name = self.tracer.agents[agent_id].get("name") agent_name = self.live_view.agents[agent_id].get("name")
if isinstance(agent_name, str): if isinstance(agent_name, str):
return agent_name return agent_name
except (KeyError, AttributeError) as e: except (KeyError, AttributeError) as e:
@@ -1822,12 +1862,12 @@ class StrixTUIApp(App): # type: ignore[misc]
agent_name = "Unknown Agent" agent_name = "Unknown Agent"
try: try:
if self.tracer and self.selected_agent_id in self.tracer.agents: if self.selected_agent_id in self.live_view.agents:
agent_data = self.tracer.agents[self.selected_agent_id] agent_data = self.live_view.agents[self.selected_agent_id]
agent_name = agent_data.get("name", "Unknown Agent") agent_name = agent_data.get("name", "Unknown Agent")
agent_status = agent_data.get("status", "running") agent_status = agent_data.get("status", "running")
if agent_status not in ["running", "waiting", "llm_failed"]: if agent_status not in ["running", "waiting"]:
return agent_name, False return agent_name, False
agent_events = self._gather_agent_events(self.selected_agent_id) agent_events = self._gather_agent_events(self.selected_agent_id)
@@ -1862,7 +1902,7 @@ class StrixTUIApp(App): # type: ignore[misc]
self._scan_thread.join(timeout=1.0) self._scan_thread.join(timeout=1.0)
self.tracer.cleanup() self.scan_store.cleanup()
self.exit() self.exit()
+14 -109
View File
@@ -23,16 +23,6 @@ from rich.text import Text
from strix.config import load_settings from strix.config import load_settings
# Token formatting utilities
def format_token_count(count: float) -> str:
count = int(count)
if count >= 1_000_000:
return f"{count / 1_000_000:.1f}M"
if count >= 1_000:
return f"{count / 1_000:.1f}K"
return str(count)
# Display utilities # Display utilities
def get_severity_color(severity: str) -> str: def get_severity_color(severity: str) -> str:
severity_colors = { severity_colors = {
@@ -206,13 +196,13 @@ def format_vulnerability_report(report: dict[str, Any]) -> Text: # noqa: PLR091
return text return text
def _build_vulnerability_stats(stats_text: Text, tracer: Any) -> None: def _build_vulnerability_stats(stats_text: Text, scan_store: Any) -> None:
"""Build vulnerability section of stats text.""" """Build vulnerability section of stats text."""
vuln_count = len(tracer.vulnerability_reports) vuln_count = len(scan_store.vulnerability_reports)
if vuln_count > 0: if vuln_count > 0:
severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0} severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
for report in tracer.vulnerability_reports: for report in scan_store.vulnerability_reports:
severity = report.get("severity", "").lower() severity = report.get("severity", "").lower()
if severity in severity_counts: if severity in severity_counts:
severity_counts[severity] += 1 severity_counts[severity] += 1
@@ -245,63 +235,20 @@ def _build_vulnerability_stats(stats_text: Text, tracer: Any) -> None:
stats_text.append("\n") stats_text.append("\n")
def _build_llm_stats(stats_text: Text, total_stats: dict[str, Any]) -> None: def build_final_stats_text(scan_store: Any) -> Text:
"""Build LLM usage section of stats text.""" """Build final stats from Strix-owned scan artifacts."""
if total_stats["requests"] > 0:
stats_text.append("\n")
stats_text.append("Input Tokens ", style="dim")
stats_text.append(format_token_count(total_stats["input_tokens"]), style="white")
if total_stats["cached_tokens"] > 0:
stats_text.append(" · ", style="dim white")
stats_text.append("Cached Tokens ", style="dim")
stats_text.append(format_token_count(total_stats["cached_tokens"]), style="white")
stats_text.append(" · ", style="dim white")
stats_text.append("Output Tokens ", style="dim")
stats_text.append(format_token_count(total_stats["output_tokens"]), style="white")
if total_stats["cost"] > 0:
stats_text.append(" · ", style="dim white")
stats_text.append("Cost ", style="dim")
stats_text.append(f"${total_stats['cost']:.4f}", style="bold #fbbf24")
else:
stats_text.append("\n")
stats_text.append("Cost ", style="dim")
stats_text.append("$0.0000 ", style="#fbbf24")
stats_text.append("· ", style="dim white")
stats_text.append("Tokens ", style="dim")
stats_text.append("0", style="white")
def build_final_stats_text(tracer: Any) -> Text:
"""Build stats text for final output with detailed messages and LLM usage."""
stats_text = Text() stats_text = Text()
if not tracer: if not scan_store:
return stats_text return stats_text
_build_vulnerability_stats(stats_text, tracer) _build_vulnerability_stats(stats_text, scan_store)
tool_count = tracer.get_real_tool_count()
agent_count = len(tracer.agents)
stats_text.append("Agents", style="dim")
stats_text.append(" ")
stats_text.append(str(agent_count), style="bold white")
stats_text.append(" · ", style="dim white")
stats_text.append("Tools", style="dim")
stats_text.append(" ")
stats_text.append(str(tool_count), style="bold white")
llm_stats = tracer.get_total_llm_stats()
_build_llm_stats(stats_text, llm_stats["total"])
return stats_text return stats_text
def build_live_stats_text(tracer: Any) -> Text: def build_live_stats_text(scan_store: Any) -> Text:
stats_text = Text() stats_text = Text()
if not tracer: if not scan_store:
return stats_text return stats_text
model = load_settings().llm.model or "unknown" model = load_settings().llm.model or "unknown"
@@ -309,16 +256,13 @@ def build_live_stats_text(tracer: Any) -> Text:
stats_text.append(str(model), style="white") stats_text.append(str(model), style="white")
stats_text.append("\n") stats_text.append("\n")
vuln_count = len(tracer.vulnerability_reports) vuln_count = len(scan_store.vulnerability_reports)
tool_count = tracer.get_real_tool_count()
agent_count = len(tracer.agents)
stats_text.append("Vulnerabilities ", style="dim") stats_text.append("Vulnerabilities ", style="dim")
stats_text.append(f"{vuln_count}", style="white") stats_text.append(f"{vuln_count}", style="white")
stats_text.append("\n") stats_text.append("\n")
if vuln_count > 0: if vuln_count > 0:
severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0} severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
for report in tracer.vulnerability_reports: for report in scan_store.vulnerability_reports:
severity = report.get("severity", "").lower() severity = report.get("severity", "").lower()
if severity in severity_counts: if severity in severity_counts:
severity_counts[severity] += 1 severity_counts[severity] += 1
@@ -340,57 +284,18 @@ def build_live_stats_text(tracer: Any) -> Text:
stats_text.append("\n") stats_text.append("\n")
stats_text.append("Agents ", style="dim")
stats_text.append(str(agent_count), style="white")
stats_text.append(" · ", style="dim white")
stats_text.append("Tools ", style="dim")
stats_text.append(str(tool_count), style="white")
llm_stats = tracer.get_total_llm_stats()
total_stats = llm_stats["total"]
stats_text.append("\n")
stats_text.append("Input Tokens ", style="dim")
stats_text.append(format_token_count(total_stats["input_tokens"]), style="white")
stats_text.append(" · ", style="dim white")
stats_text.append("Cached Tokens ", style="dim")
stats_text.append(format_token_count(total_stats["cached_tokens"]), style="white")
stats_text.append("\n")
stats_text.append("Output Tokens ", style="dim")
stats_text.append(format_token_count(total_stats["output_tokens"]), style="white")
stats_text.append(" · ", style="dim white")
stats_text.append("Cost ", style="dim")
stats_text.append(f"${total_stats['cost']:.4f}", style="#fbbf24")
return stats_text return stats_text
def build_tui_stats_text(tracer: Any) -> Text: def build_tui_stats_text(scan_store: Any) -> Text:
stats_text = Text() stats_text = Text()
if not tracer: if not scan_store:
return stats_text return stats_text
model = load_settings().llm.model or "unknown" model = load_settings().llm.model or "unknown"
stats_text.append(str(model), style="white") stats_text.append(str(model), style="white")
llm_stats = tracer.get_total_llm_stats() caido_url = getattr(scan_store, "caido_url", None)
total_stats = llm_stats["total"]
total_tokens = total_stats["input_tokens"] + total_stats["output_tokens"]
if total_tokens > 0:
stats_text.append("\n")
stats_text.append(f"{format_token_count(total_tokens)} tokens", style="white")
if total_stats["cost"] > 0:
stats_text.append(" · ", style="white")
stats_text.append(f"${total_stats['cost']:.2f}", style="white")
caido_url = getattr(tracer, "caido_url", None)
if caido_url: if caido_url:
stats_text.append("\n") stats_text.append("\n")
stats_text.append("Caido: ", style="bold white") stats_text.append("Caido: ", style="bold white")
+1 -2
View File
@@ -3,8 +3,7 @@
- :class:`AgentCoordinator` owns Strix-specific graph/status/wake state. - :class:`AgentCoordinator` owns Strix-specific graph/status/wake state.
- SDK ``SQLiteSession`` owns per-agent conversation history and message - SDK ``SQLiteSession`` owns per-agent conversation history and message
transport. transport.
- :func:`run_with_continuation` wraps SDK ``Runner.run_streamed`` only - ``runner.py`` owns SDK ``Runner.run_streamed`` and child-agent spawning.
enough to keep interactive agents addressable between bounded runs.
Import deeply (for example, ``from strix.orchestration.coordinator Import deeply (for example, ``from strix.orchestration.coordinator
import AgentCoordinator``) so ``import strix.orchestration`` doesn't import AgentCoordinator``) so ``import strix.orchestration`` doesn't
+62 -439
View File
@@ -13,39 +13,27 @@ import json
import logging import logging
import tempfile import tempfile
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal from typing import TYPE_CHECKING, Any, Literal, cast
from agents import Runner
from agents.exceptions import AgentsException, MaxTurnsExceeded, UserError
from agents.memory import SQLiteSession
from openai import APIError
if TYPE_CHECKING: if TYPE_CHECKING:
from agents.items import TResponseInputItem from agents.items import TResponseInputItem
from agents.memory import Session from agents.memory import Session
from agents.result import RunResultBase, RunResultStreaming
from agents.run_config import RunConfig
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed", "llm_failed"] Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed"]
ACTIVE_AGENT_STATUSES = {"running", "waiting", "llm_failed"}
_SNAPSHOT_VERSION = 3
_WAITING_TIMEOUT_SUBAGENT = 300.0
_TIMEOUT_RESUME_MESSAGE = "Waiting timeout reached. Resuming execution."
@dataclass(slots=True) @dataclass(slots=True)
class AgentRuntime: class AgentRuntime:
session: Session | None = None session: Session | None = None
task: asyncio.Task[Any] | None = None task: asyncio.Task[Any] | None = None
stream: RunResultStreaming | None = None stream: Any | None = None
interrupt_on_message: bool = False
wake: asyncio.Event = field(default_factory=asyncio.Event) wake: asyncio.Event = field(default_factory=asyncio.Event)
tool_calls: dict[str, str] = field(default_factory=dict)
class AgentCoordinator: class AgentCoordinator:
@@ -57,9 +45,6 @@ class AgentCoordinator:
self.names: dict[str, str] = {} self.names: dict[str, str] = {}
self.metadata: dict[str, dict[str, Any]] = {} self.metadata: dict[str, dict[str, Any]] = {}
self.pending_counts: dict[str, int] = {} self.pending_counts: dict[str, int] = {}
self.pending_user_counts: dict[str, int] = {}
self.queued_messages: dict[str, list[dict[str, Any]]] = {}
self.stats_live: dict[str, dict[str, Any]] = {}
self.runtimes: dict[str, AgentRuntime] = {} self.runtimes: dict[str, AgentRuntime] = {}
self._lock = asyncio.Lock() self._lock = asyncio.Lock()
self._snapshot_path: Path | None = None self._snapshot_path: Path | None = None
@@ -75,23 +60,15 @@ class AgentCoordinator:
*, *,
task: str | None = None, task: str | None = None,
skills: list[str] | None = None, skills: list[str] | None = None,
is_whitebox: bool = False,
scan_mode: str = "deep",
diff_scope: dict[str, Any] | None = None,
) -> None: ) -> None:
async with self._lock: async with self._lock:
self.statuses[agent_id] = "running" self.statuses[agent_id] = "running"
self.parent_of[agent_id] = parent_id self.parent_of[agent_id] = parent_id
self.names[agent_id] = name self.names[agent_id] = name
self.pending_counts.setdefault(agent_id, 0) self.pending_counts.setdefault(agent_id, 0)
self.pending_user_counts.setdefault(agent_id, 0)
self.stats_live.setdefault(agent_id, _empty_stats())
self.metadata[agent_id] = { self.metadata[agent_id] = {
"task": task or "", "task": task or "",
"skills": list(skills or []), "skills": list(skills or []),
"is_whitebox": bool(is_whitebox),
"scan_mode": scan_mode,
"diff_scope": diff_scope,
} }
self.runtimes.setdefault(agent_id, AgentRuntime()) self.runtimes.setdefault(agent_id, AgentRuntime())
logger.info("agent.register %s (%s) parent=%s", agent_id, name, parent_id or "-") logger.info("agent.register %s (%s) parent=%s", agent_id, name, parent_id or "-")
@@ -103,6 +80,7 @@ class AgentCoordinator:
*, *,
session: Session | None = None, session: Session | None = None,
task: asyncio.Task[Any] | None = None, task: asyncio.Task[Any] | None = None,
interrupt_on_message: bool | None = None,
) -> None: ) -> None:
async with self._lock: async with self._lock:
runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
@@ -110,8 +88,8 @@ class AgentCoordinator:
runtime.session = session runtime.session = session
if task is not None: if task is not None:
runtime.task = task runtime.task = task
if session is not None: if interrupt_on_message is not None:
await self.flush_queued_messages(agent_id) runtime.interrupt_on_message = interrupt_on_message
async def mark_running(self, agent_id: str) -> None: async def mark_running(self, agent_id: str) -> None:
async with self._lock: async with self._lock:
@@ -122,9 +100,6 @@ class AgentCoordinator:
async def park_waiting(self, agent_id: str) -> None: async def park_waiting(self, agent_id: str) -> None:
await self.set_status(agent_id, "waiting") await self.set_status(agent_id, "waiting")
async def mark_llm_failed(self, agent_id: str) -> None:
await self.set_status(agent_id, "llm_failed")
async def set_status(self, agent_id: str, status: Status | str) -> None: async def set_status(self, agent_id: str, status: Status | str) -> None:
async with self._lock: async with self._lock:
if agent_id not in self.statuses: if agent_id not in self.statuses:
@@ -137,17 +112,15 @@ class AgentCoordinator:
async def send(self, target_agent_id: str, message: dict[str, Any]) -> bool: async def send(self, target_agent_id: str, message: dict[str, Any]) -> bool:
"""Deliver a user/peer message by appending it to the target SDK session.""" """Deliver a user/peer message by appending it to the target SDK session."""
should_queue = False
async with self._lock: async with self._lock:
if target_agent_id not in self.statuses: if target_agent_id not in self.statuses:
logger.debug("agent.send dropped unknown target=%s", target_agent_id) logger.debug("agent.send dropped unknown target=%s", target_agent_id)
return False return False
runtime = self.runtimes.setdefault(target_agent_id, AgentRuntime()) runtime = self.runtimes.setdefault(target_agent_id, AgentRuntime())
session = runtime.session session = runtime.session
should_queue = session is None or runtime.stream is not None stream = runtime.stream
if should_queue: interrupt = runtime.interrupt_on_message
self.queued_messages.setdefault(target_agent_id, []).append(dict(message)) if session is not None:
if session is not None and not should_queue:
try: try:
await session.add_items([self._message_to_session_item(message)]) await session.add_items([self._message_to_session_item(message)])
except Exception: except Exception:
@@ -158,81 +131,55 @@ class AgentCoordinator:
return False return False
async with self._lock: async with self._lock:
self.pending_counts[target_agent_id] = self.pending_counts.get(target_agent_id, 0) + 1 self.pending_counts[target_agent_id] = self.pending_counts.get(target_agent_id, 0) + 1
if message.get("from") == "user":
self.pending_user_counts[target_agent_id] = (
self.pending_user_counts.get(target_agent_id, 0) + 1
)
self.runtimes.setdefault(target_agent_id, AgentRuntime()).wake.set() self.runtimes.setdefault(target_agent_id, AgentRuntime()).wake.set()
if should_queue: if stream is not None and interrupt:
logger.debug( stream.cancel(mode="immediate")
"agent.send %s queued until SDK session is safe to append", target_agent_id
)
await self._maybe_snapshot() await self._maybe_snapshot()
return True return True
async def flush_queued_messages(self, agent_id: str) -> None: async def wait_for_message(self, agent_id: str) -> None:
async with self._lock:
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
session = runtime.session
queued = self.queued_messages.pop(agent_id, [])
if not queued:
return
if session is None:
async with self._lock:
self.queued_messages.setdefault(agent_id, []).extend(queued)
return
try:
await session.add_items([self._message_to_session_item(msg) for msg in queued])
except Exception:
async with self._lock:
self.queued_messages.setdefault(agent_id, [])[0:0] = queued
logger.exception("agent.flush_queued_messages failed for %s", agent_id)
async def wait_for_message(self, agent_id: str, *, user_only: bool = False) -> None:
while True: while True:
async with self._lock: async with self._lock:
pending = ( if self.pending_counts.get(agent_id, 0) > 0:
self.pending_user_counts.get(agent_id, 0)
if user_only
else self.pending_counts.get(agent_id, 0)
)
if pending > 0:
return return
wake = self.runtimes.setdefault(agent_id, AgentRuntime()).wake wake = self.runtimes.setdefault(agent_id, AgentRuntime()).wake
wake.clear() wake.clear()
await wake.wait() await wake.wait()
async def wait_for_user_message(self, agent_id: str) -> None: async def consume_pending(
await self.wait_for_message(agent_id, user_only=True) self,
agent_id: str,
async def consume_wake(self, agent_id: str) -> None: *,
include_items: bool = False,
) -> tuple[int, list[Any]]:
async with self._lock: async with self._lock:
count = self.pending_counts.get(agent_id, 0)
self.pending_counts[agent_id] = 0 self.pending_counts[agent_id] = 0
self.pending_user_counts[agent_id] = 0
async def pending_count(self, agent_id: str) -> int:
async with self._lock:
return self.pending_counts.get(agent_id, 0)
async def recent_session_items(self, agent_id: str, count: int) -> list[TResponseInputItem]:
if count <= 0:
return []
async with self._lock:
session = self.runtimes.get(agent_id, AgentRuntime()).session session = self.runtimes.get(agent_id, AgentRuntime()).session
if session is None: if count <= 0:
return [] return 0, []
if not include_items or session is None:
return count, []
items = await session.get_items() items = await session.get_items()
return list(items[-count:]) return count, list(items[-count:])
async def request_interrupt(self, agent_id: str, mode: str = "after_turn") -> bool: async def session_items_snapshot(self) -> dict[str, list[Any]]:
async with self._lock: async with self._lock:
stream = self.runtimes.get(agent_id, AgentRuntime()).stream sessions = {
if stream is None: aid: runtime.session
return False for aid, runtime in self.runtimes.items()
stream.cancel(mode=mode) # type: ignore[arg-type] if runtime.session is not None
return True }
async def request_stop(self, agent_id: str, *, interrupt: bool = True) -> None: snapshots: dict[str, list[Any]] = {}
for aid, session in sessions.items():
try:
snapshots[aid] = list(await session.get_items())
except Exception:
logger.exception("failed to read SDK session items for %s", aid)
return snapshots
async def request_stop(self, agent_id: str) -> None:
async with self._lock: async with self._lock:
if agent_id not in self.statuses: if agent_id not in self.statuses:
return return
@@ -240,7 +187,7 @@ class AgentCoordinator:
runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
runtime.wake.set() runtime.wake.set()
stream = runtime.stream stream = runtime.stream
if interrupt and stream is not None: if stream is not None:
stream.cancel(mode="after_turn") stream.cancel(mode="after_turn")
await self._maybe_snapshot() await self._maybe_snapshot()
@@ -266,7 +213,7 @@ class AgentCoordinator:
async def attach_stream( async def attach_stream(
self, self,
agent_id: str, agent_id: str,
stream: RunResultStreaming, stream: Any,
) -> None: ) -> None:
async with self._lock: async with self._lock:
self.runtimes.setdefault(agent_id, AgentRuntime()).stream = stream self.runtimes.setdefault(agent_id, AgentRuntime()).stream = stream
@@ -274,7 +221,7 @@ class AgentCoordinator:
async def detach_stream( async def detach_stream(
self, self,
agent_id: str, agent_id: str,
stream: RunResultStreaming, stream: Any,
) -> None: ) -> None:
async with self._lock: async with self._lock:
runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
@@ -284,49 +231,40 @@ class AgentCoordinator:
async def active_agents_except(self, agent_id: str) -> list[dict[str, Any]]: async def active_agents_except(self, agent_id: str) -> list[dict[str, Any]]:
async with self._lock: async with self._lock:
return [ return [
self._agent_info_locked(aid) {
"agent_id": aid,
"name": self.names.get(aid, aid),
"status": status,
"parent_id": self.parent_of.get(aid),
}
for aid, status in self.statuses.items() for aid, status in self.statuses.items()
if aid != agent_id and status in ACTIVE_AGENT_STATUSES if aid != agent_id and status in {"running", "waiting"}
] ]
async def agent_info(self, agent_id: str) -> dict[str, Any] | None:
async with self._lock:
if agent_id not in self.statuses:
return None
return self._agent_info_locked(agent_id)
async def graph_snapshot( async def graph_snapshot(
self, self,
) -> tuple[dict[str, str | None], dict[str, Status], dict[str, str]]: ) -> tuple[dict[str, str | None], dict[str, Status], dict[str, str]]:
async with self._lock: async with self._lock:
return dict(self.parent_of), dict(self.statuses), dict(self.names) return dict(self.parent_of), dict(self.statuses), dict(self.names)
async def record_usage(self, agent_id: str, usage: Any) -> None:
if usage is None:
return
async with self._lock:
stats = self.stats_live.setdefault(agent_id, _empty_stats())
stats["calls"] += 1
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)
stats["cached"] += getattr(details, "cached_tokens", 0) or 0 if details else 0
def _message_to_session_item(self, message: dict[str, Any]) -> TResponseInputItem: def _message_to_session_item(self, message: dict[str, Any]) -> TResponseInputItem:
sender = str(message.get("from", "unknown")) sender = str(message.get("from", "unknown"))
content = str(message.get("content", "")) content = str(message.get("content", ""))
if sender == "user": if sender == "user":
return {"role": "user", "content": content} return cast("TResponseInputItem", {"role": "user", "content": content})
sender_name = self.names.get(sender, sender) sender_name = self.names.get(sender, sender)
msg_type = message.get("type", "information") msg_type = message.get("type", "information")
priority = message.get("priority", "normal") priority = message.get("priority", "normal")
return { return cast(
"role": "user", "TResponseInputItem",
"content": ( {
f"[Message from {sender_name} ({sender}) | type={msg_type} " "role": "user",
f"| priority={priority}]\n{content}" "content": (
), f"[Message from {sender_name} ({sender}) | type={msg_type} "
} f"| priority={priority}]\n{content}"
),
},
)
def _subtree_order_locked(self, agent_id: str) -> list[str]: def _subtree_order_locked(self, agent_id: str) -> list[str]:
queue = [agent_id] queue = [agent_id]
@@ -337,29 +275,14 @@ class AgentCoordinator:
queue.extend(child for child, parent in self.parent_of.items() if parent == aid) queue.extend(child for child, parent in self.parent_of.items() if parent == aid)
return order return order
def _agent_info_locked(self, agent_id: str) -> dict[str, Any]:
return {
"agent_id": agent_id,
"name": self.names.get(agent_id, agent_id),
"status": self.statuses.get(agent_id),
"parent_id": self.parent_of.get(agent_id),
"pending_messages": self.pending_counts.get(agent_id, 0),
}
async def snapshot(self) -> dict[str, Any]: async def snapshot(self) -> dict[str, Any]:
async with self._lock: async with self._lock:
return { return {
"version": _SNAPSHOT_VERSION,
"statuses": dict(self.statuses), "statuses": dict(self.statuses),
"parent_of": dict(self.parent_of), "parent_of": dict(self.parent_of),
"names": dict(self.names), "names": dict(self.names),
"metadata": {aid: dict(md) for aid, md in self.metadata.items()}, "metadata": {aid: dict(md) for aid, md in self.metadata.items()},
"pending_counts": dict(self.pending_counts), "pending_counts": dict(self.pending_counts),
"pending_user_counts": dict(self.pending_user_counts),
"queued_messages": {
aid: [dict(msg) for msg in msgs] for aid, msgs in self.queued_messages.items()
},
"stats_live": {aid: dict(s) for aid, s in self.stats_live.items()},
} }
async def restore(self, snap: dict[str, Any]) -> None: async def restore(self, snap: dict[str, Any]) -> None:
@@ -369,12 +292,6 @@ class AgentCoordinator:
self.names = dict(snap.get("names", {})) self.names = dict(snap.get("names", {}))
self.metadata = {aid: dict(md) for aid, md in snap.get("metadata", {}).items()} self.metadata = {aid: dict(md) for aid, md in snap.get("metadata", {}).items()}
self.pending_counts = dict(snap.get("pending_counts", {})) self.pending_counts = dict(snap.get("pending_counts", {}))
self.pending_user_counts = dict(snap.get("pending_user_counts", {}))
self.queued_messages = {
aid: [dict(msg) for msg in msgs]
for aid, msgs in snap.get("queued_messages", {}).items()
}
self.stats_live = {aid: dict(s) for aid, s in snap.get("stats_live", {}).items()}
for aid in self.statuses: for aid in self.statuses:
self.runtimes.setdefault(aid, AgentRuntime()) self.runtimes.setdefault(aid, AgentRuntime())
@@ -404,297 +321,3 @@ class AgentCoordinator:
def coordinator_from_context(ctx: dict[str, Any]) -> AgentCoordinator | None: def coordinator_from_context(ctx: dict[str, Any]) -> AgentCoordinator | None:
coordinator = ctx.get("coordinator") coordinator = ctx.get("coordinator")
return coordinator if isinstance(coordinator, AgentCoordinator) else None return coordinator if isinstance(coordinator, AgentCoordinator) else None
async def run_with_continuation(
*,
agent: Any,
initial_input: Any,
run_config: RunConfig,
context: dict[str, Any],
max_turns: int,
coordinator: AgentCoordinator,
agent_id: str,
interactive: bool,
session: Session | None = None,
start_parked: bool = False,
) -> RunResultBase | None:
await coordinator.attach_runtime(agent_id, session=session)
waiting_timeout = await _waiting_timeout(coordinator, agent_id, interactive)
result: RunResultBase | None = None
if not (start_parked and interactive):
result = await _run_cycle(
agent,
coordinator,
agent_id,
input_data=initial_input,
run_config=run_config,
context=context,
max_turns=max_turns,
session=session,
interactive=interactive,
)
if not interactive:
return result
while True:
async with coordinator._lock:
status = coordinator.statuses.get(agent_id)
try:
if status == "llm_failed":
await coordinator.wait_for_user_message(agent_id)
elif waiting_timeout is None:
await coordinator.wait_for_message(agent_id)
else:
await asyncio.wait_for(
coordinator.wait_for_message(agent_id),
timeout=waiting_timeout,
)
except asyncio.CancelledError:
return result
except TimeoutError:
result = await _run_cycle(
agent,
coordinator,
agent_id,
input_data=_TIMEOUT_RESUME_MESSAGE,
run_config=run_config,
context=context,
max_turns=max_turns,
session=session,
interactive=interactive,
)
continue
await coordinator.consume_wake(agent_id)
result = await _run_cycle(
agent,
coordinator,
agent_id,
input_data=[],
run_config=run_config,
context=context,
max_turns=max_turns,
session=session,
interactive=interactive,
)
async def _waiting_timeout(
coordinator: AgentCoordinator,
agent_id: str,
interactive: bool,
) -> float | None:
if not interactive:
return None
async with coordinator._lock:
return (
_WAITING_TIMEOUT_SUBAGENT if coordinator.parent_of.get(agent_id) is not None else None
)
async def _run_cycle(
agent: Any,
coordinator: AgentCoordinator,
agent_id: str,
*,
input_data: Any,
run_config: RunConfig,
context: dict[str, Any],
max_turns: int,
session: Session | None,
interactive: bool,
) -> RunResultBase | None:
try:
await coordinator.mark_running(agent_id)
stream = Runner.run_streamed(
agent,
input=input_data,
run_config=run_config,
context=context,
max_turns=max_turns,
session=session,
)
await coordinator.attach_stream(agent_id, stream)
try:
async for event in stream.stream_events():
await _handle_stream_event(coordinator, agent_id, context, event)
finally:
await coordinator.detach_stream(agent_id, stream)
await coordinator.flush_queued_messages(agent_id)
await _settle_run_result(coordinator, agent_id, stream.final_output, interactive, context)
return stream
except (AgentsException, APIError):
if not interactive:
raise
logger.exception("LLM/runtime failure for %s; waiting for user resume", agent_id)
await coordinator.mark_llm_failed(agent_id)
return None
except Exception as exc:
if not interactive:
raise
status: Status = "stopped" if isinstance(exc, MaxTurnsExceeded) else "crashed"
if isinstance(exc, UserError):
status = "failed"
logger.exception("agent run failed for %s; parking as %s", agent_id, status)
await coordinator.set_status(agent_id, status)
await _notify_parent_on_crash(coordinator, agent_id, status)
_mirror_tracer_status(context, coordinator, agent_id, status)
return None
async def _settle_run_result(
coordinator: AgentCoordinator,
agent_id: str,
final_output: Any,
interactive: bool,
context: dict[str, Any],
) -> None:
parsed = _parse_final_output(final_output)
async with coordinator._lock:
current_status = coordinator.statuses.get(agent_id)
if current_status == "stopped":
status: Status = "stopped"
elif parsed.get("agent_completed") or parsed.get("scan_completed"):
status = "completed"
elif parsed.get("agent_waiting") or interactive:
status = "waiting"
else:
status = "crashed"
await coordinator.set_status(agent_id, status)
await _notify_parent_on_crash(coordinator, agent_id, status)
_mirror_tracer_status(context, coordinator, agent_id, status)
def _parse_final_output(output: Any) -> dict[str, Any]:
if not isinstance(output, str):
return {}
try:
parsed = json.loads(output)
except (TypeError, ValueError):
return {}
return parsed if isinstance(parsed, dict) and parsed.get("success") else {}
async def _notify_parent_on_crash(
coordinator: AgentCoordinator,
agent_id: str,
status: str,
) -> None:
if status != "crashed":
return
async with coordinator._lock:
parent = coordinator.parent_of.get(agent_id)
name = coordinator.names.get(agent_id, agent_id)
if parent is None:
return
await coordinator.send(
parent,
{
"from": agent_id,
"type": "crash",
"priority": "high",
"content": (
f"[Agent crash] {name} ({agent_id}) terminated unexpectedly. "
"Stop waiting on this child unless you want to message it again."
),
},
)
def _mirror_tracer_status(
context: dict[str, Any],
coordinator: AgentCoordinator,
agent_id: str,
status: str,
) -> None:
tracer = context.get("tracer")
if tracer is None:
return
now = datetime.now(UTC).isoformat()
tracer.agents.setdefault(
agent_id,
{
"id": agent_id,
"name": coordinator.names.get(agent_id, agent_id),
"parent_id": coordinator.parent_of.get(agent_id),
"created_at": now,
},
)
tracer.agents[agent_id]["status"] = status
tracer.agents[agent_id]["updated_at"] = now
async def _handle_stream_event(
coordinator: AgentCoordinator,
agent_id: str,
context: dict[str, Any],
event: Any,
) -> None:
tracer = context.get("tracer")
if event.type == "raw_response_event":
response = getattr(event.data, "response", None)
usage = getattr(response, "usage", None)
if usage is not None:
await coordinator.record_usage(agent_id, usage)
if usage is not None and tracer is not None and hasattr(tracer, "record_llm_usage"):
details = getattr(usage, "input_tokens_details", None)
tracer.record_llm_usage(
input_tokens=int(getattr(usage, "input_tokens", 0) or 0),
output_tokens=int(getattr(usage, "output_tokens", 0) or 0),
cached_tokens=int(getattr(details, "cached_tokens", 0) or 0) if details else 0,
)
return
if tracer is None or event.type != "run_item_stream_event":
return
item = event.item
raw = getattr(item, "raw_item", None)
if event.name == "tool_called":
call_id = str(getattr(raw, "call_id", None) or getattr(raw, "id", ""))
tool_name = str(getattr(raw, "name", None) or getattr(raw, "type", "tool"))
args = _parse_tool_args(getattr(raw, "arguments", None))
runtime = coordinator.runtimes.setdefault(agent_id, AgentRuntime())
if call_id:
runtime.tool_calls[call_id] = tool_name
tracer.log_tool_start(agent_id, tool_name, args)
elif event.name == "tool_output":
call_id = str(getattr(raw, "call_id", None) or "")
runtime = coordinator.runtimes.setdefault(agent_id, AgentRuntime())
tool_name = runtime.tool_calls.get(call_id, "tool")
tracer.log_tool_end(agent_id, tool_name, _dump_raw(raw))
def _parse_tool_args(raw: Any) -> dict[str, Any]:
if not raw:
return {}
try:
parsed = json.loads(raw)
except (TypeError, ValueError):
return {}
return parsed if isinstance(parsed, dict) else {}
def _dump_raw(raw: Any) -> Any:
if hasattr(raw, "model_dump"):
return raw.model_dump(exclude_unset=True)
return raw
def open_agent_session(agent_id: str, path: Path) -> SQLiteSession:
path.parent.mkdir(parents=True, exist_ok=True)
return SQLiteSession(session_id=agent_id, db_path=path)
def _empty_stats() -> dict[str, Any]:
return {
"in": 0,
"out": 0,
"cached": 0,
"cost": 0.0,
"calls": 0,
}
+636
View File
@@ -0,0 +1,636 @@
"""Top-level Strix scan runner.
The SDK owns model/tool execution and per-agent sessions. This module owns
Strix-specific scan setup, child-agent startup, resume, and the small wake loop
needed to keep every agent addressable after its SDK run parks.
"""
from __future__ import annotations
import asyncio
import contextlib
import json
import logging
import uuid
from pathlib import Path
from typing import TYPE_CHECKING, Any
from agents import RunConfig, Runner
from agents.exceptions import AgentsException, MaxTurnsExceeded, UserError
from agents.memory import SQLiteSession
from agents.sandbox import SandboxRunConfig
from openai import APIError
from strix.agents.factory import build_strix_agent, make_child_factory
from strix.config import load_settings
from strix.llm.multi_provider_setup import build_multi_provider
from strix.orchestration.coordinator import AgentCoordinator, Status
from strix.orchestration.utils import (
DEFAULT_MAX_TURNS,
build_root_task,
build_scope_context,
child_initial_input,
make_model_settings,
)
from strix.runtime import session_manager
from strix.telemetry.logging import set_scan_id, setup_scan_logging
if TYPE_CHECKING:
from agents.memory import Session
from agents.result import RunResultBase
logger = logging.getLogger(__name__)
def _open_agent_session(agent_id: str, path: Path) -> SQLiteSession:
path.parent.mkdir(parents=True, exist_ok=True)
return SQLiteSession(session_id=agent_id, db_path=path)
async def _run_agent_loop(
*,
agent: Any,
initial_input: Any,
run_config: RunConfig,
context: dict[str, Any],
max_turns: int,
coordinator: AgentCoordinator,
agent_id: str,
interactive: bool,
session: Session | None = None,
start_parked: bool = False,
) -> RunResultBase | None:
await coordinator.attach_runtime(
agent_id,
session=session,
interrupt_on_message=interactive,
)
result: RunResultBase | None = None
if not (start_parked and interactive):
result = await _run_cycle(
agent,
coordinator,
agent_id,
input_data=initial_input,
run_config=run_config,
context=context,
max_turns=max_turns,
session=session,
interactive=interactive,
)
if not interactive:
return result
while True:
try:
await coordinator.wait_for_message(agent_id)
except asyncio.CancelledError:
return result
await coordinator.consume_pending(agent_id)
result = await _run_cycle(
agent,
coordinator,
agent_id,
input_data=[],
run_config=run_config,
context=context,
max_turns=max_turns,
session=session,
interactive=interactive,
)
async def _run_cycle(
agent: Any,
coordinator: AgentCoordinator,
agent_id: str,
*,
input_data: Any,
run_config: RunConfig,
context: dict[str, Any],
max_turns: int,
session: Session | None,
interactive: bool,
) -> RunResultBase | None:
try:
await coordinator.mark_running(agent_id)
stream = Runner.run_streamed(
agent,
input=input_data,
run_config=run_config,
context=context,
max_turns=max_turns,
session=session,
)
await coordinator.attach_stream(agent_id, stream)
try:
async for _event in stream.stream_events():
pass
if stream.run_loop_exception is not None:
raise stream.run_loop_exception
finally:
await coordinator.detach_stream(agent_id, stream)
except Exception as exc:
if not interactive:
raise
if isinstance(exc, MaxTurnsExceeded):
status: Status = "stopped"
elif isinstance(exc, UserError | AgentsException | APIError):
status = "failed"
else:
status = "crashed"
logger.exception("agent run failed for %s; parking as %s", agent_id, status)
await coordinator.set_status(agent_id, status)
await _notify_parent_on_crash(coordinator, agent_id, status)
return None
else:
await _settle_run_result(coordinator, agent_id, interactive)
return stream
async def _settle_run_result(
coordinator: AgentCoordinator,
agent_id: str,
interactive: bool,
) -> None:
async with coordinator._lock:
current_status = coordinator.statuses.get(agent_id)
if current_status != "running":
return
status: Status = "waiting" if interactive else "crashed"
await coordinator.set_status(agent_id, status)
await _notify_parent_on_crash(coordinator, agent_id, status)
async def _notify_parent_on_crash(
coordinator: AgentCoordinator,
agent_id: str,
status: str,
) -> None:
if status != "crashed":
return
async with coordinator._lock:
parent = coordinator.parent_of.get(agent_id)
name = coordinator.names.get(agent_id, agent_id)
if parent is None:
return
await coordinator.send(
parent,
{
"from": agent_id,
"type": "crash",
"priority": "high",
"content": (
f"[Agent crash] {name} ({agent_id}) terminated unexpectedly. "
"Stop waiting on this child unless you want to message it again."
),
},
)
async def _spawn_child_agent(
*,
coordinator: AgentCoordinator,
factory: Any,
agents_db_path: Path,
sessions_to_close: list[SQLiteSession],
run_config: RunConfig,
max_turns: int,
interactive: bool,
parent_ctx: dict[str, Any],
name: str,
task: str,
skills: list[str],
parent_history: list[Any],
) -> dict[str, Any]:
parent_id = parent_ctx.get("agent_id")
if not isinstance(parent_id, str):
raise TypeError("Parent agent_id missing from context")
child_id = uuid.uuid4().hex[:8]
child_agent = factory(name=name, skills=skills)
await coordinator.register(
child_id,
name,
parent_id,
task=task,
skills=skills,
)
await _start_child_runner(
parent_ctx=parent_ctx,
coordinator=coordinator,
agents_db_path=agents_db_path,
sessions_to_close=sessions_to_close,
run_config=run_config,
max_turns=max_turns,
interactive=interactive,
child_agent=child_agent,
child_id=child_id,
name=name,
parent_id=parent_id,
task=task,
initial_input=child_initial_input(
name=name,
child_id=child_id,
parent_id=parent_id,
task=task,
parent_history=parent_history,
),
)
return {
"success": True,
"agent_id": child_id,
"name": name,
"parent_id": parent_id,
"message": f"Spawned '{name}' ({child_id}) running in parallel.",
}
async def _start_child_runner(
*,
parent_ctx: dict[str, Any],
coordinator: AgentCoordinator,
agents_db_path: Path,
sessions_to_close: list[SQLiteSession],
run_config: RunConfig,
max_turns: int,
interactive: bool,
child_agent: Any,
child_id: str,
name: str,
parent_id: str | None,
task: str,
initial_input: Any,
start_parked: bool = False,
) -> None:
session = _open_agent_session(child_id, agents_db_path)
sessions_to_close.append(session)
await coordinator.attach_runtime(child_id, session=session)
child_ctx: dict[str, Any] = dict(parent_ctx)
child_ctx["agent_id"] = child_id
child_ctx["parent_id"] = parent_id
child_ctx["task"] = task
task_handle = asyncio.create_task(
_run_agent_loop(
agent=child_agent,
initial_input=initial_input,
run_config=run_config,
context=child_ctx,
max_turns=max_turns,
coordinator=coordinator,
agent_id=child_id,
interactive=interactive,
session=session,
start_parked=start_parked,
),
name=f"agent-{name}-{child_id}",
)
await coordinator.attach_runtime(child_id, task=task_handle)
async def run_strix_scan(
*,
scan_config: dict[str, Any],
scan_id: str | None = None,
image: str,
local_sources: list[dict[str, str]] | None = None,
coordinator: AgentCoordinator | None = None,
interactive: bool = False,
max_turns: int = DEFAULT_MAX_TURNS,
model: str | None = None,
cleanup_on_exit: bool = True,
) -> RunResultBase | None:
"""Run or resume one Strix scan against a sandbox."""
if scan_id is None:
scan_id = f"scan-{uuid.uuid4().hex[:8]}"
# Resolve run_dir before any heavy bring-up so the log file captures
# everything from sandbox start onwards.
run_dir = Path.cwd() / "strix_runs" / scan_id
run_dir.mkdir(parents=True, exist_ok=True)
teardown_logging = setup_scan_logging(run_dir)
set_scan_id(scan_id)
agents_path = run_dir / "agents.json"
agents_db = run_dir / "agents.db"
is_resume = agents_path.exists()
logger.info(
"%s Strix scan %s (image=%s, max_turns=%d, interactive=%s, run_dir=%s)",
"Resuming" if is_resume else "Starting",
scan_id,
image,
max_turns,
interactive,
run_dir,
)
settings = load_settings()
resolved_model = model or settings.llm.model
if not resolved_model:
raise RuntimeError(
"No LLM model configured. Set STRIX_LLM env or pass model= to run_strix_scan().",
)
logger.info("LLM model resolved: %s", resolved_model)
# Caller may pre-create the coordinator so it can route stop/chat
# commands while the scan loop runs in another thread.
if coordinator is None:
coordinator = AgentCoordinator()
coordinator.set_snapshot_path(agents_path)
# Wire the per-agent todo store to ``{run_dir}/todos.json`` (mirrored
# on every CRUD) and reload any prior todos so respawned subagents
# find their lists intact. Same for the shared notes store.
from strix.tools.notes.tools import hydrate_notes_from_disk
from strix.tools.todo.tools import hydrate_todos_from_disk
hydrate_todos_from_disk(run_dir)
hydrate_notes_from_disk(run_dir)
root_id: str | None = None
if is_resume:
try:
snap = json.loads(agents_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
raise RuntimeError(
f"Cannot resume scan {scan_id}: agents.json is unreadable: {exc}",
) from exc
if not agents_db.exists():
raise RuntimeError(
f"Cannot resume scan {scan_id}: missing SDK session database at {agents_db}",
)
await coordinator.restore(snap)
for aid, parent in coordinator.parent_of.items():
if parent is None:
root_id = aid
break
if root_id is None:
raise RuntimeError(
f"Cannot resume scan {scan_id}: agents.json has no root agent (parent=None)",
)
logger.info(
"Resume: restored coordinator with %d agent(s); root=%s",
len(coordinator.statuses),
root_id,
)
else:
root_id = uuid.uuid4().hex[:8]
logger.info("Bringing up sandbox session for scan %s", scan_id)
bundle = await session_manager.create_or_reuse(
scan_id,
image=image,
local_sources=local_sources or [],
)
logger.info("Sandbox ready for scan %s", scan_id)
sessions_to_close: list[SQLiteSession] = []
try:
targets = scan_config.get("targets") or []
scan_mode = str(scan_config.get("scan_mode") or "deep")
is_whitebox = any(t.get("type") == "local_code" for t in targets)
skills = list(scan_config.get("skills") or [])
root_task = build_root_task(scan_config)
model_settings = make_model_settings(settings.llm.reasoning_effort)
run_config = RunConfig(
model=resolved_model,
model_provider=build_multi_provider(),
model_settings=model_settings,
sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]),
trace_include_sensitive_data=False,
)
scope_context = build_scope_context(scan_config)
root_agent = build_strix_agent(
name="strix",
skills=skills,
is_root=True,
scan_mode=scan_mode,
is_whitebox=is_whitebox,
interactive=interactive,
system_prompt_context=scope_context,
)
if not is_resume:
await coordinator.register(
root_id,
"strix",
parent_id=None,
task=root_task,
skills=skills,
)
agent_factory = make_child_factory(
scan_mode=scan_mode,
is_whitebox=is_whitebox,
interactive=interactive,
system_prompt_context=scope_context,
)
async def spawn_child_agent(**kwargs: Any) -> dict[str, Any]:
return await _spawn_child_agent(
coordinator=coordinator,
factory=agent_factory,
agents_db_path=agents_db,
sessions_to_close=sessions_to_close,
run_config=run_config,
max_turns=max_turns,
interactive=interactive,
**kwargs,
)
context: dict[str, Any] = {
"coordinator": coordinator,
"sandbox_session": bundle["session"],
"caido_client": bundle["caido_client"],
"agent_id": root_id,
"parent_id": None,
"interactive": interactive,
"spawn_child_agent": spawn_child_agent,
}
# All agents share one SQLite database; SDK session_id separates
# each agent's conversation inside that database.
root_session = _open_agent_session(root_id, agents_db)
sessions_to_close.append(root_session)
await coordinator.attach_runtime(root_id, session=root_session)
if is_resume:
await _respawn_subagents(
coordinator=coordinator,
factory=agent_factory,
agents_db_path=agents_db,
sessions_to_close=sessions_to_close,
run_config=run_config,
max_turns=max_turns,
interactive=interactive,
parent_ctx=context,
root_id=root_id,
)
initial_input: Any = [] if is_resume else root_task
# Resume + new ``--instruction``: SDK replay drives root from
# agents.db with ``initial_input=[]``, so a brand-new instruction
# passed on the resume CLI would otherwise be silently ignored.
# Inject it as a fresh user message in root's SDK session; the
# next run cycle will replay it with the rest of the session.
resume_instruction = str(scan_config.get("resume_instruction") or "").strip()
if is_resume and resume_instruction:
await coordinator.send(
root_id,
{
"from": "user",
"type": "instruction",
"priority": "high",
"content": resume_instruction,
},
)
logger.info(
"Resume: injected new instruction into root SDK session (len=%d)",
len(resume_instruction),
)
async with coordinator._lock:
root_status = coordinator.statuses.get(root_id)
return await _run_agent_loop(
agent=root_agent,
initial_input=initial_input,
run_config=run_config,
context=context,
max_turns=max_turns,
coordinator=coordinator,
agent_id=root_id,
interactive=interactive,
session=root_session,
start_parked=bool(interactive and is_resume and root_status != "running"),
)
except BaseException:
logger.exception("Strix scan %s failed", scan_id)
# Cancel any descendant tasks the root spawned before unwinding.
# cancel_descendants is idempotent and handles the empty-tree case.
if root_id is not None:
await coordinator.cancel_descendants(root_id)
# The SDK's on_agent_end hook only fires after a successful
# ``Runner.run_streamed`` reaches the agent's first turn. A
# failure earlier (e.g., model-provider routing, sandbox
# bring-up) leaves the root stuck at status="running" — the
# TUI keeps animating "Initializing" forever. Finalize it
# here so the coordinator reflects reality.
with contextlib.suppress(Exception):
await coordinator.set_status(root_id, "failed")
raise
finally:
for s in sessions_to_close:
with contextlib.suppress(Exception):
s.close()
with contextlib.suppress(Exception):
await coordinator._maybe_snapshot()
if cleanup_on_exit:
logger.info("Tearing down sandbox session for scan %s", scan_id)
await session_manager.cleanup(scan_id)
logger.info("Strix scan %s done", scan_id)
teardown_logging()
async def _respawn_subagents(
*,
coordinator: AgentCoordinator,
factory: Any,
agents_db_path: Path,
sessions_to_close: list[SQLiteSession],
run_config: RunConfig,
max_turns: int,
interactive: bool,
parent_ctx: dict[str, Any],
root_id: str,
) -> None:
"""Re-spawn subagent runners from a restored coordinator snapshot.
Each child gets its own SDK ``session_id`` inside the shared
``agents.db`` so the SDK replays its prior conversation. Interactive
mode respawns every registered child as a parked runner unless it was
actively running before the crash. That keeps completed/stopped/
crashed/failed agents addressable: a later message wakes the SDK
session instead of being dropped into a dead inbox.
"""
async with coordinator._lock:
# Snapshot the iteration view first so we can mutate via coordinator
# below without "dict changed during iteration" trouble.
agents_snapshot = [
(aid, status, dict(coordinator.metadata.get(aid, {})))
for aid, status in coordinator.statuses.items()
]
candidates: list[tuple[str, str, str | None, dict[str, Any]]] = []
for aid, status, md in agents_snapshot:
if not interactive and status not in {"running", "waiting"}:
continue
if coordinator.parent_of.get(aid) is None or aid == root_id:
continue
md["_restored_status"] = status
candidates.append(
(
aid,
coordinator.names.get(aid, aid),
coordinator.parent_of.get(aid),
md,
)
)
for child_id, name, parent_id, md in candidates:
try:
restored_status = str(md.get("_restored_status") or "running")
start_parked = interactive and restored_status != "running"
if start_parked:
logger.warning(
"respawn %s (%s): starting parked from status=%s",
child_id,
name,
restored_status,
)
child_skills = list(md.get("skills") or [])
child_agent = factory(name=name, skills=child_skills)
await _start_child_runner(
parent_ctx=parent_ctx,
coordinator=coordinator,
agents_db_path=agents_db_path,
sessions_to_close=sessions_to_close,
run_config=run_config,
max_turns=max_turns,
interactive=interactive,
child_agent=child_agent,
child_id=child_id,
name=name,
parent_id=parent_id,
task=str(md.get("task", "")),
initial_input=[],
start_parked=start_parked,
)
logger.info(
"respawned %s (%s) parent=%s task_len=%d",
child_id,
name,
parent_id or "-",
len(md.get("task", "")),
)
except Exception:
logger.exception("respawn %s failed; marking crashed", child_id)
with contextlib.suppress(Exception):
await coordinator.set_status(child_id, "crashed")
-645
View File
@@ -1,645 +0,0 @@
"""Top-level scan entry point with auto-resume.
1. Build (or take from caller) the per-scan ``AgentCoordinator``.
2. Wire a snapshot path so lifecycle events auto-persist ``agents.json``.
3. Acquire an advisory file lock so a second ``strix`` process can't run
on the same ``scan_id`` concurrently.
4. **Resume detection**: if ``{run_dir}/agents.json`` already exists, restore
the coordinator, hydrate the tracer, reuse the persisted ``root_id`` instead
of generating a fresh one, and respawn every non-terminal subagent
from the shared SDK ``agents.db`` before starting the root.
5. Bring up (or reuse) a sandbox session for ``scan_id``.
6. Build the root ``Agent`` + child factory.
7. Open root ``SQLiteSession`` in ``agents.db`` so the SDK replays prior
turns on resume.
8. Call ``Runner.run`` (via ``run_with_continuation``).
9. ``finally``: close every per-agent session, take a final snapshot,
tear down the sandbox, release the lock.
Resume is **always on**: there is no flag — presence of ``agents.json`` is
the trigger. Fresh runs simply have no ``agents.json`` to begin with.
"""
from __future__ import annotations
import asyncio
import contextlib
import json
import logging
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
from agents import RunConfig
from agents.memory import SQLiteSession
from agents.model_settings import ModelSettings
from agents.sandbox import SandboxRunConfig
from openai.types.shared import Reasoning
from strix.agents.factory import build_strix_agent, make_child_factory
from strix.config import load_settings
from strix.llm.multi_provider_setup import build_multi_provider
from strix.llm.retry import DEFAULT_RETRY
from strix.orchestration.coordinator import (
AgentCoordinator,
open_agent_session,
run_with_continuation,
)
from strix.runtime import session_manager
from strix.telemetry.logging import set_agent_id, set_scan_id, setup_scan_logging
#: Default ``max_turns`` budget passed to ``Runner.run``.
_MAX_TURNS = 300
if TYPE_CHECKING:
from agents.result import RunResultBase
logger = logging.getLogger(__name__)
def _build_root_task(scan_config: dict[str, Any]) -> str:
"""Format the user-facing task for the root agent.
Collects each target type into a labelled section, appends
diff-scope context if active, and tacks on user_instructions. The
structured section headers are referenced by the system prompt
template, so the shape matters for prompt parity.
"""
targets = scan_config.get("targets", []) or []
diff_scope = scan_config.get("diff_scope") or {}
user_instructions = scan_config.get("user_instructions", "") or ""
repos: list[str] = []
locals_: list[str] = []
urls: list[str] = []
ips: list[str] = []
for target in targets:
ttype = target.get("type")
details = target.get("details") or {}
workspace_subdir = details.get("workspace_subdir")
workspace_path = f"/workspace/{workspace_subdir}" if workspace_subdir else "/workspace"
if ttype == "repository":
url = details.get("target_repo", "")
cloned = details.get("cloned_repo_path")
repos.append(
f"- {url} (available at: {workspace_path})" if cloned else f"- {url}",
)
elif ttype == "local_code":
path = details.get("target_path", "unknown")
locals_.append(f"- {path} (available at: {workspace_path})")
elif ttype == "web_application":
urls.append(f"- {details.get('target_url', '')}")
elif ttype == "ip_address":
ips.append(f"- {details.get('target_ip', '')}")
parts: list[str] = []
if repos:
parts.append("\n\nRepositories:")
parts.extend(repos)
if locals_:
parts.append("\n\nLocal Codebases:")
parts.extend(locals_)
if urls:
parts.append("\n\nURLs:")
parts.extend(urls)
if ips:
parts.append("\n\nIP Addresses:")
parts.extend(ips)
if diff_scope.get("active"):
parts.append("\n\nScope Constraints:")
parts.append(
"- Pull request diff-scope mode is active. Prioritize changed files "
"and use other files only for context.",
)
for repo_scope in diff_scope.get("repos", []) or []:
label = (
repo_scope.get("workspace_subdir") or repo_scope.get("source_path") or "repository"
)
changed = repo_scope.get("analyzable_files_count", 0)
deleted = repo_scope.get("deleted_files_count", 0)
parts.append(f"- {label}: {changed} changed file(s) in primary scope")
if deleted:
parts.append(f"- {label}: {deleted} deleted file(s) are context-only")
task = " ".join(parts)
if user_instructions:
task = f"{task}\n\nSpecial instructions: {user_instructions}"
return task
def _build_scope_context(scan_config: dict[str, Any]) -> dict[str, Any]:
"""Produce the system_prompt_context block used by the prompt template.
The prompt template's ``system_prompt_context.authorized_targets``
lookups expect this exact shape.
"""
authorized: list[dict[str, str]] = []
for target in scan_config.get("targets", []) or []:
ttype = target.get("type", "unknown")
details = target.get("details") or {}
if ttype == "repository":
value = details.get("target_repo", "")
elif ttype == "local_code":
value = details.get("target_path", "")
elif ttype == "web_application":
value = details.get("target_url", "")
elif ttype == "ip_address":
value = details.get("target_ip", "")
else:
value = target.get("original", "")
workspace_subdir = details.get("workspace_subdir")
workspace_path = f"/workspace/{workspace_subdir}" if workspace_subdir else ""
authorized.append(
{"type": ttype, "value": value, "workspace_path": workspace_path},
)
return {
"scope_source": "system_scan_config",
"authorization_source": "strix_platform_verified_targets",
"authorized_targets": authorized,
"user_instructions_do_not_expand_scope": True,
}
async def run_strix_scan(
*,
scan_config: dict[str, Any],
scan_id: str | None = None,
image: str,
local_sources: list[dict[str, str]] | None = None,
tracer: Any | None = None,
coordinator: AgentCoordinator | None = None,
interactive: bool = False,
max_turns: int = _MAX_TURNS,
model: str | None = None,
cleanup_on_exit: bool = True,
) -> RunResultBase | None:
"""Run one Strix scan end-to-end against a freshly-prepared sandbox.
Args:
scan_config: Per-scan configuration — ``targets``,
``user_instructions``, ``diff_scope``, ``scan_mode``,
``skills``. ``is_whitebox`` is derived from ``targets``.
scan_id: Used to key the sandbox session cache. Auto-generated
if omitted — callers that want resume-after-crash semantics
should pass a stable id.
image: Docker image tag for the sandbox (e.g.
``"strix-sandbox:0.2.0"``).
local_sources: Per-source mount specs from
:func:`strix.interface.utils.collect_local_sources` —
each entry's ``source_path`` (host) is bind-mounted at
``/workspace/<workspace_subdir>``. Pass ``None`` (or ``[]``)
for non-whitebox runs.
tracer: Optional Strix tracer. Stored in context for the
telemetry hook chain. Pass ``None`` for unit tests.
interactive: Renders the interactive-mode prompt block on the
root agent.
max_turns: Cap on root-agent LLM turns (default 300).
model: Litellm model alias. ``None`` (default) reads
:attr:`Settings.llm.model` — caller pre-validates via
:func:`validate_environment` that it's set.
cleanup_on_exit: When True (default), tears down the sandbox
session in a ``finally``. Set to False for resume scenarios
where the caller wants to preserve the container.
Returns the SDK ``RunResult`` from ``Runner.run``. Raises if the
sandbox bring-up fails or the run itself raises.
"""
if scan_id is None:
scan_id = f"scan-{uuid.uuid4().hex[:8]}"
# Resolve run_dir before any heavy bring-up so the log file captures
# everything from sandbox start onwards. Tracer (if present) owns the
# canonical path; otherwise fall back to ``./strix_runs/<scan_id>``.
run_dir = (
tracer.get_run_dir()
if tracer is not None and hasattr(tracer, "get_run_dir")
else Path.cwd() / "strix_runs" / scan_id
)
run_dir.mkdir(parents=True, exist_ok=True)
teardown_logging = setup_scan_logging(run_dir)
set_scan_id(scan_id)
agents_path = run_dir / "agents.json"
agents_db = run_dir / "agents.db"
is_resume = agents_path.exists()
lock_handle = _acquire_run_lock(run_dir)
logger.info(
"%s Strix scan %s (image=%s, max_turns=%d, interactive=%s, run_dir=%s)",
"Resuming" if is_resume else "Starting",
scan_id,
image,
max_turns,
interactive,
run_dir,
)
resolved_model = model or load_settings().llm.model
if not resolved_model:
_release_run_lock(lock_handle)
raise RuntimeError(
"No LLM model configured. Set STRIX_LLM env or pass model= to run_strix_scan().",
)
logger.info("LLM model resolved: %s", resolved_model)
# Caller may pre-create the coordinator so it can route stop/chat
# commands while the scan loop runs in another thread.
if coordinator is None:
coordinator = AgentCoordinator()
coordinator.set_snapshot_path(agents_path)
if tracer is not None and hasattr(tracer, "hydrate_from_run_dir"):
tracer.hydrate_from_run_dir()
# Wire the per-agent todo store to ``{run_dir}/todos.json`` (mirrored
# on every CRUD) and reload any prior todos so respawned subagents
# find their lists intact. Same for the shared notes store.
from strix.tools.notes.tools import hydrate_notes_from_disk
from strix.tools.todo.tools import hydrate_todos_from_disk
hydrate_todos_from_disk(run_dir)
hydrate_notes_from_disk(run_dir)
root_id: str | None = None
if is_resume:
try:
snap = json.loads(agents_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
_release_run_lock(lock_handle)
raise RuntimeError(
f"Cannot resume scan {scan_id}: agents.json is unreadable: {exc}",
) from exc
if not agents_db.exists():
_release_run_lock(lock_handle)
raise RuntimeError(
f"Cannot resume scan {scan_id}: missing SDK session database at {agents_db}",
)
await coordinator.restore(snap)
for aid, parent in coordinator.parent_of.items():
if parent is None:
root_id = aid
break
if root_id is None:
_release_run_lock(lock_handle)
raise RuntimeError(
f"Cannot resume scan {scan_id}: agents.json has no root agent (parent=None)",
)
logger.info(
"Resume: restored coordinator with %d agent(s); root=%s",
len(coordinator.statuses),
root_id,
)
else:
root_id = uuid.uuid4().hex[:8]
logger.info("Bringing up sandbox session for scan %s", scan_id)
bundle = await session_manager.create_or_reuse(
scan_id,
image=image,
local_sources=local_sources or [],
)
logger.info("Sandbox ready for scan %s", scan_id)
sessions_to_close: list[SQLiteSession] = []
try:
# Lazy: ``strix.interface`` pulls cli→tui→scan which would cycle.
from strix.interface.utils import is_whitebox_scan
scan_mode = str(scan_config.get("scan_mode") or "deep")
is_whitebox = is_whitebox_scan(scan_config.get("targets") or [])
skills = list(scan_config.get("skills") or [])
diff_scope = scan_config.get("diff_scope") or None
run_id = scan_config.get("run_id") or scan_id
scope_context = _build_scope_context(scan_config)
root_agent = build_strix_agent(
name="strix",
skills=skills,
is_root=True,
scan_mode=scan_mode,
is_whitebox=is_whitebox,
interactive=interactive,
system_prompt_context=scope_context,
)
if not is_resume:
await coordinator.register(
root_id,
"strix",
parent_id=None,
task=_build_root_task(scan_config),
skills=skills,
is_whitebox=is_whitebox,
scan_mode=scan_mode,
diff_scope=diff_scope,
)
agent_factory = make_child_factory(
scan_mode=scan_mode,
is_whitebox=is_whitebox,
interactive=interactive,
system_prompt_context=scope_context,
)
context: dict[str, Any] = {
"coordinator": coordinator,
"sandbox_session": bundle["session"],
"sandbox_client": bundle["client"],
"caido_client": bundle["caido_client"],
"agent_id": root_id,
"parent_id": None,
"tracer": tracer,
"model": resolved_model,
"model_settings": None,
"max_turns": max_turns,
"agent_finish_called": False,
"is_whitebox": is_whitebox,
"interactive": interactive,
"scan_mode": scan_mode,
"diff_scope": diff_scope,
"run_id": run_id,
"agent_factory": agent_factory,
"agents_db_path": agents_db,
"_sessions_to_close": sessions_to_close,
}
reasoning_effort: Literal["low", "medium", "high"] | None = (
load_settings().llm.reasoning_effort
)
model_settings = ModelSettings(
parallel_tool_calls=False,
tool_choice="required",
retry=DEFAULT_RETRY,
)
if reasoning_effort is not None:
model_settings = model_settings.resolve(
ModelSettings(reasoning=Reasoning(effort=reasoning_effort)),
)
run_config = RunConfig(
model=resolved_model,
model_provider=build_multi_provider(),
model_settings=model_settings,
sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]),
tracing_disabled=False,
trace_include_sensitive_data=False,
)
if is_resume:
await _respawn_subagents(
coordinator=coordinator,
agents_db_path=agents_db,
factory=agent_factory,
parent_ctx=context,
resolved_model=resolved_model,
reasoning_effort=reasoning_effort,
root_id=root_id,
sessions_to_close=sessions_to_close,
)
# All agents share one SQLite database; SDK session_id separates
# each agent's conversation inside that database.
root_session = open_agent_session(root_id, agents_db)
sessions_to_close.append(root_session)
await coordinator.attach_runtime(root_id, session=root_session)
initial_input: Any = [] if is_resume else _build_root_task(scan_config)
# Resume + new ``--instruction``: SDK replay drives root from
# agents.db with ``initial_input=[]``, so a brand-new instruction
# passed on the resume CLI would otherwise be silently ignored.
# Inject it as a fresh user message in root's SDK session; the
# next run cycle will replay it with the rest of the session.
resume_instruction = str(scan_config.get("resume_instruction") or "").strip()
if is_resume and resume_instruction:
await coordinator.send(
root_id,
{
"from": "user",
"type": "instruction",
"priority": "high",
"content": resume_instruction,
},
)
logger.info(
"Resume: injected new instruction into root SDK session (len=%d)",
len(resume_instruction),
)
async with coordinator._lock:
root_status = coordinator.statuses.get(root_id)
return await run_with_continuation(
agent=root_agent,
initial_input=initial_input,
run_config=run_config,
context=context,
max_turns=max_turns,
coordinator=coordinator,
agent_id=root_id,
interactive=interactive,
session=root_session,
start_parked=bool(interactive and is_resume and root_status != "running"),
)
except BaseException as exc:
logger.exception("Strix scan %s failed", scan_id)
# Cancel any descendant tasks the root spawned before unwinding.
# cancel_descendants is idempotent and handles the empty-tree case.
if root_id is not None:
await coordinator.cancel_descendants(root_id)
# The SDK's on_agent_end hook only fires after a successful
# ``Runner.run_streamed`` reaches the agent's first turn. A
# failure earlier (e.g., model-provider routing, sandbox
# bring-up) leaves the root stuck at status="running" — the
# TUI keeps animating "Initializing" forever. Finalize it
# here so the coordinator + tracer reflect reality, and stash the
# error message for the status-line display.
error_message = f"{type(exc).__name__}: {exc}"
if tracer is not None and root_id in getattr(tracer, "agents", {}):
tracer.agents[root_id]["status"] = "failed"
tracer.agents[root_id]["error_message"] = error_message
tracer.agents[root_id]["updated_at"] = datetime.now(UTC).isoformat()
with contextlib.suppress(Exception):
await coordinator.set_status(root_id, "failed")
set_agent_id(None)
raise
finally:
for s in sessions_to_close:
with contextlib.suppress(Exception):
s.close()
with contextlib.suppress(Exception):
await coordinator._maybe_snapshot()
if cleanup_on_exit:
logger.info("Tearing down sandbox session for scan %s", scan_id)
await session_manager.cleanup(scan_id)
_release_run_lock(lock_handle)
logger.info("Strix scan %s done", scan_id)
teardown_logging()
async def _respawn_subagents(
*,
coordinator: AgentCoordinator,
agents_db_path: Path,
factory: Any,
parent_ctx: dict[str, Any],
resolved_model: str,
reasoning_effort: Literal["low", "medium", "high"] | None,
root_id: str,
sessions_to_close: list[SQLiteSession],
) -> None:
"""Re-spawn subagent runners from a restored coordinator snapshot.
Each child gets its own SDK ``session_id`` inside the shared
``agents.db`` so the SDK replays its prior conversation. Interactive
mode respawns every registered child as a parked runner unless it was
actively running before the crash. That keeps completed/stopped/
crashed/failed agents addressable: a later message wakes the SDK
session instead of being dropped into a dead inbox.
"""
interactive = bool(parent_ctx.get("interactive", False))
async with coordinator._lock:
# Snapshot the iteration view first so we can mutate via coordinator
# below without "dict changed during iteration" trouble.
agents_snapshot = [
(aid, status, dict(coordinator.metadata.get(aid, {})))
for aid, status in coordinator.statuses.items()
]
candidates: list[tuple[str, str, str | None, dict[str, Any]]] = []
for aid, status, md in agents_snapshot:
if not interactive and status not in {"running", "waiting", "llm_failed"}:
continue
if coordinator.parent_of.get(aid) is None or aid == root_id:
continue
md["_restored_status"] = status
candidates.append(
(
aid,
coordinator.names.get(aid, aid),
coordinator.parent_of.get(aid),
md,
)
)
for child_id, name, parent_id, md in candidates:
try:
restored_status = str(md.get("_restored_status") or "running")
start_parked = interactive and restored_status != "running"
if start_parked:
logger.warning(
"respawn %s (%s): starting parked from status=%s",
child_id,
name,
restored_status,
)
child_session = open_agent_session(child_id, agents_db_path)
sessions_to_close.append(child_session)
await coordinator.attach_runtime(child_id, session=child_session)
child_skills = list(md.get("skills") or [])
child_agent = factory(name=name, skills=child_skills)
child_ctx: dict[str, Any] = dict(parent_ctx)
child_ctx["agent_id"] = child_id
child_ctx["parent_id"] = parent_id
child_ctx["agent_finish_called"] = False
child_ctx["task"] = md.get("task", "")
child_model_settings = ModelSettings(
parallel_tool_calls=False,
tool_choice="required",
retry=DEFAULT_RETRY,
)
if reasoning_effort is not None:
child_model_settings = child_model_settings.resolve(
ModelSettings(reasoning=Reasoning(effort=reasoning_effort)),
)
child_run_config = RunConfig(
model=resolved_model,
model_provider=build_multi_provider(),
model_settings=child_model_settings,
sandbox=SandboxRunConfig(
client=parent_ctx["sandbox_client"],
session=parent_ctx["sandbox_session"],
),
tracing_disabled=False,
trace_include_sensitive_data=False,
)
task_handle = asyncio.create_task(
run_with_continuation(
agent=child_agent,
initial_input=[],
run_config=child_run_config,
context=child_ctx,
max_turns=int(parent_ctx.get("max_turns", 300)),
coordinator=coordinator,
agent_id=child_id,
interactive=bool(parent_ctx.get("interactive", False)),
session=child_session,
start_parked=start_parked,
),
name=f"agent-{name}-{child_id}",
)
await coordinator.attach_runtime(child_id, task=task_handle)
logger.info(
"respawned %s (%s) parent=%s task_len=%d",
child_id,
name,
parent_id or "-",
len(md.get("task", "")),
)
except Exception:
logger.exception("respawn %s failed; marking crashed", child_id)
with contextlib.suppress(Exception):
await coordinator.set_status(child_id, "crashed")
def _acquire_run_lock(run_dir: Path) -> Any:
"""Take an exclusive flock on ``{run_dir}/.lock`` so two strix processes
can't run on the same scan_id concurrently. Raises ``RuntimeError`` if
another holder is detected. Best-effort on platforms without ``fcntl``.
"""
lock_path = run_dir / ".lock"
try:
import fcntl
except ImportError:
return None
handle = lock_path.open("a+", encoding="utf-8")
try:
fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
except OSError as exc:
handle.close()
raise RuntimeError(
f"Another strix process appears to be running on this scan "
f"(could not acquire lock at {lock_path}). Aborting.",
) from exc
return handle
def _release_run_lock(handle: Any) -> None:
if handle is None:
return
try:
import fcntl
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
except (ImportError, OSError):
pass
finally:
with contextlib.suppress(Exception):
handle.close()
+157
View File
@@ -0,0 +1,157 @@
"""Pure helper builders for Strix scan orchestration."""
from __future__ import annotations
import json
from typing import Any, Literal
from agents.model_settings import ModelSettings
from openai.types.shared import Reasoning
from strix.llm.retry import DEFAULT_RETRY
# Default max_turns budget passed to the SDK runner.
DEFAULT_MAX_TURNS = 300
def build_root_task(scan_config: dict[str, Any]) -> str:
"""Format the user-facing task for the root agent."""
targets = scan_config.get("targets", []) or []
diff_scope = scan_config.get("diff_scope") or {}
user_instructions = scan_config.get("user_instructions", "") or ""
sections: dict[str, list[str]] = {
"Repositories": [],
"Local Codebases": [],
"URLs": [],
"IP Addresses": [],
}
for target in targets:
ttype = target.get("type")
details = target.get("details") or {}
workspace_subdir = details.get("workspace_subdir")
workspace_path = f"/workspace/{workspace_subdir}" if workspace_subdir else "/workspace"
if ttype == "repository":
url = details.get("target_repo", "")
cloned = details.get("cloned_repo_path")
sections["Repositories"].append(
f"- {url} (available at: {workspace_path})" if cloned else f"- {url}",
)
elif ttype == "local_code":
path = details.get("target_path", "unknown")
sections["Local Codebases"].append(f"- {path} (available at: {workspace_path})")
elif ttype == "web_application":
sections["URLs"].append(f"- {details.get('target_url', '')}")
elif ttype == "ip_address":
sections["IP Addresses"].append(f"- {details.get('target_ip', '')}")
parts: list[str] = []
for label, items in sections.items():
if items:
parts.append(f"\n\n{label}:")
parts.extend(items)
if diff_scope.get("active"):
parts.append("\n\nScope Constraints:")
parts.append(
"- Pull request diff-scope mode is active. Prioritize changed files "
"and use other files only for context.",
)
for repo_scope in diff_scope.get("repos", []) or []:
label = (
repo_scope.get("workspace_subdir") or repo_scope.get("source_path") or "repository"
)
changed = repo_scope.get("analyzable_files_count", 0)
deleted = repo_scope.get("deleted_files_count", 0)
parts.append(f"- {label}: {changed} changed file(s) in primary scope")
if deleted:
parts.append(f"- {label}: {deleted} deleted file(s) are context-only")
task = " ".join(parts)
if user_instructions:
task = f"{task}\n\nSpecial instructions: {user_instructions}"
return task
def build_scope_context(scan_config: dict[str, Any]) -> dict[str, Any]:
"""Build the system_prompt_context block consumed by the prompt template."""
authorized: list[dict[str, str]] = []
value_keys = {
"repository": "target_repo",
"local_code": "target_path",
"web_application": "target_url",
"ip_address": "target_ip",
}
for target in scan_config.get("targets", []) or []:
ttype = target.get("type", "unknown")
details = target.get("details") or {}
key = value_keys.get(ttype)
value = details.get(key, "") if key is not None else target.get("original", "")
workspace_subdir = details.get("workspace_subdir")
workspace_path = f"/workspace/{workspace_subdir}" if workspace_subdir else ""
authorized.append(
{"type": ttype, "value": value, "workspace_path": workspace_path},
)
return {
"scope_source": "system_scan_config",
"authorization_source": "strix_platform_verified_targets",
"authorized_targets": authorized,
"user_instructions_do_not_expand_scope": True,
}
def make_model_settings(
reasoning_effort: Literal["low", "medium", "high"] | None,
) -> ModelSettings:
model_settings = ModelSettings(
parallel_tool_calls=False,
tool_choice="required",
retry=DEFAULT_RETRY,
)
if reasoning_effort is not None:
model_settings = model_settings.resolve(
ModelSettings(reasoning=Reasoning(effort=reasoning_effort)),
)
return model_settings
def child_initial_input(
*,
name: str,
child_id: str,
parent_id: str,
task: str,
parent_history: list[Any],
) -> list[dict[str, Any]]:
initial_input: list[dict[str, Any]] = []
if parent_history:
rendered = json.dumps(parent_history, ensure_ascii=False, default=str)
initial_input.append(
{
"role": "user",
"content": (
"== Inherited context from parent (background only) ==\n"
f"{rendered}\n"
"== End of inherited context ==\n"
"Use the above as background only; do not continue the "
"parent's work. Your task follows."
),
},
)
initial_input.append(
{
"role": "user",
"content": (
f"You are agent {name} ({child_id}); your parent is {parent_id}. "
"Maintain your own identity. Call agent_finish when your task "
"is complete."
),
}
)
initial_input.append({"role": "user", "content": task})
return initial_input
+1 -1
View File
@@ -67,7 +67,7 @@ async def create_or_reuse(
# Build Manifest entries keyed by ``workspace_subdir`` — the SDK # Build Manifest entries keyed by ``workspace_subdir`` — the SDK
# mounts each at ``/workspace/<key>``, which is exactly the path # mounts each at ``/workspace/<key>``, which is exactly the path
# ``_build_root_task`` puts in the agent's task prompt. Mounting # ``build_root_task`` puts in the agent's task prompt. Mounting
# only the listed source dirs (not their parent) avoids leaking # only the listed source dirs (not their parent) avoids leaking
# unrelated host content into the sandbox. # unrelated host content into the sandbox.
entries: dict[str | Path, BaseEntry] = {} entries: dict[str | Path, BaseEntry] = {}
+1 -1
View File
@@ -16,7 +16,7 @@ We collect only very **basic** usage data including:
**System Context:** OS type, architecture, Strix version\ **System Context:** OS type, architecture, Strix version\
**Scan Context:** Scan mode (quick/standard/deep), scan type (whitebox/blackbox)\ **Scan Context:** Scan mode (quick/standard/deep), scan type (whitebox/blackbox)\
**Model Usage:** Which LLM model is being used (not prompts or responses)\ **Model Usage:** Which LLM model is being used (not prompts or responses)\
**Aggregate Metrics:** Vulnerability counts by severity, agent/tool counts, token usage and cost estimates **Aggregate Metrics:** Vulnerability counts by severity
For complete transparency, you can inspect our [telemetry implementation](https://github.com/usestrix/strix/blob/main/strix/telemetry/posthog.py) to see the exact events we track. For complete transparency, you can inspect our [telemetry implementation](https://github.com/usestrix/strix/blob/main/strix/telemetry/posthog.py) to see the exact events we track.
+8 -4
View File
@@ -1,10 +1,14 @@
from . import posthog from . import posthog
from .tracer import Tracer, get_global_tracer, set_global_tracer from .scan_store import (
ScanStore,
get_global_scan_store,
set_global_scan_store,
)
__all__ = [ __all__ = [
"Tracer", "ScanStore",
"get_global_tracer", "get_global_scan_store",
"posthog", "posthog",
"set_global_tracer", "set_global_scan_store",
] ]
+6 -13
View File
@@ -11,7 +11,7 @@ from strix.config import load_settings
if TYPE_CHECKING: if TYPE_CHECKING:
from strix.telemetry.tracer import Tracer from strix.telemetry.scan_store import ScanStore
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -114,22 +114,19 @@ def finding(severity: str) -> None:
) )
def end(tracer: "Tracer", exit_reason: str = "completed") -> None: def end(scan_store: "ScanStore", exit_reason: str = "completed") -> None:
vulnerabilities_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0} vulnerabilities_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
for v in tracer.vulnerability_reports: for v in scan_store.vulnerability_reports:
sev = v.get("severity", "info").lower() sev = v.get("severity", "info").lower()
if sev in vulnerabilities_counts: if sev in vulnerabilities_counts:
vulnerabilities_counts[sev] += 1 vulnerabilities_counts[sev] += 1
llm = tracer.get_total_llm_stats()
total = llm.get("total", {})
duration = 0.0 duration = 0.0
try: try:
from datetime import datetime from datetime import datetime
start = datetime.fromisoformat(tracer.start_time.replace("Z", "+00:00")) start = datetime.fromisoformat(scan_store.start_time.replace("Z", "+00:00"))
end_iso = tracer.end_time or datetime.now(start.tzinfo).isoformat() end_iso = scan_store.end_time or datetime.now(start.tzinfo).isoformat()
duration = (datetime.fromisoformat(end_iso.replace("Z", "+00:00")) - start).total_seconds() duration = (datetime.fromisoformat(end_iso.replace("Z", "+00:00")) - start).total_seconds()
except (ValueError, TypeError, AttributeError): except (ValueError, TypeError, AttributeError):
pass pass
@@ -140,12 +137,8 @@ def end(tracer: "Tracer", exit_reason: str = "completed") -> None:
**_base_props(), **_base_props(),
"exit_reason": exit_reason, "exit_reason": exit_reason,
"duration_seconds": round(duration), "duration_seconds": round(duration),
"vulnerabilities_total": len(tracer.vulnerability_reports), "vulnerabilities_total": len(scan_store.vulnerability_reports),
**{f"vulnerabilities_{k}": v for k, v in vulnerabilities_counts.items()}, **{f"vulnerabilities_{k}": v for k, v in vulnerabilities_counts.items()},
"agent_count": len(tracer.agents),
"tool_count": tracer.get_real_tool_count(),
"llm_tokens": llm.get("total_tokens", 0),
"llm_cost": total.get("cost", 0.0),
}, },
) )
-234
View File
@@ -1,234 +0,0 @@
"""Per-scan artifact writer.
Writes the customer-facing penetration-test report and per-vulnerability
markdown + a ``vulnerabilities.csv`` index under ``strix_runs/<run>/``,
plus a machine-readable ``vulnerabilities.json`` so a resumed scan's
:class:`~strix.telemetry.tracer.Tracer` can hydrate its in-memory list
back from disk (otherwise vuln-id allocation collides post-restart).
"""
from __future__ import annotations
import csv
import json
import logging
import tempfile
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
logger = logging.getLogger(__name__)
_SEVERITY_ORDER = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
class ScanArtifactWriter:
"""Writes scan artifacts under ``run_dir``. Idempotent on repeat calls.
Tracks which vulnerability ids have already been written so that
re-saves only emit new files; the ``vulnerabilities.csv`` index is
fully rewritten each call so the displayed order stays in sync with
severity sorting.
"""
def __init__(self, run_dir: Path):
self._run_dir = run_dir
self._saved_vuln_ids: set[str] = set()
@property
def run_dir(self) -> Path:
return self._run_dir
def save(
self,
*,
vulnerability_reports: list[dict[str, Any]],
final_scan_result: str | None,
run_metadata: dict[str, Any] | None = None,
) -> None:
"""Write any new vulnerability MDs + rewrite the CSV index +
write the executive penetration-test report if available + dump
``run_metadata.json`` for resume hydration.
Tolerant of OSError / RuntimeError — logs and swallows so a
cleanup failure can't prevent the next scan from finishing.
"""
try:
self._run_dir.mkdir(parents=True, exist_ok=True)
if final_scan_result:
self._write_executive_report(final_scan_result)
if vulnerability_reports:
self._write_vulnerabilities(vulnerability_reports)
if run_metadata is not None:
_atomic_write_text(
self._run_dir / "run_metadata.json",
json.dumps(run_metadata, ensure_ascii=False, indent=2, default=str),
)
logger.info("📊 Essential scan data saved to: %s", self._run_dir)
except (OSError, RuntimeError):
logger.exception("Failed to save scan data")
# --- internals ---------------------------------------------------------
def _write_executive_report(self, body: str) -> None:
path = self._run_dir / "penetration_test_report.md"
with path.open("w", encoding="utf-8") as f:
f.write("# Security Penetration Test Report\n\n")
f.write(f"**Generated:** {datetime.now(UTC).strftime('%Y-%m-%d %H:%M:%S UTC')}\n\n")
f.write(f"{body}\n")
logger.info("Saved final penetration test report to: %s", path)
def _write_vulnerabilities(self, reports: list[dict[str, Any]]) -> None:
vuln_dir = self._run_dir / "vulnerabilities"
vuln_dir.mkdir(exist_ok=True)
new_reports = [r for r in reports if r["id"] not in self._saved_vuln_ids]
for report in new_reports:
(vuln_dir / f"{report['id']}.md").write_text(
_render_vulnerability_md(report),
encoding="utf-8",
)
self._saved_vuln_ids.add(report["id"])
sorted_reports = sorted(
reports,
key=lambda r: (_SEVERITY_ORDER.get(r["severity"], 5), r["timestamp"]),
)
csv_path = self._run_dir / "vulnerabilities.csv"
with csv_path.open("w", encoding="utf-8", newline="") as f:
fieldnames = ["id", "title", "severity", "timestamp", "file"]
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
for report in sorted_reports:
writer.writerow(
{
"id": report["id"],
"title": report["title"],
"severity": report["severity"].upper(),
"timestamp": report["timestamp"],
"file": f"vulnerabilities/{report['id']}.md",
},
)
# JSON index: machine-readable mirror used by ``Tracer.hydrate_from_run_dir``
# so a process restart (resume path in ``orchestration/scan.py``) can
# rebuild ``vulnerability_reports`` and re-establish the next id slot
# before any new ``add_vulnerability_report`` call collides on disk.
_atomic_write_text(
self._run_dir / "vulnerabilities.json",
json.dumps(reports, ensure_ascii=False, indent=2, default=str),
)
if new_reports:
logger.info(
"Saved %d new vulnerability report(s) to: %s",
len(new_reports),
vuln_dir,
)
logger.info("Updated vulnerability index: %s", csv_path)
def _atomic_write_text(path: Path, payload: str) -> None:
"""``tempfile`` + atomic rename so a crash mid-write leaves the prior file."""
path.parent.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
dir=str(path.parent),
prefix=f".{path.name}.",
suffix=".tmp",
delete=False,
) as tmp:
tmp.write(payload)
tmp_path = Path(tmp.name)
tmp_path.replace(path)
def _render_vulnerability_md(report: dict[str, Any]) -> str:
lines: list[str] = [
f"# {report.get('title', 'Untitled Vulnerability')}\n",
f"**ID:** {report.get('id', 'unknown')}",
f"**Severity:** {report.get('severity', 'unknown').upper()}",
f"**Found:** {report.get('timestamp', 'unknown')}",
]
metadata: list[tuple[str, Any]] = [
("Target", report.get("target")),
("Endpoint", report.get("endpoint")),
("Method", report.get("method")),
("CVE", report.get("cve")),
("CWE", report.get("cwe")),
]
cvss = report.get("cvss")
if cvss is not None:
metadata.append(("CVSS", cvss))
for label, value in metadata:
if value:
lines.append(f"**{label}:** {value}")
lines.append("")
lines.append("## Description\n")
lines.append(report.get("description") or "No description provided.")
lines.append("")
if report.get("impact"):
lines.append("## Impact\n")
lines.append(str(report["impact"]))
lines.append("")
if report.get("technical_analysis"):
lines.append("## Technical Analysis\n")
lines.append(str(report["technical_analysis"]))
lines.append("")
if report.get("poc_description") or report.get("poc_script_code"):
lines.append("## Proof of Concept\n")
if report.get("poc_description"):
lines.append(str(report["poc_description"]))
lines.append("")
if report.get("poc_script_code"):
lines.append("```")
lines.append(str(report["poc_script_code"]))
lines.append("```")
lines.append("")
if report.get("code_locations"):
lines.append("## Code Analysis\n")
for i, loc in enumerate(report["code_locations"]):
file_ref = loc.get("file", "unknown")
line_ref = ""
if loc.get("start_line") is not None:
if loc.get("end_line") and loc["end_line"] != loc["start_line"]:
line_ref = f" (lines {loc['start_line']}-{loc['end_line']})"
else:
line_ref = f" (line {loc['start_line']})"
lines.append(f"**Location {i + 1}:** `{file_ref}`{line_ref}")
if loc.get("label"):
lines.append(f" {loc['label']}")
if loc.get("snippet"):
lines.append(f" ```\n {loc['snippet']}\n ```")
if loc.get("fix_before") or loc.get("fix_after"):
lines.append("\n **Suggested Fix:**")
lines.append("```diff")
if loc.get("fix_before"):
for ln in str(loc["fix_before"]).splitlines():
lines.append(f"- {ln}")
if loc.get("fix_after"):
for ln in str(loc["fix_after"]).splitlines():
lines.append(f"+ {ln}")
lines.append("```")
lines.append("")
if report.get("remediation_steps"):
lines.append("## Remediation\n")
lines.append(str(report["remediation_steps"]))
lines.append("")
return "\n".join(lines)
@@ -1,5 +1,7 @@
import csv
import json import json
import logging import logging
import tempfile
from collections.abc import Callable from collections.abc import Callable
from datetime import UTC, datetime from datetime import UTC, datetime
from pathlib import Path from pathlib import Path
@@ -7,32 +9,31 @@ from typing import Any, Optional
from uuid import uuid4 from uuid import uuid4
from strix.telemetry import posthog from strix.telemetry import posthog
from strix.telemetry.scan_artifacts import ScanArtifactWriter
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_global_tracer: Optional["Tracer"] = None _global_scan_store: Optional["ScanStore"] = None
_SEVERITY_ORDER = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
def get_global_tracer() -> Optional["Tracer"]: def get_global_scan_store() -> Optional["ScanStore"]:
return _global_tracer return _global_scan_store
def set_global_tracer(tracer: "Tracer") -> None: def set_global_scan_store(scan_store: "ScanStore") -> None:
global _global_tracer # noqa: PLW0603 global _global_scan_store # noqa: PLW0603
_global_tracer = tracer _global_scan_store = scan_store
class Tracer: class ScanStore:
"""Per-scan in-memory state the TUI renders + per-scan artifact writer. """Per-scan product artifact state plus artifact writer.
Holds live state the TUI reads (chat messages, agent tree, tool The Agents SDK owns model/tool execution, tracing, and conversation
executions, vulnerability reports, LLM usage). Writes vulnerability persistence. This store keeps only Strix-owned scan artifacts and
markdown + CSV + final pentest report to ``strix_runs/<scan>/``. report metadata. Live UI projections belong to the interface layer.
Conversation history goes to the SDK's ``SQLiteSession`` instead; It is not a trace mirror and does not consume SDK tracing processors.
SDK trace events are not persisted here.
""" """
def __init__(self, run_name: str | None = None): def __init__(self, run_name: str | None = None):
@@ -41,23 +42,9 @@ class Tracer:
self.start_time = datetime.now(UTC).isoformat() self.start_time = datetime.now(UTC).isoformat()
self.end_time: str | None = None self.end_time: str | None = None
self.agents: dict[str, dict[str, Any]] = {}
self.tool_executions: dict[int, dict[str, Any]] = {}
self.chat_messages: list[dict[str, Any]] = []
self._next_exec_id = 1
self.vulnerability_reports: list[dict[str, Any]] = [] self.vulnerability_reports: list[dict[str, Any]] = []
self.final_scan_result: str | None = None self.final_scan_result: str | None = None
# LLM usage roll-up across all agents in this run.
self._llm_stats: dict[str, Any] = {
"input_tokens": 0,
"output_tokens": 0,
"cached_tokens": 0,
"cost": 0.0,
"requests": 0,
}
self.scan_results: dict[str, Any] | None = None self.scan_results: dict[str, Any] | None = None
self.scan_config: dict[str, Any] | None = None self.scan_config: dict[str, Any] | None = None
self.run_metadata: dict[str, Any] = { self.run_metadata: dict[str, Any] = {
@@ -69,8 +56,7 @@ class Tracer:
"status": "running", "status": "running",
} }
self._run_dir: Path | None = None self._run_dir: Path | None = None
self._writer: ScanArtifactWriter | None = None self._saved_vuln_ids: set[str] = set()
self._next_message_id = 1
self.caido_url: str | None = None self.caido_url: str | None = None
self.vulnerability_found_callback: Callable[[dict[str, Any]], None] | None = None self.vulnerability_found_callback: Callable[[dict[str, Any]], None] | None = None
@@ -81,7 +67,7 @@ class Tracer:
self.run_metadata["run_name"] = run_name self.run_metadata["run_name"] = run_name
self.run_metadata["run_id"] = run_name self.run_metadata["run_id"] = run_name
self._run_dir = None self._run_dir = None
self._writer = None self._saved_vuln_ids.clear()
def get_run_dir(self) -> Path: def get_run_dir(self) -> Path:
if self._run_dir is None: if self._run_dir is None:
@@ -135,11 +121,7 @@ class Tracer:
if k in {"run_id", "run_name", "start_time", "targets", "status"} if k in {"run_id", "run_name", "start_time", "targets", "status"}
}, },
) )
logger.info( logger.info("scan store hydrated run_metadata from %s", meta_path)
"tracer hydrated run_metadata from %s (start_time=%s)",
meta_path,
self.start_time,
)
json_path = run_dir / "vulnerabilities.json" json_path = run_dir / "vulnerabilities.json"
if json_path.exists(): if json_path.exists():
@@ -156,86 +138,14 @@ class Tracer:
f"vulnerabilities.json at {json_path} is not a list", f"vulnerabilities.json at {json_path} is not a list",
) )
self.vulnerability_reports = [r for r in data if isinstance(r, dict)] self.vulnerability_reports = [r for r in data if isinstance(r, dict)]
writer = self._get_writer()
for r in self.vulnerability_reports: for r in self.vulnerability_reports:
rid = r.get("id") rid = r.get("id")
if isinstance(rid, str): if isinstance(rid, str):
writer._saved_vuln_ids.add(rid) self._saved_vuln_ids.add(rid)
logger.info( logger.info(
"tracer hydrated %d vulnerability report(s) from %s", "scan store hydrated %d vulnerability report(s)", len(self.vulnerability_reports)
len(self.vulnerability_reports),
json_path,
) )
agents_path = run_dir / "agents.json"
if agents_path.exists():
try:
agents_data = json.loads(agents_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError):
# Caller will surface this same corruption via coordinator restore;
# no need to fail twice. Skip the agents/stats hydrate path.
agents_data = None
if isinstance(agents_data, dict):
self._hydrate_agents_tree(agents_data)
self._hydrate_llm_stats(agents_data)
def _hydrate_agents_tree(self, agents_data: dict[str, Any]) -> None:
"""Populate ``self.agents`` from the coordinator snapshot.
Without this, the TUI tree on resume would only show agents
currently running (mirrored by ``on_agent_start``); completed /
crashed / stopped children from the prior run would be invisible
even though the coordinator knows about them.
"""
statuses = agents_data.get("statuses") or {}
names = agents_data.get("names") or {}
parent_of = agents_data.get("parent_of") or {}
if not isinstance(statuses, dict):
return
timestamp = self.start_time
for agent_id, status in statuses.items():
if not isinstance(agent_id, str):
continue
self.agents[agent_id] = {
"id": agent_id,
"name": names.get(agent_id, agent_id) if isinstance(names, dict) else agent_id,
"parent_id": parent_of.get(agent_id) if isinstance(parent_of, dict) else None,
"status": status,
"created_at": timestamp,
"updated_at": timestamp,
}
logger.info("tracer hydrated %d agent(s) into tree", len(self.agents))
def _hydrate_llm_stats(self, agents_data: dict[str, Any]) -> None:
"""Seed ``self._llm_stats`` from the coordinator snapshot's counters.
This keeps the resumed scan's TUI footer cumulative instead of
resetting to zero.
"""
totals = {"input_tokens": 0, "output_tokens": 0, "cached_tokens": 0, "requests": 0}
bucket = agents_data.get("stats_live") or {}
if isinstance(bucket, dict):
for entry in bucket.values():
if isinstance(entry, dict):
totals["input_tokens"] += int(entry.get("in", 0) or 0)
totals["output_tokens"] += int(entry.get("out", 0) or 0)
totals["cached_tokens"] += int(entry.get("cached", 0) or 0)
totals["requests"] += int(entry.get("calls", 0) or 0)
for k, v in totals.items():
self._llm_stats[k] = v
logger.info(
"tracer hydrated llm stats from agents snapshot (in=%d out=%d cached=%d requests=%d)",
totals["input_tokens"],
totals["output_tokens"],
totals["cached_tokens"],
totals["requests"],
)
def _get_writer(self) -> ScanArtifactWriter:
if self._writer is None:
self._writer = ScanArtifactWriter(self.get_run_dir())
return self._writer
def add_vulnerability_report( def add_vulnerability_report(
self, self,
title: str, title: str,
@@ -254,6 +164,8 @@ class Tracer:
cve: str | None = None, cve: str | None = None,
cwe: str | None = None, cwe: str | None = None,
code_locations: list[dict[str, Any]] | None = None, code_locations: list[dict[str, Any]] | None = None,
agent_id: str | None = None,
agent_name: str | None = None,
) -> str: ) -> str:
report_id = f"vuln-{len(self.vulnerability_reports) + 1:04d}" report_id = f"vuln-{len(self.vulnerability_reports) + 1:04d}"
@@ -292,6 +204,10 @@ class Tracer:
report["cwe"] = cwe.strip() report["cwe"] = cwe.strip()
if code_locations: if code_locations:
report["code_locations"] = code_locations report["code_locations"] = code_locations
if agent_id:
report["agent_id"] = agent_id
if agent_name:
report["agent_name"] = agent_name
self.vulnerability_reports.append(report) self.vulnerability_reports.append(report)
logger.info(f"Added vulnerability report: {report_id} - {title}") logger.info(f"Added vulnerability report: {report_id} - {title}")
@@ -343,28 +259,6 @@ class Tracer:
self.save_run_data(mark_complete=True) self.save_run_data(mark_complete=True)
posthog.end(self, exit_reason="finished_by_tool") posthog.end(self, exit_reason="finished_by_tool")
def log_chat_message(
self,
content: str,
role: str,
agent_id: str | None = None,
metadata: dict[str, Any] | None = None,
) -> int:
message_id = self._next_message_id
self._next_message_id += 1
self.chat_messages.append(
{
"message_id": message_id,
"content": content,
"role": role,
"agent_id": agent_id,
"timestamp": datetime.now(UTC).isoformat(),
"metadata": metadata or {},
}
)
return message_id
def set_scan_config(self, config: dict[str, Any]) -> None: def set_scan_config(self, config: dict[str, Any]) -> None:
self.scan_config = config self.scan_config = config
self.run_metadata.update( self.run_metadata.update(
@@ -382,91 +276,180 @@ class Tracer:
self.run_metadata["end_time"] = self.end_time self.run_metadata["end_time"] = self.end_time
self.run_metadata["status"] = "completed" self.run_metadata["status"] = "completed"
self._get_writer().save( self._save_artifacts()
vulnerability_reports=self.vulnerability_reports,
final_scan_result=self.final_scan_result,
run_metadata=dict(self.run_metadata),
)
def log_tool_start(
self,
agent_id: str,
tool_name: str,
args: dict[str, Any] | None = None,
) -> int:
"""Record a tool invocation in flight. Returns an exec_id."""
exec_id = self._next_exec_id
self._next_exec_id += 1
self.tool_executions[exec_id] = {
"agent_id": agent_id,
"tool_name": tool_name,
"args": args or {},
"status": "running",
"result": None,
"timestamp": datetime.now(UTC).isoformat(),
}
return exec_id
def log_tool_end(self, agent_id: str, tool_name: str, result: Any) -> None:
"""Mark the most recent matching exec as completed."""
for exec_id in reversed(self.tool_executions):
entry = self.tool_executions[exec_id]
if (
entry.get("agent_id") == agent_id
and entry.get("tool_name") == tool_name
and entry.get("status") == "running"
):
entry["status"] = "completed"
entry["result"] = result
return
# No matching start (e.g. hooks added later in life) — record as completed.
exec_id = self._next_exec_id
self._next_exec_id += 1
self.tool_executions[exec_id] = {
"agent_id": agent_id,
"tool_name": tool_name,
"status": "completed",
"result": result,
"timestamp": datetime.now(UTC).isoformat(),
}
def get_real_tool_count(self) -> int:
return sum(
1
for exec_data in list(self.tool_executions.values())
if exec_data.get("tool_name") not in ["scan_start_info", "subagent_start_info"]
)
def get_total_llm_stats(self) -> dict[str, Any]:
"""Snapshot the run's aggregated LLM usage."""
stats = self._llm_stats
total = {
"input_tokens": int(stats["input_tokens"]),
"output_tokens": int(stats["output_tokens"]),
"cached_tokens": int(stats["cached_tokens"]),
"cost": round(float(stats["cost"]), 4),
"requests": int(stats["requests"]),
}
return {
"total": total,
"total_tokens": total["input_tokens"] + total["output_tokens"],
}
def record_llm_usage(
self,
*,
input_tokens: int = 0,
output_tokens: int = 0,
cached_tokens: int = 0,
cost: float = 0.0,
requests: int = 1,
) -> None:
"""Accumulate LLM usage from the orchestration hooks."""
self._llm_stats["input_tokens"] += input_tokens
self._llm_stats["output_tokens"] += output_tokens
self._llm_stats["cached_tokens"] += cached_tokens
self._llm_stats["cost"] += cost
self._llm_stats["requests"] += requests
def cleanup(self) -> None: def cleanup(self) -> None:
self.save_run_data(mark_complete=True) self.save_run_data(mark_complete=True)
def _save_artifacts(self) -> None:
"""Write scan artifacts under ``run_dir``."""
run_dir = self.get_run_dir()
try:
run_dir.mkdir(parents=True, exist_ok=True)
if self.final_scan_result:
self._write_executive_report(run_dir)
if self.vulnerability_reports:
self._write_vulnerabilities(run_dir)
_atomic_write_text(
run_dir / "run_metadata.json",
json.dumps(self.run_metadata, ensure_ascii=False, indent=2, default=str),
)
logger.info("Essential scan data saved to: %s", run_dir)
except (OSError, RuntimeError):
logger.exception("Failed to save scan data")
def _write_executive_report(self, run_dir: Path) -> None:
path = run_dir / "penetration_test_report.md"
with path.open("w", encoding="utf-8") as f:
f.write("# Security Penetration Test Report\n\n")
f.write(f"**Generated:** {datetime.now(UTC).strftime('%Y-%m-%d %H:%M:%S UTC')}\n\n")
f.write(f"{self.final_scan_result}\n")
logger.info("Saved final penetration test report to: %s", path)
def _write_vulnerabilities(self, run_dir: Path) -> None:
vuln_dir = run_dir / "vulnerabilities"
vuln_dir.mkdir(exist_ok=True)
new_reports = [r for r in self.vulnerability_reports if r["id"] not in self._saved_vuln_ids]
for report in new_reports:
(vuln_dir / f"{report['id']}.md").write_text(
_render_vulnerability_md(report),
encoding="utf-8",
)
self._saved_vuln_ids.add(report["id"])
sorted_reports = sorted(
self.vulnerability_reports,
key=lambda r: (_SEVERITY_ORDER.get(r["severity"], 5), r["timestamp"]),
)
csv_path = run_dir / "vulnerabilities.csv"
with csv_path.open("w", encoding="utf-8", newline="") as f:
fieldnames = ["id", "title", "severity", "timestamp", "file"]
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
for report in sorted_reports:
writer.writerow(
{
"id": report["id"],
"title": report["title"],
"severity": report["severity"].upper(),
"timestamp": report["timestamp"],
"file": f"vulnerabilities/{report['id']}.md",
},
)
_atomic_write_text(
run_dir / "vulnerabilities.json",
json.dumps(self.vulnerability_reports, ensure_ascii=False, indent=2, default=str),
)
if new_reports:
logger.info(
"Saved %d new vulnerability report(s) to: %s",
len(new_reports),
vuln_dir,
)
logger.info("Updated vulnerability index: %s", csv_path)
def _atomic_write_text(path: Path, payload: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
dir=str(path.parent),
prefix=f".{path.name}.",
suffix=".tmp",
delete=False,
) as tmp:
tmp.write(payload)
tmp_path = Path(tmp.name)
tmp_path.replace(path)
def _render_vulnerability_md(report: dict[str, Any]) -> str:
lines: list[str] = [
f"# {report.get('title', 'Untitled Vulnerability')}\n",
f"**ID:** {report.get('id', 'unknown')}",
f"**Severity:** {report.get('severity', 'unknown').upper()}",
f"**Found:** {report.get('timestamp', 'unknown')}",
]
metadata: list[tuple[str, Any]] = [
("Target", report.get("target")),
("Endpoint", report.get("endpoint")),
("Method", report.get("method")),
("CVE", report.get("cve")),
("CWE", report.get("cwe")),
]
cvss = report.get("cvss")
if cvss is not None:
metadata.append(("CVSS", cvss))
for label, value in metadata:
if value:
lines.append(f"**{label}:** {value}")
lines.append("")
lines.append("## Description\n")
lines.append(report.get("description") or "No description provided.")
lines.append("")
if report.get("impact"):
lines.append("## Impact\n")
lines.append(str(report["impact"]))
lines.append("")
if report.get("technical_analysis"):
lines.append("## Technical Analysis\n")
lines.append(str(report["technical_analysis"]))
lines.append("")
if report.get("poc_description") or report.get("poc_script_code"):
lines.append("## Proof of Concept\n")
if report.get("poc_description"):
lines.append(str(report["poc_description"]))
lines.append("")
if report.get("poc_script_code"):
lines.append("```")
lines.append(str(report["poc_script_code"]))
lines.append("```")
lines.append("")
if report.get("code_locations"):
lines.append("## Code Analysis\n")
for i, loc in enumerate(report["code_locations"]):
file_ref = loc.get("file", "unknown")
line_ref = ""
if loc.get("start_line") is not None:
if loc.get("end_line") and loc["end_line"] != loc["start_line"]:
line_ref = f" (lines {loc['start_line']}-{loc['end_line']})"
else:
line_ref = f" (line {loc['start_line']})"
lines.append(f"**Location {i + 1}:** `{file_ref}`{line_ref}")
if loc.get("label"):
lines.append(f" {loc['label']}")
if loc.get("snippet"):
lines.append(f" ```\n {loc['snippet']}\n ```")
if loc.get("fix_before") or loc.get("fix_after"):
lines.append("\n **Suggested Fix:**")
lines.append("```diff")
if loc.get("fix_before"):
for ln in str(loc["fix_before"]).splitlines():
lines.append(f"- {ln}")
if loc.get("fix_after"):
for ln in str(loc["fix_after"]).splitlines():
lines.append(f"+ {ln}")
lines.append("```")
lines.append("")
if report.get("remediation_steps"):
lines.append("## Remediation\n")
lines.append(str(report["remediation_steps"]))
lines.append("")
return "\n".join(lines)
+36 -220
View File
@@ -1,13 +1,10 @@
"""Multi-agent graph tools backed by the SDK-native :class:`AgentCoordinator`. """Multi-agent graph tools backed by the SDK-native :class:`AgentCoordinator`.
- ``view_agent_graph``: render the parent/child tree. - ``view_agent_graph``: render the parent/child tree.
- ``agent_status``: per-agent status + pending message count.
- ``send_message_to_agent``: append a message to another agent's SDK session. - ``send_message_to_agent``: append a message to another agent's SDK session.
- ``wait_for_message``: pause this agent until a message arrives or - ``wait_for_message``: pause this agent until a message arrives or
``timeout_seconds`` elapses. ``timeout_seconds`` elapses.
- ``create_agent``: spawn a child via - ``create_agent``: asks the scan runner to spawn an addressable child.
``asyncio.create_task(Runner.run(...))``; the task handle is stored
so a root-level cancel cascades to descendants.
- ``agent_finish``: subagents only — posts a structured completion - ``agent_finish``: subagents only — posts a structured completion
report to the parent's SDK session and returns a final-output marker. report to the parent's SDK session and returns a final-output marker.
""" """
@@ -19,27 +16,11 @@ import json
import logging import logging
import uuid import uuid
from datetime import UTC, datetime from datetime import UTC, datetime
from pathlib import Path from typing import Any, Literal
from typing import TYPE_CHECKING, Any, Literal
from agents import RunConfig, RunContextWrapper, function_tool from agents import RunContextWrapper, function_tool
from agents.items import TResponseInputItem
from agents.model_settings import ModelSettings
from agents.sandbox import SandboxRunConfig
from strix.llm.multi_provider_setup import build_multi_provider from strix.orchestration.coordinator import coordinator_from_context
from strix.llm.retry import DEFAULT_RETRY
from strix.orchestration.coordinator import (
coordinator_from_context,
open_agent_session,
run_with_continuation,
)
if TYPE_CHECKING:
from collections.abc import Callable
from agents import Agent as SDKAgent
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -145,42 +126,6 @@ async def view_agent_graph(ctx: RunContextWrapper) -> str:
) )
@function_tool(timeout=30)
async def agent_status(ctx: RunContextWrapper, agent_id: str) -> str:
"""Look up one agent's lifecycle state + pending message count.
Use when you need precise state on a specific agent (e.g., "is the
XSS specialist still going?") rather than the full tree view.
Returns ``status`` (``running`` / ``waiting`` / ``completed`` /
``crashed`` / ``stopped``), ``parent_id``, and ``pending_messages``.
Args:
agent_id: The 8-char id from ``view_agent_graph`` /
``create_agent``.
"""
inner = _ctx(ctx)
coordinator = coordinator_from_context(inner)
if coordinator is None:
return json.dumps(
{"success": False, "error": "Agent coordinator not initialized in context."},
ensure_ascii=False,
default=str,
)
info = await coordinator.agent_info(agent_id)
if info is None:
return json.dumps(
{
"success": False,
"error": f"Unknown agent_id: {agent_id}",
},
ensure_ascii=False,
default=str,
)
info["success"] = True
return json.dumps(info, ensure_ascii=False, default=str)
@function_tool(timeout=30) @function_tool(timeout=30)
async def send_message_to_agent( async def send_message_to_agent(
ctx: RunContextWrapper, ctx: RunContextWrapper,
@@ -191,8 +136,9 @@ async def send_message_to_agent(
) -> str: ) -> str:
"""Send a message to another agent's inbox — sparingly. """Send a message to another agent's inbox — sparingly.
Inter-agent messages are surfaced at the top of the target's next Inter-agent messages are appended to the target's SDK session and
LLM turn. Use only when essential: interrupt any active target turn so the next run cycle sees them.
Use only when essential:
- Sharing a discovered finding/credential another agent needs. - Sharing a discovered finding/credential another agent needs.
- Asking a specialist a focused question. - Asking a specialist a focused question.
@@ -202,8 +148,8 @@ async def send_message_to_agent(
**Don't** use for routine "hello/status" pings, for context the **Don't** use for routine "hello/status" pings, for context the
target already has (children inherit parent history), or when target already has (children inherit parent history), or when
parent/child completion via ``agent_finish`` already covers the parent/child completion via ``agent_finish`` already covers the
flow. Messages to any registered agent are queued, regardless of flow. Messages to any registered agent wake it, regardless of
status, so a follow-up can wake a completed/stopped/failed agent. status, so a follow-up can restart a completed/stopped/failed agent.
Args: Args:
target_agent_id: Recipient's 8-char id. target_agent_id: Recipient's 8-char id.
@@ -249,7 +195,7 @@ async def send_message_to_agent(
"success": True, "success": True,
"message_id": msg_id, "message_id": msg_id,
"target_agent_id": target_agent_id, "target_agent_id": target_agent_id,
"delivery_status": "queued", "delivery_status": "delivered",
}, },
ensure_ascii=False, ensure_ascii=False,
default=str, default=str,
@@ -269,7 +215,7 @@ def _session_items_payload(items: list[Any]) -> list[dict[str, Any]]:
@function_tool(timeout=601) @function_tool(timeout=601)
async def wait_for_message( async def wait_for_message( # noqa: PLR0911
ctx: RunContextWrapper, ctx: RunContextWrapper,
reason: str = "Waiting for messages from other agents", reason: str = "Waiting for messages from other agents",
timeout_seconds: int = 600, timeout_seconds: int = 600,
@@ -325,10 +271,8 @@ async def wait_for_message(
default=str, default=str,
) )
pending = await coordinator.pending_count(me) pending, items = await coordinator.consume_pending(me, include_items=True)
if pending > 0: if pending > 0:
items = await coordinator.recent_session_items(me, pending)
await coordinator.consume_wake(me)
await coordinator.mark_running(me) await coordinator.mark_running(me)
return json.dumps( return json.dumps(
{ {
@@ -344,7 +288,6 @@ async def wait_for_message(
if interactive: if interactive:
await coordinator.park_waiting(me) await coordinator.park_waiting(me)
inner["agent_waiting_called"] = True
return json.dumps( return json.dumps(
{ {
"success": True, "success": True,
@@ -388,9 +331,7 @@ async def wait_for_message(
default=str, default=str,
) )
pending = await coordinator.pending_count(me) pending, items = await coordinator.consume_pending(me, include_items=True)
items = await coordinator.recent_session_items(me, pending)
await coordinator.consume_wake(me)
await coordinator.mark_running(me) await coordinator.mark_running(me)
return json.dumps( return json.dumps(
@@ -456,7 +397,7 @@ async def create_agent(
inner = _ctx(ctx) inner = _ctx(ctx)
coordinator = coordinator_from_context(inner) coordinator = coordinator_from_context(inner)
parent_id = inner.get("agent_id") parent_id = inner.get("agent_id")
factory: Callable[..., SDKAgent] | None = inner.get("agent_factory") spawner = inner.get("spawn_child_agent")
if coordinator is None or parent_id is None: if coordinator is None or parent_id is None:
return json.dumps( return json.dumps(
@@ -464,155 +405,38 @@ async def create_agent(
ensure_ascii=False, ensure_ascii=False,
default=str, default=str,
) )
if factory is None: if not callable(spawner):
return json.dumps( return json.dumps(
{ {
"success": False, "success": False,
"error": ( "error": "Scan runner did not provide a child-agent spawner in context.",
"No agent_factory in context. "
"The root assembly must inject one when building the run context."
),
}, },
ensure_ascii=False, ensure_ascii=False,
default=str, default=str,
) )
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 json.dumps(
{
"success": False,
"error": f"agent_factory failed: {e!s}",
},
ensure_ascii=False,
default=str,
)
await coordinator.register(
child_id,
name,
parent_id,
task=task,
skills=list(skills or []),
is_whitebox=bool(inner.get("is_whitebox", False)),
scan_mode=str(inner.get("scan_mode", "deep")),
diff_scope=inner.get("diff_scope"),
)
# ``ctx.turn_input`` carries the parent's full conversation up to and # ``ctx.turn_input`` carries the parent's full conversation up to and
# including the call that's currently invoking ``create_agent`` # 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 [] parent_history = list(ctx.turn_input) if inherit_context and ctx.turn_input else []
initial_input: list[TResponseInputItem] = [] try:
if parent_history: result = await spawner(
rendered = json.dumps(parent_history, ensure_ascii=False, default=str) parent_ctx=inner,
initial_input.append( name=name,
{ task=task,
"role": "user", skills=list(skills or []),
"content": ( parent_history=parent_history,
"== Inherited context from parent (background only) ==\n" )
f"{rendered}\n" except Exception as e:
"== End of inherited context ==\n" logger.exception("create_agent: scan runner failed to spawn child '%s'", name)
"Use the above as background only; do not continue the " return json.dumps(
"parent's work. Your task follows." {"success": False, "error": f"child spawn failed: {e!s}"},
), ensure_ascii=False,
}, default=str,
) )
initial_input.append(
{
"role": "user",
"content": (
f"You are agent {name} ({child_id}); your parent is {parent_id}. "
f"Maintain your own identity. Call agent_finish when your task "
f"is complete."
),
}
)
initial_input.append({"role": "user", "content": task})
child_ctx: dict[str, Any] = {
"coordinator": coordinator,
"sandbox_session": inner.get("sandbox_session"),
"sandbox_client": inner.get("sandbox_client"),
"caido_client": inner.get("caido_client"),
"agent_id": child_id,
"parent_id": parent_id,
"tracer": inner.get("tracer"),
"model": inner["model"],
"model_settings": inner.get("model_settings"),
"max_turns": int(inner.get("max_turns", 300)),
"agent_finish_called": False,
"is_whitebox": bool(inner.get("is_whitebox", False)),
"interactive": bool(inner.get("interactive", False)),
"scan_mode": str(inner.get("scan_mode", "deep")),
"diff_scope": inner.get("diff_scope"),
"run_id": inner.get("run_id"),
"agent_factory": factory,
# Stashed for ``agent_finish`` to echo back in its completion report.
"task": task,
# Inherited so ``create_agent`` calls from this child can also
# add their per-child sessions to the same teardown list.
"_sessions_to_close": inner.get("_sessions_to_close", []),
}
# Every agent gets its own SDK session_id inside the shared agents.db.
agents_db_path = inner.get("agents_db_path")
if not isinstance(agents_db_path, Path):
run_id = inner.get("run_id") or "default"
agents_db_path = Path.cwd() / "strix_runs" / str(run_id) / "agents.db"
child_session = open_agent_session(child_id, agents_db_path)
sessions_list = child_ctx.get("_sessions_to_close")
if isinstance(sessions_list, list):
sessions_list.append(child_session)
await coordinator.attach_runtime(child_id, session=child_session)
child_model_settings = ModelSettings(
parallel_tool_calls=False,
tool_choice="required",
retry=DEFAULT_RETRY,
)
override = inner.get("model_settings")
if override is not None:
child_model_settings = child_model_settings.resolve(override)
sandbox_session = inner.get("sandbox_session")
child_run_config = RunConfig(
model=inner["model"],
model_provider=build_multi_provider(),
model_settings=child_model_settings,
sandbox=(
SandboxRunConfig(client=inner.get("sandbox_client"), session=sandbox_session)
if sandbox_session is not None
else None
),
tracing_disabled=False,
trace_include_sensitive_data=False,
)
task_handle = asyncio.create_task(
run_with_continuation(
agent=child_agent,
initial_input=initial_input,
run_config=child_run_config,
context=child_ctx,
max_turns=int(inner.get("max_turns", 300)),
coordinator=coordinator,
agent_id=child_id,
interactive=bool(inner.get("interactive", False)),
session=child_session,
),
name=f"agent-{name}-{child_id}",
)
await coordinator.attach_runtime(child_id, task=task_handle)
logger.info( logger.info(
"create_agent: spawned %s (%s) parent=%s skills=%d task_len=%d", "create_agent: spawned %s (%s) parent=%s skills=%d task_len=%d",
child_id, result.get("agent_id"),
name, name,
parent_id or "-", parent_id or "-",
len(skills or []), len(skills or []),
@@ -620,13 +444,7 @@ async def create_agent(
) )
return json.dumps( return json.dumps(
{ result,
"success": True,
"agent_id": child_id,
"name": name,
"parent_id": parent_id,
"message": f"Spawned '{name}' ({child_id}) running in parallel.",
},
ensure_ascii=False, ensure_ascii=False,
default=str, default=str,
) )
@@ -700,10 +518,6 @@ async def agent_finish(
default=str, default=str,
) )
# The coordinator settles lifecycle from this JSON final output after
# the SDK run finishes; the tool only sends the parent report.
inner["agent_finish_called"] = True
parent_notified = False parent_notified = False
if report_to_parent: if report_to_parent:
async with coordinator._lock: async with coordinator._lock:
@@ -736,6 +550,7 @@ async def agent_finish(
len(findings or []), len(findings or []),
parent_notified, parent_notified,
) )
await coordinator.set_status(me, "completed")
return json.dumps( return json.dumps(
{ {
@@ -798,7 +613,8 @@ async def stop_agent(
ensure_ascii=False, ensure_ascii=False,
default=str, default=str,
) )
if await coordinator.agent_info(target_agent_id) is None: _, statuses, _ = await coordinator.graph_snapshot()
if target_agent_id not in statuses:
return json.dumps( return json.dumps(
{"success": False, "error": f"Unknown agent_id: {target_agent_id}"}, {"success": False, "error": f"Unknown agent_id: {target_agent_id}"},
ensure_ascii=False, ensure_ascii=False,
+16 -11
View File
@@ -44,24 +44,24 @@ def _do_finish(
return {"success": False, "message": "Validation failed", "errors": errors} return {"success": False, "message": "Validation failed", "errors": errors}
try: try:
from strix.telemetry.tracer import get_global_tracer from strix.telemetry.scan_store import get_global_scan_store
tracer = get_global_tracer() scan_store = get_global_scan_store()
if tracer is None: if scan_store is None:
logger.warning("No global tracer; scan results not persisted") logger.warning("No global scan store; scan results not persisted")
return { return {
"success": True, "success": True,
"scan_completed": True, "scan_completed": True,
"message": "Scan completed (not persisted)", "message": "Scan completed (not persisted)",
"warning": "Results could not be persisted - tracer unavailable", "warning": "Results could not be persisted - scan store unavailable",
} }
tracer.update_scan_final_fields( scan_store.update_scan_final_fields(
executive_summary=executive_summary.strip(), executive_summary=executive_summary.strip(),
methodology=methodology.strip(), methodology=methodology.strip(),
technical_analysis=technical_analysis.strip(), technical_analysis=technical_analysis.strip(),
recommendations=recommendations.strip(), recommendations=recommendations.strip(),
) )
vuln_count = len(tracer.vulnerability_reports) vuln_count = len(scan_store.vulnerability_reports)
except (ImportError, AttributeError) as e: except (ImportError, AttributeError) as e:
logger.exception("finish_scan persistence failed") logger.exception("finish_scan persistence failed")
return {"success": False, "message": f"Failed to complete scan: {e!s}"} return {"success": False, "message": f"Failed to complete scan: {e!s}"}
@@ -98,8 +98,8 @@ async def finish_scan(
**Pre-flight checklist (mandatory — do not skip):** **Pre-flight checklist (mandatory — do not skip):**
1. **Call ``view_agent_graph`` first.** Inspect every entry in the 1. **Call ``view_agent_graph`` first.** Inspect every entry in the
summary. If ANY agent is in ``running`` / ``waiting`` / summary. If ANY agent is in ``running`` / ``waiting`` state,
``llm_failed`` state, you MUST NOT call ``finish_scan`` yet — you MUST NOT call ``finish_scan`` yet —
wrap them up first via ``send_message_to_agent`` (ask them to wrap them up first via ``send_message_to_agent`` (ask them to
finish), ``wait_for_message`` (block until their report finish), ``wait_for_message`` (block until their report
arrives), or ``stop_agent`` (graceful cancel). Only ``completed`` arrives), or ``stop_agent`` (graceful cancel). Only ``completed``
@@ -178,6 +178,11 @@ async def finish_scan(
technical_analysis=technical_analysis, technical_analysis=technical_analysis,
recommendations=recommendations, recommendations=recommendations,
) )
if result.get("success") and result.get("scan_completed"): if (
inner["finish_scan_called"] = True result.get("success")
and result.get("scan_completed")
and coordinator is not None
and isinstance(me, str)
):
await coordinator.set_status(me, "completed")
return json.dumps(result, ensure_ascii=False, default=str) return json.dumps(result, ensure_ascii=False, default=str)
+25 -8
View File
@@ -151,7 +151,7 @@ _REQUIRED_FIELDS = {
} }
async def _do_create( # noqa: PLR0912 async def _do_create( # noqa: PLR0912, PLR0915
*, *,
title: str, title: str,
description: str, description: str,
@@ -167,6 +167,8 @@ async def _do_create( # noqa: PLR0912
cve: str | None, cve: str | None,
cwe: str | None, cwe: str | None,
code_locations: list[dict[str, Any]] | None, code_locations: list[dict[str, Any]] | None,
agent_id: str | None = None,
agent_name: str | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
errors: list[str] = [] errors: list[str] = []
fields = { fields = {
@@ -212,20 +214,20 @@ async def _do_create( # noqa: PLR0912
cvss_score, severity, _vector = _calculate_cvss(cvss_breakdown) cvss_score, severity, _vector = _calculate_cvss(cvss_breakdown)
try: try:
from strix.telemetry.tracer import get_global_tracer from strix.telemetry.scan_store import get_global_scan_store
tracer = get_global_tracer() scan_store = get_global_scan_store()
if tracer is None: if scan_store is None:
logger.warning("No global tracer; vulnerability report not persisted") logger.warning("No global scan store; vulnerability report not persisted")
return { return {
"success": True, "success": True,
"message": f"Vulnerability report '{title}' created (not persisted)", "message": f"Vulnerability report '{title}' created (not persisted)",
"warning": "Report could not be persisted - tracer unavailable", "warning": "Report could not be persisted - scan store unavailable",
} }
from strix.llm.dedupe import check_duplicate from strix.llm.dedupe import check_duplicate
existing = tracer.get_existing_vulnerabilities() existing = scan_store.get_existing_vulnerabilities()
candidate = { candidate = {
"title": title, "title": title,
"description": description, "description": description,
@@ -256,7 +258,7 @@ async def _do_create( # noqa: PLR0912
"reason": dedupe.get("reason", ""), "reason": dedupe.get("reason", ""),
} }
report_id = tracer.add_vulnerability_report( report_id = scan_store.add_vulnerability_report(
title=title, title=title,
description=description, description=description,
severity=severity, severity=severity,
@@ -273,6 +275,8 @@ async def _do_create( # noqa: PLR0912
cve=cve, cve=cve,
cwe=cwe, cwe=cwe,
code_locations=parsed_locations, code_locations=parsed_locations,
agent_id=agent_id if isinstance(agent_id, str) else None,
agent_name=agent_name if isinstance(agent_name, str) else None,
) )
except (ImportError, AttributeError) as e: except (ImportError, AttributeError) as e:
logger.exception("create_vulnerability_report persistence failed") logger.exception("create_vulnerability_report persistence failed")
@@ -404,6 +408,17 @@ async def create_vulnerability_report(
``fix_before`` (verbatim source), ``fix_after`` (suggested ``fix_before`` (verbatim source), ``fix_after`` (suggested
replacement). replacement).
""" """
inner = ctx.context if isinstance(ctx.context, dict) else {}
raw_agent_id = inner.get("agent_id")
agent_id = raw_agent_id if isinstance(raw_agent_id, str) else None
agent_name = None
coordinator = inner.get("coordinator")
if agent_id is not None and coordinator is not None:
names = getattr(coordinator, "names", {})
if isinstance(names, dict):
raw_agent_name = names.get(agent_id)
agent_name = raw_agent_name if isinstance(raw_agent_name, str) else None
result = await _do_create( result = await _do_create(
title=title, title=title,
description=description, description=description,
@@ -419,5 +434,7 @@ async def create_vulnerability_report(
cve=cve, cve=cve,
cwe=cwe, cwe=cwe,
code_locations=code_locations, code_locations=code_locations,
agent_id=agent_id,
agent_name=agent_name,
) )
return json.dumps(result, ensure_ascii=False, default=str) return json.dumps(result, ensure_ascii=False, default=str)