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