From 6f8a42d5177a73fa82effe8c6eedcd320cc2655d Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 28 Oct 2026 15:52:39 -0700 Subject: [PATCH] fix(mcp): vouch for delegate callers with OAuth2 introspection and proxy headers On delegate servers, a bearer that isn't LiteLLM's own was only tried as an identity through custom auth or JWT auth, so proxies using enable_oauth2_auth or enable_oauth2_proxy_auth admitted those callers anonymously and lost their budgets and attribution The identity step now picks the check LiteLLM's own auth runs for that token and route, in its order (OAuth2 introspection, then reverse-proxy headers, then JWT), and follows the JWT rules. A rejected identity is anonymous, a refusal after the caller is vouched for fails, and a granted user is admitted by name with the token still forwarded. The route rule for OAuth2 introspection moves into oauth2_introspection_applies so both paths share it --- litellm/proxy/auth/user_api_key_auth.py | 42 ++++++---- .../mcp_server/auth/user_api_key_auth_mcp.py | 76 ++++++++++++++---- .../auth/test_user_api_key_auth_mcp.py | 78 +++++++++++++++++++ 3 files changed, 162 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 99edfb940d7..8db52f2353a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1102,6 +1102,24 @@ def _matches_routing_override(token_claims: dict, override: "JWTRoutingOverride" ) +def oauth2_introspection_applies( + api_key: str, route: str, *, enable_oauth2_auth: bool, enable_jwt_auth: bool, jwt_handler: JWTHandler +) -> bool: + """Whether LiteLLM validates ``api_key`` on ``route`` by OAuth2 token introspection, on LLM API, info and + lens routes only: a non-JWT token when OAuth2 auth is on globally, or a JWT matched by a routing override.""" + is_jwt: Final = enable_jwt_auth and jwt_handler.is_jwt(token=api_key) + # Routing uses unverified JWT claims only to choose auth path. + # Final authentication is enforced by the selected validator. + route_jwt_to_oauth2: Final = is_jwt and _should_route_jwt_to_oauth2_override(token=api_key, jwt_handler=jwt_handler) + covered_route: Final = ( + RouteChecks.is_llm_api_route(route=route) + or RouteChecks.is_info_route(route=route) + or route == "/lens" + or route.startswith("/lens/") + ) + return covered_route and ((enable_oauth2_auth and not is_jwt) or route_jwt_to_oauth2) + + def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> bool: routing_overrides: Final = jwt_handler.litellm_jwtauth.routing_overrides if not routing_overrides: @@ -2216,23 +2234,13 @@ async def user_api_key_auth_builder( enable_oauth2_auth: Final = general_settings.get("enable_oauth2_auth", False) is True enable_jwt_auth: Final = general_settings.get("enable_jwt_auth", False) is True - is_jwt = jwt_handler.is_jwt(token=api_key) if enable_jwt_auth else False - - # Routing uses unverified JWT claims only to choose auth path. - # Final authentication is enforced by the selected validator. - route_jwt_to_oauth2 = is_jwt and _should_route_jwt_to_oauth2_override(token=api_key, jwt_handler=jwt_handler) - - # OAuth2 applies for: - # 1) when global OAuth2 auth is enabled on LLM + info routes - # 2) JWT tokens that explicitly match routing_overrides on LLM + info routes - route_is_lens: Final = route == "/lens" or route.startswith("/lens/") - should_apply_override_oauth2: Final = route_jwt_to_oauth2 and ( - RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route) or route_is_lens - ) - should_apply_global_oauth2: Final = enable_oauth2_auth and ( - RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route) or route_is_lens - ) - if (should_apply_global_oauth2 and not is_jwt) or should_apply_override_oauth2: + if oauth2_introspection_applies( + api_key, + route, + enable_oauth2_auth=enable_oauth2_auth, + enable_jwt_auth=enable_jwt_auth, + jwt_handler=jwt_handler, + ): from litellm.proxy.proxy_server import premium_user if premium_user is not True: diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index 78ec2c1fe6e..0886054534d 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from fastapi import HTTPException from pydantic import SecretStr, TypeAdapter @@ -48,6 +48,7 @@ from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports _get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth _run_centralized_common_checks, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + oauth2_introspection_applies, run_centralized_common_checks, user_api_key_auth, ) @@ -102,6 +103,7 @@ _EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({}) OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) _PROXY_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _CAUSE_CHAIN_MAX_DEPTH: Final = 16 +_IdentityCheck: TypeAlias = Literal["oauth2", "oauth2_proxy", "jwt"] _CACHED_VALUE: Final = TypeAdapter(object) @@ -266,38 +268,78 @@ async def _bearer_is_litellm_credential(token: str) -> bool: return token.startswith(LITELLM_VIRTUAL_KEY_PREFIX) and await virtual_key_may_exist(token) -async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: Request) -> UserAPIKeyAuth | None: - """The LiteLLM user an upstream bearer also proves, or ``None`` when it proves none. - - Custom auth decides alone, so any rejection is ``None``. A JWT is tried only when JWT auth is on, its - issuer may be LiteLLM's, and its signature verifies. Past that, only a user LiteLLM does not know is - ``None``; every other refusal (budget, blocked team, disallowed route) propagates.""" +def _identity_check(token: str, route: str, settings: Mapping[str, object]) -> _IdentityCheck | None: + """The identity check LiteLLM's own auth would run on ``token`` for ``route``, in its order: OAuth2 + introspection, then trusted reverse-proxy identity headers, then a JWT whose issuer may be LiteLLM's.""" from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module - custom_auth_configured, jwt_issuer_may_be_litellm, ) from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration - general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict jwt_handler, premium_user, ) + enable_jwt_auth: Final = settings.get("enable_jwt_auth") is True + if oauth2_introspection_applies( + token, + route, + enable_oauth2_auth=settings.get("enable_oauth2_auth") is True, + enable_jwt_auth=enable_jwt_auth, + jwt_handler=jwt_handler, + ): + return "oauth2" if premium_user is True else None + if settings.get("enable_oauth2_proxy_auth") is True: + return "oauth2_proxy" + if enable_jwt_auth and premium_user is True and JWTHandler.is_jwt(token=token) and jwt_issuer_may_be_litellm(token): + return "jwt" + return None + + +async def _verify_identity(check: _IdentityCheck, token: str, request: Request) -> None: + from litellm.proxy.auth.oauth2_check import Oauth2Handler # noqa: PLC0415 # proxy import cycle + from litellm.proxy.auth.oauth2_proxy_hook import ( # noqa: PLC0415 # proxy import cycle + handle_oauth2_proxy_request, + ) + from litellm.proxy.proxy_server import jwt_handler # noqa: PLC0415 # startup owns the active auth configuration + + match check: + case "oauth2": + await Oauth2Handler.check_oauth2_token(token=token) + case "oauth2_proxy": + await handle_oauth2_proxy_request(request=request) + case "jwt": + await jwt_handler.auth_jwt(token=token) # pyright: ignore[reportUnknownMemberType] # returns a bare dict; only whether it raises matters + case _: + assert_never(check) + + +async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: Request) -> UserAPIKeyAuth | None: + """The LiteLLM user an upstream bearer also proves, or ``None`` when it proves none. + + Custom auth decides alone, so any rejection is ``None``. Otherwise the identity check LiteLLM would run + (OAuth2 introspection, reverse-proxy headers, or a JWT) must pass first. Past that, only a user LiteLLM + does not know is ``None``; every other refusal (budget, blocked team, disallowed route) propagates.""" + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 # proxy import cycle + from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module + custom_auth_configured, + ) + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration + general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict + ) + if custom_auth_configured(): try: return await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException): return None - settings: Final = _PROXY_SETTINGS_ADAPTER.validate_python(general_settings) - if ( - settings.get("enable_jwt_auth") is not True - or premium_user is not True - or not JWTHandler.is_jwt(token=token) - or not jwt_issuer_may_be_litellm(token) - ): + check: Final = _identity_check( + token, get_request_route(request), _PROXY_SETTINGS_ADAPTER.validate_python(general_settings) + ) + if check is None: return None try: - await jwt_handler.auth_jwt(token=token) # pyright: ignore[reportUnknownMemberType] # returns a bare dict; only whether it raises matters + await _verify_identity(check, token, request) except Exception: # noqa: BLE001 # any verification failure means LiteLLM does not vouch for this bearer return None try: diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index 26a5e5de988..bd252ff21fe 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -2826,6 +2826,84 @@ class TestMCPDelegateAuthToUpstream: assert exc_info.value is refusal + @staticmethod + @contextlib.contextmanager + def _idp_auth_enabled(mechanism: str, verify: AsyncMock, premium: bool = True): + from litellm.proxy import proxy_server + from litellm.proxy.auth.oauth2_check import Oauth2Handler + + settings = {"oauth2": {"enable_oauth2_auth": True}, "oauth2_proxy": {"enable_oauth2_proxy_auth": True}} + verifier = ( + patch.object(Oauth2Handler, "check_oauth2_token", verify) + if mechanism == "oauth2" + else patch("litellm.proxy.auth.oauth2_proxy_hook.handle_oauth2_proxy_request", verify) + ) + with ( + patch.dict(proxy_server.general_settings, settings[mechanism]), + patch.object(proxy_server, "premium_user", premium), + patch.object(proxy_server, "user_custom_auth", None), + verifier, + ): + yield + + @pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"]) + @pytest.mark.parametrize( + "verify_effect, allowed_server_ids, expected_user_id", + [ + (None, ["test-server-id"], "idp-user"), + (None, [], None), + (ValueError("Token is not active"), ["test-server-id"], None), + ], + ids=["granted-user", "user-not-granted", "identity-rejected"], + ) + async def test_delegate_idp_identity_follows_the_jwt_rules( + self, mechanism, verify_effect, allowed_server_ids, expected_user_id + ): + """OAuth2 introspection and trusted reverse-proxy headers admit a granted user by name, and a user + who is not granted or not vouched for anonymously. The bearer is forwarded either way.""" + admission = AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user")) + with ( + TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock(side_effect=verify_effect)), + TestMCPDelegateAuthToUpstream._delegate_admission(admission, allowed_server_ids), + ): + auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request( + TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token") + ) + + assert auth_result.user_id == expected_user_id + assert auth_result.mcp_forwards_caller_authorization is True + assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"} + + @pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"]) + async def test_delegate_idp_user_refusal_propagates(self, mechanism): + """Once the provider vouches for the caller, LiteLLM's budget and permission refusals apply.""" + refusal = ProxyException(message="Budget has been exceeded", type="budget_exceeded", param=None, code=400) + with ( + TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock()), + TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(side_effect=refusal)), + ): + with pytest.raises(ProxyException) as exc_info: + await MCPRequestHandler.process_mcp_request( + TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token") + ) + + assert exc_info.value is refusal + + async def test_delegate_oauth2_without_premium_never_calls_the_provider(self): + """OAuth2 introspection is enterprise-only, so without a license the bearer is the upstream's alone.""" + introspect = AsyncMock() + with ( + TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2", introspect, premium=False), + TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))), + ): + auth_result, *_ = await MCPRequestHandler.process_mcp_request( + TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token") + ) + + assert auth_result.user_id is None + assert auth_result.mcp_forwards_caller_authorization is True + introspect.assert_not_awaited() + async def test_delegate_jwt_from_unconfigured_issuer_is_forwarded_without_jwt_auth(self): """A JWT from an issuer LiteLLM's JWT auth is not configured for never reaches JWT auth.""" auth_jwt = AsyncMock(return_value={"sub": "jwt-user"})