diff --git a/tests/test_mcp_pool_auth_prompt_integration.py b/tests/test_mcp_pool_auth_prompt_integration.py new file mode 100644 index 00000000..a7d9b20e --- /dev/null +++ b/tests/test_mcp_pool_auth_prompt_integration.py @@ -0,0 +1,720 @@ +"""Phase 7b integration tests — real-transport prompt get 401/403/etc. + +Mirror of :mod:`tests.test_mcp_pool_auth_resource_integration` for the +prompt path (RFC §3.3). Drives through the real ``streamablehttp_client``, +real httpx response-hook plumbing, and a real upstream subprocess +(``FastMCP`` with a programmable ``BehaviorMiddleware``). Direct +``httpx.HTTPStatusError`` injection is forbidden (invariant 14). +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +import socket +import threading +import time +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any +from unittest.mock import MagicMock, patch + +import pytest +import uvicorn +from mcp.server.fastmcp import FastMCP +from starlette.middleware.base import BaseHTTPMiddleware + +from tests.conftest import make_mcp_token_cipher +from turnstone.core.mcp_client import MCPClientManager +from turnstone.core.mcp_crypto import MCPTokenStore +from turnstone.core.mcp_oauth import TokenLookupResult +from turnstone.core.storage._sqlite import SQLiteBackend + +if TYPE_CHECKING: + from collections.abc import Callable + + from starlette.requests import Request + from starlette.responses import Response + +logging.getLogger("uvicorn.error").setLevel(logging.WARNING) +logging.getLogger("uvicorn.access").setLevel(logging.WARNING) +logging.getLogger("mcp").setLevel(logging.WARNING) + + +class BehaviorMiddleware(BaseHTTPMiddleware): + """Programmable upstream behaviour — see + :mod:`tests.test_mcp_pool_auth_integration` for the semantics. This + copy serves the prompt integration tests. + """ + + def __init__(self, app: Any, behaviour: dict[str, Any]) -> None: + super().__init__(app) + self._behaviour = behaviour + + async def dispatch(self, request: Request, call_next: Callable[..., Any]) -> Response: + from starlette.responses import Response as StarletteResponse + + if request.method == "POST" and "/mcp" in str(request.url): + self._behaviour.setdefault("post_auth_headers", []).append( + request.headers.get("authorization") + ) + + mode = self._behaviour.get("mode", "never") + if mode == "once_401": + if not self._behaviour.get("_fired"): + self._behaviour["_fired"] = True + return StarletteResponse( + "unauthorized", + status_code=401, + headers={ + "www-authenticate": self._behaviour.get( + "www_authenticate", 'Bearer error="invalid_token"' + ) + }, + ) + elif mode == "always_401": + return StarletteResponse( + "unauthorized", + status_code=401, + headers={ + "www-authenticate": self._behaviour.get( + "www_authenticate", 'Bearer error="invalid_token"' + ) + }, + ) + elif mode == "once_403_insufficient": + if not self._behaviour.get("_fired"): + self._behaviour["_fired"] = True + return StarletteResponse( + "forbidden", + status_code=403, + headers={ + "www-authenticate": self._behaviour.get( + "www_authenticate", + 'Bearer error="insufficient_scope", scope="prompts:read"', + ) + }, + ) + elif mode == "once_403_generic" and not self._behaviour.get("_fired"): + self._behaviour["_fired"] = True + return StarletteResponse( + "forbidden", + status_code=403, + headers={ + "www-authenticate": self._behaviour.get("www_authenticate", "Bearer realm=mcp") + }, + ) + return await call_next(request) + + +def _find_free_port() -> int: + s = socket.socket() + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def _build_server(port: int, behaviour: dict[str, Any]) -> uvicorn.Server: + mcp = FastMCP(name="phase7b-prompt-target", streamable_http_path="/mcp") + + @mcp.prompt() + def greet(who: str = "world") -> str: + return f"Hello, {who}!" + + @mcp.prompt() + def summarize(topic: str = "today") -> str: + return f"Please summarize {topic}." + + app = mcp.streamable_http_app() + app.add_middleware(BehaviorMiddleware, behaviour=behaviour) + config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning", access_log=False) + return uvicorn.Server(config) + + +def _wait_ready(port: int, timeout: float = 5.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.5): + return + except OSError: + time.sleep(0.05) + raise TimeoutError(f"upstream at 127.0.0.1:{port} not ready after {timeout}s") + + +@pytest.fixture +def upstream(): + port = _find_free_port() + behaviour: dict[str, Any] = {} + server = _build_server(port, behaviour) + + def _run() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(server.serve()) + + t = threading.Thread(target=_run, daemon=True, name="phase7b-prompt-upstream") + t.start() + try: + _wait_ready(port) + yield f"http://127.0.0.1:{port}/mcp", behaviour + finally: + server.should_exit = True + t.join(timeout=5) + + +@pytest.fixture +def storage(tmp_path: Any) -> SQLiteBackend: + return SQLiteBackend(str(tmp_path / "test.db")) + + +def _seed_oauth_server( + storage: SQLiteBackend, + *, + name: str = "pool-srv", + server_id: str = "srv-pool", + url: str = "https://mcp.example.com/sse", +) -> None: + storage.create_mcp_server( + server_id=server_id, + name=name, + transport="streamable-http", + url=url, + auth_type="oauth_user", + oauth_client_id="client-abc", + oauth_scopes="openid", + oauth_audience=url, + ) + + +def _seed_user_token( + storage: SQLiteBackend, + cipher: Any, + *, + user_id: str = "user-1", + server_name: str = "pool-srv", + expires_in_seconds: int = 3600, + access_token: str = "access-aaa", + refresh_token: str | None = "refresh-rrr", +) -> None: + expires_at = (datetime.now(UTC) + timedelta(seconds=expires_in_seconds)).strftime( + "%Y-%m-%dT%H:%M:%S" + ) + store = MCPTokenStore(storage, cipher, node_id="test") + store.create_user_token( + user_id, + server_name, + access_token=access_token, + refresh_token=refresh_token, + expires_at=expires_at, + scopes="openid", + as_issuer="https://as.example.com", + audience="https://mcp.example.com", + ) + + +def _make_app_state(storage: SQLiteBackend, *, cipher: Any) -> SimpleNamespace: + return SimpleNamespace( + auth_storage=storage, + mcp_token_store=MCPTokenStore(storage, cipher, node_id="test"), + mcp_oauth_http_client=MagicMock(), + mcp_oauth_refresh_locks={}, + mcp_oauth_metadata_cache={}, + ) + + +@pytest.fixture +def running_loop_mgr(): + 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 _seed_pool_prompt_map( + mgr: MCPClientManager, + user_id: str, + server_name: str, + prefixed_name: str, + original_name: str, +) -> None: + """Pre-seed ``_user_prompt_map`` so ``_resolve_pool_target_prompt`` + finds the prefixed name. Production wires this through + ``_connect_one_pool``; the integration tests seed it directly so the + test focuses on the dispatch behaviour after resolution succeeds. + """ + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry((user_id, server_name)) + entry.prompts = [ + { + "name": prefixed_name, + "original_name": original_name, + "server": server_name, + "description": "", + "arguments": [], + } + ] + mgr._rebuild_user_prompt_map(user_id) + + assert mgr._loop is not None + asyncio.run_coroutine_threadsafe(_seed(), mgr._loop).result(timeout=5) + + +# --------------------------------------------------------------------------- +# I-PR-1: 401 → refresh → retry → success (prompt path) +# --------------------------------------------------------------------------- + + +def test_prompt_get_401_refresh_and_retry_succeeds( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Real upstream returns 401 once, then 200. Carrier captures 401, + force_refresh=True mints a new bearer, retry returns the prompt + messages. Hard invariant 3: breaker counter remains 0. + """ + url, behaviour = upstream + behaviour["mode"] = "once_401" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="token", token="refreshed-bearer") + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + messages = mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "everyone"}, + user_id="user-1", + timeout=15, + ) + + assert isinstance(messages, list) + assert len(messages) == 1 + assert messages[0]["role"] == "user" + assert "everyone" in messages[0]["content"] + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + post_headers = behaviour.get("post_auth_headers", []) + assert len(post_headers) >= 2, f"expected >=2 POSTs; got {len(post_headers)}" + assert post_headers[0] != post_headers[1], ( + "retry attached the same bearer as the initial; the dispatcher " + "did not pick up the refreshed token." + ) + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.session is not None + + +# --------------------------------------------------------------------------- +# I-PR-2: persistent 401 → mcp_consent_required (prompt path) → RuntimeError +# --------------------------------------------------------------------------- + + +def test_prompt_get_persistent_401_emits_consent_required( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, behaviour = upstream + behaviour["mode"] = "always_401" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="token", token="refreshed-bearer") + return TokenLookupResult(kind="token", token="access-aaa") + + with ( + patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ), + pytest.raises(RuntimeError) as excinfo, + ): + mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "world"}, + user_id="user-1", + timeout=15, + ) + + payload = json.loads(str(excinfo.value)) + assert payload["error"]["code"] == "mcp_consent_required" + assert payload["error"]["server"] == "pool-srv" + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + +# --------------------------------------------------------------------------- +# I-PR-3: 403 + insufficient_scope → mcp_insufficient_scope (prompt path) +# --------------------------------------------------------------------------- + + +def test_prompt_get_403_insufficient_scope_emits_structured_error( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, behaviour = upstream + behaviour["mode"] = "once_403_insufficient" + behaviour["www_authenticate"] = 'Bearer error="insufficient_scope", scope="prompts:read"' + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + with ( + patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ), + pytest.raises(RuntimeError) as excinfo, + ): + mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "world"}, + user_id="user-1", + timeout=15, + ) + + payload = json.loads(str(excinfo.value)) + assert payload["error"]["code"] == "mcp_insufficient_scope" + assert payload["error"]["scopes_required"] == ["prompts:read"] + post_headers = behaviour.get("post_auth_headers", []) + assert len(post_headers) == 1, ( + f"403 must NOT trigger a retry; observed {len(post_headers)} POSTs" + ) + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + +# --------------------------------------------------------------------------- +# I-PR-3b: 403 generic → mcp_prompt_get_forbidden +# --------------------------------------------------------------------------- + + +def test_prompt_get_403_generic_forbidden( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, behaviour = upstream + behaviour["mode"] = "once_403_generic" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + with ( + patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ), + pytest.raises(RuntimeError) as excinfo, + ): + mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "world"}, + user_id="user-1", + timeout=15, + ) + + payload = json.loads(str(excinfo.value)) + # Per the kind="prompt" wiring of `_handle_auth_403`, the + # operation-specific code surfaces here rather than the tool path's + # generic mcp_tool_call_forbidden. + assert payload["error"]["code"] == "mcp_prompt_get_forbidden" + assert "scopes_required" not in payload["error"] + + +# --------------------------------------------------------------------------- +# I-PR-6: breaker isolation — auth failures NEVER trip the breaker +# --------------------------------------------------------------------------- + + +def test_prompt_get_breaker_unaffected_by_auth_failures( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Repeated 401 + refresh-failed cycles leave breaker at 0 + (hard invariant 3 verified end-to-end for the prompt path).""" + url, behaviour = upstream + behaviour["mode"] = "always_401" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="refresh_failed") + return TokenLookupResult(kind="token", token="access-aaa") + + # Re-seed each iteration: symmetric eviction (Phase 7b) clears + # ``_user_prompt_map`` on auth failure so the next dispatch's + # resolver would miss without a fresh seed. Production reconnect + # repopulates this; the test simulates that out-of-band. + for _ in range(10): + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + with ( + patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ), + pytest.raises(RuntimeError) as excinfo, + ): + mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "world"}, + user_id="user-1", + timeout=15, + ) + payload = json.loads(str(excinfo.value)) + assert payload["error"]["code"] == "mcp_consent_required" + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + +# --------------------------------------------------------------------------- +# Negative tests — token lookup edge cases (prompt path) +# --------------------------------------------------------------------------- + + +def test_prompt_get_missing_token_emits_consent_required( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, _behaviour = upstream + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="missing") + + with ( + patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ), + pytest.raises(RuntimeError) as excinfo, + ): + mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "world"}, + user_id="user-1", + timeout=10, + ) + + payload = json.loads(str(excinfo.value)) + assert payload["error"]["code"] == "mcp_consent_required" + + +def test_prompt_get_decrypt_failure_emits_token_undecryptable( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, _behaviour = upstream + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="decrypt_failure") + + with ( + patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ), + pytest.raises(RuntimeError) as excinfo, + ): + mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "world"}, + user_id="user-1", + timeout=10, + ) + + payload = json.loads(str(excinfo.value)) + assert payload["error"]["code"] == "mcp_token_undecryptable_key_unknown" + + +def test_prompt_get_http_url_emits_url_insecure( + running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """An ``http://`` (non-loopback) oauth_user URL must surface + ``mcp_oauth_url_insecure`` BEFORE the bearer is attached. + """ + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url="http://example.com/mcp") + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + with ( + patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ), + pytest.raises(RuntimeError) as excinfo, + ): + mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "world"}, + user_id="user-1", + timeout=5, + ) + + payload = json.loads(str(excinfo.value)) + assert payload["error"]["code"] == "mcp_oauth_url_insecure" + + +def test_prompt_get_unknown_name_raises_value_error( + running_loop_mgr: Any, +) -> None: + """When the prefixed name doesn't resolve to either pool or static, + the static-path code raises ``ValueError``. Per-user-first + resolution (scope decision 0.1) means user_id-bearing callers still + hit this path when their pool catalog doesn't carry the name.""" + mgr, _loop, _ = running_loop_mgr + with pytest.raises(ValueError, match="Unknown MCP prompt"): + mgr.get_prompt_sync( + "mcp__nonexistent__missing", + None, + user_id="user-1", + timeout=5, + ) + + +# --------------------------------------------------------------------------- +# I-PR-E2E: real discovery + dispatch in same connect (no _seed_pool_prompt_map) +# --------------------------------------------------------------------------- + + +def test_prompt_get_e2e_discovery_then_dispatch_succeeds( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Drive REAL discovery + dispatch end-to-end through the pool path. + + Mirror of the tool path's + ``test_integration_pool_reuse_401_refresh_and_retry_succeeds``: skips + the ``_seed_pool_prompt_map`` shortcut and lets ``_connect_one_pool`` + populate ``_user_prompt_map`` from the real ``prompts/list`` + upstream response. Verifies that the entry's discovered prompts + match what the FastMCP fixture advertises AND that + ``_user_prompt_map[user_id]`` is populated with the prefixed name + after dispatch — proving the discovery path actually fired. + + This is the structural gate against a regression where prompt + dispatch silently bypasses discovery (e.g., a mis-wired resolver + that finds the (server, original) via prefix-parsing alone never + populates the per-user catalog). + """ + url, behaviour = upstream + behaviour["mode"] = "never" # passthrough — discovery + dispatch both succeed + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + # NB: no `_seed_pool_prompt_map` — the resolver finds (server, original) + # via the `mcp__{server}__{prompt}` prefix and hands off to + # ``_dispatch_pool_prompt_sync``, which lazy-connects via + # ``_connect_one_pool``. The connect runs the real ``prompts/list`` + # against the FastMCP fixture and populates the per-user catalog. + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + messages = mgr.get_prompt_sync( + "mcp__pool-srv__greet", + {"who": "world"}, + user_id="user-1", + timeout=15, + ) + + assert isinstance(messages, list) + assert len(messages) == 1 + assert messages[0]["role"] == "user" + assert "world" in messages[0]["content"] + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + # Discovery populated the entry's prompts with both fixtures + # (``greet`` and ``summarize``) — proves real ``prompts/list`` + # ran during the connect, not just the targeted ``prompts/get``. + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.session is not None + assert entry.prompts is not None + discovered_names = {p["name"] for p in entry.prompts} + assert "mcp__pool-srv__greet" in discovered_names + assert "mcp__pool-srv__summarize" in discovered_names + + # ``_rebuild_user_prompt_map`` ran during the connect, populating the + # per-user catalog. This is the signal that discovery wired into the + # routing tables — without it, a follow-up ``get_prompt_sync`` would + # need to re-resolve via prefix parsing every time. + user_prompt_map = mgr._user_prompt_map.get("user-1") or {} + assert "mcp__pool-srv__greet" in user_prompt_map + assert "mcp__pool-srv__summarize" in user_prompt_map diff --git a/tests/test_mcp_pool_auth_resource_integration.py b/tests/test_mcp_pool_auth_resource_integration.py new file mode 100644 index 00000000..27320e7f --- /dev/null +++ b/tests/test_mcp_pool_auth_resource_integration.py @@ -0,0 +1,669 @@ +"""Phase 7b integration tests — real-transport resource read 401/403/etc. + +Mirror of :mod:`tests.test_mcp_pool_auth_integration` for the resource +path (RFC §3.2). Drives through the real ``streamablehttp_client``, +real httpx response-hook plumbing, and a real upstream subprocess +(``FastMCP`` with a programmable ``BehaviorMiddleware``). Direct +``httpx.HTTPStatusError`` injection is forbidden (invariant 14). +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +import socket +import threading +import time +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any +from unittest.mock import MagicMock, patch + +import pytest +import uvicorn +from mcp.server.fastmcp import FastMCP +from starlette.middleware.base import BaseHTTPMiddleware + +from tests.conftest import make_mcp_token_cipher +from turnstone.core.mcp_client import MCPClientManager +from turnstone.core.mcp_crypto import MCPTokenStore +from turnstone.core.mcp_oauth import TokenLookupResult +from turnstone.core.storage._sqlite import SQLiteBackend + +if TYPE_CHECKING: + from collections.abc import Callable + + from starlette.requests import Request + from starlette.responses import Response + +logging.getLogger("uvicorn.error").setLevel(logging.WARNING) +logging.getLogger("uvicorn.access").setLevel(logging.WARNING) +logging.getLogger("mcp").setLevel(logging.WARNING) + + +class BehaviorMiddleware(BaseHTTPMiddleware): + """Programmable upstream behaviour — see + :mod:`tests.test_mcp_pool_auth_integration` for the semantics. This + copy serves the resource integration tests. + """ + + def __init__(self, app: Any, behaviour: dict[str, Any]) -> None: + super().__init__(app) + self._behaviour = behaviour + + async def dispatch(self, request: Request, call_next: Callable[..., Any]) -> Response: + from starlette.responses import Response as StarletteResponse + + if request.method == "POST" and "/mcp" in str(request.url): + self._behaviour.setdefault("post_auth_headers", []).append( + request.headers.get("authorization") + ) + + mode = self._behaviour.get("mode", "never") + if mode == "once_401": + if not self._behaviour.get("_fired"): + self._behaviour["_fired"] = True + return StarletteResponse( + "unauthorized", + status_code=401, + headers={ + "www-authenticate": self._behaviour.get( + "www_authenticate", 'Bearer error="invalid_token"' + ) + }, + ) + elif mode == "always_401": + return StarletteResponse( + "unauthorized", + status_code=401, + headers={ + "www-authenticate": self._behaviour.get( + "www_authenticate", 'Bearer error="invalid_token"' + ) + }, + ) + elif mode == "once_403_insufficient": + if not self._behaviour.get("_fired"): + self._behaviour["_fired"] = True + return StarletteResponse( + "forbidden", + status_code=403, + headers={ + "www-authenticate": self._behaviour.get( + "www_authenticate", + 'Bearer error="insufficient_scope", scope="files:read"', + ) + }, + ) + elif mode == "once_403_generic" and not self._behaviour.get("_fired"): + self._behaviour["_fired"] = True + return StarletteResponse( + "forbidden", + status_code=403, + headers={ + "www-authenticate": self._behaviour.get("www_authenticate", "Bearer realm=mcp") + }, + ) + return await call_next(request) + + +def _find_free_port() -> int: + s = socket.socket() + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def _build_server(port: int, behaviour: dict[str, Any]) -> uvicorn.Server: + mcp = FastMCP(name="phase7b-resource-target", streamable_http_path="/mcp") + + @mcp.resource("res://hello") + def hello() -> str: + return "world" + + @mcp.resource("res://json/data") + def jdata() -> str: + return '{"k": 1}' + + # Echo tool exists so the e2e test can trigger ``_connect_one_pool`` + # (and the full tool + resource + prompt discovery) via prefix-parsed + # ``call_tool_sync`` BEFORE the resource read. The other tests in this + # module use ``_seed_pool_resource_map`` and never invoke tools, so + # adding the tool is invisible to them. + @mcp.tool() + async def echo(payload: str = "default") -> str: + return f"echoed:{payload}" + + app = mcp.streamable_http_app() + app.add_middleware(BehaviorMiddleware, behaviour=behaviour) + config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning", access_log=False) + return uvicorn.Server(config) + + +def _wait_ready(port: int, timeout: float = 5.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.5): + return + except OSError: + time.sleep(0.05) + raise TimeoutError(f"upstream at 127.0.0.1:{port} not ready after {timeout}s") + + +@pytest.fixture +def upstream(): + port = _find_free_port() + behaviour: dict[str, Any] = {} + server = _build_server(port, behaviour) + + def _run() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(server.serve()) + + t = threading.Thread(target=_run, daemon=True, name="phase7b-resource-upstream") + t.start() + try: + _wait_ready(port) + yield f"http://127.0.0.1:{port}/mcp", behaviour + finally: + server.should_exit = True + t.join(timeout=5) + + +@pytest.fixture +def storage(tmp_path: Any) -> SQLiteBackend: + return SQLiteBackend(str(tmp_path / "test.db")) + + +def _seed_oauth_server( + storage: SQLiteBackend, + *, + name: str = "pool-srv", + server_id: str = "srv-pool", + url: str = "https://mcp.example.com/sse", +) -> None: + storage.create_mcp_server( + server_id=server_id, + name=name, + transport="streamable-http", + url=url, + auth_type="oauth_user", + oauth_client_id="client-abc", + oauth_scopes="openid", + oauth_audience=url, + ) + + +def _seed_user_token( + storage: SQLiteBackend, + cipher: Any, + *, + user_id: str = "user-1", + server_name: str = "pool-srv", + expires_in_seconds: int = 3600, + access_token: str = "access-aaa", + refresh_token: str | None = "refresh-rrr", +) -> None: + expires_at = (datetime.now(UTC) + timedelta(seconds=expires_in_seconds)).strftime( + "%Y-%m-%dT%H:%M:%S" + ) + store = MCPTokenStore(storage, cipher, node_id="test") + store.create_user_token( + user_id, + server_name, + access_token=access_token, + refresh_token=refresh_token, + expires_at=expires_at, + scopes="openid", + as_issuer="https://as.example.com", + audience="https://mcp.example.com", + ) + + +def _make_app_state(storage: SQLiteBackend, *, cipher: Any) -> SimpleNamespace: + return SimpleNamespace( + auth_storage=storage, + mcp_token_store=MCPTokenStore(storage, cipher, node_id="test"), + mcp_oauth_http_client=MagicMock(), + mcp_oauth_refresh_locks={}, + mcp_oauth_metadata_cache={}, + ) + + +@pytest.fixture +def running_loop_mgr(): + 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 _seed_pool_resource_map( + mgr: MCPClientManager, user_id: str, server_name: str, uri: str +) -> None: + """Pre-seed ``_user_resource_map`` so ``_resolve_pool_target_resource`` + finds the URI. Production wires this through ``_connect_one_pool``; + the integration tests seed it directly so the test focuses on the + dispatch behaviour after resolution succeeds. + """ + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry((user_id, server_name)) + entry.resources = [ + { + "uri": uri, + "name": "", + "description": "", + "mimeType": "", + "server": server_name, + } + ] + mgr._rebuild_user_resource_map(user_id) + + assert mgr._loop is not None + asyncio.run_coroutine_threadsafe(_seed(), mgr._loop).result(timeout=5) + + +# --------------------------------------------------------------------------- +# I-RP-1: 401 → refresh → retry → success (resource path) +# --------------------------------------------------------------------------- + + +def test_resource_read_401_refresh_and_retry_succeeds( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Real upstream returns 401 once, then 200. Carrier captures 401, + force_refresh=True mints a new bearer, retry returns the resource. + Hard invariant 3: breaker counter remains 0. + """ + url, behaviour = upstream + behaviour["mode"] = "once_401" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="token", token="refreshed-bearer") + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15) + + assert result == "world" + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + post_headers = behaviour.get("post_auth_headers", []) + assert len(post_headers) >= 2, f"expected >=2 POSTs; got {len(post_headers)}" + assert post_headers[0] != post_headers[1], ( + "retry attached the same bearer as the initial; the dispatcher " + "did not pick up the refreshed token." + ) + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.session is not None + + +# --------------------------------------------------------------------------- +# I-RP-2: persistent 401 → mcp_consent_required (resource path) +# --------------------------------------------------------------------------- + + +def test_resource_read_persistent_401_emits_consent_required( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, behaviour = upstream + behaviour["mode"] = "always_401" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="token", token="refreshed-bearer") + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_consent_required" + assert payload["error"]["server"] == "pool-srv" + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + +# --------------------------------------------------------------------------- +# I-RP-3: 403 + insufficient_scope → mcp_insufficient_scope (resource path) +# --------------------------------------------------------------------------- + + +def test_resource_read_403_insufficient_scope_emits_structured_error( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, behaviour = upstream + behaviour["mode"] = "once_403_insufficient" + behaviour["www_authenticate"] = 'Bearer error="insufficient_scope", scope="files:read"' + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_insufficient_scope" + assert payload["error"]["scopes_required"] == ["files:read"] + post_headers = behaviour.get("post_auth_headers", []) + assert len(post_headers) == 1, ( + f"403 must NOT trigger a retry; observed {len(post_headers)} POSTs" + ) + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + +# --------------------------------------------------------------------------- +# I-RP-3b: 403 generic → mcp_resource_read_forbidden +# --------------------------------------------------------------------------- + + +def test_resource_read_403_generic_forbidden( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, behaviour = upstream + behaviour["mode"] = "once_403_generic" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15) + + payload = json.loads(result) + # Per the kind="resource" wiring of `_handle_auth_403`, the + # operation-specific code surfaces here rather than the tool path's + # generic mcp_tool_call_forbidden. + assert payload["error"]["code"] == "mcp_resource_read_forbidden" + assert "scopes_required" not in payload["error"] + + +# --------------------------------------------------------------------------- +# I-RP-6: breaker isolation — auth failures NEVER trip the breaker +# --------------------------------------------------------------------------- + + +def test_resource_read_breaker_unaffected_by_auth_failures( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Repeated 401 + refresh-failed cycles leave breaker at 0 + (hard invariant 3 verified end-to-end for the resource path).""" + url, behaviour = upstream + behaviour["mode"] = "always_401" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="refresh_failed") + return TokenLookupResult(kind="token", token="access-aaa") + + # Re-seed each iteration: symmetric eviction (Phase 7b) clears + # ``_user_resource_map`` on auth failure so the next dispatch's + # resolver would miss without a fresh seed. Production reconnect + # repopulates this; the test simulates that out-of-band. + for _ in range(10): + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15) + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_consent_required" + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + +# --------------------------------------------------------------------------- +# Negative tests — token lookup edge cases (resource path) +# --------------------------------------------------------------------------- + + +def test_resource_read_missing_token_emits_consent_required( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, _behaviour = upstream + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="missing") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=10) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_consent_required" + + +def test_resource_read_decrypt_failure_emits_token_undecryptable( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + url, _behaviour = upstream + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="decrypt_failure") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=10) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_token_undecryptable_key_unknown" + + +def test_resource_read_http_url_emits_url_insecure( + running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """An ``http://`` (non-loopback) oauth_user URL must surface + ``mcp_oauth_url_insecure`` BEFORE the bearer is attached. + """ + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url="http://example.com/mcp") + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + _seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello") + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=5) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_oauth_url_insecure" + + +def test_resource_read_unknown_uri_raises_value_error( + running_loop_mgr: Any, +) -> None: + """When the URI doesn't resolve to either pool or static, the + static-path code raises ``ValueError``. Per-user-first resolution + (scope decision 0.1) means user_id-bearing callers still hit this + path when their pool catalog doesn't carry the URI.""" + mgr, _loop, _ = running_loop_mgr + with pytest.raises(ValueError, match="Unknown MCP resource"): + mgr.read_resource_sync("res://nonexistent", user_id="user-1", timeout=5) + + +# --------------------------------------------------------------------------- +# I-RP-E2E: real discovery + dispatch in same connect (no _seed_pool_resource_map) +# --------------------------------------------------------------------------- + + +def test_resource_read_e2e_discovery_then_dispatch_succeeds( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Drive REAL discovery + dispatch end-to-end through the pool path. + + Mirror of the tool path's + ``test_integration_pool_reuse_401_refresh_and_retry_succeeds``: skips + the ``_seed_pool_resource_map`` shortcut and lets ``_connect_one_pool`` + populate ``_user_resource_map`` from the real ``resources/list`` + upstream response. Verifies that the entry's discovered resources + match what the FastMCP fixture advertises AND that + ``_user_resource_map[user_id]`` is populated with the URI(s) after + discovery — proving the discovery path actually fired. + + Resource URIs do NOT carry a server-name prefix (unlike tools and + prompts), so the resource resolver cannot derive (server, uri) by + parsing alone. The test triggers the connect via a prefix-parsed + ``call_tool_sync`` first (which runs the full + tools+resources+prompts discovery against the FastMCP fixture), + then drives ``read_resource_sync`` against a URI that the + upstream advertised — proving that real discovery wired the URI + into the per-user catalog. + + Structural gate against a regression where resource discovery is + silently skipped (e.g., a capability-gating bug that drops the + ``resources/list`` call but keeps the connect succeeding). + """ + url, behaviour = upstream + behaviour["mode"] = "never" # passthrough — discovery + dispatch both succeed + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv", url=url) + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + # NB: no `_seed_pool_resource_map` — the connect runs the real + # ``resources/list`` against the FastMCP fixture and populates the + # per-user catalog. The tool call below triggers that connect because + # ``_resolve_pool_target`` derives (server, original) from the + # ``mcp__pool-srv__echo`` prefix and lazy-connects via + # ``_connect_one_pool``. + + async def _fake_classified(**_kwargs: Any) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + # Step 1: trigger the connect via prefix-parsed tool dispatch. + # Discovery (tools + resources + prompts) populates the per-user + # catalogs. + tool_result = mgr.call_tool_sync( + "mcp__pool-srv__echo", {"payload": "ignite"}, user_id="user-1", timeout=15 + ) + assert "echoed:ignite" in tool_result + + # Step 2: now that discovery has populated ``_user_resource_map``, + # the resource resolver finds ``res://hello`` and dispatches the + # read on the SAME pool entry / session. + result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15) + + assert result == "world" + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + # Discovery populated the entry's resources with both fixtures + # (``res://hello`` and ``res://json/data``) — proves real + # ``resources/list`` ran during the connect, not just the targeted + # ``resources/read``. + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.session is not None + assert entry.resources is not None + discovered_uris = {r["uri"] for r in entry.resources if not r.get("template")} + assert "res://hello" in discovered_uris + assert "res://json/data" in discovered_uris + + # ``_rebuild_user_resource_map`` ran during the connect, populating + # the per-user catalog. This is the signal that discovery wired into + # the routing tables — without it, ``read_resource_sync`` would have + # raised ValueError because the resolver had no entry for the URI. + user_resource_map = mgr._user_resource_map.get("user-1") or {} + assert "res://hello" in user_resource_map + assert "res://json/data" in user_resource_map diff --git a/tests/test_mcp_user_catalog.py b/tests/test_mcp_user_catalog.py index 19bbc1df..251c627e 100644 --- a/tests/test_mcp_user_catalog.py +++ b/tests/test_mcp_user_catalog.py @@ -119,6 +119,11 @@ def _make_jsonrpc_handler( list_tools_status: int = 200, list_tools_error_payload: dict[str, Any] | None = None, list_tools_seq: list[dict[str, Any]] | None = None, + list_resources_response: dict[str, Any] | None = None, + list_resource_templates_response: dict[str, Any] | None = None, + list_prompts_response: dict[str, Any] | None = None, + list_resources_seq: list[dict[str, Any]] | None = None, + list_prompts_seq: list[dict[str, Any]] | None = None, counter: list[int] | None = None, record_bodies: list[str] | None = None, ) -> Any: @@ -142,9 +147,49 @@ def _make_jsonrpc_handler( "id": 1, "result": {"tools": []}, } + if list_resources_response is None: + list_resources_response = { + "jsonrpc": "2.0", + "id": 2, + "result": {"resources": []}, + } + if list_resource_templates_response is None: + list_resource_templates_response = { + "jsonrpc": "2.0", + "id": 3, + "result": {"resourceTemplates": []}, + } + if list_prompts_response is None: + list_prompts_response = { + "jsonrpc": "2.0", + "id": 4, + "result": {"prompts": []}, + } counter = counter if counter is not None else [0] list_tools_seq = list_tools_seq or [] list_tools_index = [0] + list_resources_seq = list_resources_seq or [] + list_resources_index = [0] + list_prompts_seq = list_prompts_seq or [] + list_prompts_index = [0] + + def _extract_id(body: str) -> Any: + """Pull the JSON-RPC request id out of *body* via a tiny regex. + + The SDK assigns request ids sequentially per :class:`ClientSession` + (initialize=0, list_tools=1, list_resources=2, ...). The mock + responses MUST echo the same id so the SDK's correlator can route + the response back to the awaiting future. Extracting from the + request body lets a single response template work regardless of + which list method ran first. + """ + import json as _json + import re as _re + + match = _re.search(r'"id"\s*:\s*(\d+)', body) + if match: + return _json.loads(match.group(1)) + return 1 async def _handler(req: httpx.Request) -> httpx.Response: if req.method == "GET": @@ -157,11 +202,46 @@ def _make_jsonrpc_handler( if "notifications/initialized" in body: return httpx.Response(202) counter[0] += 1 + req_id = _extract_id(body) if "method" in body and '"initialize"' in body: + response_payload = dict(init_response) + response_payload["id"] = req_id return httpx.Response( 200, headers={"content-type": "application/json", "mcp-session-id": "sess-1"}, - json=init_response, + json=response_payload, + ) + if '"resources/templates/list"' in body: + payload = dict(list_resource_templates_response) + payload["id"] = req_id + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + json=payload, + ) + if '"resources/list"' in body: + if list_resources_seq and list_resources_index[0] < len(list_resources_seq): + payload = dict(list_resources_seq[list_resources_index[0]]) + list_resources_index[0] += 1 + else: + payload = dict(list_resources_response) + payload["id"] = req_id + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + json=payload, + ) + if '"prompts/list"' in body: + if list_prompts_seq and list_prompts_index[0] < len(list_prompts_seq): + payload = dict(list_prompts_seq[list_prompts_index[0]]) + list_prompts_index[0] += 1 + else: + payload = dict(list_prompts_response) + payload["id"] = req_id + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + json=payload, ) # ``tools/list`` — by request order: explicit override first, then # the staged sequence, then fall back to ``list_tools_response``. @@ -172,17 +252,15 @@ def _make_jsonrpc_handler( 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]] + payload = dict(list_tools_seq[list_tools_index[0]]) list_tools_index[0] += 1 - return httpx.Response( - 200, - headers={"content-type": "application/json"}, - json=payload, - ) + else: + payload = dict(list_tools_response) + payload["id"] = req_id return httpx.Response( 200, headers={"content-type": "application/json"}, - json=list_tools_response, + json=payload, ) return _handler @@ -197,6 +275,38 @@ def _list_tools_payload(tools: list[dict[str, Any]], req_id: int = 1) -> dict[st } +def _list_resources_payload(resources: list[dict[str, Any]], req_id: int = 2) -> dict[str, Any]: + """Build a minimal ``resources/list`` JSON-RPC result payload.""" + return {"jsonrpc": "2.0", "id": req_id, "result": {"resources": resources}} + + +def _list_prompts_payload(prompts: list[dict[str, Any]], req_id: int = 4) -> dict[str, Any]: + """Build a minimal ``prompts/list`` JSON-RPC result payload.""" + return {"jsonrpc": "2.0", "id": req_id, "result": {"prompts": prompts}} + + +def _init_response_with_caps( + *, resources: bool = False, prompts: bool = False, tools: bool = True +) -> dict[str, Any]: + """Build an ``initialize`` response that advertises selected capabilities.""" + caps: dict[str, Any] = {} + if tools: + caps["tools"] = {"listChanged": False} + if resources: + caps["resources"] = {"listChanged": False} + if prompts: + caps["prompts"] = {"listChanged": False} + return { + "jsonrpc": "2.0", + "id": 0, + "result": { + "protocolVersion": "2025-06-18", + "capabilities": caps, + "serverInfo": {"name": "fake", "version": "1"}, + }, + } + + def _tool_spec(name: str, description: str = "") -> dict[str, Any]: return { "name": name, @@ -205,6 +315,18 @@ def _tool_spec(name: str, description: str = "") -> dict[str, Any]: } +def _resource_spec(uri: str, *, name: str = "") -> dict[str, Any]: + return {"uri": uri, "name": name or uri, "mimeType": "text/plain"} + + +def _prompt_spec(name: str, *, description: str = "") -> dict[str, Any]: + return { + "name": name, + "description": description or f"prompt {name}", + "arguments": [], + } + + def _connect_pool( mgr: MCPClientManager, loop: asyncio.AbstractEventLoop, @@ -716,3 +838,861 @@ def test_listener_identity_includes_user_id() -> None: calls.clear() mgr._notify_listeners() assert calls == [], "callback fired after BOTH user_id-keyed registrations removed" + + +# --------------------------------------------------------------------------- +# Phase 7b — listener-identity widening for resources and prompts +# --------------------------------------------------------------------------- + + +def test_resource_listener_identity_includes_user_id() -> None: + """``add_resource_listener`` / ``remove_resource_listener`` use the + ``(user_id, callback)`` tuple as identity, mirroring the tool path. + + Negative-test reproducer: revert the listener storage to a flat + list of callables — this test fails because removing one + registration accidentally removes both. + """ + mgr = MCPClientManager({}) + calls = [] + + def _shared_cb() -> None: + calls.append(1) + + mgr.add_resource_listener(_shared_cb, user_id="user-A") + mgr.add_resource_listener(_shared_cb, user_id="user-B") + mgr.remove_resource_listener(_shared_cb, user_id="user-A") + mgr._notify_resource_listeners() + assert len(calls) == 1, ( + f"shared callback fired {len(calls)} times after removing user-A; expected 1" + ) + + mgr.remove_resource_listener(_shared_cb, user_id="user-B") + calls.clear() + mgr._notify_resource_listeners() + assert calls == [], "callback fired after BOTH registrations removed" + + +def test_prompt_listener_identity_includes_user_id() -> None: + """Mirror of :func:`test_resource_listener_identity_includes_user_id` + for prompt listeners — Phase 7b widens both APIs.""" + mgr = MCPClientManager({}) + calls = [] + + def _shared_cb() -> None: + calls.append(1) + + mgr.add_prompt_listener(_shared_cb, user_id="user-A") + mgr.add_prompt_listener(_shared_cb, user_id="user-B") + mgr.remove_prompt_listener(_shared_cb, user_id="user-A") + mgr._notify_prompt_listeners() + assert len(calls) == 1, ( + f"shared callback fired {len(calls)} times after removing user-A; expected 1" + ) + + mgr.remove_prompt_listener(_shared_cb, user_id="user-B") + calls.clear() + mgr._notify_prompt_listeners() + assert calls == [], "callback fired after BOTH registrations removed" + + +def test_user_resource_listeners_isolated_to_matching_user() -> None: + """``_notify_user_resource_listeners(user_id)`` fires only matching + user-id listeners + admin (None), never another user's listener. + RFC §3.3 — pool catalog change is private to its owning user. + """ + mgr = MCPClientManager({}) + user_a_calls = [0] + user_b_calls = [0] + admin_calls = [0] + + def _user_a() -> None: + user_a_calls[0] += 1 + + def _user_b() -> None: + user_b_calls[0] += 1 + + def _admin() -> None: + admin_calls[0] += 1 + + mgr.add_resource_listener(_user_a, user_id="user-A") + mgr.add_resource_listener(_user_b, user_id="user-B") + mgr.add_resource_listener(_admin) # admin / None + + mgr._notify_user_resource_listeners("user-A") + assert user_a_calls[0] == 1 + assert user_b_calls[0] == 0, ( + f"unrelated user-B listener fired {user_b_calls[0]} times — RFC §3.3 violation" + ) + assert admin_calls[0] == 1 + + +def test_user_prompt_listeners_isolated_to_matching_user() -> None: + """Mirror for prompt listeners — see resource version above.""" + mgr = MCPClientManager({}) + user_a_calls = [0] + user_b_calls = [0] + admin_calls = [0] + + def _user_a() -> None: + user_a_calls[0] += 1 + + def _user_b() -> None: + user_b_calls[0] += 1 + + def _admin() -> None: + admin_calls[0] += 1 + + mgr.add_prompt_listener(_user_a, user_id="user-A") + mgr.add_prompt_listener(_user_b, user_id="user-B") + mgr.add_prompt_listener(_admin) + + mgr._notify_user_prompt_listeners("user-A") + assert user_a_calls[0] == 1 + assert user_b_calls[0] == 0, ( + f"unrelated user-B listener fired {user_b_calls[0]} times — RFC §3.3 violation" + ) + assert admin_calls[0] == 1 + + +# --------------------------------------------------------------------------- +# Phase 7b — _rebuild_user_resource_map / _rebuild_user_prompt_map +# (sibling-cache invariant; loop-only writes; per-user isolation) +# --------------------------------------------------------------------------- + + +def test_rebuild_user_resource_map_isolates_users(running_loop_mgr: Any) -> None: + """``_rebuild_user_resource_map`` populates user-A's resource map + from user-A's pool entries only; user-B remains untouched. Empty + rebuilds drop the user_id key from all three sibling dicts. + """ + mgr, loop, _ = running_loop_mgr + + async def _seed() -> None: + a_entry = await mgr._ensure_pool_entry(("user-A", "pool-srv")) + a_entry.session = MagicMock() + a_entry.resources = [ + { + "uri": "res://a/1", + "name": "r1", + "description": "", + "mimeType": "", + "server": "pool-srv", + }, + ] + b_entry = await mgr._ensure_pool_entry(("user-B", "pool-srv")) + b_entry.session = MagicMock() + b_entry.resources = [ + { + "uri": "res://b/1", + "name": "r2", + "description": "", + "mimeType": "", + "server": "pool-srv", + }, + ] + mgr._rebuild_user_resource_map("user-A") + mgr._rebuild_user_resource_map("user-B") + + _run_on_loop(loop, _seed()) + + a_map = mgr._user_resource_map.get("user-A") or {} + b_map = mgr._user_resource_map.get("user-B") or {} + assert "res://a/1" in a_map and "res://b/1" not in a_map + assert "res://b/1" in b_map and "res://a/1" not in b_map + assert mgr._user_resources["user-A"][0]["uri"] == "res://a/1" + assert mgr._user_resources["user-B"][0]["uri"] == "res://b/1" + + # Empty rebuild after clearing entries drops the key from all three + # dicts so idle users don't retain empty-list sentinels. + async def _clear() -> None: + mgr._user_pool_entries[("user-A", "pool-srv")].resources = None + mgr._rebuild_user_resource_map("user-A") + + _run_on_loop(loop, _clear()) + assert "user-A" not in mgr._user_resource_map + assert "user-A" not in mgr._user_resources + assert "user-A" not in mgr._user_template_prefixes + + +def test_rebuild_user_resource_map_carries_template_prefixes( + running_loop_mgr: Any, +) -> None: + """Template entries (``r["template"] == True``) populate the + per-user template-prefix dict; the longer-prefix-wins collision + policy mirrors the static path.""" + mgr, loop, _ = running_loop_mgr + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-A", "pool-srv")) + entry.session = MagicMock() + entry.resources = [ + { + "uri": "res://item/{id}", + "name": "items", + "description": "", + "mimeType": "", + "server": "pool-srv", + "template": True, + }, + ] + mgr._rebuild_user_resource_map("user-A") + + _run_on_loop(loop, _seed()) + prefixes = mgr._user_template_prefixes.get("user-A") or {} + assert "res://item/" in prefixes + assert prefixes["res://item/"] == ("pool-srv", "res://item/{id}") + + +def test_rebuild_user_prompt_map_isolates_users(running_loop_mgr: Any) -> None: + """Mirror of resource isolation test for prompts.""" + mgr, loop, _ = running_loop_mgr + + async def _seed() -> None: + a_entry = await mgr._ensure_pool_entry(("user-A", "pool-srv")) + a_entry.session = MagicMock() + a_entry.prompts = [ + { + "name": "mcp__pool-srv__a_prompt", + "original_name": "a_prompt", + "server": "pool-srv", + "description": "", + "arguments": [], + } + ] + b_entry = await mgr._ensure_pool_entry(("user-B", "pool-srv")) + b_entry.session = MagicMock() + b_entry.prompts = [ + { + "name": "mcp__pool-srv__b_prompt", + "original_name": "b_prompt", + "server": "pool-srv", + "description": "", + "arguments": [], + } + ] + mgr._rebuild_user_prompt_map("user-A") + mgr._rebuild_user_prompt_map("user-B") + + _run_on_loop(loop, _seed()) + assert "mcp__pool-srv__a_prompt" in (mgr._user_prompt_map.get("user-A") or {}) + assert "mcp__pool-srv__a_prompt" not in (mgr._user_prompt_map.get("user-B") or {}) + assert "mcp__pool-srv__b_prompt" in (mgr._user_prompt_map.get("user-B") or {}) + + +def test_shutdown_clears_user_resource_and_prompt_state() -> None: + """``shutdown`` releases all per-user catalog dicts (Phase 7b + addition); without this clear, idle users would retain + map+list+template_prefix sentinels across the manager lifecycle. + """ + mgr = MCPClientManager({}) + # Seed all per-user state directly (no loop required for the dict + # mutation; ``shutdown`` is also no-op when ``_loop`` is None). + mgr._user_tool_map["user-A"] = {"x": ("s", "x")} + mgr._user_tools["user-A"] = [{"function": {"name": "x"}}] + mgr._user_resource_map["user-A"] = {"u": ("s", "u")} + mgr._user_resources["user-A"] = [{"uri": "u"}] + mgr._user_template_prefixes["user-A"] = {"u/": ("s", "u/{id}")} + mgr._user_prompt_map["user-A"] = {"p": ("s", "p")} + mgr._user_prompts["user-A"] = [{"name": "p"}] + + mgr.shutdown() + assert mgr._user_tool_map == {} + assert mgr._user_tools == {} + assert mgr._user_resource_map == {} + assert mgr._user_resources == {} + assert mgr._user_template_prefixes == {} + assert mgr._user_prompt_map == {} + assert mgr._user_prompts == {} + + +# --------------------------------------------------------------------------- +# Phase 7b — public read API widening (get_resources / get_prompts / +# is_mcp_prompt / *_count_for_user / _match_template) +# --------------------------------------------------------------------------- + + +def test_get_resources_user_id_none_returns_static_only() -> None: + """``get_resources()`` with no kwargs returns the static-only + catalog — preserves the pre-Phase-7b contract for legacy callers + (admin endpoints, boot-time logging).""" + mgr = MCPClientManager({}) + static_res = {"uri": "res://static/1", "server": "static-srv"} + pool_res = {"uri": "res://pool/1", "server": "pool-srv"} + mgr._resources = [static_res] + mgr._user_resources["user-A"] = [pool_res] + + assert mgr.get_resources() == [static_res] + # Sibling user_id-less default doesn't reach into per-user state. + assert mgr.get_resources(None) == [static_res] + + +def test_get_resources_user_id_returns_per_user_first_then_static() -> None: + """Per-user-first ordering (scope decision 0.1): pool resources + come BEFORE static in the merged list. The same user with no pool + entries sees just the static catalog.""" + mgr = MCPClientManager({}) + static_res = {"uri": "res://static/1", "server": "static-srv"} + pool_res = {"uri": "res://pool/1", "server": "pool-srv"} + mgr._resources = [static_res] + mgr._user_resources["user-A"] = [pool_res] + + merged = mgr.get_resources("user-A") + assert merged[0]["uri"] == "res://pool/1" # per-user-first + assert merged[1]["uri"] == "res://static/1" + + # User-B has no pool entries → static-only merged view. + assert mgr.get_resources("user-B") == [static_res] + + +def test_get_prompts_user_id_returns_per_user_first_then_static() -> None: + """Mirror of resource version for prompts.""" + mgr = MCPClientManager({}) + static_p = {"name": "mcp__static-srv__p1", "server": "static-srv"} + pool_p = {"name": "mcp__pool-srv__p1", "server": "pool-srv"} + mgr._prompts = [static_p] + mgr._user_prompts["user-A"] = [pool_p] + + merged = mgr.get_prompts("user-A") + assert merged[0]["name"] == "mcp__pool-srv__p1" + assert merged[1]["name"] == "mcp__static-srv__p1" + assert mgr.get_prompts() == [static_p] + + +def test_is_mcp_prompt_user_id_extends_lookup_to_pool() -> None: + """``is_mcp_prompt(name, user_id="...")`` returns True when the + name lives in either the static map OR the user's per-user prompt + map. Without the kwarg, only the static map is consulted.""" + mgr = MCPClientManager({}) + mgr._prompt_map["mcp__static-srv__p"] = ("static-srv", "p") + mgr._user_prompt_map["user-A"] = {"mcp__pool-srv__p": ("pool-srv", "p")} + + assert mgr.is_mcp_prompt("mcp__static-srv__p") is True + assert mgr.is_mcp_prompt("mcp__pool-srv__p") is False # no user_id → static only + assert mgr.is_mcp_prompt("mcp__pool-srv__p", user_id="user-A") is True + assert mgr.is_mcp_prompt("mcp__pool-srv__p", user_id="user-B") is False + + +def test_resource_count_for_user_includes_pool() -> None: + """``resource_count_for_user`` reports static + the user's pool; + the legacy ``resource_count`` property stays static-only (admin + endpoints rely on the property contract).""" + mgr = MCPClientManager({}) + mgr._resources = [{"uri": "res://static/1"}, {"uri": "res://static/2"}] + mgr._user_resources["user-A"] = [{"uri": "res://pool/1"}] + + assert mgr.resource_count == 2 # process-global, unchanged + assert mgr.resource_count_for_user(None) == 2 + assert mgr.resource_count_for_user("user-A") == 3 + assert mgr.resource_count_for_user("user-B") == 2 # no pool entries + + +def test_prompt_count_for_user_includes_pool() -> None: + """Mirror of resource version for prompts.""" + mgr = MCPClientManager({}) + mgr._prompts = [{"name": "p1"}] + mgr._user_prompts["user-A"] = [{"name": "pa"}, {"name": "pb"}] + + assert mgr.prompt_count == 1 + assert mgr.prompt_count_for_user("user-A") == 3 + assert mgr.prompt_count_for_user("user-B") == 1 + + +def test_match_template_per_user_first_resolution() -> None: + """Per-user templates win over static templates at the same prefix + (scope decision 0.1). Without ``user_id``, only static prefixes + are consulted.""" + mgr = MCPClientManager({}) + mgr._template_prefixes["res://item/"] = ("static-srv", "res://item/{id}") + mgr._user_template_prefixes["user-A"] = { + "res://item/": ("pool-srv", "res://item/{id}"), + } + + # Pre-Phase-7b call (no user_id) keeps static behaviour. + assert mgr._match_template("res://item/42") == ("static-srv", "res://item/{id}") + # User-A sees the per-user template first. + assert mgr._match_template("res://item/42", user_id="user-A") == ( + "pool-srv", + "res://item/{id}", + ) + # User-B has no per-user templates → falls through to static. + assert mgr._match_template("res://item/42", user_id="user-B") == ( + "static-srv", + "res://item/{id}", + ) + + +def test_match_template_longest_prefix_within_user_scope() -> None: + """Within the per-user index, longest-prefix-wins matching applies + just like the static path. The static fall-through MUST NOT win + when the per-user index already produced a match.""" + mgr = MCPClientManager({}) + mgr._template_prefixes["res://"] = ("static-srv", "res://{rest}") + mgr._user_template_prefixes["user-A"] = { + "res://item/": ("pool-srv", "res://item/{id}"), + } + # User-A's longer-prefix template wins even though static has a + # shorter prefix that would also match. + assert mgr._match_template("res://item/42", user_id="user-A") == ( + "pool-srv", + "res://item/{id}", + ) + + +# --------------------------------------------------------------------------- +# Phase 7b — discovery via REAL streamablehttp_client + httpx.MockTransport +# (invariant 14: boundary-crossing code drives through real SDK plumbing). +# --------------------------------------------------------------------------- + + +def test_pool_resource_discovery_on_connect_via_real_streamable_http( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """``_connect_one_pool`` discovers the user's resource catalog after + ``initialize()`` (capability-gated) and stores it on the entry plus + the per-user resource map. + + Negative-test reproducer: revert the discovery + rebuild block in + ``_connect_one_pool`` (drop the ``list_resources`` call and the + ``_rebuild_user_resource_map`` invocation): this test fails because + ``entry.resources`` stays None and ``_user_resource_map`` never + gains the user_id key. + """ + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(resources=True), + list_resources_response=_list_resources_payload( + [_resource_spec("res://test/1"), _resource_spec("res://test/2")] + ), + ) + _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.supports_resources is True + assert entry.resources is not None + uris = {r["uri"] for r in entry.resources} + assert uris == {"res://test/1", "res://test/2"} + user_map = mgr._user_resource_map.get("user-1") + assert user_map is not None + assert user_map["res://test/1"] == ("pool-srv", "res://test/1") + # Per-user resource list reflects discovery. + assert {r["uri"] for r in mgr._user_resources["user-1"]} == {"res://test/1", "res://test/2"} + + +def test_pool_prompt_discovery_on_connect_via_real_streamable_http( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """``_connect_one_pool`` discovers the user's prompt catalog + capability-gated. Mirror of resource discovery test.""" + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(prompts=True), + list_prompts_response=_list_prompts_payload( + [_prompt_spec("greet"), _prompt_spec("summary")] + ), + ) + _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.supports_prompts is True + assert entry.prompts is not None + names = {p["name"] for p in entry.prompts} + assert names == {"mcp__pool-srv__greet", "mcp__pool-srv__summary"} + user_map = mgr._user_prompt_map.get("user-1") + assert user_map is not None + assert user_map["mcp__pool-srv__greet"] == ("pool-srv", "greet") + # is_mcp_prompt with user_id surfaces the discovered prompt. + assert mgr.is_mcp_prompt("mcp__pool-srv__greet", user_id="user-1") is True + assert mgr.is_mcp_prompt("mcp__pool-srv__greet") is False + + +def test_pool_discovery_skips_resources_and_prompts_without_capability( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """When the server does NOT advertise resources / prompts in its + initialize response, ``_connect_one_pool`` MUST skip the discovery + round-trips entirely. ``entry.supports_resources`` / + ``supports_prompts`` stay False; the catalog stays None. + + Negative-test: drop the capability gates in ``_connect_one_pool`` + (force unconditional discovery) — this test fails because the + handler returns no resources/list response, the SDK times out, and + the connect coroutine raises. + """ + mgr, loop, _ = running_loop_mgr + # init_response defaults declare ONLY tools. + handler = _make_jsonrpc_handler() + _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.supports_resources is False + assert entry.supports_prompts is False + assert entry.resources is None + assert entry.prompts is None + # No per-user resource/prompt map entry for this user. + assert "user-1" not in mgr._user_resource_map + assert "user-1" not in mgr._user_prompt_map + + +def test_pool_resource_user_isolation( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """Two users connecting to the same resource-bearing server have + independent per-user resource maps; one user never sees another's + catalog.""" + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(resources=True), + list_resources_response=_list_resources_payload([_resource_spec("res://shared")]), + ) + _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_resource_map.get("user-1") or {} + map_2 = mgr._user_resource_map.get("user-2") or {} + assert map_1 is not map_2 + assert "res://shared" in map_1 and "res://shared" in map_2 + + +def test_pool_prompt_user_isolation(running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch) -> None: + """Mirror for prompts — independent per-user prompt maps.""" + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(prompts=True), + list_prompts_response=_list_prompts_payload([_prompt_spec("shared_prompt")]), + ) + _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_prompt_map.get("user-1") or {} + map_2 = mgr._user_prompt_map.get("user-2") or {} + assert map_1 is not map_2 + assert "mcp__pool-srv__shared_prompt" in map_1 + assert "mcp__pool-srv__shared_prompt" in map_2 + + +# --------------------------------------------------------------------------- +# Phase 7b — refresh on notification (resources/prompts list_changed) +# --------------------------------------------------------------------------- + + +def test_refresh_pool_server_resources_isolates_user( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """``_refresh_pool_server_resources`` refreshes THIS user's catalog only, + leaving the static path and other users' catalogs untouched. + + Negative-test: revert the refresh to call ``_rebuild_resources()`` + (the static path) instead of ``_rebuild_user_resource_map(user_id)``; + this test fails because the static map is mutated and the user map + doesn't pick up the rotated catalog. + """ + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(resources=True), + list_resources_seq=[ + _list_resources_payload([_resource_spec("res://v1")]), + _list_resources_payload([_resource_spec("res://v2")]), + ], + ) + _build_mock_transport_factory(mgr, monkeypatch, handler) + _patch_tcp_probe(mgr, monkeypatch) + + # Pre-seed a static-path resource so we can assert it's untouched. + from turnstone.core.mcp_client import StaticServerState + + sentinel_static = StaticServerState(name="static-srv", session=MagicMock()) + sentinel_static.resources = [ + {"uri": "res://static/keep", "server": "static-srv", "name": "", "mimeType": ""} + ] + mgr._static_servers["static-srv"] = sentinel_static + mgr._resource_map["res://static/keep"] = ("static-srv", "res://static/keep") + + # Pre-seed user-2's pool entry so we can assert it stays untouched. + async def _seed_other_user() -> None: + entry = await mgr._ensure_pool_entry(("user-2", "pool-srv")) + entry.session = MagicMock() + entry.supports_resources = True + entry.resources = [ + {"uri": "res://user2_only", "server": "pool-srv", "name": "", "mimeType": ""} + ] + mgr._rebuild_user_resource_map("user-2") + + _run_on_loop(loop, _seed_other_user()) + assert "res://user2_only" in (mgr._user_resource_map.get("user-2") or {}) + + _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") + assert "res://v1" in (mgr._user_resource_map.get("user-1") or {}) + + added_removed = _run_on_loop(loop, mgr._refresh_pool_server_resources(("user-1", "pool-srv"))) + added, removed = added_removed + assert added == ["res://v2"] + assert removed == ["res://v1"] + + # User-1 sees the new resource. + assert "res://v2" in (mgr._user_resource_map.get("user-1") or {}) + assert "res://v1" not in (mgr._user_resource_map.get("user-1") or {}) + # Static path untouched — invariant 1. + assert mgr._resource_map == {"res://static/keep": ("static-srv", "res://static/keep")} + # User-2's pool catalog untouched. + assert mgr._user_resource_map["user-2"] == { + "res://user2_only": ("pool-srv", "res://user2_only") + } + + +def test_refresh_pool_server_prompts_isolates_user( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """Mirror for prompts — refresh isolates to one user.""" + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(prompts=True), + list_prompts_seq=[ + _list_prompts_payload([_prompt_spec("v1")]), + _list_prompts_payload([_prompt_spec("v2")]), + ], + ) + _build_mock_transport_factory(mgr, monkeypatch, handler) + _patch_tcp_probe(mgr, monkeypatch) + + # Pre-seed a static-path prompt and another user's pool prompt. + from turnstone.core.mcp_client import StaticServerState + + sentinel = StaticServerState(name="static-srv", session=MagicMock()) + sentinel.prompts = [ + { + "name": "mcp__static-srv__keep", + "original_name": "keep", + "server": "static-srv", + "description": "", + "arguments": [], + } + ] + mgr._static_servers["static-srv"] = sentinel + mgr._prompt_map["mcp__static-srv__keep"] = ("static-srv", "keep") + + async def _seed_other_user() -> None: + entry = await mgr._ensure_pool_entry(("user-2", "pool-srv")) + entry.session = MagicMock() + entry.supports_prompts = True + entry.prompts = [ + { + "name": "mcp__pool-srv__user2", + "original_name": "user2", + "server": "pool-srv", + "description": "", + "arguments": [], + } + ] + mgr._rebuild_user_prompt_map("user-2") + + _run_on_loop(loop, _seed_other_user()) + + _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") + assert "mcp__pool-srv__v1" in (mgr._user_prompt_map.get("user-1") or {}) + + added_removed = _run_on_loop(loop, mgr._refresh_pool_server_prompts(("user-1", "pool-srv"))) + added, removed = added_removed + assert added == ["mcp__pool-srv__v2"] + assert removed == ["mcp__pool-srv__v1"] + + assert mgr._prompt_map == {"mcp__static-srv__keep": ("static-srv", "keep")} + assert mgr._user_prompt_map["user-2"] == {"mcp__pool-srv__user2": ("pool-srv", "user2")} + + +def test_refresh_pool_server_resources_skips_when_capability_unset( + running_loop_mgr: Any, +) -> None: + """If ``entry.supports_resources`` is False, the refresh returns + ``([], [])`` without making any list calls — no SDK round-trip.""" + mgr, loop, _ = running_loop_mgr + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-1", "pool-srv")) + entry.session = MagicMock() # session present but capability unset + entry.supports_resources = False + + _run_on_loop(loop, _seed()) + added, removed = _run_on_loop(loop, mgr._refresh_pool_server_resources(("user-1", "pool-srv"))) + assert added == [] + assert removed == [] + + +def test_refresh_pool_server_prompts_skips_when_capability_unset( + running_loop_mgr: Any, +) -> None: + """Mirror for prompts.""" + mgr, loop, _ = running_loop_mgr + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-1", "pool-srv")) + entry.session = MagicMock() + entry.supports_prompts = False + + _run_on_loop(loop, _seed()) + added, removed = _run_on_loop(loop, mgr._refresh_pool_server_prompts(("user-1", "pool-srv"))) + assert added == [] + assert removed == [] + + +# --------------------------------------------------------------------------- +# Phase 7b — symmetric eviction (resources/prompts cleared + listeners fire) +# --------------------------------------------------------------------------- + + +def test_eviction_clears_resource_and_prompt_catalogs_and_fires_listeners( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """``_evict_session`` clears ``entry.resources`` and ``entry.prompts``, + rebuilds both per-user maps (drops the now-empty entries), and + fires the matching user-keyed + admin listeners for ALL three + catalogs (tools, resources, prompts). + + Negative-test: drop the resource/prompt cleanup additions in + ``_evict_session``: this test fails because the user-resource map + keeps the evicted URIs and no resource listener fires. + """ + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(resources=True, prompts=True), + list_resources_response=_list_resources_payload([_resource_spec("res://r/1")]), + list_prompts_response=_list_prompts_payload([_prompt_spec("p1")]), + ) + _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 "res://r/1" in (mgr._user_resource_map.get("user-1") or {}) + assert "mcp__pool-srv__p1" in (mgr._user_prompt_map.get("user-1") or {}) + + res_calls = [0] + prompt_calls = [0] + other_res_calls = [0] + other_prompt_calls = [0] + + def _user_res_cb() -> None: + res_calls[0] += 1 + + def _user_prompt_cb() -> None: + prompt_calls[0] += 1 + + def _other_res_cb() -> None: + other_res_calls[0] += 1 + + def _other_prompt_cb() -> None: + other_prompt_calls[0] += 1 + + mgr.add_resource_listener(_user_res_cb, user_id="user-1") + mgr.add_resource_listener(_other_res_cb, user_id="user-2") + mgr.add_prompt_listener(_user_prompt_cb, user_id="user-1") + mgr.add_prompt_listener(_other_prompt_cb, user_id="user-2") + + mgr._evict_session(("user-1", "pool-srv")) + + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.session is None + assert entry.tools is None + assert entry.resources is None + assert entry.prompts is None + # Per-user maps drop the evicted entries. + assert "user-1" not in mgr._user_resource_map + assert "user-1" not in mgr._user_prompt_map + # Listeners fire for the evicted user but not for the unrelated user. + assert res_calls[0] == 1 + assert prompt_calls[0] == 1 + assert other_res_calls[0] == 0, "unrelated user-2 resource listener fired — RFC §3.3 violation" + assert other_prompt_calls[0] == 0, "unrelated user-2 prompt listener fired — RFC §3.3 violation" + + +def test_close_pool_entry_if_idle_clears_resource_and_prompt_catalogs( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """LRU/TTL eviction (``_close_pool_entry_if_idle``) symmetric + cleanup for resources & prompts. Bug-pair to the tools-only + cleanup added in Phase 7 round-2.""" + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(resources=True, prompts=True), + list_resources_response=_list_resources_payload([_resource_spec("res://r/1")]), + list_prompts_response=_list_prompts_payload([_prompt_spec("p1")]), + ) + _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") + + res_calls = [0] + prompt_calls = [0] + + def _res_cb() -> None: + res_calls[0] += 1 + + def _prompt_cb() -> None: + prompt_calls[0] += 1 + + mgr.add_resource_listener(_res_cb, user_id="user-1") + mgr.add_prompt_listener(_prompt_cb, user_id="user-1") + + _run_on_loop(loop, mgr._close_pool_entry_if_idle(key)) + + # Entry fully removed. + assert key not in mgr._user_pool_entries + # Per-user catalogs cleared. + assert "user-1" not in mgr._user_resource_map + assert "user-1" not in mgr._user_prompt_map + # Listeners fired exactly once. + assert res_calls[0] == 1 + assert prompt_calls[0] == 1 + + +def test_reconnect_after_eviction_repopulates_resource_and_prompt_catalogs( + running_loop_mgr: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """After eviction, the next ``_connect_one_pool`` re-populates the + resource and prompt catalogs from the SDK. Bug-class: an extra + glue layer would only repopulate tools.""" + mgr, loop, _ = running_loop_mgr + handler = _make_jsonrpc_handler( + init_response=_init_response_with_caps(resources=True, prompts=True), + list_resources_seq=[ + _list_resources_payload([_resource_spec("res://a")]), + _list_resources_payload([_resource_spec("res://b")]), + ], + list_prompts_seq=[ + _list_prompts_payload([_prompt_spec("p_a")]), + _list_prompts_payload([_prompt_spec("p_b")]), + ], + ) + _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 "res://a" in (mgr._user_resource_map.get("user-1") or {}) + assert "mcp__pool-srv__p_a" in (mgr._user_prompt_map.get("user-1") or {}) + + mgr._evict_session(("user-1", "pool-srv")) + assert "user-1" not in mgr._user_resource_map + assert "user-1" not in mgr._user_prompt_map + + _connect_pool(mgr, loop, user_id="user-1", server_name="pool-srv") + # New catalogs reflect the rotated payload. + assert "res://b" in (mgr._user_resource_map.get("user-1") or {}) + assert "res://a" not in (mgr._user_resource_map.get("user-1") or {}) + assert "mcp__pool-srv__p_b" in (mgr._user_prompt_map.get("user-1") or {}) + assert "mcp__pool-srv__p_a" not in (mgr._user_prompt_map.get("user-1") or {}) diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index 8c06d2d1..18289a85 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -199,10 +199,17 @@ 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. + # Phase 7b: ``_connect_one_pool`` discovers tools, resources, + # and prompts after ``initialize()`` returns (resources/prompts + # capability-gated). The capability stub returns a tools-only + # advertisement so the test can keep its narrow focus on the + # bearer-injection contract; resources/prompts paths are + # exercised by the real-transport tests in + # ``tests/test_mcp_user_catalog.py``. + fake_caps = MagicMock() + fake_caps.resources = None + fake_caps.prompts = None + fake_session.get_server_capabilities = MagicMock(return_value=fake_caps) fake_session.list_tools = AsyncMock(return_value=MagicMock(tools=[])) def _stream_factory(*, url: str, headers: dict[str, str]) -> _AsyncCM: diff --git a/tests/test_sessions.py b/tests/test_sessions.py index c6ff6d78..74eb8fc5 100644 --- a/tests/test_sessions.py +++ b/tests/test_sessions.py @@ -914,8 +914,11 @@ class TestMCPToolGating: """read_resource excluded when MCP client has no resources.""" mcp_client = MagicMock() mcp_client.get_tools.return_value = [] - mcp_client.resource_count = 0 - mcp_client.prompt_count = 2 + # Phase 7b: gating uses ``*_count_for_user`` so the test mocks + # the per-user variant (the property remains for static-only + # admin paths). Returning 0 / 2 mirrors the prior contract. + mcp_client.resource_count_for_user.return_value = 0 + mcp_client.prompt_count_for_user.return_value = 2 session = ChatSession( client=mock_openai_client, @@ -937,8 +940,8 @@ class TestMCPToolGating: """use_prompt excluded when MCP client has no prompts.""" mcp_client = MagicMock() mcp_client.get_tools.return_value = [] - mcp_client.resource_count = 3 - mcp_client.prompt_count = 0 + mcp_client.resource_count_for_user.return_value = 3 + mcp_client.prompt_count_for_user.return_value = 0 session = ChatSession( client=mock_openai_client, @@ -960,8 +963,8 @@ class TestMCPToolGating: """Both tools present when MCP client has resources and prompts.""" mcp_client = MagicMock() mcp_client.get_tools.return_value = [] - mcp_client.resource_count = 1 - mcp_client.prompt_count = 1 + mcp_client.resource_count_for_user.return_value = 1 + mcp_client.prompt_count_for_user.return_value = 1 session = ChatSession( client=mock_openai_client, @@ -983,8 +986,8 @@ class TestMCPToolGating: """Gating applies even when tool_search is active (client-side path).""" mcp_client = MagicMock() mcp_client.get_tools.return_value = [] - mcp_client.resource_count = 0 - mcp_client.prompt_count = 0 + mcp_client.resource_count_for_user.return_value = 0 + mcp_client.prompt_count_for_user.return_value = 0 session = ChatSession( client=mock_openai_client, @@ -1012,8 +1015,8 @@ class TestMCPToolGating: mcp_client = MagicMock() mcp_client.get_tools.return_value = [] - mcp_client.resource_count = 0 - mcp_client.prompt_count = 0 + mcp_client.resource_count_for_user.return_value = 0 + mcp_client.prompt_count_for_user.return_value = 0 session = ChatSession( client=mock_openai_client, @@ -1034,3 +1037,42 @@ class TestMCPToolGating: names = [t.get("function", {}).get("name") for t in tools] assert "read_resource" not in names assert "use_prompt" not in names + + def test_pool_only_user_keeps_read_resource_and_use_prompt(self, tmp_db, mock_openai_client): + """Phase 7b canary: a pool-only user (static catalog empty) still + sees ``read_resource`` and ``use_prompt`` because the gating + consults ``*_count_for_user`` (scope decision 0.2). + + Drives ``resource_count = prompt_count = 0`` (the static-only + properties are zero) but ``*_count_for_user(uid) > 0`` because + the user has pool entries; the tools must remain visible. + """ + mcp_client = MagicMock() + mcp_client.get_tools.return_value = [] + # Static catalog is empty; admin-style legacy properties say 0. + mcp_client.resource_count = 0 + mcp_client.prompt_count = 0 + # Per-user variant reports the user's pool entries. + mcp_client.resource_count_for_user.return_value = 2 + mcp_client.prompt_count_for_user.return_value = 1 + + session = ChatSession( + client=mock_openai_client, + model="local-model", + ui=MagicMock(), + instructions=None, + temperature=0.5, + max_tokens=1000, + tool_timeout=10, + mcp_client=mcp_client, + user_id="pool-only-user", + ) + + tools = session._get_active_tools() + names = [t.get("function", {}).get("name") for t in tools] + assert "read_resource" in names + assert "use_prompt" in names + # Verify the per-user gate was actually consulted with the + # session's ``user_id`` (sanity-check on the wiring). + mcp_client.resource_count_for_user.assert_any_call("pool-only-user") + mcp_client.prompt_count_for_user.assert_any_call("pool-only-user") diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index e3b90dec..bdbcf62d 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -49,12 +49,12 @@ from turnstone.core.mcp_http_parsers import ( ) from turnstone.core.mcp_oauth import ( TokenLookupResult, - emit_insufficient_scope_audit, + emit_oauth_failure_audit, get_user_access_token_classified, ) if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Awaitable, Callable log = get_logger("turnstone.mcp") @@ -129,6 +129,18 @@ _MAX_INSUFFICIENT_SCOPE_REPORTED = 32 # emit a warning so operators can investigate. _MAX_TOOLS_PER_SERVER = 1000 +# Defensive caps mirroring ``_MAX_TOOLS_PER_SERVER`` for the resource and +# prompt list paths discovered in :meth:`MCPClientManager._connect_one_pool` +# (RFC §3.2 for resources, §3.3 for prompts). Real MCP servers expose at +# most a few dozen of each; a +# hostile or misconfigured upstream returning thousands would amplify +# memory (one dict per entry) plus downstream rendering cost. Truncate +# rather than reject so partial visibility beats zero visibility, and +# emit a warning so operators can investigate. +_MAX_RESOURCES_PER_SERVER = 1000 +_MAX_RESOURCE_TEMPLATES_PER_SERVER = 1000 +_MAX_PROMPTS_PER_SERVER = 1000 + @dataclass class _AuthCapture: @@ -276,6 +288,62 @@ def _cap_server_tools(server_name: str, tools: list[Any]) -> list[Any]: return tools[:_MAX_TOOLS_PER_SERVER] +def _cap_server_resources(server_name: str, resources: list[Any]) -> list[Any]: + """Apply ``_MAX_RESOURCES_PER_SERVER`` cap with operator-visible warning. + + Identity for inputs at or below the cap (no copy); slice + warn on + overflow. Mirrors :func:`_cap_server_tools` for the resource list path. + """ + if len(resources) <= _MAX_RESOURCES_PER_SERVER: + return resources + log.warning( + "MCP server '%s' returned %d resources — truncating to %d " + "(_MAX_RESOURCES_PER_SERVER cap). Misconfigured or hostile upstream?", + server_name, + len(resources), + _MAX_RESOURCES_PER_SERVER, + ) + return resources[:_MAX_RESOURCES_PER_SERVER] + + +def _cap_server_resource_templates(server_name: str, templates: list[Any]) -> list[Any]: + """Apply ``_MAX_RESOURCE_TEMPLATES_PER_SERVER`` cap with operator-visible warning. + + Identity for inputs at or below the cap (no copy); slice + warn on + overflow. Mirrors :func:`_cap_server_resources` for the resource + template list path (RFC §3.2 templates are a separate catalog from + concrete resources but share the per-server amplification risk). + """ + if len(templates) <= _MAX_RESOURCE_TEMPLATES_PER_SERVER: + return templates + log.warning( + "MCP server '%s' returned %d resource templates — truncating to %d " + "(_MAX_RESOURCE_TEMPLATES_PER_SERVER cap). Misconfigured or hostile upstream?", + server_name, + len(templates), + _MAX_RESOURCE_TEMPLATES_PER_SERVER, + ) + return templates[:_MAX_RESOURCE_TEMPLATES_PER_SERVER] + + +def _cap_server_prompts(server_name: str, prompts: list[Any]) -> list[Any]: + """Apply ``_MAX_PROMPTS_PER_SERVER`` cap with operator-visible warning. + + Identity for inputs at or below the cap (no copy); slice + warn on + overflow. Mirrors :func:`_cap_server_tools` for the prompt list path. + """ + if len(prompts) <= _MAX_PROMPTS_PER_SERVER: + return prompts + log.warning( + "MCP server '%s' returned %d prompts — truncating to %d " + "(_MAX_PROMPTS_PER_SERVER cap). Misconfigured or hostile upstream?", + server_name, + len(prompts), + _MAX_PROMPTS_PER_SERVER, + ) + return prompts[:_MAX_PROMPTS_PER_SERVER] + + # --------------------------------------------------------------------------- # Per-server state containers # --------------------------------------------------------------------------- @@ -348,6 +416,17 @@ class PoolEntryState: # loop-binding contract); ``_ensure_pool_entry`` runs on the loop # so the dataclass default_factory is safe. auth_fired_event: asyncio.Event = field(default_factory=asyncio.Event) + # Capability flags mirror ``StaticServerState`` so the pool path can + # gate resource / prompt discovery the same way (RFC §3.2). Tools + # are always discovered, so no ``supports_tools`` flag is needed — + # but ``listChanged`` for tools is also implicit (the notification + # handler is always registered). Resources and prompts get explicit + # presence flags because we skip discovery entirely when the + # capability is absent. + supports_resources: bool = False + supports_prompts: bool = False + supports_resource_list_changed: bool = False + supports_prompt_list_changed: bool = False # --------------------------------------------------------------------------- @@ -406,13 +485,19 @@ class MCPClientManager: # Merged resource catalog self._resources: list[dict[str, Any]] = [] self._resource_map: dict[str, tuple[str, str]] = {} # uri → (server, uri) - self._resource_listeners: list[Callable[[], None]] = [] + # Phase 7b: each entry is ``(user_id, callback)`` mirroring the + # tool-listener shape (RFC §3.3) — ``user_id=None`` is the admin + # / global listener (fires on every change), a string ``user_id`` + # fires only on changes scoped to that user OR on global static + # changes. + self._resource_listeners: list[tuple[str | None, Callable[[], None]]] = [] self._resource_listeners_lock = threading.Lock() # Merged prompt catalog self._prompts: list[dict[str, Any]] = [] self._prompt_map: dict[str, tuple[str, str]] = {} # prefixed → (server, original) - self._prompt_listeners: list[Callable[[], None]] = [] + # Phase 7b: see ``_resource_listeners`` for the tuple shape rationale. + self._prompt_listeners: list[tuple[str | None, Callable[[], None]]] = [] self._prompt_listeners_lock = threading.Lock() # Template prefix → (server_name, full_template_uri) for URI expansion @@ -447,8 +532,7 @@ class MCPClientManager: # 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. + # Phase 7b lights up the resource/prompt sibling dicts below. 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 @@ -460,6 +544,26 @@ class MCPClientManager: # ``_close_pool_entry_if_idle`` / ``_evict_session``). self._user_tools: dict[str, list[dict[str, Any]]] = {} + # Per-user resource catalog. Mirrors ``_user_tool_map`` / + # ``_user_tools``: outer key is ``user_id``, inner is + # ``uri → (server_name, uri)``. Loop-only writes via + # :meth:`_rebuild_user_resource_map`; sync-thread reads via a + # single dict-get on ``_user_resources`` (atomic under GIL). + # RFC §3.2 — Phase 7b lights this up. + self._user_resource_map: dict[str, dict[str, tuple[str, str]]] = {} + self._user_resources: dict[str, list[dict[str, Any]]] = {} + # Per-user template-prefix index for resource URI expansion. + # ``_user_template_prefixes[user_id][prefix] = (server, full_template_uri)``. + # Used by :meth:`_match_template`'s per-user-first lookup branch + # (scope decision 0.1). + self._user_template_prefixes: dict[str, dict[str, tuple[str, str]]] = {} + + # Per-user prompt catalog. Mirrors ``_user_tool_map`` / + # ``_user_tools`` for prompts: outer key is ``user_id``, inner + # is ``prefixed_name → (server_name, original_name)``. + self._user_prompt_map: dict[str, dict[str, tuple[str, str]]] = {} + self._user_prompts: 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 @@ -1041,12 +1145,13 @@ class MCPClientManager: * Streamable-HTTP transport only — pool servers are remote. * ``Authorization: Bearer {access_token}`` injected into headers alongside any operator-supplied static headers. - * 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. + * Tool / resource / prompt catalog discovery runs after + ``initialize()`` (RFC §3.2 — discovery covers all three + catalogs, capability-gated). The notification handler is + bound to ``(user_id, server_name)`` so push-driven + ``*/list_changed`` updates only refresh the owning user's + catalog (the static refreshers must NEVER fire from a pool + session — they would clobber static-path state). When ``auth_capture`` is supplied, the underlying ``httpx`` client is built via a factory whose response hook records 401/403 @@ -1141,9 +1246,10 @@ class MCPClientManager: # 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``. + # other sessions. Phase 7 wired tool list-changed; Phase 7b + # extends the same shape to resources and prompts via + # ``_refresh_pool_server_resources`` / + # ``_refresh_pool_server_prompts``. async def _on_pool_notification( msg: Any, # RequestResponder | ServerNotification | Exception ) -> None: @@ -1170,17 +1276,21 @@ class MCPClientManager: 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", + log.info( + "Received resources/list_changed from pool user=%s server=%s", user_id, server_name, ) + self._last_pool_notification_refresh[key] = now + await self._refresh_pool_server_resources(key) elif isinstance(root, mcp_types.PromptListChangedNotification): - log.debug( - "pool prompts/list_changed (deferred to Phase 7b) user=%s server=%s", + log.info( + "Received prompts/list_changed from pool user=%s server=%s", user_id, server_name, ) + self._last_pool_notification_refresh[key] = now + await self._refresh_pool_server_prompts(key) except Exception as exc: # Structured fields only — ``exc_info=True`` would # serialize the chained ``httpx.Request`` whose headers @@ -1231,6 +1341,19 @@ class MCPClientManager: await self._safe_teardown_on_connect_failure(key, stack) raise + # Capability fetch — mirror the static path at mcp_client.py:920-932. + # Per RFC §3.2 the pool path now discovers resources & prompts too, + # capability-gated so a server that doesn't implement them stays + # cheap (no extra round-trips). Capabilities are populated by the + # initialize roundtrip and immutable thereafter (R13). + caps = session.get_server_capabilities() + resources_cap = getattr(caps, "resources", None) if caps else None + prompts_cap = getattr(caps, "prompts", None) if caps else None + entry.supports_resources = resources_cap is not None + entry.supports_resource_list_changed = bool(getattr(resources_cap, "listChanged", False)) + entry.supports_prompts = prompts_cap is not None + entry.supports_prompt_list_changed = bool(getattr(prompts_cap, "listChanged", False)) + # Discover this user's tool catalog. R6 verified: a 401 here # propagates through anyio TaskGroup unwinding (raises an # ``ExceptionGroup`` from the surrounding ``streamablehttp_client`` @@ -1238,8 +1361,7 @@ class MCPClientManager: # 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. + # doesn't apply to first-connect discovery. # # Why ``asyncio.timeout``, not ``asyncio.wait_for``: per # ``feedback_asyncio_timeout_vs_wait_for.md`` and the f6a3b66 @@ -1250,7 +1372,10 @@ class MCPClientManager: # ``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. + # traverse anyio cleanup. R6 (Phase 7b) verified the same hazard + # applies to ``list_resources`` / ``list_resource_templates`` / + # ``list_prompts`` — they all traverse the same + # ``mcp/shared/session.py:240-314`` send_request infrastructure. try: async with asyncio.timeout(self._CONNECT_TIMEOUT): tools_result = await session.list_tools() @@ -1273,6 +1398,108 @@ class MCPClientManager: capped_tools = _cap_server_tools(server_name, tools_result.tools) entry.tools = [_mcp_to_openai(server_name, tool) for tool in capped_tools] + + # Phase 7b — discover resources (capability-gated). Same anyio / + # ``asyncio.timeout`` invariant as the tool discovery above (R1). + server_resources: list[dict[str, Any]] = [] + if resources_cap is not None: + try: + async with asyncio.timeout(self._CONNECT_TIMEOUT): + # 1-RTT (gather) instead of 2 sequential RTTs — both + # calls share the same timeout budget and target + # disjoint catalogs (resources vs. templates), so + # ordering is irrelevant. + res_result, tmpl_result = await asyncio.gather( + session.list_resources(), + session.list_resource_templates(), + ) + 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 resource 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 resource 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 + + for r in _cap_server_resources(server_name, res_result.resources): + server_resources.append( + { + "uri": str(r.uri), + "name": r.name or "", + "description": r.description or "", + "mimeType": r.mimeType or "", + "server": server_name, + } + ) + for t in _cap_server_resource_templates(server_name, tmpl_result.resourceTemplates): + server_resources.append( + { + "uri": str(t.uriTemplate), + "name": t.name or "", + "description": t.description or "", + "mimeType": t.mimeType or "", + "server": server_name, + "template": True, + } + ) + + # Phase 7b — discover prompts (capability-gated). + server_prompts: list[dict[str, Any]] = [] + if prompts_cap is not None: + try: + async with asyncio.timeout(self._CONNECT_TIMEOUT): + prompt_result = await session.list_prompts() + 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 prompt 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 prompt 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 + + for p in _cap_server_prompts(server_name, prompt_result.prompts): + server_prompts.append( + { + "name": f"mcp__{server_name}__{p.name}", + "original_name": p.name, + "server": server_name, + "description": p.description or "", + "arguments": [ + { + "name": a.name, + "description": a.description or "", + "required": a.required or False, + } + for a in (p.arguments or []) + ], + } + ) + + entry.resources = server_resources if resources_cap is not None else None + entry.prompts = server_prompts if prompts_cap is not None else None + # 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 @@ -1282,18 +1509,24 @@ class MCPClientManager: # 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. + # recovery path. Same ordering invariant covers resources/prompts + # (R16). 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``. + # Loop-only mutation; sync-thread readers observe the new + # catalog atomically via per-user dict-gets on ``_user_tools`` / + # ``_user_resources`` / ``_user_prompts``. self._rebuild_user_tool_map(user_id) + self._rebuild_user_resource_map(user_id) + self._rebuild_user_prompt_map(user_id) # Wake user-keyed AND admin (None) listeners; per-user fan-out # ensures another user's session never observes this user's - # tool change. + # catalog change. self._notify_user_tool_listeners(user_id) + self._notify_user_resource_listeners(user_id) + self._notify_user_prompt_listeners(user_id) return entry # -- pool eviction -------------------------------------------------------- @@ -1402,14 +1635,19 @@ class MCPClientManager: 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. + # entry without rebuilding the per-user catalogs would leave + # ``is_mcp_tool`` / per-user resource & prompt maps returning + # stale entries whose backing pool is gone, and ChatSession's + # tool / resource / prompt lists 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._rebuild_user_resource_map(user_id) + self._rebuild_user_prompt_map(user_id) self._notify_user_tool_listeners(user_id) + self._notify_user_resource_listeners(user_id) + self._notify_user_prompt_listeners(user_id) evicted = True finally: lock.release() @@ -1535,6 +1773,118 @@ class MCPClientManager: self._user_tool_map.pop(user_id, None) self._user_tools.pop(user_id, None) + def _rebuild_user_resource_map(self, user_id: str) -> None: + """Rebuild the per-user resource index from pool entries. + + Mirrors :meth:`_rebuild_user_tool_map` for resources (RFC §3.2): + scan ``_user_pool_entries`` for keys whose first element matches + ``user_id``, materialize a fresh ``uri → (server, uri)`` dict, a + parallel resource list, AND a parallel template-prefix dict, + and assign each atomically. + + URI collisions across pool servers within a single user's pool + follow the same "later wins, log warning" policy as the static + :meth:`_rebuild_resources` (mcp_client.py:1689-1744). Empty + rebuilds drop the user_id key from all three dicts so idle + users don't retain empty-list sentinels. + + ``_resource_map`` / ``_template_prefixes`` (the static indices) + are NEVER mutated here, preserving invariant 1. + + MUST run on the mcp-loop. The two writes (``_user_resource_map``, + ``_user_resources``, ``_user_template_prefixes``) happen + back-to-back with no awaits between them so a sync-thread reader + cannot observe a torn state. + """ + new_map: dict[str, tuple[str, str]] = {} + new_resources: list[dict[str, Any]] = [] + new_prefixes: dict[str, tuple[str, str]] = {} + for (uid, server_name), entry in self._user_pool_entries.items(): + if uid != user_id or entry.resources is None: + continue + for res in entry.resources: + if res.get("template"): + tmpl_uri: str = res["uri"] + brace = tmpl_uri.find("{") + prefix = tmpl_uri[:brace] if brace >= 0 else tmpl_uri + if prefix: + if prefix in new_prefixes: + existing_srv, existing_tmpl = new_prefixes[prefix] + if len(tmpl_uri) > len(existing_tmpl): + log.warning( + "Pool template prefix collision: '%s' from '%s' overrides " + "'%s' (keeping more specific template) user=%s", + prefix, + server_name, + existing_srv, + user_id, + ) + new_prefixes[prefix] = (server_name, tmpl_uri) + else: + log.warning( + "Pool template prefix collision: '%s' from '%s' ignored in " + "favor of '%s' (keeping more specific template) user=%s", + prefix, + server_name, + existing_srv, + user_id, + ) + else: + new_prefixes[prefix] = (server_name, tmpl_uri) + new_resources.append(res) + continue + uri: str = res["uri"] + if uri in new_map: + log.warning( + "Pool resource URI collision: '%s' from '%s' overrides '%s' user=%s", + uri, + server_name, + new_map[uri][0], + user_id, + ) + new_map[uri] = (server_name, uri) + new_resources.append(res) + if new_map or new_resources or new_prefixes: + self._user_resource_map[user_id] = new_map + self._user_resources[user_id] = new_resources + self._user_template_prefixes[user_id] = new_prefixes + else: + self._user_resource_map.pop(user_id, None) + self._user_resources.pop(user_id, None) + self._user_template_prefixes.pop(user_id, None) + + def _rebuild_user_prompt_map(self, user_id: str) -> None: + """Rebuild the per-user prompt index from pool entries. + + Mirrors :meth:`_rebuild_user_tool_map` for prompts (RFC §3.3): + scan ``_user_pool_entries`` for keys whose first element matches + ``user_id``, materialize a fresh + ``prefixed_name → (server, original_name)`` dict and a parallel + prompt list, and assign each atomically. + + ``_prompt_map`` (the static index) is NEVER mutated here, + preserving invariant 1. + + Empty rebuilds drop the user_id key from BOTH dicts. + + MUST run on the mcp-loop. + """ + new_map: dict[str, tuple[str, str]] = {} + new_prompts: list[dict[str, Any]] = [] + for (uid, _server_name), entry in self._user_pool_entries.items(): + if uid != user_id or entry.prompts is None: + continue + for prompt in entry.prompts: + prefixed: str = prompt["name"] + new_map[prefixed] = (prompt["server"], prompt["original_name"]) + new_prompts.append(prompt) + if new_map: + self._user_prompt_map[user_id] = new_map + self._user_prompts[user_id] = new_prompts + else: + self._user_prompt_map.pop(user_id, None) + self._user_prompts.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) @@ -1611,6 +1961,138 @@ class MCPClientManager: ) return added, removed + async def _refresh_pool_server_resources( + self, key: tuple[str, str] + ) -> tuple[list[str], list[str]]: + """Re-fetch resources for one pool entry. Returns ``(added, removed)`` URIs. + + Mirror of :meth:`_refresh_pool_server_tools` for the resource + path (RFC §3.2). Skips when ``supports_resources`` is False so + a server that no longer advertises resources doesn't trigger + a list call. ``asyncio.timeout`` is mandatory — see R6. + + MUST run on the mcp-loop. Caller does not need to hold + ``open_lock``; this is invoked from the pool notification + handler running on the SDK's receive task. + """ + entry = self._user_pool_entries.get(key) + if entry is None or entry.session is None or not entry.supports_resources: + return [], [] + session = entry.session + user_id, server_name = key + old_uris = {r["uri"] for r in (entry.resources or []) if not r.get("template")} + + async with asyncio.timeout(self._CONNECT_TIMEOUT): + # 1-RTT (gather) instead of 2 sequential RTTs — both calls + # share the same timeout budget and target disjoint catalogs + # (resources vs. templates), so ordering is irrelevant. + res_result, tmpl_result = await asyncio.gather( + session.list_resources(), + session.list_resource_templates(), + ) + + server_resources: list[dict[str, Any]] = [] + for r in _cap_server_resources(server_name, res_result.resources): + server_resources.append( + { + "uri": str(r.uri), + "name": r.name or "", + "description": r.description or "", + "mimeType": r.mimeType or "", + "server": server_name, + } + ) + for t in _cap_server_resource_templates(server_name, tmpl_result.resourceTemplates): + server_resources.append( + { + "uri": str(t.uriTemplate), + "name": t.name or "", + "description": t.description or "", + "mimeType": t.mimeType or "", + "server": server_name, + "template": True, + } + ) + + new_uris = {r["uri"] for r in server_resources if not r.get("template")} + entry.resources = server_resources + self._rebuild_user_resource_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_resource_listeners(user_id) + + added = sorted(new_uris - old_uris) + removed = sorted(old_uris - new_uris) + if added or removed: + log.info( + "Refreshed pool MCP server user=%s server=%s: +%d/-%d resource(s)", + user_id, + server_name, + len(added), + len(removed), + ) + return added, removed + + async def _refresh_pool_server_prompts( + self, key: tuple[str, str] + ) -> tuple[list[str], list[str]]: + """Re-fetch prompts for one pool entry. Returns ``(added, removed)`` names. + + Mirror of :meth:`_refresh_pool_server_tools` for the prompt + path (RFC §3.3). Skips when ``supports_prompts`` is False so a + server that no longer advertises prompts doesn't trigger a list + call. ``asyncio.timeout`` is mandatory — see R6. + + MUST run on the mcp-loop. Caller does not need to hold + ``open_lock``; this is invoked from the pool notification + handler running on the SDK's receive task. + """ + entry = self._user_pool_entries.get(key) + if entry is None or entry.session is None or not entry.supports_prompts: + return [], [] + session = entry.session + user_id, server_name = key + old_names = {p["name"] for p in (entry.prompts or [])} + + async with asyncio.timeout(self._CONNECT_TIMEOUT): + prompt_result = await session.list_prompts() + + server_prompts: list[dict[str, Any]] = [] + for p in _cap_server_prompts(server_name, prompt_result.prompts): + server_prompts.append( + { + "name": f"mcp__{server_name}__{p.name}", + "original_name": p.name, + "server": server_name, + "description": p.description or "", + "arguments": [ + { + "name": a.name, + "description": a.description or "", + "required": a.required or False, + } + for a in (p.arguments or []) + ], + } + ) + + new_names = {p["name"] for p in server_prompts} + entry.prompts = server_prompts + self._rebuild_user_prompt_map(user_id) + self._notify_user_prompt_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 prompt(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. @@ -1896,40 +2378,110 @@ class MCPClientManager: except Exception: log.warning("Tool-change listener raised", exc_info=True) - def add_resource_listener(self, callback: Callable[[], None]) -> None: - """Register a callback invoked when the resource list changes.""" - with self._resource_listeners_lock: - self._resource_listeners.append(callback) + def add_resource_listener( + self, callback: Callable[[], None], *, user_id: str | None = None + ) -> None: + """Register a callback invoked when the resource list changes. - def remove_resource_listener(self, callback: Callable[[], None]) -> None: - """Unregister a resource-change callback.""" + ``user_id=None`` (default) registers a global / admin listener + that fires on every resource-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._resource_listeners_lock: + self._resource_listeners.append((user_id, callback)) + + def remove_resource_listener( + self, callback: Callable[[], None], *, user_id: str | None = None + ) -> None: + """Unregister a resource-change callback. + + ``user_id`` MUST match the value used at registration; the + ``(user_id, callback)`` pair is the listener identity. + """ with self._resource_listeners_lock, contextlib.suppress(ValueError): - self._resource_listeners.remove(callback) + self._resource_listeners.remove((user_id, callback)) def _notify_resource_listeners(self) -> None: - """Invoke all registered resource-change listeners.""" + """Static-path resource change — fires ALL registered listeners. + + Mirrors ``_notify_listeners`` for tools (RFC §3.3): the static + catalog is process-wide so the fan-out is unconditional. + """ with self._resource_listeners_lock: listeners = list(self._resource_listeners) + for _uid, cb in listeners: + try: + cb() + except Exception: + log.warning("Resource-change listener raised", exc_info=True) + + def _notify_user_resource_listeners(self, user_id: str) -> None: + """Pool-entry resource change — fires only listeners that should see it. + + Scoped fan-out targets the matching ``user_id`` AND admin + (``None``) listeners. Mirrors ``_notify_user_tool_listeners``. + RFC §3.3. + """ + with self._resource_listeners_lock: + listeners = [ + cb for uid, cb in self._resource_listeners if uid == user_id or uid is None + ] for cb in listeners: try: cb() except Exception: log.warning("Resource-change listener raised", exc_info=True) - def add_prompt_listener(self, callback: Callable[[], None]) -> None: - """Register a callback invoked when the prompt list changes.""" - with self._prompt_listeners_lock: - self._prompt_listeners.append(callback) + def add_prompt_listener( + self, callback: Callable[[], None], *, user_id: str | None = None + ) -> None: + """Register a callback invoked when the prompt list changes. - def remove_prompt_listener(self, callback: Callable[[], None]) -> None: - """Unregister a prompt-change callback.""" + ``user_id=None`` (default) registers a global / admin listener + that fires on every prompt-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._prompt_listeners_lock: + self._prompt_listeners.append((user_id, callback)) + + def remove_prompt_listener( + self, callback: Callable[[], None], *, user_id: str | None = None + ) -> None: + """Unregister a prompt-change callback. + + ``user_id`` MUST match the value used at registration; the + ``(user_id, callback)`` pair is the listener identity. + """ with self._prompt_listeners_lock, contextlib.suppress(ValueError): - self._prompt_listeners.remove(callback) + self._prompt_listeners.remove((user_id, callback)) def _notify_prompt_listeners(self) -> None: - """Invoke all registered prompt-change listeners.""" + """Static-path prompt change — fires ALL registered listeners. + + Mirrors ``_notify_listeners`` for tools (RFC §3.3): the static + catalog is process-wide so the fan-out is unconditional. + """ with self._prompt_listeners_lock: listeners = list(self._prompt_listeners) + for _uid, cb in listeners: + try: + cb() + except Exception: + log.warning("Prompt-change listener raised", exc_info=True) + + def _notify_user_prompt_listeners(self, user_id: str) -> None: + """Pool-entry prompt change — fires only listeners that should see it. + + Scoped fan-out targets the matching ``user_id`` AND admin + (``None``) listeners. Mirrors ``_notify_user_tool_listeners``. + RFC §3.3. + """ + with self._prompt_listeners_lock: + listeners = [cb for uid, cb in self._prompt_listeners if uid == user_id or uid is None] for cb in listeners: try: cb() @@ -2140,6 +2692,14 @@ class MCPClientManager: self._user_pool_entries.clear() self._user_pool_last_used.clear() self._user_pool_locks.clear() + self._user_tool_map.clear() + self._user_tools.clear() + # Phase 7b — per-user resource / prompt catalogs. + self._user_resource_map.clear() + self._user_resources.clear() + self._user_template_prefixes.clear() + self._user_prompt_map.clear() + self._user_prompts.clear() self._user_pool_eviction_task = None log.info("MCP client shut down") @@ -2489,24 +3049,84 @@ class MCPClientManager: 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).""" - return [dict(r) for r in self._resources] + def get_resources(self, user_id: str | None = None) -> list[dict[str, Any]]: + """Return discovered MCP resources (shallow-copied dicts). - def get_prompts(self) -> list[dict[str, Any]]: - """Return discovered MCP prompts (shallow-copied dicts).""" - return [dict(p) for p in self._prompts] + ``user_id=None`` returns the global static-path catalog only — + the legacy behaviour every pre-Phase-7b caller relies on. A + string ``user_id`` returns the merged view: per-user pool + resources first (per scope decision 0.1 — per-user-first + ordering), then the static catalog. Reads ``_user_resources`` + via a single dict-get (atomic under GIL). + + Returned dicts are shallow-copied; nested objects are shared + with the manager's catalog, mirroring the pre-Phase-7b contract. + """ + if user_id is None: + return [dict(r) for r in self._resources] + merged: list[dict[str, Any]] = [dict(r) for r in self._user_resources.get(user_id, [])] + merged.extend(dict(r) for r in self._resources) + return merged + + def get_prompts(self, user_id: str | None = None) -> list[dict[str, Any]]: + """Return discovered MCP prompts (shallow-copied dicts). + + ``user_id=None`` returns the global static-path catalog only. + A string ``user_id`` returns the merged view: per-user pool + prompts first (scope decision 0.1), then static. + """ + if user_id is None: + return [dict(p) for p in self._prompts] + merged: list[dict[str, Any]] = [dict(p) for p in self._user_prompts.get(user_id, [])] + merged.extend(dict(p) for p in self._prompts) + return merged @property def resource_count(self) -> int: - """Number of discovered resources (no allocation).""" + """Number of discovered static resources (no allocation). + + Process-global by design — admin endpoints rely on this contract. + Use :meth:`resource_count_for_user` for the per-user count that + includes the user's pool resources. + """ return len(self._resources) @property def prompt_count(self) -> int: - """Number of discovered prompts (no allocation).""" + """Number of discovered static prompts (no allocation). + + Process-global by design — admin endpoints rely on this contract. + Use :meth:`prompt_count_for_user` for the per-user count that + includes the user's pool prompts. + """ return len(self._prompts) + def resource_count_for_user(self, user_id: str | None = None) -> int: + """Return ``len(static) + len(user_pool)`` resources. + + ``user_id=None`` returns the static-only count (matches the + ``resource_count`` property). A string ``user_id`` adds that + user's pool resources. Used by ChatSession's ``read_resource`` + tool gating so a pool-only user still sees the tool when + their pool has resources but the static catalog is empty + (scope decision 0.2). + """ + base = len(self._resources) + if user_id is None: + return base + return base + len(self._user_resources.get(user_id, [])) + + def prompt_count_for_user(self, user_id: str | None = None) -> int: + """Return ``len(static) + len(user_pool)`` prompts. + + See :meth:`resource_count_for_user` — same shape, prompts + instead of resources. + """ + base = len(self._prompts) + if user_id is None: + return base + return base + len(self._user_prompts.get(user_id, [])) + def is_mcp_tool(self, func_name: str, *, user_id: str | None = None) -> bool: """Check whether *func_name* belongs to an MCP server. @@ -2530,9 +3150,23 @@ class MCPClientManager: 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 is_mcp_prompt(self, name: str, *, user_id: str | None = None) -> bool: + """Check whether *name* is a known MCP prompt. + + ``user_id=None`` (default) asks "is this a static-path prompt?" — + the answer is process-global. A string ``user_id`` extends the + lookup to that user's pool catalog: returns ``True`` if ``name`` + is in either the static map OR the user's per-user prompt map. + Mirrors :meth:`is_mcp_tool` (scope decision 0.1, per-user-first + is moot for prompts because their prefixed names are + per-server-disjoint by construction). + """ + if name in self._prompt_map: + return True + if user_id is None: + return False + user_map = self._user_prompt_map.get(user_id) + return user_map is not None and name in user_map def server_auth_type(self, server_name: str) -> str | None: """Return ``'oauth_user'`` for pool-backed servers, else ``None``. @@ -2770,6 +3404,69 @@ class MCPClientManager: return None return row + def _resolve_pool_target_resource( + self, uri: str, user_id: str + ) -> tuple[str, str, dict[str, Any]] | None: + """Resolve ``(server_name, uri, server_row)`` for pool resource reads. + + Per scope decision 0.1: per-user-first resolution. Consults the + user's ``_user_resource_map`` first, then the user's + ``_user_template_prefixes`` (longest-prefix match). Returns + ``None`` when the URI doesn't match a pool entry — the caller + falls through to the static path. + + Pool eligibility is gated on ``server_row.auth_type == + 'oauth_user'``: a per-user map entry that points at a static- + path server is treated as a miss (the static path will pick it + up). Defence-in-depth — _user_resource_map is populated only + from pool entries, so this case shouldn't occur in practice. + """ + user_map = self._user_resource_map.get(user_id) or {} + mapping = user_map.get(uri) + if mapping is None: + # Try per-user templates (longest-prefix wins within scope). + best: tuple[str, str] | None = None + best_len = 0 + for prefix, prefix_mapping in (self._user_template_prefixes.get(user_id) or {}).items(): + if uri.startswith(prefix) and len(prefix) > best_len: + best = prefix_mapping + best_len = len(prefix) + if best is None: + return None + server_name = best[0] + else: + server_name = mapping[0] + row = self._lookup_server_row(server_name) + if row is None or row.get("auth_type") != "oauth_user": + return None + return server_name, uri, row + + def _resolve_pool_target_prompt( + self, + prefixed_name: str, + static_server: str | None, + static_original: str | None, + ) -> tuple[str, str, dict[str, Any]] | None: + """Resolve ``(server_name, original_name, server_row)`` for pool prompts. + + Mirror of :meth:`_resolve_pool_target` for prompts: prompts + carry the ``mcp__{server}__{name}`` prefix shape so the same + parsing applies. Presence in static ``_prompt_map`` (signalled + by non-None ``static_server`` / ``static_original``) short- + circuits to None — the static path owns that name. + """ + if static_server is not None and static_original is not None: + return None + if not prefixed_name.startswith("mcp__") or prefixed_name.count("__") < 2: + return None + _, server_name, original = prefixed_name.split("__", 2) + if not server_name or not original: + return None + row = self._lookup_server_row(server_name) + if row is None or row.get("auth_type") != "oauth_user": + return None + return server_name, original, row + def _dispatch_pool_sync( self, *, @@ -2883,6 +3580,174 @@ class MCPClientManager: self._cb_record_failure(server_name) raise TimeoutError(f"MCP tool call timed out after {original_timeout}s") from None + def _dispatch_pool_resource_sync( + self, + *, + user_id: str, + server_name: str, + uri: str, + server_row: dict[str, Any], + timeout: int, + ) -> str: + """Synchronous wrapper for pool resource read. + + Mirrors :meth:`_dispatch_pool_sync` for the resource path. + Returns either the resource body or a structured-error JSON + string when token state precludes the call. + + Retry-on-401 follows the same fresh-task scheduling pattern + as the tool path — the dispatcher coroutine raises + :class:`_PoolDispatchRetryRequested` after refreshing the + bearer; we re-issue on a brand-new task so the retry's anyio + cancel-scope state is independent of the prior connect's + TaskGroup teardown. + """ + assert self._loop is not None + start = time.monotonic() + try: + return self._run_pool_dispatch_resource_attempt( + retry_count=0, + timeout=timeout, + original_timeout=timeout, + user_id=user_id, + server_name=server_name, + uri=uri, + server_row=server_row, + ) + except _PoolDispatchRetryRequested: + remaining = max(1, int(timeout - (time.monotonic() - start))) + return self._run_pool_dispatch_resource_attempt( + retry_count=1, + timeout=remaining, + original_timeout=timeout, + user_id=user_id, + server_name=server_name, + uri=uri, + server_row=server_row, + ) + + def _run_pool_dispatch_resource_attempt( + self, + *, + retry_count: int, + timeout: int, + original_timeout: int, + user_id: str, + server_name: str, + uri: str, + server_row: dict[str, Any], + ) -> str: + """Schedule one resource dispatch attempt; same shape as + :meth:`_run_pool_dispatch_attempt` for tools.""" + assert self._loop is not None + future = asyncio.run_coroutine_threadsafe( + self._dispatch_pool_resource( + retry_count=retry_count, + user_id=user_id, + server_name=server_name, + uri=uri, + server_row=server_row, + ), + self._loop, + ) + try: + return future.result(timeout=timeout) + except concurrent.futures.TimeoutError: + future.cancel() + self._cb_record_failure(server_name) + raise TimeoutError(f"MCP resource read timed out after {original_timeout}s") from None + + def _dispatch_pool_prompt_sync( + self, + *, + user_id: str, + server_name: str, + original_name: str, + arguments: dict[str, str] | None, + server_row: dict[str, Any], + timeout: int, + ) -> list[dict[str, Any]]: + """Synchronous wrapper for pool prompt invocation. + + Mirrors :meth:`_dispatch_pool_sync` for the prompt path. The + async dispatcher returns either a list of expanded messages OR + a structured-error JSON string (the failure path); to keep the + return shape consistent with the static path + (:meth:`get_prompt_sync`), errors are surfaced via + :class:`RuntimeError` so the agent-loop's existing + ``except Exception`` block at the call site renders the + structured-error message. Embedding the JSON in a single-element + message list would pollute the prompt protocol — open question 1 + in the plan. + """ + assert self._loop is not None + start = time.monotonic() + try: + result = self._run_pool_dispatch_prompt_attempt( + retry_count=0, + timeout=timeout, + original_timeout=timeout, + user_id=user_id, + server_name=server_name, + original_name=original_name, + arguments=arguments, + server_row=server_row, + ) + except _PoolDispatchRetryRequested: + remaining = max(1, int(timeout - (time.monotonic() - start))) + result = self._run_pool_dispatch_prompt_attempt( + retry_count=1, + timeout=remaining, + original_timeout=timeout, + user_id=user_id, + server_name=server_name, + original_name=original_name, + arguments=arguments, + server_row=server_row, + ) + if isinstance(result, str): + # Structured-error path — surface as RuntimeError so the + # agent-loop renders the JSON via its except-Exception + # handler. + raise RuntimeError(result) + return result + + def _run_pool_dispatch_prompt_attempt( + self, + *, + retry_count: int, + timeout: int, + original_timeout: int, + user_id: str, + server_name: str, + original_name: str, + arguments: dict[str, str] | None, + server_row: dict[str, Any], + ) -> list[dict[str, Any]] | str: + """Schedule one prompt dispatch attempt; returns either decoded + messages or a structured-error string. Mirror of + :meth:`_run_pool_dispatch_attempt` for prompts.""" + assert self._loop is not None + future = asyncio.run_coroutine_threadsafe( + self._dispatch_pool_prompt( + retry_count=retry_count, + user_id=user_id, + server_name=server_name, + original_name=original_name, + arguments=arguments, + server_row=server_row, + ), + self._loop, + ) + try: + return future.result(timeout=timeout) + except concurrent.futures.TimeoutError: + future.cancel() + self._cb_record_failure(server_name) + raise TimeoutError( + f"MCP prompt retrieval timed out after {original_timeout}s" + ) from None + async def _dispatch_pool( self, *, @@ -3080,6 +3945,290 @@ class MCPClientManager: self._cb_record_success(server_name) return result + async def _dispatch_pool_resource( + self, + *, + user_id: str, + server_name: str, + uri: str, + server_row: dict[str, Any], + retry_count: int = 0, + ) -> str: + """Pool-side coroutine for resource reads (RFC §3.2). + + Mirror of :meth:`_dispatch_pool` for the resource path. Returns + either the textualized resource content or a structured-error + JSON string. Reuses the same token-classify, breaker-gate, + URL-hygiene, carrier-race, classification and retry plumbing + as the tool dispatcher; only the SDK call differs. + """ + if self._app_state is None: + raise RuntimeError("Pool dispatch requires set_app_state() to have been called") + + lookup: TokenLookupResult = await get_user_access_token_classified( + app_state=self._app_state, + user_id=user_id, + server_name=server_name, + force_refresh=retry_count > 0, + ) + if lookup.kind == "missing": + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="No token for user. Consent flow required.", + ) + if lookup.kind == "decrypt_failure": + return _structured_error( + code="mcp_token_undecryptable_key_unknown", + server=server_name, + detail=( + "Stored token cannot be decrypted by any installed encryption key. " + "Operator action required." + ), + ) + if lookup.kind == "refresh_failed": + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="Refresh token rejected. Re-consent required.", + ) + access_token = lookup.token or "" + if not access_token: + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="No token for user. Consent flow required.", + ) + + url = str(server_row.get("url") or "") + try: + _validate_oauth_user_url(url) + except ValueError as exc: + log.warning( + "mcp_pool.url_insecure server=%s scheme=%s", + server_name, + urllib.parse.urlparse(url).scheme, + ) + return _structured_error( + code="mcp_oauth_url_insecure", + server=server_name, + detail=str(exc), + ) + + self._cb_gate(server_name) + + cfg = _pool_cfg_from_row(server_row) + key = (user_id, server_name) + entry = await self._ensure_pool_entry(key) + capture = entry.auth_capture + try: + sdk_result = await self._dispatch_pool_with_entry_call( + entry=entry, + key=key, + cfg=cfg, + access_token=access_token, + sdk_call=lambda s: s.read_resource(uri), + ) + except BaseException as exc: + classification = self._classify_failure(exc, capture=capture) + if classification == "auth_401": + self._evict_session(key) + if retry_count == 0: + log.debug( + "mcp_pool.auth_401_initial_resource server=%s user=%s exc=%s", + server_name, + user_id, + type(exc).__name__, + ) + raise _PoolDispatchRetryRequested from None + log.debug( + "mcp_pool.auth_401_retry_failed_resource server=%s user=%s exc=%s", + server_name, + user_id, + type(exc).__name__, + ) + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="Refreshed token still rejected. Re-consent required.", + ) + if classification == "auth_403": + self._evict_session(key) + log.debug( + "mcp_pool.auth_403_resource 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, + server_row=server_row, + capture=capture, + kind="resource", + ) + if classification == "transport": + self._cb_record_failure(server_name) + self._evict_session(key) + log.debug( + "mcp_pool.transport_failure_resource server=%s user=%s exc=%s", + server_name, + user_id, + type(exc).__name__, + ) + raise + raise + + self._cb_record_success(server_name) + return _decode_resource_result(sdk_result) + + async def _dispatch_pool_prompt( + self, + *, + user_id: str, + server_name: str, + original_name: str, + arguments: dict[str, str] | None, + server_row: dict[str, Any], + retry_count: int = 0, + ) -> list[dict[str, Any]] | str: + """Pool-side coroutine for prompt invocation (RFC §3.3). + + Mirror of :meth:`_dispatch_pool_resource` for the prompt path. + Returns either the decoded list of expanded messages or a + structured-error JSON string. Reuses the same token-classify, + breaker-gate, URL-hygiene, carrier-race, classification and + retry plumbing as the tool / resource dispatchers; only the SDK + call differs. + + The dual return type (``list[dict[str, Any]] | str``) is shaped + for the sync wrapper at :meth:`_dispatch_pool_prompt_sync` — + it surfaces the structured-error string via + :class:`RuntimeError` so the agent-loop's ``except Exception`` + block at the call site renders the JSON. Embedding error JSON + in a single-element message list would pollute the prompt + protocol (open question 1, plan §5). + """ + if self._app_state is None: + raise RuntimeError("Pool dispatch requires set_app_state() to have been called") + + lookup: TokenLookupResult = await get_user_access_token_classified( + app_state=self._app_state, + user_id=user_id, + server_name=server_name, + force_refresh=retry_count > 0, + ) + if lookup.kind == "missing": + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="No token for user. Consent flow required.", + ) + if lookup.kind == "decrypt_failure": + return _structured_error( + code="mcp_token_undecryptable_key_unknown", + server=server_name, + detail=( + "Stored token cannot be decrypted by any installed encryption key. " + "Operator action required." + ), + ) + if lookup.kind == "refresh_failed": + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="Refresh token rejected. Re-consent required.", + ) + access_token = lookup.token or "" + if not access_token: + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="No token for user. Consent flow required.", + ) + + url = str(server_row.get("url") or "") + try: + _validate_oauth_user_url(url) + except ValueError as exc: + log.warning( + "mcp_pool.url_insecure server=%s scheme=%s", + server_name, + urllib.parse.urlparse(url).scheme, + ) + return _structured_error( + code="mcp_oauth_url_insecure", + server=server_name, + detail=str(exc), + ) + + self._cb_gate(server_name) + + cfg = _pool_cfg_from_row(server_row) + key = (user_id, server_name) + entry = await self._ensure_pool_entry(key) + capture = entry.auth_capture + try: + sdk_result = await self._dispatch_pool_with_entry_call( + entry=entry, + key=key, + cfg=cfg, + access_token=access_token, + sdk_call=lambda s: s.get_prompt(original_name, arguments=arguments), + ) + except BaseException as exc: + classification = self._classify_failure(exc, capture=capture) + if classification == "auth_401": + self._evict_session(key) + if retry_count == 0: + log.debug( + "mcp_pool.auth_401_initial_prompt server=%s user=%s exc=%s", + server_name, + user_id, + type(exc).__name__, + ) + raise _PoolDispatchRetryRequested from None + log.debug( + "mcp_pool.auth_401_retry_failed_prompt server=%s user=%s exc=%s", + server_name, + user_id, + type(exc).__name__, + ) + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="Refreshed token still rejected. Re-consent required.", + ) + if classification == "auth_403": + self._evict_session(key) + log.debug( + "mcp_pool.auth_403_prompt 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, + server_row=server_row, + capture=capture, + kind="prompt", + ) + if classification == "transport": + self._cb_record_failure(server_name) + self._evict_session(key) + log.debug( + "mcp_pool.transport_failure_prompt server=%s user=%s exc=%s", + server_name, + user_id, + type(exc).__name__, + ) + raise + raise + + self._cb_record_success(server_name) + return _decode_prompt_result(sdk_result) + def _evict_session(self, key: tuple[str, str]) -> None: """Drop the cached session AND catalog on a pool entry. @@ -3091,30 +4240,39 @@ class MCPClientManager: 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. + Catalog cleanup (Phase 7 / 7b): clearing ``entry.tools`` / + ``entry.resources`` / ``entry.prompts`` here ensures an + evicted-then-not-yet-reconnected entry contributes no stale + entries to ``_user_tool_map`` / ``_user_resource_map`` / + ``_user_prompt_map``. Without this, ``is_mcp_tool`` / + ``is_mcp_prompt`` / pool resource resolution could return True + for names whose backing pool is gone, then the resolver 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 + evict.resources = None + evict.prompts = 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._rebuild_user_resource_map(user_id) + self._rebuild_user_prompt_map(user_id) + # Wake the user's session so its merged tool / resource / + # prompt lists shrink 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) + self._notify_user_resource_listeners(user_id) + self._notify_user_prompt_listeners(user_id) async def _handle_auth_403( self, @@ -3123,6 +4281,7 @@ class MCPClientManager: server_name: str, server_row: dict[str, Any], capture: _AuthCapture, + kind: Literal["tool", "resource", "prompt"] = "tool", ) -> str: """Map a 403 + WWW-Authenticate into a structured error. @@ -3131,32 +4290,64 @@ class MCPClientManager: can construct an authorize URL with the union of original + new scopes — re-consenting with the original scopes alone would loop because the AS would re-issue the same insufficient - token. Other 403s become a generic ``mcp_tool_call_forbidden`` - with no retry; the user lacks permission and a step-up - wouldn't help. + token. Other 403s become a generic per-operation forbidden + code (``mcp_tool_call_forbidden`` / ``mcp_resource_read_forbidden`` + / ``mcp_prompt_get_forbidden``) so the dashboard renderer can + special-case the message. The user lacks permission and a + step-up wouldn't help; no retry. + + Both branches now emit an audit event via + :func:`emit_oauth_failure_audit` carrying the ``kind`` and + resolved error ``code`` so operators can distinguish tool-call + vs resource-read vs prompt-get 403s for the same + ``(user, server)`` and detect cross-tenant probing of generic + forbidden surfaces. """ header = capture.www_authenticate or "" error_token = parse_www_authenticate_error(header) if error_token == "insufficient_scope": scopes = parse_www_authenticate_scope(header) scopes = scopes[:_MAX_INSUFFICIENT_SCOPE_REPORTED] - await emit_insufficient_scope_audit( + await emit_oauth_failure_audit( app_state=self._app_state, user_id=user_id, server_name=server_name, server_row=server_row, + kind=kind, + code="mcp_insufficient_scope", scopes=scopes, ) return _structured_error( code="mcp_insufficient_scope", server=server_name, - detail=("Tool requires elevated scopes. Re-consent flow with new scopes required."), + detail=( + f"{kind.capitalize()} requires elevated scopes. " + "Re-consent flow with new scopes required." + ), scopes_required=list(scopes), ) + forbidden_code = { + "tool": "mcp_tool_call_forbidden", + "resource": "mcp_resource_read_forbidden", + "prompt": "mcp_prompt_get_forbidden", + }[kind] + forbidden_detail = { + "tool": "Tool call forbidden by upstream policy.", + "resource": "Resource read forbidden by upstream policy.", + "prompt": "Prompt invocation forbidden by upstream policy.", + }[kind] + await emit_oauth_failure_audit( + app_state=self._app_state, + user_id=user_id, + server_name=server_name, + server_row=server_row, + kind=kind, + code=forbidden_code, + ) return _structured_error( - code="mcp_tool_call_forbidden", + code=forbidden_code, server=server_name, - detail="Tool call forbidden by upstream policy.", + detail=forbidden_detail, ) async def _dispatch_pool_with_entry( @@ -3171,24 +4362,67 @@ class MCPClientManager: ) -> str: """Hold ``entry.open_lock`` across connect-or-reuse AND ``call_tool``. - Lock held across ``call_tool`` because the entry's + Thin wrapper over :meth:`_dispatch_pool_with_entry_call` — + passes a closure that invokes ``session.call_tool`` and + decodes the SDK result. See the helper for the carrier-race + rationale (which Phase 7 verified for the tool path; Phase 7b + reuses the same shape for resources & prompts). + + **Why this wrapper exists** (revisited under Phase 7b pre-push + review): the autouse fixture + ``tests/test_mcp_pool_auth_introspection.py::_install_capture_intercept`` + monkeypatches this method to stash ``entry.auth_capture`` on + ``mgr._test_active_capture`` for downstream call_tool stubs. + Inlining the wrapper would require redirecting the monkeypatch + to ``_dispatch_pool_with_entry_call`` (different kwargs shape) + and re-validating every test that depends on the interception. + Keeping the wrapper preserves a stable interception point on + the codebase's hottest correctness path. + """ + result = await self._dispatch_pool_with_entry_call( + entry=entry, + key=key, + cfg=cfg, + access_token=access_token, + sdk_call=lambda s: s.call_tool(original_name, arguments), + ) + return _decode_tool_result(result) + + async def _dispatch_pool_with_entry_call( + self, + *, + entry: PoolEntryState, + key: tuple[str, str], + cfg: dict[str, Any], + access_token: str, + sdk_call: Callable[[Any], Awaitable[Any]], + ) -> Any: + """Hold ``entry.open_lock`` across connect-or-reuse AND ``sdk_call``. + + Generalisation of the Phase 7 ``_dispatch_pool_with_entry`` + machinery (RFC §3.2): ``sdk_call`` is a closure that takes the + SDK ``ClientSession`` and returns an awaitable. The dispatcher + passes ``lambda s: s.call_tool(name, args)`` for tools, + ``lambda s: s.read_resource(uri)`` for resources, and + ``lambda s: s.get_prompt(name, arguments=args)`` for prompts. + + Lock held across the SDK call because the entry's ``_AuthCapture`` is keyed off the httpx event hook; releasing the lock would let a concurrent same-(user, server) dispatch overwrite the carrier mid-flight and attribute one caller's auth failure to another. Holding the lock serialises same-(user, server) calls — acceptable because that contention - scenario is rare in practice (ChatSession dispatches - sequentially and per-(user, server) parallelism is not a - production requirement). + scenario is rare in practice. Reset of ``entry.auth_capture`` and ``entry.auth_fired_event`` - happens under the lock before ``call_tool`` — see + happens under the lock before the SDK call — see :class:`PoolEntryState` for why the carrier is entry-owned. ``entry.in_flight`` accounting is preserved for the eviction - interlock — it's belt-and-braces here since ``open_lock.locked()`` + interlock — belt-and-braces since ``open_lock.locked()`` already signals "do not evict", but the in-flight counter - remains the source of truth for :meth:`_close_pool_entry_if_idle`. + remains the source of truth for + :meth:`_close_pool_entry_if_idle`. """ async with entry.open_lock: entry.last_used = time.monotonic() @@ -3211,9 +4445,9 @@ class MCPClientManager: raise RuntimeError(f"Pool connect for {key!r} produced no session") entry.in_flight += 1 try: - # Race ``call_tool`` against the carrier's fired event. + # Race ``sdk_call`` against the carrier's fired event. # Without this race, an upstream 4xx on a REUSED session - # never propagates back through ``call_tool``: the SDK's + # never propagates back through the SDK call: the SDK's # ``_receive_loop`` is in BaseSession's TaskGroup, nested # inside ``streamablehttp_client``'s TaskGroup. When # the spawned ``handle_request_async`` task raises @@ -3230,10 +4464,13 @@ class MCPClientManager: # forever-hung ``response_stream_reader.receive()``. # The carrier-fired event lets us short-circuit before # the SDK's hang manifests. - call_task = asyncio.create_task(session.call_tool(original_name, arguments)) + async def _await_sdk_call() -> Any: + return await sdk_call(session) + + call_task: asyncio.Task[Any] = asyncio.create_task(_await_sdk_call()) fired_task = asyncio.create_task(entry.auth_fired_event.wait()) try: - done, pending = await asyncio.wait( + done, _pending = await asyncio.wait( {call_task, fired_task}, return_when=asyncio.FIRST_COMPLETED, ) @@ -3256,48 +4493,91 @@ class MCPClientManager: with contextlib.suppress(BaseException): await task if call_task in done: - result = call_task.result() - else: - # Hook captured 4xx before call_tool returned. The - # SDK won't propagate the failure through - # call_tool, so eagerly tear down the session and - # raise a sentinel that the dispatcher's - # ``_classify_failure`` will resolve via the - # carrier (which holds the captured status). - raise _CarrierAuthSignal() + return call_task.result() + # Hook captured 4xx before the SDK call returned. The + # SDK won't propagate the failure through the call, so + # eagerly tear down the session and raise a sentinel + # that the dispatcher's ``_classify_failure`` will + # resolve via the carrier (which holds the captured + # status). + raise _CarrierAuthSignal() finally: entry.in_flight -= 1 - return _decode_tool_result(result) # -- resource read ------------------------------------------------------- - def _match_template(self, uri: str) -> tuple[str, str] | None: + def _match_template(self, uri: str, *, user_id: str | None = None) -> tuple[str, str] | None: """Find the longest matching template prefix for an expanded URI. Returns ``(server_name, template_uri)`` or *None* if no match. - The match uses the longest static prefix stored in - ``_template_prefixes`` (the portion of each template URI before - the first ``{``), with simple ``startswith`` matching. + Uses ``startswith`` longest-prefix matching. + + Per scope decision 0.1, when ``user_id is not None`` the + per-user template index is consulted FIRST; static templates + only run as a fallback. This keeps a user's own scoped templates + authoritative within their namespace and avoids surprising + static→pool fallthrough that would attach a per-user bearer + to a server the operator believed was the same name. """ best: tuple[str, str] | None = None best_len = 0 + if user_id is not None: + for prefix, mapping in (self._user_template_prefixes.get(user_id) or {}).items(): + if uri.startswith(prefix) and len(prefix) > best_len: + best = mapping + best_len = len(prefix) + if best is not None: + return best for prefix, mapping in self._template_prefixes.items(): if uri.startswith(prefix) and len(prefix) > best_len: best = mapping best_len = len(prefix) return best - def read_resource_sync(self, uri: str, timeout: int = 120) -> str: + def read_resource_sync( + self, uri: str, *, user_id: str | None = None, timeout: int = 120 + ) -> str: """Read a resource by URI synchronously (blocks the calling thread). Returns text content for ``TextResourceContents``, or base64 data for ``BlobResourceContents``. + + When ``user_id`` is supplied AND the URI resolves to a pool + entry (per scope decision 0.1, per-user-first), dispatch goes + through the per-(user, server) pool with the same 401 / 403 / + consent-required handling as :meth:`call_tool_sync`. Otherwise + the call takes the byte-identical static path (invariant 1). """ + # Phase 7b — per-user-first pool dispatch. + if user_id and self._app_state is not None and self._storage is not None: + pool_target = self._resolve_pool_target_resource(uri, user_id) + if pool_target is not None: + return self._dispatch_pool_resource_sync( + user_id=user_id, + server_name=pool_target[0], + uri=pool_target[1], + server_row=pool_target[2], + timeout=timeout, + ) + mapping = self._resource_map.get(uri) if mapping is None: # Fall back to template prefix matching for expanded URIs mapping = self._match_template(uri) if mapping is None: + # Per-user resources whose static lookup misses fall through here + # only when the user has no pool entry for the server. A pool + # entry exists in ``_user_resource_map`` only AFTER discovery; the + # pool resolver tried it above and returned None (no oauth_user + # row). + if ( + user_id is not None + and (self._user_resource_map.get(user_id) or {}).get(uri) is not None + ): + # Per-user map carries this URI but the resolver could not + # match it to an oauth_user server row; this is an internal + # inconsistency. + raise ValueError(f"Unknown MCP resource: {uri} (per-user map / DB mismatch)") raise ValueError(f"Unknown MCP resource: {uri}") server_name, _ = mapping @@ -3329,16 +4609,7 @@ class MCPClientManager: raise self._cb_record_success(server_name) - - parts: list[str] = [] - for item in result.contents: - if hasattr(item, "text"): - parts.append(item.text) - elif hasattr(item, "blob"): - parts.append(item.blob) - else: - parts.append(str(item)) - return "\n".join(parts) if parts else "(empty resource)" + return _decode_resource_result(result) # -- prompt invocation --------------------------------------------------- @@ -3346,16 +4617,67 @@ class MCPClientManager: self, prefixed_name: str, arguments: dict[str, str] | None = None, + *, + user_id: str | None = None, timeout: int = 30, ) -> list[dict[str, Any]]: """Invoke an MCP prompt synchronously and return expanded messages. Returns a list of ``{role: str, content: str}`` dicts. + + When ``user_id`` is supplied AND ``prefixed_name`` resolves to a + pool entry, dispatch goes through the per-(user, server) pool + with the same 401 / 403 / consent-required handling as + :meth:`call_tool_sync`. Otherwise the call takes the byte- + identical static path (invariant 1). + + Pool error path: structured-error responses (consent required, + decrypt failure, insufficient scope, etc.) are surfaced via + :class:`RuntimeError` so the agent-loop's existing + ``except Exception`` block at the call site renders the error + message. The return type stays ``list[dict[str, Any]]`` because + embedding error JSON in a single-element message list would + pollute the prompt protocol. """ - mapping = self._prompt_map.get(prefixed_name) - if mapping is None: + # Phase 7b — pool dispatch (per-user-first, scope decision 0.1). + static_mapping = self._prompt_map.get(prefixed_name) + static_server: str | None = None + static_original: str | None = None + if static_mapping is not None: + static_server, static_original = static_mapping + + if user_id and self._app_state is not None and self._storage is not None: + pool_target = self._resolve_pool_target_prompt( + prefixed_name, static_server, static_original + ) + if pool_target is not None: + return self._dispatch_pool_prompt_sync( + user_id=user_id, + server_name=pool_target[0], + original_name=pool_target[1], + arguments=arguments, + server_row=pool_target[2], + timeout=timeout, + ) + + if static_mapping is None: + # Per-user prompts whose static lookup misses fall through here + # only when the user has no pool entry for the server. A pool + # entry exists in ``_user_prompt_map`` only AFTER discovery; the + # pool resolver tried it above and returned None (no oauth_user + # row). + if ( + user_id is not None + and (self._user_prompt_map.get(user_id) or {}).get(prefixed_name) is not None + ): + # Per-user map carries this prefixed name but the resolver + # could not match it to an oauth_user server row; this is + # an internal inconsistency. + raise ValueError( + f"Unknown MCP prompt: {prefixed_name} (per-user map / DB mismatch)" + ) raise ValueError(f"Unknown MCP prompt: {prefixed_name}") - server_name, original_name = mapping + server_name, original_name = static_mapping self._cb_gate(server_name) @@ -3387,13 +4709,7 @@ class MCPClientManager: raise self._cb_record_success(server_name) - - messages: list[dict[str, Any]] = [] - for msg in result.messages: - content = msg.content - text = content.text if hasattr(content, "text") else str(content) - messages.append({"role": msg.role, "content": text}) - return messages + return _decode_prompt_result(result) # --------------------------------------------------------------------------- @@ -3424,6 +4740,36 @@ def _decode_tool_result(result: Any) -> str: return output +def _decode_resource_result(result: Any) -> str: + """Render an MCP ``resources/read`` result into the string the agent sees. + + Walks ``result.contents`` collecting text parts (TextResourceContents) + and base64 data (BlobResourceContents). Shared by the static and pool + resource-read paths. + """ + parts: list[str] = [] + for item in result.contents: + if hasattr(item, "text"): + parts.append(item.text) + elif hasattr(item, "blob"): + parts.append(item.blob) + else: + parts.append(str(item)) + return "\n".join(parts) if parts else "(empty resource)" + + +def _decode_prompt_result(result: Any) -> list[dict[str, Any]]: + """Render an MCP ``prompts/get`` result into the list-of-messages + the agent sees. Shared by the static and pool prompt-get paths. + """ + messages: list[dict[str, Any]] = [] + for msg in result.messages: + content = msg.content + text = content.text if hasattr(content, "text") else str(content) + messages.append({"role": msg.role, "content": text}) + return messages + + def _structured_error( *, code: str, diff --git a/turnstone/core/mcp_oauth.py b/turnstone/core/mcp_oauth.py index aaeb0754..c8a26729 100644 --- a/turnstone/core/mcp_oauth.py +++ b/turnstone/core/mcp_oauth.py @@ -1528,21 +1528,34 @@ async def _audit_event( log.debug("mcp_server.oauth.audit_emit_failed", action=action, exc_info=True) -async def emit_insufficient_scope_audit( +async def emit_oauth_failure_audit( *, app_state: Any, user_id: str, server_name: str, server_row: dict[str, Any], - scopes: tuple[str, ...], + kind: str, + code: str, + scopes: tuple[str, ...] = (), ) -> None: """Emit ``mcp_server.oauth.insufficient_scope_emitted`` audit event. Best-effort: :func:`_audit_event` already swallows storage / write failures internally so audit emission never breaks dispatch. - Operators tracking step-up patterns consume this via the standard - audit log. Called by the pool dispatcher after classifying a 403 - ``WWW-Authenticate: error="insufficient_scope"``. + Operators tracking step-up patterns and forbidden-policy hits + consume this via the standard audit log. Called by the pool + dispatcher after classifying a 403 — both + ``WWW-Authenticate: error="insufficient_scope"`` and the generic + forbidden branch route here so cross-tenant probing leaves an + audit trail (Phase 7 left the generic 403 branch silent; Phase 7b + closes that gap). + + The ``kind`` ("tool" / "resource" / "prompt") and ``code`` + (``mcp_insufficient_scope`` / ``mcp_tool_call_forbidden`` / + ``mcp_resource_read_forbidden`` / ``mcp_prompt_get_forbidden``) + fields land in the audit detail so operators can distinguish + tool-call vs resource-read vs prompt-get 403s for the same + ``(user, server)``. """ if app_state is None: return @@ -1553,7 +1566,7 @@ async def emit_insufficient_scope_audit( user_id=user_id, action="mcp_server.oauth.insufficient_scope_emitted", server_name=server_name, - detail={"scopes_required": list(scopes)}, + detail={"scopes_required": list(scopes), "kind": kind, "code": code}, ) diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 8aca290e..e2f76d60 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -880,12 +880,16 @@ class ChatSession: # 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, user_id=self._mcp_user_id) - # Register for resource-change notifications + # Register for resource-change notifications. + # ``user_id`` scopes the listener so pool-only resource + # changes for OTHER users do not wake this session. self._mcp_resource_cb = self._on_mcp_resources_changed - mcp_client.add_resource_listener(self._mcp_resource_cb) - # Register for prompt-change notifications + mcp_client.add_resource_listener(self._mcp_resource_cb, user_id=self._mcp_user_id) + # Register for prompt-change notifications. + # ``user_id`` scopes the listener so pool-only prompt changes + # for OTHER users do not wake this session. self._mcp_prompt_cb = self._on_mcp_prompts_changed - mcp_client.add_prompt_listener(self._mcp_prompt_cb) + mcp_client.add_prompt_listener(self._mcp_prompt_cb, user_id=self._mcp_user_id) else: self._tools = INTERACTIVE_TOOLS self._task_tools = TASK_AGENT_TOOLS @@ -1445,10 +1449,15 @@ class ChatSession: 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) + # ``user_id`` MUST mirror the value passed at registration — + # the listener identity is ``(user_id, callback)`` and an + # unscoped removal would leave the registration in place. + self._mcp_client.remove_resource_listener( + self._mcp_resource_cb, user_id=self._mcp_user_id + ) self._mcp_resource_cb = None if self._mcp_client and self._mcp_prompt_cb: - self._mcp_client.remove_prompt_listener(self._mcp_prompt_cb) + self._mcp_client.remove_prompt_listener(self._mcp_prompt_cb, user_id=self._mcp_user_id) self._mcp_prompt_cb = None if self._watch_runner: self._watch_runner.remove_dispatch_fn(self._ws_id) @@ -1915,7 +1924,9 @@ class ChatSession: ) # MCP resource catalog (lets the model know what's available for read_resource) if self._mcp_client: - all_resources = self._mcp_client.get_resources() + # Per-user merge: pool entries for ``self._mcp_user_id`` are + # included; other users' pool resources are not. + all_resources = self._mcp_client.get_resources(user_id=self._mcp_user_id) concrete = [r for r in all_resources if not r.get("template")] templates = [r for r in all_resources if r.get("template")] if concrete or templates: @@ -1940,7 +1951,9 @@ class ChatSession: dev_parts.append("\n".join(lines)) # MCP prompt catalog (lets the model know what's available for use_prompt) if self._mcp_client: - prompts = self._mcp_client.get_prompts() + # Per-user merge: pool entries for ``self._mcp_user_id`` are + # included; other users' pool prompts are not. + prompts = self._mcp_client.get_prompts(user_id=self._mcp_user_id) if prompts: lines = [""] for p in prompts[:30]: @@ -2329,10 +2342,13 @@ class ChatSession: if not caps.supports_web_search and not self._resolve_search_client(): tools = _without_tool(tools, "web_search") - # Gate MCP tools: only include when relevant MCP servers are connected - if not self._mcp_client or not self._mcp_client.resource_count: + # Gate MCP tools: only include when relevant MCP servers are + # connected. Per-user variants (scope decision 0.2) keep the + # tool visible for a pool-only user even when the static catalog + # is empty. + if not self._mcp_client or not self._mcp_client.resource_count_for_user(self._mcp_user_id): tools = _without_tool(tools, "read_resource") - if not self._mcp_client or not self._mcp_client.prompt_count: + if not self._mcp_client or not self._mcp_client.prompt_count_for_user(self._mcp_user_id): tools = _without_tool(tools, "use_prompt") return tools @@ -7941,14 +7957,21 @@ class ChatSession: assert self._mcp_client is not None mcp_error = False try: - output = self._mcp_client.read_resource_sync(uri, timeout=self.tool_timeout) + # Per-user pool dispatch (Phase 7b): when ``user_id`` is set + # and the URI resolves to an oauth_user pool entry, the read + # goes through the per-(user, server) pool with token / + # 401 / 403 / consent-required handling. Otherwise the + # static path runs byte-identical (invariant 1). + output = self._mcp_client.read_resource_sync( + uri, user_id=self._mcp_user_id, timeout=self.tool_timeout + ) except TimeoutError: output = f"MCP resource read timed out after {self.tool_timeout}s" mcp_error = True self.ui.on_error(output) - except Exception: + except Exception as e: log.warning("MCP resource read failed for %s", uri, exc_info=True) - output = "MCP resource error: failed to read resource" + output = f"MCP resource error: {e}" mcp_error = True self.ui.on_error(output) @@ -7977,7 +8000,7 @@ class ChatSession: "needs_approval": False, "error": "No MCP servers configured", } - if not self._mcp_client.is_mcp_prompt(name): + if not self._mcp_client.is_mcp_prompt(name, user_id=self._mcp_user_id): return { "call_id": call_id, "func_name": "use_prompt", @@ -8023,17 +8046,25 @@ class ChatSession: assert self._mcp_client is not None mcp_error = False try: + # Per-user pool dispatch (Phase 7b): structured-error + # responses (consent required, decrypt failure, insufficient + # scope, etc.) surface here as ``RuntimeError`` carrying the + # JSON payload — caught by the broad ``except Exception`` + # below so the agent renders the error message. messages = self._mcp_client.get_prompt_sync( - name, arguments or None, timeout=self.tool_timeout + name, + arguments or None, + user_id=self._mcp_user_id, + timeout=self.tool_timeout, ) output = "\n\n".join(f"[{m['role']}]: {m['content']}" for m in messages) except TimeoutError: output = f"MCP prompt timed out after {self.tool_timeout}s" mcp_error = True self.ui.on_error(output) - except Exception: + except Exception as e: log.warning("MCP prompt invocation failed for %s", name, exc_info=True) - output = "MCP prompt error: failed to invoke prompt" + output = f"MCP prompt error: {e}" mcp_error = True self.ui.on_error(output) @@ -10192,13 +10223,12 @@ class ChatSession: elif arg and arg.split()[0] == "refresh": self._handle_mcp_refresh(arg) else: - # 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). + # Phase 7 + 7b: pass session-bound user_id so the /mcp + # listing surfaces this user's pool tools, resources, + # and prompts alongside the static catalog. tools = self._mcp_client.get_tools(user_id=self._mcp_user_id) - resources = self._mcp_client.get_resources() - prompts = self._mcp_client.get_prompts() + resources = self._mcp_client.get_resources(user_id=self._mcp_user_id) + prompts = self._mcp_client.get_prompts(user_id=self._mcp_user_id) mcp_lines = [] if tools: mcp_lines.append(f"MCP tools ({len(tools)}):")