fix(mcp): strip the revocation-generation protocol; keep the stable core

Round 4 confirmed six correctness bugs, all inside round 3's
catalog_gen machinery (a ChatSession.__init__ crash from the mirror-
race re-run, an orphaned-lock race created by the ensure-before-lookup
reorder, no generation memory across entry re-creation, gen reset on
re-ensure, a raw internal error surfacing to the session layer). Four
rounds of evidence: hardening this event-driven subsystem with new
concurrency machinery breeds interaction bugs about as fast as it
closes cosmetic races. Decision: remove the protocol, keep the core.

Stripped: PoolEntryState.catalog_gen, the expected_gen threading
through dispatch/prime/connect, the dispatcher ensure-before-lookup
reorders, _PoolGrantRevokedError, and the refresh gen-guards (the
entry-identity check stays — it protects against entry replacement
with no protocol). The publisher-suspended-across-a-drop races those
closed are now ACCEPTED RESIDUALS, documented at
_evict_session_drop_catalog: the ghost self-heals at next use via the
dead-grant drop (dispatch AND priming), and a reconnected stale bearer
dies at access-token expiry — the same bound every warm session
already rides at revocation time.

Kept from round 3 (stable, orthogonal): staged discovery publication,
prime-side dead-grant drops, the obo re-login prime, the single
_pool_lookup_verdict classification, drop_session() pairing, and the
tracked revocation drop task.

Fixed from round 4's orthogonal findings:
- ChatSession construction converges its tool lists with a bounded
  re-read loop instead of calling _on_mcp_tools_changed, which
  dereferences tool-search state initialized later in construction.
- _refresh_pool_server_tools gets the asyncio.timeout its resource and
  prompt siblings already had — a wedged server no longer hangs the
  notification-handler task.
- The OIDC capture-site prime is gated on the user having a live
  session listener (new public has_live_session_listener): routine SSO
  re-logins with nothing open no longer fan out mints and connects.
- _pool_lookup_verdict returns a Literal so a typo'd verdict
  comparison fails mypy instead of silently never matching.
- The triplicated double-401 comment blocks shrink to two-liners
  pointing at the single rationale in _evict_session's docstring.
This commit is contained in:
Patrick Buckley
2026-07-13 14:44:54 -07:00
parent 1e84f62619
commit 1c4761971c
6 changed files with 149 additions and 221 deletions
+7 -38
View File
@@ -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
+15 -40
View File
@@ -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"))
+41 -6
View File
@@ -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)
+10 -3
View File
@@ -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:
+63 -124
View File
@@ -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``
+13 -10
View File
@@ -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