fix(mcp): retain per-user catalogs when pool sessions close under live sessions

Idle-TTL eviction tore down a per-user pool entry, rebuilt the user's
tool/resource/prompt catalogs (now empty), and notified listeners — so
every live ChatSession for that user silently lost the server's tools
after 10 idle minutes, with no way back: prime_user_pools only runs at
session construction, acting-user change, and reconcile, and the
emptied catalog closes the session-side is_mcp_tool gate, so even a
history-motivated call can't reach the lazy-reconnect dispatch path.
The dispatch-failure paths (_evict_session on 401/403/transport)
cleared catalogs the same way, so a transport blip during a tool call
caused the same permanent loss with no TTL involved — and made the
breaker's half-open recovery and the consent/step-up cards unreachable.

Both now follow _on_pool_owner_death's evict-session-keep-entry shape:

- _evict_session drops only the session. The catalog stays; the next
  dispatch connect-or-reuses and re-runs discovery, so drift
  self-corrects and the refresh notification fans out then.
- TTL eviction COOLS entries of users with a live session (a
  registered user-scoped tool listener): transport closed, entry and
  catalog retained, no fan-out. Users without one keep the full drop,
  so departed users' entries don't outlive their sessions.
- The LRU cap now bounds WARM entries — the connection resources it
  exists to limit. Over the cap, live-listener users' entries are
  cooled rather than dropped; cooled catalog-only entries are bounded
  by live users x pool servers and reaped one tick after the user's
  last listener goes away.
- Explicit disconnect keeps its semantics: evict_user_session routes
  to the new _evict_session_drop_catalog (clear + rebuild + notify) —
  the user asked for the tools to leave. Clearing the catalog also
  marks the entry droppable, so it can't linger cooled.

Never-discovered stubs (no catalog) are always dropped, already-cooled
entries are skipped by later ticks, and a cooled entry keeps its
open_lock object for in-flight dispatchers.

Applies to oauth_user and oauth_obo alike: the pool and its eviction
are auth-type-agnostic, and for obo priming is the only path tools
enter a catalog at all.

Fixes #836
This commit is contained in:
Patrick Buckley
2026-07-13 07:58:38 -07:00
parent 3742e9660a
commit cb94ea349f
4 changed files with 624 additions and 117 deletions
+4 -1
View File
@@ -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
+307 -60
View File
@@ -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.
+139
View File
@@ -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
+174 -56
View File
@@ -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)