perf: bootstrap Caido concurrently with the scan start (#1143)

Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>

strix/runtime/caido_bootstrap.py

@@ -96,15 +96,17 @@ async def bootstrap_caido(
     access_token = await _login_as_guest(session, container_url=container_url)
 
     client = Client(host_url, auth=TokenAuthOptions(token=access_token))
-    await client.connect()
-
     try:
+        # connect() is inside the guard as well: a cancellation there (scan
+        # teardown while the bootstrap is still in flight) would otherwise
+        # leave the half-connected transport behind.
+        await client.connect()
         project = await client.project.create(
             CreateProjectOptions(name="sandbox", temporary=True),
         )
         await client.project.select(project.id)
     except BaseException:
-        # The connected client never reaches the session bundle if project
+        # The client never reaches the session bundle if connect or project
         # setup fails, so close it here to avoid leaking the transport.
         with contextlib.suppress(Exception):
             await client.aclose()

strix/runtime/caido_handle.py

@@ -0,0 +1,60 @@
+"""Handle for a Caido bootstrap running concurrently with the scan start.
+
+The Caido sidecar login + project setup costs a couple of seconds of
+guest-side polling, and nothing needs the client until the first proxy
+tool call (or the first traffic poll). :class:`CaidoBootstrapHandle`
+wraps the in-flight bootstrap task so session bring-up can return as
+soon as the container is up; consumers resolve the client at first use.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from caido_sdk_client import Client
+
+
+logger = logging.getLogger(__name__)
+
+
+class CaidoBootstrapHandle:
+    """Resolves to the connected Caido client once the bootstrap finishes.
+
+    A failed bootstrap is surfaced (once) to every ``get()`` caller as the
+    original exception; proxy tools degrade to their "client unavailable"
+    result instead of the failure killing the scan at bring-up.
+    """
+
+    def __init__(self, task: asyncio.Task[Client]) -> None:
+        self._task = task
+
+    async def get(self) -> Client:
+        """Wait for the bootstrap and return the client.
+
+        Shielded so one caller's cancellation (e.g. a tool timeout) does not
+        cancel the shared bootstrap for everyone else.
+        """
+        return await asyncio.shield(self._task)
+
+    def peek(self) -> Client | None:
+        """Return the client if the bootstrap already finished cleanly."""
+        if self._task.done() and not self._task.cancelled() and self._task.exception() is None:
+            return self._task.result()
+        return None
+
+    async def aclose(self) -> None:
+        """Cancel an in-flight bootstrap or close the finished client."""
+        if not self._task.done():
+            self._task.cancel()
+            with contextlib.suppress(asyncio.CancelledError, Exception):
+                await self._task
+            return
+        client = self.peek()
+        if client is not None:
+            with contextlib.suppress(Exception):
+                await client.aclose()

strix/runtime/session_manager.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
 import sys
@@ -15,6 +16,7 @@ from strix.config import load_settings
 from strix.core.paths import run_dir_for, runtime_state_dir
 from strix.runtime.backends import backend_supports_bind_mounts, get_backend
 from strix.runtime.caido_bootstrap import bootstrap_caido
+from strix.runtime.caido_handle import CaidoBootstrapHandle
 
 
 if TYPE_CHECKING:
@@ -333,10 +335,19 @@ async def create_or_reuse(
     host_caido_url = f"{scheme}://{caido_endpoint.host}:{caido_endpoint.port}"
     logger.debug("Caido host endpoint resolved: %s", host_caido_url)
 
-    caido_client = await bootstrap_caido(
-        session,
-        host_url=host_caido_url,
-        container_url=container_caido_url,
+    # The Caido login + project setup polls the guest for a couple of seconds
+    # and nothing needs the client before the first proxy tool call, so it
+    # runs concurrently with the rest of scan start; consumers resolve the
+    # handle at first use (see CaidoBootstrapHandle).
+    caido_client = CaidoBootstrapHandle(
+        asyncio.create_task(
+            bootstrap_caido(
+                session,
+                host_url=host_caido_url,
+                container_url=container_caido_url,
+            ),
+            name=f"caido-bootstrap-{scan_id}",
+        )
     )
 
     bundle = {

strix/tools/proxy/tools.py

@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Literal
 
 from agents import RunContextWrapper, function_tool
 
+from strix.runtime.caido_handle import CaidoBootstrapHandle
 from strix.tools.proxy import caido_api
 
 
@@ -47,9 +48,16 @@ ScopeAction = Literal["get", "list", "create", "update", "delete"]
 _CAIDO_CALL_LOCK = asyncio.Lock()
 
 
-def _ctx_client(ctx: RunContextWrapper) -> Client | None:
-    inner = ctx.context if isinstance(ctx.context, dict) else {}
-    return inner.get("caido_client")
+async def _ctx_client(ctx: RunContextWrapper) -> Client | None:
+    inner: dict[str, Any] = ctx.context if isinstance(ctx.context, dict) else {}
+    client: Client | CaidoBootstrapHandle | None = inner.get("caido_client")
+    if isinstance(client, CaidoBootstrapHandle):
+        try:
+            return await client.get()
+        except Exception:  # noqa: BLE001
+            logger.warning("Caido bootstrap failed; proxy tools unavailable", exc_info=True)
+            return None
+    return client
 
 
 async def _call[T](client: Client, fn: Callable[[Client], Awaitable[T]]) -> T:
@@ -155,7 +163,7 @@ async def list_requests(
         sort_order: ``asc`` or ``desc``.
         scope_id: Restrict to a Caido scope (managed via ``scope_rules``).
     """
-    client = _ctx_client(ctx)
+    client = await _ctx_client(ctx)
     if client is None:
         return _no_client()
 
@@ -261,7 +269,7 @@ async def view_request(
         page: 1-indexed page number (only when no ``search_pattern``).
         page_size: Lines per page.
     """
-    client = _ctx_client(ctx)
+    client = await _ctx_client(ctx)
     if client is None:
         return _no_client()
 
@@ -379,7 +387,7 @@ async def repeat_request(
             - ``body`` — replace the body string entirely.
             - ``cookies`` — dict of cookies to add/update.
     """
-    client = _ctx_client(ctx)
+    client = await _ctx_client(ctx)
     if client is None:
         return _no_client()
     mods = modifications or {}
@@ -461,7 +469,7 @@ async def list_sitemap(
             (recursive subtree). Only meaningful with ``parent_id``.
         page: 1-indexed page (30 entries per page).
     """
-    client = _ctx_client(ctx)
+    client = await _ctx_client(ctx)
     if client is None:
         return _no_client()
     try:
@@ -495,7 +503,7 @@ async def view_sitemap_entry(
     Args:
         entry_id: ID from ``list_sitemap`` (or any nested entry).
     """
-    client = _ctx_client(ctx)
+    client = await _ctx_client(ctx)
     if client is None:
         return _no_client()
     try:
@@ -554,7 +562,7 @@ async def scope_rules(
         scope_id: Required for ``get`` / ``update`` / ``delete``.
         scope_name: Required for ``create`` / ``update``.
     """
-    client = _ctx_client(ctx)
+    client = await _ctx_client(ctx)
     if client is None:
         return _no_client()
 

tests/test_caido_bootstrap.py

@@ -0,0 +1,81 @@
+"""A bootstrap that dies mid-setup must not leave its transport behind.
+
+The bootstrap now runs concurrently with the scan start, so teardown can
+cancel it at any await — including inside ``Client.connect()``, where the
+client exists but no caller will ever see it to close it.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+import types
+from typing import Any
+
+import pytest
+
+from strix.runtime.caido_bootstrap import bootstrap_caido
+
+
+class _FakeExecResult:
+    stderr = b""
+    exit_code = 0
+
+    def __init__(self, stdout: str) -> None:
+        self.stdout = stdout
+
+    def ok(self) -> bool:
+        return True
+
+
+class _FakeSession:
+    async def exec(self, *_args: Any, **_kwargs: Any) -> _FakeExecResult:
+        return _FakeExecResult('{"data":{"loginAsGuest":{"token":{"accessToken":"t"}}}}')
+
+
+class _FakeClient:
+    def __init__(self, connect_error: BaseException) -> None:
+        self.connect_error = connect_error
+        self.closed = False
+
+    async def connect(self) -> None:
+        raise self.connect_error
+
+    async def aclose(self) -> None:
+        self.closed = True
+
+
+async def _bootstrap_expecting(
+    monkeypatch: pytest.MonkeyPatch, error: BaseException
+) -> _FakeClient:
+    """Run a bootstrap whose ``connect()`` fails with ``error``."""
+    client = _FakeClient(error)
+    # The SDK is imported inside bootstrap_caido (it is slow to import), so the
+    # fakes are injected as the modules it imports.
+    sdk = types.ModuleType("caido_sdk_client")
+    sdk.Client = lambda *_a, **_k: client  # type: ignore[attr-defined]
+    sdk.TokenAuthOptions = lambda token: token  # type: ignore[attr-defined]
+    sdk_types = types.ModuleType("caido_sdk_client.types")
+    sdk_types.CreateProjectOptions = lambda **_k: None  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "caido_sdk_client", sdk)
+    monkeypatch.setitem(sys.modules, "caido_sdk_client.types", sdk_types)
+
+    with pytest.raises(type(error)):
+        await bootstrap_caido(
+            _FakeSession(),  # type: ignore[arg-type]
+            host_url="http://host",
+            container_url="http://container",
+        )
+    return client
+
+
+async def test_cancellation_during_connect_closes_the_client(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    client = await _bootstrap_expecting(monkeypatch, asyncio.CancelledError())
+    assert client.closed
+
+
+async def test_failed_connect_closes_the_client(monkeypatch: pytest.MonkeyPatch) -> None:
+    client = await _bootstrap_expecting(monkeypatch, RuntimeError("no listener"))
+    assert client.closed

tests/test_caido_handle.py

@@ -0,0 +1,103 @@
+"""Tests for the concurrent Caido bootstrap handle."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import pytest
+
+from strix.runtime.caido_handle import CaidoBootstrapHandle
+
+
+class _FakeClient:
+    def __init__(self) -> None:
+        self.closed = False
+
+    async def aclose(self) -> None:
+        self.closed = True
+
+
+def _handle(coro: Any) -> CaidoBootstrapHandle:
+    return CaidoBootstrapHandle(asyncio.ensure_future(coro))
+
+
+async def test_get_waits_for_the_bootstrap() -> None:
+    client = _FakeClient()
+    started = asyncio.Event()
+
+    async def _bootstrap() -> Any:
+        started.set()
+        await asyncio.sleep(0.01)
+        return client
+
+    handle = _handle(_bootstrap())
+    await started.wait()
+    assert handle.peek() is None
+    assert await handle.get() is client
+    assert handle.peek() is client
+
+
+async def test_get_reraises_bootstrap_failure_to_every_caller() -> None:
+    async def _bootstrap() -> Any:
+        raise RuntimeError("caido never came up")
+
+    handle = _handle(_bootstrap())
+    for _ in range(2):
+        with pytest.raises(RuntimeError, match="caido never came up"):
+            await handle.get()
+    assert handle.peek() is None
+
+
+async def test_caller_cancellation_does_not_cancel_the_shared_bootstrap() -> None:
+    client = _FakeClient()
+
+    async def _bootstrap() -> Any:
+        await asyncio.sleep(0.05)
+        return client
+
+    handle = _handle(_bootstrap())
+
+    with pytest.raises(TimeoutError):
+        await asyncio.wait_for(handle.get(), timeout=0.01)
+
+    assert await handle.get() is client
+
+
+async def test_aclose_closes_a_finished_client() -> None:
+    client = _FakeClient()
+
+    async def _bootstrap() -> Any:
+        return client
+
+    handle = _handle(_bootstrap())
+    await handle.get()
+    await handle.aclose()
+    assert client.closed is True
+
+
+async def test_aclose_cancels_an_in_flight_bootstrap() -> None:
+    cancelled = asyncio.Event()
+
+    async def _bootstrap() -> Any:
+        try:
+            await asyncio.sleep(10)
+        except asyncio.CancelledError:
+            cancelled.set()
+            raise
+        return _FakeClient()
+
+    handle = _handle(_bootstrap())
+    await asyncio.sleep(0)
+    await handle.aclose()
+    assert cancelled.is_set()
+
+
+async def test_aclose_swallows_a_failed_bootstrap() -> None:
+    async def _bootstrap() -> Any:
+        raise RuntimeError("boom")
+
+    handle = _handle(_bootstrap())
+    with pytest.raises(RuntimeError, match="boom"):
+        await handle.get()
+    await handle.aclose()

tests/test_proxy_client.py

@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, cast
 
 import pytest
 
+from strix.runtime.caido_handle import CaidoBootstrapHandle
 from strix.tools.proxy import caido_api, tools
 
 
@@ -198,12 +199,31 @@ class _Ctx:
         self.context = context
 
 
-def test_ctx_client_returns_client_when_present() -> None:
+async def test_ctx_client_returns_client_when_present() -> None:
     client = _FakeClient("host")
-    got = tools._ctx_client(cast("Any", _Ctx({"caido_client": client})))
+    got = await tools._ctx_client(cast("Any", _Ctx({"caido_client": client})))
     assert got is client
 
 
-def test_ctx_client_returns_none_without_client() -> None:
-    assert tools._ctx_client(cast("Any", _Ctx({}))) is None
-    assert tools._ctx_client(cast("Any", _Ctx(None))) is None
+async def test_ctx_client_returns_none_without_client() -> None:
+    assert await tools._ctx_client(cast("Any", _Ctx({}))) is None
+    assert await tools._ctx_client(cast("Any", _Ctx(None))) is None
+
+
+async def test_ctx_client_resolves_bootstrap_handle() -> None:
+    client = _FakeClient("host")
+
+    async def _bootstrap() -> Any:
+        return client
+
+    handle = CaidoBootstrapHandle(asyncio.ensure_future(_bootstrap()))
+    got = await tools._ctx_client(cast("Any", _Ctx({"caido_client": handle})))
+    assert got is client
+
+
+async def test_ctx_client_degrades_when_bootstrap_failed() -> None:
+    async def _bootstrap() -> Any:
+        raise RuntimeError("caido never came up")
+
+    handle = CaidoBootstrapHandle(asyncio.ensure_future(_bootstrap()))
+    assert await tools._ctx_client(cast("Any", _Ctx({"caido_client": handle}))) is None