Retry the Caido client connection

strix/runtime/caido_bootstrap.py

@@ -28,7 +28,7 @@ _LOGIN_AS_GUEST_BODY = (
     '{"query":"mutation LoginAsGuest { loginAsGuest { token { accessToken } } }"}'
 )
 _PROJECT_SETUP_TIMEOUT_MS = 45_000
-_PROJECT_SETUP_ATTEMPTS = 3
+_BOOTSTRAP_ATTEMPTS = 3
 
 
 async def _login_as_guest(
@@ -103,7 +103,7 @@ async def _setup_project(host_url: str, access_token: str) -> None:
 
     project_id: str | None = None
     last_exc: Exception | None = None
-    for attempt in range(1, _PROJECT_SETUP_ATTEMPTS + 1):
+    for attempt in range(1, _BOOTSTRAP_ATTEMPTS + 1):
         client = Client(
             host_url,
             auth=TokenAuthOptions(token=access_token),
@@ -127,10 +127,10 @@ async def _setup_project(host_url: str, access_token: str) -> None:
             logger.warning(
                 "Caido project setup attempt %d/%d failed: %s",
                 attempt,
-                _PROJECT_SETUP_ATTEMPTS,
+                _BOOTSTRAP_ATTEMPTS,
                 exc,
             )
-            if attempt < _PROJECT_SETUP_ATTEMPTS:
+            if attempt < _BOOTSTRAP_ATTEMPTS:
                 await asyncio.sleep(min(2.0 * attempt, 8.0))
         else:
             logger.info("Caido project selected: %s", project_id)
@@ -139,7 +139,7 @@ async def _setup_project(host_url: str, access_token: str) -> None:
             with contextlib.suppress(Exception):
                 await client.aclose()
     raise RuntimeError(
-        f"Caido project setup failed after {_PROJECT_SETUP_ATTEMPTS} attempts"
+        f"Caido project setup failed after {_BOOTSTRAP_ATTEMPTS} attempts"
     ) from last_exc
 
 
@@ -161,16 +161,32 @@ async def bootstrap_caido(
 
     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()
-    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
-    return client
+    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

@@ -11,13 +11,17 @@ import asyncio
 import sys
 import types
 from dataclasses import dataclass
-from typing import Any
+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
@@ -129,33 +133,55 @@ def _setup_clients(
 
 
 async def _bootstrap_expecting(
-    monkeypatch: pytest.MonkeyPatch, error: BaseException
-) -> _FakeClient:
-    """Run a bootstrap whose ``connect()`` fails with ``error``."""
+    monkeypatch: pytest.MonkeyPatch, errors: Sequence[BaseException]
+) -> tuple[list[_FakeClient], list[float], BaseException]:
+    """Run a bootstrap whose scan-client connections fail."""
     setup_client = _FakeClient()
-    client = _FakeClient(error)
-    _install_sdk(monkeypatch, [setup_client, client])
+    scan_clients = [_FakeClient(error) for error in errors]
+    _install_sdk(monkeypatch, [setup_client, *scan_clients])
+    sleep_calls: list[float] = []
 
-    with pytest.raises(type(error)):
+    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",
         )
     assert setup_client.closed
-    return client
+    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(