From dad98c062d75d4beaec35c246df03e1d50bdcbf3 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Wed, 6 May 2026 12:52:41 -0700 Subject: [PATCH] =?UTF-8?q?feat(mcp):=20per-user=20catalog=20scoping=20(Ph?= =?UTF-8?q?ase=207=20=E2=80=94=20tools)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- tests/test_mcp_client.py | 358 ++++++++++++++++ tests/test_mcp_user_catalog.py | 718 +++++++++++++++++++++++++++++++++ tests/test_mcp_user_pool.py | 5 + tests/test_web_search.py | 47 +++ turnstone/core/mcp_client.py | 546 +++++++++++++++++++++++-- turnstone/core/session.py | 62 ++- turnstone/core/web_search.py | 20 + 7 files changed, 1702 insertions(+), 54 deletions(-) create mode 100644 tests/test_mcp_user_catalog.py diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 378314fa..9b9cc85c 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -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): diff --git a/tests/test_mcp_user_catalog.py b/tests/test_mcp_user_catalog.py new file mode 100644 index 00000000..19bbc1df --- /dev/null +++ b/tests/test_mcp_user_catalog.py @@ -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" diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index ffeef835..8c06d2d1 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -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 diff --git a/tests/test_web_search.py b/tests/test_web_search.py index 936eaed7..dd120cb2 100644 --- a/tests/test_web_search.py +++ b/tests/test_web_search.py @@ -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) diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 86b387bc..1aa16c80 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -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 diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 1f943d98..f62827f5 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -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 = [] diff --git a/turnstone/core/web_search.py b/turnstone/core/web_search.py index 45fc2d18..8fb022bf 100644 --- a/turnstone/core/web_search.py +++ b/turnstone/core/web_search.py @@ -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