diff --git a/pyproject.toml b/pyproject.toml index 6dd6251..66e2016 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,6 +139,7 @@ module = [ "opentelemetry.*", "scrubadub.*", "traceloop.*", + "docker.*", ] ignore_missing_imports = true @@ -299,6 +300,9 @@ ignore = [ "strix/tools/python/python_sdk_tool.py" = ["TC002"] "strix/tools/proxy/proxy_sdk_tools.py" = ["TC002"] "strix/tools/agents_graph/agents_graph_sdk_tools.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"] # 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/sandbox/caido_capability.py b/strix/sandbox/caido_capability.py new file mode 100644 index 0000000..a3c5e28 --- /dev/null +++ b/strix/sandbox/caido_capability.py @@ -0,0 +1,212 @@ +"""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 from Phase 2.5 are returned here. The SDK runtime collects + tools from every capability and merges them with the agent's + ``tools=[...]`` declaration, so individual 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. + +References: + - PLAYBOOK.md §3.2 + - AUDIT.md §2.5 (C5 — healthcheck wired to RunHooks) +""" + +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.proxy_sdk_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 (legacy and SDK both bind to 127.0.0.1). +_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/healthcheck.py b/strix/sandbox/healthcheck.py new file mode 100644 index 0000000..a69f606 --- /dev/null +++ b/strix/sandbox/healthcheck.py @@ -0,0 +1,121 @@ +"""Sandbox port readiness probes used during session bring-up. + +The in-container tool server (FastAPI) takes a few seconds to start +listening after the Docker container is created, and Caido's HTTPS +proxy takes a similar window. The session manager waits for both +before returning a session bundle so that the first tool call from +an agent doesn't hit a connection refused. + +Two helpers are exposed: + +- :func:`wait_for_http_ready` for the FastAPI tool server, whose + ``/health`` endpoint returns ``{"status": "healthy"}`` once the + process is up. We don't require the JSON shape exactly — any 2xx + is treated as ready, mirroring the legacy ``_wait_for_tool_server`` + but more lenient (the legacy version checked the JSON body too, + which made test images without that handler fail spuriously). + +- :func:`wait_for_tcp_ready` for Caido, which serves an HTTP forward + proxy on its port and does *not* expose ``/health``. A TCP connect + is the most we can probe without sending real proxy traffic. + +References: + - PLAYBOOK.md §3.1 + - HARNESS_WIKI.md §6.4 (legacy ``_wait_for_tool_server`` pattern) +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging + +import httpx + + +logger = logging.getLogger(__name__) + + +class SandboxNotReadyError(Exception): + """Raised when a sandbox port doesn't accept connections in time.""" + + +# Default per-attempt HTTP timeout. The legacy harness used 5s; we +# match it so a slow first request (image still warming up) doesn't +# misfire as a hard failure on a single attempt. +_DEFAULT_HTTP_PROBE_TIMEOUT = 5.0 + +# Default polling cadence between attempts. Balanced for CI-style +# fast bring-up (sub-second) without burning CPU when the port is +# legitimately taking a few seconds. +_DEFAULT_POLL_INTERVAL = 0.5 + + +async def wait_for_http_ready( + url: str, + *, + timeout: float = 30.0, + poll_interval: float = _DEFAULT_POLL_INTERVAL, + probe_timeout: float = _DEFAULT_HTTP_PROBE_TIMEOUT, +) -> None: + """Poll ``url`` until any 2xx response, or raise after ``timeout``. + + Network errors (ConnectError / TimeoutException / RequestError) + are treated as "not ready yet" — the loop continues. Any other + exception class will surface immediately so a programmer error + (bad URL, etc.) doesn't get silently retried for 30 seconds. + """ + deadline = asyncio.get_event_loop().time() + timeout + last_error: str | None = None + async with httpx.AsyncClient(timeout=probe_timeout, trust_env=False) as client: + while asyncio.get_event_loop().time() < deadline: + try: + response = await client.get(url) + if 200 <= response.status_code < 300: + return + last_error = f"HTTP {response.status_code}" + except (httpx.ConnectError, httpx.TimeoutException, httpx.RequestError) as e: + last_error = type(e).__name__ + await asyncio.sleep(poll_interval) + + raise SandboxNotReadyError( + f"HTTP probe of {url} did not return 2xx within {timeout}s (last error: {last_error})", + ) + + +async def wait_for_tcp_ready( + host: str, + port: int, + *, + timeout: float = 30.0, + poll_interval: float = _DEFAULT_POLL_INTERVAL, +) -> None: + """Poll ``host:port`` until a TCP connect succeeds, or raise after ``timeout``. + + Used for ports that don't expose an HTTP health endpoint (Caido's + forward proxy). We open the socket and immediately close it — the + handshake completing is enough to confirm readiness. + """ + deadline = asyncio.get_event_loop().time() + timeout + last_error: str | None = None + while asyncio.get_event_loop().time() < deadline: + try: + reader, writer = await asyncio.wait_for( + asyncio.open_connection(host, port), + timeout=poll_interval * 4, + ) + except (TimeoutError, OSError) as e: + last_error = type(e).__name__ + else: + writer.close() + # Some servers close hard immediately after accept; we only + # care that the connect itself succeeded. + with contextlib.suppress(OSError): + await writer.wait_closed() + del reader + return + await asyncio.sleep(poll_interval) + + raise SandboxNotReadyError( + f"TCP probe of {host}:{port} did not connect within {timeout}s (last error: {last_error})", + ) diff --git a/strix/sandbox/session_manager.py b/strix/sandbox/session_manager.py new file mode 100644 index 0000000..739aa96 --- /dev/null +++ b/strix/sandbox/session_manager.py @@ -0,0 +1,204 @@ +"""Per-scan sandbox session lifecycle. + +Replaces the legacy ``DockerRuntime`` (``strix/runtime/docker_runtime.py``) +with the SDK-native session model. One session per scan, reused across +every agent in that scan's tree. + +The bundle returned by :func:`create_or_reuse` is what the per-agent +context dict reads from in ``run_config_factory.make_agent_context`` — +``client``, ``session``, ``tool_server_host_port``, ``caido_host_port``, +and ``bearer`` for authenticating to the in-container FastAPI tool server. + +Cache strategy: a module-level dict keyed by ``scan_id``. The same scan +issuing multiple ``create_or_reuse`` calls (e.g., resume after a crash +on the host side) gets the same bundle back. ``cleanup`` is best-effort +— a leaked container is preferable to a stuck cleanup that prevents the +next scan from starting. + +References: + - PLAYBOOK.md §3.3 + - HARNESS_WIKI.md §6 (legacy Docker runtime) +""" + +from __future__ import annotations + +import logging +import secrets +import socket +from typing import TYPE_CHECKING, Any + +import docker +from agents.sandbox.entries import LocalDir +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: + from pathlib import Path + + +logger = logging.getLogger(__name__) + + +# In-container ports (must match the image's tool server + Caido sidecar +# binds). Defined here as a single source of truth for both the +# capability and the manifest env vars. +_CONTAINER_TOOL_SERVER_PORT = 48081 +_CONTAINER_CAIDO_PORT = 48080 + + +# Per-scan session cache. Module-level so a scan that bounces through +# multiple host-side processes (e.g., re-imports the module) doesn't +# spin up a second container — though in practice we expect one +# Strix process per scan. +_SESSION_CACHE: dict[str, dict[str, Any]] = {} + + +def _alloc_loopback_port() -> int: + """Reserve a free 127.0.0.1 port via ephemeral socket bind. + + Used only as a fallback when the SDK doesn't return a resolved + host port (older SDK versions before ``_resolve_exposed_port`` + existed). Modern path uses the SDK's resolution. + """ + sock = socket.socket() + try: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + finally: + sock.close() + + +async def create_or_reuse( + scan_id: str, + *, + image: str, + sources_path: Path, + execution_timeout: int = 120, +) -> dict[str, Any]: + """Return the existing bundle for ``scan_id`` or create a new one. + + Args: + scan_id: Caller-provided scan identifier (used as cache key). + image: Docker image tag (e.g. ``"strix-sandbox:0.1.13"``). + sources_path: Host directory mounted into the container's + ``/workspace/sources`` so the agent can read user code. + execution_timeout: ``STRIX_SANDBOX_EXECUTION_TIMEOUT`` env var + inside the container — caps how long the in-container tool + server waits for a tool to finish before responding 504. + 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). + """ + cached = _SESSION_CACHE.get(scan_id) + if cached is not None: + logger.info("Reusing existing sandbox session for scan %s", scan_id) + return cached + + bearer = secrets.token_urlsafe(32) + + capability = CaidoCapability() + # ``Manifest.environment`` is an ``Environment`` model — a bare dict + # is silently dropped by Pydantic, so wrap explicitly. + manifest = Manifest( + entries={"sources": LocalDir(src=sources_path)}, + environment=Environment( + value={ + "TOOL_SERVER_TOKEN": bearer, + "TOOL_SERVER_PORT": str(_CONTAINER_TOOL_SERVER_PORT), + "CAIDO_PORT": str(_CONTAINER_CAIDO_PORT), + "STRIX_SANDBOX_EXECUTION_TIMEOUT": str(execution_timeout), + "PYTHONUNBUFFERED": "1", + "HOST_GATEWAY": "host.docker.internal", + }, + ), + ) + + # The SDK's DockerSandboxClient requires a docker.DockerClient instance + # at construction time (since openai-agents 0.14.x). We use the + # caller's environment to find the daemon — same as the legacy + # DockerRuntime did via ``docker.from_env()``. + client = StrixDockerSandboxClient(docker.from_env()) + options = DockerSandboxClientOptions( + image=image, + exposed_ports=(_CONTAINER_TOOL_SERVER_PORT, _CONTAINER_CAIDO_PORT), + ) + + logger.info( + "Creating sandbox session for scan %s (image=%s, exec_timeout=%ds)", + scan_id, + image, + execution_timeout, + ) + 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, + } + _SESSION_CACHE[scan_id] = bundle + return bundle + + +async def cleanup(scan_id: str) -> None: + """Tear down ``scan_id``'s container and drop its cache entry. + + Best-effort: any error during ``client.delete`` is logged and + swallowed. We never want a cleanup failure to prevent the next + scan from starting; the worst case is a stranded container that + Docker's normal reaping will catch on next ``docker prune``. + """ + bundle = _SESSION_CACHE.pop(scan_id, None) + if bundle is 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) + except Exception: + logger.exception( + "cleanup(%s): client.delete raised; container may need manual reaping", + scan_id, + ) + + +def cached_scan_ids() -> list[str]: + """Snapshot of currently-cached scan ids. Used by the TUI / CLI.""" + return list(_SESSION_CACHE.keys()) + + +def _reset_cache_for_tests() -> None: + """Test helper — clears the module cache between unit tests.""" + _SESSION_CACHE.clear() diff --git a/tests/sandbox/__init__.py b/tests/sandbox/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/sandbox/test_caido_capability.py b/tests/sandbox/test_caido_capability.py new file mode 100644 index 0000000..fb5227e --- /dev/null +++ b/tests/sandbox/test_caido_capability.py @@ -0,0 +1,156 @@ +"""Phase 4 tests for CaidoCapability. + +The capability has three observable behaviors that need parity with the +PLAYBOOK contract: + +1. ``process_manifest`` injects http_proxy / https_proxy / ALL_PROXY env + vars into the manifest's ``Environment.value`` dict. +2. ``tools()`` returns the seven Caido SDK function tools we wrapped in + Phase 2.5 — same instances, in the documented order. +3. ``bind`` schedules an aggregated healthcheck task; the orchestration + hook later awaits it on first agent start. + +The healthcheck task itself is exercised by the healthcheck unit tests; +here we only verify the wiring (task created, name set, points at the +right ports). +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest +from agents.sandbox.entries import LocalDir +from agents.sandbox.manifest import Environment, Manifest + +from strix.sandbox.caido_capability import CaidoCapability + + +def test_capability_type_and_default_state() -> None: + cap = CaidoCapability() + assert cap.type == "caido" + assert cap._healthcheck_task is None + assert cap._tool_server_host_port is None + assert cap._caido_host_port is None + + +def test_process_manifest_injects_proxy_env_vars(tmp_path: object) -> None: + """Existing env vars must be preserved; proxy keys are added.""" + cap = CaidoCapability() + manifest = Manifest( + environment=Environment( + value={"PYTHONUNBUFFERED": "1", "TOOL_SERVER_TOKEN": "abc"}, + ), + ) + out = cap.process_manifest(manifest) + env = out.environment.value + # Pre-existing entries preserved. + assert env["PYTHONUNBUFFERED"] == "1" + assert env["TOOL_SERVER_TOKEN"] == "abc" + # Proxy entries injected, all pointing at the in-container Caido port. + assert env["http_proxy"] == "http://127.0.0.1:48080" + assert env["https_proxy"] == "http://127.0.0.1:48080" + assert env["ALL_PROXY"] == "http://127.0.0.1:48080" + + +def test_process_manifest_handles_missing_environment() -> None: + """A manifest without env entries should still get the proxy block.""" + cap = CaidoCapability() + # ``LocalDir`` requires a real path on disk; use a temp one to satisfy + # the validator without actually mounting anything. + manifest = Manifest(entries={"src": LocalDir(src="/tmp")}) + out = cap.process_manifest(manifest) + env = out.environment.value + assert env["http_proxy"] == "http://127.0.0.1:48080" + + +def test_tools_returns_seven_caido_tools_in_order() -> None: + cap = CaidoCapability() + names = [t.name for t in cap.tools()] + assert names == [ + "list_requests", + "view_request", + "send_request", + "repeat_request", + "scope_rules", + "list_sitemap", + "view_sitemap_entry", + ] + + +def test_tools_returns_a_fresh_list_per_call() -> None: + """SDK convention — caller may mutate the returned list.""" + cap = CaidoCapability() + a = cap.tools() + b = cap.tools() + assert a == b + assert a is not b + + +@pytest.mark.asyncio +async def test_instructions_mentions_caido_and_tools() -> None: + cap = CaidoCapability() + out = await cap.instructions(Manifest()) + assert out is not None + assert "" in out + # Every tool name appears verbatim so the model knows what's available. + for name in ( + "list_requests", + "view_request", + "send_request", + "repeat_request", + "scope_rules", + "list_sitemap", + "view_sitemap_entry", + ): + assert name in out + + +def test_configure_host_ports_stores_both() -> None: + cap = CaidoCapability() + cap.configure_host_ports(tool_server_host_port=12345, caido_host_port=12346) + assert cap._tool_server_host_port == 12345 + assert cap._caido_host_port == 12346 + + +@pytest.mark.asyncio +async def test_bind_without_configured_ports_skips_healthcheck() -> None: + """If the session manager forgets to configure ports, bind shouldn't + schedule a probe against ``None`` — it should warn and no-op. + """ + cap = CaidoCapability() + fake_session = MagicMock() + cap.bind(fake_session) + assert cap._healthcheck_task is None + assert cap.session is fake_session + + +@pytest.mark.asyncio +async def test_bind_schedules_healthcheck_task_when_ports_configured() -> None: + """The hook chain (StrixOrchestrationHooks.on_agent_start) awaits this + task — it must exist as an asyncio.Task with a useful name. + """ + cap = CaidoCapability() + cap.configure_host_ports(tool_server_host_port=54321, caido_host_port=54322) + fake_session = MagicMock() + + # Patch the actual probes so we don't try to connect for real. + async def _fake_probe(*args: object, **kwargs: object) -> None: + return None + + with ( + patch( + "strix.sandbox.caido_capability.wait_for_http_ready", + side_effect=_fake_probe, + ), + patch( + "strix.sandbox.caido_capability.wait_for_tcp_ready", + side_effect=_fake_probe, + ), + ): + cap.bind(fake_session) + assert cap._healthcheck_task is not None + assert isinstance(cap._healthcheck_task, asyncio.Task) + assert "caido-healthcheck-54321" in cap._healthcheck_task.get_name() + await cap._healthcheck_task # must complete without error diff --git a/tests/sandbox/test_healthcheck.py b/tests/sandbox/test_healthcheck.py new file mode 100644 index 0000000..6ee1734 --- /dev/null +++ b/tests/sandbox/test_healthcheck.py @@ -0,0 +1,171 @@ +"""Phase 4 tests for the sandbox port readiness probes. + +The two helpers (``wait_for_http_ready`` and ``wait_for_tcp_ready``) +gate session bring-up, so a regression here would mean every fresh +scan hits a connection-refused on its first tool call. Tests cover: + +- Happy path returns when the probe succeeds. +- Polling continues across transient failures. +- Timeout raises ``SandboxNotReadyError`` with a useful last-error. +- Real ``asyncio.open_connection`` against a local listener verifies + the TCP probe end-to-end (no mocking — the helper is small enough + that a real socket is the cheaper test). +""" + +from __future__ import annotations + +import asyncio +import contextlib +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from strix.sandbox.healthcheck import ( + SandboxNotReadyError, + wait_for_http_ready, + wait_for_tcp_ready, +) + + +# --- HTTP probe ---------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_wait_for_http_ready_returns_immediately_on_2xx() -> None: + response = MagicMock(spec=httpx.Response) + response.status_code = 200 + client = AsyncMock() + client.get = AsyncMock(return_value=response) + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) + + with patch("strix.sandbox.healthcheck.httpx.AsyncClient", return_value=client): + await wait_for_http_ready("http://localhost:9999/health", timeout=1) + + assert client.get.await_count == 1 + + +@pytest.mark.asyncio +async def test_wait_for_http_ready_polls_through_connect_errors() -> None: + """Two connect errors followed by a 200 — the helper should keep going.""" + response_ok = MagicMock(spec=httpx.Response) + response_ok.status_code = 200 + + side_effects: list[Any] = [ + httpx.ConnectError("conn refused"), + httpx.ConnectError("conn refused"), + response_ok, + ] + + client = AsyncMock() + client.get = AsyncMock(side_effect=side_effects) + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) + + with patch("strix.sandbox.healthcheck.httpx.AsyncClient", return_value=client): + await wait_for_http_ready( + "http://localhost:9999/health", + timeout=5, + poll_interval=0.01, + ) + assert client.get.await_count == 3 + + +@pytest.mark.asyncio +async def test_wait_for_http_ready_raises_after_timeout() -> None: + client = AsyncMock() + client.get = AsyncMock(side_effect=httpx.ConnectError("nope")) + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( + "strix.sandbox.healthcheck.httpx.AsyncClient", + return_value=client, + ), + pytest.raises(SandboxNotReadyError) as exc_info, + ): + await wait_for_http_ready( + "http://localhost:9999/health", + timeout=0.3, + poll_interval=0.05, + ) + + err = str(exc_info.value) + assert "http://localhost:9999/health" in err + assert "ConnectError" in err + + +@pytest.mark.asyncio +async def test_wait_for_http_ready_treats_5xx_as_not_ready() -> None: + response_500 = MagicMock(spec=httpx.Response) + response_500.status_code = 500 + response_ok = MagicMock(spec=httpx.Response) + response_ok.status_code = 200 + + client = AsyncMock() + client.get = AsyncMock(side_effect=[response_500, response_ok]) + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) + + with patch("strix.sandbox.healthcheck.httpx.AsyncClient", return_value=client): + await wait_for_http_ready( + "http://localhost:9999/health", + timeout=2, + poll_interval=0.01, + ) + assert client.get.await_count == 2 + + +# --- TCP probe ----------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_wait_for_tcp_ready_against_real_listener() -> None: + """Spin up a local TCP echo server and verify the probe connects.""" + + async def _server_handler( + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + # Drain any bytes the test sends, then close. + await reader.read(0) + writer.close() + with contextlib.suppress(OSError): + await writer.wait_closed() + + server = await asyncio.start_server(_server_handler, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + try: + await wait_for_tcp_ready("127.0.0.1", port, timeout=2, poll_interval=0.05) + finally: + server.close() + await server.wait_closed() + + +@pytest.mark.asyncio +async def test_wait_for_tcp_ready_raises_when_port_closed() -> None: + async def _no_handler( + _reader: asyncio.StreamReader, + _writer: asyncio.StreamWriter, + ) -> None: + return + + # Bind and immediately close to claim a definitely-unused port number. + server = await asyncio.start_server(_no_handler, "127.0.0.1", 0) + closed_port = server.sockets[0].getsockname()[1] + server.close() + await server.wait_closed() + + with pytest.raises(SandboxNotReadyError) as exc_info: + await wait_for_tcp_ready( + "127.0.0.1", + closed_port, + timeout=0.3, + poll_interval=0.05, + ) + + err = str(exc_info.value) + assert f"127.0.0.1:{closed_port}" in err diff --git a/tests/sandbox/test_session_manager.py b/tests/sandbox/test_session_manager.py new file mode 100644 index 0000000..bacdbbe --- /dev/null +++ b/tests/sandbox/test_session_manager.py @@ -0,0 +1,206 @@ +"""Phase 4 tests for the per-scan sandbox session manager. + +We don't spin up real Docker here — the ``StrixDockerSandboxClient`` is +patched and we assert on the manifest / options / bundle shape. Goals: + +- Cache hit: a second ``create_or_reuse(scan_id, ...)`` returns the same + bundle without calling client.create twice. +- Manifest carries the right env vars (TOOL_SERVER_TOKEN, container ports, + STRIX_SANDBOX_EXECUTION_TIMEOUT, PYTHONUNBUFFERED). +- The Docker client options request both container ports be exposed. +- Capability is configured with the resolved host ports *before* bind, + so its healthcheck task probes the right ones. +- Bundle is cached and surfaces in ``cached_scan_ids``. +- ``cleanup`` cancels the healthcheck task and calls ``client.delete``; + errors during delete are swallowed. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from strix.sandbox import session_manager +from strix.sandbox.caido_capability import CaidoCapability + + +@pytest.fixture(autouse=True) +def _isolate_cache() -> Iterator[None]: + session_manager._reset_cache_for_tests() + yield + session_manager._reset_cache_for_tests() + + +def _noop_bind(_self: Any, _session: Any) -> None: + """Stand-in for CaidoCapability.bind that skips the healthcheck task.""" + + +def _make_endpoint(port: int) -> Any: + ep = MagicMock() + ep.port = port + ep.host = "127.0.0.1" + ep.tls = False + return ep + + +def _make_client_and_session( + *, + tool_port: int = 12001, + caido_port: int = 12002, +) -> tuple[Any, Any]: + """Build a fake DockerSandboxClient and session pair.""" + session = MagicMock() + session._resolve_exposed_port = AsyncMock( + side_effect=lambda port: _make_endpoint( + tool_port if port == 48081 else caido_port, + ), + ) + client = MagicMock() + client.create = AsyncMock(return_value=session) + client.delete = AsyncMock() + return client, session + + +@pytest.mark.asyncio +async def test_create_or_reuse_creates_new_session(tmp_path: Any) -> None: + client, session = _make_client_and_session() + # Patch the capability's bind to a no-op so we don't spin up the + # healthcheck task in unit tests. + with ( + patch( + "strix.sandbox.session_manager.StrixDockerSandboxClient", + return_value=client, + ), + patch.object(CaidoCapability, "bind", _noop_bind), + ): + bundle = await session_manager.create_or_reuse( + "scan-1", + image="strix-sandbox:test", + sources_path=tmp_path, + ) + + # Bundle shape. + assert bundle["client"] is client + assert bundle["session"] is session + assert bundle["tool_server_host_port"] == 12001 + assert bundle["caido_host_port"] == 12002 + assert isinstance(bundle["bearer"], str) and len(bundle["bearer"]) >= 32 + assert isinstance(bundle["capability"], CaidoCapability) + # Capability got the resolved host ports BEFORE bind would have run. + assert bundle["capability"]._tool_server_host_port == 12001 + assert bundle["capability"]._caido_host_port == 12002 + + # client.create called exactly once with manifest + exposed ports. + assert client.create.await_count == 1 + options = client.create.await_args.kwargs["options"] + assert options.image == "strix-sandbox:test" + assert set(options.exposed_ports) == {48080, 48081} + + manifest = client.create.await_args.kwargs["manifest"] + env = manifest.environment.value + assert env["TOOL_SERVER_TOKEN"] == bundle["bearer"] + assert env["TOOL_SERVER_PORT"] == "48081" + assert env["CAIDO_PORT"] == "48080" + assert env["STRIX_SANDBOX_EXECUTION_TIMEOUT"] == "120" + assert env["PYTHONUNBUFFERED"] == "1" + assert env["HOST_GATEWAY"] == "host.docker.internal" + + +@pytest.mark.asyncio +async def test_create_or_reuse_returns_cached_bundle(tmp_path: Any) -> None: + client, _ = _make_client_and_session() + with ( + patch( + "strix.sandbox.session_manager.StrixDockerSandboxClient", + return_value=client, + ), + patch.object(CaidoCapability, "bind", _noop_bind), + ): + first = await session_manager.create_or_reuse( + "scan-X", + image="i", + sources_path=tmp_path, + ) + second = await session_manager.create_or_reuse( + "scan-X", + image="i", + sources_path=tmp_path, + ) + + assert first is second + assert client.create.await_count == 1 + assert "scan-X" in session_manager.cached_scan_ids() + + +@pytest.mark.asyncio +async def test_create_or_reuse_passes_custom_execution_timeout(tmp_path: Any) -> None: + client, _ = _make_client_and_session() + with ( + patch( + "strix.sandbox.session_manager.StrixDockerSandboxClient", + return_value=client, + ), + patch.object(CaidoCapability, "bind", _noop_bind), + ): + await session_manager.create_or_reuse( + "scan-2", + image="i", + sources_path=tmp_path, + execution_timeout=300, + ) + + manifest = client.create.await_args.kwargs["manifest"] + assert manifest.environment.value["STRIX_SANDBOX_EXECUTION_TIMEOUT"] == "300" + + +@pytest.mark.asyncio +async def test_cleanup_calls_delete_and_drops_cache(tmp_path: Any) -> None: + client, session = _make_client_and_session() + with ( + patch( + "strix.sandbox.session_manager.StrixDockerSandboxClient", + return_value=client, + ), + patch.object(CaidoCapability, "bind", _noop_bind), + ): + await session_manager.create_or_reuse( + "scan-3", + image="i", + sources_path=tmp_path, + ) + assert "scan-3" in session_manager.cached_scan_ids() + await session_manager.cleanup("scan-3") + + client.delete.assert_awaited_once_with(session) + assert "scan-3" not in session_manager.cached_scan_ids() + + +@pytest.mark.asyncio +async def test_cleanup_swallows_delete_errors(tmp_path: Any) -> None: + """A flaky Docker daemon shouldn't prevent cache eviction.""" + client, _ = _make_client_and_session() + client.delete = AsyncMock(side_effect=RuntimeError("docker daemon went away")) + with ( + patch( + "strix.sandbox.session_manager.StrixDockerSandboxClient", + return_value=client, + ), + patch.object(CaidoCapability, "bind", _noop_bind), + ): + await session_manager.create_or_reuse( + "scan-4", + image="i", + sources_path=tmp_path, + ) + await session_manager.cleanup("scan-4") # must not raise + + assert "scan-4" not in session_manager.cached_scan_ids() + + +@pytest.mark.asyncio +async def test_cleanup_unknown_scan_is_noop() -> None: + """No cached entry → cleanup is a quiet no-op.""" + await session_manager.cleanup("never-existed") # must not raise