diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index 33bc8dbb..7700abde 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -642,7 +642,7 @@ class TestEviction: entry.tools = _fake_pool_tools("pool-srv", "t") mgr._rebuild_user_tool_map("u0") entry.session = MagicMock() # warm BEFORE the failed lookup - mgr._schedule_dead_grant_drop(dead, key) + mgr._schedule_dead_grant_drop(dead, key, observed_session=entry.session) _run_on_loop(loop, _warm_then_schedule()) _drain_background(mgr, loop) @@ -660,7 +660,7 @@ class TestEviction: fresh.tools = _fake_pool_tools("pool-srv", "t") mgr._rebuild_user_tool_map("u0") fresh.session = MagicMock() # observed at schedule time - mgr._schedule_dead_grant_drop(dead, key) + mgr._schedule_dead_grant_drop(dead, key, observed_session=fresh.session) fresh.session = MagicMock() # re-consent connect lands first _run_on_loop(loop, _reseed_and_swap()) @@ -668,6 +668,34 @@ class TestEviction: entry = mgr._user_pool_entries[key] assert entry.tools is not None, "drop must spare a session newer than its observation" + def test_dead_grant_drop_skips_when_nothing_to_converge(self, running_loop_mgr) -> None: + """At scale every unconsented server classifies as a dead grant + on every prime; spawning a tracked no-op drop task per (user, + unconsented server) is pure mcp-loop overhead. The schedule-side + guard skips when there is no entry, or a session-less entry with + no catalog — a catalog appearing later implies a lookup that + succeeded after this one failed.""" + mgr, loop, _ = running_loop_mgr + mgr._app_state = SimpleNamespace(mcp_token_store=object()) + mgr._storage = object() + dead = SimpleNamespace(kind="missing", token=None) + key = ("u0", "pool-srv") + + async def _schedule_no_entry() -> int: + before = len(mgr._background_tasks) + mgr._schedule_dead_grant_drop(dead, key, observed_session=None) + return len(mgr._background_tasks) - before + + assert _run_on_loop(loop, _schedule_no_entry()) == 0 + + async def _schedule_bare_stub() -> int: + await mgr._ensure_pool_entry(key) + before = len(mgr._background_tasks) + mgr._schedule_dead_grant_drop(dead, key, observed_session=None) + return len(mgr._background_tasks) - before + + assert _run_on_loop(loop, _schedule_bare_stub()) == 0 + def test_notification_refresh_failure_keeps_debounce_stamp(self, running_loop_mgr) -> None: """A failed spawned refresh must KEEP the debounce stamp: popping it re-arms the handler on every notification, so a @@ -679,7 +707,12 @@ class TestEviction: notification refreshes immediately).""" mgr, loop, _ = running_loop_mgr key = ("u0", "pool-srv") - _run_on_loop(loop, mgr._ensure_pool_entry(key)) + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() # runner refreshes only live sessions + + _run_on_loop(loop, _seed()) mgr._last_pool_notification_refresh[key] = 123.0 async def _boom(_key: tuple[str, str]) -> tuple[list[str], list[str]]: @@ -707,7 +740,12 @@ class TestEviction: catch) — the anyio stray-cancel shape.""" mgr, loop, _ = running_loop_mgr key = ("u0", "pool-srv") - _run_on_loop(loop, mgr._ensure_pool_entry(key)) + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() # runner refreshes only live sessions + + _run_on_loop(loop, _seed()) mgr._last_pool_notification_refresh[key] = 123.0 async def _wedge(_key: tuple[str, str]) -> tuple[list[str], list[str]]: @@ -732,6 +770,7 @@ class TestEviction: async def _scenario() -> bool: entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() # runner refreshes only live sessions await entry.open_lock.acquire() task = asyncio.ensure_future(mgr._run_notification_refresh(key, "tools", _rec)) for _ in range(3): @@ -775,6 +814,89 @@ class TestEviction: _run_on_loop(loop, _scenario()) assert refreshed == [] + def test_notification_coalesces_when_refresh_already_queued(self, running_loop_mgr) -> None: + """While a runner is queued for a key+kind (coalesce marker + set), further notifications spawn NOTHING — the parked runner's + fresh list observes their change when it acquires the lock. + This bounds the ``open_lock`` waiter queue at one parked runner + per key+kind: without it a notifying-but-slow server accretes + waiters (admitted 1/5s, drained 1/30s) that starve same-key + dispatches past their wall-clock budget and starve idle + eviction, while ``_background_tasks`` grows without bound.""" + mgr, loop, _ = running_loop_mgr + key = ("u0", "pool-srv") + refreshed: list[tuple[str, str]] = [] + + async def _fake_refresh(k: tuple[str, str]) -> tuple[list[str], list[str]]: + refreshed.append(k) + return [], [] + + mgr._refresh_pool_server_tools = _fake_refresh # type: ignore[method-assign] + handler = mgr._make_pool_notification_handler(key) + + async def _fire_twice() -> int: + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() + await entry.open_lock.acquire() # park the first runner + note = mcp_types.ServerNotification( + mcp_types.ToolListChangedNotification(method="notifications/tools/list_changed") + ) + before = len(mgr._background_tasks) + await handler(note) + # Age the stamp past the debounce window: the MARKER (not + # the stamp) must do the suppression for the second fire. + mgr._last_pool_notification_refresh[key] = time.monotonic() - 10.0 + await handler(note) + spawned = len(mgr._background_tasks) - before + entry.open_lock.release() + return spawned + + spawned = _run_on_loop(loop, _fire_twice()) + _drain_background(mgr, loop) + assert spawned == 1, "second notification must coalesce into the parked runner" + assert refreshed == [key] + # The runner released its marker (at lock acquire + finally). + assert not mgr._pool_refresh_pending + + def test_teardown_pool_entry_pops_debounce_stamp(self, running_loop_mgr) -> None: + """Every teardown path must pop the debounce stamp — the + keep-stamp-on-failure design leans on it: a reconnect's first + ``list_changed`` refreshes immediately, so a change announced + in a failed window converges at the next reconnect.""" + mgr, loop, _ = running_loop_mgr + key = ("u0", "pool-srv") + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() + + _run_on_loop(loop, _seed()) + mgr._last_pool_notification_refresh[key] = 123.0 + _run_on_loop(loop, mgr._teardown_pool_entry(key)) + assert key not in mgr._last_pool_notification_refresh + + def test_owner_death_pops_debounce_stamp(self, running_loop_mgr) -> None: + """The unrequested-collapse path (owner done-callback) is a + teardown too: the stamp must not outlive the transport.""" + mgr, loop, _ = running_loop_mgr + key = ("u0", "pool-srv") + + async def _scenario() -> None: + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() + + async def _noop() -> None: + pass + + task = asyncio.ensure_future(_noop()) + await task + entry.owner_task = task + mgr._last_pool_notification_refresh[key] = 123.0 + mgr._on_pool_owner_death(key, task) + + _run_on_loop(loop, _scenario()) + assert key not in mgr._last_pool_notification_refresh + def test_list_changed_handler_spawns_refresh_off_receive_loop(self, running_loop_mgr) -> None: """The notification handler must SPAWN the refresh, not await it: the SDK awaits handlers inline in its receive loop, so an @@ -793,8 +915,10 @@ class TestEviction: async def _fire() -> bool: # The runner refreshes through the entry's ``open_lock``; - # without an entry it (correctly) discards the notification. - await mgr._ensure_pool_entry(key) + # without an entry (or a live session) it correctly + # discards the notification. + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() note = mcp_types.ServerNotification( mcp_types.ToolListChangedNotification(method="notifications/tools/list_changed") ) @@ -1038,6 +1162,55 @@ class TestDispatchStateMachine: assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is False assert fired[0] == 1 + def test_mid_lookup_reconsent_survives_dead_grant_drop( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """The dead-grant observation is snapshotted BEFORE the + classified lookup's awaits: a consent-completion prime that + connects and publishes DURING the lookup (its awaits can park on + executor hops) must read as re-consent evidence, not as the + pre-observation session — or the scheduled drop evicts the + just-restored catalog with no remaining re-prime path.""" + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + self._wire_pool(mgr, storage, cipher) + self._seed_cooled_catalog(mgr, loop) + key = ("user-1", "pool-srv") + fresh_session = MagicMock() + + async def _lookup_with_race( + server_row: dict[str, Any], + user_id: str, + server_name: str, + force_refresh: bool = False, + ) -> Any: + # The consent callback completes mid-lookup: prime connects + # and publishes the fresh catalog before the lookup returns + # its (now stale) dead-grant verdict. + entry = mgr._user_pool_entries[key] + entry.session = fresh_session + entry.tools = _fake_pool_tools("pool-srv", "do_thing") + mgr._rebuild_user_tool_map("user-1") + return SimpleNamespace(kind="missing", token=None) + + mgr._pool_token_lookup = _lookup_with_race # type: ignore[method-assign] + + with pytest.raises(RuntimeError) as exc_info: + mgr.call_tool_sync( + "mcp__pool-srv__do_thing", + {}, + user_id="user-1", + timeout=5, + ) + payload = json.loads(str(exc_info.value)) + assert payload["error"]["code"] == "mcp_consent_required" + _drain_background(mgr, loop) + entry = mgr._user_pool_entries[key] + assert entry.tools is not None, "drop must spare the mid-lookup re-consent" + assert entry.session is fresh_session + assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True + def test_decrypt_failure_does_not_emit_consent( self, running_loop_mgr, storage: SQLiteBackend ) -> None: diff --git a/turnstone/core/auth.py b/turnstone/core/auth.py index 7d1b0d45..8a5e99b7 100644 --- a/turnstone/core/auth.py +++ b/turnstone/core/auth.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from turnstone.core.oidc import OIDCConfig from turnstone.core.log import get_logger +from turnstone.core.mcp_client import try_prime_user_pools from turnstone.core.oidc import ( OIDC_STATE_TTL_SECONDS, OIDCError, @@ -2160,25 +2161,15 @@ async def handle_oidc_callback(request: Request, audience: str, cookie_name: str # 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") - and hasattr(mcp_client, "has_live_session_listener") - ): - # The gate call sits INSIDE the try: nothing on this - # best-effort path may fail the login (the credential - # is already persisted; the JWT is not yet issued). - try: - if mcp_client.has_live_session_listener(user["user_id"]): - mcp_client.prime_user_pools(user["user_id"]) - except Exception: - log.debug( - "oidc.capture: post-capture pool prime scheduling failed user=%s", - user["user_id"], - exc_info=True, - ) + # Fire-and-forget; a failure changes nothing about login + # (the credential is already persisted; the JWT is not + # yet issued) — the helper owns the swallow. + try_prime_user_pools( + getattr(request.app.state, "mcp_client", None), + user["user_id"], + require_live_listener=True, + context="oidc-capture", + ) # Load permissions and issue Turnstone JWT perms = await asyncio.to_thread(_load_user_permissions, storage, user["user_id"]) diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index df7dfd9c..81bf06c4 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -148,18 +148,59 @@ _MAX_PROMPTS_PER_SERVER = 1000 # thundering herd of connects on every session start. _PRIME_MAX_CONCURRENCY = 4 -# One row per pool ``*/list_changed`` kind: notification type → (kind label, -# refresh-method name). The pool notification handler drives all three catalog -# kinds through this table so the debounce-stamp/log/spawn protocol exists -# exactly once — a protocol change cannot silently diverge between kinds. -# Exact-type lookup is safe: the SDK's discriminated-union parser instantiates -# these concrete classes, never subclasses. -_POOL_LIST_CHANGED_REFRESHERS: dict[type, tuple[str, str]] = { - mcp_types.ToolListChangedNotification: ("tools", "_refresh_pool_server_tools"), - mcp_types.ResourceListChangedNotification: ("resources", "_refresh_pool_server_resources"), - mcp_types.PromptListChangedNotification: ("prompts", "_refresh_pool_server_prompts"), +# One row per pool ``*/list_changed`` kind: notification type → kind label. +# The pool notification handler drives all three catalog kinds through this +# table (plus its kind → bound-method map, which stays mypy-checked attribute +# access rather than a name-string table) so the debounce/coalesce/log/spawn +# protocol exists exactly once — a protocol change cannot silently diverge +# between kinds. Exact-type lookup is safe: the SDK's discriminated-union +# parser instantiates these concrete classes, never subclasses. +_POOL_LIST_CHANGED_KINDS: dict[type, str] = { + mcp_types.ToolListChangedNotification: "tools", + mcp_types.ResourceListChangedNotification: "resources", + mcp_types.PromptListChangedNotification: "prompts", } + +def try_prime_user_pools( + mcp_client: Any, + user_id: str | None, + *, + require_live_listener: bool = False, + context: str = "prime", +) -> None: + """Best-effort per-user pool prime: schedule and swallow, never raise. + + The ONE copy of the fire-and-forget prime idiom shared by the + session-construction, acting-user-change, and OIDC capture-success + call sites — three hand-synced copies had already drifted. + ``mcp_client`` is duck-typed (session and auth tests stub it): a + client without the prime surface, or a falsy ``user_id``, is a + silent no-op, matching the guards this replaces. Failures never + propagate — a prime is an optimisation, and neither login nor + session construction may fail on it. ``require_live_listener`` + gates on an open session to heal (the OIDC capture site: priming on + every routine SSO re-login would warm transports and mint tokens + for users with nothing open, at deployment scale). + """ + if not user_id or not hasattr(mcp_client, "prime_user_pools"): + return + try: + if require_live_listener and not ( + hasattr(mcp_client, "has_live_session_listener") + and mcp_client.has_live_session_listener(user_id) + ): + return + mcp_client.prime_user_pools(user_id) + except Exception: + log.debug( + "MCP pool prime scheduling failed (%s) user=%s", + context, + user_id, + exc_info=True, + ) + + # How long a node trusts its own pending-consent DELETE before re-running it on # the next dispatch success. Bounds two things at once: the hot-path SQL rate # (at most one DELETE per (user, server) per window) and the staleness of a @@ -751,6 +792,13 @@ class MCPClientManager: # so a noisy server in one user's pool doesn't suppress a refresh # in another user's pool of the same server. self._last_pool_notification_refresh: dict[tuple[str, str], float] = {} + # Coalescing markers for spawned ``list_changed`` refreshes, keyed + # ``((user_id, server), kind)``: set at spawn, cleared the moment + # the runner acquires ``open_lock`` (before its list call). While + # set, further notifications for the key+kind are dropped — the + # parked runner's fresh list will observe their change — bounding + # the lock's waiter queue at ONE parked runner per key+kind. + self._pool_refresh_pending: set[tuple[tuple[str, str], str]] = set() # Pool tuning (read at construction; falls back to defaults). mcp_cfg = load_config("mcp") @@ -1866,10 +1914,9 @@ class MCPClientManager: ) -> None: if not isinstance(msg, mcp_types.ServerNotification): return - spec = _POOL_LIST_CHANGED_REFRESHERS.get(type(msg.root)) - if spec is None: + kind = _POOL_LIST_CHANGED_KINDS.get(type(msg.root)) + if kind is None: return - kind, refresh_name = spec now = time.monotonic() last = self._last_pool_notification_refresh.get(key, 0.0) if now - last < self._NOTIFICATION_DEBOUNCE: @@ -1880,6 +1927,21 @@ class MCPClientManager: now - last, ) return + marker = (key, kind) + if marker in self._pool_refresh_pending: + # A runner for this key+kind is queued but has not yet + # issued its list call — it will observe this change + # when it runs. Skipping bounds ``open_lock``'s waiter + # queue at one parked runner per key+kind, so a + # notifying-but-slow server cannot accrete waiters that + # starve same-key dispatches and idle eviction. + log.debug( + "Coalescing pool %s notification user=%s server=%s (refresh queued)", + kind, + user_id, + server_name, + ) + return try: # SPAWNED, never awaited: the SDK awaits notification # handlers inline in its receive loop, so a handler that @@ -1888,6 +1950,18 @@ class MCPClientManager: # parked, every in-flight call on the session stalls # with it, and the refresh only ever exits via its # timeout. The refresh runs as its own tracked task. + # + # Bound at dispatch time (not a module-level name table) + # so instance-level overrides keep working and mypy + # checks the attribute references. + refreshers: dict[ + str, + Callable[[tuple[str, str]], Awaitable[tuple[list[str], list[str]]]], + ] = { + "tools": self._refresh_pool_server_tools, + "resources": self._refresh_pool_server_resources, + "prompts": self._refresh_pool_server_prompts, + } log.info( "Received %s/list_changed from pool user=%s server=%s", kind, @@ -1895,15 +1969,18 @@ class MCPClientManager: server_name, ) self._last_pool_notification_refresh[key] = now + self._pool_refresh_pending.add(marker) self._spawn_background( - self._run_notification_refresh(key, kind, getattr(self, refresh_name)), + self._run_notification_refresh(key, kind, refreshers[kind]), f"pool {kind} refresh for '{server_name}'", ) except Exception as exc: - # Scheduling failed (loop shutting down). Structured - # fields only — ``exc_info=True`` would serialize the - # chained ``httpx.Request`` whose headers carry - # ``Authorization: Bearer ``. + # Scheduling failed (loop shutting down) — release the + # coalesce marker or this key+kind never refreshes again. + # Structured fields only — ``exc_info=True`` would + # serialize the chained ``httpx.Request`` whose headers + # carry ``Authorization: Bearer ``. + self._pool_refresh_pending.discard(marker) log.warning( "Pool refresh scheduling failed user=%s server=%s exc=%s", user_id, @@ -1928,14 +2005,26 @@ class MCPClientManager: refresh for the same key, where the slower list call can publish the older catalog last. Under the lock each refresh issues its list call only after the previous publisher - finished, so the last publish is always the freshest. The wait - is bounded: lock holders are a connect/dispatch (one SDK call) - or another refresh (capped by ``_CONNECT_TIMEOUT``). + finished, so the last publish is always the freshest. + + The coalesce marker (set by the handler at spawn) is cleared + the moment the lock is ACQUIRED, before the list call: a + notification landing during the in-flight list may announce a + change that list already missed, so it must spawn exactly one + successor — which parks behind this lock. Together with the + handler's marker check this bounds the waiter queue at one + parked runner per key+kind: a same-key dispatch waits at most + two refresh timeouts, not an unbounded runner FIFO. The + ``finally`` discard covers cancellation while parked. The same-entry recheck under the lock discards a refresh whose entry was replaced (full drop + re-create) while it waited — the notification belonged to the old transport and the - replacement published its own discovery. + replacement published its own discovery. A session evicted + while we were parked returns quietly for the same reason: + the reconnect's discovery republishes, and every teardown path + pops the debounce stamp, so the reconnected transport's first + notification refreshes immediately. The debounce stamp (set at schedule time) deliberately SURVIVES a failed refresh: popping it re-armed the handler on every @@ -1943,9 +2032,7 @@ class MCPClientManager: its notification rate. Keeping it caps attempts at one per debounce window; a change announced during the remainder of a failed window converges on the server's next ``list_changed`` - or the entry's next reconnect (connects run full discovery, and - teardown pops the stamp so a reconnect's first notification - refreshes immediately). + or the entry's next reconnect. ``BaseExceptionGroup`` is caught alongside ``Exception``: a wedged anyio transport surfaces session-op failures as groups, @@ -1953,15 +2040,30 @@ class MCPClientManager: ``_spawn_background``'s failure log, whose ``exc_info`` serializes the chained ``httpx.Request`` — headers carrying ``Authorization: Bearer ``. + + Bounded residual, accepted: a server that keeps notifying while + every list call hangs to ``_CONNECT_TIMEOUT`` keeps THIS + entry's ``open_lock`` near-continuously occupied (one active + + one parked runner), deferring idle eviction of the entry (the + eviction pass skips a contested lock). The churn ends at the + first real dispatch (whose transport failure evicts the + session, after which parked runners bail on the session check), + at server recovery, or when the notifications stop; other + entries are unaffected (per-entry lock). """ user_id, server_name = key + marker = (key, kind) entry = self._user_pool_entries.get(key) if entry is None: + self._pool_refresh_pending.discard(marker) return try: async with entry.open_lock: + self._pool_refresh_pending.discard(marker) if self._user_pool_entries.get(key) is not entry: return + if entry.session is None: + return await refresh(key) except (Exception, BaseExceptionGroup) as exc: # Structured fields only — ``exc_info`` would serialize @@ -1974,6 +2076,8 @@ class MCPClientManager: server_name, type(exc).__name__, ) + finally: + self._pool_refresh_pending.discard(marker) async def _pool_transport_owner( self, @@ -2108,6 +2212,9 @@ class MCPClientManager: if entry is None: return entry.drop_session() + # The debounce stamp must not outlive the transport: a + # reconnect's first ``list_changed`` refreshes immediately. + self._last_pool_notification_refresh.pop(key, None) owner = entry.owner_task close_requested = entry.close_requested entry.owner_task = None @@ -2157,6 +2264,9 @@ class MCPClientManager: return entry.owner_task = None entry.close_requested = None + # The debounce stamp must not outlive the transport: the + # reconnect's first ``list_changed`` refreshes immediately. + self._last_pool_notification_refresh.pop(key, None) if entry.session is not None: entry.drop_session() user_id, server_name = key @@ -2657,6 +2767,10 @@ class MCPClientManager: entry = self._user_pool_entries.get(key) if entry is not None and entry.session is not None: return # already connected — nothing to do + # Pre-lookup observation for the dead-grant drop below: at + # this point the session is None; one connected during the + # lookup's awaits must read as re-consent evidence. + observed_session = self._observed_pool_session(user_id, server_name) 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- @@ -2716,7 +2830,9 @@ class MCPClientManager: # is durably gone must drop the retained # catalog other live sessions still serve — # e.g. a disconnect made on another node. - self._schedule_dead_grant_drop(lookup, key) + self._schedule_dead_grant_drop( + lookup, key, observed_session=observed_session + ) return # not consented / no credential / refresh failed — lazy paths handle it if server_row is None: server_row = await asyncio.to_thread( @@ -2751,6 +2867,13 @@ class MCPClientManager: # kind="missing". The oauth_user branch keeps its own deferred-row # optimisation. On any read hiccup, fail open and let the # per-server mint path classify it. + # Pre-read observation for the dead-grant drops below: a + # session connected DURING the credential read (a re-login + # capture prime racing this one) must read as re-consent + # evidence at drop time, not as the pre-observation session. + observed_obo = { + name: self._observed_pool_session(user_id, name) for name in self._obo_server_names + } issuer = str(getattr(getattr(self._app_state, "oidc_config", None), "issuer", "") or "") has_credential = False if issuer: @@ -2766,9 +2889,11 @@ class MCPClientManager: if not has_credential: # The credential is GONE — fire the same dead-grant # convergence the per-server lookup would have - # (kind="missing") for any retained obo catalogs, or + # (kind="missing") for any retained obo entries, or # ghosts survive every new-session prime: this gate - # skips exactly the lookup that would drop them. + # skips exactly the lookup that would drop them (the + # schedule-side guard skips entries with nothing to + # converge). # Contract: the synthesized kind="missing" mirrors # get_obo_access_token_classified's durably-absent- # credential verdict (its missing-credential return @@ -2776,10 +2901,11 @@ class MCPClientManager: # ever splits — say a transient sub-case — update this # synthesis with it. for name in self._obo_server_names: - obo_key = (user_id, name) - obo_entry = self._user_pool_entries.get(obo_key) - if obo_entry is not None and self._entry_has_catalog(obo_entry): - self._schedule_dead_grant_drop(TokenLookupResult(kind="missing"), obo_key) + self._schedule_dead_grant_drop( + TokenLookupResult(kind="missing"), + (user_id, name), + observed_session=observed_obo.get(name), + ) prime_names -= self._obo_server_names await asyncio.gather(*(_prime_one(s) for s in list(prime_names))) @@ -3072,33 +3198,28 @@ class MCPClientManager: # — a transport-layer group must not kill the loop. log.warning("MCP pool eviction iteration failed", exc_info=True) - def _user_has_live_listener(self, user_id: str) -> bool: + def has_live_session_listener(self, user_id: str) -> bool: """True when *user_id* has a registered user-scoped tool listener. A live listener means a live ChatSession whose merged tool list - is derived from this user's pool catalog — the signal the - eviction passes use to decide cool-vs-drop (#836). Admin - (``None``) listeners don't count: they track catalog state for - operator tooling, not a user's model-visible tool list. + is derived from this user's pool catalog. The ONE liveness + predicate: the eviction passes use it to decide cool-vs-drop + (#836), and outside callers (the OIDC capture-site prime) use + it to 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. Admin (``None``) listeners don't count: they + track catalog state for operator tooling, not a user's + model-visible tool list. """ 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). The TTL pass checks retention for every idle entry every tick; - a per-entry ``_user_has_live_listener`` scan would be + a per-entry ``has_live_session_listener`` scan would be O(entries × listeners) under ``_listeners_lock`` on the mcp-loop. Snapshot staleness is bounded by one tick and benign: a listener added after the snapshot re-primes at session start @@ -3176,7 +3297,7 @@ class MCPClientManager: return False if live_uids is not None: return user_id in live_uids - return self._user_has_live_listener(user_id) + return self.has_live_session_listener(user_id) def _rebuild_and_notify_user_catalogs(self, user_id: str) -> None: """Rebuild all three per-user catalog maps, then fan out to all @@ -3320,10 +3441,6 @@ class MCPClientManager: if entry.in_flight > 0: return await self._teardown_pool_entry(key) - # Prune the push-notification debounce stamp with the - # transport either way — a future reconnect's first - # ``list_changed`` must refresh immediately. - self._last_pool_notification_refresh.pop(key, None) user_id, _server_name = key if self._retain_cooled(key, entry): # Cooled: entry + catalog stay, so the live session's @@ -3682,11 +3799,14 @@ class MCPClientManager: Mirror of :meth:`_refresh_pool_server_tools` for the resource path (RFC §3.2). Skips when ``supports_resources`` is False so a server that no longer advertises resources doesn't trigger - a list call. ``asyncio.timeout`` is mandatory — see R6. + a list call. ``asyncio.timeout`` is mandatory — a wedged server + must not hang the refresh (and the lock it holds) forever. - MUST run on the mcp-loop. Caller does not need to hold - ``open_lock``; this is invoked from the pool notification - handler running on the SDK's receive task. + MUST run on the mcp-loop, with the entry's ``open_lock`` HELD + (:meth:`_run_notification_refresh` acquires it): an unlocked + refresh races a connect's discovery wiring and sibling + refreshes, either of which can publish an older catalog over + this one — see the runner's ordering rationale. """ entry = self._user_pool_entries.get(key) if entry is None or entry.session is None or not entry.supports_resources: @@ -3757,11 +3877,14 @@ class MCPClientManager: Mirror of :meth:`_refresh_pool_server_tools` for the prompt path (RFC §3.3). Skips when ``supports_prompts`` is False so a server that no longer advertises prompts doesn't trigger a list - call. ``asyncio.timeout`` is mandatory — see R6. + call. ``asyncio.timeout`` is mandatory — a wedged server must + not hang the refresh (and the lock it holds) forever. - MUST run on the mcp-loop. Caller does not need to hold - ``open_lock``; this is invoked from the pool notification - handler running on the SDK's receive task. + MUST run on the mcp-loop, with the entry's ``open_lock`` HELD + (:meth:`_run_notification_refresh` acquires it): an unlocked + refresh races a connect's discovery wiring and sibling + refreshes, either of which can publish an older catalog over + this one — see the runner's ordering rationale. """ entry = self._user_pool_entries.get(key) if entry is None or entry.session is None or not entry.supports_prompts: @@ -6525,10 +6648,16 @@ 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. + # Observe the session BEFORE the lookup's awaits: a session the + # consent-completion prime connects DURING the lookup must read + # as re-consent evidence at drop time, not as pre-observation. + observed_session = self._observed_pool_session(user_id, server_name) lookup: TokenLookupResult = await self._pool_token_lookup( server_row, user_id, server_name, force_refresh=retry_count > 0 ) - lookup_error = self._pool_lookup_failure(lookup, server_name, server_row, user_id) + lookup_error = self._pool_lookup_failure( + lookup, server_name, server_row, user_id, observed_session=observed_session + ) if lookup_error is not None: return lookup_error access_token = lookup.token or "" @@ -6671,10 +6800,16 @@ class MCPClientManager: if self._app_state is None: raise RuntimeError("Pool dispatch requires set_app_state() to have been called") + # Observe the session BEFORE the lookup's awaits: a session the + # consent-completion prime connects DURING the lookup must read + # as re-consent evidence at drop time, not as pre-observation. + observed_session = self._observed_pool_session(user_id, server_name) lookup: TokenLookupResult = await self._pool_token_lookup( server_row, user_id, server_name, force_refresh=retry_count > 0 ) - lookup_error = self._pool_lookup_failure(lookup, server_name, server_row, user_id) + lookup_error = self._pool_lookup_failure( + lookup, server_name, server_row, user_id, observed_session=observed_session + ) if lookup_error is not None: return lookup_error access_token = lookup.token or "" @@ -6795,10 +6930,16 @@ class MCPClientManager: if self._app_state is None: raise RuntimeError("Pool dispatch requires set_app_state() to have been called") + # Observe the session BEFORE the lookup's awaits: a session the + # consent-completion prime connects DURING the lookup must read + # as re-consent evidence at drop time, not as pre-observation. + observed_session = self._observed_pool_session(user_id, server_name) lookup: TokenLookupResult = await self._pool_token_lookup( server_row, user_id, server_name, force_refresh=retry_count > 0 ) - lookup_error = self._pool_lookup_failure(lookup, server_name, server_row, user_id) + lookup_error = self._pool_lookup_failure( + lookup, server_name, server_row, user_id, observed_session=observed_session + ) if lookup_error is not None: return lookup_error access_token = lookup.token or "" @@ -6991,12 +7132,28 @@ class MCPClientManager: return False return _pool_lookup_verdict(lookup) == "mcp_consent_required" + def _observed_pool_session(self, user_id: str, server_name: str) -> Any: + """Snapshot the entry's session for a dead-grant observation. + + MUST be called BEFORE the classified lookup's first await (the + callers' contract with :meth:`_schedule_dead_grant_drop`): a + snapshot taken after those awaits can capture a session the + consent-completion prime connected mid-lookup, and the drop + would then evict the just-restored catalog it was designed to + spare — the session has to predate the failure to count as the + pre-observation transport. + """ + entry = self._user_pool_entries.get((user_id, server_name)) + return entry.session if entry is not None else None + def _pool_lookup_failure( self, lookup: TokenLookupResult, server_name: str, server_row: dict[str, Any], user_id: str, + *, + observed_session: Any, ) -> str | None: """Render a failed pool lookup AND pair it with its convergence drop. @@ -7005,24 +7162,34 @@ class MCPClientManager: would keep offering revoked tools behind a consent card — the cross-node non-convergence #836 fixes, silently split by dispatch surface. Returns ``None`` exactly when *lookup* - carries a usable bearer. + carries a usable bearer. ``observed_session`` is the caller's + PRE-lookup snapshot (:meth:`_observed_pool_session`). """ lookup_error = _pool_lookup_error(lookup, server_name, server_row) if lookup_error is not None: - self._schedule_dead_grant_drop(lookup, (user_id, server_name)) + self._schedule_dead_grant_drop( + lookup, (user_id, server_name), observed_session=observed_session + ) return lookup_error - def _schedule_dead_grant_drop(self, lookup: TokenLookupResult, key: tuple[str, str]) -> None: + def _schedule_dead_grant_drop( + self, + lookup: TokenLookupResult, + key: tuple[str, str], + *, + observed_session: Any, + ) -> None: """Converge live sessions with a durably-gone grant (#836). Called on the mcp-loop wherever a classified lookup fails — the - three dispatchers and session-start priming. When - :meth:`_lookup_grant_dead` confirms the grant is gone (revoked, - disconnected, or unlinked — not an infrastructure blip), the - (user, server) catalog is dropped so the tools leave the user's - live sessions instead of dangling behind a consent card for - access that no longer exists. Re-consent restores them via the - consent-completion prime; obo re-login via the capture prime. + three dispatchers, session-start priming, and the obo + credential gate. When :meth:`_lookup_grant_dead` confirms the + grant is gone (revoked, disconnected, or unlinked — not an + infrastructure blip), the (user, server) catalog is dropped so + the tools leave the user's live sessions instead of dangling + behind a consent card for access that no longer exists. + Re-consent restores them via the consent-completion prime; obo + re-login via the capture prime. SCHEDULED, never awaited: the locked drop can park behind a same-key dispatch holding ``open_lock`` across its entire SDK @@ -7030,21 +7197,35 @@ class MCPClientManager: timeout and charged the breaker they are documented to bypass. ``_spawn_background`` tracks the task so shutdown cancels it. - The entry's CURRENT session is snapshotted here, at the moment - the dead grant was observed, and handed to the drop as - ``observed_session``: a warm session that already existed when - the lookup failed predates the revocation and must be evicted - (the #836 warm case — dispatches short-circuit on the failed - lookup without ever touching the session, so no 401 path would - converge it). Only a session that CHANGED after this snapshot - proves a later successful connect, i.e. a live grant. + ``observed_session`` is REQUIRED and must be snapshotted by the + caller BEFORE the classified lookup's first await + (:meth:`_observed_pool_session`): a warm session that already + existed at observation time predates the revocation and must be + evicted (the #836 warm case — dispatches short-circuit on the + failed lookup without ever touching the session, so no 401 path + would converge it), while only a session that CHANGED after the + observation proves a later successful connect. Snapshotting + here, after the lookup returned, would capture a session the + consent-completion prime connected mid-lookup and evict the + just-restored catalog. + + Skips when there is provably nothing to converge — no entry, or + a session-less entry with no catalog. At scale every + unconsented server classifies as a dead grant on every prime; + spawning a tracked no-op task per (user, unconsented server) + for that is pure mcp-loop overhead. A catalog appearing AFTER + this check implies a successful connect, i.e. a lookup that + succeeded after this one failed. """ if not self._lookup_grant_dead(lookup): return entry = self._user_pool_entries.get(key) - observed = entry.session if entry is not None else None + if entry is None or (entry.session is None and not self._entry_has_catalog(entry)): + return self._spawn_background( - self._drop_catalog_locked(key, skip_if_connected=True, observed_session=observed), + self._drop_catalog_locked( + key, skip_if_connected=True, observed_session=observed_session + ), f"dead-grant catalog drop for '{key[1]}'", ) diff --git a/turnstone/core/session.py b/turnstone/core/session.py index d8eec52b..c0889180 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -70,6 +70,7 @@ from turnstone.core.lowering import ( tool_args_preview, wire_valid_arguments, ) +from turnstone.core.mcp_client import try_prime_user_pools from turnstone.core.memory import ( count_messages, count_structured_memories, @@ -1689,6 +1690,10 @@ class ChatSession: # listeners register so tool/resource/prompt refreshes flow # through to this session. # * interactive (no mcp) — INTERACTIVE_TOOLS. + # Construction-scoped (a local, NOT instance state): the seq + # value at the authoritative merged read, compared once at the + # end of tool setup. Only ``_mcp_tools_change_seq`` lives on. + mcp_tools_seq_at_read = 0 if kind == WorkstreamKind.COORDINATOR: self._tools = list(COORDINATOR_TOOLS) self._task_tools = [] @@ -1727,7 +1732,7 @@ class ChatSession: # attributes and be swallowed by the fan-out, losing its # effect, and one that succeeds here would be clobbered by # the tool-search construction below reading mixed state. - self._mcp_tools_seq_at_read = self._mcp_tools_change_seq + mcp_tools_seq_at_read = 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) @@ -1737,15 +1742,7 @@ class ChatSession: # listeners registered just above deliver the catalog to this # session once each prime completes. No-op for users with no # consented oauth_user servers. - if self._mcp_user_id and hasattr(self._mcp_client, "prime_user_pools"): - try: - self._mcp_client.prime_user_pools(self._mcp_user_id) - except Exception: - log.debug( - "mcp prime_user_pools scheduling failed user=%s", - self._mcp_user_id, - exc_info=True, - ) + try_prime_user_pools(self._mcp_client, self._mcp_user_id, context="session-start") else: self._tools = INTERACTIVE_TOOLS self._task_tools = TASK_AGENT_TOOLS @@ -1803,7 +1800,7 @@ class ChatSession: if ( self._kind != WorkstreamKind.COORDINATOR and self._mcp_client - and self._mcp_tools_change_seq != self._mcp_tools_seq_at_read + and self._mcp_tools_change_seq != mcp_tools_seq_at_read ): self._on_mcp_tools_changed() # Skill: explicit name overrides is_default skills. ``skill_arguments`` @@ -5652,15 +5649,7 @@ class ChatSession: mcp.remove_prompt_listener(self._mcp_prompt_cb, user_id=old_listener_uid) mcp.add_prompt_listener(self._mcp_prompt_cb, user_id=new_listener_uid) self._mcp_listener_user_id = new_listener_uid - if new_listener_uid and hasattr(mcp, "prime_user_pools"): - try: - mcp.prime_user_pools(new_listener_uid) - except Exception: - log.debug( - "mcp prime_user_pools scheduling failed user=%s", - new_listener_uid, - exc_info=True, - ) + try_prime_user_pools(mcp, new_listener_uid, context="acting-user-change") # Rebuild the merged tool list and resource/prompt-dependent # state under the new identity NOW — the prime above completes # asynchronously and only notifies on catalog changes, while