From c4dd06a0bb95776e998195cfb1a8adbb9321bccb Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 11 Jul 2026 19:16:58 -0700 Subject: [PATCH] feat(mcp): client-held refresh envelope for the dcr_bridge oauth_delegate flow A dcr_bridge oauth_delegate access envelope is capped at one hour, and until now the mode had no refresh at all: when the envelope expired the client had to re-run the interactive authorization_code flow. This adds a second client-held credential, the refresh envelope, so the client renews on a back channel and only re-authenticates when the refresh envelope expires or the upstream refresh token dies. The refresh envelope is a distinct llm_refresh_ credential that seals only the upstream refresh token (never the access token) bound to the same litellm identity and MCP server as the access envelope, under the same master-key-derived keys, with nothing stored server-side. Both envelopes now carry a signed kind claim ("access" or "refresh") that open() requires to match, so a refresh envelope can never open as an access credential even if its wire prefix is swapped (the prefix is not signed; the claim is). A refresh envelope presented at the MCP tool-call edge is not an access envelope, so admission fails it closed the same way it already fails any non-access bearer. At the token endpoint the authorization_code mint now returns a refresh envelope alongside the access envelope whenever the upstream returned a refresh token, and the refresh_token grant is supported for bridge servers: the client presents its refresh envelope, the endpoint opens it, re-validates the sealed litellm key so a revoked key cannot keep refreshing, unwraps the real upstream refresh token, exchanges it with the upstream IdP, and returns a fresh access envelope. Because the endpoint re-seals a refresh envelope only when the upstream returns a new refresh token, the design mirrors the upstream's own rotation policy rather than reinventing it: with a rotating upstream the client rotates and reuse is detected upstream; with a non-rotating upstream the original refresh envelope stands until its bounded 14-day TTL. Both preconditions and the unwrap run before the exchange, so a rejected refresh never consumes or rotates an upstream token. The pure envelope and credential layers stay side-effect free: mint/open share one signing, size, and kind gate across both envelope kinds, and every failure is a value. Tests cover the refresh round-trip, the kind-claim and server-id bindings, the revoked-key gate, upstream rotation carried through, the unwrap sending the real upstream token upstream, and edge rejection of a refresh envelope; the three security bindings are mutation-checked. Limitation documented in the PR: gateway-enforced refresh rotation with reuse detection would require server-side state, which this zero-custody mode omits by design, so the refresh envelope inherits the upstream's rotation posture plus gateway identity binding and a bounded TTL. --- .../mcp_server/discoverable_endpoints.py | 281 +++++++++++++-- .../bridge_credentials.py | 66 ++++ .../outbound_credentials/envelope.py | 244 +++++++++++-- .../test_bridge_credentials.py | 71 ++++ .../outbound_credentials/test_envelope.py | 90 ++++- .../mcp_server/test_discoverable_endpoints.py | 330 +++++++++++++++++- 6 files changed, 998 insertions(+), 84 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index bd26abe0b26..e56210ca393 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -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,67 @@ 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. Classification matches :func:`_reload_active_key_by_hash`: no DB connection or an + unexpected error is a gateway fault, a ``ProxyException`` / ``HTTPException`` from ``get_user_object`` + or a missing / deactivated user is ``no_active_key`` (the caller maps it to invalid_grant on refresh), + and a database-service-unavailable error is a retryable outage.""" + 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 # classify: a DB outage is retryable, anything else is an opaque gateway fault + if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): + return "unavailable" + verbose_logger.debug("_reload_active_user_by_id: unexpected user-resolution error (%s)", type(exc).__name__) + return "unresolvable" + 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 _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 (a blocked key, a deactivated user) gates renewal + for either identity source, the same way admission gates the egress.""" + match identity.subject_type: + case "key_hash": + reloaded = await _reload_active_key_by_hash(identity.subject) + return None if isinstance(reloaded, _ResolvedKey) else reloaded + 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 +978,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 +1016,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 +1094,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 +1120,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 +1133,77 @@ 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, and the upstream + refresh token (unwrapped from the client's refresh envelope) to exchange with the upstream IdP.""" + + ready: "_BridgeMintReady" + upstream_refresh_token: str + + +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( + request: Request, mcp_server: MCPServer, refresh_value: Optional[str] +) -> "_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. 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.get_secret_value(), + ) + + 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 +1221,69 @@ 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.""" + 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 + scope = token_response.get("scope") + lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in")) + return RefreshCredential( + refresh_token=SecretStr(refresh), + scope=scope if isinstance(scope, str) and scope else None, + expires_in=lifetime if isinstance(lifetime, int) else None, + ) + + +def _mint_refresh_envelope_value( + identity: "EnvelopeIdentity", token_response: object, keys: "EnvelopeKeys", now: datetime, mcp_server: MCPServer +) -> str | None: + """Seal the upstream refresh grant (if any) into a refresh envelope and return its bearer string, or + ``None`` when the upstream returned no refresh token or the refresh token is too large to seal. A + too-large refresh token degrades to an access-only response (logged) rather than failing an exchange + that already succeeded upstream: the client simply re-authenticates when the access envelope expires.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import + build_bridge_refresh_token_response, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import + SealedEnvelope, + ) + + refresh_credential = _upstream_refresh_credential(token_response) + if refresh_credential is None: + return None + sealed = build_bridge_refresh_token_response(identity, refresh_credential, keys, now) + if isinstance(sealed, SealedEnvelope): + return sealed.token.get_secret_value() + verbose_logger.warning( + "bridge mint: the upstream refresh token is too large to seal into a refresh envelope for " + "server=%s; issuing an access-only response, so the client re-authenticates at access expiry", + mcp_server.server_id, + ) + return None + + async def exchange_token_with_server( request: Request, mcp_server: MCPServer, @@ -1124,15 +1319,31 @@ 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: 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(request, 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 + # 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 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: @@ -1175,17 +1386,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, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py index 5530fbc46fd..e6506ce27b8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -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.""" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index 783e64d13e2..9118a3e129d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -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() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index ecea86bbed4..37fe5704db5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -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,72 @@ 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. It is not an access envelope, so the + admission consumer returns NotBridgeEnvelope, which admission fails closed (401): a refresh + credential can only ever be presented back to the token endpoint.""" + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = resolve_bridge_envelope(_sealed_refresh(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, NotBridgeEnvelope) + + +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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py index 7a2b51c2a95..ae196c9080b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -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", [ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 966619ee6b8..0505268b2b9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -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,313 @@ 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 +): + """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).""" + 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)), 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_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() + + @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