From 76e9d0a4f110f3aca05eb299b2f5443de804da90 Mon Sep 17 00:00:00 2001 From: metaclassing <1027733+metaclassing@users.noreply.github.com> Date: Thu, 25 Jun 2026 21:26:47 -0500 Subject: [PATCH] Address remaining blockers to entra id on behalf of flow for user impersonation to mcp servers (#706) * This is a collection of little snippits to resolve all the OBO flow problems required to get this talking to entra id for on behalf of user impersonating to protected mcp servers. we make sure turnstone checks these mcp servers on startup, and address some of microsofts opinionated implementations of oauth2/oidc and metadata provided by the identity provider. * minor token timeout bugfix --------- Co-authored-by: root --- turnstone/core/mcp_client.py | 162 +++++++++++++++++++++++++++++++++++ turnstone/core/mcp_oauth.py | 61 +++++++++++-- turnstone/core/session.py | 11 +++ 3 files changed, 226 insertions(+), 8 deletions(-) diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index dcd22cf6..83075bda 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -406,6 +406,12 @@ class PoolEntryState: prompts: list[dict[str, Any]] | None = None last_used: float = 0.0 in_flight: int = 0 + # Access token this session's httpx client was connected with. The bearer is + # frozen into the client's STATIC headers at connect (_connect_one_pool), so + # when the stored token later refreshes we compare against this to detect a + # stale session and reconnect (rebind the current token) proactively — + # instead of replaying the stale bearer and eating a guaranteed upstream 401. + bound_token: str | None = None auth_capture: _AuthCapture = field(default_factory=_AuthCapture) # Set by the response hook when the carrier captures a 4xx; awaited # by ``_dispatch_pool_with_entry``'s race against ``call_tool``. @@ -1568,6 +1574,7 @@ class MCPClientManager: # recovery path. Same ordering invariant covers resources/prompts # (R16). entry.session = session + entry.bound_token = access_token # remember the bearer this session carries entry.last_used = time.monotonic() self._user_pool_last_used[key] = entry.last_used @@ -1585,6 +1592,137 @@ class MCPClientManager: self._notify_user_prompt_listeners(user_id) return entry + # -- pool priming --------------------------------------------------------- + + async def _prime_user_server( + self, key: tuple[str, str], cfg: dict[str, Any], access_token: str + ) -> int: + """Proactively connect a pool entry so its catalog populates into + ``get_tools(user_id)`` WITHOUT waiting for a tool dispatch. + + ``auth_type='oauth_user'`` tools are per-user and were previously + discovered ONLY lazily, on first dispatch (``_dispatch_pool_with_entry`` + at the ``session is None`` branch). That creates a chicken-and-egg: + the model can't emit a call for a tool it can't see, but the tool + only appears after a call connects the pool — so the per-user + catalog stays empty and the server is stuck "connecting". This + connects at a known-good moment (OAuth consent completion), commits + the catalog and fires the per-user tool/resource/prompt listeners so + any live :class:`ChatSession` refreshes its tool list. + + Idempotent: a no-op if a dispatch (or an earlier prime) already + established the session. MUST run on the mcp-loop; takes + ``entry.open_lock`` exactly like dispatch so it can't race a + concurrent connect/eviction. Returns the discovered tool count. + """ + entry = await self._ensure_pool_entry(key) + async with entry.open_lock: + if entry.session is not None: + return len(entry.tools or []) + fresh = await self._connect_one_pool( + key, + cfg, + access_token, + auth_capture=entry.auth_capture, + auth_fired_event=entry.auth_fired_event, + ) + return len(fresh.tools or []) + + async def 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). + + 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. + """ + if server_name not in self._oauth_user_server_names: + return False + loop = self._loop + if loop is None or not access_token: + return False + 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 + ) + 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: + log.warning( + "mcp pool prime failed user=%s server=%s: %s", user_id, server_name, exc + ) + return False + + def prime_user_pools(self, user_id: str) -> None: + """Fire-and-forget: warm THIS user's consented ``oauth_user`` pools. + + Called at ChatSession start so a per-user OAuth server's tools are + present automatically — no manual reconnect after a reboot/upgrade, and + no chicken-and-egg (the model can't dispatch a tool it can't see). Only + touches servers the user already has a stored token for; skips servers + already connected. Non-blocking: schedules onto the mcp-loop and returns + immediately. The per-user tool listeners (registered by ChatSession) + deliver the catalog to the live session when each prime completes. + """ + if not user_id or self._loop is None or not self._oauth_user_server_names: + return + if self._app_state is None or self._storage is None: + 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) + + 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 + + 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, + ) + # -- pool eviction -------------------------------------------------------- async def _user_pool_eviction_loop(self) -> None: @@ -4714,6 +4852,30 @@ class MCPClientManager: entry.auth_capture.www_authenticate = None entry.auth_fired_event.clear() session = entry.session + if ( + session is not None + and entry.bound_token is not None + and entry.bound_token != access_token + ): + # The stored token refreshed since this session connected. The + # pooled httpx client's Authorization header is frozen at connect + # (_connect_one_pool), so a warm session would replay the now- + # stale bearer and eat a guaranteed upstream 401 before the + # auth_401 retry could heal it. Proactively reconnect to rebind + # the current token: _connect_one_pool pre-closes the old + # stack/streams, and entry.tools is retained (catalog intact), so + # this is a transparent in-place token rotation — attempt #0 now + # carries a valid bearer. + # + # ``bound_token is not None`` gate: only rotate when this session + # was established through ``_connect_one_pool`` (which records the + # 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] + ) + session = None if session is None: # Lazy connect — also covers post-eviction recovery. fresh = await self._connect_one_pool( diff --git a/turnstone/core/mcp_oauth.py b/turnstone/core/mcp_oauth.py index 8e378642..50cdb233 100644 --- a/turnstone/core/mcp_oauth.py +++ b/turnstone/core/mcp_oauth.py @@ -249,14 +249,29 @@ async def _fetch_as_metadata( except OAuthSSRFError as exc: raise MCPOAuthDiscoveryError(f"AS issuer URL rejected: {exc}") from exc - metadata_url = issuer.rstrip("/") + "/.well-known/oauth-authorization-server" - try: - resp = 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 resp.status_code != 200: - raise MCPOAuthDiscoveryError(f"AS metadata returned HTTP {resp.status_code}") + # MCP auth permits OpenID Connect discovery as a fallback to RFC 8414. + # Major IdPs (notably Microsoft Entra) serve ONLY the OIDC document + # (.well-known/openid-configuration) and 404 the oauth-authorization-server + # path — refusing to support them would lock out the most common + # enterprise AS. Try RFC 8414 first (preferred), then OIDC. + base = issuer.rstrip("/") + metadata_candidates = ( + base + "/.well-known/oauth-authorization-server", + base + "/.well-known/openid-configuration", + ) + resp = None + last_status: int | None = None + for 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 + break + last_status = r.status_code + if resp is None: + raise MCPOAuthDiscoveryError(f"AS metadata returned HTTP {last_status}") if len(resp.content) > _MAX_DISCOVERY_BODY_BYTES: raise MCPOAuthDiscoveryError("AS metadata response body exceeds size limit") @@ -321,6 +336,14 @@ 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. + log.info("mcp_server.oauth.s256_assumed_absent_advertisement") + code_methods = ("S256",) if "S256" not in code_methods: raise MCPOAuthDiscoveryError( "AS metadata does not advertise S256 PKCE — refusing to proceed" @@ -2350,6 +2373,28 @@ async def _handle_mcp_oauth_callback_inner(request: Request) -> Response: }, ) + # Prime the per-user pool so the just-consented server's tools populate + # into this user's catalog immediately. Without this, oauth_user tools + # are discovered only lazily on first dispatch — but the agent can't + # emit a call for a tool it can't yet see, so the catalog stays empty + # 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, + ) + # Phase 9 — clear any deferred-consent records for this (user, # server) now that consent has completed. Best-effort: a storage # failure here doesn't change the user-observable callback success; diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 105ec477..3d133675 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1203,6 +1203,17 @@ class ChatSession: # for OTHER users do not wake this session. self._mcp_prompt_cb = self._on_mcp_prompts_changed mcp_client.add_prompt_listener(self._mcp_prompt_cb, user_id=self._mcp_user_id) + # Proactively warm this user's per-user OAuth (oauth_user) pools so + # their tools are present without a manual reconnect (e.g. after a + # reboot/upgrade, or right after consent). Fire-and-forget — the + # listeners registered just above deliver the catalog to this + # session once each prime completes. No-op for users with no + # consented oauth_user servers. + if self._mcp_user_id and hasattr(mcp_client, "prime_user_pools"): + try: + mcp_client.prime_user_pools(self._mcp_user_id) + except Exception: + pass else: self._tools = INTERACTIVE_TOOLS self._task_tools = TASK_AGENT_TOOLS