From efcef9fc1a18c7a0ad94de1c89404fc9ee952661 Mon Sep 17 00:00:00 2001 From: Alex Schapiro Date: Mon, 31 Aug 2026 14:26:05 +0000 Subject: [PATCH] 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 ca143ce4..c1b96c7f 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 0a4680e6..655ca3ce 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(