mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-13 23:42:25 -06:00
feat(mcp): per-user catalog scoping (Phase 7 — tools)
Light up production reachability of pool dispatch (RFC §3, invariant 8)
by widening the public catalog API to optionally take a ``user_id``:
- ``MCPClientManager.get_tools(user_id=None)`` returns the merged
static + per-user pool view when ``user_id`` is supplied; the default
preserves the legacy global-only contract.
- ``is_mcp_tool(name, *, user_id=None)`` extends the lookup to the
per-user ``_user_tool_map``. Pool tools become reachable from
``ChatSession._prepare_tool`` only when the session-bound user_id
flows through — flipping invariant 8 from "must hold" to "satisfied".
- Listener identity becomes ``(user_id, callback)``. Static-path
changes fire ALL listeners (admin + every user); pool-entry
changes fire only matching-user + admin (``None``) listeners.
RFC §3.3.
- Pool sessions discover their tool list on first connect
(``_connect_one_pool`` → ``await session.list_tools()``); the
notification closure binds to ``(user_id, server_name)`` so
push-driven ``list_changed`` updates target the correct user's
catalog. R6 verified empirically: ``list_tools()`` 401 propagates
through anyio TaskGroup unwinding, no hang — plain ``await`` is
fine, no carrier-race shape needed for discovery.
- ``_evict_session`` drops ``entry.tools`` and rebuilds the user's
index so an evicted-then-reconnected session doesn't carry
stale catalog state.
- ``web_search.resolve_web_search_client`` refuses
``auth_type=oauth_user`` backends (per-node web search can't
carry per-user tokens).
Resources / prompts pool dispatch deferred to Phase 7b — invariant 8
is satisfied by the tool path alone, and the resource/prompt path
needs sibling ``_dispatch_pool_resource_sync`` /
``_dispatch_pool_prompt_sync`` helpers each with their own
carrier-race plumbing (~400 LOC). Phase 7b will follow the patterns
established here.
CLI sessions default ``user_id=""`` and so cannot use oauth_user
MCP servers — documented limitation; users must use the web UI.
Round-1 review fixes (4-finder review applied, no push yet):
- bug-1: get_tools(user_id) was iterating _user_pool_entries from sync
threads while the mcp-loop concurrently mutated it (RuntimeError:
dictionary changed size during iteration). Now reads from a sibling
_user_tools dict updated atomically by _rebuild_user_tool_map.
- bug-2: _close_pool_entry_if_idle (LRU/TTL eviction) skipped the
catalog cleanup that _evict_session does — stale tools persisted
in _user_tool_map and ChatSession's tool list never rebuilt. Now
mirrors _evict_session.
- perf-1: _last_pool_notification_refresh debounce dict was never
pruned in either eviction path. Now popped alongside the entry.
- perf-3: web_search resolver was issuing a sync SQL query per LLM
turn to gate oauth_user backends. Now reads from the cached
in-memory config.
- sec-1: bearer token could leak into exc_info-rendered tracebacks
via Sentry/faulthandler. log.debug now uses structured fields,
not exc_info.
- sec-2: tools-per-server response now capped at 1000 (defensive,
mirrors _MAX_ERROR_LEN / _MAX_INSUFFICIENT_SCOPE_REPORTED).
- Test cleanup: dropped two listener fan-out tests duplicating
test_mcp_client.py coverage; renamed test_pool_session_notification_handler
to match its actual scope (_refresh_pool_server_tools); removed
stale comments referencing /tmp/r6-spike*.py scratchpads and a
misleading "copy-on-write" comment.
Round-2 pre-push review fixes (focused single-pass review applied):
- round2-1: bug-2's catalog-cleanup block in _close_pool_entry_if_idle
had no integration test (exactly the failure mode flagged in
feedback_tests_through_boundaries.md). Added
test_close_pool_entry_if_idle_clears_catalog_and_fires_listener
driving the LRU/TTL eviction path through real streamablehttp_client +
MockTransport. Negative-test verified: reverting the
_rebuild_user_tool_map / _notify_user_tool_listeners calls makes
the new test fail.
- round2-3: documented the _oauth_user_server_names cache invariant
in add_server_sync / remove_server_sync docstrings. Cache is
reconcile_sync's sole owner — direct callers leave it stale, but
_db_servers_to_config strips oauth_user rows so production paths
are unaffected. Static→oauth_user transitions correctly leave the
name in the cache because remove_server_sync drops the static
connection, not the cache identity.
- round2-6: strengthened test_rebuild_user_tool_map_populates and
test_rebuild_user_tool_map_drops_empty_user to assert on the
_user_tools sibling cache (bug-1 fix). Without this, a future
revert dropping the sibling write would still pass the unit
tests because get_tools coverage lives in separate tests.
Round-3 full-stack review fixes (multi-stage review on the final
state caught what the layered apply passes missed):
- q-1 REGRESSION: pool tool-discovery used asyncio.wait_for around
session.list_tools(), the exact pattern the f6a3b66 fix (and
feedback_asyncio_timeout_vs_wait_for.md) put in place to avoid.
Python 3.11's asyncio.wait_for wraps the inner coroutine in a
fresh task → cross-task scope-exit when the SDK's anyio TaskGroup
unwinds on a 401. Switched to `async with asyncio.timeout(...):`
pattern used by _safe_close_stack.
- sec-2: TOCTOU in _connect_one_pool — entry.tools was published
(via _rebuild_user_tool_map + listener fan-out) BEFORE entry.session
was assigned. A sync-thread reader could observe a tool whose
backing entry has session=None. Defence-in-depth — dispatch
re-fetches its own token and lazy-reconnects on session=None — but
reordering catches the race at the source. entry.session now
publishes BEFORE catalog visibility.
- bug-1: _close_pool_entry_if_idle's _user_pool_locks.pop ran
unconditionally after the try/finally, but the early-return
branches (entry None on re-check, in_flight > 0 under lock) skip
it via Python's return-through-finally semantics. The lock was
never popped on those paths. Now gated behind an `evicted` flag
set only on the success path; in_flight > 0 leaves the lock for
the active dispatcher to reuse, entry-None races leave the lock
for re-allocation by _ensure_pool_entry. Comment now describes
the actual semantics, not the original promise.
- bug-2: softened the _rebuild_user_tool_map docstring's atomicity
claim. The two-dict write is technically non-atomic across Python
statements; in practice the window is sub-microsecond on the
mcp-loop with no awaits between writes, and the listener fan-out
fires AFTER both writes complete. Docstring now says "back-to-back
on the mcp-loop" instead of "atomically alongside".
- q-3: dropped `hasattr(mcp_client, "server_auth_type")` defensive
check in web_search.py. The method ships in this commit; the
hasattr created a silent fallthrough that would let a future
rename silently re-enable oauth_user backends.
- q-4: surfaced the CLI / empty-user_id limitation in a docstring
comment at ChatSession.__init__'s self._user_id assignment. The
note previously lived only inside is_mcp_tool's docstring — a
future maintainer wiring CLI features against MCP pool servers
wouldn't think to read is_mcp_tool to find the constraint.
- q-2 + q-5: deleted a tautological duplicate test in
test_mcp_user_catalog.py whose docstring claimed to test
ChatSession.close but never instantiated a ChatSession (the
manager-level identity semantics are already covered by
test_listener_identity_includes_user_id in the same file and by
test_session_close_removes_listener_with_same_user_id in
test_mcp_client.py which DOES drive a ChatSession). Reworded a
misleading "fixture provides only 5s" comment to point at the
actual `_run_on_loop(..., timeout=5)` site.
- q-6: the `self._user_id or None` collapse repeated at 8 sites
across session.py. Cached once at __init__ as
``self._mcp_user_id`` (since ``_user_id`` is set once and never
mutated); 8 call sites now read the cached value. The empty-
string-is-CLI-sentinel invariant is documented at the assignment
site, not re-asserted at each consumer.
Deferred to follow-up:
- sec-1: a hostile MCP server bound to user-A could craft a
tool.name containing `__` to synthesize a prefixed-name collision
in user-A's own catalog. Bounded impact: cross-tenant dispatch is
prevented by the per-tenant token gate in _dispatch_pool, and
user-B's get_tools(user_id="B") never includes user-A's pool
entries. The fix needs policy decisions (reject vs. sanitize)
and touches _mcp_to_openai which is shared between static and
pool paths; better discussed in its own follow-up where the
policy applies uniformly to static-path servers too. The threat
model already requires user-A to have consented to a malicious
server, who has many more dangerous vectors than tool-name
shenanigans.
Test count delta: +31 tests (5435 → 5466, ``-m "not live"``; one
test deleted in round-3 apply per q-2):
- ``tests/test_mcp_client.py`` +20 (per-user catalog state, listener
identity, session thread-through)
- ``tests/test_mcp_user_catalog.py`` +9 NEW (integration tests
driving real ``streamablehttp_client`` + ``httpx.MockTransport`` per
invariant 14: discovery on connect, user isolation, eviction +
reconnect, LRU/TTL eviction (round2-1), R6 401-propagation
regression, static byte-identical canonical regression; review
passes dropped duplicate listener fan-out tests from earlier
drafts whose coverage lived in test_mcp_client.py)
- ``tests/test_web_search.py`` +2 (oauth_user backend rejection +
static backend acceptance regression; updated to use the new
``server_auth_type`` in-memory accessor)
This commit is contained in:
Patrick Buckley
@@ -397,6 +397,217 @@ class TestMCPClientManager:
|
||||
mgr = MCPClientManager({})
|
||||
mgr.shutdown() # should be a no-op
|
||||
|
||||
# -- Phase 7: per-user catalog scoping ---------------------------------
|
||||
|
||||
def test_is_mcp_tool_user_id_default_none_unchanged(self):
|
||||
"""Sanity: default ``user_id=None`` answers static-only.
|
||||
|
||||
The legacy single-arg call still works, and unknown names still
|
||||
return False — Phase 7 adds an optional keyword without
|
||||
rewriting the static-path semantics.
|
||||
"""
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tool_map["mcp__static__list"] = ("static", "list")
|
||||
# Legacy single-arg call still works.
|
||||
assert mgr.is_mcp_tool("mcp__static__list") is True
|
||||
assert mgr.is_mcp_tool("mcp__static__list", user_id=None) is True
|
||||
assert mgr.is_mcp_tool("nonexistent") is False
|
||||
assert mgr.is_mcp_tool("nonexistent", user_id=None) is False
|
||||
|
||||
def test_is_mcp_tool_user_keyed_pool_tool(self):
|
||||
"""A name visible only via ``_user_tool_map`` resolves only for
|
||||
the matching ``user_id``.
|
||||
|
||||
Verifies the new branch: ``_tool_map`` miss + ``user_id`` hit.
|
||||
"""
|
||||
mgr = MCPClientManager({})
|
||||
mgr._user_tool_map["user-1"] = {
|
||||
"mcp__pool-srv__do": ("pool-srv", "do"),
|
||||
}
|
||||
# Visible to user-1.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id="user-1") is True
|
||||
# Invisible to None caller (admin / web-search backend resolution).
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id=None) is False
|
||||
# Invisible to a different user.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id="user-2") is False
|
||||
|
||||
def test_is_mcp_tool_static_wins_for_any_user(self):
|
||||
"""Static-path tools are visible regardless of ``user_id`` —
|
||||
the merged view is ``static ∪ user-pool``."""
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tool_map["mcp__static__list"] = ("static", "list")
|
||||
# Even for an unknown user, a static tool is still reachable —
|
||||
# static-path is process-global.
|
||||
assert mgr.is_mcp_tool("mcp__static__list", user_id="user-1") is True
|
||||
assert mgr.is_mcp_tool("mcp__static__list", user_id="anybody") is True
|
||||
|
||||
def test_get_tools_user_id_none_returns_static_only(self):
|
||||
"""``get_tools(user_id=None)`` returns the global static catalog.
|
||||
|
||||
Pool tools are NEVER included in the default-arg view — that's
|
||||
the legacy contract every pre-Phase-7 caller relies on.
|
||||
"""
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tools = [_fake_openai_tool("mcp__static__list")]
|
||||
# Seed a pool entry that should NOT appear in the default view.
|
||||
from turnstone.core.mcp_client import PoolEntryState
|
||||
|
||||
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
|
||||
entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
|
||||
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
|
||||
|
||||
tools = mgr.get_tools()
|
||||
names = [t["function"]["name"] for t in tools]
|
||||
assert names == ["mcp__static__list"]
|
||||
|
||||
def test_get_tools_user_id_merges_pool(self):
|
||||
"""``get_tools(user_id='user-1')`` merges static + that user's pool tools.
|
||||
|
||||
Other users' pool entries MUST NOT leak into the result —
|
||||
privacy / RBAC invariant.
|
||||
"""
|
||||
from turnstone.core.mcp_client import PoolEntryState
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tools = [_fake_openai_tool("mcp__static__list")]
|
||||
e1 = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
|
||||
e1.tools = [_fake_openai_tool("mcp__pool-srv__do")]
|
||||
e2 = PoolEntryState(key=("user-2", "pool-srv"), open_lock=MagicMock())
|
||||
e2.tools = [_fake_openai_tool("mcp__pool-srv__other")]
|
||||
mgr._user_pool_entries[("user-1", "pool-srv")] = e1
|
||||
mgr._user_pool_entries[("user-2", "pool-srv")] = e2
|
||||
# Production invariant: ``_connect_one_pool`` /
|
||||
# ``_refresh_pool_server_tools`` / ``_evict_session`` /
|
||||
# ``_close_pool_entry_if_idle`` all call ``_rebuild_user_tool_map``
|
||||
# immediately after mutating ``_user_pool_entries``. Tests that
|
||||
# seed pool entries directly must mirror that invariant —
|
||||
# ``get_tools(user_id=...)`` reads from the ``_user_tools``
|
||||
# snapshot (built by ``_rebuild_user_tool_map``), never iterating
|
||||
# ``_user_pool_entries`` directly.
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
mgr._rebuild_user_tool_map("user-2")
|
||||
|
||||
u1_names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-1")]
|
||||
assert sorted(u1_names) == ["mcp__pool-srv__do", "mcp__static__list"]
|
||||
|
||||
u2_names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-2")]
|
||||
assert sorted(u2_names) == ["mcp__pool-srv__other", "mcp__static__list"]
|
||||
|
||||
# Default still global-only — unaffected by either user's entries.
|
||||
default_names = [t["function"]["name"] for t in mgr.get_tools()]
|
||||
assert default_names == ["mcp__static__list"]
|
||||
|
||||
def test_get_tools_user_id_returns_copies(self):
|
||||
"""Mirror existing ``test_get_tools_returns_copy``: caller mutation
|
||||
of the returned list MUST NOT affect the manager's catalog."""
|
||||
from turnstone.core.mcp_client import PoolEntryState
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tools = [_fake_openai_tool("mcp__static__a")]
|
||||
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
|
||||
entry.tools = [_fake_openai_tool("mcp__pool-srv__b")]
|
||||
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
|
||||
# See note in ``test_get_tools_user_id_merges_pool``.
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
|
||||
tools = mgr.get_tools(user_id="user-1")
|
||||
assert len(tools) == 2
|
||||
tools.clear()
|
||||
# Re-fetch — original catalog unchanged.
|
||||
assert len(mgr.get_tools(user_id="user-1")) == 2
|
||||
|
||||
def test_get_tools_user_with_none_tools_skipped(self):
|
||||
"""A pool entry that hasn't completed discovery (``entry.tools is None``)
|
||||
contributes no tools — the merged view skips it cleanly."""
|
||||
from turnstone.core.mcp_client import PoolEntryState
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tools = [_fake_openai_tool("mcp__static__a")]
|
||||
# Brand-new pool entry, discovery not yet run.
|
||||
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
|
||||
assert entry.tools is None
|
||||
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
|
||||
# Rebuild observes ``entry.tools is None`` and skips this entry.
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
|
||||
names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-1")]
|
||||
assert names == ["mcp__static__a"]
|
||||
|
||||
def test_rebuild_user_tool_map_populates(self):
|
||||
"""``_rebuild_user_tool_map`` materializes the per-user index from
|
||||
pool entries owned by that user."""
|
||||
from turnstone.core.mcp_client import PoolEntryState
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
|
||||
entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
|
||||
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
|
||||
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
assert mgr._user_tool_map["user-1"] == {"mcp__pool-srv__do": ("pool-srv", "do")}
|
||||
# Sibling _user_tools cache (bug-1 fix) MUST be populated alongside
|
||||
# the map — otherwise get_tools(user_id="user-1") would silently
|
||||
# return the static-only view despite is_mcp_tool returning True.
|
||||
assert mgr._user_tools["user-1"] == [_fake_openai_tool("mcp__pool-srv__do")]
|
||||
|
||||
def test_rebuild_user_tool_map_drops_empty_user(self):
|
||||
"""Rebuilding for a user with no pool entries removes the key
|
||||
rather than retaining an empty-dict sentinel."""
|
||||
from turnstone.core.mcp_client import PoolEntryState
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
|
||||
entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
|
||||
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
assert "user-1" in mgr._user_tool_map
|
||||
assert "user-1" in mgr._user_tools
|
||||
|
||||
# Drop the entry, rebuild — user_id key should be removed from BOTH
|
||||
# the map and the sibling tool list (bug-1 fix). A drop in only one
|
||||
# would leave get_tools and is_mcp_tool out of sync.
|
||||
mgr._user_pool_entries.clear()
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
assert "user-1" not in mgr._user_tool_map
|
||||
assert "user-1" not in mgr._user_tools
|
||||
|
||||
def test_rebuild_user_tool_map_isolates_users(self):
|
||||
"""Rebuilding for ``user-1`` MUST NOT touch ``user-2``'s entry."""
|
||||
from turnstone.core.mcp_client import PoolEntryState
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
e1 = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
|
||||
e1.tools = [_fake_openai_tool("mcp__pool-srv__one")]
|
||||
e2 = PoolEntryState(key=("user-2", "pool-srv"), open_lock=MagicMock())
|
||||
e2.tools = [_fake_openai_tool("mcp__pool-srv__two")]
|
||||
mgr._user_pool_entries[("user-1", "pool-srv")] = e1
|
||||
mgr._user_pool_entries[("user-2", "pool-srv")] = e2
|
||||
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
mgr._rebuild_user_tool_map("user-2")
|
||||
assert mgr._user_tool_map["user-1"] == {"mcp__pool-srv__one": ("pool-srv", "one")}
|
||||
assert mgr._user_tool_map["user-2"] == {"mcp__pool-srv__two": ("pool-srv", "two")}
|
||||
|
||||
# Clear user-1's entry only; rebuild user-1; user-2 must remain.
|
||||
mgr._user_pool_entries.pop(("user-1", "pool-srv"))
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
assert "user-1" not in mgr._user_tool_map
|
||||
assert mgr._user_tool_map["user-2"] == {"mcp__pool-srv__two": ("pool-srv", "two")}
|
||||
|
||||
def test_rebuild_user_tool_map_does_not_touch_static(self):
|
||||
"""Invariant 1: per-user rebuild must NOT mutate ``_tool_map``."""
|
||||
from turnstone.core.mcp_client import PoolEntryState
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tool_map["mcp__static__list"] = ("static", "list")
|
||||
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
|
||||
entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
|
||||
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
|
||||
|
||||
before = dict(mgr._tool_map)
|
||||
mgr._rebuild_user_tool_map("user-1")
|
||||
assert mgr._tool_map == before
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session integration (mock MCP client)
|
||||
@@ -597,6 +808,84 @@ class TestSessionIntegration:
|
||||
assert "MCP tool error" in output
|
||||
assert "server crashed" in output
|
||||
|
||||
# -- Phase 7: per-user catalog scoping ---------------------------------
|
||||
|
||||
def test_session_passes_user_id_to_get_tools(self, tmp_db):
|
||||
"""ChatSession threads its ``user_id`` into ``get_tools`` so the
|
||||
merged static + pool view is scoped to the session's user."""
|
||||
mock_mcp = MagicMock()
|
||||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||||
self._make_session(mcp_client=mock_mcp, user_id="user-7")
|
||||
mock_mcp.get_tools.assert_called_with(user_id="user-7")
|
||||
|
||||
def test_session_get_tools_empty_user_id_passes_none(self, tmp_db):
|
||||
"""Sentinel ``user_id=""`` (CLI / service / unknown) collapses to
|
||||
``user_id=None`` so the static-only view is returned."""
|
||||
mock_mcp = MagicMock()
|
||||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||||
self._make_session(mcp_client=mock_mcp, user_id="")
|
||||
mock_mcp.get_tools.assert_called_with(user_id=None)
|
||||
|
||||
def test_session_passes_user_id_to_add_listener(self, tmp_db):
|
||||
"""ChatSession registers its tool-change listener under its own
|
||||
``user_id`` so pool-only changes for OTHER users do not fire it."""
|
||||
mock_mcp = MagicMock()
|
||||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||||
self._make_session(mcp_client=mock_mcp, user_id="user-7")
|
||||
# ``add_listener`` was called with ``user_id="user-7"``.
|
||||
listener_calls = mock_mcp.add_listener.call_args_list
|
||||
assert listener_calls, "ChatSession did not register a tool listener"
|
||||
first_call = listener_calls[0]
|
||||
assert first_call.kwargs.get("user_id") == "user-7"
|
||||
|
||||
def test_session_close_removes_listener_with_same_user_id(self, tmp_db):
|
||||
"""R4 critical: register and remove MUST agree on ``user_id`` —
|
||||
the listener identity is ``(user_id, callback)``."""
|
||||
mock_mcp = MagicMock()
|
||||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||||
session = self._make_session(mcp_client=mock_mcp, user_id="user-7")
|
||||
|
||||
session.close()
|
||||
# ``remove_listener`` must be called with the same ``user_id``.
|
||||
remove_calls = mock_mcp.remove_listener.call_args_list
|
||||
assert remove_calls, "ChatSession.close did not unregister a tool listener"
|
||||
first_remove = remove_calls[0]
|
||||
assert first_remove.kwargs.get("user_id") == "user-7"
|
||||
# And the callback identity must match what was registered.
|
||||
registered_cb = mock_mcp.add_listener.call_args_list[0].args[0]
|
||||
removed_cb = first_remove.args[0]
|
||||
assert registered_cb is removed_cb
|
||||
|
||||
def test_session_unknown_tool_lists_user_scoped_catalog(self, tmp_db):
|
||||
"""The "Unknown tool" error message lists tools the session can
|
||||
actually invoke — drawn from the merged user-scoped catalog,
|
||||
not the manager's private static-only ``_tool_map``."""
|
||||
mock_mcp = MagicMock()
|
||||
# Pretend the user's merged view contains a static + pool entry.
|
||||
mock_mcp.get_tools.return_value = [
|
||||
_fake_openai_tool("mcp__static__list"),
|
||||
_fake_openai_tool("mcp__pool-srv__do"),
|
||||
]
|
||||
mock_mcp.is_mcp_tool.return_value = False
|
||||
session = self._make_session(mcp_client=mock_mcp, user_id="user-7")
|
||||
# Reset the call counter so we observe only the _prepare_tool call.
|
||||
mock_mcp.get_tools.reset_mock()
|
||||
|
||||
tc = {
|
||||
"id": "call_unknown",
|
||||
"function": {"name": "no_such_tool", "arguments": "{}"},
|
||||
}
|
||||
prepared = session._prepare_tool(tc)
|
||||
assert "error" in prepared
|
||||
# The error mentions both static and pool tools — proves we're
|
||||
# consulting the merged catalog rather than ``_tool_map``.
|
||||
assert "mcp__static__list" in prepared["error"]
|
||||
assert "mcp__pool-srv__do" in prepared["error"]
|
||||
# And the catalog request was scoped to this session's user.
|
||||
assert any(
|
||||
call.kwargs.get("user_id") == "user-7" for call in mock_mcp.get_tools.call_args_list
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Server name validation
|
||||
@@ -794,6 +1083,75 @@ class TestListeners:
|
||||
mgr.add_listener(lambda: 1 / 0) # will raise ZeroDivisionError
|
||||
mgr._rebuild_tools() # should not raise
|
||||
|
||||
# -- Phase 7: user-keyed listener fan-out ------------------------------
|
||||
|
||||
def test_add_listener_records_user_id(self):
|
||||
"""``add_listener`` stores ``(user_id, callback)`` tuples — the
|
||||
listener identity carries the user_id."""
|
||||
mgr = MCPClientManager({})
|
||||
cb_admin = lambda: None # noqa: E731
|
||||
cb_user = lambda: None # noqa: E731
|
||||
mgr.add_listener(cb_admin) # default: user_id=None (admin)
|
||||
mgr.add_listener(cb_user, user_id="user-1")
|
||||
assert (None, cb_admin) in mgr._listeners
|
||||
assert ("user-1", cb_user) in mgr._listeners
|
||||
|
||||
def test_remove_listener_requires_matching_user_id(self):
|
||||
"""Removing with a different ``user_id`` must NOT remove the
|
||||
original registration — listener identity is the pair."""
|
||||
mgr = MCPClientManager({})
|
||||
calls: list[int] = []
|
||||
cb = lambda: calls.append(1) # noqa: E731
|
||||
mgr.add_listener(cb, user_id="user-1")
|
||||
|
||||
# Try to remove with the wrong user_id — should be a no-op.
|
||||
mgr.remove_listener(cb, user_id="user-2")
|
||||
# The user-1 listener should still be live.
|
||||
mgr._notify_user_tool_listeners("user-1")
|
||||
assert calls == [1]
|
||||
|
||||
# Now remove with the right user_id.
|
||||
mgr.remove_listener(cb, user_id="user-1")
|
||||
mgr._notify_user_tool_listeners("user-1")
|
||||
assert calls == [1] # not invoked again
|
||||
|
||||
def test_static_change_fires_all_listeners(self):
|
||||
"""``_rebuild_tools`` (static-path change) fires ALL registered
|
||||
listeners — admin + every user. RFC §3.3."""
|
||||
mgr = MCPClientManager({})
|
||||
admin_calls: list[int] = []
|
||||
u1_calls: list[int] = []
|
||||
u2_calls: list[int] = []
|
||||
mgr.add_listener(lambda: admin_calls.append(1))
|
||||
mgr.add_listener(lambda: u1_calls.append(1), user_id="user-1")
|
||||
mgr.add_listener(lambda: u2_calls.append(1), user_id="user-2")
|
||||
_seed_static_state(mgr, "a", tools=[_fake_openai_tool("mcp__a__x")])
|
||||
mgr._rebuild_tools()
|
||||
assert admin_calls == [1]
|
||||
assert u1_calls == [1]
|
||||
assert u2_calls == [1]
|
||||
|
||||
def test_user_tool_listeners_only_fire_for_matching_user(self):
|
||||
"""``_notify_user_tool_listeners('user-1')`` fires admin (None)
|
||||
and user-1 listeners; user-2's listener is silent."""
|
||||
mgr = MCPClientManager({})
|
||||
admin_calls: list[int] = []
|
||||
u1_calls: list[int] = []
|
||||
u2_calls: list[int] = []
|
||||
mgr.add_listener(lambda: admin_calls.append(1))
|
||||
mgr.add_listener(lambda: u1_calls.append(1), user_id="user-1")
|
||||
mgr.add_listener(lambda: u2_calls.append(1), user_id="user-2")
|
||||
|
||||
mgr._notify_user_tool_listeners("user-1")
|
||||
assert admin_calls == [1]
|
||||
assert u1_calls == [1]
|
||||
assert u2_calls == []
|
||||
|
||||
mgr._notify_user_tool_listeners("user-2")
|
||||
assert admin_calls == [1, 1]
|
||||
assert u1_calls == [1]
|
||||
assert u2_calls == [1]
|
||||
|
||||
|
||||
class TestServerNames:
|
||||
def test_server_names_property(self):
|
||||
|
||||
@@ -0,0 +1,718 @@
|
||||
"""Phase 7 integration tests — per-user catalog scoping.
|
||||
|
||||
These tests drive ``_connect_one_pool`` against a real
|
||||
``streamablehttp_client`` with the SDK's response hook bound to an
|
||||
``httpx.MockTransport`` (R6 verified the happy path closes cleanly and
|
||||
a mid-discovery 401 propagates through anyio TaskGroup unwinding).
|
||||
Direct method injection on ``MagicMock`` as the SOLE gate is
|
||||
forbidden — invariant 14 of the OAuth-MCP RFC.
|
||||
|
||||
The test scaffolding mirrors ``tests/test_mcp_pool_auth_introspection.py``
|
||||
(per-test ``running_loop_mgr`` fixture, ``_seed_oauth_server`` helper,
|
||||
``_run_on_loop``) and adds a tool-discovery transport that returns
|
||||
``initialize`` + ``tools/list`` responses programmatically.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from turnstone.core.mcp_client import (
|
||||
MCPClientManager,
|
||||
PoolEntryState,
|
||||
_AuthCapture,
|
||||
_make_capturing_http_factory,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures and helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def running_loop_mgr() -> Any:
|
||||
"""Background mcp-loop fixture. Mirrors test_mcp_pool_auth_introspection."""
|
||||
cfg: dict[str, Any] = {}
|
||||
mgr = MCPClientManager(cfg)
|
||||
loop = asyncio.new_event_loop()
|
||||
thread = threading.Thread(target=loop.run_forever, daemon=True, name="mcp-pool-test-loop")
|
||||
thread.start()
|
||||
mgr._loop = loop
|
||||
try:
|
||||
yield mgr, loop, thread
|
||||
finally:
|
||||
|
||||
async def _drain(m: MCPClientManager) -> None:
|
||||
task = m._user_pool_eviction_task
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
with contextlib.suppress(BaseException):
|
||||
await task
|
||||
m._user_pool_eviction_task = None
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
asyncio.run_coroutine_threadsafe(_drain(mgr), loop).result(timeout=2)
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
thread.join(timeout=2)
|
||||
|
||||
|
||||
def _run_on_loop(loop: asyncio.AbstractEventLoop, coro: Any, timeout: float = 10) -> Any:
|
||||
fut = asyncio.run_coroutine_threadsafe(coro, loop)
|
||||
return fut.result(timeout=timeout)
|
||||
|
||||
|
||||
def _build_mock_transport_factory(
|
||||
mgr: MCPClientManager,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
handler: Any,
|
||||
) -> None:
|
||||
"""Wrap ``_make_capturing_http_factory`` so the resulting httpx
|
||||
client uses an ``httpx.MockTransport(handler)``.
|
||||
|
||||
The production capturing factory is preserved (we still get the
|
||||
response-hook plumbing); only the underlying transport is swapped.
|
||||
Mirrors the spike v3 pattern that R6 used to verify discovery
|
||||
behaviour against the real SDK.
|
||||
"""
|
||||
real_factory = _make_capturing_http_factory
|
||||
|
||||
def _wrapped_factory(
|
||||
capture: _AuthCapture,
|
||||
fired_event: asyncio.Event | None = None,
|
||||
) -> Any:
|
||||
inner = real_factory(capture, fired_event=fired_event)
|
||||
|
||||
def _factory(*args: Any, **kwargs: Any) -> httpx.AsyncClient:
|
||||
client = inner(*args, **kwargs)
|
||||
client._transport = httpx.MockTransport(handler)
|
||||
return client
|
||||
|
||||
return _factory
|
||||
|
||||
monkeypatch.setattr(
|
||||
"turnstone.core.mcp_client._make_capturing_http_factory",
|
||||
_wrapped_factory,
|
||||
)
|
||||
|
||||
|
||||
def _patch_tcp_probe(mgr: MCPClientManager, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Skip the TCP probe — MockTransport never opens a real socket."""
|
||||
|
||||
async def _noop_probe(*_args: Any, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(MCPClientManager, "_tcp_probe", _noop_probe)
|
||||
|
||||
|
||||
def _make_jsonrpc_handler(
|
||||
*,
|
||||
init_response: dict[str, Any] | None = None,
|
||||
list_tools_response: dict[str, Any] | None = None,
|
||||
list_tools_status: int = 200,
|
||||
list_tools_error_payload: dict[str, Any] | None = None,
|
||||
list_tools_seq: list[dict[str, Any]] | None = None,
|
||||
counter: list[int] | None = None,
|
||||
record_bodies: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Build an httpx async handler for the ``streamable-http`` shape.
|
||||
|
||||
Returns a coroutine the MockTransport invokes per request.
|
||||
"""
|
||||
if init_response is None:
|
||||
init_response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 0,
|
||||
"result": {
|
||||
"protocolVersion": "2025-06-18",
|
||||
"capabilities": {"tools": {"listChanged": False}},
|
||||
"serverInfo": {"name": "fake", "version": "1"},
|
||||
},
|
||||
}
|
||||
if list_tools_response is None:
|
||||
list_tools_response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"result": {"tools": []},
|
||||
}
|
||||
counter = counter if counter is not None else [0]
|
||||
list_tools_seq = list_tools_seq or []
|
||||
list_tools_index = [0]
|
||||
|
||||
async def _handler(req: httpx.Request) -> httpx.Response:
|
||||
if req.method == "GET":
|
||||
return httpx.Response(405)
|
||||
if req.method == "DELETE":
|
||||
return httpx.Response(200)
|
||||
body = req.content.decode() if req.content else ""
|
||||
if record_bodies is not None:
|
||||
record_bodies.append(body)
|
||||
if "notifications/initialized" in body:
|
||||
return httpx.Response(202)
|
||||
counter[0] += 1
|
||||
if "method" in body and '"initialize"' in body:
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "application/json", "mcp-session-id": "sess-1"},
|
||||
json=init_response,
|
||||
)
|
||||
# ``tools/list`` — by request order: explicit override first, then
|
||||
# the staged sequence, then fall back to ``list_tools_response``.
|
||||
if list_tools_status != 200:
|
||||
return httpx.Response(
|
||||
list_tools_status,
|
||||
headers={"www-authenticate": 'Bearer error="invalid_token"'},
|
||||
json=list_tools_error_payload or {"error": "unauthorized"},
|
||||
)
|
||||
if list_tools_seq and list_tools_index[0] < len(list_tools_seq):
|
||||
payload = list_tools_seq[list_tools_index[0]]
|
||||
list_tools_index[0] += 1
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "application/json"},
|
||||
json=payload,
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "application/json"},
|
||||
json=list_tools_response,
|
||||
)
|
||||
|
||||
return _handler
|
||||
|
||||
|
||||
def _list_tools_payload(tools: list[dict[str, Any]], req_id: int = 1) -> dict[str, Any]:
|
||||
"""Build a minimal ``tools/list`` JSON-RPC result payload."""
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {"tools": tools},
|
||||
}
|
||||
|
||||
|
||||
def _tool_spec(name: str, description: str = "") -> dict[str, Any]:
|
||||
return {
|
||||
"name": name,
|
||||
"description": description or f"do {name}",
|
||||
"inputSchema": {"type": "object", "properties": {}},
|
||||
}
|
||||
|
||||
|
||||
def _connect_pool(
|
||||
mgr: MCPClientManager,
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
*,
|
||||
user_id: str,
|
||||
server_name: str,
|
||||
url: str = "https://mcp.example.com/sse",
|
||||
access_token: str = "access-aaa",
|
||||
) -> PoolEntryState:
|
||||
"""Drive ``_connect_one_pool`` once; returns the resulting entry."""
|
||||
cfg: dict[str, Any] = {"type": "streamable-http", "url": url, "headers": {}}
|
||||
capture = _AuthCapture()
|
||||
fired = asyncio.Event()
|
||||
|
||||
async def _go() -> PoolEntryState:
|
||||
return await mgr._connect_one_pool(
|
||||
(user_id, server_name),
|
||||
cfg,
|
||||
access_token,
|
||||
auth_capture=capture,
|
||||
auth_fired_event=fired,
|
||||
)
|
||||
|
||||
return _run_on_loop(loop, _go())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Discovery on first connect — drives REAL streamablehttp_client + httpx
|
||||
# MockTransport per invariant 14.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_discovery_on_pool_connect_via_real_streamable_http(
|
||||
running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""``_connect_one_pool`` discovers the user's tool catalog after
|
||||
``initialize()`` and stores it on the entry.
|
||||
|
||||
Drives through the real ``streamablehttp_client`` and the real
|
||||
SDK ``ClientSession`` (only the underlying httpx transport is
|
||||
swapped). Asserts on observable state rather than mock call count
|
||||
so a refactor that produces the same end state still passes.
|
||||
|
||||
Verified by reverting the discovery block in ``_connect_one_pool``
|
||||
(the ``await session.list_tools()`` + ``_rebuild_user_tool_map``
|
||||
calls): this test fails because ``entry.tools`` stays ``None`` and
|
||||
``_user_tool_map`` never gains the user_id key.
|
||||
"""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
handler = _make_jsonrpc_handler(
|
||||
list_tools_response=_list_tools_payload(
|
||||
[_tool_spec("do_thing"), _tool_spec("other")],
|
||||
),
|
||||
)
|
||||
_build_mock_transport_factory(mgr, monkeypatch, handler)
|
||||
_patch_tcp_probe(mgr, monkeypatch)
|
||||
|
||||
entry = _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv")
|
||||
|
||||
assert entry.session is not None
|
||||
assert entry.tools is not None
|
||||
tool_names = {t["function"]["name"] for t in entry.tools}
|
||||
assert tool_names == {"mcp__pool-srv__do_thing", "mcp__pool-srv__other"}
|
||||
# Per-user index reflects the discovery.
|
||||
user_map = mgr._user_tool_map.get("user-1")
|
||||
assert user_map is not None
|
||||
assert "mcp__pool-srv__do_thing" in user_map
|
||||
assert user_map["mcp__pool-srv__do_thing"] == ("pool-srv", "do_thing")
|
||||
# is_mcp_tool with the user's id sees the discovered tool.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True
|
||||
# Default (None) caller does NOT — pool is per-user.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do_thing") is False
|
||||
|
||||
|
||||
def test_pool_tool_visibility_user_isolation(
|
||||
running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Two users connecting to the same server have independent
|
||||
``_user_tool_map`` entries; one user never sees another's catalog."""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
# Both users see the same set, but the per-user maps must stay
|
||||
# independent — privacy is a structural property even when the
|
||||
# contents happen to match.
|
||||
handler = _make_jsonrpc_handler(
|
||||
list_tools_response=_list_tools_payload([_tool_spec("shared_tool")]),
|
||||
)
|
||||
_build_mock_transport_factory(mgr, monkeypatch, handler)
|
||||
_patch_tcp_probe(mgr, monkeypatch)
|
||||
|
||||
_connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv")
|
||||
_connect_pool(mgr, loop, user_id="user-2", server_name="pool-srv")
|
||||
|
||||
map_1 = mgr._user_tool_map.get("user-1")
|
||||
map_2 = mgr._user_tool_map.get("user-2")
|
||||
assert map_1 is not None and map_2 is not None
|
||||
assert map_1 is not map_2
|
||||
assert "mcp__pool-srv__shared_tool" in map_1
|
||||
assert "mcp__pool-srv__shared_tool" in map_2
|
||||
# Cross-user visibility is forbidden via is_mcp_tool: user-1's name
|
||||
# presence implies nothing about user-3.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__shared_tool", user_id="user-3") is False
|
||||
|
||||
|
||||
def test_eviction_drops_catalog_and_fires_listener(
|
||||
running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""``_evict_session`` clears ``entry.tools``, rebuilds the user's
|
||||
map (drops the now-empty entry), and fires user + admin listeners.
|
||||
|
||||
Verified by reverting the catalog-cleanup block in ``_evict_session``
|
||||
(drop the ``evict.tools = None`` / ``_rebuild_user_tool_map`` /
|
||||
``_notify_user_tool_listeners`` calls): the test fails because the
|
||||
listener never fires AND ``is_mcp_tool`` keeps returning True for
|
||||
the now-evicted tool.
|
||||
"""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
handler = _make_jsonrpc_handler(
|
||||
list_tools_response=_list_tools_payload([_tool_spec("do_thing")]),
|
||||
)
|
||||
_build_mock_transport_factory(mgr, monkeypatch, handler)
|
||||
_patch_tcp_probe(mgr, monkeypatch)
|
||||
|
||||
_connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv")
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True
|
||||
|
||||
# Register one user-keyed and one admin (None) listener; both
|
||||
# MUST fire on this user's eviction.
|
||||
user_calls = [0]
|
||||
admin_calls = [0]
|
||||
other_calls = [0]
|
||||
|
||||
def _user_cb() -> None:
|
||||
user_calls[0] += 1
|
||||
|
||||
def _admin_cb() -> None:
|
||||
admin_calls[0] += 1
|
||||
|
||||
def _other_cb() -> None:
|
||||
other_calls[0] += 1
|
||||
|
||||
mgr.add_listener(_user_cb, user_id="user-1")
|
||||
mgr.add_listener(_admin_cb) # admin / None
|
||||
mgr.add_listener(_other_cb, user_id="user-2")
|
||||
|
||||
mgr._evict_session(("user-1", "pool-srv"))
|
||||
|
||||
# Catalog cleared.
|
||||
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
|
||||
assert entry.session is None
|
||||
assert entry.tools is None
|
||||
# User map dropped (no remaining pool entries for this user).
|
||||
assert "user-1" not in mgr._user_tool_map
|
||||
# is_mcp_tool no longer surfaces the evicted name.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is False
|
||||
# Listener fan-out: matching user + admin fire, OTHER user does not.
|
||||
assert user_calls[0] == 1, f"user-keyed listener fired {user_calls[0]} times; expected 1"
|
||||
assert admin_calls[0] == 1, f"admin listener fired {admin_calls[0]} times; expected 1"
|
||||
assert other_calls[0] == 0, (
|
||||
f"unrelated user-2 listener fired {other_calls[0]} times; expected 0 — "
|
||||
"RFC §3.3 privacy violation."
|
||||
)
|
||||
|
||||
|
||||
def test_close_pool_entry_if_idle_clears_catalog_and_fires_listener(
|
||||
running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""LRU/TTL eviction (``_close_pool_entry_if_idle``) mirrors
|
||||
``_evict_session``'s catalog-cleanup contract: drops the entry,
|
||||
prunes the notification-debounce dict, rebuilds the user's tool
|
||||
map, and fires user + admin listeners.
|
||||
|
||||
Phase 7 round-2 review hardening (round2-1): the bug-2 fix added
|
||||
the catalog-cleanup block to this method but no integration test
|
||||
drove it — exactly the failure mode flagged in
|
||||
``feedback_tests_through_boundaries.md``. Negative test: drop the
|
||||
``_rebuild_user_tool_map`` / ``_notify_user_tool_listeners`` calls
|
||||
from ``_close_pool_entry_if_idle``'s post-pop block; this test
|
||||
fails because ``is_mcp_tool`` keeps returning ``True`` for the
|
||||
evicted name AND no listener fires.
|
||||
"""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
handler = _make_jsonrpc_handler(
|
||||
list_tools_response=_list_tools_payload([_tool_spec("do_thing")]),
|
||||
)
|
||||
_build_mock_transport_factory(mgr, monkeypatch, handler)
|
||||
_patch_tcp_probe(mgr, monkeypatch)
|
||||
|
||||
_connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv")
|
||||
key = ("user-1", "pool-srv")
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is True
|
||||
# Sanity: in_flight must be 0 for the eviction path to proceed.
|
||||
assert mgr._user_pool_entries[key].in_flight == 0
|
||||
# Seed the debounce dict so the perf-1 prune is observable.
|
||||
mgr._last_pool_notification_refresh[key] = 0.0
|
||||
|
||||
user_calls = [0]
|
||||
admin_calls = [0]
|
||||
other_calls = [0]
|
||||
|
||||
def _user_cb() -> None:
|
||||
user_calls[0] += 1
|
||||
|
||||
def _admin_cb() -> None:
|
||||
admin_calls[0] += 1
|
||||
|
||||
def _other_cb() -> None:
|
||||
other_calls[0] += 1
|
||||
|
||||
mgr.add_listener(_user_cb, user_id="user-1")
|
||||
mgr.add_listener(_admin_cb) # admin / None
|
||||
mgr.add_listener(_other_cb, user_id="user-2")
|
||||
|
||||
# Drive the LRU/TTL eviction path directly. ``open_lock`` is
|
||||
# uncontested (no concurrent dispatch) and ``in_flight`` is 0,
|
||||
# so the close proceeds without retry.
|
||||
_run_on_loop(loop, mgr._close_pool_entry_if_idle(key))
|
||||
|
||||
# Entry fully removed (LRU eviction pops the dict — unlike
|
||||
# ``_evict_session`` which keeps the entry as a ``session=None``
|
||||
# phantom for the next dispatch to re-connect).
|
||||
assert key not in mgr._user_pool_entries
|
||||
assert key not in mgr._user_pool_last_used
|
||||
assert key not in mgr._user_pool_locks
|
||||
# perf-1 prune: debounce dict no longer carries the key.
|
||||
assert key not in mgr._last_pool_notification_refresh
|
||||
# Catalog cleanup ran in BOTH dicts (bug-1 sibling + bug-2 cleanup).
|
||||
assert "user-1" not in mgr._user_tool_map
|
||||
assert "user-1" not in mgr._user_tools
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__do_thing", user_id="user-1") is False
|
||||
# Listener fan-out: matching user + admin fire, OTHER user does not.
|
||||
assert user_calls[0] == 1, f"user-keyed listener fired {user_calls[0]} times; expected 1"
|
||||
assert admin_calls[0] == 1, f"admin listener fired {admin_calls[0]} times; expected 1"
|
||||
assert other_calls[0] == 0, (
|
||||
f"unrelated user-2 listener fired {other_calls[0]} times; expected 0 — "
|
||||
"RFC §3.3 privacy violation."
|
||||
)
|
||||
|
||||
|
||||
def test_reconnect_after_eviction_repopulates_catalog(
|
||||
running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""After eviction, the next ``_connect_one_pool`` re-populates the
|
||||
catalog from the SDK by construction — no extra glue required.
|
||||
|
||||
Verified by reverting the discovery block in ``_connect_one_pool``:
|
||||
the reconnected entry's ``tools`` stays ``None`` and the user map
|
||||
never re-emerges.
|
||||
"""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
# Sequence: first connect sees [tool_a]; eviction clears; second
|
||||
# connect (after backend rotates) sees [tool_b].
|
||||
handler = _make_jsonrpc_handler(
|
||||
list_tools_seq=[
|
||||
_list_tools_payload([_tool_spec("tool_a")]),
|
||||
_list_tools_payload([_tool_spec("tool_b")], req_id=1),
|
||||
],
|
||||
)
|
||||
_build_mock_transport_factory(mgr, monkeypatch, handler)
|
||||
_patch_tcp_probe(mgr, monkeypatch)
|
||||
|
||||
_connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv")
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__tool_a", user_id="user-1") is True
|
||||
|
||||
mgr._evict_session(("user-1", "pool-srv"))
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__tool_a", user_id="user-1") is False
|
||||
|
||||
_connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv")
|
||||
# Reconnect picked up the rotated catalog.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__tool_b", user_id="user-1") is True
|
||||
# Old name was correctly purged — not still hanging around.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__tool_a", user_id="user-1") is False
|
||||
|
||||
|
||||
def test_refresh_pool_server_tools_isolates_user(
|
||||
running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""``_refresh_pool_server_tools`` refreshes THIS user's catalog only,
|
||||
leaving the static path and other users' catalogs untouched.
|
||||
|
||||
Drives ``_refresh_pool_server_tools`` directly (which is also what
|
||||
the pool session's notification handler invokes) and asserts:
|
||||
* The pool entry's ``tools`` field reflects the new catalog.
|
||||
* ``_user_tool_map`` for the owning user is updated.
|
||||
* ``_tool_map`` (static) is untouched — invariant 1.
|
||||
* Other users' ``_user_tool_map`` entries are untouched.
|
||||
|
||||
Verified by reverting ``_refresh_pool_server_tools`` to call
|
||||
``_rebuild_tools()`` (the static path) instead of
|
||||
``_rebuild_user_tool_map(user_id)``: this test fails because
|
||||
``_tool_map`` is mutated and the user map doesn't pick up the
|
||||
rotated catalog.
|
||||
"""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
handler = _make_jsonrpc_handler(
|
||||
list_tools_seq=[
|
||||
_list_tools_payload([_tool_spec("v1")]),
|
||||
_list_tools_payload([_tool_spec("v2")], req_id=2),
|
||||
],
|
||||
)
|
||||
_build_mock_transport_factory(mgr, monkeypatch, handler)
|
||||
_patch_tcp_probe(mgr, monkeypatch)
|
||||
|
||||
# Pre-seed a static-path entry to assert it stays untouched.
|
||||
from turnstone.core.mcp_client import StaticServerState
|
||||
|
||||
sentinel_static = StaticServerState(name="static-srv", session=MagicMock())
|
||||
sentinel_static.tools = [
|
||||
{"type": "function", "function": {"name": "mcp__static-srv__static_one"}}
|
||||
]
|
||||
mgr._static_servers["static-srv"] = sentinel_static
|
||||
mgr._tool_map["mcp__static-srv__static_one"] = ("static-srv", "static_one")
|
||||
|
||||
# Pre-seed an unrelated user's pool entry so we can assert it's
|
||||
# untouched too.
|
||||
async def _seed_other_user() -> None:
|
||||
entry = await mgr._ensure_pool_entry(("user-2", "pool-srv"))
|
||||
entry.session = MagicMock()
|
||||
entry.tools = [{"type": "function", "function": {"name": "mcp__pool-srv__user2_only"}}]
|
||||
mgr._rebuild_user_tool_map("user-2")
|
||||
|
||||
_run_on_loop(loop, _seed_other_user())
|
||||
# Sanity: user-2's map is set.
|
||||
assert "mcp__pool-srv__user2_only" in (mgr._user_tool_map.get("user-2") or {})
|
||||
|
||||
_connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv")
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__v1", user_id="user-1") is True
|
||||
|
||||
# Trigger a refresh (simulating the notification handler invoking it).
|
||||
added_removed = _run_on_loop(loop, mgr._refresh_pool_server_tools(("user-1", "pool-srv")))
|
||||
added, removed = added_removed
|
||||
assert added == ["mcp__pool-srv__v2"]
|
||||
assert removed == ["mcp__pool-srv__v1"]
|
||||
|
||||
# User-1 sees the new tool.
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__v2", user_id="user-1") is True
|
||||
assert mgr.is_mcp_tool("mcp__pool-srv__v1", user_id="user-1") is False
|
||||
|
||||
# Static path untouched — hard invariant 1.
|
||||
assert mgr._tool_map == {"mcp__static-srv__static_one": ("static-srv", "static_one")}
|
||||
assert mgr._static_servers["static-srv"].tools == [
|
||||
{"type": "function", "function": {"name": "mcp__static-srv__static_one"}}
|
||||
]
|
||||
# User-2's pool catalog untouched.
|
||||
assert mgr._user_tool_map["user-2"] == {
|
||||
"mcp__pool-srv__user2_only": ("pool-srv", "user2_only"),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Invariant 1 — canonical regression: static path byte-identical when
|
||||
# oauth_user is disabled.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_static_path_byte_identical_with_oauth_user_disabled(
|
||||
running_loop_mgr: Any,
|
||||
) -> None:
|
||||
"""In a static-only deployment (no oauth_user config rows), the
|
||||
Phase 7 changes MUST leave every static-path catalog API byte-
|
||||
identical to pre-Phase-7 behaviour. This is invariant 1.
|
||||
|
||||
Drives ``_connect_one`` (static path) and asserts:
|
||||
* ``get_tools()`` (no user_id) returns the static catalog only.
|
||||
* ``is_mcp_tool(name)`` (no user_id) works as before.
|
||||
* ``_tool_map`` is the only catalog index touched.
|
||||
* No ``_user_*`` state exists (no pool entries / no per-user maps).
|
||||
|
||||
Static-path connect uses ``StaticServerState`` instead of pool
|
||||
entries, so this test seeds it directly rather than driving through
|
||||
``_connect_one`` (which is exercised by the existing integration
|
||||
suite). The point is to assert the catalog API contract stays
|
||||
intact under Phase 7's signature widening.
|
||||
|
||||
Verified by widening the default of ``get_tools(user_id)`` to
|
||||
``user_id=""`` (instead of ``None``) — the test fails because
|
||||
``get_tools()`` now goes through the per-user merge branch and
|
||||
starts paying for an empty-loop iteration that pre-Phase-7
|
||||
callers never saw.
|
||||
"""
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
# Seed the static-path catalog directly (no pool plumbing needed).
|
||||
from turnstone.core.mcp_client import StaticServerState
|
||||
|
||||
sentinel = StaticServerState(name="static-srv", session=MagicMock())
|
||||
static_tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "mcp__static-srv__list",
|
||||
"description": "list",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
sentinel.tools = [static_tool]
|
||||
mgr._static_servers["static-srv"] = sentinel
|
||||
mgr._tools = [static_tool]
|
||||
mgr._tool_map["mcp__static-srv__list"] = ("static-srv", "list")
|
||||
|
||||
# Pre-Phase-7 caller pattern: no kwargs.
|
||||
pre_phase7_tools = mgr.get_tools()
|
||||
assert pre_phase7_tools == [static_tool]
|
||||
assert mgr.is_mcp_tool("mcp__static-srv__list") is True
|
||||
assert mgr.is_mcp_tool("nonexistent") is False
|
||||
# Per-user state does NOT exist when no pool entries are present.
|
||||
assert mgr._user_pool_entries == {}
|
||||
assert mgr._user_tool_map == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# R6 verification (production form): a 401 mid-discovery propagates
|
||||
# rather than hangs. This test pins the empirical conclusion so a future
|
||||
# SDK upgrade that re-introduces the hang surfaces here.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_pool_connect_list_tools_401_propagates(
|
||||
running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""A 401 returned to the discovery POST surfaces as an exception
|
||||
from ``_connect_one_pool`` within a bounded timeout (no hang).
|
||||
|
||||
The SDK propagates the 401 through anyio TaskGroup unwinding (raises
|
||||
an ``ExceptionGroup`` from the surrounding ``streamablehttp_client``
|
||||
context); the production reading is "discovery does not need the
|
||||
carrier-race shape from ``_dispatch_pool_with_entry``."
|
||||
|
||||
If a future SDK bump silently re-introduces the hang, this test
|
||||
will time out — ``_run_on_loop`` is invoked below with a 5s budget.
|
||||
"""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
handler = _make_jsonrpc_handler(list_tools_status=401)
|
||||
_build_mock_transport_factory(mgr, monkeypatch, handler)
|
||||
_patch_tcp_probe(mgr, monkeypatch)
|
||||
|
||||
cfg: dict[str, Any] = {
|
||||
"type": "streamable-http",
|
||||
"url": "https://mcp.example.com/sse",
|
||||
"headers": {},
|
||||
}
|
||||
capture = _AuthCapture()
|
||||
fired = asyncio.Event()
|
||||
t0 = time.monotonic()
|
||||
# The 401 surfaces as ``ExceptionGroup`` (anyio TaskGroup unwinding
|
||||
# wraps the underlying ``HTTPStatusError``); both are ``Exception``
|
||||
# subclasses. The point of the assertion is "the call returned/raised
|
||||
# within the timeout instead of hanging" — the type-discrimination
|
||||
# tests cover the carrier classification path.
|
||||
with pytest.raises(Exception): # noqa: B017 — generic catch documents the no-hang assertion
|
||||
_run_on_loop(
|
||||
loop,
|
||||
mgr._connect_one_pool(
|
||||
("user-1", "pool-srv"),
|
||||
cfg,
|
||||
"access-aaa",
|
||||
auth_capture=capture,
|
||||
auth_fired_event=fired,
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
# Hard upper bound — the spike's happy path closed in ~0.05s, so a
|
||||
# 401 propagation that takes longer than 5s is the hang regression.
|
||||
assert elapsed < 5.0, f"_connect_one_pool 401 hang regression: {elapsed:.2f}s"
|
||||
# Carrier captured the 401 — the response hook fired before the
|
||||
# SDK propagated the failure, so the dispatcher would have the
|
||||
# auth signal even though discovery aborted.
|
||||
assert capture.status == 401, (
|
||||
f"401 mid-discovery did NOT surface to carrier; capture.status={capture.status}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Listener fan-out — RFC §3.3 verification (Chunk 3)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_listener_identity_includes_user_id() -> None:
|
||||
"""Listener identity is the ``(user_id, callback)`` tuple, not the
|
||||
callback alone — removing for user A leaves user B's registration
|
||||
intact even when the same callable was registered for both.
|
||||
|
||||
Negative-test reproducer: revert ``add_listener`` /
|
||||
``remove_listener`` to keep storage as a flat list of callables —
|
||||
this test fails because removing one registration accidentally
|
||||
removes BOTH (callable identity is shared).
|
||||
"""
|
||||
mgr = MCPClientManager({})
|
||||
|
||||
calls = []
|
||||
|
||||
def _shared_cb() -> None:
|
||||
calls.append(time.monotonic())
|
||||
|
||||
mgr.add_listener(_shared_cb, user_id="user-A")
|
||||
mgr.add_listener(_shared_cb, user_id="user-B")
|
||||
|
||||
# Removing user-A's registration leaves user-B's intact.
|
||||
mgr.remove_listener(_shared_cb, user_id="user-A")
|
||||
|
||||
# Static-path notify fires both registrations; with user-A removed,
|
||||
# only user-B fires now.
|
||||
mgr._notify_listeners()
|
||||
assert len(calls) == 1, (
|
||||
f"shared callback fired {len(calls)} times after removing user-A's "
|
||||
"registration; expected 1 (user-B's registration intact)"
|
||||
)
|
||||
|
||||
# Now remove user-B's; static-path notify fires nobody.
|
||||
mgr.remove_listener(_shared_cb, user_id="user-B")
|
||||
calls.clear()
|
||||
mgr._notify_listeners()
|
||||
assert calls == [], "callback fired after BOTH user_id-keyed registrations removed"
|
||||
@@ -199,6 +199,11 @@ class TestLazyConnect:
|
||||
|
||||
fake_session = MagicMock()
|
||||
fake_session.initialize = AsyncMock(return_value=None)
|
||||
# Phase 7: ``_connect_one_pool`` discovers the user's tool
|
||||
# catalog after ``initialize()`` returns. This stub returns a
|
||||
# zero-tool result so the test can keep its narrow focus on
|
||||
# the bearer-injection contract.
|
||||
fake_session.list_tools = AsyncMock(return_value=MagicMock(tools=[]))
|
||||
|
||||
def _stream_factory(*, url: str, headers: dict[str, str]) -> _AsyncCM:
|
||||
observed_kwargs["url"] = url
|
||||
|
||||
@@ -173,3 +173,50 @@ class TestResolveClient:
|
||||
def test_unknown_backend_returns_none(self):
|
||||
client = resolve_web_search_client("typo_backend", tavily_key="key")
|
||||
assert client is None
|
||||
|
||||
def test_resolve_web_search_client_rejects_oauth_user_backend(self):
|
||||
"""A web_search backend pointing at an ``auth_type=oauth_user``
|
||||
MCP server MUST be rejected at boot — per-node web_search
|
||||
cannot carry per-user tokens, so resolving the backend would
|
||||
guarantee a 401-on-call instead of a clean disablement.
|
||||
|
||||
Phase 7 invariant 8 corollary: pool tools are user-scoped;
|
||||
every entry point that lacks per-user identity (web_search
|
||||
boot resolver, eval harness, CLI default) MUST refuse them
|
||||
rather than silently produce a broken client.
|
||||
|
||||
Verified by reverting the ``server_auth_type(...) == 'oauth_user'``
|
||||
guard in ``resolve_web_search_client``: the resolver returns
|
||||
an ``MCPSearchClient`` whose ``call_tool_sync`` would surface
|
||||
a 401 / consent_required structured error on every search.
|
||||
"""
|
||||
mcp = MagicMock()
|
||||
mcp.is_mcp_tool.return_value = True # name resolves
|
||||
mcp.server_auth_type.return_value = "oauth_user"
|
||||
client = resolve_web_search_client(
|
||||
"mcp:oauth-search:search", tavily_key=None, mcp_client=mcp
|
||||
)
|
||||
assert client is None, (
|
||||
"oauth_user-backed web_search backend resolved to a non-None client; "
|
||||
"boot-time guard missing or regressed."
|
||||
)
|
||||
# Per-turn callers must read from the in-memory cache, never
|
||||
# the SQL helper — perf regression guard.
|
||||
mcp.server_auth_type.assert_called_with("oauth-search")
|
||||
assert not mcp._lookup_server_row.called, (
|
||||
"resolver issued a SQL roundtrip via _lookup_server_row; "
|
||||
"per-turn web_search backend resolution must use the "
|
||||
"in-memory server_auth_type accessor."
|
||||
)
|
||||
|
||||
def test_resolve_web_search_client_accepts_static_backend(self):
|
||||
"""Static-path (``auth_type=none`` or ``static``) MCP backends
|
||||
still resolve cleanly — the new guard ONLY rejects oauth_user.
|
||||
"""
|
||||
mcp = MagicMock()
|
||||
mcp.is_mcp_tool.return_value = True
|
||||
mcp.server_auth_type.return_value = None
|
||||
client = resolve_web_search_client(
|
||||
"mcp:static-search:search", tavily_key=None, mcp_client=mcp
|
||||
)
|
||||
assert isinstance(client, MCPSearchClient)
|
||||
|
||||
+502
-44
@@ -119,6 +119,16 @@ def _validate_oauth_user_url(url: str) -> None:
|
||||
# bloating either surface via a thousand-token ``scope=`` value.
|
||||
_MAX_INSUFFICIENT_SCOPE_REPORTED = 32
|
||||
|
||||
# Defensive cap on the number of tools we accept from any single MCP
|
||||
# server's ``tools/list`` response. Real servers expose at most a few
|
||||
# dozen tools; a misconfigured or hostile upstream returning thousands
|
||||
# would amplify both memory (one OpenAI tool dict per entry) and
|
||||
# downstream BM25 reindex cost. Mirrors the ``_MAX_ERROR_LEN`` /
|
||||
# ``_MAX_INSUFFICIENT_SCOPE_REPORTED`` defensive ceilings: we truncate
|
||||
# rather than reject so partial visibility beats zero visibility, and
|
||||
# emit a warning so operators can investigate.
|
||||
_MAX_TOOLS_PER_SERVER = 1000
|
||||
|
||||
|
||||
@dataclass
|
||||
class _AuthCapture:
|
||||
@@ -247,6 +257,25 @@ def _mcp_to_openai(server_name: str, tool: Any) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _cap_server_tools(server_name: str, tools: list[Any]) -> list[Any]:
|
||||
"""Apply ``_MAX_TOOLS_PER_SERVER`` cap with operator-visible warning.
|
||||
|
||||
Identity for inputs at or below the cap (no copy); slice + warn on
|
||||
overflow. Caller is responsible for converting the returned list to
|
||||
the OpenAI shape via :func:`_mcp_to_openai`.
|
||||
"""
|
||||
if len(tools) <= _MAX_TOOLS_PER_SERVER:
|
||||
return tools
|
||||
log.warning(
|
||||
"MCP server '%s' returned %d tools — truncating to %d "
|
||||
"(_MAX_TOOLS_PER_SERVER cap). Misconfigured or hostile upstream?",
|
||||
server_name,
|
||||
len(tools),
|
||||
_MAX_TOOLS_PER_SERVER,
|
||||
)
|
||||
return tools[:_MAX_TOOLS_PER_SERVER]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-server state containers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -366,8 +395,12 @@ class MCPClientManager:
|
||||
self._last_error: dict[str, str] = {}
|
||||
self._MAX_ERROR_LEN = 256
|
||||
|
||||
# Listener infrastructure (tool-change callbacks for ChatSession)
|
||||
self._listeners: list[Callable[[], None]] = []
|
||||
# Listener infrastructure (tool-change callbacks for ChatSession).
|
||||
# Each entry is ``(user_id, callback)``: ``user_id=None`` is the
|
||||
# admin / global listener (fires on every tool-change), a string
|
||||
# ``user_id`` fires only on changes scoped to that user OR on
|
||||
# global static-path changes. RFC §3.3.
|
||||
self._listeners: list[tuple[str | None, Callable[[], None]]] = []
|
||||
self._listeners_lock = threading.Lock()
|
||||
|
||||
# Merged resource catalog
|
||||
@@ -408,6 +441,31 @@ class MCPClientManager:
|
||||
# so each ``asyncio.Lock`` binds to the correct loop on first use.
|
||||
self._user_pool_locks: dict[tuple[str, str], asyncio.Lock] = {}
|
||||
|
||||
# Per-(user, server) catalog state. Tools live on
|
||||
# ``PoolEntryState.tools`` (the dataclass already carries the
|
||||
# field — it's populated lazily by per-user discovery on first
|
||||
# connect). ``_user_tool_map`` is the per-user prefixed-name
|
||||
# index, mirroring static ``_tool_map``: outer key is ``user_id``,
|
||||
# inner is ``prefixed_name → (server_name, original_name)``.
|
||||
# Resource/prompt mirrors are deferred to Phase 7b — Phase 7
|
||||
# only lights up the tool path needed for invariant 8.
|
||||
self._user_tool_map: dict[str, dict[str, tuple[str, str]]] = {}
|
||||
# Per-user merged tool list (one snapshot per user_id), updated
|
||||
# atomically alongside ``_user_tool_map`` in
|
||||
# ``_rebuild_user_tool_map``. ``get_tools(user_id=...)`` reads
|
||||
# this via a single dict-get (atomic under GIL) so sync-thread
|
||||
# callers (ChatSession) never iterate ``_user_pool_entries``
|
||||
# concurrently with the mcp-loop's mutations of the same dict
|
||||
# (insert in ``_ensure_pool_entry`` / pop in
|
||||
# ``_close_pool_entry_if_idle`` / ``_evict_session``).
|
||||
self._user_tools: dict[str, list[dict[str, Any]]] = {}
|
||||
|
||||
# Notification debounce for pool sessions, keyed ``(user_id, server)``.
|
||||
# Mirrors ``_last_notification_refresh`` (static) but per-pool-key
|
||||
# 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] = {}
|
||||
|
||||
# Pool tuning (read at construction; falls back to defaults).
|
||||
mcp_cfg = load_config("mcp")
|
||||
self._user_pool_idle_ttl_s = float(mcp_cfg.get("user_session_idle_ttl_seconds", 600))
|
||||
@@ -419,6 +477,15 @@ class MCPClientManager:
|
||||
# never hit it.
|
||||
self._app_state: Any = None
|
||||
|
||||
# In-memory cache of server names whose ``auth_type='oauth_user'``.
|
||||
# ``_db_servers_to_config`` strips oauth_user rows on the way into
|
||||
# ``_server_configs``, so neither that dict nor ``_static_servers``
|
||||
# carries auth_type. This set is populated alongside ``reconcile_sync``
|
||||
# / ``set_oauth_user_servers`` so per-turn callers (web_search
|
||||
# backend resolution) can answer "is this server pool-backed?"
|
||||
# without a SQL roundtrip.
|
||||
self._oauth_user_server_names: set[str] = set()
|
||||
|
||||
# Idle-eviction task handle. Scheduled lazily on the mcp-loop the
|
||||
# first time a pool entry is created (start() runs before pool
|
||||
# rows exist, so deferring keeps the task count at zero in
|
||||
@@ -866,9 +933,8 @@ class MCPClientManager:
|
||||
|
||||
# Discover tools
|
||||
result = await session.list_tools()
|
||||
server_tools: list[dict[str, Any]] = []
|
||||
for tool in result.tools:
|
||||
server_tools.append(_mcp_to_openai(name, tool))
|
||||
capped = _cap_server_tools(name, result.tools)
|
||||
server_tools: list[dict[str, Any]] = [_mcp_to_openai(name, tool) for tool in capped]
|
||||
|
||||
state.tools = server_tools
|
||||
self._rebuild_tools()
|
||||
@@ -975,9 +1041,12 @@ class MCPClientManager:
|
||||
* Streamable-HTTP transport only — pool servers are remote.
|
||||
* ``Authorization: Bearer {access_token}`` injected into headers
|
||||
alongside any operator-supplied static headers.
|
||||
* No catalog discovery (tools / resources / prompts) and no
|
||||
notification handler — those land later once per-user catalog
|
||||
state is in place.
|
||||
* Tool catalog discovery runs after ``initialize()``; the
|
||||
notification handler is bound to ``(user_id, server_name)`` so
|
||||
push-driven ``tools/list_changed`` updates only refresh the
|
||||
owning user's catalog (the static refresher must NEVER fire
|
||||
from a pool session — it would clobber static-path state).
|
||||
Resource / prompt discovery is deferred to Phase 7b.
|
||||
|
||||
When ``auth_capture`` is supplied, the underlying ``httpx``
|
||||
client is built via a factory whose response hook records 401/403
|
||||
@@ -1064,8 +1133,66 @@ class MCPClientManager:
|
||||
await self._safe_teardown_on_connect_failure(key, stack)
|
||||
raise
|
||||
|
||||
# Pool-scoped notification handler — bound to ``(user_id,
|
||||
# server_name)`` via closure so a push-driven
|
||||
# ``tools/list_changed`` only refreshes THIS user's catalog,
|
||||
# never the static path's. Calling the static
|
||||
# ``_refresh_server_tools(server_name)`` from a pool session
|
||||
# would mutate ``_static_servers[server_name].tools`` and
|
||||
# ``_tool_map`` — breaking invariant 1 (static-path
|
||||
# byte-identical) and broadcasting one user's tool view to all
|
||||
# other sessions. Phase 7 ToolListChangedNotification only;
|
||||
# resource / prompt branches log + skip until Phase 7b adds
|
||||
# ``_refresh_pool_server_resources`` / ``_refresh_pool_server_prompts``.
|
||||
async def _on_pool_notification(
|
||||
msg: Any, # RequestResponder | ServerNotification | Exception
|
||||
) -> None:
|
||||
if not isinstance(msg, mcp_types.ServerNotification):
|
||||
return
|
||||
root = msg.root
|
||||
now = time.monotonic()
|
||||
last = self._last_pool_notification_refresh.get(key, 0.0)
|
||||
if now - last < self._NOTIFICATION_DEBOUNCE:
|
||||
log.debug(
|
||||
"Debouncing pool notification user=%s server=%s (%.1fs since last refresh)",
|
||||
user_id,
|
||||
server_name,
|
||||
now - last,
|
||||
)
|
||||
return
|
||||
try:
|
||||
if isinstance(root, mcp_types.ToolListChangedNotification):
|
||||
log.info(
|
||||
"Received tools/list_changed from pool user=%s server=%s",
|
||||
user_id,
|
||||
server_name,
|
||||
)
|
||||
self._last_pool_notification_refresh[key] = now
|
||||
await self._refresh_pool_server_tools(key)
|
||||
elif isinstance(root, mcp_types.ResourceListChangedNotification):
|
||||
log.debug(
|
||||
"pool resources/list_changed (deferred to Phase 7b) user=%s server=%s",
|
||||
user_id,
|
||||
server_name,
|
||||
)
|
||||
elif isinstance(root, mcp_types.PromptListChangedNotification):
|
||||
log.debug(
|
||||
"pool prompts/list_changed (deferred to Phase 7b) user=%s server=%s",
|
||||
user_id,
|
||||
server_name,
|
||||
)
|
||||
except Exception:
|
||||
log.warning(
|
||||
"Pool refresh after notification failed user=%s server=%s",
|
||||
user_id,
|
||||
server_name,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
try:
|
||||
session = await stack.enter_async_context(ClientSession(read, write))
|
||||
session = await stack.enter_async_context(
|
||||
ClientSession(read, write, message_handler=_on_pool_notification) # type: ignore[arg-type]
|
||||
)
|
||||
except Exception:
|
||||
await self._safe_teardown_on_connect_failure(key, stack)
|
||||
raise
|
||||
@@ -1089,9 +1216,70 @@ class MCPClientManager:
|
||||
entry.stack = None
|
||||
await self._safe_teardown_on_connect_failure(key, stack)
|
||||
raise
|
||||
|
||||
# Discover this user's tool catalog. R6 verified: a 401 here
|
||||
# propagates through anyio TaskGroup unwinding (raises an
|
||||
# ``ExceptionGroup`` from the surrounding ``streamablehttp_client``
|
||||
# context) — plain ``await`` under ``asyncio.timeout`` is
|
||||
# sufficient. The carrier-race shape used by
|
||||
# ``_dispatch_pool_with_entry`` defends a different scenario
|
||||
# (reused-session 401 from inside a SECOND dispatch) that
|
||||
# doesn't apply to first-connect discovery. Resource / prompt
|
||||
# discovery deferred to Phase 7b.
|
||||
#
|
||||
# Why ``asyncio.timeout``, not ``asyncio.wait_for``: per
|
||||
# ``feedback_asyncio_timeout_vs_wait_for.md`` and the f6a3b66
|
||||
# fix, Python 3.11's ``asyncio.wait_for`` wraps the inner
|
||||
# coroutine in a fresh task. When the SDK's ``streamablehttp_client``
|
||||
# TaskGroup unwinds (e.g. on a 401), ``aclose`` on the surrounding
|
||||
# anyio scope runs from a different task than entered it →
|
||||
# ``RuntimeError("Attempted to exit cancel scope in a different
|
||||
# task")``. ``asyncio.timeout`` runs the inner coroutine in the
|
||||
# current task and is the safe shape for any await that may
|
||||
# traverse anyio cleanup.
|
||||
try:
|
||||
async with asyncio.timeout(self._CONNECT_TIMEOUT):
|
||||
tools_result = await session.list_tools()
|
||||
except asyncio.CancelledError:
|
||||
entry.stack = None
|
||||
task = asyncio.current_task()
|
||||
if task is not None and task.cancelling():
|
||||
await self._safe_teardown_on_connect_failure(key, stack)
|
||||
raise
|
||||
await self._safe_teardown_on_connect_failure(key, stack)
|
||||
raise TimeoutError(f"Pool discovery failed for '{server_name}'") from None
|
||||
except TimeoutError:
|
||||
entry.stack = None
|
||||
await self._safe_teardown_on_connect_failure(key, stack)
|
||||
raise TimeoutError(f"Pool discovery timed out after {self._CONNECT_TIMEOUT}s") from None
|
||||
except Exception:
|
||||
entry.stack = None
|
||||
await self._safe_teardown_on_connect_failure(key, stack)
|
||||
raise
|
||||
|
||||
capped_tools = _cap_server_tools(server_name, tools_result.tools)
|
||||
entry.tools = [_mcp_to_openai(server_name, tool) for tool in capped_tools]
|
||||
# Publish session readiness BEFORE catalog visibility.
|
||||
# ``_rebuild_user_tool_map`` makes ``is_mcp_tool(name, user_id=U)``
|
||||
# return True for the discovered names; if catalog visibility
|
||||
# preceded ``entry.session`` assignment, a sync-thread reader
|
||||
# racing with this coroutine could observe a tool whose backing
|
||||
# entry has ``session=None``. Defence-in-depth — dispatch
|
||||
# re-fetches its own token through ``get_user_access_token_classified``
|
||||
# and lazy-reconnects on session=None — but ordering catches the
|
||||
# race at the source rather than relying on the dispatch-time
|
||||
# recovery path.
|
||||
entry.session = session
|
||||
entry.last_used = time.monotonic()
|
||||
self._user_pool_last_used[key] = entry.last_used
|
||||
|
||||
# Loop-only mutation; sync-thread readers observe the new tool
|
||||
# list atomically via the per-user dict-get on ``_user_tools``.
|
||||
self._rebuild_user_tool_map(user_id)
|
||||
# Wake user-keyed AND admin (None) listeners; per-user fan-out
|
||||
# ensures another user's session never observes this user's
|
||||
# tool change.
|
||||
self._notify_user_tool_listeners(user_id)
|
||||
return entry
|
||||
|
||||
# -- pool eviction --------------------------------------------------------
|
||||
@@ -1181,6 +1369,7 @@ class MCPClientManager:
|
||||
)
|
||||
except (TimeoutError, asyncio.CancelledError):
|
||||
return
|
||||
evicted = False
|
||||
try:
|
||||
entry = self._user_pool_entries.get(key)
|
||||
if entry is None:
|
||||
@@ -1198,11 +1387,27 @@ class MCPClientManager:
|
||||
await self._safe_close_stack(stack)
|
||||
self._user_pool_entries.pop(key, None)
|
||||
self._user_pool_last_used.pop(key, None)
|
||||
# Mirror ``_evict_session``'s catalog cleanup: dropping the
|
||||
# entry without rebuilding ``_user_tool_map`` would leave
|
||||
# ``is_mcp_tool`` returning True for tools whose backing
|
||||
# pool is gone, and ChatSession's ``_tools`` would never
|
||||
# rebuild because no listener fires.
|
||||
self._last_pool_notification_refresh.pop(key, None)
|
||||
user_id, _server_name = key
|
||||
self._rebuild_user_tool_map(user_id)
|
||||
self._notify_user_tool_listeners(user_id)
|
||||
evicted = True
|
||||
finally:
|
||||
lock.release()
|
||||
# Drop the now-orphaned lock so the dict doesn't grow without
|
||||
# bound across the process lifetime. A new key will reallocate.
|
||||
self._user_pool_locks.pop(key, None)
|
||||
# bound across the process lifetime — but ONLY on the success
|
||||
# path. The early-return branches (entry already removed by a
|
||||
# racing path, or in_flight > 0) leave the lock in place: an
|
||||
# in-flight dispatcher needs the same lock object on its next
|
||||
# acquire, and a key whose entry races a concurrent eviction
|
||||
# will be re-allocated by ``_ensure_pool_entry`` next time.
|
||||
if evicted:
|
||||
self._user_pool_locks.pop(key, None)
|
||||
|
||||
# -- failure classification (pool dispatch) ------------------------------
|
||||
|
||||
@@ -1270,6 +1475,52 @@ class MCPClientManager:
|
||||
self._tool_map = new_map
|
||||
self._notify_listeners()
|
||||
|
||||
def _rebuild_user_tool_map(self, user_id: str) -> None:
|
||||
"""Rebuild the per-user prefixed-name index from pool entries.
|
||||
|
||||
Mirrors :meth:`_rebuild_tools` for one user's pool entries:
|
||||
scan ``_user_pool_entries`` for keys whose first element matches
|
||||
``user_id``, materialize a fresh ``prefixed_name → (server, original)``
|
||||
dict and a parallel tool-list, assign each to its dict in
|
||||
``_user_tool_map`` and ``_user_tools`` respectively. Each
|
||||
per-key write is individually atomic under the GIL; the two
|
||||
writes happen back-to-back on the mcp-loop with no awaits
|
||||
between them, so a sync-thread reader cannot interleave at the
|
||||
Python statement level — but the cross-dict write is not a
|
||||
single atomic operation. In practice the window is sub-microsecond
|
||||
and the listener fan-out (which fires AFTER both writes
|
||||
complete) is the trigger for any session-side rebuild that
|
||||
might re-read both dicts.
|
||||
|
||||
``_tool_map`` (the static index) is NEVER mutated here, so
|
||||
invariant 1 (static path byte-identical) is preserved.
|
||||
|
||||
Empty rebuilds drop the user_id key from BOTH dicts so an idle
|
||||
user with no pool entries doesn't retain permanent empty-list
|
||||
sentinels.
|
||||
|
||||
MUST run on the mcp-loop. The pool dict scan here cannot race
|
||||
with sync-thread reads because sync threads never touch
|
||||
``_user_pool_entries``; they read ``_user_tools`` instead.
|
||||
"""
|
||||
new_map: dict[str, tuple[str, str]] = {}
|
||||
new_tools: list[dict[str, Any]] = []
|
||||
for (uid, _server_name), entry in self._user_pool_entries.items():
|
||||
if uid != user_id or entry.tools is None:
|
||||
continue
|
||||
for tool in entry.tools:
|
||||
prefixed: str = tool["function"]["name"]
|
||||
# Extract original name from the mcp__server__original pattern.
|
||||
original = prefixed.split("__", 2)[2] if prefixed.count("__") >= 2 else prefixed
|
||||
new_map[prefixed] = (_server_name, original)
|
||||
new_tools.append(tool)
|
||||
if new_map:
|
||||
self._user_tool_map[user_id] = new_map
|
||||
self._user_tools[user_id] = new_tools
|
||||
else:
|
||||
self._user_tool_map.pop(user_id, None)
|
||||
self._user_tools.pop(user_id, None)
|
||||
|
||||
async def _refresh_server_tools(self, name: str) -> tuple[list[str], list[str]]:
|
||||
"""Re-fetch tools for one server. Returns ``(added, removed)`` names."""
|
||||
state = self._static_servers.get(name)
|
||||
@@ -1283,7 +1534,8 @@ class MCPClientManager:
|
||||
old_names = {t["function"]["name"] for t in state.tools}
|
||||
|
||||
result = await session.list_tools()
|
||||
server_tools = [_mcp_to_openai(name, tool) for tool in result.tools]
|
||||
capped = _cap_server_tools(name, result.tools)
|
||||
server_tools = [_mcp_to_openai(name, tool) for tool in capped]
|
||||
new_names = {t["function"]["name"] for t in server_tools}
|
||||
|
||||
state.tools = server_tools
|
||||
@@ -1300,6 +1552,51 @@ class MCPClientManager:
|
||||
)
|
||||
return added, removed
|
||||
|
||||
async def _refresh_pool_server_tools(self, key: tuple[str, str]) -> tuple[list[str], list[str]]:
|
||||
"""Re-fetch tools for one pool entry. Returns ``(added, removed)`` names.
|
||||
|
||||
Mirror of :meth:`_refresh_server_tools` for the pool path:
|
||||
targets a single ``(user_id, server_name)`` entry, mutates ONLY
|
||||
``entry.tools`` + the per-user index, fires the user-scoped
|
||||
listener fan-out. Static-path catalogs are NEVER touched, so
|
||||
invariant 1 (static path byte-identical) is preserved.
|
||||
|
||||
MUST run on the mcp-loop. Caller does not need to hold
|
||||
``open_lock`` — ``_refresh_pool_server_tools`` is invoked from
|
||||
the pool session's notification handler, which already runs in
|
||||
the SDK's receive task on the loop.
|
||||
"""
|
||||
entry = self._user_pool_entries.get(key)
|
||||
if entry is None or entry.session is None:
|
||||
raise RuntimeError(f"Pool entry {key!r} is not connected")
|
||||
# Snapshot session locally — a concurrent transport-error
|
||||
# eviction can clear ``entry.session`` after our reads, so a
|
||||
# post-await ``entry.session.<...>`` call would raise
|
||||
# ``AttributeError``. Mirrors the static path's pattern.
|
||||
session = entry.session
|
||||
user_id, server_name = key
|
||||
old_names = {t["function"]["name"] for t in (entry.tools or [])}
|
||||
result = await session.list_tools()
|
||||
capped = _cap_server_tools(server_name, result.tools)
|
||||
server_tools = [_mcp_to_openai(server_name, tool) for tool in capped]
|
||||
new_names = {t["function"]["name"] for t in server_tools}
|
||||
entry.tools = server_tools
|
||||
self._rebuild_user_tool_map(user_id)
|
||||
# Fire user-keyed AND admin (None) listeners. Other users'
|
||||
# listeners do NOT see this change — the pool catalog is private.
|
||||
self._notify_user_tool_listeners(user_id)
|
||||
added = sorted(new_names - old_names)
|
||||
removed = sorted(old_names - new_names)
|
||||
if added or removed:
|
||||
log.info(
|
||||
"Refreshed pool MCP server user=%s server=%s: +%d/-%d tool(s)",
|
||||
user_id,
|
||||
server_name,
|
||||
len(added),
|
||||
len(removed),
|
||||
)
|
||||
return added, removed
|
||||
|
||||
async def _refresh_server(self, name: str) -> tuple[list[str], list[str]]:
|
||||
"""Re-fetch tools, resources, and prompts for one server.
|
||||
|
||||
@@ -1533,20 +1830,52 @@ class MCPClientManager:
|
||||
|
||||
# -- listener infrastructure ---------------------------------------------
|
||||
|
||||
def add_listener(self, callback: Callable[[], None]) -> None:
|
||||
"""Register a callback invoked when the tool list changes."""
|
||||
with self._listeners_lock:
|
||||
self._listeners.append(callback)
|
||||
def add_listener(self, callback: Callable[[], None], *, user_id: str | None = None) -> None:
|
||||
"""Register a callback invoked when the tool list changes.
|
||||
|
||||
def remove_listener(self, callback: Callable[[], None]) -> None:
|
||||
"""Unregister a tool-change callback."""
|
||||
``user_id=None`` (default) registers a global / admin listener
|
||||
that fires on every tool-change (static or pool). A string
|
||||
``user_id`` scopes the listener: it fires on global static-path
|
||||
changes AND on pool changes for that user only — never on
|
||||
another user's pool change. RFC §3.3.
|
||||
"""
|
||||
with self._listeners_lock:
|
||||
self._listeners.append((user_id, callback))
|
||||
|
||||
def remove_listener(self, callback: Callable[[], None], *, user_id: str | None = None) -> None:
|
||||
"""Unregister a tool-change callback.
|
||||
|
||||
``user_id`` MUST match the value used at registration; the
|
||||
``(user_id, callback)`` pair is the listener identity.
|
||||
"""
|
||||
with self._listeners_lock, contextlib.suppress(ValueError):
|
||||
self._listeners.remove(callback)
|
||||
self._listeners.remove((user_id, callback))
|
||||
|
||||
def _notify_listeners(self) -> None:
|
||||
"""Invoke all registered tool-change listeners."""
|
||||
"""Static-path tool change — fires ALL registered listeners.
|
||||
|
||||
The static catalog is a process-wide concern: a static-server
|
||||
notification (or reconcile) updates ``_tool_map`` which every
|
||||
user-scoped session relies on, so the fan-out is unconditional.
|
||||
"""
|
||||
with self._listeners_lock:
|
||||
listeners = list(self._listeners)
|
||||
for _uid, cb in listeners:
|
||||
try:
|
||||
cb()
|
||||
except Exception:
|
||||
log.warning("Tool-change listener raised", exc_info=True)
|
||||
|
||||
def _notify_user_tool_listeners(self, user_id: str) -> None:
|
||||
"""Pool-entry tool change — fires only listeners that should see it.
|
||||
|
||||
A pool catalog change touches one user's view; firing every
|
||||
registered listener would broadcast that user's tool set to
|
||||
unrelated sessions. Scoped fan-out targets the matching
|
||||
``user_id`` AND admin (``None``) listeners. RFC §3.3.
|
||||
"""
|
||||
with self._listeners_lock:
|
||||
listeners = [cb for uid, cb in self._listeners if uid == user_id or uid is None]
|
||||
for cb in listeners:
|
||||
try:
|
||||
cb()
|
||||
@@ -1807,6 +2136,16 @@ class MCPClientManager:
|
||||
"""Connect a new MCP server at runtime (blocks the calling thread).
|
||||
|
||||
Returns status dict with keys: connected, tools, resources, prompts, error.
|
||||
|
||||
Note on ``_oauth_user_server_names``: this method does NOT update
|
||||
the oauth_user-name cache. ``_db_servers_to_config`` strips
|
||||
``auth_type='oauth_user'`` rows before this method is reached, so
|
||||
production callers (only :meth:`reconcile_sync`) never pass an
|
||||
oauth_user cfg here — the cache is rebuilt wholesale by
|
||||
``reconcile_sync`` at the top of every reconcile, which is the
|
||||
canonical update point. Direct test callers passing an oauth_user
|
||||
cfg would leave the cache stale; route through ``reconcile_sync``
|
||||
instead.
|
||||
"""
|
||||
if "__" in name:
|
||||
return {
|
||||
@@ -1929,6 +2268,17 @@ class MCPClientManager:
|
||||
with notification handlers and refresh tasks.
|
||||
|
||||
Returns True if the server was connected and successfully removed.
|
||||
|
||||
Note on ``_oauth_user_server_names``: this method does NOT discard
|
||||
``name`` from the oauth_user-name cache. The cache is the source of
|
||||
truth for "this server exists in the DB as oauth_user", not "this
|
||||
server is connected on the static path" — those are separate
|
||||
concerns. A static→oauth_user transition during reconcile_sync
|
||||
keeps ``name`` in the cache (the new identity) AND calls this
|
||||
method to drop the old static-path connection; discarding here
|
||||
would silently make web_search resolve to a now-pool-only server.
|
||||
The cache is updated only by :meth:`reconcile_sync`'s wholesale
|
||||
rebuild at line 2340.
|
||||
"""
|
||||
existing = self._static_servers.get(name)
|
||||
was_connected = existing is not None and existing.session is not None
|
||||
@@ -2039,6 +2389,14 @@ class MCPClientManager:
|
||||
log.warning("reconcile_sync: failed to read mcp_servers table", exc_info=True)
|
||||
return {"added": [], "removed": [], "updated": []}
|
||||
|
||||
# Refresh the in-memory oauth_user name cache from the rows we
|
||||
# just read — feeds :meth:`server_auth_type` so callers (e.g.
|
||||
# ``web_search.resolve_web_search_client``) avoid a per-turn SQL
|
||||
# roundtrip.
|
||||
self._oauth_user_server_names = {
|
||||
row["name"] for row in rows if row.get("auth_type") == "oauth_user"
|
||||
}
|
||||
|
||||
desired = _db_servers_to_config(rows)
|
||||
desired_names = set(desired)
|
||||
|
||||
@@ -2093,9 +2451,29 @@ class MCPClientManager:
|
||||
|
||||
# -- query methods -------------------------------------------------------
|
||||
|
||||
def get_tools(self) -> list[dict[str, Any]]:
|
||||
"""Return MCP tools in OpenAI function-calling format."""
|
||||
return [dict(t) for t in self._tools]
|
||||
def get_tools(self, user_id: str | None = None) -> list[dict[str, Any]]:
|
||||
"""Return MCP tools in OpenAI function-calling format.
|
||||
|
||||
``user_id=None`` returns the global static-path catalog only —
|
||||
the legacy behaviour every pre-Phase-7 caller relies on. A
|
||||
string ``user_id`` returns the merged view: static catalog
|
||||
followed by that user's cached pool tools, read from
|
||||
``_user_tools`` (a single dict-get, atomic under GIL). The
|
||||
list snapshot is materialised on the mcp-loop by
|
||||
:meth:`_rebuild_user_tool_map` and replaced atomically — sync
|
||||
threads never iterate ``_user_pool_entries`` directly, so the
|
||||
mcp-loop is free to insert/pop entries concurrently.
|
||||
Callers that don't carry a session-bound user (boot-time
|
||||
logging, web_search backend resolution) MUST use the default.
|
||||
|
||||
Returned dicts are shallow-copied; nested objects are shared
|
||||
with the manager's catalog, mirroring the pre-Phase-7 contract.
|
||||
"""
|
||||
base = [dict(t) for t in self._tools]
|
||||
if user_id is None:
|
||||
return base
|
||||
base.extend(dict(t) for t in self._user_tools.get(user_id, []))
|
||||
return base
|
||||
|
||||
def get_resources(self) -> list[dict[str, Any]]:
|
||||
"""Return discovered MCP resources (shallow-copied dicts)."""
|
||||
@@ -2115,24 +2493,47 @@ class MCPClientManager:
|
||||
"""Number of discovered prompts (no allocation)."""
|
||||
return len(self._prompts)
|
||||
|
||||
def is_mcp_tool(self, func_name: str) -> bool:
|
||||
def is_mcp_tool(self, func_name: str, *, user_id: str | None = None) -> bool:
|
||||
"""Check whether *func_name* belongs to an MCP server.
|
||||
|
||||
Caveat: only static-path tools (``auth_type ∈ {none, static}``)
|
||||
populate ``_tool_map``; pool tools become reachable via this
|
||||
method only once per-user catalog scoping lands. Until then,
|
||||
pool dispatch is reachable from production callers like
|
||||
``ChatSession._exec_mcp_tool`` only when the LLM produces a
|
||||
``mcp__{server}__{tool}`` name that bypasses ``is_mcp_tool``-
|
||||
style gating, or via direct ``call_tool_sync`` with a known
|
||||
prefixed name (the path the new pool tests use).
|
||||
``user_id=None`` (default) asks "is this a static-path tool?" —
|
||||
the answer is process-global. Boot-time and per-node callers
|
||||
(e.g. ``resolve_web_search_client``) MUST use the default
|
||||
because they don't carry a session-bound user identity.
|
||||
|
||||
A string ``user_id`` extends the lookup to that user's pool
|
||||
catalog: returns ``True`` if ``func_name`` is in either the
|
||||
static map OR the user's per-user tool map. Pool tools become
|
||||
reachable via this gate ONLY once a session-bound caller
|
||||
threads its ``user_id`` through; CLI sessions default
|
||||
``user_id=""`` and so cannot see pool tools — documented
|
||||
limitation.
|
||||
"""
|
||||
return func_name in self._tool_map
|
||||
if func_name in self._tool_map:
|
||||
return True
|
||||
if user_id is None:
|
||||
return False
|
||||
user_map = self._user_tool_map.get(user_id)
|
||||
return user_map is not None and func_name in user_map
|
||||
|
||||
def is_mcp_prompt(self, name: str) -> bool:
|
||||
"""Check whether *name* is a known MCP prompt."""
|
||||
return name in self._prompt_map
|
||||
|
||||
def server_auth_type(self, server_name: str) -> str | None:
|
||||
"""Return ``'oauth_user'`` for pool-backed servers, else ``None``.
|
||||
|
||||
In-memory accessor for the per-turn callers that need to
|
||||
distinguish pool-backed servers from static-path ones without a
|
||||
SQL roundtrip. ``None`` means "either static-path or unknown" —
|
||||
the boot-time / per-node web_search resolver only uses this as
|
||||
a defence-in-depth gate, so a missing-cache miss is safe (the
|
||||
outer ``is_mcp_tool`` check already proves the server is in
|
||||
``_tool_map``, which by construction excludes oauth_user).
|
||||
Populated by ``reconcile_sync`` and ``create_mcp_client``.
|
||||
"""
|
||||
return "oauth_user" if server_name in self._oauth_user_server_names else None
|
||||
|
||||
@property
|
||||
def server_count(self) -> int:
|
||||
return sum(1 for s in self._static_servers.values() if s.session is not None)
|
||||
@@ -2605,11 +3006,28 @@ class MCPClientManager:
|
||||
# ``streamablehttp_client`` TaskGroup — see
|
||||
# :meth:`_dispatch_pool_sync` for the architectural
|
||||
# rationale.
|
||||
log.debug("mcp_pool.auth_401_initial", exc_info=exc)
|
||||
#
|
||||
# exc_info=False: the underlying httpx exception
|
||||
# carries ``request.headers["authorization"]`` with
|
||||
# the rejected bearer; standard tracebacks don't
|
||||
# render locals but Sentry / faulthandler hooks
|
||||
# capture frame state. Structured fields below
|
||||
# provide the diagnostic signal without the secret.
|
||||
log.debug(
|
||||
"mcp_pool.auth_401_initial server=%s user=%s exc=%s",
|
||||
server_name,
|
||||
user_id,
|
||||
type(exc).__name__,
|
||||
)
|
||||
raise _PoolDispatchRetryRequested from None
|
||||
# retry_count == 1 — refreshed bearer also rejected;
|
||||
# emit consent_required so the user/operator re-grants.
|
||||
log.debug("mcp_pool.auth_401_retry_failed", exc_info=exc)
|
||||
log.debug(
|
||||
"mcp_pool.auth_401_retry_failed server=%s user=%s exc=%s",
|
||||
server_name,
|
||||
user_id,
|
||||
type(exc).__name__,
|
||||
)
|
||||
return _structured_error(
|
||||
code="mcp_consent_required",
|
||||
server=server_name,
|
||||
@@ -2617,7 +3035,12 @@ class MCPClientManager:
|
||||
)
|
||||
if classification == "auth_403":
|
||||
self._evict_session(key)
|
||||
log.debug("mcp_pool.auth_403", exc_info=exc)
|
||||
log.debug(
|
||||
"mcp_pool.auth_403 server=%s user=%s exc=%s",
|
||||
server_name,
|
||||
user_id,
|
||||
type(exc).__name__,
|
||||
)
|
||||
return await self._handle_auth_403(
|
||||
user_id=user_id,
|
||||
server_name=server_name,
|
||||
@@ -2627,7 +3050,15 @@ class MCPClientManager:
|
||||
if classification == "transport":
|
||||
self._cb_record_failure(server_name)
|
||||
self._evict_session(key)
|
||||
log.debug("mcp_pool.transport_failure", exc_info=exc)
|
||||
# See auth_401 branch for why exc_info=False — pool
|
||||
# dispatch errors all share the same request object
|
||||
# whose headers carry the bearer.
|
||||
log.debug(
|
||||
"mcp_pool.transport_failure server=%s user=%s exc=%s",
|
||||
server_name,
|
||||
user_id,
|
||||
type(exc).__name__,
|
||||
)
|
||||
raise
|
||||
# protocol / other — don't trip the breaker.
|
||||
raise
|
||||
@@ -2636,18 +3067,40 @@ class MCPClientManager:
|
||||
return result
|
||||
|
||||
def _evict_session(self, key: tuple[str, str]) -> None:
|
||||
"""Drop the cached session on a pool entry. Stack/streams left for reconnect.
|
||||
"""Drop the cached session AND catalog on a pool entry.
|
||||
|
||||
Auth/transport branches both call this — the next connect's
|
||||
``_connect_one_pool`` tears down the stale stack lazily via the
|
||||
stale-entry guard at the top of the method. Closing eagerly
|
||||
from here is incorrect under cancellation: ``stack.aclose()``
|
||||
must run inside the same anyio scope it was entered in, which
|
||||
the next connect arranges.
|
||||
Stack/streams left for reconnect. Auth/transport branches both
|
||||
call this — the next connect's ``_connect_one_pool`` tears
|
||||
down the stale stack lazily via the stale-entry guard at the
|
||||
top of the method. Closing eagerly from here is incorrect
|
||||
under cancellation: ``stack.aclose()`` must run inside the
|
||||
same anyio scope it was entered in, which the next connect
|
||||
arranges.
|
||||
|
||||
Catalog cleanup (Phase 7): clearing ``entry.tools`` here
|
||||
ensures an evicted-then-not-yet-reconnected entry contributes
|
||||
no stale entries to ``_user_tool_map``. Without this,
|
||||
``is_mcp_tool(name, user_id=user_id)`` could return True for a
|
||||
name whose backing pool is gone, then ``_resolve_pool_target``
|
||||
would dispatch into a session-less entry and the next call
|
||||
would surface as a generic transport error instead of a
|
||||
clean reconnect path.
|
||||
"""
|
||||
evict = self._user_pool_entries.get(key)
|
||||
if evict is not None:
|
||||
evict.session = None
|
||||
evict.tools = None
|
||||
# Prune the debounce stamp in lockstep with the entry — the
|
||||
# dict grows otherwise (slow leak) across (user, server)
|
||||
# churn. Mirrors the cleanup in ``_close_pool_entry_if_idle``.
|
||||
self._last_pool_notification_refresh.pop(key, None)
|
||||
user_id, _server_name = key
|
||||
self._rebuild_user_tool_map(user_id)
|
||||
# Wake the user's session so its merged tool list shrinks
|
||||
# back to static-only until the next connect populates the
|
||||
# entry. Admin (None) listeners also fire — operator
|
||||
# tooling tracking pool catalog state observes the drop.
|
||||
self._notify_user_tool_listeners(user_id)
|
||||
|
||||
async def _handle_auth_403(
|
||||
self,
|
||||
@@ -3129,11 +3582,15 @@ def create_mcp_client(
|
||||
"""
|
||||
# Check DB first to know which servers are DB-managed
|
||||
db_names: set[str] = set()
|
||||
oauth_user_names: set[str] = set()
|
||||
if storage is not None:
|
||||
try:
|
||||
rows = storage.list_mcp_servers(enabled_only=True)
|
||||
if rows:
|
||||
db_names = {r["name"] for r in rows}
|
||||
# Cache oauth_user names so per-turn callers (web_search
|
||||
# backend resolution) can answer auth_type without SQL.
|
||||
oauth_user_names = {r["name"] for r in rows if r.get("auth_type") == "oauth_user"}
|
||||
except Exception:
|
||||
log.warning("Failed to load DB-managed MCP servers", exc_info=True)
|
||||
|
||||
@@ -3144,5 +3601,6 @@ def create_mcp_client(
|
||||
mgr = MCPClientManager(servers)
|
||||
# Mark DB-sourced servers so reconcile_sync won't remove config-file servers
|
||||
mgr._db_managed = {name for name in servers if name in db_names}
|
||||
mgr._oauth_user_server_names = oauth_user_names
|
||||
mgr.start()
|
||||
return mgr
|
||||
|
||||
+52
-10
@@ -712,7 +712,24 @@ class ChatSession:
|
||||
self.debug = False
|
||||
self.auto_approve = False
|
||||
self._node_id = node_id
|
||||
# ``user_id`` is the authenticated principal's UUID for HTTP-borne
|
||||
# sessions (server, console) and the empty string ``""`` for CLI
|
||||
# / eval / unauthenticated callers. The empty string is a SENTINEL
|
||||
# that collapses to ``None`` for MCPClientManager's optional
|
||||
# ``user_id`` arguments (cached once below as ``_mcp_user_id``),
|
||||
# which short-circuits per-user OAuth pool lookup → CLI / eval
|
||||
# sessions cannot use ``auth_type=oauth_user`` MCP servers
|
||||
# (no per-request user identity to attach a bearer to). End
|
||||
# users running such servers must pre-link via the web UI;
|
||||
# static-path MCP servers continue to work in all session
|
||||
# contexts.
|
||||
self._user_id = user_id
|
||||
# ``_user_id`` is set once here and never reassigned, so we
|
||||
# compute the MCP-API form (empty-string-to-None collapse) once
|
||||
# rather than re-asserting the invariant at each call site.
|
||||
# Listener identity, catalog merge, and dispatch all consume
|
||||
# this through ``_mcp_user_id``.
|
||||
self._mcp_user_id: str | None = user_id or None
|
||||
self._username = username
|
||||
self._client_type = client_type
|
||||
self._config_store = config_store
|
||||
@@ -841,13 +858,15 @@ class ChatSession:
|
||||
self._task_tools = []
|
||||
self._agent_tools = []
|
||||
elif mcp_client:
|
||||
mcp_tools = mcp_client.get_tools()
|
||||
mcp_tools = 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)
|
||||
self._agent_tools = merge_mcp_tools(AGENT_TOOLS, mcp_tools)
|
||||
# Register for tool-change notifications from MCP servers
|
||||
# Register for tool-change notifications from MCP servers.
|
||||
# ``user_id`` is the listener identity component — pool-only
|
||||
# changes for OTHER users must not fire this callback.
|
||||
self._mcp_refresh_cb = self._on_mcp_tools_changed
|
||||
mcp_client.add_listener(self._mcp_refresh_cb)
|
||||
mcp_client.add_listener(self._mcp_refresh_cb, user_id=self._mcp_user_id)
|
||||
# Register for resource-change notifications
|
||||
self._mcp_resource_cb = self._on_mcp_resources_changed
|
||||
mcp_client.add_resource_listener(self._mcp_resource_cb)
|
||||
@@ -1203,7 +1222,11 @@ class ChatSession:
|
||||
# is fixed at COORDINATOR_TOOLS. Ignore MCP server changes.
|
||||
if self._kind == WorkstreamKind.COORDINATOR:
|
||||
return
|
||||
mcp_tools = self._mcp_client.get_tools()
|
||||
# Phase 7: pass session-bound user_id so the merged tool list
|
||||
# includes this user's pool catalog. The static path is included
|
||||
# by ``get_tools`` regardless; ``user_id=None`` would silently
|
||||
# drop pool tools that the LLM is allowed to call.
|
||||
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)
|
||||
self._agent_tools = merge_mcp_tools(AGENT_TOOLS, mcp_tools)
|
||||
@@ -1369,7 +1392,11 @@ class ChatSession:
|
||||
if self._judge_cancel_event is not None:
|
||||
self._judge_cancel_event.set()
|
||||
if self._mcp_client and self._mcp_refresh_cb:
|
||||
self._mcp_client.remove_listener(self._mcp_refresh_cb)
|
||||
# ``user_id`` MUST match the value used at registration —
|
||||
# the listener identity is ``(user_id, callback)``, not
|
||||
# callback alone. ``self._user_id`` is set once in
|
||||
# ``__init__`` and never mutated, so identity is stable.
|
||||
self._mcp_client.remove_listener(self._mcp_refresh_cb, user_id=self._mcp_user_id)
|
||||
self._mcp_refresh_cb = None
|
||||
if self._mcp_client and self._mcp_resource_cb:
|
||||
self._mcp_client.remove_resource_listener(self._mcp_resource_cb)
|
||||
@@ -4725,13 +4752,24 @@ class ChatSession:
|
||||
}
|
||||
preparer = preparers.get(func_name)
|
||||
if not preparer:
|
||||
# Check if this is an MCP tool
|
||||
if self._mcp_client and self._mcp_client.is_mcp_tool(func_name):
|
||||
# Check if this is an MCP tool. Phase 7: pass session-bound
|
||||
# ``user_id`` so per-user pool tools become reachable here —
|
||||
# without this kwarg the gate stays static-only and pool
|
||||
# dispatch is structurally unreachable from
|
||||
# ``ChatSession._prepare_tool`` (RFC §3, invariant 8).
|
||||
if self._mcp_client and self._mcp_client.is_mcp_tool(
|
||||
func_name, user_id=self._mcp_user_id
|
||||
):
|
||||
return self._prepare_mcp_tool(call_id, func_name, args)
|
||||
self.ui.on_error(f"Model called unknown tool: {func_name!r}")
|
||||
available = list(preparers)
|
||||
if self._mcp_client:
|
||||
available.extend(sorted(self._mcp_client._tool_map))
|
||||
available.extend(
|
||||
sorted(
|
||||
t["function"]["name"]
|
||||
for t in self._mcp_client.get_tools(user_id=self._mcp_user_id)
|
||||
)
|
||||
)
|
||||
return {
|
||||
"call_id": call_id,
|
||||
"func_name": func_name,
|
||||
@@ -7782,7 +7820,7 @@ class ChatSession:
|
||||
output = self._mcp_client.call_tool_sync(
|
||||
func_name,
|
||||
args,
|
||||
user_id=self._user_id or None,
|
||||
user_id=self._mcp_user_id,
|
||||
timeout=self.tool_timeout,
|
||||
)
|
||||
except TimeoutError:
|
||||
@@ -10148,7 +10186,11 @@ class ChatSession:
|
||||
elif arg and arg.split()[0] == "refresh":
|
||||
self._handle_mcp_refresh(arg)
|
||||
else:
|
||||
tools = self._mcp_client.get_tools()
|
||||
# Phase 7: pass session-bound user_id so the /mcp listing
|
||||
# surfaces this user's pool tools alongside the static
|
||||
# catalog. Resource / prompt query stays user_id-less
|
||||
# (deferred to Phase 7b).
|
||||
tools = self._mcp_client.get_tools(user_id=self._mcp_user_id)
|
||||
resources = self._mcp_client.get_resources()
|
||||
prompts = self._mcp_client.get_prompts()
|
||||
mcp_lines = []
|
||||
|
||||
@@ -168,7 +168,27 @@ def resolve_web_search_client(
|
||||
if len(parts) == 3 and mcp_client is not None:
|
||||
_, server, tool = parts
|
||||
prefixed = f"mcp__{server}__{tool}"
|
||||
# Boot-time gate: ``is_mcp_tool`` without ``user_id`` returns
|
||||
# True only for static-path catalogs. Pool-backed
|
||||
# (``auth_type=oauth_user``) servers are NEVER reachable via
|
||||
# the per-node web_search client because the boot-time
|
||||
# resolver has no per-user identity to attach a bearer to —
|
||||
# the resolved client would be shared across requests, but
|
||||
# the bearer can't be (RFC §3, invariant 8 corollary).
|
||||
if mcp_client.is_mcp_tool(prefixed):
|
||||
# Defence-in-depth: even if a future change widens
|
||||
# ``_tool_map`` to include oauth_user names by accident,
|
||||
# refuse the backend explicitly. ``server_auth_type``
|
||||
# is an in-memory accessor — this resolver is invoked
|
||||
# per LLM turn, so a SQL hop here would amplify token
|
||||
# cost on every chat round.
|
||||
if mcp_client.server_auth_type(server) == "oauth_user":
|
||||
log.warning(
|
||||
"web_search_backend %r points at oauth_user MCP server; "
|
||||
"per-node web search cannot use per-user tokens — disabling",
|
||||
backend,
|
||||
)
|
||||
return None
|
||||
return MCPSearchClient(mcp_client, prefixed, timeout=timeout)
|
||||
return None
|
||||
|
||||
|
||||
Reference in New Issue
Block a user