Retry Caido project setup with isolated client

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
+_PROJECT_SETUP_ATTEMPTS = 3
 
 
 async def _login_as_guest(
@@ -78,6 +80,55 @@ async def _login_as_guest(
     raise RuntimeError(f"loginAsGuest failed after {attempts} attempts: {last_err}")
 
 
+async def _setup_project(host_url: str, access_token: str) -> None:
+    """Connect with a longer deadline and select the sandbox project."""
+    from caido_sdk_client import Client, TokenAuthOptions
+    from caido_sdk_client.types import CreateProjectOptions
+
+    client = Client(
+        host_url,
+        auth=TokenAuthOptions(token=access_token),
+        timeout_ms=_PROJECT_SETUP_TIMEOUT_MS,
+    )
+    try:
+        await client.connect()
+        project = None
+        last_exc: Exception | None = None
+        for i in range(1, _PROJECT_SETUP_ATTEMPTS + 1):
+            try:
+                if project is None:
+                    project = await client.project.create(
+                        CreateProjectOptions(name="sandbox", temporary=True),
+                    )
+                await client.project.select(project.id)
+            except Exception as exc:  # noqa: BLE001
+                last_exc = exc
+                if project is None:
+                    with contextlib.suppress(Exception):
+                        projects = await client.project.list()
+                        sandbox_projects = [item for item in projects if item.name == "sandbox"]
+                        if sandbox_projects:
+                            project = max(sandbox_projects, key=lambda item: item.id)
+                logger.warning(
+                    "Caido project setup attempt %d/%d failed: %s",
+                    i,
+                    _PROJECT_SETUP_ATTEMPTS,
+                    exc,
+                )
+                if i < _PROJECT_SETUP_ATTEMPTS:
+                    await asyncio.sleep(min(2.0 * i, 8.0))
+            else:
+                logger.info("Caido project selected: %s", project.id)
+                return
+        assert last_exc is not None
+        raise RuntimeError(
+            f"Caido project setup failed after {_PROJECT_SETUP_ATTEMPTS} attempts"
+        ) from last_exc
+    finally:
+        with contextlib.suppress(Exception):
+            await client.aclose()
+
+
 async def bootstrap_caido(
     session: BaseSandboxSession,
     *,
@@ -89,27 +140,23 @@ 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)
 
+    await _setup_project(host_url, access_token)
+
     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

tests/test_caido_bootstrap.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 import asyncio
 import sys
 import types
+from dataclasses import dataclass
 from typing import Any
 
 import pytest
@@ -33,32 +34,95 @@ 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
 
 
+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 = 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 **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
+
+
 async def _bootstrap_expecting(
     monkeypatch: pytest.MonkeyPatch, error: BaseException
 ) -> _FakeClient:
     """Run a bootstrap whose ``connect()`` fails with ``error``."""
+    setup_client = _FakeClient()
     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)
+    _install_sdk(monkeypatch, [setup_client, client])
 
     with pytest.raises(type(error)):
         await bootstrap_caido(
@@ -66,6 +130,7 @@ async def _bootstrap_expecting(
             host_url="http://host",
             container_url="http://container",
         )
+    assert setup_client.closed
     return client
 
 
@@ -79,3 +144,130 @@ async def test_cancellation_during_connect_closes_the_client(
 async def test_failed_connect_closes_the_client(monkeypatch: pytest.MonkeyPatch) -> None:
     client = await _bootstrap_expecting(monkeypatch, RuntimeError("no listener"))
     assert client.closed
+
+
+async def test_select_retries_without_creating_another_project(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    setup_project = _FakeProjectSDK(select_errors=[RuntimeError("not ready")])
+    setup_client = _FakeClient(project=setup_project)
+    returned_client = _FakeClient()
+    factory = _install_sdk(monkeypatch, [setup_client, 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 setup_client.closed
+    assert factory.calls[0][1]["timeout_ms"] == 45_000
+    assert factory.calls[1][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_client = _FakeClient(project=setup_project)
+    returned_client = _FakeClient()
+    _install_sdk(monkeypatch, [setup_client, 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 setup_client.closed
+
+
+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_client = _FakeClient(project=setup_project)
+    _install_sdk(monkeypatch, [setup_client])
+
+    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 setup_client.closed
+    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_client = _FakeClient(project=setup_project)
+    _install_sdk(monkeypatch, [setup_client])
+    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 setup_client.closed