From 1e84f62619a4763fe84e2c3c06bbfd9f00a7c3f0 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 13 Jul 2026 11:38:05 -0700 Subject: [PATCH] =?UTF-8?q?fix(mcp):=20round-3=20review=20fixes=20?= =?UTF-8?q?=E2=80=94=20revocation=20generation=20for=20catalog=20publisher?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 3 identified the class behind the remaining bugs: catalog PUBLISHERS never re-validate revocation state, so anything that read a token or suspended before a drop could republish (resurrect) a revoked catalog that retention then keeps forever. One primitive closes the class: - PoolEntryState.catalog_gen, bumped by _evict_session_drop_catalog. The three list_changed refreshes snapshot it before their awaits and discard results if it moved; dispatch and priming snapshot it before their token reads, and _connect_one_pool refuses to connect (raising _PoolGrantRevokedError, a non-breaker failure) when the generation moved past the caller's snapshot — the bearer in hand predates a disconnect. - _connect_one_pool stages all three discovery results locally and publishes them together in the final wiring block: a mid-discovery failure now leaves the retained catalog exactly as it was instead of a torn half-update diverging from the per-user maps. - Priming converges dead grants too: _prime_one schedules the same catalog drop the dispatchers use, so a NEW session's prime clears ghosts left by a disconnect made on another node. - obo re-login is the obo restore moment: a successful credential capture at the OIDC callback now schedules prime_user_pools, so a previously dropped obo catalog returns to LIVE sessions (obo has no consent flow to heal through). - ChatSession construction re-runs its tool rebuild when the change marker advanced during its authoritative read — the mirror race where a fresher listener update was clobbered by the constructor's staler snapshot. - evict_user_session's drop task is now tracked (_spawn_background) so shutdown cancels it instead of abandoning a parked task. Dedup/altitude from the round: _schedule_dead_grant_drop is the single drop block (was three byte-identical copies); _pool_lookup_verdict is the single lookup classification — rendering and _lookup_grant_dead both derive from it, with literal code strings kept so the consent-url sibling audit still sees the sites (expected count 7 -> 5 after the collapse); PoolEntryState.drop_session() pairs session/bound_token clearing structurally (owner-death was missing the bearer clear). --- tests/test_mcp_consent_url_sibling_audit.py | 22 +- tests/test_mcp_oauth_connections.py | 1 + tests/test_mcp_pool_auth_introspection.py | 45 ++- tests/test_mcp_user_pool.py | 114 ++++++- tests/test_oidc_handlers.py | 62 ++++ turnstone/core/auth.py | 17 + turnstone/core/mcp_client.py | 343 +++++++++++++------- turnstone/core/session.py | 14 + 8 files changed, 486 insertions(+), 132 deletions(-) diff --git a/tests/test_mcp_consent_url_sibling_audit.py b/tests/test_mcp_consent_url_sibling_audit.py index 1b3ef281..cd02585b 100644 --- a/tests/test_mcp_consent_url_sibling_audit.py +++ b/tests/test_mcp_consent_url_sibling_audit.py @@ -97,23 +97,23 @@ def test_every_user_actionable_structured_error_passes_consent_url() -> None: def test_audit_finds_all_known_user_actionable_sites() -> None: """Lock the count so accidental deletions are caught. - There are 7 user-actionable ``_structured_error`` call sites today: - 3 in ``_pool_lookup_error`` (the single lookup-kind → error mapping - shared by the tool / resource / prompt dispatchers — it replaced the - 9 hand-synced per-dispatcher copies), 3 in the post-retry-failed - branches, and 1 in ``_handle_auth_403``'s insufficient-scope branch. - If a new exec path is added the count can rise; if a branch is - removed the count can fall — both are fine, but require an - intentional bump of this number to confirm the change went through - review. + There are 5 user-actionable ``_structured_error`` call sites today: + 1 in ``_pool_lookup_error`` (the ``_pool_lookup_verdict`` split + collapsed the former 3 consent-required renderings into one site — + the per-kind branches now select only the DETAIL copy), 3 in the + post-retry-failed branches, and 1 in ``_handle_auth_403``'s + insufficient-scope branch. If a new exec path is added the count + can rise; if a branch is removed the count can fall — both are + fine, but require an intentional bump of this number to confirm + the change went through review. """ source = _read_source() blocks = _find_structured_error_blocks(source) user_actionable_count = sum( 1 for _, blk in blocks if any(f'code="{code}"' in blk for code in _USER_ACTIONABLE_CODES) ) - assert user_actionable_count == 7, ( - f"Expected 7 user-actionable _structured_error sites, got " + assert user_actionable_count == 5, ( + f"Expected 5 user-actionable _structured_error sites, got " f"{user_actionable_count}. If this is intentional, bump the " f"expected count and document why in the commit message." ) diff --git a/tests/test_mcp_oauth_connections.py b/tests/test_mcp_oauth_connections.py index 84ef3193..b1f249f8 100644 --- a/tests/test_mcp_oauth_connections.py +++ b/tests/test_mcp_oauth_connections.py @@ -804,6 +804,7 @@ class TestEvictUserSession: mgr._loop = loop # type: ignore[attr-defined] mgr._user_pool_entries = {} # type: ignore[attr-defined] mgr._last_pool_notification_refresh = {} # type: ignore[attr-defined] + mgr._background_tasks = set() # type: ignore[attr-defined] evicted: list[tuple[str, str]] = [] async def _fake_evict(key: tuple[str, str]) -> None: diff --git a/tests/test_mcp_pool_auth_introspection.py b/tests/test_mcp_pool_auth_introspection.py index e61b08ad..6dc2e2ab 100644 --- a/tests/test_mcp_pool_auth_introspection.py +++ b/tests/test_mcp_pool_auth_introspection.py @@ -571,6 +571,7 @@ class TestDispatcherAuthFlows: *, auth_capture: Any = None, auth_fired_event: Any = None, + expected_gen: Any = None, ) -> Any: entry = await self_inner._ensure_pool_entry(key) sess = MagicMock() @@ -980,6 +981,7 @@ class TestBreakerInvariant: *, auth_capture: Any = None, auth_fired_event: Any = None, + expected_gen: Any = None, ) -> Any: entry = await self_inner._ensure_pool_entry(key) sess = MagicMock() @@ -1600,7 +1602,11 @@ class TestPoolPrimingAndTokenRotation: primed: list[tuple[tuple[str, str], str]] = [] async def _fake_prime( - self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + token: str, + expected_gen: Any = None, ) -> int: primed.append((key, token)) return 3 @@ -1631,7 +1637,11 @@ class TestPoolPrimingAndTokenRotation: primed: list[tuple[tuple[str, str], str]] = [] async def _fake_prime( - self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + token: str, + expected_gen: Any = None, ) -> int: primed.append((key, token)) return 3 @@ -1671,7 +1681,11 @@ class TestPoolPrimingAndTokenRotation: primed: list[tuple[str, str]] = [] async def _fake_prime( - self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + token: str, + expected_gen: Any = None, ) -> int: primed.append(key) return 0 @@ -1806,7 +1820,11 @@ class TestPoolPrimingAndTokenRotation: primed: list[tuple[str, str]] = [] async def _fake_prime( - self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + token: str, + expected_gen: Any = None, ) -> int: primed.append(key) return 0 @@ -1858,7 +1876,11 @@ class TestPoolPrimingAndTokenRotation: done = threading.Event() async def _fake_prime( - self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + token: str, + expected_gen: Any = None, ) -> int: captured["key"] = key captured["token"] = token @@ -1918,6 +1940,7 @@ class TestPoolPrimingAndTokenRotation: *, auth_capture: Any = None, auth_fired_event: Any = None, + expected_gen: Any = None, ) -> Any: reconnect_tokens.append(access_token) entry = await self_inner._ensure_pool_entry(key) @@ -1954,7 +1977,11 @@ class TestPoolPrimingAndTokenRotation: primed: list[tuple[str, str]] = [] async def _fake_prime( - self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + token: str, + expected_gen: Any = None, ) -> int: primed.append(key) return 0 @@ -1978,7 +2005,11 @@ class TestPoolPrimingAndTokenRotation: self._wire(mgr, storage, cipher) async def _fake_prime( - self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + token: str, + expected_gen: Any = None, ) -> int: return 1 diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index 0f39c245..2c585102 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -521,6 +521,62 @@ class TestEviction: assert blocked, "drop must wait for the in-flight connect's open_lock" assert cleared, "drop must win once the connect completes — no resurrection" + def test_connect_refuses_when_generation_moved(self, running_loop_mgr) -> None: + """``_connect_one_pool`` must refuse to connect (and so to + publish) when the entry's revocation generation moved past the + caller's pre-token-read snapshot — the bearer in hand predates + a disconnect, and publishing would resurrect the dropped + catalog with nothing left to clear it.""" + from turnstone.core.mcp_client import _PoolGrantRevokedError + + mgr, loop, _ = running_loop_mgr + key = ("u0", "pool-srv") + + async def _scenario() -> None: + entry = await mgr._ensure_pool_entry(key) + gen0 = entry.catalog_gen + entry.catalog_gen += 1 # a revocation drop landed in the window + with pytest.raises(_PoolGrantRevokedError): + await mgr._connect_one_pool( + key, + {"type": "streamable-http", "url": "https://mcp.example.com/sse"}, + "stale-bearer", + expected_gen=gen0, + ) + + _run_on_loop(loop, _scenario()) + + def test_refresh_discards_result_when_generation_moved(self, running_loop_mgr) -> None: + """A ``list_changed`` refresh whose await straddles a revocation + drop must DISCARD its result — publishing would resurrect the + revoked catalog, which retention then keeps alive forever.""" + mgr, loop, _ = running_loop_mgr + mgr._oauth_user_server_names = {"pool-srv"} + key = ("u0", "pool-srv") + + async def _scenario() -> tuple[list, list, Any, bool]: + entry = await mgr._ensure_pool_entry(key) + entry.tools = self._fake_tools("pool-srv", 0) + mgr._rebuild_user_tool_map("u0") + + class _RacingSession: + async def list_tools(self) -> Any: + # The disconnect lands while list_tools is in flight. + mgr._evict_session_drop_catalog(key) + res = MagicMock() + res.tools = [] + return res + + entry.session = _RacingSession() + added, removed = await mgr._refresh_pool_server_tools(key) + return added, removed, entry.tools, "u0" in mgr._user_tool_map + + added, removed, tools, in_map = _run_on_loop(loop, _scenario()) + assert (added, removed) == ([], []) + # The drop's clear stands; the refresh did not republish. + assert tools is None + assert in_map is False + def test_lookup_grant_dead_requires_wired_infrastructure(self, running_loop_mgr) -> None: """kind='missing' is authoritative only when the stores that could know are wired: the obo lookup returns 'missing' for an @@ -1369,7 +1425,9 @@ class TestOboPriming: warmed: list[tuple[Any, Any, str]] = [] - async def _fake_prime_server(key: Any, cfg: Any, token: str) -> None: + async def _fake_prime_server( + key: Any, cfg: Any, token: str, expected_gen: Any = None + ) -> None: warmed.append((key, cfg, token)) obo_lookup = AsyncMock(return_value=SimpleNamespace(kind="token", token="minted-at")) @@ -1388,6 +1446,60 @@ class TestOboPriming: assert len(warmed) == 1 assert warmed[0][2] == "minted-at" + def test_prime_drops_retained_catalog_on_dead_grant(self, running_loop_mgr, storage) -> None: + """Priming is a convergence point (#836): a NEW session's prime + that finds the grant durably GONE must drop the retained catalog + that other live sessions still serve (e.g. a disconnect made on + another node) — not just skip the server.""" + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + mgr.set_storage(storage) + app_state = _make_app_state(storage, cipher=cipher) + mgr.set_app_state(app_state) + storage.create_mcp_server( + server_id="srv-o", + name="pool-srv", + transport="streamable-http", + url="https://mcp.example.com/sse", + auth_type="oauth_user", + ) + mgr._oauth_user_server_names = {"pool-srv"} + key = ("user-1", "pool-srv") + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(key) + entry.tools = [ + { + "type": "function", + "function": { + "name": "mcp__pool-srv__t", + "description": "", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + mgr._rebuild_user_tool_map("user-1") + + _run_on_loop(loop, _seed()) + fired = [0] + + def _cb() -> None: + fired[0] += 1 + + mgr.add_listener(_cb, user_id="user-1") + assert mgr.is_mcp_tool("mcp__pool-srv__t", user_id="user-1") is True + + dead = AsyncMock(return_value=SimpleNamespace(kind="missing", token=None)) + with patch("turnstone.core.mcp_client.get_user_access_token_classified", new=dead): + _run_on_loop(loop, mgr._prime_user_pools("user-1")) + # The drop is scheduled — flush the loop before asserting. + _run_on_loop(loop, asyncio.sleep(0.05)) + + entry = mgr._user_pool_entries[key] + assert entry.tools is None + assert mgr.is_mcp_tool("mcp__pool-srv__t", user_id="user-1") is False + assert fired[0] == 1 + def test_prime_skips_obo_server_when_no_credential(self, running_loop_mgr, storage) -> None: """A user without a captured credential is skipped BEFORE any per-server obo work: one existence SELECT decides all obo servers diff --git a/tests/test_oidc_handlers.py b/tests/test_oidc_handlers.py index 9d4b0a39..0d8de0b8 100644 --- a/tests/test_oidc_handlers.py +++ b/tests/test_oidc_handlers.py @@ -8,6 +8,7 @@ on the HTTP handler logic, request/response wiring, and storage side-effects. from __future__ import annotations import urllib.parse +from types import SimpleNamespace from typing import TYPE_CHECKING, Any from unittest.mock import AsyncMock, patch @@ -832,6 +833,67 @@ class TestOIDCCallbackCapture: assert plain is not None assert plain["refresh_token"] == "rt-1" + @patch("turnstone.core.auth.provision_oidc_user") + @patch("turnstone.core.auth.validate_id_token") + @patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock) + def test_capture_success_primes_user_pools( + self, + mock_exchange: AsyncMock, + mock_validate: Any, + mock_provision: Any, + storage: SQLiteBackend, + oidc_config: OIDCConfig, + ) -> None: + """Re-login is the OBO restore moment (#836): a successful + credential capture schedules a pool prime so a previously + dropped obo catalog returns to the user's LIVE sessions — obo + has no consent flow, so nothing else re-primes an open + workstream after re-login.""" + client, store, cfg = self._capture_client(storage, oidc_config) + primed: list[str] = [] + client.app.state.mcp_client = SimpleNamespace( # type: ignore[attr-defined] + prime_user_pools=primed.append + ) + resp = self._login( + client, + storage, + mock_exchange, + mock_validate, + mock_provision, + tokens={"id_token": "fake.jwt.token", "access_token": "at", "refresh_token": "rt-1"}, + ) + assert resp.status_code == 302 + assert primed == ["test-admin"] + + @patch("turnstone.core.auth.provision_oidc_user") + @patch("turnstone.core.auth.validate_id_token") + @patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock) + def test_no_capture_no_prime( + self, + mock_exchange: AsyncMock, + mock_validate: Any, + mock_provision: Any, + storage: SQLiteBackend, + oidc_config: OIDCConfig, + ) -> None: + """No refresh token in the response → no capture → no prime + (the prime is gated on a persisted credential, not on login).""" + client, store, cfg = self._capture_client(storage, oidc_config) + primed: list[str] = [] + client.app.state.mcp_client = SimpleNamespace( # type: ignore[attr-defined] + prime_user_pools=primed.append + ) + resp = self._login( + client, + storage, + mock_exchange, + mock_validate, + mock_provision, + tokens={"id_token": "fake.jwt.token", "access_token": "at"}, + ) + assert resp.status_code == 302 + assert primed == [] + @patch("turnstone.core.auth.provision_oidc_user") @patch("turnstone.core.auth.validate_id_token") @patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock) diff --git a/turnstone/core/auth.py b/turnstone/core/auth.py index ba17b060..1f04e899 100644 --- a/turnstone/core/auth.py +++ b/turnstone/core/auth.py @@ -2152,6 +2152,23 @@ async def handle_oidc_callback(request: Request, audience: str, cookie_name: str ) except Exception: log.exception("oidc.capture: failed to persist credential") + else: + # Re-login is the OBO restore moment (#836): a dropped + # obo catalog (credential unlinked / mint rejected) has + # no consent flow to heal through, so warm this user's + # pools now — live sessions pick the tools back up via + # their listeners. Fire-and-forget; a failure changes + # nothing about login. + mcp_client = getattr(request.app.state, "mcp_client", None) + if mcp_client is not None and hasattr(mcp_client, "prime_user_pools"): + try: + mcp_client.prime_user_pools(user["user_id"]) + except Exception: + log.debug( + "oidc.capture: post-capture pool prime scheduling failed user=%s", + user["user_id"], + exc_info=True, + ) # Load permissions and issue Turnstone JWT perms = await asyncio.to_thread(_load_user_permissions, storage, user["user_id"]) diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index ecbe4f47..7702ba12 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -278,6 +278,17 @@ class _CarrierAuthSignal(Exception): # noqa: N818 """ +class _PoolGrantRevokedError(Exception): + """The grant was revoked between a caller's token read and its locked connect. + + Raised by ``_connect_one_pool`` when the entry's ``catalog_gen`` + moved past the caller's pre-token-read snapshot: a disconnect ran + in that window, and connecting with the pre-revocation bearer would + publish (resurrect) a catalog for a dead grant. Classified as a + non-breaker failure — the server is healthy; the grant is not. + """ + + def _make_capturing_http_factory( capture: _AuthCapture, fired_event: asyncio.Event | None = None ) -> McpHttpClientFactory: @@ -538,6 +549,14 @@ class PoolEntryState: tools: list[dict[str, Any]] | None = None resources: list[dict[str, Any]] | None = None prompts: list[dict[str, Any]] | None = None + # Revocation generation. Bumped by ``_evict_session_drop_catalog``; + # every catalog PUBLISHER (connect wiring, the three list_changed + # refreshes) compares against the generation it snapshotted before + # its unlocked window and discards its publish on mismatch — + # otherwise a publisher that read state before the revocation + # republishes (resurrects) a revoked catalog after the drop, and + # the retention rules then keep the ghost alive indefinitely. + catalog_gen: int = 0 last_used: float = 0.0 in_flight: int = 0 # Access token this session's httpx client was connected with. The bearer is @@ -565,6 +584,20 @@ class PoolEntryState: supports_resource_list_changed: bool = False supports_prompt_list_changed: bool = False + def drop_session(self) -> None: + """Drop the cached session AND its paired plaintext bearer copy. + + The ONE way to null ``session``: ``bound_token`` is read only + under a live session (stale-rebind detection) and overwritten + at reconnect, so it is dead the moment the session goes — and + an entry may cool indefinitely after any drop, which must not + retain a plaintext credential for the life of the user's + sessions. A site that nulled ``session`` directly would + silently re-open that hole. + """ + self.session = None + self.bound_token = None + # --------------------------------------------------------------------------- # Client manager @@ -2025,12 +2058,7 @@ class MCPClientManager: entry = self._user_pool_entries.get(key) if entry is None: return - entry.session = None - # The bearer copy is dead once the session is gone (read only - # under a live session for stale-rebind detection; overwritten - # at reconnect) — don't retain a plaintext credential on a - # cooled entry for the life of the user's sessions. - entry.bound_token = None + entry.drop_session() owner = entry.owner_task close_requested = entry.close_requested entry.owner_task = None @@ -2081,11 +2109,7 @@ class MCPClientManager: entry.owner_task = None entry.close_requested = None if entry.session is not None: - entry.session = None - # Third session-drop site of the bearer-clearing sweep (with - # _teardown_pool_entry and _evict_session): the copy is dead - # without a session, and the entry may now cool indefinitely. - entry.bound_token = None + entry.drop_session() user_id, server_name = key log.info( "MCP pool transport terminated user=%s server=%s; session evicted for reconnect", @@ -2145,6 +2169,7 @@ class MCPClientManager: *, auth_capture: _AuthCapture | None = None, auth_fired_event: asyncio.Event | None = None, + expected_gen: int | None = None, ) -> PoolEntryState: """Connect a single per-(user, server) pool entry. @@ -2180,6 +2205,15 @@ class MCPClientManager: ) entry = await self._ensure_pool_entry(key) + if expected_gen is not None and entry.catalog_gen != expected_gen: + # The grant was revoked between the caller's token read and + # this locked connect (the drop bumped the generation while + # we queued on open_lock). The bearer in hand predates the + # revocation — connecting would publish a catalog the drop + # can no longer clear (#836). + raise _PoolGrantRevokedError( + f"grant for user={user_id!r} server={server_name!r} revoked during connect window" + ) # Guard: close any stale owner/session so we don't leak, the same way # ``_connect_one_locked`` does for the static path (cf. PR #296 @@ -2309,7 +2343,12 @@ class MCPClientManager: raise capped_tools = _cap_server_tools(server_name, tools_result.tools) - entry.tools = [_mcp_to_openai(server_name, tool) for tool in capped_tools] + # STAGED — published to the entry only in the final wiring block + # below, together with resources/prompts: a mid-discovery + # failure tears the transport down and must leave the entry's + # (retained) catalog exactly as it was, never half-updated with + # the per-user maps still holding the old view. + server_tools = [_mcp_to_openai(server_name, tool) for tool in capped_tools] # Phase 7b — discover resources (capability-gated). Same anyio / # ``asyncio.timeout`` invariant as the tool discovery above (R1). @@ -2406,6 +2445,7 @@ class MCPClientManager: } ) + entry.tools = server_tools entry.resources = server_resources if resources_cap is not None else None entry.prompts = server_prompts if prompts_cap is not None else None @@ -2432,7 +2472,11 @@ class MCPClientManager: # -- pool priming --------------------------------------------------------- async def _prime_user_server( - self, key: tuple[str, str], cfg: dict[str, Any], access_token: str + self, + key: tuple[str, str], + cfg: dict[str, Any], + access_token: str, + expected_gen: int | None = None, ) -> int: """Proactively connect a pool entry so its catalog populates into ``get_tools(user_id)`` WITHOUT waiting for a tool dispatch. @@ -2462,6 +2506,7 @@ class MCPClientManager: access_token, auth_capture=entry.auth_capture, auth_fired_event=entry.auth_fired_event, + expected_gen=expected_gen, ) return len(fresh.tools or []) @@ -2512,8 +2557,19 @@ class MCPClientManager: failure must not change the user-observable consent outcome. """ try: - count = await self._prime_user_server(key, cfg, access_token) + # Snapshot the revocation generation on-loop, before the + # locked connect: a disconnect racing this consent-time + # prime must not have its drop republished by us. + pre = self._user_pool_entries.get(key) + expected_gen = pre.catalog_gen if pre is not None else None + count = await self._prime_user_server(key, cfg, access_token, expected_gen) log.info("mcp pool primed user=%s server=%s tools=%d", user_id, server_name, count) + except _PoolGrantRevokedError: + log.info( + "mcp pool prime skipped: grant revoked mid-prime user=%s server=%s", + user_id, + server_name, + ) except Exception: log.warning( "mcp pool prime failed user=%s server=%s", @@ -2578,6 +2634,10 @@ class MCPClientManager: entry = self._user_pool_entries.get(key) if entry is not None and entry.session is not None: return # already connected — nothing to do + # Snapshot the revocation generation before the token read + # (see _dispatch_pool); None = no entry existed yet, so + # there is no retained catalog a racer could resurrect. + expected_gen = entry.catalog_gen if entry is not None else None if key in self._priming_keys: return # a concurrent prime for this (user, server) is in flight # Claim synchronously before any await — the mcp-loop is single- @@ -2632,6 +2692,12 @@ class MCPClientManager: revoke_ambiguous_escalation=False, ) if lookup.kind != "token" or not lookup.token: + # Priming is also a convergence point (#836): + # a NEW session's prime discovering the grant + # is durably gone must drop the retained + # catalog other live sessions still serve — + # e.g. a disconnect made on another node. + self._schedule_dead_grant_drop(lookup, key) return # not consented / no credential / refresh failed — lazy paths handle it if server_row is None: server_row = await asyncio.to_thread( @@ -2640,12 +2706,18 @@ class MCPClientManager: if not server_row: return cfg = _pool_cfg_from_row(server_row) - await self._prime_user_server(key, cfg, lookup.token) + await self._prime_user_server(key, cfg, lookup.token, expected_gen) log.info( "mcp pool auto-primed at session start user=%s server=%s", user_id, server_name, ) + except _PoolGrantRevokedError: + log.info( + "mcp pool prime skipped: grant revoked mid-prime user=%s server=%s", + user_id, + server_name, + ) except Exception: log.debug( "mcp pool auto-prime failed user=%s server=%s", @@ -3532,7 +3604,13 @@ class MCPClientManager: session = entry.session user_id, server_name = key old_names = {t["function"]["name"] for t in (entry.tools or [])} + gen = entry.catalog_gen result = await session.list_tools() + if self._user_pool_entries.get(key) is not entry or entry.catalog_gen != gen: + # A revocation drop landed while list_tools was in flight — + # publishing now would resurrect the revoked catalog with + # nothing left to ever clear it (retention keeps it). Discard. + return [], [] capped = _cap_server_tools(server_name, result.tools) server_tools = [_mcp_to_openai(server_name, tool) for tool in capped] new_names = {t["function"]["name"] for t in server_tools} @@ -3573,6 +3651,7 @@ class MCPClientManager: session = entry.session user_id, server_name = key old_uris = {r["uri"] for r in (entry.resources or []) if not r.get("template")} + gen = entry.catalog_gen async with asyncio.timeout(self._CONNECT_TIMEOUT): # 1-RTT (gather) instead of 2 sequential RTTs — both calls @@ -3582,6 +3661,9 @@ class MCPClientManager: session.list_resources(), session.list_resource_templates(), ) + if self._user_pool_entries.get(key) is not entry or entry.catalog_gen != gen: + # Revocation drop landed mid-flight — discard, don't resurrect. + return [], [] server_resources: list[dict[str, Any]] = [] for r in _cap_server_resources(server_name, res_result.resources): @@ -3645,9 +3727,13 @@ class MCPClientManager: session = entry.session user_id, server_name = key old_names = {p["name"] for p in (entry.prompts or [])} + gen = entry.catalog_gen async with asyncio.timeout(self._CONNECT_TIMEOUT): prompt_result = await session.list_prompts() + if self._user_pool_entries.get(key) is not entry or entry.catalog_gen != gen: + # Revocation drop landed mid-flight — discard, don't resurrect. + return [], [] server_prompts: list[dict[str, Any]] = [] for p in _cap_server_prompts(server_name, prompt_result.prompts): @@ -4342,7 +4428,7 @@ class MCPClientManager: entry = self._user_pool_entries.get(key) if entry is None: continue - entry.session = None + entry.drop_session() owner = entry.owner_task close_requested = entry.close_requested entry.owner_task = None @@ -6399,26 +6485,19 @@ class MCPClientManager: # this retry; the local cached token is the one the AS just # rejected, so reading it back without ``force_refresh=True`` # would re-attempt with the same (rejected) bearer. + # Ensure the entry and snapshot its revocation generation BEFORE + # the token read: a disconnect landing between the row read and + # the locked connect bumps the generation, and _connect_one_pool + # then refuses to publish for the revoked grant (#836). + key = (user_id, server_name) + entry = await self._ensure_pool_entry(key) + expected_gen = entry.catalog_gen lookup: TokenLookupResult = await self._pool_token_lookup( server_row, user_id, server_name, force_refresh=retry_count > 0 ) lookup_error = _pool_lookup_error(lookup, server_name, server_row) if lookup_error is not None: - if self._lookup_grant_dead(lookup): - # The grant is GONE (revoked / permanently rejected) — - # converge this node's live sessions now instead of - # leaving tools dangling behind a consent card for - # access the user no longer holds (#836 cross-node - # disconnect). Re-consent restores them via the - # consent-completion prime. SCHEDULED, not awaited: the - # drop waits on open_lock, which a same-key dispatch - # may hold across its entire SDK call — awaiting here - # let a token-side error stall past the sync timeout - # and charge the breaker it is documented to bypass. - self._spawn_background( - self._drop_catalog_locked((user_id, server_name)), - f"dead-grant catalog drop for '{server_name}'", - ) + self._schedule_dead_grant_drop(lookup, key) return lookup_error access_token = lookup.token or "" @@ -6442,8 +6521,6 @@ class MCPClientManager: self._cb_gate(server_name) cfg = _pool_cfg_from_row(server_row) - key = (user_id, server_name) - entry = await self._ensure_pool_entry(key) # See PoolEntryState.auth_capture for why the carrier is # entry-owned; _dispatch_pool_with_entry resets it under @@ -6457,6 +6534,7 @@ class MCPClientManager: access_token=access_token, original_name=original_name, arguments=arguments, + expected_gen=expected_gen, ) except BaseException as exc: classification = self._classify_failure(exc, capture=capture) @@ -6564,26 +6642,17 @@ class MCPClientManager: if self._app_state is None: raise RuntimeError("Pool dispatch requires set_app_state() to have been called") + # Snapshot the revocation generation before the token read — + # see _dispatch_pool. + key = (user_id, server_name) + entry = await self._ensure_pool_entry(key) + expected_gen = entry.catalog_gen lookup: TokenLookupResult = await self._pool_token_lookup( server_row, user_id, server_name, force_refresh=retry_count > 0 ) lookup_error = _pool_lookup_error(lookup, server_name, server_row) if lookup_error is not None: - if self._lookup_grant_dead(lookup): - # The grant is GONE (revoked / permanently rejected) — - # converge this node's live sessions now instead of - # leaving tools dangling behind a consent card for - # access the user no longer holds (#836 cross-node - # disconnect). Re-consent restores them via the - # consent-completion prime. SCHEDULED, not awaited: the - # drop waits on open_lock, which a same-key dispatch - # may hold across its entire SDK call — awaiting here - # let a token-side error stall past the sync timeout - # and charge the breaker it is documented to bypass. - self._spawn_background( - self._drop_catalog_locked((user_id, server_name)), - f"dead-grant catalog drop for '{server_name}'", - ) + self._schedule_dead_grant_drop(lookup, key) return lookup_error access_token = lookup.token or "" @@ -6605,8 +6674,6 @@ class MCPClientManager: self._cb_gate(server_name) cfg = _pool_cfg_from_row(server_row) - key = (user_id, server_name) - entry = await self._ensure_pool_entry(key) capture = entry.auth_capture try: sdk_result = await self._dispatch_pool_with_entry_call( @@ -6615,6 +6682,7 @@ class MCPClientManager: cfg=cfg, access_token=access_token, sdk_call=lambda s: s.read_resource(uri), + expected_gen=expected_gen, ) except BaseException as exc: classification = self._classify_failure(exc, capture=capture) @@ -6707,26 +6775,17 @@ class MCPClientManager: if self._app_state is None: raise RuntimeError("Pool dispatch requires set_app_state() to have been called") + # Snapshot the revocation generation before the token read — + # see _dispatch_pool. + key = (user_id, server_name) + entry = await self._ensure_pool_entry(key) + expected_gen = entry.catalog_gen lookup: TokenLookupResult = await self._pool_token_lookup( server_row, user_id, server_name, force_refresh=retry_count > 0 ) lookup_error = _pool_lookup_error(lookup, server_name, server_row) if lookup_error is not None: - if self._lookup_grant_dead(lookup): - # The grant is GONE (revoked / permanently rejected) — - # converge this node's live sessions now instead of - # leaving tools dangling behind a consent card for - # access the user no longer holds (#836 cross-node - # disconnect). Re-consent restores them via the - # consent-completion prime. SCHEDULED, not awaited: the - # drop waits on open_lock, which a same-key dispatch - # may hold across its entire SDK call — awaiting here - # let a token-side error stall past the sync timeout - # and charge the breaker it is documented to bypass. - self._spawn_background( - self._drop_catalog_locked((user_id, server_name)), - f"dead-grant catalog drop for '{server_name}'", - ) + self._schedule_dead_grant_drop(lookup, key) return lookup_error access_token = lookup.token or "" @@ -6748,8 +6807,6 @@ class MCPClientManager: self._cb_gate(server_name) cfg = _pool_cfg_from_row(server_row) - key = (user_id, server_name) - entry = await self._ensure_pool_entry(key) capture = entry.auth_capture try: sdk_result = await self._dispatch_pool_with_entry_call( @@ -6758,6 +6815,7 @@ class MCPClientManager: cfg=cfg, access_token=access_token, sdk_call=lambda s: s.get_prompt(original_name, arguments=arguments), + expected_gen=expected_gen, ) except BaseException as exc: classification = self._classify_failure(exc, capture=capture) @@ -6850,9 +6908,7 @@ class MCPClientManager: """ evict = self._user_pool_entries.get(key) if evict is not None: - evict.session = None - # Dead once the session is gone (see _teardown_pool_entry). - evict.bound_token = None + evict.drop_session() # Prune the push-notification debounce stamp so a # reconnect's first ``list_changed`` refreshes immediately. self._last_pool_notification_refresh.pop(key, None) @@ -6880,6 +6936,11 @@ class MCPClientManager: evict = self._user_pool_entries.get(key) if evict is None: return + # Invalidate in-flight publishers FIRST — even when the entry is + # catalog-less, a racer that read pre-revocation state (a token + # row, a live session) must find the generation moved and + # discard its publish, or it resurrects the revoked catalog. + evict.catalog_gen += 1 if not self._entry_has_catalog(evict): # Already catalog-less (repeat drop from a racing dispatch, # or a never-discovered stub) — nothing to clear, and a @@ -6894,27 +6955,48 @@ class MCPClientManager: self._rebuild_and_notify_user_catalogs(key[0]) def _lookup_grant_dead(self, lookup: TokenLookupResult) -> bool: - """Single source of truth: does this failed lookup prove the grant is GONE? + """Does this failed lookup prove the grant is GONE? - True only when the infrastructure that could know is actually - wired — the obo lookup returns kind="missing" for an - unconfigured token store / storage too, and dropping a catalog - on a boot-ordering window would strand live sessions until a - fresh prime. With the stores present, "missing" (row / - credential absent — revoked, disconnected, or unlinked) and - "refresh_failed" (the AS / mint durably rejected the grant) are - authoritative; the barely-reachable empty-token fallback maps - to the same consent_required verdict and classifies the same - way. Transient refresh failures and decrypt failures keep the + Derived from :func:`_pool_lookup_verdict` — the classification + has ONE source, so the consent card the user sees and the + catalog-drop convergence can never disagree by kind. True only + when the infrastructure that could know is actually wired: the + obo lookup returns kind="missing" for an unconfigured token + store / storage too, and dropping a catalog on a boot-ordering + window would strand live sessions until a fresh prime. + Transient refresh failures and decrypt failures keep the catalog: access still exists. """ if getattr(self._app_state, "mcp_token_store", None) is None: return False if self._storage is None: return False - if lookup.kind in ("missing", "refresh_failed"): - return True - return lookup.kind == "token" and not (lookup.token or "") + return _pool_lookup_verdict(lookup) == "mcp_consent_required" + + def _schedule_dead_grant_drop(self, lookup: TokenLookupResult, key: tuple[str, str]) -> None: + """Converge live sessions with a durably-gone grant (#836). + + Called on the mcp-loop wherever a classified lookup fails — the + three dispatchers and session-start priming. When + :meth:`_lookup_grant_dead` confirms the grant is gone (revoked, + disconnected, or unlinked — not an infrastructure blip), the + (user, server) catalog is dropped so the tools leave the user's + live sessions instead of dangling behind a consent card for + access that no longer exists. Re-consent restores them via the + consent-completion prime; obo re-login via the capture prime. + + SCHEDULED, never awaited: the locked drop can park behind a + same-key dispatch holding ``open_lock`` across its entire SDK + call — awaiting here stalled token-side errors past the sync + timeout and charged the breaker they are documented to bypass. + ``_spawn_background`` tracks the task so shutdown cancels it. + """ + if not self._lookup_grant_dead(lookup): + return + self._spawn_background( + self._drop_catalog_locked(key), + f"dead-grant catalog drop for '{key[1]}'", + ) async def _drop_catalog_locked(self, key: tuple[str, str]) -> None: """Serialize a catalog drop against an in-flight connect. @@ -7015,6 +7097,7 @@ class MCPClientManager: access_token: str, original_name: str, arguments: dict[str, Any], + expected_gen: int | None = None, ) -> str: """Hold ``entry.open_lock`` across connect-or-reuse AND ``call_tool``. @@ -7041,6 +7124,7 @@ class MCPClientManager: cfg=cfg, access_token=access_token, sdk_call=lambda s: s.call_tool(original_name, arguments), + expected_gen=expected_gen, ) return _decode_tool_result(result) @@ -7052,6 +7136,7 @@ class MCPClientManager: cfg: dict[str, Any], access_token: str, sdk_call: Callable[[Any], Awaitable[Any]], + expected_gen: int | None = None, ) -> Any: """Hold ``entry.open_lock`` across connect-or-reuse AND ``sdk_call``. @@ -7117,6 +7202,7 @@ class MCPClientManager: access_token, auth_capture=entry.auth_capture, auth_fired_event=entry.auth_fired_event, + expected_gen=expected_gen, ) session = fresh.session if session is None: @@ -7403,11 +7489,23 @@ class MCPClientManager: return key = (user_id, server_name) + async def _spawn_tracked() -> None: + # ``_spawn_background`` is loop-thread-only — hop first. + # Tracking matters: the locked drop can park behind a + # same-key dispatch's long SDK call, and an untracked task + # still pending at shutdown is abandoned with "Task was + # destroyed but it is pending!" noise; tracked tasks are + # cancelled by shutdown(). + self._spawn_background( + self._drop_catalog_locked(key), + f"revocation catalog drop for '{server_name}'", + ) + try: # Locked: an in-flight connect completing its discovery # after an unserialized drop would republish (resurrect) # the revoked catalog with nothing left to clear it. - asyncio.run_coroutine_threadsafe(self._drop_catalog_locked(key), self._loop) + asyncio.run_coroutine_threadsafe(_spawn_tracked(), self._loop) except RuntimeError as exc: log.info( "mcp_pool.evict_user_session_failed server=%s user=%s error=%s", @@ -7712,19 +7810,20 @@ def _pool_lookup_error( ) -> str | None: """Map a failed pool token lookup to its structured error; ``None`` on success. - Single copy of the lookup-kind → structured-error mapping shared by the - tool / resource / prompt dispatchers — previously three hand-synced + Single copy of the lookup-kind → structured-error RENDERING shared by + the tool / resource / prompt dispatchers — previously three hand-synced copies that every auth-model change had to edit in lockstep. Returns - ``None`` exactly when *lookup* carries a non-empty bearer. + ``None`` exactly when *lookup* carries a non-empty bearer. The + CLASSIFICATION lives in :func:`_pool_lookup_verdict`; this function + only chooses per-kind detail copy for the code it returns. """ - if lookup.kind == "missing": - return _structured_error( - code="mcp_consent_required", - server=server_name, - detail=_consent_missing_detail(server_row), - consent_url=_build_consent_url(server_row), - ) - if lookup.kind == "decrypt_failure": + # Literal ``code=`` strings in each branch (not ``code=verdict``) so the + # consent-url sibling audit (tests/test_mcp_consent_url_sibling_audit.py) + # keeps seeing these sites — a variable code is invisible to its scan. + verdict = _pool_lookup_verdict(lookup) + if verdict is None: + return None + if verdict == "mcp_token_undecryptable_key_unknown": return _structured_error( code="mcp_token_undecryptable_key_unknown", server=server_name, @@ -7733,7 +7832,7 @@ def _pool_lookup_error( "Operator action required." ), ) - if lookup.kind == "refresh_failed_transient": + if verdict == "mcp_refresh_unavailable": # Transient refresh failure (AS/network blip) — the token was kept; # a retry may succeed once the AS recovers. Retryable, NOT re-consent. return _structured_error( @@ -7741,27 +7840,45 @@ def _pool_lookup_error( server=server_name, detail="Token refresh temporarily failed; please retry.", ) + # verdict == "mcp_consent_required": pick the detail by kind. if lookup.kind == "refresh_failed": # The ``mcp_server.oauth.token_revoked`` audit was already emitted by # the classified lookup when it deleted the row; no second audit here. - return _structured_error( - code="mcp_consent_required", - server=server_name, - detail=_pool_error_detail(server_row, "refresh_failed"), - consent_url=_build_consent_url(server_row), - ) + detail = _pool_error_detail(server_row, "refresh_failed") + else: + # kind == "missing", or the barely-reachable empty-token fallback — + # auth-model-aware so an obo row never shows per-server-consent copy + # + a null consent_url. + detail = _consent_missing_detail(server_row) + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail=detail, + consent_url=_build_consent_url(server_row), + ) + + +def _pool_lookup_verdict(lookup: TokenLookupResult) -> str | None: + """Classify a pool token lookup outcome as a structured-error code. + + THE single classification of failed lookups: ``_pool_lookup_error`` + renders it, and ``MCPClientManager._lookup_grant_dead`` derives the + catalog-drop decision from it (``mcp_consent_required`` == the grant + is gone) — so the card the user sees and the convergence behavior + can never disagree by kind. Returns ``None`` exactly when *lookup* + carries a non-empty bearer. + """ + if lookup.kind == "missing": + return "mcp_consent_required" + if lookup.kind == "decrypt_failure": + return "mcp_token_undecryptable_key_unknown" + if lookup.kind == "refresh_failed_transient": + return "mcp_refresh_unavailable" + if lookup.kind == "refresh_failed": + return "mcp_consent_required" # kind == "token" if not (lookup.token or ""): - # A classified lookup returns kind=="token" only with a real token, so - # this empty-token fallback is barely reachable — but keep it - # auth-model-aware like the sibling ``missing`` branch above, so an obo - # row never shows per-server-consent copy + a null consent_url. - return _structured_error( - code="mcp_consent_required", - server=server_name, - detail=_consent_missing_detail(server_row), - consent_url=_build_consent_url(server_row), - ) + return "mcp_consent_required" return None diff --git a/turnstone/core/session.py b/turnstone/core/session.py index a4eb8b1a..bce28607 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1700,6 +1700,7 @@ class ChatSession: # merged read runs after the registrations below. self._tools = list(INTERACTIVE_TOOLS) self._task_tools = list(TASK_AGENT_TOOLS) + self._mcp_tools_change_seq = 0 # Register for tool-change notifications from MCP servers. # ``user_id`` is the listener identity component — pool-only # changes for OTHER users must not fire this callback. @@ -1719,9 +1720,16 @@ class ChatSession: # catalog change firing earlier than this fans out to the # listeners just registered, so nothing can slip between # read and register with its only notification unheard. + seq_before = self._mcp_tools_change_seq mcp_tools = self._mcp_client.get_tools(user_id=self._mcp_user_id) self._tools = merge_mcp_tools(INTERACTIVE_TOOLS, mcp_tools) self._task_tools = merge_mcp_tools(TASK_AGENT_TOOLS, mcp_tools) + if self._mcp_tools_change_seq != seq_before: + # The mirror race: a listener callback fired between the + # read above and these assignments — its fresher merge + # was just clobbered by our staler snapshot. Re-running + # the callback converges on the current maps. + self._on_mcp_tools_changed() # Proactively warm this user's per-user OAuth (oauth_user) pools so # their tools are present without a manual reconnect (e.g. after a # reboot/upgrade, or right after consent). Fire-and-forget — the @@ -2475,6 +2483,12 @@ class ChatSession: # is fixed at COORDINATOR_TOOLS. Ignore MCP server changes. if self._kind == WorkstreamKind.COORDINATOR: return + # Monotonic change marker: the constructor snapshots this around + # its authoritative post-registration read and re-runs this + # callback if it advanced — otherwise a notification landing + # between that read and its assignments is clobbered by the + # staler snapshot (its only notification already consumed). + self._mcp_tools_change_seq += 1 # Pass the effective user_id (acting user on shared workstreams, # owner otherwise) so the merged tool list includes that user's # pool catalog. The static path is included by ``get_tools``