diff --git a/tests/test_mcp_oauth_connections.py b/tests/test_mcp_oauth_connections.py index f792e73b..6f3ce322 100644 --- a/tests/test_mcp_oauth_connections.py +++ b/tests/test_mcp_oauth_connections.py @@ -809,7 +809,10 @@ class TestEvictUserSession: def _fake_evict(key: tuple[str, str]) -> None: evicted.append(key) - mgr._evict_session = _fake_evict # type: ignore[method-assign] + # 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] # 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 251c627e..e38e3e1c 100644 --- a/tests/test_mcp_user_catalog.py +++ b/tests/test_mcp_user_catalog.py @@ -430,17 +430,75 @@ def test_pool_tool_visibility_user_isolation( assert mgr.is_mcp_tool("mcp__pool-srv__shared_tool", user_id="user-3") is False -def test_eviction_drops_catalog_and_fires_listener( +def test_evict_session_keeps_catalog_and_fires_no_listener( running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """``_evict_session`` clears ``entry.tools``, rebuilds the user's - map (drops the now-empty entry), and fires user + admin listeners. + """``_evict_session`` (dispatch-failure eviction: 401 / 403 / + transport blip) drops ONLY the session — the catalog and per-user + maps stay, no listener fires, and ``is_mcp_tool`` keeps resolving + the name (#836). - Verified by reverting the catalog-cleanup block in ``_evict_session`` - (drop the ``evict.tools = None`` / ``_rebuild_user_tool_map`` / - ``_notify_user_tool_listeners`` calls): the test fails because the + Clearing the catalog here silently removed the server's tools from + the user's live sessions on the first failed dispatch: the maps + rebuilt empty, the session-side ``is_mcp_tool`` gate closed, and + with no re-prime path the tools never came back — which also made + the breaker's half-open recovery and the consent / step-up cards + unreachable. + + Verified by restoring the old catalog-cleanup block in + ``_evict_session`` (``evict.tools = None`` + rebuild + notify): + this test fails because the listener fires AND ``is_mcp_tool`` + flips to False. + """ + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + list_tools_response=_list_tools_payload([_tool_spec("do_thing")]), + ) + _build_mock_transport_factory(mgr, monkeypatch, handler) + _patch_tcp_probe(mgr, monkeypatch) + + _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") + assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True + + user_calls = [0] + admin_calls = [0] + + def _user_cb() -> None: + user_calls[0] += 1 + + def _admin_cb() -> None: + admin_calls[0] += 1 + + mgr.add_listener(_user_cb, user_id="user-1") + mgr.add_listener(_admin_cb) # admin / None + + mgr._evict_session(("user-1", "pool-srv")) + + # Session dropped; catalog RETAINED. + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.session 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 + assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True + # Nothing changed catalog-wise → NO fan-out. + assert user_calls[0] == 0, f"user-keyed listener fired {user_calls[0]} times; expected 0" + assert admin_calls[0] == 0, f"admin listener fired {admin_calls[0]} times; expected 0" + + +def test_evict_user_session_drops_catalog_and_fires_listener( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """``_evict_session_drop_catalog`` (the explicit-revocation flavor + behind ``evict_user_session``) clears ``entry.tools``, rebuilds the + user's map (drops the now-empty entry), and fires user + admin + listeners — the user asked for the disconnect, so their live + sessions SHOULD see the tools leave. + + Verified by reverting the catalog-cleanup block in + ``_evict_session_drop_catalog``: the test fails because the listener never fires AND ``is_mcp_tool`` keeps returning True for - the now-evicted tool. + the now-revoked tool. """ mgr, loop, _ = running_loop_mgr handler = _make_jsonrpc_handler( @@ -453,7 +511,7 @@ def test_eviction_drops_catalog_and_fires_listener( assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True # Register one user-keyed and one admin (None) listener; both - # MUST fire on this user's eviction. + # MUST fire on this user's revocation. user_calls = [0] admin_calls = [0] other_calls = [0] @@ -471,7 +529,7 @@ def test_eviction_drops_catalog_and_fires_listener( mgr.add_listener(_admin_cb) # admin / None mgr.add_listener(_other_cb, user_id="user-2") - mgr._evict_session(("user-1", "pool-srv")) + mgr._evict_session_drop_catalog(("user-1", "pool-srv")) # Catalog cleared. entry = mgr._user_pool_entries[("user-1", "pool-srv")] @@ -479,7 +537,7 @@ def test_eviction_drops_catalog_and_fires_listener( assert entry.tools is None # User map dropped (no remaining pool entries for this user). assert "user-1" not in mgr._user_tool_map - # is_mcp_tool no longer surfaces the evicted name. + # is_mcp_tool no longer surfaces the revoked name. assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is False # Listener fan-out: matching user + admin fire, OTHER user does not. assert user_calls[0] == 1, f"user-keyed listener fired {user_calls[0]} times; expected 1" @@ -490,22 +548,19 @@ def test_eviction_drops_catalog_and_fires_listener( ) -def test_close_pool_entry_if_idle_clears_catalog_and_fires_listener( +def test_close_pool_entry_if_idle_cools_entry_for_live_session_user( running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """LRU/TTL eviction (``_close_pool_entry_if_idle``) mirrors - ``_evict_session``'s catalog-cleanup contract: drops the entry, - prunes the notification-debounce dict, rebuilds the user's tool - map, and fires user + admin listeners. + """TTL/LRU eviction COOLS the entry when the user has a live + session (a registered user-scoped tool listener): the transport is + torn down but the entry, its catalog, the per-user maps, and the + entry's ``open_lock`` all survive, and NO listener fires — the live + session's model-visible tool list must not shrink because the user + went 10 minutes without an MCP dispatch (#836). - Phase 7 round-2 review hardening (round2-1): the bug-2 fix added - the catalog-cleanup block to this method but no integration test - drove it — exactly the failure mode flagged in - ``feedback_tests_through_boundaries.md``. Negative test: drop the - ``_rebuild_user_tool_map`` / ``_notify_user_tool_listeners`` calls - from ``_close_pool_entry_if_idle``'s post-pop block; this test - fails because ``is_mcp_tool`` keeps returning ``True`` for the - evicted name AND no listener fires. + Negative test: restore the unconditional pop + rebuild + notify in + ``_close_pool_entry_if_idle``: this test fails because the entry + vanishes, ``is_mcp_tool`` flips to False, and the listener fires. """ mgr, loop, _ = running_loop_mgr handler = _make_jsonrpc_handler( @@ -517,36 +572,100 @@ def test_close_pool_entry_if_idle_clears_catalog_and_fires_listener( _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") key = ("user-1", "pool-srv") assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True - # Sanity: in_flight must be 0 for the eviction path to proceed. assert mgr._user_pool_entries[key].in_flight == 0 - # Seed the debounce dict so the perf-1 prune is observable. + # Seed the debounce dict so the prune-on-close is observable. mgr._last_pool_notification_refresh[key] = 0.0 user_calls = [0] admin_calls = [0] - other_calls = [0] def _user_cb() -> None: user_calls[0] += 1 + def _admin_cb() -> None: + admin_calls[0] += 1 + + # The user-scoped tool listener is the liveness signal (#836). + mgr.add_listener(_user_cb, user_id="user-1") + mgr.add_listener(_admin_cb) # admin / None + + _run_on_loop(loop, mgr._close_pool_entry_if_idle(key)) + + # Entry cooled, not dropped: transport gone, catalog intact. + 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.tools is not None + # The lock object must survive with the entry — an in-flight + # dispatcher's next acquire needs the same lock. + assert key in mgr._user_pool_locks + # Debounce stamp pruned with the transport. + assert key not in mgr._last_pool_notification_refresh + # Per-user catalog view untouched. + assert "user-1" in mgr._user_tool_map + assert "user-1" in mgr._user_tools + assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True + # Nothing changed catalog-wise → NO fan-out (a fan-out here would + # make every live session rebuild its tool list every idle TTL). + assert user_calls[0] == 0, f"user-keyed listener fired {user_calls[0]} times; expected 0" + assert admin_calls[0] == 0, f"admin listener fired {admin_calls[0]} times; expected 0" + + # A second close on the already-cooled entry is a no-op — no drop, + # no fan-out (the eviction loop additionally skips cooled entries + # before even getting here). + _run_on_loop(loop, mgr._close_pool_entry_if_idle(key)) + assert key in mgr._user_pool_entries + assert user_calls[0] == 0 and admin_calls[0] == 0 + + +def test_close_pool_entry_if_idle_drops_entry_without_live_listener( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """TTL/LRU eviction FULLY drops the entry when the user has no live + session: entry popped, notification-debounce dict pruned, per-user + maps rebuilt, and the (admin-only) fan-out fires — departed users' + entries must not outlive their sessions. + + Phase 7 round-2 review hardening (round2-1) heritage: the + catalog-cleanup block needs an integration test driving it — the + failure mode flagged in ``feedback_tests_through_boundaries.md``. + Negative test: drop the ``_rebuild_user_tool_map`` / + ``_notify_user_tool_listeners`` calls from the post-pop block; this + test fails because ``is_mcp_tool`` keeps returning ``True`` for the + evicted name AND the admin listener never fires. + """ + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + list_tools_response=_list_tools_payload([_tool_spec("do_thing")]), + ) + _build_mock_transport_factory(mgr, monkeypatch, handler) + _patch_tcp_probe(mgr, monkeypatch) + + _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") + key = ("user-1", "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. + mgr._last_pool_notification_refresh[key] = 0.0 + + admin_calls = [0] + other_calls = [0] + def _admin_cb() -> None: admin_calls[0] += 1 def _other_cb() -> None: other_calls[0] += 1 - mgr.add_listener(_user_cb, user_id="user-1") + # NO user-1 tool listener — user-1 has no live session. The admin + # (None) and unrelated-user listeners don't count as liveness. mgr.add_listener(_admin_cb) # admin / None mgr.add_listener(_other_cb, user_id="user-2") - # Drive the LRU/TTL eviction path directly. ``open_lock`` is - # uncontested (no concurrent dispatch) and ``in_flight`` is 0, - # so the close proceeds without retry. _run_on_loop(loop, mgr._close_pool_entry_if_idle(key)) - # Entry fully removed (LRU eviction pops the dict — unlike - # ``_evict_session`` which keeps the entry as a ``session=None`` - # phantom for the next dispatch to re-connect). + # Entry fully removed. assert key not in mgr._user_pool_entries assert key not in mgr._user_pool_last_used assert key not in mgr._user_pool_locks @@ -556,8 +675,8 @@ def test_close_pool_entry_if_idle_clears_catalog_and_fires_listener( assert "user-1" not in mgr._user_tool_map assert "user-1" not in mgr._user_tools assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is False - # Listener fan-out: matching user + admin fire, OTHER user does not. - assert user_calls[0] == 1, f"user-keyed listener fired {user_calls[0]} times; expected 1" + # Admin fan-out fires (operator tooling observes the drop); the + # unrelated user's listener does not. assert admin_calls[0] == 1, f"admin listener fired {admin_calls[0]} times; expected 1" assert other_calls[0] == 0, ( f"unrelated user-2 listener fired {other_calls[0]} times; expected 0 — " @@ -565,18 +684,25 @@ def test_close_pool_entry_if_idle_clears_catalog_and_fires_listener( ) -def test_reconnect_after_eviction_repopulates_catalog( +def test_reconnect_after_eviction_corrects_catalog_drift( running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """After eviction, the next ``_connect_one_pool`` re-populates the - catalog from the SDK by construction — no extra glue required. + """The retained catalog self-corrects at reconnect: the next + ``_connect_one_pool`` re-runs discovery and REPLACES the stale + snapshot — no extra glue required (#836). + + Between eviction and reconnect the OLD names stay visible (that is + the point of retention — the model can still emit them, and the + dispatch that follows performs this reconnect); a name the server + dropped in the meantime dies at the server as a per-call error + while the refreshed catalog fans out. Verified by reverting the discovery block in ``_connect_one_pool``: - the reconnected entry's ``tools`` stays ``None`` and the user map - never re-emerges. + the reconnected entry keeps serving the stale ``tool_a`` and + ``tool_b`` never appears. """ mgr, loop, _ = running_loop_mgr - # Sequence: first connect sees [tool_a]; eviction clears; second + # Sequence: first connect sees [tool_a]; the session dies; second # connect (after backend rotates) sees [tool_b]. handler = _make_jsonrpc_handler( list_tools_seq=[ @@ -591,7 +717,9 @@ def test_reconnect_after_eviction_repopulates_catalog( assert mgr.is_mcp_tool("mcp__pool-srv__tool_a", user_id="user-1") is True mgr._evict_session(("user-1", "pool-srv")) - assert mgr.is_mcp_tool("mcp__pool-srv__tool_a", user_id="user-1") is False + # Retention: the stale snapshot keeps serving the live session + # until the reconnect refreshes it. + assert mgr.is_mcp_tool("mcp__pool-srv__tool_a", user_id="user-1") is True _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") # Reconnect picked up the rotated catalog. @@ -1555,17 +1683,14 @@ def test_refresh_pool_server_prompts_skips_when_capability_unset( # --------------------------------------------------------------------------- -def test_eviction_clears_resource_and_prompt_catalogs_and_fires_listeners( +def test_evict_session_keeps_resource_and_prompt_catalogs( running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """``_evict_session`` clears ``entry.resources`` and ``entry.prompts``, - rebuilds both per-user maps (drops the now-empty entries), and - fires the matching user-keyed + admin listeners for ALL three - catalogs (tools, resources, prompts). - - Negative-test: drop the resource/prompt cleanup additions in - ``_evict_session``: this test fails because the user-resource map - keeps the evicted URIs and no resource listener fires. + """``_evict_session`` retains ``entry.resources`` / ``entry.prompts`` + and their per-user maps, firing NO listeners — symmetric with the + tool-catalog retention (#836). The revocation flavor + (``_evict_session_drop_catalog``) is what clears and notifies; see + the sibling test below. """ mgr, loop, _ = running_loop_mgr handler = _make_jsonrpc_handler( @@ -1604,15 +1729,81 @@ def test_eviction_clears_resource_and_prompt_catalogs_and_fires_listeners( mgr._evict_session(("user-1", "pool-srv")) + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.session is None + # All three catalogs retained. + assert entry.tools is not None + assert entry.resources is not None + assert entry.prompts is not None + # Per-user maps keep the entries. + assert "res://r/1" in (mgr._user_resource_map.get("user-1") or {}) + assert "mcp__pool-srv__p1" in (mgr._user_prompt_map.get("user-1") or {}) + # Nothing changed catalog-wise → NO fan-out for anyone. + assert res_calls[0] == 0 + assert prompt_calls[0] == 0 + assert other_res_calls[0] == 0, "unrelated user-2 resource listener fired — RFC §3.3 violation" + assert other_prompt_calls[0] == 0, "unrelated user-2 prompt listener fired — RFC §3.3 violation" + + +def test_evict_user_session_clears_resource_and_prompt_catalogs( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """``_evict_session_drop_catalog`` (explicit revocation) clears + ``entry.resources`` / ``entry.prompts``, rebuilds both per-user + maps (drops the now-empty entries), and fires the matching + user-keyed listeners for resources and prompts. + + Negative-test: drop the resource/prompt cleanup in + ``_evict_session_drop_catalog``: this test fails because the + user-resource map keeps the revoked URIs and no resource listener + fires. + """ + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(resources=True, prompts=True), + list_resources_response=_list_resources_payload([_resource_spec("res://r/1")]), + list_prompts_response=_list_prompts_payload([_prompt_spec("p1")]), + ) + _build_mock_transport_factory(mgr, monkeypatch, handler) + _patch_tcp_probe(mgr, monkeypatch) + + _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") + assert "res://r/1" in (mgr._user_resource_map.get("user-1") or {}) + assert "mcp__pool-srv__p1" in (mgr._user_prompt_map.get("user-1") or {}) + + res_calls = [0] + prompt_calls = [0] + other_res_calls = [0] + other_prompt_calls = [0] + + def _user_res_cb() -> None: + res_calls[0] += 1 + + def _user_prompt_cb() -> None: + prompt_calls[0] += 1 + + def _other_res_cb() -> None: + other_res_calls[0] += 1 + + def _other_prompt_cb() -> None: + other_prompt_calls[0] += 1 + + mgr.add_resource_listener(_user_res_cb, user_id="user-1") + mgr.add_resource_listener(_other_res_cb, user_id="user-2") + mgr.add_prompt_listener(_user_prompt_cb, user_id="user-1") + mgr.add_prompt_listener(_other_prompt_cb, user_id="user-2") + + mgr._evict_session_drop_catalog(("user-1", "pool-srv")) + entry = mgr._user_pool_entries[("user-1", "pool-srv")] assert entry.session is None assert entry.tools is None assert entry.resources is None assert entry.prompts is None - # Per-user maps drop the evicted entries. + # Per-user maps drop the revoked entries. assert "user-1" not in mgr._user_resource_map assert "user-1" not in mgr._user_prompt_map - # Listeners fire for the evicted user but not for the unrelated user. + # Listeners fire for the revoked user but not for the unrelated user. assert res_calls[0] == 1 assert prompt_calls[0] == 1 assert other_res_calls[0] == 0, "unrelated user-2 resource listener fired — RFC §3.3 violation" @@ -1623,8 +1814,13 @@ def test_close_pool_entry_if_idle_clears_resource_and_prompt_catalogs( running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch ) -> None: """LRU/TTL eviction (``_close_pool_entry_if_idle``) symmetric - cleanup for resources & prompts. Bug-pair to the tools-only - cleanup added in Phase 7 round-2.""" + cleanup for resources & prompts on the FULL-DROP path. Bug-pair to + the tools-only cleanup added in Phase 7 round-2. + + Deliberately registers NO user-1 TOOL listener: the user-scoped + tool listener is the liveness signal (#836) — resource/prompt + listeners alone do not mark a user live (every real ChatSession + registers the tool listener), so this drives the full drop.""" mgr, loop, _ = running_loop_mgr handler = _make_jsonrpc_handler( init_response=_init_response_with_caps(resources=True, prompts=True), @@ -1661,12 +1857,62 @@ def test_close_pool_entry_if_idle_clears_resource_and_prompt_catalogs( assert prompt_calls[0] == 1 +def test_close_pool_entry_if_idle_cooling_keeps_resources_and_prompts( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """The cooled path retains resources & prompts symmetrically with + tools: a user-1 TOOL listener (the liveness signal, #836) makes + TTL/LRU eviction keep the entry, both per-user maps, and fire no + resource/prompt listener.""" + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(resources=True, prompts=True), + list_resources_response=_list_resources_payload([_resource_spec("res://r/1")]), + list_prompts_response=_list_prompts_payload([_prompt_spec("p1")]), + ) + _build_mock_transport_factory(mgr, monkeypatch, handler) + _patch_tcp_probe(mgr, monkeypatch) + + _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") + key = ("user-1", "pool-srv") + + res_calls = [0] + prompt_calls = [0] + + def _res_cb() -> None: + res_calls[0] += 1 + + def _prompt_cb() -> None: + prompt_calls[0] += 1 + + # The TOOL listener marks user-1 live; the resource/prompt + # listeners observe (non-)fan-out. + mgr.add_listener(lambda: None, user_id="user-1") + mgr.add_resource_listener(_res_cb, user_id="user-1") + mgr.add_prompt_listener(_prompt_cb, user_id="user-1") + + _run_on_loop(loop, mgr._close_pool_entry_if_idle(key)) + + # Entry cooled: transport gone, catalogs and maps intact. + entry = mgr._user_pool_entries.get(key) + assert entry is not None + assert entry.session is None + assert entry.resources is not None + assert entry.prompts is not None + assert "res://r/1" in (mgr._user_resource_map.get("user-1") or {}) + assert "mcp__pool-srv__p1" in (mgr._user_prompt_map.get("user-1") or {}) + # No fan-out — nothing changed catalog-wise. + assert res_calls[0] == 0 + assert prompt_calls[0] == 0 + + def test_reconnect_after_eviction_repopulates_resource_and_prompt_catalogs( running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """After eviction, the next ``_connect_one_pool`` re-populates the - resource and prompt catalogs from the SDK. Bug-class: an extra - glue layer would only repopulate tools.""" + """After a session eviction the next ``_connect_one_pool`` REPLACES + the retained resource and prompt catalogs from the SDK — drift + self-corrects for all three catalogs, not just tools (#836). + Bug-class: an extra glue layer would only repopulate tools.""" mgr, loop, _ = running_loop_mgr handler = _make_jsonrpc_handler( init_response=_init_response_with_caps(resources=True, prompts=True), @@ -1687,8 +1933,9 @@ def test_reconnect_after_eviction_repopulates_resource_and_prompt_catalogs( assert "mcp__pool-srv__p_a" in (mgr._user_prompt_map.get("user-1") or {}) mgr._evict_session(("user-1", "pool-srv")) - assert "user-1" not in mgr._user_resource_map - assert "user-1" not in mgr._user_prompt_map + # Retention: the stale snapshots keep serving until reconnect. + assert "res://a" in (mgr._user_resource_map.get("user-1") or {}) + assert "mcp__pool-srv__p_a" in (mgr._user_prompt_map.get("user-1") or {}) _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") # New catalogs reflect the rotated payload. diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index cf9ff650..fbe9c974 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -400,6 +400,145 @@ class TestEviction: # All entries removed from the dict regardless. assert mgr._user_pool_entries == {} + @staticmethod + def _fake_tools(server_name: str, i: int) -> list[dict]: + return [ + { + "type": "function", + "function": { + "name": f"mcp__{server_name}__t{i}", + "description": "", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def test_idle_eviction_cools_entries_for_live_listener_users(self, running_loop_mgr) -> None: + """TTL eviction cools (retains) a live-listener user's + catalog-bearing entry and full-drops a listener-less user's — + the #836 split. A second tick must not disturb the cooled entry + (the already-cooled skip).""" + mgr, loop, _ = running_loop_mgr + mgr._user_pool_idle_ttl_s = 0.0 # everything is stale + + async def _seed() -> None: + for i in range(2): + entry = await mgr._ensure_pool_entry((f"u{i}", "pool-srv")) + entry.session = MagicMock() + entry.tools = self._fake_tools("pool-srv", i) + + _run_on_loop(loop, _seed()) + # 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. + 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 + # u1: full drop. + assert ("u1", "pool-srv") not in mgr._user_pool_entries + + # Second tick: the cooled entry is skipped, not re-processed. + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + assert ("u0", "pool-srv") in mgr._user_pool_entries + + def test_idle_eviction_drops_catalogless_stub_despite_live_listener( + self, running_loop_mgr + ) -> None: + """A cold, never-discovered stub (``_ensure_pool_entry`` + allocated, connect failed before discovery) carries no catalog + worth retaining — TTL eviction drops it even for a live-listener + user, so revoke-cleared and connect-failed entries can't + accumulate as zombies behind an open session.""" + mgr, loop, _ = running_loop_mgr + mgr._user_pool_idle_ttl_s = 0.0 + + async def _seed() -> None: + await mgr._ensure_pool_entry(("u-stub", "pool-srv")) + + _run_on_loop(loop, _seed()) + mgr.add_listener(lambda: None, user_id="u-stub") + + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + assert ("u-stub", "pool-srv") not in mgr._user_pool_entries + + def test_lru_cap_ignores_cooled_entries(self, running_loop_mgr) -> None: + """The LRU cap bounds WARM entries (connection resources), not + cooled catalog-only ones — cooled entries neither count toward + the cap nor get evicted by it.""" + mgr, loop, _ = running_loop_mgr + mgr._user_pool_idle_ttl_s = 999_999.0 # TTL effectively disabled + mgr._user_pool_lru_max = 2 + + async def _seed() -> None: + base = time.monotonic() + # Three cooled entries (no session/owner, catalog present). + for i in range(3): + key = (f"cool{i}", "pool-srv") + entry = await mgr._ensure_pool_entry(key) + entry.tools = self._fake_tools("pool-srv", i) + entry.last_used = base + i + mgr._user_pool_last_used[key] = base + i + # Two warm entries, newer than the cooled ones. + for i in range(2): + key = (f"warm{i}", "pool-srv") + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() + entry.tools = self._fake_tools("pool-srv", 10 + i) + entry.last_used = base + 10 + i + mgr._user_pool_last_used[key] = base + 10 + i + + _run_on_loop(loop, _seed()) + for i in range(3): + mgr.add_listener(lambda: None, user_id=f"cool{i}") + + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + + # Warm count (2) is at the cap — nothing evicted, cooled + # entries (which would be "oldest" by last_used) untouched. + assert len(mgr._user_pool_entries) == 5 + assert all((f"cool{i}", "pool-srv") in mgr._user_pool_entries for i in range(3)) + + def test_lru_cap_cools_live_listener_entries(self, running_loop_mgr) -> None: + """Over the cap, warm entries of live-listener users are COOLED + (transport closed, entry + catalog retained) oldest-first until + the warm count meets the cap — cap pressure must not reintroduce + the #836 tool loss for live sessions.""" + mgr, loop, _ = running_loop_mgr + mgr._user_pool_idle_ttl_s = 999_999.0 + mgr._user_pool_lru_max = 1 + + async def _seed() -> None: + base = time.monotonic() + for i in range(3): + key = (f"u{i}", "pool-srv") + entry = await mgr._ensure_pool_entry(key) + entry.session = MagicMock() + entry.tools = self._fake_tools("pool-srv", i) + entry.last_used = base + i + mgr._user_pool_last_used[key] = base + i + + _run_on_loop(loop, _seed()) + for i in range(3): + mgr.add_listener(lambda: None, user_id=f"u{i}") + + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + + # All three entries survive with catalogs; only the newest is + # still warm. + assert len(mgr._user_pool_entries) == 3 + warm = [ + key + for key, e in mgr._user_pool_entries.items() + if e.session is not None or e.owner_task is not None + ] + assert warm == [("u2", "pool-srv")] + for i in range(3): + assert mgr._user_pool_entries[(f"u{i}", "pool-srv")].tools is not None + # --------------------------------------------------------------------------- # Dispatch state machine diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index dbdac57f..6dd52040 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -697,7 +697,7 @@ class MCPClientManager: # 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``). + # ``_close_pool_entry_if_idle`` / ``_evict_session_drop_catalog``). self._user_tools: dict[str, list[dict[str, Any]]] = {} # Per-user resource catalog. Mirrors ``_user_tool_map`` / @@ -2965,12 +2965,61 @@ 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: + """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. + """ + with self._listeners_lock: + return any(uid == user_id for uid, _cb in self._listeners) + + @staticmethod + def _entry_has_catalog(entry: PoolEntryState) -> bool: + """True when the entry carries a discovered catalog worth retaining. + + Never-connected stubs (``_ensure_pool_entry`` allocated, connect + failed before discovery) and revoke-cleared entries + (:meth:`_evict_session_drop_catalog`) carry none — retaining + those serves nobody, so the eviction passes full-drop them even + when the user has a live session. + """ + 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. + + 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. + """ + return sum( + 1 + for e in self._user_pool_entries.values() + if e.session is not None or e.owner_task is not None + ) + async def _evict_idle_pool_entries(self) -> None: """Evict pool entries past the idle TTL or above the LRU cap. Skips any key whose ``open_lock`` is currently held or whose ``in_flight`` counter is non-zero — eviction never blocks on a contested lock or an active dispatch; the next tick retries. + + Both passes close TRANSPORTS. Whether the ENTRY (and with it the + user's catalog contribution) survives is decided per user in + :meth:`_close_pool_entry_if_idle`: a user with a live session + keeps a cooled catalog-only entry so their model-visible tool + list never silently shrinks (#836); a user without one gets the + full drop. The LRU cap therefore bounds WARM entries — the + connection resources (transport, httpx client, owner task) are + what the cap exists to limit. Cooled entries hold none of + those; they are bounded by live-session users × pool servers + and reaped by the TTL pass within a tick of their user's last + listener going away. """ if not self._user_pool_entries: return @@ -2982,37 +3031,70 @@ class MCPClientManager: ttl_targets: list[tuple[str, str]] = [] for key, entry in list(self._user_pool_entries.items()): last = self._user_pool_last_used.get(key, entry.last_used) - if (now - last) >= ttl: - ttl_targets.append(key) + 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]) + ): + # 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. + continue + ttl_targets.append(key) if ttl_targets: await asyncio.gather( *(self._close_pool_entry_if_idle(k) for k in ttl_targets), return_exceptions=True, ) - # Second pass: LRU cap. Iterate the canonical entry map (not the - # last_used view) so brand-new entries that were created via - # ``_ensure_pool_entry`` but haven't dispatched yet are still - # eviction-eligible. - if len(self._user_pool_entries) <= self._user_pool_lru_max: + # Second pass: LRU cap over warm entries. Iterate the canonical + # entry map (not the last_used view) so brand-new entries that + # were created via ``_ensure_pool_entry`` but haven't dispatched + # yet are still eviction-eligible. + 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 len(warm) <= self._user_pool_lru_max: return ordered = sorted( - self._user_pool_entries.items(), + 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). for key, _entry in ordered: - if len(self._user_pool_entries) <= self._user_pool_lru_max: + if self._warm_pool_count() <= self._user_pool_lru_max: break await self._close_pool_entry_if_idle(key) async def _close_pool_entry_if_idle(self, key: tuple[str, str]) -> None: - """Close ``key`` iff its open_lock is uncontested AND in_flight==0. + """Close ``key``'s transport iff its open_lock is uncontested AND in_flight==0. Best-effort: a contested lock or an active dispatch causes the 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`): + + - 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 + ``_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. """ entry = self._user_pool_entries.get(key) if entry is None: @@ -3043,16 +3125,26 @@ class MCPClientManager: if entry.in_flight > 0: return await self._teardown_pool_entry(key) - self._user_pool_entries.pop(key, None) - self._user_pool_last_used.pop(key, None) - # Mirror ``_evict_session``'s catalog cleanup: 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, and ChatSession's - # tool / resource / prompt lists would never rebuild because - # no listener fires. + # 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._entry_has_catalog(entry) and self._user_has_live_listener(user_id): + # 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 + # entry's ``open_lock`` must survive with it (an + # in-flight dispatcher's next acquire needs the same + # lock object), so skip the ``evicted`` cleanup too. + return + self._user_pool_entries.pop(key, None) + self._user_pool_last_used.pop(key, None) + # 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) @@ -6573,7 +6665,7 @@ class MCPClientManager: return _decode_prompt_result(sdk_result) def _evict_session(self, key: tuple[str, str]) -> None: - """Drop the cached session AND catalog on a pool entry. + """Drop the cached session on a pool entry; KEEP the catalog. Owner/streams left for reconnect. Auth/transport branches both call this — the next connect's ``_connect_one_pool`` tears down the stale @@ -6584,39 +6676,63 @@ class MCPClientManager: following reconnect is reaped by idle eviction, which drives the same protocol. - Catalog cleanup (Phase 7 / 7b): clearing ``entry.tools`` / - ``entry.resources`` / ``entry.prompts`` here ensures an - evicted-then-not-yet-reconnected entry contributes no stale - entries to ``_user_tool_map`` / ``_user_resource_map`` / - ``_user_prompt_map``. Without this, ``is_mcp_tool`` / - ``is_mcp_prompt`` / pool resource resolution could return True - for names whose backing pool is gone, then the resolver would - dispatch into a session-less entry and the next call would - surface as a generic transport error instead of a clean - reconnect path. + The catalog (``entry.tools`` / ``resources`` / ``prompts``) is + deliberately RETAINED — the evict-session-keep-entry shape of + ``_on_pool_owner_death`` (#836). Clearing it here removed the + server's tools from the user's live sessions on the first failed + dispatch (a transport blip, a 403, a double-401): the per-user + maps rebuilt empty, the session-side ``is_mcp_tool`` gate + closed, and with no re-prime path for a live session the tools + never came back — which also made the breaker's half-open + recovery and the consent / step-up cards unreachable (the model + could no longer emit the name they need to fire). A session-less + entry stays fully dispatchable: tool-name resolution never reads + the catalog (``_resolve_pool_target`` is name + server-row + based) and ``_dispatch_pool_with_entry`` connect-or-reuses, + re-running discovery — post-reconnect drift self-corrects and + the catalog-refresh notification fans out then. """ evict = self._user_pool_entries.get(key) if evict is not None: evict.session = None - evict.tools = None - evict.resources = None - evict.prompts = None - # Prune the debounce stamp in lockstep with the entry — the - # dict grows otherwise (slow leak) across (user, server) - # churn. Mirrors the cleanup in ``_close_pool_entry_if_idle``. + # Prune the push-notification debounce stamp so a + # reconnect's first ``list_changed`` refreshes immediately. 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 session so its merged tool / resource / - # prompt lists shrink back to static-only until the next - # connect populates the entry. 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) + + 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). + 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). + """ + 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) async def _handle_auth_403( self, @@ -7073,22 +7189,24 @@ class MCPClientManager: return _decode_prompt_result(result) def evict_user_session(self, user_id: str, server_name: str) -> None: - """Drop the cached pool session for ``(user_id, server_name)``. + """Drop the cached pool session AND catalog for ``(user_id, server_name)``. - Sync entry point for callers (e.g. the OAuth revoke handler) - that mutate token state from outside the mcp-loop and need the - next dispatch to reconnect with fresh credentials. Idempotent — + Sync entry point for the OAuth revoke handler: the user + explicitly disconnected the server, so the session is dropped + and — unlike the dispatch-failure eviction (#836) — the user's + catalog view of the server is removed and their live sessions + notified (:meth:`_evict_session_drop_catalog`). Idempotent — a missing key is a silent no-op. Fire-and-forget: schedules - :meth:`_evict_session` on the mcp-loop and returns immediately - without waiting for the future. Best-effort: a closed loop or - scheduling failure logs at info level; never raises. + onto the mcp-loop and returns immediately without waiting for + the future. Best-effort: a closed loop or scheduling failure + logs at info level; never raises. """ if self._loop is None: return key = (user_id, server_name) async def _do_evict() -> None: - self._evict_session(key) + self._evict_session_drop_catalog(key) try: asyncio.run_coroutine_threadsafe(_do_evict(), self._loop)