feat(mcp): per-user resource + prompt pool dispatch (Phase 7b)

Extends the Phase 7 per-(user, server) ClientSession pool to cover
RFC §3.2 (resources/read) and §3.3 (prompts/get) on the same shape
already proven for tools/call. Pool discovery is capability-gated so
servers without resources/ or prompts/ stay free of extra round-trips.

API additions / widenings (MCPClientManager):
- ``read_resource_sync(uri, *, user_id=None, timeout=120)`` —
  per-user-first dispatch; falls through to the byte-identical static
  path when ``user_id`` is None or the URI doesn't resolve to an
  ``oauth_user`` pool entry.
- ``get_prompt_sync(prefixed_name, arguments=None, *, user_id=None,
  timeout=30)`` — same dispatch shape; structured-error responses
  surface via ``RuntimeError`` so the agent-loop's ``except Exception``
  block renders the JSON without polluting the prompt-protocol return
  shape.
- ``get_resources(user_id=None)`` / ``get_prompts(user_id=None)`` —
  per-user merged catalogs (admin/global call still passes None).
- ``add_{resource,prompt}_listener`` /
  ``remove_{resource,prompt}_listener`` —  ``user_id`` keyword scopes
  the listener so a pool-only catalog change for one user does not
  wake another user's session.
- ``resource_count_for_user(user_id=None)`` /
  ``prompt_count_for_user(user_id=None)`` — method-form variants used
  by ChatSession's ``read_resource`` / ``use_prompt`` tool gating; the
  legacy ``resource_count`` / ``prompt_count`` properties remain
  static-only for admin paths.
- ``_dispatch_pool_resource`` / ``_dispatch_pool_prompt`` async coros
  — mirror ``_dispatch_pool`` for the new SDK calls; share the
  carrier-race-and-cancel core via ``_dispatch_pool_with_entry_call``.
- ``_handle_auth_403`` extended with ``kind=Literal["tool",
  "resource", "prompt"]`` so the per-operation ``mcp_*_forbidden``
  code surfaces (kind="tool" remains the default for back-compat).
- Pool notification handler now refreshes resources / prompts on
  ``ResourceListChangedNotification`` / ``PromptListChangedNotification``
  via ``_refresh_pool_server_resources`` / ``_refresh_pool_server_prompts``.

ChatSession (``turnstone/core/session.py``) call-site updates:
- 12 sites threaded the session-bound ``user_id`` through
  ``add_*_listener`` / ``remove_*_listener``, ``get_resources`` /
  ``get_prompts``, gating, ``read_resource_sync`` /
  ``get_prompt_sync``, and ``is_mcp_prompt`` so the per-user merged
  catalog drives both the visible-tool set and dispatch.
- ``/mcp`` slash command now lists this user's pool resources and
  prompts alongside tools (Phase 7 already scoped tools).

Scope decisions:
- Per-user-first URI ordering (decision 0.1): the dispatcher attempts
  the user's pool catalog first, falling back to the static catalog
  only when no pool entry resolves the URI / prefixed name. Pool-only
  users never see the static catalog leak into their resolution.
- Method-form ``*_count_for_user`` (vs property) keeps the legacy
  ``resource_count`` / ``prompt_count`` properties intact for admin
  endpoints whose contract is "static catalog size only".
- Shared ``_dispatch_pool_with_entry_call`` helper accepts an
  ``sdk_call: Callable[[ClientSession], Awaitable[Any]]`` closure,
  keeping the entry-locked carrier-race / classification / retry
  plumbing single-source instead of a 3x copy across tool / resource
  / prompt paths.

R6 (anyio uniformity): every pool-side list / read / get path uses
``async with asyncio.timeout(...)`` — ``asyncio.wait_for`` is
forbidden in those paths because it wraps the inner awaitable in a
fresh task and surfaces ``CancelledError`` from inside
``streamablehttp_client``'s anyio TaskGroup on Python 3.11
(per ``feedback_asyncio_timeout_vs_wait_for.md``).

