scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 213

feat(jwt): allow virtual_key_claim_field per issuer

Browse source 
Multi-IdP deployments can now set virtual_key_claim_field and
unregistered_jwt_client_behavior on a JWTIssuerConfig entry. Tokens from
that issuer use the issuer-specific claim path and no-match policy for the
virtual key mapping lookup; issuers that omit them keep the global values.
The auth flow now enters the mapping lookup when any issuer configures the
field, not only when the global field is set.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-13 00:05:50 +00:00
parent d4a72e7372
commit f41c8556b5
4 changed files with 324 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/proxy/_types.py
View file

@ -4837,6 +4837,14 @@ class JWTIssuerConfig(BaseModel):
default=None,
description="Issuer-specific claim path to normalize into LiteLLM's end-user id.",
)
virtual_key_claim_field: str | None = Field(
default=None,
description="Issuer-specific claim path used for the virtual key mapping lookup. Falls back to the global field.",
)
unregistered_jwt_client_behavior: UnregisteredJWTClientBehavior | None = Field(
default=None,
description="Issuer-specific policy when the virtual key claim has no mapping. Falls back to the global policy.",
)
model_config = {
"extra": "forbid",
@ -5063,6 +5071,28 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
super().__init__(**kwargs)
def get_issuer_config(self, issuer: str | None) -> JWTIssuerConfig | None:
if issuer is None or self.issuers is None:
return None
return next((config for config in self.issuers if config.issuer == issuer), None)
def is_virtual_key_mapping_configured(self) -> bool:
if self.virtual_key_claim_field is not None:
return True
return any(config.virtual_key_claim_field is not None for config in self.issuers or ())
def get_virtual_key_claim_field(self, issuer: str | None) -> str | None:
issuer_config: Final = self.get_issuer_config(issuer)
if issuer_config is not None and issuer_config.virtual_key_claim_field is not None:
return issuer_config.virtual_key_claim_field
return self.virtual_key_claim_field
def get_unregistered_jwt_client_behavior(self, issuer: str | None) -> UnregisteredJWTClientBehavior:
issuer_config: Final = self.get_issuer_config(issuer)
if issuer_config is not None and issuer_config.unregistered_jwt_client_behavior is not None:
return issuer_config.unregistered_jwt_client_behavior
return self.unregistered_jwt_client_behavior
class PrismaCompatibleUpdateDBModel(TypedDict, total=False):
model_name: str

19
litellm/proxy/auth/user_api_key_auth.py
View file

@ -987,9 +987,12 @@ async def _resolve_jwt_to_virtual_key(
- Raises HTTPException: REJECT policy hit, missing claim under
REJECT/AUTO_REGISTER, or other policy violations.
"""
virtual_key_claim_field: Final = jwt_handler.litellm_jwtauth.virtual_key_claim_field
raw_issuer: Final = jwt_claims.get(JWTHandler.LITELLM_JWT_ISSUER_CLAIM)
normalized_issuer: Final = raw_issuer if isinstance(raw_issuer, str) else None
virtual_key_claim_field: Final = jwt_handler.litellm_jwtauth.get_virtual_key_claim_field(normalized_issuer)
if virtual_key_claim_field is None:
return None
behavior: Final = jwt_handler.litellm_jwtauth.get_unregistered_jwt_client_behavior(normalized_issuer)
claim_value: Final = get_nested_value(
data=jwt_claims,
@ -1006,7 +1009,6 @@ async def _resolve_jwt_to_virtual_key(
# simply by presenting a JWT that omits the configured field. For
# AUTO_REGISTER there is no stable identity to map without a claim
# value, so we deny rather than create a sentinel-keyed record.
behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
if behavior in (
UnregisteredJWTClientBehavior.REJECT,
UnregisteredJWTClientBehavior.AUTO_REGISTER,
@ -1021,7 +1023,13 @@ async def _resolve_jwt_to_virtual_key(
return None
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER
cached_mapping: Final = (
None
if raw_cached_mapping == _JWT_PROXY_ADMIN_SENTINEL and not sentinel_written_by_this_policy
else raw_cached_mapping
)
if cached_mapping == _JWT_PROXY_ADMIN_SENTINEL:
# Previously resolved to a proxy admin via auth_builder; skip the
@ -1030,7 +1038,6 @@ async def _resolve_jwt_to_virtual_key(
return None
if cached_mapping == "__NO_MAPPING__":
behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
if behavior == UnregisteredJWTClientBehavior.REJECT:
raise HTTPException(
status_code=403,
@ -1093,8 +1100,6 @@ async def _resolve_jwt_to_virtual_key(
)
# No mapping found (DB miss or no DB) — apply no-match policy.
behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
if behavior == UnregisteredJWTClientBehavior.REJECT:
# Cache the miss before raising so repeated rejections are served from
# cache and don't re-query the DB on every request.
@ -1428,7 +1433,7 @@ async def _user_api_key_auth_builder(
# unnecessary DB queries in auth_builder
do_standard_jwt_auth = True
pending_auto_register: _PendingAutoRegister | None = None
if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None:
if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured():
# Decode JWT to get claims without running full auth_builder
jwt_claims: dict | None
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not is_jwt:

225
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -13,7 +13,7 @@ from unittest.mock import ANY, AsyncMock, MagicMock, patch
import pytest
from fastapi import status
from fastapi import HTTPException, status
import litellm
import litellm.proxy.proxy_server
@ -7442,3 +7442,226 @@ async def test_claude_view_never_reinterprets_explicit_names(monkeypatch, layer)
assert data["model"] == ("foo" if layer == "unclaimed" else encoded)
await _normalize_claude_model(data, token, request, "/v1/messages")
assert data["model"] == ("foo" if layer == "unclaimed" else encoded)
ISSUER_ONE = "https://issuer-one.example.com"
ISSUER_TWO = "https://issuer-two.example.com"
def _per_issuer_virtual_key_jwt_handler(
global_claim_field: str | None, global_behavior: str = "fallback_team_mapping"
) -> MagicMock:
jwt_handler = MagicMock()
jwt_handler.is_jwt.return_value = True
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field=global_claim_field,
unregistered_jwt_client_behavior=global_behavior,
issuers=[
{
"issuer": ISSUER_ONE,
"jwks_url": f"{ISSUER_ONE}/keys",
"audience": "audience-one",
"team_id_jwt_field": "sub",
},
{
"issuer": ISSUER_TWO,
"jwks_url": f"{ISSUER_TWO}/keys",
"audience": "audience-two",
"virtual_key_claim_field": "sub",
"unregistered_jwt_client_behavior": "reject",
},
],
)
return jwt_handler
def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]:
find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token))
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
return prisma_client, find_first
def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]:
return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True}
@pytest.mark.asyncio
async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for_the_db_lookup():
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None)
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping("hashed-mapped-key")
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-mapped-key",
value=UserAPIKeyAuth(token="hashed-mapped-key", api_key="hashed-mapped-key", team_id="svc-team"),
)
resolved = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "svc-account-7"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7"))
assert isinstance(resolved, UserAPIKeyAuth)
assert resolved.token == "hashed-mapped-key"
assert resolved.team_id == "svc-team"
assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key"
@pytest.mark.asyncio
async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer():
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None)
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None)
team_issuer_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert team_issuer_result is None
find_first.assert_not_awaited()
with pytest.raises(HTTPException) as exc:
await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "unknown-svc"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert exc.value.status_code == 403
assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc"))
@pytest.mark.asyncio
async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_reject():
from litellm.proxy.auth.user_api_key_auth import _JWT_PROXY_ADMIN_SENTINEL, _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register")
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL)
auto_register_issuer_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert auto_register_issuer_result is None
find_first.assert_not_awaited()
with pytest.raises(HTTPException) as exc:
await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "admin-7"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert exc.value.status_code == 403
assert "No registered mapping for sub='admin-7'" in str(exc.value.detail)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7"))
@pytest.mark.asyncio
async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_field():
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="client_id")
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None)
with_claim = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha", "client_id": "app-9"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
without_claim = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert with_claim is None
assert without_claim is None
find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9"))
@pytest.mark.asyncio
async def test_auth_flow_enters_virtual_key_mapping_when_only_an_issuer_configures_the_claim_field():
jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtYWNjb3VudC03In0.signature"
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None)
jwt_handler.auth_jwt = AsyncMock(
return_value={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "svc-account-7"}
)
mapped_key = UserAPIKeyAuth(token="hashed-mapped-key", api_key="hashed-mapped-key", team_id="svc-team")
mock_request = MagicMock()
mock_request.url.path = "/v1/chat/completions"
mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
mock_request.query_params = {}
mock_request.state = SimpleNamespace()
with (
patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam
"litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}
),
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: module-global proxy state
patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state
patch( # test-quality-ok: module-global proxy state
"litellm.proxy.proxy_server.user_api_key_cache", DualCache()
),
patch( # test-quality-ok: module-global proxy state
"litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()
),
patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), # test-quality-ok: module-global proxy state
patch( # test-quality-ok: the regression is whether the builder reaches this seam at all
"litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key",
new_callable=AsyncMock,
return_value=mapped_key,
) as resolve_mock,
patch( # test-quality-ok: a mapped key must short-circuit standard JWT auth; reaching it is the failure
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
new_callable=AsyncMock,
side_effect=AssertionError("standard JWT auth must not run for a mapped virtual key"),
),
):
result = await _user_api_key_auth_builder(
request=mock_request,
api_key=jwt_token,
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={"model": "gpt-4o-mini"},
)
resolve_mock.assert_awaited_once()
assert resolve_mock.await_args.kwargs["jwt_claims"][JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == ISSUER_TWO
assert result.api_key == "hashed-mapped-key"
assert result.team_id == "svc-team"

58
tests/test_litellm/proxy/test__types.py
View file

@ -277,3 +277,61 @@ def test_team_membership_budget_table_present_still_works():
}
result = LiteLLM_TeamMembership.model_validate(data)
assert result.litellm_budget_table is None
def test_a_jwt_issuer_can_override_the_virtual_key_claim_field_while_other_issuers_keep_the_global_one():
from litellm.proxy._types import LiteLLM_JWTAuth, UnregisteredJWTClientBehavior
jwt_auth = LiteLLM_JWTAuth(
virtual_key_claim_field="client_id",
issuers=[
{
"issuer": "https://team-idp.example.com",
"jwks_url": "https://team-idp.example.com/keys",
"audience": "litellm",
"team_id_jwt_field": "sub",
},
{
"issuer": "https://service-idp.example.com",
"jwks_url": "https://service-idp.example.com/keys",
"audience": "litellm",
"virtual_key_claim_field": "sub",
"unregistered_jwt_client_behavior": "reject",
},
],
)
assert jwt_auth.get_virtual_key_claim_field("https://service-idp.example.com") == "sub"
assert jwt_auth.get_unregistered_jwt_client_behavior("https://service-idp.example.com") is (
UnregisteredJWTClientBehavior.REJECT
)
assert jwt_auth.get_virtual_key_claim_field("https://team-idp.example.com") == "client_id"
assert jwt_auth.get_unregistered_jwt_client_behavior("https://team-idp.example.com") is (
UnregisteredJWTClientBehavior.FALLBACK_TEAM_MAPPING
)
assert jwt_auth.get_virtual_key_claim_field(None) == "client_id"
assert jwt_auth.get_virtual_key_claim_field("https://unknown-idp.example.com") == "client_id"
@pytest.mark.parametrize(
("global_field", "issuer_field", "is_configured"),
((None, None, False), ("sub", None, True), (None, "sub", True)),
)
def test_virtual_key_mapping_counts_as_configured_when_any_issuer_sets_the_claim_field(
global_field, issuer_field, is_configured
):
from litellm.proxy._types import LiteLLM_JWTAuth
jwt_auth = LiteLLM_JWTAuth(
virtual_key_claim_field=global_field,
issuers=[
{
"issuer": "https://idp.example.com",
"jwks_url": "https://idp.example.com/keys",
"audience": "litellm",
"virtual_key_claim_field": issuer_field,
}
],
)
assert jwt_auth.is_virtual_key_mapping_configured() is is_configured