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