Tests:
- ``test_mcp_pool_auth_resource_integration.py`` — 9 real-transport
  resource tests (FastMCP upstream + ``BehaviorMiddleware``):
  401-refresh-retry success, persistent 401 -> consent_required,
  403+insufficient_scope, 403 generic -> mcp_resource_read_forbidden,
  breaker-isolation under repeated auth failures, missing-token,
  decrypt-failure, http:// URL guard, unknown-URI ValueError.
- ``test_mcp_pool_auth_prompt_integration.py`` — 9 mirror tests for
  the prompt path; structured-error responses verified via
  ``RuntimeError`` payload shape.
- ``test_mcp_user_catalog.py`` — extended unit coverage for per-user
  resource / prompt rebuild + collision policy + symmetric eviction.
- ``test_sessions.py::TestMCPToolGating`` — pool-only-user canary
  asserts ``read_resource`` / ``use_prompt`` stay visible when the
  static catalog is empty but the user has pool entries.

Round-1 review fixes (4-finder review applied, no push yet):
- bug-1: ``_exec_use_prompt`` was hardcoding ``"MCP prompt error: failed
  to invoke prompt"`` — discarding the structured-error JSON that
  ``_dispatch_pool_prompt_sync`` raises via ``RuntimeError``. Now uses
  ``f"MCP prompt error: {e}"`` mirroring ``_exec_mcp_tool``; pool-prompt
  consent_required / insufficient_scope / forbidden errors now reach
  the LLM as intended.
- bug-2 + bug-3: resource template discovery was uncapped —
  ``_cap_server_resources`` covered ``res_result.resources`` but the
  separate ``tmpl_result.resourceTemplates`` loop appended every
  template a server returned. Added ``_MAX_RESOURCE_TEMPLATES_PER_SERVER``
  (1000) + ``_cap_server_resource_templates`` helper, applied at both
  the initial discovery site (``_connect_one_pool``) and the refresh
  site (``_refresh_pool_server_resources``). Mirrors the existing
  ``_MAX_TOOLS_PER_SERVER`` / ``_MAX_PROMPTS_PER_SERVER`` defensive
  ceilings.
- sec-1 + sec-2: ``emit_insufficient_scope_audit`` generalized to
  ``emit_oauth_failure_audit(kind, code, ...)``, called from both the
  insufficient_scope branch AND the previously-silent generic 403
  branch. Audit detail now records ``{"kind": kind, "code": code,
  "scopes_required": [...]}`` so operators can distinguish tool-call
  vs resource-read vs prompt-get 403s in audit logs and so cross-
  tenant probing on the generic 403 path leaves a trail. The Phase 7
  inherited gap (``mcp_tool_call_forbidden`` had the same silence) is
  closed in the same refactor.
- perf-1: pool resource discovery now uses ``asyncio.gather(
  list_resources, list_resource_templates)`` inside the existing
  ``async with asyncio.timeout(...)`` budget — disjoint catalogs, no
  ordering dependency. Typical-case 2-RTT cold-connect resource block
  collapses to 1-RTT. Same change applied at ``_refresh_pool_server_resources``.
- q-1: ``_rebuild_user_prompt_map`` docstring corrected RFC §3.2 →
  §3.3 (resources are §3.2; prompts are §3.3).
- q-2: ``_refresh_pool_server_prompts`` docstring now carries the
  R6 / mcp-loop note that the resource sibling already had — both
  refresh paths now declare the asyncio.timeout invariant explicitly.
- q-5: added the ``_user_resource_map`` / DB-mismatch guard to
  ``read_resource_sync`` for parity with ``get_prompt_sync``. A stale
  per-user map entry with no matching oauth_user row now raises a
  specific ValueError instead of silently falling through to a
  generic ``Unknown MCP resource``.
- q-6: ``_dispatch_pool_with_entry`` (now a single-caller wrapper
  after the ``_dispatch_pool_with_entry_call`` extraction) gains a
  one-line docstring explaining why the wrapper is preserved
  (tool-decode localization + stack-trace identity for debugging).
