diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index a2026c5d..e158c053 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -2487,6 +2487,133 @@ class TestCircuitBreaker: # Circuit should NOT have recorded a failure assert mgr._consecutive_failures.get("test", 0) == 0 + def test_closed_resource_error_evicts_session_and_trips_circuit(self): + """Regression: the MCP SDK's streamable-http transport raises + ``anyio.ClosedResourceError`` (NOT BrokenPipeError) when its write + stream is dead. That must evict the session AND trip the breaker — + otherwise the corpse session is re-used on every call forever and + only a full process restart recovers it.""" + import anyio + + mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}}) + mock_session = MagicMock() + mock_session.call_tool = MagicMock(return_value="sentinel") + _seed_static_state(mgr, "test", session=mock_session) + mgr._loop = MagicMock() + mgr._tool_map["mcp__test__ping"] = ("test", "ping") + mock_future = MagicMock() + mock_future.result.side_effect = anyio.ClosedResourceError() + with ( + patch("asyncio.run_coroutine_threadsafe", new=_dispatch_stub(mock_future)), + pytest.raises(anyio.ClosedResourceError), + ): + mgr.call_tool_sync("mcp__test__ping", {}, timeout=5) + assert mgr._static_servers["test"].session is None + assert mgr._consecutive_failures.get("test", 0) == 1 + + def test_session_terminated_mcperror_evicts_and_trips_circuit(self): + """Regression: when the MCP SERVER restarts and loses its session map, our + held mcp-session-id is stale; the server returns HTTP 404 and the SDK + surfaces McpError(code=32600, 'Session terminated'). That is NOT a healthy + protocol rejection — the session must be evicted so the next dispatch + reconnects with a fresh initialize; reusing it 404s forever (restart-hang).""" + from mcp import McpError + from mcp.types import ErrorData + + mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}}) + mock_session = MagicMock() + mock_session.call_tool = MagicMock(return_value="sentinel") + _seed_static_state(mgr, "test", session=mock_session) + mgr._loop = MagicMock() + mgr._tool_map["mcp__test__ping"] = ("test", "ping") + mock_future = MagicMock() + # Exactly what the streamable-http SDK injects on a 404 stale session. + mock_future.result.side_effect = McpError( + ErrorData(code=32600, message="Session terminated") + ) + with ( + patch("asyncio.run_coroutine_threadsafe", new=_dispatch_stub(mock_future)), + pytest.raises(McpError), + ): + mgr.call_tool_sync("mcp__test__ping", {}, timeout=5) + assert mgr._static_servers["test"].session is None + assert mgr._consecutive_failures.get("test", 0) == 1 + + def test_httpx_connect_error_evicts_session(self): + """A dead underlying httpx connection (server down mid-call) is transport + death, not a protocol rejection — evict so the next call reconnects.""" + import httpx + + mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}}) + mock_session = MagicMock() + mock_session.call_tool = MagicMock(return_value="sentinel") + _seed_static_state(mgr, "test", session=mock_session) + mgr._loop = MagicMock() + mgr._tool_map["mcp__test__ping"] = ("test", "ping") + mock_future = MagicMock() + mock_future.result.side_effect = httpx.ConnectError("connection refused") + with ( + patch("asyncio.run_coroutine_threadsafe", new=_dispatch_stub(mock_future)), + pytest.raises(httpx.ConnectError), + ): + mgr.call_tool_sync("mcp__test__ping", {}, timeout=5) + assert mgr._static_servers["test"].session is None + assert mgr._consecutive_failures.get("test", 0) == 1 + + def test_connection_closed_mcperror_evicts_and_trips_circuit(self): + """Regression: when the SDK's ``post_writer`` swallows the transport + error, a dead connection surfaces as ``McpError(CONNECTION_CLOSED)``. + Unlike a genuine protocol rejection, this MUST evict + trip the + breaker so the next dispatch reconnects instead of looping.""" + from mcp import McpError + from mcp.types import CONNECTION_CLOSED, ErrorData + + mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}}) + mock_session = MagicMock() + mock_session.call_tool = MagicMock(return_value="sentinel") + _seed_static_state(mgr, "test", session=mock_session) + mgr._loop = MagicMock() + mgr._tool_map["mcp__test__ping"] = ("test", "ping") + mock_future = MagicMock() + mock_future.result.side_effect = McpError( + ErrorData(code=CONNECTION_CLOSED, message="connection closed") + ) + with ( + patch("asyncio.run_coroutine_threadsafe", new=_dispatch_stub(mock_future)), + pytest.raises(McpError), + ): + mgr.call_tool_sync("mcp__test__ping", {}, timeout=5) + assert mgr._static_servers["test"].session is None + assert mgr._consecutive_failures.get("test", 0) == 1 + + def test_refresh_all_evicts_dead_session_so_next_tick_reconnects(self): + """Regression: a periodic refresh that hits a dead-but-non-None + session must null the session so the reconnect branch (gated on + ``session is None``) fires on the NEXT tick. Without this the + refresh re-probes the corpse forever — the bug that required a + full restart.""" + import anyio + + async def _run() -> None: + mgr = MCPClientManager({}) + mgr._server_configs["test"] = {"type": "stdio", "command": "x"} + dead = anyio.ClosedResourceError() + mock_session = MagicMock() + mock_session.list_tools = AsyncMock(side_effect=dead) + mock_session.list_resources = AsyncMock(side_effect=dead) + mock_session.list_resource_templates = AsyncMock(side_effect=dead) + mock_session.list_prompts = AsyncMock(side_effect=dead) + _seed_static_state(mgr, "test", session=mock_session) + + await mgr._refresh_all("test") + + # Dead session evicted → next refresh tick / dispatch reconnects. + assert mgr._static_servers["test"].session is None + ts, outcome = mgr._last_refresh["test"] + assert outcome == "error:ClosedResourceError" + + asyncio.run(_run()) + # --------------------------------------------------------------------------- # Fix 3: Safe transport stream pre-close diff --git a/tests/test_mcp_pool_auth_introspection.py b/tests/test_mcp_pool_auth_introspection.py index e40b5667..f891de72 100644 --- a/tests/test_mcp_pool_auth_introspection.py +++ b/tests/test_mcp_pool_auth_introspection.py @@ -1609,16 +1609,60 @@ class TestPoolPrimingAndTokenRotation: assert primed == [(("user-1", "pool-srv"), "bearer-fresh")] - def test_prime_user_pools_skips_near_expiry_without_revoking( + def test_prime_user_pools_refreshes_expired_token_and_warms( self, running_loop_mgr, storage: SQLiteBackend ) -> None: - """bug-1 regression: a near-expiry token is skipped (not refreshed), so a - transient refresh failure during priming can never revoke the token.""" + """An expired/near-expiry token is now REFRESHED (via the guarded + classified resolver) and the pool is warmed with the fresh token — + closing the chicken-and-egg where an expired token left the pool + permanently cold ("connecting" / no tools / never-refreshed).""" + from unittest.mock import patch + + from turnstone.core.mcp_oauth import TokenLookupResult + + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher, expires_in_seconds=5, access_token="bearer-stale") + self._wire(mgr, storage, cipher) + + primed: list[tuple[tuple[str, str], str]] = [] + + async def _fake_prime( + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + ) -> int: + primed.append((key, token)) + return 3 + + mgr._prime_user_server = _fake_prime.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + # The resolver refreshed the expired token and returns the fresh one. + return TokenLookupResult(kind="token", token="bearer-refreshed") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + _run_on_loop(loop, mgr._prime_user_pools("user-1")) + + assert primed == [(("user-1", "pool-srv"), "bearer-refreshed")], ( + "expired token must be refreshed and the pool warmed with the fresh token" + ) + + def test_prime_user_pools_transient_refresh_failure_skips_without_revoking( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """Safety invariant preserved: a TRANSIENT refresh failure during priming + does not warm the pool AND does not revoke — the classified resolver keeps + the token (kind=refresh_failed_transient) and lazy dispatch retries later.""" + from unittest.mock import patch + + from turnstone.core.mcp_oauth import TokenLookupResult + mgr, loop, _ = running_loop_mgr cipher = make_mcp_token_cipher() _seed_oauth_server(storage, name="pool-srv") - # Inside the 60s refresh-skew window -> the refreshing lookup would have - # driven a refresh here. _seed_user_token(storage, cipher, expires_in_seconds=5, access_token="bearer-stale") self._wire(mgr, storage, cipher) @@ -1632,10 +1676,17 @@ class TestPoolPrimingAndTokenRotation: mgr._prime_user_server = _fake_prime.__get__(mgr, type(mgr)) # type: ignore[method-assign] - _run_on_loop(loop, mgr._prime_user_pools("user-1")) + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="refresh_failed_transient") - assert primed == [], "near-expiry token must be skipped, not primed (no refresh driven)" - # The token row must survive — priming must never revoke. + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + _run_on_loop(loop, mgr._prime_user_pools("user-1")) + + assert primed == [], "transient refresh failure must not warm the pool" + # The token row must survive — priming must never revoke on a transient blip. store = MCPTokenStore(storage, cipher, node_id="test") assert store.get_user_token("user-1", "pool-srv") is not None @@ -1840,5 +1891,40 @@ class TestPoolPrimingAndTokenRotation: assert mgr._priming_keys == set(), "in-flight marker must be cleared in finally" +class TestOAuthUserServerStatus: + """``get_server_status`` for ``auth_type='oauth_user'`` servers reflects + per-user pool warmth instead of a permanent global "connecting" — so the + console pill flips to connected once a pool is primed (the third leg of the + OBO fix set).""" + + def test_oauth_user_status_connected_when_pool_warm(self) -> None: + from turnstone.core.mcp_client import PoolEntryState + + mgr = MCPClientManager({}) + mgr._oauth_user_server_names = {"pool-srv"} + entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock()) + entry.session = MagicMock() + entry.tools = [{"function": {"name": "mcp__pool-srv__do"}}] + mgr._user_pool_entries[("user-1", "pool-srv")] = entry + + st = mgr.get_server_status("pool-srv") + assert st["connected"] is True + assert st["tools"] == 1 + assert st["auth_type"] == "oauth_user" + assert st["user_pools"] == 1 + # Also surfaced in the all-servers map (oauth_user is absent from + # _server_configs, so this exercises the explicit union). + assert "pool-srv" in mgr.get_all_server_status() + + def test_oauth_user_status_connecting_when_no_warm_pool(self) -> None: + mgr = MCPClientManager({}) + mgr._oauth_user_server_names = {"pool-srv"} + st = mgr.get_server_status("pool-srv") + assert st["connected"] is False + assert st["tools"] == 0 + assert st["user_pools"] == 0 + assert st["auth_type"] == "oauth_user" + + # Suppress unused-import warning for AsyncMock. _ = AsyncMock diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 29330067..eeb51c4b 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -31,6 +31,7 @@ from datetime import UTC, datetime from pathlib import Path from typing import TYPE_CHECKING, Any, Literal +import anyio import httpx import mcp.types as mcp_types from mcp import ClientSession, McpError, StdioServerParameters @@ -145,6 +146,56 @@ _MAX_PROMPTS_PER_SERVER = 1000 _PRIME_MAX_CONCURRENCY = 4 +def _is_dead_transport(exc: BaseException) -> bool: + """True when *exc* means the MCP session's transport is dead and the + session must be torn down and rebuilt (vs a protocol-level rejection + from a still-healthy connection). + + The streamable-http SDK holds the session over anyio in-memory streams. + Three distinct death modes all mean "reconnect me", not "the server + rejected my request": + + 1. **Local stream torn down** — the GET/SSE or POST stream died (idle close, + peer reset, keep-alive expiry) and the ``ClientSession`` object survives + with a closed write stream, so ``list_tools``/``call_tool`` raises + :class:`anyio.ClosedResourceError` / :class:`anyio.BrokenResourceError`. + 2. **Transport-swallowed** — the SDK's ``post_writer`` swallows the upstream + error and the caller sees ``McpError(CONNECTION_CLOSED)`` (-32000); or the + underlying httpx connection is simply gone (server down / peer closed + mid-response → ``httpx.ConnectError`` / ``httpx.RemoteProtocolError``). + 3. **Server-side session lost** — the MCP server RESTARTED and dropped its + session map, so our held ``mcp-session-id`` is unknown. The server returns + HTTP 404 and the SDK surfaces ``McpError(code=32600, "Session terminated")`` + (see ``streamable_http._send_session_terminated_error``). The session-id is + dead server-side; only a fresh ``initialize`` (full reconnect) recovers it — + reusing it 404s forever (the restart-hang). Matched by message because the + SDK's code (32600) is not a documented MCP constant. + + The raw-socket variants (``BrokenPipeError`` / ``ConnectionResetError`` / + ``EOFError``) are kept for the stdio transport and defense-in-depth. + """ + if isinstance( + exc, + anyio.ClosedResourceError + | anyio.BrokenResourceError + | BrokenPipeError + | ConnectionResetError + | EOFError + | httpx.ConnectError + | httpx.ConnectTimeout + | httpx.RemoteProtocolError, + ): + return True + if isinstance(exc, McpError): + err = exc.error + if getattr(err, "code", None) == mcp_types.CONNECTION_CLOSED: + return True + msg = (getattr(err, "message", "") or "").lower() + if "session terminated" in msg or "session not found" in msg: + return True + return False + + @dataclass class _AuthCapture: """Carrier populated by the response hook on 4xx upstream responses.""" @@ -1722,16 +1773,20 @@ class MCPClientManager: async def _prime_user_pools(self, user_id: str) -> None: """Warm THIS user's consented ``oauth_user`` pools (runs on the mcp-loop). - Best-effort and NON-DESTRUCTIVE: each token is read directly (NOT via the - refresh state machine) and missing/near-expiry tokens are skipped, so a - transient AS/network failure during priming can never delete a token and - force re-consent — a refresh that may fail belongs on the lazy dispatch - path, driven by actual use. Servers are primed concurrently under - ``_PRIME_MAX_CONCURRENCY`` so one slow/unreachable upstream can't stall - the rest. + Best-effort and transient-safe: each token is resolved via the SAME + guarded refresh state machine the lazy-dispatch path uses + (:func:`get_user_access_token_classified`). That refreshes an expired / + near-expiry access token and persists it, but a TRANSIENT AS/network + failure keeps the token (``kind=refresh_failed_transient``, no revoke) + and is simply skipped here — only a genuinely PERMANENT rejection (the + user must re-consent anyway) clears it. This closes the chicken-and-egg + where an expired token made priming skip the server, its tools never + entered the per-user catalog, and lazy dispatch — the only OTHER refresh + trigger — could therefore never fire, leaving the pool permanently cold + and stuck "connecting" with no token. Servers are primed concurrently + under ``_PRIME_MAX_CONCURRENCY`` so one slow/unreachable upstream can't + stall the rest. """ - from turnstone.core.mcp_oauth import _token_needs_refresh - token_store = getattr(self._app_state, "mcp_token_store", None) if token_store is None: return @@ -1750,22 +1805,24 @@ class MCPClientManager: try: async with sem: try: - # Non-refreshing read — priming must not drive a refresh - # whose transient failure would revoke the token. - plain = await asyncio.to_thread( - token_store.get_user_token, user_id, server_name + # Guarded resolve: refreshes an expired/near-expiry token + # and persists it, but a transient failure keeps the token + # (kind=refresh_failed_transient) so priming can never + # revoke a live grant. Only kind == "token" warms the pool. + lookup = await get_user_access_token_classified( + app_state=self._app_state, + user_id=user_id, + server_name=server_name, ) - if plain is None or not plain.get("access_token"): - return # no usable token (not consented) — lazy paths handle it - if _token_needs_refresh(plain.get("expires_at")): - return # near expiry — let lazy dispatch refresh on actual use + if lookup.kind != "token" or not lookup.token: + return # not consented / undecryptable / refresh failed — lazy paths handle it server_row = await asyncio.to_thread( self._storage.get_mcp_server_by_name, server_name ) if not server_row: return cfg = _pool_cfg_from_row(server_row) - await self._prime_user_server(key, cfg, plain["access_token"]) + await self._prime_user_server(key, cfg, lookup.token) log.info( "mcp pool auto-primed at session start user=%s server=%s", user_id, @@ -1948,9 +2005,15 @@ class MCPClientManager: return "auth_401" if status == 403: return "auth_403" + # A closed/broken transport must be classified BEFORE the McpError + # branch: the SDK surfaces a dead connection as McpError(CONNECTION_CLOSED), + # which would otherwise be mistaken for a healthy protocol rejection and + # leave the dead pool entry in place forever (no rebuild). + if _is_dead_transport(exc): + return "transport" if isinstance(exc, McpError): return "protocol" - if isinstance(exc, BrokenPipeError | ConnectionResetError | EOFError | TimeoutError): + if isinstance(exc, TimeoutError): return "transport" return "other" @@ -2434,6 +2497,17 @@ class MCPClientManager: log.warning("Refresh failed for MCP server '%s'", name, exc_info=True) self._set_error(name, f"Refresh failed: {exc}") results[name] = ([], []) + # A dead transport leaves a non-None but unusable session, so + # the reconnect branch at the top of this loop (gated on + # ``session is None``) would never fire and we would re-probe + # the corpse on every tick forever — the exact failure that + # required a full process restart to clear. Evicting the + # session here makes the NEXT refresh tick reconnect, turning + # this periodic refresh into a self-healing liveness probe. + if _is_dead_transport(exc): + dead_state = self._static_servers.get(name) + if dead_state is not None: + dead_state.session = None # Overwrite unconditionally with the freshest observed # outcome. Two cases produce the write: # (1) Reconnect branch: ``_connect_one`` raised before @@ -3256,6 +3330,13 @@ class MCPClientManager: def get_server_status(self, name: str) -> dict[str, Any]: """Return live status for a single server, including config details.""" + # auth_type='oauth_user' servers hold NO process-global session — they + # are warmed per-user into the pool — so the static-session check below + # would always report them "connecting". Derive their status from warm + # per-user pool entries instead, so the console pill reflects real + # per-user reachability once a pool is primed. + if name in self._oauth_user_server_names: + return self._oauth_user_server_status(name) state = self._static_servers.get(name) connected = state is not None and state.session is not None cfg = self._server_configs.get(name, {}) @@ -3284,11 +3365,55 @@ class MCPClientManager: "last_refresh_outcome": last_refresh[1] if last_refresh is not None else None, } + def _oauth_user_server_status(self, name: str) -> dict[str, Any]: + """Live status for an ``auth_type='oauth_user'`` server. + + These have no global session (stripped from ``_server_configs`` / + ``_static_servers``); they connect per-user into ``_user_pool_entries``. + Report ``connected=True`` when at least one user has a warm pool entry, + with a representative catalog count and the number of warm user pools — + so the console stops showing a permanent "connecting"/``---`` for a + server that is in fact reachable and in use. + """ + warm = [ + entry + for (uid, sname), entry in self._user_pool_entries.items() + if sname == name and entry.session is not None + ] + rep = warm[0] if warm else None + cb_deadline = self._circuit_open_until.get(name) + cb_open = cb_deadline is not None and time.monotonic() < cb_deadline + last_refresh = self._last_refresh.get(name) + return { + "connected": bool(warm), + "tools": len(rep.tools) if rep is not None and rep.tools else 0, + "resources": len(rep.resources) if rep is not None and rep.resources else 0, + "prompts": len(rep.prompts) if rep is not None and rep.prompts else 0, + "error": self._last_error.get(name, ""), + "transport": "streamable-http", + "command": "", + "url": "", + "circuit_open": cb_open, + "consecutive_failures": self._consecutive_failures.get(name, 0), + "auth_type": "oauth_user", + "user_pools": len(warm), + "last_refresh_at": last_refresh[0] if last_refresh is not None else None, + "last_refresh_outcome": last_refresh[1] if last_refresh is not None else None, + } + def get_all_server_status(self) -> dict[str, dict[str, Any]]: - """Return live status for all configured servers.""" + """Return live status for all configured servers. + + Includes ``oauth_user`` servers (which are absent from + ``_server_configs``) so the console list reports their real per-user + pool status instead of falling back to a DB-only "connecting" default. + """ result: dict[str, dict[str, Any]] = {} for name in list(self._server_configs): result[name] = self.get_server_status(name) + for name in list(self._oauth_user_server_names): + if name not in result: + result[name] = self.get_server_status(name) return result def reconcile_sync(self, storage: Any, timeout: int = 30) -> dict[str, Any]: @@ -3722,12 +3847,19 @@ class MCPClientManager: except Exception as exc: # Protocol errors (McpError) come from a healthy connection that # rejected the request — only transport errors trip the breaker. - if not isinstance(exc, McpError): + # A dead transport (anyio Closed/BrokenResourceError, or the + # SDK-swallowed McpError(CONNECTION_CLOSED)) IS a transport failure + # even though it is an McpError, so it must trip the breaker AND + # evict the session. + dead = _is_dead_transport(exc) + if dead or not isinstance(exc, McpError): self._cb_record_failure(server_name) - if isinstance(exc, BrokenPipeError | ConnectionResetError | EOFError): + if dead: # Evict the session only — leave stack/streams behind so the # stale-session-and-stack guard in _connect_one cleans them up - # on the next connect attempt. + # on the next connect attempt. Eviction is what lets the next + # dispatch's ``session is None`` check fire _cb_auto_reconnect + # instead of re-using the corpse forever. evict = self._static_servers.get(server_name) if evict is not None: evict.session = None