diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ffb27d5f92e..ebd5f43bf02 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + TOKEN_EXCHANGE_GRANT_TYPE, VendorCredentialState, aggregate_authorize, aggregate_token, @@ -60,6 +61,9 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( relative_request_url, revoke_refresh_token, ) +from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( + exchange_idp_subject_token, +) from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, RefreshTokenPresented, @@ -1980,6 +1984,9 @@ async def token_endpoint( refresh_token: str | None = Form(None), scope: str | None = Form(None), resource: str | None = Form(None), + subject_token: str | None = Form(None), + subject_token_type: str | None = Form(None), + requested_token_type: str | None = Form(None), mcp_server_name: str | None = None, ): """ @@ -2010,6 +2017,10 @@ async def token_endpoint( cache=user_api_key_cache, resource=resource, mint_proxy_credential=mint_proxy_credential, + subject_token=subject_token, + subject_token_type=subject_token_type, + requested_token_type=requested_token_type, + exchange_subject_token=exchange_idp_subject_token, ) lookup_name: Final = mcp_server_name or client_id @@ -2638,7 +2649,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: "registration_endpoint": f"{request_base_url}/register", "response_types_supported": ["code"], "scopes_supported": [], - "grant_types_supported": ["authorization_code", "refresh_token"], + "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], } diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index f3fdd54b39d..9bdde3c5edc 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -51,7 +51,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ConfigDict, Field, ValidationError -from typing_extensions import ReadOnly, TypedDict, assert_never +from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._logging import verbose_logger from litellm.caching.caching import DualCache @@ -187,6 +187,41 @@ class MintProxyCredential(Protocol): ) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ... +TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" +"""RFC 8693: a native client that already holds a token from the customer's identity +provider trades it for the proxy-API credential without a browser round trip.""" + +_IssuedTokenType = Literal["urn:ietf:params:oauth:token-type:access_token"] +ACCESS_TOKEN_TOKEN_TYPE: Final[_IssuedTokenType] = "urn:ietf:params:oauth:token-type:access_token" +SUBJECT_TOKEN_TYPES: Final = frozenset( + { + "urn:ietf:params:oauth:token-type:jwt", + "urn:ietf:params:oauth:token-type:id_token", + ACCESS_TOKEN_TOKEN_TYPE, + } +) + + +class SubjectIdentity(BaseModel): + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + team_id: str | None = None + + +class SubjectTokenRefusal(BaseModel): + model_config = ConfigDict(frozen=True) + error: Literal["unsupported_grant_type", "invalid_request"] + description: str = Field(min_length=1) + + +class ExchangeSubjectToken(Protocol): + """Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the + IdP token the way the proxy's own JWT auth does and names the litellm user and team it + stands for, or says why this gateway will not take it.""" + + def __call__(self, subject_token: str, request: Request, /) -> Awaitable[SubjectIdentity | SubjectTokenRefusal]: ... + + class ConsentTeam(BaseModel): model_config = ConfigDict(frozen=True) team_id: str = Field(min_length=1) @@ -213,6 +248,12 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr return "unresolvable" +async def _refuse_subject_token(subject_token: str, request: Request) -> SubjectTokenRefusal: + return SubjectTokenRefusal( + error="unsupported_grant_type", description="this gateway is not configured to exchange IdP tokens" + ) + + async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: return "unavailable" @@ -382,7 +423,7 @@ async def register_aggregate_client(request: Request, request_body: Mapping[str, "client_id_issued_at": int(now.timestamp()), "redirect_uris": list(raw_uris), "token_endpoint_auth_method": "none", - "grant_types": ["authorization_code", "refresh_token"], + "grant_types": ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE], "response_types": ["code"], }, ) @@ -595,7 +636,7 @@ def native_client_auth_contract(request: Request) -> NativeClientAuthContract: "revocation_endpoint": f"{base_url}/revoke", "resource": base_url, "response_types_supported": ("code",), - "grant_types_supported": ("authorization_code", "refresh_token"), + "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), "code_challenge_methods_supported": ("S256",), "token_endpoint_auth_methods_supported": ("none",), "revocation_endpoint_auth_methods_supported": ("none",), @@ -1033,20 +1074,26 @@ class _ProxyCredentialTokenResponse(TypedDict): refresh_token: ReadOnly[str] user_id: ReadOnly[str] team_id: ReadOnly[str | None] + issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]] def _proxy_credential_response( - minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime + minted: MintedProxyCredential, + principal: SessionPrincipal, + keys: SessionSigningKeys, + now: datetime, + issued_token_type: _IssuedTokenType | None = None, ) -> Response: """The proxy-API token response: the access token is the very credential ``lite login`` stores (accepted on every proxy route with user and team attribution), and the refresh token is a gateway-sealed rotating token bound to the team the credential - was minted for, so a renewal keeps the team the user consented to.""" + was minted for, so a renewal keeps the team the user consented to. A token exchange + also states ``issued_token_type``, which RFC 8693 section 2.2.1 requires.""" bound_principal: Final = principal.model_copy(update=MappingProxyType({"team_id": minted.team_id})) refresh: Final = mint_session_refresh_token(bound_principal, keys, now) if not isinstance(refresh, MintedSessionToken): return _oauth_error(500, "server_error", "failed to mint the session credential") - body: Final[_ProxyCredentialTokenResponse] = { + credential: Final[_ProxyCredentialTokenResponse] = { "access_token": minted.key, "token_type": "Bearer", "expires_in": minted.expires_in, @@ -1054,7 +1101,10 @@ def _proxy_credential_response( "user_id": minted.user_id, "team_id": minted.team_id, } - return JSONResponse(status_code=200, content=body, headers=TOKEN_NO_CACHE_HEADERS) + if issued_token_type is None: + return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS) + exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type} + return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS) def _reload_failure_response(failure: ReloadUserFailure) -> Response: @@ -1116,11 +1166,16 @@ async def aggregate_token( cache: DualCache, resource: str | None = None, mint_proxy_credential: MintProxyCredential = _refuse_proxy_credential, + subject_token: str | None = None, + subject_token_type: str | None = None, + requested_token_type: str | None = None, + exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token, ) -> Response: """The aggregate token verb: authorization_code and refresh_token grants for the identity-only session pair, or for the proxy-API credential when the grant was issued - with that audience. Every path re-validates the litellm user live before minting, so a - deactivated user cannot obtain or renew a session.""" + with that audience, and the RFC 8693 token exchange that turns an IdP token straight + into the proxy-API credential. Every path re-validates the litellm user live before + minting, so a deactivated user cannot obtain or renew a session.""" if master_key is None: verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured") return _oauth_error(500, "server_error", "the gateway has no master key configured") @@ -1159,7 +1214,20 @@ async def aggregate_token( now=now, issue=issue, ) - return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token") + if grant_type == TOKEN_EXCHANGE_GRANT_TYPE: + return await _token_exchange_grant( + subject_token=subject_token, + subject_token_type=subject_token_type, + requested_token_type=requested_token_type, + client_id=client_id, + exchange_subject_token=exchange_subject_token, + issue=issue, + ) + return _oauth_error( + 400, + "unsupported_grant_type", + f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}", + ) class _GrantIssuer: @@ -1211,10 +1279,9 @@ class _GrantIssuer: async def _issue_proxy_credential( self, principal: SessionPrincipal, claim_key: str, claim_ttl_seconds: int, replayed: str ) -> Response: - if self._resource is not None and not is_proxy_api_resource(self._request, self._resource): - return _oauth_error( - 400, "invalid_target", "resource does not match the proxy API this grant was issued for" - ) + target_refusal: Final = self._proxy_api_target_refusal() + if target_refusal is not None: + return target_refusal minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id) if not isinstance(minted, MintedProxyCredential): return _mint_failure_response(minted) @@ -1223,6 +1290,33 @@ class _GrantIssuer: return refusal return _proxy_credential_response(minted, principal, self._keys, self._now) + async def exchange( + self, subject_token: str, client_id: str, exchange_subject_token: ExchangeSubjectToken + ) -> Response: + """The RFC 8693 tail: prove the IdP token, then mint. No single-use marker, because + the subject token stays a valid proof for as long as the IdP says it is and every + exchange mints a fresh credential and refresh token of its own.""" + target_refusal: Final = self._proxy_api_target_refusal() + if target_refusal is not None: + return target_refusal + identity: Final = await exchange_subject_token(subject_token, self._request) + if isinstance(identity, SubjectTokenRefusal): + return _oauth_error(400, identity.error, identity.description) + principal: Final = SessionPrincipal( + user_id=identity.user_id, client_id=client_id, audience=PROXY_API_AUDIENCE, team_id=identity.team_id + ) + minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id) + if not isinstance(minted, MintedProxyCredential): + return _mint_failure_response(minted) + return _proxy_credential_response( + minted, principal, self._keys, self._now, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE + ) + + def _proxy_api_target_refusal(self) -> Response | None: + if self._resource is None or is_proxy_api_resource(self._request, self._resource): + return None + return _oauth_error(400, "invalid_target", "resource does not match the proxy API this grant was issued for") + async def _claim_refusal(self, claim_key: str, claim_ttl_seconds: int, replayed: str) -> Response | None: return _claim_refusal( await self._guard.claim(claim_key, claim_ttl_seconds), replayed=_oauth_error(400, "invalid_grant", replayed) @@ -1297,6 +1391,32 @@ async def _refresh_token_grant( ) +async def _token_exchange_grant( + subject_token: str | None, + subject_token_type: str | None, + requested_token_type: str | None, + client_id: str, + exchange_subject_token: ExchangeSubjectToken, + issue: _GrantIssuer, +) -> Response: + """RFC 8693 token exchange for a registered native client that already holds an IdP + token: the gateway proves the token the way its JWT auth does and answers with the + proxy-API credential, so a fresh laptop with only an IdP login gets a gateway key + without a browser round trip. The client must be registered because the refresh token + in the answer is bound to it.""" + if not is_gateway_dcr_client_id(client_id) or open_gateway_dcr_client(client_id) is None: + return _oauth_error(401, "invalid_client", "unknown or malformed client_id") + if not subject_token or not subject_token_type: + return _oauth_error(400, "invalid_request", "subject_token and subject_token_type are required") + if subject_token_type not in SUBJECT_TOKEN_TYPES: + return _oauth_error( + 400, "invalid_request", f"subject_token_type must be one of {', '.join(sorted(SUBJECT_TOKEN_TYPES))}" + ) + if requested_token_type is not None and requested_token_type != ACCESS_TOKEN_TOKEN_TYPE: + return _oauth_error(400, "invalid_request", f"requested_token_type must be {ACCESS_TOKEN_TOKEN_TYPE}") + return await issue.exchange(subject_token, client_id, exchange_subject_token) + + async def revoke_refresh_token(token: str, client_id: str, master_key: str | None, cache: DualCache) -> Response: """RFC 7009 revocation for the gateway's refresh tokens: burn the presented token's ``jti`` so neither the holder nor a thief can rotate it again. Access tokens are diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py new file mode 100644 index 00000000000..69b81713810 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -0,0 +1,105 @@ +"""The identity-provider side of the RFC 8693 token exchange on ``POST /token``: a native +client that already holds a JWT from the customer's IdP trades it for the same proxy-API +credential ``lite login`` stores, proven by the proxy's own JWT auth (signature, claims, +and the user and team sync it performs), so no browser round trip is needed.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from typing import Final, Protocol + +from fastapi import HTTPException, Request + +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal +from litellm.proxy._types import JWTAuthBuilderResult, ProxyException +from litellm.proxy.auth.handle_jwt import JWTAuthManager + +EXCHANGE_ROUTE: Final = "/token" + + +class AuthorizeSubjectToken(Protocol): + """Injected JWT authorization ``(subject_token, request_headers)``: the proxy's + ``JWTAuthManager.auth_builder`` in production, which raises when the token is not + acceptable and otherwise names the user and team it resolved.""" + + def __call__( + self, subject_token: str, request_headers: Mapping[str, str], / + ) -> Awaitable[JWTAuthBuilderResult]: ... + + +async def exchange_idp_subject_token(subject_token: str, request: Request) -> SubjectIdentity | SubjectTokenRefusal: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + general_settings, + jwt_handler, + premium_user, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + async def authorize(token: str, request_headers: Mapping[str, str]) -> JWTAuthBuilderResult: + return await JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=jwt_handler, + request_data={}, + general_settings=general_settings, + route=EXCHANGE_ROUTE, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + request_headers=request_headers, + request_method="POST", + ) + + return await identity_from_subject_token( + subject_token, + request_headers=request.headers, + jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, + has_database=prisma_client is not None, + licensed=premium_user is True, + is_jwt=jwt_handler.is_jwt, + authorize=authorize, + ) + + +async def identity_from_subject_token( + subject_token: str, + request_headers: Mapping[str, str], + jwt_auth_enabled: bool, + has_database: bool, + licensed: bool, + is_jwt: Callable[[str], bool], + authorize: AuthorizeSubjectToken, +) -> SubjectIdentity | SubjectTokenRefusal: + """Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the + proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which + RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token.""" + if not jwt_auth_enabled: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens", + ) + if not has_database: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no database, so it cannot exchange IdP tokens", + ) + if not is_jwt(subject_token): + return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT") + if not licensed: + return SubjectTokenRefusal( + error="unsupported_grant_type", description="JWT auth is an enterprise only feature; no license is set" + ) + try: + result: Final = await authorize(subject_token, request_headers) + except HTTPException as denied: + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.detail}") + except ProxyException as denied: + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.message}") + except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied}") + user_id: Final = result["user_id"] + if user_id is None: + return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows") + return SubjectIdentity(user_id=user_id, team_id=result["team_id"]) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 74f38b3ca6d..6893c06d2d7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19346,7 +19346,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { @@ -23115,6 +23115,17 @@ ], "title": "Refresh Token" }, + "requested_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Requested Token Type" + }, "resource": { "anyOf": [ { @@ -23136,6 +23147,28 @@ } ], "title": "Scope" + }, + "subject_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" } }, "required": [ @@ -23189,6 +23222,17 @@ ], "title": "Refresh Token" }, + "requested_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Requested Token Type" + }, "resource": { "anyOf": [ { @@ -23210,6 +23254,28 @@ } ], "title": "Scope" + }, + "subject_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" } }, "required": [ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 63d0bfcc5b8..5f7365cedd5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -525,6 +525,7 @@ class LiteLLMRoutes(enum.Enum): "/mcp-rest/tools/call", "/v1/mcp/tools", "/introspect", + "/token", ] # MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS. diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 6a28cd7ff99..803093ff93a 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1867,7 +1867,7 @@ class JWTAuthManager: @staticmethod def get_team_id_from_header( - request_headers: dict | None, + request_headers: Mapping[str, str] | None, allowed_team_ids: set[str], fallback_to_db_teams: bool = False, ) -> str | None: @@ -2037,7 +2037,7 @@ class JWTAuthManager: async def _attach_team_from_header_for_admin( admin_result: JWTAuthBuilderResult, route: str, - request_headers: dict | None, + request_headers: Mapping[str, str] | None, jwt_handler: JWTHandler, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, @@ -2293,7 +2293,7 @@ class JWTAuthManager: user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging, - request_headers: dict | None = None, + request_headers: Mapping[str, str] | None = None, request_method: str | None = None, ) -> JWTAuthBuilderResult: return await JWTAuthManager.authorize_jwt( @@ -2390,7 +2390,7 @@ class JWTAuthManager: user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging, - request_headers: dict[str, str] | None = None, + request_headers: Mapping[str, str] | None = None, request_method: str | None = None, provisioning: _JWTProvisioning | None = None, ) -> JWTAuthBuilderResult: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 7c80ee77cd7..1be6ebb6e22 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -15,13 +15,18 @@ from starlette.requests import Request from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( _AUTH_CODE_DEBUG_KEY, + ACCESS_TOKEN_TOKEN_TYPE, CONNECT_FLOW_COOKIE_PREFIX, GATEWAY_AUTH_CODE_PREFIX, GATEWAY_AUTH_CODE_TTL_SECONDS, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, MAX_CLIENT_ID_LENGTH, + SUBJECT_TOKEN_TYPES, + TOKEN_EXCHANGE_GRANT_TYPE, ConsentTeam, MintedProxyCredential, + SubjectIdentity, + SubjectTokenRefusal, _GatewayAuthCode, _open_sealed, _seal, @@ -105,6 +110,7 @@ async def _reload_user_active(user_id: str): async def test_register_mints_stateless_public_client(): body = await _register([REDIRECT_URI]) assert body["token_endpoint_auth_method"] == "none" + assert body["grant_types"] == ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE] assert "client_secret" not in body assert body["redirect_uris"] == [REDIRECT_URI] assert is_gateway_dcr_client_id(body["client_id"]) @@ -1957,7 +1963,11 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): "revocation_endpoint": "https://llm.example.com/revoke", "resource": "https://llm.example.com", "response_types_supported": ["code"], - "grant_types_supported": ["authorization_code", "refresh_token"], + "grant_types_supported": [ + "authorization_code", + "refresh_token", + "urn:ietf:params:oauth:grant-type:token-exchange", + ], "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none"], "revocation_endpoint_auth_methods_supported": ["none"], @@ -2148,3 +2158,177 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a ) assert renewed.status_code == 200 assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id" + + +JWT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +IDP_TOKEN = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" + + +class _Exchanger: + def __init__(self, result=None): + self.calls = [] + self.result = result + + async def __call__(self, subject_token, request): + self.calls.append((subject_token, request.url.path)) + if self.result is not None: + return self.result + return SubjectIdentity(user_id="u1", team_id="team-b") + + +async def _exchange_native(client_id, minter, exchanger, cache=None, **overrides): + arguments = { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "subject_token": IDP_TOKEN, + "subject_token_type": JWT_SUBJECT_TOKEN_TYPE, + "exchange_subject_token": exchanger, + } + return await _redeem_native(None, client_id, minter, cache=cache, **{**arguments, **overrides}) + + +@pytest.mark.asyncio +async def test_token_exchange_mints_the_proxy_credential_for_the_idp_subject(): + """RFC 8693: a registered native client trades the IdP token it already holds for the + same credential the consent flow mints, attributed to the user and team the gateway's + JWT auth resolved, with a rotating refresh token bound to that team and the client. + The exchange can be repeated while the IdP token lives; nothing is burned.""" + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, exchanger, cache = _Minter(), _Exchanger(), DualCache() + response = await _exchange_native(client_id, minter, exchanger, cache=cache) + assert response.status_code == 200 + assert response.headers["cache-control"] == "no-store" + body = json.loads(response.body) + assert exchanger.calls == [(IDP_TOKEN, "/token")] + assert minter.calls == [("u1", "team-b")] + assert body["issued_token_type"] == ACCESS_TOKEN_TOKEN_TYPE + assert body["access_token"] == "sk-cli-u1" + assert body["token_type"] == "Bearer" + assert body["expires_in"] == 3600 + assert (body["user_id"], body["team_id"]) == ("u1", "team-b") + principal = _opened_refresh(body["refresh_token"], client_id) + assert (principal.user_id, principal.client_id, principal.audience, principal.team_id) == ( + "u1", + client_id, + "proxy_api", + "team-b", + ) + again = await _exchange_native(client_id, minter, exchanger, cache=cache) + assert again.status_code == 200 + assert json.loads(again.body)["refresh_token"] != body["refresh_token"] + assert minter.calls == [("u1", "team-b"), ("u1", "team-b")] + + +@pytest.mark.asyncio +async def test_exchanged_credential_refreshes_and_rotates_like_a_consented_one(): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, cache = _Minter(), DualCache() + exchanged = json.loads((await _exchange_native(client_id, minter, _Exchanger(), cache=cache)).body) + refreshed = await _refresh_native(exchanged["refresh_token"], client_id, minter, cache) + assert refreshed.status_code == 200 + body = json.loads(refreshed.body) + assert "issued_token_type" not in body + assert (body["access_token"], body["user_id"], body["team_id"]) == ("sk-cli-u1", "u1", "team-b") + assert body["refresh_token"] != exchanged["refresh_token"] + assert minter.calls == [("u1", "team-b"), ("u1", "team-b")] + replay = await _refresh_native(exchanged["refresh_token"], client_id, minter, cache) + assert replay.status_code == 400 + assert json.loads(replay.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_token_exchange_for_a_teamless_subject_mints_a_teamless_credential(): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + response = await _exchange_native(client_id, minter, _Exchanger(SubjectIdentity(user_id="u2"))) + assert response.status_code == 200 + body = json.loads(response.body) + assert minter.calls == [("u2", None)] + assert (body["user_id"], body["team_id"]) == ("u2", None) + assert _opened_refresh(body["refresh_token"], client_id).team_id is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("subject_token_type", sorted(SUBJECT_TOKEN_TYPES)) +async def test_token_exchange_accepts_every_advertised_subject_token_type(subject_token_type): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + response = await _exchange_native(client_id, _Minter(), _Exchanger(), subject_token_type=subject_token_type) + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_token_exchange_without_an_idp_exchanger_is_unsupported(): + """A gateway that wires no IdP verifier into the endpoint answers the way it always + answered an unknown grant, and never reaches the minter.""" + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + response = await _redeem_native( + None, + client_id, + minter, + grant_type=TOKEN_EXCHANGE_GRANT_TYPE, + subject_token=IDP_TOKEN, + subject_token_type=JWT_SUBJECT_TOKEN_TYPE, + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "unsupported_grant_type" + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, status, error", + [ + ({"subject_token": None}, 400, "invalid_request"), + ({"subject_token": ""}, 400, "invalid_request"), + ({"subject_token_type": None}, 400, "invalid_request"), + ({"subject_token_type": "urn:ietf:params:oauth:token-type:saml2"}, 400, "invalid_request"), + ({"requested_token_type": "urn:ietf:params:oauth:token-type:refresh_token"}, 400, "invalid_request"), + ({"resource": "https://other.example.com"}, 400, "invalid_target"), + ({"resource": "https://llm.example.com/mcp"}, 400, "invalid_target"), + ({"client_id": "llm_dcrc_forged"}, 401, "invalid_client"), + ({"client_id": "not-a-gateway-client"}, 401, "invalid_client"), + ], +) +async def test_token_exchange_refuses_a_malformed_request_before_touching_the_idp_token(overrides, status, error): + registered = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, exchanger = _Minter(), _Exchanger() + response = await _exchange_native( + overrides.get("client_id", registered), + minter, + exchanger, + **{name: value for name, value in overrides.items() if name != "client_id"}, + ) + assert response.status_code == status + assert json.loads(response.body)["error"] == error + assert exchanger.calls == [] + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", ["unsupported_grant_type", "invalid_request"]) +async def test_token_exchange_relays_the_idp_refusal_and_never_mints(error): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + exchanger = _Exchanger(SubjectTokenRefusal(error=error, description="subject_token was rejected: bad signature")) + response = await _exchange_native(client_id, minter, exchanger) + assert response.status_code == 400 + body = json.loads(response.body) + assert (body["error"], body["error_description"]) == (error, "subject_token was rejected: bad signature") + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure, status, error", + [ + ("not_a_member", 400, "invalid_grant"), + ("team_required", 400, "invalid_grant"), + ("no_active_key", 400, "invalid_grant"), + ("unavailable", 503, "temporarily_unavailable"), + ], +) +async def test_token_exchange_relays_a_mint_refusal(failure, status, error): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + response = await _exchange_native(client_id, _Minter(failure), _Exchanger()) + assert response.status_code == status + assert json.loads(response.body)["error"] == error diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py new file mode 100644 index 00000000000..00440614120 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -0,0 +1,115 @@ +import pytest +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal +from litellm.proxy._experimental.mcp_server.idp_token_exchange import identity_from_subject_token +from litellm.proxy._types import ProxyException +from litellm.proxy.auth.handle_jwt import JWTHandler + +IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" +REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} + + +def _authorized(user_id="u1", team_id="team-b"): + return { + "is_proxy_admin": False, + "team_object": None, + "user_object": None, + "end_user_object": None, + "org_object": None, + "token": IDP_JWT, + "team_id": team_id, + "user_id": user_id, + "user_email": None, + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": user_id}, + "agent_id": None, + } + + +class _Authorizer: + def __init__(self, result=None, raises=None): + self.calls = [] + self.result = result if result is not None else _authorized() + self.raises = raises + + async def __call__(self, subject_token, request_headers): + self.calls.append((subject_token, dict(request_headers))) + if self.raises is not None: + raise self.raises + return self.result + + +async def _identity(authorizer, subject_token=IDP_JWT, **overrides): + arguments = { + "request_headers": REQUEST_HEADERS, + "jwt_auth_enabled": True, + "has_database": True, + "licensed": True, + "is_jwt": JWTHandler.is_jwt, + "authorize": authorizer, + } + return await identity_from_subject_token(subject_token, **{**arguments, **overrides}) + + +@pytest.mark.asyncio +async def test_a_jwt_the_proxy_accepts_names_its_user_and_team(): + """The subject token goes to the proxy's own JWT auth with the caller's headers (that is + where the team header is read), and the identity it resolved is what gets minted.""" + authorizer = _Authorizer() + assert await _identity(authorizer) == SubjectIdentity(user_id="u1", team_id="team-b") + assert authorizer.calls == [(IDP_JWT, REQUEST_HEADERS)] + + +@pytest.mark.asyncio +async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): + assert await _identity(_Authorizer(_authorized(team_id=None))) == SubjectIdentity(user_id="u1", team_id=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, subject_token, error, mentions", + [ + ({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"), + ({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"), + ({"licensed": False}, IDP_JWT, "unsupported_grant_type", "enterprise"), + ({}, "sk-litellm-virtual-key", "invalid_request", "not a JWT"), + ], +) +async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verification( + overrides, subject_token, error, mentions +): + authorizer = _Authorizer() + refusal = await _identity(authorizer, subject_token=subject_token, **overrides) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == error + assert mentions in refusal.description + assert authorizer.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raised, mentions", + [ + (HTTPException(status_code=403, detail="User not allowed to access this route"), "not allowed"), + (ProxyException(message="Token expired", type="auth_error", param="token", code=401), "Token expired"), + (Exception("Validation fails: signature verification failed"), "signature verification failed"), + (Exception("Invalid JWT Submitted"), "Invalid JWT"), + ], +) +async def test_a_jwt_the_proxy_rejects_is_an_invalid_subject_token(raised, mentions): + refusal = await _identity(_Authorizer(raises=raised)) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == "invalid_request" + assert refusal.description.startswith("subject_token was rejected: ") + assert mentions in refusal.description + + +@pytest.mark.asyncio +async def test_a_jwt_that_resolves_no_user_cannot_be_exchanged(): + refusal = await _identity(_Authorizer(_authorized(user_id=None))) + assert refusal == SubjectTokenRefusal( + error="invalid_request", description="subject_token names no user the gateway knows" + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26ae28a57d2..e4c03399b30 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8388,3 +8388,21 @@ async def test_access_group_model_fallback_uses_the_injected_database(channel: s llm_router=None, prisma_client=client, ) is True reader.assert_awaited_once_with(where={"access_group_id": "group-a"}) + + +def test_jwt_team_role_reaches_the_gateway_token_endpoint_by_default(): + """The RFC 8693 token exchange authorizes the IdP JWT against ``POST /token`` itself, and JWT + auth only binds a team from a multi-team claim when that team may call the route, so the + default team allowlist has to cover the gateway's token endpoint or the exchange would mint + teamless credentials for every ``team_ids_jwt_field`` deployment.""" + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + assert allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/token", litellm_proxy_roles=LiteLLM_JWTAuth() + ) + assert not allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route="/token", + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=[]), + ) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 806c55d51ce..630f69bf0f1 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -627,6 +627,7 @@ def test_virtual_key_llm_api_routes_denies_spend_logs_v2(): "/mcp/tools/call", "/mcp-rest/tools/call", "/mcp/tools/list", + "/token", ], ) def test_mcp_inference_routes_classified_as_llm_api(route): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 17ec8367324..b896f6901cb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24505,10 +24505,16 @@ export interface components { redirect_uri?: string; /** Refresh Token */ refresh_token?: string | null; + /** Requested Token Type */ + requested_token_type?: string | null; /** Resource */ resource?: string | null; /** Scope */ scope?: string | null; + /** Subject Token */ + subject_token?: string | null; + /** Subject Token Type */ + subject_token_type?: string | null; }; /** Body_token_endpoint_token_post */ Body_token_endpoint_token_post: { @@ -24526,10 +24532,16 @@ export interface components { redirect_uri?: string; /** Refresh Token */ refresh_token?: string | null; + /** Requested Token Type */ + requested_token_type?: string | null; /** Resource */ resource?: string | null; /** Scope */ scope?: string | null; + /** Subject Token */ + subject_token?: string | null; + /** Subject Token Type */ + subject_token_type?: string | null; }; /** Body_upload_logo_upload_logo_post */ Body_upload_logo_upload_logo_post: {