diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ebceb320906..8d1713a5911 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,8 +1,10 @@ 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 @@ -10,7 +12,8 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, SecretStr, ValidationError +from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -37,6 +40,10 @@ 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 ( + EnvelopeKeys, + UpstreamTokenGrant, + ) from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. @@ -326,66 +333,136 @@ def _litellm_key_from_request(request: Request) -> Optional[str]: return None -def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]: - """The key's ``user_id``, or ``None`` if the key is blocked or expired. +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 its - identity is trusted to key a stored credential; a revoked or expired key must not be able to - write or overwrite the per-user OAuth token. ``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. + 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 None + return False expires = key_obj.expires if expires is not None: - expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) + 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 None - return key_obj.user_id + return False + return True -async def _extract_user_id_from_request(request: Request) -> Optional[str]: - """Resolve the LiteLLM ``user_id`` at the OAuth token endpoint so a per-user token is stored - under the same identity the egress later reads it by (``user_api_key_auth.user_id``). +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 - Resolves authoritatively via ``get_key_object`` (cache first, then DB) instead of a raw cache - peek. On a multi-replica gateway the token-exchange request can land on a worker whose in-memory - cache never saw the key, and a cross-replica Redis hit deserializes to a plain ``dict`` rather - than a ``UserAPIKeyAuth``; the previous code read only ``Authorization`` and did - ``getattr(cached, "user_id")`` with no ``model_type`` rehydration and no DB fallback, so it - silently returned ``None`` and the token was never persisted, which makes the egress 401 on every - reconnect. The resolved key is validated (``_active_key_user_id``) before its identity is trusted, - so a blocked or expired key cannot write. Returns ``None`` when no key is present, the key cannot - be resolved, or it is blocked/expired. - """ + +@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 None - try: - from litellm.proxy._types import hash_token # noqa: PLC0415 - from litellm.proxy.auth.auth_checks import get_key_object # noqa: PLC0415 - from litellm.proxy.proxy_server import ( # noqa: PLC0415 - prisma_client, - user_api_key_cache, - ) + return "no_active_key" + from litellm.proxy._types import ( # noqa: PLC0415 # inline import avoids a module-load circular import + ProxyException, + hash_token, + ) + 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" + key_hash = hash_token(token) + try: key_obj = await get_key_object( - hashed_token=hash_token(token), + hashed_token=key_hash, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, ) - return _active_key_user_id(key_obj) - except Exception as exc: + 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( - "_extract_user_id_from_request: could not resolve a LiteLLM user_id for the presented " - "key (%s); per-user token will not be stored server-side.", + "_resolve_active_litellm_key: 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 _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( @@ -654,6 +731,255 @@ 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", + "unsupported_grant", + "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 authorizing key hash and the + master-key-derived envelope keys. Passing this forward means identity resolution and key derivation + happen exactly once, and ``_finish_bridge_mint`` has no preconditions left that could fail.""" + + key_hash: str + 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; send a litellm credential " + "(x-litellm-api-key or Authorization) on the token request", + ) + case "unsupported_grant": + status, code, desc = ( + 400, + "unsupported_grant_type", + "this server issues a gateway-bound credential and supports only the authorization_code " + "grant; re-run authorization_code to renew rather than refresh_token", + ) + 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, grant_type: str) -> "_BridgeMintReady | _BridgeMintError": + """Phase 1, BEFORE the upstream exchange: reject a grant this mint does not support, 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 or rotating a refresh token. A bridge + server issues only envelopes and seals no upstream refresh_token, so the client holds none to + present: the refresh_token grant is rejected up front rather than exchanged (which could rotate the + upstream credential) and its result then discarded. Identity-resolution failures keep their origin + so the mapper statuses each truthfully.""" + 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.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import + master_key, + ) + + if grant_type != "authorization_code": + return "unsupported_grant" + if not master_key: + return "not_configured" + resolved = await _resolve_active_litellm_key(request) + if not isinstance(resolved, _ResolvedKey): + return _key_resolution_failure_to_mint_error(resolved) + return _BridgeMintReady(key_hash=resolved.key_hash, keys=envelope_keys_from_master_key(master_key)) + + +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 envelope using + the pre-resolved identity and keys, so the client holds one bearer that admits it and forwards the + upstream token with nothing stored server-side. The only failures here are properties of the + upstream response (no usable token, an already-expired lifetime, or a token too large to seal), + returned as values.""" + 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 + EnvelopeIdentity, + SealedEnvelope, + UpstreamTokenGrant, + ) + + grant = _bridge_grant_from_token_response(token_response) + if not isinstance(grant, UpstreamTokenGrant): + return _upstream_rejection_to_mint_error(grant) + identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=ready.key_hash) + sealed = build_bridge_token_response(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())) + body = {"access_token": sealed.token.get_secret_value(), "token_type": "Bearer", "expires_in": expires_in} + return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS) + + async def exchange_token_with_server( request: Request, mcp_server: MCPServer, @@ -727,6 +1053,16 @@ async def exchange_token_with_server( if code_verifier: token_data["code_verifier"] = code_verifier + # Phase 1: for a bridge oauth_delegate mint, validate all preconditions and resolve identity+keys + # BEFORE the exchange below consumes the single-use upstream code, and carry the ready context to + # phase 3. A failure here returns without ever touching the upstream credential. + bridge_mint_ready: _BridgeMintReady | None = None + if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: + prepared = await _prepare_bridge_mint(request, grant_type) + if not isinstance(prepared, _BridgeMintReady): + return _bridge_mint_error_response(prepared) + bridge_mint_ready = prepared + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response = await async_client.post( mcp_server.token_url, @@ -751,7 +1087,6 @@ async def exchange_token_with_server( ) raise token_response = response.json() - access_token = token_response["access_token"] # Validate token response against server-configured rules before any storage. # This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc. @@ -791,8 +1126,17 @@ async def exchange_token_with_server( mcp_server.server_id, ) + # A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the + # upstream token) instead of the raw upstream token, so the one bearer both admits the caller and + # forwards the upstream credential. Only this mode mints; every other server returns the raw token. + if bridge_mint_ready is not None: + # Phase 3: seal the upstream grant into the client-held envelope; failures map through the same + # OAuth-shaped response as the phase-1 preconditions. + minted = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc)) + return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted) + result = { - "access_token": access_token, + "access_token": token_response["access_token"], "token_type": token_response.get("token_type", "Bearer"), } 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 c55a631c7b3..68466e624ec 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 @@ -4362,6 +4362,501 @@ async def test_register_bridge_relay_never_persists(): mock_persist.assert_not_called() +_BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" + + +async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_client_out=None): + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _ResolvedKey, + exchange_token_with_server, + ) + + fake_http_response = MagicMock() + fake_http_response.json.return_value = upstream_body + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + # The mint consumes _resolve_active_litellm_key's tagged result: an active key resolves to a + # _ResolvedKey carrying its hash; a request with no usable credential resolves to "no_active_key". + resolution = _ResolvedKey(key_hash=key_hash, key=MagicMock()) if key_hash is not None else "no_active_key" + key_resolver = AsyncMock(return_value=resolution) + if fake_client_out is not None: + fake_client_out["client"] = fake_http_client + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_active_litellm_key", + new=key_resolver, + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + if server.is_oauth_delegate and server.is_dcr_bridge: + key_resolver.assert_awaited_once() + else: + key_resolver.assert_not_awaited() + return response + + +@pytest.mark.asyncio +async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token(): + """A dcr_bridge oauth_delegate token exchange returns a gateway-bound envelope, not the raw + upstream token: the response access_token opens (under the same master-key-derived keys and the + server_id) to the caller's identity and the upstream Authorization, and the raw upstream token + never appears in the bearer the client receives.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + envelope_keys_from_master_key, + resolve_bridge_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + token = body["access_token"] + assert body["token_type"] == "Bearer" + assert body["expires_in"] > 0 + assert token.startswith("llm_env_") + assert "UPSTREAM-SECRET-TOKEN" not in token + assert "refresh_token" not in body + + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id) + assert isinstance(opened, BridgeEnvelopeAdmitted) + assert opened.identity.key_hash == "hashed-litellm-key-77" + assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" + + +@pytest.mark.asyncio +async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm_identity(): + """Without a resolvable litellm identity on the token request, the exchange must not mint an + identity-less envelope. It returns an RFC 6749 §5.2-shaped invalid_request (error at the top + level, not wrapped in detail) BEFORE exchanging the upstream code, so the single-use code is not + burned and the client can retry.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + captured: dict = {} + response = await _exchange_for_bridge_server(server, upstream, key_hash=None, fake_client_out=captured) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + # identity resolution failed first, so the upstream single-use code was never exchanged (not burned) + captured["client"].post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_envelope_too_large_upstream_token_is_502(): + """An upstream token too large to seal into the envelope is an upstream-payload condition, so the + mint surfaces a 502 (as an RFC 6749 §5.2 error body, not a raised HTTPException) rather than a 500: + build_bridge_token_response returns EnvelopeTooLarge as a value, _finish_bridge_mint returns the + "too_large" failure, and _bridge_mint_error_response maps it to a truthful status.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "x" * 40000, "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +async def test_bridge_envelope_does_not_seal_upstream_refresh_token(): + """The upstream refresh_token is never sealed into the client-held envelope: the edge never + consumes it and a long-lived upstream credential should not live in the client bearer. The opened + envelope's grant carries no refresh token even when the upstream returned one, and neither does + the response body.""" + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + envelope_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + OpenedEnvelope, + open_envelope, + ) + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = { + "access_token": "UP", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "UPSTREAM-REFRESH", + } + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + assert "refresh_token" not in body + assert "UPSTREAM-REFRESH" not in body["access_token"] + keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY) + opened = open_envelope(body["access_token"], keys, datetime.now(timezone.utc)) + assert isinstance(opened, OpenedEnvelope) + assert opened.grant.refresh_token is None + + +@pytest.mark.asyncio +async def test_bridge_refresh_grant_is_rejected_before_upstream(): + """A bridge oauth_delegate server issues only envelopes and seals no upstream refresh_token, so the + client never holds one to present. _prepare_bridge_mint rejects the refresh_token grant up front + with unsupported_grant_type, BEFORE any upstream exchange, so a stray refresh request can never + rotate or consume the client's upstream refresh credential; renewal is re-running + authorization_code. This is checked before identity resolution, so it holds even with a valid key.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="dcr-client-123", + client_secret=None, + code_verifier=None, + refresh_token="client-refresh-token", + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "unsupported_grant_type" + fake_http_client.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset(): + """master_key is validated BEFORE the upstream exchange (in _prepare_bridge_mint), so a + misconfigured gateway returns a 500 server_error without consuming the single-use code, avoiding + the burn-then-fail the pre-exchange phase exists to prevent. The failure is returned as an RFC 6749 + error body, not raised.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_active_litellm_key", + new=AsyncMock(return_value="no_active_key"), + ), + patch("litellm.proxy.proxy_server.master_key", None), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + + assert response.status_code == 500 + assert json.loads(response.body)["error"] == "server_error" + fake_http_client.post.assert_not_called() + + +async def _prepare_only_bridge_exchange(resolver_result): + """Drive exchange_token_with_server for a bridge oauth_delegate authorization_code request with the + identity resolver stubbed to a given tagged result, returning (response, post_mock) so a test can + assert the mapped status and that the single-use code was never exchanged.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_active_litellm_key", + new=AsyncMock(return_value=resolver_result), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + response = await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://claude.ai/api/mcp/auth_callback", + client_id="dcr-client-123", + client_secret=None, + code_verifier="verifier", + ) + return response, fake_http_client.post + + +@pytest.mark.asyncio +async def test_bridge_mint_db_outage_is_503_before_upstream(): + """A DB outage while resolving identity is a retryable gateway failure, so the mint returns 503 + temporarily_unavailable WITHOUT consuming the single-use code, matching how admission statuses the + same outage on the egress side. Collapsing every resolution failure to None used to blame the + client with 400 invalid_request for an infrastructure problem.""" + response, post = await _prepare_only_bridge_exchange("unavailable") + assert response.status_code == 503 + assert json.loads(response.body)["error"] == "temporarily_unavailable" + post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_unresolvable_identity_is_500_before_upstream(): + """An unresolvable identity (no DB connection, or an unexpected resolution error) is a gateway + fault, so the mint returns 500 server_error before the exchange, a status distinct from both the + caller's 400 and the transient 503, matching admission's 500-vs-503 split for the same conditions.""" + response, post = await _prepare_only_bridge_exchange("unresolvable") + assert response.status_code == 500 + assert json.loads(response.body)["error"] == "server_error" + post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_upstream_expired_lifetime_is_502(): + """An upstream token response reporting an already-elapsed lifetime (a parseable non-positive + expires_in) is rejected with 502 rather than sealed into an hour-long envelope around a dead + bearer. Regression for expires_in<=0 silently falling through to the 1h cap.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 0} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +async def test_bridge_mint_positive_sub_second_lifetime_mints_not_502(): + """A positive fractional expires_in in (0, 1) is a live token, not an elapsed one, so it mints a + (1s-floored) envelope rather than being truncated to 0 and rejected with 502 after the single-use + code was already consumed. Regression for classifying a sub-second remaining lifetime as expired.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 0.5} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert body["expires_in"] >= 0 + + +@pytest.mark.asyncio +async def test_bridge_mint_unknown_lifetime_is_capped_not_rejected(): + """An absent or unparseable expires_in leaves the lifetime unknown, which the envelope caps (never + inventing a longer life than the upstream stated); it is NOT rejected. Only an explicitly-dead + lifetime fails, so a metadata glitch on an otherwise-valid token still mints a bounded envelope.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": "not-a-number"} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert 0 < body["expires_in"] <= 3600 + + +@pytest.mark.asyncio +async def test_bridge_reported_expires_in_does_not_overstate_jwt_exp(): + """The reported expires_in is derived from the envelope JWT's second-truncated exp (rounding the + elapsed portion up), so the client is never told the bearer lives past the point admission expires + it. Regression for the sub-second overstatement of the raw (expires_at - now) delta.""" + import time + + import jwt as _jwt + + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 300} + before = int(time.time()) + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + body = json.loads(response.body) + claims = _jwt.decode(body["access_token"].removeprefix("llm_env_"), options={"verify_signature": False}) + # projecting the reported lifetime from a time no later than the mint must not exceed the JWT exp + assert before + body["expires_in"] <= claims["exp"] + + +def test_bridge_reported_expires_in_can_be_zero_at_jwt_exp_boundary(): + from datetime import datetime, timezone + + from fastapi.responses import JSONResponse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _BridgeMintReady, + _finish_bridge_mint, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + envelope_keys_from_master_key, + ) + from litellm.types.mcp import MCPAuth + + ready = _BridgeMintReady( + key_hash="hashed-litellm-key-77", + keys=envelope_keys_from_master_key(_BRIDGE_MASTER_KEY), + ) + response = _finish_bridge_mint( + ready=ready, + mcp_server=_bridge_server(auth_type=MCPAuth.oauth_delegate), + token_response={"access_token": "UP", "expires_in": 1}, + now=datetime.fromtimestamp(100.25, tz=timezone.utc), + ) + + assert isinstance(response, JSONResponse) + assert json.loads(response.body)["expires_in"] == 0 + + +def test_classify_upstream_lifetime(): + """expires_in from an IdP may be an int, a float (3600.0), or a numeric string ("3600"); each + coerces to a positive number of seconds. Absent or unparseable input (bool, non-numeric, NaN/inf, + 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 + + assert _classify_upstream_lifetime(300) == 300 + assert _classify_upstream_lifetime(300.0) == 300 + assert _classify_upstream_lifetime("300") == 300 + assert _classify_upstream_lifetime(" 300 ") == 300 + # explicit, parseable, non-positive -> the upstream says the token is already dead + assert _classify_upstream_lifetime(0) == "expired" + assert _classify_upstream_lifetime(-5) == "expired" + assert _classify_upstream_lifetime(-0.5) == "expired" + # a positive sub-second lifetime is alive, not elapsed; it clamps up to the envelope's 1s floor + # rather than truncating to 0 and being misread as expired + assert _classify_upstream_lifetime(0.5) == 1 + assert _classify_upstream_lifetime(0.001) == 1 + # a positive value >= 1 truncates toward zero (never overstating the stated lifetime) + assert _classify_upstream_lifetime(1.9) == 1 + # unknown lifetime -> cap (never invent a longer life than the upstream stated) + assert _classify_upstream_lifetime(None) == "unspecified" + assert _classify_upstream_lifetime(True) == "unspecified" + assert _classify_upstream_lifetime("nope") == "unspecified" + # hostile numerics must not raise (int(float(...)) can OverflowError) -> unspecified + assert _classify_upstream_lifetime("inf") == "unspecified" + assert _classify_upstream_lifetime("1e999") == "unspecified" + assert _classify_upstream_lifetime("-inf") == "unspecified" + assert _classify_upstream_lifetime("nan") == "unspecified" + assert _classify_upstream_lifetime(float("inf")) == "unspecified" + assert _classify_upstream_lifetime(10**400) == "unspecified" + + +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 + + def grant(v): + return _bridge_grant_from_token_response({"access_token": "x", "expires_in": v}) + + assert grant(300).expires_in == 300 + assert grant(120.0).expires_in == 120 + # unknown lifetime backs a grant whose expires_in the envelope caps; it is not a rejection + assert grant("nope").expires_in is None + assert _bridge_grant_from_token_response({"access_token": "x"}).expires_in is None + # an explicitly already-dead lifetime is rejected, not silently capped at 1h + assert grant(0) == "expired_lifetime" + assert grant(-5) == "expired_lifetime" + + +@pytest.mark.asyncio +async def test_bridge_token_exchange_honors_short_float_expires_in_ttl(): + """A short float expires_in from the upstream caps the envelope TTL, so the client-held envelope + does not outlive the upstream token. Before coercion a float was dropped and the envelope + defaulted to the 1h cap (3600), which would forward a stale bearer after the upstream token + expired.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 120.0} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert json.loads(response.body)["expires_in"] <= 120 + + +@pytest.mark.asyncio +async def test_oauth_delegate_bridge_token_exchange_missing_access_token_is_502_not_keyerror(): + """When the upstream token response has no access_token, a dcr_bridge oauth_delegate exchange + returns a clean 502 error body rather than raising a KeyError. _finish_bridge_mint asks + _bridge_grant_from_token_response for a typed grant, gets None, and returns the "no_upstream_token" + failure, which maps to 502; nothing indexes token_response["access_token"] on the bridge path.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"token_type": "Bearer", "expires_in": 3600} + + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" + + +@pytest.mark.asyncio +async def test_true_passthrough_bridge_token_exchange_returns_raw_upstream_token(): + """Only oauth_delegate mints. A true_passthrough dcr_bridge server relays the raw upstream token + to the client, since that mode has no litellm identity to bind and the caller owns the token.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.true_passthrough) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + assert body["access_token"] == "UPSTREAM-SECRET-TOKEN" + assert not body["access_token"].startswith("llm_env_") + + +@pytest.mark.asyncio +async def test_non_bridge_oauth_delegate_token_exchange_returns_raw_upstream_token(): + """An oauth_delegate server without dcr_bridge keeps the pre-change contract: the raw upstream + token is returned, so flag-off behavior is byte-identical.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate, dcr_bridge=None) + upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + body = json.loads(response.body) + assert body["access_token"] == "UPSTREAM-SECRET-TOKEN" + + async def _exchange_persistence_attempted_for_auth_type(auth_type) -> bool: """Run exchange_token_with_server for a server of ``auth_type`` and report whether it attempted to persist the exchanged token server-side. The client-forwarded token modes must not persist: @@ -4706,6 +5201,165 @@ async def test_extract_user_id_rejects_expired_key(proxy_globals): assert await _extract_user_id_from_request(request) is None +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_returns_resolved_key_for_active_key(proxy_globals): + """The dcr_bridge mint seals the hash of the authorizing key so admission can reload the live + 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 ( + _resolve_active_litellm_key, + _ResolvedKey, + ) + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-alice-key" + cache = UserApiKeyCache() + await cache.async_set_cache( + hash_token(key), + UserAPIKeyAuth(token=hash_token(key), user_id="alice"), + model_type=UserAPIKeyAuth, + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = object() + + request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) + resolved = await _resolve_active_litellm_key(request) + assert isinstance(resolved, _ResolvedKey) + assert resolved.key_hash == hash_token(key) + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_resolves_key_without_user_id(proxy_globals): + """A valid team-scoped or service-account key has no user_id but is a legitimate credential, so it + must still resolve to a hash and be able to mint a bridge envelope. Gating the resolver on user_id + 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 ( + _extract_user_id_from_request, + _resolve_active_litellm_key, + _ResolvedKey, + ) + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-team-scoped-key" + cache = UserApiKeyCache() + await cache.async_set_cache( + hash_token(key), + UserAPIKeyAuth(token=hash_token(key), user_id=None, team_id="team-x"), + model_type=UserAPIKeyAuth, + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = object() + + request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) + resolved = await _resolve_active_litellm_key(request) + assert isinstance(resolved, _ResolvedKey) + assert resolved.key_hash == hash_token(key) + assert await _extract_user_id_from_request(request) is None + + +@pytest.mark.asyncio +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 ( + _resolve_active_litellm_key, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="blocked-user", blocked=True) + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": "sk-blocked-key"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_fails_closed_on_malformed_expiry(proxy_globals): + """A key whose stored expires string does not parse must fail closed to no-hash (the mint then + 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 ( + _resolve_active_litellm_key, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="u", expires="not-a-parseable-date") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": "sk-bad-expiry-key"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@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 ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + request = _token_request({"content-type": "application/json"}) + assert await _resolve_active_litellm_key(request) == "no_active_key" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_db_outage_is_unavailable(proxy_globals): + """A database outage while resolving the presented key is a retryable infrastructure failure, not + 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 ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _OutagePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + raise ConnectionError("connection refused") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _OutagePrisma() + + request = _token_request({"x-litellm-api-key": "sk-during-outage"}) + assert await _resolve_active_litellm_key(request) == "unavailable" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_no_database_is_unresolvable(proxy_globals): + """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 ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = None + + request = _token_request({"x-litellm-api-key": "sk-no-db"}) + assert await _resolve_active_litellm_key(request) == "unresolvable" + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the