- q-7: added 1 resource + 1 prompt end-to-end integration test that
  drive REAL discovery + dispatch in the same connect (no
  ``_seed_pool_*_map`` shortcuts), mirroring the tool path's
  ``test_integration_pool_reuse_401_refresh_and_retry_succeeds``.
  The seeded-map tests stay (faster, focused on dispatch); the new
  e2e tests cover the connect-discover-dispatch composition that
  caught Phase 6's carrier-on-entry bug.

Pre-push round-1 review fixes (3-finder review on the final state —
the lesson from Phase 7 round-3's q-1 regression: round-2 catches
what the round-1 apply pass missed):
- q-1 (MAJOR): the bug-1 sibling that round-1 missed —
  ``_exec_read_resource`` was hardcoding ``"MCP resource error: failed
  to read resource"`` while ``_exec_use_prompt`` (post-bug-1) preserved
  the structured-error JSON via ``f"... error: {e}"``. The round-1
  apply pass patched the prompt side but not the resource side. q-5's
  per-user-map / DB-mismatch ValueError was being swallowed at the
  agent loop boundary, defeating the operator-diagnostic intent. Now
  ``_exec_read_resource`` mirrors ``_exec_mcp_tool`` and ``_exec_use_prompt``.
- q-6 (nit): defensive-cap comment block at module-level cited
  "(RFC §3.2)" while covering both resource and prompt list paths;
  prompts are §3.3. Now reads "(RFC §3.2 for resources, §3.3 for
  prompts)" matching the convention the q-1 apply established.
- q-5 (rejected with better justification): the reviewer flagged
  ``_dispatch_pool_with_entry`` as a single-caller wrapper that should
  be inlined. After examination — the autouse fixture
  ``tests/test_mcp_pool_auth_introspection.py::_install_capture_intercept``
  monkeypatches this method to stash ``entry.auth_capture`` for the
  fake call_tool stubs in dispatcher-asserting tests. Inlining would
  redirect the patch to ``_dispatch_pool_with_entry_call`` (different
  kwargs shape) and require re-validating every test that depends on
  the interception. The wrapper IS load-bearing; q-6 docstring updated
  to cite the test-fixture rationale instead of the thin "stack-trace
  identity" claim.

Deferred to follow-up (documented rationale):
- perf-2: single-pass partition for system-message resource list
  (concrete vs templates). Sub-microsecond at expected scale;
  opportunistic-only.
- q-2 (pre-push): ~200 lines of fixture infrastructure
  (``BehaviorMiddleware``, ``_build_server``, ``_seed_oauth_server``,
  ``running_loop_mgr``, etc.) duplicated across three pool-integration
  test files. Real maintenance cost, but a 200-line conftest extraction
  is a focused refactor that earns its own commit / PR. Tracking as
  follow-up rather than balloon Phase 7b's diff further.
- q-3 / q-4 (refactor): extract shared dispatcher / scheduler
  helpers to compress three near-identical 90-line bodies (round-1
  q-3 was the same root cause; the pre-push q-3/q-4 reviewer
  reaffirmed it concretely). Three named methods preserve readability
  for the codebase's hottest correctness path; follow-up if
  duplication grows further or if a per-path divergence ships.
- q-4 (round-1, distinct from pre-push q-4): split pool concerns
  into ``mcp_pool.py``. Out-of-scope per finder; future refactor as
  the file approaches the navigation/merge-conflict threshold.

3.13: 5590 passed (5541 baseline -> +49 net; pre-review +47, q-7
e2e tests added +2). Existing audit-detail tests updated in-place
to expect the new ``kind`` and ``code`` fields.
3.11: 5590 passed (parity gate per ``feedback_pytest_env_parity.md``).
This commit is contained in:
Patrick Buckley
2026-05-06 16:19:09 -07:00
parent c757c22f55
commit 124615cce0
8 changed files with 3981 additions and 174 deletions
@@ -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
@@ -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
File diff suppressed because it is too large Load Diff
+11 -4
View File
@@ -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:
+52 -10
View File
@@ -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")
+1468 -122
View File
File diff suppressed because it is too large Load Diff
+19 -6
View File
@@ -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},
)
+54 -24
View File
@@ -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 = ["<mcp-prompts>"]
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)}):")