mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-04 23:07:50 +00:00
feat: allow JWT and OAuth2 auth to coexist on the same instance (#23153)
When both enable_jwt_auth and enable_oauth2_auth are True, the proxy now routes tokens based on their format: - JWT tokens (3 dot-separated parts) -> JWT auth handler - Opaque tokens -> OAuth2 auth handler This enables using JWT for human users and OAuth2 for M2M (machine) clients on the same LiteLLM instance. Previously, enabling OAuth2 would intercept all tokens on LLM API routes before JWT auth could run. When only one auth method is enabled, behavior is unchanged (backward compatible).
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user