From ad0dc4f5e91444f58a5949aec1609e603cbfe3eb Mon Sep 17 00:00:00 2001 From: yoni Date: Fri, 28 Aug 2026 03:44:04 +0000 Subject: [PATCH] Fix MCP resilience follow-up defects --- strix/tools/mcp/client.py | 43 ++++++--- strix/tools/mcp/config.py | 5 +- strix/tools/mcp/failures.py | 2 +- strix/tools/mcp/session.py | 177 ++++++++++++++++++++--------------- tests/test_mcp_client.py | 83 ++++++++++------ tests/test_mcp_resilience.py | 66 +++++++++++-- 6 files changed, 243 insertions(+), 133 deletions(-) diff --git a/strix/tools/mcp/client.py b/strix/tools/mcp/client.py index db5c3d3f..7335ffd4 100644 --- a/strix/tools/mcp/client.py +++ b/strix/tools/mcp/client.py @@ -75,6 +75,13 @@ class ConnectedMcpServer(NamedTuple): notes: str | None = None +class BuiltMcpServer(NamedTuple): + """A constructed SDK server and its optional HTTP failure recorder.""" + + server: MCPServer + recorder: HttpStatusRecorder | None + + def _auth_headers(config: McpConnectionConfig) -> dict[str, str]: """Build the per-server request headers from the connection's auth.""" auth = config.auth @@ -112,9 +119,12 @@ class _QuietMCPServerStdio(MCPServerStdio): return _quiet_stdio_streams(self.params) -def _build_server(config: McpConnectionConfig) -> MCPServer: +def _build_server(config: McpConnectionConfig) -> BuiltMcpServer: """Construct (but do not connect) the SDK server for one connection. + The returned tuple carries the server and, for HTTP connections, a recorder + that retains sanitized response metadata for the owning session. + When ``allowed_tools`` is a list the static filter means the server will not even list tools outside it, so it is the authoritative gate on what ``describe_mcp`` and ``call_mcp`` can see. When it is ``None`` no filter is @@ -132,11 +142,14 @@ def _build_server(config: McpConnectionConfig) -> MCPServer: "args": config.args, "env": config.env, } - return _QuietMCPServerStdio( - params=stdio_params, - name=config.name, - tool_filter=tool_filter, - cache_tools_list=True, + return BuiltMcpServer( + _QuietMCPServerStdio( + params=stdio_params, + name=config.name, + tool_filter=tool_filter, + cache_tools_list=True, + ), + None, ) recorder = HttpStatusRecorder() @@ -157,16 +170,16 @@ def _build_server(config: McpConnectionConfig) -> MCPServer: "sse_read_timeout": config.sse_read_timeout_seconds, "httpx_client_factory": httpx_client_factory, } - server = MCPServerStreamableHttp( - params=http_params, - name=config.name, - tool_filter=tool_filter, - cache_tools_list=True, - client_session_timeout_seconds=config.session_timeout_seconds, + return BuiltMcpServer( + MCPServerStreamableHttp( + params=http_params, + name=config.name, + tool_filter=tool_filter, + cache_tools_list=True, + client_session_timeout_seconds=config.session_timeout_seconds, + ), + recorder, ) - # The recorder is intentionally private and contains only status metadata. - server._strix_http_status_recorder = recorder # type: ignore[attr-defined] - return server def _mcp_result_to_tool_output(server: MCPServer, result: Any) -> Any: diff --git a/strix/tools/mcp/config.py b/strix/tools/mcp/config.py index e822b598..795a58a3 100644 --- a/strix/tools/mcp/config.py +++ b/strix/tools/mcp/config.py @@ -12,6 +12,9 @@ from typing import Annotated, Literal from pydantic import BaseModel, ConfigDict, Field, model_validator +DEFAULT_MAX_CONCURRENT_CALLS = 4 + + class BearerAuth(BaseModel): """Header-token auth, sent as ``Authorization: Bearer ``.""" @@ -74,7 +77,7 @@ class McpConnectionConfig(BaseModel): session_timeout_seconds: float = Field(default=60.0, gt=0) """MCP operation timeout for SQL queries and cloud describe fan-outs.""" - max_concurrent_calls: int = Field(default=4, ge=1) + max_concurrent_calls: int = Field(default=DEFAULT_MAX_CONCURRENT_CALLS, ge=1) """Maximum concurrent calls for this connection name across sessions.""" @model_validator(mode="after") diff --git a/strix/tools/mcp/failures.py b/strix/tools/mcp/failures.py index a9a0c886..05953851 100644 --- a/strix/tools/mcp/failures.py +++ b/strix/tools/mcp/failures.py @@ -129,7 +129,7 @@ class HttpStatusRecorder: def __init__(self) -> None: self._failure: FailureInfo | None = None - def __call__(self, response: httpx.Response) -> None: + async def __call__(self, response: httpx.Response) -> None: if not 200 <= response.status_code < 300: request = response.request self._failure = _from_status( diff --git a/strix/tools/mcp/session.py b/strix/tools/mcp/session.py index 07c73f32..d0f4a8d7 100644 --- a/strix/tools/mcp/session.py +++ b/strix/tools/mcp/session.py @@ -47,8 +47,10 @@ import dataclasses import logging import secrets import time +import weakref from typing import TYPE_CHECKING, Any, cast +from strix.tools.mcp.config import DEFAULT_MAX_CONCURRENT_CALLS from strix.tools.mcp.failures import FailureInfo, HttpStatusRecorder, classify @@ -75,7 +77,9 @@ logger = logging.getLogger(__name__) _SHUTDOWN_TIMEOUT = 10.0 _MAX_ATTEMPTS = 3 _SETTLE_DELAY = 0.05 -_SEMAPHORES: dict[str, asyncio.Semaphore] = {} +_SEMAPHORES: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, asyncio.Semaphore]] = ( + weakref.WeakKeyDictionary() +) _JITTER = secrets.SystemRandom() @@ -86,6 +90,12 @@ def _retry_delay(attempt: int, retry_after: float | None) -> float: return base + _JITTER.uniform(0.0, base * 0.1) # type: ignore[no-any-return] +def _call_semaphore(name: str, limit: int) -> asyncio.Semaphore: + loop = asyncio.get_running_loop() + semaphores = _SEMAPHORES.setdefault(loop, {}) + return semaphores.setdefault(name, asyncio.Semaphore(limit)) + + class McpConnectionUnavailableError(RuntimeError): """A dead MCP connection could not be reached and did not come back. @@ -148,14 +158,7 @@ class SupervisedMcpSession: self._quarantine_count = 0 self._last_failure = FailureInfo("unknown", reason="connection unavailable") self._reconnect_lock = asyncio.Lock() - self._call_semaphore = _SEMAPHORES.setdefault( - self._name, asyncio.Semaphore(config.max_concurrent_calls) - ) - # Guards the idle-death self-heal against a flapping server: set after an - # idle reconnect, cleared once a real call runs. If the session dies idle - # again before serving anything, we give up instead of reconnecting in a - # tight loop. - self._healed_without_progress = False + self._call_semaphore: asyncio.Semaphore | None = None @classmethod def adopt( @@ -169,7 +172,7 @@ class SupervisedMcpSession: Calls run inline against ``server`` on the caller's task, matching the old direct-dispatch behavior. Reconnect is available only when ``config`` is - given; otherwise a failed call marks the connection dead. + given; otherwise a failed call can be quarantined but cannot be revived. """ self = cls.__new__(cls) self._name = name @@ -188,10 +191,7 @@ class SupervisedMcpSession: self._quarantine_count = 0 self._last_failure = FailureInfo("unknown", reason="connection unavailable") self._reconnect_lock = asyncio.Lock() - self._call_semaphore = _SEMAPHORES.setdefault( - name, asyncio.Semaphore(config.max_concurrent_calls if config else 4) - ) - self._healed_without_progress = False + self._call_semaphore = None return self # -- read-only accessors -------------------------------------------------- @@ -411,13 +411,27 @@ class SupervisedMcpSession: await self._safe_cleanup() self._fail_pending() return - except BaseException as exc: # noqa: BLE001 - classify cancellation groups + except BaseExceptionGroup as exc: failure = classify(exc) logger.warning( "Skipping MCP connection %r kind=%s status=%s attempt=1 delay=0", self._name, failure.kind, failure.status, + exc_info=True, + ) + self._report_ready(value=False) + await self._safe_cleanup() + self._fail_pending() + return + except Exception as exc: # noqa: BLE001 - classify ordinary connect failures + failure = classify(exc) + logger.warning( + "Skipping MCP connection %r kind=%s status=%s attempt=1 delay=0", + self._name, + failure.kind, + failure.status, + exc_info=True, ) self._report_ready(value=False) await self._safe_cleanup() @@ -440,55 +454,23 @@ class SupervisedMcpSession: # A cancellation while idle is the transport's task group cancelling # this supervising task because a background session task failed. # Contained here. If we are closing, this is an ordinary shutdown, - # so let it propagate. Otherwise try to self-heal once: reconnect a - # fresh session and keep serving. The flag stops a flapping server - # (one that dies again before serving any call) from reconnecting in - # a tight loop; there we give up and mark the connection dead. Later - # calls then short-circuit to the dead output without this task. + # so let it propagate. Otherwise quarantine the failed session and + # keep serving requests so a later call can revive it. if self._closing: raise failure = self._recorder.take() if self._recorder is not None else None failure = failure or FailureInfo("transport", reason="session cancelled") self._last_failure = failure - if not self._healed_without_progress: - logger.warning( - "MCP connection %r session died while idle kind=%s status=%s " - "attempt=1 delay=0; reconnecting", - self._name, - failure.kind, - failure.status, - ) - self._healed_without_progress = True - reconnected, reconnect_failure = await self._reconnect() - if reconnected: - logger.info( - "MCP connection %r revived kind=%s status=%s attempt=1 delay=%.2f", - self._name, - failure.kind, - failure.status, - _SETTLE_DELAY, - ) - continue - if reconnect_failure is not None: - failure = reconnect_failure - else: - logger.warning( - "MCP connection %r died again before serving a call; " - "marking it permanently unavailable kind=%s status=%s " - "attempt=1 delay=0", - self._name, - failure.kind, - failure.status, - ) - self._mark_dead(failure) - await self._safe_cleanup() - return + if failure.kind == "auth": + self._mark_dead(failure, attempt=1) + return + await self._quarantine(failure, attempt=1) + if self._dead: + return + continue if request is None: # shutdown sentinel return outcome = await self._execute(request.job) - # A served call is real progress: clear the idle-heal guard so a future - # idle death is again allowed one reconnect. - self._healed_without_progress = False if not request.future.done(): request.future.set_result(outcome) self._pending.discard(request.future) @@ -497,7 +479,7 @@ class SupervisedMcpSession: # -- run one job with bounded classified retries -------------------------- - async def _execute(self, job: Job) -> _Outcome: # noqa: PLR0911, PLR0912 + async def _execute(self, job: Job) -> _Outcome: # noqa: PLR0911, PLR0912, PLR0915 """Run one job with classified retries and temporary quarantine.""" if self._dead: return _Outcome(dead=True) @@ -511,11 +493,20 @@ class SupervisedMcpSession: self._name, self._last_failure.kind, self._last_failure.status, - remaining, + 0.0, ) failure: FailureInfo | None = None for attempt in range(1, _MAX_ATTEMPTS + 1): + if self._call_semaphore is None: + self._call_semaphore = _call_semaphore( + self._name, + ( + self._config.max_concurrent_calls + if self._config is not None + else DEFAULT_MAX_CONCURRENT_CALLS + ), + ) if self._server is None: reconnected, reconnect_failure = await self._reconnect() if not reconnected: @@ -527,23 +518,29 @@ class SupervisedMcpSession: self._mark_dead(failure, attempt=attempt) return _Outcome(dead=True) if attempt == _MAX_ATTEMPTS: - self._quarantine(failure, attempt=attempt) + await self._quarantine(failure, attempt=attempt) return _Outcome(dead=True) delay = _retry_delay(attempt, failure.retry_after) self._log_retry(failure, attempt, delay) await asyncio.sleep(delay) continue assert self._server is not None + call_semaphore = self._call_semaphore + assert call_semaphore is not None try: - async with self._call_semaphore: - return _Outcome(value=await self._run_job_call(job)) + async with call_semaphore: + return _Outcome(value=await job(self._server)) except asyncio.CancelledError: if not self._supervised or self._closing: raise failure = ( self._recorder.take() if self._recorder is not None else None ) or FailureInfo("transport", reason="session cancelled") - except BaseException as exc: # noqa: BLE001 - classify SDK groups + except BaseExceptionGroup as exc: + failure = classify(exc) + if failure.kind == "unknown" and self._recorder is not None: + failure = self._recorder.take() or failure + except Exception as exc: # noqa: BLE001 - classify ordinary call failures failure = classify(exc) if failure.kind == "unknown" and self._recorder is not None: failure = self._recorder.take() or failure @@ -553,20 +550,26 @@ class SupervisedMcpSession: self._mark_dead(failure, attempt=attempt) return _Outcome(dead=True) if attempt == _MAX_ATTEMPTS: - self._quarantine(failure, attempt=attempt) + await self._quarantine(failure, attempt=attempt) return _Outcome(dead=True) delay = _retry_delay(attempt, failure.retry_after) self._log_retry(failure, attempt, delay) await asyncio.sleep(delay) reconnected, reconnect_failure = await self._reconnect() - if not reconnected and reconnect_failure is not None: - self._mark_dead(reconnect_failure, attempt=attempt) - return _Outcome(dead=True) + if not reconnected: + failure = reconnect_failure or FailureInfo("transport", reason="reconnect failed") + self._last_failure = failure + if failure.kind == "auth": + self._mark_dead(failure, attempt=attempt) + return _Outcome(dead=True) + if attempt == _MAX_ATTEMPTS: + await self._quarantine(failure, attempt=attempt) + return _Outcome(dead=True) + delay = _retry_delay(attempt, failure.retry_after) + self._log_retry(failure, attempt, delay) + await asyncio.sleep(delay) return _Outcome(dead=True) - async def _run_job_call(self, job: Job) -> Any: - return await job(self._server) # type: ignore[arg-type] - def _log_retry(self, failure: FailureInfo, attempt: int, delay: float) -> None: logger.warning( "MCP connection %r retryable failure kind=%s status=%s attempt=%d delay=%.2f", @@ -577,7 +580,8 @@ class SupervisedMcpSession: delay, ) - def _quarantine(self, failure: FailureInfo, *, attempt: int) -> None: + async def _quarantine(self, failure: FailureInfo, *, attempt: int) -> None: + await self._safe_cleanup() self._quarantine_count += 1 if self._quarantine_count >= 3: self._mark_dead(failure, attempt=attempt) @@ -606,7 +610,13 @@ class SupervisedMcpSession: raise self._server = None return False, FailureInfo("transport", reason="reconnect cancelled") - except BaseException as exc: # noqa: BLE001 - classify SDK groups + except BaseExceptionGroup as exc: + self._server = None + failure = classify(exc) + if failure.kind == "unknown" and self._recorder is not None: + failure = self._recorder.take() or failure + return False, failure + except Exception as exc: # noqa: BLE001 - classify ordinary reconnect failures self._server = None failure = classify(exc) if failure.kind == "unknown" and self._recorder is not None: @@ -618,7 +628,7 @@ class SupervisedMcpSession: await asyncio.sleep(_SETTLE_DELAY) except asyncio.CancelledError: self._server = None - with contextlib.suppress(BaseException): + with contextlib.suppress(Exception): await server.cleanup() # type: ignore[no-untyped-call] if self._closing: raise @@ -632,15 +642,30 @@ class SupervisedMcpSession: same task before the error propagates, so a failed connect never orphans an MCP subprocess or half-open HTTP session. """ - from strix.tools.mcp.client import _build_server + from strix.tools.mcp.client import BuiltMcpServer, _build_server if self._config is None: raise RuntimeError(f"MCP connection {self._name!r} has no config to connect") - server = _build_server(self._config) - self._recorder = getattr(server, "_strix_http_status_recorder", None) + built = _build_server(self._config) + if isinstance(built, BuiltMcpServer): + server = built.server + self._recorder = built.recorder + else: + # Test and private consumers may still provide a connected server + # directly when replacing the private builder. + server = cast("MCPServer", built) # type: ignore[unreachable] + self._recorder = None try: await server.connect() # type: ignore[no-untyped-call] - except BaseException: + except asyncio.CancelledError: + with contextlib.suppress(Exception): + await server.cleanup() # type: ignore[no-untyped-call] + raise + except BaseExceptionGroup: + with contextlib.suppress(Exception): + await server.cleanup() # type: ignore[no-untyped-call] + raise + except Exception: with contextlib.suppress(Exception): await server.cleanup() # type: ignore[no-untyped-call] raise diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index c75efa4d..0ce12780 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -11,6 +11,8 @@ import asyncio import contextlib import json import re +import time +from functools import partial from typing import TYPE_CHECKING, Any import pytest @@ -160,7 +162,7 @@ def _config(name: str, allowed_tools: list[str] | None) -> McpConnectionConfig: return McpConnectionConfig( name=name, url="https://mcp.example.com", - auth=BearerAuth(token="run-token"), + auth=BearerAuth(token="run-token"), # nosec B106 allowed_tools=allowed_tools, ) @@ -204,7 +206,7 @@ def test_bearer_config_parses_from_dict() -> None: ) assert isinstance(config.auth, BearerAuth) - assert config.auth.token == "abc" + assert config.auth.token == "abc" # nosec B105 assert config.allowed_tools == ["list_files"] @@ -356,7 +358,7 @@ def test_build_server_stdio_branch() -> None: env={"TOKEN": "x"}, ) - server = mcp_client._build_server(config) + server = mcp_client._build_server(config).server assert isinstance(server, MCPServerStdio) assert server.name == "local_fs" @@ -366,7 +368,7 @@ def test_build_server_stdio_branch() -> None: def test_build_server_http_branch() -> None: - server = mcp_client._build_server(_config("files_main", ["list_files"])) + server = mcp_client._build_server(_config("files_main", ["list_files"])).server assert isinstance(server, MCPServerStreamableHttp) assert server.name == "files_main" @@ -1213,7 +1215,7 @@ def _secret_config(name: str) -> McpConnectionConfig: return McpConnectionConfig( name=name, url="https://mcp.example.com", - auth=BearerAuth(token="super-secret-bearer-token-42"), + auth=BearerAuth(token="super-secret-bearer-token-42"), # nosec B106 allowed_tools=["read_file"], ) @@ -1257,8 +1259,8 @@ async def test_call_mcp_reconnects_and_retries_after_a_session_death( async def test_call_mcp_marks_connection_dead_when_reconnect_keeps_failing( monkeypatch: pytest.MonkeyPatch, ) -> None: - # The session dies and the reconnect attempt also fails: the connection is - # marked dead and the call returns the standard failed-tool output. + # Reconnect failures are retried and then quarantine the connection rather + # than permanently retiring it on the first failed reconnect. first = _DyingHttpServer("fs", [_mcp_tool("read_file")], death=ConnectionError("403")) built = {"n": 0} @@ -1282,9 +1284,12 @@ async def test_call_mcp_marks_connection_dead_when_reconnect_keeps_failing( assert isinstance(out, dict) assert out["success"] is False assert "unavailable" in out["content"] - assert session.is_dead is True + assert session.is_dead is False + assert session.is_unavailable is True + assert session.server is None + assert session._task is not None and not session._task.done() - # A later call short-circuits to the same failed output without a new attempt. + # A later call during cooldown short-circuits to the same failed output. again = await call_mcp.on_invoke_tool( _ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"}) ) @@ -1309,14 +1314,17 @@ async def _pump_until(predicate: Callable[[], bool], *, limit: int = 100) -> Non raise AssertionError("condition not reached") +def _quarantine_reached(session: SupervisedMcpSession, count: int) -> bool: + return session.is_dead or session._quarantine_count >= count + + @pytest.mark.asyncio async def test_idle_session_death_self_heals_on_reconnect( monkeypatch: pytest.MonkeyPatch, ) -> None: # A session that dies while idle (its supervising task cancelled between calls, - # modeling the transport scope dying with no call in flight) reconnects once on - # its own and keeps serving, rather than staying dead until a later call would - # have triggered a reconnect. + # modeling the transport scope dying with no call in flight) is quarantined and + # keeps serving, rather than ending its supervising task. first = FakeMCPServer("fs", [_mcp_tool("read_file")]) second = FakeMCPServer("fs", [_mcp_tool("read_file")]) built = iter([first, second]) @@ -1328,15 +1336,16 @@ async def test_idle_session_death_self_heals_on_reconnect( assert session._task is not None session._task.cancel() # idle transport death: no call in flight - await _pump_until(lambda: session.server is second) + await _pump_until(lambda: session.is_unavailable) assert session.is_dead is False + assert session.server is None - # The reconnected session serves calls normally. + # Once the cooldown expires, the next call reconnects onto a fresh session. + session._unavailable_until = time.monotonic() - 1 out = await call_mcp.on_invoke_tool( _ctx(registry), json.dumps({"connection": "fs", "tool": "read_file"}) ) assert out == {"type": "text", "text": "routed:read_file"} - await session.aclose() @@ -1344,9 +1353,8 @@ async def test_idle_session_death_self_heals_on_reconnect( async def test_flapping_idle_session_is_marked_dead_without_looping( monkeypatch: pytest.MonkeyPatch, ) -> None: - # If a session reconnects after an idle death but dies again before serving any - # call, the supervisor stops reconnecting and marks the connection dead, so a - # server that instantly drops on connect cannot spin in a reconnect loop. + # Repeated idle deaths consume quarantine slots; the supervisor stays alive + # until the configured permanent-death threshold is reached. first = FakeMCPServer("fs", [_mcp_tool("read_file")]) second = FakeMCPServer("fs", [_mcp_tool("read_file")]) built = iter([first, second]) @@ -1357,13 +1365,12 @@ async def test_flapping_idle_session_is_marked_dead_without_looping( registry.add(name="fs", session=session, tool_count=1) assert session._task is not None - # First idle death heals onto the second server (only two builds ever happen). - session._task.cancel() - await _pump_until(lambda: session.server is second) - assert session.is_dead is False - - # Second idle death before any call is served: give up rather than reconnect. - session._task.cancel() + # Three idle deaths exhaust the quarantine budget. + for count in range(1, 4): + session._task.cancel() + await _pump_until(partial(_quarantine_reached, session, count)) + if session.is_dead: + break await _pump_until(lambda: session._task is not None and session._task.done()) assert session.is_dead is True @@ -1566,23 +1573,37 @@ class _RaisingMCPServer(FakeMCPServer): @pytest.mark.asyncio -async def test_session_on_dead_fires_once_on_the_death_transition() -> None: - # An adopted session with no config cannot reconnect, so the first failed - # call marks it dead; the on-dead callback fires exactly once, on the edge. +async def test_session_on_dead_fires_once_on_the_death_transition( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # An adopted session with no config cannot reconnect, so repeated transient + # exhaustion eventually marks it dead; the callback fires once on that edge. server = _RaisingMCPServer("db", [_mcp_tool("read")]) session = SupervisedMcpSession.adopt(server, name="db") fires: list[int] = [] session.set_on_dead(lambda: fires.append(1)) + clock = [100.0] + monkeypatch.setattr("strix.tools.mcp.session.time.monotonic", lambda: clock[0]) + monkeypatch.setattr(mcp_session_mod, "_retry_delay", lambda _attempt, _retry_after: 0) + + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(asyncio, "sleep", no_sleep) out = await session.dispatch("read", {}, label="db_read") - assert session.is_dead is True + assert session.is_dead is False assert isinstance(out, dict) and out.get("success") is False - assert fires == [1] + assert fires == [] - # A later call to the already-dead session must not fire the callback again. + clock[0] += 31 + await session.dispatch("read", {}, label="db_read") + assert session.is_dead is False + clock[0] += 61 await session.dispatch("read", {}, label="db_read") assert fires == [1] + assert session.is_dead is True def test_registry_statuses_report_the_live_dead_flag_and_provider() -> None: diff --git a/tests/test_mcp_resilience.py b/tests/test_mcp_resilience.py index bd91bf83..42e29dfb 100644 --- a/tests/test_mcp_resilience.py +++ b/tests/test_mcp_resilience.py @@ -5,7 +5,7 @@ from __future__ import annotations import asyncio import importlib from datetime import UTC, datetime, timedelta -from typing import Any +from typing import Any, cast import httpx import pytest @@ -67,15 +67,16 @@ def test_classifies_nested_exception_groups_by_specificity() -> None: assert info.retryable is False -def test_retry_after_parses_seconds_and_http_date() -> None: +@pytest.mark.asyncio +async def test_retry_after_parses_seconds_and_http_date() -> None: seconds = HttpStatusRecorder() - seconds(_http_error(429, retry_after="12").response) + await seconds(_http_error(429, retry_after="12").response) assert seconds.take() is not None assert seconds.take() is None date = (datetime.now(UTC) + timedelta(seconds=20)).strftime("%a, %d %b %Y %H:%M:%S GMT") recorder = HttpStatusRecorder() - recorder(_http_error(429, retry_after=date).response) + await recorder(_http_error(429, retry_after=date).response) info = recorder.take() assert info is not None retry_after = info.retry_after @@ -83,10 +84,11 @@ def test_retry_after_parses_seconds_and_http_date() -> None: assert 0 <= retry_after <= 20 -def test_recorder_only_keeps_non_sensitive_request_metadata() -> None: +@pytest.mark.asyncio +async def test_recorder_only_keeps_non_sensitive_request_metadata() -> None: recorder = HttpStatusRecorder() response = _http_error(500, retry_after="3").response - recorder(response) + await recorder(response) info = recorder.take() assert info == FailureInfo( "server", @@ -174,6 +176,7 @@ async def test_server_exhaustion_quarantines_then_revives( assert result["success"] is False assert session.is_dead is False assert session.is_unavailable is True + assert session.server is None clock[0] += 31 result = await session.dispatch("read", {}, label="quarantine_read") assert result == {"type": "text", "text": "routed:read"} @@ -199,16 +202,20 @@ async def test_cancelled_call_uses_recorded_status( ) -> None: recorder = HttpStatusRecorder() first = _sequence_server("cancelled", asyncio.CancelledError()) - first._strix_http_status_recorder = recorder second = _sequence_server("cancelled") - builds = iter([first, second]) + builds = iter( + [ + mcp_client.BuiltMcpServer(first, recorder), + mcp_client.BuiltMcpServer(second, None), + ] + ) monkeypatch.setattr(mcp_client, "_build_server", lambda _config: next(builds)) monkeypatch.setattr(mcp_session, "_retry_delay", _zero_delay) monkeypatch.setattr(asyncio, "sleep", _no_sleep) session = mcp_session.SupervisedMcpSession(_config("cancelled")) assert await session.start() - recorder(_http_error(503).response) + await recorder(_http_error(503).response) result = await session.dispatch("read", {}, label="cancelled_read") assert result == {"type": "text", "text": "routed:read"} await session.aclose() @@ -241,6 +248,29 @@ async def test_build_server_passes_explicit_http_values( await client.aclose() +@pytest.mark.asyncio +async def test_http_factory_awaits_response_recorder() -> None: + built = mcp_client._build_server(_config("hook")) + assert built.recorder is not None + factory = cast("Any", built.server).params["httpx_client_factory"] + client = factory(headers={}, timeout=httpx.Timeout(1), auth=None) + + def response(request: httpx.Request) -> httpx.Response: + return httpx.Response(429, headers={"Retry-After": "7"}, request=request) + + client._transport = httpx.MockTransport(response) + result = await client.get("https://provider.example/mcp?token=secret-query") + assert result.status_code == 429 + info = built.recorder.take() + assert info is not None + assert info.kind == "rate_limit" + assert info.status == 429 + assert info.retry_after == 7 + assert info.request_method == "GET" + assert info.request_path == "/mcp" + await client.aclose() + + @pytest.mark.asyncio async def test_same_name_sessions_share_concurrency_cap() -> None: active = 0 @@ -277,6 +307,24 @@ async def test_same_name_sessions_share_concurrency_cap() -> None: await right.aclose() +def test_same_name_semaphore_works_across_event_loops() -> None: + async def run_once() -> None: + server = FakeMCPServer("loop-cap", [_mcp_tool("read")]) + session = mcp_session.SupervisedMcpSession.adopt( + server, + name="loop-cap", + config=_config("loop-cap", max_concurrent_calls=1), + ) + assert await session.dispatch("read", {}, label="loop_cap_read") == { + "type": "text", + "text": "routed:read", + } + await session.aclose() + + asyncio.run(run_once()) + asyncio.run(run_once()) + + @pytest.mark.asyncio async def test_resilience_logs_do_not_include_request_secrets( monkeypatch: pytest.MonkeyPatch,