From df51eeedd09fadf14b65cc59ca13f37605950a17 Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sat, 25 Apr 2026 13:32:11 -0700 Subject: [PATCH] refactor: flatten CaidoCapability into direct wiring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The custom ``Capability`` subclass was 207 LoC bundling four tiny concerns (env-var injection, tool exposure, system-prompt block, healthcheck) — and three of them were dead code: the SDK's ``SandboxRunConfig`` doesn't accept capabilities, so ``process_manifest``, ``tools()``, and ``instructions()`` were never called. Only ``bind()`` ran, because we invoked it manually. Replace each piece with the obvious direct equivalent: - **Env vars**: inject ``http_proxy`` / ``https_proxy`` / ``ALL_PROXY`` directly into the manifest in ``session_manager.create_or_reuse``. This *also fixes a latent bug* — the proxy env vars in ``CaidoCapability.process_manifest`` weren't being applied to live containers, so shelled-out HTTP traffic from terminal/python tools wasn't actually flowing through Caido. - **Tool exposure**: add the seven Caido tools (``list_requests``, ``view_request``, ``send_request``, ``repeat_request``, ``scope_rules``, ``list_sitemap``, ``view_sitemap_entry``) to ``_BASE_TOOLS`` in ``agents/factory.py`` like every other sandbox tool. They were already defined in ``tools/proxy/tools.py``. - **Healthcheck**: ``entry.py`` now ``await``s ``wait_for_http_ready`` + ``wait_for_tcp_ready`` inline after ``session_manager.create_or_reuse`` returns, before any agent runs. No more capability state, ``configure_host_ports`` plumbing, or ``on_agent_start`` await-the-task indirection. - **Instructions block**: dropped. The seven proxy tools' docstrings cover the HTTPQL syntax and usage already; the duplicate prompt fragment was overhead. Cascade cleanups: - Drop ``caido_capability`` from the agent context (was passed to every ``make_agent_context`` call but only used by the now-deleted ``on_agent_start`` await). - Strip the capability await branch from ``StrixOrchestrationHooks.on_agent_start``; that hook now does only the ``tracer.agents`` mirroring it always should have. - Drop the ``capability`` key from the session bundle. - Drop ``strix/sandbox/caido_capability.py`` — entire file (207 LoC). - Drop the per-file ruff ignore for the deleted file. mypy clean on every touched file. Net -217 LoC. Co-Authored-By: Claude Opus 4.7 (1M context) --- pyproject.toml | 3 - strix/agents/factory.py | 31 +++-- strix/entry.py | 17 ++- strix/orchestration/hooks.py | 29 +---- strix/run_config_factory.py | 2 - strix/sandbox/__init__.py | 4 +- strix/sandbox/caido_capability.py | 207 ------------------------------ strix/sandbox/session_manager.py | 33 ++--- strix/tools/agents_graph/tools.py | 1 - 9 files changed, 55 insertions(+), 272 deletions(-) delete mode 100644 strix/sandbox/caido_capability.py diff --git a/pyproject.toml b/pyproject.toml index 30142fa..b01a853 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -249,9 +249,6 @@ ignore = [ "strix/tools/proxy/tools.py" = ["TC002"] "strix/tools/agents_graph/tools.py" = ["TC002"] "strix/agents/factory.py" = ["TC002"] -# 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"] # 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 diff --git a/strix/agents/factory.py b/strix/agents/factory.py index e03d56c..521e349 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -1,8 +1,7 @@ """``build_strix_agent`` — assemble an ``agents.Agent`` for root or child. -Wires the SDK function tools, multi-agent graph tools, -``CaidoCapability``, and the rendered Jinja prompt into one -``agents.Agent`` ready for ``Runner.run``. +Wires the SDK function tools, multi-agent graph tools, and the rendered +Jinja prompt into one ``agents.Agent`` ready for ``Runner.run``. Two flavors: @@ -13,9 +12,8 @@ Two flavors: there — without ``stop_at_tool_names`` the SDK loop would keep running to ``max_turns`` even after the child reported back. -Caido tools come from ``CaidoCapability.tools()`` via the SDK's -capability merge — we don't list them here. Skills are baked into the -system prompt at scan bring-up; there's no runtime skill-loading tool. +Skills are baked into the system prompt at scan bring-up; there's no +runtime skill-loading tool. """ from __future__ import annotations @@ -50,6 +48,15 @@ from strix.tools.notes.tools import ( list_notes, update_note, ) +from strix.tools.proxy.tools import ( + list_requests, + list_sitemap, + repeat_request, + scope_rules, + send_request, + view_request, + view_sitemap_entry, +) from strix.tools.python.tool import python_action from strix.tools.reporting.tool import create_vulnerability_report from strix.tools.terminal.tool import terminal_execute @@ -68,9 +75,7 @@ from strix.tools.web_search.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. +# Tools every Strix agent has, root or child. _BASE_TOOLS: tuple[Tool, ...] = ( # Thinking + planning think, @@ -101,6 +106,14 @@ _BASE_TOOLS: tuple[Tool, ...] = ( browser_action, terminal_execute, python_action, + # Caido HTTP/HTTPS proxy + list_requests, + view_request, + send_request, + repeat_request, + scope_rules, + list_sitemap, + view_sitemap_entry, # Multi-agent graph tools (the bus is in ctx.context) view_agent_graph, agent_status, diff --git a/strix/entry.py b/strix/entry.py index c00408a..8874d10 100644 --- a/strix/entry.py +++ b/strix/entry.py @@ -15,6 +15,7 @@ from __future__ import annotations +import asyncio import logging import uuid from pathlib import Path @@ -34,6 +35,7 @@ from strix.run_config_factory import ( make_run_config, ) from strix.sandbox import session_manager +from strix.sandbox.healthcheck import wait_for_http_ready, wait_for_tcp_ready from strix.telemetry.strix_processor import StrixTracingProcessor @@ -219,6 +221,20 @@ async def run_strix_scan( sources_path=sources_path, ) + # Wait for the in-container FastAPI tool server + Caido sidecar to + # come up before any agent fires its first tool call. + await asyncio.gather( + wait_for_http_ready( + f"http://127.0.0.1:{bundle['tool_server_host_port']}/health", + timeout=60.0, + ), + wait_for_tcp_ready( + "127.0.0.1", + int(bundle["caido_host_port"]), + timeout=60.0, + ), + ) + try: scan_mode = str(scan_config.get("scan_mode") or "deep") is_whitebox = bool(scan_config.get("is_whitebox", False)) @@ -254,7 +270,6 @@ async def run_strix_scan( sandbox_token=bundle["bearer"], tool_server_host_port=bundle["tool_server_host_port"], caido_host_port=bundle["caido_host_port"], - caido_capability=bundle.get("capability"), agent_id=root_id, parent_id=None, tracer=tracer, diff --git a/strix/orchestration/hooks.py b/strix/orchestration/hooks.py index 7c8ae26..8576b49 100644 --- a/strix/orchestration/hooks.py +++ b/strix/orchestration/hooks.py @@ -17,16 +17,13 @@ logger = logging.getLogger(__name__) class StrixOrchestrationHooks(RunHooks[Any]): """Lifecycle hooks for Strix multi-agent runs. - Wires four concerns: + Wires three concerns: 1. Turn-budget warnings injected into ``input_items`` at 85% and ``N - 3`` of ``max_turns``. - 2. LLM usage recording into the bus + tracer. - 3. Sandbox readiness: awaits the - ``CaidoCapability._healthcheck_task`` on first agent start so - the agent doesn't fire tools before Caido and the tool server - are ready. - 4. Subagent crash detection: if ``on_agent_end`` fires without + 2. LLM usage recording into the bus + tracer, plus mirroring the + bus's agent tree into ``tracer.agents`` for the TUI. + 3. Subagent crash detection: if ``on_agent_end`` fires without ``agent_finish_called`` being set, posts a crash message to the parent's inbox so the parent learns on its next turn instead of waiting forever. @@ -106,26 +103,14 @@ class StrixOrchestrationHooks(RunHooks[Any]): context: AgentHookContext[Any], agent: Any, ) -> None: - # Two concerns wired together because they fire at the same - # moment in the lifecycle: - # - # 1. CaidoCapability is bound to the sandbox session, not the - # Agent (we use plain ``Agent``, not ``SandboxAgent``), so - # the healthcheck task gets stashed in the context dict at - # scan bring-up and we await it on first agent start. - # 2. The TUI reads ``tracer.agents`` to render the agent tree. - # We mirror the bus state into the tracer here so the tree - # actually shows live and historical agents. + # The TUI reads ``tracer.agents`` to render the agent tree; + # mirror the bus state into the tracer here so the tree + # populates as agents come online. del agent try: ctx = context.context if not isinstance(ctx, dict): return - cap = ctx.get("caido_capability") - task = getattr(cap, "_healthcheck_task", None) - if task is not None: - await task - tracer = ctx.get("tracer") bus = ctx.get("bus") me = ctx.get("agent_id") diff --git a/strix/run_config_factory.py b/strix/run_config_factory.py index 5a82a28..5c059c9 100644 --- a/strix/run_config_factory.py +++ b/strix/run_config_factory.py @@ -126,7 +126,6 @@ def make_agent_context( run_id: str | None = None, sandbox_client: Any | None = None, agent_factory: Any | None = None, - caido_capability: Any | None = None, ) -> dict[str, Any]: """Build the per-agent ``context`` dict passed to ``Runner.run(context=...)``. @@ -146,7 +145,6 @@ def make_agent_context( "sandbox_token": sandbox_token, "tool_server_host_port": tool_server_host_port, "caido_host_port": caido_host_port, - "caido_capability": caido_capability, "agent_id": agent_id, "parent_id": parent_id, "tracer": tracer, diff --git a/strix/sandbox/__init__.py b/strix/sandbox/__init__.py index 37788a1..5ceb951 100644 --- a/strix/sandbox/__init__.py +++ b/strix/sandbox/__init__.py @@ -1,8 +1,6 @@ """Strix sandbox layer on top of OpenAI Agents SDK SandboxAgent / Manifest. -- :mod:`.caido_capability` — Caido proxy + 7 GraphQL function tools - + system prompt block. -- :mod:`.healthcheck` — ``wait_for_ports_ready``. +- :mod:`.healthcheck` — ``wait_for_http_ready`` / ``wait_for_tcp_ready``. - :mod:`.session_manager` — ``create_or_reuse`` / ``cleanup`` keyed by scan id. """ diff --git a/strix/sandbox/caido_capability.py b/strix/sandbox/caido_capability.py deleted file mode 100644 index 0ac6f28..0000000 --- a/strix/sandbox/caido_capability.py +++ /dev/null @@ -1,207 +0,0 @@ -"""CaidoCapability — sandbox capability for the Caido HTTP/HTTPS proxy. - -Three concerns wired into the SDK's capability lifecycle: - -1. **Manifest mutation** (``process_manifest``): inject ``http_proxy`` / - ``https_proxy`` / ``ALL_PROXY`` env vars pointing at the in-container - Caido listener. Any tool that ultimately shells out (curl, requests, - etc.) now flows through the proxy automatically. - -2. **Tool exposure** (``tools``): the seven Caido SDK function-tool - wrappers are returned here. The SDK runtime collects tools from - every capability and merges them with the agent's ``tools=[...]`` - declaration, so agents don't have to redeclare them. - -3. **Healthcheck task** (``bind``): when a session binds, we kick off - :func:`wait_for_http_ready` against the FastAPI tool server's - ``/health`` endpoint and :func:`wait_for_tcp_ready` against the - Caido proxy port. The aggregated task handle is stored on - ``self._healthcheck_task``, which the - :class:`StrixOrchestrationHooks.on_agent_start` hook awaits before - the first LLM call so the agent never hits a connection-refused - on its very first tool invocation. -""" - -from __future__ import annotations - -import asyncio -import logging -from typing import TYPE_CHECKING, ClassVar, Literal - -from agents.sandbox.capabilities.capability import Capability -from agents.tool import Tool -from pydantic import PrivateAttr - -from strix.sandbox.healthcheck import wait_for_http_ready, wait_for_tcp_ready -from strix.tools.proxy.tools import ( - list_requests, - list_sitemap, - repeat_request, - scope_rules, - send_request, - view_request, - view_sitemap_entry, -) - - -if TYPE_CHECKING: - from agents.sandbox.manifest import Manifest - from agents.sandbox.session.base_sandbox_session import BaseSandboxSession - - -logger = logging.getLogger(__name__) - - -# Container-internal Caido listener. The in-container Caido sidecar binds -# on this port; the host gets a randomly mapped port we resolve at -# session create time and pass into the per-agent context as -# ``caido_host_port`` for the proxy SDK tools' dispatcher. -_CAIDO_INTERNAL_PORT = 48080 - -# Container-internal FastAPI tool server. Same shape as Caido — host -# port is resolved at session create. -_TOOL_SERVER_INTERNAL_PORT = 48081 - -# Probe URLs used inside ``bind``. ``host=127.0.0.1`` because the host -# port mapping is loopback-only. -_PROBE_HOST = "127.0.0.1" - - -# Cached tool list — building Tool instances has side effects via the -# function_tool decorator and we don't want re-instantiation each time -# the SDK calls ``tools()``. -_CAIDO_TOOLS: tuple[Tool, ...] = ( - list_requests, - view_request, - send_request, - repeat_request, - scope_rules, - list_sitemap, - view_sitemap_entry, -) - - -class CaidoCapability(Capability): - """Caido HTTP/HTTPS forward proxy + 7 GraphQL function tools. - - Lifetime: one instance per scan. The SDK clones capabilities - per-run (see ``Capability.clone``); we accept that — each cloned - instance opens its own healthcheck task on ``bind``, which is - cheap and idempotent. - """ - - type: Literal["caido"] = "caido" - - # Pydantic ``PrivateAttr`` for runtime-only state. Pydantic forbids - # underscore-prefixed *fields*, but private attributes are first-class - # and cleanly excluded from model dumps and serialization. - _healthcheck_task: asyncio.Task[None] | None = PrivateAttr(default=None) - - # The two ports the host needs to reach. Populated by the session - # manager *after* the SDK creates the container and we've resolved - # the random host-side mappings via ``session._resolve_exposed_port``. - _tool_server_host_port: int | None = PrivateAttr(default=None) - _caido_host_port: int | None = PrivateAttr(default=None) - - # Per-capability healthcheck timeout. Long enough to cover image - # pulls on a cold cache plus tool-server boot, short enough that a - # mis-configured image fails the run inside a few minutes. - _HEALTHCHECK_TIMEOUT: ClassVar[float] = 60.0 - - def process_manifest(self, manifest: Manifest) -> Manifest: - """Inject proxy env vars into the manifest's environment. - - Mutates in place; returns the same manifest. Mirrors the SDK's - Capability protocol where ``process_manifest`` is the single - synchronous hook for changing what the container sees. - """ - env = dict(manifest.environment.value or {}) - env.update( - { - "http_proxy": f"http://127.0.0.1:{_CAIDO_INTERNAL_PORT}", - "https_proxy": f"http://127.0.0.1:{_CAIDO_INTERNAL_PORT}", - "ALL_PROXY": f"http://127.0.0.1:{_CAIDO_INTERNAL_PORT}", - }, - ) - manifest.environment.value = env - return manifest - - def tools(self) -> list[Tool]: - """Return the seven Caido function tools. - - The SDK runtime calls this at agent-build time and merges the - result with the agent's own tool list. Returning a fresh list - each call (rather than yielding the cached tuple directly) is - SDK convention. - """ - return list(_CAIDO_TOOLS) - - async def instructions(self, manifest: Manifest) -> str | None: # noqa: ARG002 - """System-prompt fragment appended for every Caido-equipped agent.""" - return ( - "\n" - "All HTTP/HTTPS traffic in this sandbox is automatically captured " - f"by Caido (in-container at 127.0.0.1:{_CAIDO_INTERNAL_PORT}; " - "host_proxy / https_proxy env vars are pre-set).\n" - "Tools: list_requests, view_request, send_request, repeat_request, " - "scope_rules, list_sitemap, view_sitemap_entry.\n" - "HTTPQL filter examples: " - "'request.method == \"POST\"', " - "'response.status >= 400', " - "'request.host == \"target.com\"'.\n" - "" - ) - - def configure_host_ports( - self, - *, - tool_server_host_port: int, - caido_host_port: int, - ) -> None: - """Record the resolved host-side ports. - - Called by the session manager after ``client.create(...)`` - returns, before binding the session. The healthcheck task - reads these to know which mapped ports to probe. - """ - self._tool_server_host_port = tool_server_host_port - self._caido_host_port = caido_host_port - - def bind(self, session: BaseSandboxSession) -> None: - """Schedule a healthcheck task on session bind. - - Stores the task handle so :class:`StrixOrchestrationHooks` can - await it on the first agent start. We never raise from here — - the healthcheck failure surfaces inside on_agent_start, which - is the right place to fail the run because by then we have a - live RunContextWrapper to log against. - """ - super().bind(session) - if self._tool_server_host_port is None or self._caido_host_port is None: - logger.warning( - "CaidoCapability.bind called before configure_host_ports; " - "skipping healthcheck task scheduling.", - ) - return - self._healthcheck_task = asyncio.create_task( - self._run_healthcheck(), - name=f"caido-healthcheck-{self._tool_server_host_port}", - ) - - async def _run_healthcheck(self) -> None: - """Probe both ports concurrently; raise on first failure.""" - # Mypy sees these as Optional, but ``bind`` checks both before - # creating the task. - assert self._tool_server_host_port is not None - assert self._caido_host_port is not None - await asyncio.gather( - wait_for_http_ready( - f"http://{_PROBE_HOST}:{self._tool_server_host_port}/health", - timeout=self._HEALTHCHECK_TIMEOUT, - ), - wait_for_tcp_ready( - _PROBE_HOST, - self._caido_host_port, - timeout=self._HEALTHCHECK_TIMEOUT, - ), - ) diff --git a/strix/sandbox/session_manager.py b/strix/sandbox/session_manager.py index 3f558ad..9b0d924 100644 --- a/strix/sandbox/session_manager.py +++ b/strix/sandbox/session_manager.py @@ -27,7 +27,6 @@ from agents.sandbox.manifest import Environment, Manifest from agents.sandbox.sandboxes.docker import DockerSandboxClientOptions from strix.runtime.strix_docker_client import StrixDockerSandboxClient -from strix.sandbox.caido_capability import CaidoCapability if TYPE_CHECKING: @@ -86,8 +85,7 @@ async def create_or_reuse( Defaults to 120s, matching the legacy harness. Returns the bundle dict containing ``client``, ``session``, - ``tool_server_host_port``, ``caido_host_port``, ``bearer``, - and ``capability`` (the live CaidoCapability instance). + ``tool_server_host_port``, ``caido_host_port``, and ``bearer``. """ cached = _SESSION_CACHE.get(scan_id) if cached is not None: @@ -96,9 +94,11 @@ async def create_or_reuse( bearer = secrets.token_urlsafe(32) - capability = CaidoCapability() - # ``Manifest.environment`` is an ``Environment`` model — a bare dict - # is silently dropped by Pydantic, so wrap explicitly. + # Caido runs as an in-container sidecar on _CONTAINER_CAIDO_PORT and + # all HTTP(S) traffic from shelled-out tools (curl, Python requests, + # etc.) needs to flow through it — set the conventional env vars so + # standard libraries pick them up automatically. + caido_proxy_url = f"http://127.0.0.1:{_CONTAINER_CAIDO_PORT}" manifest = Manifest( entries={"sources": LocalDir(src=sources_path)}, environment=Environment( @@ -109,6 +109,9 @@ async def create_or_reuse( "STRIX_SANDBOX_EXECUTION_TIMEOUT": str(execution_timeout), "PYTHONUNBUFFERED": "1", "HOST_GATEWAY": "host.docker.internal", + "http_proxy": caido_proxy_url, + "https_proxy": caido_proxy_url, + "ALL_PROXY": caido_proxy_url, }, ), ) @@ -130,26 +133,14 @@ async def create_or_reuse( ) session = await client.create(options=options, manifest=manifest) - # Resolve the host-side mapped ports the SDK assigned. The capability - # needs these *before* it binds, so its healthcheck task probes the - # right ports. tool_server_endpoint = await session._resolve_exposed_port( _CONTAINER_TOOL_SERVER_PORT, ) caido_endpoint = await session._resolve_exposed_port(_CONTAINER_CAIDO_PORT) - capability.configure_host_ports( - tool_server_host_port=tool_server_endpoint.port, - caido_host_port=caido_endpoint.port, - ) - # Bind the capability against the live session — this schedules the - # healthcheck task that on_agent_start awaits. - capability.bind(session) - bundle = { "client": client, "session": session, - "capability": capability, "tool_server_host_port": tool_server_endpoint.port, "caido_host_port": caido_endpoint.port, "bearer": bearer, @@ -171,12 +162,6 @@ async def cleanup(scan_id: str) -> None: logger.debug("cleanup(%s): no cached session", scan_id) return - capability = bundle.get("capability") - if isinstance(capability, CaidoCapability): - task = capability._healthcheck_task - if task is not None and not task.done(): - task.cancel() - try: await bundle["client"].delete(bundle["session"]) logger.info("Cleaned up sandbox session for scan %s", scan_id) diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index 7207db2..7ac938f 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -406,7 +406,6 @@ async def create_agent( sandbox_token=inner.get("sandbox_token"), tool_server_host_port=inner.get("tool_server_host_port"), caido_host_port=inner.get("caido_host_port"), - caido_capability=inner.get("caido_capability"), agent_id=child_id, parent_id=parent_id, tracer=inner.get("tracer"),