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