mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
Merge pull request #32980 from BerriAI/litellm_bridge_refresh_envelope
feat(mcp): client-held refresh envelope for the dcr_bridge oauth_delegate flow
This commit is contained in:
tin-berri
GitHub
commit
53aaabba5e
6 changed files with 1405 additions and 90 deletions
|
|
@ -43,6 +43,7 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
|
||||
EnvelopeIdentity,
|
||||
EnvelopeKeys,
|
||||
RefreshCredential,
|
||||
UpstreamTokenGrant,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth
|
||||
|
|
@ -486,9 +487,23 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR
|
|||
token = _litellm_key_from_request(request)
|
||||
if not token:
|
||||
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._types import hash_token # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
|
||||
return await _reload_active_key_by_hash(hash_token(token))
|
||||
|
||||
|
||||
async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResolutionFailure":
|
||||
"""Reload the live key record for ``key_hash`` (cache first, then DB) and gate it on active state,
|
||||
returning the resolved key or a precise failure. Shared by the token request's presented-key
|
||||
resolution (:func:`_resolve_active_litellm_key`, which hashes the presented key) and the refresh
|
||||
path (which already holds the hash sealed in the refresh envelope), so both re-validate identity
|
||||
through one active-key gate and one failure classification. Classification mirrors admission's
|
||||
``_reload_admitted_key``: no DB connection is a gateway fault, a ``ProxyException`` / ``HTTPException``
|
||||
from ``get_key_object`` is an unknown or invalid key, a database-service-unavailable error is a
|
||||
retryable outage, and anything else is an unexpected gateway fault. A blocked or expired key is
|
||||
``no_active_key``, so a revoked key can neither mint nor refresh a bridge envelope."""
|
||||
from litellm.proxy._types import (
|
||||
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
get_key_object,
|
||||
|
|
@ -503,7 +518,6 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR
|
|||
|
||||
if prisma_client is None:
|
||||
return "unresolvable"
|
||||
key_hash = hash_token(token)
|
||||
try:
|
||||
key_obj = await get_key_object(
|
||||
hashed_token=key_hash,
|
||||
|
|
@ -516,7 +530,7 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR
|
|||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc):
|
||||
return "unavailable"
|
||||
verbose_logger.debug(
|
||||
"_resolve_active_litellm_key: unexpected key-resolution error (%s)",
|
||||
"_reload_active_key_by_hash: unexpected key-resolution error (%s)",
|
||||
type(exc).__name__,
|
||||
)
|
||||
return "unresolvable"
|
||||
|
|
@ -525,6 +539,107 @@ async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyR
|
|||
return _ResolvedKey(key_hash=key_hash, key=key_obj)
|
||||
|
||||
|
||||
async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | None":
|
||||
"""Re-validate a live litellm user by id, returning ``None`` when the user is active or a precise
|
||||
failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a
|
||||
user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a
|
||||
deactivated user cannot keep refreshing, mirroring how admission re-validates the same user subject on
|
||||
the egress side. No DB connection is a gateway fault (``unresolvable``) and a
|
||||
database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails
|
||||
closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` /
|
||||
``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object``
|
||||
catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look
|
||||
identical, the original error surviving only as ``__context__``), so the outage check walks the cause
|
||||
chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault."""
|
||||
from litellm.proxy._types import (
|
||||
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
get_user_object,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
PrismaDBExceptionHandler,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
return "unresolvable"
|
||||
try:
|
||||
user_object = await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
)
|
||||
except (ProxyException, HTTPException):
|
||||
return "no_active_key"
|
||||
except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc):
|
||||
return "unavailable"
|
||||
verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__)
|
||||
return "no_active_key"
|
||||
if user_object is None:
|
||||
return "no_active_key"
|
||||
if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False:
|
||||
return "no_active_key"
|
||||
return None
|
||||
|
||||
|
||||
async def _key_owner_scim_deactivated(key: "UserAPIKeyAuth") -> bool:
|
||||
"""True only when the key's owning user was explicitly SCIM-deactivated, so a refresh revokes an
|
||||
offboarded owner's key exactly as admission does via ``_reject_if_admitted_owner_scim_deactivated``.
|
||||
A key with no owner, a missing owner record, or a failed lookup fails OPEN (returns ``False``),
|
||||
matching admission and the standard builder: a key may outlive its owner record, and a transient DB
|
||||
blip must not revoke a live key. Only an explicit ``scim_active`` of ``False`` gates renewal."""
|
||||
if key.user_id is None:
|
||||
return False
|
||||
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
get_user_object,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
return False
|
||||
try:
|
||||
owner = await get_user_object(
|
||||
user_id=key.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # fail open: a missing owner (get_user_object's wrapped ValueError) or a DB blip must not revoke a live key
|
||||
verbose_logger.debug("refresh: key-owner SCIM lookup failed, not revoking (%s)", type(exc).__name__)
|
||||
return False
|
||||
return owner is not None and isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False
|
||||
|
||||
|
||||
async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResolutionFailure | None":
|
||||
"""Re-validate that the subject sealed in a refresh envelope is still live, dispatching on its type:
|
||||
a key_hash reloads the virtual key, a user_id reloads the user. Returns ``None`` when the subject is
|
||||
active or a precise failure otherwise, so revocation gates renewal for either identity source the same
|
||||
way admission gates the egress: a blocked or expired key, a SCIM-deactivated key owner (mirroring
|
||||
admission's owner check, so an offboarded user cannot keep renewing a still-active key), and a
|
||||
deactivated or deleted user all fail closed to ``no_active_key``."""
|
||||
match identity.subject_type:
|
||||
case "key_hash":
|
||||
reloaded = await _reload_active_key_by_hash(identity.subject)
|
||||
if not isinstance(reloaded, _ResolvedKey):
|
||||
return reloaded
|
||||
if await _key_owner_scim_deactivated(reloaded.key):
|
||||
return "no_active_key"
|
||||
return None
|
||||
case "user_id":
|
||||
return await _reload_active_user_by_id(identity.subject)
|
||||
case _:
|
||||
assert_never(identity.subject_type)
|
||||
|
||||
|
||||
async def _extract_user_id_from_request(request: Request) -> str | None:
|
||||
"""The litellm ``user_id`` for the token request, so a per-user token is stored under the same
|
||||
identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome
|
||||
|
|
@ -903,7 +1018,7 @@ def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenG
|
|||
|
||||
_BridgeMintError = Literal[
|
||||
"no_identity",
|
||||
"unsupported_grant",
|
||||
"invalid_refresh",
|
||||
"identity_unavailable",
|
||||
"identity_unresolvable",
|
||||
"not_configured",
|
||||
|
|
@ -941,12 +1056,12 @@ def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse:
|
|||
"this server issues a gateway-bound credential; complete the interactive sign-in, or "
|
||||
"send a litellm credential (x-litellm-api-key or Authorization) on the token request",
|
||||
)
|
||||
case "unsupported_grant":
|
||||
case "invalid_refresh":
|
||||
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",
|
||||
"invalid_grant",
|
||||
"the refresh credential is not a valid, live refresh envelope for this server; "
|
||||
"re-run authorization_code to obtain a new one",
|
||||
)
|
||||
case "identity_unavailable":
|
||||
status, code, desc = (
|
||||
|
|
@ -1019,21 +1134,21 @@ def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _Br
|
|||
|
||||
async def _prepare_bridge_mint(
|
||||
request: Request,
|
||||
grant_type: str,
|
||||
mcp_server: MCPServer,
|
||||
bridge_identity: _BridgeAuthorizationCode | None = None,
|
||||
) -> "_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.
|
||||
"""Phase 1 for the authorization_code grant, BEFORE the upstream exchange: confirm the gateway can
|
||||
mint (master_key set), resolve the litellm identity, and derive the envelope keys. Returns a ready
|
||||
context or a precise failure value. Running before the exchange is what makes every failure here fail
|
||||
closed without consuming the single-use code.
|
||||
|
||||
Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged
|
||||
authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway
|
||||
authorization code) and mints a user subject. The scripted two-header client presents a litellm key
|
||||
on the token request instead, so its identity is the active key's hash and mints a key_hash subject.
|
||||
A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully;
|
||||
neither source present is ``no_identity``."""
|
||||
neither source present is ``no_identity``. The refresh_token grant has its own phase-1
|
||||
(:func:`_prepare_bridge_refresh`), which recovers identity from the presented refresh envelope."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
envelope_keys_from_master_key,
|
||||
)
|
||||
|
|
@ -1045,8 +1160,6 @@ async def _prepare_bridge_mint(
|
|||
master_key,
|
||||
)
|
||||
|
||||
if grant_type != "authorization_code":
|
||||
return "unsupported_grant"
|
||||
if not master_key:
|
||||
return "not_configured"
|
||||
keys = envelope_keys_from_master_key(master_key)
|
||||
|
|
@ -1060,14 +1173,85 @@ async def _prepare_bridge_mint(
|
|||
return _BridgeMintReady(identity=identity, keys=keys)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeRefreshReady:
|
||||
"""A validated refresh request: the identity+keys to mint the renewed pair under, the upstream refresh
|
||||
token (unwrapped from the client's refresh envelope) to exchange with the upstream IdP, and the scope
|
||||
sealed alongside it at mint. The upstream refresh token is a ``SecretStr`` like every other credential
|
||||
in this layer, so a repr or a traceback that captures this value never exposes the raw upstream refresh
|
||||
token in plaintext. ``upstream_scope`` carries the originally-granted scope so the renewal re-requests
|
||||
it when the client (a DCR/MCP client that typically omits scope on refresh) sends none, keeping the
|
||||
renewed token's scope stable against an upstream that would otherwise narrow or drop it."""
|
||||
|
||||
ready: "_BridgeMintReady"
|
||||
upstream_refresh_token: SecretStr
|
||||
upstream_scope: str | None = None
|
||||
|
||||
|
||||
def _refresh_key_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError:
|
||||
"""Lift an identity-resolution failure on the refresh path into the mint taxonomy. Unlike the mint
|
||||
path, a resolved-but-inactive (or unknown) key is ``invalid_grant`` rather than ``invalid_request``:
|
||||
the client did present an identity (sealed in the refresh envelope), but it is no longer live, so the
|
||||
refresh is invalid and the client must re-authenticate. A transient outage is still 503 and a gateway
|
||||
fault still 500, matching the mint path and admission."""
|
||||
match failure:
|
||||
case "no_active_key":
|
||||
return "invalid_refresh"
|
||||
case "unavailable":
|
||||
return "identity_unavailable"
|
||||
case "unresolvable":
|
||||
return "identity_unresolvable"
|
||||
case _:
|
||||
assert_never(failure)
|
||||
|
||||
|
||||
async def _prepare_bridge_refresh(
|
||||
mcp_server: MCPServer, refresh_value: str | None
|
||||
) -> "_BridgeRefreshReady | _BridgeMintError":
|
||||
"""Phase 1 for the refresh_token grant, BEFORE the upstream exchange: open the client's refresh
|
||||
envelope, re-validate the sealed litellm identity so a revoked key cannot keep refreshing, and
|
||||
recover the upstream refresh token to exchange. Identity comes entirely from the sealed envelope, not
|
||||
the HTTP request, so the request object is not needed here. The client presents a refresh envelope,
|
||||
never a raw upstream refresh token, so a missing value, a non-envelope, an unopenable envelope, or one
|
||||
minted for another server is ``invalid_grant``. Running before the exchange means a rejected refresh
|
||||
never consumes or rotates the upstream refresh token."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
BridgeRefreshOpened,
|
||||
envelope_keys_from_master_key,
|
||||
open_bridge_refresh_envelope,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
master_key,
|
||||
)
|
||||
|
||||
if not master_key:
|
||||
return "not_configured"
|
||||
if not refresh_value:
|
||||
return "invalid_refresh"
|
||||
keys = envelope_keys_from_master_key(master_key)
|
||||
opened = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id)
|
||||
if not isinstance(opened, BridgeRefreshOpened):
|
||||
return "invalid_refresh"
|
||||
failure = await _revalidate_active_subject(opened.identity)
|
||||
if failure is not None:
|
||||
return _refresh_key_failure_to_mint_error(failure)
|
||||
return _BridgeRefreshReady(
|
||||
ready=_BridgeMintReady(identity=opened.identity, keys=keys),
|
||||
upstream_refresh_token=opened.refresh.refresh_token,
|
||||
upstream_scope=opened.refresh.scope,
|
||||
)
|
||||
|
||||
|
||||
def _finish_bridge_mint(
|
||||
ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime
|
||||
) -> "JSONResponse | _BridgeMintError":
|
||||
"""Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held envelope under
|
||||
the pre-resolved identity and keys, so the client holds one bearer that admits it and forwards the
|
||||
upstream token with nothing stored server-side. The only failures here are properties of the
|
||||
upstream response (no usable token, an already-expired lifetime, or a token too large to seal),
|
||||
returned as values."""
|
||||
"""Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held access envelope
|
||||
using the pre-resolved identity and keys, and, when the upstream returned a refresh token, seal a
|
||||
long-lived refresh envelope alongside it so the client can renew without re-authenticating. Shared by
|
||||
the authorization_code and refresh_token paths, so a renewal that the upstream rotates re-issues a
|
||||
fresh refresh envelope. The only hard failures here are properties of the upstream access token (no
|
||||
usable token, an already-expired lifetime, or a token too large to seal); a refresh token that cannot
|
||||
be sealed degrades to an access-only response rather than failing the whole exchange."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
build_bridge_token_response,
|
||||
)
|
||||
|
|
@ -1085,10 +1269,90 @@ def _finish_bridge_mint(
|
|||
# Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the
|
||||
# client is never told the bearer lives past the point admission (which uses that exp) rejects it.
|
||||
expires_in = max(0, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp()))
|
||||
body = {"access_token": sealed.token.get_secret_value(), "token_type": "Bearer", "expires_in": expires_in}
|
||||
refresh_envelope = _mint_refresh_envelope_value(ready.identity, token_response, ready.keys, now, mcp_server)
|
||||
body = {
|
||||
"access_token": sealed.token.get_secret_value(),
|
||||
"token_type": "Bearer",
|
||||
"expires_in": expires_in,
|
||||
# A refresh envelope rides along only when the upstream returned a refresh token to seal; when it
|
||||
# rotates on renewal, the client receives the new one and the old envelope's upstream token dies.
|
||||
**({"refresh_token": refresh_envelope} if refresh_envelope is not None else {}),
|
||||
}
|
||||
return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS)
|
||||
|
||||
|
||||
def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | None":
|
||||
"""Extract the upstream refresh grant from a token response, or ``None`` when there is none to seal.
|
||||
Each field is isinstance-checked so nothing untyped reaches the refresh envelope; ``refresh_expires_in``
|
||||
(the refresh token's own lifetime, when the upstream reports it) is classified like ``expires_in`` and
|
||||
bounds the refresh envelope's TTL. An upstream that reports the refresh token itself as already elapsed
|
||||
(``refresh_expires_in`` non-positive) yields ``None`` rather than a refresh envelope: sealing a dead
|
||||
token would hand the client a full-TTL-capped envelope the IdP will reject, so the exchange degrades to
|
||||
an access-only response (the client re-authenticates at access expiry), mirroring how
|
||||
:func:`_bridge_grant_from_token_response` refuses an already-elapsed access token instead of capping it."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
RefreshCredential,
|
||||
)
|
||||
|
||||
if not isinstance(token_response, dict):
|
||||
return None
|
||||
refresh = token_response.get("refresh_token")
|
||||
if not isinstance(refresh, str) or not refresh:
|
||||
return None
|
||||
lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in"))
|
||||
if lifetime == "expired":
|
||||
return None
|
||||
scope = token_response.get("scope")
|
||||
return RefreshCredential(
|
||||
refresh_token=SecretStr(refresh),
|
||||
scope=scope if isinstance(scope, str) and scope else None,
|
||||
expires_in=lifetime if isinstance(lifetime, int) else None,
|
||||
)
|
||||
|
||||
|
||||
def _mint_refresh_envelope_value(
|
||||
identity: "EnvelopeIdentity", token_response: object, keys: "EnvelopeKeys", now: datetime, mcp_server: MCPServer
|
||||
) -> str | None:
|
||||
"""Seal the upstream refresh grant (if any) into a refresh envelope and return its bearer string, or
|
||||
``None`` when the upstream returned no refresh token or the refresh token is too large to seal. A
|
||||
too-large refresh token degrades to an access-only response (logged) rather than failing an exchange
|
||||
that already succeeded upstream: the client simply re-authenticates when the access envelope expires."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
build_bridge_refresh_token_response,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
SealedEnvelope,
|
||||
)
|
||||
|
||||
refresh_credential = _upstream_refresh_credential(token_response)
|
||||
if refresh_credential is None:
|
||||
return None
|
||||
sealed = build_bridge_refresh_token_response(identity, refresh_credential, keys, now)
|
||||
if isinstance(sealed, SealedEnvelope):
|
||||
return sealed.token.get_secret_value()
|
||||
verbose_logger.warning(
|
||||
"bridge mint: the upstream refresh token is too large to seal into a refresh envelope for "
|
||||
"server=%s; issuing an access-only response, so the client re-authenticates at access expiry",
|
||||
mcp_server.server_id,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _upstream_oauth_error(response: httpx.Response) -> str | None:
|
||||
"""The RFC 6749 5.2 ``error`` code from an upstream token-endpoint error body, or ``None`` when the
|
||||
body is not a JSON object carrying a string ``error``. Reading the field beats substring-matching the
|
||||
raw text, which would false-match a code that only appears inside ``error_description`` (a false
|
||||
invalid_grant would trigger a needless authorization_code re-run)."""
|
||||
try:
|
||||
body = json.loads(response.text)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
error = body.get("error")
|
||||
return error if isinstance(error, str) else None
|
||||
|
||||
|
||||
async def exchange_token_with_server(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -1124,19 +1388,41 @@ async def exchange_token_with_server(
|
|||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
||||
bridge_identity: _BridgeAuthorizationCode | None = None
|
||||
bridge_mint_ready: _BridgeMintReady | None = None
|
||||
bridge_upstream_refresh: SecretStr | None = None
|
||||
bridge_upstream_scope: str | None = None
|
||||
refresh_request_scope: str | None = None
|
||||
is_bridge = mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge
|
||||
|
||||
if grant_type == "refresh_token":
|
||||
if not refresh_token:
|
||||
# Phase 1 for a bridge refresh: open the client's refresh envelope, re-validate the sealed
|
||||
# identity, and unwrap the real upstream refresh token BEFORE building token_data, so the exchange
|
||||
# sends the upstream token and never the envelope. A failure returns without touching the upstream.
|
||||
if is_bridge:
|
||||
prepared_refresh = await _prepare_bridge_refresh(mcp_server, refresh_token)
|
||||
if not isinstance(prepared_refresh, _BridgeRefreshReady):
|
||||
return _bridge_mint_error_response(prepared_refresh)
|
||||
bridge_mint_ready = prepared_refresh.ready
|
||||
bridge_upstream_refresh = prepared_refresh.upstream_refresh_token
|
||||
bridge_upstream_scope = prepared_refresh.upstream_scope
|
||||
# A bridge server sends the unwrapped upstream refresh token recovered from the client's refresh
|
||||
# envelope above; every other server sends the client's own refresh token verbatim.
|
||||
upstream_refresh_token = (
|
||||
bridge_upstream_refresh.get_secret_value() if bridge_upstream_refresh is not None else refresh_token
|
||||
)
|
||||
if not upstream_refresh_token:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="refresh_token is required for refresh_token grant",
|
||||
)
|
||||
token_data: dict = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
"refresh_token": upstream_refresh_token,
|
||||
**client_auth.body,
|
||||
}
|
||||
if scope:
|
||||
token_data["scope"] = scope
|
||||
refresh_request_scope = scope or bridge_upstream_scope
|
||||
if refresh_request_scope:
|
||||
token_data["scope"] = refresh_request_scope
|
||||
else:
|
||||
if not code:
|
||||
raise HTTPException(
|
||||
|
|
@ -1175,17 +1461,13 @@ async def exchange_token_with_server(
|
|||
}
|
||||
if code_verifier:
|
||||
token_data["code_verifier"] = code_verifier
|
||||
|
||||
# Phase 1: for a bridge oauth_delegate mint, validate all preconditions and resolve identity+keys
|
||||
# BEFORE the exchange below consumes the single-use upstream code, and carry the ready context to
|
||||
# phase 3. A failure here returns without ever touching the upstream credential.
|
||||
bridge_mint_ready: _BridgeMintReady | None = None
|
||||
if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge:
|
||||
prepared = await _prepare_bridge_mint(request, grant_type, mcp_server, bridge_identity)
|
||||
if not isinstance(prepared, _BridgeMintReady):
|
||||
return _bridge_mint_error_response(prepared)
|
||||
bridge_mint_ready = prepared
|
||||
|
||||
# Phase 1 for a bridge authorization_code mint: resolve identity (the SSO user recovered above, or
|
||||
# the presented litellm key) and the envelope keys BEFORE the exchange consumes the single-use code.
|
||||
if is_bridge:
|
||||
prepared = await _prepare_bridge_mint(request, mcp_server, bridge_identity)
|
||||
if not isinstance(prepared, _BridgeMintReady):
|
||||
return _bridge_mint_error_response(prepared)
|
||||
bridge_mint_ready = prepared
|
||||
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
|
||||
response = await async_client.post(
|
||||
mcp_server.token_url,
|
||||
|
|
@ -1208,6 +1490,20 @@ async def exchange_token_with_server(
|
|||
"does not send yet (tracked as LIT-4339)",
|
||||
mcp_server.server_id,
|
||||
)
|
||||
upstream_rejected_bridge_refresh = (
|
||||
is_bridge
|
||||
and grant_type == "refresh_token"
|
||||
and exc.response.status_code == 400
|
||||
and _upstream_oauth_error(exc.response) == "invalid_grant"
|
||||
)
|
||||
if upstream_rejected_bridge_refresh:
|
||||
verbose_logger.info(
|
||||
"bridge refresh: the upstream rejected the sealed refresh token for server=%s with "
|
||||
"invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client "
|
||||
"re-runs authorization_code rather than an opaque upstream error",
|
||||
mcp_server.server_id,
|
||||
)
|
||||
return _bridge_mint_error_response("invalid_refresh")
|
||||
raise
|
||||
token_response = response.json()
|
||||
|
||||
|
|
@ -1253,6 +1549,8 @@ async def exchange_token_with_server(
|
|||
# upstream token) instead of the raw upstream token, so the one bearer both admits the caller and
|
||||
# forwards the upstream credential. Only this mode mints; every other server returns the raw token.
|
||||
if bridge_mint_ready is not None:
|
||||
if refresh_request_scope and isinstance(token_response, dict) and not token_response.get("scope"):
|
||||
token_response = {**token_response, "scope": refresh_request_scope}
|
||||
# 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))
|
||||
|
|
|
|||
|
|
@ -21,11 +21,16 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
|
|||
EnvelopeKeys,
|
||||
EnvelopeMintError,
|
||||
OpenedEnvelope,
|
||||
OpenedRefreshEnvelope,
|
||||
RefreshCredential,
|
||||
SealedEnvelope,
|
||||
UpstreamTokenGrant,
|
||||
is_envelope,
|
||||
is_refresh_envelope,
|
||||
mint_envelope,
|
||||
mint_refresh_envelope,
|
||||
open_envelope,
|
||||
open_refresh_envelope,
|
||||
)
|
||||
|
||||
_SIGNING_KEY_DOMAIN = b"litellm-mcp-bridge:envelope-signing:"
|
||||
|
|
@ -92,6 +97,67 @@ def build_bridge_token_response(
|
|||
return mint_envelope(identity, grant, keys, now)
|
||||
|
||||
|
||||
def build_bridge_refresh_token_response(
|
||||
identity: EnvelopeIdentity,
|
||||
refresh: RefreshCredential,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
) -> SealedEnvelope | EnvelopeMintError:
|
||||
"""Seal ``refresh`` for ``identity`` into the long-lived refresh envelope the token endpoint returns
|
||||
alongside the access envelope, so the client can renew without re-authenticating. A thin, pure
|
||||
wrapper over :func:`mint_refresh_envelope`; returns the mint error as a value for the caller to map.
|
||||
"""
|
||||
return mint_refresh_envelope(identity, refresh, keys, now)
|
||||
|
||||
|
||||
class BridgeRefreshOpened(BaseModel):
|
||||
"""A valid refresh envelope presented to the token endpoint: the identity to re-validate and renew
|
||||
under, and the upstream refresh grant to exchange."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["opened"] = "opened"
|
||||
identity: EnvelopeIdentity
|
||||
refresh: RefreshCredential
|
||||
|
||||
|
||||
class BridgeRefreshInvalid(BaseModel):
|
||||
"""The presented refresh grant is not a valid refresh envelope for this server (not refresh-shaped,
|
||||
will not open, or minted for a different server); the token endpoint fails the refresh closed."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["invalid"] = "invalid"
|
||||
|
||||
|
||||
BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid
|
||||
|
||||
|
||||
def open_bridge_refresh_envelope(
|
||||
refresh_value: str,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
expected_server_id: str,
|
||||
) -> BridgeRefreshResult:
|
||||
"""Open a refresh envelope a bridge ``oauth_delegate`` client presented on a refresh_token grant.
|
||||
|
||||
The token-endpoint mirror of :func:`resolve_bridge_envelope`: strips an optional ``Bearer`` scheme,
|
||||
then returns ``BridgeRefreshOpened`` with the recovered identity and upstream refresh grant, or
|
||||
``BridgeRefreshInvalid`` for anything that is not a valid refresh envelope for this server. Never
|
||||
raises; total over hostile input via :func:`open_refresh_envelope`. ``expected_server_id`` binds the
|
||||
envelope to the server the request targets, so a refresh envelope minted for one server cannot renew
|
||||
against another. A raw upstream refresh token (not envelope-shaped) is ``BridgeRefreshInvalid``: this
|
||||
mode never hands the client a bare upstream refresh token, so it must never accept one.
|
||||
"""
|
||||
candidate = _strip_bearer(refresh_value)
|
||||
if not is_refresh_envelope(candidate):
|
||||
return BridgeRefreshInvalid()
|
||||
opened = open_refresh_envelope(candidate, keys, now)
|
||||
if not isinstance(opened, OpenedRefreshEnvelope):
|
||||
return BridgeRefreshInvalid()
|
||||
if opened.identity.server_id != expected_server_id:
|
||||
return BridgeRefreshInvalid()
|
||||
return BridgeRefreshOpened(identity=opened.identity, refresh=opened.refresh)
|
||||
|
||||
|
||||
class NotBridgeEnvelope(BaseModel):
|
||||
"""The bearer is not an envelope; admission continues on its normal path."""
|
||||
|
||||
|
|
@ -128,10 +194,12 @@ def _strip_bearer(value: str) -> str:
|
|||
|
||||
|
||||
def is_bridge_envelope_shaped(authorization_value: str) -> bool:
|
||||
"""Cheap, keyless test that an ``Authorization`` value carries an envelope (optional
|
||||
``Bearer`` scheme stripped). The admission edge engages the bridge arm only for an
|
||||
envelope, so a plain upstream bearer falls through to normal oauth2 admission."""
|
||||
return is_envelope(_strip_bearer(authorization_value))
|
||||
"""Cheap, keyless test that an ``Authorization`` value carries an envelope of either kind (optional
|
||||
``Bearer`` scheme stripped). The admission edge engages the bridge arm for an access envelope (to
|
||||
admit) and for a refresh envelope (to reject it explicitly, since a refresh credential is never
|
||||
usable at the tool-call edge); a plain upstream bearer falls through to normal oauth2 admission."""
|
||||
candidate = _strip_bearer(authorization_value)
|
||||
return is_envelope(candidate) or is_refresh_envelope(candidate)
|
||||
|
||||
|
||||
def resolve_bridge_envelope(
|
||||
|
|
@ -148,6 +216,10 @@ def resolve_bridge_envelope(
|
|||
envelope, and ``BridgeEnvelopeInvalid`` for an envelope-shaped bearer that will not
|
||||
open. Never raises: it is total over hostile input via :func:`open_envelope`.
|
||||
|
||||
A refresh envelope is ``BridgeEnvelopeInvalid`` here: it is a valid gateway credential but only ever
|
||||
presented back to the token endpoint, never usable to authenticate a tool call, so admission must
|
||||
fail it closed rather than let it fall through to another arm.
|
||||
|
||||
``expected_server_id`` is the ``server_id`` of the MCP server the request targets; an
|
||||
opened envelope whose sealed ``server_id`` does not match is rejected as
|
||||
``BridgeEnvelopeInvalid``. Binding here (rather than leaving it to the caller) prevents
|
||||
|
|
@ -157,6 +229,8 @@ def resolve_bridge_envelope(
|
|||
unlike ``hmac.compare_digest`` on ``str``, does not raise on a non-ASCII server_id.
|
||||
"""
|
||||
candidate = _strip_bearer(authorization_value)
|
||||
if is_refresh_envelope(candidate):
|
||||
return BridgeEnvelopeInvalid()
|
||||
if not is_envelope(candidate):
|
||||
return NotBridgeEnvelope()
|
||||
opened = open_envelope(candidate, keys, now)
|
||||
|
|
|
|||
|
|
@ -44,18 +44,33 @@ from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
|
|||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value
|
||||
|
||||
ENVELOPE_PREFIX = "llm_env_"
|
||||
"""Marker prefix on every serialized envelope so the edge can cheaply tell an envelope
|
||||
"""Marker prefix on every serialized ACCESS envelope so the edge can cheaply tell an envelope
|
||||
from a raw upstream token before doing any cryptography."""
|
||||
|
||||
REFRESH_ENVELOPE_PREFIX = "llm_refresh_"
|
||||
"""Marker prefix on every serialized REFRESH envelope. A distinct prefix keeps the two credentials
|
||||
routable without crypto and, together with the signed ``kind`` claim, stops one from being presented
|
||||
where the other is expected: a refresh envelope carries a long-lived upstream refresh token and is only
|
||||
ever presented back to the token endpoint, never forwarded upstream on a tool call."""
|
||||
|
||||
ENVELOPE_ISSUER = "litellm-mcp-bridge"
|
||||
"""``iss`` claim stamped into every envelope and required back on open."""
|
||||
|
||||
MAX_ENVELOPE_TTL_SECONDS = 3600
|
||||
"""Hard ceiling on envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)``
|
||||
"""Hard ceiling on ACCESS envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)``
|
||||
(the cap alone when the upstream omits ``expires_in``), matching the 1h lifetime of the
|
||||
BYOK session bearer this module's signing approach is borrowed from: a client-held
|
||||
credential should never outlive a bounded window even when the upstream token does."""
|
||||
|
||||
MAX_REFRESH_ENVELOPE_TTL_SECONDS = 1209600
|
||||
"""Hard ceiling on REFRESH envelope lifetime (14 days). A refresh envelope only renews the short-lived
|
||||
access envelope, and each renewal re-validates the sealed litellm key (revocation gates it) and is
|
||||
re-minted with a fresh window, so the practical bound is idle time, not a fixed session. ``exp`` is
|
||||
``min(upstream refresh_expires_in, this cap)`` (the cap alone when the upstream omits it); if the
|
||||
upstream refresh token dies first, the next renewal simply fails at the upstream and the client
|
||||
re-authenticates. The value is deliberately far shorter than a typical upstream refresh-token lifetime
|
||||
so a leaked refresh envelope is bounded even if the upstream would have honoured it for longer."""
|
||||
|
||||
MAX_ENVELOPE_BYTES = 12288
|
||||
"""Size cap on the final serialized envelope (prefix + JWT, in bytes). Upstream JWTs
|
||||
commonly run 2-4KB; base64 plus encryption overhead roughly doubles that inside the
|
||||
|
|
@ -66,6 +81,11 @@ typed error, never truncated."""
|
|||
|
||||
_ENVELOPE_JWT_ALGORITHM = "HS256"
|
||||
|
||||
EnvelopeKind = Literal["access", "refresh"]
|
||||
"""Which credential an envelope is. Stamped into the signed claims and required to match on open, so a
|
||||
signature-valid envelope of one kind cannot be replayed as the other even if its wire prefix is swapped
|
||||
(the prefix is not part of the signed payload; this claim is)."""
|
||||
|
||||
|
||||
EnvelopeSubjectType: TypeAlias = Literal["key_hash", "user_id"]
|
||||
"""Discriminator for what litellm principal the envelope binds the grant to.
|
||||
|
|
@ -121,6 +141,21 @@ class UpstreamTokenGrant(BaseModel):
|
|||
expires_in: int | None = Field(default=None, gt=0)
|
||||
|
||||
|
||||
class RefreshCredential(BaseModel):
|
||||
"""The upstream refresh grant sealed inside a refresh envelope.
|
||||
|
||||
Only the refresh token (plus the scope to re-request and the refresh token's own lifetime, when the
|
||||
upstream reports it) is sealed; the access token is never in a refresh envelope. ``refresh_token`` is
|
||||
a ``SecretStr`` so reprs never leak it, and ``expires_in`` (the refresh token's lifetime, not the
|
||||
access token's) must be positive when present.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
refresh_token: SecretStr = Field(min_length=1)
|
||||
scope: str | None = None
|
||||
expires_in: int | None = Field(default=None, gt=0)
|
||||
|
||||
|
||||
class EnvelopeKeys(BaseModel):
|
||||
"""Injected key material: the HS256 signing key and the symmetric encryption key.
|
||||
|
||||
|
|
@ -143,13 +178,21 @@ class SealedEnvelope(BaseModel):
|
|||
|
||||
|
||||
class OpenedEnvelope(BaseModel):
|
||||
"""A validated envelope: the identity it was minted for and the recovered grant."""
|
||||
"""A validated access envelope: the identity it was minted for and the recovered grant."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
identity: EnvelopeIdentity
|
||||
grant: UpstreamTokenGrant
|
||||
|
||||
|
||||
class OpenedRefreshEnvelope(BaseModel):
|
||||
"""A validated refresh envelope: the identity it was minted for and the recovered refresh grant."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
identity: EnvelopeIdentity
|
||||
refresh: RefreshCredential
|
||||
|
||||
|
||||
class EnvelopeTooLarge(BaseModel):
|
||||
"""The serialized envelope exceeded ``MAX_ENVELOPE_BYTES``; carries sizes only."""
|
||||
|
||||
|
|
@ -221,6 +264,7 @@ class _EnvelopeClaims(BaseModel):
|
|||
iss: str
|
||||
iat: int
|
||||
exp: int
|
||||
kind: EnvelopeKind
|
||||
server_id: str = Field(min_length=1)
|
||||
subject_type: EnvelopeSubjectType
|
||||
subject: str = Field(min_length=1)
|
||||
|
|
@ -236,11 +280,25 @@ class _GrantWire(BaseModel):
|
|||
expires_in: int | None = None
|
||||
|
||||
|
||||
class _RefreshWire(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
refresh_token: str
|
||||
scope: str | None = None
|
||||
expires_in: int | None = None
|
||||
|
||||
|
||||
def is_envelope(candidate: str) -> bool:
|
||||
"""Cheap prefix check so the edge can route envelopes vs raw tokens without crypto."""
|
||||
"""Cheap prefix check for an ACCESS envelope so the edge can route envelopes vs raw tokens without
|
||||
crypto. A refresh envelope has a different prefix and is not an access envelope."""
|
||||
return candidate.startswith(ENVELOPE_PREFIX)
|
||||
|
||||
|
||||
def is_refresh_envelope(candidate: str) -> bool:
|
||||
"""Cheap prefix check for a REFRESH envelope so the token endpoint can route a refresh grant that
|
||||
carries an envelope vs a raw upstream refresh token without crypto."""
|
||||
return candidate.startswith(REFRESH_ENVELOPE_PREFIX)
|
||||
|
||||
|
||||
def mint_envelope(
|
||||
identity: EnvelopeIdentity,
|
||||
grant: UpstreamTokenGrant,
|
||||
|
|
@ -254,24 +312,15 @@ def mint_envelope(
|
|||
serialized envelope exceeds ``MAX_ENVELOPE_BYTES``.
|
||||
"""
|
||||
expires_at = now + timedelta(seconds=_envelope_ttl_seconds(grant.expires_in))
|
||||
claims = _EnvelopeClaims(
|
||||
iss=ENVELOPE_ISSUER,
|
||||
iat=int(now.timestamp()),
|
||||
exp=int(expires_at.timestamp()),
|
||||
server_id=identity.server_id,
|
||||
subject_type=identity.subject_type,
|
||||
subject=identity.subject,
|
||||
grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key),
|
||||
return _seal(
|
||||
kind="access",
|
||||
prefix=ENVELOPE_PREFIX,
|
||||
identity=identity,
|
||||
grant_blob=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key),
|
||||
expires_at=expires_at,
|
||||
signing_key=keys.signing_key,
|
||||
now=now,
|
||||
)
|
||||
token = ENVELOPE_PREFIX + jwt.encode(
|
||||
claims.model_dump(),
|
||||
keys.signing_key.get_secret_value(),
|
||||
algorithm=_ENVELOPE_JWT_ALGORITHM,
|
||||
)
|
||||
size_bytes = len(token.encode("utf-8"))
|
||||
if size_bytes > MAX_ENVELOPE_BYTES:
|
||||
return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES)
|
||||
return SealedEnvelope(token=SecretStr(token), expires_at=expires_at)
|
||||
|
||||
|
||||
def open_envelope(
|
||||
|
|
@ -287,20 +336,9 @@ def open_envelope(
|
|||
re-derived, so it is stale by up to the envelope's lifetime; callers that need a
|
||||
live remaining lifetime should use ``now`` against the upstream, not this field.
|
||||
"""
|
||||
if not is_envelope(candidate):
|
||||
return NotAnEnvelope()
|
||||
# UTF-8 byte length is never below character length, so a character count already over the
|
||||
# cap rejects an oversize candidate in O(1) without encoding it; the exact byte check then
|
||||
# runs only on candidates already bounded to <= MAX_ENVELOPE_BYTES characters.
|
||||
if len(candidate) > MAX_ENVELOPE_BYTES:
|
||||
return MalformedPayload()
|
||||
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES:
|
||||
return MalformedPayload()
|
||||
claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key)
|
||||
claims = _open_claims(candidate, prefix=ENVELOPE_PREFIX, expected_kind="access", keys=keys, now=now)
|
||||
if not isinstance(claims, _EnvelopeClaims):
|
||||
return claims
|
||||
if now.timestamp() >= claims.exp:
|
||||
return Expired()
|
||||
grant = _decrypt_grant(claims.grant, keys.encryption_key)
|
||||
if not isinstance(grant, UpstreamTokenGrant):
|
||||
return grant
|
||||
|
|
@ -310,12 +348,124 @@ def open_envelope(
|
|||
)
|
||||
|
||||
|
||||
def mint_refresh_envelope(
|
||||
identity: EnvelopeIdentity,
|
||||
refresh: RefreshCredential,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
) -> SealedEnvelope | EnvelopeMintError:
|
||||
"""Seal ``refresh`` for ``identity`` into a long-lived, client-held refresh envelope.
|
||||
|
||||
``exp`` is ``min(refresh.expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS)`` seconds from ``now`` (the
|
||||
cap alone when the upstream omits the refresh lifetime). Sealing a distinct ``kind="refresh"`` claim
|
||||
is what keeps a refresh envelope from ever opening as an access credential at the MCP edge. Returns
|
||||
``EnvelopeTooLarge`` when the serialized envelope exceeds ``MAX_ENVELOPE_BYTES``.
|
||||
"""
|
||||
expires_at = now + timedelta(seconds=_refresh_ttl_seconds(refresh.expires_in))
|
||||
return _seal(
|
||||
kind="refresh",
|
||||
prefix=REFRESH_ENVELOPE_PREFIX,
|
||||
identity=identity,
|
||||
grant_blob=_encrypt_grant_blob(_refresh_plaintext(refresh), keys.encryption_key),
|
||||
expires_at=expires_at,
|
||||
signing_key=keys.signing_key,
|
||||
now=now,
|
||||
)
|
||||
|
||||
|
||||
def open_refresh_envelope(
|
||||
candidate: str,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
) -> OpenedRefreshEnvelope | EnvelopeOpenError:
|
||||
"""Validate a refresh ``candidate`` and recover the identity and inner refresh grant.
|
||||
|
||||
Total over hostile input exactly like :func:`open_envelope`: every invalid, expired, tampered,
|
||||
wrong-kind, or undecryptable candidate maps to a distinct ``EnvelopeOpenError`` variant, never a
|
||||
raise. The ``kind="refresh"`` claim is required, so an access envelope re-prefixed as a refresh one
|
||||
is rejected as ``MalformedPayload``.
|
||||
"""
|
||||
claims = _open_claims(candidate, prefix=REFRESH_ENVELOPE_PREFIX, expected_kind="refresh", keys=keys, now=now)
|
||||
if not isinstance(claims, _EnvelopeClaims):
|
||||
return claims
|
||||
refresh = _decrypt_refresh(claims.grant, keys.encryption_key)
|
||||
if not isinstance(refresh, RefreshCredential):
|
||||
return refresh
|
||||
return OpenedRefreshEnvelope(
|
||||
identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject),
|
||||
refresh=refresh,
|
||||
)
|
||||
|
||||
|
||||
def _seal(
|
||||
kind: EnvelopeKind,
|
||||
prefix: str,
|
||||
identity: EnvelopeIdentity,
|
||||
grant_blob: str,
|
||||
expires_at: datetime,
|
||||
signing_key: SecretStr,
|
||||
now: datetime,
|
||||
) -> SealedEnvelope | EnvelopeTooLarge:
|
||||
"""Sign the claims for either envelope kind and enforce the size cap. Shared by both mints so the
|
||||
JWT shape, issuer, and size guard cannot drift between access and refresh envelopes."""
|
||||
claims = _EnvelopeClaims(
|
||||
iss=ENVELOPE_ISSUER,
|
||||
iat=int(now.timestamp()),
|
||||
exp=int(expires_at.timestamp()),
|
||||
kind=kind,
|
||||
server_id=identity.server_id,
|
||||
subject_type=identity.subject_type,
|
||||
subject=identity.subject,
|
||||
grant=grant_blob,
|
||||
)
|
||||
token = prefix + jwt.encode(claims.model_dump(), signing_key.get_secret_value(), algorithm=_ENVELOPE_JWT_ALGORITHM)
|
||||
size_bytes = len(token.encode("utf-8"))
|
||||
if size_bytes > MAX_ENVELOPE_BYTES:
|
||||
return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES)
|
||||
return SealedEnvelope(token=SecretStr(token), expires_at=expires_at)
|
||||
|
||||
|
||||
def _open_claims(
|
||||
candidate: str,
|
||||
prefix: str,
|
||||
expected_kind: EnvelopeKind,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
) -> _EnvelopeClaims | EnvelopeOpenError:
|
||||
"""Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an attacker-controlled
|
||||
candidate, shared by both openers so the security gate is identical for access and refresh. Returns
|
||||
the validated claims or a distinct ``EnvelopeOpenError``; never raises."""
|
||||
if not candidate.startswith(prefix):
|
||||
return NotAnEnvelope()
|
||||
# UTF-8 byte length is never below character length, so a character count already over the cap
|
||||
# rejects an oversize candidate in O(1) without encoding it; the exact byte check then runs only on
|
||||
# candidates already bounded to <= MAX_ENVELOPE_BYTES characters.
|
||||
if len(candidate) > MAX_ENVELOPE_BYTES:
|
||||
return MalformedPayload()
|
||||
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES:
|
||||
return MalformedPayload()
|
||||
claims = _decode_claims(candidate.removeprefix(prefix), keys.signing_key)
|
||||
if not isinstance(claims, _EnvelopeClaims):
|
||||
return claims
|
||||
if claims.kind != expected_kind:
|
||||
return MalformedPayload()
|
||||
if now.timestamp() >= claims.exp:
|
||||
return Expired()
|
||||
return claims
|
||||
|
||||
|
||||
def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int:
|
||||
if upstream_expires_in is None:
|
||||
return MAX_ENVELOPE_TTL_SECONDS
|
||||
return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS)
|
||||
|
||||
|
||||
def _refresh_ttl_seconds(upstream_refresh_expires_in: int | None) -> int:
|
||||
if upstream_refresh_expires_in is None:
|
||||
return MAX_REFRESH_ENVELOPE_TTL_SECONDS
|
||||
return min(upstream_refresh_expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS)
|
||||
|
||||
|
||||
def _grant_plaintext(grant: UpstreamTokenGrant) -> str:
|
||||
wire = _GrantWire(
|
||||
access_token=grant.access_token.get_secret_value(),
|
||||
|
|
@ -327,6 +477,15 @@ def _grant_plaintext(grant: UpstreamTokenGrant) -> str:
|
|||
return wire.model_dump_json(exclude_none=True)
|
||||
|
||||
|
||||
def _refresh_plaintext(refresh: RefreshCredential) -> str:
|
||||
wire = _RefreshWire(
|
||||
refresh_token=refresh.refresh_token.get_secret_value(),
|
||||
scope=refresh.scope,
|
||||
expires_in=refresh.expires_in,
|
||||
)
|
||||
return wire.model_dump_json(exclude_none=True)
|
||||
|
||||
|
||||
def _decode_claims(
|
||||
compact: str,
|
||||
signing_key: SecretStr,
|
||||
|
|
@ -388,3 +547,22 @@ def _decrypt_grant(
|
|||
return UpstreamTokenGrant.model_validate_json(plaintext)
|
||||
except ValidationError:
|
||||
return MalformedPayload()
|
||||
|
||||
|
||||
def _decrypt_refresh(
|
||||
blob: str,
|
||||
encryption_key: SecretStr,
|
||||
) -> RefreshCredential | DecryptFailed | MalformedPayload:
|
||||
from nacl.exceptions import CryptoError
|
||||
|
||||
try:
|
||||
plaintext = decrypt_value(
|
||||
value=base64.urlsafe_b64decode(blob),
|
||||
signing_key=encryption_key.get_secret_value(),
|
||||
)
|
||||
except (CryptoError, ValueError):
|
||||
return DecryptFailed()
|
||||
try:
|
||||
return RefreshCredential.model_validate_json(plaintext)
|
||||
except ValidationError:
|
||||
return MalformedPayload()
|
||||
|
|
|
|||
|
|
@ -15,10 +15,14 @@ from pydantic import SecretStr
|
|||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
|
||||
BridgeEnvelopeAdmitted,
|
||||
BridgeEnvelopeInvalid,
|
||||
BridgeRefreshInvalid,
|
||||
BridgeRefreshOpened,
|
||||
NotBridgeEnvelope,
|
||||
build_bridge_refresh_token_response,
|
||||
build_bridge_token_response,
|
||||
envelope_keys_from_master_key,
|
||||
is_bridge_envelope_shaped,
|
||||
open_bridge_refresh_envelope,
|
||||
resolve_bridge_envelope,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
|
||||
|
|
@ -26,6 +30,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
|
|||
EnvelopeIdentity,
|
||||
EnvelopeKeys,
|
||||
EnvelopeTooLarge,
|
||||
RefreshCredential,
|
||||
SealedEnvelope,
|
||||
UpstreamTokenGrant,
|
||||
key_hash_identity,
|
||||
|
|
@ -49,6 +54,76 @@ def _sealed_token(keys: EnvelopeKeys, now: datetime = _NOW, identity: EnvelopeId
|
|||
return sealed.token.get_secret_value()
|
||||
|
||||
|
||||
_UPSTREAM_REFRESH = "upstream-refresh-do-not-leak-9b2c"
|
||||
|
||||
|
||||
def _sealed_refresh(keys: EnvelopeKeys, now: datetime = _NOW, identity: EnvelopeIdentity = _IDENTITY) -> str:
|
||||
sealed = build_bridge_refresh_token_response(
|
||||
identity, RefreshCredential(refresh_token=SecretStr(_UPSTREAM_REFRESH)), keys, now
|
||||
)
|
||||
assert isinstance(sealed, SealedEnvelope)
|
||||
return sealed.token.get_secret_value()
|
||||
|
||||
|
||||
def test_open_bridge_refresh_envelope_round_trips_identity_and_refresh():
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
result = open_bridge_refresh_envelope(_sealed_refresh(keys), keys, _NOW, _SERVER_ID)
|
||||
assert isinstance(result, BridgeRefreshOpened)
|
||||
assert result.identity == _IDENTITY
|
||||
assert result.refresh.refresh_token.get_secret_value() == _UPSTREAM_REFRESH
|
||||
|
||||
|
||||
def test_open_bridge_refresh_envelope_strips_bearer_scheme():
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
result = open_bridge_refresh_envelope(f"Bearer {_sealed_refresh(keys)}", keys, _NOW, _SERVER_ID)
|
||||
assert isinstance(result, BridgeRefreshOpened)
|
||||
|
||||
|
||||
def test_open_bridge_refresh_envelope_rejects_wrong_server():
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
result = open_bridge_refresh_envelope(_sealed_refresh(keys), keys, _NOW, "a-different-server")
|
||||
assert isinstance(result, BridgeRefreshInvalid)
|
||||
|
||||
|
||||
def test_open_bridge_refresh_envelope_rejects_non_refresh_bearers():
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
# an access envelope is not a refresh envelope; a raw upstream refresh token is not one either
|
||||
assert isinstance(open_bridge_refresh_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID), BridgeRefreshInvalid)
|
||||
assert isinstance(open_bridge_refresh_envelope("raw-refresh-token", keys, _NOW, _SERVER_ID), BridgeRefreshInvalid)
|
||||
|
||||
|
||||
def test_open_bridge_refresh_envelope_rejects_under_wrong_master_key():
|
||||
minted = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
other = envelope_keys_from_master_key(_MASTER_KEY + "-rotated")
|
||||
result = open_bridge_refresh_envelope(_sealed_refresh(minted), other, _NOW, _SERVER_ID)
|
||||
assert isinstance(result, BridgeRefreshInvalid)
|
||||
|
||||
|
||||
def test_refresh_envelope_is_never_admitted_at_the_tool_call_edge():
|
||||
"""A refresh envelope must never authenticate a tool call. The admission edge engages the bridge arm
|
||||
for it (is_bridge_envelope_shaped is true for either envelope kind), and the consumer rejects it as
|
||||
BridgeEnvelopeInvalid, which admission fails closed (401): a refresh credential is only ever
|
||||
presented back to the token endpoint."""
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
refresh = _sealed_refresh(keys)
|
||||
assert is_bridge_envelope_shaped(refresh) is True
|
||||
assert is_bridge_envelope_shaped(f"Bearer {refresh}") is True
|
||||
result = resolve_bridge_envelope(refresh, keys, _NOW, _SERVER_ID)
|
||||
assert isinstance(result, BridgeEnvelopeInvalid)
|
||||
|
||||
|
||||
def test_refresh_jwt_wearing_the_access_prefix_is_rejected_at_the_edge():
|
||||
"""Belt-and-suspenders against a swapped wire prefix: a refresh JWT re-prefixed as an access envelope
|
||||
opens far enough to hit the signed kind claim, which rejects it, so admission fails closed rather
|
||||
than forwarding a refresh credential's contents upstream."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import REFRESH_ENVELOPE_PREFIX
|
||||
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
swapped = ENVELOPE_PREFIX + _sealed_refresh(keys).removeprefix(REFRESH_ENVELOPE_PREFIX)
|
||||
result = resolve_bridge_envelope(swapped, keys, _NOW, _SERVER_ID)
|
||||
assert isinstance(result, BridgeEnvelopeInvalid)
|
||||
|
||||
|
||||
def test_key_derivation_is_deterministic():
|
||||
assert envelope_keys_from_master_key(_MASTER_KEY) == envelope_keys_from_master_key(_MASTER_KEY)
|
||||
|
||||
|
|
|
|||
|
|
@ -24,6 +24,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
|
|||
ENVELOPE_PREFIX,
|
||||
MAX_ENVELOPE_BYTES,
|
||||
MAX_ENVELOPE_TTL_SECONDS,
|
||||
MAX_REFRESH_ENVELOPE_TTL_SECONDS,
|
||||
REFRESH_ENVELOPE_PREFIX,
|
||||
BadSignature,
|
||||
DecryptFailed,
|
||||
EnvelopeIdentity,
|
||||
|
|
@ -33,12 +35,17 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
|
|||
MalformedPayload,
|
||||
NotAnEnvelope,
|
||||
OpenedEnvelope,
|
||||
OpenedRefreshEnvelope,
|
||||
RefreshCredential,
|
||||
SealedEnvelope,
|
||||
UpstreamTokenGrant,
|
||||
is_envelope,
|
||||
is_refresh_envelope,
|
||||
key_hash_identity,
|
||||
mint_envelope,
|
||||
mint_refresh_envelope,
|
||||
open_envelope,
|
||||
open_refresh_envelope,
|
||||
user_identity,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value
|
||||
|
|
@ -139,10 +146,11 @@ def test_minimal_grant_round_trips_without_none_leakage_into_claims():
|
|||
def test_claim_layout_and_no_plaintext_token_in_envelope():
|
||||
token = _sealed_token(_full_grant())
|
||||
claims = _unverified_claims(token)
|
||||
assert set(claims) == {"iss", "iat", "exp", "server_id", "subject_type", "subject", "grant"}
|
||||
assert set(claims) == {"iss", "iat", "exp", "kind", "server_id", "subject_type", "subject", "grant"}
|
||||
assert claims["iss"] == ENVELOPE_ISSUER
|
||||
assert claims["iat"] == int(_NOW.timestamp())
|
||||
assert claims["exp"] == int(_NOW.timestamp()) + 600
|
||||
assert claims["kind"] == "access"
|
||||
assert claims["server_id"] == "srv-456"
|
||||
assert claims["subject_type"] == "key_hash"
|
||||
assert claims["subject"] == "hashed-key-123"
|
||||
|
|
@ -151,6 +159,86 @@ def test_claim_layout_and_no_plaintext_token_in_envelope():
|
|||
assert _REFRESH_TOKEN not in json.dumps(claims)
|
||||
|
||||
|
||||
def _refresh_credential() -> RefreshCredential:
|
||||
return RefreshCredential(refresh_token=SecretStr(_REFRESH_TOKEN), scope="read:tools", expires_in=None)
|
||||
|
||||
|
||||
def _sealed_refresh_token(refresh: RefreshCredential | None = None, keys: EnvelopeKeys = _KEYS) -> str:
|
||||
sealed = mint_refresh_envelope(_IDENTITY, refresh or _refresh_credential(), keys, _NOW)
|
||||
assert isinstance(sealed, SealedEnvelope)
|
||||
return sealed.token.get_secret_value()
|
||||
|
||||
|
||||
def test_refresh_envelope_round_trips_identity_and_refresh_token():
|
||||
token = _sealed_refresh_token()
|
||||
assert is_refresh_envelope(token)
|
||||
assert not is_envelope(token)
|
||||
opened = open_refresh_envelope(token, _KEYS, _NOW)
|
||||
assert isinstance(opened, OpenedRefreshEnvelope)
|
||||
assert opened.identity == _IDENTITY
|
||||
assert opened.refresh.refresh_token.get_secret_value() == _REFRESH_TOKEN
|
||||
assert opened.refresh.scope == "read:tools"
|
||||
|
||||
|
||||
def test_refresh_envelope_ttl_is_min_of_upstream_refresh_lifetime_and_cap():
|
||||
short = mint_refresh_envelope(
|
||||
_IDENTITY, RefreshCredential(refresh_token=SecretStr("r"), expires_in=120), _KEYS, _NOW
|
||||
)
|
||||
assert isinstance(short, SealedEnvelope)
|
||||
assert short.expires_at == _NOW + timedelta(seconds=120)
|
||||
capped = mint_refresh_envelope(
|
||||
_IDENTITY,
|
||||
RefreshCredential(refresh_token=SecretStr("r"), expires_in=MAX_REFRESH_ENVELOPE_TTL_SECONDS + 86400),
|
||||
_KEYS,
|
||||
_NOW,
|
||||
)
|
||||
assert isinstance(capped, SealedEnvelope)
|
||||
assert capped.expires_at == _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS)
|
||||
default = mint_refresh_envelope(_IDENTITY, RefreshCredential(refresh_token=SecretStr("r")), _KEYS, _NOW)
|
||||
assert isinstance(default, SealedEnvelope)
|
||||
assert default.expires_at == _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS)
|
||||
|
||||
|
||||
def test_access_and_refresh_envelopes_do_not_cross_open():
|
||||
access = _sealed_token(_full_grant())
|
||||
refresh = _sealed_refresh_token()
|
||||
# each opener rejects the other kind's prefix outright
|
||||
assert isinstance(open_refresh_envelope(access, _KEYS, _NOW), NotAnEnvelope)
|
||||
assert isinstance(open_envelope(refresh, _KEYS, _NOW), NotAnEnvelope)
|
||||
|
||||
|
||||
def test_prefix_swap_is_rejected_by_the_signed_kind_claim():
|
||||
# the wire prefix is not signed, so swap it; the signed kind claim must still reject the cross-use
|
||||
refresh = _sealed_refresh_token()
|
||||
swapped_to_access = ENVELOPE_PREFIX + refresh.removeprefix(REFRESH_ENVELOPE_PREFIX)
|
||||
assert isinstance(open_envelope(swapped_to_access, _KEYS, _NOW), MalformedPayload)
|
||||
access = _sealed_token(_full_grant())
|
||||
swapped_to_refresh = REFRESH_ENVELOPE_PREFIX + access.removeprefix(ENVELOPE_PREFIX)
|
||||
assert isinstance(open_refresh_envelope(swapped_to_refresh, _KEYS, _NOW), MalformedPayload)
|
||||
|
||||
|
||||
def test_refresh_envelope_total_over_hostile_input():
|
||||
token = _sealed_refresh_token()
|
||||
# expired against the injected clock
|
||||
assert isinstance(
|
||||
open_refresh_envelope(token, _KEYS, _NOW + timedelta(seconds=MAX_REFRESH_ENVELOPE_TTL_SECONDS)), Expired
|
||||
)
|
||||
# wrong signing key
|
||||
assert isinstance(open_refresh_envelope(token, _WRONG_SIGNING, _NOW), BadSignature)
|
||||
# right signature, wrong encryption key
|
||||
assert isinstance(open_refresh_envelope(token, _WRONG_ENCRYPTION, _NOW), DecryptFailed)
|
||||
# not an envelope at all
|
||||
assert isinstance(open_refresh_envelope("raw-upstream-refresh-token", _KEYS, _NOW), NotAnEnvelope)
|
||||
|
||||
|
||||
def test_refresh_envelope_never_leaks_the_refresh_token_in_plaintext():
|
||||
token = _sealed_refresh_token()
|
||||
assert _REFRESH_TOKEN not in token
|
||||
claims = jwt.decode(token.removeprefix(REFRESH_ENVELOPE_PREFIX), options={"verify_signature": False})
|
||||
assert claims["kind"] == "refresh"
|
||||
assert _REFRESH_TOKEN not in json.dumps(claims)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expires_in, expected_ttl",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -4672,11 +4672,11 @@ async def test_bridge_envelope_too_large_upstream_token_is_502():
|
|||
|
||||
|
||||
@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."""
|
||||
async def test_bridge_access_envelope_never_carries_upstream_refresh_token():
|
||||
"""The upstream refresh token is never sealed into the ACCESS envelope, the bearer forwarded upstream
|
||||
on every tool call: the opened access grant carries no refresh token even when the upstream returned
|
||||
one, and the raw refresh token never appears in the access envelope. It rides only in the separate
|
||||
refresh envelope returned as the response's refresh_token, encrypted, never in plaintext."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
|
||||
|
|
@ -4698,21 +4698,22 @@ async def test_bridge_envelope_does_not_seal_upstream_refresh_token():
|
|||
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
|
||||
# the refresh token rides only in the separate, encrypted refresh envelope, never in plaintext
|
||||
assert body["refresh_token"].startswith("llm_refresh_")
|
||||
assert "UPSTREAM-REFRESH" not in body["refresh_token"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_grant_is_rejected_before_upstream():
|
||||
"""A bridge oauth_delegate server issues only envelopes and seals no upstream refresh_token, so the
|
||||
client never holds one to present. _prepare_bridge_mint rejects the refresh_token grant up front
|
||||
with unsupported_grant_type, BEFORE any upstream exchange, so a stray refresh request can never
|
||||
rotate or consume the client's upstream refresh credential; renewal is re-running
|
||||
authorization_code. This is checked before identity resolution, so it holds even with a valid key."""
|
||||
async def test_bridge_refresh_grant_with_non_envelope_is_invalid_grant_before_upstream():
|
||||
"""A bridge oauth_delegate client only ever holds a refresh envelope, never a raw upstream refresh
|
||||
token, so a refresh_token grant carrying a bare (non-envelope) value is invalid_grant, rejected in
|
||||
_prepare_bridge_refresh BEFORE any upstream exchange. Rejecting before the exchange means a bad
|
||||
refresh request can never consume or rotate an upstream refresh token."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
|
|
@ -4739,10 +4740,557 @@ async def test_bridge_refresh_grant_is_rejected_before_upstream():
|
|||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "unsupported_grant_type"
|
||||
assert json.loads(response.body)["error"] == "invalid_grant"
|
||||
fake_http_client.post.assert_not_called()
|
||||
|
||||
|
||||
def _mint_test_refresh_envelope(
|
||||
server_id="bridge_srv", key_hash="hashed-litellm-key-77", upstream_refresh="UPSTREAM-REFRESH", identity=None,
|
||||
scope=None,
|
||||
):
|
||||
"""Mint a refresh envelope the way the producer does, for driving the refresh_token grant in tests.
|
||||
Defaults to a key_hash subject; pass ``identity`` to seal a specific subject (e.g. a user_id), and
|
||||
``scope`` to seal the scope to re-request on refresh."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from pydantic import SecretStr
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
|
||||
build_bridge_refresh_token_response,
|
||||
envelope_keys_from_master_key,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
|
||||
RefreshCredential,
|
||||
SealedEnvelope,
|
||||
key_hash_identity,
|
||||
)
|
||||
|
||||
keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY)
|
||||
identity = identity if identity is not None else key_hash_identity(server_id=server_id, key_hash=key_hash)
|
||||
sealed = build_bridge_refresh_token_response(
|
||||
identity, RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope=scope), keys,
|
||||
datetime.now(timezone.utc),
|
||||
)
|
||||
assert isinstance(sealed, SealedEnvelope)
|
||||
return sealed.token.get_secret_value()
|
||||
|
||||
|
||||
async def _refresh_for_bridge_server(
|
||||
server, refresh_envelope_value, upstream_body, revalidate_result=None, fake_client_out=None
|
||||
):
|
||||
"""Drive a refresh_token grant for a bridge server: the client presents ``refresh_envelope_value``,
|
||||
the sealed subject re-validates to ``revalidate_result`` (``None`` when the key or user is still
|
||||
active, or a failure literal like "no_active_key" when revoked/deactivated), and the upstream returns
|
||||
``upstream_body``. Patching the single subject-revalidation dispatch covers both a key_hash and a
|
||||
user_id refresh envelope. Returns the response; the captured client exposes the POST call so a test
|
||||
can assert what refresh token was actually sent upstream."""
|
||||
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)
|
||||
if fake_client_out is not None:
|
||||
fake_client_out["client"] = fake_http_client
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=fake_http_client,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._revalidate_active_subject",
|
||||
new=AsyncMock(return_value=revalidate_result),
|
||||
),
|
||||
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="refresh_token",
|
||||
code=None,
|
||||
redirect_uri=None,
|
||||
client_id="dcr-client-123",
|
||||
client_secret=None,
|
||||
code_verifier=None,
|
||||
refresh_token=refresh_envelope_value,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_mint_returns_refresh_envelope_that_opens_to_upstream_refresh():
|
||||
"""When the upstream returns a refresh token, the authorization_code mint returns a refresh envelope
|
||||
alongside the access envelope. The refresh envelope is a distinct llm_refresh_ credential that opens
|
||||
(under the same keys and server_id) to the upstream refresh token, so the client can renew later."""
|
||||
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 (
|
||||
OpenedRefreshEnvelope,
|
||||
open_refresh_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": "R-UP"}
|
||||
response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
|
||||
|
||||
body = json.loads(response.body)
|
||||
refresh_env = body["refresh_token"]
|
||||
assert refresh_env.startswith("llm_refresh_")
|
||||
keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY)
|
||||
opened = open_refresh_envelope(refresh_env, keys, datetime.now(timezone.utc))
|
||||
assert isinstance(opened, OpenedRefreshEnvelope)
|
||||
assert opened.identity.server_id == server.server_id
|
||||
assert opened.refresh.refresh_token.get_secret_value() == "R-UP"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_mint_omits_refresh_envelope_when_upstream_has_no_refresh():
|
||||
"""No refresh envelope is issued when the upstream returns no refresh token, so the response carries
|
||||
only the access envelope; the client re-authenticates at access expiry (nothing to renew with)."""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
upstream = {"access_token": "UP", "token_type": "Bearer", "expires_in": 3600}
|
||||
response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
|
||||
|
||||
body = json.loads(response.body)
|
||||
assert body["access_token"].startswith("llm_env_")
|
||||
assert "refresh_token" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_grant_sends_unwrapped_upstream_token_and_renews():
|
||||
"""A refresh_token grant carrying a valid refresh envelope renews: the exchange unwraps the envelope
|
||||
and sends the REAL upstream refresh token upstream (never the envelope), then returns a fresh access
|
||||
envelope. This is the flow that lets the client renew without re-authenticating."""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="UPSTREAM-REFRESH")
|
||||
upstream = {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600}
|
||||
captured: dict = {}
|
||||
response = await _refresh_for_bridge_server(server, refresh_env, upstream, None, fake_client_out=captured)
|
||||
|
||||
assert response.status_code == 200
|
||||
body = json.loads(response.body)
|
||||
assert body["access_token"].startswith("llm_env_")
|
||||
# the upstream exchange received the unwrapped upstream refresh token, never the client's envelope
|
||||
sent = captured["client"].post.call_args.kwargs["data"]
|
||||
assert sent["grant_type"] == "refresh_token"
|
||||
assert sent["refresh_token"] == "UPSTREAM-REFRESH"
|
||||
assert not sent["refresh_token"].startswith("llm_refresh_")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_grant_rotates_refresh_envelope_wrapping_new_upstream_token():
|
||||
"""When the upstream rotates the refresh token on renewal, the client receives a new refresh envelope
|
||||
that wraps the NEW upstream refresh token, so the rotation is carried through faithfully."""
|
||||
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 (
|
||||
OpenedRefreshEnvelope,
|
||||
open_refresh_envelope,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="OLD-UP-REFRESH")
|
||||
upstream = {
|
||||
"access_token": "NEW-ACCESS",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"refresh_token": "NEW-UP-REFRESH",
|
||||
}
|
||||
response = await _refresh_for_bridge_server(server, refresh_env, upstream, None)
|
||||
|
||||
body = json.loads(response.body)
|
||||
keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY)
|
||||
opened = open_refresh_envelope(body["refresh_token"], keys, datetime.now(timezone.utc))
|
||||
assert isinstance(opened, OpenedRefreshEnvelope)
|
||||
assert opened.refresh.refresh_token.get_secret_value() == "NEW-UP-REFRESH"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_grant_with_revoked_key_is_invalid_grant_before_upstream():
|
||||
"""A valid refresh envelope whose sealed litellm key has since been revoked cannot keep refreshing:
|
||||
the reload gate reports no_active_key and the refresh is invalid_grant, returned BEFORE the upstream
|
||||
exchange so the upstream refresh token is never consumed. Revocation kills renewal."""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
refresh_env = _mint_test_refresh_envelope(server_id=server.server_id)
|
||||
captured: dict = {}
|
||||
response = await _refresh_for_bridge_server(
|
||||
server, refresh_env, {"access_token": "NEW"}, "no_active_key", fake_client_out=captured
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "invalid_grant"
|
||||
captured["client"].post.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_envelope_for_another_server_is_invalid_grant():
|
||||
"""A refresh envelope minted for one server cannot renew against another: the sealed server_id must
|
||||
match the server the refresh targets, so a cross-server refresh envelope is invalid_grant and never
|
||||
reaches the upstream exchange."""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
foreign_env = _mint_test_refresh_envelope(server_id="some-other-server")
|
||||
captured: dict = {}
|
||||
response = await _refresh_for_bridge_server(
|
||||
server, foreign_env, {"access_token": "NEW"}, None, fake_client_out=captured
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "invalid_grant"
|
||||
captured["client"].post.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_grant_renews_a_user_subject_envelope():
|
||||
"""The interactive SSO client mints a user_id-subject envelope, so its refresh envelope carries a
|
||||
user subject too. Renewing it re-validates the user (still active here), unwraps the upstream refresh
|
||||
token, and returns a fresh access envelope that opens back to the same user_id subject; the upstream
|
||||
exchange received the real upstream refresh token, not the client's envelope."""
|
||||
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.proxy._experimental.mcp_server.outbound_credentials.envelope import user_identity
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
user_env = _mint_test_refresh_envelope(
|
||||
identity=user_identity(server_id=server.server_id, user_id="sso-user-42"), upstream_refresh="UP-REFRESH-USER"
|
||||
)
|
||||
upstream = {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600}
|
||||
captured: dict = {}
|
||||
response = await _refresh_for_bridge_server(server, user_env, upstream, None, fake_client_out=captured)
|
||||
|
||||
assert response.status_code == 200
|
||||
body = json.loads(response.body)
|
||||
keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY)
|
||||
opened = resolve_bridge_envelope(body["access_token"], keys, datetime.now(timezone.utc), server.server_id)
|
||||
assert isinstance(opened, BridgeEnvelopeAdmitted)
|
||||
assert opened.identity.subject_type == "user_id"
|
||||
assert opened.identity.subject == "sso-user-42"
|
||||
assert captured["client"].post.call_args.kwargs["data"]["refresh_token"] == "UP-REFRESH-USER"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_re_requests_the_sealed_scope_when_client_omits_it():
|
||||
"""A DCR/MCP client omits scope on the refresh request, so the gateway must re-request the scope sealed
|
||||
at mint; dropping it lets a stricter upstream narrow the renewed token. The upstream POST must carry
|
||||
the sealed scope even though the client sent none. Regression for the dropped sealed refresh scope."""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
refresh_env = _mint_test_refresh_envelope(
|
||||
server_id=server.server_id, upstream_refresh="UP-REFRESH", scope="read:tools write:tools"
|
||||
)
|
||||
captured: dict = {}
|
||||
response = await _refresh_for_bridge_server(
|
||||
server, refresh_env, {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600}, None,
|
||||
fake_client_out=captured,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured["client"].post.call_args.kwargs["data"]["scope"] == "read:tools write:tools"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_re_seals_scope_when_upstream_omits_it_so_the_chain_keeps_it():
|
||||
"""RFC 6749 5.1 lets an upstream omit scope in a refresh response when it is unchanged. The re-minted
|
||||
refresh envelope must still seal the scope that was requested, otherwise the NEXT refresh loses it and
|
||||
a stricter upstream could narrow the token. The returned refresh envelope carries the scope even though
|
||||
the upstream response had none, and a second refresh off it still re-requests the scope."""
|
||||
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 (
|
||||
OpenedRefreshEnvelope,
|
||||
open_refresh_envelope,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
refresh_env = _mint_test_refresh_envelope(
|
||||
server_id=server.server_id, upstream_refresh="UP-1", scope="mcp:read mcp:write"
|
||||
)
|
||||
# the upstream rotates the refresh token but OMITS scope (valid when unchanged)
|
||||
upstream_no_scope = {"access_token": "NEW", "token_type": "Bearer", "expires_in": 3600, "refresh_token": "UP-2"}
|
||||
|
||||
captured: dict = {}
|
||||
r1 = await _refresh_for_bridge_server(server, refresh_env, upstream_no_scope, None, fake_client_out=captured)
|
||||
assert r1.status_code == 200
|
||||
assert captured["client"].post.call_args.kwargs["data"]["scope"] == "mcp:read mcp:write"
|
||||
|
||||
keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY)
|
||||
new_env = json.loads(r1.body)["refresh_token"]
|
||||
opened = open_refresh_envelope(new_env, keys, datetime.now(timezone.utc))
|
||||
assert isinstance(opened, OpenedRefreshEnvelope)
|
||||
assert opened.refresh.scope == "mcp:read mcp:write"
|
||||
|
||||
captured2: dict = {}
|
||||
r2 = await _refresh_for_bridge_server(server, new_env, upstream_no_scope, None, fake_client_out=captured2)
|
||||
assert r2.status_code == 200
|
||||
assert captured2["client"].post.call_args.kwargs["data"]["scope"] == "mcp:read mcp:write"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_grant_with_deactivated_user_is_invalid_grant_before_upstream():
|
||||
"""A user_id-subject refresh envelope whose user has since been deactivated (SCIM offboarding, or
|
||||
the user no longer exists) cannot keep refreshing: subject re-validation reports no_active_key and
|
||||
the refresh is invalid_grant, returned BEFORE the upstream exchange. Revocation kills renewal for the
|
||||
user subject exactly as it does for the key subject."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import user_identity
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
user_env = _mint_test_refresh_envelope(identity=user_identity(server_id=server.server_id, user_id="gone-user"))
|
||||
captured: dict = {}
|
||||
response = await _refresh_for_bridge_server(
|
||||
server, user_env, {"access_token": "NEW"}, "no_active_key", fake_client_out=captured
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "invalid_grant"
|
||||
captured["client"].post.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revalidate_active_subject_dispatches_on_subject_type():
|
||||
"""Subject re-validation routes a key_hash envelope to the key reload and a user_id envelope to the
|
||||
user reload, so revocation gates renewal for either identity source through one dispatch point."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
_ResolvedKey,
|
||||
_revalidate_active_subject,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity, user_identity
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash",
|
||||
new=AsyncMock(return_value=_ResolvedKey(key_hash="kh", key=MagicMock())),
|
||||
) as key_reload,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_user_by_id",
|
||||
new=AsyncMock(return_value=None),
|
||||
) as user_reload,
|
||||
):
|
||||
assert await _revalidate_active_subject(key_hash_identity(server_id="s", key_hash="kh")) is None
|
||||
key_reload.assert_awaited_once_with("kh")
|
||||
user_reload.assert_not_awaited()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash",
|
||||
new=AsyncMock(),
|
||||
) as key_reload2,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_user_by_id",
|
||||
new=AsyncMock(return_value="no_active_key"),
|
||||
) as user_reload2,
|
||||
):
|
||||
assert await _revalidate_active_subject(user_identity(server_id="s", user_id="u42")) == "no_active_key"
|
||||
user_reload2.assert_awaited_once_with("u42")
|
||||
key_reload2.assert_not_awaited()
|
||||
|
||||
|
||||
def test_upstream_refresh_credential_expired_refresh_token_is_not_sealed():
|
||||
"""An upstream that reports its refresh token already elapsed (refresh_expires_in non-positive) must
|
||||
not be sealed: _upstream_refresh_credential returns None so the exchange degrades to an access-only
|
||||
response, mirroring how the access grant refuses an already-elapsed access token rather than capping a
|
||||
dead token to the full refresh TTL. A live or unspecified lifetime still yields a credential."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _upstream_refresh_credential
|
||||
|
||||
assert _upstream_refresh_credential({"access_token": "A", "refresh_token": "R", "refresh_expires_in": 0}) is None
|
||||
assert _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": -5}) is None
|
||||
live = _upstream_refresh_credential({"refresh_token": "R", "refresh_expires_in": 1800})
|
||||
assert live is not None and live.expires_in == 1800
|
||||
unspecified = _upstream_refresh_credential({"refresh_token": "R"})
|
||||
assert unspecified is not None and unspecified.expires_in is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_upstream_invalid_grant_maps_to_invalid_grant():
|
||||
"""When the sealed upstream refresh token has been revoked or expired at the IdP, the upstream returns
|
||||
400 invalid_grant. The bridge refresh path maps that to an RFC 6749 invalid_grant response so the OAuth
|
||||
client re-runs authorization_code, rather than surfacing the opaque upstream error it cannot act on."""
|
||||
import httpx
|
||||
|
||||
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)
|
||||
refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="LIVE-ENVELOPE-REFRESH")
|
||||
|
||||
error_response = MagicMock()
|
||||
error_response.status_code = 400
|
||||
error_response.text = '{"error": "invalid_grant", "error_description": "refresh token expired"}'
|
||||
error_response.raise_for_status = MagicMock(
|
||||
side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response)
|
||||
)
|
||||
fake_http_client = MagicMock()
|
||||
fake_http_client.post = AsyncMock(return_value=error_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._revalidate_active_subject",
|
||||
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=refresh_env,
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "invalid_grant"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_refresh_upstream_error_detection_parses_json_not_substring():
|
||||
"""The upstream invalid_grant detection parses the RFC 6749 5.2 error field, not a substring of the
|
||||
body. An upstream error whose code is not invalid_grant (here invalid_client, with the string
|
||||
invalid_grant only inside error_description) must NOT be mistaken for a dead refresh token, so it
|
||||
propagates as the upstream error rather than triggering a spurious authorization_code re-run."""
|
||||
import httpx
|
||||
|
||||
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)
|
||||
refresh_env = _mint_test_refresh_envelope(server_id=server.server_id, upstream_refresh="UP")
|
||||
|
||||
error_response = MagicMock()
|
||||
error_response.status_code = 400
|
||||
error_response.text = '{"error": "invalid_client", "error_description": "this is not an invalid_grant problem"}'
|
||||
error_response.raise_for_status = MagicMock(
|
||||
side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response)
|
||||
)
|
||||
fake_http_client = MagicMock()
|
||||
fake_http_client.post = AsyncMock(return_value=error_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._revalidate_active_subject",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY),
|
||||
):
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
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=refresh_env,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revalidate_key_subject_revoked_when_owner_scim_deactivated(proxy_globals):
|
||||
"""A key_hash refresh envelope whose key is still active but whose OWNING user was SCIM-deactivated must
|
||||
fail closed to no_active_key, mirroring how admission's _reject_if_admitted_owner_scim_deactivated
|
||||
revokes an offboarded owner's key. Without this, an offboarded user keeps renewing a live key."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _ResolvedKey, _revalidate_active_subject
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
proxy_globals.user_api_key_cache = UserApiKeyCache()
|
||||
proxy_globals.prisma_client = object()
|
||||
|
||||
resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="offboarded-owner"))
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash",
|
||||
new=AsyncMock(return_value=resolved),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_checks.get_user_object",
|
||||
new=AsyncMock(return_value=MagicMock(metadata={"scim_active": False})),
|
||||
),
|
||||
):
|
||||
result = await _revalidate_active_subject(key_hash_identity(server_id="s", key_hash="kh"))
|
||||
|
||||
assert result == "no_active_key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revalidate_key_subject_active_owner_renews_and_missing_owner_fails_open(proxy_globals):
|
||||
"""The key-owner SCIM gate blocks only an explicit scim_active False: an active owner renews (None), and
|
||||
a missing owner (get_user_object's wrapped ValueError) fails OPEN, since a key may outlive its owner
|
||||
record and a transient blip must not revoke a live key."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _ResolvedKey, _revalidate_active_subject
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
proxy_globals.user_api_key_cache = UserApiKeyCache()
|
||||
proxy_globals.prisma_client = object()
|
||||
resolved = _ResolvedKey(key_hash="kh", key=MagicMock(user_id="live-owner"))
|
||||
identity = key_hash_identity(server_id="s", key_hash="kh")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash",
|
||||
new=AsyncMock(return_value=resolved),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_checks.get_user_object",
|
||||
new=AsyncMock(return_value=MagicMock(metadata={"scim_active": True})),
|
||||
),
|
||||
):
|
||||
assert await _revalidate_active_subject(identity) is None
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_key_by_hash",
|
||||
new=AsyncMock(return_value=resolved),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_checks.get_user_object",
|
||||
new=AsyncMock(side_effect=_wrapped_user_lookup_error(Exception())),
|
||||
),
|
||||
):
|
||||
assert await _revalidate_active_subject(identity) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset():
|
||||
"""master_key is validated BEFORE the upstream exchange (in _prepare_bridge_mint), so a
|
||||
|
|
@ -5553,6 +6101,60 @@ async def test_resolve_active_litellm_key_no_database_is_unresolvable(proxy_glob
|
|||
assert await _resolve_active_litellm_key(request) == "unresolvable"
|
||||
|
||||
|
||||
def _wrapped_user_lookup_error(original: BaseException) -> ValueError:
|
||||
"""Reproduce get_user_object's real exception contract (litellm/proxy/auth/auth_checks.py): it
|
||||
catches every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the original
|
||||
error (a missing-user Exception or a real outage) survives only as ``__context__``. Injecting a raw
|
||||
ConnectionError/Exception instead would exercise a shape production never produces and let a
|
||||
chain-blind outage classifier pass. The wrapping fidelity is pinned by
|
||||
test_get_user_object_wraps_db_outage_as_valueerror_preserving_context in test_auth_checks."""
|
||||
try:
|
||||
raise original
|
||||
except BaseException:
|
||||
try:
|
||||
raise ValueError(f"User doesn't exist in db. Got error - {original}")
|
||||
except ValueError as wrapped:
|
||||
return wrapped
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reload_active_user_by_id_missing_user_is_no_active_key(proxy_globals):
|
||||
"""A user_id refresh envelope whose user has been deleted must fail closed to no_active_key (the
|
||||
refresh path maps it to invalid_grant), not unresolvable/500. get_user_object catches the missing row
|
||||
and re-raises a bare ValueError, so a missing user must not be misclassified as a DB outage or an
|
||||
opaque gateway fault."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
proxy_globals.user_api_key_cache = UserApiKeyCache()
|
||||
proxy_globals.prisma_client = object()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.auth_checks.get_user_object",
|
||||
new=AsyncMock(side_effect=_wrapped_user_lookup_error(Exception())),
|
||||
):
|
||||
assert await _reload_active_user_by_id("gone-user") == "no_active_key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reload_active_user_by_id_db_outage_is_unavailable(proxy_globals):
|
||||
"""A transient DB outage while re-validating the user on refresh is a retryable outage, distinct from
|
||||
a missing user, so the refresh path surfaces "unavailable" (a 503) rather than blaming the caller.
|
||||
get_user_object wraps the outage in a bare ValueError, so this exercises the chain-aware classifier; a
|
||||
raw ConnectionError would falsely pass even a chain-blind check because it is an OSError."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _reload_active_user_by_id
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
proxy_globals.user_api_key_cache = UserApiKeyCache()
|
||||
proxy_globals.prisma_client = object()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.auth_checks.get_user_object",
|
||||
new=AsyncMock(side_effect=_wrapped_user_lookup_error(ConnectionError("user database unreachable"))),
|
||||
):
|
||||
assert await _reload_active_user_by_id("sso-user-7") == "unavailable"
|
||||
|
||||
|
||||
@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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue