diff --git a/strix/core/runner.py b/strix/core/runner.py index 9f1be995..1e5d66b4 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import contextlib import json import logging @@ -283,12 +282,11 @@ async def run_strix_scan( context: dict[str, Any] = { "coordinator": coordinator, "sandbox_session": bundle["session"], + # One ``SharedCaidoClient`` is reused by every agent in the scan + # (child contexts are shallow copies via ``dict(parent_ctx)``). It + # serializes access to the non-concurrency-safe GraphQL transport + # and rebuilds it if it dies mid-scan. "caido_client": bundle["caido_client"], - # One shared Caido client is reused by every agent in the scan; its - # GraphQL transport is not safe for concurrent use. Child contexts - # are shallow copies (``dict(parent_ctx)``) so they inherit this - # same lock object, serializing all proxy calls scan-wide. - "caido_lock": asyncio.Lock(), "agent_id": root_id, "parent_id": None, "interactive": interactive, diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index 7057d228..bd8423e1 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/strix/runtime/caido_bootstrap.py @@ -79,23 +79,56 @@ async def _login_as_guest( raise RuntimeError(f"loginAsGuest failed after {attempts} attempts: {last_err}") -async def bootstrap_caido( +async def _connect_client( session: BaseSandboxSession, *, host_url: str, container_url: str, ) -> Client: - """Connect to the in-container Caido sidecar and select a fresh project.""" - 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)) await client.connect() + return client + + +async def bootstrap_caido( + session: BaseSandboxSession, + *, + host_url: str, + container_url: str, +) -> tuple[Client, str]: + """Connect to the in-container Caido sidecar and select a fresh project. + + Returns the connected client and the id of the temporary project it + selected. The project id lets :func:`reconnect_caido` rebuild a dead + transport while staying on the *same* project (and its captured traffic) + instead of creating a new empty one. + """ + logger.info("Bootstrapping Caido client (host=%s, container=%s)", host_url, container_url) + + client = await _connect_client(session, host_url=host_url, container_url=container_url) project = await client.project.create( CreateProjectOptions(name="sandbox", temporary=True), ) await client.project.select(project.id) logger.info("Caido project selected: %s", project.id) + return client, str(project.id) + + +async def reconnect_caido( + session: BaseSandboxSession, + *, + host_url: str, + container_url: str, + project_id: str, +) -> Client: + """Rebuild a Caido client after its transport died, keeping the project. + + Re-authenticates, reconnects, and re-selects the existing project so the + caller keeps access to the traffic captured before the disconnect. + """ + logger.info("Reconnecting Caido client (host=%s, project=%s)", host_url, project_id) + client = await _connect_client(session, host_url=host_url, container_url=container_url) + await client.project.select(project_id) return client diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py index cdc6d955..84950666 100644 --- a/strix/runtime/session_manager.py +++ b/strix/runtime/session_manager.py @@ -5,15 +5,20 @@ from __future__ import annotations import logging import shutil from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from agents.sandbox.entries import BaseEntry, LocalDir from agents.sandbox.manifest import Environment, Manifest from strix.config import load_settings from strix.runtime.backends import get_backend -from strix.runtime.caido_bootstrap import bootstrap_caido +from strix.runtime.caido_bootstrap import bootstrap_caido, reconnect_caido from strix.runtime.local_dir_staging import stage_symlink_safe_dir +from strix.tools.proxy.caido_api import SharedCaidoClient + + +if TYPE_CHECKING: + from caido_sdk_client import Client as CaidoClient logger = logging.getLogger(__name__) @@ -131,16 +136,24 @@ async def create_or_reuse( host_caido_url = f"{scheme}://{caido_endpoint.host}:{caido_endpoint.port}" logger.debug("Caido host endpoint resolved: %s", host_caido_url) - caido_client = await bootstrap_caido( + caido_client, caido_project_id = await bootstrap_caido( session, host_url=host_caido_url, container_url=container_caido_url, ) + async def _reconnect_caido() -> CaidoClient: + return await reconnect_caido( + session, + host_url=host_caido_url, + container_url=container_caido_url, + project_id=caido_project_id, + ) + bundle = { "client": client, "session": session, - "caido_client": caido_client, + "caido_client": SharedCaidoClient(caido_client, _reconnect_caido), } _SESSION_CACHE[scan_id] = bundle logger.info("Sandbox session for scan %s ready and cached", scan_id) diff --git a/strix/tools/proxy/caido_api.py b/strix/tools/proxy/caido_api.py index cf5aef23..11e26859 100644 --- a/strix/tools/proxy/caido_api.py +++ b/strix/tools/proxy/caido_api.py @@ -3,7 +3,9 @@ from __future__ import annotations import asyncio +import contextlib import json +import logging import os import time import urllib.request @@ -26,6 +28,9 @@ if TYPE_CHECKING: from caido_sdk_client import Client as CaidoClient +logger = logging.getLogger(__name__) + + RequestPart = Literal["request", "response"] SortBy = Literal[ "timestamp", @@ -103,6 +108,14 @@ async def _new_client() -> Client: return client +async def _safe_aclose(client: Client | None) -> None: + """Close a (possibly dead) client without letting teardown errors escape.""" + if client is None: + return + with contextlib.suppress(Exception): + await client.aclose() + + def _is_connection_error(exc: BaseException) -> bool: message = str(exc).lower() if any(marker in message for marker in _CONNECTION_ERROR_MARKERS): @@ -126,15 +139,22 @@ async def get_client() -> Client: return client -async def call_with_client[T](fn: Callable[[Client], Awaitable[T]]) -> T: +async def call_with_client[T]( + fn: Callable[[Client], Awaitable[T]], *, idempotent: bool = True +) -> T: """Run ``fn`` against the shared client, serialized and reconnect-safe. The Caido GraphQL transport is not safe for concurrent use: two in-flight requests race and raise "Transport is already connected". All proxy calls are therefore serialized through ``_CLIENT_LOCK``. If the cached client's transport has since died ("Connector is closed" / "Server disconnected"), - the stale client is rebuilt and the call retried once, instead of every - subsequent call in the run failing against a dead client. + the stale client is closed and rebuilt so subsequent calls stop failing + against a dead client. + + ``fn`` is only re-run automatically when ``idempotent`` is true. For + mutations (replay, scope create/update/delete) a connection error may + arrive *after* Caido applied the change, so we heal the client for future + calls but re-raise instead of risking a double-apply. """ async with _CLIENT_LOCK: client = _CLIENT_CACHE.get("default") @@ -146,9 +166,46 @@ async def call_with_client[T](fn: Callable[[Client], Awaitable[T]]) -> T: except Exception as exc: if not _is_connection_error(exc): raise - client = await _new_client() - _CLIENT_CACHE["default"] = client - return await fn(client) + new_client = await _new_client() + _CLIENT_CACHE["default"] = new_client + await _safe_aclose(client) + if not idempotent: + raise + return await fn(new_client) + + +class SharedCaidoClient: + """Serialized, reconnect-safe wrapper around one host-side Caido client. + + Every agent in a scan shares a single instance (propagated through the + shallow-copied run context). ``call`` serializes access — the SDK transport + is not concurrency-safe — and, when the transport dies, rebuilds the client + via ``reconnect`` (which preserves the Caido project) and closes the dead + one, so a transient Caido restart no longer disables proxy tools for the + rest of the scan. + """ + + def __init__(self, client: Client, reconnect: Callable[[], Awaitable[Client]]) -> None: + self._client = client + self._reconnect = reconnect + self._lock = asyncio.Lock() + + async def call[T](self, fn: Callable[[Client], Awaitable[T]], *, idempotent: bool = True) -> T: + async with self._lock: + try: + return await fn(self._client) + except Exception as exc: + if not _is_connection_error(exc): + raise + dead, self._client = self._client, await self._reconnect() + await _safe_aclose(dead) + if not idempotent: + raise + return await fn(self._client) + + async def aclose(self) -> None: + async with self._lock: + await _safe_aclose(self._client) async def close_client() -> None: @@ -489,7 +546,9 @@ async def repeat_request( ) return await replay_send_raw(client, raw=raw, connection=connection) - return await call_with_client(_run) + # A replay mutates server state; don't auto-retry if the transport dies + # mid-send (the request may already have been sent). + return await call_with_client(_run, idempotent=False) async def scope_rules( @@ -510,7 +569,8 @@ async def scope_rules( scope_name=scope_name, ) - return await call_with_client(_run) + # get/list are read-only and safe to retry; create/update/delete mutate. + return await call_with_client(_run, idempotent=action in {"get", "list"}) async def _scope_rules_with_client( @@ -759,6 +819,7 @@ async def view_sitemap_entry(entry_id: str) -> dict[str, Any]: __all__ = [ "RequestPart", "ScopeAction", + "SharedCaidoClient", "SitemapDepth", "SortBy", "SortOrder", diff --git a/strix/tools/proxy/tools.py b/strix/tools/proxy/tools.py index 1610dfde..4fdb850a 100644 --- a/strix/tools/proxy/tools.py +++ b/strix/tools/proxy/tools.py @@ -2,8 +2,6 @@ from __future__ import annotations -import asyncio -import contextlib import dataclasses import json import logging @@ -15,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Literal from agents import RunContextWrapper, function_tool from strix.tools.proxy import caido_api +from strix.tools.proxy.caido_api import SharedCaidoClient logger = logging.getLogger(__name__) @@ -30,7 +29,7 @@ if TYPE_CHECKING: SortOrder, ) else: - from strix.tools.proxy.caido_api import ( # noqa: TC001 + from strix.tools.proxy.caido_api import ( RequestPart, SitemapDepth, SortBy, @@ -41,26 +40,18 @@ else: ScopeAction = Literal["get", "list", "create", "update", "delete"] -def _ctx_client(ctx: RunContextWrapper) -> Client | None: - inner = ctx.context if isinstance(ctx.context, dict) else {} - return inner.get("caido_client") +def _ctx_proxy(ctx: RunContextWrapper) -> SharedCaidoClient | None: + """Return the scan-wide serialized, reconnect-safe Caido client holder. - -def _ctx_lock(ctx: RunContextWrapper) -> contextlib.AbstractAsyncContextManager[None]: - """Return the scan-wide lock serializing access to the shared Caido client. - - All agents in a scan share one ``caido_client`` whose GraphQL transport is - not concurrency-safe (parallel calls raise "Transport is already - connected", and racing session teardown yields "Connector is closed" / - "Server disconnected"). Holding this lock around every proxy call serializes - them. Falls back to a no-op context when no lock is present (e.g. standalone - tool invocation outside a scan run). + All agents in a scan share one :class:`SharedCaidoClient` whose GraphQL + transport is not concurrency-safe (parallel calls raise "Transport is + already connected"). ``SharedCaidoClient.call`` serializes access and + rebuilds the transport if it dies mid-scan. Returns ``None`` when no holder + is present (e.g. standalone tool invocation outside a scan run). """ inner = ctx.context if isinstance(ctx.context, dict) else {} - lock = inner.get("caido_lock") - if isinstance(lock, asyncio.Lock): - return lock - return contextlib.nullcontext() + proxy = inner.get("caido_client") + return proxy if isinstance(proxy, SharedCaidoClient) else None def _to_tool_json(value: Any) -> Any: @@ -193,13 +184,13 @@ async def list_requests( sort_order: ``asc`` or ``desc``. scope_id: Restrict to a Caido scope (managed via ``scope_rules``). """ - client = _ctx_client(ctx) - if client is None: + proxy = _ctx_proxy(ctx) + if proxy is None: return _no_client() try: - async with _ctx_lock(ctx): - connection = await caido_api.list_requests_with_client( + connection = await proxy.call( + lambda client: caido_api.list_requests_with_client( client, httpql_filter=httpql_filter, first=first, @@ -208,6 +199,7 @@ async def list_requests( sort_order=sort_order, scope_id=scope_id, ) + ) entries = [] for edge in connection.edges: @@ -299,13 +291,14 @@ async def view_request( page: 1-indexed page number (only when no ``search_pattern``). page_size: Lines per page. """ - client = _ctx_client(ctx) - if client is None: + proxy = _ctx_proxy(ctx) + if proxy is None: return _no_client() try: - async with _ctx_lock(ctx): - result = await caido_api.get_request_with_client(client, request_id, part=part) + result = await proxy.call( + lambda client: caido_api.get_request_with_client(client, request_id, part=part) + ) if result is None: return json.dumps( {"success": False, "error": f"Request {request_id} not found"}, @@ -415,33 +408,38 @@ async def repeat_request( - ``body`` — replace the body string entirely. - ``cookies`` — dict of cookies to add/update. """ - client = _ctx_client(ctx) - if client is None: + proxy = _ctx_proxy(ctx) + if proxy is None: return _no_client() mods = modifications or {} - try: - async with _ctx_lock(ctx): - result = await caido_api.get_request_with_client(client, request_id, part="request") - if result is None or result.request.raw is None: - return json.dumps( - {"success": False, "error": f"Request {request_id} not found"}, - ensure_ascii=False, - default=str, - ) + async def _do(client: Client) -> dict[str, Any] | None: + result = await caido_api.get_request_with_client(client, request_id, part="request") + if result is None or result.request.raw is None: + return None + original = result.request + raw_str = result.request.raw.decode("utf-8", errors="replace") + components = caido_api.parse_raw_request(raw_str) + full_url = caido_api.full_url_from_components(original, components, mods) + modified = caido_api.apply_modifications(components, mods, full_url) + connection, raw = caido_api.build_raw_request( + method=modified["method"], + url=modified["url"], + headers=modified["headers"], + body=modified["body"], + ) + return await caido_api.replay_send_raw(client, raw=raw, connection=connection) - original = result.request - raw_str = result.request.raw.decode("utf-8", errors="replace") - components = caido_api.parse_raw_request(raw_str) - full_url = caido_api.full_url_from_components(original, components, mods) - modified = caido_api.apply_modifications(components, mods, full_url) - connection, raw = caido_api.build_raw_request( - method=modified["method"], - url=modified["url"], - headers=modified["headers"], - body=modified["body"], + try: + # A replay mutates target state, so don't auto-retry on a mid-send + # transport failure (the request may already have been sent). + replay = await proxy.call(_do, idempotent=False) + if replay is None: + return json.dumps( + {"success": False, "error": f"Request {request_id} not found"}, + ensure_ascii=False, + default=str, ) - replay = await caido_api.replay_send_raw(client, raw=raw, connection=connection) return _format_replay_tool_result(replay) except Exception as exc: # noqa: BLE001 return _err("repeat_request", exc) @@ -494,18 +492,19 @@ async def list_sitemap( (recursive subtree). Only meaningful with ``parent_id``. page: 1-indexed page (30 entries per page). """ - client = _ctx_client(ctx) - if client is None: + proxy = _ctx_proxy(ctx) + if proxy is None: return _no_client() try: - async with _ctx_lock(ctx): - payload = await caido_api.list_sitemap_with_client( + payload = await proxy.call( + lambda client: caido_api.list_sitemap_with_client( client, scope_id=scope_id, parent_id=parent_id, depth=depth, page=page, ) + ) return json.dumps(payload, ensure_ascii=False, default=str) except Exception as exc: # noqa: BLE001 return _err("list_sitemap", exc) @@ -526,12 +525,13 @@ async def view_sitemap_entry( Args: entry_id: ID from ``list_sitemap`` (or any nested entry). """ - client = _ctx_client(ctx) - if client is None: + proxy = _ctx_proxy(ctx) + if proxy is None: return _no_client() try: - async with _ctx_lock(ctx): - payload = await caido_api.view_sitemap_entry_with_client(client, entry_id) + payload = await proxy.call( + lambda client: caido_api.view_sitemap_entry_with_client(client, entry_id) + ) return json.dumps(payload, ensure_ascii=False, default=str) except Exception as exc: # noqa: BLE001 return _err("view_sitemap_entry", exc) @@ -583,80 +583,85 @@ async def scope_rules( scope_id: Required for ``get`` / ``update`` / ``delete``. scope_name: Required for ``create`` / ``update``. """ - client = _ctx_client(ctx) - if client is None: + proxy = _ctx_proxy(ctx) + if proxy is None: return _no_client() try: - async with _ctx_lock(ctx): - if action == "list": - scopes = await caido_api.scope_list(client) - return json.dumps( - {"success": True, "scopes": [_to_tool_json(s) for s in scopes]}, - ensure_ascii=False, - default=str, - ) - if action == "get": - if not scope_id: - return json.dumps( - {"success": False, "error": "Scope_id is required for action='get'"}, - ensure_ascii=False, - default=str, - ) - scope = await caido_api.scope_get(client, scope_id) - return json.dumps( - {"success": True, "scope": _to_tool_json(scope)}, - ensure_ascii=False, - default=str, - ) - if action == "create": - if not scope_name: - return json.dumps( - {"success": False, "error": "Scope_name is required for action='create'"}, - ensure_ascii=False, - default=str, - ) - scope = await caido_api.scope_create( - client, name=scope_name, allowlist=allowlist, denylist=denylist - ) - return json.dumps( - {"success": True, "scope": _to_tool_json(scope)}, - ensure_ascii=False, - default=str, - ) - if action == "update": - if not scope_id or not scope_name: - return json.dumps( - { - "success": False, - "error": "Scope_id and scope_name are required for action='update'", - }, - ensure_ascii=False, - default=str, - ) - scope = await caido_api.scope_update( - client, scope_id, name=scope_name, allowlist=allowlist, denylist=denylist - ) - return json.dumps( - {"success": True, "scope": _to_tool_json(scope)}, - ensure_ascii=False, - default=str, - ) - if not scope_id: - return json.dumps( - {"success": False, "error": "Scope_id is required for action='delete'"}, - ensure_ascii=False, - default=str, - ) - await caido_api.scope_delete(client, scope_id) + if action == "list": + scopes = await proxy.call(caido_api.scope_list) return json.dumps( - { - "success": True, - "deleted": scope_id, - "message": f"Scope {scope_id} deleted", - }, + {"success": True, "scopes": [_to_tool_json(s) for s in scopes]}, ensure_ascii=False, default=str, ) + if action == "get": + if not scope_id: + return json.dumps( + {"success": False, "error": "Scope_id is required for action='get'"}, + ensure_ascii=False, + default=str, + ) + scope = await proxy.call(lambda client: caido_api.scope_get(client, scope_id)) + return json.dumps( + {"success": True, "scope": _to_tool_json(scope)}, + ensure_ascii=False, + default=str, + ) + if action == "create": + if not scope_name: + return json.dumps( + {"success": False, "error": "Scope_name is required for action='create'"}, + ensure_ascii=False, + default=str, + ) + scope = await proxy.call( + lambda client: caido_api.scope_create( + client, name=scope_name, allowlist=allowlist, denylist=denylist + ), + idempotent=False, + ) + return json.dumps( + {"success": True, "scope": _to_tool_json(scope)}, + ensure_ascii=False, + default=str, + ) + if action == "update": + if not scope_id or not scope_name: + return json.dumps( + { + "success": False, + "error": "Scope_id and scope_name are required for action='update'", + }, + ensure_ascii=False, + default=str, + ) + scope = await proxy.call( + lambda client: caido_api.scope_update( + client, scope_id, name=scope_name, allowlist=allowlist, denylist=denylist + ), + idempotent=False, + ) + return json.dumps( + {"success": True, "scope": _to_tool_json(scope)}, + ensure_ascii=False, + default=str, + ) + if not scope_id: + return json.dumps( + {"success": False, "error": "Scope_id is required for action='delete'"}, + ensure_ascii=False, + default=str, + ) + await proxy.call(lambda client: caido_api.scope_delete(client, scope_id), idempotent=False) + return json.dumps( + { + "success": True, + "deleted": scope_id, + "message": f"Scope {scope_id} deleted", + }, + ensure_ascii=False, + default=str, + ) except Exception as exc: # noqa: BLE001 return _err("scope_rules", exc) diff --git a/tests/test_proxy_client.py b/tests/test_proxy_client.py index 6022f5a2..64436583 100644 --- a/tests/test_proxy_client.py +++ b/tests/test_proxy_client.py @@ -1,20 +1,20 @@ """Tests for the shared Caido client lifecycle and proxy error handling. Covers the concurrency/reconnect guarantees of ``caido_api.call_with_client`` -(the sandbox-imported path) and the host-side helpers in ``proxy.tools`` -(scan-wide lock + actionable HTTPQL errors). +(the sandbox-imported path) and ``caido_api.SharedCaidoClient`` (the host-side +holder), plus the actionable HTTPQL errors in ``proxy.tools``. """ from __future__ import annotations import asyncio -import contextlib import json from typing import TYPE_CHECKING, Any, cast import pytest from strix.tools.proxy import caido_api, tools +from strix.tools.proxy.caido_api import SharedCaidoClient if TYPE_CHECKING: @@ -24,6 +24,10 @@ if TYPE_CHECKING: class _FakeClient: def __init__(self, name: str) -> None: self.name = name + self.closed = False + + async def aclose(self) -> None: + self.closed = True @pytest.fixture(autouse=True) @@ -35,7 +39,7 @@ def _clear_cache() -> Iterator[None]: async def test_call_with_client_reuses_cached_client(monkeypatch: pytest.MonkeyPatch) -> None: cached = _FakeClient("cached") - caido_api._CLIENT_CACHE["default"] = cached + caido_api._CLIENT_CACHE["default"] = cast("Any", cached) async def _new() -> Any: raise AssertionError("_new_client must not run when a client is cached") @@ -87,12 +91,12 @@ async def test_failed_init_does_not_poison_cache(monkeypatch: pytest.MonkeyPatch assert "default" not in caido_api._CLIENT_CACHE -async def test_call_with_client_reconnects_on_dead_transport( +async def test_call_with_client_reconnects_and_closes_dead_transport( monkeypatch: pytest.MonkeyPatch, ) -> None: dead = _FakeClient("dead") fresh = _FakeClient("fresh") - caido_api._CLIENT_CACHE["default"] = dead + caido_api._CLIENT_CACHE["default"] = cast("Any", dead) new_calls = {"n": 0} @@ -114,13 +118,41 @@ async def test_call_with_client_reconnects_on_dead_transport( assert attempts == [dead, fresh] assert new_calls["n"] == 1 assert caido_api._CLIENT_CACHE["default"] is fresh + assert dead.closed is True # stale transport is not leaked + + +async def test_call_with_client_non_idempotent_rebuilds_but_reraises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + dead = _FakeClient("dead") + fresh = _FakeClient("fresh") + caido_api._CLIENT_CACHE["default"] = cast("Any", dead) + + async def _new() -> Any: + return fresh + + monkeypatch.setattr(caido_api, "_new_client", _new) + + calls = {"n": 0} + + async def fn(_client: Any) -> str: + calls["n"] += 1 + raise RuntimeError("Server disconnected") + + # A mutation must not be auto-retried (it may already have applied), but the + # dead client is still healed so later calls succeed. + with pytest.raises(RuntimeError, match="Server disconnected"): + await caido_api.call_with_client(fn, idempotent=False) + assert calls["n"] == 1 + assert caido_api._CLIENT_CACHE["default"] is fresh + assert dead.closed is True async def test_call_with_client_does_not_retry_application_errors( monkeypatch: pytest.MonkeyPatch, ) -> None: cached = _FakeClient("cached") - caido_api._CLIENT_CACHE["default"] = cached + caido_api._CLIENT_CACHE["default"] = cast("Any", cached) async def _new() -> Any: raise AssertionError("deterministic errors must not trigger a reconnect") @@ -142,7 +174,7 @@ async def test_call_with_client_does_not_retry_application_errors( async def test_call_with_client_serializes_concurrent_calls( monkeypatch: pytest.MonkeyPatch, ) -> None: - caido_api._CLIENT_CACHE["default"] = _FakeClient("shared") + caido_api._CLIENT_CACHE["default"] = cast("Any", _FakeClient("shared")) async def _new() -> Any: raise AssertionError("no reconnect expected") @@ -162,6 +194,87 @@ async def test_call_with_client_serializes_concurrent_calls( assert state["max"] == 1 +async def test_shared_client_reconnects_and_closes_dead_transport() -> None: + dead = _FakeClient("dead") + fresh = _FakeClient("fresh") + + async def _reconnect() -> Any: + return fresh + + holder = SharedCaidoClient(cast("Any", dead), _reconnect) + + attempts: list[Any] = [] + + async def fn(client: Any) -> str: + attempts.append(client) + if len(attempts) == 1: + raise RuntimeError("Connector is closed") + return "ok" + + assert await holder.call(fn) == "ok" + assert attempts == [dead, fresh] + assert dead.closed is True + + +async def test_shared_client_non_idempotent_rebuilds_but_reraises() -> None: + dead = _FakeClient("dead") + fresh = _FakeClient("fresh") + + async def _reconnect() -> Any: + return fresh + + holder = SharedCaidoClient(cast("Any", dead), _reconnect) + + calls = {"n": 0} + + async def fn(_client: Any) -> str: + calls["n"] += 1 + raise RuntimeError("Server disconnected") + + with pytest.raises(RuntimeError, match="Server disconnected"): + await holder.call(fn, idempotent=False) + assert calls["n"] == 1 + assert dead.closed is True + # The healthy client remains for the next call. + assert await holder.call(lambda _c: _ok()) == "ok" + + +async def _ok() -> str: + return "ok" + + +async def test_shared_client_serializes_concurrent_calls() -> None: + async def _reconnect() -> Any: + raise AssertionError("no reconnect expected") + + holder = SharedCaidoClient(cast("Any", _FakeClient("shared")), _reconnect) + + state = {"active": 0, "max": 0} + + async def fn(_client: Any) -> str: + state["active"] += 1 + state["max"] = max(state["max"], state["active"]) + await asyncio.sleep(0.01) + state["active"] -= 1 + return "ok" + + await asyncio.gather(*(holder.call(fn) for _ in range(6))) + assert state["max"] == 1 + + +async def test_shared_client_passes_through_application_errors() -> None: + async def _reconnect() -> Any: + raise AssertionError("deterministic errors must not trigger a reconnect") + + holder = SharedCaidoClient(cast("Any", _FakeClient("c")), _reconnect) + + async def fn(_client: Any) -> str: + raise ValueError("Invalid HTTPQL filter") + + with pytest.raises(ValueError, match="Invalid HTTPQL"): + await holder.call(fn) + + def test_is_connection_error_matches_markers_and_causes() -> None: assert caido_api._is_connection_error(RuntimeError("Transport is already connected")) assert caido_api._is_connection_error(RuntimeError("Connector is closed")) @@ -178,17 +291,19 @@ class _Ctx: self.context = context -def test_ctx_lock_returns_lock_when_present() -> None: - lock = asyncio.Lock() - got = tools._ctx_lock(cast("Any", _Ctx({"caido_lock": lock}))) - assert got is lock +def test_ctx_proxy_returns_holder_when_present() -> None: + async def _reconnect() -> Any: + raise AssertionError("unused") + + holder = SharedCaidoClient(cast("Any", _FakeClient("c")), _reconnect) + got = tools._ctx_proxy(cast("Any", _Ctx({"caido_client": holder}))) + assert got is holder -def test_ctx_lock_falls_back_to_noop_without_lock() -> None: - got = tools._ctx_lock(cast("Any", _Ctx({}))) - assert isinstance(got, contextlib.nullcontext) - got_non_dict = tools._ctx_lock(cast("Any", _Ctx(None))) - assert isinstance(got_non_dict, contextlib.nullcontext) +def test_ctx_proxy_returns_none_without_holder() -> None: + assert tools._ctx_proxy(cast("Any", _Ctx({}))) is None + assert tools._ctx_proxy(cast("Any", _Ctx(None))) is None + assert tools._ctx_proxy(cast("Any", _Ctx({"caido_client": object()}))) is None def test_is_httpql_error_detection() -> None: