diff --git a/tests/test_mcp_pool_auth_introspection.py b/tests/test_mcp_pool_auth_introspection.py index 6dc2e2ab..e61b08ad 100644 --- a/tests/test_mcp_pool_auth_introspection.py +++ b/tests/test_mcp_pool_auth_introspection.py @@ -571,7 +571,6 @@ 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() @@ -981,7 +980,6 @@ 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() @@ -1602,11 +1600,7 @@ 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, - expected_gen: Any = None, + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str ) -> int: primed.append((key, token)) return 3 @@ -1637,11 +1631,7 @@ 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, - expected_gen: Any = None, + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str ) -> int: primed.append((key, token)) return 3 @@ -1681,11 +1671,7 @@ class TestPoolPrimingAndTokenRotation: primed: list[tuple[str, str]] = [] async def _fake_prime( - self_inner: MCPClientManager, - key: tuple[str, str], - cfg: dict[str, Any], - token: str, - expected_gen: Any = None, + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str ) -> int: primed.append(key) return 0 @@ -1820,11 +1806,7 @@ class TestPoolPrimingAndTokenRotation: primed: list[tuple[str, str]] = [] async def _fake_prime( - self_inner: MCPClientManager, - key: tuple[str, str], - cfg: dict[str, Any], - token: str, - expected_gen: Any = None, + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str ) -> int: primed.append(key) return 0 @@ -1876,11 +1858,7 @@ class TestPoolPrimingAndTokenRotation: done = threading.Event() async def _fake_prime( - self_inner: MCPClientManager, - key: tuple[str, str], - cfg: dict[str, Any], - token: str, - expected_gen: Any = None, + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str ) -> int: captured["key"] = key captured["token"] = token @@ -1940,7 +1918,6 @@ 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) @@ -1977,11 +1954,7 @@ class TestPoolPrimingAndTokenRotation: primed: list[tuple[str, str]] = [] async def _fake_prime( - self_inner: MCPClientManager, - key: tuple[str, str], - cfg: dict[str, Any], - token: str, - expected_gen: Any = None, + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str ) -> int: primed.append(key) return 0 @@ -2005,11 +1978,7 @@ class TestPoolPrimingAndTokenRotation: self._wire(mgr, storage, cipher) async def _fake_prime( - self_inner: MCPClientManager, - key: tuple[str, str], - cfg: dict[str, Any], - token: str, - expected_gen: Any = None, + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str ) -> int: return 1 diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index 2c585102..17ea7601 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -521,61 +521,38 @@ 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.""" + def test_refresh_discards_result_when_entry_replaced(self, running_loop_mgr) -> None: + """A ``list_changed`` refresh whose await straddles a full drop + and re-creation must DISCARD its result — it belongs to the old + entry object and would clobber the replacement's state.""" mgr, loop, _ = running_loop_mgr mgr._oauth_user_server_names = {"pool-srv"} key = ("u0", "pool-srv") - async def _scenario() -> tuple[list, list, Any, bool]: + async def _scenario() -> tuple[list, list, Any]: 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) + # The entry is fully dropped and re-created while + # list_tools is in flight. + mgr._user_pool_entries.pop(key, None) + await mgr._ensure_pool_entry(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 + fresh = mgr._user_pool_entries[key] + return added, removed, fresh.tools - added, removed, tools, in_map = _run_on_loop(loop, _scenario()) + added, removed, fresh_tools = _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 + # The stale result was not published onto the replacement entry. + assert fresh_tools is None def test_lookup_grant_dead_requires_wired_infrastructure(self, running_loop_mgr) -> None: """kind='missing' is authoritative only when the stores that @@ -1425,9 +1402,7 @@ class TestOboPriming: warmed: list[tuple[Any, Any, str]] = [] - async def _fake_prime_server( - key: Any, cfg: Any, token: str, expected_gen: Any = None - ) -> None: + async def _fake_prime_server(key: Any, cfg: Any, token: str) -> None: warmed.append((key, cfg, token)) obo_lookup = AsyncMock(return_value=SimpleNamespace(kind="token", token="minted-at")) diff --git a/tests/test_oidc_handlers.py b/tests/test_oidc_handlers.py index 0d8de0b8..44c94631 100644 --- a/tests/test_oidc_handlers.py +++ b/tests/test_oidc_handlers.py @@ -845,14 +845,15 @@ class TestOIDCCallbackCapture: 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.""" + credential capture for a user with a LIVE session schedules a + pool prime so a previously dropped obo catalog returns to their + open workstreams — obo has no consent flow, so nothing else + re-primes them 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 + prime_user_pools=primed.append, + has_live_session_listener=lambda _uid: True, ) resp = self._login( client, @@ -881,7 +882,8 @@ class TestOIDCCallbackCapture: 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 + prime_user_pools=primed.append, + has_live_session_listener=lambda _uid: True, ) resp = self._login( client, @@ -894,6 +896,39 @@ class TestOIDCCallbackCapture: 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) + def test_capture_without_live_session_does_not_prime( + self, + mock_exchange: AsyncMock, + mock_validate: Any, + mock_provision: Any, + storage: SQLiteBackend, + oidc_config: OIDCConfig, + ) -> None: + """Routine SSO re-login with nothing open must not fan out pool + warms — the prime exists to heal LIVE sessions only.""" + 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, + has_live_session_listener=lambda _uid: False, + ) + 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 + # Credential captured, but no live session → no prime. + assert store is not None + assert store.get_oidc_credential("test-admin", cfg.issuer) is not None + 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 1f04e899..4317fd3d 100644 --- a/turnstone/core/auth.py +++ b/turnstone/core/auth.py @@ -2157,10 +2157,17 @@ async def handle_oidc_callback(request: Request, audience: str, cookie_name: str # 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. + # their listeners. Gated on a LIVE session: routine SSO + # re-logins by users with nothing open must not fan out + # mints and transport connects at deployment scale. + # 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"): + if ( + mcp_client is not None + and hasattr(mcp_client, "prime_user_pools") + and hasattr(mcp_client, "has_live_session_listener") + and mcp_client.has_live_session_listener(user["user_id"]) + ): try: mcp_client.prime_user_pools(user["user_id"]) except Exception: diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 7702ba12..c1c22002 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -278,17 +278,6 @@ 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: @@ -549,14 +538,6 @@ 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 @@ -2169,7 +2150,6 @@ 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. @@ -2205,15 +2185,6 @@ 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 @@ -2472,11 +2443,7 @@ class MCPClientManager: # -- pool priming --------------------------------------------------------- async def _prime_user_server( - self, - key: tuple[str, str], - cfg: dict[str, Any], - access_token: str, - expected_gen: int | None = None, + self, key: tuple[str, str], cfg: dict[str, Any], access_token: str ) -> int: """Proactively connect a pool entry so its catalog populates into ``get_tools(user_id)`` WITHOUT waiting for a tool dispatch. @@ -2506,7 +2473,6 @@ 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 []) @@ -2557,19 +2523,8 @@ class MCPClientManager: failure must not change the user-observable consent outcome. """ try: - # 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) + count = await self._prime_user_server(key, cfg, access_token) 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", @@ -2634,10 +2589,6 @@ 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- @@ -2706,18 +2657,12 @@ class MCPClientManager: if not server_row: return cfg = _pool_cfg_from_row(server_row) - await self._prime_user_server(key, cfg, lookup.token, expected_gen) + await self._prime_user_server(key, cfg, lookup.token) 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", @@ -3055,6 +3000,16 @@ class MCPClientManager: with self._listeners_lock: return any(uid == user_id for uid, _cb in self._listeners) + def has_live_session_listener(self, user_id: str) -> bool: + """Public: does *user_id* have a live session's tool listener? + + Lets outside callers (the OIDC capture-site prime) fan out work + only for users who actually have an open session to heal — + priming on every routine SSO re-login would warm transports and + mint tokens for users with nothing open, at deployment scale. + """ + return self._user_has_live_listener(user_id) + def _live_listener_uids(self) -> set[str]: """Snapshot the user ids with a live tool listener (one lock take). @@ -3604,12 +3559,14 @@ 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. + # ``asyncio.timeout`` mandatory (R6) — a wedged server must not + # hang the notification-handler task; siblings already comply. + async with asyncio.timeout(self._CONNECT_TIMEOUT): + result = await session.list_tools() + if self._user_pool_entries.get(key) is not entry: + # The entry was replaced (full drop + re-create) while + # list_tools was in flight — this result belongs to the + # old entry; publishing it would clobber the new one. return [], [] capped = _cap_server_tools(server_name, result.tools) server_tools = [_mcp_to_openai(server_name, tool) for tool in capped] @@ -3651,7 +3608,6 @@ 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 @@ -3661,8 +3617,8 @@ 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. + if self._user_pool_entries.get(key) is not entry: + # Entry replaced mid-flight — stale result, discard. return [], [] server_resources: list[dict[str, Any]] = [] @@ -3727,12 +3683,11 @@ 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. + if self._user_pool_entries.get(key) is not entry: + # Entry replaced mid-flight — stale result, discard. return [], [] server_prompts: list[dict[str, Any]] = [] @@ -6485,19 +6440,12 @@ 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: - self._schedule_dead_grant_drop(lookup, key) + self._schedule_dead_grant_drop(lookup, (user_id, server_name)) return lookup_error access_token = lookup.token or "" @@ -6521,6 +6469,8 @@ 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 @@ -6534,7 +6484,6 @@ 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) @@ -6576,13 +6525,9 @@ class MCPClientManager: user_id, type(exc).__name__, ) - # No catalog drop here: the forced refresh SUCCEEDED, - # so the grant is alive at the AS — this second 401 is - # the resource server rejecting a fresh bearer (JWKS - # lag, audience misconfig, clock skew). Retention lets - # the next dispatch self-heal once the RS recovers - # (#836); a genuinely revoked grant converges via the - # token-lookup drop instead (the row is gone by then). + # No catalog drop: refresh SUCCEEDED, so this 401 is + # RS-side (JWKS lag / audience / skew), not a dead + # grant — see _evict_session's docstring. return _structured_error( code="mcp_consent_required", server=server_name, @@ -6642,17 +6587,12 @@ 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: - self._schedule_dead_grant_drop(lookup, key) + self._schedule_dead_grant_drop(lookup, (user_id, server_name)) return lookup_error access_token = lookup.token or "" @@ -6674,6 +6614,8 @@ 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( @@ -6682,7 +6624,6 @@ 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) @@ -6702,13 +6643,9 @@ class MCPClientManager: user_id, type(exc).__name__, ) - # No catalog drop here: the forced refresh SUCCEEDED, - # so the grant is alive at the AS — this second 401 is - # the resource server rejecting a fresh bearer (JWKS - # lag, audience misconfig, clock skew). Retention lets - # the next dispatch self-heal once the RS recovers - # (#836); a genuinely revoked grant converges via the - # token-lookup drop instead (the row is gone by then). + # No catalog drop: refresh SUCCEEDED, so this 401 is + # RS-side (JWKS lag / audience / skew), not a dead + # grant — see _evict_session's docstring. return _structured_error( code="mcp_consent_required", server=server_name, @@ -6775,17 +6712,12 @@ 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: - self._schedule_dead_grant_drop(lookup, key) + self._schedule_dead_grant_drop(lookup, (user_id, server_name)) return lookup_error access_token = lookup.token or "" @@ -6807,6 +6739,8 @@ 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( @@ -6815,7 +6749,6 @@ 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) @@ -6835,13 +6768,9 @@ class MCPClientManager: user_id, type(exc).__name__, ) - # No catalog drop here: the forced refresh SUCCEEDED, - # so the grant is alive at the AS — this second 401 is - # the resource server rejecting a fresh bearer (JWKS - # lag, audience misconfig, clock skew). Retention lets - # the next dispatch self-heal once the RS recovers - # (#836); a genuinely revoked grant converges via the - # token-lookup drop instead (the row is gone by then). + # No catalog drop: refresh SUCCEEDED, so this 401 is + # RS-side (JWKS lag / audience / skew), not a dead + # grant — see _evict_session's docstring. return _structured_error( code="mcp_consent_required", server=server_name, @@ -6931,16 +6860,23 @@ class MCPClientManager: one-cancel close protocol). Callers that can race an in-flight connect must serialize via :meth:`_drop_catalog_locked` so a completing discovery can't republish the cleared catalog. + + ACCEPTED RESIDUAL (deliberate, after a rev-generation protocol + proved buggier than the race it closed): a publisher already + suspended across this drop — a ``list_changed`` refresh in + flight, or a connect whose token was read pre-revocation and + queued on ``open_lock`` ahead of us — can republish the catalog + after the drop. The ghost is bounded and self-healing: any use + of it fails the token lookup and re-schedules this drop + (:meth:`_schedule_dead_grant_drop`, from dispatch AND priming), + and a reconnected stale bearer dies at access-token expiry — + the same bound every warm session already rides at revocation + time. """ self._evict_session(key) 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 @@ -7097,7 +7033,6 @@ 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``. @@ -7124,7 +7059,6 @@ 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) @@ -7136,7 +7070,6 @@ 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``. @@ -7202,7 +7135,6 @@ 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: @@ -7858,7 +7790,14 @@ def _pool_lookup_error( ) -def _pool_lookup_verdict(lookup: TokenLookupResult) -> str | None: +_PoolLookupVerdict = Literal[ + "mcp_consent_required", + "mcp_token_undecryptable_key_unknown", + "mcp_refresh_unavailable", +] + + +def _pool_lookup_verdict(lookup: TokenLookupResult) -> _PoolLookupVerdict | None: """Classify a pool token lookup outcome as a structured-error code. THE single classification of failed lookups: ``_pool_lookup_error`` diff --git a/turnstone/core/session.py b/turnstone/core/session.py index bce28607..b3b76796 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1720,16 +1720,19 @@ 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() + seq_before = self._mcp_tools_change_seq - 1 + while self._mcp_tools_change_seq != seq_before: + # Re-read until stable: a listener callback firing + # between a read and its assignments would have its + # fresher merge clobbered by our staler snapshot (its + # only notification already consumed). NEVER call + # _on_mcp_tools_changed here — it dereferences tool- + # search state initialized later in construction; the + # downstream init consumes these converged lists. + 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) # 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