From b9399195608429e86242b130a8707dccefceb410 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Thu, 25 Jun 2026 20:28:51 -0700 Subject: [PATCH] fix(mcp): harden Entra OBO OAuth review follow-ups for #706 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up review of the #706 on-behalf-of / Entra ID MCP changes (#682). security (PKCE downgrade): the AS-metadata "assume S256 when code_challenge_methods_supported is absent" relaxation applied to BOTH the RFC 8414 oauth-authorization-server document and the OIDC openid-configuration document. Per RFC 8414 an omitted field on the oauth-authorization-server document means the AS does NOT support PKCE, so this was fail-open. The client always sends code_challenge_method=S256, making this discovery check the only pre-flight that the AS enforces PKCE. Track which document won discovery and assume S256 only for the OIDC document; the RFC 8414 document now fails closed. Also log which discovery profile (rfc8414 vs oidc) answered, for operators debugging an enterprise AS. bug (consent loss): session-start pool priming called the refreshing token lookup for every cold oauth_user server. A near-expiry token triggered a refresh, and a transient refresh failure (network/5xx/429) deletes the token and emits token_revoked — so a blip during a cold-pool warm (e.g. after a reboot) silently revoked consent across servers the user wasn't even using. Priming now reads the token directly and skips missing/near-expiry tokens; a refresh that may fail stays on the lazy dispatch path. perf/UX (blocking redirect): the OAuth callback awaited prime_user_server (default 20s timeout), holding the consent redirect on a slow/unreachable MCP server. Replaced with fire-and-forget schedule_prime_user_server that schedules onto the mcp-loop (GC-safe, no unreferenced request-loop task) and returns at once. perf: prime a user's pools concurrently under a bound instead of serially, so one slow upstream can't stall the rest. hygiene: log (not silently swallow) prime scheduling failures at session start; add exc_info to the prime-failure warning; guard run_coroutine_threadsafe against a closed mcp-loop. tests: per-document S256 + OIDC-fallback discovery cases; pool priming (non-destructive on near-expiry, skips connected) and bound-token rotation reconnect. --- tests/test_mcp_oauth_discovery.py | 113 ++++++++++++ tests/test_mcp_pool_auth_introspection.py | 215 ++++++++++++++++++++++ turnstone/core/mcp_client.py | 163 ++++++++++------ turnstone/core/mcp_oauth.py | 61 +++--- turnstone/core/session.py | 6 +- 5 files changed, 475 insertions(+), 83 deletions(-) diff --git a/tests/test_mcp_oauth_discovery.py b/tests/test_mcp_oauth_discovery.py index 8321c735..1abe073c 100644 --- a/tests/test_mcp_oauth_discovery.py +++ b/tests/test_mcp_oauth_discovery.py @@ -359,6 +359,119 @@ class TestASMetadataValidation: client.get.assert_not_called() +class TestS256PerDocumentAndOIDCFallback: + """PKCE S256 defaulting is per-discovery-document, and OIDC discovery is a + fallback to RFC 8414 (PR #706 follow-up). + + The client always sends ``code_challenge_method=S256``, so the AS-metadata + check is the only PKCE-enforcement pre-flight. An ABSENT + ``code_challenge_methods_supported`` is treated as "S256 supported" ONLY for + the OIDC ``openid-configuration`` document (where the field is optional and + Entra omits it); for the RFC 8414 ``oauth-authorization-server`` document an + absent field fails closed. + """ + + @staticmethod + def _doc_without_code_challenge() -> dict[str, Any]: + doc = _good_as_metadata_doc() + del doc["code_challenge_methods_supported"] + return doc + + def test_absent_field_on_oidc_doc_assumes_s256(self) -> None: + # RFC 8414 path 404s; the OIDC doc omits code_challenge_methods_supported + # -> assume S256 (Entra's shape) and discovery succeeds. + async def _get(url, *args, **kwargs): + if url.endswith("/oauth-authorization-server"): + return _mk_response(404, json_body=None) + if url.endswith("/openid-configuration"): + return _mk_response(200, self._doc_without_code_challenge()) + raise AssertionError(f"unexpected URL: {url}") + + client = MagicMock(spec=httpx.AsyncClient) + client.get = AsyncMock(side_effect=_get) + storage = _mk_storage_mock() + + async def _run(): + with _public_addr_patch(): + return await discover_authorization_server( + server_name="srv-x", + server_url="https://mcp.example.com/sse", + override_url="https://as.example.com", + cached_issuer=None, + http_client=client, + storage=storage, + server_id="srv-id", + trusted_hosts=frozenset(), + ) + + meta = asyncio.run(_run()) + assert isinstance(meta, ASMetadata) + assert meta.token_endpoint == "https://as.example.com/token" + + def test_absent_field_on_rfc8414_doc_fails_closed(self) -> None: + # The RFC 8414 doc is served (200) but omits the field — must NOT assume + # S256. Per RFC 8414 an omitted field means "no PKCE advertised", so + # discovery fails closed rather than silently downgrading. + async def _get(url, *args, **kwargs): + if url.endswith("/oauth-authorization-server"): + return _mk_response(200, self._doc_without_code_challenge()) + raise AssertionError(f"unexpected URL: {url}") + + client = MagicMock(spec=httpx.AsyncClient) + client.get = AsyncMock(side_effect=_get) + storage = _mk_storage_mock() + + async def _run(): + with _public_addr_patch(): + await discover_authorization_server( + server_name="srv-x", + server_url="https://mcp.example.com/sse", + override_url="https://as.example.com", + cached_issuer=None, + http_client=client, + storage=storage, + server_id="srv-id", + trusted_hosts=frozenset(), + ) + + with pytest.raises(MCPOAuthDiscoveryError, match="S256"): + asyncio.run(_run()) + + def test_rfc8414_404_falls_back_to_openid_configuration(self) -> None: + # RFC 8414 path 404s; the OIDC doc (advertising S256) is parsed instead. + async def _get(url, *args, **kwargs): + if url.endswith("/oauth-authorization-server"): + return _mk_response(404, json_body=None) + if url.endswith("/openid-configuration"): + return _mk_response(200, _good_as_metadata_doc()) + raise AssertionError(f"unexpected URL: {url}") + + client = MagicMock(spec=httpx.AsyncClient) + client.get = AsyncMock(side_effect=_get) + storage = _mk_storage_mock() + + async def _run(): + with _public_addr_patch(): + return await discover_authorization_server( + server_name="srv-x", + server_url="https://mcp.example.com/sse", + override_url="https://as.example.com", + cached_issuer=None, + http_client=client, + storage=storage, + server_id="srv-id", + trusted_hosts=frozenset(), + ) + + meta = asyncio.run(_run()) + assert meta.issuer == "https://as.example.com" + assert meta.token_endpoint == "https://as.example.com/token" + # Both candidate URLs were tried, RFC 8414 first then OIDC. + called = [c.args[0] for c in client.get.call_args_list] + assert any("oauth-authorization-server" in u for u in called) + assert any("openid-configuration" in u for u in called) + + # --------------------------------------------------------------------------- # Caching # --------------------------------------------------------------------------- diff --git a/tests/test_mcp_pool_auth_introspection.py b/tests/test_mcp_pool_auth_introspection.py index a123581c..3dd9ccbd 100644 --- a/tests/test_mcp_pool_auth_introspection.py +++ b/tests/test_mcp_pool_auth_introspection.py @@ -1535,5 +1535,220 @@ def test_call_tool_sync_does_not_wrap_non_structured_string( assert result == payload +class TestPoolPrimingAndTokenRotation: + """Per-user pool priming (PR #706 follow-up) and the bound-token rotation + reconnect. Priming must be NON-DESTRUCTIVE — it must never drive a token + refresh whose transient failure would revoke consent.""" + + def _wire(self, mgr: MCPClientManager, storage: SQLiteBackend, cipher: Any) -> None: + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + mgr._oauth_user_server_names = {"pool-srv"} + + def test_prime_user_pools_warms_fresh_token_server( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher, expires_in_seconds=3600, access_token="bearer-fresh") + self._wire(mgr, storage, cipher) + + primed: list[tuple[tuple[str, str], str]] = [] + + async def _fake_prime( + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + ) -> int: + primed.append((key, token)) + return 3 + + mgr._prime_user_server = _fake_prime.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + _run_on_loop(loop, mgr._prime_user_pools("user-1")) + + assert primed == [(("user-1", "pool-srv"), "bearer-fresh")] + + def test_prime_user_pools_skips_near_expiry_without_revoking( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """bug-1 regression: a near-expiry token is skipped (not refreshed), so a + transient refresh failure during priming can never revoke the token.""" + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + # Inside the 60s refresh-skew window -> the refreshing lookup would have + # driven a refresh here. + _seed_user_token(storage, cipher, expires_in_seconds=5, access_token="bearer-stale") + self._wire(mgr, storage, cipher) + + primed: list[tuple[str, str]] = [] + + async def _fake_prime( + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + ) -> int: + primed.append(key) + return 0 + + mgr._prime_user_server = _fake_prime.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + _run_on_loop(loop, mgr._prime_user_pools("user-1")) + + assert primed == [], "near-expiry token must be skipped, not primed (no refresh driven)" + # The token row must survive — priming must never revoke. + store = MCPTokenStore(storage, cipher, node_id="test") + assert store.get_user_token("user-1", "pool-srv") is not None + + def test_prime_user_pools_skips_already_connected( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + _seed_user_token(storage, cipher, expires_in_seconds=3600) + self._wire(mgr, storage, cipher) + + async def _seed() -> None: + entry = await mgr._ensure_pool_entry(("user-1", "pool-srv")) + entry.session = MagicMock() # already connected + + _run_on_loop(loop, _seed()) + + primed: list[tuple[str, str]] = [] + + async def _fake_prime( + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + ) -> int: + primed.append(key) + return 0 + + mgr._prime_user_server = _fake_prime.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + _run_on_loop(loop, mgr._prime_user_pools("user-1")) + + assert primed == [], "already-connected pool entry must be skipped" + + def test_schedule_prime_user_server_noop_for_non_oauth_user( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + self._wire(mgr, storage, cipher) + mgr._oauth_user_server_names = set() # nothing registered as oauth_user + + ran = threading.Event() + + async def _fake_logged( + self_inner: MCPClientManager, + key: tuple[str, str], + cfg: dict[str, Any], + token: str, + user_id: str, + server_name: str, + ) -> None: + ran.set() + + mgr._prime_user_server_logged = _fake_logged.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + mgr.schedule_prime_user_server( + user_id="user-1", server_name="not-oauth", access_token="t", server_row={} + ) + # Give any erroneously-scheduled coroutine a chance to run. + _run_on_loop(loop, asyncio.sleep(0.05)) + assert not ran.is_set(), "non-oauth_user server must not schedule a prime" + + def test_schedule_prime_user_server_runs_for_oauth_user( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + self._wire(mgr, storage, cipher) + + captured: dict[str, Any] = {} + done = threading.Event() + + async def _fake_prime( + self_inner: MCPClientManager, key: tuple[str, str], cfg: dict[str, Any], token: str + ) -> int: + captured["key"] = key + captured["token"] = token + done.set() + return 5 + + mgr._prime_user_server = _fake_prime.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + server_row = storage.get_mcp_server_by_name("pool-srv") + mgr.schedule_prime_user_server( + user_id="user-1", + server_name="pool-srv", + access_token="bearer-x", + server_row=server_row, + ) + assert done.wait(timeout=5), "scheduled prime did not run on the mcp-loop" + assert captured["key"] == ("user-1", "pool-srv") + assert captured["token"] == "bearer-x" + + def test_dispatch_reconnects_when_bound_token_rotated( + self, running_loop_mgr, storage: SQLiteBackend + ) -> None: + """A warm session bound to a stale bearer is transparently reconnected + with the current token; the discovered catalog is retained.""" + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + _seed_oauth_server(storage, name="pool-srv") + # The CURRENT stored token the dispatch will resolve. + _seed_user_token(storage, cipher, expires_in_seconds=3600, access_token="bearer-new") + self._wire(mgr, storage, cipher) + + reconnect_tokens: list[str] = [] + + async def _ok_call_tool(name: str, args: dict[str, Any]) -> Any: + 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 = _ok_call_tool + entry.session = sess + entry.bound_token = "bearer-old" # connected with the OLD token + entry.tools = [{"name": "do_thing"}] # catalog already discovered + + _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, + auth_fired_event: Any = None, + ) -> Any: + reconnect_tokens.append(access_token) + entry = await self_inner._ensure_pool_entry(key) + sess = MagicMock() + sess.call_tool = _ok_call_tool + entry.session = sess + entry.bound_token = access_token + return entry + + mgr._connect_one_pool = _fake_connect.__get__(mgr, type(mgr)) # type: ignore[method-assign] + + result = mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5) + + assert result == "ok" + # Stale bound token (bearer-old) != resolved token (bearer-new) -> exactly + # one reconnect carrying the current bearer. + assert reconnect_tokens == ["bearer-new"] + # Catalog retained across the in-place rotation. + entry = mgr._user_pool_entries[("user-1", "pool-srv")] + assert entry.tools == [{"name": "do_thing"}] + + # Suppress unused-import warning for AsyncMock. _ = AsyncMock diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 83075bda..e67ae9ee 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -138,6 +138,12 @@ _MAX_RESOURCES_PER_SERVER = 1000 _MAX_RESOURCE_TEMPLATES_PER_SERVER = 1000 _MAX_PROMPTS_PER_SERVER = 1000 +# Upper bound on concurrent per-user pool primes at ChatSession start. Servers +# are warmed in parallel (so one slow/unreachable upstream can't stall the rest) +# but capped so a deployment with many oauth_user servers doesn't fire a +# thundering herd of connects on every session start. +_PRIME_MAX_CONCURRENCY = 4 + @dataclass class _AuthCapture: @@ -1628,46 +1634,62 @@ class MCPClientManager: ) return len(fresh.tools or []) - async def prime_user_server( + def schedule_prime_user_server( self, *, user_id: str, server_name: str, access_token: str, server_row: dict[str, Any], - timeout: float = 20.0, - ) -> bool: - """Best-effort warm of a ``(user, server)`` pool so the user's tool - catalog populates immediately (e.g. right after OAuth consent). + ) -> None: + """Fire-and-forget warm of a ``(user, server)`` pool — e.g. right after + OAuth consent — so the user's tool catalog populates immediately WITHOUT + holding the consent redirect on a slow/unreachable MCP server. - Returns ``False`` (no-op) for non-``oauth_user`` servers, before the - mcp-loop is running, or with no token. NEVER raises: a prime failure - must not change the user-observable consent/redirect outcome, and the - lazy dispatch path remains the backstop. Schedules the connect onto - the mcp-loop (this is called from the request loop) and awaits it - without blocking that loop. + No-op for non-``oauth_user`` servers, before the mcp-loop is running, or + with no token. Schedules the connect onto the mcp-loop and returns at + once; the per-user tool listeners deliver the catalog to live sessions + when the prime completes, and lazy dispatch remains the backstop. """ if server_name not in self._oauth_user_server_names: - return False + return loop = self._loop if loop is None or not access_token: - return False + return cfg = _pool_cfg_from_row(server_row) key = (user_id, server_name) try: - fut = asyncio.run_coroutine_threadsafe( - self._prime_user_server(key, cfg, access_token), loop + asyncio.run_coroutine_threadsafe( + self._prime_user_server_logged(key, cfg, access_token, user_id, server_name), + loop, ) - count = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout) - log.info( - "mcp pool primed user=%s server=%s tools=%d", user_id, server_name, count - ) - return True - except Exception as exc: + except RuntimeError: + # mcp-loop is shutting down — skip; lazy dispatch is the backstop. + log.debug("mcp pool prime skipped: loop closed user=%s server=%s", user_id, server_name) + + async def _prime_user_server_logged( + self, + key: tuple[str, str], + cfg: dict[str, Any], + access_token: str, + user_id: str, + server_name: str, + ) -> None: + """Best-effort body scheduled by :meth:`schedule_prime_user_server`. + + Runs on the mcp-loop and never lets an exception escape onto it: a prime + failure must not change the user-observable consent outcome. + """ + try: + count = await self._prime_user_server(key, cfg, access_token) + log.info("mcp pool primed user=%s server=%s tools=%d", user_id, server_name, count) + except Exception: log.warning( - "mcp pool prime failed user=%s server=%s: %s", user_id, server_name, exc + "mcp pool prime failed user=%s server=%s", + user_id, + server_name, + exc_info=True, ) - return False def prime_user_pools(self, user_id: str) -> None: """Fire-and-forget: warm THIS user's consented ``oauth_user`` pools. @@ -1686,42 +1708,67 @@ class MCPClientManager: return # run_coroutine_threadsafe keeps the task referenced by the loop while # it runs, so no strong-ref bookkeeping is needed here. - asyncio.run_coroutine_threadsafe(self._prime_user_pools(user_id), self._loop) + try: + asyncio.run_coroutine_threadsafe(self._prime_user_pools(user_id), self._loop) + except RuntimeError: + # mcp-loop is shutting down — skip; lazy dispatch is the backstop. + log.debug("mcp pool prime skipped: loop closed user=%s", user_id) async def _prime_user_pools(self, user_id: str) -> None: - """Per-server body of :meth:`prime_user_pools` (runs on the mcp-loop).""" - from turnstone.core.mcp_oauth import get_user_access_token_classified + """Warm THIS user's consented ``oauth_user`` pools (runs on the mcp-loop). - for server_name in list(self._oauth_user_server_names): - try: - key = (user_id, server_name) - entry = self._user_pool_entries.get(key) - if entry is not None and entry.session is not None: - continue # already connected — nothing to do - lookup = await get_user_access_token_classified( - app_state=self._app_state, user_id=user_id, server_name=server_name - ) - if lookup.kind != "token" or not lookup.token: - continue # no usable token (not consented) — lazy paths handle it - server_row = await asyncio.to_thread( - self._storage.get_mcp_server_by_name, server_name - ) - if not server_row: - continue - cfg = _pool_cfg_from_row(server_row) - await self._prime_user_server(key, cfg, lookup.token) - log.info( - "mcp pool auto-primed at session start user=%s server=%s", - user_id, - server_name, - ) - except Exception: - log.debug( - "mcp pool auto-prime failed user=%s server=%s", - user_id, - server_name, - exc_info=True, - ) + Best-effort and NON-DESTRUCTIVE: each token is read directly (NOT via the + refresh state machine) and missing/near-expiry tokens are skipped, so a + transient AS/network failure during priming can never delete a token and + force re-consent — a refresh that may fail belongs on the lazy dispatch + path, driven by actual use. Servers are primed concurrently under + ``_PRIME_MAX_CONCURRENCY`` so one slow/unreachable upstream can't stall + the rest. + """ + from turnstone.core.mcp_oauth import _token_needs_refresh + + token_store = getattr(self._app_state, "mcp_token_store", None) + if token_store is None: + return + sem = asyncio.Semaphore(_PRIME_MAX_CONCURRENCY) + + async def _prime_one(server_name: str) -> None: + async with sem: + try: + key = (user_id, server_name) + entry = self._user_pool_entries.get(key) + if entry is not None and entry.session is not None: + return # already connected — nothing to do + # Non-refreshing read — priming must not drive a refresh whose + # transient failure would revoke the token (see docstring). + plain = await asyncio.to_thread( + token_store.get_user_token, user_id, server_name + ) + if plain is None or not plain.get("access_token"): + return # no usable token (not consented) — lazy paths handle it + if _token_needs_refresh(plain.get("expires_at")): + return # near expiry — let lazy dispatch refresh on actual use + server_row = await asyncio.to_thread( + self._storage.get_mcp_server_by_name, server_name + ) + if not server_row: + return + cfg = _pool_cfg_from_row(server_row) + await self._prime_user_server(key, cfg, plain["access_token"]) + log.info( + "mcp pool auto-primed at session start user=%s server=%s", + user_id, + server_name, + ) + except Exception: + log.debug( + "mcp pool auto-prime failed user=%s server=%s", + user_id, + server_name, + exc_info=True, + ) + + await asyncio.gather(*(_prime_one(s) for s in list(self._oauth_user_server_names))) # -- pool eviction -------------------------------------------------------- @@ -4872,9 +4919,7 @@ class MCPClientManager: # bearer). A directly-injected session (None bind token) is left # to the existing auth_401 retry path — and is the shape unit # tests use, so this avoids spurious reconnects there. - log.debug( - "mcp_pool.token_rotated_reconnect user=%s server=%s", key[0], key[1] - ) + log.debug("mcp_pool.token_rotated_reconnect user=%s server=%s", key[0], key[1]) session = None if session is None: # Lazy connect — also covers post-eviction recovery. diff --git a/turnstone/core/mcp_oauth.py b/turnstone/core/mcp_oauth.py index 50cdb233..dafbf60f 100644 --- a/turnstone/core/mcp_oauth.py +++ b/turnstone/core/mcp_oauth.py @@ -255,23 +255,32 @@ async def _fetch_as_metadata( # path — refusing to support them would lock out the most common # enterprise AS. Try RFC 8414 first (preferred), then OIDC. base = issuer.rstrip("/") + # (profile, url): ``profile`` records WHICH discovery document each candidate + # is so the S256 PKCE check below can apply the correct per-document + # defaulting rule — an absent ``code_challenge_methods_supported`` is only + # treated as "S256 supported" for the OIDC document (see below). metadata_candidates = ( - base + "/.well-known/oauth-authorization-server", - base + "/.well-known/openid-configuration", + ("rfc8414", base + "/.well-known/oauth-authorization-server"), + ("oidc", base + "/.well-known/openid-configuration"), ) resp = None + winning_profile: str | None = None last_status: int | None = None - for metadata_url in metadata_candidates: + for profile, metadata_url in metadata_candidates: try: r = await http_client.get(metadata_url, timeout=_DEFAULT_HTTP_TIMEOUT) except httpx.HTTPError as exc: raise MCPOAuthDiscoveryError(f"AS metadata fetch failed: {exc}") from exc if r.status_code == 200: resp = r + winning_profile = profile break last_status = r.status_code if resp is None: raise MCPOAuthDiscoveryError(f"AS metadata returned HTTP {last_status}") + # Which discovery profile answered (rfc8414 vs oidc) is the load-bearing + # detail when debugging an enterprise AS (e.g. Entra serves only OIDC). + log.debug("mcp_server.oauth.as_metadata_discovered", profile=winning_profile) if len(resp.content) > _MAX_DISCOVERY_BODY_BYTES: raise MCPOAuthDiscoveryError("AS metadata response body exceeds size limit") @@ -336,12 +345,22 @@ async def _fetch_as_metadata( if not isinstance(code_methods_raw, list): code_methods_raw = [] code_methods = tuple(str(m) for m in code_methods_raw) - if not code_methods: - # The OIDC discovery document does not require advertising - # code_challenge_methods_supported, and some IdPs (Entra) omit it - # despite fully supporting S256. Absence is not a denial — assume - # S256 (mandated by OAuth 2.1 / MCP auth) rather than locking the - # AS out. A NON-empty list missing S256 is still a hard refusal. + if not code_methods and winning_profile == "oidc": + # OIDC document (openid-configuration) only: it does not require + # advertising code_challenge_methods_supported, and some IdPs (Entra) + # omit it despite fully supporting S256, so treat absence as "S256 + # supported" (mandated by OAuth 2.1 / MCP auth) rather than locking the + # AS out. + # + # For the RFC 8414 oauth-authorization-server document we deliberately + # do NOT assume: an omitted field there is taken at face value as "no + # PKCE advertised", so code_methods stays empty and the check below + # fails closed. The client always sends code_challenge_method=S256, so + # this guard is the ONLY pre-flight that the AS actually enforces PKCE; + # assuming S256 on a document that omitted it would silently admit a + # non-enforcing AS and forfeit code-interception protection on the + # on-behalf-of bearer. A NON-empty list missing S256 is always a hard + # refusal, for both documents. log.info("mcp_server.oauth.s256_assumed_absent_advertisement") code_methods = ("S256",) if "S256" not in code_methods: @@ -2380,20 +2399,16 @@ async def _handle_mcp_oauth_callback_inner(request: Request) -> Response: # and the server is stuck "connecting". Best-effort: a prime failure # does not change consent success; lazy dispatch remains the backstop. mcp_client = getattr(request.app.state, "mcp_client", None) - if mcp_client is not None and hasattr(mcp_client, "prime_user_server"): - try: - await mcp_client.prime_user_server( - user_id=user_id, - server_name=server_name, - access_token=access_token, - server_row=server_row, - ) - except Exception: - log.debug( - "mcp_server.oauth.pool_prime_failed", - server_name=server_name, - exc_info=True, - ) + if mcp_client is not None and hasattr(mcp_client, "schedule_prime_user_server"): + # Fire-and-forget so the consent redirect is not held on a slow or + # unreachable MCP server; the warm runs on the mcp-loop in the + # background and live sessions pick up the catalog via the listeners. + mcp_client.schedule_prime_user_server( + user_id=user_id, + server_name=server_name, + access_token=access_token, + server_row=server_row, + ) # Phase 9 — clear any deferred-consent records for this (user, # server) now that consent has completed. Best-effort: a storage diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 3d133675..4727c2ff 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1213,7 +1213,11 @@ class ChatSession: try: mcp_client.prime_user_pools(self._mcp_user_id) except Exception: - pass + log.debug( + "mcp prime_user_pools scheduling failed user=%s", + self._mcp_user_id, + exc_info=True, + ) else: self._tools = INTERACTIVE_TOOLS self._task_tools = TASK_AGENT_TOOLS