mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
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:
@@ -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
@@ -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"))
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user