From 34c6cce70562ded634b27771a11f698c46c184dd Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 10 Jul 2026 14:36:20 -0700 Subject: [PATCH 01/12] feat(mcp): mint gateway-bound envelope at the token endpoint for dcr_bridge oauth_delegate --- .../mcp_server/discoverable_endpoints.py | 91 +++++++++++++- .../mcp_server/test_discoverable_endpoints.py | 116 ++++++++++++++++++ 2 files changed, 206 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ebceb320906..41ed49a7508 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -10,7 +10,7 @@ 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 litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -37,6 +37,9 @@ 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 ( + UpstreamTokenGrant, + ) from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. @@ -654,6 +657,86 @@ async def authorize_with_server( return response +def _bridge_grant_from_token_response(token_response: object) -> Optional["UpstreamTokenGrant"]: + """Validate an upstream OAuth token response into a typed grant, or None when it lacks a usable + access token. Each field is isinstance-checked so nothing untyped from ``response.json()`` flows + into the grant.""" + 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 None + access = token_response.get("access_token") + if not isinstance(access, str) or not access: + return None + token_type = token_response.get("token_type") + refresh = token_response.get("refresh_token") + scope = token_response.get("scope") + expires_in = token_response.get("expires_in") + return UpstreamTokenGrant( + access_token=SecretStr(access), + token_type=token_type if isinstance(token_type, str) and token_type else "Bearer", + refresh_token=SecretStr(refresh) if isinstance(refresh, str) and refresh else None, + scope=scope if isinstance(scope, str) and scope else None, + expires_in=expires_in if isinstance(expires_in, int) and expires_in > 0 else None, + ) + + +async def _mint_bridge_delegate_token_response( + request: Request, mcp_server: MCPServer, token_response: object +) -> JSONResponse: + """Return the client-held envelope bearer for a DCR-bridge ``oauth_delegate`` token exchange. + + The envelope binds the caller's litellm identity (resolved from the token request) to the + upstream grant, so the client holds one bearer that later admits it and forwards the upstream + token, with nothing stored server-side. Fails closed with an OAuth ``invalid_request`` when no + litellm identity accompanies the token request rather than minting an identity-less credential. + """ + 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, + 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 + EnvelopeIdentity, + SealedEnvelope, + ) + from litellm.proxy.proxy_server import ( + master_key, # noqa: PLC0415 # inline import avoids a module-load circular import + ) + + if not master_key: + raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") + + user_id = await _extract_user_id_from_request(request) + if not user_id: + raise HTTPException( + status_code=400, + detail={ + "error": "invalid_request", + "error_description": ( + "this server issues a gateway-bound credential; send a litellm credential " + "(x-litellm-api-key or Authorization) on the token request" + ), + }, + ) + + grant = _bridge_grant_from_token_response(token_response) + if grant is None: + raise HTTPException(status_code=502, detail="Upstream token response has no usable access_token") + + now = datetime.now(timezone.utc) + keys = envelope_keys_from_master_key(master_key) + identity = EnvelopeIdentity(user_id=user_id, server_id=mcp_server.server_id) + sealed = build_bridge_token_response(identity, grant, keys, now) + if not isinstance(sealed, SealedEnvelope): + raise HTTPException(status_code=500, detail="Failed to mint the gateway-bound credential") + + expires_in = max(1, int((sealed.expires_at - now).total_seconds())) + 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, @@ -791,6 +874,12 @@ 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 mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: + return await _mint_bridge_delegate_token_response(request, mcp_server, token_response) + result = { "access_token": 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..a2e8d693fab 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,122 @@ 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, user_id): + 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 + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + + 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._extract_user_id_from_request", + new=AsyncMock(return_value=user_id), + ), + patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), + ): + return 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", + ) + + +@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, user_id="user-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.user_id == "user-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 OAuth invalid_request so the client sends a credential.""" + 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} + + with pytest.raises(HTTPException) as exc: + await _exchange_for_bridge_server(server, upstream, user_id=None) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"] == "invalid_request" + + +@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, user_id="user-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, user_id="user-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: From 85255c96fb45fc2473fa5b4ad1939102c8ab9db0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 10 Jul 2026 17:22:59 -0700 Subject: [PATCH 02/12] feat(mcp): seal the authorizing key hash in the dcr_bridge envelope The mint bound only user_id/server_id into the envelope, which gave admission no way to reload the caller's key and enforce its current restrictions. Seal the hashed authorizing key instead (a one-way digest, not a usable credential), so admission reloads the live UserAPIKeyAuth by it and the key's team/org/tool permissions and revocation apply per request. Extract the token endpoint's key resolution into a shared _resolve_active_litellm_key so the per-user token store (user_id) and the bridge mint (key hash) derive from one active-key-gated path, and fail the mint closed with invalid_request when no active key accompanies the request. --- .../mcp_server/discoverable_endpoints.py | 83 +++++++++++++------ .../mcp_server/test_discoverable_endpoints.py | 78 +++++++++++++++-- 2 files changed, 127 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 41ed49a7508..eeadeb290f3 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -351,44 +351,73 @@ def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]: return key_obj.user_id -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``). +async def _resolve_active_litellm_key(request: Request) -> Optional[Tuple[str, "UserAPIKeyAuth"]]: + """Resolve the presented litellm key to ``(its hash, the live active key record)``, or ``None`` + when the key is absent, unresolvable, or blocked/expired. - 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. + Single resolution path the OAuth token endpoint reuses. 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``. The resolved key + is validated (``_active_key_user_id``) before it is trusted, so a blocked or expired key resolves + to ``None``. The returned hash is the value ``get_key_object`` and the cache/DB layer key the + record by. Callers derive the ``user_id`` (per-user token store) or seal the hash (dcr_bridge + envelope) from the result. """ 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 + from litellm.proxy._types import ( # noqa: PLC0415 # inline import avoids a module-load circular import + 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.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import prisma_client, user_api_key_cache, ) + key_hash = hash_token(token) 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 Exception as exc: # noqa: BLE001 # fail closed to None on any key-resolution error 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: could not resolve the presented key (%s)", type(exc).__name__, ) return None + if _active_key_user_id(key_obj) is None: + return None + return key_hash, key_obj + + +async def _extract_user_id_from_request(request: Request) -> Optional[str]: + """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 (``user_api_key_auth.user_id``). ``None`` when no active + key is present. See :func:`_resolve_active_litellm_key` for the resolution and active-key gate. + """ + resolved = await _resolve_active_litellm_key(request) + return _active_key_user_id(resolved[1]) if resolved else None + + +async def _extract_active_key_hash_from_request(request: Request) -> Optional[str]: + """The hash of the litellm key that authorized the token request, when it maps to an active key. + + A DCR-bridge envelope seals this hash so admission can reload the live ``UserAPIKeyAuth`` record + and enforce the key's current team/org/tool restrictions and revocation, rather than trusting a + frozen identity. The hash is a one-way digest, not a usable credential (the edge rejects a bare + hash presented as a bearer). ``None`` when no active key is present, so no envelope is minted for + a missing, unresolvable, or revoked key. + """ + resolved = await _resolve_active_litellm_key(request) + return resolved[0] if resolved else None async def _store_per_user_token_server_side( @@ -688,10 +717,12 @@ async def _mint_bridge_delegate_token_response( ) -> JSONResponse: """Return the client-held envelope bearer for a DCR-bridge ``oauth_delegate`` token exchange. - The envelope binds the caller's litellm identity (resolved from the token request) to the + The envelope binds the authorizing litellm key (its hash, resolved from the token request) to the upstream grant, so the client holds one bearer that later admits it and forwards the upstream - token, with nothing stored server-side. Fails closed with an OAuth ``invalid_request`` when no - litellm identity accompanies the token request rather than minting an identity-less credential. + token, with nothing stored server-side. Admission reloads the live key by that hash, so the key's + current restrictions and revocation gate the request. Fails closed with an OAuth + ``invalid_request`` when no active litellm key accompanies the token request rather than minting + an unbound credential. """ 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, @@ -708,8 +739,8 @@ async def _mint_bridge_delegate_token_response( if not master_key: raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") - user_id = await _extract_user_id_from_request(request) - if not user_id: + key_hash = await _extract_active_key_hash_from_request(request) + if not key_hash: raise HTTPException( status_code=400, detail={ @@ -727,7 +758,7 @@ async def _mint_bridge_delegate_token_response( now = datetime.now(timezone.utc) keys = envelope_keys_from_master_key(master_key) - identity = EnvelopeIdentity(user_id=user_id, server_id=mcp_server.server_id) + identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=key_hash) sealed = build_bridge_token_response(identity, grant, keys, now) if not isinstance(sealed, SealedEnvelope): raise HTTPException(status_code=500, detail="Failed to mint the gateway-bound credential") 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 a2e8d693fab..e69d94d9615 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 @@ -4365,7 +4365,7 @@ async def test_register_bridge_relay_never_persists(): _BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" -async def _exchange_for_bridge_server(server, upstream_body, user_id): +async def _exchange_for_bridge_server(server, upstream_body, key_hash): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( exchange_token_with_server, ) @@ -4382,8 +4382,8 @@ async def _exchange_for_bridge_server(server, upstream_body, user_id): return_value=fake_http_client, ), patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", - new=AsyncMock(return_value=user_id), + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_active_key_hash_from_request", + new=AsyncMock(return_value=key_hash), ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), ): @@ -4416,7 +4416,7 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token 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, user_id="user-77") + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") body = json.loads(response.body) token = body["access_token"] @@ -4429,7 +4429,7 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token 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.user_id == "user-77" + assert opened.identity.key_hash == "hashed-litellm-key-77" assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN" @@ -4443,7 +4443,7 @@ async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600} with pytest.raises(HTTPException) as exc: - await _exchange_for_bridge_server(server, upstream, user_id=None) + await _exchange_for_bridge_server(server, upstream, key_hash=None) assert exc.value.status_code == 400 assert exc.value.detail["error"] == "invalid_request" @@ -4457,7 +4457,7 @@ async def test_true_passthrough_bridge_token_exchange_returns_raw_upstream_token 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, user_id="user-77") + 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" @@ -4472,7 +4472,7 @@ async def test_non_bridge_oauth_delegate_token_exchange_returns_raw_upstream_tok 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, user_id="user-77") + 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" @@ -4822,6 +4822,68 @@ 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_extract_active_key_hash_returns_hash_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 ( + _extract_active_key_hash_from_request, + ) + 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}"}) + assert await _extract_active_key_hash_from_request(request) == hash_token(key) + + +@pytest.mark.asyncio +async def test_extract_active_key_hash_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 ( + _extract_active_key_hash_from_request, + ) + 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 _extract_active_key_hash_from_request(request) is None + + +@pytest.mark.asyncio +async def test_extract_active_key_hash_none_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 ( + _extract_active_key_hash_from_request, + ) + 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 _extract_active_key_hash_from_request(request) is None + + @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 From 7df848aa6c6efe2fd32ea9174cd2e5bc51ed479f Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 14:41:44 -0700 Subject: [PATCH 03/12] fix(mcp): return 502 not KeyError when a bridge upstream response lacks access_token The eager access_token = token_response["access_token"] extraction ran before the dcr_bridge branch, so a missing upstream access_token raised an unhandled KeyError and _bridge_grant_from_token_response's nil guard (which maps to a clean 502) was dead code. Move the extraction onto the non-bridge result path so the bridge branch reaches its 502 guard. --- .../mcp_server/discoverable_endpoints.py | 3 +-- .../mcp_server/test_discoverable_endpoints.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index eeadeb290f3..ef42fef312d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -865,7 +865,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. @@ -912,7 +911,7 @@ async def exchange_token_with_server( return await _mint_bridge_delegate_token_response(request, mcp_server, token_response) 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 e69d94d9615..5826d1f28b6 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 @@ -4449,6 +4449,23 @@ async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm assert exc.value.detail["error"] == "invalid_request" +@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 rather than raising a KeyError. The eager access_token extraction used to run + before the bridge branch, so a missing token raised KeyError and _bridge_grant_from_token_response's + nil guard (which maps to 502) was dead code; the extraction now lives on the non-bridge path only.""" + from litellm.types.mcp import MCPAuth + + server = _bridge_server(auth_type=MCPAuth.oauth_delegate) + upstream = {"token_type": "Bearer", "expires_in": 3600} + + with pytest.raises(HTTPException) as exc: + await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + + assert exc.value.status_code == 502 + + @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 From 362c78e30864f5826a3dc826dc5428c84a7d4f5e Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 15:05:30 -0700 Subject: [PATCH 04/12] fix(mcp): let a keyless-user active key mint a bridge envelope _resolve_active_litellm_key gated on _active_key_user_id, which returns None both for blocked/expired keys AND for valid keys with no user_id, so a team-scoped or service-account key was wrongly rejected with invalid_request at bridge token exchange. Split the active-state gate (_key_is_active: blocked/expiry only) from the user_id extraction; the mint seals the key hash, not the user, and admission already handles a keyless-user key. The per-user token store still gets no user for such a key, as there is none to key a stored credential by. --- .../mcp_server/discoverable_endpoints.py | 44 ++++++++++++------- .../mcp_server/test_discoverable_endpoints.py | 29 ++++++++++++ 2 files changed, 57 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ef42fef312d..df1615aff99 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -329,26 +329,37 @@ 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`. """ 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 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 + + +def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]: + """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 async def _resolve_active_litellm_key(request: Request) -> Optional[Tuple[str, "UserAPIKeyAuth"]]: @@ -361,10 +372,11 @@ async def _resolve_active_litellm_key(request: Request) -> Optional[Tuple[str, " 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``. The resolved key - is validated (``_active_key_user_id``) before it is trusted, so a blocked or expired key resolves - to ``None``. The returned hash is the value ``get_key_object`` and the cache/DB layer key the - record by. Callers derive the ``user_id`` (per-user token store) or seal the hash (dcr_bridge - envelope) from the result. + is validated (``_key_is_active``) before it is trusted, so a blocked or expired key resolves to + ``None``, while a valid team-scoped or service-account key (no ``user_id``) still resolves so it + can mint a bridge envelope. The returned hash is the value ``get_key_object`` and the cache/DB + layer key the record by. Callers derive the ``user_id`` (per-user token store) or seal the hash + (dcr_bridge envelope) from the result. """ token = _litellm_key_from_request(request) if not token: @@ -393,7 +405,7 @@ async def _resolve_active_litellm_key(request: Request) -> Optional[Tuple[str, " type(exc).__name__, ) return None - if _active_key_user_id(key_obj) is None: + if not _key_is_active(key_obj): return None return key_hash, key_obj 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 5826d1f28b6..2bf5e49ec79 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 @@ -4865,6 +4865,35 @@ async def test_extract_active_key_hash_returns_hash_for_active_key(proxy_globals assert await _extract_active_key_hash_from_request(request) == hash_token(key) +@pytest.mark.asyncio +async def test_extract_active_key_hash_returns_hash_for_active_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_active_key_hash_from_request, + _extract_user_id_from_request, + ) + 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}"}) + assert await _extract_active_key_hash_from_request(request) == hash_token(key) + assert await _extract_user_id_from_request(request) is None + + @pytest.mark.asyncio async def test_extract_active_key_hash_rejects_blocked_key(proxy_globals): """A blocked key must not yield a hash, so no gateway-bound envelope is minted for a revoked key; From ceff2d1f3c99f2e314e497e86c302372ba26e64b Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 15:47:27 -0700 Subject: [PATCH 05/12] style(mcp): use X | None annotations on the touched key-resolution helpers The keyless-user fix moved these signatures, so their pre-existing Optional[...] annotations counted against the diff and tripped the UP045 strict-budget gate. Modernize the four touched return annotations to the X | None form the gate wants; runtime behavior is unchanged. --- .../_experimental/mcp_server/discoverable_endpoints.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index df1615aff99..c5d7c3b40fd 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -355,14 +355,14 @@ def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool: return True -def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]: +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 -async def _resolve_active_litellm_key(request: Request) -> Optional[Tuple[str, "UserAPIKeyAuth"]]: +async def _resolve_active_litellm_key(request: Request) -> Tuple[str, "UserAPIKeyAuth"] | None: """Resolve the presented litellm key to ``(its hash, the live active key record)``, or ``None`` when the key is absent, unresolvable, or blocked/expired. @@ -410,7 +410,7 @@ async def _resolve_active_litellm_key(request: Request) -> Optional[Tuple[str, " return key_hash, key_obj -async def _extract_user_id_from_request(request: Request) -> Optional[str]: +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 (``user_api_key_auth.user_id``). ``None`` when no active key is present. See :func:`_resolve_active_litellm_key` for the resolution and active-key gate. @@ -419,7 +419,7 @@ async def _extract_user_id_from_request(request: Request) -> Optional[str]: return _active_key_user_id(resolved[1]) if resolved else None -async def _extract_active_key_hash_from_request(request: Request) -> Optional[str]: +async def _extract_active_key_hash_from_request(request: Request) -> str | None: """The hash of the litellm key that authorized the token request, when it maps to an active key. A DCR-bridge envelope seals this hash so admission can reload the live ``UserAPIKeyAuth`` record From 2f349f6cd18194a67b7c5985e989de5009bec4c9 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 16:16:52 -0700 Subject: [PATCH 06/12] fix(mcp): coerce numeric expires_in and make the active-key check total Two correctness gaps in the bridge mint. _bridge_grant_from_token_response only accepted an int expires_in, dropping a float (3600.0) or numeric-string ('3600') lifetime to None so the envelope fell back to its 1h cap and could outlive a shorter-lived upstream token; coerce it to a positive int (bool excluded). And _key_is_active called datetime.fromisoformat on the str|datetime expires outside the resolver's try, so a malformed stored expiry raised an unhandled 500 instead of the fail-closed invalid_request; it now fails closed (inactive) on an unparseable expiry. Regression tests cover int/float/string/bool coercion, the short-float TTL, and the malformed-expiry fail-closed path. --- .../mcp_server/discoverable_endpoints.py | 39 ++++++++++-- .../mcp_server/test_discoverable_endpoints.py | 60 +++++++++++++++++++ 2 files changed, 95 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index c5d7c3b40fd..00586a6afb3 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -342,12 +342,23 @@ def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool: 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: - 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): @@ -698,10 +709,31 @@ async def authorize_with_server( return response +def _coerce_positive_expires_in(value: object) -> int | None: + """Coerce an upstream ``expires_in`` to a positive int, or ``None`` when it is absent or not a + usable number. IdPs return it as an int, a float (``3600.0``), or a numeric string (``"3600"``); + accepting only ``int`` would drop the float/string cases to ``None`` and fall back to the + envelope's 1h cap, which can outlive a shorter-lived upstream token and forward a stale bearer. + ``bool`` is excluded (it is an ``int`` subclass but never a real lifetime).""" + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + seconds = int(value) + return seconds if seconds > 0 else None + if isinstance(value, str): + try: + seconds = int(float(value.strip())) + except (ValueError, TypeError): + return None + return seconds if seconds > 0 else None + return None + + def _bridge_grant_from_token_response(token_response: object) -> Optional["UpstreamTokenGrant"]: """Validate an upstream OAuth token response into a typed grant, or None when it lacks a usable access token. Each field is isinstance-checked so nothing untyped from ``response.json()`` flows - into the grant.""" + into the grant; ``expires_in`` is numerically coerced so a float/string lifetime is honored + rather than dropped to the envelope's default cap.""" from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import UpstreamTokenGrant, ) @@ -714,13 +746,12 @@ def _bridge_grant_from_token_response(token_response: object) -> Optional["Upstr token_type = token_response.get("token_type") refresh = token_response.get("refresh_token") scope = token_response.get("scope") - expires_in = token_response.get("expires_in") return UpstreamTokenGrant( access_token=SecretStr(access), token_type=token_type if isinstance(token_type, str) and token_type else "Bearer", refresh_token=SecretStr(refresh) if isinstance(refresh, str) and refresh else None, scope=scope if isinstance(scope, str) and scope else None, - expires_in=expires_in if isinstance(expires_in, int) and expires_in > 0 else None, + expires_in=_coerce_positive_expires_in(token_response.get("expires_in")), ) 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 2bf5e49ec79..0dc8d8116df 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 @@ -4449,6 +4449,43 @@ async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm assert exc.value.detail["error"] == "invalid_request" +def test_bridge_grant_coerces_numeric_expires_in(): + """expires_in from an IdP may be an int, a float (3600.0), or a numeric string ("3600"); coerce + it to a positive int so the envelope TTL honors the real lifetime instead of dropping a non-int + value and defaulting to the 1h cap (which can outlive a shorter-lived upstream token). bool and + non-numeric values become None.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _bridge_grant_from_token_response, + ) + + def ei(v): + return _bridge_grant_from_token_response({"access_token": "x", "expires_in": v}).expires_in + + assert ei(300) == 300 + assert ei(300.0) == 300 + assert ei("300") == 300 + assert ei(" 300 ") == 300 + assert ei(True) is None + assert ei("nope") is None + assert ei(0) is None + assert ei(-5) is None + assert ei(None) is None + + +@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 @@ -4915,6 +4952,29 @@ async def test_extract_active_key_hash_rejects_blocked_key(proxy_globals): assert await _extract_active_key_hash_from_request(request) is None +@pytest.mark.asyncio +async def test_extract_active_key_hash_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 ( + _extract_active_key_hash_from_request, + ) + 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 _extract_active_key_hash_from_request(request) is None + + @pytest.mark.asyncio async def test_extract_active_key_hash_none_without_litellm_key(proxy_globals): """No LiteLLM key on the request yields no hash without consulting the resolver.""" From 7a63e516252a5fc78a6150da4b3fc6cd03bab1c6 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 16:36:08 -0700 Subject: [PATCH 07/12] fix(mcp): harden the bridge token mint (multi-lens review pass) Findings from a full adversarial review of the mint path across security, correctness, error-handling, concurrency, and OAuth-protocol dimensions. - expires_in coercion is now total: int(float(...)) can raise OverflowError on Infinity / a giant numeric string, which escaped the ValueError/TypeError catch and 500'd the token endpoint. Unified to catch OverflowError too. - Resolve the litellm identity BEFORE exchanging the single-use upstream code, so a missing or transiently-unresolvable identity fails closed with invalid_request without burning the code (the mint re-resolves via a cache hit). - The no-identity failure is now an RFC 6749 5.2-shaped invalid_request (JSONResponse, top-level error, no-store) instead of a detail-wrapped HTTPException, matching the BYOK OAuth endpoint. - EnvelopeTooLarge (upstream token too big to seal) surfaces a 502, not a 500. - The upstream refresh_token is no longer sealed into the envelope: the edge never consumes it, so it was dead weight embedding a long-lived upstream credential in the client bearer and enlarging the envelope; refresh is a follow-up (a dedicated refresh-envelope). Security review found no exploitable defect (forgery, cross-server/user replay, leakage, confused-deputy all closed). Regression tests cover the OverflowError, the code-not-burned path, the RFC-shaped error, the 502, and the dropped refresh. --- .../mcp_server/discoverable_endpoints.py | 76 ++++++++++++------- .../mcp_server/test_discoverable_endpoints.py | 71 +++++++++++++++-- 2 files changed, 116 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 00586a6afb3..4f72c2cb8c2 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -373,7 +373,7 @@ def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> str | None: return key_obj.user_id if _key_is_active(key_obj) else None -async def _resolve_active_litellm_key(request: Request) -> Tuple[str, "UserAPIKeyAuth"] | None: +async def _resolve_active_litellm_key(request: Request) -> tuple[str, "UserAPIKeyAuth"] | None: """Resolve the presented litellm key to ``(its hash, the live active key record)``, or ``None`` when the key is absent, unresolvable, or blocked/expired. @@ -714,19 +714,17 @@ def _coerce_positive_expires_in(value: object) -> int | None: usable number. IdPs return it as an int, a float (``3600.0``), or a numeric string (``"3600"``); accepting only ``int`` would drop the float/string cases to ``None`` and fall back to the envelope's 1h cap, which can outlive a shorter-lived upstream token and forward a stale bearer. - ``bool`` is excluded (it is an ``int`` subclass but never a real lifetime).""" - if isinstance(value, bool): + ``bool`` is excluded (it is an ``int`` subclass but never a real lifetime). Total over hostile + input: a non-numeric string, ``NaN``, ``Infinity``, or an over-large value all resolve to + ``None`` rather than raising (``int(float(...))`` can raise ``ValueError`` or ``OverflowError``), + so a malformed upstream ``expires_in`` never surfaces as a 500 from the token endpoint.""" + if isinstance(value, bool) or not isinstance(value, (int, float, str)): return None - if isinstance(value, (int, float)): - seconds = int(value) - return seconds if seconds > 0 else None - if isinstance(value, str): - try: - seconds = int(float(value.strip())) - except (ValueError, TypeError): - return None - return seconds if seconds > 0 else None - return None + try: + seconds = int(float(value)) + except (ValueError, TypeError, OverflowError): + return None + return seconds if seconds > 0 else None def _bridge_grant_from_token_response(token_response: object) -> Optional["UpstreamTokenGrant"]: @@ -744,17 +742,38 @@ def _bridge_grant_from_token_response(token_response: object) -> Optional["Upstr if not isinstance(access, str) or not access: return None token_type = token_response.get("token_type") - refresh = token_response.get("refresh_token") 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", - refresh_token=SecretStr(refresh) if isinstance(refresh, str) and refresh else None, + # 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=_coerce_positive_expires_in(token_response.get("expires_in")), ) +def _bridge_invalid_request_response() -> JSONResponse: + """RFC 6749 §5.2-shaped ``invalid_request`` for a bridge token exchange that carries no resolvable + litellm identity. Returned (not raised) so the OAuth error members sit at the top level rather than + wrapped in FastAPI's ``detail``, with the no-store token-endpoint headers, matching the BYOK OAuth + endpoint and what a strict DCR client parses per RFC 6749 §5.2.""" + return JSONResponse( + status_code=400, + content={ + "error": "invalid_request", + "error_description": ( + "this server issues a gateway-bound credential; send a litellm credential " + "(x-litellm-api-key or Authorization) on the token request" + ), + }, + headers=TOKEN_NO_CACHE_HEADERS, + ) + + async def _mint_bridge_delegate_token_response( request: Request, mcp_server: MCPServer, token_response: object ) -> JSONResponse: @@ -784,16 +803,7 @@ async def _mint_bridge_delegate_token_response( key_hash = await _extract_active_key_hash_from_request(request) if not key_hash: - raise HTTPException( - status_code=400, - detail={ - "error": "invalid_request", - "error_description": ( - "this server issues a gateway-bound credential; send a litellm credential " - "(x-litellm-api-key or Authorization) on the token request" - ), - }, - ) + return _bridge_invalid_request_response() grant = _bridge_grant_from_token_response(token_response) if grant is None: @@ -804,7 +814,11 @@ async def _mint_bridge_delegate_token_response( identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=key_hash) sealed = build_bridge_token_response(identity, grant, keys, now) if not isinstance(sealed, SealedEnvelope): - raise HTTPException(status_code=500, detail="Failed to mint the gateway-bound credential") + # build_bridge_token_response returns EnvelopeTooLarge as a value when the upstream token is + # too large to seal; that is an upstream-payload condition, so surface a 502, not a 500. + raise HTTPException( + status_code=502, detail="Upstream token is too large to seal into a gateway-bound credential" + ) expires_in = max(1, int((sealed.expires_at - now).total_seconds())) body = {"access_token": sealed.token.get_secret_value(), "token_type": "Bearer", "expires_in": expires_in} @@ -884,6 +898,16 @@ async def exchange_token_with_server( if code_verifier: token_data["code_verifier"] = code_verifier + # For a bridge oauth_delegate mint, resolve the litellm identity BEFORE exchanging the + # single-use upstream code. A missing or transiently-unresolvable identity then fails closed + # with invalid_request without consuming the code, so the client can retry the same code + # instead of being forced back through the full interactive authorize. The mint below + # re-resolves authoritatively; get_key_object is cache-first, so that second call is a cache + # hit and this adds no extra database round-trip. + if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: + if not await _extract_active_key_hash_from_request(request): + return _bridge_invalid_request_response() + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response = await async_client.post( mcp_server.token_url, 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 0dc8d8116df..8aad8b24e88 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 @@ -4365,7 +4365,7 @@ async def test_register_bridge_relay_never_persists(): _BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" -async def _exchange_for_bridge_server(server, upstream_body, key_hash): +async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_client_out=None): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( exchange_token_with_server, ) @@ -4375,6 +4375,8 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash): fake_http_response.raise_for_status = MagicMock() fake_http_client = MagicMock() fake_http_client.post = AsyncMock(return_value=fake_http_response) + if fake_client_out is not None: + fake_client_out["client"] = fake_http_client with ( patch( @@ -4436,17 +4438,69 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_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 OAuth invalid_request so the client sends a credential.""" + 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 rather than a 500 (build_bridge_token_response returns EnvelopeTooLarge as a + value, and the caller 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} with pytest.raises(HTTPException) as exc: - await _exchange_for_bridge_server(server, upstream, key_hash=None) + await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + assert exc.value.status_code == 502 - assert exc.value.status_code == 400 - assert exc.value.detail["error"] == "invalid_request" + +@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 def test_bridge_grant_coerces_numeric_expires_in(): @@ -4470,6 +4524,13 @@ def test_bridge_grant_coerces_numeric_expires_in(): assert ei(0) is None assert ei(-5) is None assert ei(None) is None + # hostile numerics must not raise (int(float(...)) can OverflowError) -> None + assert ei("inf") is None + assert ei("1e999") is None + assert ei("-inf") is None + assert ei("nan") is None + assert ei(float("inf")) is None + assert ei(10**400) is None @pytest.mark.asyncio From e16ad044c3773ac958b1cffff0ad6d15bb5e0296 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 16:58:14 -0700 Subject: [PATCH 08/12] fix(mcp): close the burn-before-check gate for both grants and validate master_key first Follow-up to the pre-exchange identity gate, which I had only added to the authorization_code branch and which left the master_key check inside the mint (after the upstream exchange) - so the very burn-then-fail pattern it was meant to prevent still applied to refresh_token grants and to a misconfigured gateway. - Hoist a single pre-exchange gate above the upstream call that covers BOTH grant types: it fails closed (invalid_request) on an unresolvable litellm identity and 500s on an unset master_key BEFORE the single-use code or refresh token is exchanged/rotated, so a bad key or a misconfigured gateway never burns the upstream credential. - Report expires_in from the envelope JWT's own second-truncated exp (rounding the elapsed portion up) instead of the raw expires_at - now delta, so the client is never told the bearer is valid past the ~1s point admission already expires it. Regression tests assert the upstream exchange is never called on the no-identity refresh grant and the master_key-unset path, and that the reported expires_in does not overstate the JWT exp. --- .../mcp_server/discoverable_endpoints.py | 28 ++++-- .../mcp_server/test_discoverable_endpoints.py | 97 +++++++++++++++++++ 2 files changed, 115 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 4f72c2cb8c2..5bf4de09bba 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,6 +1,7 @@ import asyncio import html as _html import json +import math import secrets import time from datetime import datetime, timezone @@ -820,7 +821,10 @@ async def _mint_bridge_delegate_token_response( status_code=502, detail="Upstream token is too large to seal into a gateway-bound credential" ) - expires_in = max(1, int((sealed.expires_at - now).total_seconds())) + # The JWT exp is int(expires_at.timestamp()) (second-truncated), and admission expires the envelope + # against that exp. Report expires_in from the same truncated exp, rounding the elapsed portion up, + # so the client is never told the bearer lives past the point admission already rejects it. + expires_in = max(1, 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) @@ -898,15 +902,19 @@ async def exchange_token_with_server( if code_verifier: token_data["code_verifier"] = code_verifier - # For a bridge oauth_delegate mint, resolve the litellm identity BEFORE exchanging the - # single-use upstream code. A missing or transiently-unresolvable identity then fails closed - # with invalid_request without consuming the code, so the client can retry the same code - # instead of being forced back through the full interactive authorize. The mint below - # re-resolves authoritatively; get_key_object is cache-first, so that second call is a cache - # hit and this adds no extra database round-trip. - if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: - if not await _extract_active_key_hash_from_request(request): - return _bridge_invalid_request_response() + # A bridge oauth_delegate mint must fail closed BEFORE the upstream exchange consumes or rotates the + # single-use code (or refresh token): confirm the gateway can mint at all (master_key set) and that + # the request carries a resolvable litellm identity. Applies to both grant types, so an invalid key + # or a misconfigured gateway never burns the upstream credential. The mint below re-checks + # authoritatively; get_key_object is cache-first, so the identity re-resolution is a cache hit and + # adds no extra database round-trip. + if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: + from litellm.proxy.proxy_server import master_key as _bridge_master_key # noqa: PLC0415 + + if not _bridge_master_key: + raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") + if not await _extract_active_key_hash_from_request(request): + return _bridge_invalid_request_response() async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response = await async_client.post( 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 8aad8b24e88..4a4ff398915 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 @@ -4503,6 +4503,103 @@ async def test_bridge_envelope_does_not_seal_upstream_refresh_token(): assert opened.grant.refresh_token is None +@pytest.mark.asyncio +async def test_bridge_refresh_grant_fails_closed_before_upstream_when_no_identity(): + """The pre-exchange identity gate covers the refresh_token grant, not just authorization_code: an + unresolvable litellm identity fails closed with invalid_request BEFORE the upstream refresh is + exchanged, so the client's refresh token is not rotated/consumed on a rejected request.""" + 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._extract_active_key_hash_from_request", + new=AsyncMock(return_value=None), + ), + 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"] == "invalid_request" + 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, so a misconfigured gateway 500s without + consuming the single-use code, avoiding the burn-then-fail the pre-exchange gate exists to prevent.""" + 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._extract_active_key_hash_from_request", + new=AsyncMock(return_value="hashed-litellm-key-77"), + ), + patch("litellm.proxy.proxy_server.master_key", None), + ): + with pytest.raises(HTTPException) as exc: + 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 exc.value.status_code == 500 + fake_http_client.post.assert_not_called() + + +@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_grant_coerces_numeric_expires_in(): """expires_in from an IdP may be an int, a float (3600.0), or a numeric string ("3600"); coerce it to a positive int so the envelope TTL honors the real lifetime instead of dropping a non-int From 2f0ddc82f7f38b282dff5e299436cd99285e1ba7 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 17:22:01 -0700 Subject: [PATCH 09/12] refactor(mcp): make the bridge delegate mint a phased failures-as-values pipeline The dcr_bridge oauth_delegate token mint validated its preconditions in two places: a pre-exchange guard inside exchange_token_with_server (master_key set, resolvable litellm identity) and an authoritative re-check inside the post-exchange _mint_bridge_delegate_token_response. Keeping the two in step by hand is what kept producing the same class of finding: a precondition guarded on one grant branch but not the other, master_key checked after the exchange on one path, identity resolved twice, and each failure raising an ad-hoc HTTPException with its own status and body shape. Model the mint as three phases whose failures are values. _prepare_bridge_mint runs before the exchange, checks every precondition once (master_key, then identity), and returns either a frozen _BridgeMintReady carrying the resolved key hash and the master-key-derived envelope keys, or a _BridgeMintError literal. Because every precondition lives in prepare, and prepare runs before the upstream POST, no failure can burn the single-use code or rotate a refresh token, for either grant type, by construction rather than by a guard we have to remember to keep in sync. _finish_bridge_mint runs after the exchange and has no preconditions left that can fail; its only failure values are properties of the upstream response itself (no usable access_token, or a token too large to seal). One mapper, _bridge_mint_error_response, turns each _BridgeMintError into an RFC 6749 section 5.2-shaped body with a status truthful about where the failure is (400 for the caller, 500 for gateway config, 502 for the upstream), with an exhaustive match plus assert_never so a new failure mode cannot be added without a matching status. Behavior is unchanged for the client. Every failure that previously raised now returns the same status as an OAuth error body, which is the correct token-endpoint contract; the three tests that asserted a raised HTTPException now assert the returned response. _exchange_for_bridge_server additionally asserts the identity resolver is awaited exactly once for a bridge server and never for a non-bridge one. --- .../mcp_server/discoverable_endpoints.py | 167 +++++++++++------- .../mcp_server/test_discoverable_endpoints.py | 63 ++++--- 2 files changed, 140 insertions(+), 90 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 5bf4de09bba..1bdbc8ea8a4 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -4,6 +4,7 @@ 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 @@ -12,6 +13,7 @@ import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response 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 ( @@ -39,6 +41,7 @@ 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 @@ -757,73 +760,112 @@ def _bridge_grant_from_token_response(token_response: object) -> Optional["Upstr ) -def _bridge_invalid_request_response() -> JSONResponse: - """RFC 6749 §5.2-shaped ``invalid_request`` for a bridge token exchange that carries no resolvable - litellm identity. Returned (not raised) so the OAuth error members sit at the top level rather than - wrapped in FastAPI's ``detail``, with the no-store token-endpoint headers, matching the BYOK OAuth - endpoint and what a strict DCR client parses per RFC 6749 §5.2.""" - return JSONResponse( - status_code=400, - content={ - "error": "invalid_request", - "error_description": ( +# --------------------------------------------------------------------------- +# 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["not_configured", "no_identity", "no_upstream_token", "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) for every case, with a status truthful about where the failure is: + the caller's request (400), the gateway config (500), or the upstream (502).""" + if error == "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" ), - }, - headers=TOKEN_NO_CACHE_HEADERS, + ) + elif error == "not_configured": + status, code, desc = ( + 500, + "server_error", + ("the gateway is not configured to mint a gateway-bound credential (master_key is not set)"), + ) + elif error == "no_upstream_token": + status, code, desc = 502, "server_error", "the upstream token response has no usable access_token" + elif error == "too_large": + status, code, desc = ( + 502, + "server_error", + ("the upstream token is too large to seal into a gateway-bound credential"), + ) + else: + assert_never(error) + return JSONResponse( + status_code=status, content={"error": code, "error_description": desc}, headers=TOKEN_NO_CACHE_HEADERS ) -async def _mint_bridge_delegate_token_response( - request: Request, mcp_server: MCPServer, token_response: object -) -> JSONResponse: - """Return the client-held envelope bearer for a DCR-bridge ``oauth_delegate`` token exchange. +async def _prepare_bridge_mint(request: Request, mcp_server: MCPServer) -> "_BridgeMintReady | _BridgeMintError": + """Phase 1, BEFORE the upstream exchange: validate that the gateway can mint (master_key set) and + that the request carries a resolvable litellm identity, and derive the envelope keys. Returns a + ready context or a failure value. Running before the exchange is what makes a missing master_key or + an unresolvable identity fail closed without consuming the single-use code / rotating a refresh + token, for both grant types.""" + 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, + ) - The envelope binds the authorizing litellm key (its hash, resolved from the token request) to the - upstream grant, so the client holds one bearer that later admits it and forwards the upstream - token, with nothing stored server-side. Admission reloads the live key by that hash, so the key's - current restrictions and revocation gate the request. Fails closed with an OAuth - ``invalid_request`` when no active litellm key accompanies the token request rather than minting - an unbound credential. - """ + if not master_key: + return "not_configured" + key_hash = await _extract_active_key_hash_from_request(request) + if not key_hash: + return "no_identity" + return _BridgeMintReady(key_hash=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, 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, - 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 EnvelopeIdentity, SealedEnvelope, ) - from litellm.proxy.proxy_server import ( - master_key, # noqa: PLC0415 # inline import avoids a module-load circular import - ) - - if not master_key: - raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") - - key_hash = await _extract_active_key_hash_from_request(request) - if not key_hash: - return _bridge_invalid_request_response() grant = _bridge_grant_from_token_response(token_response) if grant is None: - raise HTTPException(status_code=502, detail="Upstream token response has no usable access_token") - - now = datetime.now(timezone.utc) - keys = envelope_keys_from_master_key(master_key) - identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=key_hash) - sealed = build_bridge_token_response(identity, grant, keys, now) + return "no_upstream_token" + 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): - # build_bridge_token_response returns EnvelopeTooLarge as a value when the upstream token is - # too large to seal; that is an upstream-payload condition, so surface a 502, not a 500. - raise HTTPException( - status_code=502, detail="Upstream token is too large to seal into a gateway-bound credential" - ) - - # The JWT exp is int(expires_at.timestamp()) (second-truncated), and admission expires the envelope - # against that exp. Report expires_in from the same truncated exp, rounding the elapsed portion up, - # so the client is never told the bearer lives past the point admission already rejects it. + 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(1, 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) @@ -902,19 +944,15 @@ async def exchange_token_with_server( if code_verifier: token_data["code_verifier"] = code_verifier - # A bridge oauth_delegate mint must fail closed BEFORE the upstream exchange consumes or rotates the - # single-use code (or refresh token): confirm the gateway can mint at all (master_key set) and that - # the request carries a resolvable litellm identity. Applies to both grant types, so an invalid key - # or a misconfigured gateway never burns the upstream credential. The mint below re-checks - # authoritatively; get_key_object is cache-first, so the identity re-resolution is a cache hit and - # adds no extra database round-trip. + # 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: - from litellm.proxy.proxy_server import master_key as _bridge_master_key # noqa: PLC0415 - - if not _bridge_master_key: - raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set") - if not await _extract_active_key_hash_from_request(request): - return _bridge_invalid_request_response() + prepared = await _prepare_bridge_mint(request, mcp_server) + 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( @@ -982,8 +1020,11 @@ async def exchange_token_with_server( # 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 mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge: - return await _mint_bridge_delegate_token_response(request, mcp_server, token_response) + 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": token_response["access_token"], 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 4a4ff398915..4bbaf7b0b76 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 @@ -4375,6 +4375,7 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie fake_http_response.raise_for_status = MagicMock() fake_http_client = MagicMock() fake_http_client.post = AsyncMock(return_value=fake_http_response) + key_resolver = AsyncMock(return_value=key_hash) if fake_client_out is not None: fake_client_out["client"] = fake_http_client @@ -4385,11 +4386,11 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie ), patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_active_key_hash_from_request", - new=AsyncMock(return_value=key_hash), + new=key_resolver, ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), ): - return await exchange_token_with_server( + response = await exchange_token_with_server( request=_bridge_mock_request(), mcp_server=server, grant_type="authorization_code", @@ -4399,6 +4400,11 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie 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 @@ -4457,15 +4463,16 @@ async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm @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 rather than a 500 (build_bridge_token_response returns EnvelopeTooLarge as a - value, and the caller maps it to a truthful status).""" + 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} - with pytest.raises(HTTPException) as exc: - await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") - assert exc.value.status_code == 502 + 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 @@ -4544,8 +4551,10 @@ async def test_bridge_refresh_grant_fails_closed_before_upstream_when_no_identit @pytest.mark.asyncio async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset(): - """master_key is validated BEFORE the upstream exchange, so a misconfigured gateway 500s without - consuming the single-use code, avoiding the burn-then-fail the pre-exchange gate exists to prevent.""" + """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 @@ -4563,19 +4572,19 @@ async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset(): ), patch("litellm.proxy.proxy_server.master_key", None), ): - with pytest.raises(HTTPException) as exc: - 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", - ) + 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 exc.value.status_code == 500 + assert response.status_code == 500 + assert json.loads(response.body)["error"] == "server_error" fake_http_client.post.assert_not_called() @@ -4647,18 +4656,18 @@ async def test_bridge_token_exchange_honors_short_float_expires_in_ttl(): @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 rather than raising a KeyError. The eager access_token extraction used to run - before the bridge branch, so a missing token raised KeyError and _bridge_grant_from_token_response's - nil guard (which maps to 502) was dead code; the extraction now lives on the non-bridge path only.""" + 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} - with pytest.raises(HTTPException) as exc: - await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") + response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77") - assert exc.value.status_code == 502 + assert response.status_code == 502 + assert json.loads(response.body)["error"] == "server_error" @pytest.mark.asyncio From 4ba7221b7a1ae717d55ff548247aa8416dbe8232 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 17:47:54 -0700 Subject: [PATCH 10/12] fix(mcp): let the bridge envelope report expires_in 0 at the jwt exp boundary _finish_bridge_mint floored the reported expires_in at 1. Admission expires the envelope against the JWT's second-truncated exp, so when the mint lands in the same second that exp falls on (a sub-second upstream lifetime, for instance), the true remaining life is 0 and reporting 1 tells the client the bearer lives one second past the point admission already rejects it. Floor at 0 instead so the reported lifetime never overstates the exp; the value still cannot go negative. The regression pins the boundary directly: minting at now=100.25 with a 1s upstream token seals exp=101, and the reported expires_in is max(0, 101 - ceil(100.25)) = 0. Under the old floor of 1 it reads 1, so the test fails on that mutation. Also drops the unused mcp_server parameter from _prepare_bridge_mint; identity and key derivation there never referenced the server. --- .../mcp_server/discoverable_endpoints.py | 6 ++-- .../mcp_server/test_discoverable_endpoints.py | 29 +++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 1bdbc8ea8a4..abb375b5b6a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -821,7 +821,7 @@ def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: ) -async def _prepare_bridge_mint(request: Request, mcp_server: MCPServer) -> "_BridgeMintReady | _BridgeMintError": +async def _prepare_bridge_mint(request: Request) -> "_BridgeMintReady | _BridgeMintError": """Phase 1, BEFORE the upstream exchange: validate that the gateway can mint (master_key set) and that the request carries a resolvable litellm identity, and derive the envelope keys. Returns a ready context or a failure value. Running before the exchange is what makes a missing master_key or @@ -866,7 +866,7 @@ def _finish_bridge_mint( 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(1, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp())) + 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) @@ -949,7 +949,7 @@ async def exchange_token_with_server( # 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, mcp_server) + prepared = await _prepare_bridge_mint(request) if not isinstance(prepared, _BridgeMintReady): return _bridge_mint_error_response(prepared) bridge_mint_ready = prepared 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 4bbaf7b0b76..4aa23247249 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 @@ -4609,6 +4609,35 @@ async def test_bridge_reported_expires_in_does_not_overstate_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_bridge_grant_coerces_numeric_expires_in(): """expires_in from an IdP may be an int, a float (3600.0), or a numeric string ("3600"); coerce it to a positive int so the envelope TTL honors the real lifetime instead of dropping a non-int From a07aba05798941c75e369789831e701801355daa Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 18:28:32 -0700 Subject: [PATCH 11/12] refactor(mcp): make bridge-mint resolvers return tagged unions so status is truthful by construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three findings landed together, all one defect: a resolution step crushed several distinct outcomes into a single None or a silent default, so the mint's error mapper could not tell them apart and assigned the wrong status. Identity resolution mapped a database outage to the same None as a missing credential, which the mint reported as 400 invalid_request, blaming the caller for a gateway outage while admission statuses the same outage 503/500. Lifetime coercion mapped an explicit non-positive expires_in to the same None as an absent one, so an upstream token the IdP reports as already dead was sealed into an hour-long envelope. And the refresh_token grant was run through the upstream exchange (which can rotate the client's upstream refresh credential) and its result then discarded, even though a bridge server seals no refresh_token and the client never holds one to present. Rather than add a mapping branch per finding, the fix changes the return types so a wrong status is not representable. Each resolution step now returns a precise tagged value instead of None: identity resolution returns a _ResolvedKey or one of no_active_key / unavailable / unresolvable, classified the same way admission's _reload_admitted_key classifies the same conditions; upstream-lifetime classification returns a positive number of seconds, "unspecified" (absent or unparseable, which the envelope caps), or "expired" (a parseable non-positive value, an already-dead token); and upstream-grant validation returns a typed grant or one of no_access_token / expired_lifetime. Thin exhaustive mappers (match plus assert_never) lift each vocabulary into one bridge-mint taxonomy of eight named failures, and a single _bridge_mint_error_response gives each its truthful RFC 6749 §5.2 status: 400 for the caller's missing credential or an unsupported grant, 503 for a transient auth-DB outage, 500 for a gateway that cannot resolve identity or is not configured, and 502 for an upstream response with no usable token, an already-expired lifetime, or a token too large to seal. Adding a failure mode now requires a new literal and a match arm the type checker forces, so the class of wrong-status bug cannot recur silently. The refresh_token grant is rejected in _prepare_bridge_mint before the exchange with unsupported_grant_type, so it can never rotate or consume the client's upstream refresh credential; renewal is re-running authorization_code, as the sealed refresh_token=None already intends. An absent or unparseable expires_in still mints a capped envelope (the by-design behaviour for an upstream that omits the field); only an explicitly-dead lifetime is rejected. Tests cover the resolver's three failure classes (including a real connection-error outage and a missing prisma_client), the mint statuses for each (503 before the upstream exchange, 500, 502 on an expired upstream lifetime, and a capped mint on an unknown one), and the refresh-grant rejection before any exchange. The three findings are mutation-checked: reverting each fix turns its regression test red. --- .../mcp_server/discoverable_endpoints.py | 334 ++++++++++++------ .../mcp_server/test_discoverable_endpoints.py | 259 +++++++++++--- 2 files changed, 424 insertions(+), 169 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index abb375b5b6a..b8d10335182 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -377,74 +377,92 @@ def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> str | None: return key_obj.user_id if _key_is_active(key_obj) else None -async def _resolve_active_litellm_key(request: Request) -> tuple[str, "UserAPIKeyAuth"] | None: - """Resolve the presented litellm key to ``(its hash, the live active key record)``, or ``None`` - when the key is absent, unresolvable, or 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.""" - Single resolution path the OAuth token endpoint reuses. 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``. The resolved key - is validated (``_key_is_active``) before it is trusted, so a blocked or expired key resolves to - ``None``, while a valid team-scoped or service-account key (no ``user_id``) still resolves so it - can mint a bridge envelope. The returned hash is the value ``get_key_object`` and the cache/DB - layer key the record by. Callers derive the ``user_id`` (per-user token store) or seal the hash - (dcr_bridge envelope) from the result. - """ + 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 ( # noqa: PLC0415 # inline import avoids a module-load circular import - 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.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import - 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, + ) - key_hash = hash_token(token) + if prisma_client is None: + return "unresolvable" + key_hash = hash_token(token) + try: key_obj = await get_key_object( hashed_token=key_hash, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, ) - except Exception as exc: # noqa: BLE001 # fail closed to None on any key-resolution error + 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( - "_resolve_active_litellm_key: could not resolve the presented key (%s)", + "_resolve_active_litellm_key: unexpected key-resolution error (%s)", type(exc).__name__, ) - return None + return "unresolvable" if not _key_is_active(key_obj): - return None - return key_hash, 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 (``user_api_key_auth.user_id``). ``None`` when no active - key is present. See :func:`_resolve_active_litellm_key` for the resolution and active-key gate. - """ + """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) - return _active_key_user_id(resolved[1]) if resolved else None - - -async def _extract_active_key_hash_from_request(request: Request) -> str | None: - """The hash of the litellm key that authorized the token request, when it maps to an active key. - - A DCR-bridge envelope seals this hash so admission can reload the live ``UserAPIKeyAuth`` record - and enforce the key's current team/org/tool restrictions and revocation, rather than trusting a - frozen identity. The hash is a one-way digest, not a usable credential (the edge rejects a bare - hash presented as a bearer). ``None`` when no active key is present, so no envelope is minted for - a missing, unresolvable, or revoked key. - """ - resolved = await _resolve_active_litellm_key(request) - return resolved[0] if resolved else None + if not isinstance(resolved, _ResolvedKey): + return None + return _active_key_user_id(resolved.key) async def _store_per_user_token_server_side( @@ -713,38 +731,51 @@ async def authorize_with_server( return response -def _coerce_positive_expires_in(value: object) -> int | None: - """Coerce an upstream ``expires_in`` to a positive int, or ``None`` when it is absent or not a - usable number. IdPs return it as an int, a float (``3600.0``), or a numeric string (``"3600"``); - accepting only ``int`` would drop the float/string cases to ``None`` and fall back to the - envelope's 1h cap, which can outlive a shorter-lived upstream token and forward a stale bearer. - ``bool`` is excluded (it is an ``int`` subclass but never a real lifetime). Total over hostile - input: a non-numeric string, ``NaN``, ``Infinity``, or an over-large value all resolve to - ``None`` rather than raising (``int(float(...))`` can raise ``ValueError`` or ``OverflowError``), - so a malformed upstream ``expires_in`` never surfaces as a 500 from the token endpoint.""" - if isinstance(value, bool) or not isinstance(value, (int, float, str)): - return None +_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 parseable 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. ``bool`` is excluded (an ``int`` subclass but never a real lifetime), and + ``int(float(...))`` 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: - seconds = int(float(value)) + seconds = int(float(raw_expires_in)) except (ValueError, TypeError, OverflowError): - return None - return seconds if seconds > 0 else None + return "unspecified" + return seconds if seconds > 0 else "expired" -def _bridge_grant_from_token_response(token_response: object) -> Optional["UpstreamTokenGrant"]: - """Validate an upstream OAuth token response into a typed grant, or None when it lacks a usable - access token. Each field is isinstance-checked so nothing untyped from ``response.json()`` flows - into the grant; ``expires_in`` is numerically coerced so a float/string lifetime is honored - rather than dropped to the envelope's default cap.""" +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 None + return "no_access_token" access = token_response.get("access_token") if not isinstance(access, str) or not access: - return None + 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( @@ -756,7 +787,7 @@ def _bridge_grant_from_token_response(token_response: object) -> Optional["Upstr # 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=_coerce_positive_expires_in(token_response.get("expires_in")), + expires_in=lifetime if isinstance(lifetime, int) else None, ) @@ -774,7 +805,16 @@ def _bridge_grant_from_token_response(token_response: object) -> Optional["Upstr # shape are uniform. Adding a failure mode is a new literal plus a match arm the type checker forces. # --------------------------------------------------------------------------- -_BridgeMintError = Literal["not_configured", "no_identity", "no_upstream_token", "too_large"] +_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) @@ -788,45 +828,105 @@ class _BridgeMintReady: 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) for every case, with a status truthful about where the failure is: - the caller's request (400), the gateway config (500), or the upstream (502).""" - if error == "no_identity": - status, code, desc = ( - 400, - "invalid_request", - ( + """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" - ), - ) - elif error == "not_configured": - status, code, desc = ( - 500, - "server_error", - ("the gateway is not configured to mint a gateway-bound credential (master_key is not set)"), - ) - elif error == "no_upstream_token": - status, code, desc = 502, "server_error", "the upstream token response has no usable access_token" - elif error == "too_large": - status, code, desc = ( - 502, - "server_error", - ("the upstream token is too large to seal into a gateway-bound credential"), - ) - else: - assert_never(error) + "(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 ) -async def _prepare_bridge_mint(request: Request) -> "_BridgeMintReady | _BridgeMintError": - """Phase 1, BEFORE the upstream exchange: validate that the gateway can mint (master_key set) and - that the request carries a resolvable litellm identity, and derive the envelope keys. Returns a - ready context or a failure value. Running before the exchange is what makes a missing master_key or - an unresolvable identity fail closed without consuming the single-use code / rotating a refresh - token, for both grant types.""" +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, ) @@ -834,12 +934,14 @@ async def _prepare_bridge_mint(request: Request) -> "_BridgeMintReady | _BridgeM master_key, ) + if grant_type != "authorization_code": + return "unsupported_grant" if not master_key: return "not_configured" - key_hash = await _extract_active_key_hash_from_request(request) - if not key_hash: - return "no_identity" - return _BridgeMintReady(key_hash=key_hash, keys=envelope_keys_from_master_key(master_key)) + 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( @@ -848,18 +950,20 @@ def _finish_bridge_mint( """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, or a token too large to seal), returned as values.""" + 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 grant is None: - return "no_upstream_token" + 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): @@ -949,7 +1053,7 @@ async def exchange_token_with_server( # 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) + prepared = await _prepare_bridge_mint(request, grant_type) if not isinstance(prepared, _BridgeMintReady): return _bridge_mint_error_response(prepared) bridge_mint_ready = prepared 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 4aa23247249..f611d2f6a24 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 @@ -4367,6 +4367,7 @@ _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, ) @@ -4375,7 +4376,10 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie fake_http_response.raise_for_status = MagicMock() fake_http_client = MagicMock() fake_http_client.post = AsyncMock(return_value=fake_http_response) - key_resolver = AsyncMock(return_value=key_hash) + # 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 @@ -4385,7 +4389,7 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie return_value=fake_http_client, ), patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_active_key_hash_from_request", + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._resolve_active_litellm_key", new=key_resolver, ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), @@ -4511,10 +4515,12 @@ async def test_bridge_envelope_does_not_seal_upstream_refresh_token(): @pytest.mark.asyncio -async def test_bridge_refresh_grant_fails_closed_before_upstream_when_no_identity(): - """The pre-exchange identity gate covers the refresh_token grant, not just authorization_code: an - unresolvable litellm identity fails closed with invalid_request BEFORE the upstream refresh is - exchanged, so the client's refresh token is not rotated/consumed on a rejected request.""" +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 @@ -4526,10 +4532,6 @@ async def test_bridge_refresh_grant_fails_closed_before_upstream_when_no_identit "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", return_value=fake_http_client, ), - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_active_key_hash_from_request", - new=AsyncMock(return_value=None), - ), patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY), ): response = await exchange_token_with_server( @@ -4545,7 +4547,7 @@ async def test_bridge_refresh_grant_fails_closed_before_upstream_when_no_identit ) assert response.status_code == 400 - assert json.loads(response.body)["error"] == "invalid_request" + assert json.loads(response.body)["error"] == "unsupported_grant_type" fake_http_client.post.assert_not_called() @@ -4567,8 +4569,8 @@ 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._extract_active_key_hash_from_request", - new=AsyncMock(return_value="hashed-litellm-key-77"), + "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), ): @@ -4588,6 +4590,93 @@ async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset(): 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_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 @@ -4638,34 +4727,51 @@ def test_bridge_reported_expires_in_can_be_zero_at_jwt_exp_boundary(): assert json.loads(response.body)["expires_in"] == 0 -def test_bridge_grant_coerces_numeric_expires_in(): - """expires_in from an IdP may be an int, a float (3600.0), or a numeric string ("3600"); coerce - it to a positive int so the envelope TTL honors the real lifetime instead of dropping a non-int - value and defaulting to the 1h cap (which can outlive a shorter-lived upstream token). bool and - non-numeric values become None.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _bridge_grant_from_token_response, - ) +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 - def ei(v): - return _bridge_grant_from_token_response({"access_token": "x", "expires_in": v}).expires_in + 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" + # 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" - assert ei(300) == 300 - assert ei(300.0) == 300 - assert ei("300") == 300 - assert ei(" 300 ") == 300 - assert ei(True) is None - assert ei("nope") is None - assert ei(0) is None - assert ei(-5) is None - assert ei(None) is None - # hostile numerics must not raise (int(float(...)) can OverflowError) -> None - assert ei("inf") is None - assert ei("1e999") is None - assert ei("-inf") is None - assert ei("nan") is None - assert ei(float("inf")) is None - assert ei(10**400) is None + +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 @@ -5073,13 +5179,14 @@ async def test_extract_user_id_rejects_expired_key(proxy_globals): @pytest.mark.asyncio -async def test_extract_active_key_hash_returns_hash_for_active_key(proxy_globals): +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 ( - _extract_active_key_hash_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 @@ -5095,19 +5202,22 @@ async def test_extract_active_key_hash_returns_hash_for_active_key(proxy_globals proxy_globals.prisma_client = object() request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) - assert await _extract_active_key_hash_from_request(request) == hash_token(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_extract_active_key_hash_returns_hash_for_active_key_without_user_id(proxy_globals): +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_active_key_hash_from_request, _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 @@ -5123,16 +5233,18 @@ async def test_extract_active_key_hash_returns_hash_for_active_key_without_user_ proxy_globals.prisma_client = object() request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) - assert await _extract_active_key_hash_from_request(request) == hash_token(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_extract_active_key_hash_rejects_blocked_key(proxy_globals): +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 ( - _extract_active_key_hash_from_request, + _resolve_active_litellm_key, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -5145,17 +5257,17 @@ async def test_extract_active_key_hash_rejects_blocked_key(proxy_globals): proxy_globals.prisma_client = _FakePrisma() request = _token_request({"x-litellm-api-key": "sk-blocked-key"}) - assert await _extract_active_key_hash_from_request(request) is None + assert await _resolve_active_litellm_key(request) == "no_active_key" @pytest.mark.asyncio -async def test_extract_active_key_hash_fails_closed_on_malformed_expiry(proxy_globals): +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 ( - _extract_active_key_hash_from_request, + _resolve_active_litellm_key, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -5168,14 +5280,14 @@ async def test_extract_active_key_hash_fails_closed_on_malformed_expiry(proxy_gl proxy_globals.prisma_client = _FakePrisma() request = _token_request({"x-litellm-api-key": "sk-bad-expiry-key"}) - assert await _extract_active_key_hash_from_request(request) is None + assert await _resolve_active_litellm_key(request) == "no_active_key" @pytest.mark.asyncio -async def test_extract_active_key_hash_none_without_litellm_key(proxy_globals): +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 ( - _extract_active_key_hash_from_request, + _resolve_active_litellm_key, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -5183,7 +5295,46 @@ async def test_extract_active_key_hash_none_without_litellm_key(proxy_globals): proxy_globals.prisma_client = object() request = _token_request({"content-type": "application/json"}) - assert await _extract_active_key_hash_from_request(request) is None + 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 From 55ff3a242c269228dc1d54d0ffaffa40fb66ca5d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 18:39:00 -0700 Subject: [PATCH 12/12] fix(mcp): treat a positive sub-second upstream lifetime as alive, not expired _classify_upstream_lifetime decided "expired" from int(float(expires_in)), which truncates toward zero, so a positive fractional lifetime in (0, 1) became 0 and was misread as already elapsed. That rejected the mint with 502 in _finish_bridge_mint after the single-use upstream code had already been consumed, even though the upstream reported a positive remaining lifetime. Decide expired on the parsed numeric value rather than its truncated int, so only a genuinely non-positive value is expired. The envelope works in whole seconds and cannot represent a sub-second lifetime, so a positive value that truncates to 0 clamps up to the 1s floor instead of being rejected. Values >= 1 still truncate toward zero so the envelope never claims more life than the upstream stated, and NaN / Infinity / oversized input still read as unparseable ("unspecified"). Regression covers the classifier (0.5 and 0.001 clamp to 1, 1.9 truncates to 1, -0.5 stays expired) and the mint (a 0.5s upstream lifetime mints a 200 envelope rather than a 502); reverting to the truncate-then-check reddens both. --- .../mcp_server/discoverable_endpoints.py | 21 ++++++++++------- .../mcp_server/test_discoverable_endpoints.py | 23 +++++++++++++++++++ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index b8d10335182..8d1713a5911 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -742,19 +742,24 @@ 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 parseable 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. ``bool`` is excluded (an ``int`` subclass but never a real lifetime), and - ``int(float(...))`` can raise on ``NaN`` / ``Infinity`` / oversized input, which reads as - unparseable rather than surfacing as a 500.""" + 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: - seconds = int(float(raw_expires_in)) + numeric = float(raw_expires_in) + seconds = int(numeric) except (ValueError, TypeError, OverflowError): return "unspecified" - return seconds if seconds > 0 else "expired" + if numeric <= 0: + return "expired" + return max(1, seconds) def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenGrant | _UpstreamGrantRejection": 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 f611d2f6a24..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 @@ -4661,6 +4661,22 @@ async def test_bridge_mint_upstream_expired_lifetime_is_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 @@ -4742,6 +4758,13 @@ def test_classify_upstream_lifetime(): # 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"