diff --git a/strix/tools/mcp/session.py b/strix/tools/mcp/session.py index 2d16566a7..8b123b523 100644 --- a/strix/tools/mcp/session.py +++ b/strix/tools/mcp/session.py @@ -28,9 +28,13 @@ lifetime, and ``cleanup()``. Three consequences: "connection unavailable" value instead of a cancellation propagating into the agent loop. -When a call fails the supervisor classifies it, retries boundedly, and temporarily -quarantines transient failures before lazily reviving the session. Authentication -failures and repeated transient exhaustion permanently retire a connection. +Failure handling follows connection-pool discipline: discard on error, rebuild on +next use. The instant a call errors, the supervisor disposes that session on its own +task and never awaits on it again; the next call lazily builds a fresh one. A +classified transient failure is retried once on the rebuilt session and then, if it +keeps failing, temporarily quarantines the connection so the circuit breaker plus the +next call can recover it. Authentication failures and repeated transient exhaustion +permanently retire a connection. Security: the connection's :class:`~strix.tools.mcp.config.McpConnectionConfig` holds a live bearer credential and is kept here in memory only, on the same @@ -472,8 +476,33 @@ class SupervisedMcpSession: # -- run one job with bounded classified retries -------------------------- - async def _execute(self, job: Job) -> _Outcome: # noqa: PLR0911, PLR0912 - """Run one job with classified retries and temporary quarantine.""" + async def _execute(self, job: Job) -> _Outcome: # noqa: PLR0912 + """Run one job on a healthy session, disposing it the instant it errors. + + Discard-on-error, rebuild-on-next-use is the whole discipline here, and it + rests on one invariant: **a session object is only ever awaited while + healthy.** The moment a call fails, the very next thing this method does, + before any other ``await`` including the backoff sleep inside + :meth:`_handle_failure`, is dispose that session on this task + (:meth:`_safe_cleanup` runs the transport teardown and clears ``_server``). + + Why the ordering is the crux, not a nicety: when a provider returns a non-2xx + status mid-call, the streamable-HTTP transport's task group cancels its scope, + which cancels this supervising task; the failure surfaces as a + ``CancelledError`` and the scope keeps firing (re-raising on every subsequent + ``await``) until the session is torn down. Disposing closes the transport's + AsyncExitStack, which exits that firing scope. If instead we slept for backoff + first, the sleep would re-raise the firing ``CancelledError``, escape this + method, and kill the supervising task, leaving the slot wedged with + ``is_dead`` False forever. Disposing first is what turns a failure into a + returned value and keeps the task alive to rebuild on the next attempt. + + The rebuild itself happens lazily at the top of the loop: once a failure has + set ``_server`` to None, the next iteration builds a fresh session (guarded by + :meth:`_reconnect`) and retries the call on it. A genuine shutdown + (``_closing``) and a real external cancellation still propagate; only the + transport's teardown cancellation is contained. + """ if self._dead: return _Outcome(dead=True) if self._unavailable_until is not None: @@ -499,6 +528,10 @@ class SupervisedMcpSession: ) failure: FailureInfo | None = None for attempt in range(1, _MAX_ATTEMPTS + 1): + # Lazy, atomic rebuild: a prior failure disposed the session, so build a + # fresh one here. The rebuild lock lets concurrent callers (adopted + # sessions dispatched from several agent tasks) share one rebuild rather + # than each building their own. if self._server is None: reconnected, reconnect_failure = await self._reconnect() if not reconnected: @@ -514,27 +547,38 @@ class SupervisedMcpSession: assert call_semaphore is not None try: async with call_semaphore: - return _Outcome(value=await job(self._server)) + result = await job(self._server) + # A success clears the quarantine strikes. A connection that + # recovered and served a call is healthy again, so transient + # failure bursts separated by successful revivals must not + # accumulate toward permanent retirement; only sustained failure + # with no success in between should retire the connection. + self._quarantine_count = 0 + return _Outcome(value=result) 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") + # Dispose BEFORE any other await. The transport's cancel scope may be + # firing right now; _safe_cleanup exits it so the backoff sleep below + # cannot re-raise the cancellation and kill this task. See the + # method docstring for why this ordering is load-bearing. + await self._safe_cleanup() except _CLASSIFIABLE as exc: failure = classify(exc) if failure.kind == "unknown" and self._recorder is not None: failure = self._recorder.take() or failure + # Dispose BEFORE any other await, same reason as the branch above: + # never await on a session that has already errored. + await self._safe_cleanup() + # Session is disposed and _server is None; _handle_failure may sleep for + # backoff safely, and the next loop iteration rebuilds and retries. outcome = await self._handle_failure(failure, attempt) if outcome is not None: return outcome - reconnected, reconnect_failure = await self._reconnect() - if not reconnected: - failure = reconnect_failure or FailureInfo("transport", reason="reconnect failed") - outcome = await self._handle_failure(failure, attempt) - if outcome is not None: - return outcome return _Outcome(dead=True) async def _handle_failure(self, failure: FailureInfo, attempt: int) -> _Outcome | None: @@ -578,9 +622,19 @@ class SupervisedMcpSession: ) async def _reconnect(self) -> tuple[bool, FailureInfo | None]: - """Rebuild and reconnect, reusing the stored config and settling briefly.""" + """Build a fresh session under the rebuild lock, so concurrent callers share one. + + Called only when ``_server`` is None (a prior failure already disposed the old + session). The lock serializes rebuilds; a caller that finds the session already + rebuilt by whoever held the lock first reuses it instead of building a second + one. There is deliberately no cleanup of an existing ``_server`` here: this + method never runs against a live session, because the failure path disposes + before it ever reaches a rebuild. + """ async with self._reconnect_lock: - await self._safe_cleanup() + if self._server is not None: + # Another caller rebuilt while we waited for the lock; share it. + return True, None if self._config is None: return False, FailureInfo("transport", reason="no reconnect config") try: @@ -601,9 +655,9 @@ class SupervisedMcpSession: try: await asyncio.sleep(_SETTLE_DELAY) except asyncio.CancelledError: - self._server = None - with contextlib.suppress(Exception): - await server.cleanup() # type: ignore[no-untyped-call] + # Dispose the just-built session before returning; _safe_cleanup + # re-raises when we are shutting down and absorbs otherwise. + await self._safe_cleanup() if self._closing: raise return False, FailureInfo("transport", reason="reconnect cancelled") @@ -638,12 +692,48 @@ class SupervisedMcpSession: # -- helpers -------------------------------------------------------------- async def _safe_cleanup(self) -> None: + """Dispose the live session on this task, completing teardown even under a + firing cancel scope. + + Why this is delicate: the streamable-HTTP transport holds an anyio task group + whose cancel scope was entered on this supervising task. When a background POST + got a non-2xx status the SDK cancelled that scope, and until the scope is + exited every ``await`` on this task re-raises ``CancelledError``. + ``server.cleanup()`` closes the AsyncExitStack that runs the task group's + ``__aexit__``, and that ``__aexit__`` is exactly what exits the scope and stops + the firing; it also absorbs the scope's own cancellation internally, so the + common case returns cleanly. A stray ``CancelledError`` can still surface, + though, and ``contextlib.suppress(Exception)`` would let it through because + ``CancelledError`` is a ``BaseException``, not an ``Exception``. + + So we catch ``CancelledError`` explicitly. During a real shutdown + (``_closing``) that cancellation is the run going down and must propagate, so + we re-raise it. Otherwise we absorb it and retry the close a bounded number of + times: if a cleanup was interrupted before the exit stack finished unwinding, + closing again continues from where it left off (the stack pops one callback at + a time), so the scope still ends up exited and this task stays runnable for the + next rebuild. + """ server = self._server self._server = None if server is None: return - with contextlib.suppress(Exception): - await server.cleanup() # type: ignore[no-untyped-call] + for _ in range(_MAX_ATTEMPTS): + try: + # suppress(Exception) absorbs an ordinary cleanup error but lets a + # CancelledError through, because it is a BaseException; the outer + # handler below is what decides whether to propagate or retry it. + with contextlib.suppress(Exception): + await server.cleanup() # type: ignore[no-untyped-call] + except asyncio.CancelledError: + if self._closing: + raise + # Firing scope hit the cleanup await before the stack finished + # unwinding; swallow this cancellation and close again to complete + # the teardown. A fully-closed stack makes the retry a clean no-op. + continue + else: + return def _report_ready(self, value: bool) -> None: if self._ready is not None and not self._ready.done(): diff --git a/tests/test_mcp_resilience.py b/tests/test_mcp_resilience.py index e1b35104d..cd79caa23 100644 --- a/tests/test_mcp_resilience.py +++ b/tests/test_mcp_resilience.py @@ -201,6 +201,44 @@ async def test_server_exhaustion_quarantines_then_revives( await session.aclose() +@pytest.mark.asyncio +async def test_success_resets_quarantine_strikes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # A quarantine strike must be cleared by a successful revival, so transient + # failure bursts separated by successes do not accumulate toward permanent + # retirement. Without the reset, three such bursts would mark the connection + # dead even though it recovered between each one. + monkeypatch.setattr(mcp_session, "_retry_delay", _zero_delay) + monkeypatch.setattr(asyncio, "sleep", _no_sleep) + clock = [100.0] + monkeypatch.setattr("strix.tools.mcp.session.time.monotonic", lambda: clock[0]) + builds = iter( + [ + _sequence_server("strikes", _http_error(500)), + _sequence_server("strikes", _http_error(500)), + _sequence_server("strikes", _http_error(500)), + _sequence_server("strikes"), + ] + ) + monkeypatch.setattr(mcp_client, "_build_server", lambda _config: _built_server(next(builds))) + session = mcp_session.SupervisedMcpSession(_config("strikes")) + assert await session.start() + + # First burst exhausts three attempts and quarantines: one strike. + result = await session.dispatch("read", {}, label="strikes_read") + assert result["success"] is False + assert session._quarantine_count == 1 + + # The revive succeeds, which must clear the strike back to zero. + clock[0] += 31 + result = await session.dispatch("read", {}, label="strikes_read") + assert result == {"type": "text", "text": "routed:read"} + assert session._quarantine_count == 0 + assert session.is_dead is False + await session.aclose() + + @pytest.mark.asyncio async def test_auth_failure_dies_without_retry(monkeypatch: pytest.MonkeyPatch) -> None: builds = [_sequence_server("auth", _http_error(401))]