mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
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:
@@ -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
@@ -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
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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
@@ -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)}):")
|
||||
|
||||
Reference in New Issue
Block a user