diff --git a/pyproject.toml b/pyproject.toml index 66e2016..f31a8cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -303,6 +303,15 @@ ignore = [ # CaidoCapability uses agents.tool.Tool at runtime — pydantic Field # annotations and the cached _CAIDO_TOOLS tuple need it eagerly. "strix/sandbox/caido_capability.py" = ["TC002"] +# Agent factory: agents.tool.Tool is used at runtime in the _BASE_TOOLS +# tuple type, not just for annotations. +"strix/agents/sdk_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 +# walks every supported target type — splitting it into per-type +# helpers would add indirection without simplifying anything. +"strix/sdk_entry.py" = ["TC003", "PLR0912"] # Sandbox dispatch helper has many short-circuit error returns (auth fail, # size cap, decode fail, etc). Each is a distinct, documented failure mode # the model needs to see verbatim — collapsing them harms readability. diff --git a/strix/agents/sdk_factory.py b/strix/agents/sdk_factory.py new file mode 100644 index 0000000..1f62213 --- /dev/null +++ b/strix/agents/sdk_factory.py @@ -0,0 +1,221 @@ +"""build_strix_agent — assemble an ``agents.Agent`` for root or child runs. + +This is the keystone that links Phase 2's SDK function tools, Phase 3's +graph tools, Phase 4's CaidoCapability, and the rendered Jinja prompt +from :mod:`strix.agents.sdk_prompt` into a single ``agents.Agent`` +instance ready for ``Runner.run``. + +Two flavors: + +- **Root** (``is_root=True``): the top-level scan agent. Carries + ``finish_scan`` (terminates the scan), no ``agent_finish`` (that's + for subagents). ``tool_use_behavior`` stops on ``finish_scan`` so + the model can't accidentally keep talking after marking the scan + complete. + +- **Child** (``is_root=False``): subagents spawned by the + ``create_agent`` graph tool. Carries ``agent_finish``, no + ``finish_scan``. ``tool_use_behavior`` stops on ``agent_finish`` + (C4 — without this, the SDK loop would keep going to ``max_turns`` + even after the child reported back to its parent). + +Caido tools come from ``CaidoCapability.tools()`` automatically via +the SDK's capability merge — we don't include them here. Skills are +injected via the prompt; the model can also load more at runtime via +the ``load_skill`` tool. + +References: + - PLAYBOOK.md §4.3 (graph tool wiring) + - AUDIT.md §2.4 (C4 — stop_at_tool_names is required for subagents) +""" + +from __future__ import annotations + +import logging +from typing import Any + +from agents import Agent +from agents.agent import StopAtTools +from agents.tool import Tool + +from strix.agents.sdk_prompt import render_system_prompt +from strix.tools.agents_graph.agents_graph_sdk_tools import ( + agent_finish, + agent_status, + create_agent, + send_message_to_agent, + view_agent_graph, + wait_for_message, +) +from strix.tools.browser.browser_sdk_tool import browser_action +from strix.tools.file_edit.file_edit_sdk_tools import ( + list_files, + search_files, + str_replace_editor, +) +from strix.tools.finish.finish_sdk_tool import finish_scan +from strix.tools.load_skill.load_skill_sdk_tool import load_skill +from strix.tools.notes.notes_sdk_tools import ( + create_note, + delete_note, + get_note, + list_notes, + update_note, +) +from strix.tools.python.python_sdk_tool import python_action +from strix.tools.reporting.reporting_sdk_tools import create_vulnerability_report +from strix.tools.terminal.terminal_sdk_tool import terminal_execute +from strix.tools.thinking.thinking_sdk_tools import think +from strix.tools.todo.todo_sdk_tools import ( + create_todo, + delete_todo, + list_todos, + mark_todo_done, + mark_todo_pending, + update_todo, +) +from strix.tools.web_search.web_search_sdk_tool import web_search + + +logger = logging.getLogger(__name__) + + +# Tools every Strix agent has, root or child. The Caido proxy tools +# (list_requests, view_request, send_request, ...) are NOT here — +# CaidoCapability.tools() returns them and the SDK merges them in. +_BASE_TOOLS: tuple[Tool, ...] = ( + # Thinking + planning + think, + # Per-agent todos + create_todo, + list_todos, + update_todo, + mark_todo_done, + mark_todo_pending, + delete_todo, + # Shared notes (per-run JSONL store) + create_note, + list_notes, + get_note, + update_note, + delete_note, + # Web search (only registered if PERPLEXITY_API_KEY is set; the + # tool itself returns a structured error when not configured, so + # always exposing it is safe) + web_search, + # File edit (sandbox-bound) + str_replace_editor, + list_files, + search_files, + # Reporting + create_vulnerability_report, + # Skill loading + load_skill, + # Sandbox primitives + browser_action, + terminal_execute, + python_action, + # Multi-agent graph tools (the bus is in ctx.context) + view_agent_graph, + agent_status, + send_message_to_agent, + wait_for_message, + create_agent, +) + + +def build_strix_agent( + *, + name: str = "strix", + skills: list[str] | None = None, + is_root: bool, + scan_mode: str = "deep", + is_whitebox: bool = False, + interactive: bool = False, + system_prompt_context: dict[str, Any] | None = None, +) -> Agent[Any]: + """Build an ``agents.Agent`` configured for either root or child use. + + Args: + name: Agent name. Surfaces in traces and the bus's ``names`` map. + Defaults to ``"strix"`` for the root; create_agent passes + distinct names per child. + skills: Skills to preload into the system prompt. The agent can + also load more at runtime via the ``load_skill`` tool. + is_root: Selects the tool list and ``tool_use_behavior``. + Root carries ``finish_scan`` and stops there; child carries + ``agent_finish`` and stops there. + scan_mode: ``"deep"`` etc.; routes the scan-mode skill section + of the prompt template. + is_whitebox: Whitebox source-aware mode toggle. Adds two extra + skills to the prompt and gates whitebox-only behavior in + the create_agent / wiki integration. + interactive: Renders the interactive-mode communication block + in the system prompt. + system_prompt_context: Free-form dict the prompt template + renders into the ``system_prompt_context`` variable — + today carries the scan scope / authorization block. + + Returns the ``Agent`` instance with ``model=None`` so the + ``RunConfig.model`` (built by ``make_run_config``) drives provider + selection. ``agents.Agent`` is generic on context type; we let + the caller's ``Runner.run(context=...)`` typing determine that. + """ + instructions = render_system_prompt( + skills=skills, + scan_mode=scan_mode, + is_whitebox=is_whitebox, + interactive=interactive, + system_prompt_context=system_prompt_context, + ) + + # Tool list + termination tool depend on is_root. The tuple-then- + # list dance keeps _BASE_TOOLS immutable so concurrent agent builds + # can't accidentally mutate each other's tool list. + if is_root: + tools: list[Tool] = [*_BASE_TOOLS, finish_scan] + stop_at = ("finish_scan",) + else: + tools = [*_BASE_TOOLS, agent_finish] + stop_at = ("agent_finish",) + + return Agent( + name=name, + instructions=instructions, + tools=tools, + tool_use_behavior=StopAtTools(stop_at_tool_names=list(stop_at)), + # model=None so ``RunConfig.model`` (e.g. ``strix/claude-sonnet-4.6``) + # routes through MultiProvider rather than the SDK's default. + model=None, + ) + + +def make_child_factory( + *, + scan_mode: str = "deep", + is_whitebox: bool = False, + interactive: bool = False, + system_prompt_context: dict[str, Any] | None = None, +) -> Any: + """Return a callable suitable for ``ctx.context['agent_factory']``. + + The Phase 3 ``create_agent`` graph tool reads + ``ctx.context['agent_factory']`` and calls it with ``name=`` and + ``skills=`` to build a child Agent. We snapshot the run-level + arguments (scan_mode, is_whitebox, etc.) into a closure so each + child inherits the right scan-level configuration without the + create_agent tool having to know about them. + """ + + def _factory(*, name: str, skills: list[str]) -> Agent[Any]: + return build_strix_agent( + name=name, + skills=skills, + is_root=False, + scan_mode=scan_mode, + is_whitebox=is_whitebox, + interactive=interactive, + system_prompt_context=system_prompt_context, + ) + + return _factory diff --git a/strix/agents/sdk_prompt.py b/strix/agents/sdk_prompt.py new file mode 100644 index 0000000..f77e964 --- /dev/null +++ b/strix/agents/sdk_prompt.py @@ -0,0 +1,130 @@ +"""Standalone Jinja-based system-prompt renderer for SDK agents. + +The legacy ``LLM._load_system_prompt`` couples prompt rendering to the +LLM client class. The SDK migration owns the model client through +``MultiProvider`` instead, so we extract the rendering logic into a +plain function that the SDK agent factory can call without pulling in +the legacy ``LLM`` instance. + +Reuses the existing Jinja template at +``strix/agents/StrixAgent/system_prompt.jinja`` (508 lines, expanding +into the multi-section prompt with skills, tools, scan modes, etc.) so +behavior parity is preserved verbatim — only the call site changes. + +References: + - HARNESS_WIKI.md §4.1 (system prompt assembly) + - PLAYBOOK.md §4 (per-tool migration contracts) +""" + +from __future__ import annotations + +import logging +from typing import Any + +from jinja2 import Environment, FileSystemLoader, select_autoescape + +from strix.skills import load_skills +from strix.tools import get_tools_prompt +from strix.utils.resource_paths import get_strix_resource_path + + +logger = logging.getLogger(__name__) + + +# Hard-coded to the StrixAgent template since it's the only agent type +# under the SDK migration. The legacy harness supported multiple agent +# names but in practice only StrixAgent ships. +_AGENT_NAME = "StrixAgent" + + +def _resolve_skills( + *, + requested: list[str] | None, + scan_mode: str = "deep", + is_whitebox: bool = False, +) -> list[str]: + """Build the deduped, ordered skills list for the prompt render. + + Mirrors :py:meth:`LLM._get_skills_to_load` exactly so the rendered + prompt is byte-identical to the legacy path: + + 1. Whatever the caller asked for, in order. + 2. ``scan_modes/`` (always). + 3. Whitebox-specific skills if applicable. + """ + ordered: list[str] = list(requested or []) + ordered.append(f"scan_modes/{scan_mode}") + if is_whitebox: + ordered.append("coordination/source_aware_whitebox") + ordered.append("custom/source_aware_sast") + + deduped: list[str] = [] + seen: set[str] = set() + for skill in ordered: + if skill and skill not in seen: + deduped.append(skill) + seen.add(skill) + return deduped + + +def render_system_prompt( + *, + skills: list[str] | None = None, + scan_mode: str = "deep", + is_whitebox: bool = False, + interactive: bool = False, + system_prompt_context: dict[str, Any] | None = None, +) -> str: + """Render the StrixAgent system prompt. + + Args: + skills: Skills the caller wants preloaded into the prompt + context (the agent can also load more at runtime via the + ``load_skill`` tool). + scan_mode: ``"deep" | "fast" | ...``. Maps to ``scan_modes/`` + skill. + is_whitebox: When True, the source-aware whitebox skill stack + is loaded too. + interactive: When True, the prompt renders the interactive-mode + communication rules block. + system_prompt_context: Free-form dict that the template's + ``system_prompt_context`` variable receives — used today for + the scan-scope authorization block from + :py:meth:`StrixAgent._build_system_scope_context`. + + Returns the rendered prompt string. If anything goes wrong (template + missing, render failure), returns an empty string and logs — same + fail-soft posture as the legacy method, because a missing prompt is + survivable but a hard failure during agent construction is not. + """ + try: + prompt_dir = get_strix_resource_path("agents", _AGENT_NAME) + skills_dir = get_strix_resource_path("skills") + env = Environment( + loader=FileSystemLoader([prompt_dir, skills_dir]), + autoescape=select_autoescape( + enabled_extensions=(), + default_for_string=False, + ), + ) + + skills_to_load = _resolve_skills( + requested=skills, + scan_mode=scan_mode, + is_whitebox=is_whitebox, + ) + skill_content = load_skills(skills_to_load) + env.globals["get_skill"] = lambda name: skill_content.get(name, "") + + rendered = env.get_template("system_prompt.jinja").render( + get_tools_prompt=get_tools_prompt, + loaded_skill_names=list(skill_content.keys()), + interactive=interactive, + system_prompt_context=system_prompt_context or {}, + **skill_content, + ) + except Exception: + logger.exception("render_system_prompt failed; returning empty prompt") + return "" + else: + return str(rendered) diff --git a/strix/sdk_entry.py b/strix/sdk_entry.py new file mode 100644 index 0000000..0967386 --- /dev/null +++ b/strix/sdk_entry.py @@ -0,0 +1,282 @@ +"""Top-level SDK scan entry point. + +Replaces the legacy ``strix.cli.main → StrixAgent.execute_scan`` +pipeline with the SDK-native equivalent: + +1. Build the per-scan ``AgentMessageBus``. +2. Bring up (or reuse) a sandbox session for ``scan_id`` via the + :mod:`strix.sandbox.session_manager`. +3. Build the root ``Agent`` via :func:`build_strix_agent` and a + matching child factory via :func:`make_child_factory`. +4. Build the root context dict (bus + sandbox bundle + agent_factory). +5. Register the root in the bus. +6. Build the ``RunConfig`` via the factory. +7. Call ``Runner.run(...)`` and surface the result. +8. ``finally`` cleanup the sandbox session. + +Phase 5 lands the wiring; the streaming accumulator + TUI integration +land in Phase 5b. The entry point is intentionally not wired to the +CLI yet — that's a follow-up under ``STRIX_USE_SDK_HARNESS=1`` (see +PLAYBOOK §7.1 cutover plan). + +References: + - PLAYBOOK.md §3.3 (session manager), §4.3 (graph tools), §7.1 + - AUDIT_R3.md C9 (cancel_descendants on cleanup) +""" + +from __future__ import annotations + +import logging +import uuid +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from agents import Runner + +from strix.agents.sdk_factory import build_strix_agent, make_child_factory +from strix.orchestration.bus import AgentMessageBus +from strix.orchestration.hooks import StrixOrchestrationHooks +from strix.run_config_factory import ( + STRIX_DEFAULT_MAX_TURNS, + make_agent_context, + make_run_config, +) +from strix.sandbox import session_manager + + +if TYPE_CHECKING: + from agents.result import RunResult + + +logger = logging.getLogger(__name__) + + +def _build_root_task(scan_config: dict[str, Any]) -> str: + """Format the user-facing task for the root agent. + + Mirrors :py:meth:`StrixAgent.execute_scan` (legacy) — collects each + target type into a labelled section, appends diff-scope context if + active, and tacks on user_instructions. The structured shape is + important for prompt parity: the system prompt template references + these section headers. + """ + 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. + + Same shape as the legacy + :py:meth:`StrixAgent._build_system_scope_context` so the prompt + template's ``system_prompt_context.authorized_targets`` lookups + stay byte-identical. + """ + 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, + sources_path: Path, + tracer: Any | None = None, + interactive: bool = False, + max_turns: int = STRIX_DEFAULT_MAX_TURNS, + cleanup_on_exit: bool = True, +) -> RunResult: + """Run one Strix scan end-to-end against a freshly-prepared sandbox. + + Args: + scan_config: Same shape the legacy ``StrixAgent.execute_scan`` + takes (targets, user_instructions, diff_scope, scan_mode, + is_whitebox, skills). + 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.1.13"``). + sources_path: Host directory mounted into ``/workspace/sources``. + 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. Mirrors legacy + ``AgentState.max_iterations`` (300). + 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]}" + logger.info("Starting Strix scan %s", scan_id) + + bus = AgentMessageBus() + root_id = uuid.uuid4().hex[:8] + + bundle = await session_manager.create_or_reuse( + scan_id, + image=image, + sources_path=sources_path, + ) + + try: + scan_mode = str(scan_config.get("scan_mode") or "deep") + is_whitebox = bool(scan_config.get("is_whitebox", False)) + 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, + ) + + await bus.register(root_id, "strix", parent_id=None) + + agent_factory = make_child_factory( + scan_mode=scan_mode, + is_whitebox=is_whitebox, + interactive=interactive, + system_prompt_context=scope_context, + ) + + context = make_agent_context( + bus=bus, + sandbox_session=bundle["session"], + sandbox_client=bundle["client"], + sandbox_token=bundle["bearer"], + tool_server_host_port=bundle["tool_server_host_port"], + caido_host_port=bundle["caido_host_port"], + agent_id=root_id, + agent_name="strix", + parent_id=None, + tracer=tracer, + max_turns=max_turns, + is_whitebox=is_whitebox, + diff_scope=diff_scope, + run_id=run_id, + agent_factory=agent_factory, + ) + + run_config = make_run_config( + sandbox_session=bundle["session"], + sandbox_client=bundle["client"], + ) + + task_text = _build_root_task(scan_config) + return await Runner.run( + root_agent, + input=task_text, + run_config=run_config, + context=context, + hooks=StrixOrchestrationHooks(), + max_turns=max_turns, + ) + except BaseException: + # Cancel any descendant tasks the root spawned before unwinding. + # cancel_descendants is idempotent and handles the empty-tree case. + await bus.cancel_descendants(root_id) + raise + finally: + if cleanup_on_exit: + await session_manager.cleanup(scan_id) diff --git a/tests/agents/test_sdk_factory.py b/tests/agents/test_sdk_factory.py new file mode 100644 index 0000000..452dcb9 --- /dev/null +++ b/tests/agents/test_sdk_factory.py @@ -0,0 +1,201 @@ +"""Phase 5 tests for the SDK agent factory + prompt renderer. + +These two modules are the keystone wiring between Phases 2-4 and an +actual ``Runner.run`` invocation. The tests verify: + +- The prompt renderer reuses the existing Jinja template (parity with + legacy LLM._load_system_prompt) and degrades gracefully when the + template isn't available. +- ``build_strix_agent(is_root=True)`` carries ``finish_scan`` and + stops on it; child agents carry ``agent_finish`` and stop on it. +- ``make_child_factory`` snapshots scan-level config into a closure + so each spawned child inherits the right scan_mode / is_whitebox / + prompt context without create_agent having to re-derive it. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import patch + +from agents import Agent +from agents.tool import FunctionTool + +from strix.agents.sdk_factory import build_strix_agent, make_child_factory +from strix.agents.sdk_prompt import _resolve_skills, render_system_prompt + + +# --- prompt renderer ---------------------------------------------------- + + +def test_resolve_skills_deduplicates_and_orders() -> None: + out = _resolve_skills( + requested=["recon", "xss", "recon"], + scan_mode="deep", + is_whitebox=False, + ) + assert out == ["recon", "xss", "scan_modes/deep"] + + +def test_resolve_skills_adds_whitebox_pair() -> None: + out = _resolve_skills(requested=None, scan_mode="fast", is_whitebox=True) + # The whitebox pair sits at the tail; scan_modes goes in the middle + # because callers can append more skills after it via the requested arg. + assert out == [ + "scan_modes/fast", + "coordination/source_aware_whitebox", + "custom/source_aware_sast", + ] + + +def test_render_system_prompt_returns_string() -> None: + """Smoke: the StrixAgent template is on disk and renders to non-empty.""" + out = render_system_prompt(skills=[], scan_mode="deep") + assert isinstance(out, str) + # The first line of the template starts with 'You are Strix'. + assert out.startswith("You are Strix") + + +def test_render_system_prompt_swallows_template_errors() -> None: + """If the template path can't be resolved, return an empty string + (not raise) — agent construction must never blow up on prompt load.""" + with patch( + "strix.agents.sdk_prompt.get_strix_resource_path", + side_effect=RuntimeError("missing"), + ): + out = render_system_prompt(skills=[]) + assert out == "" + + +# --- factory: shape + tools -------------------------------------------- + + +def test_root_agent_carries_finish_scan_and_stops_there() -> None: + agent = build_strix_agent(name="strix", is_root=True) + assert isinstance(agent, Agent) + tool_names = {t.name for t in agent.tools if isinstance(t, FunctionTool)} + assert "finish_scan" in tool_names + assert "agent_finish" not in tool_names + behavior = agent.tool_use_behavior + # StopAtTools is a TypedDict at runtime → behavior is a dict. + assert isinstance(behavior, dict) + assert behavior["stop_at_tool_names"] == ["finish_scan"] + + +def test_child_agent_carries_agent_finish_and_stops_there() -> None: + agent = build_strix_agent(name="recon-bot", is_root=False) + tool_names = {t.name for t in agent.tools if isinstance(t, FunctionTool)} + assert "agent_finish" in tool_names + assert "finish_scan" not in tool_names + behavior = agent.tool_use_behavior + assert isinstance(behavior, dict) + assert behavior["stop_at_tool_names"] == ["agent_finish"] + + +def test_root_and_child_share_base_tool_set() -> None: + """The base tool set (think/todo/notes/file_edit/web_search/etc) is + identical between root and child — only the terminator differs.""" + root = build_strix_agent(is_root=True) + child = build_strix_agent(is_root=False) + root_names = {t.name for t in root.tools if isinstance(t, FunctionTool)} + child_names = {t.name for t in child.tools if isinstance(t, FunctionTool)} + # Drop the terminators and compare. + assert root_names - {"finish_scan"} == child_names - {"agent_finish"} + + +def test_agent_includes_graph_and_sandbox_tools() -> None: + """The graph + sandbox tool families are required for parity with + legacy. Spot-check the ones most likely to be forgotten in a refactor.""" + agent = build_strix_agent(is_root=True) + names = {t.name for t in agent.tools if isinstance(t, FunctionTool)} + expected = { + "think", + "create_todo", + "create_note", + "web_search", + "str_replace_editor", + "create_vulnerability_report", + "load_skill", + "browser_action", + "terminal_execute", + "python_action", + "view_agent_graph", + "agent_status", + "send_message_to_agent", + "wait_for_message", + "create_agent", + } + missing = expected - names + assert not missing, f"missing tools: {missing}" + + +def test_agent_does_not_include_caido_tools() -> None: + """Caido tools come from CaidoCapability.tools(); the agent doesn't + declare them directly to avoid double-registration when the SDK + runtime merges capability tools.""" + agent = build_strix_agent(is_root=True) + names = {t.name for t in agent.tools if isinstance(t, FunctionTool)} + caido = { + "list_requests", + "view_request", + "send_request", + "repeat_request", + "scope_rules", + "list_sitemap", + "view_sitemap_entry", + } + overlap = names & caido + assert overlap == set(), f"unexpected Caido tools in agent.tools: {overlap}" + + +def test_agent_uses_run_config_model() -> None: + """``model=None`` so the RunConfig drives the model alias through + MultiProvider rather than an SDK default like gpt-4.1.""" + agent = build_strix_agent(is_root=True) + assert agent.model is None + + +def test_agent_instructions_contain_rendered_prompt() -> None: + """The factory must wire the rendered prompt into ``instructions``.""" + agent = build_strix_agent(is_root=True, scan_mode="deep") + assert isinstance(agent.instructions, str) + assert agent.instructions.startswith("You are Strix") + + +# --- child factory ------------------------------------------------------ + + +def test_make_child_factory_returns_callable_that_builds_child() -> None: + factory = make_child_factory(scan_mode="deep", is_whitebox=False) + assert callable(factory) + child = factory(name="sub-1", skills=["recon"]) + assert isinstance(child, Agent) + assert child.name == "sub-1" + behavior = child.tool_use_behavior + assert isinstance(behavior, dict) + assert behavior["stop_at_tool_names"] == ["agent_finish"] + + +def test_make_child_factory_passes_scan_level_config() -> None: + """Verify scan_mode + is_whitebox flow into the rendered prompt + via the closure rather than the create_agent call site.""" + captured: dict[str, Any] = {} + + def fake_render(**kwargs: Any) -> str: + captured.update(kwargs) + return "stub-prompt" + + factory = make_child_factory( + scan_mode="fast", + is_whitebox=True, + interactive=True, + system_prompt_context={"scope_source": "test"}, + ) + with patch("strix.agents.sdk_factory.render_system_prompt", side_effect=fake_render): + factory(name="child", skills=["xss"]) + + assert captured["scan_mode"] == "fast" + assert captured["is_whitebox"] is True + assert captured["interactive"] is True + assert captured["system_prompt_context"] == {"scope_source": "test"} + assert captured["skills"] == ["xss"] diff --git a/tests/test_sdk_entry.py b/tests/test_sdk_entry.py new file mode 100644 index 0000000..c9e383f --- /dev/null +++ b/tests/test_sdk_entry.py @@ -0,0 +1,319 @@ +"""Phase 5 tests for the top-level SDK scan entry point. + +We never spin up a real Docker container or hit a real LLM here. The +tests patch ``session_manager.create_or_reuse``, ``Runner.run``, and +the agent factory so we can verify the wiring shape: + +- The bus is registered with a root agent before Runner.run. +- The context dict carries every field downstream code (tools, hooks, + filter) reads. +- The session manager's bundle flows through to the context (host + ports, bearer, sandbox session/client). +- ``cleanup_on_exit=True`` always cleans up, even when Runner.run + raises. +- ``cleanup_on_exit=False`` preserves the cached session. +- Cancellation propagates: if Runner.run raises, descendants are + cancelled before re-raising. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from strix.orchestration.bus import AgentMessageBus +from strix.sdk_entry import _build_root_task, _build_scope_context, run_strix_scan + + +# --- helpers ------------------------------------------------------------ + + +def _bundle_for_test() -> dict[str, Any]: + return { + "client": MagicMock(name="docker_client"), + "session": MagicMock(name="sandbox_session"), + "capability": MagicMock(), + "tool_server_host_port": 12001, + "caido_host_port": 12002, + "bearer": "test-bearer-token-1234567890", + } + + +def _scan_config(**overrides: Any) -> dict[str, Any]: + base: dict[str, Any] = { + "targets": [ + { + "type": "web_application", + "details": {"target_url": "https://example.com"}, + }, + ], + "user_instructions": "find xss", + "scan_mode": "deep", + "is_whitebox": False, + } + base.update(overrides) + return base + + +# --- task / scope builders --------------------------------------------- + + +def test_build_root_task_groups_targets_and_appends_instructions() -> None: + config = _scan_config( + targets=[ + { + "type": "repository", + "details": { + "target_repo": "https://github.com/x/y", + "cloned_repo_path": "/tmp/y", + "workspace_subdir": "y", + }, + }, + { + "type": "ip_address", + "details": {"target_ip": "10.0.0.1"}, + }, + ], + user_instructions="report only critical issues", + ) + task = _build_root_task(config) + assert "Repositories:" in task + assert "https://github.com/x/y (available at: /workspace/y)" in task + assert "IP Addresses:" in task + assert "10.0.0.1" in task + assert "Special instructions: report only critical issues" in task + + +def test_build_root_task_renders_diff_scope_block() -> None: + config = _scan_config( + diff_scope={ + "active": True, + "repos": [ + { + "workspace_subdir": "service-x", + "analyzable_files_count": 7, + "deleted_files_count": 2, + }, + ], + }, + ) + task = _build_root_task(config) + assert "Scope Constraints:" in task + assert "service-x: 7 changed file(s)" in task + assert "service-x: 2 deleted file(s)" in task + + +def test_build_scope_context_marks_authorization_source() -> None: + config = _scan_config( + targets=[ + { + "type": "web_application", + "details": {"target_url": "https://target.test"}, + }, + ], + ) + ctx = _build_scope_context(config) + assert ctx["scope_source"] == "system_scan_config" + assert ctx["authorization_source"] == "strix_platform_verified_targets" + assert ctx["user_instructions_do_not_expand_scope"] is True + assert ctx["authorized_targets"] == [ + {"type": "web_application", "value": "https://target.test", "workspace_path": ""}, + ] + + +# --- run_strix_scan wiring --------------------------------------------- + + +@pytest.mark.asyncio +async def test_run_strix_scan_wires_context_and_calls_runner(tmp_path: Path) -> None: + """End-to-end (mocked) — assert every downstream consumer of context + sees the bundle's bearer + host ports.""" + bundle = _bundle_for_test() + captured_context: dict[str, Any] = {} + + async def fake_runner_run(*args: Any, **kwargs: Any) -> Any: + captured_context.update(kwargs.get("context", {})) + return MagicMock(name="run_result") + + with ( + patch( + "strix.sdk_entry.session_manager.create_or_reuse", + new=AsyncMock(return_value=bundle), + ) as create_mock, + patch( + "strix.sdk_entry.session_manager.cleanup", + new=AsyncMock(), + ) as cleanup_mock, + patch("strix.sdk_entry.Runner.run", side_effect=fake_runner_run) as runner_mock, + # Stub the factory to avoid rendering the 158k-char prompt for + # every test (it's covered by sdk_prompt tests). + patch( + "strix.sdk_entry.build_strix_agent", + return_value=MagicMock(name="root_agent"), + ) as factory_mock, + ): + await run_strix_scan( + scan_config=_scan_config(), + scan_id="scan-test", + image="strix-sandbox:test", + sources_path=tmp_path, + ) + + # Session manager calls. + create_mock.assert_awaited_once() + create_args = create_mock.await_args + assert create_args is not None + assert create_args.args == ("scan-test",) + assert create_args.kwargs["image"] == "strix-sandbox:test" + assert create_args.kwargs["sources_path"] == tmp_path + cleanup_mock.assert_awaited_once_with("scan-test") + + # Factory called with is_root=True. + factory_mock.assert_called_once() + assert factory_mock.call_args.kwargs["is_root"] is True + + # Runner.run called once with the root agent. + assert runner_mock.call_count == 1 + + # Context shape passed into Runner.run. + assert captured_context["sandbox_session"] is bundle["session"] + assert captured_context["sandbox_client"] is bundle["client"] + assert captured_context["sandbox_token"] == bundle["bearer"] + assert captured_context["tool_server_host_port"] == bundle["tool_server_host_port"] + assert captured_context["caido_host_port"] == bundle["caido_host_port"] + # Bus is registered and root agent_id is populated. + bus = captured_context["bus"] + assert isinstance(bus, AgentMessageBus) + assert captured_context["agent_id"] in bus.statuses + assert bus.parent_of[captured_context["agent_id"]] is None + # Child factory wired through so create_agent works. + assert callable(captured_context["agent_factory"]) + + +@pytest.mark.asyncio +async def test_run_strix_scan_cleans_up_on_runner_failure(tmp_path: Path) -> None: + """If Runner.run raises, cleanup must still fire (the finally branch).""" + bundle = _bundle_for_test() + + with ( + patch( + "strix.sdk_entry.session_manager.create_or_reuse", + new=AsyncMock(return_value=bundle), + ), + patch( + "strix.sdk_entry.session_manager.cleanup", + new=AsyncMock(), + ) as cleanup_mock, + patch( + "strix.sdk_entry.Runner.run", + side_effect=RuntimeError("simulated LLM blow-up"), + ), + patch("strix.sdk_entry.build_strix_agent", return_value=MagicMock()), + pytest.raises(RuntimeError, match="simulated LLM"), + ): + await run_strix_scan( + scan_config=_scan_config(), + scan_id="scan-fail", + image="i", + sources_path=tmp_path, + ) + + cleanup_mock.assert_awaited_once_with("scan-fail") + + +@pytest.mark.asyncio +async def test_run_strix_scan_skips_cleanup_when_disabled(tmp_path: Path) -> None: + bundle = _bundle_for_test() + + async def fake_runner_run(*args: Any, **kwargs: Any) -> Any: + return MagicMock() + + with ( + patch( + "strix.sdk_entry.session_manager.create_or_reuse", + new=AsyncMock(return_value=bundle), + ), + patch( + "strix.sdk_entry.session_manager.cleanup", + new=AsyncMock(), + ) as cleanup_mock, + patch("strix.sdk_entry.Runner.run", side_effect=fake_runner_run), + patch("strix.sdk_entry.build_strix_agent", return_value=MagicMock()), + ): + await run_strix_scan( + scan_config=_scan_config(), + scan_id="scan-keep", + image="i", + sources_path=tmp_path, + cleanup_on_exit=False, + ) + + cleanup_mock.assert_not_called() + + +@pytest.mark.asyncio +async def test_run_strix_scan_auto_generates_scan_id(tmp_path: Path) -> None: + """A caller without a stable id should still get a valid scan_id + flowing into create_or_reuse.""" + bundle = _bundle_for_test() + captured_scan_id: list[str] = [] + + async def fake_create(scan_id: str, **_kwargs: Any) -> Any: + captured_scan_id.append(scan_id) + return bundle + + with ( + patch( + "strix.sdk_entry.session_manager.create_or_reuse", + new=AsyncMock(side_effect=fake_create), + ), + patch("strix.sdk_entry.session_manager.cleanup", new=AsyncMock()), + patch("strix.sdk_entry.Runner.run", new=AsyncMock(return_value=MagicMock())), + patch("strix.sdk_entry.build_strix_agent", return_value=MagicMock()), + ): + await run_strix_scan( + scan_config=_scan_config(), + image="i", + sources_path=tmp_path, + ) + + assert len(captured_scan_id) == 1 + assert captured_scan_id[0].startswith("scan-") + assert len(captured_scan_id[0]) > len("scan-") + + +@pytest.mark.asyncio +async def test_run_strix_scan_passes_scan_level_config_into_factory( + tmp_path: Path, +) -> None: + """scan_mode / is_whitebox flow from scan_config into both the + root factory call and the child factory closure.""" + bundle = _bundle_for_test() + factory_calls: list[dict[str, Any]] = [] + + def fake_factory(**kwargs: Any) -> Any: + factory_calls.append(kwargs) + return MagicMock() + + with ( + patch( + "strix.sdk_entry.session_manager.create_or_reuse", + new=AsyncMock(return_value=bundle), + ), + patch("strix.sdk_entry.session_manager.cleanup", new=AsyncMock()), + patch("strix.sdk_entry.Runner.run", new=AsyncMock(return_value=MagicMock())), + patch("strix.sdk_entry.build_strix_agent", side_effect=fake_factory), + ): + await run_strix_scan( + scan_config=_scan_config(scan_mode="fast", is_whitebox=True), + scan_id="s", + image="i", + sources_path=tmp_path, + ) + + assert factory_calls[0]["scan_mode"] == "fast" + assert factory_calls[0]["is_whitebox"] is True + assert factory_calls[0]["is_root"] is True