scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

refactor(mcp): make bridge-mint resolvers return tagged unions so status is truthful by construction

Browse source 
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.
This commit is contained in:
Tin Chi Lo 2026-07-11 18:28:32 -07:00
parent 4ba7221b7a
commit a07aba0579
2 changed files with 424 additions and 169 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

334
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -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

259
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -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