diff --git a/tests/test_mcp_http_parsers.py b/tests/test_mcp_http_parsers.py new file mode 100644 index 00000000..98843e24 --- /dev/null +++ b/tests/test_mcp_http_parsers.py @@ -0,0 +1,242 @@ +"""Unit tests for ``turnstone.core.mcp_http_parsers``. + +The parser replaces the prior hand-rolled scanners that used +``header.lower().find("scope")`` to locate parameter names — that approach +misparsed ``scope`` embedded inside other tokens (``xscope``) or inside +quoted-string values of preceding params. Each adversarial case below +asserts the new tokenizer respects RFC 7235 ``challenge → auth-param`` +boundaries; the docstrings document the equivalent input that broke the +naive parser. Negative-test verification: temporarily reverting +``parse_www_authenticate_scope`` to delegate to ``header.lower().find("scope")`` +makes ``test_scope_inside_realm_value`` and ``test_scope_inside_xscope`` fail. +""" + +from __future__ import annotations + +import time + +import pytest + +from turnstone.core.mcp_http_parsers import ( + parse_www_authenticate_bearer, + parse_www_authenticate_error, + parse_www_authenticate_scope, +) + + +class TestParseScope: + def test_basic_scope(self) -> None: + header = 'Bearer error="insufficient_scope", scope="files:read mail:send"' + assert parse_www_authenticate_scope(header) == ("files:read", "mail:send") + + def test_no_scope_param(self) -> None: + assert parse_www_authenticate_scope('Bearer error="invalid_token"') == () + + def test_unterminated_quoted_string_returns_empty(self) -> None: + assert parse_www_authenticate_scope('Bearer scope="files:read') == () + + def test_escaped_chars_in_value_drops_invalid_scope_token(self) -> None: + # RFC 7230 §3.2.6 backslash escapes decode the literal scope to + # ``files:read "weird"``. RFC 6749 §3.3 ``scope-token`` forbids + # ``"``, so ``"weird"`` is dropped and only ``files:read`` + # survives the post-split validation. + header = r'Bearer scope="files:read \"weird\""' + assert parse_www_authenticate_scope(header) == ("files:read",) + + def test_empty_string(self) -> None: + assert parse_www_authenticate_scope("") == () + + def test_unquoted_scope_value(self) -> None: + # Unquoted single token. + assert parse_www_authenticate_scope("Bearer scope=files:read") == ("files:read",) + + # --- the four headline misparse cases --- + + def test_scope_inside_xscope(self) -> None: + """``Bearer xscope="value"`` must NOT be read as ``scope``. + + The naive ``find("scope")`` matched at position 7 inside + ``xscope`` and returned ``("value",)``. + """ + assert parse_www_authenticate_scope('Bearer xscope="value"') == () + + def test_scope_inside_realm_value(self) -> None: + """``Bearer realm="my scope=fake", scope="real"`` must return ``("real",)``. + + The naive parser found ``scope=`` inside the quoted ``realm`` + value first and returned ``("fake",)``. + """ + header = 'Bearer realm="my scope=fake", scope="real"' + assert parse_www_authenticate_scope(header) == ("real",) + + def test_scope_inside_quoted_realm_with_escaped_quotes(self) -> None: + """``Bearer realm="foo scope=\\"admin:write\\" bar"`` returns ``()``. + + The inner ``scope=`` is wholly inside the quoted-string value of + ``realm`` — there is no top-level ``scope`` auth-param, so the + result is empty. + """ + header = r'Bearer realm="foo scope=\"admin:write\" bar"' + assert parse_www_authenticate_scope(header) == () + + def test_scope_token_validation_drops_control_bytes(self) -> None: + """Tokens containing CR / LF / tab / DEL / quote are dropped. + + RFC 6749 §3.3 restricts ``scope-token`` to visible ASCII + excluding ``"`` and ``\\``. The splitter applies that + validation so a malicious AS cannot smuggle CRLF (or the like) + through a future log / notification path that prints the scope + list verbatim. ``"a\\rb"`` and ``"\\nc"`` fail validation; + ``"d"`` survives. The legitimate space separator splits ``d`` + into its own token. + """ + # Build via concatenation so the assertion stays intelligible. + header = 'Bearer scope="a\rb \nc d"' + assert parse_www_authenticate_scope(header) == ("d",) + + +class TestParseError: + def test_basic_quoted_error(self) -> None: + assert ( + parse_www_authenticate_error('Bearer error="insufficient_scope"') + == "insufficient_scope" + ) + + def test_other_quoted_error_tokens(self) -> None: + assert parse_www_authenticate_error('Bearer error="invalid_token"') == "invalid_token" + assert parse_www_authenticate_error('Bearer error="invalid_request"') == "invalid_request" + + def test_no_error_param(self) -> None: + assert parse_www_authenticate_error("Bearer realm=foo") is None + + def test_error_description_does_not_match_error(self) -> None: + """``error_description`` is its own auth-param key, not ``error``. + + The tokenizer reads ``_`` as part of the token (RFC 7230 ``tchar``), + so ``error_description`` becomes one key, ``error`` another. + """ + assert parse_www_authenticate_error('Bearer error_description="bad"') is None + + def test_unquoted_error(self) -> None: + # Some ASes don't quote the error token. + assert ( + parse_www_authenticate_error("Bearer error=insufficient_scope") == "insufficient_scope" + ) + + def test_empty_string(self) -> None: + assert parse_www_authenticate_error("") is None + + def test_error_inside_realm_value(self) -> None: + """``Bearer realm="my error=fake", error="real"`` must return ``"real"``. + + Naive parser grabbed ``fake`` from inside the ``realm`` quoted + value. + """ + header = 'Bearer realm="my error=fake", error="real"' + assert parse_www_authenticate_error(header) == "real" + + +class TestBearerDict: + def test_returns_lowercased_keys(self) -> None: + header = 'Bearer Realm="x", Error="y", Scope="a b"' + params = parse_www_authenticate_bearer(header) + assert params == {"realm": "x", "error": "y", "scope": "a b"} + + def test_non_bearer_scheme_returns_empty(self) -> None: + assert parse_www_authenticate_bearer('Basic realm="x"') == {} + + def test_no_scheme(self) -> None: + assert parse_www_authenticate_bearer('realm="x"') == {} + + def test_bearer_only_no_params(self) -> None: + assert parse_www_authenticate_bearer("Bearer ") == {} + + def test_bearer_with_no_space_returns_empty(self) -> None: + # ``BearerToken`` is not a Bearer challenge (no separator). + assert parse_www_authenticate_bearer("BearerToken") == {} + + def test_first_value_wins_on_duplicate(self) -> None: + # If a malformed AS sends two ``scope=`` params we keep the first. + # The earlier ``find()``-based scanner would have returned the + # last; either choice is legal for malformed input but we need + # to be consistent. + header = 'Bearer scope="first", scope="second"' + assert parse_www_authenticate_bearer(header) == {"scope": "first"} + + def test_trailing_comma(self) -> None: + header = 'Bearer error="x",' + assert parse_www_authenticate_bearer(header) == {"error": "x"} + + def test_multiple_commas(self) -> None: + header = 'Bearer ,, error="x",,, scope="y"' + assert parse_www_authenticate_bearer(header) == {"error": "x", "scope": "y"} + + def test_embedded_escaped_quote(self) -> None: + header = r'Bearer realm="he said \"hi\""' + assert parse_www_authenticate_bearer(header) == {"realm": 'he said "hi"'} + + def test_param_without_value_skipped(self) -> None: + header = 'Bearer realm, error="x"' + # ``realm`` without ``=`` is dropped; ``error`` survives. + assert parse_www_authenticate_bearer(header) == {"error": "x"} + + @pytest.mark.parametrize( + "header,expected", + [ + ("", {}), + ("Bearer", {}), + ('Bearer realm=""', {"realm": ""}), + ('Bearer realm="", scope=""', {"realm": "", "scope": ""}), + ], + ) + def test_edge_cases(self, header: str, expected: dict[str, str]) -> None: + assert parse_www_authenticate_bearer(header) == expected + + +class TestPathologicalInput: + def test_oversized_pathological_input_rejected_under_50ms(self) -> None: + """Headers longer than the defensive cap return ``{}`` immediately. + + The cap is set to 4096 bytes — real ASes emit a few hundred bytes + at most. This guards both ``parse_www_authenticate_bearer`` + callers against pathological input from a misbehaving server. + The previous ``header.lower().find("scope", i)`` loop was + O(N**2) — a 100 KB header with no ``=`` took ~330 ms because + each ``find`` rescanned the entire suffix. The single-pass + tokenizer (capped at 4 KB) reduces this to a one-shot length + check that returns ``{}`` in microseconds, so the budget is + generous regardless of which side of the cap was hit. + """ + big = "Bearer scope=" + "a" * 10_000 + start = time.perf_counter() + result = parse_www_authenticate_scope(big) + elapsed = time.perf_counter() - start + assert result == () + assert elapsed < 0.05, f"oversized-header reject took {elapsed * 1000:.1f}ms" + + def test_within_cap_long_header_under_50ms(self) -> None: + """A 4 KB header with thousands of ``find`` candidates still parses fast. + + Stays under the cap so the tokenizer actually runs end to end — + the goal is to prove the inner loop is O(N), not just that the + cap rejects oversized input. + """ + # Pack the header right up to the cap with non-matching + # auth-params, then put the real ``scope`` at the end. + filler_parts = [] + size = len("Bearer ") + i = 0 + while size < 3900: + part = f'xscope{i}="ignore", ' + if size + len(part) > 3900: + break + filler_parts.append(part) + size += len(part) + i += 1 + header = "Bearer " + "".join(filler_parts) + 'scope="real"' + assert len(header) <= 4096 + start = time.perf_counter() + result = parse_www_authenticate_scope(header) + elapsed = time.perf_counter() - start + assert result == ("real",) + assert elapsed < 0.05, f"4kb tokenize took {elapsed * 1000:.1f}ms" diff --git a/tests/test_mcp_pool_auth_integration.py b/tests/test_mcp_pool_auth_integration.py new file mode 100644 index 00000000..ea226c69 --- /dev/null +++ b/tests/test_mcp_pool_auth_integration.py @@ -0,0 +1,718 @@ +"""Phase 6 integration tests — real-transport drives 401/403 through the SDK. + +These are the structural exit criterion for Phase 6. They MUST drive +through the real ``streamablehttp_client``, the real httpx response-hook +path, and a REAL upstream MCP server (a ``FastMCP`` in-process subprocess +with a starlette middleware that programmatically returns 401/403 with +crafted ``WWW-Authenticate`` headers). + +Direct ``httpx.HTTPStatusError`` injection is FORBIDDEN here — Phase 5 +bug-1 was masked precisely by that pattern (the production code path +was structurally unreachable, but the unit-test injection bypassed the +SDK's swallow). The integration tests gate that the production path +actually receives the carrier signal end-to-end. + +The fixture upstream is built in-thread (uvicorn on its own asyncio +loop in a background thread) — same pattern as +``tests/spike_sdk_concurrency.py``. Per the orchestrator's startup-cost +note, measured at ~0.05s per fixture spin-up locally; well under the +2s threshold for default-collection inclusion. +""" + +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 + +# Quiet noisy logs during tests. +logging.getLogger("uvicorn.error").setLevel(logging.WARNING) +logging.getLogger("uvicorn.access").setLevel(logging.WARNING) +logging.getLogger("mcp").setLevel(logging.WARNING) + + +# --------------------------------------------------------------------------- +# Fixture upstream — programmable BehaviorMiddleware +# --------------------------------------------------------------------------- + + +class BehaviorMiddleware(BaseHTTPMiddleware): + """Inspects per-request behaviour state and returns 401/403 on demand. + + The behaviour is steered by a mutable ``behaviour`` dict on the + middleware instance; tests mutate it via the fixture handle. + Records every request's Authorization header for assertion. + + Behaviour semantics: + * ``"once_401"``: return 401 once, then 200 thereafter. + * ``"always_401"``: always return 401. + * ``"once_403_insufficient"``: return 403 with insufficient_scope once. + * ``"once_403_generic"``: return 403 without error param once. + * ``"once_multi_www_auth_403"``: return 403 with TWO + ``WWW-Authenticate`` headers — first ``Bearer`` challenge + carries the SAFE scopes, second carries INJECTED scopes. The + dispatcher must report only the first. + * ``"never"`` (default): pass through to the real handler. + + ``www_authenticate`` overrides the default header crafted per shape. + """ + + 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 + + # Record the Authorization header for assertion. POST is the + # tools/call request the dispatcher sends. + 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:write mail:send"', + ) + }, + ) + 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") + }, + ) + elif mode == "once_multi_www_auth_403" and not self._behaviour.get("_fired"): + self._behaviour["_fired"] = True + # Two ``WWW-Authenticate: Bearer ...`` challenges. The + # first carries ``error=insufficient_scope`` but NO + # ``scope=`` parameter; the second carries the INJECTED + # scopes the dispatcher must NOT report. The first + # challenge intentionally lacks ``scope`` because + # ``parse_www_authenticate_bearer`` uses ``setdefault`` — + # if the first challenge HAD a scope, ``setdefault`` would + # already win on first-occurrence. The vector this test + # guards is the case where a defended absence becomes a + # silent presence: a hook regression to ``get(...)`` joins + # repeated headers with ``, `` and the parser then folds + # the second challenge's scope into the first challenge's + # params dict because there is no first-occurrence to + # protect. + response = StarletteResponse("forbidden", status_code=403) + response.headers.append( + "www-authenticate", + 'Bearer realm="legit", error="insufficient_scope"', + ) + response.headers.append( + "www-authenticate", + 'Bearer error="insufficient_scope", scope="org:admin db:write"', + ) + return response + 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="phase6-target", streamable_http_path="/mcp") + + @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(): + """Boot a FastMCP fixture upstream in a background thread. + + Yields ``(url, behaviour)`` where ``behaviour`` is a mutable dict + the test mutates to steer the middleware (set ``mode`` to one of + the BehaviorMiddleware shapes). + """ + 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="phase6-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) + + +# --------------------------------------------------------------------------- +# Test 21: 401 → refresh-and-retry → success +# --------------------------------------------------------------------------- + + +def test_integration_401_refresh_and_retry_succeeds( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Real upstream returns 401 once with ``WWW-Authenticate: Bearer + error="invalid_token"``, then 200. Dispatcher carrier captures the + 401, ``force_refresh=True`` mints a new bearer (stubbed), retry + succeeds. Hard invariant 3: breaker counter remains 0. + + Drives through the REAL ``streamablehttp_client`` and a REAL + upstream subprocess (no ``httpx.HTTPStatusError`` injection). This + is the structural exit gate for Phase 6 — the equivalent unit + tests CANNOT prove the production wiring works because the SDK + swallows the underlying exception. + """ + url, behaviour = upstream + behaviour["mode"] = "once_401" + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + # Override URL to point at the local upstream (loopback http:// is + # exempt from the URL-validator). + _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)) + + 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.call_tool_sync( + "mcp__pool-srv__echo", {"payload": "hi"}, user_id="user-1", timeout=15 + ) + + assert "echoed:hi" in result + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + # Server saw at least 2 POSTs to /mcp (initial + retry). + post_headers = behaviour.get("post_auth_headers", []) + assert len(post_headers) >= 2, f"expected >=2 POSTs; got {len(post_headers)}" + # Retry carries a different bearer than the initial. + initial = post_headers[0] + retry = post_headers[1] + assert initial != retry, ( + "retry attached the same bearer as the initial; the dispatcher " + "did not pick up the refreshed token." + ) + # Pool entry has a session after the successful retry. + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.session is not None + + +# --------------------------------------------------------------------------- +# Test 22: 401 + refresh failure → mcp_consent_required +# --------------------------------------------------------------------------- + + +def test_integration_401_with_refresh_failure_emits_consent_required( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + 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)) + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="refresh_failed") + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.call_tool_sync( + "mcp__pool-srv__echo", {"payload": "x"}, 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 + + +# --------------------------------------------------------------------------- +# Test 23: 403 + insufficient_scope → mcp_insufficient_scope with parsed scopes +# --------------------------------------------------------------------------- + + +def test_integration_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:write mail:send"' + ) + + 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)) + + 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.call_tool_sync( + "mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15 + ) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_insufficient_scope" + assert payload["error"]["scopes_required"] == ["files:write", "mail:send"] + # No retry — exactly ONE POST attempted before the structured error. + 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 + + +# --------------------------------------------------------------------------- +# Test 24: 403 without insufficient_scope → generic forbidden +# --------------------------------------------------------------------------- + + +def test_integration_403_no_insufficient_scope_emits_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)) + + 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.call_tool_sync( + "mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15 + ) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_tool_call_forbidden" + assert "scopes_required" not in payload["error"] + post_headers = behaviour.get("post_auth_headers", []) + assert len(post_headers) == 1, ( + f"403 must NOT trigger a retry; observed {len(post_headers)} POSTs" + ) + + +# --------------------------------------------------------------------------- +# sec-1: multi-WWW-Authenticate header injection — only the FIRST +# Bearer challenge feeds the structured-error / audit emission. +# --------------------------------------------------------------------------- + + +def test_integration_403_multi_www_authenticate_drops_injected_scopes( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Upstream returns a 403 with TWO ``WWW-Authenticate: Bearer ...`` + challenges. The first carries ``error=insufficient_scope`` but NO + ``scope=`` parameter; the second carries INJECTED scopes + (``["org:admin", "db:write"]``). The dispatcher must report + ``scopes_required == []`` — derived from the first challenge alone + — never the second challenge's injected scopes. + + Two layers of defence cooperate (either alone neutralises the + vector; both run together so a regression in one cannot silently + re-open it): + + 1. ``_make_capturing_http_factory._hook`` reads + ``response.headers.get_list("www-authenticate")[0]`` rather than + ``response.headers.get(...)`` — the latter joins repeated + headers with ``", "`` which the RFC 7235 tokenizer would + otherwise consume as a continuation of the first challenge. + 2. ``parse_www_authenticate_bearer`` stops at the first ``Bearer`` + challenge boundary even if the input was already joined, so a + hook regression to ``get(...)`` would NOT re-open the vector. + + The first challenge intentionally lacks ``scope=`` — the parser + uses ``setdefault`` so a first-occurrence ``scope`` would already + win and mask a single-layer regression. The undefended-absence + case is what proves both layers actually do their job. + + Negative-test (CRITICAL — Phase 5 lesson): verified by reverting + the hook to ``response.headers.get("www-authenticate")`` AND + removing the ``_looks_like_bearer_challenge_start`` guard in + ``parse_www_authenticate_bearer``. The test then fails because + ``scopes_required`` becomes ``["org:admin", "db:write"]`` — the + injected scopes from the second challenge silently fold into the + first challenge's params dict via httpx's comma-joined header + value (the absence of a first-occurrence scope means nothing + blocks the fold). + """ + url, behaviour = upstream + behaviour["mode"] = "once_multi_www_auth_403" + + 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)) + + 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.call_tool_sync( + "mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15 + ) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_insufficient_scope", ( + f"expected mcp_insufficient_scope; got {payload!r}" + ) + # ``scopes_required`` derives from the FIRST challenge alone, which + # carries no ``scope=`` parameter. The injected second challenge + # MUST NOT appear here. + assert payload["error"]["scopes_required"] == [], ( + "Multi-header injection slipped through: dispatcher reported " + "scopes from the SECOND Bearer challenge. Got " + f"{payload['error']['scopes_required']!r}; expected []." + ) + + +# --------------------------------------------------------------------------- +# Test 25: 401 retry ceiling — never recurse +# --------------------------------------------------------------------------- + + +def test_integration_401_retry_ceiling( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Upstream always returns 401; refresh stub keeps minting tokens. + After exactly ONE retry, dispatcher emits ``mcp_consent_required``. + """ + 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)) + + refresh_count = 0 + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + nonlocal refresh_count + if kwargs.get("force_refresh"): + refresh_count += 1 + return TokenLookupResult(kind="token", token=f"refreshed-{refresh_count}") + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.call_tool_sync( + "mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15 + ) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_consent_required" + # Exactly ONE refresh round-trip. + assert refresh_count == 1, f"expected exactly 1 refresh round-trip; got {refresh_count}" + # Server saw EXACTLY 2 POSTs (initial + 1 retry). + post_headers = behaviour.get("post_auth_headers", []) + assert len(post_headers) == 2, ( + f"expected exactly 2 POSTs (initial + 1 retry); got {len(post_headers)}" + ) + + +# --------------------------------------------------------------------------- +# Test 26: breaker unaffected by repeated auth failures (slow — 50 cycles) +# --------------------------------------------------------------------------- + + +def test_integration_breaker_unaffected_by_auth_failures( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """50 sequential dispatches all hit 401 with refresh-failed → 50 + cycles of ``mcp_consent_required``. ``_consecutive_failures`` MUST + stay at 0 throughout (hard invariant 3 verified end-to-end). + """ + 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)) + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="refresh_failed") + return TokenLookupResult(kind="token", token="access-aaa") + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + for _ in range(50): + result = mgr.call_tool_sync( + "mcp__pool-srv__echo", {"payload": "x"}, 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 + + +# --------------------------------------------------------------------------- +# Test 27: static path unaffected by Phase 6 changes +# --------------------------------------------------------------------------- + + +def test_integration_static_path_unaffected( + upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend +) -> None: + """Static-path connect against an unauthed upstream succeeds without + going through the capturing factory. This is the integration-level + mirror of ``test_reconnect_preserves_static_state_identity``. + + Drives the static path against the same fixture upstream (with + ``behaviour={}`` so middleware passes through) — confirms the + static path's session lifecycle is byte-identical even when the + pool path's auth introspection is wired up. + """ + url, _behaviour = upstream + # No mode → middleware passes through to FastMCP. + + mgr, loop, _ = running_loop_mgr + + # Manually configure mgr with a static-path server pointing at the + # fixture upstream. Use _connect_one (not the pool path). + cfg = {"type": "streamable-http", "url": url} + + async def _connect_static() -> None: + await mgr._connect_one("static-srv", cfg) + + fut = asyncio.run_coroutine_threadsafe(_connect_static(), loop) + fut.result(timeout=15) + + state_before = mgr._static_servers.get("static-srv") + assert state_before is not None + assert state_before.session is not None + # Snapshot identity. + state_id_before = id(state_before) + session_before = state_before.session + + # Reconnect — the canonical regression check is that the + # StaticServerState object identity is preserved. + fut = asyncio.run_coroutine_threadsafe(_connect_static(), loop) + fut.result(timeout=15) + + state_after = mgr._static_servers.get("static-srv") + assert state_after is not None + assert id(state_after) == state_id_before, ( + "Static path StaticServerState identity changed across reconnect; " + "hard invariant 1 violated." + ) + assert state_after.session is not None + assert state_after.session is not session_before, ( + "Reconnect did not actually replace the session" + ) diff --git a/tests/test_mcp_pool_auth_introspection.py b/tests/test_mcp_pool_auth_introspection.py new file mode 100644 index 00000000..351b4952 --- /dev/null +++ b/tests/test_mcp_pool_auth_introspection.py @@ -0,0 +1,1291 @@ +"""Phase 6 unit tests — pool dispatch auth introspection (carrier + retry). + +Covers the ``_classify_failure`` split into ``auth_401`` / ``auth_403``, +the response-hook carrier reset / leak guarantees, the +``force_refresh=True`` semantics in +:func:`get_user_access_token_classified`, the dispatcher's +refresh-and-retry-once policy on 401, and the +``mcp_insufficient_scope`` emission on 403. Parser-helper coverage lives +in ``tests/test_mcp_http_parsers.py``. + +Negative-test verifications (run + revert + run + restore + restore): +several tests note explicit "verified by reverting [production line] to +[no-op]" lines. This bakes the Phase 5 fix-up workflow into the test +authoring discipline so future readers can re-verify the assertions +hold for the right reason. + +Direct ``httpx.HTTPStatusError`` injection appears ONLY in the +classification tests for defense-in-depth coverage of the non-SDK +refresh path. The dispatcher-asserting tests drive the carrier through +the production path. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import threading +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from tests.conftest import make_mcp_token_cipher +from turnstone.core.mcp_client import ( + MCPClientManager, + _AuthCapture, + _make_capturing_http_factory, +) +from turnstone.core.mcp_crypto import MCPTokenStore +from turnstone.core.storage._sqlite import SQLiteBackend + +# --------------------------------------------------------------------------- +# Fixtures and helpers (mirror conventions in test_mcp_user_pool.py) +# --------------------------------------------------------------------------- + + +@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(): + """Background mcp-loop fixture mirroring test_mcp_user_pool.py.""" + cfg: dict[str, Any] = {} + mgr = MCPClientManager(cfg) + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True, name="mcp-pool-test-loop") + thread.start() + mgr._loop = loop + try: + yield mgr, loop, thread + finally: + + async def _drain(m: MCPClientManager) -> None: + task = m._user_pool_eviction_task + if task is not None: + task.cancel() + with contextlib.suppress(BaseException): + await task + m._user_pool_eviction_task = None + + with contextlib.suppress(Exception): + asyncio.run_coroutine_threadsafe(_drain(mgr), loop).result(timeout=2) + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=2) + + +def _run_on_loop(loop: asyncio.AbstractEventLoop, coro: Any) -> Any: + fut = asyncio.run_coroutine_threadsafe(coro, loop) + return fut.result(timeout=5) + + +# --------------------------------------------------------------------------- +# _classify_failure with capture vs legacy fallback +# --------------------------------------------------------------------------- + + +class TestClassifyFailureWithCapture: + def test_classify_failure_auth_401_with_capture(self) -> None: + mgr = MCPClientManager({}) + capture = _AuthCapture(status=401) + # Exception type doesn't matter — capture wins. + assert mgr._classify_failure(RuntimeError("any"), capture=capture) == "auth_401" + + def test_classify_failure_auth_403_with_capture(self) -> None: + mgr = MCPClientManager({}) + capture = _AuthCapture(status=403) + assert mgr._classify_failure(RuntimeError("any"), capture=capture) == "auth_403" + + def test_classify_failure_no_capture_falls_through_to_legacy(self) -> None: + """Defense-in-depth: ``_refresh_and_persist`` raises ``HTTPStatusError`` + directly without going through the SDK swallow, so the legacy branch + must still classify.""" + mgr = MCPClientManager({}) + req = httpx.Request("POST", "https://mcp.example.com/sse") + for status, label in ((401, "auth_401"), (403, "auth_403")): + resp = httpx.Response(status, request=req) + exc = httpx.HTTPStatusError("err", request=req, response=resp) + assert mgr._classify_failure(exc, capture=None) == label + + def test_classify_failure_capture_with_unrelated_status_falls_through(self) -> None: + """Carrier with status=500 (not 401/403) doesn't classify as auth.""" + mgr = MCPClientManager({}) + capture = _AuthCapture(status=500) + # The carrier's status isn't 401 or 403, and the exception is generic. + assert mgr._classify_failure(ValueError("nope"), capture=capture) == "other" + + +# --------------------------------------------------------------------------- +# Tests 9-10: carrier per-dispatch isolation + hook only fires on 4xx +# --------------------------------------------------------------------------- + + +class TestCarrierLifecycle: + def test_capture_resets_per_dispatch(self, running_loop_mgr, storage: SQLiteBackend) -> None: + """Each ``_dispatch_pool`` allocates a fresh ``_AuthCapture`` so a + prior dispatch's 401 cannot leak into the next dispatch. + + The autouse ``_install_capture_intercept`` fixture stashes the + per-dispatch carrier on the manager. Two consecutive dispatches + must observe DIFFERENT carrier identities — if the dispatcher + reused one capture, the second dispatch would see the same + ``id()``. + + Verified by reverting ``_dispatch_pool`` to allocate + ``_AuthCapture`` once (e.g., set it on ``self._shared_capture`` + in ``__init__`` and reuse) and confirming this test fails + because the two dispatches observe the same carrier identity. + """ + from unittest.mock import patch + + from turnstone.core.mcp_oauth import TokenLookupResult + + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + + observed_capture_ids: list[int] = [] + + async def _call_tool(name: str, args: dict[str, Any]) -> Any: + cap = getattr(mgr, "_test_active_capture", None) + if cap is not None: + observed_capture_ids.append(id(cap)) + content = MagicMock() + content.text = "ok" + res = MagicMock() + res.content = [content] + res.isError = False + return res + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-1", "pool-srv")) + sess = MagicMock() + sess.call_tool = _call_tool + entry.session = sess + + _run_on_loop(loop, _seed()) + + 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, + ): + mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5) + mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5) + + assert len(observed_capture_ids) == 2 + assert observed_capture_ids[0] != observed_capture_ids[1], ( + "Both dispatches saw the same _AuthCapture object; the dispatcher " + "is reusing the carrier across calls instead of allocating fresh." + ) + + def test_capture_dataclass_isolated_when_constructed_separately(self) -> None: + """Sanity check: independent ``_AuthCapture`` instances do not + share mutable state. Pure dataclass shape verification.""" + cap_a = _AuthCapture() + cap_b = _AuthCapture() + cap_a.status = 401 + cap_a.www_authenticate = 'Bearer error="invalid_token"' + assert cap_b.status is None + assert cap_b.www_authenticate is None + cap_b.status = 403 + assert cap_a.status == 401 + + @pytest.mark.asyncio + async def test_response_hook_captures_4xx_only(self) -> None: + """Hook records on 401 and 403 only; 200/201/202/500 are ignored. + + Verified by removing the ``status in (401, 403)`` guard in + ``_make_capturing_http_factory._hook`` and confirming that + ``status=200`` populates the carrier (test fails because we + assert ``capture.status is None`` for 200). + + The hook is ``async`` (httpx invokes ``await hook(response)``), + so the test awaits it directly via the ``event_hooks`` slot. + """ + capture = _AuthCapture() + factory = _make_capturing_http_factory(capture) + client = factory() + try: + hooks = client.event_hooks["response"] + assert len(hooks) == 1 + hook = hooks[0] + + req = httpx.Request("POST", "https://mcp.example.com/") + for ignored_status in (200, 201, 202, 500): + resp = httpx.Response( + ignored_status, + request=req, + headers={"www-authenticate": "Bearer should-be-ignored"}, + ) + await hook(resp) + assert capture.status is None, ( + f"hook recorded on {ignored_status}; expected only 401/403" + ) + + for tracked_status in (401, 403): + capture.status = None + capture.www_authenticate = None + resp = httpx.Response( + tracked_status, + request=req, + headers={"www-authenticate": f'Bearer error="x{tracked_status}"'}, + ) + await hook(resp) + assert capture.status == tracked_status + assert capture.www_authenticate == f'Bearer error="x{tracked_status}"' + finally: + await client.aclose() + + +# --------------------------------------------------------------------------- +# Test 11-12: force_refresh=True semantics in get_user_access_token_classified +# --------------------------------------------------------------------------- + + +class TestForceRefresh: + @pytest.mark.asyncio + async def test_force_refresh_bypasses_token_freshness_check( + self, storage: SQLiteBackend + ) -> None: + """``force_refresh=True`` MUST go through the lock + AS round-trip + even when ``_token_needs_refresh`` returns False. + + Verified by replacing the ``force_refresh`` parameter with a + no-op default (i.e., dropping the ``not force_refresh and`` + guard so the fast path always wins) and confirming this test + fails — the call returns the original cached token instead of + the refreshed value. + """ + from unittest.mock import patch + + from turnstone.core.mcp_oauth import get_user_access_token_classified + + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="srv-oauth") + state = _make_app_state(storage, cipher=cipher) + # Token expires 1 hour from now — _token_needs_refresh returns False. + _seed_user_token( + storage, cipher, user_id="user-1", server_name="srv-oauth", expires_in_seconds=3600 + ) + + async def _fake_refresh_and_persist(**_kwargs: Any) -> tuple[str, str | None, str | None]: + return ("refreshed-token", "rotated-rrr", None) + + with patch( + "turnstone.core.mcp_oauth._refresh_and_persist", + side_effect=_fake_refresh_and_persist, + ): + # Without force_refresh: returns the cached token. + cached = await get_user_access_token_classified( + app_state=state, user_id="user-1", server_name="srv-oauth" + ) + assert cached.kind == "token" + assert cached.token == "access-aaa" + # With force_refresh: forces the AS round-trip. + forced = await get_user_access_token_classified( + app_state=state, + user_id="user-1", + server_name="srv-oauth", + force_refresh=True, + ) + assert forced.kind == "token" + assert forced.token == "refreshed-token" + + @pytest.mark.asyncio + async def test_force_refresh_collapses_concurrent_callers_via_lock( + self, storage: SQLiteBackend + ) -> None: + """Two concurrent ``force_refresh=True`` callers MUST collapse to + one AS round-trip via the dual-layer lock + the timestamp-comparison + guard inside the locked block. + + Verified by removing the timestamp-comparison guard inside the + ``async with lock, pg_lock:`` block (i.e., letting the second + caller refresh again because its ``not _token_needs_refresh`` + check still has ``force_refresh=True``) and confirming this + test counts 2 AS round-trips instead of 1. + """ + from unittest.mock import patch + + from turnstone.core.mcp_oauth import get_user_access_token_classified + + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="srv-oauth") + state = _make_app_state(storage, cipher=cipher) + _seed_user_token( + storage, cipher, user_id="user-1", server_name="srv-oauth", expires_in_seconds=3600 + ) + + call_count = 0 + call_count_lock = threading.Lock() + + async def _fake_refresh_and_persist(**kwargs: Any) -> tuple[str, str | None, str | None]: + nonlocal call_count + with call_count_lock: + call_count += 1 + # Yield once so the second concurrent caller can hit the + # lock while the first is still inside. + await asyncio.sleep(0.05) + # Mirror what the real ``_refresh_and_persist`` does so the + # timestamp-comparison guard inside the locked block sees a + # fresh ``last_refreshed`` and the second caller short-circuits. + token_store: MCPTokenStore = kwargs["token_store"] + user_id = kwargs["user_id"] + server_name = kwargs["server_name"] + future_expires = (datetime.now(UTC) + timedelta(seconds=3600)).strftime( + "%Y-%m-%dT%H:%M:%S" + ) + await asyncio.to_thread( + token_store.update_user_token_after_refresh, + user_id, + server_name, + access_token="refreshed-token", + refresh_token="rotated-rrr", + expires_at=future_expires, + ) + return ("refreshed-token", "rotated-rrr", future_expires) + + with patch( + "turnstone.core.mcp_oauth._refresh_and_persist", + side_effect=_fake_refresh_and_persist, + ): + results = await asyncio.gather( + get_user_access_token_classified( + app_state=state, + user_id="user-1", + server_name="srv-oauth", + force_refresh=True, + ), + get_user_access_token_classified( + app_state=state, + user_id="user-1", + server_name="srv-oauth", + force_refresh=True, + ), + ) + + for r in results: + assert r.kind == "token" + assert call_count == 1, ( + f"Expected the in-process lock + timestamp guard to collapse two " + f"concurrent force_refresh callers to one AS round-trip, but " + f"observed {call_count} round-trips." + ) + + @pytest.mark.asyncio + async def test_force_refresh_goes_through_pg_refresh_lock_path( + self, storage: SQLiteBackend + ) -> None: + """``force_refresh=True`` paths MUST still flow through the + ``_PgRefreshLock`` infrastructure (hard invariants 11-13 unchanged). + """ + from unittest.mock import patch + + from turnstone.core.mcp_oauth import get_user_access_token_classified + + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="srv-oauth") + state = _make_app_state(storage, cipher=cipher) + _seed_user_token( + storage, cipher, user_id="user-1", server_name="srv-oauth", expires_in_seconds=3600 + ) + + acquired = False + + original_acquire = storage.acquire_advisory_lock_sync + + @contextlib.contextmanager + def _tracking_acquire(key_text: str = "") -> Any: + nonlocal acquired + with original_acquire(key_text): + acquired = True + yield + + async def _fake_refresh_and_persist(**_kwargs: Any) -> tuple[str, str | None, str | None]: + return ("refreshed-token", "rotated-rrr", None) + + with ( + patch.object(storage, "acquire_advisory_lock_sync", side_effect=_tracking_acquire), + patch( + "turnstone.core.mcp_oauth._refresh_and_persist", + side_effect=_fake_refresh_and_persist, + ), + ): + result = await get_user_access_token_classified( + app_state=state, + user_id="user-1", + server_name="srv-oauth", + force_refresh=True, + ) + + assert result.kind == "token" + assert acquired, ( + "force_refresh=True bypassed the _PgRefreshLock advisory-lock path; " + "hard invariant 13 violated." + ) + + +# --------------------------------------------------------------------------- +# Test 13-17: dispatcher refresh-and-retry + 403 step-up emission +# --------------------------------------------------------------------------- + + +class TestDispatcherAuthFlows: + """The dispatcher consults the carrier and routes to refresh-and-retry + (401) or structured-error emission (403).""" + + def _wire_pool( + self, + mgr: MCPClientManager, + storage: SQLiteBackend, + cipher: Any, + ) -> SimpleNamespace: + mgr.set_storage(storage) + state = _make_app_state(storage, cipher=cipher) + mgr.set_app_state(state) + return state + + def _seed_pool_entry_with_call_tool( + self, + mgr: MCPClientManager, + loop: asyncio.AbstractEventLoop, + call_tool: Any, + ) -> None: + """Seed a pool entry with a fake session whose call_tool is *call_tool*. + + Also patches ``_connect_one_pool`` so the retry path (which the + dispatcher exercises after dropping the session on auth failure) + can re-install the same fake session without going through the + real TCP probe / SDK handshake. The patch reinstates the same + fake on every reconnect — the test's ``call_tool`` stub is + responsible for varying behaviour across calls. + """ + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-1", "pool-srv")) + sess = MagicMock() + sess.call_tool = call_tool + entry.session = sess + + _run_on_loop(loop, _seed()) + + async def _fake_connect( + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + access_token: str, + *, + auth_capture: Any = None, + ) -> Any: + entry = await self_inner._ensure_pool_entry(key) + sess = MagicMock() + sess.call_tool = call_tool + entry.session = sess + return entry + + # Install via setattr — monkeypatching the bound method on the + # instance avoids leaking into other tests that share the class. + mgr._connect_one_pool = _fake_connect.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + def test_dispatch_pool_401_refreshes_and_retries( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """First call sees carrier=401, dispatcher signals retry; the + retry's ``_dispatch_pool`` reads the token with ``force_refresh=True`` + (via ``retry_count > 0``) and the second call returns success. + + Verified by reverting ``_dispatch_pool_sync`` to remove the + ``except _PoolDispatchRetryRequested`` block (so the signal + propagates instead of triggering the retry on a fresh task) + and confirming this test fails because the retry never fires. + """ + from unittest.mock import patch + + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + self._wire_pool(mgr, storage, cipher) + + from turnstone.core.mcp_oauth import TokenLookupResult + + call_count = 0 + + # Stub force_refresh to return a fresh token. + 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") + + # The first call_tool populates the carrier with 401; the + # second call_tool returns success. + async def _call_tool(name: str, args: dict[str, Any]) -> Any: + nonlocal call_count + call_count += 1 + # Find the calling _dispatch_pool's capture by walking the + # most recently allocated entry session — the test bridges + # via the dispatch path, so we mutate the per-dispatch + # carrier directly via the patched factory's hook semantics. + # Easier: use a sentinel exception that the dispatcher's + # capture pre-populates via the hook test path. But for + # the dispatcher unit test we drive the carrier through + # an exception-injection that ALSO populates the captured + # carrier — see the test fixtures below. + if call_count == 1: + # Simulate the SDK swallow path: carrier was populated + # by the hook, but the exception that surfaces is a + # generic CONNECTION_CLOSED-shaped error. + _populate_active_capture(mgr, status=401, header='Bearer error="invalid_token"') + raise RuntimeError("upstream 401 (SDK-swallow shape)") + content = MagicMock() + content.text = "ok" + res = MagicMock() + res.content = [content] + res.isError = False + return res + + self._seed_pool_entry_with_call_tool(mgr, loop, _call_tool) + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.call_tool_sync( + "mcp__pool-srv__do_thing", + {}, + user_id="user-1", + timeout=5, + ) + + assert result == "ok" + assert call_count == 2, ( + f"Expected exactly 2 call_tool invocations (initial + retry); got {call_count}" + ) + # Auth failures must not affect the breaker (hard invariant 3). + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + def test_dispatch_pool_401_retry_with_refresh_failure_emits_consent_required( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """When the AS rejects the refresh, the dispatcher emits + ``mcp_consent_required`` (no exception, no breaker tick).""" + from unittest.mock import patch + + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + self._wire_pool(mgr, storage, cipher) + + from turnstone.core.mcp_oauth import TokenLookupResult + + async def _fake_classified( + **kwargs: Any, + ) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="refresh_failed") + return TokenLookupResult(kind="token", token="access-aaa") + + async def _call_tool(name: str, args: dict[str, Any]) -> Any: + _populate_active_capture(mgr, status=401, header='Bearer error="invalid_token"') + raise RuntimeError("upstream 401") + + self._seed_pool_entry_with_call_tool(mgr, loop, _call_tool) + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.call_tool_sync( + "mcp__pool-srv__do_thing", + {}, + user_id="user-1", + timeout=5, + ) + + 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 + + def test_dispatch_pool_401_retry_ceiling_caps_at_one( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """Even when the refresh succeeds, if the retry ALSO 401s, the + dispatcher emits ``mcp_consent_required`` rather than recursing. + + Verified by editing ``_dispatch_pool`` to relax the + ``if retry_count == 0`` guard (e.g. to ``if retry_count <= 1``) + so the auth_401 branch raises ``_PoolDispatchRetryRequested`` + even at the ceiling. ``_dispatch_pool_sync`` only catches the + signal once, so the second raise propagates back to + ``call_tool_sync`` and the test fails on the missing + ``mcp_consent_required`` payload (``json.loads`` on a + non-JSON / raised result). The point of the negative-test is + to prove the ceiling (``retry_count == 0`` guard) is what + bounds the retry loop — without it, the dispatcher would loop + the bearer-rejection cycle indefinitely. + """ + from unittest.mock import patch + + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + self._wire_pool(mgr, storage, cipher) + + from turnstone.core.mcp_oauth import TokenLookupResult + + async def _fake_classified( + **kwargs: Any, + ) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="bearer-XYZ") + + call_count = 0 + + async def _call_tool(name: str, args: dict[str, Any]) -> Any: + nonlocal call_count + call_count += 1 + _populate_active_capture(mgr, status=401, header='Bearer error="invalid_token"') + raise RuntimeError("upstream 401") + + self._seed_pool_entry_with_call_tool(mgr, loop, _call_tool) + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.call_tool_sync( + "mcp__pool-srv__do_thing", + {}, + user_id="user-1", + timeout=5, + ) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_consent_required" + # Exactly TWO calls: initial + retry. No recursion. + assert call_count == 2, ( + f"Expected exactly 2 call_tool invocations (initial + 1 retry); got {call_count}" + ) + + def test_dispatch_pool_403_emits_insufficient_scope_with_parsed_scopes( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """403 + ``error="insufficient_scope"`` emits ``mcp_insufficient_scope`` + with the parsed scope set; no retry.""" + from unittest.mock import patch + + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + self._wire_pool(mgr, storage, cipher) + + from turnstone.core.mcp_oauth import TokenLookupResult + + async def _fake_classified( + **_kwargs: Any, + ) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + call_count = 0 + + async def _call_tool(name: str, args: dict[str, Any]) -> Any: + nonlocal call_count + call_count += 1 + _populate_active_capture( + mgr, + status=403, + header='Bearer error="insufficient_scope", scope="files:write mail:send"', + ) + raise RuntimeError("upstream 403") + + self._seed_pool_entry_with_call_tool(mgr, loop, _call_tool) + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.call_tool_sync( + "mcp__pool-srv__do_thing", + {}, + user_id="user-1", + timeout=5, + ) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_insufficient_scope" + assert payload["error"]["scopes_required"] == ["files:write", "mail:send"] + assert call_count == 1, "403 must NOT trigger a retry" + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + def test_dispatch_pool_403_without_insufficient_scope_emits_generic_forbidden( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """403 without ``error="insufficient_scope"`` → generic forbidden.""" + from unittest.mock import patch + + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + self._wire_pool(mgr, storage, cipher) + + from turnstone.core.mcp_oauth import TokenLookupResult + + async def _fake_classified( + **_kwargs: Any, + ) -> TokenLookupResult: + return TokenLookupResult(kind="token", token="access-aaa") + + async def _call_tool(name: str, args: dict[str, Any]) -> Any: + _populate_active_capture(mgr, status=403, header="Bearer realm=mcp") + raise RuntimeError("upstream 403") + + self._seed_pool_entry_with_call_tool(mgr, loop, _call_tool) + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result = mgr.call_tool_sync( + "mcp__pool-srv__do_thing", + {}, + user_id="user-1", + timeout=5, + ) + + payload = json.loads(result) + assert payload["error"]["code"] == "mcp_tool_call_forbidden" + assert "scopes_required" not in payload["error"] + + +# --------------------------------------------------------------------------- +# Test 18: hard invariant 3 — auth failures never trip the breaker +# --------------------------------------------------------------------------- + + +class TestBreakerInvariant: + def test_auth_failures_do_not_trip_breaker( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """After 401-then-success and after 403-emit, the per-server breaker + counter MUST remain 0 (hard invariant 3). + + Verified by adding ``self._cb_record_failure(server_name)`` to + the ``auth_401`` and ``auth_403`` branches of ``_dispatch_pool`` + and confirming this test fails because the counter advances. + """ + from unittest.mock import patch + + from turnstone.core.mcp_oauth import TokenLookupResult + + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + + async def _fake_classified(**kwargs: Any) -> TokenLookupResult: + if kwargs.get("force_refresh"): + return TokenLookupResult(kind="token", token="bearer-refreshed") + return TokenLookupResult(kind="token", token="bearer-original") + + # Multi-stage call_tool: 401 (initial) → ok (retry) → 403 (next dispatch). + # The 403 step doesn't retry so it's the third element. + async def _call_tool_401(name: str, args: dict[str, Any]) -> Any: + _populate_active_capture(mgr, status=401, header='Bearer error="invalid_token"') + raise RuntimeError("upstream 401") + + async def _call_tool_success(name: str, args: dict[str, Any]) -> Any: + content = MagicMock() + content.text = "ok" + res = MagicMock() + res.content = [content] + res.isError = False + return res + + async def _call_tool_403(name: str, args: dict[str, Any]) -> Any: + _populate_active_capture( + mgr, status=403, header='Bearer error="insufficient_scope", scope="x:y"' + ) + raise RuntimeError("upstream 403") + + seq = iter([_call_tool_401, _call_tool_success, _call_tool_403]) + + async def _staged_call_tool(name: str, args: dict[str, Any]) -> Any: + fn = next(seq) + return await fn(name, args) + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-1", "pool-srv")) + sess = MagicMock() + sess.call_tool = _staged_call_tool + entry.session = sess + + _run_on_loop(loop, _seed()) + + # Patch _connect_one_pool to re-install the staged session on + # reconnect (the dispatcher drops the session after auth failure). + async def _fake_connect( + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + access_token: str, + *, + auth_capture: Any = None, + ) -> Any: + entry = await self_inner._ensure_pool_entry(key) + sess = MagicMock() + sess.call_tool = _staged_call_tool + entry.session = sess + return entry + + mgr._connect_one_pool = _fake_connect.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result_401_retry = mgr.call_tool_sync( + "mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5 + ) + assert result_401_retry == "ok" + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + with patch( + "turnstone.core.mcp_client.get_user_access_token_classified", + side_effect=_fake_classified, + ): + result_403 = mgr.call_tool_sync( + "mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5 + ) + payload = json.loads(result_403) + assert payload["error"]["code"] == "mcp_insufficient_scope" + # STILL zero after the 403 cycle. + assert mgr._consecutive_failures.get("pool-srv", 0) == 0 + + +# --------------------------------------------------------------------------- +# Test 19: hard invariant 1 — static path does NOT receive the capture factory +# --------------------------------------------------------------------------- + + +class TestStaticPathUnchanged: + def test_static_path_does_not_pass_capturing_factory(self) -> None: + """``_connect_one`` (static path) must NEVER pass + ``httpx_client_factory`` to ``streamablehttp_client``. + + Hard invariant 1: any change to the static-path connect plumbing + is potentially breaking. Verified by source inspection rather + than runtime mocking — the static path's call site is the only + non-test ``streamablehttp_client(...)`` invocation that must + omit the factory parameter. + """ + import inspect + + from turnstone.core import mcp_client + + source = inspect.getsource(mcp_client.MCPClientManager._connect_one) + + # The static path's streamablehttp_client invocation should NOT + # mention ``httpx_client_factory``. Pool path keeps it. + # Find the streamablehttp_client(...) call inside _connect_one. + assert "streamablehttp_client" in source + # The call site in _connect_one is bare — no factory keyword. + # We grep by line: the factory keyword must not appear in the + # static-path source. + for line in source.splitlines(): + if "httpx_client_factory" in line: + pytest.fail( + "_connect_one (static path) passes httpx_client_factory to " + "streamablehttp_client; hard invariant 1 violated." + ) + + +# --------------------------------------------------------------------------- +# Test 20: hold open_lock across call_tool (Phase 6 multiplex revert) +# --------------------------------------------------------------------------- + + +class TestOpenLockHeldAcrossCallTool: + def test_pool_dispatch_holds_open_lock_across_call_tool( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """Phase 6 reverts Phase 5 perf-1 for the auth-aware path: hold + ``open_lock`` across ``call_tool`` so concurrent dispatches can't + race on the response-hook carrier. + + This test mirrors + ``test_pool_concurrent_dispatch_to_same_user_server_is_serialized`` + in ``test_mcp_user_pool.py`` but pins the assertion to the carrier + isolation rationale specifically. + + Verified by reverting ``_dispatch_pool_with_entry`` to release + ``open_lock`` before ``call_tool`` (move the + ``in_flight += 1`` / ``call_tool`` / decrement out of the + ``async with`` body) and confirming this test observes + ``observed_max_concurrency == 2``. + """ + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + + observed_max = 0 + in_flight = 0 + in_flight_lock = threading.Lock() + + async def _call_tool(name: str, args: dict[str, Any]) -> Any: + nonlocal observed_max, in_flight + with in_flight_lock: + in_flight += 1 + observed_max = max(observed_max, in_flight) + try: + await asyncio.sleep(0.1) + content = MagicMock() + content.text = "ok" + res = MagicMock() + res.content = [content] + res.isError = False + return res + finally: + with in_flight_lock: + in_flight -= 1 + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-1", "pool-srv")) + sess = MagicMock() + sess.call_tool = _call_tool + entry.session = sess + + _run_on_loop(loop, _seed()) + + results: list[str] = [] + + def _dispatch() -> None: + results.append( + mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5) + ) + + t1 = threading.Thread(target=_dispatch) + t2 = threading.Thread(target=_dispatch) + t1.start() + t2.start() + t1.join(timeout=5) + t2.join(timeout=5) + + assert results == ["ok", "ok"] + assert observed_max == 1, ( + "open_lock was released before call_tool, allowing concurrent " + "carrier crosstalk between same-(user, server) dispatches." + ) + + +class TestCrossTaskRetryIsolation: + """The auth_401 retry MUST run on a fresh asyncio.Task. + + Phase 6's cross-task hop (``_dispatch_pool_sync`` schedules the + retry via a SECOND ``asyncio.run_coroutine_threadsafe`` call so the + retry's ``streamablehttp_client`` TaskGroup gets a clean anyio + cancel-scope state, free of the prior connect's teardown + pollution). An in-task retry inherits the prior anyio scope across + ``task.uncancel()`` and surfaces ``CancelledError`` from inside the + retry's own ``streamablehttp_client`` scope (verified empirically + by the prior implementer; cross-task is the architectural fix). + """ + + def test_dispatch_pool_sync_retries_on_fresh_task( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """Mock ``_dispatch_pool`` to raise ``_PoolDispatchRetryRequested`` + once, then succeed. Assert the two invocations ran on different + ``asyncio.Task`` instances (proving the retry got a fresh task). + + Verified by reverting ``_dispatch_pool_sync`` to remove the + ``except _PoolDispatchRetryRequested`` block (so the signal + propagates instead of being caught and re-issued via a second + ``_run_pool_dispatch_attempt``) and confirming the test fails — + the retry never fires, so only one ``_dispatch_pool`` invocation + is recorded and the expected success result never returns. + """ + from turnstone.core.mcp_client import _PoolDispatchRetryRequested + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + + # Capture each invocation's task identity. + invocations: list[dict[str, Any]] = [] + + async def _fake_dispatch_pool(**kwargs: Any) -> str: + task = asyncio.current_task() + invocations.append( + { + "retry_count": kwargs.get("retry_count"), + "task_id": id(task), + "task_name": task.get_name() if task else None, + } + ) + if kwargs.get("retry_count") == 0: + raise _PoolDispatchRetryRequested + return "RETRY_OK" + + mgr._dispatch_pool = _fake_dispatch_pool # type: ignore[method-assign] + + result = mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5) + + assert result == "RETRY_OK" + assert len(invocations) == 2, ( + f"expected 2 _dispatch_pool invocations (initial + retry); got {len(invocations)}" + ) + assert invocations[0]["retry_count"] == 0 + assert invocations[1]["retry_count"] == 1 + assert invocations[0]["task_id"] != invocations[1]["task_id"], ( + "retry ran on the SAME asyncio.Task as the initial attempt; " + "cross-task isolation broken — the retry's anyio cancel-scope " + "would inherit teardown state from the prior connect." + ) + + +class TestRetryTimeoutBudget: + """The retry's ``future.result(timeout=...)`` window MUST be reduced + by however long the first attempt consumed before raising + ``_PoolDispatchRetryRequested``. + + The earlier ``for retry_count in (0, 1):`` loop passed the full + ``timeout`` to BOTH ``future.result`` calls. A first attempt that + consumed almost the entire budget could therefore double the + caller-observed timeout window (initial budget + retry budget) when + the retry stalled. The wall-clock budget collapses both attempts + into one ``timeout``-bounded window. + + Negative-test verification: reverting ``_dispatch_pool_sync`` to + pass ``timeout`` (instead of ``remaining``) on the second attempt + makes ``test_retry_attempt_timeout_reduced_by_first_attempt_duration`` + fail — the captured second-attempt timeout would equal the original. + """ + + def test_retry_attempt_timeout_reduced_by_first_attempt_duration( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """First attempt sleeps ~1s before raising the retry signal; the + second attempt's ``future.result(timeout=...)`` MUST receive a + value strictly less than the original ``timeout``. + """ + from turnstone.core.mcp_client import _PoolDispatchRetryRequested + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + + observed_timeouts: list[int] = [] + + async def _slow_then_succeed(**kwargs: Any) -> str: + if kwargs.get("retry_count") == 0: + # Burn ~1s of wall-clock on the first attempt before + # signalling the retry. + await asyncio.sleep(1.0) + raise _PoolDispatchRetryRequested + return "RETRY_OK" + + mgr._dispatch_pool = _slow_then_succeed # type: ignore[method-assign] + + original_run = mgr._run_pool_dispatch_attempt + + def _spy_run(**kwargs: Any) -> str: + observed_timeouts.append(kwargs["timeout"]) + return original_run(**kwargs) + + mgr._run_pool_dispatch_attempt = _spy_run # type: ignore[method-assign] + + result = mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=10) + + assert result == "RETRY_OK" + assert len(observed_timeouts) == 2, ( + f"expected 2 attempts (initial + retry); got {len(observed_timeouts)} " + f"timeouts={observed_timeouts!r}" + ) + # Initial attempt sees the full budget. + assert observed_timeouts[0] == 10 + # Retry's window is reduced by the first attempt's ~1s sleep. + assert observed_timeouts[1] < 10, ( + "retry attempt received the full original timeout instead of " + f"the wall-clock remainder; observed_timeouts={observed_timeouts!r}" + ) + + def test_timeout_message_reports_original_budget( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """``TimeoutError`` message uses the caller's original budget, + even when the retry's trimmed window is what actually expired. + """ + from turnstone.core.mcp_client import _PoolDispatchRetryRequested + + mgr, _loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher) + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + + async def _retry_then_hang(**kwargs: Any) -> str: + if kwargs.get("retry_count") == 0: + raise _PoolDispatchRetryRequested + # Hang past whatever timeout we get. + await asyncio.sleep(60) + return "never" + + mgr._dispatch_pool = _retry_then_hang # type: ignore[method-assign] + + with pytest.raises(TimeoutError) as exc_info: + mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=2) + assert "timed out after 2s" in str(exc_info.value) + + +# --------------------------------------------------------------------------- +# Test helpers — bridge the test stand-in for the real httpx event hook. +# +# The dispatcher allocates a fresh ``_AuthCapture`` per dispatch and threads +# it into ``_dispatch_pool_with_entry`` via the ``auth_capture`` kwarg. To +# unit-test the dispatcher's reaction without a real upstream, the fake +# ``call_tool`` populates the carrier directly via this helper, simulating +# what the production response hook would do. +# --------------------------------------------------------------------------- + + +def _populate_active_capture(mgr: MCPClientManager, *, status: int, header: str) -> None: + """Mutate the per-dispatch ``_AuthCapture`` stashed on the manager + by the autouse ``_install_capture_intercept`` fixture. + + Used by stub ``call_tool`` implementations to simulate what the + production response hook would record on a 4xx upstream response. + Raises ``RuntimeError`` if invoked outside the autouse fixture's + intercept window — the carrier is per-dispatch, not per-test, so + it only exists while a dispatch is in flight. + """ + capture = getattr(mgr, "_test_active_capture", None) + if capture is None: + # Allocate a fresh sentinel that the dispatcher's auth-classification + # branch will see. The fake call_tool mutates THIS object; the + # dispatcher consults its own per-dispatch capture, but since we + # also patch _classify_failure to consult the test sentinel... + raise RuntimeError( + "test setup error: _populate_active_capture called before " + "_install_capture_intercept; the dispatcher's per-dispatch " + "_AuthCapture isn't observable from this stub." + ) + capture.status = status + capture.www_authenticate = header + + +@pytest.fixture(autouse=True) +def _install_capture_intercept(monkeypatch: pytest.MonkeyPatch) -> None: + """Wrap ``_dispatch_pool_with_entry`` so the test's fake ``call_tool`` + can populate the per-dispatch ``_AuthCapture`` via ``mgr._test_active_capture``. + + The wrap is a no-op for tests that don't call ``_populate_active_capture``; + they simply never read the attribute. Dispatcher-asserting tests rely + on this so the fake call_tool stub can mutate the same carrier object + the dispatcher inspects after raising. + """ + from turnstone.core import mcp_client as mcp_client_mod + + original = mcp_client_mod.MCPClientManager._dispatch_pool_with_entry + + async def _wrapped(self: MCPClientManager, **kwargs: Any) -> str: + # Stash the carrier so the test's call_tool stub can populate it. + self._test_active_capture = kwargs.get("auth_capture") # type: ignore[attr-defined] + try: + return await original(self, **kwargs) + finally: + self._test_active_capture = None # type: ignore[attr-defined] + + monkeypatch.setattr( + mcp_client_mod.MCPClientManager, + "_dispatch_pool_with_entry", + _wrapped, + ) + + +# Suppress unused-import warning for AsyncMock. +_ = AsyncMock diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index 496c2ddc..ffeef835 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -517,23 +517,29 @@ class TestClassifyFailure: mgr = MCPClientManager({}) assert mgr._classify_failure(ValueError("nope")) == "other" - def test_http_401_classified_as_auth(self) -> None: + def test_http_401_classified_as_auth_401(self) -> None: + """Defense-in-depth: ``HTTPStatusError`` classification still works + even though Phase 6 normally consults the carrier instead. + + Phase 6 split ``"auth"`` into ``"auth_401"`` / ``"auth_403"`` + so the dispatcher can refresh-and-retry only on 401. + """ import httpx mgr = MCPClientManager({}) req = httpx.Request("POST", "https://mcp.example.com/sse") resp = httpx.Response(401, request=req) exc = httpx.HTTPStatusError("unauthorized", request=req, response=resp) - assert mgr._classify_failure(exc) == "auth" + assert mgr._classify_failure(exc) == "auth_401" - def test_http_403_classified_as_auth(self) -> None: + def test_http_403_classified_as_auth_403(self) -> None: import httpx mgr = MCPClientManager({}) req = httpx.Request("POST", "https://mcp.example.com/sse") resp = httpx.Response(403, request=req) exc = httpx.HTTPStatusError("forbidden", request=req, response=resp) - assert mgr._classify_failure(exc) == "auth" + assert mgr._classify_failure(exc) == "auth_403" def test_http_500_not_classified_as_auth(self) -> None: import httpx @@ -577,37 +583,6 @@ class TestDispatchFailureWiring: _run_on_loop(loop, _seed()) - def test_dispatch_pool_401_does_not_trip_breaker( - self, running_loop_mgr, storage: SQLiteBackend - ) -> None: - import httpx - - mgr, loop, _ = running_loop_mgr - cipher = make_mcp_token_cipher() - _seed_oauth_server(storage, name="pool-srv") - _seed_user_token(storage, cipher) - self._wire_pool(mgr, storage, cipher) - - req = httpx.Request("POST", "https://mcp.example.com/sse") - resp = httpx.Response(401, request=req) - self._seed_connected_session( - mgr, loop, httpx.HTTPStatusError("unauthorized", request=req, response=resp) - ) - - with pytest.raises(httpx.HTTPStatusError): - mgr.call_tool_sync( - "mcp__pool-srv__do_thing", - {}, - user_id="user-1", - timeout=5, - ) - # Auth failure must NOT tick the per-server breaker. - assert mgr._consecutive_failures.get("pool-srv", 0) == 0 - # But the pool entry's session must be cleared so the next call - # re-authenticates with a fresh bearer. - entry = mgr._user_pool_entries[("user-1", "pool-srv")] - assert entry.session is None - def test_dispatch_pool_transport_failure_trips_breaker( self, running_loop_mgr, storage: SQLiteBackend ) -> None: @@ -811,11 +786,27 @@ class TestLruInterlock: class TestConcurrentDispatch: - def test_pool_concurrent_dispatch_to_same_user_server_is_concurrent( + def test_pool_concurrent_dispatch_to_same_user_server_is_serialized( self, running_loop_mgr, storage: SQLiteBackend ) -> None: - """Two tool calls on the SAME (user, server) must overlap in flight, - not serialize on ``open_lock``.""" + """Phase 6: two tool calls on the SAME (user, server) MUST serialize + on ``open_lock`` so the auth-introspection carrier never crosses + between concurrent dispatches. + + Phase 5 perf-1 released ``open_lock`` before ``call_tool`` so two + concurrent same-key calls multiplexed on a shared + ``ClientSession``. Phase 6 reverts that for the auth-aware path + because the per-dispatch ``_AuthCapture`` is keyed off the + ``httpx.AsyncClient`` event hook — releasing the lock would let + a concurrent dispatch overwrite the carrier mid-flight, + attributing one caller's 401 to another (a security bug). + + Verified by reverting ``_dispatch_pool_with_entry`` to the + Phase 5 shape (release ``open_lock`` before ``call_tool`` — + i.e. move the ``in_flight += 1`` / ``call_tool`` / decrement + block out of the ``async with`` body) and confirming this test + observes ``max_concurrency == 2``. + """ mgr, loop, _ = running_loop_mgr cipher = make_mcp_token_cipher() _seed_oauth_server(storage, name="pool-srv") @@ -834,7 +825,8 @@ class TestConcurrentDispatch: in_flight += 1 observed_max_concurrency = max(observed_max_concurrency, in_flight) try: - # Hold a moment so concurrent calls overlap. + # Hold a moment so concurrent calls would overlap if + # they weren't serialized on ``open_lock``. await asyncio.sleep(0.1) content = MagicMock() content.text = "ok" @@ -880,9 +872,9 @@ class TestConcurrentDispatch: assert errors == [] assert results == ["ok", "ok"] - # Both dispatches were in flight at the same time — open_lock did - # not serialize them. - assert observed_max_concurrency == 2 + # ``open_lock`` held across ``call_tool`` — the second dispatch + # waits for the first to release before entering call_tool. + assert observed_max_concurrency == 1 # --------------------------------------------------------------------------- diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index be747fdd..6c35e932 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -35,11 +35,21 @@ import mcp.types as mcp_types from mcp import ClientSession, McpError, StdioServerParameters from mcp.client.stdio import stdio_client from mcp.client.streamable_http import streamablehttp_client +from mcp.shared._httpx_utils import ( + MCP_DEFAULT_SSE_READ_TIMEOUT, + MCP_DEFAULT_TIMEOUT, + McpHttpClientFactory, +) from turnstone.core.config import load_config from turnstone.core.log import get_logger +from turnstone.core.mcp_http_parsers import ( + parse_www_authenticate_error, + parse_www_authenticate_scope, +) from turnstone.core.mcp_oauth import ( TokenLookupResult, + emit_insufficient_scope_audit, get_user_access_token_classified, ) @@ -81,6 +91,116 @@ def _validate_oauth_user_url(url: str) -> None: ) +# --------------------------------------------------------------------------- +# Pool dispatch auth introspection (response-hook carrier) +# --------------------------------------------------------------------------- +# +# The MCP SDK's ``streamable_http`` transport raises +# :class:`httpx.HTTPStatusError` inside ``_handle_post_request`` and the +# enclosing ``post_writer`` swallows it (``mcp/client/streamable_http.py`` +# logger.exception path). The dispatcher then sees only +# ``McpError(CONNECTION_CLOSED)`` with no status / headers preserved. +# To recover the upstream 401/403 we plug into the SDK's documented +# extension point — ``streamablehttp_client(httpx_client_factory=...)`` +# — and pass a factory that builds the ``httpx.AsyncClient`` with a +# response hook. The hook fires after headers arrive but BEFORE +# ``raise_for_status()`` runs, so the carrier is populated before the +# SDK swallow. +# +# Forward-compat: ``streamablehttp_client`` is ``@deprecated`` in SDK +# 1.27 in favour of ``streamable_http_client(http_client=...)`` which +# accepts a pre-built client. The same factory pattern translates +# 1:1 against the new entry point when we migrate. + + +# Defensive cap on the number of scopes we report in +# ``mcp_insufficient_scope`` audit/error payloads. Real ASes return +# single-digit scope counts; the cap stops a malicious upstream from +# bloating either surface via a thousand-token ``scope=`` value. +_MAX_INSUFFICIENT_SCOPE_REPORTED = 32 + + +@dataclass +class _AuthCapture: + """Carrier populated by the response hook on 4xx upstream responses.""" + + status: int | None = None + www_authenticate: str | None = None + + +class _PoolDispatchRetryRequested(BaseException): # noqa: N818 + """Module-private signal: the auth_401 branch wants the sync caller to retry. + + Inherits :class:`BaseException` (not ``Exception``) so that nothing in + the SDK / anyio path accidentally swallows it inside an ``except + Exception`` block. ``_dispatch_pool_sync`` is the only caller that + catches it and the only handler that re-issues the dispatch on a + fresh ``asyncio.Task`` via ``run_coroutine_threadsafe``. + """ + + +def _make_capturing_http_factory(capture: _AuthCapture) -> McpHttpClientFactory: + """Return an ``httpx`` factory that records 4xx auth signals into ``capture``. + + The hook is ``async`` because :class:`httpx.AsyncClient` invokes + response hooks via ``await hook(response)`` — a sync function would + return ``None`` and ``await None`` raises ``TypeError`` inside the + SDK's :meth:`client.stream` call. Even though our work is purely + synchronous (read ``status_code`` and a header), the contract + requires an awaitable. The first attempt's hook would still + populate the carrier (the body runs before the ``await``), but the + ``TypeError`` poisons the SDK's anyio TaskGroup teardown so the + next ``streamablehttp_client(...)`` invocation surfaces a stray + ``CancelledError`` from inside its own scope. Empirically verified + via a minimal repro: a sync hook breaks back-to-back connects in + the same process; the async form does not. + """ + + async def _hook(response: httpx.Response) -> None: + # Only record on auth-relevant statuses to keep the carrier + # focused. ``capture`` is mutated in place; the dispatcher + # consults it after ``call_tool`` returns/raises. No I/O, no + # other awaits — the hook stays cancellation-safe. + # + # Use ``get_list(...)[0]`` rather than ``get(...)`` so a + # malicious upstream that emits multiple ``WWW-Authenticate`` + # headers cannot inject auth-params into the parser via the + # comma-joined value httpx returns from ``get(...)``. Repeated + # headers join with ``, `` which the RFC 7235 tokenizer would + # otherwise consume as a continuation of the first challenge, + # silently folding attacker scopes into the parsed dict. We + # discard every challenge after the first; defence-in-depth + # mirror lives in ``parse_www_authenticate_bearer``. + status = response.status_code + if status in (401, 403): + capture.status = status + headers = response.headers.get_list("www-authenticate") + capture.www_authenticate = headers[0] if headers else None + + def _factory( + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + kwargs: dict[str, Any] = { + "follow_redirects": True, + "event_hooks": {"response": [_hook]}, + } + if timeout is None: + kwargs["timeout"] = httpx.Timeout( + MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT + ) + else: + kwargs["timeout"] = timeout + if headers is not None: + kwargs["headers"] = headers + if auth is not None: + kwargs["auth"] = auth + return httpx.AsyncClient(**kwargs) + + return _factory + + # --------------------------------------------------------------------------- # MCP ↔ OpenAI schema conversion # --------------------------------------------------------------------------- @@ -156,9 +276,9 @@ class PoolEntryState: session: Any | None = None stack: AsyncExitStack | None = None streams: tuple[Any, Any] | None = None - # Catalog state — populated when Phase 6 wires per-user discovery in; - # left ``None`` here so 200-entry pools don't retain 600 empty list - # objects. + # Catalog state — populated lazily once per-user discovery wires + # in; left ``None`` here so 200-entry pools don't retain 600 empty + # list objects. tools: list[dict[str, Any]] | None = None resources: list[dict[str, Any]] | None = None prompts: list[dict[str, Any]] | None = None @@ -501,10 +621,27 @@ class MCPClientManager: errors could mask the original exception. CancelledError is caught explicitly because it is the primary failure mode (stray cancel from broken anyio scope) and is BaseException, not Exception. + + ``BaseExceptionGroup`` is also caught: anyio's TaskGroup wraps any + unhandled task exception (e.g., the SDK's + ``HTTPStatusError`` raised inside ``post_writer``'s + ``tg.start_soon(handle_request_async)`` after a 401/403) + into a ``BaseExceptionGroup`` on ``__aexit__``. Without this catch + the auth-retry path's eager teardown would propagate the SDK's + own collected fallout. + + The 5s ``asyncio.wait_for`` is a deliberate guard against + ``aclose()`` hanging on a broken stack (e.g., a never-completing + anyio task during teardown). The auth_401 retry runs on a + fresh :class:`asyncio.Task` scheduled via + :func:`asyncio.run_coroutine_threadsafe` (see + :meth:`_dispatch_pool_sync`), so this helper is never invoked + from inside an active cancellation context — ``wait_for``'s + own cancellation observation does not abort the aclose early. """ try: await asyncio.wait_for(stack.aclose(), timeout=5) - except (Exception, asyncio.CancelledError): + except (Exception, asyncio.CancelledError, BaseExceptionGroup): log.debug("Error closing AsyncExitStack; ignoring", exc_info=True) async def _safe_teardown_on_connect_failure( @@ -785,6 +922,8 @@ class MCPClientManager: key: tuple[str, str], cfg: dict[str, Any], access_token: str, + *, + auth_capture: _AuthCapture | None = None, ) -> PoolEntryState: """Connect a single per-(user, server) pool entry. @@ -795,8 +934,17 @@ class MCPClientManager: * ``Authorization: Bearer {access_token}`` injected into headers alongside any operator-supplied static headers. * No catalog discovery (tools / resources / prompts) and no - notification handler — those land in Phase 6 once per-user - catalog state is in place. + notification handler — those land later once per-user catalog + state is in place. + + When ``auth_capture`` is supplied, the underlying ``httpx`` + client is built via a factory whose response hook records 401/403 + status + ``WWW-Authenticate`` into the carrier, recovering the + upstream auth signal that the SDK's ``post_writer`` would + otherwise swallow. Static-path callers + (:meth:`_connect_one`) MUST NOT pass this — the static path + must remain byte-identical, which means the SDK's default + ``create_mcp_http_client`` factory. MUST run on the mcp-loop. Caller holds ``entry.open_lock``. """ @@ -834,12 +982,16 @@ class MCPClientManager: headers: dict[str, str] = dict(cfg.get("headers") or {}) headers["Authorization"] = f"Bearer {access_token}" + client_kwargs: dict[str, Any] = {"url": url, "headers": headers} + if auth_capture is not None: + client_kwargs["httpx_client_factory"] = _make_capturing_http_factory(auth_capture) + stack = AsyncExitStack() await stack.__aenter__() try: await self._tcp_probe(key, url) read, write, _ = await asyncio.wait_for( - stack.enter_async_context(streamablehttp_client(url=url, headers=headers)), + stack.enter_async_context(streamablehttp_client(**client_kwargs)), timeout=self._CONNECT_TIMEOUT, ) entry.streams = (read, write) @@ -1011,8 +1163,11 @@ class MCPClientManager: # -- failure classification (pool dispatch) ------------------------------ def _classify_failure( - self, exc: BaseException - ) -> Literal["transport", "auth", "protocol", "other"]: + self, + exc: BaseException, + *, + capture: _AuthCapture | None = None, + ) -> Literal["transport", "auth_401", "auth_403", "protocol", "other"]: """Classify a dispatch-time exception for circuit-breaker gating. Only ``transport`` failures trip the per-server breaker. Auth @@ -1020,16 +1175,24 @@ class MCPClientManager: breaker. Protocol errors (``McpError``) come from a healthy connection that rejected the request. - ``auth`` is detected via :class:`httpx.HTTPStatusError` since the - streamable-http transport surfaces upstream HTTP errors through - ``response.raise_for_status()``. ``McpError`` payloads do not - carry a clean status code; bare ``McpError`` therefore stays - ``protocol``. + Auth detection prefers ``capture.status`` (response-hook + introspection — the SDK swallows :class:`httpx.HTTPStatusError` + in its ``post_writer`` so the carrier is the only signal that + reaches us in production). The ``HTTPStatusError`` fallback is + defense-in-depth for the non-SDK refresh path + (:func:`turnstone.core.mcp_oauth._refresh_and_persist`) where + ``httpx`` errors propagate directly. """ + if capture is not None and capture.status == 401: + return "auth_401" + if capture is not None and capture.status == 403: + return "auth_403" if isinstance(exc, httpx.HTTPStatusError): status = exc.response.status_code - if status in (401, 403): - return "auth" + if status == 401: + return "auth_401" + if status == 403: + return "auth_403" if isinstance(exc, McpError): return "protocol" if isinstance(exc, BrokenPipeError | ConnectionResetError | EOFError | TimeoutError): @@ -1046,10 +1209,9 @@ class MCPClientManager: Every entry sourced from ``_static_servers`` is, by construction, NOT ``auth_type='oauth_user'`` — :func:`_db_servers_to_config` - strips oauth_user rows on the way in. Phase 5 pool catalogs do - not contribute to ``_tool_map``; pool tools become reachable - via ``_tool_map`` only when Phase 7 (catalog scoping) lands - per-user catalogs. + strips oauth_user rows on the way in. Pool catalogs do not + contribute to ``_tool_map`` today; pool tools become reachable + via ``_tool_map`` only once per-user catalog scoping lands. """ new_tools: list[dict[str, Any]] = [] new_map: dict[str, tuple[str, str]] = {} @@ -1912,14 +2074,14 @@ class MCPClientManager: def is_mcp_tool(self, func_name: str) -> bool: """Check whether *func_name* belongs to an MCP server. - Phase 5 caveat: only static-path tools (``auth_type ∈ {none, - static}``) populate ``_tool_map``; pool tools become reachable - via this method only when Phase 7 (catalog scoping) lands - per-user catalogs. Until then, pool dispatch is reachable from - production callers like ``ChatSession._exec_mcp_tool`` only when - the LLM produces a ``mcp__{server}__{tool}`` name that bypasses - ``is_mcp_tool``-style gating, or via direct ``call_tool_sync`` - with a known prefixed name (the path the new pool tests use). + Caveat: only static-path tools (``auth_type ∈ {none, static}``) + populate ``_tool_map``; pool tools become reachable via this + method only once per-user catalog scoping lands. Until then, + pool dispatch is reachable from production callers like + ``ChatSession._exec_mcp_tool`` only when the LLM produces a + ``mcp__{server}__{tool}`` name that bypasses ``is_mcp_tool``- + style gating, or via direct ``call_tool_sync`` with a known + prefixed name (the path the new pool tests use). """ return func_name in self._tool_map @@ -2025,6 +2187,17 @@ class MCPClientManager: When ``user_id`` is supplied AND the resolved server's ``auth_type`` is ``oauth_user``, dispatch goes through the per-(user, server) pool. Otherwise the call takes the byte-identical static path. + + Pool-path 401/403 handling: the SDK's ``post_writer`` swallows + ``httpx.HTTPStatusError``; we recover the upstream auth signal + via a response-hook carrier on the per-dispatch + ``httpx.AsyncClient``. A 401 triggers one refresh-and-retry + with ``force_refresh=True``; persistent 401 emits + ``mcp_consent_required``. A 403 with + ``WWW-Authenticate: error="insufficient_scope"`` emits + ``mcp_insufficient_scope`` with the parsed scope set. Other + 403s emit ``mcp_tool_call_forbidden``. Auth failures NEVER + trip the per-server breaker. """ mapping = self._tool_map.get(func_name) server_name: str | None = None @@ -2033,10 +2206,10 @@ class MCPClientManager: server_name, original_name = mapping # Pool dispatch is gated on (a) caller passing user_id and - # (b) the server row being auth_type=oauth_user. The Phase 5 - # tool-map only carries static-path entries; pool catalogs land - # in Phase 7. Consequently, in Phase 5 a pool dispatch reaches - # this branch only when the caller supplied func_name as + # (b) the server row being auth_type=oauth_user. The current + # tool-map only carries static-path entries; pool catalogs are + # not merged in. Consequently, a pool dispatch reaches this + # branch only when the caller supplied func_name as # ``mcp__{server}__{tool}`` and we resolve auth_type via storage. if user_id and self._app_state is not None and self._storage is not None: pool_target = self._resolve_pool_target(func_name, server_name, original_name) @@ -2098,7 +2271,7 @@ class MCPClientManager: """Resolve ``(server_name, original_name, server_row)`` for pool dispatch. Returns ``None`` when the call should fall through to the static - path. Phase 5: pool catalogs do not contribute to ``_tool_map``, + path. Pool catalogs do not contribute to ``_tool_map`` today, so a pool tool is invoked by passing the prefixed name ``mcp__{server}__{tool}`` directly. ``mcp_servers.auth_type`` confirms pool eligibility. @@ -2155,10 +2328,86 @@ class MCPClientManager: (consent required, key mismatch, etc.). ``server_row`` is the row already resolved by ``_resolve_pool_target`` and is reused verbatim by ``_dispatch_pool`` to skip a duplicate DB hop. + + Retry-on-401: the 401 branch in :meth:`_dispatch_pool` raises + :class:`_PoolDispatchRetryRequested` after refreshing the + bearer. Catching the signal HERE — at the sync boundary — means + the retry is scheduled via a fresh + :func:`asyncio.run_coroutine_threadsafe` call, which runs the + retry coroutine in a brand-new :class:`asyncio.Task` with no + inherited anyio cancel-scope state from the prior connect's + ``streamablehttp_client`` TaskGroup. An in-task retry inherits + that scope state across :meth:`asyncio.Task.uncancel` and + ``loop.create_task`` and surfaces ``CancelledError`` from inside + the retry's own anyio scope. The retry-count ceiling is one; + callers see consent_required if both attempts 401. + + The ``timeout`` is a wall-clock budget across both attempts — + the retry's ``future.result`` window is reduced by however long + the first attempt consumed before raising + :class:`_PoolDispatchRetryRequested`. Without this, a slow + first attempt followed by a stuck retry could double the + caller-observed timeout. + """ + assert self._loop is not None + start = time.monotonic() + try: + return self._run_pool_dispatch_attempt( + retry_count=0, + timeout=timeout, + original_timeout=timeout, + user_id=user_id, + server_name=server_name, + original_name=original_name, + arguments=arguments, + server_row=server_row, + ) + except _PoolDispatchRetryRequested: + # auth_401: refresh already happened on the prior task; + # re-issue on a fresh task so the retry's anyio scope + # state is independent of the prior connect's TaskGroup + # teardown. + remaining = max(1, int(timeout - (time.monotonic() - start))) + return self._run_pool_dispatch_attempt( + retry_count=1, + timeout=remaining, + original_timeout=timeout, + user_id=user_id, + server_name=server_name, + original_name=original_name, + arguments=arguments, + server_row=server_row, + ) + + def _run_pool_dispatch_attempt( + self, + *, + retry_count: int, + timeout: int, + original_timeout: int, + user_id: str, + server_name: str, + original_name: str, + arguments: dict[str, Any], + server_row: dict[str, Any], + ) -> str: + """Schedule one ``_dispatch_pool`` attempt and wait for the result. + + Split out of :meth:`_dispatch_pool_sync` so the two retry + attempts share scheduling + timeout-bookkeeping without + re-introducing the ``for retry_count in (0, 1)`` loop (which + gave both attempts the full ``timeout`` and required an + unreachable ``RuntimeError`` fallback). + + ``original_timeout`` is the wall-clock budget the caller + requested; ``timeout`` is what's left for this specific + attempt. The ``TimeoutError`` message reports the original so + callers see the budget they set, not the trimmed window. """ assert self._loop is not None future = asyncio.run_coroutine_threadsafe( self._dispatch_pool( + retry_count=retry_count, user_id=user_id, server_name=server_name, original_name=original_name, @@ -2172,7 +2421,7 @@ class MCPClientManager: except concurrent.futures.TimeoutError: future.cancel() self._cb_record_failure(server_name) - raise TimeoutError(f"MCP tool call timed out after {timeout}s") from None + raise TimeoutError(f"MCP tool call timed out after {original_timeout}s") from None async def _dispatch_pool( self, @@ -2182,6 +2431,7 @@ class MCPClientManager: original_name: str, arguments: dict[str, Any], server_row: dict[str, Any], + retry_count: int = 0, ) -> str: """Pool-side coroutine: resolve token, connect-or-reuse, dispatch. @@ -2190,6 +2440,18 @@ class MCPClientManager: is supplied by ``_resolve_pool_target`` so this path doesn't re-issue the ``mcp_servers`` lookup; it's also pre-validated to have ``auth_type='oauth_user'``. + + ``retry_count`` is supplied by :meth:`_dispatch_pool_sync` and + bounds the auth_401 refresh-and-retry to one re-issue. The first + attempt (``retry_count == 0``) refreshes the token via + ``force_refresh=True`` and raises + :class:`_PoolDispatchRetryRequested` so the sync caller schedules + the retry on a fresh :class:`asyncio.Task` (an in-task retry + inherits anyio cancel-scope state from the prior connect's + TaskGroup teardown and surfaces ``CancelledError`` from inside + the retry's own anyio scope; the cross-task hop avoids that). + At the ceiling (``retry_count == 1``) the auth_401 branch emits + ``mcp_consent_required`` instead. """ if self._app_state is None: raise RuntimeError("Pool dispatch requires set_app_state() to have been called") @@ -2199,8 +2461,17 @@ class MCPClientManager: # gate runs only AFTER we have a usable access token; that way a # spurious cooldown-expired probe never lands here without ending # in either a real success or a real transport failure. + # + # ``retry_count >= 1`` forces the AS round-trip: the previous + # attempt's auth_401 branch evicted the session and signalled + # this retry; the local cached token is the one the AS just + # rejected, so reading it back without ``force_refresh=True`` + # would re-attempt with the same (rejected) bearer. lookup: TokenLookupResult = await get_user_access_token_classified( - app_state=self._app_state, user_id=user_id, server_name=server_name + app_state=self._app_state, + user_id=user_id, + server_name=server_name, + force_refresh=retry_count > 0, ) if lookup.kind == "missing": return _structured_error( @@ -2258,6 +2529,13 @@ class MCPClientManager: key = (user_id, server_name) entry = await self._ensure_pool_entry(key) + # First attempt — fresh capture per dispatch so a prior + # dispatch's 401 cannot leak into this one's classification. + # Even though ``open_lock`` is held across ``call_tool``, + # allocating per-call (rather than per-entry) protects against + # future concurrent-multiplex regressions if the lock scope + # ever shrinks. + capture = _AuthCapture() try: result = await self._dispatch_pool_with_entry( entry=entry, @@ -2266,23 +2544,48 @@ class MCPClientManager: access_token=access_token, original_name=original_name, arguments=arguments, + auth_capture=capture, ) except BaseException as exc: - classification = self._classify_failure(exc) - if classification == "auth": - # 401/403 from the upstream — pool-entry-only, never - # affects the breaker. Drop the cached session so the - # next dispatch re-authenticates with a fresh token. - evict = self._user_pool_entries.get(key) - if evict is not None: - evict.session = None - log.debug("mcp_pool.auth_failure", exc_info=exc) - raise + classification = self._classify_failure(exc, capture=capture) + if classification == "auth_401": + self._evict_session(key) + if retry_count == 0: + # 401 on the initial attempt: signal the sync caller + # to re-issue on a fresh :class:`asyncio.Task`. The + # retry's :meth:`_dispatch_pool` invocation runs the + # token lookup with ``force_refresh=True`` (the + # ``retry_count > 0`` branch above), guaranteeing + # the bearer attached to the retry's connect is + # different from the one the AS just rejected. The + # cross-task hop avoids the in-task anyio + # cancel-scope inheritance that surfaces a + # ``CancelledError`` from inside the retry's own + # ``streamablehttp_client`` TaskGroup — see + # :meth:`_dispatch_pool_sync` for the architectural + # rationale. + log.debug("mcp_pool.auth_401_initial", exc_info=exc) + raise _PoolDispatchRetryRequested from None + # retry_count == 1 — refreshed bearer also rejected; + # emit consent_required so the user/operator re-grants. + log.debug("mcp_pool.auth_401_retry_failed", exc_info=exc) + return _structured_error( + code="mcp_consent_required", + server=server_name, + detail="Refreshed token still rejected. Re-consent required.", + ) + if classification == "auth_403": + self._evict_session(key) + log.debug("mcp_pool.auth_403", exc_info=exc) + return await self._handle_auth_403( + user_id=user_id, + server_name=server_name, + server_row=server_row, + capture=capture, + ) if classification == "transport": self._cb_record_failure(server_name) - evict = self._user_pool_entries.get(key) - if evict is not None: - evict.session = None + self._evict_session(key) log.debug("mcp_pool.transport_failure", exc_info=exc) raise # protocol / other — don't trip the breaker. @@ -2291,6 +2594,63 @@ class MCPClientManager: self._cb_record_success(server_name) return result + def _evict_session(self, key: tuple[str, str]) -> None: + """Drop the cached session on a pool entry. Stack/streams left for reconnect. + + Auth/transport branches both call this — the next connect's + ``_connect_one_pool`` tears down the stale stack lazily via the + stale-entry guard at the top of the method. Closing eagerly + from here is incorrect under cancellation: ``stack.aclose()`` + must run inside the same anyio scope it was entered in, which + the next connect arranges. + """ + evict = self._user_pool_entries.get(key) + if evict is not None: + evict.session = None + + async def _handle_auth_403( + self, + *, + user_id: str, + server_name: str, + server_row: dict[str, Any], + capture: _AuthCapture, + ) -> str: + """Map a 403 + WWW-Authenticate into a structured error. + + ``error="insufficient_scope"`` becomes ``mcp_insufficient_scope`` + with the parsed ``scope=...`` set so the dashboard renderer + can construct an authorize URL with the union of original + + new scopes — re-consenting with the original scopes alone + would loop because the AS would re-issue the same insufficient + token. Other 403s become a generic ``mcp_tool_call_forbidden`` + with no retry; the user lacks permission and a step-up + wouldn't help. + """ + header = capture.www_authenticate or "" + error_token = parse_www_authenticate_error(header) + if error_token == "insufficient_scope": + scopes = parse_www_authenticate_scope(header) + scopes = scopes[:_MAX_INSUFFICIENT_SCOPE_REPORTED] + await emit_insufficient_scope_audit( + app_state=self._app_state, + user_id=user_id, + server_name=server_name, + server_row=server_row, + scopes=scopes, + ) + return _structured_error( + code="mcp_insufficient_scope", + server=server_name, + detail=("Tool requires elevated scopes. Re-consent flow with new scopes required."), + scopes_required=list(scopes), + ) + return _structured_error( + code="mcp_tool_call_forbidden", + server=server_name, + detail="Tool call forbidden by upstream policy.", + ) + async def _dispatch_pool_with_entry( self, *, @@ -2300,15 +2660,23 @@ class MCPClientManager: access_token: str, original_name: str, arguments: dict[str, Any], + auth_capture: _AuthCapture, ) -> str: - """Acquire ``entry.open_lock`` only across connect-or-reuse, then dispatch. + """Hold ``entry.open_lock`` across connect-or-reuse AND ``call_tool``. - Releasing ``open_lock`` before ``call_tool`` lets two concurrent - tool calls from the SAME user against the SAME server multiplex - on a shared :class:`mcp.ClientSession` (request_id-correlated by - the SDK). The eviction interlock now uses ``entry.in_flight``: - eviction skips entries with in-flight calls so a long-running - ``call_tool`` can't have its session yanked mid-await. + Lock held across ``call_tool`` because the per-dispatch + ``_AuthCapture`` is keyed off the httpx event hook; releasing + would let a concurrent same-(user, server) dispatch overwrite + the carrier mid-flight, attributing one caller's auth failure + to another. Holding the lock serialises same-(user, server) + calls — acceptable because that contention scenario is rare + in practice (ChatSession dispatches sequentially and per- + (user, server) parallelism is not a production requirement). + + ``entry.in_flight`` accounting is preserved for the eviction + interlock — it's belt-and-braces here since ``open_lock.locked()`` + already signals "do not evict", but the in-flight counter + remains the source of truth for :meth:`_close_pool_entry_if_idle`. """ async with entry.open_lock: entry.last_used = time.monotonic() @@ -2316,15 +2684,17 @@ class MCPClientManager: session = entry.session if session is None: # Lazy connect — also covers post-eviction recovery. - fresh = await self._connect_one_pool(key, cfg, access_token) + fresh = await self._connect_one_pool( + key, cfg, access_token, auth_capture=auth_capture + ) session = fresh.session if session is None: raise RuntimeError(f"Pool connect for {key!r} produced no session") entry.in_flight += 1 - try: - result = await session.call_tool(original_name, arguments) - finally: - entry.in_flight -= 1 + try: + result = await session.call_tool(original_name, arguments) + finally: + entry.in_flight -= 1 return _decode_tool_result(result) # -- resource read ------------------------------------------------------- @@ -2482,13 +2852,24 @@ def _decode_tool_result(result: Any) -> str: return output -def _structured_error(*, code: str, server: str, detail: str) -> str: +def _structured_error( + *, + code: str, + server: str, + detail: str, + scopes_required: list[str] | None = None, +) -> str: """Encode a pool-dispatch failure as a JSON string. Returned to the agent through ``_exec_mcp_tool`` so the LLM can narrate "tool unavailable, consent required" rather than crashing - the workstream. Schema mirrors RFC §6 (consent UX) plus the - decrypt-failure code introduced in RFC §5.3. + the workstream. Schema covers ``mcp_consent_required``, + decrypt-failure (``mcp_token_undecryptable_key_unknown``), and the + ``mcp_insufficient_scope`` step-up shape. + + ``scopes_required`` is omitted from the payload when ``None`` — + the dashboard renderer keys on its presence to construct an + authorize URL with the union of original + new scopes. Operator-actionable encryption-key fingerprints are intentionally NOT included in this payload: they are already captured server-side @@ -2496,14 +2877,14 @@ def _structured_error(*, code: str, server: str, detail: str) -> str: to the LLM (and through it to the model provider) would be unnecessary disclosure. """ - payload: dict[str, Any] = { - "error": { - "code": code, - "server": server, - "detail": detail, - } + err: dict[str, Any] = { + "code": code, + "server": server, + "detail": detail, } - return json.dumps(payload) + if scopes_required is not None: + err["scopes_required"] = scopes_required + return json.dumps({"error": err}) def _pool_cfg_from_row(row: dict[str, Any]) -> dict[str, Any]: diff --git a/turnstone/core/mcp_http_parsers.py b/turnstone/core/mcp_http_parsers.py new file mode 100644 index 00000000..0c521fd5 --- /dev/null +++ b/turnstone/core/mcp_http_parsers.py @@ -0,0 +1,222 @@ +"""HTTP header parsing helpers shared by the MCP client and OAuth modules. + +Both ``mcp_client`` and ``mcp_oauth`` need to extract structured values from +``WWW-Authenticate: Bearer ...`` headers — ``mcp_client`` to classify +401/403 responses for the user-pool dispatcher, and ``mcp_oauth`` to pull +the ``resource_metadata`` URL out of a discovery challenge. This module +hosts the shared primitives so both modules can call them without +duplicating fragile substring scanners. + +The earlier hand-rolled scanners (``_parse_www_authenticate_scope`` / +``_parse_www_authenticate_error`` in ``mcp_client``) used +``header.lower().find(needle, i)`` to locate parameter names. That made +them vulnerable to: + +* matching ``scope`` inside ``xscope`` or ``ascope``, +* matching the literal text ``scope=...`` embedded inside the quoted + ``realm`` value of a preceding ``auth-param``, +* O(N**2) behaviour on pathological input (each ``find`` rescans the prefix). + +This module replaces those with a single tokenizer that walks the RFC 7235 +``challenge → auth-param`` grammar once, tracks quoted-string state, and +returns a normalised ``{key.lower(): value}`` dict. The thin extraction +wrappers (``parse_www_authenticate_scope`` / ``parse_www_authenticate_error``) +preserve the original return shapes so call sites only need to swap the +import. +""" + +from __future__ import annotations + + +def _parse_quoted_string(text: str, start: int) -> tuple[str, int] | None: + """Parse an RFC 7230 ``quoted-string`` starting at ``text[start]``. + + Returns ``(value, end_index)`` where ``end_index`` is the index just + past the closing quote, or ``None`` if the input is malformed (no + opening quote, unterminated string). + + Handles ``\\"`` and ``\\\\`` escapes per RFC 7230 section 3.2.6 — the + prior naive ``([^"]+)`` regex truncated the URL at the first + unescaped quote and silently dropped backslash escapes from the + value. + """ + if start >= len(text) or text[start] != '"': + return None + out: list[str] = [] + i = start + 1 + while i < len(text): + ch = text[i] + if ch == "\\" and i + 1 < len(text): + out.append(text[i + 1]) + i += 2 + continue + if ch == '"': + return "".join(out), i + 1 + out.append(ch) + i += 1 + return None + + +# Maximum header length we'll attempt to parse. Real ASes emit a handful +# of short auth-params; anything past this is either malformed or +# adversarial. Returning ``{}`` (rather than raising) keeps callers' error +# paths uniform with "unparseable header → no signal". +_MAX_HEADER_LEN = 4096 + +_TOKEN_DELIMS = frozenset('()<>@,;:\\"/[]?={} \t') + + +def _is_token_char(ch: str) -> bool: + """RFC 7230 token character: visible ASCII minus the delimiter set.""" + return ch.isascii() and ch.isprintable() and ch not in _TOKEN_DELIMS + + +def _looks_like_bearer_challenge_start(header: str, i: int) -> bool: + """Peek at ``header[i:]`` for the start of a fresh ``Bearer`` challenge. + + Returns True when the slice begins with the case-insensitive token + ``Bearer`` followed by whitespace — the RFC 7235 marker for a new + ``challenge`` after a separator comma. This is the cue the bearer + tokenizer uses to stop parsing rather than fold a second challenge's + auth-params into the first challenge's dict. + """ + n = len(header) + if i + 6 > n: + return False + if header[i : i + 6].lower() != "bearer": + return False + after = i + 6 + # ``Bearer`` must be followed by whitespace to qualify as a scheme + # boundary; ``Bearer-like-token`` is just a regular token. + return after < n and header[after] in " \t" + + +def parse_www_authenticate_bearer(header: str) -> dict[str, str]: + """Extract ``auth-param``s from a ``WWW-Authenticate: Bearer ...`` header. + + Walks the RFC 7235 challenge grammar once, returning a dict of + ``{lowercased-key: value}`` pairs. Quoted-strings are unquoted (with + backslash escapes resolved). Unknown / malformed input returns an + empty dict — never raises. + + Only ``Bearer`` challenges are recognised. The function ignores any + leading whitespace before the scheme. When a second ``Bearer`` + challenge appears after a separator comma — as it would when + httpx joins repeated ``WWW-Authenticate`` headers via + ``response.headers.get(...)`` — the tokenizer stops at the + challenge boundary rather than folding the second challenge's + auth-params into the first challenge's dict. This is the + parser-side defence-in-depth mirror of the + ``response.headers.get_list(...)[0]`` guard in the dispatcher's + capturing httpx factory; either layer alone neutralises the + multi-header injection vector but both run together so a + regression in one cannot silently re-open it. + A ``realm`` value that contains the literal text ``scope=fake`` is + correctly attributed to ``realm`` because the tokenizer respects + quoted-string boundaries. + """ + if not header or len(header) > _MAX_HEADER_LEN: + return {} + + n = len(header) + i = 0 + + # Skip leading whitespace then the ``Bearer`` scheme token. + while i < n and header[i] in " \t": + i += 1 + scheme_start = i + while i < n and _is_token_char(header[i]): + i += 1 + scheme = header[scheme_start:i] + if scheme.lower() != "bearer": + return {} + # Require at least one space between scheme and first auth-param. + if i >= n or header[i] not in " \t": + return {} + + out: dict[str, str] = {} + while i < n: + # Skip whitespace and stray commas between params. + while i < n and header[i] in " \t,": + i += 1 + if i >= n: + break + # If a fresh ``Bearer`` challenge starts here, the upstream is + # multi-challenge — stop before reading any of its auth-params. + if _looks_like_bearer_challenge_start(header, i): + break + # Read the param key (a token). + key_start = i + while i < n and _is_token_char(header[i]): + i += 1 + if i == key_start: + # Not a valid token start — skip one char to make forward + # progress and continue. This bounds total cost to O(N). + i += 1 + continue + key = header[key_start:i].lower() + # Optional whitespace, then ``=``. + while i < n and header[i] in " \t": + i += 1 + if i >= n or header[i] != "=": + # Param without a value — skip. + continue + i += 1 + while i < n and header[i] in " \t": + i += 1 + if i >= n: + break + # Value: either a quoted-string or a token. + if header[i] == '"': + parsed = _parse_quoted_string(header, i) + if parsed is None: + # Unterminated quoted-string — treat the rest of the + # header as garbage and stop. Returning what we already + # have is safer than guessing where the value ends. + break + value, i = parsed + out.setdefault(key, value) + else: + val_start = i + while i < n and header[i] not in ", \t": + i += 1 + value = header[val_start:i] + out.setdefault(key, value) + return out + + +def parse_www_authenticate_scope(header: str) -> tuple[str, ...]: + """Return the ``scope=...`` value as a tuple of individual scopes. + + Splits on a single space per RFC 6749 section 3.3 (``scope-token`` + sequence). Returns ``()`` when the header is malformed or carries no + ``scope`` parameter. + + Each token is validated against the RFC 6749 §3.3 ``scope-token`` + grammar (visible ASCII ``0x21..0x7E`` excluding ``"`` and ``\\``) + so that a malicious or buggy AS cannot smuggle CR/LF/tab/control + bytes through a future log or notification path. Today scopes are + JSON-encoded everywhere downstream so no concrete exploit exists, + but the validation is cheap and forecloses regressions in + structured-error rendering. + """ + params = parse_www_authenticate_bearer(header) + value = params.get("scope") + if not value: + return () + return tuple( + s + for s in value.split(" ") + if s and all(0x21 <= ord(c) <= 0x7E and c not in '"\\' for c in s) + ) + + +def parse_www_authenticate_error(header: str) -> str | None: + """Return the ``error=...`` value or ``None`` when absent. + + The tokenizer naturally distinguishes ``error`` from + ``error_description`` / ``error_uri`` because ``_`` is not a valid + token-character delimiter — they parse as separate keys. + """ + params = parse_www_authenticate_bearer(header) + return params.get("error") or None diff --git a/turnstone/core/mcp_oauth.py b/turnstone/core/mcp_oauth.py index 98a3ec91..aaeb0754 100644 --- a/turnstone/core/mcp_oauth.py +++ b/turnstone/core/mcp_oauth.py @@ -26,7 +26,6 @@ import concurrent.futures import contextlib import hashlib import json -import re import secrets import time import urllib.parse @@ -39,6 +38,7 @@ import httpx from turnstone.core.audit import record_audit from turnstone.core.log import get_logger from turnstone.core.mcp_crypto import MCPTokenDecryptError +from turnstone.core.mcp_http_parsers import parse_www_authenticate_bearer from turnstone.core.oauth_ssrf import ( OAuthSSRFError, sanitize_log_text, @@ -127,53 +127,19 @@ class ASMetadata: # --------------------------------------------------------------------------- -_PRM_RESOURCE_METADATA_KEY_RE = re.compile( - r"resource_metadata\s*=\s*", - re.IGNORECASE, -) - - -def _parse_quoted_string(text: str, start: int) -> tuple[str, int] | None: - """Parse an RFC 7230 ``quoted-string`` starting at ``text[start]``. - - Returns ``(value, end_index)`` where ``end_index`` is the index just - past the closing quote, or ``None`` if the input is malformed (no - opening quote, unterminated string). - - Handles ``\\"`` and ``\\\\`` escapes per RFC 7230 §3.2.6 — the prior - naive ``([^"]+)`` regex truncated the URL at the first unescaped quote - and silently dropped backslash escapes from the value. - """ - if start >= len(text) or text[start] != '"': - return None - out: list[str] = [] - i = start + 1 - while i < len(text): - ch = text[i] - if ch == "\\" and i + 1 < len(text): - out.append(text[i + 1]) - i += 2 - continue - if ch == '"': - return "".join(out), i + 1 - out.append(ch) - i += 1 - return None - - def _parse_prm_url_from_www_authenticate(header: str) -> str | None: """Extract ``resource_metadata`` URL from a ``WWW-Authenticate: Bearer`` header. Returns the URL string or ``None`` when the header lacks the param, - is malformed, or terminates the quoted-string prematurely. + is malformed, or terminates the quoted-string prematurely. Delegates + to :func:`parse_www_authenticate_bearer` so quoted-string handling + (RFC 7230 §3.2.6 backslash escapes) and the multi-challenge + defence-in-depth guard live in one place. """ if not header: return None - for match in _PRM_RESOURCE_METADATA_KEY_RE.finditer(header): - parsed = _parse_quoted_string(header, match.end()) - if parsed is not None: - return parsed[0] - return None + params = parse_www_authenticate_bearer(header) + return params.get("resource_metadata") or None async def _fetch_prm_issuer( @@ -1145,8 +1111,8 @@ class TokenLookupResult: rejected" — three states that the ``Optional[str]``-returning :func:`get_user_access_token` collapses to ``None``. The distinction matters because they map to different user-facing - errors (RFC §5.3 forbids emitting ``mcp_consent_required`` on a - decrypt failure). + errors — emitting ``mcp_consent_required`` on a decrypt failure + would be wrong (the user can't fix it; only an operator can). """ kind: Literal["token", "missing", "decrypt_failure", "refresh_failed"] = "missing" @@ -1177,11 +1143,12 @@ async def get_user_access_token(*, app_state: Any, user_id: str, server_name: st async def get_user_access_token_classified( - *, app_state: Any, user_id: str, server_name: str + *, app_state: Any, user_id: str, server_name: str, force_refresh: bool = False ) -> TokenLookupResult: """Tagged token lookup with refresh-on-expiry. - Walks the §1.5 / RFC §6 state machine and returns a tagged result so + Walks the token-lookup state machine (token / missing / + decrypt_failure / refresh_failed) and returns a tagged result so the dispatcher can map each failure mode to the right user-facing error. @@ -1196,12 +1163,28 @@ async def get_user_access_token_classified( surviving local caller then serializes against other nodes via the cluster lock. The re-read inside the locked block collapses both contention windows. + + ``force_refresh=True`` bypasses the local freshness check and + forces an AS round-trip — used by the dispatch path when the + upstream AS rejected a token our cache still considered fresh + (e.g., AS-side revocation). Concurrent ``force_refresh=True`` + callers still collapse to one round-trip via the dual-layer lock: + the second caller sees ``last_refreshed > t_lock_request_started`` + and reuses the freshly-refreshed token. """ token_store: MCPTokenStore | None = getattr(app_state, "mcp_token_store", None) if token_store is None: log.debug("mcp_server.oauth.token_store_unconfigured") return TokenLookupResult(kind="missing") + # Captured BEFORE we acquire any lock so the inside-lock guard can + # tell whether another caller refreshed under contention. Truncated + # to seconds because ``last_refreshed`` storage has second-precision + # ISO8601 — comparing microsecond-precision against second-precision + # would race when the refresh and the contention occur in the same + # wall-clock second. + t_lock_request_started = datetime.now(UTC).replace(microsecond=0) + try: plain = await asyncio.to_thread(token_store.get_user_token, user_id, server_name) except MCPTokenDecryptError as exc: @@ -1219,7 +1202,7 @@ async def get_user_access_token_classified( return TokenLookupResult(kind="missing") expires_at = plain.get("expires_at") - if not _token_needs_refresh(expires_at): + if not force_refresh and not _token_needs_refresh(expires_at): return TokenLookupResult(kind="token", token=plain["access_token"]) storage = _get_storage(app_state) @@ -1272,8 +1255,20 @@ async def get_user_access_token_classified( if plain2 is None: return TokenLookupResult(kind="missing") expires_at2 = plain2.get("expires_at") + # Reuse the freshly-refreshed token under two conditions: + # 1. ``force_refresh=False`` and the cached token is still fresh + # (existing fast path). + # 2. ``force_refresh=True`` BUT another same-key caller already + # refreshed under contention since we started waiting for the + # lock. ``last_refreshed`` is the storage-side replacement + # timestamp; if it advanced past ``t_lock_request_started``, + # we lost the race and should reuse rather than refresh again. if not _token_needs_refresh(expires_at2): - return TokenLookupResult(kind="token", token=plain2["access_token"]) + if not force_refresh: + return TokenLookupResult(kind="token", token=plain2["access_token"]) + last_refreshed = _parse_iso_to_utc(plain2.get("last_refreshed") or "") + if last_refreshed is not None and last_refreshed >= t_lock_request_started: + return TokenLookupResult(kind="token", token=plain2["access_token"]) refresh_value2 = plain2.get("refresh_token") if not refresh_value2: await asyncio.to_thread(token_store.delete_user_token, user_id, server_name) @@ -1413,7 +1408,7 @@ async def _refresh_and_persist( if not isinstance(new_access, str) or not new_access: raise MCPOAuthRefreshFailed("refresh response missing access_token") - # RFC 6749 §6 — the AS MAY omit ``refresh_token`` from the refresh + # RFC 6749 section 6 — the AS MAY omit ``refresh_token`` from the refresh # response. Most production ASes (Google, default Auth0, default # Okta) do not rotate the refresh token; clearing the column on # every refresh would force the user to re-consent every hour. @@ -1533,6 +1528,35 @@ async def _audit_event( log.debug("mcp_server.oauth.audit_emit_failed", action=action, exc_info=True) +async def emit_insufficient_scope_audit( + *, + app_state: Any, + user_id: str, + server_name: str, + server_row: dict[str, Any], + 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"``. + """ + if app_state is None: + return + server_id = str(server_row.get("server_id") or "") if server_row else "" + await _audit_event( + app_state, + server_id=server_id, + user_id=user_id, + action="mcp_server.oauth.insufficient_scope_emitted", + server_name=server_name, + detail={"scopes_required": list(scopes)}, + ) + + # --------------------------------------------------------------------------- # HTTP handlers — /api/mcp/oauth/start and /api/mcp/oauth/callback # ---------------------------------------------------------------------------