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