From dd02fede169acbb1059925f55b2a3edc9b3dcee8 Mon Sep 17 00:00:00 2001 From: Alex Schapiro Date: Mon, 31 Aug 2026 13:06:25 +0000 Subject: [PATCH 1/4] Retry Caido project setup with isolated client --- strix/runtime/caido_bootstrap.py | 59 ++++++++- tests/test_caido_bootstrap.py | 214 +++++++++++++++++++++++++++++-- 2 files changed, 256 insertions(+), 17 deletions(-) diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index a9c7c82ae..80a5815a3 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/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 diff --git a/tests/test_caido_bootstrap.py b/tests/test_caido_bootstrap.py index 98e42a94d..f4c68c884 100644 --- a/tests/test_caido_bootstrap.py +++ b/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 From 0e3d8e506eb1ab8c34aeada6e73c6e82010c58f7 Mon Sep 17 00:00:00 2001 From: Alex Schapiro Date: Mon, 31 Aug 2026 13:07:41 +0000 Subject: [PATCH 2/4] Avoid assert in Caido setup failure path --- strix/runtime/caido_bootstrap.py | 1 - 1 file changed, 1 deletion(-) diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index 80a5815a3..9f93cecec 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/strix/runtime/caido_bootstrap.py @@ -120,7 +120,6 @@ async def _setup_project(host_url: str, access_token: str) -> None: 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 From 5ea5129b6732c645378ae59bb77e6f3dc6677416 Mon Sep 17 00:00:00 2001 From: Alex Schapiro Date: Mon, 31 Aug 2026 14:15:13 +0000 Subject: [PATCH 3/4] Retry the Caido setup connection --- strix/runtime/caido_bootstrap.py | 95 ++++++++++++++++++-------------- tests/test_caido_bootstrap.py | 69 ++++++++++++++++++----- 2 files changed, 111 insertions(+), 53 deletions(-) diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index 9f93cecec..ca143ce4d 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/strix/runtime/caido_bootstrap.py @@ -80,52 +80,67 @@ 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: - """Connect with a longer deadline and select the sandbox project.""" + """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 - 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( + project_id: str | None = None + last_exc: Exception | None = None + for attempt in range(1, _PROJECT_SETUP_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), ) - 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 - raise RuntimeError( - f"Caido project setup failed after {_PROJECT_SETUP_ATTEMPTS} attempts" - ) from last_exc - finally: - with contextlib.suppress(Exception): - await client.aclose() + 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, + _PROJECT_SETUP_ATTEMPTS, + exc, + ) + if attempt < _PROJECT_SETUP_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 {_PROJECT_SETUP_ATTEMPTS} attempts" + ) from last_exc async def bootstrap_caido( diff --git a/tests/test_caido_bootstrap.py b/tests/test_caido_bootstrap.py index f4c68c884..0a4680e6d 100644 --- a/tests/test_caido_bootstrap.py +++ b/tests/test_caido_bootstrap.py @@ -116,6 +116,18 @@ def _install_sdk( return factory +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, error: BaseException ) -> _FakeClient: @@ -150,9 +162,9 @@ 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) + setup_clients = _setup_clients(setup_project, 2) returned_client = _FakeClient() - factory = _install_sdk(monkeypatch, [setup_client, returned_client]) + factory = _install_sdk(monkeypatch, [*setup_clients, returned_client]) sleep_calls: list[float] = [] async def _sleep(delay: float) -> None: @@ -172,9 +184,10 @@ async def test_select_retries_without_creating_another_project( assert result is returned_client assert setup_project.create_calls == 1 assert setup_project.selected_ids == ["created", "created"] - assert setup_client.closed + assert all(client.closed for client in setup_clients) assert factory.calls[0][1]["timeout_ms"] == 45_000 - assert factory.calls[1][1].get("timeout_ms") is None + assert factory.calls[1][1]["timeout_ms"] == 45_000 + assert factory.calls[2][1].get("timeout_ms") is None assert sleep_calls == [2.0] @@ -189,9 +202,9 @@ async def test_create_failure_reuses_the_most_recent_sandbox_project( _FakeProject("other", name="other"), ], ) - setup_client = _FakeClient(project=setup_project) + setup_clients = _setup_clients(setup_project, 2) returned_client = _FakeClient() - _install_sdk(monkeypatch, [setup_client, returned_client]) + _install_sdk(monkeypatch, [*setup_clients, returned_client]) async def _sleep(_delay: float) -> None: pass @@ -208,7 +221,7 @@ async def test_create_failure_reuses_the_most_recent_sandbox_project( assert setup_project.create_calls == 1 assert setup_project.list_calls == 1 assert setup_project.selected_ids == ["project-2"] - assert setup_client.closed + assert all(client.closed for client in setup_clients) async def test_project_setup_failure_chains_last_error_and_closes_setup_client( @@ -220,8 +233,8 @@ async def test_project_setup_failure_chains_last_error_and_closes_setup_client( RuntimeError("last"), ] setup_project = _FakeProjectSDK(select_errors=errors) - setup_client = _FakeClient(project=setup_project) - _install_sdk(monkeypatch, [setup_client]) + setup_clients = _setup_clients(setup_project, 3) + _install_sdk(monkeypatch, setup_clients) async def _sleep(_delay: float) -> None: pass @@ -239,7 +252,7 @@ async def test_project_setup_failure_chains_last_error_and_closes_setup_client( ) assert exc_info.value.__cause__ is errors[-1] - assert setup_client.closed + assert all(client.closed for client in setup_clients) assert setup_project.create_calls == 1 assert setup_project.selected_ids == ["created", "created", "created"] @@ -248,8 +261,8 @@ 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]) + setup_clients = _setup_clients(setup_project, 1) + _install_sdk(monkeypatch, setup_clients) sleep_calls: list[float] = [] async def _sleep(delay: float) -> None: @@ -270,4 +283,34 @@ async def test_cancelled_project_setup_is_not_retried( assert setup_project.create_calls == 1 assert setup_project.selected_ids == ["created"] assert sleep_calls == [] - assert setup_client.closed + 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) From efcef9fc1a18c7a0ad94de1c89404fc9ee952661 Mon Sep 17 00:00:00 2001 From: Alex Schapiro Date: Mon, 31 Aug 2026 14:26:05 +0000 Subject: [PATCH 4/4] Retry the Caido client connection --- strix/runtime/caido_bootstrap.py | 52 +++++++++++++++++++++----------- tests/test_caido_bootstrap.py | 50 ++++++++++++++++++++++-------- 2 files changed, 72 insertions(+), 30 deletions(-) diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index ca143ce4d..c1b96c7f9 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/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 diff --git a/tests/test_caido_bootstrap.py b/tests/test_caido_bootstrap.py index 0a4680e6d..655ca3cef 100644 --- a/tests/test_caido_bootstrap.py +++ b/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(