diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 94cb7510a5..c992cfb53e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -615,17 +615,23 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # This allows UI SSO to work separately from API M2M authentication # Note: Info routes are already scoped to the user if RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route): - # return UserAPIKeyAuth object - # helper to check if the api_key is a valid oauth2 token - from litellm.proxy.proxy_server import premium_user + # When both OAuth2 and JWT auth are enabled, use token format to decide: + # - JWT tokens (3 dot-separated parts) -> skip OAuth2, fall through to JWT handler + # - Opaque tokens -> use OAuth2 handler + # This allows JWT for users and OAuth2 for M2M on the same instance + is_jwt_token = jwt_handler.is_jwt(token=api_key) if general_settings.get("enable_jwt_auth", False) is True else False + if not is_jwt_token: + # return UserAPIKeyAuth object + # helper to check if the api_key is a valid oauth2 token + from litellm.proxy.proxy_server import premium_user - if premium_user is not True: - raise ValueError( - "Oauth2 token validation is only available for premium users" - + CommonProxyErrors.not_premium_user.value - ) + if premium_user is not True: + raise ValueError( + "Oauth2 token validation is only available for premium users" + + CommonProxyErrors.not_premium_user.value + ) - return await Oauth2Handler.check_oauth2_token(token=api_key) + return await Oauth2Handler.check_oauth2_token(token=api_key) if general_settings.get("enable_oauth2_proxy_auth", False) is True: return await handle_oauth2_proxy_request(request=request) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 79c2ed4158..f3f0ba56cb 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -13,8 +13,12 @@ from unittest.mock import MagicMock import pytest +import litellm.proxy.proxy_server +from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import LiteLLM_JWTAuth, UserAPIKeyAuth +from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.auth.user_api_key_auth import get_api_key +from litellm.proxy.auth.user_api_key_auth import get_api_key, user_api_key_auth def test_get_api_key(): @@ -515,3 +519,169 @@ def test_proxy_admin_jwt_auth_handles_no_team_object(): assert result.team_metadata is None assert result.org_id is None assert result.end_user_id is None + + +class TestJWTOAuth2Coexistence: + """ + Test that JWT and OAuth2 auth can coexist on the same instance. + + When both enable_jwt_auth and enable_oauth2_auth are True, the proxy should + route tokens based on their format: + - JWT tokens (3 dot-separated parts) -> JWT auth handler + - Opaque tokens -> OAuth2 auth handler + """ + + def test_is_jwt_detects_jwt_tokens(self): + """JWT tokens have 3 dot-separated parts.""" + assert JWTHandler.is_jwt("header.payload.signature") is True + assert JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123") is True + + def test_is_jwt_rejects_opaque_tokens(self): + """Opaque OAuth2 tokens do not have 3 dot-separated parts.""" + assert JWTHandler.is_jwt("some-opaque-oauth2-token") is False + assert JWTHandler.is_jwt("sk-12345678") is False + assert JWTHandler.is_jwt("Bearer token") is False + assert JWTHandler.is_jwt("two.parts") is False + + @pytest.mark.asyncio + async def test_both_enabled_opaque_token_uses_oauth2(self): + """ + When both enable_jwt_auth and enable_oauth2_auth are True, + an opaque token should be handled by OAuth2 auth (not JWT). + """ + opaque_token = "some-opaque-m2m-oauth2-token" + + general_settings = { + "enable_oauth2_auth": True, + "enable_jwt_auth": True, + } + + mock_oauth2_response = UserAPIKeyAuth( + api_key=opaque_token, + user_id="machine-client-1", + team_id="m2m-team", + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.headers = {"authorization": f"Bearer {opaque_token}"} + mock_request.query_params = {} + + with patch("litellm.proxy.proxy_server.general_settings", general_settings), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch("litellm.proxy.proxy_server.master_key", "sk-master"), \ + patch("litellm.proxy.proxy_server.prisma_client", None), \ + patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock, return_value=mock_oauth2_response) as mock_oauth2, \ + patch("litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", new_callable=AsyncMock) as mock_jwt_auth: + + litellm.proxy.proxy_server.jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + result = await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {opaque_token}", + ) + + # OAuth2 SHOULD be called for opaque tokens + mock_oauth2.assert_called_once_with(token=opaque_token) + # JWT auth should NOT be called + mock_jwt_auth.assert_not_called() + assert result.user_id == "machine-client-1" + + @pytest.mark.asyncio + async def test_both_enabled_jwt_token_skips_oauth2(self): + """ + When both enable_jwt_auth and enable_oauth2_auth are True, + a JWT-formatted token should skip OAuth2 and reach the JWT handler. + """ + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + + general_settings = { + "enable_oauth2_auth": True, + "enable_jwt_auth": True, + } + + mock_jwt_result = { + "is_proxy_admin": True, + "team_object": None, + "user_object": None, + "end_user_object": None, + "org_object": None, + "token": jwt_token, + "team_id": "jwt-team", + "user_id": "jwt-human-user", + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": "user1"}, + } + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + + with patch("litellm.proxy.proxy_server.general_settings", general_settings), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch("litellm.proxy.proxy_server.master_key", "sk-master"), \ + patch("litellm.proxy.proxy_server.prisma_client", None), \ + patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock) as mock_oauth2, \ + patch("litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", new_callable=AsyncMock, return_value=mock_jwt_result) as mock_jwt_auth: + + litellm.proxy.proxy_server.jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + result = await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {jwt_token}", + ) + + # OAuth2 should NOT be called for JWT tokens + mock_oauth2.assert_not_called() + # JWT auth SHOULD be called + mock_jwt_auth.assert_called_once() + assert result.user_id == "jwt-human-user" + + @pytest.mark.asyncio + async def test_only_oauth2_enabled_handles_all_tokens(self): + """ + When only enable_oauth2_auth is True (no JWT), all LLM API tokens + should go through OAuth2 - backward compatible behavior. + """ + jwt_like_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + + general_settings = { + "enable_oauth2_auth": True, + "enable_jwt_auth": False, + } + + mock_oauth2_response = UserAPIKeyAuth( + api_key=jwt_like_token, + user_id="oauth2-user", + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.headers = {"authorization": f"Bearer {jwt_like_token}"} + mock_request.query_params = {} + + with patch("litellm.proxy.proxy_server.general_settings", general_settings), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch("litellm.proxy.proxy_server.master_key", "sk-master"), \ + patch("litellm.proxy.proxy_server.prisma_client", None), \ + patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock, return_value=mock_oauth2_response) as mock_oauth2: + + result = await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {jwt_like_token}", + ) + + # OAuth2 should handle it since JWT auth is disabled + mock_oauth2.assert_called_once_with(token=jwt_like_token) + assert result.user_id == "oauth2-user"