From a30c25a1216b2ad42131cc13b9bc37f380ba2d45 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 17:14:57 -0700 Subject: [PATCH] refactor(mcp): extract the dcr_bridge token flow into bridge_token_flow.py discoverable_endpoints.py had grown to 2695 lines mixing FastAPI route handlers with the dcr_bridge token-flow logic, against the no-monster-files convention. This moves the bridge token flow (the litellm-key/user resolution, the SCIM revalidation gate, and the mint/refresh envelope logic with their types and error mappers) into a dedicated bridge_token_flow.py, leaving the route handlers and the shared exchange_token_with_server orchestrator in discoverable_endpoints.py importing from it Pure relocation, zero behavior change. The moved code is byte-verbatim except one type annotation quoted as a forward reference (_BridgeAuthorizationCode is used only for typing and imported under TYPE_CHECKING to avoid a cycle), and the new module imports nothing from discoverable_endpoints at runtime. 275 tests pass unchanged; the test patch targets for moved internals were repointed to the new module and verified to still apply --- .../mcp_server/bridge_token_flow.py | 694 ++++++++++++++++++ .../mcp_server/discoverable_endpoints.py | 689 +---------------- .../mcp_server/test_discoverable_endpoints.py | 76 +- 3 files changed, 741 insertions(+), 718 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/bridge_token_flow.py diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py new file mode 100644 index 00000000000..19048e2eb7c --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -0,0 +1,694 @@ +"""Bridge token flow: litellm identity resolution and the DCR-bridge oauth_delegate mint/refresh pipeline.""" + +import math +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Literal, Optional + +from fastapi import HTTPException, Request +from fastapi.responses import JSONResponse +from pydantic import SecretStr +from typing_extensions import assert_never + +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _BridgeAuthorizationCode + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, + EnvelopeKeys, + RefreshCredential, + UpstreamTokenGrant, + ) + from litellm.proxy._types import UserAPIKeyAuth + + +def _litellm_key_from_request(request: Request) -> Optional[str]: + """Return the LiteLLM API key presented on the request, or ``None``. + + Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code + send) as well as ``Authorization``; either may carry a bare token or ``Bearer ``. + ``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry + an OAuth/upstream bearer. + """ + for header_value in ( + request.headers.get("x-litellm-api-key"), + request.headers.get("Authorization") or request.headers.get("authorization"), + ): + if not header_value: + continue + value = header_value.strip() + if value.lower().startswith("bearer "): + value = value[7:].strip() + if value: + return value + return None + + +def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool: + """``True`` when the presented key is neither blocked nor past its expiry. + + The OAuth token endpoint is unauthenticated, so the presented key is validated here before it is + trusted; a revoked or expired key must not mint a bridge envelope or write a stored credential. + ``get_key_object`` resolves a row without these checks (the main ``user_api_key_auth`` pipeline + enforces them downstream, which this endpoint bypasses), so they are applied here. Deleted keys + are already rejected upstream, where ``get_key_object`` raises on a row that no longer exists. + + This is an active-state gate only; it deliberately does not require a ``user_id``. A valid + team-scoped or service-account key has no ``user_id`` yet is a legitimate credential, so gating + on ``user_id`` presence would wrongly reject it. Callers that need the user (the per-user token + store) derive it separately via :func:`_active_key_user_id`. + + Total by design: ``expires`` is typed ``str | datetime``, and an unparseable string would make + ``datetime.fromisoformat`` raise. Since the callers run this outside their key-resolution + ``try``, an uncaught parse error would surface as a 500 instead of the endpoint's fail-closed + behavior, so a malformed expiry is treated as inactive (return ``False``) rather than raising. + """ + if key_obj.blocked is True: + return False + expires = key_obj.expires + if expires is not None: + if isinstance(expires, datetime): + expiry = expires + else: + try: + expiry = datetime.fromisoformat(expires) + except (ValueError, TypeError): + return False + if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: + expiry = expiry.replace(tzinfo=timezone.utc) + if expiry < datetime.now(timezone.utc): + return False + return True + + +def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> str | None: + """The active key's ``user_id``, or ``None`` when the key is blocked/expired or simply has no + ``user_id`` (a team-scoped or service-account key). Used only by the per-user token store, which + needs a user to key the stored credential; the bridge mint uses the key hash and does not.""" + return key_obj.user_id if _key_is_active(key_obj) else None + + +@dataclass(frozen=True, slots=True) +class _ResolvedKey: + """An active litellm key resolved from the token request: its hash (the value ``get_key_object`` + and the cache/DB layer key the record by) and the live record.""" + + key_hash: str + key: "UserAPIKeyAuth" + + +_KeyResolutionFailure = Literal["no_active_key", "unavailable", "unresolvable"] +"""Why a token request yielded no active litellm key, kept distinct so a caller statuses each truthfully +instead of blaming the client for a gateway problem: +- ``no_active_key``: none was presented, or the presented key is unknown / blocked / expired (the + caller's request is at fault) +- ``unavailable``: the auth database was transiently unreachable while resolving (retryable) +- ``unresolvable``: the gateway cannot resolve identity right now (no DB connection, or an unexpected + error) -- a gateway fault, not the caller's +The classification mirrors admission's ``_reload_admitted_key`` so the mint (ingress) and admission +(egress) never disagree on the status of the same outage.""" + + +async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyResolutionFailure": + """Resolve the presented litellm key to an active key record, or say precisely why not. + + Single resolution path the OAuth token endpoint reuses, resolving authoritatively via + ``get_key_object`` (cache first, then DB). The failure is a value, not a bare ``None``, so a caller + can tell "the client sent no usable credential" (a request error) apart from "the gateway could not + check" (an infrastructure error) and status each truthfully; collapsing both to ``None`` is what let + a DB outage read as a 400. A resolved key is still gated by ``_key_is_active``, so a blocked or + expired key is ``no_active_key`` while a valid team-scoped or service-account key (no ``user_id``) + resolves. Classification mirrors admission's ``_reload_admitted_key``: no DB connection is a gateway + fault, a ``ProxyException`` / ``HTTPException`` from ``get_key_object`` is an unknown or invalid key, + a database-service-unavailable error is a retryable outage, and anything else is an unexpected + gateway fault.""" + token = _litellm_key_from_request(request) + if not token: + return "no_active_key" + from litellm.proxy._types import hash_token # noqa: PLC0415 # inline import avoids a module-load circular import + + return await _reload_active_key_by_hash(hash_token(token)) + + +async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResolutionFailure": + """Reload the live key record for ``key_hash`` (cache first, then DB) and gate it on active state, + returning the resolved key or a precise failure. Shared by the token request's presented-key + resolution (:func:`_resolve_active_litellm_key`, which hashes the presented key) and the refresh + path (which already holds the hash sealed in the refresh envelope), so both re-validate identity + through one active-key gate and one failure classification. Classification mirrors admission's + ``_reload_admitted_key``: no DB connection is a gateway fault, a ``ProxyException`` / ``HTTPException`` + from ``get_key_object`` is an unknown or invalid key, a database-service-unavailable error is a + retryable outage, and anything else is an unexpected gateway fault. A blocked or expired key is + ``no_active_key``, so a revoked key can neither mint nor refresh a bridge envelope.""" + from litellm.proxy._types import ( + ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import + ) + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_key_object, + ) + from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import + PrismaDBExceptionHandler, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return "unresolvable" + try: + key_obj = await get_key_object( + hashed_token=key_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except (ProxyException, HTTPException): + return "no_active_key" + except Exception as exc: # noqa: BLE001 # classify: a DB outage is retryable, anything else is an opaque gateway fault + if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): + return "unavailable" + verbose_logger.debug( + "_reload_active_key_by_hash: unexpected key-resolution error (%s)", + type(exc).__name__, + ) + return "unresolvable" + if not _key_is_active(key_obj): + return "no_active_key" + return _ResolvedKey(key_hash=key_hash, key=key_obj) + + +async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | None": + """Re-validate a live litellm user by id, returning ``None`` when the user is active or a precise + failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a + user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a + deactivated user cannot keep refreshing, mirroring how admission re-validates the same user subject on + the egress side. No DB connection is a gateway fault (``unresolvable``) and a + database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails + closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` / + ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` + catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look + identical, the original error surviving only as ``__context__``), so the outage check walks the cause + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.""" + from litellm.proxy._types import ( + ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import + ) + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_user_object, + ) + from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import + PrismaDBExceptionHandler, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return "unresolvable" + try: + user_object = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except (ProxyException, HTTPException): + return "no_active_key" + except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500 + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc): + return "unavailable" + verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__) + return "no_active_key" + if user_object is None: + return "no_active_key" + if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: + return "no_active_key" + return None + + +async def _key_owner_scim_deactivated(key: "UserAPIKeyAuth") -> bool: + """True only when the key's owning user was explicitly SCIM-deactivated, so a refresh revokes an + offboarded owner's key exactly as admission does via ``_reject_if_admitted_owner_scim_deactivated``. + A key with no owner, a missing owner record, or a failed lookup fails OPEN (returns ``False``), + matching admission and the standard builder: a key may outlive its owner record, and a transient DB + blip must not revoke a live key. Only an explicit ``scim_active`` of ``False`` gates renewal.""" + if key.user_id is None: + return False + from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import + get_user_object, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + return False + try: + owner = await get_user_object( + user_id=key.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except Exception as exc: # noqa: BLE001 # fail open: a missing owner (get_user_object's wrapped ValueError) or a DB blip must not revoke a live key + verbose_logger.debug("refresh: key-owner SCIM lookup failed, not revoking (%s)", type(exc).__name__) + return False + return owner is not None and isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False + + +async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResolutionFailure | None": + """Re-validate that the subject sealed in a refresh envelope is still live, dispatching on its type: + a key_hash reloads the virtual key, a user_id reloads the user. Returns ``None`` when the subject is + active or a precise failure otherwise, so revocation gates renewal for either identity source the same + way admission gates the egress: a blocked or expired key, a SCIM-deactivated key owner (mirroring + admission's owner check, so an offboarded user cannot keep renewing a still-active key), and a + deactivated or deleted user all fail closed to ``no_active_key``.""" + match identity.subject_type: + case "key_hash": + reloaded = await _reload_active_key_by_hash(identity.subject) + if not isinstance(reloaded, _ResolvedKey): + return reloaded + if await _key_owner_scim_deactivated(reloaded.key): + return "no_active_key" + return None + case "user_id": + return await _reload_active_user_by_id(identity.subject) + case _: + assert_never(identity.subject_type) + + +async def _extract_user_id_from_request(request: Request) -> str | None: + """The litellm ``user_id`` for the token request, so a per-user token is stored under the same + identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome + (including a transient DB outage) collapses to ``None`` here and the caller simply skips the store; + the bridge mint, which must status those outcomes differently, consumes + :func:`_resolve_active_litellm_key` directly.""" + resolved = await _resolve_active_litellm_key(request) + if not isinstance(resolved, _ResolvedKey): + return None + return _active_key_user_id(resolved.key) + + +_UpstreamGrantRejection = Literal["no_access_token", "expired_lifetime"] +"""Why an upstream token response cannot back a bridge envelope: +- ``no_access_token``: the response carries no usable ``access_token`` +- ``expired_lifetime``: the response reports a parseable, non-positive ``expires_in``, i.e. an upstream + token that is already dead, so sealing it would forward a bearer the edge cannot use +An absent or unparseable ``expires_in`` is NOT a rejection; the lifetime is merely unknown and the +envelope caps it, the by-design behaviour for an upstream that omits the field.""" + + +def _classify_upstream_lifetime(raw_expires_in: object) -> "int | Literal['unspecified', 'expired']": + """Classify an upstream ``expires_in`` into a positive number of seconds, ``"unspecified"`` (absent + or unparseable, so the envelope caps it), or ``"expired"`` (a non-positive value the upstream reports + as already elapsed). Telling "we do not know the lifetime" apart from "the upstream says it is + already dead" is what stops an explicitly-expired token from silently receiving the envelope's 1h + cap. The expired decision is made on the parsed numeric value, not on ``int(...)`` of it, so a + positive sub-second lifetime in ``(0, 1)`` is not truncated to ``0`` and misread as elapsed; the + envelope works in whole seconds, so such a lifetime clamps up to its 1s floor. ``bool`` is excluded + (an ``int`` subclass but never a real lifetime), and the conversions can raise on ``NaN`` / + ``Infinity`` / oversized input, which reads as unparseable rather than surfacing as a 500.""" + if raw_expires_in is None or isinstance(raw_expires_in, bool) or not isinstance(raw_expires_in, (int, float, str)): + return "unspecified" + try: + numeric = float(raw_expires_in) + seconds = int(numeric) + except (ValueError, TypeError, OverflowError): + return "unspecified" + if numeric <= 0: + return "expired" + return max(1, seconds) + + +def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenGrant | _UpstreamGrantRejection": + """Validate an upstream OAuth token response into a typed grant, or say why it cannot back an + envelope. Each field is isinstance-checked so nothing untyped from ``response.json()`` reaches the + grant. ``expires_in`` is read three ways (see :func:`_classify_upstream_lifetime`): an unknown + lifetime leaves the grant ``expires_in`` ``None`` for the envelope to cap, a positive value is + honoured, and an explicit already-elapsed value is a rejection rather than a silent fall-through to + the cap.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + UpstreamTokenGrant, + ) + + if not isinstance(token_response, dict): + return "no_access_token" + access = token_response.get("access_token") + if not isinstance(access, str) or not access: + return "no_access_token" + lifetime = _classify_upstream_lifetime(token_response.get("expires_in")) + if lifetime == "expired": + return "expired_lifetime" + token_type = token_response.get("token_type") + scope = token_response.get("scope") + return UpstreamTokenGrant( + access_token=SecretStr(access), + token_type=token_type if isinstance(token_type, str) and token_type else "Bearer", + # The upstream refresh_token is deliberately NOT sealed: the edge never consumes it (it forwards + # only token_type + access_token), so it would be dead weight embedding a long-lived upstream + # credential in the client-held bearer, and it enlarges the envelope. Refresh support is a + # follow-up (a dedicated refresh-envelope); the client re-runs authorization_code at the cap. + refresh_token=None, + scope=scope if isinstance(scope, str) and scope else None, + expires_in=lifetime if isinstance(lifetime, int) else None, + ) + + +# --------------------------------------------------------------------------- +# DCR-bridge oauth_delegate mint: a three-phase pipeline whose failures are values. +# +# prepare (before the upstream exchange) -> validate every precondition and resolve identity+keys +# exchange (the single-use upstream code is consumed here, in exchange_token_with_server) +# finish (after the exchange) -> seal the upstream grant into the client-held envelope +# +# Every precondition lives in ``prepare``, which runs BEFORE the exchange, so no failure can burn the +# single-use code or rotate a refresh token, for either grant type -- that whole class of bug is gone +# by construction rather than guarded case by case. Failures are values mapped to an OAuth-shaped +# response in one place (``_bridge_mint_error_response``), so status codes and the RFC 6749 §5.2 body +# shape are uniform. Adding a failure mode is a new literal plus a match arm the type checker forces. +# --------------------------------------------------------------------------- + +_BridgeMintError = Literal[ + "no_identity", + "invalid_refresh", + "identity_unavailable", + "identity_unresolvable", + "not_configured", + "no_upstream_token", + "upstream_token_expired", + "too_large", +] + + +@dataclass(frozen=True, slots=True) +class _BridgeMintReady: + """Everything the seal needs, resolved once before the exchange: the identity to bind the envelope + to and the master-key-derived envelope keys. The identity is a key_hash subject for the scripted + two-header client (resolved from the litellm key it presents) or a user_id subject for the + interactive SSO client (the user recovered from the gateway authorization code), so one phase-3 seal + serves both. Resolving identity here means ``_finish_bridge_mint`` has no preconditions left to + fail.""" + + identity: "EnvelopeIdentity" + keys: "EnvelopeKeys" + + +def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: + """Map a bridge-mint failure value to its token-endpoint response: one place, RFC 6749 §5.2 shape + (top-level ``error``, no-store headers) for every case, with a status truthful about where the + failure is. The caller's request is 400, a transient gateway outage is 503, a gateway + misconfiguration is 500, and an upstream problem is 502. The identity-resolution statuses match how + admission statuses the same conditions on the egress side, so mint and admit never disagree under + one outage.""" + match error: + case "no_identity": + status, code, desc = ( + 400, + "invalid_request", + "this server issues a gateway-bound credential; complete the interactive sign-in, or " + "send a litellm credential (x-litellm-api-key or Authorization) on the token request", + ) + case "invalid_refresh": + status, code, desc = ( + 400, + "invalid_grant", + "the refresh credential is not a valid, live refresh envelope for this server; " + "re-run authorization_code to obtain a new one", + ) + case "identity_unavailable": + status, code, desc = ( + 503, + "temporarily_unavailable", + "the authentication database is temporarily unreachable; retry shortly", + ) + case "identity_unresolvable": + status, code, desc = ( + 500, + "server_error", + "the gateway could not resolve the litellm identity for this request", + ) + case "not_configured": + status, code, desc = ( + 500, + "server_error", + "the gateway is not configured to mint a gateway-bound credential (master_key is not set)", + ) + case "no_upstream_token": + status, code, desc = ( + 502, + "server_error", + "the upstream token response has no usable access_token", + ) + case "upstream_token_expired": + status, code, desc = ( + 502, + "server_error", + "the upstream token response reports an already-expired lifetime", + ) + case "too_large": + status, code, desc = ( + 502, + "server_error", + "the upstream token is too large to seal into a gateway-bound credential", + ) + case _: + assert_never(error) + return JSONResponse( + status_code=status, content={"error": code, "error_description": desc}, headers=TOKEN_NO_CACHE_HEADERS + ) + + +def _key_resolution_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError: + """Lift an identity-resolution failure into the mint taxonomy, preserving origin so the status stays + truthful: the caller's missing credential is 400, a transient DB outage is 503, and a gateway that + cannot resolve identity is 500.""" + match failure: + case "no_active_key": + return "no_identity" + case "unavailable": + return "identity_unavailable" + case "unresolvable": + return "identity_unresolvable" + case _: + assert_never(failure) + + +def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _BridgeMintError: + """Lift an upstream-response rejection into the mint taxonomy; both are upstream faults (502).""" + match rejection: + case "no_access_token": + return "no_upstream_token" + case "expired_lifetime": + return "upstream_token_expired" + case _: + assert_never(rejection) + + +async def _prepare_bridge_mint( + request: Request, + mcp_server: MCPServer, + bridge_identity: "_BridgeAuthorizationCode | None" = None, +) -> "_BridgeMintReady | _BridgeMintError": + """Phase 1 for the authorization_code grant, BEFORE the upstream exchange: confirm the gateway can + mint (master_key set), resolve the litellm identity, and derive the envelope keys. Returns a ready + context or a precise failure value. Running before the exchange is what makes every failure here fail + closed without consuming the single-use code. + + Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged + authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway + authorization code) and mints a user subject. The scripted two-header client presents a litellm key + on the token request instead, so its identity is the active key's hash and mints a key_hash subject. + A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully; + neither source present is ``no_identity``. The refresh_token grant has its own phase-1 + (:func:`_prepare_bridge_refresh`), which recovers identity from the presented refresh envelope.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + key_hash_identity, + user_identity, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + master_key, + ) + + if not master_key: + return "not_configured" + keys = envelope_keys_from_master_key(master_key) + if bridge_identity is not None: + identity = user_identity(server_id=mcp_server.server_id, user_id=bridge_identity.litellm_user_id) + return _BridgeMintReady(identity=identity, keys=keys) + resolved = await _resolve_active_litellm_key(request) + if not isinstance(resolved, _ResolvedKey): + return _key_resolution_failure_to_mint_error(resolved) + identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved.key_hash) + return _BridgeMintReady(identity=identity, keys=keys) + + +@dataclass(frozen=True, slots=True) +class _BridgeRefreshReady: + """A validated refresh request: the identity+keys to mint the renewed pair under, the upstream refresh + token (unwrapped from the client's refresh envelope) to exchange with the upstream IdP, and the scope + sealed alongside it at mint. The upstream refresh token is a ``SecretStr`` like every other credential + in this layer, so a repr or a traceback that captures this value never exposes the raw upstream refresh + token in plaintext. ``upstream_scope`` carries the originally-granted scope so the renewal re-requests + it when the client (a DCR/MCP client that typically omits scope on refresh) sends none, keeping the + renewed token's scope stable against an upstream that would otherwise narrow or drop it.""" + + ready: "_BridgeMintReady" + upstream_refresh_token: SecretStr + upstream_scope: str | None = None + + +def _refresh_key_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError: + """Lift an identity-resolution failure on the refresh path into the mint taxonomy. Unlike the mint + path, a resolved-but-inactive (or unknown) key is ``invalid_grant`` rather than ``invalid_request``: + the client did present an identity (sealed in the refresh envelope), but it is no longer live, so the + refresh is invalid and the client must re-authenticate. A transient outage is still 503 and a gateway + fault still 500, matching the mint path and admission.""" + match failure: + case "no_active_key": + return "invalid_refresh" + case "unavailable": + return "identity_unavailable" + case "unresolvable": + return "identity_unresolvable" + case _: + assert_never(failure) + + +async def _prepare_bridge_refresh( + mcp_server: MCPServer, refresh_value: str | None +) -> "_BridgeRefreshReady | _BridgeMintError": + """Phase 1 for the refresh_token grant, BEFORE the upstream exchange: open the client's refresh + envelope, re-validate the sealed litellm identity so a revoked key cannot keep refreshing, and + recover the upstream refresh token to exchange. Identity comes entirely from the sealed envelope, not + the HTTP request, so the request object is not needed here. The client presents a refresh envelope, + never a raw upstream refresh token, so a missing value, a non-envelope, an unopenable envelope, or one + minted for another server is ``invalid_grant``. Running before the exchange means a rejected refresh + never consumes or rotates the upstream refresh token.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + BridgeRefreshOpened, + envelope_keys_from_master_key, + open_bridge_refresh_envelope, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + master_key, + ) + + if not master_key: + return "not_configured" + if not refresh_value: + return "invalid_refresh" + keys = envelope_keys_from_master_key(master_key) + opened = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id) + if not isinstance(opened, BridgeRefreshOpened): + return "invalid_refresh" + failure = await _revalidate_active_subject(opened.identity) + if failure is not None: + return _refresh_key_failure_to_mint_error(failure) + return _BridgeRefreshReady( + ready=_BridgeMintReady(identity=opened.identity, keys=keys), + upstream_refresh_token=opened.refresh.refresh_token, + upstream_scope=opened.refresh.scope, + ) + + +def _finish_bridge_mint( + ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime +) -> "JSONResponse | _BridgeMintError": + """Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held access envelope + using the pre-resolved identity and keys, and, when the upstream returned a refresh token, seal a + long-lived refresh envelope alongside it so the client can renew without re-authenticating. Shared by + the authorization_code and refresh_token paths, so a renewal that the upstream rotates re-issues a + fresh refresh envelope. The only hard failures here are properties of the upstream access token (no + usable token, an already-expired lifetime, or a token too large to seal); a refresh token that cannot + be sealed degrades to an access-only response rather than failing the whole exchange.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + build_bridge_token_response, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + SealedEnvelope, + UpstreamTokenGrant, + ) + + grant = _bridge_grant_from_token_response(token_response) + if not isinstance(grant, UpstreamTokenGrant): + return _upstream_rejection_to_mint_error(grant) + sealed = build_bridge_token_response(ready.identity, grant, ready.keys, now) + if not isinstance(sealed, SealedEnvelope): + return "too_large" + # Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the + # client is never told the bearer lives past the point admission (which uses that exp) rejects it. + expires_in = max(0, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp())) + refresh_envelope = _mint_refresh_envelope_value(ready.identity, token_response, ready.keys, now, mcp_server) + body = { + "access_token": sealed.token.get_secret_value(), + "token_type": "Bearer", + "expires_in": expires_in, + # A refresh envelope rides along only when the upstream returned a refresh token to seal; when it + # rotates on renewal, the client receives the new one and the old envelope's upstream token dies. + **({"refresh_token": refresh_envelope} if refresh_envelope is not None else {}), + } + return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS) + + +def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | None": + """Extract the upstream refresh grant from a token response, or ``None`` when there is none to seal. + Each field is isinstance-checked so nothing untyped reaches the refresh envelope; ``refresh_expires_in`` + (the refresh token's own lifetime, when the upstream reports it) is classified like ``expires_in`` and + bounds the refresh envelope's TTL. An upstream that reports the refresh token itself as already elapsed + (``refresh_expires_in`` non-positive) yields ``None`` rather than a refresh envelope: sealing a dead + token would hand the client a full-TTL-capped envelope the IdP will reject, so the exchange degrades to + an access-only response (the client re-authenticates at access expiry), mirroring how + :func:`_bridge_grant_from_token_response` refuses an already-elapsed access token instead of capping it.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + RefreshCredential, + ) + + if not isinstance(token_response, dict): + return None + refresh = token_response.get("refresh_token") + if not isinstance(refresh, str) or not refresh: + return None + lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in")) + if lifetime == "expired": + return None + scope = token_response.get("scope") + return RefreshCredential( + refresh_token=SecretStr(refresh), + scope=scope if isinstance(scope, str) and scope else None, + expires_in=lifetime if isinstance(lifetime, int) else None, + ) + + +def _mint_refresh_envelope_value( + identity: "EnvelopeIdentity", token_response: object, keys: "EnvelopeKeys", now: datetime, mcp_server: MCPServer +) -> str | None: + """Seal the upstream refresh grant (if any) into a refresh envelope and return its bearer string, or + ``None`` when the upstream returned no refresh token or the refresh token is too large to seal. A + too-large refresh token degrades to an access-only response (logged) rather than failing an exchange + that already succeeded upstream: the client simply re-authenticates when the access envelope expires.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + build_bridge_refresh_token_response, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + SealedEnvelope, + ) + + refresh_credential = _upstream_refresh_credential(token_response) + if refresh_credential is None: + return None + sealed = build_bridge_refresh_token_response(identity, refresh_credential, keys, now) + if isinstance(sealed, SealedEnvelope): + return sealed.token.get_secret_value() + verbose_logger.warning( + "bridge mint: the upstream refresh token is too large to seal into a refresh envelope for " + "server=%s; issuing an access-only response, so the client re-authenticates at access expiry", + mcp_server.server_id, + ) + return None diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index b6e9a094cd5..54aff86aab2 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,10 +1,8 @@ import asyncio import html as _html import json -import math import secrets import time -from dataclasses import dataclass from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse @@ -13,7 +11,6 @@ import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError -from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -24,6 +21,15 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( TokenEndpointAuthConfigError, build_token_endpoint_client_auth, ) +from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _bridge_mint_error_response, + _BridgeMintReady, + _BridgeRefreshReady, + _extract_user_id_from_request, + _finish_bridge_mint, + _prepare_bridge_mint, + _prepare_bridge_refresh, +) from litellm.proxy._experimental.mcp_server.faults import ( CallerRejected, CredentialSource, @@ -49,13 +55,7 @@ from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer if TYPE_CHECKING: - from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( - EnvelopeIdentity, - EnvelopeKeys, - RefreshCredential, - UpstreamTokenGrant, - ) - from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth + from litellm.proxy._types import LiteLLM_MCPServerTable # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. # Keeps us from hammering the upstream IdP on each discovery request. @@ -393,274 +393,6 @@ def _validate_token_response( ) -def _litellm_key_from_request(request: Request) -> Optional[str]: - """Return the LiteLLM API key presented on the request, or ``None``. - - Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code - send) as well as ``Authorization``; either may carry a bare token or ``Bearer ``. - ``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry - an OAuth/upstream bearer. - """ - for header_value in ( - request.headers.get("x-litellm-api-key"), - request.headers.get("Authorization") or request.headers.get("authorization"), - ): - if not header_value: - continue - value = header_value.strip() - if value.lower().startswith("bearer "): - value = value[7:].strip() - if value: - return value - return None - - -def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool: - """``True`` when the presented key is neither blocked nor past its expiry. - - The OAuth token endpoint is unauthenticated, so the presented key is validated here before it is - trusted; a revoked or expired key must not mint a bridge envelope or write a stored credential. - ``get_key_object`` resolves a row without these checks (the main ``user_api_key_auth`` pipeline - enforces them downstream, which this endpoint bypasses), so they are applied here. Deleted keys - are already rejected upstream, where ``get_key_object`` raises on a row that no longer exists. - - This is an active-state gate only; it deliberately does not require a ``user_id``. A valid - team-scoped or service-account key has no ``user_id`` yet is a legitimate credential, so gating - on ``user_id`` presence would wrongly reject it. Callers that need the user (the per-user token - store) derive it separately via :func:`_active_key_user_id`. - - Total by design: ``expires`` is typed ``str | datetime``, and an unparseable string would make - ``datetime.fromisoformat`` raise. Since the callers run this outside their key-resolution - ``try``, an uncaught parse error would surface as a 500 instead of the endpoint's fail-closed - behavior, so a malformed expiry is treated as inactive (return ``False``) rather than raising. - """ - if key_obj.blocked is True: - return False - expires = key_obj.expires - if expires is not None: - if isinstance(expires, datetime): - expiry = expires - else: - try: - expiry = datetime.fromisoformat(expires) - except (ValueError, TypeError): - return False - if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: - expiry = expiry.replace(tzinfo=timezone.utc) - if expiry < datetime.now(timezone.utc): - return False - return True - - -def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> str | None: - """The active key's ``user_id``, or ``None`` when the key is blocked/expired or simply has no - ``user_id`` (a team-scoped or service-account key). Used only by the per-user token store, which - needs a user to key the stored credential; the bridge mint uses the key hash and does not.""" - return key_obj.user_id if _key_is_active(key_obj) else None - - -@dataclass(frozen=True, slots=True) -class _ResolvedKey: - """An active litellm key resolved from the token request: its hash (the value ``get_key_object`` - and the cache/DB layer key the record by) and the live record.""" - - key_hash: str - key: "UserAPIKeyAuth" - - -_KeyResolutionFailure = Literal["no_active_key", "unavailable", "unresolvable"] -"""Why a token request yielded no active litellm key, kept distinct so a caller statuses each truthfully -instead of blaming the client for a gateway problem: -- ``no_active_key``: none was presented, or the presented key is unknown / blocked / expired (the - caller's request is at fault) -- ``unavailable``: the auth database was transiently unreachable while resolving (retryable) -- ``unresolvable``: the gateway cannot resolve identity right now (no DB connection, or an unexpected - error) -- a gateway fault, not the caller's -The classification mirrors admission's ``_reload_admitted_key`` so the mint (ingress) and admission -(egress) never disagree on the status of the same outage.""" - - -async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyResolutionFailure": - """Resolve the presented litellm key to an active key record, or say precisely why not. - - Single resolution path the OAuth token endpoint reuses, resolving authoritatively via - ``get_key_object`` (cache first, then DB). The failure is a value, not a bare ``None``, so a caller - can tell "the client sent no usable credential" (a request error) apart from "the gateway could not - check" (an infrastructure error) and status each truthfully; collapsing both to ``None`` is what let - a DB outage read as a 400. A resolved key is still gated by ``_key_is_active``, so a blocked or - expired key is ``no_active_key`` while a valid team-scoped or service-account key (no ``user_id``) - resolves. Classification mirrors admission's ``_reload_admitted_key``: no DB connection is a gateway - fault, a ``ProxyException`` / ``HTTPException`` from ``get_key_object`` is an unknown or invalid key, - a database-service-unavailable error is a retryable outage, and anything else is an unexpected - gateway fault.""" - token = _litellm_key_from_request(request) - if not token: - return "no_active_key" - from litellm.proxy._types import hash_token # noqa: PLC0415 # inline import avoids a module-load circular import - - return await _reload_active_key_by_hash(hash_token(token)) - - -async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResolutionFailure": - """Reload the live key record for ``key_hash`` (cache first, then DB) and gate it on active state, - returning the resolved key or a precise failure. Shared by the token request's presented-key - resolution (:func:`_resolve_active_litellm_key`, which hashes the presented key) and the refresh - path (which already holds the hash sealed in the refresh envelope), so both re-validate identity - through one active-key gate and one failure classification. Classification mirrors admission's - ``_reload_admitted_key``: no DB connection is a gateway fault, a ``ProxyException`` / ``HTTPException`` - from ``get_key_object`` is an unknown or invalid key, a database-service-unavailable error is a - retryable outage, and anything else is an unexpected gateway fault. A blocked or expired key is - ``no_active_key``, so a revoked key can neither mint nor refresh a bridge envelope.""" - from litellm.proxy._types import ( - ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import - ) - from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import - get_key_object, - ) - from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import - PrismaDBExceptionHandler, - ) - from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import - prisma_client, - user_api_key_cache, - ) - - if prisma_client is None: - return "unresolvable" - try: - key_obj = await get_key_object( - hashed_token=key_hash, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - except (ProxyException, HTTPException): - return "no_active_key" - except Exception as exc: # noqa: BLE001 # classify: a DB outage is retryable, anything else is an opaque gateway fault - if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): - return "unavailable" - verbose_logger.debug( - "_reload_active_key_by_hash: unexpected key-resolution error (%s)", - type(exc).__name__, - ) - return "unresolvable" - if not _key_is_active(key_obj): - return "no_active_key" - return _ResolvedKey(key_hash=key_hash, key=key_obj) - - -async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | None": - """Re-validate a live litellm user by id, returning ``None`` when the user is active or a precise - failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a - user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a - deactivated user cannot keep refreshing, mirroring how admission re-validates the same user subject on - the egress side. No DB connection is a gateway fault (``unresolvable``) and a - database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails - closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` / - ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` - catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look - identical, the original error surviving only as ``__context__``), so the outage check walks the cause - chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.""" - from litellm.proxy._types import ( - ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import - ) - from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import - get_user_object, - ) - from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import - PrismaDBExceptionHandler, - ) - from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import - prisma_client, - user_api_key_cache, - ) - - if prisma_client is None: - return "unresolvable" - try: - user_object = await get_user_object( - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - ) - except (ProxyException, HTTPException): - return "no_active_key" - except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500 - if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc): - return "unavailable" - verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__) - return "no_active_key" - if user_object is None: - return "no_active_key" - if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: - return "no_active_key" - return None - - -async def _key_owner_scim_deactivated(key: "UserAPIKeyAuth") -> bool: - """True only when the key's owning user was explicitly SCIM-deactivated, so a refresh revokes an - offboarded owner's key exactly as admission does via ``_reject_if_admitted_owner_scim_deactivated``. - A key with no owner, a missing owner record, or a failed lookup fails OPEN (returns ``False``), - matching admission and the standard builder: a key may outlive its owner record, and a transient DB - blip must not revoke a live key. Only an explicit ``scim_active`` of ``False`` gates renewal.""" - if key.user_id is None: - return False - from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import - get_user_object, - ) - from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import - prisma_client, - user_api_key_cache, - ) - - if prisma_client is None: - return False - try: - owner = await get_user_object( - user_id=key.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - ) - except Exception as exc: # noqa: BLE001 # fail open: a missing owner (get_user_object's wrapped ValueError) or a DB blip must not revoke a live key - verbose_logger.debug("refresh: key-owner SCIM lookup failed, not revoking (%s)", type(exc).__name__) - return False - return owner is not None and isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False - - -async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResolutionFailure | None": - """Re-validate that the subject sealed in a refresh envelope is still live, dispatching on its type: - a key_hash reloads the virtual key, a user_id reloads the user. Returns ``None`` when the subject is - active or a precise failure otherwise, so revocation gates renewal for either identity source the same - way admission gates the egress: a blocked or expired key, a SCIM-deactivated key owner (mirroring - admission's owner check, so an offboarded user cannot keep renewing a still-active key), and a - deactivated or deleted user all fail closed to ``no_active_key``.""" - match identity.subject_type: - case "key_hash": - reloaded = await _reload_active_key_by_hash(identity.subject) - if not isinstance(reloaded, _ResolvedKey): - return reloaded - if await _key_owner_scim_deactivated(reloaded.key): - return "no_active_key" - return None - case "user_id": - return await _reload_active_user_by_id(identity.subject) - case _: - assert_never(identity.subject_type) - - -async def _extract_user_id_from_request(request: Request) -> str | None: - """The litellm ``user_id`` for the token request, so a per-user token is stored under the same - identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome - (including a transient DB outage) collapses to ``None`` here and the caller simply skips the store; - the bridge mint, which must status those outcomes differently, consumes - :func:`_resolve_active_litellm_key` directly.""" - resolved = await _resolve_active_litellm_key(request) - if not isinstance(resolved, _ResolvedKey): - return None - return _active_key_user_id(resolved.key) - - async def _store_per_user_token_server_side( server: MCPServer, user_id: str, @@ -946,350 +678,6 @@ async def authorize_with_server( return response -_UpstreamGrantRejection = Literal["no_access_token", "expired_lifetime"] -"""Why an upstream token response cannot back a bridge envelope: -- ``no_access_token``: the response carries no usable ``access_token`` -- ``expired_lifetime``: the response reports a parseable, non-positive ``expires_in``, i.e. an upstream - token that is already dead, so sealing it would forward a bearer the edge cannot use -An absent or unparseable ``expires_in`` is NOT a rejection; the lifetime is merely unknown and the -envelope caps it, the by-design behaviour for an upstream that omits the field.""" - - -def _classify_upstream_lifetime(raw_expires_in: object) -> "int | Literal['unspecified', 'expired']": - """Classify an upstream ``expires_in`` into a positive number of seconds, ``"unspecified"`` (absent - or unparseable, so the envelope caps it), or ``"expired"`` (a non-positive value the upstream reports - as already elapsed). Telling "we do not know the lifetime" apart from "the upstream says it is - already dead" is what stops an explicitly-expired token from silently receiving the envelope's 1h - cap. The expired decision is made on the parsed numeric value, not on ``int(...)`` of it, so a - positive sub-second lifetime in ``(0, 1)`` is not truncated to ``0`` and misread as elapsed; the - envelope works in whole seconds, so such a lifetime clamps up to its 1s floor. ``bool`` is excluded - (an ``int`` subclass but never a real lifetime), and the conversions can raise on ``NaN`` / - ``Infinity`` / oversized input, which reads as unparseable rather than surfacing as a 500.""" - if raw_expires_in is None or isinstance(raw_expires_in, bool) or not isinstance(raw_expires_in, (int, float, str)): - return "unspecified" - try: - numeric = float(raw_expires_in) - seconds = int(numeric) - except (ValueError, TypeError, OverflowError): - return "unspecified" - if numeric <= 0: - return "expired" - return max(1, seconds) - - -def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenGrant | _UpstreamGrantRejection": - """Validate an upstream OAuth token response into a typed grant, or say why it cannot back an - envelope. Each field is isinstance-checked so nothing untyped from ``response.json()`` reaches the - grant. ``expires_in`` is read three ways (see :func:`_classify_upstream_lifetime`): an unknown - lifetime leaves the grant ``expires_in`` ``None`` for the envelope to cap, a positive value is - honoured, and an explicit already-elapsed value is a rejection rather than a silent fall-through to - the cap.""" - from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import - UpstreamTokenGrant, - ) - - if not isinstance(token_response, dict): - return "no_access_token" - access = token_response.get("access_token") - if not isinstance(access, str) or not access: - return "no_access_token" - lifetime = _classify_upstream_lifetime(token_response.get("expires_in")) - if lifetime == "expired": - return "expired_lifetime" - token_type = token_response.get("token_type") - scope = token_response.get("scope") - return UpstreamTokenGrant( - access_token=SecretStr(access), - token_type=token_type if isinstance(token_type, str) and token_type else "Bearer", - # The upstream refresh_token is deliberately NOT sealed: the edge never consumes it (it forwards - # only token_type + access_token), so it would be dead weight embedding a long-lived upstream - # credential in the client-held bearer, and it enlarges the envelope. Refresh support is a - # follow-up (a dedicated refresh-envelope); the client re-runs authorization_code at the cap. - refresh_token=None, - scope=scope if isinstance(scope, str) and scope else None, - expires_in=lifetime if isinstance(lifetime, int) else None, - ) - - -# --------------------------------------------------------------------------- -# DCR-bridge oauth_delegate mint: a three-phase pipeline whose failures are values. -# -# prepare (before the upstream exchange) -> validate every precondition and resolve identity+keys -# exchange (the single-use upstream code is consumed here, in exchange_token_with_server) -# finish (after the exchange) -> seal the upstream grant into the client-held envelope -# -# Every precondition lives in ``prepare``, which runs BEFORE the exchange, so no failure can burn the -# single-use code or rotate a refresh token, for either grant type -- that whole class of bug is gone -# by construction rather than guarded case by case. Failures are values mapped to an OAuth-shaped -# response in one place (``_bridge_mint_error_response``), so status codes and the RFC 6749 §5.2 body -# shape are uniform. Adding a failure mode is a new literal plus a match arm the type checker forces. -# --------------------------------------------------------------------------- - -_BridgeMintError = Literal[ - "no_identity", - "invalid_refresh", - "identity_unavailable", - "identity_unresolvable", - "not_configured", - "no_upstream_token", - "upstream_token_expired", - "too_large", -] - - -@dataclass(frozen=True, slots=True) -class _BridgeMintReady: - """Everything the seal needs, resolved once before the exchange: the identity to bind the envelope - to and the master-key-derived envelope keys. The identity is a key_hash subject for the scripted - two-header client (resolved from the litellm key it presents) or a user_id subject for the - interactive SSO client (the user recovered from the gateway authorization code), so one phase-3 seal - serves both. Resolving identity here means ``_finish_bridge_mint`` has no preconditions left to - fail.""" - - identity: "EnvelopeIdentity" - keys: "EnvelopeKeys" - - -def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: - """Map a bridge-mint failure value to its token-endpoint response: one place, RFC 6749 §5.2 shape - (top-level ``error``, no-store headers) for every case, with a status truthful about where the - failure is. The caller's request is 400, a transient gateway outage is 503, a gateway - misconfiguration is 500, and an upstream problem is 502. The identity-resolution statuses match how - admission statuses the same conditions on the egress side, so mint and admit never disagree under - one outage.""" - match error: - case "no_identity": - status, code, desc = ( - 400, - "invalid_request", - "this server issues a gateway-bound credential; complete the interactive sign-in, or " - "send a litellm credential (x-litellm-api-key or Authorization) on the token request", - ) - case "invalid_refresh": - status, code, desc = ( - 400, - "invalid_grant", - "the refresh credential is not a valid, live refresh envelope for this server; " - "re-run authorization_code to obtain a new one", - ) - case "identity_unavailable": - status, code, desc = ( - 503, - "temporarily_unavailable", - "the authentication database is temporarily unreachable; retry shortly", - ) - case "identity_unresolvable": - status, code, desc = ( - 500, - "server_error", - "the gateway could not resolve the litellm identity for this request", - ) - case "not_configured": - status, code, desc = ( - 500, - "server_error", - "the gateway is not configured to mint a gateway-bound credential (master_key is not set)", - ) - case "no_upstream_token": - status, code, desc = ( - 502, - "server_error", - "the upstream token response has no usable access_token", - ) - case "upstream_token_expired": - status, code, desc = ( - 502, - "server_error", - "the upstream token response reports an already-expired lifetime", - ) - case "too_large": - status, code, desc = ( - 502, - "server_error", - "the upstream token is too large to seal into a gateway-bound credential", - ) - case _: - assert_never(error) - return JSONResponse( - status_code=status, content={"error": code, "error_description": desc}, headers=TOKEN_NO_CACHE_HEADERS - ) - - -def _key_resolution_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError: - """Lift an identity-resolution failure into the mint taxonomy, preserving origin so the status stays - truthful: the caller's missing credential is 400, a transient DB outage is 503, and a gateway that - cannot resolve identity is 500.""" - match failure: - case "no_active_key": - return "no_identity" - case "unavailable": - return "identity_unavailable" - case "unresolvable": - return "identity_unresolvable" - case _: - assert_never(failure) - - -def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _BridgeMintError: - """Lift an upstream-response rejection into the mint taxonomy; both are upstream faults (502).""" - match rejection: - case "no_access_token": - return "no_upstream_token" - case "expired_lifetime": - return "upstream_token_expired" - case _: - assert_never(rejection) - - -async def _prepare_bridge_mint( - request: Request, - mcp_server: MCPServer, - bridge_identity: _BridgeAuthorizationCode | None = None, -) -> "_BridgeMintReady | _BridgeMintError": - """Phase 1 for the authorization_code grant, BEFORE the upstream exchange: confirm the gateway can - mint (master_key set), resolve the litellm identity, and derive the envelope keys. Returns a ready - context or a precise failure value. Running before the exchange is what makes every failure here fail - closed without consuming the single-use code. - - Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged - authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway - authorization code) and mints a user subject. The scripted two-header client presents a litellm key - on the token request instead, so its identity is the active key's hash and mints a key_hash subject. - A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully; - neither source present is ``no_identity``. The refresh_token grant has its own phase-1 - (:func:`_prepare_bridge_refresh`), which recovers identity from the presented refresh envelope.""" - from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import - envelope_keys_from_master_key, - ) - from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import - key_hash_identity, - user_identity, - ) - from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import - master_key, - ) - - if not master_key: - return "not_configured" - keys = envelope_keys_from_master_key(master_key) - if bridge_identity is not None: - identity = user_identity(server_id=mcp_server.server_id, user_id=bridge_identity.litellm_user_id) - return _BridgeMintReady(identity=identity, keys=keys) - resolved = await _resolve_active_litellm_key(request) - if not isinstance(resolved, _ResolvedKey): - return _key_resolution_failure_to_mint_error(resolved) - identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved.key_hash) - return _BridgeMintReady(identity=identity, keys=keys) - - -@dataclass(frozen=True, slots=True) -class _BridgeRefreshReady: - """A validated refresh request: the identity+keys to mint the renewed pair under, the upstream refresh - token (unwrapped from the client's refresh envelope) to exchange with the upstream IdP, and the scope - sealed alongside it at mint. The upstream refresh token is a ``SecretStr`` like every other credential - in this layer, so a repr or a traceback that captures this value never exposes the raw upstream refresh - token in plaintext. ``upstream_scope`` carries the originally-granted scope so the renewal re-requests - it when the client (a DCR/MCP client that typically omits scope on refresh) sends none, keeping the - renewed token's scope stable against an upstream that would otherwise narrow or drop it.""" - - ready: "_BridgeMintReady" - upstream_refresh_token: SecretStr - upstream_scope: str | None = None - - -def _refresh_key_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError: - """Lift an identity-resolution failure on the refresh path into the mint taxonomy. Unlike the mint - path, a resolved-but-inactive (or unknown) key is ``invalid_grant`` rather than ``invalid_request``: - the client did present an identity (sealed in the refresh envelope), but it is no longer live, so the - refresh is invalid and the client must re-authenticate. A transient outage is still 503 and a gateway - fault still 500, matching the mint path and admission.""" - match failure: - case "no_active_key": - return "invalid_refresh" - case "unavailable": - return "identity_unavailable" - case "unresolvable": - return "identity_unresolvable" - case _: - assert_never(failure) - - -async def _prepare_bridge_refresh( - mcp_server: MCPServer, refresh_value: str | None -) -> "_BridgeRefreshReady | _BridgeMintError": - """Phase 1 for the refresh_token grant, BEFORE the upstream exchange: open the client's refresh - envelope, re-validate the sealed litellm identity so a revoked key cannot keep refreshing, and - recover the upstream refresh token to exchange. Identity comes entirely from the sealed envelope, not - the HTTP request, so the request object is not needed here. The client presents a refresh envelope, - never a raw upstream refresh token, so a missing value, a non-envelope, an unopenable envelope, or one - minted for another server is ``invalid_grant``. Running before the exchange means a rejected refresh - never consumes or rotates the upstream refresh token.""" - from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import - BridgeRefreshOpened, - envelope_keys_from_master_key, - open_bridge_refresh_envelope, - ) - from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import - master_key, - ) - - if not master_key: - return "not_configured" - if not refresh_value: - return "invalid_refresh" - keys = envelope_keys_from_master_key(master_key) - opened = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id) - if not isinstance(opened, BridgeRefreshOpened): - return "invalid_refresh" - failure = await _revalidate_active_subject(opened.identity) - if failure is not None: - return _refresh_key_failure_to_mint_error(failure) - return _BridgeRefreshReady( - ready=_BridgeMintReady(identity=opened.identity, keys=keys), - upstream_refresh_token=opened.refresh.refresh_token, - upstream_scope=opened.refresh.scope, - ) - - -def _finish_bridge_mint( - ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime -) -> "JSONResponse | _BridgeMintError": - """Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held access envelope - using the pre-resolved identity and keys, and, when the upstream returned a refresh token, seal a - long-lived refresh envelope alongside it so the client can renew without re-authenticating. Shared by - the authorization_code and refresh_token paths, so a renewal that the upstream rotates re-issues a - fresh refresh envelope. The only hard failures here are properties of the upstream access token (no - usable token, an already-expired lifetime, or a token too large to seal); a refresh token that cannot - be sealed degrades to an access-only response rather than failing the whole exchange.""" - from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import - build_bridge_token_response, - ) - from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import - SealedEnvelope, - UpstreamTokenGrant, - ) - - grant = _bridge_grant_from_token_response(token_response) - if not isinstance(grant, UpstreamTokenGrant): - return _upstream_rejection_to_mint_error(grant) - sealed = build_bridge_token_response(ready.identity, grant, ready.keys, now) - if not isinstance(sealed, SealedEnvelope): - return "too_large" - # Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the - # client is never told the bearer lives past the point admission (which uses that exp) rejects it. - expires_in = max(0, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp())) - refresh_envelope = _mint_refresh_envelope_value(ready.identity, token_response, ready.keys, now, mcp_server) - body = { - "access_token": sealed.token.get_secret_value(), - "token_type": "Bearer", - "expires_in": expires_in, - # A refresh envelope rides along only when the upstream returned a refresh token to seal; when it - # rotates on renewal, the client receives the new one and the old envelope's upstream token dies. - **({"refresh_token": refresh_envelope} if refresh_envelope is not None else {}), - } - return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS) - - def _token_credential_source(mcp_server: MCPServer) -> CredentialSource: """Mirrors the resolved-client rule in :func:`exchange_token_with_server`: when the server has a stored client_id the gateway presents its own credentials upstream, so a credential rejection is @@ -1297,63 +685,6 @@ def _token_credential_source(mcp_server: MCPServer) -> CredentialSource: return "gateway_stored" if mcp_server.client_id else "caller_supplied" -def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | None": - """Extract the upstream refresh grant from a token response, or ``None`` when there is none to seal. - Each field is isinstance-checked so nothing untyped reaches the refresh envelope; ``refresh_expires_in`` - (the refresh token's own lifetime, when the upstream reports it) is classified like ``expires_in`` and - bounds the refresh envelope's TTL. An upstream that reports the refresh token itself as already elapsed - (``refresh_expires_in`` non-positive) yields ``None`` rather than a refresh envelope: sealing a dead - token would hand the client a full-TTL-capped envelope the IdP will reject, so the exchange degrades to - an access-only response (the client re-authenticates at access expiry), mirroring how - :func:`_bridge_grant_from_token_response` refuses an already-elapsed access token instead of capping it.""" - from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import - RefreshCredential, - ) - - if not isinstance(token_response, dict): - return None - refresh = token_response.get("refresh_token") - if not isinstance(refresh, str) or not refresh: - return None - lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in")) - if lifetime == "expired": - return None - scope = token_response.get("scope") - return RefreshCredential( - refresh_token=SecretStr(refresh), - scope=scope if isinstance(scope, str) and scope else None, - expires_in=lifetime if isinstance(lifetime, int) else None, - ) - - -def _mint_refresh_envelope_value( - identity: "EnvelopeIdentity", token_response: object, keys: "EnvelopeKeys", now: datetime, mcp_server: MCPServer -) -> str | None: - """Seal the upstream refresh grant (if any) into a refresh envelope and return its bearer string, or - ``None`` when the upstream returned no refresh token or the refresh token is too large to seal. A - too-large refresh token degrades to an access-only response (logged) rather than failing an exchange - that already succeeded upstream: the client simply re-authenticates when the access envelope expires.""" - from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import - build_bridge_refresh_token_response, - ) - from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import - SealedEnvelope, - ) - - refresh_credential = _upstream_refresh_credential(token_response) - if refresh_credential is None: - return None - sealed = build_bridge_refresh_token_response(identity, refresh_credential, keys, now) - if isinstance(sealed, SealedEnvelope): - return sealed.token.get_secret_value() - verbose_logger.warning( - "bridge mint: the upstream refresh token is too large to seal into a refresh envelope for " - "server=%s; issuing an access-only response, so the client re-authenticates at access expiry", - mcp_server.server_id, - ) - return None - - async def exchange_token_with_server( request: Request, mcp_server: MCPServer, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 47a49b96327..d1aceffa968 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4374,10 +4374,8 @@ _BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" async def _exchange_for_bridge_server(server, upstream_body, key_hash, code="auth-code", fake_client_out=None): - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _ResolvedKey, - exchange_token_with_server, - ) + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _ResolvedKey + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server fake_http_response = MagicMock() fake_http_response.json.return_value = upstream_body @@ -4397,7 +4395,7 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, code="aut return_value=fake_http_client, ), patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_active_litellm_key", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_active_litellm_key", new=key_resolver, ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), @@ -4808,7 +4806,7 @@ async def _refresh_for_bridge_server( return_value=fake_http_client, ), patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._revalidate_active_subject", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._revalidate_active_subject", new=AsyncMock(return_value=revalidate_result), ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), @@ -5085,7 +5083,7 @@ async def test_bridge_refresh_grant_with_deactivated_user_is_invalid_grant_befor async def test_revalidate_active_subject_dispatches_on_subject_type(): """Subject re-validation routes a key_hash envelope to the key reload and a user_id envelope to the user reload, so revocation gates renewal for either identity source through one dispatch point.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _ResolvedKey, _revalidate_active_subject, ) @@ -5093,11 +5091,11 @@ async def test_revalidate_active_subject_dispatches_on_subject_type(): with ( patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", new=AsyncMock(return_value=_ResolvedKey(key_hash="kh", key=MagicMock())), ) as key_reload, patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_user_by_id", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_user_by_id", new=AsyncMock(return_value=None), ) as user_reload, ): @@ -5107,11 +5105,11 @@ async def test_revalidate_active_subject_dispatches_on_subject_type(): with ( patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", new=AsyncMock(), ) as key_reload2, patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_user_by_id", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_user_by_id", new=AsyncMock(return_value="no_active_key"), ) as user_reload2, ): @@ -5125,7 +5123,7 @@ def test_upstream_refresh_credential_expired_refresh_token_is_not_sealed(): not be sealed: _upstream_refresh_credential returns None so the exchange degrades to an access-only response, mirroring how the access grant refuses an already-elapsed access token rather than capping a dead token to the full refresh TTL. A live or unspecified lifetime still yields a credential.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _upstream_refresh_credential + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _upstream_refresh_credential assert _upstream_refresh_credential({"access_token": "A", "refresh_token": "R", "refresh_expires_in": 0}) is None assert _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": -5}) is None @@ -5164,7 +5162,7 @@ async def test_bridge_refresh_upstream_invalid_grant_maps_to_invalid_grant(): return_value=fake_http_client, ), patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._revalidate_active_subject", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._revalidate_active_subject", new=AsyncMock(return_value=None), ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), @@ -5217,7 +5215,7 @@ async def test_bridge_refresh_upstream_error_detection_parses_json_not_substring return_value=fake_http_client, ), patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._revalidate_active_subject", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._revalidate_active_subject", new=AsyncMock(return_value=None), ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), @@ -5244,7 +5242,7 @@ async def test_revalidate_key_subject_revoked_when_owner_scim_deactivated(proxy_ """A key_hash refresh envelope whose key is still active but whose OWNING user was SCIM-deactivated must fail closed to no_active_key, mirroring how admission's _reject_if_admitted_owner_scim_deactivated revokes an offboarded owner's key. Without this, an offboarded user keeps renewing a live key.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _ResolvedKey, _revalidate_active_subject + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _ResolvedKey, _revalidate_active_subject from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -5254,7 +5252,7 @@ async def test_revalidate_key_subject_revoked_when_owner_scim_deactivated(proxy_ resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="offboarded-owner")) with ( patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", new=AsyncMock(return_value=resolved), ), patch( @@ -5272,7 +5270,7 @@ async def test_revalidate_key_subject_active_owner_renews_and_missing_owner_fail """The key-owner SCIM gate blocks only an explicit scim_active False: an active owner renews (None), and a missing owner (get_user_object's wrapped ValueError) fails OPEN, since a key may outlive its owner record and a transient blip must not revoke a live key.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _ResolvedKey, _revalidate_active_subject + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _ResolvedKey, _revalidate_active_subject from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -5283,7 +5281,7 @@ async def test_revalidate_key_subject_active_owner_renews_and_missing_owner_fail with ( patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", new=AsyncMock(return_value=resolved), ), patch( @@ -5295,7 +5293,7 @@ async def test_revalidate_key_subject_active_owner_renews_and_missing_owner_fail with ( patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._reload_active_key_by_hash", new=AsyncMock(return_value=resolved), ), patch( @@ -5324,7 +5322,7 @@ async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset(): return_value=fake_http_client, ), patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_active_litellm_key", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_active_litellm_key", new=AsyncMock(return_value="no_active_key"), ), patch("litellm.proxy.proxy_server.master_key", None), @@ -5361,7 +5359,7 @@ async def _prepare_only_bridge_exchange(resolver_result): return_value=fake_http_client, ), patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_active_litellm_key", + "litellm.proxy._experimental.mcp_server.bridge_token_flow._resolve_active_litellm_key", new=AsyncMock(return_value=resolver_result), ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), @@ -5505,7 +5503,7 @@ def test_classify_upstream_lifetime(): oversized) is "unspecified" so the envelope caps it, while a parseable non-positive value is "expired": the upstream reporting an already-dead token, which the mint must reject rather than silently give the 1h cap.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _classify_upstream_lifetime + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _classify_upstream_lifetime assert _classify_upstream_lifetime(300) == 300 assert _classify_upstream_lifetime(300.0) == 300 @@ -5538,7 +5536,7 @@ def test_bridge_grant_honors_and_rejects_upstream_lifetime(): """The grant validator honors a positive lifetime, leaves an unknown one None for the envelope to cap, and rejects an explicitly-expired one with "expired_lifetime" so a dead upstream token is never sealed into an hour-long envelope.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _bridge_grant_from_token_response + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _bridge_grant_from_token_response def grant(v): return _bridge_grant_from_token_response({"access_token": "x", "expires_in": v}) @@ -5831,7 +5829,7 @@ async def test_extract_user_id_reads_x_litellm_api_key_header(proxy_globals): """The LiteLLM key arrives on x-litellm-api-key (what Claude Desktop/Code send), not Authorization. Reading only Authorization dropped the identity, so the per-user token was never stored and the egress 401'd forever. Resolution must honor x-litellm-api-key.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth, hash_token @@ -5856,7 +5854,7 @@ async def test_extract_user_id_rehydrates_cross_replica_dict_cache(proxy_globals """Cross-replica, async_get_cache hands back a serialized dict, not a UserAPIKeyAuth. Resolution must rehydrate it; the old getattr(cached, "user_id") returned None on a dict, which is exactly why a multi-replica gateway never found the stored token.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import hash_token @@ -5877,7 +5875,7 @@ async def test_extract_user_id_falls_back_to_db_on_cache_miss(proxy_globals): """A cache miss must read the key from the DB rather than returning None; the old code did a cache-only peek and skipped the DB, so any replica that hadn't just authenticated the key failed to store the token.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -5899,7 +5897,7 @@ async def test_extract_user_id_falls_back_to_db_on_cache_miss(proxy_globals): @pytest.mark.asyncio async def test_extract_user_id_none_without_litellm_key(proxy_globals): """No LiteLLM key on the request resolves to None without consulting the resolver.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -5916,7 +5914,7 @@ async def test_extract_user_id_rejects_blocked_key(proxy_globals): """A blocked LiteLLM key must not resolve an identity. get_key_object returns the DB row without checking blocked/expiry (the main auth pipeline does, and the public token endpoint bypasses it), so a revoked key could otherwise overwrite the stored per-user OAuth token for its user.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -5938,7 +5936,7 @@ async def test_extract_user_id_rejects_expired_key(proxy_globals): """An expired LiteLLM key must not resolve an identity, for the same reason as a blocked key.""" from datetime import datetime, timedelta, timezone - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth @@ -5963,7 +5961,7 @@ async def test_resolve_active_litellm_key_returns_resolved_key_for_active_key(pr record. For an active key the resolver returns exactly hash_token(key), the same value get_key_object and the whole cache/DB layer key the record by, so the sealed reference resolves back to this key at admission.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _resolve_active_litellm_key, _ResolvedKey, ) @@ -5993,10 +5991,10 @@ async def test_resolve_active_litellm_key_resolves_key_without_user_id(proxy_glo presence wrongly rejected these keys with invalid_request; the active-state gate now checks only blocked and expiry, and the key hash (not the user) is what the mint seals. The per-user token store still gets no user for such a key, since there is none to key a stored credential by.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _extract_user_id_from_request, - _resolve_active_litellm_key, _ResolvedKey, + _resolve_active_litellm_key, ) from litellm.proxy._types import UserAPIKeyAuth, hash_token from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -6022,7 +6020,7 @@ async def test_resolve_active_litellm_key_resolves_key_without_user_id(proxy_glo async def test_resolve_active_litellm_key_rejects_blocked_key(proxy_globals): """A blocked key must not yield a hash, so no gateway-bound envelope is minted for a revoked key; the mint fails closed with invalid_request instead.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _resolve_active_litellm_key, ) from litellm.proxy._types import UserAPIKeyAuth @@ -6045,7 +6043,7 @@ async def test_resolve_active_litellm_key_fails_closed_on_malformed_expiry(proxy returns invalid_request), not surface an unhandled 500. The active-state check runs outside the resolver's try, so it must be total over a bad expires rather than letting datetime.fromisoformat raise. Before the fix this raised a ValueError instead of returning None.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _resolve_active_litellm_key, ) from litellm.proxy._types import UserAPIKeyAuth @@ -6065,7 +6063,7 @@ async def test_resolve_active_litellm_key_fails_closed_on_malformed_expiry(proxy @pytest.mark.asyncio async def test_resolve_active_litellm_key_no_active_key_without_litellm_key(proxy_globals): """No LiteLLM key on the request yields no hash without consulting the resolver.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _resolve_active_litellm_key, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -6083,7 +6081,7 @@ async def test_resolve_active_litellm_key_db_outage_is_unavailable(proxy_globals the caller's fault, so the resolver reports "unavailable" (the mint statuses it 503) rather than collapsing it to the same value as a missing credential. is_database_service_unavailable_error classifies a connection error (an OSError) as an outage, matching admission's egress-side handling.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _resolve_active_litellm_key, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -6104,7 +6102,7 @@ async def test_resolve_active_litellm_key_no_database_is_unresolvable(proxy_glob """With no database connection configured the gateway cannot verify the presented key at all, so the resolver reports "unresolvable" (the mint statuses it 500) instead of blaming the caller. Mirrors admission, which 500s a missing prisma_client on the egress side.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( _resolve_active_litellm_key, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -6138,7 +6136,7 @@ async def test_reload_active_user_by_id_missing_user_is_no_active_key(proxy_glob refresh path maps it to invalid_grant), not unresolvable/500. get_user_object catches the missing row and re-raises a bare ValueError, so a missing user must not be misclassified as a DB outage or an opaque gateway fault.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache proxy_globals.user_api_key_cache = UserApiKeyCache() @@ -6157,7 +6155,7 @@ async def test_reload_active_user_by_id_db_outage_is_unavailable(proxy_globals): a missing user, so the refresh path surfaces "unavailable" (a 503) rather than blaming the caller. get_user_object wraps the outage in a bare ValueError, so this exercises the chain-aware classifier; a raw ConnectionError would falsely pass even a chain-blind check because it is an OSError.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache proxy_globals.user_api_key_cache = UserApiKeyCache()