mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-28 06:44:51 -06:00
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.
This commit is contained in:
@@ -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
|
||||
|
||||
+100
-42
@@ -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
|
||||
|
||||
@@ -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; "
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user