fix(mcp): observe before the lookup; coalesce queued refreshes; pop stamps on every teardown

Close the round-7 review findings:

- The dead-grant observation is now snapshotted BEFORE the classified
  lookup's first await, by the callers (the three dispatchers via
  _pool_lookup_failure, _prime_one, and the obo credential gate), and
  _schedule_dead_grant_drop requires it as a parameter: snapshotting
  after the lookup returned could capture a session the
  consent-completion prime connected mid-lookup — its awaits can park
  on executor hops — and the drop then evicted the just-restored
  catalog it exists to spare, with no remaining re-prime path.

- Spawned list_changed refreshes coalesce on a per-(key, kind) marker:
  set at spawn, cleared the moment the runner acquires open_lock
  (before its list call, so a change the in-flight list missed spawns
  exactly one successor). Admission was one per 5s debounce window
  while each runner can hold the lock up to the 30s refresh timeout,
  so a notifying-but-slow server accreted lock waiters without bound —
  FIFO dispatch waits past the 120s budget, idle eviction starved by
  the contested lock, and background tasks growing for as long as the
  server kept notifying. The runner also returns quietly for an
  evicted session instead of failing through the log. The residual
  duty-cycle case (a wedged-but-notifying server defers idle eviction
  of its own entry until the first dispatch, recovery, or silence) is
  documented at the runner.

- Every teardown path now pops the notification debounce stamp:
  _teardown_pool_entry and _on_pool_owner_death left it in place, so
  the keep-stamp design's documented reconnect backstop did not exist
  on the idle-collapse and connect-failure paths — a change announced
  in a failed window could be debounced against a pre-collapse stamp
  after reconnect and never land. The idle-close path's own pop is
  now owned by _teardown_pool_entry.

- Cleanups: the notification table maps type to kind label only, with
  the kind-to-refresher map bound at dispatch time (mypy-checked
  attribute references, instance overrides keep working) instead of
  getattr on a name string; _schedule_dead_grant_drop skips when there
  is provably nothing to converge (no entry, or a session-less
  catalog-less stub), sparing a tracked no-op task per unconsented
  server per prime at scale; the fire-and-forget prime idiom's three
  hand-synced copies collapse into try_prime_user_pools (session
  construction, acting-user change, OIDC capture); the stale
  lock-contract docstrings on the resources/prompts refreshers now
  state the held-lock requirement; has_live_session_listener is the
  sole listener-liveness predicate (the private alias is gone); the
  construction-scoped tools-seq read is a constructor local instead of
  a persistent ChatSession attribute.
