scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83
litellm/tests/test_litellm/proxy/auth/test_handle_jwt.py
ryan-crabbe-berri b76def0e5d
test: require a match= on broad pytest.raises, and drop duplicate parametrize cases (#37769) 
`pytest.raises(Exception)` with no `match=` passes on any error that broad. A
TypeError from a refactor, a botched fixture, an import that moved: all of them
read as the rejection the test claims to police, so the test goes green for the
wrong reason and stays green after the behaviour it guards is gone.

PT011 closes that gap for the 317 sites B017 could not reach, because B017 only
fires on a single-statement body with no `as e` binding. Each pattern here is the
message the code actually raised, recorded by running the sites under a plugin
that logged the concrete type and text per call site, so the assertions describe
observed behaviour rather than a guess. Where a site raises more than one message
across its parametrize cases, the pattern is an alternation of what was seen;
where the exception carries an empty `str()` and puts the text on `.message`, the
site keeps a narrow `noqa` with the reason.

PT014 removes four parametrize cases that were listed twice. The duplicate re-runs
an assertion that already passed, and it usually marks a case someone meant to
vary and forgot to edit.
2026-08-20 20:24:49 -07:00

6736 lines
241 KiB
Python
Raw Permalink Blame History

import asyncio
import re
import time
from collections.abc import Mapping, Sequence
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi import HTTPException
import httpx
import pytest
from litellm.proxy._types import (
DEFAULT_JWKS_STALE_TTL,
JWTLiteLLMRoleMap,
LiteLLM_JWTAuth,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
Member,
ProxyErrorTypes,
ProxyException,
)
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.handle_jwt import (
JWKS_FETCH_ATTEMPTS,
STALE_CACHE_KEY_PREFIX,
STALE_WRITTEN_AT_CACHE_KEY_PREFIX,
JWKSUnreachableError,
JWTAuthManager,
JWTHandler,
NoMatchingJWTPublicKeyError,
)
@pytest.mark.asyncio
async def test_map_user_to_teams_user_already_in_team():
"""Test that no action is taken when user is already in team"""
# Setup test data
user = LiteLLM_UserTable(user_id="test_user_1")
team = LiteLLM_TeamTable(
team_id="test_team_1",
members_with_roles=[Member(user_id="test_user_1", role="user")],
)
# Mock team_member_add to ensure it's not called
with patch(
"litellm.proxy.management_endpoints.team_endpoints.team_member_add",
new_callable=AsyncMock,
) as mock_add:
await JWTAuthManager.map_user_to_teams(user_object=user, team_object=team)
mock_add.assert_not_called()
@pytest.mark.asyncio
async def test_map_user_to_teams_add_new_user():
"""Test that new user is added to team"""
# Setup test data
user = LiteLLM_UserTable(user_id="test_user_1")
team = LiteLLM_TeamTable(team_id="test_team_1", members_with_roles=[])
# Mock team_member_add
with patch(
"litellm.proxy.management_endpoints.team_endpoints.team_member_add",
new_callable=AsyncMock,
) as mock_add:
await JWTAuthManager.map_user_to_teams(user_object=user, team_object=team)
mock_add.assert_called_once()
# Verify the correct data was passed to team_member_add
call_args = mock_add.call_args[1]["data"]
assert call_args.member.user_id == "test_user_1"
assert call_args.member.role == "user"
assert call_args.team_id == "test_team_1"
@pytest.mark.asyncio
async def test_map_user_to_teams_handles_already_in_team_exception():
"""Test that team_member_already_in_team exception is handled gracefully"""
# Setup test data
user = LiteLLM_UserTable(user_id="test_user_1")
team = LiteLLM_TeamTable(team_id="test_team_1", members_with_roles=[])
# Create a ProxyException with team_member_already_in_team error type
already_in_team_exception = ProxyException(
message="User test_user_1 already in team",
type=ProxyErrorTypes.team_member_already_in_team,
param="user_id",
code="400",
)
# Mock team_member_add to raise the exception
with patch(
"litellm.proxy.management_endpoints.team_endpoints.team_member_add",
new_callable=AsyncMock,
side_effect=already_in_team_exception,
) as mock_add:
with patch("litellm.proxy.auth.handle_jwt.verbose_proxy_logger") as mock_logger:
# This should not raise an exception
result = await JWTAuthManager.map_user_to_teams(
user_object=user, team_object=team
)
# Verify the method completed successfully
assert result is None
mock_add.assert_called_once()
@pytest.mark.asyncio
async def test_map_user_to_teams_reraises_other_proxy_exceptions():
"""Test that other ProxyException types are re-raised"""
# Setup test data
user = LiteLLM_UserTable(user_id="test_user_1")
team = LiteLLM_TeamTable(team_id="test_team_1", members_with_roles=[])
# Create a ProxyException with a different error type
other_exception = ProxyException(
message="Some other error",
type=ProxyErrorTypes.internal_server_error,
param="some_param",
code="500",
)
# Mock team_member_add to raise the exception
with patch(
"litellm.proxy.management_endpoints.team_endpoints.team_member_add",
new_callable=AsyncMock,
side_effect=other_exception,
) as mock_add:
# This should re-raise the exception
with pytest.raises(ProxyException) as exc_info:
await JWTAuthManager.map_user_to_teams(user_object=user, team_object=team)
@pytest.mark.asyncio
async def test_map_user_to_teams_null_inputs():
"""Test that method handles null inputs gracefully"""
# Test with null user
await JWTAuthManager.map_user_to_teams(
user_object=None, team_object=LiteLLM_TeamTable(team_id="test_team_1")
)
# Test with null team
await JWTAuthManager.map_user_to_teams(
user_object=LiteLLM_UserTable(user_id="test_user_1"), team_object=None
)
# Test with both null
await JWTAuthManager.map_user_to_teams(user_object=None, team_object=None)
@pytest.mark.asyncio
async def test_find_team_with_model_access_reports_passthrough_allowlist_denial():
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
team = LiteLLM_TeamTable(
team_id="team-a",
models=["gpt-4"],
metadata={},
)
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=team,
),
patch(
"litellm.proxy.auth.handle_jwt.can_team_access_model",
new_callable=AsyncMock,
return_value=True,
),
patch(
"litellm.proxy.auth.handle_jwt.allowed_routes_check",
return_value=True,
),
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
return_value=True,
) as mock_is_auth_enforced_pass_through_route,
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
return_value=False,
) as mock_passthrough_check,
):
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.find_team_with_model_access(
team_ids={"team-a"},
requested_model="gpt-4",
route="/my-pass-through",
request_method="POST",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert exc_info.value.status_code == 403
assert "allowed_passthrough_routes" in exc_info.value.detail
assert "requested model" not in exc_info.value.detail
mock_is_auth_enforced_pass_through_route.assert_called_once_with(
route="/my-pass-through", method="POST"
)
user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"]
assert user_api_key_dict.metadata == {}
assert user_api_key_dict.team_metadata == {}
@pytest.mark.asyncio
async def test_find_team_with_model_access_uses_request_method_for_passthrough_auth():
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
team = LiteLLM_TeamTable(
team_id="team-a",
models=["gpt-4"],
metadata={},
)
mock_registered_routes = {
"test-uuid-1:exact:/custom:GET": {
"endpoint_id": "test-uuid-1",
"path": "/custom",
"type": "exact",
"methods": ["GET"],
"auth": False,
},
"test-uuid-2:exact:/custom:POST": {
"endpoint_id": "test-uuid-2",
"path": "/custom",
"type": "exact",
"methods": ["POST"],
"auth": True,
},
}
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=team,
),
patch(
"litellm.proxy.auth.handle_jwt.allowed_routes_check",
return_value=True,
),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
mock_registered_routes,
),
patch(
"litellm.proxy.utils.get_server_root_path",
return_value="/",
),
):
team_id, team_obj = await JWTAuthManager.find_team_with_model_access(
team_ids={"team-a"},
requested_model=None,
route="/custom",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
request_method="GET",
)
assert team_id == "team-a"
assert team_obj == team
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.find_team_with_model_access(
team_ids={"team-a"},
requested_model=None,
route="/custom",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
request_method="POST",
)
assert exc_info.value.status_code == 403
assert "allowed_passthrough_routes" in exc_info.value.detail
@pytest.mark.asyncio
async def test_auth_builder_proxy_admin_user_role():
"""Test that is_proxy_admin is True when user_object.user_role is PROXY_ADMIN"""
# Setup test data
api_key = "test_jwt_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
# Create user object with PROXY_ADMIN role
user_object = LiteLLM_UserTable(
user_id="test_user_1", user_role=LitellmUserRoles.PROXY_ADMIN
)
# Create mock JWT handler
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
# Mock all the dependencies and method calls
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(
JWTAuthManager, "check_rbac_role", new_callable=AsyncMock
) as mock_check_rbac,
patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac,
patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes,
patch.object(
jwt_handler, "get_object_id", return_value=None
) as mock_get_object_id,
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=("test_user_1", "test@example.com", True),
) as mock_get_user_info,
patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id,
patch.object(
jwt_handler, "get_end_user_id", return_value=None
) as mock_get_end_user_id,
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
) as mock_check_admin,
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team,
patch.object(
JWTAuthManager, "get_all_team_ids", return_value=set()
) as mock_get_all_team_ids,
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team_access,
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
) as mock_get_objects,
patch.object(
JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock
) as mock_map_user,
patch.object(
JWTAuthManager, "validate_object_id", return_value=True
) as mock_validate_object,
):
# Set up the mock return values
mock_auth_jwt.return_value = {"sub": "test_user_1", "scope": ""}
# Call the auth_builder method
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
# Verify that is_proxy_admin is True
assert result["is_proxy_admin"] is True
assert result["user_object"] == user_object
assert result["user_id"] == "test_user_1"
@pytest.mark.asyncio
async def test_auth_builder_non_proxy_admin_user_role():
"""Test that is_proxy_admin is False when user_object.user_role is not PROXY_ADMIN"""
# Setup test data
api_key = "test_jwt_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
# Create user object with regular USER role
user_object = LiteLLM_UserTable(
user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER
)
# Create mock JWT handler
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
# Mock all the dependencies and method calls
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(
JWTAuthManager, "check_rbac_role", new_callable=AsyncMock
) as mock_check_rbac,
patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac,
patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes,
patch.object(
jwt_handler, "get_object_id", return_value=None
) as mock_get_object_id,
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=("test_user_1", "test@example.com", True),
) as mock_get_user_info,
patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id,
patch.object(
jwt_handler, "get_end_user_id", return_value=None
) as mock_get_end_user_id,
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
) as mock_check_admin,
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team,
patch.object(
JWTAuthManager, "get_all_team_ids", return_value=set()
) as mock_get_all_team_ids,
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team_access,
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
) as mock_get_objects,
patch.object(
JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock
) as mock_map_user,
patch.object(
JWTAuthManager, "validate_object_id", return_value=True
) as mock_validate_object,
):
# Set up the mock return values
mock_auth_jwt.return_value = {"sub": "test_user_1", "scope": ""}
# Call the auth_builder method
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
# Verify that is_proxy_admin is False
assert result["is_proxy_admin"] is False
assert result["user_object"] == user_object
assert result["user_id"] == "test_user_1"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"row_email,expected_email",
[
("row@example.com", "row@example.com"),
(None, "claim@example.com"),
("", "claim@example.com"),
],
)
async def test_auth_builder_result_includes_user_email(row_email, expected_email):
"""LIT-4238: auth_builder must return user_email (user row wins, JWT claim
is the fallback) so the auth object and metrics get the email."""
api_key = "test_jwt_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
user_object = LiteLLM_UserTable(
user_id="test_user_1",
user_email=row_email,
user_role=LitellmUserRoles.INTERNAL_USER,
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=("test_user_1", "claim@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
) as mock_check_admin,
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
):
mock_auth_jwt.return_value = {"sub": "test_user_1", "scope": ""}
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result["user_email"] == expected_email
assert mock_check_admin.call_args.kwargs["user_email"] == "claim@example.com"
@pytest.mark.asyncio
async def test_check_admin_access_result_includes_user_email():
"""LIT-4238: the scope-based admin path has no user row, so the JWT claim
email must ride the JWTAuthBuilderResult."""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
admin_jwt_scope="litellm_proxy_admin",
admin_allowed_routes=["/chat/completions"],
)
result = await JWTAuthManager.check_admin_access(
jwt_handler=jwt_handler,
scopes=["litellm_proxy_admin"],
route="/chat/completions",
user_id="admin-user",
user_email="admin@example.com",
org_id=None,
api_key="test_jwt_token",
jwt_valid_token={"sub": "admin-user"},
)
assert result is not None
assert result["is_proxy_admin"] is True
assert result["user_email"] == "admin@example.com"
@pytest.mark.asyncio
async def test_sync_user_role_and_teams():
from unittest.mock import MagicMock
# Create mock objects for required types
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=mock_user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
jwt_litellm_role_map=[
JWTLiteLLMRoleMap(
jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN
)
],
roles_jwt_field="roles",
team_ids_jwt_field="my_id_teams",
sync_user_role_and_teams=True,
),
)
token = {"roles": ["ADMIN"], "my_id_teams": ["team1", "team2"]}
user = LiteLLM_UserTable(
user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=["team2"]
)
prisma = AsyncMock()
prisma.db.litellm_usertable.update = AsyncMock()
with patch(
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
new_callable=AsyncMock,
) as mock_patch:
await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, prisma)
prisma.db.litellm_usertable.update.assert_called_once()
mock_patch.assert_called_once()
assert user.user_role == LitellmUserRoles.PROXY_ADMIN.value
assert set(user.teams) == {"team1", "team2"}
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_cache_invalidation_on_role_change():
"""Test that user cache is updated when role changes."""
mock_cache = AsyncMock()
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=AsyncMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
jwt_litellm_role_map=[
JWTLiteLLMRoleMap(
jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN
)
],
roles_jwt_field="roles",
team_ids_jwt_field="my_id_teams",
sync_user_role_and_teams=True,
),
)
token = {"roles": ["ADMIN"], "my_id_teams": ["team1"]}
user = LiteLLM_UserTable(
user_id="u1",
user_role=LitellmUserRoles.INTERNAL_USER.value,
teams=["team1"], # teams already match — only role differs
)
prisma = AsyncMock()
prisma.db.litellm_usertable.update = AsyncMock()
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler, token, user, prisma, user_api_key_cache=mock_cache
)
mock_cache.async_set_cache.assert_called_once()
call_kwargs = mock_cache.async_set_cache.call_args
assert call_kwargs.kwargs["key"] == "u1"
assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable)
assert call_kwargs.kwargs["value"].user_role == LitellmUserRoles.PROXY_ADMIN.value
assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_cache_invalidation_on_team_change():
"""Test that user cache is updated when team memberships change."""
mock_cache = AsyncMock()
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=AsyncMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
jwt_litellm_role_map=[
JWTLiteLLMRoleMap(
jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN
)
],
roles_jwt_field="roles",
team_ids_jwt_field="my_id_teams",
sync_user_role_and_teams=True,
),
)
token = {"roles": ["ADMIN"], "my_id_teams": ["team1", "team2"]}
user = LiteLLM_UserTable(
user_id="u1",
user_role=LitellmUserRoles.PROXY_ADMIN.value, # role already matches
teams=["team2"], # teams differ
)
prisma = AsyncMock()
prisma.db.litellm_usertable.update = AsyncMock()
with patch(
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
new_callable=AsyncMock,
):
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler, token, user, prisma, user_api_key_cache=mock_cache
)
mock_cache.async_set_cache.assert_called_once()
call_kwargs = mock_cache.async_set_cache.call_args
assert call_kwargs.kwargs["key"] == "u1"
assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable)
assert set(call_kwargs.kwargs["value"].teams) == {"team1", "team2"}
assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes():
"""Test that cache is NOT written when role and teams already match."""
mock_cache = AsyncMock()
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=AsyncMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
jwt_litellm_role_map=[
JWTLiteLLMRoleMap(
jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN
)
],
roles_jwt_field="roles",
team_ids_jwt_field="my_id_teams",
sync_user_role_and_teams=True,
),
)
token = {"roles": ["ADMIN"], "my_id_teams": ["team1"]}
user = LiteLLM_UserTable(
user_id="u1",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
teams=["team1"],
)
prisma = AsyncMock()
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler, token, user, prisma, user_api_key_cache=mock_cache
)
mock_cache.async_set_cache.assert_not_called()
def test_get_all_jwt_team_ids_unions_singular_and_plural():
"""get_all_jwt_team_ids must include the singular team_id_jwt_field claim
in addition to the plural team_ids_jwt_field, deduplicated."""
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=MagicMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
team_id_jwt_field="team_id",
team_ids_jwt_field="teams",
),
)
# singular only — Okta/Auth0 default shape
assert jwt_handler.get_all_jwt_team_ids({"team_id": "team-low"}) == ["team-low"]
# plural only — pre-fix shape
assert jwt_handler.get_all_jwt_team_ids({"teams": ["a", "b"]}) == ["a", "b"]
# both populated, no overlap
assert jwt_handler.get_all_jwt_team_ids(
{"team_id": "primary", "teams": ["a", "b"]}
) == ["a", "b", "primary"]
# both populated with overlap — singular dedup'd
assert jwt_handler.get_all_jwt_team_ids({"team_id": "a", "teams": ["a", "b"]}) == [
"a",
"b",
]
# singular field as multi-element list (some IdPs) — merge all, preserve plural-first order
assert jwt_handler.get_all_jwt_team_ids(
{"team_id": ["primary", "secondary"], "teams": ["a"]}
) == [
"a",
"primary",
"secondary",
]
# neither populated
assert jwt_handler.get_all_jwt_team_ids({}) == []
def test_get_all_jwt_team_ids_does_not_use_team_id_default():
"""team_id_default is a JWT-bearer-flow auth-builder fallback, not a token
claim. It must NOT leak into get_all_jwt_team_ids — otherwise SSO logins
would silently start adding users to the default team for any tenant that
has team_id_default configured."""
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=MagicMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
team_id_jwt_field="team_id",
team_ids_jwt_field="teams",
team_id_default="default-team",
),
)
# team_id claim missing — must not fall back to default-team
assert jwt_handler.get_all_jwt_team_ids({"teams": []}) == []
assert jwt_handler.get_all_jwt_team_ids({}) == []
# only the plural is populated — default still must not be added
assert jwt_handler.get_all_jwt_team_ids({"teams": ["a"]}) == ["a"]
# team_id_jwt_field unset entirely + only default configured: still no default
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=MagicMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="teams",
team_id_default="default-team",
),
)
assert jwt_handler.get_all_jwt_team_ids({"teams": []}) == []
@pytest.mark.asyncio
async def test_map_jwt_role_to_litellm_role():
"""Test JWT role mapping to LiteLLM roles with various patterns"""
from unittest.mock import MagicMock
# Create mock objects for required types
mock_user_api_key_cache = MagicMock()
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=mock_user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
jwt_litellm_role_map=[
# Exact match
JWTLiteLLMRoleMap(
jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN
),
# Wildcard patterns
JWTLiteLLMRoleMap(
jwt_role="user_*", litellm_role=LitellmUserRoles.INTERNAL_USER
),
JWTLiteLLMRoleMap(
jwt_role="team_?", litellm_role=LitellmUserRoles.TEAM
),
JWTLiteLLMRoleMap(
jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER
),
],
roles_jwt_field="roles",
),
)
# Test exact match
token = {"roles": ["ADMIN"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result == LitellmUserRoles.PROXY_ADMIN
# Test wildcard match with *
token = {"roles": ["user_manager"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result == LitellmUserRoles.INTERNAL_USER
token = {"roles": ["user_"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result == LitellmUserRoles.INTERNAL_USER
# Test wildcard match with ?
token = {"roles": ["team_1"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result == LitellmUserRoles.TEAM
token = {"roles": ["team_a"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result == LitellmUserRoles.TEAM
# Test character class match
token = {"roles": ["dev_1"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result == LitellmUserRoles.INTERNAL_USER
token = {"roles": ["dev_2"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result == LitellmUserRoles.INTERNAL_USER
# Test no match
token = {"roles": ["unknown_role"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result is None
# Test multiple roles - should return first mapping match
token = {"roles": ["user_test", "ADMIN"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result == LitellmUserRoles.PROXY_ADMIN # ADMIN matches first mapping
# Test empty roles
token = {"roles": []}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result is None
# Test no roles field
token = {}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result is None
# Test no role mappings configured
jwt_handler.litellm_jwtauth.jwt_litellm_role_map = None
token = {"roles": ["ADMIN"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result is None
# Test empty role mappings
jwt_handler.litellm_jwtauth.jwt_litellm_role_map = []
token = {"roles": ["ADMIN"]}
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result is None
# Test patterns that don't match character classes
jwt_handler.litellm_jwtauth.jwt_litellm_role_map = [
JWTLiteLLMRoleMap(
jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER
),
]
token = {"roles": ["dev_4"]} # 4 is not in [123]
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result is None
# Test ? pattern that requires exactly one character
jwt_handler.litellm_jwtauth.jwt_litellm_role_map = [
JWTLiteLLMRoleMap(jwt_role="team_?", litellm_role=LitellmUserRoles.TEAM),
]
token = {"roles": ["team_12"]} # More than one character after underscore
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result is None
token = {"roles": ["team_"]} # No character after underscore
result = jwt_handler.map_jwt_role_to_litellm_role(token)
assert result is None
@pytest.mark.asyncio
async def test_nested_jwt_field_access():
"""
Test that all JWT fields support dot notation for nested access
This test verifies that:
1. All JWT field methods can access nested values using dot notation
2. Backward compatibility is maintained for flat field names
3. Missing nested paths return appropriate defaults
"""
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
# Create JWT handler
jwt_handler = JWTHandler()
# Test token with nested claims
nested_token = {
"user": {"sub": "u123", "email": "user@example.com"},
"resource_access": {"my-client": {"roles": ["admin", "user"]}},
"groups": ["team1", "team2"],
"organization": {"id": "org456"},
"profile": {"object_id": "obj789"},
"customer": {"end_user_id": "customer123"},
"tenant": {"team_id": "team456"},
}
# Test flat token for backward compatibility
flat_token = {
"sub": "u123",
"email": "user@example.com",
"roles": ["admin", "user"],
"groups": ["team1", "team2"],
"org_id": "org456",
"object_id": "obj789",
"end_user_id": "customer123",
"team_id": "team456",
}
# Test 1: user_id_jwt_field with nested access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_id_jwt_field="user.sub")
assert jwt_handler.get_user_id(nested_token, None) == "u123"
# Test 1b: user_id_jwt_field with flat access (backward compatibility)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_id_jwt_field="sub")
assert jwt_handler.get_user_id(flat_token, None) == "u123"
# Test 2: user_email_jwt_field with nested access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="user.email")
assert jwt_handler.get_user_email(nested_token, None) == "user@example.com"
# Test 2b: user_email_jwt_field with flat access (backward compatibility)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="email")
assert jwt_handler.get_user_email(flat_token, None) == "user@example.com"
# Test 3: team_ids_jwt_field with nested access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups")
assert jwt_handler.get_team_ids_from_jwt(nested_token) == ["team1", "team2"]
# Test 3b: team_ids_jwt_field with flat access (backward compatibility)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups")
assert jwt_handler.get_team_ids_from_jwt(flat_token) == ["team1", "team2"]
# Test 4: org_id_jwt_field with nested access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_id_jwt_field="organization.id")
assert jwt_handler.get_org_id(nested_token, None) == "org456"
# Test 4b: org_id_jwt_field with flat access (backward compatibility)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_id_jwt_field="org_id")
assert jwt_handler.get_org_id(flat_token, None) == "org456"
# Test 5: object_id_jwt_field with nested access (requires role_mappings)
from litellm.proxy._types import LitellmUserRoles, RoleMapping
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
object_id_jwt_field="profile.object_id",
role_mappings=[
RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)
],
)
assert jwt_handler.get_object_id(nested_token, None) == "obj789"
# Test 5b: object_id_jwt_field with flat access (backward compatibility)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
object_id_jwt_field="object_id",
role_mappings=[
RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)
],
)
assert jwt_handler.get_object_id(flat_token, None) == "obj789"
# Test 6: end_user_id_jwt_field with nested access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
end_user_id_jwt_field="customer.end_user_id"
)
assert jwt_handler.get_end_user_id(nested_token, None) == "customer123"
# Test 6b: end_user_id_jwt_field with flat access (backward compatibility)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(end_user_id_jwt_field="end_user_id")
assert jwt_handler.get_end_user_id(flat_token, None) == "customer123"
# Test 7: team_id_jwt_field with nested access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="tenant.team_id")
assert jwt_handler.get_team_id(nested_token, None) == "team456"
# Test 7b: team_id_jwt_field with flat access (backward compatibility)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id")
assert jwt_handler.get_team_id(flat_token, None) == "team456"
# Test 8: roles_jwt_field with deeply nested access (already supported, but testing)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
roles_jwt_field="resource_access.my-client.roles"
)
assert jwt_handler.get_jwt_role(nested_token, []) == ["admin", "user"]
# Test 9: user_roles_jwt_field with nested access (already supported, but testing)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
user_roles_jwt_field="resource_access.my-client.roles",
user_allowed_roles=["admin", "user"],
)
assert jwt_handler.get_user_roles(nested_token, []) == ["admin", "user"]
@pytest.mark.asyncio
async def test_nested_jwt_field_missing_paths():
"""
Test handling of missing nested paths in JWT tokens
This test verifies that:
1. Missing nested paths return appropriate defaults
2. Partial paths that exist but don't have the final key return defaults
3. team_id_default fallback works with nested fields
"""
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
# Create JWT handler
jwt_handler = JWTHandler()
# Test token with missing nested paths
incomplete_token = {
"user": {
"name": "test user"
# missing "sub" and "email"
},
"resource_access": {
"other-client": {"roles": ["viewer"]}
# missing "my-client"
},
# missing "organization", "profile", "customer", "tenant", "groups"
}
# Test 1: Missing user.sub should return default
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_id_jwt_field="user.sub")
assert jwt_handler.get_user_id(incomplete_token, "default_user") == "default_user"
# Test 2: Missing user.email should return default
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="user.email")
assert (
jwt_handler.get_user_email(incomplete_token, "default@example.com")
== "default@example.com"
)
# Test 3: Missing groups should return empty list
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups")
assert jwt_handler.get_team_ids_from_jwt(incomplete_token) == []
# Test 4: Missing organization.id should return default
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_id_jwt_field="organization.id")
assert jwt_handler.get_org_id(incomplete_token, "default_org") == "default_org"
# Test 5: Missing profile.object_id should return default (requires role_mappings)
from litellm.proxy._types import LitellmUserRoles, RoleMapping
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
object_id_jwt_field="profile.object_id",
role_mappings=[
RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)
],
)
assert jwt_handler.get_object_id(incomplete_token, "default_obj") == "default_obj"
# Test 6: Missing customer.end_user_id should return default
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
end_user_id_jwt_field="customer.end_user_id"
)
assert (
jwt_handler.get_end_user_id(incomplete_token, "default_customer")
== "default_customer"
)
# Test 7: Missing tenant.team_id should use team_id_default fallback
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_id_jwt_field="tenant.team_id", team_id_default="fallback_team"
)
assert jwt_handler.get_team_id(incomplete_token, "default_team") == "fallback_team"
# Test 8: Missing resource_access.my-client.roles should return default
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
roles_jwt_field="resource_access.my-client.roles"
)
assert jwt_handler.get_jwt_role(incomplete_token, ["default_role"]) == [
"default_role"
]
# Test 9: Missing nested user roles should return default
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
user_roles_jwt_field="resource_access.my-client.roles",
user_allowed_roles=["admin", "user"],
)
assert jwt_handler.get_user_roles(incomplete_token, ["default_user_role"]) == [
"default_user_role"
]
@pytest.mark.asyncio
async def test_metadata_prefix_handling_in_nested_fields():
"""
Test that metadata. prefix is properly handled in nested JWT field access
The get_nested_value function should remove metadata. prefix before traversing
"""
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
# Create JWT handler
jwt_handler = JWTHandler()
# Test token with proper structure for metadata prefix removal
token = {
"user": {
"email": "user@example.com" # This will be accessed when metadata.user.email is used
},
"sub": "u123",
}
# Test 1: metadata.user.email should access user.email after prefix removal
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
user_email_jwt_field="metadata.user.email"
)
# The get_nested_value function removes "metadata." prefix, so "metadata.user.email" becomes "user.email"
assert jwt_handler.get_user_email(token, None) == "user@example.com"
# Test 2: user.sub should work normally without metadata prefix
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_id_jwt_field="sub")
assert jwt_handler.get_user_id(token, None) == "u123"
@pytest.mark.asyncio
async def test_find_team_with_model_access_model_group(monkeypatch):
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
router = Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "gpt-4o-mini"},
"model_info": {"access_groups": ["test-group"]},
}
]
)
import sys
import types
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
team = LiteLLM_TeamTable(team_id="team-1", models=["test-group"])
async def mock_get_team_object(*args, **kwargs): # type: ignore
return team
monkeypatch.setattr(
"litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
team_id, team_obj = await JWTAuthManager.find_team_with_model_access(
team_ids={"team-1"},
requested_model="gpt-4o-mini",
route="/chat/completions",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
assert team_id == "team-1"
assert team_obj.team_id == "team-1"
@pytest.mark.asyncio
async def test_find_team_with_model_access_v1_messages_default_routes(monkeypatch):
"""Regression for #31189: a single-team JWT that grants the requested model
through an access group must resolve on /v1/messages without an explicit
x-litellm-team-id header. /v1/messages lives in `anthropic_routes`, so when a
team has no `team_allowed_routes` configured the default allowlist must cover
it just like /chat/completions and /v1/responses; otherwise the internal route
check fails and surfaces a misleading "No team has access to the requested
model" 403."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
router = Router(
model_list=[
{
"model_name": "claude-sonnet-4-6",
"litellm_params": {"model": "claude-sonnet-4-6"},
"model_info": {"access_groups": ["coding_only_models"]},
}
]
)
import sys
import types
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
team = LiteLLM_TeamTable(team_id="coding-team", models=["coding_only_models"])
async def mock_get_team_object(*args, **kwargs): # type: ignore
return team
monkeypatch.setattr(
"litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
team_id, team_obj = await JWTAuthManager.find_team_with_model_access(
team_ids={"coding-team"},
requested_model="claude-sonnet-4-6",
route="/v1/messages",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
assert team_id == "coding-team"
assert team_obj.team_id == "coding-team"
@pytest.mark.parametrize(
"route,expected",
[
("/v1/messages", True),
("/v1/messages/count_tokens", True),
("/v1/skills", False),
("/v1/skills/skill_abc123", False),
],
)
def test_default_team_allowed_routes_cover_messages_but_not_skills(route, expected):
from litellm.proxy.auth.auth_checks import allowed_routes_check
assert (
allowed_routes_check(
user_role=LitellmUserRoles.TEAM,
user_route=route,
litellm_proxy_roles=LiteLLM_JWTAuth(),
)
is expected
)
@pytest.mark.asyncio
async def test_auth_builder_returns_team_membership_object():
"""
Test that auth_builder returns the team_membership_object when user is a member of a team.
"""
# Setup test data
api_key = "test_jwt_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
_team_id = "test_team_1"
_user_id = "test_user_1"
# Create mock objects
from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership
mock_team_membership = LiteLLM_TeamMembership(
user_id=_user_id,
team_id=_team_id,
budget_id="budget_123",
spend=10.5,
litellm_budget_table=LiteLLM_BudgetTable(
budget_id="budget_123", rpm_limit=100, tpm_limit=5000
),
)
user_object = LiteLLM_UserTable(
user_id=_user_id, user_role=LitellmUserRoles.INTERNAL_USER
)
team_object = LiteLLM_TeamTable(team_id=_team_id)
# Create mock JWT handler
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
# Mock all the dependencies and method calls
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(
JWTAuthManager, "check_rbac_role", new_callable=AsyncMock
) as mock_check_rbac,
patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac,
patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes,
patch.object(
jwt_handler, "get_object_id", return_value=None
) as mock_get_object_id,
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(_user_id, "test@example.com", True),
) as mock_get_user_info,
patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id,
patch.object(
jwt_handler, "get_end_user_id", return_value=None
) as mock_get_end_user_id,
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
) as mock_check_admin,
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(_team_id, team_object),
) as mock_find_team,
patch.object(
JWTAuthManager, "get_all_team_ids", return_value=set()
) as mock_get_all_team_ids,
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team_access,
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(
user_object,
None,
None,
mock_team_membership,
user_object.user_id,
),
) as mock_get_objects,
patch.object(
JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock
) as mock_map_user,
patch.object(
JWTAuthManager, "validate_object_id", return_value=True
) as mock_validate_object,
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
) as mock_sync_user,
):
# Set up the mock return values
mock_auth_jwt.return_value = {"sub": _user_id, "scope": ""}
# Call the auth_builder method
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
# Verify that team_membership_object is returned
assert result["team_membership"] is not None, (
"team_membership should be present"
)
assert result["team_membership"] == mock_team_membership, (
"team_membership should match the mock object"
)
assert result["team_membership"].user_id == _user_id, (
"team_membership user_id should match"
)
assert result["team_membership"].team_id == _team_id, (
"team_membership team_id should match"
)
assert result["team_membership"].budget_id == "budget_123", (
"team_membership budget_id should match"
)
assert result["team_membership"].spend == 10.5, (
"team_membership spend should match"
)
@pytest.mark.asyncio
async def test_auth_builder_with_oidc_userinfo_enabled():
"""Test that auth_builder uses OIDC UserInfo endpoint when enabled"""
from unittest.mock import MagicMock
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
# Setup test data
api_key = "test_access_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
user_object = LiteLLM_UserTable(
user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER
)
# Create JWT handler with OIDC UserInfo enabled
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
oidc_userinfo_enabled=True,
oidc_userinfo_endpoint="https://example.com/oauth2/userinfo",
user_id_jwt_field="sub",
user_email_jwt_field="email",
),
)
# Mock OIDC UserInfo response
userinfo_response = {
"sub": "test_user_1",
"email": "test@example.com",
"scope": "",
}
# Mock all the dependencies
with (
patch.object(
jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock
) as mock_get_userinfo,
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(
JWTAuthManager, "check_rbac_role", new_callable=AsyncMock
) as mock_check_rbac,
patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac,
patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes,
patch.object(
jwt_handler, "get_object_id", return_value=None
) as mock_get_object_id,
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=("test_user_1", "test@example.com", True),
) as mock_get_user_info,
patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id,
patch.object(
jwt_handler, "get_end_user_id", return_value=None
) as mock_get_end_user_id,
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
) as mock_check_admin,
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team,
patch.object(
JWTAuthManager, "get_all_team_ids", return_value=set()
) as mock_get_all_team_ids,
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team_access,
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
) as mock_get_objects,
patch.object(
JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock
) as mock_map_user,
patch.object(
JWTAuthManager, "validate_object_id", return_value=True
) as mock_validate_object,
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
) as mock_sync_user,
):
# Set up mock return values
mock_get_userinfo.return_value = userinfo_response
# Call auth_builder
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# Verify that get_oidc_userinfo was called instead of auth_jwt
mock_get_userinfo.assert_called_once_with(token=api_key)
mock_auth_jwt.assert_not_called() # Should not be called when OIDC is enabled
# Verify the result
assert result["user_id"] == "test_user_1"
assert result["user_object"] == user_object
@pytest.mark.asyncio
async def test_auth_builder_with_oidc_userinfo_disabled():
"""Test that auth_builder uses JWT validation when OIDC UserInfo is disabled"""
from unittest.mock import MagicMock
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
# Setup test data
api_key = "test_jwt_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
user_object = LiteLLM_UserTable(
user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER
)
# Create JWT handler with OIDC UserInfo disabled
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
oidc_userinfo_enabled=False, # Disabled
user_id_jwt_field="sub",
),
)
# Mock JWT validation response
jwt_response = {
"sub": "test_user_1",
"scope": "",
}
# Mock all the dependencies
with (
patch.object(
jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock
) as mock_get_userinfo,
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(
JWTAuthManager, "check_rbac_role", new_callable=AsyncMock
) as mock_check_rbac,
patch.object(jwt_handler, "get_rbac_role", return_value=None) as mock_get_rbac,
patch.object(jwt_handler, "get_scopes", return_value=[]) as mock_get_scopes,
patch.object(
jwt_handler, "get_object_id", return_value=None
) as mock_get_object_id,
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=("test_user_1", None, None),
) as mock_get_user_info,
patch.object(jwt_handler, "get_org_id", return_value=None) as mock_get_org_id,
patch.object(
jwt_handler, "get_end_user_id", return_value=None
) as mock_get_end_user_id,
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
) as mock_check_admin,
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team,
patch.object(
JWTAuthManager, "get_all_team_ids", return_value=set()
) as mock_get_all_team_ids,
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
) as mock_find_team_access,
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
) as mock_get_objects,
patch.object(
JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock
) as mock_map_user,
patch.object(
JWTAuthManager, "validate_object_id", return_value=True
) as mock_validate_object,
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
) as mock_sync_user,
):
# Set up mock return values
mock_auth_jwt.return_value = jwt_response
# Call auth_builder
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# Verify that auth_jwt was called instead of get_oidc_userinfo
mock_auth_jwt.assert_called_once_with(token=api_key)
mock_get_userinfo.assert_not_called() # Should not be called when OIDC is disabled
# Verify the result
assert result["user_id"] == "test_user_1"
assert result["user_object"] == user_object
@pytest.mark.asyncio
async def test_auth_builder_oidc_enabled_falls_back_to_jwt_auth_for_jwt_tokens():
"""
Regression test for the is_jwt routing fix.
When oidc_userinfo_enabled=True but the supplied token is a well-formed
JWT (three dot-separated parts), auth_builder must call auth_jwt and skip
get_oidc_userinfo. Sending a standard JWT to the OIDC UserInfo endpoint
is incorrect — the endpoint expects an opaque access token.
"""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
# Three-part token: recognised as a JWT by is_jwt()
api_key = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ0ZXN0X3VzZXIifQ.some_signature"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
user_object = LiteLLM_UserTable(
user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER
)
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
oidc_userinfo_enabled=True,
oidc_userinfo_endpoint="https://example.com/oauth2/userinfo",
user_id_jwt_field="sub",
),
)
jwt_response = {"sub": "test_user_1", "scope": ""}
with (
patch.object(
jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock
) as mock_get_userinfo,
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=("test_user_1", None, None),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
):
mock_auth_jwt.return_value = jwt_response
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# Token is a JWT, so standard JWT auth must be used even when
# oidc_userinfo_enabled is True.
mock_auth_jwt.assert_called_once_with(token=api_key)
mock_get_userinfo.assert_not_called()
assert result["user_id"] == "test_user_1"
assert result["user_object"] == user_object
def test_get_team_id_from_header():
"""Test get_team_id_from_header returns team when valid, None when missing, raises on invalid."""
from fastapi import HTTPException
# Valid team in allowed list
result = JWTAuthManager.get_team_id_from_header(
request_headers={"x-litellm-team-id": "team-1"},
allowed_team_ids={"team-1", "team-2"},
)
assert result == "team-1"
# No header returns None
result = JWTAuthManager.get_team_id_from_header(
request_headers={"authorization": "Bearer token"},
allowed_team_ids={"team-1"},
)
assert result is None
# Invalid team raises 403
with pytest.raises(HTTPException) as exc_info:
JWTAuthManager.get_team_id_from_header(
request_headers={"x-litellm-team-id": "invalid-team"},
allowed_team_ids={"team-1", "team-2"},
)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_auth_builder_uses_team_from_header_e2e():
"""Test auth_builder e2e flow: selects team from x-litellm-team-id header."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
user_id_jwt_field="sub",
),
)
team_object = LiteLLM_TeamTable(team_id="team-2")
user_object = LiteLLM_UserTable(
user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER
)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_team,
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
):
mock_auth_jwt.return_value = {
"sub": "user-1",
"scope": "",
"groups": ["team-1", "team-2"],
}
mock_get_team.return_value = team_object
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_headers={"x-litellm-team-id": "team-2"},
)
assert result["team_id"] == "team-2"
assert result["team_object"] == team_object
@pytest.mark.asyncio
async def test_auth_builder_header_team_denies_auth_passthrough_without_allowlist():
"""Header-selected JWT teams must enforce team allowed_passthrough_routes."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
user_id_jwt_field="sub",
),
)
team_object = LiteLLM_TeamTable(team_id="team-2", metadata={})
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=team_object,
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
) as mock_get_objects,
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
return_value=True,
),
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
return_value=False,
) as mock_passthrough_check,
):
mock_auth_jwt.return_value = {
"sub": "user-1",
"scope": "",
"groups": ["team-1", "team-2"],
}
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={},
route="/my-pass-through",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_headers={"x-litellm-team-id": "team-2"},
request_method="POST",
)
assert exc_info.value.status_code == 403
assert "allowed_passthrough_routes" in exc_info.value.detail
mock_get_objects.assert_not_called()
user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"]
assert user_api_key_dict.team_metadata == {}
@pytest.mark.asyncio
async def test_auth_builder_specific_team_denies_auth_passthrough_without_allowlist():
"""JWT-field-selected teams must enforce team allowed_passthrough_routes."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_id_jwt_field="team_id",
user_id_jwt_field="sub",
),
)
team_object = LiteLLM_TeamTable(team_id="team-1", metadata={})
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=team_object,
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
) as mock_get_objects,
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
return_value=True,
),
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
return_value=False,
) as mock_passthrough_check,
):
mock_auth_jwt.return_value = {
"sub": "user-1",
"scope": "",
"team_id": "team-1",
}
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={},
route="/my-pass-through",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_method="POST",
)
assert exc_info.value.status_code == 403
assert "allowed_passthrough_routes" in exc_info.value.detail
mock_get_objects.assert_not_called()
user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"]
assert user_api_key_dict.team_metadata == {}
@pytest.mark.asyncio
async def test_auth_builder_rbac_team_loads_team_for_passthrough_allowlist():
"""RBAC role-claim teams (team_object unset) must load team metadata before gating."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(),
)
team_object = LiteLLM_TeamTable(
team_id="team-rbac",
metadata={"allowed_passthrough_routes": ["/my-pass-through"]},
)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(jwt_handler, "get_rbac_role", return_value=LitellmUserRoles.TEAM),
patch.object(jwt_handler, "get_object_id", return_value="team-rbac"),
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=team_object,
) as mock_get_team,
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(None, None, None, None, None),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
return_value=True,
),
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
return_value=True,
) as mock_passthrough_check,
):
mock_auth_jwt.return_value = {"scope": ""}
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={},
route="/my-pass-through",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_method="POST",
)
assert result["team_id"] == "team-rbac"
mock_get_team.assert_awaited_once()
assert mock_get_team.await_args.kwargs["team_id"] == "team-rbac"
user_api_key_dict = mock_passthrough_check.call_args.kwargs["user_api_key_dict"]
assert user_api_key_dict.team_metadata == {
"allowed_passthrough_routes": ["/my-pass-through"]
}
@pytest.mark.asyncio
async def test_auth_builder_rbac_team_denies_passthrough_without_allowlist():
"""RBAC role-claim teams without an allowlist are still denied for passthrough."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(),
)
team_object = LiteLLM_TeamTable(team_id="team-rbac", metadata={})
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(jwt_handler, "get_rbac_role", return_value=LitellmUserRoles.TEAM),
patch.object(jwt_handler, "get_object_id", return_value="team-rbac"),
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=team_object,
) as mock_get_team,
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
return_value=True,
),
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.check_passthrough_route_access",
return_value=False,
),
):
mock_auth_jwt.return_value = {"scope": ""}
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={},
route="/my-pass-through",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_method="POST",
)
assert exc_info.value.status_code == 403
assert "allowed_passthrough_routes" in exc_info.value.detail
mock_get_team.assert_awaited_once()
@pytest.mark.asyncio
async def test_auth_builder_admin_on_llm_route_honors_team_header():
"""JWT proxy_admin + x-litellm-team-id on an LLM API route -> team context is
attached to the admin result so team TPM/RPM limits and attribution apply."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
user_id_jwt_field="sub",
admin_allowed_routes=[
"management_routes",
"info_routes",
"openai_routes",
],
),
)
team_object = LiteLLM_TeamTable(team_id="team-low", tpm_limit=100, rpm_limit=2)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "is_admin", return_value=True),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_team,
):
mock_auth_jwt.return_value = {
"sub": "admin-user",
"scope": "",
"groups": [],
}
mock_get_team.return_value = team_object
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_headers={"x-litellm-team-id": "team-low"},
)
assert result["is_proxy_admin"] is True
assert result["team_id"] == "team-low"
assert result["team_object"] == team_object
mock_get_team.assert_called_once()
@pytest.mark.asyncio
async def test_auth_builder_admin_on_mgmt_route_ignores_team_header():
"""JWT proxy_admin + x-litellm-team-id on an admin management route -> header
is ignored; no team fetch. Preserves pre-existing bypass behavior and avoids
phantom team creation when team_id_upsert is enabled."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
user_id_jwt_field="sub",
team_id_upsert=True,
),
)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "is_admin", return_value=True),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_team,
):
mock_auth_jwt.return_value = {
"sub": "admin-user",
"scope": "",
"groups": [],
}
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={},
general_settings={},
route="/user/info",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_headers={"x-litellm-team-id": "totally-made-up-team"},
)
assert result["is_proxy_admin"] is True
assert result["team_id"] is None
assert result["team_object"] is None
mock_get_team.assert_not_called()
@pytest.mark.asyncio
async def test_auth_builder_admin_on_llm_route_without_header_unchanged():
"""JWT proxy_admin on an LLM API route without x-litellm-team-id -> no team
context (team limits not applied, admin keeps unrestricted access)."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
user_id_jwt_field="sub",
admin_allowed_routes=[
"management_routes",
"info_routes",
"openai_routes",
],
),
)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "is_admin", return_value=True),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_team,
):
mock_auth_jwt.return_value = {
"sub": "admin-user",
"scope": "",
"groups": [],
}
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_headers={},
)
assert result["is_proxy_admin"] is True
assert result["team_id"] is None
assert result["team_object"] is None
mock_get_team.assert_not_called()
@pytest.mark.asyncio
async def test_get_team_alias_with_nested_fields():
"""
Test get_team_alias() method with nested JWT fields
"""
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
jwt_handler = JWTHandler()
# Test token with nested team name
nested_token = {
"organization": {"team": {"name": "engineering-team"}},
"team_name": "flat-team",
}
# Test nested access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_alias_jwt_field="organization.team.name"
)
assert jwt_handler.get_team_alias(nested_token, None) == "engineering-team"
# Test flat access (backward compatibility)
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="team_name")
assert jwt_handler.get_team_alias(nested_token, None) == "flat-team"
# Test missing field returns default
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_alias_jwt_field="nonexistent.field"
)
assert jwt_handler.get_team_alias(nested_token, "default-team") == "default-team"
# Test with team_alias_jwt_field not configured
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() # team_alias_jwt_field is None
assert jwt_handler.get_team_alias(nested_token, "default") is None
@pytest.mark.asyncio
async def test_is_required_team_id_with_team_alias_field():
"""
Test that is_required_team_id() returns True when team_alias_jwt_field is set
"""
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
jwt_handler = JWTHandler()
# Neither field set - should return False
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
assert jwt_handler.is_required_team_id() is False
# Only team_id_jwt_field set - should return True
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id")
assert jwt_handler.is_required_team_id() is True
# Only team_alias_jwt_field set - should return True
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="team_name")
assert jwt_handler.is_required_team_id() is True
# Both fields set - should return True
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_id_jwt_field="team_id", team_alias_jwt_field="team_name"
)
assert jwt_handler.is_required_team_id() is True
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_with_team_alias():
"""
Test that find_and_validate_specific_team_id resolves team by name when team_id is not found
"""
from unittest.mock import MagicMock
from litellm.caching import DualCache
from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_TeamTable
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(team_alias_jwt_field="team_alias"),
)
# Token with team name (no team_id)
jwt_token = {"sub": "user-1", "team_alias": "my-team"}
# Mock team object returned by get_team_object_by_alias
team_object = LiteLLM_TeamTable(team_id="resolved-team-id", team_alias="my-team")
with patch(
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias", new_callable=AsyncMock
) as mock_get_by_alias:
mock_get_by_alias.return_value = team_object
team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_token,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# Should have resolved team_id from team name
assert team_id == "resolved-team-id"
assert result_team == team_object
mock_get_by_alias.assert_called_once_with(
team_alias="my-team",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
@pytest.mark.asyncio
async def test_find_and_validate_team_id_takes_precedence_over_name():
"""
Test that team_id_jwt_field takes precedence over team_alias_jwt_field
"""
from unittest.mock import MagicMock
from litellm.caching import DualCache
from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_TeamTable
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_id_jwt_field="team_id", team_alias_jwt_field="team_alias"
),
)
# Token with both team_id and team name
jwt_token = {"sub": "user-1", "team_id": "direct-team-id", "team_alias": "my-team"}
# Mock team object returned by get_team_object (by ID)
team_object = LiteLLM_TeamTable(team_id="direct-team-id")
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_by_id,
patch(
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
new_callable=AsyncMock,
) as mock_get_by_alias,
):
mock_get_by_id.return_value = team_object
team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_token,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# Should use team_id directly, not resolve by name
assert team_id == "direct-team-id"
assert result_team == team_object
mock_get_by_id.assert_called_once()
mock_get_by_alias.assert_not_called()
@pytest.mark.asyncio
async def test_find_and_validate_raises_when_required_team_not_found():
"""
Test that an exception is raised when team is required but neither team_id nor team_name is found
"""
from litellm.caching import DualCache
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_alias_jwt_field="team_alias" # Required, but not in token
),
)
# Token without team info
jwt_token = {"sub": "user-1"}
with pytest.raises(Exception, match="No team found in token\\. Checked team_id field 'None' and") as exc_info:
await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_token,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
assert "No team found in token" in str(exc_info.value)
assert "team_alias field 'team_alias'" in str(exc_info.value)
@pytest.mark.asyncio
async def test_get_org_alias_with_nested_fields():
"""
Test get_org_alias() method with nested JWT fields
"""
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
jwt_handler = JWTHandler()
# Test token with nested org name
nested_token = {
"company": {"organization": {"name": "acme-corp"}},
"org_name": "flat-org",
}
# Test nested access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
org_alias_jwt_field="company.organization.name"
)
assert jwt_handler.get_org_alias(nested_token, None) == "acme-corp"
# Test flat access
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="org_name")
assert jwt_handler.get_org_alias(nested_token, None) == "flat-org"
# Test missing field returns default
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
org_alias_jwt_field="nonexistent.field"
)
assert jwt_handler.get_org_alias(nested_token, "default-org") == "default-org"
# Test with org_alias_jwt_field not configured
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
assert jwt_handler.get_org_alias(nested_token, "default") is None
@pytest.mark.asyncio
async def test_get_objects_resolves_org_by_name():
"""
Test that get_objects resolves organization by name when org_id is not provided
"""
from litellm.caching import DualCache
from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_OrganizationTable
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(org_alias_jwt_field="org_alias"),
)
# Mock org object returned by get_org_object_by_alias
org_object = LiteLLM_OrganizationTable(
organization_id="resolved-org-id",
organization_alias="my-org",
budget_id="budget-1",
created_by="admin",
updated_by="admin",
models=[],
)
with patch(
"litellm.proxy.auth.handle_jwt.get_org_object_by_alias", new_callable=AsyncMock
) as mock_get_by_alias:
mock_get_by_alias.return_value = org_object
(
result_user_obj,
result_org_obj,
result_end_user_obj,
result_team_membership,
_result_user_id,
) = await JWTAuthManager.get_objects(
user_id=None,
user_email=None,
org_id=None, # No org_id provided
end_user_id=None,
team_id=None,
valid_user_email=None,
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
route="/chat/completions",
org_alias="my-org",
)
# Should resolve org by alias - org_id can be derived from org_object.organization_id
assert result_org_obj == org_object
assert result_org_obj.organization_id == "resolved-org-id"
mock_get_by_alias.assert_called_once_with(
org_alias="my-org",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# ---------------------------------------------------------------------------
# Fix 1: OIDC discovery URL resolution
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_resolve_jwks_url_passthrough_for_direct_jwks_url():
"""Non-discovery URLs are returned unchanged."""
from unittest.mock import AsyncMock, MagicMock
from litellm.caching.dual_cache import DualCache
handler = JWTHandler()
handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
litellm_jwtauth=LiteLLM_JWTAuth(),
)
url = "https://login.microsoftonline.com/common/discovery/keys"
result = await handler._resolve_jwks_url(url)
assert result == url
@pytest.mark.asyncio
async def test_resolve_jwks_url_resolves_oidc_discovery_document():
"""
A .well-known/openid-configuration URL should be fetched and its
jwks_uri returned.
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.caching.dual_cache import DualCache
handler = JWTHandler()
cache = DualCache()
handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(),
)
discovery_url = (
"https://login.microsoftonline.com/tenant/.well-known/openid-configuration"
)
jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys"
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"jwks_uri": jwks_url, "issuer": "https://..."}
mock_get = AsyncMock(return_value=mock_response)
handler.http_handler.get = mock_get
result = await handler._resolve_jwks_url(discovery_url)
assert result == jwks_url
mock_get.assert_called_once_with(discovery_url)
@pytest.mark.asyncio
async def test_resolve_jwks_url_caches_resolved_jwks_uri():
"""Resolved jwks_uri is cached — second call does not hit the network."""
from unittest.mock import AsyncMock, MagicMock
from litellm.caching.dual_cache import DualCache
handler = JWTHandler()
cache = DualCache()
handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(),
)
discovery_url = (
"https://login.microsoftonline.com/tenant/.well-known/openid-configuration"
)
jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys"
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"jwks_uri": jwks_url}
mock_get = AsyncMock(return_value=mock_response)
handler.http_handler.get = mock_get
first = await handler._resolve_jwks_url(discovery_url)
second = await handler._resolve_jwks_url(discovery_url)
assert first == jwks_url
assert second == jwks_url
# Network should only be hit once
assert mock_get.call_count == 1
@pytest.mark.asyncio
async def test_resolve_jwks_url_raises_if_no_jwks_uri_in_discovery_doc():
"""Raise a helpful error if the discovery document has no jwks_uri."""
from unittest.mock import AsyncMock, MagicMock
from litellm.caching.dual_cache import DualCache
handler = JWTHandler()
handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
litellm_jwtauth=LiteLLM_JWTAuth(),
)
discovery_url = "https://example.com/.well-known/openid-configuration"
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"issuer": "https://example.com"} # no jwks_uri
handler.http_handler.get = AsyncMock(return_value=mock_response)
with pytest.raises(Exception, match="jwks_uri"):
await handler._resolve_jwks_url(discovery_url)
# ---------------------------------------------------------------------------
# Fix 2: handle array values in team_id_jwt_field (e.g. AAD "roles" claim)
# ---------------------------------------------------------------------------
def _make_jwt_handler(team_id_jwt_field: str) -> JWTHandler:
from litellm.caching.dual_cache import DualCache
handler = JWTHandler()
handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
litellm_jwtauth=LiteLLM_JWTAuth(team_id_jwt_field=team_id_jwt_field),
)
return handler
def test_get_team_id_returns_first_element_when_roles_is_list():
"""
AAD sends roles as a list. get_team_id() must return the first string
element rather than the raw list (which would later crash with
'unhashable type: list').
"""
handler = _make_jwt_handler("roles")
token = {"oid": "user-oid", "roles": ["team1"]}
result = handler.get_team_id(token=token, default_value=None)
assert result == "team1"
def test_get_team_id_returns_first_element_from_multi_value_roles_list():
"""When roles has multiple entries, the first one is used."""
handler = _make_jwt_handler("roles")
token = {"roles": ["team2", "team1"]}
result = handler.get_team_id(token=token, default_value=None)
assert result == "team2"
def test_get_team_id_returns_default_when_roles_list_is_empty():
"""Empty list should fall back to default_value."""
handler = _make_jwt_handler("roles")
token = {"roles": []}
result = handler.get_team_id(token=token, default_value="fallback")
assert result == "fallback"
def test_get_team_id_still_works_with_string_value():
"""String values (non-array) continue to work as before."""
handler = _make_jwt_handler("appid")
token = {"appid": "my-team-id"}
result = handler.get_team_id(token=token, default_value=None)
assert result == "my-team-id"
def test_get_team_id_list_result_is_hashable():
"""
The value returned by get_team_id() must be hashable so it can be
added to a set (the operation that previously crashed).
"""
handler = _make_jwt_handler("roles")
token = {"roles": ["team1"]}
result = handler.get_team_id(token=token, default_value=None)
# This must not raise TypeError
s: set = set()
s.add(result)
assert "team1" in s
# ---------------------------------------------------------------------------
# Fix 3: helpful error message for dot-notation array indexing (roles.0)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_hints_bracket_notation():
"""
When team_id_jwt_field is set to 'roles.0' (unsupported dot-notation for
array indexing) and no team is found, the exception message should suggest
using 'roles' instead (and explain LiteLLM auto-unwraps list values).
"""
from unittest.mock import MagicMock
from litellm.caching.dual_cache import DualCache
handler = _make_jwt_handler("roles.0")
# token has roles as a list — dot-notation won't find anything
token = {"roles": ["team1"]}
with pytest.raises(Exception, match="is not supported\\. Use 'roles' instead — LiteLLM") as exc_info:
await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=handler,
jwt_valid_token=token,
prisma_client=None,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
error_msg = str(exc_info.value)
# Should mention the bad field name and suggest the fix
assert "roles.0" in error_msg, f"Expected field name in: {error_msg}"
assert "roles" in error_msg and "list" in error_msg, (
f"Expected hint about using 'roles' instead: {error_msg}"
)
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_hints_bracket_index_notation():
"""
When team_id_jwt_field is set to 'roles[0]' (bracket indexing, also unsupported
in get_nested_value) the error message should suggest using 'roles' instead.
"""
from unittest.mock import MagicMock
from litellm.caching.dual_cache import DualCache
handler = _make_jwt_handler("roles[0]")
token = {"roles": ["team1"]}
with pytest.raises(Exception, match="is not supported in team_id_jwt_field\\. Use 'roles' instead") as exc_info:
await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=handler,
jwt_valid_token=token,
prisma_client=None,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
error_msg = str(exc_info.value)
assert "roles[0]" in error_msg, f"Expected field name in: {error_msg}"
assert "roles" in error_msg and "list" in error_msg, (
f"Expected hint about using 'roles' instead: {error_msg}"
)
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_no_hint_for_valid_field():
"""
When team_id_jwt_field is a normal field name (no dot-notation) the
error message should not contain a spurious bracket-notation hint.
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.caching.dual_cache import DualCache
handler = _make_jwt_handler("appid")
token = {} # no appid — triggers the "no team found" path
with pytest.raises(Exception, match="No team found in token\\. Checked team_id field 'appid' and") as exc_info:
await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=handler,
jwt_valid_token=token,
prisma_client=None,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
error_msg = str(exc_info.value)
assert "Hint" not in error_msg
# ---------------------------------------------------------------------------
# Single-team DB fallback when JWT does not resolve team_id
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
(
"user_id",
"user_teams",
"get_team_object_return",
"expected_team_id",
"expect_get_team_called",
"expect_get_membership_called",
),
[
pytest.param(
"user_single_team_fb",
["team_only_fb"],
"resolved_row",
"team_only_fb",
True,
True,
id="one_db_team_resolves_team_and_membership",
),
pytest.param(
"user_multi_team_fb",
["team_a", "team_b"],
"unused",
None,
False,
False,
id="two_db_teams_ambiguous_no_fallback",
),
pytest.param(
"user_zero_teams_fb",
[],
"unused",
None,
False,
False,
id="zero_db_teams_no_fallback",
),
pytest.param(
"user_orphan_team_fb",
["team_missing_in_db"],
"http_404",
None,
True,
False,
id="one_team_id_but_row_missing_in_db",
),
pytest.param(
"user_orphan_team_non404_fb",
["team_err"],
"http_500",
None,
True,
False,
id="get_team_object_raises_non404_still_no_raise",
),
],
)
@pytest.mark.asyncio
async def test_auth_builder_single_team_db_fallback_when_jwt_has_no_team(
user_id: str,
user_teams: list,
get_team_object_return: Optional[str],
expected_team_id: Optional[str],
expect_get_team_called: bool,
expect_get_membership_called: bool,
) -> None:
"""
JWT does not set team_id (mocks return no team from token/header/routing). Behavior:
- exactly one team on user + get_team_object returns a row -> set team + membership
- two+ teams, or zero teams -> no get_team_object / no membership
- one team id but get_team_object raises (e.g. 404/500) -> skip fallback, no team, no error
"""
if len(user_teams) == 1 and get_team_object_return == "resolved_row":
only = user_teams[0]
team_table = LiteLLM_TeamTable(team_id=only)
membership = LiteLLM_TeamMembership(
user_id=user_id, team_id=only, litellm_budget_table=None
)
get_team_return_value = team_table
membership_return_value = membership
else:
team_table = None
membership = None
get_team_return_value = None
membership_return_value = None
# "http_404" / "http_500" use get_team_object.side_effect, not return_value
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=user_teams,
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
) as mock_get_team,
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
) as mock_get_membership,
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
if get_team_object_return in ("http_404", "http_500"):
from fastapi import HTTPException
code = 404 if get_team_object_return == "http_404" else 500
mock_get_team.side_effect = HTTPException(
status_code=code,
detail={
"error": f"Team doesn't exist in db. Team={user_teams[0]}. Create team via `/team/new` call."
},
)
else:
mock_get_team.return_value = get_team_return_value
if membership_return_value is not None:
mock_get_membership.return_value = membership_return_value
result = await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result["team_id"] == expected_team_id
if expected_team_id is not None:
assert result["team_object"] == team_table
assert result["team_membership"] == membership
else:
assert result["team_object"] is None
if not expect_get_membership_called:
assert result["team_membership"] is None
if expect_get_team_called:
mock_get_team.assert_called()
else:
mock_get_team.assert_not_called()
if expect_get_membership_called:
mock_get_membership.assert_called_once()
else:
mock_get_membership.assert_not_called()
@pytest.mark.asyncio
async def test_auth_builder_single_team_fallback_membership_error_skips_no_raise():
"""
get_team_object succeeds but get_team_membership raises — do not set team; no exception.
"""
from fastapi import HTTPException
user_id = "u_mem_fail"
team_id_val = "team_mem_fail"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=[team_id_val],
)
team_table = LiteLLM_TeamTable(team_id=team_id_val)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
) as mock_get_team,
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
) as mock_get_membership,
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
mock_get_team.return_value = team_table
mock_get_membership.side_effect = HTTPException(
status_code=500, detail="membership lookup failed"
)
result = await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result["team_id"] is None
assert result["team_object"] is None
assert result["team_membership"] is None
mock_get_team.assert_called()
mock_get_membership.assert_called_once()
# ---------------------------------------------------------------------------
# JWTHandler._build_decode_kwargs — VERIA-27 (audience + issuer verification)
# ---------------------------------------------------------------------------
@pytest.fixture(autouse=False)
def _reset_unscoped_warning_flag():
"""Reset the once-per-process warning sentinel so each test sees a fresh
state."""
JWTHandler._unscoped_jwt_warning_emitted = False
yield
JWTHandler._unscoped_jwt_warning_emitted = False
def test_build_decode_kwargs_no_env_disables_both_verifications(
monkeypatch, _reset_unscoped_warning_flag
):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_ISSUER", raising=False)
kwargs = JWTHandler._build_decode_kwargs()
assert kwargs["audience"] is None
assert kwargs["issuer"] is None
assert kwargs["options"] == {"verify_aud": False, "verify_iss": False}
def test_build_decode_kwargs_audience_only_enables_aud_verification(
monkeypatch, _reset_unscoped_warning_flag
):
monkeypatch.setenv("JWT_AUDIENCE", "my-proxy")
monkeypatch.delenv("JWT_ISSUER", raising=False)
kwargs = JWTHandler._build_decode_kwargs()
assert kwargs["audience"] == "my-proxy"
assert kwargs["issuer"] is None
# verify_aud not in options means PyJWT will verify audience
assert kwargs["options"] == {"verify_iss": False}
def test_build_decode_kwargs_issuer_only_enables_iss_verification(
monkeypatch, _reset_unscoped_warning_flag
):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.setenv("JWT_ISSUER", "https://idp.example.com/")
kwargs = JWTHandler._build_decode_kwargs()
assert kwargs["audience"] is None
assert kwargs["issuer"] == "https://idp.example.com/"
assert kwargs["options"] == {"verify_aud": False}
def test_build_decode_kwargs_both_set_enables_full_verification(
monkeypatch, _reset_unscoped_warning_flag
):
monkeypatch.setenv("JWT_AUDIENCE", "my-proxy")
monkeypatch.setenv("JWT_ISSUER", "https://idp.example.com/")
kwargs = JWTHandler._build_decode_kwargs()
assert kwargs["audience"] == "my-proxy"
assert kwargs["issuer"] == "https://idp.example.com/"
# No verification opt-outs — PyJWT verifies both claims by default.
assert kwargs["options"] is None
def test_build_decode_kwargs_warns_once_when_unscoped(
monkeypatch, _reset_unscoped_warning_flag, caplog
):
"""The warning about unscoped JWT auth should fire on the first call but
not on every subsequent decode."""
import logging
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_ISSUER", raising=False)
caplog.set_level(logging.WARNING)
JWTHandler._build_decode_kwargs()
JWTHandler._build_decode_kwargs()
JWTHandler._build_decode_kwargs()
matching = [
r
for r in caplog.records
if "JWT auth is enabled" in r.getMessage()
and "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage()
]
assert len(matching) == 1, (
f"Expected exactly one warning across 3 calls, got {len(matching)}"
)
def test_build_decode_kwargs_no_warning_when_scoped(
monkeypatch, _reset_unscoped_warning_flag, caplog
):
import logging
monkeypatch.setenv("JWT_AUDIENCE", "my-proxy")
monkeypatch.delenv("JWT_ISSUER", raising=False)
caplog.set_level(logging.WARNING)
JWTHandler._build_decode_kwargs()
matching = [
r
for r in caplog.records
if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage()
]
assert matching == []
# ---------------------------------------------------------------------------
# Defer to single-team DB fallback (PR #26418) when JWT claims are present
# but do not resolve to a LiteLLM team.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_unresolved_claim_returns_none():
"""With `team_claim_fallback=True`: team_id claim is present in the JWT
but the team is missing in the DB — return (None, None) so the
auth_builder single-team fallback can run, instead of raising and
failing auth."""
from fastapi import HTTPException
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_id_jwt_field="team_id",
team_claim_fallback=True,
)
token = {"sub": "user-1", "team_id": "claim-team-not-in-db"}
with patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
) as mock_get_team:
mock_get_team.side_effect = HTTPException(status_code=404, detail="missing")
team_id, team_object = await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=token,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert team_id is None
assert team_object is None
@pytest.mark.asyncio
async def test_find_team_with_model_access_unresolved_group_claim_returns_none(
monkeypatch,
):
"""With `team_claim_fallback=True`: group claim resolves to team_ids that
don't exist in the DB — return (None, None) instead of raising 403, so
the single-team fallback can run."""
import sys
import types
from fastapi import HTTPException
from litellm.router import Router
router = Router(
model_list=[
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}
]
)
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
async def raise_404(*_args, **_kwargs):
raise HTTPException(status_code=404, detail="missing")
monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", raise_404)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_claim_fallback=True)
team_id, team_object = await JWTAuthManager.find_team_with_model_access(
team_ids={"idp-group-a", "idp-group-b"},
requested_model="gpt-4o-mini",
route="/chat/completions",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert team_id is None
assert team_object is None
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_non_http_exception_still_propagates():
"""Regression guard: only the 404 HTTPException raised by
`get_team_object` ("team doesn't exist in db") is softened. Other
errors — e.g. "No DB Connected" — must still propagate so operator-side
problems are loud."""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id")
token = {"sub": "user-1", "team_id": "some-claim-team"}
with patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
) as mock_get_team:
mock_get_team.side_effect = RuntimeError("simulated infrastructure error")
with pytest.raises(RuntimeError, match="simulated infrastructure error"):
await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=token,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_non_404_http_exception_propagates():
"""Regression guard: only 404 HTTPException is softened. If
`get_team_object` is ever updated to raise a different HTTP status code
(e.g. 403 for a blocked team), that error must still propagate rather
than silently fall through to the single-team DB fallback."""
from fastapi import HTTPException
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id")
token = {"sub": "user-1", "team_id": "some-claim-team"}
for status_code in (400, 403, 500):
with patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
) as mock_get_team:
mock_get_team.side_effect = HTTPException(
status_code=status_code, detail="non-404 failure"
)
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=token,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert exc_info.value.status_code == status_code
@pytest.mark.asyncio
async def test_find_team_with_model_access_enforce_team_based_access_still_raises():
"""Regression guard: when no group claims are present and
`enforce_team_based_model_access` is on, the original 403 still fires —
the new soft-fail only applies to the unresolved-claim path inside the
loop, not to the no-team-claims-at-all path at the top."""
from fastapi import HTTPException
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(enforce_team_based_model_access=True)
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.find_team_with_model_access(
team_ids=set(),
requested_model="gpt-4o-mini",
route="/chat/completions",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert exc_info.value.status_code == 403
assert "enforce_team_based_model_access" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_find_team_with_model_access_resolved_team_without_model_still_raises_403(
monkeypatch,
):
"""Regression guard: when the JWT group claim DOES resolve to a real
LiteLLM team but that team does not grant the requested model, keep the
original 403. Only the unresolved-claim case is softened."""
import sys
import types
from fastapi import HTTPException
from litellm.router import Router
router = Router(
model_list=[
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "gpt-3.5-turbo"},
},
]
)
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
team = LiteLLM_TeamTable(team_id="real-team", models=["gpt-3.5-turbo"])
async def mock_get_team_object(*_args, **_kwargs):
return team
monkeypatch.setattr(
"litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.find_team_with_model_access(
team_ids={"real-team"},
requested_model="gpt-4o-mini",
route="/chat/completions",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert exc_info.value.status_code == 403
assert "No team has access to the requested model" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_unresolved_claim_default_raises():
"""Default `team_claim_fallback=False`: unresolved team_id claim must
still raise — preserves the strict claim-based authorization boundary
when the operator has not opted in to the fallback."""
from fastapi import HTTPException
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id")
token = {"sub": "user-1", "team_id": "claim-team-not-in-db"}
with patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
) as mock_get_team:
mock_get_team.side_effect = HTTPException(status_code=404, detail="missing")
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=token,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert exc_info.value.status_code == 404
@pytest.mark.asyncio
async def test_find_team_with_model_access_unresolved_group_claim_default_raises(
monkeypatch,
):
"""Default `team_claim_fallback=False`: group claims that don't resolve
to any LiteLLM team must still raise 403 — preserves the strict
claim-based authorization boundary."""
import sys
import types
from fastapi import HTTPException
from litellm.router import Router
router = Router(
model_list=[
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}
]
)
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
async def raise_404(*_args, **_kwargs):
raise HTTPException(status_code=404, detail="missing")
monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", raise_404)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.find_team_with_model_access(
team_ids={"idp-group-a", "idp-group-b"},
requested_model="gpt-4o-mini",
route="/chat/completions",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert exc_info.value.status_code == 403
# GH #26789: JWT claim user_id must rebind to legacy DB row after fuzzy match.
def test_canonical_user_id_rebinds_to_legacy_uuid():
"""JWT email resolves to a legacy UUID row -> use the UUID for attribution."""
legacy_uuid = "bb8ab11f-09aa-47ae-b063-6e80506ac3bc"
jwt_email = "matt@example.com"
user_object = LiteLLM_UserTable(user_id=legacy_uuid, user_email=jwt_email)
assert (
JWTAuthManager._canonical_user_id_from_db(
user_id=jwt_email, user_object=user_object
)
== legacy_uuid
)
def test_canonical_user_id_no_change_when_ids_match():
"""Fresh upserted user (row.user_id == claim) -> claim returned unchanged."""
same = "alice@example.com"
user_object = LiteLLM_UserTable(user_id=same, user_email=same)
assert (
JWTAuthManager._canonical_user_id_from_db(user_id=same, user_object=user_object)
== same
)
def test_canonical_user_id_returns_claim_when_no_user_object():
"""No resolved row (e.g. upsert disabled / brand new) -> keep the claim."""
assert (
JWTAuthManager._canonical_user_id_from_db(
user_id="newcomer@example.com", user_object=None
)
== "newcomer@example.com"
)
def test_canonical_user_id_returns_none_when_claim_none_and_no_object():
"""Defensive: no claim and no row -> stays None, never invents an id."""
assert (
JWTAuthManager._canonical_user_id_from_db(user_id=None, user_object=None)
is None
)
def test_canonical_user_id_no_change_when_db_user_id_falsy():
"""Defensive: an empty user_object.user_id must not clobber the claim."""
class _Stub:
user_id = ""
assert (
JWTAuthManager._canonical_user_id_from_db(
user_id="jwt@example.com", user_object=_Stub()
)
== "jwt@example.com"
)
@pytest.mark.asyncio
async def test_auth_jwt_expired_token_raises_401_jwk_path():
"""An expired JWT (access token) decoded via the JWK/dict public-key path
must raise a ProxyException carrying a 401 status code so the status is
preserved end-to-end (client response + OTel traces).
"""
import jwt as jwt_lib
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
with (
patch.object(
jwt_handler, "get_public_key", new_callable=AsyncMock
) as mock_get_public_key,
patch(
"litellm.proxy.auth.handle_jwt.jwt.get_unverified_header",
return_value={"kid": "test-kid"},
),
patch(
"litellm.proxy.auth.handle_jwt.PyJWK.from_dict",
return_value=MagicMock(key="fake-key"),
),
patch(
"litellm.proxy.auth.handle_jwt.jwt.decode",
side_effect=jwt_lib.ExpiredSignatureError("Signature has expired"),
),
):
mock_get_public_key.return_value = {"kty": "RSA", "kid": "test-kid"}
with pytest.raises(ProxyException) as exc_info:
await jwt_handler.auth_jwt(token="expired.jwt.token")
assert exc_info.value.code == str(401)
assert exc_info.value.type == ProxyErrorTypes.expired_key.value
assert "Token Expired" in exc_info.value.message
@pytest.mark.asyncio
async def test_auth_jwt_expired_token_raises_401_pem_cert_path():
"""Same as above but for the PEM-certificate (string public-key) decode path."""
import jwt as jwt_lib
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
mock_cert = MagicMock()
mock_cert.public_key.return_value.public_bytes.return_value = b"fake-key"
with (
patch.object(
jwt_handler, "get_public_key", new_callable=AsyncMock
) as mock_get_public_key,
patch(
"litellm.proxy.auth.handle_jwt.jwt.get_unverified_header",
return_value={"kid": "test-kid"},
),
patch(
"litellm.proxy.auth.handle_jwt.x509.load_pem_x509_certificate",
return_value=mock_cert,
),
patch(
"litellm.proxy.auth.handle_jwt.jwt.decode",
side_effect=jwt_lib.ExpiredSignatureError("Signature has expired"),
),
):
mock_get_public_key.return_value = (
"-----BEGIN CERTIFICATE-----\nfake\n-----END CERTIFICATE-----"
)
with pytest.raises(ProxyException) as exc_info:
await jwt_handler.auth_jwt(token="expired.jwt.token")
assert exc_info.value.code == str(401)
assert exc_info.value.type == ProxyErrorTypes.expired_key.value
assert "Token Expired" in exc_info.value.message
def _base64url_encode_int(value: int) -> str:
import base64
value_bytes = value.to_bytes((value.bit_length() + 7) // 8, "big")
return base64.urlsafe_b64encode(value_bytes).decode("utf-8").rstrip("=")
def _get_rsa_key_and_jwk(kid: str):
from cryptography.hazmat.primitives.asymmetric import rsa
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
public_numbers = private_key.public_key().public_numbers()
jwk = {
"kty": "RSA",
"n": _base64url_encode_int(value=public_numbers.n),
"e": _base64url_encode_int(value=public_numbers.e),
"kid": kid,
"alg": "RS256",
"use": "sig",
}
return private_key, jwk
def _encode_rsa_jwt(
private_key,
issuer: str,
audience: str,
kid: str,
extra_claims: Optional[dict] = None,
) -> str:
import time
import jwt
from cryptography.hazmat.primitives import serialization
private_key_pem = private_key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
current_time = int(time.time())
claims = {
"sub": "test-subject",
"iss": issuer,
"aud": audience,
"iat": current_time,
"exp": current_time + 300,
}
if extra_claims:
claims.update(extra_claims)
return jwt.encode(
claims,
private_key_pem,
algorithm="RS256",
headers={"kid": kid},
)
def _get_jwt_handler_with_issuer_keys(issuers: list, keys_by_url: dict) -> JWTHandler:
from litellm.caching.dual_cache import DualCache
cache = DualCache()
for jwks_url, keys in keys_by_url.items():
cache.set_cache(
key=f"litellm_jwt_auth_keys_{jwks_url}",
value=keys,
)
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(issuers=issuers),
)
return jwt_handler
@pytest.mark.asyncio
async def test_get_public_key_fetches_and_caches_jwks_response():
from unittest.mock import AsyncMock, MagicMock
from litellm.caching.dual_cache import DualCache
jwt_handler = JWTHandler()
cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(public_key_ttl=123),
)
expected_key_id = "cached-key"
_, jwk = _get_rsa_key_and_jwk(kid=expected_key_id)
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"keys": [jwk]}
jwt_handler.http_handler.get = AsyncMock(return_value=mock_response)
public_key = await jwt_handler._get_public_key_from_jwks_url(
jwks_url="https://issuer.example.com/keys",
kid=expected_key_id,
)
assert public_key == jwk
cached_keys = await cache.async_get_cache(
key="litellm_jwt_auth_keys_https://issuer.example.com/keys"
)
assert cached_keys == [jwk]
class _ScriptedJWKSEndpoint:
"""Injected stand-in for ``JWTHandler.http_handler`` with scripted per-call outcomes.
Each outcome is either an exception to raise or a JSON body to return; the
last outcome repeats for any further calls.
"""
def __init__(
self,
outcomes: Sequence[Exception | Mapping[str, object] | MagicMock],
delay: float = 0.0,
) -> None:
self.outcomes = outcomes
self.delay = delay
self.call_count = 0
async def get(
self,
url: str,
params: Mapping[str, str] | None = None,
headers: Mapping[str, str] | None = None,
) -> MagicMock:
self.call_count += 1
if self.delay:
await asyncio.sleep(self.delay)
outcome = self.outcomes[min(self.call_count - 1, len(self.outcomes) - 1)]
if isinstance(outcome, Exception):
raise outcome
if isinstance(outcome, MagicMock):
return outcome
response = MagicMock()
response.status_code = 200
response.json.return_value = outcome
return response
def _get_jwt_handler_with_scripted_endpoint(
cache: "DualCache",
endpoint: _ScriptedJWKSEndpoint,
public_key_ttl: float = 600,
public_key_stale_ttl: float = DEFAULT_JWKS_STALE_TTL,
) -> JWTHandler:
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(
public_key_ttl=public_key_ttl,
public_key_stale_ttl=public_key_stale_ttl,
),
)
jwt_handler.http_handler = endpoint
return jwt_handler
@pytest.mark.asyncio
async def test_get_public_key_retries_transient_jwks_fetch_failure():
"""A single connect timeout to the IdP must be retried, not surfaced to the caller."""
from litellm.caching.dual_cache import DualCache
_, jwk = _get_rsa_key_and_jwk(kid="retried-key")
endpoint = _ScriptedJWKSEndpoint((httpx.ConnectTimeout("connect timed out"), {"keys": [jwk]}))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(DualCache(), endpoint)
public_key = await jwt_handler._get_public_key_from_jwks_url(
jwks_url="https://issuer.example.com/keys",
kid="retried-key",
)
assert public_key == jwk
assert endpoint.call_count == 2
@pytest.mark.asyncio
async def test_get_public_key_serves_stale_keys_when_jwks_refresh_fails():
"""Once the TTL lapses, an unreachable IdP must not invalidate a still-valid signing key."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://issuer.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="stale-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="stale-key") == jwk
await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}")
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
public_key = await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="stale-key")
assert public_key == jwk
@pytest.mark.asyncio
async def test_stale_jwks_window_is_the_configured_grace_past_a_long_public_key_ttl():
"""The stale window is `public_key_stale_ttl` past the active entry, whatever `public_key_ttl` is set to.
Deriving the window from `public_key_ttl` instead would collapse it to nothing on the long TTLs that
make the fallback worth having.
"""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://long-ttl-issuer.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="long-ttl-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(
cache,
endpoint,
public_key_ttl=90000,
public_key_stale_ttl=3600,
)
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="long-ttl-key")
active_key = f"litellm_jwt_auth_keys_{jwks_url}"
active_deadline = cache.in_memory_cache.ttl_dict[active_key]
stale_deadline = cache.in_memory_cache.ttl_dict[f"{STALE_CACHE_KEY_PREFIX}{active_key}"]
assert stale_deadline - active_deadline == pytest.approx(3600, abs=1)
@pytest.mark.asyncio
async def test_long_public_key_ttl_still_serves_stale_keys_when_the_idp_is_unreachable():
"""A long `public_key_ttl` must not leave the stale fallback inert once that TTL finally lapses."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://long-ttl-fallback.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="long-ttl-fallback-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=604800)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="long-ttl-fallback-key") == jwk
await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}")
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="long-ttl-fallback-key") == jwk
@pytest.mark.asyncio
async def test_removed_signing_key_stops_being_trusted_once_the_stale_window_expires(monkeypatch):
"""The stale fallback is bounded: past its window a key the IdP dropped is no longer served."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://revoking-issuer.example.com/keys"
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url)
_, jwk = _get_rsa_key_and_jwk(kid="revoked-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
assert await jwt_handler.get_public_key(kid="revoked-key") == jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
assert await jwt_handler.get_public_key(kid="revoked-key") == jwk
await cache.async_delete_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}")
with pytest.raises(ProxyException) as exc_info:
await jwt_handler.get_public_key(kid="revoked-key")
assert exc_info.value.code == "503"
assert exc_info.value.type == ProxyErrorTypes.auth_provider_unavailable
@pytest.mark.asyncio
async def test_key_removed_from_a_reachable_jwks_is_rejected_without_consulting_the_stale_copy():
"""A reachable IdP always wins: dropping a key revokes it immediately, stale copy included."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://rotating-issuer.example.com/keys"
_, retired_jwk = _get_rsa_key_and_jwk(kid="retired-key")
_, current_jwk = _get_rsa_key_and_jwk(kid="current-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [retired_jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="retired-key") == retired_jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
endpoint.outcomes = ({"keys": [current_jwk]},)
with pytest.raises(NoMatchingJWTPublicKeyError):
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="retired-key")
assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [current_jwk]
@pytest.mark.asyncio
async def test_zero_public_key_stale_ttl_fails_closed_instead_of_serving_stale_keys():
"""`public_key_stale_ttl=0` is the escape hatch for deployments that cannot trust an unrefreshed key."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://fail-closed-issuer.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="fail-closed-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_stale_ttl=0)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="fail-closed-key") == jwk
assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}litellm_jwt_auth_keys_{jwks_url}") is None
await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}")
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
with pytest.raises(JWKSUnreachableError):
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="fail-closed-key")
@pytest.mark.asyncio
@pytest.mark.parametrize("lowered_stale_ttl", [0, 30])
async def test_lowering_public_key_stale_ttl_stops_serving_a_copy_cached_under_the_old_setting(lowered_stale_ttl):
"""Lowering the window has to bite immediately: an operator does this mid-incident, on a shared cache.
The stale entry keeps whatever expiry it was written with, so enforcing the bound only at write time would
leave a copy taken under the old, longer setting servable until it aged out on its own.
"""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://relaxed-then-tightened.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="tightened-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
generous = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_stale_ttl=86400)
assert await generous._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="tightened-key") == jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [jwk]
# The operator tightens the window and restarts; the cache, and its long-lived copy, survive.
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
tightened = _get_jwt_handler_with_scripted_endpoint(
cache, endpoint, public_key_stale_ttl=lowered_stale_ttl
)
await cache.async_set_cache(
key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}",
value=time.time() - 7200,
ttl=86400,
)
with pytest.raises(JWKSUnreachableError):
await tightened._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="tightened-key")
@pytest.mark.asyncio
async def test_zero_public_key_stale_ttl_fails_closed_even_for_a_freshly_written_copy():
"""`0` must fail closed on its own, not merely because the copy happens to be older than `public_key_ttl`.
The active entry can disappear before it expires, through cache eviction or a flush, which leaves a stale
copy younger than `public_key_ttl`. Bounding only on age would still serve it.
"""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://evicted-active-entry.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="fresh-copy-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
generous = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=600, public_key_stale_ttl=3600)
assert await generous._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="fresh-copy-key") == jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
written_at = await cache.async_get_cache(key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}")
assert time.time() - written_at < 600
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
fail_closed = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=600, public_key_stale_ttl=0)
with pytest.raises(JWKSUnreachableError):
await fail_closed._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="fresh-copy-key")
@pytest.mark.asyncio
async def test_stale_copy_with_no_recorded_write_time_is_not_served():
"""The bound is enforced from the recorded write time, so losing it must fail closed, never open."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://undated-copy.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="undated-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="undated-key") == jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
await cache.async_delete_cache(key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}")
assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [jwk]
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
with pytest.raises(JWKSUnreachableError):
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="undated-key")
@pytest.mark.asyncio
async def test_increasing_public_key_stale_ttl_only_extends_within_the_new_bound():
"""Raising the window re-measures from the copy's refresh time; it does not bless whatever is cached."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://widened-window.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="widened-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
narrow = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=600, public_key_stale_ttl=60)
assert await narrow._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="widened-key") == jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
written_at_key = f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}"
await cache.async_delete_cache(key=active_cache_key)
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
widened = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=600, public_key_stale_ttl=3600)
# Older than the widened bound of 600 + 3600, so widening must not revive it.
await cache.async_set_cache(key=written_at_key, value=time.time() - 5000, ttl=86400)
with pytest.raises(JWKSUnreachableError):
await widened._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="widened-key")
# Inside the widened bound, so it is servable again.
await cache.async_set_cache(key=written_at_key, value=time.time() - 1000, ttl=86400)
assert await widened._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="widened-key") == jwk
@pytest.mark.asyncio
async def test_stale_copy_written_at_survives_a_whole_number_epoch():
"""A Redis JSON round-trip can return the epoch as an int, and that must not read as a missing timestamp.
Rejecting it would fail closed on a copy that is well inside the window, in the shared-cache deployment
the stale fallback exists to serve.
"""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://int-epoch.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="int-epoch-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="int-epoch-key") == jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
await cache.async_set_cache(
key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}",
value=int(time.time()) - 60,
ttl=86400,
)
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="int-epoch-key") == jwk
@pytest.mark.asyncio
async def test_stale_copy_with_a_malformed_write_time_is_not_served():
"""An unreadable refresh timestamp is indistinguishable from an unbounded one, so it fails closed."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://malformed-timestamp.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="malformed-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="malformed-key") == jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
await cache.async_set_cache(
key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}",
value="whenever",
ttl=86400,
)
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
with pytest.raises(JWKSUnreachableError):
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="malformed-key")
@pytest.mark.asyncio
async def test_public_key_stale_ttl_defaults_to_one_hour():
"""The default is the exposure bound for a key the IdP revoked mid-outage, so it stays short deliberately."""
assert LiteLLM_JWTAuth().public_key_stale_ttl == 3600
@pytest.mark.asyncio
async def test_stale_fallback_warns_with_the_kid_and_how_stale_the_jwks_copy_is(caplog):
"""Serving an unrefreshed signing key is a security-relevant event, so it must be legible in the logs."""
import logging
from litellm.caching.dual_cache import DualCache
jwks_url = "https://warned-issuer.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="warned-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_stale_ttl=1800)
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="warned-key")
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
await cache.async_set_cache(
key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}",
value=time.time() - 120,
ttl=600,
)
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
caplog.set_level(logging.WARNING)
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="warned-key")
warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING]
stale_warnings = [m for m in warnings if "stale JWKS copy" in m]
assert len(stale_warnings) == 1
assert "kid=warned-key" in stale_warnings[0]
assert jwks_url in stale_warnings[0]
freshness = re.search(r"last refreshed (\d+)s ago, stops being trusted in (\d+)s", stale_warnings[0])
assert freshness is not None
age, remaining = int(freshness.group(1)), int(freshness.group(2))
assert age == pytest.approx(120, abs=2)
assert remaining == pytest.approx(600 + 1800 - 120, abs=2)
@pytest.mark.asyncio
async def test_unparseable_jwks_response_does_not_fall_back_to_the_stale_copy():
"""Only an unreachable IdP unlocks the stale copy. A reachable one that answers badly must surface the error."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://garbled-issuer.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="garbled-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="garbled-key") == jwk
await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}")
garbled = MagicMock()
garbled.status_code = 200
garbled.text = "<html>not json</html>"
garbled.json.side_effect = ValueError("Expecting value: line 1 column 1")
endpoint.outcomes = (garbled,)
with pytest.raises(Exception, match="Error parsing response"):
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="garbled-key")
@pytest.mark.asyncio
async def test_jwks_error_response_is_not_cached_over_the_last_known_good_keys():
"""An IdP error body must never be stored as the key set, least of all as the stale copy."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://erroring-issuer.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="erroring-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="erroring-key") == jwk
active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}"
await cache.async_delete_cache(key=active_cache_key)
server_error = MagicMock()
server_error.status_code = 503
server_error.text = '{"error": "upstream unavailable"}'
server_error.json.return_value = {"error": "upstream unavailable"}
endpoint.outcomes = (server_error,)
with pytest.raises(Exception, match="returned status 503"):
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="erroring-key")
assert await cache.async_get_cache(key=active_cache_key) is None
assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [jwk]
@pytest.mark.asyncio
async def test_sustained_jwks_outage_refetches_once_per_backoff_window_not_once_per_request():
"""Without a backoff, every request during an outage pays three timeouts serialised behind the refresh lock."""
from litellm.caching.dual_cache import DualCache
jwks_url = "https://flooded-issuer.example.com/keys"
_, jwk = _get_rsa_key_and_jwk(kid="flooded-key")
cache = DualCache()
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint)
await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="flooded-key")
await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}")
endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),)
calls_before_outage = endpoint.call_count
public_keys = await asyncio.gather(
*[jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="flooded-key") for _ in range(6)]
)
assert public_keys == [jwk] * 6
assert endpoint.call_count - calls_before_outage == JWKS_FETCH_ATTEMPTS
@pytest.mark.asyncio
async def test_get_public_key_raises_503_when_jwks_unreachable_and_no_cached_keys(monkeypatch):
"""An unreachable IdP is an infra failure: 503, never a 401 that clients read as bad credentials."""
from litellm.caching.dual_cache import DualCache
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://issuer.example.com/keys")
endpoint = _ScriptedJWKSEndpoint((httpx.ConnectTimeout("connect timed out"),))
jwt_handler = _get_jwt_handler_with_scripted_endpoint(DualCache(), endpoint)
with pytest.raises(ProxyException) as exc_info:
await jwt_handler.get_public_key(kid="any-key")
assert exc_info.value.code == "503"
assert exc_info.value.type == ProxyErrorTypes.auth_provider_unavailable
assert "ConnectTimeout" in exc_info.value.message
assert endpoint.call_count == 3
@pytest.mark.asyncio
async def test_get_public_key_coalesces_concurrent_jwks_refreshes():
"""Concurrent requests in the TTL-expiry window share one JWKS fetch."""
from litellm.caching.dual_cache import DualCache
_, jwk = _get_rsa_key_and_jwk(kid="coalesced-key")
endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},), delay=0.05)
jwt_handler = _get_jwt_handler_with_scripted_endpoint(DualCache(), endpoint)
public_keys = await asyncio.gather(
*[
jwt_handler._get_public_key_from_jwks_url(
jwks_url="https://coalesce.example.com/keys",
kid="coalesced-key",
)
for _ in range(5)
]
)
assert public_keys == [jwk] * 5
assert endpoint.call_count == 1
@pytest.mark.asyncio
async def test_get_public_key_tries_next_jwks_url_when_kid_missing(monkeypatch):
from litellm.caching.dual_cache import DualCache
first_jwks_url = "https://first.example.com/keys"
second_jwks_url = "https://second.example.com/keys"
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", f"{first_jwks_url}, {second_jwks_url},,")
_, first_jwk = _get_rsa_key_and_jwk(kid="first-key")
_, second_jwk = _get_rsa_key_and_jwk(kid="second-key")
cache = DualCache()
cache.set_cache(key=f"litellm_jwt_auth_keys_{first_jwks_url}", value=[first_jwk])
cache.set_cache(key=f"litellm_jwt_auth_keys_{second_jwks_url}", value=[second_jwk])
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(),
)
public_key = await jwt_handler.get_public_key(kid="second-key")
assert public_key == second_jwk
def test_get_jwks_url_for_issuer_falls_back_to_discovery_document():
jwt_handler = JWTHandler()
issuer_config = LiteLLM_JWTAuth(
issuers=[
{
"issuer": "https://issuer.example.com/tenant/",
"disable_audience_validation": True,
}
]
).issuers[0]
jwks_url = jwt_handler._get_jwks_url_for_issuer(issuer_config=issuer_config)
assert (
jwks_url == "https://issuer.example.com/tenant/.well-known/openid-configuration"
)
@pytest.mark.asyncio
async def test_get_objects_team_membership_uses_rebound_user_id():
"""team_membership lookup uses resolved DB user_id, not JWT email claim."""
from litellm.caching.caching import DualCache
legacy_uuid = "bb8ab11f-09aa-47ae-b063-6e80506ac3bc"
jwt_email = "matt@example.com"
team_id = "team-1"
resolved_user = LiteLLM_UserTable(user_id=legacy_uuid, user_email=jwt_email)
captured = {}
async def fake_get_user_object(*args, **kwargs):
return resolved_user
async def fake_get_team_membership(user_id, team_id, *args, **kwargs):
captured["user_id"] = user_id
captured["team_id"] = team_id
return None
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
user_id_jwt_field="email", user_id_upsert=True
)
with (
patch(
"litellm.proxy.auth.handle_jwt.get_user_object",
side_effect=fake_get_user_object,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
side_effect=fake_get_team_membership,
),
):
(
user_object,
_org_object,
_end_user_object,
_team_membership_object,
effective_user_id,
) = await JWTAuthManager.get_objects(
user_id=jwt_email,
user_email=jwt_email,
org_id=None,
end_user_id=None,
team_id=team_id,
valid_user_email=None,
jwt_handler=jwt_handler,
prisma_client=MagicMock(),
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
route="/chat/completions",
)
assert user_object is not None and user_object.user_id == legacy_uuid
assert effective_user_id == legacy_uuid
assert captured["user_id"] == legacy_uuid, (
"team_membership lookup must use the resolved DB user_id, not the JWT "
f"email claim (got {captured['user_id']!r})"
)
assert captured["team_id"] == team_id
@pytest.mark.asyncio
async def test_multi_issuer_jwt_validates_selected_issuer_and_maps_claims(
monkeypatch,
):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer_one = "https://issuer-one.example.com"
issuer_two = "https://issuer-two.example.com"
issuer_one_jwks_url = f"{issuer_one}/keys"
issuer_two_jwks_url = f"{issuer_two}/keys"
shared_kid = "shared-kid"
_, issuer_one_jwk = _get_rsa_key_and_jwk(kid=shared_kid)
issuer_two_private_key, issuer_two_jwk = _get_rsa_key_and_jwk(kid=shared_kid)
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer_one,
"jwks_url": issuer_one_jwks_url,
"audience": "audience-one",
"user_id_jwt_field": "email",
"user_email_jwt_field": "email",
},
{
"issuer": issuer_two,
"jwks_url": issuer_two_jwks_url,
"audience": "audience-two",
"user_id_jwt_field": "repository_owner",
"team_id_jwt_field": "repository",
},
],
keys_by_url={
issuer_one_jwks_url: [issuer_one_jwk],
issuer_two_jwks_url: [issuer_two_jwk],
},
)
token = _encode_rsa_jwt(
private_key=issuer_two_private_key,
issuer=issuer_two,
audience="audience-two",
kid=shared_kid,
extra_claims={
"repository_owner": "example-org",
"repository": "example-org/litellm-fork",
},
)
claims = await jwt_handler.auth_jwt(token=token)
assert claims[JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == issuer_two
assert jwt_handler.get_user_id(token=claims, default_value=None) == "example-org"
assert jwt_handler.get_team_id(token=claims, default_value=None) == (
"example-org/litellm-fork"
)
@pytest.mark.asyncio
async def test_auth_jwt_issuer_path_expired_token_raises_401(monkeypatch):
"""An expired JWT validated through the issuer-scoped path
(_auth_jwt_with_issuer) must raise a ProxyException carrying a 401 so the
status is preserved end-to-end, just like the non-issuer path.
"""
import time
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
kid = "expired-kid"
private_key, jwk = _get_rsa_key_and_jwk(kid=kid)
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[{"issuer": issuer, "jwks_url": jwks_url, "audience": "my-audience"}],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="my-audience",
kid=kid,
extra_claims={"exp": int(time.time()) - 100},
)
with pytest.raises(ProxyException) as exc_info:
await jwt_handler.auth_jwt(token=token)
assert exc_info.value.code == str(401)
assert exc_info.value.type == ProxyErrorTypes.expired_key.value
assert "Token Expired" in exc_info.value.message
@pytest.mark.asyncio
async def test_auth_jwt_issuer_path_unreachable_jwks_raises_503(monkeypatch):
"""The issuer-scoped path must report an unreachable IdP as 503, not as a credential failure."""
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://unreachable-issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, _ = _get_rsa_key_and_jwk(kid="unreachable-kid")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[{"issuer": issuer, "jwks_url": jwks_url, "audience": "my-audience"}],
keys_by_url={},
)
endpoint = _ScriptedJWKSEndpoint((httpx.ConnectTimeout("connect timed out"),))
jwt_handler.http_handler = endpoint
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="my-audience",
kid="unreachable-kid",
)
with pytest.raises(ProxyException) as exc_info:
await jwt_handler.auth_jwt(token=token)
assert exc_info.value.code == "503"
assert exc_info.value.type == ProxyErrorTypes.auth_provider_unavailable
assert endpoint.call_count == 3
@pytest.mark.asyncio
async def test_multi_issuer_jwt_maps_kubernetes_namespace_claim(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://oidc.eks.eu-west-1.amazonaws.com/id/test-cluster"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="k8s-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": None,
"disable_audience_validation": True,
"user_id_jwt_field": "kubernetes\\.io.namespace",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="kubernetes.default.svc",
kid="k8s-key",
extra_claims={"kubernetes.io": {"namespace": "example-namespace"}},
)
claims = await jwt_handler.auth_jwt(token=token)
assert (
jwt_handler.get_user_id(token=claims, default_value=None) == "example-namespace"
)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_unknown_issuer_falls_back_to_global_jwks(monkeypatch):
"""Tokens whose ``iss`` is not in the configured issuers list fall through
to the legacy ``JWT_PUBLIC_KEY_URL`` path so operators can add the new
``issuers`` list to a live deployment without breaking existing tokens
minted by non-configured IdPs. With no global JWKS configured, the legacy
path surfaces a ``Missing JWT Public Key URL from environment.`` error.
"""
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
configured_issuer = "https://issuer.example.com"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": configured_issuer,
"jwks_url": f"{configured_issuer}/keys",
"audience": "expected-audience",
}
],
keys_by_url={f"{configured_issuer}/keys": [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer="https://unknown-issuer.example.com",
audience="expected-audience",
kid="issuer-key",
)
with pytest.raises(Exception, match='Missing JWT Public Key URL from environment\\.') as exc:
await jwt_handler.auth_jwt(token=token)
assert "Missing JWT Public Key URL from environment." in str(exc.value)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_rejects_wrong_audience(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "expected-audience",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="wrong-audience",
kid="issuer-key",
)
with pytest.raises(Exception, match="Validation fails: Audience doesn't match") as exc:
await jwt_handler.auth_jwt(token=token)
assert "Validation fails" in str(exc.value)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer_one = "https://issuer-one.example.com"
issuer_two = "https://issuer-two.example.com"
issuer_one_jwks_url = f"{issuer_one}/keys"
issuer_two_jwks_url = f"{issuer_two}/keys"
shared_kid = "shared-kid"
issuer_one_private_key, issuer_one_jwk = _get_rsa_key_and_jwk(kid=shared_kid)
_, issuer_two_jwk = _get_rsa_key_and_jwk(kid=shared_kid)
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer_one,
"jwks_url": issuer_one_jwks_url,
"audience": "audience-one",
},
{
"issuer": issuer_two,
"jwks_url": issuer_two_jwks_url,
"audience": "audience-two",
},
],
keys_by_url={
issuer_one_jwks_url: [issuer_one_jwk],
issuer_two_jwks_url: [issuer_two_jwk],
},
)
token = _encode_rsa_jwt(
private_key=issuer_one_private_key,
issuer=issuer_two,
audience="audience-two",
kid=shared_kid,
)
with pytest.raises(Exception, match='Validation fails: Signature verification failed') as exc:
await jwt_handler.auth_jwt(token=token)
assert "Validation fails" in str(exc.value)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_missing_mapped_claim_leaves_user_id_unset(
monkeypatch,
):
"""Mapped issuer claims behave like the global ``litellm_jwtauth`` path —
present claims override the normalised value, missing ones simply leave
the corresponding LiteLLM-internal claim absent (rather than failing the
JWT outright). This keeps multi-issuer auth tolerant of tokens that omit
optional fields like email or org id.
"""
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "expected-audience",
"user_id_jwt_field": "email",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="expected-audience",
kid="issuer-key",
)
claims = await jwt_handler.auth_jwt(token=token)
assert claims[jwt_handler.LITELLM_JWT_ISSUER_CLAIM] == issuer
assert jwt_handler.LITELLM_USER_ID_CLAIM not in claims
def test_multi_issuer_jwt_requires_audience_unless_explicitly_disabled(
monkeypatch,
):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
with pytest.raises(Exception, match='must configure audience or set') as exc:
LiteLLM_JWTAuth(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
}
]
)
assert "must configure audience" in str(exc.value)
def test_multi_issuer_jwt_rejects_audience_with_disable_audience_validation():
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
with pytest.raises(Exception, match='cannot set audience and disable_audience_validation=True') as exc:
LiteLLM_JWTAuth(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "some-audience",
"disable_audience_validation": True,
}
]
)
assert "cannot set audience and disable_audience_validation=True together" in str(
exc.value
)
@pytest.mark.asyncio
async def test_global_jwt_ignores_user_supplied_internal_claims(monkeypatch):
from litellm.caching.dual_cache import DualCache
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_ISSUER", raising=False)
jwks_url = "https://global-issuer.example.com/keys"
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url)
private_key, jwk = _get_rsa_key_and_jwk(kid="global-key")
cache = DualCache()
cache.set_cache(key=f"litellm_jwt_auth_keys_{jwks_url}", value=[jwk])
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(
user_id_jwt_field="email",
user_email_jwt_field="email",
team_id_jwt_field="team.id",
team_ids_jwt_field="teams",
org_id_jwt_field="org.id",
end_user_id_jwt_field="end_user.id",
),
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer="https://global-issuer.example.com",
audience="some-other-client",
kid="global-key",
extra_claims={
"email": "real-user@example.com",
"team": {"id": "real-team"},
"teams": ["real-team", "secondary-team"],
"org": {"id": "real-org"},
"end_user": {"id": "real-end-user"},
JWTHandler.LITELLM_JWT_ISSUER_CLAIM: "https://issuer.example.com",
JWTHandler.LITELLM_USER_ID_CLAIM: "victim-user",
JWTHandler.LITELLM_USER_EMAIL_CLAIM: "victim@example.com",
JWTHandler.LITELLM_TEAM_ID_CLAIM: "victim-team",
JWTHandler.LITELLM_TEAM_IDS_CLAIM: ["victim-team"],
JWTHandler.LITELLM_ORG_ID_CLAIM: "victim-org",
JWTHandler.LITELLM_END_USER_ID_CLAIM: "victim-end-user",
},
)
claims = await jwt_handler.auth_jwt(token=token)
assert jwt_handler.get_user_id(token=claims, default_value=None) == (
"real-user@example.com"
)
assert jwt_handler.get_user_email(token=claims, default_value=None) == (
"real-user@example.com"
)
assert jwt_handler.get_team_id(token=claims, default_value=None) == "real-team"
assert jwt_handler.get_team_ids_from_jwt(token=claims) == [
"real-team",
"secondary-team",
]
assert jwt_handler.get_org_id(token=claims, default_value=None) == "real-org"
assert jwt_handler.get_end_user_id(token=claims, default_value=None) == (
"real-end-user"
)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_strips_unmapped_internal_claims(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "expected-audience",
"user_email_jwt_field": "email",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="expected-audience",
kid="issuer-key",
extra_claims={
"email": "real-user@example.com",
JWTHandler.LITELLM_USER_ID_CLAIM: "victim-user",
JWTHandler.LITELLM_TEAM_ID_CLAIM: "victim-team",
},
)
claims = await jwt_handler.auth_jwt(token=token)
assert JWTHandler.LITELLM_USER_ID_CLAIM not in claims
assert JWTHandler.LITELLM_TEAM_ID_CLAIM not in claims
assert jwt_handler.get_user_id(token=claims, default_value=None) is None
assert jwt_handler.get_team_id(token=claims, default_value=None) is None
assert jwt_handler.get_user_email(token=claims, default_value=None) == (
"real-user@example.com"
)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_does_not_emit_unscoped_global_warning(
monkeypatch, caplog
):
import logging
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_ISSUER", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
JWTHandler._unscoped_jwt_warning_emitted = False
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "expected-audience",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="expected-audience",
kid="issuer-key",
)
with caplog.at_level(logging.WARNING):
await jwt_handler.auth_jwt(token=token)
assert "Tokens minted by any application" not in caplog.text
assert JWTHandler._unscoped_jwt_warning_emitted is False
def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deployment(
monkeypatch, _reset_unscoped_warning_flag, caplog
):
"""The unscoped-fallback warning must fire even when per-issuer configs
are set. In mixed deployments, tokens whose ``iss`` does not match any
configured issuer fall through to the global path; if env-var scoping is
absent that fallback IS unscoped, and the operator needs to be told."""
import logging
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_ISSUER", raising=False)
caplog.set_level(logging.WARNING)
JWTHandler._build_decode_kwargs()
matching = [
r
for r in caplog.records
if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage()
]
assert len(matching) == 1
# ---------------------------------------------------------------------------
# fallback_to_db_teams: resolve team from DB memberships when JWT has no team
# claims (config flag on LiteLLM_JWTAuth)
# ---------------------------------------------------------------------------
def test_get_team_id_from_header_defers_to_db_membership_only_without_jwt_claims():
"""With fallback_to_db_teams=True, an x-litellm-team-id header is accepted
provisionally only when the JWT carries no team claims (allowed set empty).
When the JWT does carry team claims, the header must still be validated
against them, and the flag-off behavior must keep rejecting unknown teams."""
deferred = JWTAuthManager.get_team_id_from_header(
request_headers={"x-litellm-team-id": "team-from-db"},
allowed_team_ids=set(),
fallback_to_db_teams=True,
)
assert deferred == "team-from-db"
with pytest.raises(HTTPException) as exc_info:
JWTAuthManager.get_team_id_from_header(
request_headers={"x-litellm-team-id": "team-x"},
allowed_team_ids={"team-1", "team-2"},
fallback_to_db_teams=True,
)
assert exc_info.value.status_code == 403
with pytest.raises(HTTPException):
JWTAuthManager.get_team_id_from_header(
request_headers={"x-litellm-team-id": "team-from-db"},
allowed_team_ids=set(),
fallback_to_db_teams=False,
)
@pytest.mark.asyncio
async def test_find_team_with_model_access_defers_no_team_403_under_db_fallback():
"""find_team_with_model_access raises the early "no teams in token" 403 when
enforcement is on, but defers (returns no team) so auth_builder's DB fallback
can run when fallback_to_db_teams is enabled."""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
enforce_team_based_model_access=True,
fallback_to_db_teams=False,
)
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.find_team_with_model_access(
team_ids=set(),
requested_model="gpt-4",
route="/chat/completions",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert exc_info.value.status_code == 403
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
enforce_team_based_model_access=True,
fallback_to_db_teams=True,
)
team_id, team_object = await JWTAuthManager.find_team_with_model_access(
team_ids=set(),
requested_model="gpt-4",
route="/chat/completions",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert team_id is None
assert team_object is None
def _db_fallback_handler(litellm_jwtauth: Optional[LiteLLM_JWTAuth] = None) -> JWTHandler:
handler = JWTHandler()
handler.litellm_jwtauth = litellm_jwtauth or LiteLLM_JWTAuth()
return handler
@pytest.mark.asyncio
async def test_resolve_db_team_fallback_skips_unresolvable_membership():
"""An orphaned membership (team row missing/erroring) is skipped and the next
resolvable DB team is selected instead of aborting the fallback."""
user_object = LiteLLM_UserTable(
user_id="u_skip",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["ghost_team", "real_team"],
)
resolved = LiteLLM_TeamTable(team_id="real_team")
async def fake_get_team(team_id, **kwargs):
if team_id == "ghost_team":
raise HTTPException(status_code=404, detail="missing")
return resolved
with patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
):
(
team_id,
team_object,
_membership,
) = await JWTAuthManager._resolve_db_team_fallback(
user_object=user_object,
user_id=None,
requested_model=None,
route="/chat/completions",
jwt_handler=_db_fallback_handler(),
enforce_team_based_model_access=True,
team_id_upsert=False,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert team_id == "real_team"
assert team_object is resolved
@pytest.mark.parametrize(
(
"fallback_to_db_teams",
"user_teams",
"header_team_id",
"expected_team_id",
"expect_403",
),
[
pytest.param(
True, ["team_solo"], None, "team_solo", False, id="flag_on_single_db_team"
),
pytest.param(
True,
["team_a", "team_b"],
None,
"team_a",
False,
id="flag_on_multi_db_team_picks_first",
),
pytest.param(
True,
["team_a", "team_b"],
"team_b",
"team_b",
False,
id="flag_on_header_team_in_membership",
),
pytest.param(
True,
["team_a", "team_b"],
"team_x",
None,
True,
id="flag_on_header_team_not_in_membership_403",
),
pytest.param(True, [], None, None, True, id="flag_on_no_db_team_enforced_403"),
pytest.param(
False,
["team_a", "team_b"],
None,
None,
False,
id="flag_off_multi_db_team_no_fallback",
),
pytest.param(
False,
["team_solo"],
None,
"team_solo",
False,
id="flag_off_single_db_team_upstream_fallback",
),
],
)
@pytest.mark.asyncio
async def test_auth_builder_db_team_fallback_when_jwt_has_no_team(
fallback_to_db_teams: bool,
user_teams: list,
header_team_id: Optional[str],
expected_team_id: Optional[str],
expect_403: bool,
) -> None:
"""End-to-end auth_builder behavior with no JWT team claims.
fallback_to_db_teams=True attributes usage to the user's first resolvable DB
team, honors a valid x-litellm-team-id header, and rejects a header team the
user does not belong to. The default (flag off) preserves the upstream
single-team fallback: a lone DB team is resolved, multiple are ambiguous.
"""
user_id = "u_db_fallback"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=user_teams,
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
enforce_team_based_model_access=True,
fallback_to_db_teams=fallback_to_db_teams,
)
request_headers = {"x-litellm-team-id": header_team_id} if header_team_id else None
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
async def call_auth_builder():
with (
patch.object(
jwt_handler, "auth_jwt", new_callable=AsyncMock
) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(
JWTAuthManager,
"find_and_validate_specific_team_id",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"find_team_with_model_access",
new_callable=AsyncMock,
return_value=(None, None),
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
return_value=LiteLLM_TeamMembership(
user_id=user_id,
team_id=user_teams[0] if user_teams else "none",
litellm_budget_table=None,
),
),
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
return await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers=request_headers,
)
if expect_403:
with pytest.raises(HTTPException) as exc_info:
await call_auth_builder()
assert exc_info.value.status_code == 403
else:
result = await call_auth_builder()
assert result["team_id"] == expected_team_id
@pytest.mark.parametrize(
"fallback_to_db_teams, expect_teams_stripped",
[
pytest.param(True, False, id="fallback_on_preserves_db_teams"),
pytest.param(False, True, id="fallback_off_strips_db_teams"),
],
)
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_no_claim_team_preservation(
fallback_to_db_teams: bool,
expect_teams_stripped: bool,
) -> None:
"""A no-team-claim JWT must not permanently strip a user's DB team memberships
when fallback_to_db_teams is enabled — otherwise the DB fallback that runs
right after has nothing to resolve and every request silently wipes the user
out of their teams. With the flag off, the legacy mirror-the-IdP behavior
(remove teams absent from the token) is preserved."""
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=AsyncMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="teams",
sync_user_role_and_teams=True,
fallback_to_db_teams=fallback_to_db_teams,
),
)
token = {"sub": "u1"}
user = LiteLLM_UserTable(
user_id="u1",
user_role=LitellmUserRoles.INTERNAL_USER.value,
teams=["team_a", "team_b"],
)
prisma = AsyncMock()
with patch(
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
new_callable=AsyncMock,
) as mock_patch:
await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, prisma)
if expect_teams_stripped:
mock_patch.assert_awaited_once()
assert set(mock_patch.call_args.kwargs["teams_ids_to_remove_user_from"]) == {
"team_a",
"team_b",
}
assert user.teams == []
else:
mock_patch.assert_not_called()
assert user.teams == ["team_a", "team_b"]
@pytest.mark.asyncio
async def test_resolve_db_team_fallback_skips_team_without_model_access():
"""The DB-team fallback must apply the same per-team model-access check as the
claim-based path: a DB team that cannot access the requested model is skipped
in favor of one that can, instead of selecting the first membership blindly."""
user_object = LiteLLM_UserTable(
user_id="u_model_access",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["restricted_team", "allowed_team"],
)
teams = {
"restricted_team": LiteLLM_TeamTable(
team_id="restricted_team", models=["claude-3"]
),
"allowed_team": LiteLLM_TeamTable(team_id="allowed_team", models=["gpt-4"]),
}
async def fake_get_team(team_id, **kwargs):
return teams[team_id]
async def fake_can_access(model, team_object, llm_router, team_model_aliases=None):
if model in (team_object.models or []):
return True
raise ProxyException(
message="team not allowed to access model",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=403,
)
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.can_team_access_model",
new_callable=AsyncMock,
side_effect=fake_can_access,
),
):
(
team_id,
team_object,
_membership,
) = await JWTAuthManager._resolve_db_team_fallback(
user_object=user_object,
user_id=None,
requested_model="gpt-4",
route="/chat/completions",
jwt_handler=_db_fallback_handler(),
enforce_team_based_model_access=True,
team_id_upsert=False,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert team_id == "allowed_team"
assert team_object is teams["allowed_team"]
@pytest.mark.asyncio
async def test_resolve_db_team_fallback_enforces_team_allowed_routes():
"""The DB-team fallback must apply the same team_allowed_routes gate as the
claim-based path: a route the JWT config excludes for team-role callers must
not become reachable by selecting a DB team, even when that team can access
the requested model. Without the gate, a teamless JWT could reach the
info/management routes an admin narrowed team_allowed_routes to exclude."""
user_object = LiteLLM_UserTable(
user_id="u_routes",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["team_a"],
)
team = LiteLLM_TeamTable(team_id="team_a", models=["gpt-4"])
handler = _db_fallback_handler(LiteLLM_JWTAuth(team_allowed_routes=["openai_routes"]))
async def fake_get_team(team_id, **kwargs):
return team
async def fake_can_access(model, team_object, llm_router, team_model_aliases=None):
return True
async def resolve(route):
return await JWTAuthManager._resolve_db_team_fallback(
user_object=user_object,
user_id=None,
requested_model="gpt-4",
route=route,
jwt_handler=handler,
enforce_team_based_model_access=False,
team_id_upsert=False,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.can_team_access_model",
new_callable=AsyncMock,
side_effect=fake_can_access,
),
):
excluded_team_id, excluded_team_object, _ = await resolve("/key/info")
allowed_team_id, allowed_team_object, _ = await resolve("/chat/completions")
assert excluded_team_id is None
assert excluded_team_object is None
assert allowed_team_id == "team_a"
assert allowed_team_object is team
def test_validate_header_team_in_db_membership_does_not_leak_team_ids():
"""The 403 raised for an x-litellm-team-id header outside the user's DB
memberships must not enumerate the user's team IDs back to the caller; any
valid-JWT caller could otherwise probe header values to discover team IDs."""
user_object = LiteLLM_UserTable(
user_id="u_leak",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["secret_team_alpha", "secret_team_beta"],
)
with pytest.raises(HTTPException) as exc_info:
JWTAuthManager._validate_header_team_in_db_membership(
team_id="outsider_team",
user_object=user_object,
)
detail = exc_info.value.detail
assert exc_info.value.status_code == 403
assert "secret_team_alpha" not in detail
assert "secret_team_beta" not in detail
assert "outsider_team" in detail
async def _run_auth_builder_with_header_team(
jwt_auth_config: LiteLLM_JWTAuth,
token: dict,
header_team_id: str,
user_object: LiteLLM_UserTable,
fake_get_team,
allowed_team_ids: set,
):
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = jwt_auth_config
with (
patch.object(
jwt_handler, "auth_jwt", new_callable=AsyncMock, return_value=token
),
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_object.user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(
JWTAuthManager, "get_all_team_ids", return_value=allowed_team_ids
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_object.user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
):
return await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers={"x-litellm-team-id": header_team_id},
)
async def _team_lookup_404(team_id, **kwargs):
raise HTTPException(
status_code=404,
detail=f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call.",
)
@pytest.mark.asyncio
async def test_auth_builder_header_team_not_found_matches_non_membership_denial() -> (
None
):
"""A provisional x-litellm-team-id naming a nonexistent team must produce
the exact same 403 shape as one naming an existing team outside the
caller's memberships. Letting get_team_object's 404 surface would give any
valid-JWT caller an oracle to probe which team ids exist."""
user_object = LiteLLM_UserTable(
user_id="u_oracle",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["team_member"],
)
config = LiteLLM_JWTAuth(
enforce_team_based_model_access=True,
fallback_to_db_teams=True,
)
token = {"sub": "u_oracle", "scope": ""}
async def team_exists(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
with pytest.raises(HTTPException) as missing_exc:
await _run_auth_builder_with_header_team(
config, token, "team_ghost", user_object, _team_lookup_404, set()
)
with pytest.raises(HTTPException) as outsider_exc:
await _run_auth_builder_with_header_team(
config, token, "team_other", user_object, team_exists, set()
)
assert missing_exc.value.status_code == 403
assert outsider_exc.value.status_code == 403
assert missing_exc.value.detail.replace(
"team_ghost", "<team>"
) == outsider_exc.value.detail.replace("team_other", "<team>")
assert "exist" not in missing_exc.value.detail
@pytest.mark.asyncio
async def test_auth_builder_claim_backed_header_team_lookup_error_propagates() -> None:
"""When the JWT carries team claims the header team is not provisional, so
a failed team lookup keeps the upstream contract: get_team_object's 404
surfaces unchanged instead of being rewritten into the membership 403."""
user_object = LiteLLM_UserTable(
user_id="u_claimed",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["team_member"],
)
config = LiteLLM_JWTAuth(
enforce_team_based_model_access=True,
fallback_to_db_teams=True,
team_ids_jwt_field="team_ids",
)
token = {"sub": "u_claimed", "scope": "", "team_ids": ["team_claimed"]}
with pytest.raises(HTTPException) as exc_info:
await _run_auth_builder_with_header_team(
config, token, "team_claimed", user_object, _team_lookup_404, {"team_claimed"}
)
assert exc_info.value.status_code == 404
@pytest.mark.asyncio
async def test_resolve_db_team_fallback_loads_team_membership():
"""The DB-team fallback must load the resolved team's membership row (when a
user_id is known) so per-team membership budget limits are enforced on the
fallback path the same as on the claim-based path; returning a None membership
would silently skip LiteLLM_TeamMembership budget checks for every request."""
user_object = LiteLLM_UserTable(
user_id="u_membership",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["team_with_budget"],
)
membership = LiteLLM_TeamMembership(
user_id="u_membership",
team_id="team_with_budget",
budget_id="budget_xyz",
litellm_budget_table=None,
)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
async def fake_get_membership(user_id, team_id, **kwargs):
assert user_id == "u_membership"
assert team_id == "team_with_budget"
return membership
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
side_effect=fake_get_membership,
),
):
(
team_id,
team_object,
team_membership,
) = await JWTAuthManager._resolve_db_team_fallback(
user_object=user_object,
user_id="u_membership",
requested_model=None,
route="/chat/completions",
jwt_handler=_db_fallback_handler(),
enforce_team_based_model_access=True,
team_id_upsert=False,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert team_id == "team_with_budget"
assert team_object is not None
assert team_membership is membership
assert team_membership.budget_id == "budget_xyz"
@pytest.mark.asyncio
async def test_resolve_db_team_fallback_survives_membership_lookup_error():
"""A transient membership-lookup failure must not deny an otherwise-authorized
request. get_team_membership swallows DB errors internally and returns None, so
the fallback must return the resolved team with a None membership (budget
enforcement degrades gracefully) instead of treating it as a denial."""
user_object = LiteLLM_UserTable(
user_id="u_flaky",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["team_flaky"],
)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
async def none_on_db_error_membership(user_id, team_id, **kwargs):
return None
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
side_effect=none_on_db_error_membership,
),
):
(
team_id,
team_object,
team_membership,
) = await JWTAuthManager._resolve_db_team_fallback(
user_object=user_object,
user_id="u_flaky",
requested_model=None,
route="/chat/completions",
jwt_handler=_db_fallback_handler(),
enforce_team_based_model_access=True,
team_id_upsert=False,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert team_id == "team_flaky"
assert team_object is not None
assert team_membership is None
@pytest.mark.asyncio
async def test_auth_builder_db_fallback_does_not_validate_rbac_team_against_db_membership():
"""When fallback_to_db_teams is on and the JWT carries an RBAC team role but no
group/team claims, team_id is set from the RBAC object_id (not the provisional
x-litellm-team-id header). That RBAC-asserted team must not be re-validated
against the user's DB memberships; only a team that actually came from the
header is provisional. Without the team_id == header_team_id guard, every such
RBAC request 403s when the RBAC team is not also a DB membership."""
rbac_team = "rbac_asserted_team"
user_id = "u_rbac"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["unrelated_db_team"],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
enforce_team_based_model_access=True,
fallback_to_db_teams=True,
)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=LitellmUserRoles.TEAM),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=rbac_team),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
result = await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers=None,
)
assert result["team_id"] == rbac_team
@pytest.mark.asyncio
async def test_resolve_db_team_fallback_distinguishes_no_membership_vs_model_denied():
"""When enforce_team_based_model_access is on, a user with no DB memberships
and a user with memberships that all fail the model-access check must surface
different 403s; collapsing both into the no-membership message hides the real
cause and diverges from find_team_with_model_access's claim-based message."""
membership_user = LiteLLM_UserTable(
user_id="u_no_model",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["only_team"],
)
no_membership_user = LiteLLM_UserTable(
user_id="u_empty",
user_role=LitellmUserRoles.INTERNAL_USER,
teams=[],
)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id, models=["other"])
async def fake_can_access(model, team_object, llm_router, team_model_aliases=None):
raise ProxyException(
message="team not allowed to access model",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=403,
)
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.can_team_access_model",
new_callable=AsyncMock,
side_effect=fake_can_access,
),
):
with pytest.raises(HTTPException) as model_denied:
await JWTAuthManager._resolve_db_team_fallback(
user_object=membership_user,
user_id=None,
requested_model="gpt-4",
route="/chat/completions",
jwt_handler=_db_fallback_handler(),
enforce_team_based_model_access=True,
team_id_upsert=False,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
with pytest.raises(HTTPException) as no_member:
await JWTAuthManager._resolve_db_team_fallback(
user_object=no_membership_user,
user_id=None,
requested_model="gpt-4",
route="/chat/completions",
jwt_handler=_db_fallback_handler(),
enforce_team_based_model_access=True,
team_id_upsert=False,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert model_denied.value.status_code == 403
assert "requested model" in model_denied.value.detail
assert "gpt-4" in model_denied.value.detail
assert "only_team" not in model_denied.value.detail
assert no_member.value.status_code == 403
assert "not a member of any team" in no_member.value.detail
@pytest.mark.asyncio
async def test_auth_builder_db_fallback_runs_when_only_team_id_default_set():
"""team_id_default makes JWTHandler.get_team_id return a non-None team for a
claimless token. The fallback gate must look at real JWT team claims (not the
operator-configured default) so fallback_to_db_teams still attributes to the
user's DB memberships instead of silently routing to the default team."""
user_id = "u_default_token"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["db_team_for_user"],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
fallback_to_db_teams=True,
team_id_default="config_default_team",
)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
return_value=None,
),
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
result = await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers=None,
)
assert result["team_id"] == "db_team_for_user"
@pytest.mark.asyncio
async def test_auth_builder_alias_only_token_resolves_alias_not_db_fallback():
"""An alias-only JWT (team_alias_jwt_field set, no team-id claims) must resolve
its alias via find_and_validate_specific_team_id, not fall into the DB-membership
fallback. get_all_jwt_team_ids ignores aliases, so without the get_team_alias
clause in the db_team_fallback gate the alias is silently dropped and the request
is mis-attributed to the user's first DB team instead of the alias-named team."""
user_id = "u_alias_only"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["db_membership_team"],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
fallback_to_db_teams=True,
team_alias_jwt_field="team_name",
)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
async def fake_get_team_by_alias(team_alias, **kwargs):
return LiteLLM_TeamTable(team_id="alias_resolved_team", team_alias=team_alias)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
new_callable=AsyncMock,
side_effect=fake_get_team_by_alias,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
return_value=None,
),
):
mock_auth_jwt.return_value = {"sub": user_id, "team_name": "resolvable_alias"}
result = await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers=None,
)
assert result["team_id"] == "alias_resolved_team"
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_alias_wins_over_team_id_default():
"""When the JWT carries only an alias claim (no team_id claim) and
team_id_default is configured, alias resolution must win. get_team_id
silently substitutes team_id_default for a missing claim, which would
otherwise mask the alias-resolved team and mis-attribute spend/access
to the configured default team."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_alias_jwt_field="team_alias",
team_id_default="config_default_team",
),
)
jwt_token = {"sub": "user-1", "team_alias": "my-team"}
alias_team = LiteLLM_TeamTable(
team_id="alias_resolved_team", team_alias="my-team"
)
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_by_id,
patch(
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
new_callable=AsyncMock,
) as mock_get_by_alias,
):
mock_get_by_alias.return_value = alias_team
team_id, team_obj = await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_token,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
assert team_id == "alias_resolved_team"
assert team_obj == alias_team
mock_get_by_id.assert_not_called()
mock_get_by_alias.assert_called_once()
@pytest.mark.asyncio
async def test_find_and_validate_specific_team_id_team_id_default_used_without_alias():
"""When the token carries neither a team_id nor an alias claim and
team_id_default is configured, the default still resolves the team. The
alias-precedence fix must not regress this baseline fallback behavior."""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_alias_jwt_field="team_alias",
team_id_default="config_default_team",
),
)
jwt_token = {"sub": "user-1"}
default_team = LiteLLM_TeamTable(team_id="config_default_team")
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_by_id,
patch(
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
new_callable=AsyncMock,
) as mock_get_by_alias,
):
mock_get_by_id.return_value = default_team
team_id, team_obj = await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_token,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
assert team_id == "config_default_team"
assert team_obj == default_team
mock_get_by_alias.assert_not_called()
@pytest.mark.asyncio
async def test_auth_builder_db_fallback_enforces_passthrough_route_access():
"""A team selected only via _resolve_db_team_fallback must still pass the
auth-enforced passthrough route check; previously the earlier gate ran while
team_id was None and the fallback-resolved team bypassed it entirely."""
user_id = "u_passthrough"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["team_no_passthrough"],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(fallback_to_db_teams=True)
passthrough_route = "/vertex_ai/v1/projects/p/locations/us/publishers/google/models/gemini:generateContent"
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id, metadata={})
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.auth.handle_jwt.RouteChecks.is_auth_enforced_pass_through_route",
return_value=True,
),
patch.object(
JWTAuthManager,
"_team_has_passthrough_route_access",
return_value=False,
),
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gemini"},
general_settings={"enforce_rbac": False},
route=passthrough_route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers=None,
request_method="POST",
)
assert exc_info.value.status_code == 403
assert "passthrough route" in exc_info.value.detail
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_singular_claim_reconciles_memberships():
"""When fallback_to_db_teams is on but the JWT carries a singular team claim
(Okta/Auth0 default for users with one primary team), sync must treat it as a
real claim and reconcile DB memberships against it. Otherwise stale DB teams
persist and a subsequent claimless JWT for the same user is silently attributed
to a team the IdP never asserted on the singular-claim login."""
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=AsyncMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
team_id_jwt_field="primary_team",
sync_user_role_and_teams=True,
fallback_to_db_teams=True,
),
)
token = {"sub": "u_singular", "primary_team": "team_primary"}
user = LiteLLM_UserTable(
user_id="u_singular",
user_role=LitellmUserRoles.INTERNAL_USER.value,
teams=["team_stale_a", "team_stale_b"],
)
with patch(
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
new_callable=AsyncMock,
) as mock_patch:
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler, token, user, AsyncMock()
)
mock_patch.assert_awaited_once()
assert set(mock_patch.call_args.kwargs["teams_ids_to_remove_user_from"]) == {
"team_stale_a",
"team_stale_b",
}
assert set(mock_patch.call_args.kwargs["teams_ids_to_add_user_to"]) == {
"team_primary"
}
assert user.teams == ["team_primary"]
@pytest.mark.asyncio
async def test_auth_builder_provisional_header_team_is_not_upserted():
"""A provisional x-litellm-team-id (accepted only because the JWT carries no
team claims) must not be upserted even when team_id_upsert is enabled: it is
validated against DB membership afterwards, so upserting first would let an
attacker-supplied header create an orphaned team row. A genuine membership
team already exists, so the resolved request still succeeds."""
user_id = "u_no_upsert"
header_team = "header_supplied_team"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=[header_team],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
fallback_to_db_teams=True,
team_id_upsert=True,
)
upsert_by_team: dict[str, Optional[bool]] = {}
async def spy_get_team(team_id, **kwargs):
upsert_by_team[team_id] = kwargs.get("team_id_upsert")
return LiteLLM_TeamTable(team_id=team_id)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=spy_get_team,
),
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
result = await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers={"x-litellm-team-id": header_team},
)
assert result["team_id"] == header_team
assert upsert_by_team[header_team] is False
@pytest.mark.asyncio
async def test_auth_builder_header_cannot_override_rbac_team_under_db_fallback():
"""An RBAC team-role JWT already pins team_id to the asserted team. With
fallback_to_db_teams on, a caller must not be able to substitute that team
by sending x-litellm-team-id for any other team they happen to belong to:
the provisional-header path is only for tokens with no team identity at all,
so an RBAC token plus a non-claim header team is rejected with 403."""
user_id = "u_rbac_override"
rbac_team = "rbac_pinned_team"
other_team = "other_db_team"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=[other_team],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(fallback_to_db_teams=True)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=LitellmUserRoles.TEAM),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=rbac_team),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route="/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers={"x-litellm-team-id": other_team},
)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_auth_builder_header_team_enforces_team_allowed_routes_under_db_fallback():
"""A claimless JWT with x-litellm-team-id under fallback_to_db_teams must
obey the same team_allowed_routes gate as the auto-pick fallback path.
Otherwise the header bypasses the route gate the JWT config narrows for
team-role callers, letting management/info routes be reached with a
team_id the auto-pick path would silently refuse to set."""
user_id = "u_header_routes"
header_team = "header_supplied_team"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=[header_team],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
fallback_to_db_teams=True,
team_allowed_routes=["openai_routes"],
)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id)
async def call(route: str):
with (
patch.object(
jwt_handler, "auth_jwt", new_callable=AsyncMock
) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(
JWTAuthManager,
"check_admin_access",
new_callable=AsyncMock,
return_value=None,
),
patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(
JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
):
mock_auth_jwt.return_value = {"sub": user_id, "scope": ""}
return await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={"enforce_rbac": False},
route=route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers={"x-litellm-team-id": header_team},
)
with pytest.raises(HTTPException) as exc_info:
await call("/key/info")
assert exc_info.value.status_code == 403
assert "not allowed to access route" in exc_info.value.detail
assert "/key/info" in exc_info.value.detail
result = await call("/chat/completions")
assert result["team_id"] == header_team
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag():
"""Reading the singular team claim during sync is scoped to fallback_to_db_teams.
With the flag off, sync keeps the upstream plural-only reconciliation, so a
singular-only token is treated as claimless and existing DB teams are removed
exactly as before this PR; the new dual-claim behavior must not silently change
membership reconciliation for deployments that never opted in."""
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=AsyncMock(),
litellm_jwtauth=LiteLLM_JWTAuth(
team_id_jwt_field="primary_team",
sync_user_role_and_teams=True,
fallback_to_db_teams=False,
),
)
token = {"sub": "u_flag_off", "primary_team": "team_primary"}
user = LiteLLM_UserTable(
user_id="u_flag_off",
user_role=LitellmUserRoles.INTERNAL_USER.value,
teams=["team_existing"],
)
with patch(
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
new_callable=AsyncMock,
) as mock_patch:
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler, token, user, AsyncMock()
)
mock_patch.assert_awaited_once()
assert set(mock_patch.call_args.kwargs["teams_ids_to_remove_user_from"]) == {
"team_existing"
}
assert mock_patch.call_args.kwargs["teams_ids_to_add_user_to"] == []
assert user.teams == []
Reference in a new issue View git blame Copy permalink