mcp dead transport fix and token refresh

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

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

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