This commit is contained in:
Patrick Buckley
2026-07-13 20:14:41 -07:00
parent 7186e1e709
commit e9ecf91c07
4 changed files with 463 additions and 129 deletions
+179 -6
View File
@@ -642,7 +642,7 @@ class TestEviction:
entry.tools = _fake_pool_tools("pool-srv", "t")
mgr._rebuild_user_tool_map("u0")
entry.session = MagicMock() # warm BEFORE the failed lookup
mgr._schedule_dead_grant_drop(dead, key)
mgr._schedule_dead_grant_drop(dead, key, observed_session=entry.session)
_run_on_loop(loop, _warm_then_schedule())
_drain_background(mgr, loop)
@@ -660,7 +660,7 @@ class TestEviction:
fresh.tools = _fake_pool_tools("pool-srv", "t")
mgr._rebuild_user_tool_map("u0")
fresh.session = MagicMock() # observed at schedule time
mgr._schedule_dead_grant_drop(dead, key)
mgr._schedule_dead_grant_drop(dead, key, observed_session=fresh.session)
fresh.session = MagicMock() # re-consent connect lands first
_run_on_loop(loop, _reseed_and_swap())
@@ -668,6 +668,34 @@ class TestEviction:
entry = mgr._user_pool_entries[key]
assert entry.tools is not None, "drop must spare a session newer than its observation"
def test_dead_grant_drop_skips_when_nothing_to_converge(self, running_loop_mgr) -> None:
"""At scale every unconsented server classifies as a dead grant
on every prime; spawning a tracked no-op drop task per (user,
unconsented server) is pure mcp-loop overhead. The schedule-side
guard skips when there is no entry, or a session-less entry with
no catalog — a catalog appearing later implies a lookup that
succeeded after this one failed."""
mgr, loop, _ = running_loop_mgr
mgr._app_state = SimpleNamespace(mcp_token_store=object())
mgr._storage = object()
dead = SimpleNamespace(kind="missing", token=None)
key = ("u0", "pool-srv")
async def _schedule_no_entry() -> int:
before = len(mgr._background_tasks)
mgr._schedule_dead_grant_drop(dead, key, observed_session=None)
return len(mgr._background_tasks) - before
assert _run_on_loop(loop, _schedule_no_entry()) == 0
async def _schedule_bare_stub() -> int:
await mgr._ensure_pool_entry(key)
before = len(mgr._background_tasks)
mgr._schedule_dead_grant_drop(dead, key, observed_session=None)
return len(mgr._background_tasks) - before
assert _run_on_loop(loop, _schedule_bare_stub()) == 0
def test_notification_refresh_failure_keeps_debounce_stamp(self, running_loop_mgr) -> None:
"""A failed spawned refresh must KEEP the debounce stamp:
popping it re-arms the handler on every notification, so a
@@ -679,7 +707,12 @@ class TestEviction:
notification refreshes immediately)."""
mgr, loop, _ = running_loop_mgr
key = ("u0", "pool-srv")
_run_on_loop(loop, mgr._ensure_pool_entry(key))
async def _seed() -> None:
entry = await mgr._ensure_pool_entry(key)
entry.session = MagicMock() # runner refreshes only live sessions
_run_on_loop(loop, _seed())
mgr._last_pool_notification_refresh[key] = 123.0
async def _boom(_key: tuple[str, str]) -> tuple[list[str], list[str]]:
@@ -707,7 +740,12 @@ class TestEviction:
catch) — the anyio stray-cancel shape."""
mgr, loop, _ = running_loop_mgr
key = ("u0", "pool-srv")
_run_on_loop(loop, mgr._ensure_pool_entry(key))
async def _seed() -> None:
entry = await mgr._ensure_pool_entry(key)
entry.session = MagicMock() # runner refreshes only live sessions
_run_on_loop(loop, _seed())
mgr._last_pool_notification_refresh[key] = 123.0
async def _wedge(_key: tuple[str, str]) -> tuple[list[str], list[str]]:
@@ -732,6 +770,7 @@ class TestEviction:
async def _scenario() -> bool:
entry = await mgr._ensure_pool_entry(key)
entry.session = MagicMock() # runner refreshes only live sessions
await entry.open_lock.acquire()
task = asyncio.ensure_future(mgr._run_notification_refresh(key, "tools", _rec))
for _ in range(3):
@@ -775,6 +814,89 @@ class TestEviction:
_run_on_loop(loop, _scenario())
assert refreshed == []
def test_notification_coalesces_when_refresh_already_queued(self, running_loop_mgr) -> None:
"""While a runner is queued for a key+kind (coalesce marker
set), further notifications spawn NOTHING — the parked runner's
fresh list observes their change when it acquires the lock.
This bounds the ``open_lock`` waiter queue at one parked runner
per key+kind: without it a notifying-but-slow server accretes
waiters (admitted 1/5s, drained 1/30s) that starve same-key
dispatches past their wall-clock budget and starve idle
eviction, while ``_background_tasks`` grows without bound."""
mgr, loop, _ = running_loop_mgr
key = ("u0", "pool-srv")
refreshed: list[tuple[str, str]] = []
async def _fake_refresh(k: tuple[str, str]) -> tuple[list[str], list[str]]:
refreshed.append(k)
return [], []
mgr._refresh_pool_server_tools = _fake_refresh # type: ignore[method-assign]
handler = mgr._make_pool_notification_handler(key)
async def _fire_twice() -> int:
entry = await mgr._ensure_pool_entry(key)
entry.session = MagicMock()
await entry.open_lock.acquire() # park the first runner
note = mcp_types.ServerNotification(
mcp_types.ToolListChangedNotification(method="notifications/tools/list_changed")
)
before = len(mgr._background_tasks)
await handler(note)
# Age the stamp past the debounce window: the MARKER (not
# the stamp) must do the suppression for the second fire.
mgr._last_pool_notification_refresh[key] = time.monotonic() - 10.0
await handler(note)
spawned = len(mgr._background_tasks) - before
entry.open_lock.release()
return spawned
spawned = _run_on_loop(loop, _fire_twice())
_drain_background(mgr, loop)
assert spawned == 1, "second notification must coalesce into the parked runner"
assert refreshed == [key]
# The runner released its marker (at lock acquire + finally).
assert not mgr._pool_refresh_pending
def test_teardown_pool_entry_pops_debounce_stamp(self, running_loop_mgr) -> None:
"""Every teardown path must pop the debounce stamp — the
keep-stamp-on-failure design leans on it: a reconnect's first
``list_changed`` refreshes immediately, so a change announced
in a failed window converges at the next reconnect."""
mgr, loop, _ = running_loop_mgr
key = ("u0", "pool-srv")
async def _seed() -> None:
entry = await mgr._ensure_pool_entry(key)
entry.session = MagicMock()
_run_on_loop(loop, _seed())
mgr._last_pool_notification_refresh[key] = 123.0
_run_on_loop(loop, mgr._teardown_pool_entry(key))
assert key not in mgr._last_pool_notification_refresh
def test_owner_death_pops_debounce_stamp(self, running_loop_mgr) -> None:
"""The unrequested-collapse path (owner done-callback) is a
teardown too: the stamp must not outlive the transport."""
mgr, loop, _ = running_loop_mgr
key = ("u0", "pool-srv")
async def _scenario() -> None:
entry = await mgr._ensure_pool_entry(key)
entry.session = MagicMock()
async def _noop() -> None:
pass
task = asyncio.ensure_future(_noop())
await task
entry.owner_task = task
mgr._last_pool_notification_refresh[key] = 123.0
mgr._on_pool_owner_death(key, task)
_run_on_loop(loop, _scenario())
assert key not in mgr._last_pool_notification_refresh
def test_list_changed_handler_spawns_refresh_off_receive_loop(self, running_loop_mgr) -> None:
"""The notification handler must SPAWN the refresh, not await it:
the SDK awaits handlers inline in its receive loop, so an
@@ -793,8 +915,10 @@ class TestEviction:
async def _fire() -> bool:
# The runner refreshes through the entry's ``open_lock``;
# without an entry it (correctly) discards the notification.
await mgr._ensure_pool_entry(key)
# without an entry (or a live session) it correctly
# discards the notification.
entry = await mgr._ensure_pool_entry(key)
entry.session = MagicMock()
note = mcp_types.ServerNotification(
mcp_types.ToolListChangedNotification(method="notifications/tools/list_changed")
)
@@ -1038,6 +1162,55 @@ class TestDispatchStateMachine:
assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is False
assert fired[0] == 1
def test_mid_lookup_reconsent_survives_dead_grant_drop(
self, running_loop_mgr, storage: SQLiteBackend
) -> None:
"""The dead-grant observation is snapshotted BEFORE the
classified lookup's awaits: a consent-completion prime that
connects and publishes DURING the lookup (its awaits can park on
executor hops) must read as re-consent evidence, not as the
pre-observation session — or the scheduled drop evicts the
just-restored catalog with no remaining re-prime path."""
mgr, loop, _ = running_loop_mgr
cipher = make_mcp_token_cipher()
_seed_oauth_server(storage, name="pool-srv")
self._wire_pool(mgr, storage, cipher)
self._seed_cooled_catalog(mgr, loop)
key = ("user-1", "pool-srv")
fresh_session = MagicMock()
async def _lookup_with_race(
server_row: dict[str, Any],
user_id: str,
server_name: str,
force_refresh: bool = False,
) -> Any:
# The consent callback completes mid-lookup: prime connects
# and publishes the fresh catalog before the lookup returns
# its (now stale) dead-grant verdict.
entry = mgr._user_pool_entries[key]
entry.session = fresh_session
entry.tools = _fake_pool_tools("pool-srv", "do_thing")
mgr._rebuild_user_tool_map("user-1")
return SimpleNamespace(kind="missing", token=None)
mgr._pool_token_lookup = _lookup_with_race # type: ignore[method-assign]
with pytest.raises(RuntimeError) as exc_info:
mgr.call_tool_sync(
"mcp__pool-srv__do_thing",
{},
user_id="user-1",
timeout=5,
)
payload = json.loads(str(exc_info.value))
assert payload["error"]["code"] == "mcp_consent_required"
_drain_background(mgr, loop)
entry = mgr._user_pool_entries[key]
assert entry.tools is not None, "drop must spare the mid-lookup re-consent"
assert entry.session is fresh_session
assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True
def test_decrypt_failure_does_not_emit_consent(
self, running_loop_mgr, storage: SQLiteBackend
) -> None:
+10 -19
View File
@@ -36,6 +36,7 @@ if TYPE_CHECKING:
from turnstone.core.oidc import OIDCConfig
from turnstone.core.log import get_logger
from turnstone.core.mcp_client import try_prime_user_pools
from turnstone.core.oidc import (
OIDC_STATE_TTL_SECONDS,
OIDCError,
@@ -2160,25 +2161,15 @@ async def handle_oidc_callback(request: Request, audience: str, cookie_name: str
# their listeners. Gated on a LIVE session: routine SSO
# re-logins by users with nothing open must not fan out
# mints and transport connects at deployment scale.
# Fire-and-forget; a failure changes nothing about login.
mcp_client = getattr(request.app.state, "mcp_client", None)
if (
mcp_client is not None
and hasattr(mcp_client, "prime_user_pools")
and hasattr(mcp_client, "has_live_session_listener")
):
# The gate call sits INSIDE the try: nothing on this
# best-effort path may fail the login (the credential
# is already persisted; the JWT is not yet issued).
try:
if mcp_client.has_live_session_listener(user["user_id"]):
mcp_client.prime_user_pools(user["user_id"])
except Exception:
log.debug(
"oidc.capture: post-capture pool prime scheduling failed user=%s",
user["user_id"],
exc_info=True,
)
# Fire-and-forget; a failure changes nothing about login
# (the credential is already persisted; the JWT is not
# yet issued) — the helper owns the swallow.
try_prime_user_pools(
getattr(request.app.state, "mcp_client", None),
user["user_id"],
require_live_listener=True,
context="oidc-capture",
)
# Load permissions and issue Turnstone JWT
perms = await asyncio.to_thread(_load_user_permissions, storage, user["user_id"])
+265 -84
View File
@@ -148,18 +148,59 @@ _MAX_PROMPTS_PER_SERVER = 1000
# thundering herd of connects on every session start.
_PRIME_MAX_CONCURRENCY = 4
# One row per pool ``*/list_changed`` kind: notification type → (kind label,
# refresh-method name). The pool notification handler drives all three catalog
# kinds through this table so the debounce-stamp/log/spawn protocol exists
# exactly once — a protocol change cannot silently diverge between kinds.
# Exact-type lookup is safe: the SDK's discriminated-union parser instantiates
# these concrete classes, never subclasses.
_POOL_LIST_CHANGED_REFRESHERS: dict[type, tuple[str, str]] = {
mcp_types.ToolListChangedNotification: ("tools", "_refresh_pool_server_tools"),
mcp_types.ResourceListChangedNotification: ("resources", "_refresh_pool_server_resources"),
mcp_types.PromptListChangedNotification: ("prompts", "_refresh_pool_server_prompts"),
# One row per pool ``*/list_changed`` kind: notification type → kind label.
# The pool notification handler drives all three catalog kinds through this
# table (plus its kind → bound-method map, which stays mypy-checked attribute
# access rather than a name-string table) so the debounce/coalesce/log/spawn
# protocol exists exactly once — a protocol change cannot silently diverge
# between kinds. Exact-type lookup is safe: the SDK's discriminated-union
# parser instantiates these concrete classes, never subclasses.
_POOL_LIST_CHANGED_KINDS: dict[type, str] = {
mcp_types.ToolListChangedNotification: "tools",
mcp_types.ResourceListChangedNotification: "resources",
mcp_types.PromptListChangedNotification: "prompts",
}
def try_prime_user_pools(
mcp_client: Any,
user_id: str | None,
*,
require_live_listener: bool = False,
context: str = "prime",
) -> None:
"""Best-effort per-user pool prime: schedule and swallow, never raise.
The ONE copy of the fire-and-forget prime idiom shared by the
session-construction, acting-user-change, and OIDC capture-success
call sites three hand-synced copies had already drifted.
``mcp_client`` is duck-typed (session and auth tests stub it): a
client without the prime surface, or a falsy ``user_id``, is a
silent no-op, matching the guards this replaces. Failures never
propagate a prime is an optimisation, and neither login nor
session construction may fail on it. ``require_live_listener``
gates on an open session to heal (the OIDC capture site: priming on
every routine SSO re-login would warm transports and mint tokens
for users with nothing open, at deployment scale).
"""
if not user_id or not hasattr(mcp_client, "prime_user_pools"):
return
try:
if require_live_listener and not (
hasattr(mcp_client, "has_live_session_listener")
and mcp_client.has_live_session_listener(user_id)
):
return
mcp_client.prime_user_pools(user_id)
except Exception:
log.debug(
"MCP pool prime scheduling failed (%s) user=%s",
context,
user_id,
exc_info=True,
)
# How long a node trusts its own pending-consent DELETE before re-running it on
# the next dispatch success. Bounds two things at once: the hot-path SQL rate
# (at most one DELETE per (user, server) per window) and the staleness of a
@@ -751,6 +792,13 @@ class MCPClientManager:
# so a noisy server in one user's pool doesn't suppress a refresh
# in another user's pool of the same server.
self._last_pool_notification_refresh: dict[tuple[str, str], float] = {}
# Coalescing markers for spawned ``list_changed`` refreshes, keyed
# ``((user_id, server), kind)``: set at spawn, cleared the moment
# the runner acquires ``open_lock`` (before its list call). While
# set, further notifications for the key+kind are dropped — the
# parked runner's fresh list will observe their change — bounding
# the lock's waiter queue at ONE parked runner per key+kind.
self._pool_refresh_pending: set[tuple[tuple[str, str], str]] = set()
# Pool tuning (read at construction; falls back to defaults).
mcp_cfg = load_config("mcp")
@@ -1866,10 +1914,9 @@ class MCPClientManager:
) -> None:
if not isinstance(msg, mcp_types.ServerNotification):
return
spec = _POOL_LIST_CHANGED_REFRESHERS.get(type(msg.root))
if spec is None:
kind = _POOL_LIST_CHANGED_KINDS.get(type(msg.root))
if kind is None:
return
kind, refresh_name = spec
now = time.monotonic()
last = self._last_pool_notification_refresh.get(key, 0.0)
if now - last < self._NOTIFICATION_DEBOUNCE:
@@ -1880,6 +1927,21 @@ class MCPClientManager:
now - last,
)
return
marker = (key, kind)
if marker in self._pool_refresh_pending:
# A runner for this key+kind is queued but has not yet
# issued its list call — it will observe this change
# when it runs. Skipping bounds ``open_lock``'s waiter
# queue at one parked runner per key+kind, so a
# notifying-but-slow server cannot accrete waiters that
# starve same-key dispatches and idle eviction.
log.debug(
"Coalescing pool %s notification user=%s server=%s (refresh queued)",
kind,
user_id,
server_name,
)
return
try:
# SPAWNED, never awaited: the SDK awaits notification
# handlers inline in its receive loop, so a handler that
@@ -1888,6 +1950,18 @@ class MCPClientManager:
# parked, every in-flight call on the session stalls
# with it, and the refresh only ever exits via its
# timeout. The refresh runs as its own tracked task.
#
# Bound at dispatch time (not a module-level name table)
# so instance-level overrides keep working and mypy
# checks the attribute references.
refreshers: dict[
str,
Callable[[tuple[str, str]], Awaitable[tuple[list[str], list[str]]]],
] = {
"tools": self._refresh_pool_server_tools,
"resources": self._refresh_pool_server_resources,
"prompts": self._refresh_pool_server_prompts,
}
log.info(
"Received %s/list_changed from pool user=%s server=%s",
kind,
@@ -1895,15 +1969,18 @@ class MCPClientManager:
server_name,
)
self._last_pool_notification_refresh[key] = now
self._pool_refresh_pending.add(marker)
self._spawn_background(
self._run_notification_refresh(key, kind, getattr(self, refresh_name)),
self._run_notification_refresh(key, kind, refreshers[kind]),
f"pool {kind} refresh for '{server_name}'",
)
except Exception as exc:
# Scheduling failed (loop shutting down). Structured
# fields only — ``exc_info=True`` would serialize the
# chained ``httpx.Request`` whose headers carry
# ``Authorization: Bearer <token>``.
# Scheduling failed (loop shutting down) — release the
# coalesce marker or this key+kind never refreshes again.
# Structured fields only — ``exc_info=True`` would
# serialize the chained ``httpx.Request`` whose headers
# carry ``Authorization: Bearer <token>``.
self._pool_refresh_pending.discard(marker)
log.warning(
"Pool refresh scheduling failed user=%s server=%s exc=%s",
user_id,
@@ -1928,14 +2005,26 @@ class MCPClientManager:
refresh for the same key, where the slower list call can
publish the older catalog last. Under the lock each refresh
issues its list call only after the previous publisher
finished, so the last publish is always the freshest. The wait
is bounded: lock holders are a connect/dispatch (one SDK call)
or another refresh (capped by ``_CONNECT_TIMEOUT``).
finished, so the last publish is always the freshest.
The coalesce marker (set by the handler at spawn) is cleared
the moment the lock is ACQUIRED, before the list call: a
notification landing during the in-flight list may announce a
change that list already missed, so it must spawn exactly one
successor which parks behind this lock. Together with the
handler's marker check this bounds the waiter queue at one
parked runner per key+kind: a same-key dispatch waits at most
two refresh timeouts, not an unbounded runner FIFO. The
``finally`` discard covers cancellation while parked.
The same-entry recheck under the lock discards a refresh whose
entry was replaced (full drop + re-create) while it waited
the notification belonged to the old transport and the
replacement published its own discovery.
replacement published its own discovery. A session evicted
while we were parked returns quietly for the same reason:
the reconnect's discovery republishes, and every teardown path
pops the debounce stamp, so the reconnected transport's first
notification refreshes immediately.
The debounce stamp (set at schedule time) deliberately SURVIVES
a failed refresh: popping it re-armed the handler on every
@@ -1943,9 +2032,7 @@ class MCPClientManager:
its notification rate. Keeping it caps attempts at one per
debounce window; a change announced during the remainder of a
failed window converges on the server's next ``list_changed``
or the entry's next reconnect (connects run full discovery, and
teardown pops the stamp so a reconnect's first notification
refreshes immediately).
or the entry's next reconnect.
``BaseExceptionGroup`` is caught alongside ``Exception``: a
wedged anyio transport surfaces session-op failures as groups,
@@ -1953,15 +2040,30 @@ class MCPClientManager:
``_spawn_background``'s failure log, whose ``exc_info``
serializes the chained ``httpx.Request`` headers carrying
``Authorization: Bearer <token>``.
Bounded residual, accepted: a server that keeps notifying while
every list call hangs to ``_CONNECT_TIMEOUT`` keeps THIS
entry's ``open_lock`` near-continuously occupied (one active +
one parked runner), deferring idle eviction of the entry (the
eviction pass skips a contested lock). The churn ends at the
first real dispatch (whose transport failure evicts the
session, after which parked runners bail on the session check),
at server recovery, or when the notifications stop; other
entries are unaffected (per-entry lock).
"""
user_id, server_name = key
marker = (key, kind)
entry = self._user_pool_entries.get(key)
if entry is None:
self._pool_refresh_pending.discard(marker)
return
try:
async with entry.open_lock:
self._pool_refresh_pending.discard(marker)
if self._user_pool_entries.get(key) is not entry:
return
if entry.session is None:
return
await refresh(key)
except (Exception, BaseExceptionGroup) as exc:
# Structured fields only — ``exc_info`` would serialize
@@ -1974,6 +2076,8 @@ class MCPClientManager:
server_name,
type(exc).__name__,
)
finally:
self._pool_refresh_pending.discard(marker)
async def _pool_transport_owner(
self,
@@ -2108,6 +2212,9 @@ class MCPClientManager:
if entry is None:
return
entry.drop_session()
# The debounce stamp must not outlive the transport: a
# reconnect's first ``list_changed`` refreshes immediately.
self._last_pool_notification_refresh.pop(key, None)
owner = entry.owner_task
close_requested = entry.close_requested
entry.owner_task = None
@@ -2157,6 +2264,9 @@ class MCPClientManager:
return
entry.owner_task = None
entry.close_requested = None
# The debounce stamp must not outlive the transport: the
# reconnect's first ``list_changed`` refreshes immediately.
self._last_pool_notification_refresh.pop(key, None)
if entry.session is not None:
entry.drop_session()
user_id, server_name = key
@@ -2657,6 +2767,10 @@ class MCPClientManager:
entry = self._user_pool_entries.get(key)
if entry is not None and entry.session is not None:
return # already connected — nothing to do
# Pre-lookup observation for the dead-grant drop below: at
# this point the session is None; one connected during the
# lookup's awaits must read as re-consent evidence.
observed_session = self._observed_pool_session(user_id, server_name)
if key in self._priming_keys:
return # a concurrent prime for this (user, server) is in flight
# Claim synchronously before any await — the mcp-loop is single-
@@ -2716,7 +2830,9 @@ class MCPClientManager:
# is durably gone must drop the retained
# catalog other live sessions still serve —
# e.g. a disconnect made on another node.
self._schedule_dead_grant_drop(lookup, key)
self._schedule_dead_grant_drop(
lookup, key, observed_session=observed_session
)
return # not consented / no credential / refresh failed — lazy paths handle it
if server_row is None:
server_row = await asyncio.to_thread(
@@ -2751,6 +2867,13 @@ class MCPClientManager:
# kind="missing". The oauth_user branch keeps its own deferred-row
# optimisation. On any read hiccup, fail open and let the
# per-server mint path classify it.
# Pre-read observation for the dead-grant drops below: a
# session connected DURING the credential read (a re-login
# capture prime racing this one) must read as re-consent
# evidence at drop time, not as the pre-observation session.
observed_obo = {
name: self._observed_pool_session(user_id, name) for name in self._obo_server_names
}
issuer = str(getattr(getattr(self._app_state, "oidc_config", None), "issuer", "") or "")
has_credential = False
if issuer:
@@ -2766,9 +2889,11 @@ class MCPClientManager:
if not has_credential:
# The credential is GONE — fire the same dead-grant
# convergence the per-server lookup would have
# (kind="missing") for any retained obo catalogs, or
# (kind="missing") for any retained obo entries, or
# ghosts survive every new-session prime: this gate
# skips exactly the lookup that would drop them.
# skips exactly the lookup that would drop them (the
# schedule-side guard skips entries with nothing to
# converge).
# Contract: the synthesized kind="missing" mirrors
# get_obo_access_token_classified's durably-absent-
# credential verdict (its missing-credential return
@@ -2776,10 +2901,11 @@ class MCPClientManager:
# ever splits — say a transient sub-case — update this
# synthesis with it.
for name in self._obo_server_names:
obo_key = (user_id, name)
obo_entry = self._user_pool_entries.get(obo_key)
if obo_entry is not None and self._entry_has_catalog(obo_entry):
self._schedule_dead_grant_drop(TokenLookupResult(kind="missing"), obo_key)
self._schedule_dead_grant_drop(
TokenLookupResult(kind="missing"),
(user_id, name),
observed_session=observed_obo.get(name),
)
prime_names -= self._obo_server_names
await asyncio.gather(*(_prime_one(s) for s in list(prime_names)))
@@ -3072,33 +3198,28 @@ 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:
def has_live_session_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.
is derived from this user's pool catalog. The ONE liveness
predicate: the eviction passes use it to decide cool-vs-drop
(#836), and outside callers (the OIDC capture-site prime) use
it to fan out work only for users who actually have an open
session to heal priming on every routine SSO re-login would
warm transports and mint tokens for users with nothing open, at
deployment scale. 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)
def has_live_session_listener(self, user_id: str) -> bool:
"""Public: does *user_id* have a live session's tool listener?
Lets outside callers (the OIDC capture-site prime) fan out work
only for users who actually have an open session to heal
priming on every routine SSO re-login would warm transports and
mint tokens for users with nothing open, at deployment scale.
"""
return self._user_has_live_listener(user_id)
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
a per-entry ``has_live_session_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
@@ -3176,7 +3297,7 @@ class MCPClientManager:
return False
if live_uids is not None:
return user_id in live_uids
return self._user_has_live_listener(user_id)
return self.has_live_session_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
@@ -3320,10 +3441,6 @@ class MCPClientManager:
if entry.in_flight > 0:
return
await self._teardown_pool_entry(key)
# 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._retain_cooled(key, entry):
# Cooled: entry + catalog stay, so the live session's
@@ -3682,11 +3799,14 @@ class MCPClientManager:
Mirror of :meth:`_refresh_pool_server_tools` for the resource
path (RFC §3.2). Skips when ``supports_resources`` is False so
a server that no longer advertises resources doesn't trigger
a list call. ``asyncio.timeout`` is mandatory see R6.
a list call. ``asyncio.timeout`` is mandatory a wedged server
must not hang the refresh (and the lock it holds) forever.
MUST run on the mcp-loop. Caller does not need to hold
``open_lock``; this is invoked from the pool notification
handler running on the SDK's receive task.
MUST run on the mcp-loop, with the entry's ``open_lock`` HELD
(:meth:`_run_notification_refresh` acquires it): an unlocked
refresh races a connect's discovery wiring and sibling
refreshes, either of which can publish an older catalog over
this one see the runner's ordering rationale.
"""
entry = self._user_pool_entries.get(key)
if entry is None or entry.session is None or not entry.supports_resources:
@@ -3757,11 +3877,14 @@ class MCPClientManager:
Mirror of :meth:`_refresh_pool_server_tools` for the prompt
path (RFC §3.3). Skips when ``supports_prompts`` is False so a
server that no longer advertises prompts doesn't trigger a list
call. ``asyncio.timeout`` is mandatory see R6.
call. ``asyncio.timeout`` is mandatory a wedged server must
not hang the refresh (and the lock it holds) forever.
MUST run on the mcp-loop. Caller does not need to hold
``open_lock``; this is invoked from the pool notification
handler running on the SDK's receive task.
MUST run on the mcp-loop, with the entry's ``open_lock`` HELD
(:meth:`_run_notification_refresh` acquires it): an unlocked
refresh races a connect's discovery wiring and sibling
refreshes, either of which can publish an older catalog over
this one see the runner's ordering rationale.
"""
entry = self._user_pool_entries.get(key)
if entry is None or entry.session is None or not entry.supports_prompts:
@@ -6525,10 +6648,16 @@ class MCPClientManager:
# this retry; the local cached token is the one the AS just
# rejected, so reading it back without ``force_refresh=True``
# would re-attempt with the same (rejected) bearer.
# Observe the session BEFORE the lookup's awaits: a session the
# consent-completion prime connects DURING the lookup must read
# as re-consent evidence at drop time, not as pre-observation.
observed_session = self._observed_pool_session(user_id, server_name)
lookup: TokenLookupResult = await self._pool_token_lookup(
server_row, user_id, server_name, force_refresh=retry_count > 0
)
lookup_error = self._pool_lookup_failure(lookup, server_name, server_row, user_id)
lookup_error = self._pool_lookup_failure(
lookup, server_name, server_row, user_id, observed_session=observed_session
)
if lookup_error is not None:
return lookup_error
access_token = lookup.token or ""
@@ -6671,10 +6800,16 @@ class MCPClientManager:
if self._app_state is None:
raise RuntimeError("Pool dispatch requires set_app_state() to have been called")
# Observe the session BEFORE the lookup's awaits: a session the
# consent-completion prime connects DURING the lookup must read
# as re-consent evidence at drop time, not as pre-observation.
observed_session = self._observed_pool_session(user_id, server_name)
lookup: TokenLookupResult = await self._pool_token_lookup(
server_row, user_id, server_name, force_refresh=retry_count > 0
)
lookup_error = self._pool_lookup_failure(lookup, server_name, server_row, user_id)
lookup_error = self._pool_lookup_failure(
lookup, server_name, server_row, user_id, observed_session=observed_session
)
if lookup_error is not None:
return lookup_error
access_token = lookup.token or ""
@@ -6795,10 +6930,16 @@ class MCPClientManager:
if self._app_state is None:
raise RuntimeError("Pool dispatch requires set_app_state() to have been called")
# Observe the session BEFORE the lookup's awaits: a session the
# consent-completion prime connects DURING the lookup must read
# as re-consent evidence at drop time, not as pre-observation.
observed_session = self._observed_pool_session(user_id, server_name)
lookup: TokenLookupResult = await self._pool_token_lookup(
server_row, user_id, server_name, force_refresh=retry_count > 0
)
lookup_error = self._pool_lookup_failure(lookup, server_name, server_row, user_id)
lookup_error = self._pool_lookup_failure(
lookup, server_name, server_row, user_id, observed_session=observed_session
)
if lookup_error is not None:
return lookup_error
access_token = lookup.token or ""
@@ -6991,12 +7132,28 @@ class MCPClientManager:
return False
return _pool_lookup_verdict(lookup) == "mcp_consent_required"
def _observed_pool_session(self, user_id: str, server_name: str) -> Any:
"""Snapshot the entry's session for a dead-grant observation.
MUST be called BEFORE the classified lookup's first await (the
callers' contract with :meth:`_schedule_dead_grant_drop`): a
snapshot taken after those awaits can capture a session the
consent-completion prime connected mid-lookup, and the drop
would then evict the just-restored catalog it was designed to
spare the session has to predate the failure to count as the
pre-observation transport.
"""
entry = self._user_pool_entries.get((user_id, server_name))
return entry.session if entry is not None else None
def _pool_lookup_failure(
self,
lookup: TokenLookupResult,
server_name: str,
server_row: dict[str, Any],
user_id: str,
*,
observed_session: Any,
) -> str | None:
"""Render a failed pool lookup AND pair it with its convergence drop.
@@ -7005,24 +7162,34 @@ class MCPClientManager:
would keep offering revoked tools behind a consent card the
cross-node non-convergence #836 fixes, silently split by
dispatch surface. Returns ``None`` exactly when *lookup*
carries a usable bearer.
carries a usable bearer. ``observed_session`` is the caller's
PRE-lookup snapshot (:meth:`_observed_pool_session`).
"""
lookup_error = _pool_lookup_error(lookup, server_name, server_row)
if lookup_error is not None:
self._schedule_dead_grant_drop(lookup, (user_id, server_name))
self._schedule_dead_grant_drop(
lookup, (user_id, server_name), observed_session=observed_session
)
return lookup_error
def _schedule_dead_grant_drop(self, lookup: TokenLookupResult, key: tuple[str, str]) -> None:
def _schedule_dead_grant_drop(
self,
lookup: TokenLookupResult,
key: tuple[str, str],
*,
observed_session: Any,
) -> None:
"""Converge live sessions with a durably-gone grant (#836).
Called on the mcp-loop wherever a classified lookup fails the
three dispatchers and session-start priming. When
:meth:`_lookup_grant_dead` confirms the grant is gone (revoked,
disconnected, or unlinked not an infrastructure blip), the
(user, server) catalog is dropped so the tools leave the user's
live sessions instead of dangling behind a consent card for
access that no longer exists. Re-consent restores them via the
consent-completion prime; obo re-login via the capture prime.
three dispatchers, session-start priming, and the obo
credential gate. When :meth:`_lookup_grant_dead` confirms the
grant is gone (revoked, disconnected, or unlinked not an
infrastructure blip), the (user, server) catalog is dropped so
the tools leave the user's live sessions instead of dangling
behind a consent card for access that no longer exists.
Re-consent restores them via the consent-completion prime; obo
re-login via the capture prime.
SCHEDULED, never awaited: the locked drop can park behind a
same-key dispatch holding ``open_lock`` across its entire SDK
@@ -7030,21 +7197,35 @@ class MCPClientManager:
timeout and charged the breaker they are documented to bypass.
``_spawn_background`` tracks the task so shutdown cancels it.
The entry's CURRENT session is snapshotted here, at the moment
the dead grant was observed, and handed to the drop as
``observed_session``: a warm session that already existed when
the lookup failed predates the revocation and must be evicted
(the #836 warm case — dispatches short-circuit on the failed
lookup without ever touching the session, so no 401 path would
converge it). Only a session that CHANGED after this snapshot
proves a later successful connect, i.e. a live grant.
``observed_session`` is REQUIRED and must be snapshotted by the
caller BEFORE the classified lookup's first await
(:meth:`_observed_pool_session`): a warm session that already
existed at observation time predates the revocation and must be
evicted (the #836 warm case — dispatches short-circuit on the
failed lookup without ever touching the session, so no 401 path
would converge it), while only a session that CHANGED after the
observation proves a later successful connect. Snapshotting
here, after the lookup returned, would capture a session the
consent-completion prime connected mid-lookup and evict the
just-restored catalog.
Skips when there is provably nothing to converge no entry, or
a session-less entry with no catalog. At scale every
unconsented server classifies as a dead grant on every prime;
spawning a tracked no-op task per (user, unconsented server)
for that is pure mcp-loop overhead. A catalog appearing AFTER
this check implies a successful connect, i.e. a lookup that
succeeded after this one failed.
"""
if not self._lookup_grant_dead(lookup):
return
entry = self._user_pool_entries.get(key)
observed = entry.session if entry is not None else None
if entry is None or (entry.session is None and not self._entry_has_catalog(entry)):
return
self._spawn_background(
self._drop_catalog_locked(key, skip_if_connected=True, observed_session=observed),
self._drop_catalog_locked(
key, skip_if_connected=True, observed_session=observed_session
),
f"dead-grant catalog drop for '{key[1]}'",
)
+9 -20
View File
@@ -70,6 +70,7 @@ from turnstone.core.lowering import (
tool_args_preview,
wire_valid_arguments,
)
from turnstone.core.mcp_client import try_prime_user_pools
from turnstone.core.memory import (
count_messages,
count_structured_memories,
@@ -1689,6 +1690,10 @@ class ChatSession:
# listeners register so tool/resource/prompt refreshes flow
# through to this session.
# * interactive (no mcp) — INTERACTIVE_TOOLS.
# Construction-scoped (a local, NOT instance state): the seq
# value at the authoritative merged read, compared once at the
# end of tool setup. Only ``_mcp_tools_change_seq`` lives on.
mcp_tools_seq_at_read = 0
if kind == WorkstreamKind.COORDINATOR:
self._tools = list(COORDINATOR_TOOLS)
self._task_tools = []
@@ -1727,7 +1732,7 @@ class ChatSession:
# attributes and be swallowed by the fan-out, losing its
# effect, and one that succeeds here would be clobbered by
# the tool-search construction below reading mixed state.
self._mcp_tools_seq_at_read = self._mcp_tools_change_seq
mcp_tools_seq_at_read = self._mcp_tools_change_seq
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)
@@ -1737,15 +1742,7 @@ class ChatSession:
# listeners registered just above deliver the catalog to this
# session once each prime completes. No-op for users with no
# consented oauth_user servers.
if self._mcp_user_id and hasattr(self._mcp_client, "prime_user_pools"):
try:
self._mcp_client.prime_user_pools(self._mcp_user_id)
except Exception:
log.debug(
"mcp prime_user_pools scheduling failed user=%s",
self._mcp_user_id,
exc_info=True,
)
try_prime_user_pools(self._mcp_client, self._mcp_user_id, context="session-start")
else:
self._tools = INTERACTIVE_TOOLS
self._task_tools = TASK_AGENT_TOOLS
@@ -1803,7 +1800,7 @@ class ChatSession:
if (
self._kind != WorkstreamKind.COORDINATOR
and self._mcp_client
and self._mcp_tools_change_seq != self._mcp_tools_seq_at_read
and self._mcp_tools_change_seq != mcp_tools_seq_at_read
):
self._on_mcp_tools_changed()
# Skill: explicit name overrides is_default skills. ``skill_arguments``
@@ -5652,15 +5649,7 @@ class ChatSession:
mcp.remove_prompt_listener(self._mcp_prompt_cb, user_id=old_listener_uid)
mcp.add_prompt_listener(self._mcp_prompt_cb, user_id=new_listener_uid)
self._mcp_listener_user_id = new_listener_uid
if new_listener_uid and hasattr(mcp, "prime_user_pools"):
try:
mcp.prime_user_pools(new_listener_uid)
except Exception:
log.debug(
"mcp prime_user_pools scheduling failed user=%s",
new_listener_uid,
exc_info=True,
)
try_prime_user_pools(mcp, new_listener_uid, context="acting-user-change")
# Rebuild the merged tool list and resource/prompt-dependent
# state under the new identity NOW — the prime above completes
# asynchronously and only notifies on catalog changes, while