From bd40884fcff3e08d3d01076fdb12f50745c581a2 Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sun, 26 Apr 2026 11:30:00 -0700 Subject: [PATCH] Simplify SDK-native orchestration --- pyproject.toml | 19 +- strix.spec | 40 +- strix/agents/factory.py | 2 - strix/interface/cli.py | 24 +- strix/interface/main.py | 22 +- strix/interface/tui.py | 394 ++++++----- strix/interface/utils.py | 123 +--- strix/orchestration/__init__.py | 3 +- strix/orchestration/coordinator.py | 501 ++------------ strix/orchestration/runner.py | 636 ++++++++++++++++++ strix/orchestration/scan.py | 645 ------------------- strix/orchestration/utils.py | 157 +++++ strix/runtime/session_manager.py | 2 +- strix/telemetry/README.md | 2 +- strix/telemetry/__init__.py | 12 +- strix/telemetry/posthog.py | 19 +- strix/telemetry/scan_artifacts.py | 234 ------- strix/telemetry/{tracer.py => scan_store.py} | 417 ++++++------ strix/tools/agents_graph/tools.py | 256 ++------ strix/tools/finish/tool.py | 27 +- strix/tools/reporting/tool.py | 33 +- 21 files changed, 1435 insertions(+), 2133 deletions(-) create mode 100644 strix/orchestration/runner.py delete mode 100644 strix/orchestration/scan.py create mode 100644 strix/orchestration/utils.py delete mode 100644 strix/telemetry/scan_artifacts.py rename strix/telemetry/{tracer.py => scan_store.py} (51%) diff --git a/pyproject.toml b/pyproject.toml index 87cb0ab..b06371d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -194,11 +194,6 @@ ignore = [ "strix/tools/**/*.py" = [ "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`. "strix/llm/anthropic_cache_wrapper.py" = [ "A002", # Argument shadows builtin (parent signature uses `input`) @@ -209,9 +204,6 @@ ignore = [ # Backend factories import their backend's deps lazily so deployments # that pick a different backend don't need every backend's libs installed. "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" = [ "TC002", # Manifest, Container imported for annotations "TC003", # uuid imported for annotation @@ -234,13 +226,13 @@ ignore = [ "strix/agents/factory.py" = ["TC002"] # Entry point: ``Path`` is used at runtime by the typing of the # 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 # helpers would add indirection without simplifying anything. -"strix/orchestration/scan.py" = ["TC003", "PLR0912", "PLR0915", "PLC0415"] -# Tracer carries a long event surface and a runtime ``Callable`` -# annotation on ``vulnerability_found_callback``. -"strix/telemetry/tracer.py" = ["TC003", "PLR0912", "PLR0915", "E501"] +"strix/orchestration/runner.py" = ["TC003", "PLR0912", "PLR0915", "PLC0415"] +# ScanStore carries scan artifact/report fields, artifact rendering, and +# a runtime ``Callable`` annotation on ``vulnerability_found_callback``. +"strix/telemetry/scan_store.py" = ["TC003", "PLR0912", "PLR0915", "E501", "PERF401"] # Interface utility branches per scope-mode / target-type combination; # splitting would obscure the decision tree without simplifying it. "strix/interface/utils.py" = ["PLR0912", "BLE001", "PLC0415"] @@ -249,7 +241,6 @@ ignore = [ "strix/interface/cli.py" = ["BLE001", "PLC0415"] "strix/interface/tui.py" = ["BLE001", "PLC0415", "PLR0912", "PLR0915", "SIM105"] "strix/interface/main.py" = ["BLE001", "PLC0415", "PLR0912", "PLR0915"] -"strix/interface/streaming_parser.py" = ["PLC0415"] "strix/interface/tool_components/agent_message_renderer.py" = ["PLC0415"] [tool.ruff.lint.isort] diff --git a/strix.spec b/strix.spec index 425219a..445988e 100644 --- a/strix.spec +++ b/strix.spec @@ -119,23 +119,37 @@ hiddenimports = [ 'strix.interface.utils', 'strix.interface.tool_components', 'strix.agents', - 'strix.agents.base_agent', - 'strix.agents.state', - 'strix.agents.StrixAgent', + 'strix.agents.factory', + 'strix.agents.prompt', 'strix.llm', - 'strix.llm.llm', - 'strix.llm.config', - 'strix.llm.utils', - 'strix.llm.memory_compressor', + 'strix.llm.anthropic_cache_wrapper', + 'strix.llm.dedupe', + 'strix.llm.multi_provider_setup', + 'strix.llm.retry', + 'strix.orchestration', + 'strix.orchestration.coordinator', + 'strix.orchestration.runner', + 'strix.orchestration.utils', 'strix.runtime', - 'strix.runtime.runtime', - 'strix.runtime.docker_runtime', + 'strix.runtime.backends', + 'strix.runtime.caido_bootstrap', + 'strix.runtime.docker_client', + 'strix.runtime.session_manager', 'strix.telemetry', - 'strix.telemetry.tracer', + 'strix.telemetry.logging', + 'strix.telemetry.posthog', + 'strix.telemetry.scan_store', 'strix.tools', - 'strix.tools.registry', - 'strix.tools.executor', - 'strix.tools.argument_parser', + 'strix.tools.agents_graph.tools', + 'strix.tools.finish.tool', + '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', ] diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 54aeb0f..50ec0e2 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -29,7 +29,6 @@ from agents.tool import Tool from strix.agents.prompt import render_system_prompt from strix.tools.agents_graph.tools import ( agent_finish, - agent_status, create_agent, send_message_to_agent, stop_agent, @@ -163,7 +162,6 @@ _BASE_TOOLS: tuple[Tool, ...] = ( python_action, # Multi-agent graph tools (the coordinator is in ctx.context) view_agent_graph, - agent_status, send_message_to_agent, wait_for_message, create_agent, diff --git a/strix/interface/cli.py b/strix/interface/cli.py index 4846241..acdcbbb 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -13,9 +13,9 @@ from rich.panel import Panel from rich.text import Text 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.telemetry.tracer import Tracer, set_global_tracer +from strix.telemetry.scan_store import ScanStore, set_global_scan_store from .utils import ( 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 "", } - tracer = Tracer(args.run_name) - tracer.set_scan_config(scan_config) + scan_store = ScanStore(args.run_name) + scan_store.set_scan_config(scan_config) + scan_store.hydrate_from_run_dir() def display_vulnerability(report: dict[str, Any]) -> None: 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() - tracer.vulnerability_found_callback = display_vulnerability + scan_store.vulnerability_found_callback = display_vulnerability def cleanup_on_exit() -> None: - tracer.cleanup() + scan_store.cleanup() def signal_handler(_signum: int, _frame: Any) -> None: - tracer.cleanup() + scan_store.cleanup() sys.exit(1) atexit.register(cleanup_on_exit) @@ -129,14 +130,14 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 if hasattr(signal, "SIGHUP"): signal.signal(signal.SIGHUP, signal_handler) - set_global_tracer(tracer) + set_global_scan_store(scan_store) def create_live_status() -> Panel: status_text = Text() status_text.append("Penetration test in progress", style="bold #22c55e") status_text.append("\n\n") - stats_text = build_live_stats_text(tracer) + stats_text = build_live_stats_text(scan_store) if 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, image=_resolve_sandbox_image(), local_sources=getattr(args, "local_sources", None) or [], - tracer=tracer, interactive=bool(getattr(args, "interactive", False)), ) finally: @@ -196,7 +196,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 console.print(f"[bold red]Error during penetration test:[/] {e}") raise - if tracer.final_scan_result: + if scan_store.final_scan_result: console.print() final_report_text = Text() @@ -206,7 +206,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 Text.assemble( final_report_text, "\n\n", - tracer.final_scan_result, + scan_store.final_scan_result, ), title="[bold white]STRIX", title_align="left", diff --git a/strix/interface/main.py b/strix/interface/main.py index b500b5a..bc37a35 100644 --- a/strix/interface/main.py +++ b/strix/interface/main.py @@ -37,7 +37,7 @@ from strix.interface.utils import ( validate_llm_response, ) 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" @@ -47,7 +47,7 @@ import logging # noqa: E402 # 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 # work (``main()``, env validation, image pull) emits via the module # 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: console = Console() - tracer = get_global_tracer() + scan_store = get_global_scan_store() scan_completed = False - if tracer and tracer.scan_results: - scan_completed = tracer.scan_results.get("scan_completed", False) + if scan_store and scan_store.scan_results: + scan_completed = scan_store.scan_results.get("scan_completed", False) completion_text = Text() if scan_completed: @@ -575,7 +575,7 @@ def display_completion_message(args: argparse.Namespace, results_path: Path) -> target_text.append("\n ") 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] @@ -753,16 +753,16 @@ def main() -> None: posthog.error("unhandled_exception", str(e)) raise finally: - tracer = get_global_tracer() - if tracer: - posthog.end(tracer, exit_reason=exit_reason) + scan_store = get_global_scan_store() + if scan_store: + posthog.end(scan_store, exit_reason=exit_reason) results_path = Path("strix_runs") / args.run_name display_completion_message(args, results_path) if args.non_interactive: - tracer = get_global_tracer() - if tracer and tracer.vulnerability_reports: + scan_store = get_global_scan_store() + if scan_store and scan_store.vulnerability_reports: sys.exit(2) diff --git a/strix/interface/tui.py b/strix/interface/tui.py index 0b0ba09..fb9a7d2 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -2,13 +2,16 @@ import argparse import asyncio import atexit import contextlib +import json import logging import signal import sys import threading from collections.abc import Callable +from datetime import UTC, datetime from importlib.metadata import PackageNotFoundError from importlib.metadata import version as pkg_version +from pathlib import Path 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.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.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.telemetry.tracer import Tracer, set_global_tracer +from strix.telemetry.scan_store import ScanStore, set_global_scan_store logger = logging.getLogger(__name__) @@ -639,6 +641,139 @@ class VulnerabilitiesPanel(VerticalScroll): # type: ignore[misc] 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] def compose(self) -> ComposeResult: yield Grid( @@ -704,9 +839,13 @@ class StrixTUIApp(App): # type: ignore[misc] self.args = args self.scan_config = self._build_scan_config(args) - self.tracer = Tracer(self.scan_config["run_name"]) - self.tracer.set_scan_config(self.scan_config) - set_global_tracer(self.tracer) + self.scan_store = ScanStore(self.scan_config["run_name"]) + self.scan_store.set_scan_config(self.scan_config) + 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 # 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 cleanup_on_exit() -> None: - self.tracer.cleanup() + self.scan_store.cleanup() def signal_handler(_signum: int, _frame: Any) -> None: - self.tracer.cleanup() + self.scan_store.cleanup() sys.exit(0) atexit.register(cleanup_on_exit) @@ -881,9 +1020,9 @@ class StrixTUIApp(App): # type: ignore[misc] 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: return @@ -902,8 +1041,10 @@ class StrixTUIApp(App): # type: ignore[misc] except (ValueError, Exception): return + self._sync_live_view() + 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: self._add_agent_node(agent_data) self._displayed_agents.add(agent_id) @@ -922,6 +1063,43 @@ class StrixTUIApp(App): # type: ignore[misc] 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: if agent_id not in self.agent_nodes: return False @@ -937,7 +1115,6 @@ class StrixTUIApp(App): # type: ignore[misc] "completed": "๐ŸŸข", "failed": "๐Ÿ”ด", "stopped": "โ– ", - "llm_failed": "๐Ÿ”ด", } status_icon = status_indicators.get(status, "โ—‹") @@ -1074,8 +1251,6 @@ class StrixTUIApp(App): # type: ignore[misc] if event["type"] == "chat": content = self._render_chat_content(event["data"]) - elif event["type"] == "tool": - content = self._render_tool_content_simple(event["data"]) if content: if renderables: @@ -1090,7 +1265,7 @@ class StrixTUIApp(App): # type: ignore[misc] 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] ) -> tuple[Text | None, Text, bool]: status = agent_data.get("status", "running") @@ -1116,18 +1291,6 @@ class StrixTUIApp(App): # type: ignore[misc] text.append(msg) 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": error_msg = agent_data.get("error_message", "") text = Text() @@ -1173,7 +1336,7 @@ class StrixTUIApp(App): # type: ignore[misc] return 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( self.selected_agent_id, agent_data ) @@ -1206,7 +1369,7 @@ class StrixTUIApp(App): # type: ignore[misc] stats_content = Text() - stats_text = build_tui_stats_text(self.tracer) + stats_text = build_tui_stats_text(self.scan_store) if 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): return - vulnerabilities = self.tracer.vulnerability_reports + vulnerabilities = self.scan_store.vulnerability_reports if not vulnerabilities: self._safe_widget_operation(vuln_panel.add_class, "hidden") @@ -1234,8 +1397,10 @@ class StrixTUIApp(App): # type: ignore[misc] enriched_vulns = [] for vuln in vulnerabilities: enriched = dict(vuln) - report_id = vuln.get("id", "") - agent_name = self._get_agent_name_for_vulnerability(report_id) + agent_name = enriched.get("agent_name") + 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: enriched["agent_name"] = agent_name enriched_vulns.append(enriched) @@ -1243,18 +1408,6 @@ class StrixTUIApp(App): # type: ignore[misc] self._safe_widget_operation(vuln_panel.remove_class, "hidden") 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: text = Text() num_squares = self._sweep_num_squares @@ -1307,8 +1460,8 @@ class StrixTUIApp(App): # type: ignore[misc] def _animate_dots(self) -> None: has_active_agents = False - if self.selected_agent_id and self.selected_agent_id in self.tracer.agents: - agent_data = self.tracer.agents[self.selected_agent_id] + if self.selected_agent_id and self.selected_agent_id in self.live_view.agents: + agent_data = self.live_view.agents[self.selected_agent_id] status = agent_data.get("status", "running") if status in ["running", "waiting"]: has_active_agents = True @@ -1323,7 +1476,7 @@ class StrixTUIApp(App): # type: ignore[misc] if not has_active_agents: has_active_agents = any( 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: @@ -1331,28 +1484,12 @@ class StrixTUIApp(App): # type: ignore[misc] self._spinner_frame_index = 0 def _agent_has_real_activity(self, agent_id: str) -> bool: - initial_tools = {"scan_start_info", "subagent_start_info"} - - 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 + return any(msg.get("agent_id") == agent_id for msg in self.live_view.chat_messages) def _agent_vulnerability_count(self, agent_id: str) -> int: - count = 0 - 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 == "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 + return sum( + 1 for vuln in self.scan_store.vulnerability_reports if vuln.get("agent_id") == agent_id + ) def _gather_agent_events(self, agent_id: str) -> list[dict[str, Any]]: chat_events = [ @@ -1362,24 +1499,12 @@ class StrixTUIApp(App): # type: ignore[misc] "id": f"chat_{msg['message_id']}", "data": msg, } - for msg in self.tracer.chat_messages + for msg in self.live_view.chat_messages if msg.get("agent_id") == agent_id ] - tool_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 + chat_events.sort(key=lambda e: (e["timestamp"], e["id"])) + return chat_events def watch_selected_agent_id(self, _agent_id: str | None) -> None: 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"], image=str(image), local_sources=getattr(self.args, "local_sources", None) or [], - tracer=self.tracer, coordinator=self.coordinator, interactive=True, ), @@ -1470,7 +1594,6 @@ class StrixTUIApp(App): # type: ignore[misc] "completed": "๐ŸŸข", "failed": "๐Ÿ”ด", "stopped": "โ– ", - "llm_failed": "๐Ÿ”ด", } 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: 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") status = agent_data.get("status", "running") @@ -1542,7 +1665,6 @@ class StrixTUIApp(App): # type: ignore[misc] "completed": "๐ŸŸข", "failed": "๐Ÿ”ด", "stopped": "โ– ", - "llm_failed": "๐Ÿ”ด", } 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: 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 ( agent_data.get("parent_id") == new_parent_id and agent_id in self.agent_nodes @@ -1608,71 +1730,6 @@ class StrixTUIApp(App): # type: ignore[misc] 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] def handle_tree_highlight(self, event: Tree.NodeHighlighted) -> None: if len(self.screen_stack) > 1 or self.show_splash: @@ -1718,21 +1775,8 @@ class StrixTUIApp(App): # type: ignore[misc] self.selected_agent_id, len(message), ) - if self.tracer: - self.tracer.log_chat_message( - 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")``. + # Route to the agent's SDK session. The coordinator also interrupts + # any active stream so the message is picked up on the next run cycle. if self._scan_loop is not None and not self._scan_loop.is_closed(): target_agent_id = self.selected_agent_id asyncio.run_coroutine_threadsafe( @@ -1742,10 +1786,6 @@ class StrixTUIApp(App): # type: ignore[misc] ), 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._update_chat_view() @@ -1754,8 +1794,8 @@ class StrixTUIApp(App): # type: ignore[misc] def _get_agent_name(self, agent_id: str) -> str: try: - if self.tracer and agent_id in self.tracer.agents: - agent_name = self.tracer.agents[agent_id].get("name") + if agent_id in self.live_view.agents: + agent_name = self.live_view.agents[agent_id].get("name") if isinstance(agent_name, str): return agent_name except (KeyError, AttributeError) as e: @@ -1822,12 +1862,12 @@ class StrixTUIApp(App): # type: ignore[misc] agent_name = "Unknown Agent" try: - if self.tracer and self.selected_agent_id in self.tracer.agents: - agent_data = self.tracer.agents[self.selected_agent_id] + if self.selected_agent_id in self.live_view.agents: + agent_data = self.live_view.agents[self.selected_agent_id] agent_name = agent_data.get("name", "Unknown Agent") 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 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.tracer.cleanup() + self.scan_store.cleanup() self.exit() diff --git a/strix/interface/utils.py b/strix/interface/utils.py index 1acd4a7..98a8c7d 100644 --- a/strix/interface/utils.py +++ b/strix/interface/utils.py @@ -23,16 +23,6 @@ from rich.text import Text 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 def get_severity_color(severity: str) -> str: severity_colors = { @@ -206,13 +196,13 @@ def format_vulnerability_report(report: dict[str, Any]) -> Text: # noqa: PLR091 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.""" - vuln_count = len(tracer.vulnerability_reports) + vuln_count = len(scan_store.vulnerability_reports) if vuln_count > 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() if severity in severity_counts: severity_counts[severity] += 1 @@ -245,63 +235,20 @@ def _build_vulnerability_stats(stats_text: Text, tracer: Any) -> None: stats_text.append("\n") -def _build_llm_stats(stats_text: Text, total_stats: dict[str, Any]) -> None: - """Build LLM usage section of stats text.""" - 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.""" +def build_final_stats_text(scan_store: Any) -> Text: + """Build final stats from Strix-owned scan artifacts.""" stats_text = Text() - if not tracer: + if not scan_store: return stats_text - _build_vulnerability_stats(stats_text, tracer) - - 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"]) + _build_vulnerability_stats(stats_text, scan_store) return stats_text -def build_live_stats_text(tracer: Any) -> Text: +def build_live_stats_text(scan_store: Any) -> Text: stats_text = Text() - if not tracer: + if not scan_store: return stats_text 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("\n") - vuln_count = len(tracer.vulnerability_reports) - tool_count = tracer.get_real_tool_count() - agent_count = len(tracer.agents) - + vuln_count = len(scan_store.vulnerability_reports) stats_text.append("Vulnerabilities ", style="dim") stats_text.append(f"{vuln_count}", style="white") stats_text.append("\n") if vuln_count > 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() if severity in severity_counts: severity_counts[severity] += 1 @@ -340,57 +284,18 @@ def build_live_stats_text(tracer: Any) -> Text: 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 -def build_tui_stats_text(tracer: Any) -> Text: +def build_tui_stats_text(scan_store: Any) -> Text: stats_text = Text() - if not tracer: + if not scan_store: return stats_text model = load_settings().llm.model or "unknown" stats_text.append(str(model), style="white") - llm_stats = tracer.get_total_llm_stats() - 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) + caido_url = getattr(scan_store, "caido_url", None) if caido_url: stats_text.append("\n") stats_text.append("Caido: ", style="bold white") diff --git a/strix/orchestration/__init__.py b/strix/orchestration/__init__.py index bf6a962..7dab3be 100644 --- a/strix/orchestration/__init__.py +++ b/strix/orchestration/__init__.py @@ -3,8 +3,7 @@ - :class:`AgentCoordinator` owns Strix-specific graph/status/wake state. - SDK ``SQLiteSession`` owns per-agent conversation history and message transport. -- :func:`run_with_continuation` wraps SDK ``Runner.run_streamed`` only - enough to keep interactive agents addressable between bounded runs. +- ``runner.py`` owns SDK ``Runner.run_streamed`` and child-agent spawning. Import deeply (for example, ``from strix.orchestration.coordinator import AgentCoordinator``) so ``import strix.orchestration`` doesn't diff --git a/strix/orchestration/coordinator.py b/strix/orchestration/coordinator.py index a6f96cb..9e1e3cc 100644 --- a/strix/orchestration/coordinator.py +++ b/strix/orchestration/coordinator.py @@ -13,39 +13,27 @@ import json import logging import tempfile from dataclasses import dataclass, field -from datetime import UTC, datetime from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal - -from agents import Runner -from agents.exceptions import AgentsException, MaxTurnsExceeded, UserError -from agents.memory import SQLiteSession -from openai import APIError +from typing import TYPE_CHECKING, Any, Literal, cast if TYPE_CHECKING: from agents.items import TResponseInputItem from agents.memory import Session - from agents.result import RunResultBase, RunResultStreaming - from agents.run_config import RunConfig logger = logging.getLogger(__name__) -Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed", "llm_failed"] -ACTIVE_AGENT_STATUSES = {"running", "waiting", "llm_failed"} -_SNAPSHOT_VERSION = 3 -_WAITING_TIMEOUT_SUBAGENT = 300.0 -_TIMEOUT_RESUME_MESSAGE = "Waiting timeout reached. Resuming execution." +Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed"] @dataclass(slots=True) class AgentRuntime: session: Session | 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) - tool_calls: dict[str, str] = field(default_factory=dict) class AgentCoordinator: @@ -57,9 +45,6 @@ class AgentCoordinator: self.names: dict[str, str] = {} self.metadata: dict[str, dict[str, Any]] = {} 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._lock = asyncio.Lock() self._snapshot_path: Path | None = None @@ -75,23 +60,15 @@ class AgentCoordinator: *, task: str | None = None, skills: list[str] | None = None, - is_whitebox: bool = False, - scan_mode: str = "deep", - diff_scope: dict[str, Any] | None = None, ) -> None: async with self._lock: self.statuses[agent_id] = "running" self.parent_of[agent_id] = parent_id self.names[agent_id] = name 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] = { "task": task or "", "skills": list(skills or []), - "is_whitebox": bool(is_whitebox), - "scan_mode": scan_mode, - "diff_scope": diff_scope, } self.runtimes.setdefault(agent_id, AgentRuntime()) logger.info("agent.register %s (%s) parent=%s", agent_id, name, parent_id or "-") @@ -103,6 +80,7 @@ class AgentCoordinator: *, session: Session | None = None, task: asyncio.Task[Any] | None = None, + interrupt_on_message: bool | None = None, ) -> None: async with self._lock: runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) @@ -110,8 +88,8 @@ class AgentCoordinator: runtime.session = session if task is not None: runtime.task = task - if session is not None: - await self.flush_queued_messages(agent_id) + if interrupt_on_message is not None: + runtime.interrupt_on_message = interrupt_on_message async def mark_running(self, agent_id: str) -> None: async with self._lock: @@ -122,9 +100,6 @@ class AgentCoordinator: async def park_waiting(self, agent_id: str) -> None: 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 with self._lock: 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: """Deliver a user/peer message by appending it to the target SDK session.""" - should_queue = False async with self._lock: if target_agent_id not in self.statuses: logger.debug("agent.send dropped unknown target=%s", target_agent_id) return False runtime = self.runtimes.setdefault(target_agent_id, AgentRuntime()) session = runtime.session - should_queue = session is None or runtime.stream is not None - if should_queue: - self.queued_messages.setdefault(target_agent_id, []).append(dict(message)) - if session is not None and not should_queue: + stream = runtime.stream + interrupt = runtime.interrupt_on_message + if session is not None: try: await session.add_items([self._message_to_session_item(message)]) except Exception: @@ -158,81 +131,55 @@ class AgentCoordinator: return False async with self._lock: 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() - if should_queue: - logger.debug( - "agent.send %s queued until SDK session is safe to append", target_agent_id - ) + if stream is not None and interrupt: + stream.cancel(mode="immediate") await self._maybe_snapshot() return True - async def flush_queued_messages(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: + async def wait_for_message(self, agent_id: str) -> None: while True: async with self._lock: - pending = ( - self.pending_user_counts.get(agent_id, 0) - if user_only - else self.pending_counts.get(agent_id, 0) - ) - if pending > 0: + if self.pending_counts.get(agent_id, 0) > 0: return wake = self.runtimes.setdefault(agent_id, AgentRuntime()).wake wake.clear() await wake.wait() - async def wait_for_user_message(self, agent_id: str) -> None: - await self.wait_for_message(agent_id, user_only=True) - - async def consume_wake(self, agent_id: str) -> None: + async def consume_pending( + self, + agent_id: str, + *, + include_items: bool = False, + ) -> tuple[int, list[Any]]: async with self._lock: + count = self.pending_counts.get(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 - if session is None: - return [] + if count <= 0: + return 0, [] + if not include_items or session is None: + return count, [] 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: - stream = self.runtimes.get(agent_id, AgentRuntime()).stream - if stream is None: - return False - stream.cancel(mode=mode) # type: ignore[arg-type] - return True + sessions = { + aid: runtime.session + for aid, runtime in self.runtimes.items() + if runtime.session is not None + } - 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: if agent_id not in self.statuses: return @@ -240,7 +187,7 @@ class AgentCoordinator: runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) runtime.wake.set() stream = runtime.stream - if interrupt and stream is not None: + if stream is not None: stream.cancel(mode="after_turn") await self._maybe_snapshot() @@ -266,7 +213,7 @@ class AgentCoordinator: async def attach_stream( self, agent_id: str, - stream: RunResultStreaming, + stream: Any, ) -> None: async with self._lock: self.runtimes.setdefault(agent_id, AgentRuntime()).stream = stream @@ -274,7 +221,7 @@ class AgentCoordinator: async def detach_stream( self, agent_id: str, - stream: RunResultStreaming, + stream: Any, ) -> None: async with self._lock: 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 with self._lock: 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() - 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( self, ) -> tuple[dict[str, str | None], dict[str, Status], dict[str, str]]: async with self._lock: 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: sender = str(message.get("from", "unknown")) content = str(message.get("content", "")) if sender == "user": - return {"role": "user", "content": content} + return cast("TResponseInputItem", {"role": "user", "content": content}) sender_name = self.names.get(sender, sender) msg_type = message.get("type", "information") priority = message.get("priority", "normal") - return { - "role": "user", - "content": ( - f"[Message from {sender_name} ({sender}) | type={msg_type} " - f"| priority={priority}]\n{content}" - ), - } + return cast( + "TResponseInputItem", + { + "role": "user", + "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]: queue = [agent_id] @@ -337,29 +275,14 @@ class AgentCoordinator: queue.extend(child for child, parent in self.parent_of.items() if parent == aid) 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 with self._lock: return { - "version": _SNAPSHOT_VERSION, "statuses": dict(self.statuses), "parent_of": dict(self.parent_of), "names": dict(self.names), "metadata": {aid: dict(md) for aid, md in self.metadata.items()}, "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: @@ -369,12 +292,6 @@ class AgentCoordinator: self.names = dict(snap.get("names", {})) self.metadata = {aid: dict(md) for aid, md in snap.get("metadata", {}).items()} 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: self.runtimes.setdefault(aid, AgentRuntime()) @@ -404,297 +321,3 @@ class AgentCoordinator: def coordinator_from_context(ctx: dict[str, Any]) -> AgentCoordinator | None: coordinator = ctx.get("coordinator") 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, - } diff --git a/strix/orchestration/runner.py b/strix/orchestration/runner.py new file mode 100644 index 0000000..ee6a507 --- /dev/null +++ b/strix/orchestration/runner.py @@ -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") diff --git a/strix/orchestration/scan.py b/strix/orchestration/scan.py deleted file mode 100644 index d8d9755..0000000 --- a/strix/orchestration/scan.py +++ /dev/null @@ -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/``. 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/``. - 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() diff --git a/strix/orchestration/utils.py b/strix/orchestration/utils.py new file mode 100644 index 0000000..cc24baa --- /dev/null +++ b/strix/orchestration/utils.py @@ -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 diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py index 1394495..85a8f54 100644 --- a/strix/runtime/session_manager.py +++ b/strix/runtime/session_manager.py @@ -67,7 +67,7 @@ async def create_or_reuse( # Build Manifest entries keyed by ``workspace_subdir`` โ€” the SDK # mounts each at ``/workspace/``, 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 # unrelated host content into the sandbox. entries: dict[str | Path, BaseEntry] = {} diff --git a/strix/telemetry/README.md b/strix/telemetry/README.md index 5771f26..bb90fe1 100644 --- a/strix/telemetry/README.md +++ b/strix/telemetry/README.md @@ -16,7 +16,7 @@ We collect only very **basic** usage data including: **System Context:** OS type, architecture, Strix version\ **Scan Context:** Scan mode (quick/standard/deep), scan type (whitebox/blackbox)\ **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. diff --git a/strix/telemetry/__init__.py b/strix/telemetry/__init__.py index 0537f61..a82122b 100644 --- a/strix/telemetry/__init__.py +++ b/strix/telemetry/__init__.py @@ -1,10 +1,14 @@ 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__ = [ - "Tracer", - "get_global_tracer", + "ScanStore", + "get_global_scan_store", "posthog", - "set_global_tracer", + "set_global_scan_store", ] diff --git a/strix/telemetry/posthog.py b/strix/telemetry/posthog.py index 19656fb..4561d47 100644 --- a/strix/telemetry/posthog.py +++ b/strix/telemetry/posthog.py @@ -11,7 +11,7 @@ from strix.config import load_settings if TYPE_CHECKING: - from strix.telemetry.tracer import Tracer + from strix.telemetry.scan_store import ScanStore 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} - for v in tracer.vulnerability_reports: + for v in scan_store.vulnerability_reports: sev = v.get("severity", "info").lower() if sev in vulnerabilities_counts: vulnerabilities_counts[sev] += 1 - llm = tracer.get_total_llm_stats() - total = llm.get("total", {}) - duration = 0.0 try: from datetime import datetime - start = datetime.fromisoformat(tracer.start_time.replace("Z", "+00:00")) - end_iso = tracer.end_time or datetime.now(start.tzinfo).isoformat() + start = datetime.fromisoformat(scan_store.start_time.replace("Z", "+00:00")) + end_iso = scan_store.end_time or datetime.now(start.tzinfo).isoformat() duration = (datetime.fromisoformat(end_iso.replace("Z", "+00:00")) - start).total_seconds() except (ValueError, TypeError, AttributeError): pass @@ -140,12 +137,8 @@ def end(tracer: "Tracer", exit_reason: str = "completed") -> None: **_base_props(), "exit_reason": exit_reason, "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()}, - "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), }, ) diff --git a/strix/telemetry/scan_artifacts.py b/strix/telemetry/scan_artifacts.py deleted file mode 100644 index dee9b8e..0000000 --- a/strix/telemetry/scan_artifacts.py +++ /dev/null @@ -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//``, -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) diff --git a/strix/telemetry/tracer.py b/strix/telemetry/scan_store.py similarity index 51% rename from strix/telemetry/tracer.py rename to strix/telemetry/scan_store.py index f6dbde2..1170eed 100644 --- a/strix/telemetry/tracer.py +++ b/strix/telemetry/scan_store.py @@ -1,5 +1,7 @@ +import csv import json import logging +import tempfile from collections.abc import Callable from datetime import UTC, datetime from pathlib import Path @@ -7,32 +9,31 @@ from typing import Any, Optional from uuid import uuid4 from strix.telemetry import posthog -from strix.telemetry.scan_artifacts import ScanArtifactWriter 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"]: - return _global_tracer +def get_global_scan_store() -> Optional["ScanStore"]: + return _global_scan_store -def set_global_tracer(tracer: "Tracer") -> None: - global _global_tracer # noqa: PLW0603 - _global_tracer = tracer +def set_global_scan_store(scan_store: "ScanStore") -> None: + global _global_scan_store # noqa: PLW0603 + _global_scan_store = scan_store -class Tracer: - """Per-scan in-memory state the TUI renders + per-scan artifact writer. +class ScanStore: + """Per-scan product artifact state plus artifact writer. - Holds live state the TUI reads (chat messages, agent tree, tool - executions, vulnerability reports, LLM usage). Writes vulnerability - markdown + CSV + final pentest report to ``strix_runs//``. + The Agents SDK owns model/tool execution, tracing, and conversation + persistence. This store keeps only Strix-owned scan artifacts and + report metadata. Live UI projections belong to the interface layer. - Conversation history goes to the SDK's ``SQLiteSession`` instead; - SDK trace events are not persisted here. + It is not a trace mirror and does not consume SDK tracing processors. """ def __init__(self, run_name: str | None = None): @@ -41,23 +42,9 @@ class Tracer: self.start_time = datetime.now(UTC).isoformat() 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.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_config: dict[str, Any] | None = None self.run_metadata: dict[str, Any] = { @@ -69,8 +56,7 @@ class Tracer: "status": "running", } self._run_dir: Path | None = None - self._writer: ScanArtifactWriter | None = None - self._next_message_id = 1 + self._saved_vuln_ids: set[str] = set() self.caido_url: str | 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_id"] = run_name self._run_dir = None - self._writer = None + self._saved_vuln_ids.clear() def get_run_dir(self) -> Path: if self._run_dir is None: @@ -135,11 +121,7 @@ class Tracer: if k in {"run_id", "run_name", "start_time", "targets", "status"} }, ) - logger.info( - "tracer hydrated run_metadata from %s (start_time=%s)", - meta_path, - self.start_time, - ) + logger.info("scan store hydrated run_metadata from %s", meta_path) json_path = run_dir / "vulnerabilities.json" if json_path.exists(): @@ -156,86 +138,14 @@ class Tracer: f"vulnerabilities.json at {json_path} is not a list", ) self.vulnerability_reports = [r for r in data if isinstance(r, dict)] - writer = self._get_writer() for r in self.vulnerability_reports: rid = r.get("id") if isinstance(rid, str): - writer._saved_vuln_ids.add(rid) + self._saved_vuln_ids.add(rid) logger.info( - "tracer hydrated %d vulnerability report(s) from %s", - len(self.vulnerability_reports), - json_path, + "scan store hydrated %d vulnerability report(s)", len(self.vulnerability_reports) ) - 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( self, title: str, @@ -254,6 +164,8 @@ class Tracer: cve: str | None = None, cwe: str | None = None, code_locations: list[dict[str, Any]] | None = None, + agent_id: str | None = None, + agent_name: str | None = None, ) -> str: report_id = f"vuln-{len(self.vulnerability_reports) + 1:04d}" @@ -292,6 +204,10 @@ class Tracer: report["cwe"] = cwe.strip() if 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) logger.info(f"Added vulnerability report: {report_id} - {title}") @@ -343,28 +259,6 @@ class Tracer: self.save_run_data(mark_complete=True) 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: self.scan_config = config self.run_metadata.update( @@ -382,91 +276,180 @@ class Tracer: self.run_metadata["end_time"] = self.end_time self.run_metadata["status"] = "completed" - self._get_writer().save( - 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 + self._save_artifacts() def cleanup(self) -> None: 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) diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index 38527c2..50684b1 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -1,13 +1,10 @@ """Multi-agent graph tools backed by the SDK-native :class:`AgentCoordinator`. - ``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. - ``wait_for_message``: pause this agent until a message arrives or ``timeout_seconds`` elapses. -- ``create_agent``: spawn a child via - ``asyncio.create_task(Runner.run(...))``; the task handle is stored - so a root-level cancel cascades to descendants. +- ``create_agent``: asks the scan runner to spawn an addressable child. - ``agent_finish``: subagents only โ€” posts a structured completion report to the parent's SDK session and returns a final-output marker. """ @@ -19,27 +16,11 @@ import json import logging import uuid from datetime import UTC, datetime -from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal -from agents import RunConfig, RunContextWrapper, function_tool -from agents.items import TResponseInputItem -from agents.model_settings import ModelSettings -from agents.sandbox import SandboxRunConfig +from agents import RunContextWrapper, function_tool -from strix.llm.multi_provider_setup import build_multi_provider -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 +from strix.orchestration.coordinator import coordinator_from_context 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) async def send_message_to_agent( ctx: RunContextWrapper, @@ -191,8 +136,9 @@ async def send_message_to_agent( ) -> str: """Send a message to another agent's inbox โ€” sparingly. - Inter-agent messages are surfaced at the top of the target's next - LLM turn. Use only when essential: + Inter-agent messages are appended to the target's SDK session and + interrupt any active target turn so the next run cycle sees them. + Use only when essential: - Sharing a discovered finding/credential another agent needs. - Asking a specialist a focused question. @@ -202,8 +148,8 @@ async def send_message_to_agent( **Don't** use for routine "hello/status" pings, for context the target already has (children inherit parent history), or when parent/child completion via ``agent_finish`` already covers the - flow. Messages to any registered agent are queued, regardless of - status, so a follow-up can wake a completed/stopped/failed agent. + flow. Messages to any registered agent wake it, regardless of + status, so a follow-up can restart a completed/stopped/failed agent. Args: target_agent_id: Recipient's 8-char id. @@ -249,7 +195,7 @@ async def send_message_to_agent( "success": True, "message_id": msg_id, "target_agent_id": target_agent_id, - "delivery_status": "queued", + "delivery_status": "delivered", }, ensure_ascii=False, default=str, @@ -269,7 +215,7 @@ def _session_items_payload(items: list[Any]) -> list[dict[str, Any]]: @function_tool(timeout=601) -async def wait_for_message( +async def wait_for_message( # noqa: PLR0911 ctx: RunContextWrapper, reason: str = "Waiting for messages from other agents", timeout_seconds: int = 600, @@ -325,10 +271,8 @@ async def wait_for_message( default=str, ) - pending = await coordinator.pending_count(me) + pending, items = await coordinator.consume_pending(me, include_items=True) if pending > 0: - items = await coordinator.recent_session_items(me, pending) - await coordinator.consume_wake(me) await coordinator.mark_running(me) return json.dumps( { @@ -344,7 +288,6 @@ async def wait_for_message( if interactive: await coordinator.park_waiting(me) - inner["agent_waiting_called"] = True return json.dumps( { "success": True, @@ -388,9 +331,7 @@ async def wait_for_message( default=str, ) - pending = await coordinator.pending_count(me) - items = await coordinator.recent_session_items(me, pending) - await coordinator.consume_wake(me) + pending, items = await coordinator.consume_pending(me, include_items=True) await coordinator.mark_running(me) return json.dumps( @@ -456,7 +397,7 @@ async def create_agent( inner = _ctx(ctx) coordinator = coordinator_from_context(inner) 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: return json.dumps( @@ -464,155 +405,38 @@ async def create_agent( ensure_ascii=False, default=str, ) - if factory is None: + if not callable(spawner): return json.dumps( { "success": False, - "error": ( - "No agent_factory in context. " - "The root assembly must inject one when building the run context." - ), + "error": "Scan runner did not provide a child-agent spawner in context.", }, ensure_ascii=False, 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 - # 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. + # including the call that's currently invoking ``create_agent``. parent_history = list(ctx.turn_input) if inherit_context and ctx.turn_input else [] - initial_input: list[TResponseInputItem] = [] - if parent_history: - rendered = json.dumps(parent_history, ensure_ascii=False, default=str) - initial_input.append( - { - "role": "user", - "content": ( - "== Inherited context from parent (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." - ), - }, + try: + result = await spawner( + parent_ctx=inner, + name=name, + task=task, + skills=list(skills or []), + parent_history=parent_history, + ) + except Exception as e: + logger.exception("create_agent: scan runner failed to spawn child '%s'", name) + return json.dumps( + {"success": False, "error": f"child spawn failed: {e!s}"}, + ensure_ascii=False, + default=str, ) - 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( "create_agent: spawned %s (%s) parent=%s skills=%d task_len=%d", - child_id, + result.get("agent_id"), name, parent_id or "-", len(skills or []), @@ -620,13 +444,7 @@ async def create_agent( ) return json.dumps( - { - "success": True, - "agent_id": child_id, - "name": name, - "parent_id": parent_id, - "message": f"Spawned '{name}' ({child_id}) running in parallel.", - }, + result, ensure_ascii=False, default=str, ) @@ -700,10 +518,6 @@ async def agent_finish( 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 if report_to_parent: async with coordinator._lock: @@ -736,6 +550,7 @@ async def agent_finish( len(findings or []), parent_notified, ) + await coordinator.set_status(me, "completed") return json.dumps( { @@ -798,7 +613,8 @@ async def stop_agent( ensure_ascii=False, 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( {"success": False, "error": f"Unknown agent_id: {target_agent_id}"}, ensure_ascii=False, diff --git a/strix/tools/finish/tool.py b/strix/tools/finish/tool.py index e0207c6..dc6d456 100644 --- a/strix/tools/finish/tool.py +++ b/strix/tools/finish/tool.py @@ -44,24 +44,24 @@ def _do_finish( return {"success": False, "message": "Validation failed", "errors": errors} try: - from strix.telemetry.tracer import get_global_tracer + from strix.telemetry.scan_store import get_global_scan_store - tracer = get_global_tracer() - if tracer is None: - logger.warning("No global tracer; scan results not persisted") + scan_store = get_global_scan_store() + if scan_store is None: + logger.warning("No global scan store; scan results not persisted") return { "success": True, "scan_completed": True, "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(), methodology=methodology.strip(), technical_analysis=technical_analysis.strip(), recommendations=recommendations.strip(), ) - vuln_count = len(tracer.vulnerability_reports) + vuln_count = len(scan_store.vulnerability_reports) except (ImportError, AttributeError) as e: logger.exception("finish_scan persistence failed") 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):** 1. **Call ``view_agent_graph`` first.** Inspect every entry in the - summary. If ANY agent is in ``running`` / ``waiting`` / - ``llm_failed`` state, you MUST NOT call ``finish_scan`` yet โ€” + summary. If ANY agent is in ``running`` / ``waiting`` state, + you MUST NOT call ``finish_scan`` yet โ€” wrap them up first via ``send_message_to_agent`` (ask them to finish), ``wait_for_message`` (block until their report arrives), or ``stop_agent`` (graceful cancel). Only ``completed`` @@ -178,6 +178,11 @@ async def finish_scan( technical_analysis=technical_analysis, recommendations=recommendations, ) - if result.get("success") and result.get("scan_completed"): - inner["finish_scan_called"] = True + if ( + 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) diff --git a/strix/tools/reporting/tool.py b/strix/tools/reporting/tool.py index 4b01787..68d9f97 100644 --- a/strix/tools/reporting/tool.py +++ b/strix/tools/reporting/tool.py @@ -151,7 +151,7 @@ _REQUIRED_FIELDS = { } -async def _do_create( # noqa: PLR0912 +async def _do_create( # noqa: PLR0912, PLR0915 *, title: str, description: str, @@ -167,6 +167,8 @@ async def _do_create( # noqa: PLR0912 cve: str | None, cwe: str | None, code_locations: list[dict[str, Any]] | None, + agent_id: str | None = None, + agent_name: str | None = None, ) -> dict[str, Any]: errors: list[str] = [] fields = { @@ -212,20 +214,20 @@ async def _do_create( # noqa: PLR0912 cvss_score, severity, _vector = _calculate_cvss(cvss_breakdown) try: - from strix.telemetry.tracer import get_global_tracer + from strix.telemetry.scan_store import get_global_scan_store - tracer = get_global_tracer() - if tracer is None: - logger.warning("No global tracer; vulnerability report not persisted") + scan_store = get_global_scan_store() + if scan_store is None: + logger.warning("No global scan store; vulnerability report not persisted") return { "success": True, "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 - existing = tracer.get_existing_vulnerabilities() + existing = scan_store.get_existing_vulnerabilities() candidate = { "title": title, "description": description, @@ -256,7 +258,7 @@ async def _do_create( # noqa: PLR0912 "reason": dedupe.get("reason", ""), } - report_id = tracer.add_vulnerability_report( + report_id = scan_store.add_vulnerability_report( title=title, description=description, severity=severity, @@ -273,6 +275,8 @@ async def _do_create( # noqa: PLR0912 cve=cve, cwe=cwe, 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: logger.exception("create_vulnerability_report persistence failed") @@ -404,6 +408,17 @@ async def create_vulnerability_report( ``fix_before`` (verbatim source), ``fix_after`` (suggested 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( title=title, description=description, @@ -419,5 +434,7 @@ async def create_vulnerability_report( cve=cve, cwe=cwe, code_locations=code_locations, + agent_id=agent_id, + agent_name=agent_name, ) return json.dumps(result, ensure_ascii=False, default=str)