From 4d89efa7b86c16ba4d173c4b6bdc50b7bd30d83d Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 13 Jul 2026 09:49:07 -0700 Subject: [PATCH] fix(mcp): registry-liveness for cooled entries, dead-grant convergence, revoke interlock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix round for the review of the #836 catalog-retention change (14 findings: 9 correctness, 5 cleanup): - Cooled retention now requires the server to still exist in the pool registries (_retain_cooled — ONE policy shared by the TTL skip and the close path): an admin delete/disable/rename/auth-flip drops the ghost catalog within one eviction tick. Pre-#836 the idle TTL bounded such ghosts to ~10 minutes; retention made them immortal, including a disabled server that stayed dispatchable and duplicate tool names after a flip to static. - A dispatch that learns the grant is durably GONE (token row missing or refresh permanently rejected — the mcp_consent_required class) drops that (user, server) catalog, so a disconnect made on another node converges here at first touch instead of re-offering revoked tools behind a consent card. Re-consent restores the tools through the existing consent-completion single-server prime. - Revocation drops serialize against an in-flight connect via the entry's open_lock (_drop_catalog_locked): an unserialized drop was republished (resurrected) by the connect's completing discovery, with nothing left to ever clear it. - The LRU pass counts closes incrementally and the TTL pass checks a once-per-tick listener snapshot instead of scanning the listener registry per entry under its lock. - bound_token (a plaintext bearer) is cleared whenever the session is dropped — it is dead on a session-less entry, and cooling otherwise retained it for the life of the user's sessions. - Per-user status falls back to the cooled catalog for its counts and reports the idle pool separately (user_pools_idle): cooled is the steady state now, and the warm-only view said '0 tools' for a catalog the same user's chat was actively offered. - Session construction re-reads the merged tool lists after listener registration, closing the read-then-register window that missed a concurrent drop's only notification. - Dedup: one retention policy, one warm predicate, one rebuild+notify sequence (was three copies), and the drop-catalog path now layers on _evict_session instead of copying its prologue. Known limits, deliberately deferred: shared-workstream participants who are not the acting user still lose their catalogs at TTL (not a regression — the next send re-primes), and the pre-existing orphaned-lock race on full-drop is unchanged. --- tests/test_mcp_oauth_connections.py | 12 +- tests/test_mcp_user_catalog.py | 13 +- tests/test_mcp_user_pool.py | 158 ++++++++++++++- turnstone/core/mcp_client.py | 293 ++++++++++++++++++++-------- turnstone/core/session.py | 8 + 5 files changed, 392 insertions(+), 92 deletions(-) diff --git a/tests/test_mcp_oauth_connections.py b/tests/test_mcp_oauth_connections.py index 6f3ce322..84ef3193 100644 --- a/tests/test_mcp_oauth_connections.py +++ b/tests/test_mcp_oauth_connections.py @@ -806,13 +806,15 @@ class TestEvictUserSession: mgr._last_pool_notification_refresh = {} # type: ignore[attr-defined] evicted: list[tuple[str, str]] = [] - def _fake_evict(key: tuple[str, str]) -> None: + async def _fake_evict(key: tuple[str, str]) -> None: evicted.append(key) - # The revoke entry point must take the DROP-CATALOG flavor — - # the user asked for the disconnect, so their live sessions - # see the tools leave (unlike dispatch-failure eviction, #836). - mgr._evict_session_drop_catalog = _fake_evict # type: ignore[method-assign] + # The revoke entry point must take the LOCKED drop-catalog + # path — the user asked for the disconnect, so their live + # sessions see the tools leave (unlike dispatch-failure + # eviction, #836), serialized against an in-flight connect + # so a completing discovery can't resurrect the catalog. + mgr._drop_catalog_locked = _fake_evict # type: ignore[method-assign] # Run the dispatch on a separate thread so the loop can drain. import threading diff --git a/tests/test_mcp_user_catalog.py b/tests/test_mcp_user_catalog.py index e38e3e1c..04844df2 100644 --- a/tests/test_mcp_user_catalog.py +++ b/tests/test_mcp_user_catalog.py @@ -474,9 +474,10 @@ def test_evict_session_keeps_catalog_and_fires_no_listener( mgr._evict_session(("user-1", "pool-srv")) - # Session dropped; catalog RETAINED. + # Session dropped; catalog RETAINED; dead bearer copy cleared. entry = mgr._user_pool_entries[("user-1", "pool-srv")] assert entry.session is None + assert entry.bound_token is None assert entry.tools is not None # User map intact — the live session's merged tool list is untouched. assert "user-1" in mgr._user_tool_map @@ -571,6 +572,8 @@ def test_close_pool_entry_if_idle_cools_entry_for_live_session_user( _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") key = ("user-1", "pool-srv") + # Retention also requires the server to still exist as a pool server. + mgr._oauth_user_server_names = {"pool-srv"} assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True assert mgr._user_pool_entries[key].in_flight == 0 # Seed the debounce dict so the prune-on-close is observable. @@ -591,11 +594,13 @@ def test_close_pool_entry_if_idle_cools_entry_for_live_session_user( _run_on_loop(loop, mgr._close_pool_entry_if_idle(key)) - # Entry cooled, not dropped: transport gone, catalog intact. + # Entry cooled, not dropped: transport gone, catalog intact, dead + # bearer copy cleared with the transport. entry = mgr._user_pool_entries.get(key) assert entry is not None, "cooled entry must survive TTL eviction" assert entry.session is None assert entry.owner_task is None + assert entry.bound_token is None assert entry.tools is not None # The lock object must survive with the entry — an in-flight # dispatcher's next acquire needs the same lock. @@ -644,6 +649,9 @@ def test_close_pool_entry_if_idle_drops_entry_without_live_listener( _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") key = ("user-1", "pool-srv") + # Registry seeded so the DROP below is attributable to the missing + # listener alone, not to registry-liveness. + mgr._oauth_user_server_names = {"pool-srv"} assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True assert mgr._user_pool_entries[key].in_flight == 0 # Seed the debounce dict so the perf-1 prune is observable. @@ -1875,6 +1883,7 @@ def test_close_pool_entry_if_idle_cooling_keeps_resources_and_prompts( _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") key = ("user-1", "pool-srv") + mgr._oauth_user_server_names = {"pool-srv"} res_calls = [0] prompt_calls = [0] diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index fbe9c974..b1021574 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -428,16 +428,20 @@ class TestEviction: entry.tools = self._fake_tools("pool-srv", i) _run_on_loop(loop, _seed()) + # The server must exist in the pool registries for retention. + mgr._oauth_user_server_names = {"pool-srv"} # u0 has a live session (tool listener); u1 does not. mgr.add_listener(lambda: None, user_id="u0") _run_on_loop(loop, mgr._evict_idle_pool_entries()) - # u0: cooled — retained without a session, catalog intact. + # u0: cooled — retained without a session, catalog intact, and + # the dead bearer copy cleared with the transport. cooled = mgr._user_pool_entries.get(("u0", "pool-srv")) assert cooled is not None assert cooled.session is None assert cooled.tools is not None + assert cooled.bound_token is None # u1: full drop. assert ("u1", "pool-srv") not in mgr._user_pool_entries @@ -445,6 +449,106 @@ class TestEviction: _run_on_loop(loop, mgr._evict_idle_pool_entries()) assert ("u0", "pool-srv") in mgr._user_pool_entries + def test_eviction_drops_cooled_entry_when_server_leaves_registry( + self, running_loop_mgr + ) -> None: + """Registry-liveness: a cooled entry is retained ONLY while its + server still exists as a pool server. Admin delete / disable / + rename / flip-to-static all remove the name from the pool + registries at reconcile — the ghost catalog must leave live + sessions within one tick, not survive for the session's life + (pre-#836-fix the TTL bounded this to ~10 minutes).""" + mgr, loop, _ = running_loop_mgr + mgr._user_pool_idle_ttl_s = 0.0 + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("u0", "pool-srv")) + entry.session = MagicMock() + entry.tools = self._fake_tools("pool-srv", 0) + + _run_on_loop(loop, _seed()) + mgr._oauth_user_server_names = {"pool-srv"} + mgr.add_listener(lambda: None, user_id="u0") + + # While registered: cooled + retained. + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + assert ("u0", "pool-srv") in mgr._user_pool_entries + + # Auth-flip between POOL types keeps retention (still pool-backed; + # the reconcile flip self-heal re-primes the real catalog). + mgr._oauth_user_server_names = set() + mgr._obo_server_names = {"pool-srv"} + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + assert ("u0", "pool-srv") in mgr._user_pool_entries + + # Server leaves the pool registries entirely (deleted / disabled / + # renamed / flipped to static): full drop despite the live listener. + mgr._obo_server_names = set() + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + assert ("u0", "pool-srv") not in mgr._user_pool_entries + assert "u0" not in mgr._user_tool_map + + def test_drop_catalog_locked_serializes_with_inflight_connect(self, running_loop_mgr) -> None: + """A revocation drop must wait for an in-flight connect holding + ``open_lock`` — an unserialized drop is republished (resurrected) + by the connect's discovery, with nothing left to ever clear it.""" + mgr, loop, _ = running_loop_mgr + mgr._oauth_user_server_names = {"pool-srv"} + key = ("u0", "pool-srv") + + async def _scenario() -> tuple[bool, bool]: + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() + entry.tools = self._fake_tools("pool-srv", 0) + mgr._rebuild_user_tool_map("u0") + # Simulate an in-flight connect: open_lock held while the + # "discovery" publishes the catalog. + await entry.open_lock.acquire() + drop_task = asyncio.create_task(mgr._drop_catalog_locked(key)) + await asyncio.sleep(0) + # Drop is parked on the lock — catalog still published. + blocked = entry.tools is not None and not drop_task.done() + # Connect finishes its publish, then releases the lock. + mgr._rebuild_user_tool_map("u0") + entry.open_lock.release() + await drop_task + cleared = ( + entry.tools is None and entry.session is None and "u0" not in mgr._user_tool_map + ) + return blocked, cleared + + blocked, cleared = _run_on_loop(loop, _scenario()) + 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_status_reports_cooled_catalog_as_idle(self, running_loop_mgr) -> None: + """A cooled entry's catalog is still model-visible, so status + must not report '0 tools' for it: counts fall back to the + cooled catalog, ``connected`` stays transport-truthful, and the + idle pool is surfaced separately.""" + mgr, loop, _ = running_loop_mgr + mgr._oauth_user_server_names = {"pool-srv"} + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("u0", "pool-srv")) + entry.tools = self._fake_tools("pool-srv", 0) + + _run_on_loop(loop, _seed()) + + status = mgr._oauth_user_server_status("pool-srv", "u0") + assert status["connected"] is False + assert status["tools"] == 1 + assert status["user_pools"] == 0 + assert status["user_pools_idle"] == 1 + # Another user sees nothing (per-user scoping unchanged). + other = mgr._oauth_user_server_status("pool-srv", "u-other") + assert other["tools"] == 0 + assert other["user_pools_idle"] == 0 + # Aggregate operator view counts the idle catalog too. + agg = mgr._oauth_user_server_status("pool-srv", None, aggregate=True) + assert agg["tools"] == 1 + assert agg["user_pools_idle"] == 1 + def test_idle_eviction_drops_catalogless_stub_despite_live_listener( self, running_loop_mgr ) -> None: @@ -460,6 +564,7 @@ class TestEviction: await mgr._ensure_pool_entry(("u-stub", "pool-srv")) _run_on_loop(loop, _seed()) + mgr._oauth_user_server_names = {"pool-srv"} mgr.add_listener(lambda: None, user_id="u-stub") _run_on_loop(loop, mgr._evict_idle_pool_entries()) @@ -492,6 +597,7 @@ class TestEviction: mgr._user_pool_last_used[key] = base + 10 + i _run_on_loop(loop, _seed()) + mgr._oauth_user_server_names = {"pool-srv"} for i in range(3): mgr.add_listener(lambda: None, user_id=f"cool{i}") @@ -522,6 +628,7 @@ class TestEviction: mgr._user_pool_last_used[key] = base + i _run_on_loop(loop, _seed()) + mgr._oauth_user_server_names = {"pool-srv"} for i in range(3): mgr.add_listener(lambda: None, user_id=f"u{i}") @@ -556,13 +663,48 @@ class TestDispatchStateMachine: mgr.set_app_state(state) return state + def _seed_cooled_catalog(self, mgr: MCPClientManager, loop) -> list[int]: + """Seed a cooled catalog-bearing entry + live listener for user-1. + + Returns the listener's fire counter. Used by the dead-grant rows: + a dispatch that learns the grant is GONE must drop this catalog so + live sessions converge with the revocation (#836 cross-node + disconnect) instead of re-offering the revoked tools. + """ + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-1", "pool-srv")) + entry.tools = [ + { + "type": "function", + "function": { + "name": "mcp__pool-srv__do_thing", + "description": "", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + mgr._rebuild_user_tool_map("user-1") + + _run_on_loop(loop, _seed()) + mgr._oauth_user_server_names = {"pool-srv"} + fired = [0] + + def _cb() -> None: + fired[0] += 1 + + mgr.add_listener(_cb, user_id="user-1") + assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True + return fired + def test_no_token_emits_consent_required( self, running_loop_mgr, storage: SQLiteBackend ) -> None: - mgr, _loop, _ = running_loop_mgr + mgr, loop, _ = running_loop_mgr cipher = make_mcp_token_cipher() _seed_oauth_server(storage, name="pool-srv") self._wire_pool(mgr, storage, cipher) + fired = self._seed_cooled_catalog(mgr, loop) with pytest.raises(RuntimeError) as exc_info: mgr.call_tool_sync( @@ -574,6 +716,12 @@ class TestDispatchStateMachine: payload = json.loads(str(exc_info.value)) assert payload["error"]["code"] == "mcp_consent_required" assert payload["error"]["server"] == "pool-srv" + # kind="missing" is a DEAD grant: the retained catalog drops and + # the live session is notified — tools leave instead of dangling + # behind a consent card for access the user no longer holds. + assert mgr._user_pool_entries[("user-1", "pool-srv")].tools is None + assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is False + assert fired[0] == 1 def test_decrypt_failure_does_not_emit_consent( self, running_loop_mgr, storage: SQLiteBackend @@ -609,7 +757,7 @@ class TestDispatchStateMachine: assert "key_fingerprints_attempted" not in payload["error"] def test_refresh_failure_emits_consent(self, running_loop_mgr, storage: SQLiteBackend) -> None: - mgr, _loop, _ = running_loop_mgr + mgr, loop, _ = running_loop_mgr cipher = make_mcp_token_cipher() _seed_oauth_server(storage, name="pool-srv") # Seed an expired token with no refresh — the classified getter @@ -629,6 +777,7 @@ class TestDispatchStateMachine: as_issuer="https://as.example.com", audience="https://mcp.example.com", ) + fired = self._seed_cooled_catalog(mgr, loop) with pytest.raises(RuntimeError) as exc_info: mgr.call_tool_sync( @@ -639,6 +788,9 @@ class TestDispatchStateMachine: ) payload = json.loads(str(exc_info.value)) assert payload["error"]["code"] == "mcp_consent_required" + # kind="refresh_failed" is likewise a DEAD grant → catalog drops. + assert mgr._user_pool_entries[("user-1", "pool-srv")].tools is None + assert fired[0] == 1 def test_token_present_dispatches_to_session( self, running_loop_mgr, storage: SQLiteBackend diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 6dd52040..b64a8957 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -696,8 +696,8 @@ class MCPClientManager: # this via a single dict-get (atomic under GIL) so sync-thread # callers (ChatSession) never iterate ``_user_pool_entries`` # concurrently with the mcp-loop's mutations of the same dict - # (insert in ``_ensure_pool_entry`` / pop in - # ``_close_pool_entry_if_idle`` / ``_evict_session_drop_catalog``). + # (insert in ``_ensure_pool_entry`` / the SOLE pop in + # ``_close_pool_entry_if_idle``'s full-drop path). self._user_tools: dict[str, list[dict[str, Any]]] = {} # Per-user resource catalog. Mirrors ``_user_tool_map`` / @@ -2026,6 +2026,11 @@ class MCPClientManager: if entry is None: return entry.session = None + # The bearer copy is dead once the session is gone (read only + # under a live session for stale-rebind detection; overwritten + # at reconnect) — don't retain a plaintext credential on a + # cooled entry for the life of the user's sessions. + entry.bound_token = None owner = entry.owner_task close_requested = entry.close_requested entry.owner_task = None @@ -2415,16 +2420,9 @@ class MCPClientManager: # Loop-only mutation; sync-thread readers observe the new # catalog atomically via per-user dict-gets on ``_user_tools`` / - # ``_user_resources`` / ``_user_prompts``. - self._rebuild_user_tool_map(user_id) - self._rebuild_user_resource_map(user_id) - self._rebuild_user_prompt_map(user_id) - # Wake user-keyed AND admin (None) listeners; per-user fan-out - # ensures another user's session never observes this user's - # catalog change. - self._notify_user_tool_listeners(user_id) - self._notify_user_resource_listeners(user_id) - self._notify_user_prompt_listeners(user_id) + # ``_user_resources`` / ``_user_prompts``. Per-user fan-out + # ensures another user's session never observes this change. + self._rebuild_and_notify_user_catalogs(user_id) return entry # -- pool priming --------------------------------------------------------- @@ -2951,8 +2949,12 @@ class MCPClientManager: Note: the loop wakes on a fixed tick rather than a condition variable. Event-driven evictions would be more efficient on large idle pools but add complexity (tracking per-entry - deadlines + a wakeup ``asyncio.Event``); at the design cap of - 200 entries the unconditional 30 s wake is negligible. + deadlines + a wakeup ``asyncio.Event``). Warm entries are + bounded by the LRU cap (default 200); cooled catalog-only + entries are bounded by live-session users × registered pool + servers, and each tick's work over them is a per-entry set + lookup against a once-per-tick listener snapshot — negligible + either way. """ while True: try: @@ -2977,6 +2979,19 @@ class MCPClientManager: with self._listeners_lock: return any(uid == user_id for uid, _cb in self._listeners) + 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 + 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 + anyway, and one removed after it is caught next tick. + """ + with self._listeners_lock: + return {uid for uid, _cb in self._listeners if uid} + @staticmethod def _entry_has_catalog(entry: PoolEntryState) -> bool: """True when the entry carries a discovered catalog worth retaining. @@ -2989,18 +3004,70 @@ class MCPClientManager: """ return entry.tools is not None or entry.resources is not None or entry.prompts is not None - def _warm_pool_count(self) -> int: - """Count pool entries still holding connection resources. + @staticmethod + def _entry_is_warm(entry: PoolEntryState) -> bool: + """True when the entry holds connection resources. Warm = an open session OR a live owner task (an owner parked after ``_evict_session`` still holds the transport until the close protocol runs). Cooled catalog-only entries hold neither. + This is THE definition the LRU cap bounds; every warm check in + the eviction passes must go through it. """ - return sum( - 1 - for e in self._user_pool_entries.values() - if e.session is not None or e.owner_task is not None - ) + return entry.session is not None or entry.owner_task is not None + + def _retain_cooled( + self, + key: tuple[str, str], + entry: PoolEntryState, + live_uids: set[str] | None = None, + ) -> bool: + """Retention policy: keep this entry cooled for a live session? + + True iff ALL of: + - the entry carries a discovered catalog (stubs and + revoke-cleared entries retain nothing); + - the server still exists as a pool server in the in-memory + registries (:meth:`_is_pool_server`, rebuilt wholesale by + every reconcile). Without this, an admin delete / disable / + rename / auth-flip left an IMMORTAL cooled entry serving + ghost tools to live sessions — pre-#836-fix those aged out + with the idle TTL; with the registry check they full-drop + within one eviction tick, faster than before; + - the user has a live session (a registered user-scoped tool + listener). ``live_uids`` is the per-tick snapshot; ``None`` + consults the registry directly (authoritative, single key). + + Single copy of the policy — the TTL pass's already-cooled skip + and :meth:`_close_pool_entry_if_idle`'s cool-vs-drop decision + must never diverge. + """ + user_id, server_name = key + if not self._entry_has_catalog(entry): + return False + if not self._is_pool_server(server_name): + return False + if live_uids is not None: + return user_id in live_uids + return self._user_has_live_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 + three listener classes (user-keyed + admin ``None``). + + MUST run on the mcp-loop. Rebuild-before-notify is the + invariant (listeners re-read the maps), and tools / resources / + prompts move together — the Phase 7 round-2 "bug-pair" was + exactly a tools-only cleanup missing its resource/prompt half. + Shared by connect wiring, the full-drop eviction path, and the + revocation drop. + """ + self._rebuild_user_tool_map(user_id) + self._rebuild_user_resource_map(user_id) + self._rebuild_user_prompt_map(user_id) + self._notify_user_tool_listeners(user_id) + self._notify_user_resource_listeners(user_id) + self._notify_user_prompt_listeners(user_id) async def _evict_idle_pool_entries(self) -> None: """Evict pool entries past the idle TTL or above the LRU cap. @@ -3025,6 +3092,7 @@ class MCPClientManager: return now = time.monotonic() ttl = self._user_pool_idle_ttl_s + live_uids = self._live_listener_uids() # First pass: TTL-based eviction. Run closes in parallel so a tick # that needs to evict many entries doesn't block on serial teardowns. @@ -3033,16 +3101,12 @@ class MCPClientManager: last = self._user_pool_last_used.get(key, entry.last_used) if (now - last) < ttl: continue - if ( - entry.session is None - and entry.owner_task is None - and self._entry_has_catalog(entry) - and self._user_has_live_listener(key[0]) - ): + if not self._entry_is_warm(entry) and self._retain_cooled(key, entry, live_uids): # Already cooled — nothing to close. Retained for the - # live session's tool list; the tick after that user's - # last listener is removed, this stops matching and the - # entry takes the full-drop path below. + # live session's tool list; once the user's last + # listener goes away OR the server leaves the pool + # registries, this stops matching and the entry takes + # the full-drop path below. continue ttl_targets.append(key) if ttl_targets: @@ -3058,7 +3122,7 @@ class MCPClientManager: warm = [ (key, entry) for key, entry in self._user_pool_entries.items() - if entry.session is not None or entry.owner_task is not None + if self._entry_is_warm(entry) ] if len(warm) <= self._user_pool_lru_max: return @@ -3066,12 +3130,17 @@ class MCPClientManager: warm, key=lambda kv: self._user_pool_last_used.get(kv[0], kv[1].last_used), ) - # Compute the eviction batch up front; we re-check the cap after - # each close (in-flight skips can leave us still over). + # Count actual closes instead of rescanning the whole map per + # iteration; a skip (contested lock / in-flight) leaves the + # entry warm, so re-read the entry itself to learn the outcome. + over = len(warm) - self._user_pool_lru_max for key, _entry in ordered: - if self._warm_pool_count() <= self._user_pool_lru_max: + if over <= 0: break await self._close_pool_entry_if_idle(key) + current = self._user_pool_entries.get(key) + if current is None or not self._entry_is_warm(current): + over -= 1 async def _close_pool_entry_if_idle(self, key: tuple[str, str]) -> None: """Close ``key``'s transport iff its open_lock is uncontested AND in_flight==0. @@ -3080,21 +3149,21 @@ class MCPClientManager: function to return without mutation; the next eviction tick retries. - What happens to the ENTRY depends on whether the user still has - a live session (:meth:`_user_has_live_listener`) and the entry - carries a discovered catalog (:meth:`_entry_has_catalog`): + What happens to the ENTRY is :meth:`_retain_cooled` — ONE + policy shared with the TTL pass's already-cooled skip: - - live listener + catalog → the entry is COOLED: transport torn down, - catalog kept, no rebuild, no listener fan-out. The user's - merged tool list is untouched and the next dispatch or prime - reconnects — the evict-session-keep-entry shape of + - retained → the entry is COOLED: transport torn down, catalog + kept, no rebuild, no listener fan-out. The user's merged tool + list is untouched and the next dispatch or prime reconnects — + the evict-session-keep-entry shape of ``_on_pool_owner_death``. Dropping the catalog here instead silently removed the server's tools from live sessions with no re-prime path (#836). - otherwise → full drop: entry popped, per-user catalogs - rebuilt, listeners notified (reaching only admin/``None`` - listeners — operator tooling tracking catalog state). This - keeps departed users' entries from outliving their sessions. + rebuilt, listeners notified. Covers departed users (no live + listener), catalog-less stubs, and servers that left the pool + registries (deleted / disabled / renamed / auth-flipped) — + whose ghost tools must leave live sessions within a tick. """ entry = self._user_pool_entries.get(key) if entry is None: @@ -3130,7 +3199,7 @@ class MCPClientManager: # ``list_changed`` must refresh immediately. self._last_pool_notification_refresh.pop(key, None) user_id, _server_name = key - if self._entry_has_catalog(entry) and self._user_has_live_listener(user_id): + if self._retain_cooled(key, entry): # Cooled: entry + catalog stay, so the live session's # tool list and ``is_mcp_tool`` are untouched. Nothing # changed catalog-wise → no rebuild, no fan-out. The @@ -3143,14 +3212,9 @@ class MCPClientManager: # Dropping the entry without rebuilding the per-user # catalogs would leave ``is_mcp_tool`` / per-user resource & # prompt maps returning stale entries whose backing pool is - # gone for good (this user has no live session left to - # re-warm it lazily). - self._rebuild_user_tool_map(user_id) - self._rebuild_user_resource_map(user_id) - self._rebuild_user_prompt_map(user_id) - self._notify_user_tool_listeners(user_id) - self._notify_user_resource_listeners(user_id) - self._notify_user_prompt_listeners(user_id) + # gone for good (departed user, or a server no longer in + # the pool registries). + self._rebuild_and_notify_user_catalogs(user_id) evicted = True finally: lock.release() @@ -4727,18 +4791,22 @@ class MCPClientManager: # and the eviction loop snapshot the same way. entries = list(self._user_pool_entries.items()) if aggregate: - warm = [e for (uid, sname), e in entries if sname == name and e.session is not None] + scoped = [e for (uid, sname), e in entries if sname == name] elif user_id: - warm = [ - e - for (uid, sname), e in entries - if sname == name and uid == user_id and e.session is not None - ] + scoped = [e for (uid, sname), e in entries if sname == name and uid == user_id] else: - warm = [] + scoped = [] + warm = [e for e in scoped if e.session is not None] + # Cooled entries (transport idled out, catalog retained for the + # user's live sessions, #836) still SERVE their tools — a + # warm-only view told the user "0 tools" for a catalog their + # own chat was actively offered. ``connected`` stays + # transport-truthful; the catalog counts must not. + idle = [e for e in scoped if e.session is None and self._entry_has_catalog(e)] # rep is the first warm entry (insertion order): the requester's own pool - # when scoped, or a representative catalog for the aggregate operator view. - rep = warm[0] if warm else None + # when scoped, or a representative catalog for the aggregate operator + # view — falling back to a cooled catalog when nothing is warm. + rep = warm[0] if warm else (idle[0] if idle else None) cb_deadline = self._circuit_open_until.get(name) cb_open = cb_deadline is not None and time.monotonic() < cb_deadline last_refresh = self._last_refresh.get(name) @@ -4755,6 +4823,7 @@ class MCPClientManager: "consecutive_failures": self._consecutive_failures.get(name, 0), "auth_type": self.server_auth_type(name) or "oauth_user", "user_pools": len(warm), + "user_pools_idle": len(idle), "last_refresh_at": last_refresh[0] if last_refresh is not None else None, "last_refresh_outcome": last_refresh[1] if last_refresh is not None else None, } @@ -6314,6 +6383,14 @@ class MCPClientManager: ) lookup_error = _pool_lookup_error(lookup, server_name, server_row) if lookup_error is not None: + if lookup.kind in _DEAD_GRANT_LOOKUP_KINDS: + # The grant is GONE (revoked / permanently rejected) — + # converge this node's live sessions now instead of + # leaving tools dangling behind a consent card for + # access the user no longer holds (#836 cross-node + # disconnect). Re-consent restores them via the + # consent-completion prime. + await self._drop_catalog_locked((user_id, server_name)) return lookup_error access_token = lookup.token or "" @@ -6393,6 +6470,10 @@ class MCPClientManager: user_id, type(exc).__name__, ) + # Refreshed bearer also rejected — the grant is dead at + # the AS; drop the catalog so this node's live sessions + # converge (#836) instead of re-offering revoked tools. + await self._drop_catalog_locked(key) return _structured_error( code="mcp_consent_required", server=server_name, @@ -6457,6 +6538,14 @@ class MCPClientManager: ) lookup_error = _pool_lookup_error(lookup, server_name, server_row) if lookup_error is not None: + if lookup.kind in _DEAD_GRANT_LOOKUP_KINDS: + # The grant is GONE (revoked / permanently rejected) — + # converge this node's live sessions now instead of + # leaving tools dangling behind a consent card for + # access the user no longer holds (#836 cross-node + # disconnect). Re-consent restores them via the + # consent-completion prime. + await self._drop_catalog_locked((user_id, server_name)) return lookup_error access_token = lookup.token or "" @@ -6507,6 +6596,10 @@ class MCPClientManager: user_id, type(exc).__name__, ) + # Refreshed bearer also rejected — the grant is dead at + # the AS; drop the catalog so this node's live sessions + # converge (#836) instead of re-offering revoked tools. + await self._drop_catalog_locked(key) return _structured_error( code="mcp_consent_required", server=server_name, @@ -6578,6 +6671,14 @@ class MCPClientManager: ) lookup_error = _pool_lookup_error(lookup, server_name, server_row) if lookup_error is not None: + if lookup.kind in _DEAD_GRANT_LOOKUP_KINDS: + # The grant is GONE (revoked / permanently rejected) — + # converge this node's live sessions now instead of + # leaving tools dangling behind a consent card for + # access the user no longer holds (#836 cross-node + # disconnect). Re-consent restores them via the + # consent-completion prime. + await self._drop_catalog_locked((user_id, server_name)) return lookup_error access_token = lookup.token or "" @@ -6628,6 +6729,10 @@ class MCPClientManager: user_id, type(exc).__name__, ) + # Refreshed bearer also rejected — the grant is dead at + # the AS; drop the catalog so this node's live sessions + # converge (#836) instead of re-offering revoked tools. + await self._drop_catalog_locked(key) return _structured_error( code="mcp_consent_required", server=server_name, @@ -6695,6 +6800,8 @@ class MCPClientManager: evict = self._user_pool_entries.get(key) if evict is not None: evict.session = None + # Dead once the session is gone (see _teardown_pool_entry). + evict.bound_token = None # Prune the push-notification debounce stamp so a # reconnect's first ``list_changed`` refreshes immediately. self._last_pool_notification_refresh.pop(key, None) @@ -6702,37 +6809,48 @@ class MCPClientManager: def _evict_session_drop_catalog(self, key: tuple[str, str]) -> None: """Drop the cached session AND the entry's catalog contribution. - The explicit-revocation flavor of :meth:`_evict_session`, used - by :meth:`evict_user_session` (the OAuth disconnect handler): - the user asked for the server to be disconnected, so their live - sessions SHOULD see the tools leave — the opposite of the - dispatch-failure paths, where the catalog is retained (#836). + The dead-grant flavor of :meth:`_evict_session`, used when the + user's grant for this server is GONE — explicit disconnect + (:meth:`evict_user_session`) or a dispatch that failed with the + token-missing / grant-rejected class — so their live sessions + SHOULD see the tools leave rather than keep offering a card for + access that no longer exists (#836 cross-node convergence). Clearing the catalog also makes the eviction passes treat the entry as a droppable stub (``_entry_has_catalog`` is False), so it doesn't linger cooled behind a live listener. - Owner/streams left for reconnect teardown exactly as in - :meth:`_evict_session` (the one-cancel close protocol). + Session-drop prologue is delegated to :meth:`_evict_session` — + owner/streams left for reconnect teardown exactly as there (the + 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. """ + self._evict_session(key) evict = self._user_pool_entries.get(key) if evict is None: return - evict.session = None evict.tools = None evict.resources = None evict.prompts = None - self._last_pool_notification_refresh.pop(key, None) - user_id, _server_name = key - self._rebuild_user_tool_map(user_id) - self._rebuild_user_resource_map(user_id) - self._rebuild_user_prompt_map(user_id) # Wake the user's sessions so their merged tool / resource / - # prompt lists shrink now, not at the next turn boundary. Admin - # (None) listeners also fire — operator tooling tracking pool - # catalog state observes the drop. - self._notify_user_tool_listeners(user_id) - self._notify_user_resource_listeners(user_id) - self._notify_user_prompt_listeners(user_id) + # prompt lists shrink now, not at the next turn boundary. + self._rebuild_and_notify_user_catalogs(key[0]) + + async def _drop_catalog_locked(self, key: tuple[str, str]) -> None: + """Serialize a catalog drop against an in-flight connect. + + ``_evict_session_drop_catalog`` alone races + ``_connect_one_pool``: a connect already past the token check + publishes its discovery AFTER the drop, resurrecting a revoked + server's catalog with no remaining path that ever clears it. + ``open_lock`` is held for the connect/reuse window only, so the + wait is bounded by connect+discovery. Runs on the mcp-loop. + """ + entry = self._user_pool_entries.get(key) + if entry is None: + return + async with entry.open_lock: + self._evict_session_drop_catalog(key) async def _handle_auth_403( self, @@ -7206,7 +7324,10 @@ class MCPClientManager: key = (user_id, server_name) async def _do_evict() -> None: - self._evict_session_drop_catalog(key) + # Locked: an in-flight connect completing its discovery + # after an unserialized drop would republish (resurrect) + # the revoked catalog with nothing left to clear it. + await self._drop_catalog_locked(key) try: asyncio.run_coroutine_threadsafe(_do_evict(), self._loop) @@ -7509,6 +7630,14 @@ def _token_rejected_detail(server_row: dict[str, Any]) -> str: return _pool_error_detail(server_row, "token_rejected") +# Lookup kinds meaning the user's grant is durably gone (revoked / +# permanently rejected) — the dispatchers drop the (user, server) +# catalog on these so live sessions converge with the revocation +# (#836). Transient refresh failures and operator-actionable decrypt +# failures deliberately keep the catalog: access still exists. +_DEAD_GRANT_LOOKUP_KINDS: frozenset[str] = frozenset({"missing", "refresh_failed"}) + + def _pool_lookup_error( lookup: TokenLookupResult, server_name: str, server_row: dict[str, Any] ) -> str | None: diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 0bb66d40..0108871a 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1713,6 +1713,14 @@ class ChatSession: # for OTHER users do not wake this session. self._mcp_prompt_cb = self._on_mcp_prompts_changed self._mcp_client.add_prompt_listener(self._mcp_prompt_cb, user_id=self._mcp_user_id) + # Re-read now that the listeners exist: a pool-catalog change + # firing between the get_tools above and the registrations just + # made fanned out before this session could hear it, and its + # only notification is gone. One re-read converges the merged + # lists with the post-registration state. + 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