From e38e573f7cb91caceb8e6017d518772cd219fbc7 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Sat, 11 Jul 2026 21:55:41 -0700 Subject: [PATCH] feat(mcp): route oauth_obo servers through the per-user pool Gate sweep of the pool-backed class: oauth_obo joins oauth_user at every pool-keying site, judged individually - - _obo_server_names sibling registry (reconcile + boot); priming, keep-alive sweep, and consent-flow sites deliberately keep iterating _oauth_user_server_names only (obo has no per-server consent; its keep-alive lands with the credential lifecycle work) - pool routing/status/static-health/tool-resolve gates use the shared is_user_scoped_auth predicate; status reports the real auth_type - dispatch: _pool_token_lookup routes oauth_obo to the mint engine; 'missing' detail becomes a re-login message (no per-server Connect URL is advertised - _build_consent_url already returns None) - _db_servers_to_config skips obo rows from static auto-connect (would handshake-fail with empty headers and trip the breaker) - web_search backend refusal covers both per-user auth types - console: oauth_obo in _MCP_AUTH_TYPES, https enforcement extended; startup key requirement counts obo rows (encrypted mint cache) Refs #551. --- turnstone/console/server.py | 12 +-- turnstone/core/mcp_client.py | 142 ++++++++++++++++++++++++----------- turnstone/core/mcp_crypto.py | 10 ++- turnstone/core/web_search.py | 10 ++- 4 files changed, 121 insertions(+), 53 deletions(-) diff --git a/turnstone/console/server.py b/turnstone/console/server.py index c051e2d3..e1776fe6 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -9682,7 +9682,7 @@ async def admin_registry_install(request: Request) -> JSONResponse: _MCP_NAME_RE = re.compile(r"^[a-zA-Z0-9._-]+$") _MCP_MAX_SERVERS = 200 # fallback; prefer cluster.mcp_max_servers from storage -_MCP_AUTH_TYPES = frozenset({"none", "static", "oauth_user"}) +_MCP_AUTH_TYPES = frozenset({"none", "static", "oauth_user", "oauth_obo"}) def _clean_oauth_text(value: Any, *, max_length: int = 512) -> str | None: @@ -9714,20 +9714,22 @@ def _parse_auth_type(body: dict[str, Any]) -> tuple[str | None, JSONResponse | N auth_type = str(body["auth_type"]).strip() if auth_type not in _MCP_AUTH_TYPES: return None, JSONResponse( - {"error": "auth_type must be 'none', 'static', or 'oauth_user'"}, + {"error": "auth_type must be 'none', 'static', 'oauth_user', or 'oauth_obo'"}, status_code=400, ) return auth_type, None def _enforce_oauth_user_https(auth_type: str, url: str | None) -> JSONResponse | None: - """Reject oauth_user MCP server URLs that aren't https (or loopback http). + """Reject per-user-auth MCP server URLs that aren't https (or loopback http). - Belt-and-braces with :func:`turnstone.core.mcp_client._validate_oauth_user_url` + Applies to both pool-backed types — ``oauth_user`` and ``oauth_obo`` — + since both transmit per-user bearers to the server URL. Belt-and-braces + with :func:`turnstone.core.mcp_client._validate_oauth_user_url` so a misconfigured row never persists. Returns the error response or ``None`` when the input is acceptable. """ - if auth_type != "oauth_user": + if auth_type not in ("oauth_user", "oauth_obo"): return None if not url: return None # presence checked elsewhere; this helper only checks scheme diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 2b2af2ca..84bda7db 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -55,7 +55,9 @@ from turnstone.core.mcp_http_parsers import ( from turnstone.core.mcp_oauth import ( TokenLookupResult, emit_oauth_failure_audit, + get_obo_access_token_classified, get_user_access_token_classified, + is_user_scoped_auth, ) if TYPE_CHECKING: @@ -766,6 +768,12 @@ class MCPClientManager: # backend resolution) can answer "is this server pool-backed?" # without a SQL roundtrip. self._oauth_user_server_names: set[str] = set() + # Sibling registry for auth_type='oauth_obo' servers (issue #551): + # also pool-backed (per-user sessions), but with NO per-server + # consent flow — tokens mint from the user's single captured + # credential, so the priming / keep-alive-sweep / consent sites that + # iterate ``_oauth_user_server_names`` deliberately exclude these. + self._obo_server_names: set[str] = set() # Idle-eviction task handle. Scheduled lazily on the mcp-loop the # first time a pool entry is created (start() runs before pool @@ -4470,12 +4478,13 @@ class MCPClientManager: :meth:`_oauth_user_server_status`). Both are ignored for static servers, whose session is process-global. """ - # auth_type='oauth_user' servers hold NO process-global session — they - # are warmed per-user into the pool — so the static-session check below - # would always report them "connecting". Derive their status from the - # REQUESTING user's warm pool entry instead, so the console pill reflects - # that user's real reachability once their pool is primed. - if name in self._oauth_user_server_names: + # Pool-backed servers (oauth_user / oauth_obo) hold NO process-global + # session — they are warmed per-user into the pool — so the + # static-session check below would always report them "connecting". + # Derive their status from the REQUESTING user's warm pool entry + # instead, so the console pill reflects that user's real reachability + # once their pool is primed. + if self._is_pool_server(name): return self._oauth_user_server_status(name, user_id, aggregate=aggregate) state = self._static_servers.get(name) connected = state is not None and state.session is not None @@ -4560,7 +4569,7 @@ class MCPClientManager: "url": "", "circuit_open": cb_open, "consecutive_failures": self._consecutive_failures.get(name, 0), - "auth_type": "oauth_user", + "auth_type": self.server_auth_type(name) or "oauth_user", "user_pools": len(warm), "last_refresh_at": last_refresh[0] if last_refresh is not None else None, "last_refresh_outcome": last_refresh[1] if last_refresh is not None else None, @@ -4580,7 +4589,7 @@ class MCPClientManager: result: dict[str, dict[str, Any]] = {} for name in list(self._server_configs): result[name] = self.get_server_status(name, user_id, aggregate=aggregate) - for name in list(self._oauth_user_server_names): + for name in list(self._oauth_user_server_names | self._obo_server_names): if name not in result: result[name] = self.get_server_status(name, user_id, aggregate=aggregate) return result @@ -4612,6 +4621,9 @@ class MCPClientManager: self._oauth_user_server_names = { row["name"] for row in rows if row.get("auth_type") == "oauth_user" } + self._obo_server_names = { + row["name"] for row in rows if row.get("auth_type") == "oauth_obo" + } desired = _db_servers_to_config(rows) desired_names = set(desired) @@ -4811,18 +4823,27 @@ class MCPClientManager: return user_map is not None and name in user_map def server_auth_type(self, server_name: str) -> str | None: - """Return ``'oauth_user'`` for pool-backed servers, else ``None``. + """Return the pool-backed auth type (``'oauth_user'`` / ``'oauth_obo'``), else ``None``. In-memory accessor for the per-turn callers that need to distinguish pool-backed servers from static-path ones without a SQL roundtrip. ``None`` means "either static-path or unknown" — the boot-time / per-node web_search resolver only uses this as - a defence-in-depth gate, so a missing-cache miss is safe (the - outer ``is_mcp_tool`` check already proves the server is in - ``_tool_map``, which by construction excludes oauth_user). + a defence-in-depth gate (via :func:`is_user_scoped_auth`), so a + missing-cache miss is safe (the outer ``is_mcp_tool`` check + already proves the server is in ``_tool_map``, which by + construction excludes pool-backed servers). Populated by ``reconcile_sync`` and ``create_mcp_client``. """ - return "oauth_user" if server_name in self._oauth_user_server_names else None + if server_name in self._oauth_user_server_names: + return "oauth_user" + if server_name in self._obo_server_names: + return "oauth_obo" + return None + + def _is_pool_server(self, server_name: str) -> bool: + """True when *server_name* uses per-user pool sessions (no static session).""" + return server_name in self._oauth_user_server_names or server_name in self._obo_server_names @property def server_count(self) -> int: @@ -4962,7 +4983,7 @@ class MCPClientManager: # Per-user pools are managed separately; ``__`` names can never # connect (``_connect_one``'s reserved-delimiter guard), so retrying # them forever would only spam ``log.error`` every interval. - if name not in self._oauth_user_server_names and "__" not in name + if not self._is_pool_server(name) and "__" not in name ] if not names: return self._static_health_check_s @@ -5422,10 +5443,40 @@ class MCPClientManager: if not server_name or not original: return None row = self._lookup_server_row(server_name) - if row is None or row.get("auth_type") != "oauth_user": + if row is None or not is_user_scoped_auth(row.get("auth_type")): return None return server_name, original, row + async def _pool_token_lookup( + self, + server_row: dict[str, Any], + user_id: str, + server_name: str, + *, + force_refresh: bool, + ) -> TokenLookupResult: + """Classified token lookup routed by the server's auth model. + + ``oauth_obo`` servers mint from the user's single captured + credential (:func:`get_obo_access_token_classified`); everything + else keeps the per-(user, server) refresh-grant path. Both share + the ``TokenLookupResult`` vocabulary, so the dispatcher's error + mapping below is auth-model-agnostic. + """ + if str(server_row.get("auth_type") or "") == "oauth_obo": + return await get_obo_access_token_classified( + app_state=self._app_state, + user_id=user_id, + server_name=server_name, + force_refresh=force_refresh, + ) + return await get_user_access_token_classified( + app_state=self._app_state, + user_id=user_id, + server_name=server_name, + force_refresh=force_refresh, + ) + def _lookup_server_row(self, server_name: str) -> dict[str, Any] | None: """Return the ``mcp_servers`` row for *server_name*, or None.""" if self._storage is None: @@ -5470,7 +5521,7 @@ class MCPClientManager: else: server_name = mapping[0] row = self._lookup_server_row(server_name) - if row is None or row.get("auth_type") != "oauth_user": + if row is None or not is_user_scoped_auth(row.get("auth_type")): return None return server_name, uri, row @@ -5496,7 +5547,7 @@ class MCPClientManager: if not server_name or not original: return None row = self._lookup_server_row(server_name) - if row is None or row.get("auth_type") != "oauth_user": + if row is None or not is_user_scoped_auth(row.get("auth_type")): return None return server_name, original, row @@ -5905,17 +5956,14 @@ class MCPClientManager: # 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, - force_refresh=retry_count > 0, + lookup: TokenLookupResult = await self._pool_token_lookup( + server_row, user_id, server_name, force_refresh=retry_count > 0 ) if lookup.kind == "missing": return _structured_error( code="mcp_consent_required", server=server_name, - detail="No token for user. Consent flow required.", + detail=_consent_missing_detail(server_row), consent_url=_build_consent_url(server_row), ) if lookup.kind == "decrypt_failure": @@ -6090,17 +6138,14 @@ class MCPClientManager: if self._app_state is None: raise RuntimeError("Pool dispatch requires set_app_state() to have been called") - lookup: TokenLookupResult = await get_user_access_token_classified( - app_state=self._app_state, - user_id=user_id, - server_name=server_name, - force_refresh=retry_count > 0, + lookup: TokenLookupResult = await self._pool_token_lookup( + server_row, user_id, server_name, force_refresh=retry_count > 0 ) if lookup.kind == "missing": return _structured_error( code="mcp_consent_required", server=server_name, - detail="No token for user. Consent flow required.", + detail=_consent_missing_detail(server_row), consent_url=_build_consent_url(server_row), ) if lookup.kind == "decrypt_failure": @@ -6249,17 +6294,14 @@ class MCPClientManager: if self._app_state is None: raise RuntimeError("Pool dispatch requires set_app_state() to have been called") - lookup: TokenLookupResult = await get_user_access_token_classified( - app_state=self._app_state, - user_id=user_id, - server_name=server_name, - force_refresh=retry_count > 0, + lookup: TokenLookupResult = await self._pool_token_lookup( + server_row, user_id, server_name, force_refresh=retry_count > 0 ) if lookup.kind == "missing": return _structured_error( code="mcp_consent_required", server=server_name, - detail="No token for user. Consent flow required.", + detail=_consent_missing_detail(server_row), consent_url=_build_consent_url(server_row), ) if lookup.kind == "decrypt_failure": @@ -7126,6 +7168,19 @@ def _build_consent_url( return f"/v1/api/mcp/oauth/start?{qs}" +def _consent_missing_detail(server_row: dict[str, Any]) -> str: + """User-facing detail for a ``missing`` token lookup, per auth model. + + ``oauth_obo`` has no per-server consent flow — the missing thing is the + user's captured sign-in credential, so the affordance is a re-login + (``_build_consent_url`` already returns None for these servers, so no + per-server Connect button is advertised). + """ + if server_row.get("auth_type") == "oauth_obo": + return "No sign-in credential for this account. Sign in to Turnstone again to reconnect." + return "No token for user. Consent flow required." + + def _pool_cfg_from_row(row: dict[str, Any]) -> dict[str, Any]: """Build a streamable-http MCP-client cfg from an ``mcp_servers`` row. @@ -7164,15 +7219,15 @@ def _pool_cfg_from_row(row: dict[str, Any]) -> dict[str, Any]: def _db_servers_to_config(rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]: """Convert mcp_servers DB rows to the config dict format. - Skips ``auth_type='oauth_user'`` rows: they need per-user bearer - tokens fetched at dispatch time via the OAuth flow, so auto-connecting - them at startup with empty headers fails handshake and trips the - circuit breaker. The upcoming per-user pool integration brings them - online lazily once a user has consented. + Skips pool-backed rows (``auth_type='oauth_user'`` / ``'oauth_obo'``): + they need per-user bearer tokens fetched at dispatch time, so + auto-connecting them at startup with empty headers fails handshake and + trips the circuit breaker. The per-user pool brings them online lazily — + on consent for oauth_user, on first dispatch for oauth_obo. """ result: dict[str, dict[str, Any]] = {} for row in rows: - if row.get("auth_type") == "oauth_user": + if is_user_scoped_auth(row.get("auth_type")): continue name = row["name"] cfg: dict[str, Any] = {"type": row["transport"]} @@ -7264,14 +7319,16 @@ def create_mcp_client( # Check DB first to know which servers are DB-managed db_names: set[str] = set() oauth_user_names: set[str] = set() + obo_names: set[str] = set() if storage is not None: try: rows = storage.list_mcp_servers(enabled_only=True) if rows: db_names = {r["name"] for r in rows} - # Cache oauth_user names so per-turn callers (web_search + # Cache pool-backed names so per-turn callers (web_search # backend resolution) can answer auth_type without SQL. oauth_user_names = {r["name"] for r in rows if r.get("auth_type") == "oauth_user"} + obo_names = {r["name"] for r in rows if r.get("auth_type") == "oauth_obo"} except Exception: log.warning("Failed to load DB-managed MCP servers", exc_info=True) @@ -7283,5 +7340,6 @@ def create_mcp_client( # Mark DB-sourced servers so reconcile_sync won't remove config-file servers mgr._db_managed = {name for name in servers if name in db_names} mgr._oauth_user_server_names = oauth_user_names + mgr._obo_server_names = obo_names mgr.start() return mgr diff --git a/turnstone/core/mcp_crypto.py b/turnstone/core/mcp_crypto.py index 4c00c636..74af1a46 100644 --- a/turnstone/core/mcp_crypto.py +++ b/turnstone/core/mcp_crypto.py @@ -573,13 +573,19 @@ def initialize_mcp_crypto_state(app_state: object, *, node_id: str = "") -> None raise SystemExit(1) from exc storage = get_storage() + # Both pool-backed types persist encrypted per-user rows (oauth_user: + # tokens + refresh; oauth_obo: minted-token cache), so both force the + # key requirement. Literal tuple rather than mcp_oauth's + # USER_SCOPED_AUTH_TYPES: importing mcp_oauth here would be a cycle. oauth_user_count = sum( - 1 for row in storage.list_mcp_servers() if row.get("auth_type") == "oauth_user" + 1 + for row in storage.list_mcp_servers() + if row.get("auth_type") in ("oauth_user", "oauth_obo") ) if oauth_user_count > 0 and cipher_cfg is None: log.error( - "mcp.oauth: %d server(s) configured with auth_type='oauth_user' but no " + "mcp.oauth: %d server(s) configured with auth_type='oauth_user'/'oauth_obo' but no " "[security] mcp_token_encryption_keys (rotation list) or " "mcp_token_encryption_key (single) in config.toml. Generate a key with: " "python -c 'from cryptography.fernet import Fernet; " diff --git a/turnstone/core/web_search.py b/turnstone/core/web_search.py index 50cac666..f7cac17a 100644 --- a/turnstone/core/web_search.py +++ b/turnstone/core/web_search.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Any, Protocol import httpx from turnstone.core.log import get_logger +from turnstone.core.mcp_oauth import is_user_scoped_auth if TYPE_CHECKING: from turnstone.core.mcp_client import MCPClientManager @@ -237,14 +238,15 @@ def resolve_web_search_client( # the bearer can't be (RFC §3, invariant 8 corollary). if mcp_client.is_mcp_tool(prefixed): # Defence-in-depth: even if a future change widens - # ``_tool_map`` to include oauth_user names by accident, + # ``_tool_map`` to include pool-backed names by accident, # refuse the backend explicitly. ``server_auth_type`` # is an in-memory accessor — this resolver is invoked # per LLM turn, so a SQL hop here would amplify token - # cost on every chat round. - if mcp_client.server_auth_type(server) == "oauth_user": + # cost on every chat round. oauth_obo is equally + # unusable here: minting needs a signed-in user. + if is_user_scoped_auth(mcp_client.server_auth_type(server)): log.warning( - "web_search_backend %r points at oauth_user MCP server; " + "web_search_backend %r points at a per-user-auth MCP server; " "per-node web search cannot use per-user tokens — disabling", backend, )