mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
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.
This commit is contained in:
Tin Chi Lo
parent
c0ff81a947
commit
c4dd06a0bb
6 changed files with 998 additions and 84 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,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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue