diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ebd5f43bf02..873c3d1baec 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -46,7 +46,6 @@ 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,9 +59,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( register_aggregate_client, relative_request_url, revoke_refresh_token, + supported_grant_types, ) from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( exchange_idp_subject_token, + token_exchange_available, ) from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, @@ -2142,7 +2143,9 @@ async def introspect_endpoint(token: str = Form(...)) -> Response: async def native_client_auth_discovery(request: Request) -> JSONResponse: """The versioned contract a native client (``lite login --pkce``, or a CLI in any other language) reads to sign a user in through the browser and obtain a proxy credential.""" - return JSONResponse(native_client_auth_contract(request), headers=TOKEN_NO_CACHE_HEADERS) + return JSONResponse( + native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS + ) # Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request @@ -2630,7 +2633,7 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict: } -def _build_aggregate_authorization_server_response(request: Request) -> dict: +def _build_aggregate_authorization_server_response(request: Request, token_exchange_available: bool) -> dict: """RFC 8414 metadata for the gateway as the aggregate authorization server. The issuer is ``{base}/mcp`` and must stay equal to the value the @@ -2649,7 +2652,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", TOKEN_EXCHANGE_GRANT_TYPE), + "grant_types_supported": supported_grant_types(token_exchange_available), "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], } @@ -2687,7 +2690,7 @@ async def oauth_authorization_server_aggregate(request: Request): per-server row win here instead would serve an issuer of {base} against a resource that advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door. """ - return _build_aggregate_authorization_server_response(request) + return _build_aggregate_authorization_server_response(request, token_exchange_available()) # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} @@ -2913,7 +2916,9 @@ async def register_client(request: Request, mcp_server_name: str | None = None): # advertises that), so this does not affect it. A request without redirect_uris is not # a DCR request, so the legacy single-server-or-dummy fallback is kept for it. if data.get("redirect_uris"): - return await register_aggregate_client(request=request, request_body=data) + return await register_aggregate_client( + request=request, request_body=data, token_exchange_available=token_exchange_available() + ) resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 9bdde3c5edc..ba24e861d6e 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -188,6 +188,17 @@ class MintProxyCredential(Protocol): TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" + + +def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]: + """The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed + only where the JWT auth that proves a subject token is on, backed by a database, and + licensed, so a client never selects a grant the gateway would then refuse.""" + if token_exchange_available: + return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE) + return ("authorization_code", "refresh_token") + + """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.""" @@ -359,7 +370,9 @@ def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None: return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY) -async def register_aggregate_client(request: Request, request_body: Mapping[str, object]) -> Response: +async def register_aggregate_client( + request: Request, request_body: Mapping[str, object], token_exchange_available: bool +) -> Response: """RFC 7591 dynamic registration against the gateway itself, statelessly. Only ``redirect_uris`` is authoritative; every client is registered as a public @@ -423,7 +436,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", TOKEN_EXCHANGE_GRANT_TYPE], + "grant_types": list(supported_grant_types(token_exchange_available)), "response_types": ["code"], }, ) @@ -621,7 +634,7 @@ class NativeClientAuthContract(TypedDict): revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]] -def native_client_auth_contract(request: Request) -> NativeClientAuthContract: +def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract: """The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a native client (in any language) needs to run the sign-in without reading LiteLLM source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter @@ -636,7 +649,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", TOKEN_EXCHANGE_GRANT_TYPE), + "grant_types_supported": supported_grant_types(token_exchange_available), "code_challenge_methods_supported": ("S256",), "token_endpoint_auth_methods_supported": ("none",), "revocation_endpoint_auth_methods_supported": ("none",), diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index 69b81713810..cdefaf76d49 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -6,15 +6,69 @@ 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 dataclasses import dataclass from typing import Final, Protocol from fastapi import HTTPException, Request +from litellm._logging import verbose_proxy_logger 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" +REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" + + +@dataclass(frozen=True, slots=True) +class TokenExchangePrerequisites: + """The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT + bearer. Discovery and registration advertise the exchange grant only when every one of + them holds, and an exchange attempt is refused naming the first one that does not.""" + + jwt_auth_enabled: bool + has_database: bool + licensed: bool + + @property + def available(self) -> bool: + return self.jwt_auth_enabled and self.has_database and self.licensed + + def refusal(self) -> SubjectTokenRefusal | None: + if not self.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 self.has_database: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no database, so it cannot exchange IdP tokens", + ) + if not self.licensed: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is an enterprise only feature; no license is set", + ) + return None + + +def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + general_settings, + premium_user, + prisma_client, + ) + + return TokenExchangePrerequisites( + jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, + has_database=prisma_client is not None, + licensed=premium_user is True, + ) + + +def token_exchange_available() -> bool: + return read_token_exchange_prerequisites().available class AuthorizeSubjectToken(Protocol): @@ -31,7 +85,6 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su 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, @@ -55,9 +108,7 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su 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, + prerequisites=read_token_exchange_prerequisites(), is_jwt=jwt_handler.is_jwt, authorize=authorize, ) @@ -66,40 +117,34 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su async def identity_from_subject_token( subject_token: str, request_headers: Mapping[str, str], - jwt_auth_enabled: bool, - has_database: bool, - licensed: bool, + prerequisites: TokenExchangePrerequisites, 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", - ) + RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token. The + reason stays in the proxy log: this endpoint is public and JWT auth's own wording can + name the JWKS URL it fetched or quote the IdP's response.""" + unmet: Final = prerequisites.refusal() + if unmet is not None: + return unmet 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}") + return _rejected_by_jwt_auth(denied.detail) except ProxyException as denied: - return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.message}") + return _rejected_by_jwt_auth(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}") + return _rejected_by_jwt_auth(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"]) + + +def _rejected_by_jwt_auth(reason: object) -> SubjectTokenRefusal: + verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason) + return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) diff --git a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py index 27d0ebbd5e6..a34119edf10 100644 --- a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py @@ -16,7 +16,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( ReloadUserFailure, ) from litellm.proxy._types import LiteLLM_UserTable -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, effective_user_role from litellm.proxy.management_endpoints.ui_sso import ( CliSsoTeamDetail, fetch_cli_sso_team_details, @@ -51,12 +51,12 @@ async def mint_proxy_credential( posting the consent form without one. Memberships whose team rows are gone count as no team at all, the way ``lite login`` treats them, so they can never lock a user out. The user row handed to the minter carries no team list, exactly like ``lite login``'s, so - the minter's own first-team fallback stays inert.""" + the minter's own first-team fallback stays inert. The credential carries the role the + proxy already enforces for the user on every request, so a row with no role (JWT auth's + upsert writes none) mints as an internal user instead of being refused.""" user: Final = await load_active_user_by_id(user_id) if isinstance(user, str): return user - if user.user_role is None: - return "no_active_key" if team_id is not None and team_id not in user.teams: return "not_a_member" details: Final = await _team_details(user.teams) if user.teams else () @@ -68,7 +68,9 @@ async def mint_proxy_credential( if selected is None: return "not_a_member" key: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info=LiteLLM_UserTable(user_id=user.user_id, user_role=user.user_role, models=user.models), + user_info=LiteLLM_UserTable( + user_id=user.user_id, user_role=effective_user_role(user.user_role).value, models=user.models + ), team_id=team_id, team_alias=selected.team_alias, team_models=selected.team_models, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index ba68dc8a17f..d90553db72b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1207,21 +1207,19 @@ async def common_checks( return True +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + try: + return LitellmUserRoles(user_role) + except ValueError: + return LitellmUserRoles.INTERNAL_USER + + def _get_user_role( user_obj: LiteLLM_UserTable | None, ) -> LitellmUserRoles | None: if user_obj is None: return None - - _user: Final = user_obj - - _user_role: Final = _user.user_role - try: - role: Final = LitellmUserRoles(_user_role) - except ValueError: - return LitellmUserRoles.INTERNAL_USER - - return role + return effective_user_role(user_obj.user_role) def _is_api_route_allowed( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index aa45b2f6793..e4edb5fb4dd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11048,6 +11048,25 @@ def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(mo assert stranger.json()["error"] == "invalid_client" +@pytest.mark.parametrize("exchange_servable", [True, False]) +def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it(monkeypatch, exchange_servable): + """Every document a native client reads before it picks a grant (the versioned contract, the + aggregate authorization-server metadata, and the registration response) lists the RFC 8693 + exchange exactly when the running proxy can serve it: JWT auth on, a database, and a license.""" + client, _session_cookie, _minted = _native_client_app(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": exchange_servable}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else [] + expected = ["authorization_code", "refresh_token", *exchange_grant] + + assert client.get("/.well-known/litellm-cli-auth").json()["grant_types_supported"] == expected + assert client.get("/.well-known/oauth-authorization-server/mcp").json()["grant_types_supported"] == expected + registered = client.post("/register", json={"redirect_uris": ["http://127.0.0.1:51234/callback"]}) + assert registered.status_code == 201 + assert registered.json()["grant_types"] == expected + + def test_native_client_authorize_without_the_proxy_resource_keeps_the_mcp_flow(monkeypatch): """A registered client asking for the MCP resource (or no resource) never sees the consent page, so existing MCP clients are untouched by the native-client arm.""" 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 1be6ebb6e22..c42f8763e74 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 @@ -95,9 +95,11 @@ def _request(path="/authorize", query="", cookies=None, method="GET"): ) -async def _register(redirect_uris) -> dict: +async def _register(redirect_uris, token_exchange_available=True) -> dict: response = await register_aggregate_client( - request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris} + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": redirect_uris}, + token_exchange_available=token_exchange_available, ) return json.loads(response.body) @@ -119,11 +121,18 @@ async def test_register_mints_stateless_public_client(): assert record.redirect_uris == (REDIRECT_URI,) +@pytest.mark.asyncio +async def test_register_omits_the_exchange_grant_where_the_gateway_cannot_serve_it(): + body = await _register([REDIRECT_URI], token_exchange_available=False) + assert body["grant_types"] == ["authorization_code", "refresh_token"] + + @pytest.mark.asyncio @pytest.mark.parametrize("redirect_uris", [VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[str, ...]) -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={ "client_name": "Visual Studio Code", "client_uri": "https://code.visualstudio.com", @@ -149,6 +158,7 @@ async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[s async def test_register_rejects_five_valid_callbacks() -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={"redirect_uris": [*VSCODE_REDIRECT_URIS, "http://127.0.0.1:33419/"]}, ) assert response.status_code == 400 @@ -162,6 +172,7 @@ async def test_register_rejects_five_valid_callbacks() -> None: async def test_register_four_callbacks_preserves_encoded_size_guard() -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={"redirect_uris": [f"https://client.example/{index}/".ljust(256, "é") for index in range(4)]}, ) assert response.status_code == 400 @@ -214,6 +225,7 @@ async def test_register_rejects_userinfo_spoofed_origin(): response = await register_aggregate_client( request=_request(path="/register", method="POST"), request_body={"redirect_uris": ["https://claude.ai@attacker.example/callback"]}, + token_exchange_available=True, ) assert response.status_code == 400 assert json.loads(response.body)["error"] == "invalid_redirect_uri" @@ -234,7 +246,9 @@ async def test_register_rejects_userinfo_spoofed_origin(): ) async def test_register_rejects_bad_redirect_uris(redirect_uris): response = await register_aggregate_client( - request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris} + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": redirect_uris}, + token_exchange_available=True, ) assert response.status_code == 400 assert json.loads(response.body)["error"] in ("invalid_redirect_uri", "invalid_client_metadata") @@ -1954,7 +1968,7 @@ async def test_revoke_refuses_unknown_clients_and_a_missing_master_key(): def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): - assert json.loads(json.dumps(native_client_auth_contract(_request("/.well-known/litellm-cli-auth")))) == { + assert json.loads(json.dumps(native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), True))) == { "contract_version": 1, "issuer": "https://llm.example.com", "authorization_endpoint": "https://llm.example.com/authorize", @@ -1974,6 +1988,11 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): } +def test_native_client_auth_contract_omits_the_exchange_grant_where_the_gateway_cannot_serve_it(): + contract = native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), False) + assert list(contract["grant_types_supported"]) == ["authorization_code", "refresh_token"] + + @pytest.mark.parametrize( "resource, expected", [ 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 index 00440614120..d1b049dddd5 100644 --- 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 @@ -1,13 +1,22 @@ +import logging + 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._experimental.mcp_server.idp_token_exchange import ( + REJECTED_SUBJECT_TOKEN, + TokenExchangePrerequisites, + identity_from_subject_token, + token_exchange_available, +) 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"} +EVERY_GATE_HOLDS = {"jwt_auth_enabled": True, "has_database": True, "licensed": True} +JWKS_URL = "https://idp.example.com/.well-known/jwks.json" def _authorized(user_id="u1", team_id="team-b"): @@ -42,16 +51,14 @@ class _Authorizer: 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}) +async def _identity(authorizer, subject_token=IDP_JWT, **unmet): + return await identity_from_subject_token( + subject_token, + request_headers=REQUEST_HEADERS, + prerequisites=TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}), + is_jwt=JWTHandler.is_jwt, + authorize=authorizer, + ) @pytest.mark.asyncio @@ -70,7 +77,7 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): @pytest.mark.asyncio @pytest.mark.parametrize( - "overrides, subject_token, error, mentions", + "unmet, 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"), @@ -79,32 +86,59 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): ], ) async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verification( - overrides, subject_token, error, mentions + unmet, subject_token, error, mentions ): authorizer = _Authorizer() - refusal = await _identity(authorizer, subject_token=subject_token, **overrides) + refusal = await _identity(authorizer, subject_token=subject_token, **unmet) assert isinstance(refusal, SubjectTokenRefusal) assert refusal.error == error assert mentions in refusal.description assert authorizer.calls == [] +@pytest.mark.parametrize("unmet", [{}, {"jwt_auth_enabled": False}, {"has_database": False}, {"licensed": False}]) +def test_the_grant_is_available_exactly_when_every_gate_holds(unmet): + prerequisites = TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}) + assert prerequisites.available is (unmet == {}) + assert (prerequisites.refusal() is None) is prerequisites.available + + +@pytest.mark.parametrize( + "general_settings, prisma_client, premium_user, expected", + [ + ({"enable_jwt_auth": True}, object(), True, True), + ({}, object(), True, False), + ({"enable_jwt_auth": True}, None, True, False), + ({"enable_jwt_auth": True}, object(), False, False), + ], +) +def test_availability_is_read_from_the_running_proxy( + monkeypatch, general_settings, prisma_client, premium_user, expected +): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user) + assert token_exchange_available() is expected + + @pytest.mark.asyncio @pytest.mark.parametrize( - "raised, mentions", + "raised, reason", [ (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"), + (Exception(f"Failed to fetch keys from {JWKS_URL}: 502 Bad Gateway from the IdP"), JWKS_URL), ], ) -async def test_a_jwt_the_proxy_rejects_is_an_invalid_subject_token(raised, mentions): +async def test_a_jwt_the_proxy_rejects_is_refused_with_the_reason_kept_in_the_log(raised, reason, caplog): + """The endpoint is public, so the response never quotes JWT auth's wording (it can name + the JWKS URL or relay the IdP's reply); the operator reads the reason in the proxy log.""" + caplog.set_level(logging.WARNING, logger="LiteLLM Proxy") 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 + assert refusal == SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) + assert reason in caplog.text @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 8bb8bdada7d..85650c6a05a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -8,6 +8,7 @@ from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.models.user import LiteLLM_UserTable from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, MintedProxyCredential from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential +from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail @@ -67,10 +68,21 @@ async def test_mint_passes_user_lookup_failures_through(failure, load_user, fetc @pytest.mark.asyncio -async def test_mint_refuses_a_user_without_a_role(load_user, fetch_teams): - load_user.return_value = _user(user_role=None) - assert await mint_proxy_credential("u1", None) == "no_active_key" - fetch_teams.assert_not_awaited() +@pytest.mark.parametrize( + "stored_role, minted_role", + [ + (None, LitellmUserRoles.INTERNAL_USER), + ("made_up_role", LitellmUserRoles.INTERNAL_USER), + ("proxy_admin", LitellmUserRoles.PROXY_ADMIN), + ], +) +async def test_mint_carries_the_role_the_proxy_enforces_for_the_user(load_user, fetch_teams, stored_role, minted_role): + """A user JWT auth upserted has no role in the database, and the proxy already treats + such a user as an internal user on every request, so the credential says the same.""" + load_user.return_value = _user(user_role=stored_role) + minted = await mint_proxy_credential("u1", "team-a") + assert isinstance(minted, MintedProxyCredential) + assert _decoded(minted).user_role == minted_role @pytest.mark.asyncio