Retry the Caido client connection

strix/runtime/caido_bootstrap.py

@@ -27,6 +27,8 @@ logger = logging.getLogger(__name__)
 _LOGIN_AS_GUEST_BODY = (
     '{"query":"mutation LoginAsGuest { loginAsGuest { token { accessToken } } }"}'
 )
+_PROJECT_SETUP_TIMEOUT_MS = 45_000
+_BOOTSTRAP_ATTEMPTS = 3
 
 
 async def _login_as_guest(
@@ -78,6 +80,69 @@ async def _login_as_guest(
     raise RuntimeError(f"loginAsGuest failed after {attempts} attempts: {last_err}")
 
 
+async def _find_sandbox_project(client: Client) -> str | None:
+    """Look for a project a create that timed out client-side may have left behind."""
+    with contextlib.suppress(Exception):
+        projects = [item for item in await client.project.list() if item.name == "sandbox"]
+        if projects:
+            return str(max(projects, key=lambda item: item.id).id)
+    return None
+
+
+async def _setup_project(host_url: str, access_token: str) -> None:
+    """Select the sandbox project, retrying the whole connect/create/select sequence.
+
+    Each attempt gets a fresh client with a deadline well past the SDK default:
+    these mutations are slow on a cold Caido, while the long-lived client the
+    scan uses keeps the short default so a traffic poll cannot stall on it.
+    Until a project is selected Caido answers every proxied request with a 500,
+    so giving up here costs the whole run, not just the traffic capture.
+    """
+    from caido_sdk_client import Client, TokenAuthOptions
+    from caido_sdk_client.types import CreateProjectOptions
+
+    project_id: str | None = None
+    last_exc: Exception | None = None
+    for attempt in range(1, _BOOTSTRAP_ATTEMPTS + 1):
+        client = Client(
+            host_url,
+            auth=TokenAuthOptions(token=access_token),
+            timeout_ms=_PROJECT_SETUP_TIMEOUT_MS,
+        )
+        try:
+            await client.connect()
+            if project_id is None:
+                try:
+                    created = await client.project.create(
+                        CreateProjectOptions(name="sandbox", temporary=True),
+                    )
+                except Exception:
+                    # A create that timed out client-side may still have landed.
+                    project_id = await _find_sandbox_project(client)
+                    raise
+                project_id = created.id
+            await client.project.select(project_id)
+        except Exception as exc:  # noqa: BLE001
+            last_exc = exc
+            logger.warning(
+                "Caido project setup attempt %d/%d failed: %s",
+                attempt,
+                _BOOTSTRAP_ATTEMPTS,
+                exc,
+            )
+            if attempt < _BOOTSTRAP_ATTEMPTS:
+                await asyncio.sleep(min(2.0 * attempt, 8.0))
+        else:
+            logger.info("Caido project selected: %s", project_id)
+            return
+        finally:
+            with contextlib.suppress(Exception):
+                await client.aclose()
+    raise RuntimeError(
+        f"Caido project setup failed after {_BOOTSTRAP_ATTEMPTS} attempts"
+    ) from last_exc
+
+
 async def bootstrap_caido(
     session: BaseSandboxSession,
     *,
@@ -89,27 +154,39 @@ async def bootstrap_caido(
     # only needed once a sandbox is actually being bootstrapped, so it is
     # imported here rather than at module scope.
     from caido_sdk_client import Client, TokenAuthOptions
-    from caido_sdk_client.types import CreateProjectOptions
 
     logger.info("Bootstrapping Caido client (host=%s, container=%s)", host_url, container_url)
 
     access_token = await _login_as_guest(session, container_url=container_url)
 
-    client = Client(host_url, auth=TokenAuthOptions(token=access_token))
-    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 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()
-        raise
-    logger.info("Caido project selected: %s", project.id)
-    return client
+    await _setup_project(host_url, access_token)
+
+    last_exc: Exception | None = None
+    for attempt in range(1, _BOOTSTRAP_ATTEMPTS + 1):
+        client = Client(host_url, auth=TokenAuthOptions(token=access_token))
+        try:
+            # A cancellation while connecting can leave a half-connected
+            # transport behind, so close the client before propagating it.
+            await client.connect()
+        except Exception as exc:  # noqa: BLE001
+            with contextlib.suppress(Exception):
+                await client.aclose()
+            last_exc = exc
+            logger.warning(
+                "Caido client connect attempt %d/%d failed: %s",
+                attempt,
+                _BOOTSTRAP_ATTEMPTS,
+                exc,
+            )
+            if attempt < _BOOTSTRAP_ATTEMPTS:
+                await asyncio.sleep(min(2.0 * attempt, 8.0))
+        except BaseException:
+            # Teardown can cancel the bootstrap at any await; do not retry.
+            with contextlib.suppress(Exception):
+                await client.aclose()
+            raise
+        else:
+            return client
+    raise RuntimeError(
+        f"Caido client connect failed after {_BOOTSTRAP_ATTEMPTS} attempts"
+    ) from last_exc

tests/test_caido_bootstrap.py

@@ -10,13 +10,18 @@ from __future__ import annotations
 import asyncio
 import sys
 import types
-from typing import Any
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
 
 import pytest
 
 from strix.runtime.caido_bootstrap import bootstrap_caido
 
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
 class _FakeExecResult:
     stderr = b""
     exit_code = 0
@@ -33,49 +38,305 @@ class _FakeSession:
         return _FakeExecResult('{"data":{"loginAsGuest":{"token":{"accessToken":"t"}}}}')
 
 
+@dataclass
+class _FakeProject:
+    id: str
+    name: str = "sandbox"
+
+
+class _FakeProjectSDK:
+    def __init__(
+        self,
+        *,
+        create_errors: list[BaseException] | None = None,
+        select_errors: list[BaseException] | None = None,
+        projects: list[_FakeProject] | None = None,
+    ) -> None:
+        self.create_errors = list(create_errors or [])
+        self.select_errors = list(select_errors or [])
+        self.projects = projects or []
+        self.create_calls = 0
+        self.selected_ids: list[str] = []
+        self.list_calls = 0
+
+    async def create(self, _options: Any) -> _FakeProject:
+        self.create_calls += 1
+        if self.create_errors:
+            raise self.create_errors.pop(0)
+        return _FakeProject("created")
+
+    async def select(self, project_id: str) -> _FakeProject:
+        self.selected_ids.append(project_id)
+        if self.select_errors:
+            raise self.select_errors.pop(0)
+        return _FakeProject(project_id)
+
+    async def list(self) -> list[_FakeProject]:
+        self.list_calls += 1
+        return self.projects
+
+
 class _FakeClient:
-    def __init__(self, connect_error: BaseException) -> None:
+    def __init__(
+        self,
+        connect_error: BaseException | None = None,
+        *,
+        project: _FakeProjectSDK | None = None,
+    ) -> None:
         self.connect_error = connect_error
+        self.project = project or _FakeProjectSDK()
         self.closed = False
 
     async def connect(self) -> None:
-        raise self.connect_error
+        if self.connect_error is not 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.
+class _FakeClientFactory:
+    def __init__(self, clients: list[_FakeClient]) -> None:
+        self.clients = iter(clients)
+        self.calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
+
+    def __call__(self, *args: Any, **kwargs: Any) -> _FakeClient:
+        self.calls.append((args, kwargs))
+        return next(self.clients)
+
+
+def _install_sdk(
+    monkeypatch: pytest.MonkeyPatch,
+    clients: list[_FakeClient],
+) -> _FakeClientFactory:
+    factory = _FakeClientFactory(clients)
     sdk = types.ModuleType("caido_sdk_client")
-    sdk.Client = lambda *_a, **_k: client  # type: ignore[attr-defined]
+    sdk.Client = factory  # 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]
+    sdk_types.CreateProjectOptions = lambda **kwargs: kwargs  # type: ignore[attr-defined]
     monkeypatch.setitem(sys.modules, "caido_sdk_client", sdk)
     monkeypatch.setitem(sys.modules, "caido_sdk_client.types", sdk_types)
+    return factory
 
-    with pytest.raises(type(error)):
+
+def _setup_clients(
+    project: _FakeProjectSDK,
+    count: int,
+    *,
+    connect_errors: list[BaseException | None] | None = None,
+) -> list[_FakeClient]:
+    """One client per setup attempt, all sharing the same server-side project state."""
+    errors: list[BaseException | None] = list(connect_errors or [])
+    errors += [None] * (count - len(errors))
+    return [_FakeClient(errors[i], project=project) for i in range(count)]
+
+
+async def _bootstrap_expecting(
+    monkeypatch: pytest.MonkeyPatch, errors: Sequence[BaseException]
+) -> tuple[list[_FakeClient], list[float], BaseException]:
+    """Run a bootstrap whose scan-client connections fail."""
+    setup_client = _FakeClient()
+    scan_clients = [_FakeClient(error) for error in errors]
+    _install_sdk(monkeypatch, [setup_client, *scan_clients])
+    sleep_calls: list[float] = []
+
+    async def _sleep(delay: float) -> None:
+        sleep_calls.append(delay)
+
+    monkeypatch.setattr("strix.runtime.caido_bootstrap.asyncio.sleep", _sleep)
+
+    with pytest.raises(BaseException) as exc_info:
         await bootstrap_caido(
             _FakeSession(),  # type: ignore[arg-type]
             host_url="http://host",
             container_url="http://container",
         )
-    return client
+    assert setup_client.closed
+    return [setup_client, *scan_clients], sleep_calls, exc_info.value
 
 
 async def test_cancellation_during_connect_closes_the_client(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    client = await _bootstrap_expecting(monkeypatch, asyncio.CancelledError())
-    assert client.closed
+    clients, sleep_calls, error = await _bootstrap_expecting(
+        monkeypatch,
+        [asyncio.CancelledError()],
+    )
+    assert isinstance(error, asyncio.CancelledError)
+    assert len(clients) == 2
+    assert all(client.closed for client in clients)
+    assert sleep_calls == []
 
 
 async def test_failed_connect_closes_the_client(monkeypatch: pytest.MonkeyPatch) -> None:
-    client = await _bootstrap_expecting(monkeypatch, RuntimeError("no listener"))
-    assert client.closed
+    errors = [
+        RuntimeError("first"),
+        RuntimeError("second"),
+        RuntimeError("last"),
+    ]
+    clients, sleep_calls, error = await _bootstrap_expecting(monkeypatch, errors)
+    assert isinstance(error, RuntimeError)
+    assert str(error) == "Caido client connect failed after 3 attempts"
+    assert error.__cause__ is errors[-1]
+    assert len(clients) == 4
+    assert all(client.closed for client in clients)
+    assert sleep_calls == [2.0, 4.0]
+
+
+async def test_select_retries_without_creating_another_project(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    setup_project = _FakeProjectSDK(select_errors=[RuntimeError("not ready")])
+    setup_clients = _setup_clients(setup_project, 2)
+    returned_client = _FakeClient()
+    factory = _install_sdk(monkeypatch, [*setup_clients, returned_client])
+    sleep_calls: list[float] = []
+
+    async def _sleep(delay: float) -> None:
+        sleep_calls.append(delay)
+
+    monkeypatch.setattr(
+        "strix.runtime.caido_bootstrap.asyncio.sleep",
+        _sleep,
+    )
+
+    result = await bootstrap_caido(
+        _FakeSession(),  # type: ignore[arg-type]
+        host_url="http://host",
+        container_url="http://container",
+    )
+
+    assert result is returned_client
+    assert setup_project.create_calls == 1
+    assert setup_project.selected_ids == ["created", "created"]
+    assert all(client.closed for client in setup_clients)
+    assert factory.calls[0][1]["timeout_ms"] == 45_000
+    assert factory.calls[1][1]["timeout_ms"] == 45_000
+    assert factory.calls[2][1].get("timeout_ms") is None
+    assert sleep_calls == [2.0]
+
+
+async def test_create_failure_reuses_the_most_recent_sandbox_project(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    setup_project = _FakeProjectSDK(
+        create_errors=[RuntimeError("create timed out")],
+        projects=[
+            _FakeProject("project-1"),
+            _FakeProject("project-2"),
+            _FakeProject("other", name="other"),
+        ],
+    )
+    setup_clients = _setup_clients(setup_project, 2)
+    returned_client = _FakeClient()
+    _install_sdk(monkeypatch, [*setup_clients, returned_client])
+
+    async def _sleep(_delay: float) -> None:
+        pass
+
+    monkeypatch.setattr("strix.runtime.caido_bootstrap.asyncio.sleep", _sleep)
+
+    result = await bootstrap_caido(
+        _FakeSession(),  # type: ignore[arg-type]
+        host_url="http://host",
+        container_url="http://container",
+    )
+
+    assert result is returned_client
+    assert setup_project.create_calls == 1
+    assert setup_project.list_calls == 1
+    assert setup_project.selected_ids == ["project-2"]
+    assert all(client.closed for client in setup_clients)
+
+
+async def test_project_setup_failure_chains_last_error_and_closes_setup_client(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    errors: list[BaseException] = [
+        RuntimeError("first"),
+        RuntimeError("second"),
+        RuntimeError("last"),
+    ]
+    setup_project = _FakeProjectSDK(select_errors=errors)
+    setup_clients = _setup_clients(setup_project, 3)
+    _install_sdk(monkeypatch, setup_clients)
+
+    async def _sleep(_delay: float) -> None:
+        pass
+
+    monkeypatch.setattr("strix.runtime.caido_bootstrap.asyncio.sleep", _sleep)
+
+    with pytest.raises(
+        RuntimeError,
+        match="Caido project setup failed after 3 attempts",
+    ) as exc_info:
+        await bootstrap_caido(
+            _FakeSession(),  # type: ignore[arg-type]
+            host_url="http://host",
+            container_url="http://container",
+        )
+
+    assert exc_info.value.__cause__ is errors[-1]
+    assert all(client.closed for client in setup_clients)
+    assert setup_project.create_calls == 1
+    assert setup_project.selected_ids == ["created", "created", "created"]
+
+
+async def test_cancelled_project_setup_is_not_retried(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    setup_project = _FakeProjectSDK(select_errors=[asyncio.CancelledError()])
+    setup_clients = _setup_clients(setup_project, 1)
+    _install_sdk(monkeypatch, setup_clients)
+    sleep_calls: list[float] = []
+
+    async def _sleep(delay: float) -> None:
+        sleep_calls.append(delay)
+
+    monkeypatch.setattr(
+        "strix.runtime.caido_bootstrap.asyncio.sleep",
+        _sleep,
+    )
+
+    with pytest.raises(asyncio.CancelledError):
+        await bootstrap_caido(
+            _FakeSession(),  # type: ignore[arg-type]
+            host_url="http://host",
+            container_url="http://container",
+        )
+
+    assert setup_project.create_calls == 1
+    assert setup_project.selected_ids == ["created"]
+    assert sleep_calls == []
+    assert all(client.closed for client in setup_clients)
+
+
+async def test_setup_connect_failure_is_retried(monkeypatch: pytest.MonkeyPatch) -> None:
+    setup_project = _FakeProjectSDK()
+    setup_clients = _setup_clients(
+        setup_project,
+        2,
+        connect_errors=[RuntimeError("gateway not ready")],
+    )
+    returned_client = _FakeClient()
+    _install_sdk(monkeypatch, [*setup_clients, returned_client])
+    sleep_calls: list[float] = []
+
+    async def _sleep(delay: float) -> None:
+        sleep_calls.append(delay)
+
+    monkeypatch.setattr("strix.runtime.caido_bootstrap.asyncio.sleep", _sleep)
+
+    result = await bootstrap_caido(
+        _FakeSession(),  # type: ignore[arg-type]
+        host_url="http://host",
+        container_url="http://container",
+    )
+
+    assert result is returned_client
+    assert setup_project.create_calls == 1
+    assert setup_project.selected_ids == ["created"]
+    assert setup_project.list_calls == 0
+    assert sleep_calls == [2.0]
+    assert all(client.closed for client in setup_clients)