scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

Merge pull request #24879 from BerriAI/worktree-eager-kindling-pnueli

Browse source 
feat(proxy): JWT-to-virtual-key mapping P0/P1/P2 improvements
This commit is contained in:
ishaan-berri 2026-03-31 15:58:16 -07:00 • committed by GitHub
commit 968e1eb105
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 1184 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
litellm/proxy/_types.py
View file

@ -938,6 +938,7 @@ class LiteLLMKeyType(str, enum.Enum):
MANAGEMENT = "management" # Can call management routes (user/team/key management)
READ_ONLY = "read_only" # Can only call info/read routes
DEFAULT = "default" # Uses default allowed routes
JWT_CLIENT = "jwt_client" # JWT-mapped service identity — LLM API routes only, no key management
class GenerateKeyRequest(KeyRequestBase):
@ -3848,6 +3849,7 @@ class KeyHealthResponse(TypedDict, total=False):
class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
jwt_claim_name: str
jwt_claim_value: str
issuer: str = ""
key: str
description: Optional[str] = None
@ -3867,12 +3869,77 @@ class JWTKeyMappingResponse(LiteLLMPydanticObjectBase):
id: str
jwt_claim_name: str
jwt_claim_value: str
issuer: str = ""
description: Optional[str] = None
is_active: bool
created_at: datetime
updated_at: datetime
created_by: Optional[str] = None
updated_by: Optional[str] = None
# Virtual key fields — populated by info/unified endpoints
models: Optional[list] = None
max_budget: Optional[float] = None
budget_duration: Optional[str] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
team_id: Optional[str] = None
key_alias: Optional[str] = None
spend: Optional[float] = None
expires: Optional[datetime] = None
# Cleartext key — only populated on creation (never returned again)
key: Optional[str] = None
class UpdateJWTClientRequest(LiteLLMPydanticObjectBase):
"""Request body for POST /jwt_client/update.
Uses a request body (not query params) so that callers can explicitly set
fields to null/unlimited — e.g. ``{"max_budget": null}`` clears the budget.
Fields absent from the body are left unchanged (exclude_unset semantics).
"""
id: str
# Mapping-level fields
description: Optional[str] = None
is_active: Optional[bool] = None
# Virtual key fields — null means "clear / set unlimited"
models: Optional[list] = None
max_budget: Optional[float] = None
budget_duration: Optional[str] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
model_config = {"extra": "ignore"}
class CreateJWTClientRequest(LiteLLMPydanticObjectBase):
"""Single-call request to create a virtual key + JWT mapping atomically."""
jwt_claim_name: str
jwt_claim_value: str
issuer: str = ""
description: Optional[str] = None
# Virtual key configuration
models: Optional[list] = []
max_budget: Optional[float] = None
budget_duration: Optional[str] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
team_id: Optional[str] = None
key_alias: Optional[str] = None
duration: Optional[str] = None
metadata: Optional[dict] = {}
class JWTClientAutoRegisterDefaults(LiteLLMPydanticObjectBase):
"""Default virtual key settings applied when auto-registering unknown JWT clients."""
models: Optional[list] = None
max_budget: Optional[float] = None
budget_duration: Optional[str] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
team_id: Optional[str] = None
class SpecialHeaders(enum.Enum):
@ -4200,6 +4267,21 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
default=300,
description="TTL (seconds) for caching JWT-to-virtual-key mapping lookups.",
)
unregistered_jwt_client_behavior: Literal[
"reject", "fallback_team_mapping", "auto_register"
] = Field(
default="fallback_team_mapping",
description=(
"Behavior when a JWT arrives with no virtual-key mapping. "
"'reject' → 403. "
"'fallback_team_mapping' → standard JWT team auth (default). "
"'auto_register' → create a mapping + virtual key on first encounter."
),
)
auto_register_defaults: Optional["JWTClientAutoRegisterDefaults"] = Field(
default=None,
description="Default virtual key settings used when auto_register creates keys for unknown JWT clients.",
)
#########################################################
def __init__(self, **kwargs: Any) -> None:

11
litellm/proxy/auth/auth_checks.py
View file

@ -8,6 +8,7 @@ Run checks for:
2. If user is in budget
3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
"""
import asyncio
import re
import time
@ -414,9 +415,9 @@ async def common_checks( # noqa: PLR0915
model=_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=valid_token.team_model_aliases
if valid_token
else None,
team_model_aliases=(
valid_token.team_model_aliases if valid_token else None
),
):
raise ProxyException(
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
@ -2146,16 +2147,20 @@ async def get_jwt_key_mapping_object(
jwt_claim_name: str,
jwt_claim_value: str,
prisma_client: PrismaClient,
issuer: str = "",
) -> Optional[str]:
"""
Lookup a JWT-to-virtual-key mapping from the database.
Returns the hashed token (str) if a matching active mapping is found, else None.
The `issuer` parameter (JWT "iss" claim) is used to disambiguate identical
claim values across different identity providers.
"""
mapping = await prisma_client.db.litellm_jwtkeymapping.find_first(
where={
"jwt_claim_name": jwt_claim_name,
"jwt_claim_value": jwt_claim_value,
"issuer": issuer,
"is_active": True,
}
)

8
litellm/proxy/auth/route_checks.py
View file

@ -209,7 +209,13 @@ class RouteChecks:
route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value
)
):
pass
# JWT-bound sessions (resolved from a JWT-to-virtual-key mapping) are
# restricted to LLM API routes only. Non-LLM routes reach this branch
# only because they passed the is_llm_api_route check above, so deny here.
if valid_token.jwt_claims is not None:
RouteChecks._raise_admin_only_route_exception(
user_obj=user_obj, route=route
)
elif _user_is_org_admin(
request_data=request_data, user_object=user_obj
) and RouteChecks.check_route_access(

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

@ -471,7 +471,8 @@ async def _resolve_jwt_to_virtual_key(
)
return None
cache_key = f"jwt_key_mapping:{virtual_key_claim_field}:{claim_value}"
issuer = str(jwt_claims.get("iss", ""))
cache_key = f"jwt_key_mapping:{virtual_key_claim_field}:{claim_value}:{issuer}"
cached_mapping = await user_api_key_cache.async_get_cache(cache_key)
if cached_mapping == "__NO_MAPPING__":
@ -492,6 +493,7 @@ async def _resolve_jwt_to_virtual_key(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=str(claim_value),
prisma_client=prisma_client,
issuer=issuer,
)
if token_hash is not None:
@ -507,7 +509,27 @@ async def _resolve_jwt_to_virtual_key(
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
if behavior == "reject":
raise HTTPException(
status_code=403,
detail=f"JWT client not registered. No mapping found for claim '{virtual_key_claim_field}' = '{claim_value}'.",
)
elif behavior == "auto_register":
return await _auto_register_jwt_client(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=str(claim_value),
issuer=issuer,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
cache_key=cache_key,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
else:
# "fallback_team_mapping" — default: fall through to standard JWT auth
await user_api_key_cache.async_set_cache(
key=cache_key,
value="__NO_MAPPING__",
@ -516,6 +538,101 @@ async def _resolve_jwt_to_virtual_key(
return None
async def _auto_register_jwt_client(
jwt_claim_name: str,
jwt_claim_value: str,
issuer: str,
jwt_handler: "JWTHandler",
prisma_client: "PrismaClient",
user_api_key_cache: DualCache,
cache_key: str,
parent_otel_span: Optional["Span"],
proxy_logging_obj: "ProxyLogging",
) -> Optional["UserAPIKeyAuth"]:
"""
Auto-register a new JWT client by creating a virtual key + mapping on first encounter.
Called when unregistered_jwt_client_behavior == 'auto_register'.
"""
from litellm.proxy._types import hash_token
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
)
defaults = jwt_handler.litellm_jwtauth.auto_register_defaults
jwt_metadata = {
"jwt_bound": True,
"jwt_claim_name": jwt_claim_name,
"jwt_claim_value": jwt_claim_value,
}
key_data = await generate_key_helper_fn(
request_type="key",
models=list(defaults.models or []) if defaults else [],
metadata=jwt_metadata,
key_max_budget=defaults.max_budget if defaults else None,
key_budget_duration=defaults.budget_duration if defaults else None,
tpm_limit=defaults.tpm_limit if defaults else None,
rpm_limit=defaults.rpm_limit if defaults else None,
team_id=defaults.team_id if defaults else None,
allowed_routes=["llm_api_routes"],
)
cleartext_token: str = key_data["token"]
hashed_key = hash_token(cleartext_token)
try:
await prisma_client.db.litellm_jwtkeymapping.create(
data={
"jwt_claim_name": jwt_claim_name,
"jwt_claim_value": jwt_claim_value,
"issuer": issuer,
"token": hashed_key,
}
)
except Exception as create_exc:
# Another concurrent request won the race and already inserted the mapping.
# Clean up the orphaned key we just created, then fetch the winning mapping.
error_str = str(create_exc).lower()
if "unique" in error_str or "p2002" in error_str:
try:
await prisma_client.db.litellm_verificationtoken.delete(
where={"token": hashed_key}
)
except Exception:
pass
winner = await get_jwt_key_mapping_object(
jwt_claim_name=jwt_claim_name,
jwt_claim_value=jwt_claim_value,
prisma_client=prisma_client,
issuer=issuer,
)
if winner is not None:
hashed_key = winner
else:
verbose_proxy_logger.warning(
f"JWT auto-register race: lost insert but winner mapping not found for "
f"{jwt_claim_name}={jwt_claim_value} issuer={issuer!r}"
)
return None
else:
verbose_proxy_logger.warning(
f"JWT auto-register: failed to create mapping for "
f"{jwt_claim_name}={jwt_claim_value} issuer={issuer!r}: {create_exc}"
)
await user_api_key_cache.async_set_cache(
key=cache_key,
value=hashed_key,
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
)
return await get_key_object(
hashed_token=hashed_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
async def _user_api_key_auth_builder( # noqa: PLR0915
request: Request,
api_key: str,

240
litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py
View file

@ -1,10 +1,15 @@
import json
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from litellm.proxy._types import (
CreateJWTClientRequest,
CreateJWTKeyMappingRequest,
DeleteJWTKeyMappingRequest,
JWTKeyMappingResponse,
LitellmUserRoles,
UpdateJWTClientRequest,
UpdateJWTKeyMappingRequest,
UserAPIKeyAuth,
hash_token,
@ -14,12 +19,17 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
router = APIRouter()
def _to_response(mapping) -> JWTKeyMappingResponse:
"""Convert a Prisma mapping object to a safe response (no hashed token)."""
return JWTKeyMappingResponse(
def _to_response(mapping, key_row=None) -> JWTKeyMappingResponse:
"""Convert a Prisma mapping object to a safe response (no hashed token).
Optionally accepts the linked LiteLLM_VerificationToken row to populate
virtual key fields (models, budget, rate limits, etc.).
"""
resp = JWTKeyMappingResponse(
id=mapping.id,
jwt_claim_name=mapping.jwt_claim_name,
jwt_claim_value=mapping.jwt_claim_value,
issuer=getattr(mapping, "issuer", ""),
description=mapping.description,
is_active=mapping.is_active,
created_at=mapping.created_at,
@ -27,6 +37,17 @@ def _to_response(mapping) -> JWTKeyMappingResponse:
created_by=mapping.created_by,
updated_by=mapping.updated_by,
)
if key_row is not None:
resp.models = key_row.models or []
resp.max_budget = key_row.max_budget
resp.budget_duration = key_row.budget_duration
resp.tpm_limit = key_row.tpm_limit
resp.rpm_limit = key_row.rpm_limit
resp.team_id = key_row.team_id
resp.key_alias = key_row.key_alias
resp.spend = key_row.spend
resp.expires = key_row.expires
return resp
@router.post(
@ -53,6 +74,7 @@ async def create_jwt_key_mapping(
create_data = {
"jwt_claim_name": data.jwt_claim_name,
"jwt_claim_value": data.jwt_claim_value,
"issuer": data.issuer,
"token": hashed_key,
"created_by": user_api_key_dict.user_id,
"updated_by": user_api_key_dict.user_id,
@ -64,8 +86,34 @@ async def create_jwt_key_mapping(
data=create_data
)
# Invalidate cache
cache_key = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}"
# Stamp jwt_bound metadata and restrict routes on the mapped key
existing_key = await prisma_client.db.litellm_verificationtoken.find_first(
where={"token": hashed_key}
)
if existing_key is not None:
existing_metadata = {}
if existing_key.metadata:
raw = existing_key.metadata
if isinstance(raw, str):
existing_metadata = json.loads(raw)
elif isinstance(raw, dict):
existing_metadata = raw
jwt_metadata = {
**existing_metadata,
"jwt_bound": True,
"jwt_claim_name": data.jwt_claim_name,
"jwt_claim_value": data.jwt_claim_value,
}
await prisma_client.db.litellm_verificationtoken.update(
where={"token": hashed_key},
data={
"metadata": json.dumps(jwt_metadata),
"allowed_routes": ["llm_api_routes"],
},
)
# Invalidate cache (include issuer in cache key)
cache_key = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}:{data.issuer}"
await user_api_key_cache.async_delete_cache(cache_key)
return _to_response(new_mapping)
@ -119,7 +167,8 @@ async def update_jwt_key_mapping(
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
old_issuer = getattr(old_mapping, "issuer", "")
cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}:{old_issuer}"
await user_api_key_cache.async_delete_cache(cache_key)
updated_mapping = await prisma_client.db.litellm_jwtkeymapping.update(
@ -127,7 +176,8 @@ async def update_jwt_key_mapping(
)
# Invalidate new cache key if claim fields changed
cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}"
new_issuer = getattr(updated_mapping, "issuer", "")
cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}:{new_issuer}"
await user_api_key_cache.async_delete_cache(cache_key)
return _to_response(updated_mapping)
@ -172,7 +222,8 @@ async def delete_jwt_key_mapping(
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
del_issuer = getattr(old_mapping, "issuer", "")
cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}:{del_issuer}"
await user_api_key_cache.async_delete_cache(cache_key)
await prisma_client.db.litellm_jwtkeymapping.delete(where={"id": data.id})
@ -247,10 +298,181 @@ async def info_jwt_key_mapping(
)
if mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
return _to_response(mapping)
key_row = await prisma_client.db.litellm_verificationtoken.find_first(
where={"token": mapping.token}
)
return _to_response(mapping, key_row=key_row)
except HTTPException:
raise
except Exception:
raise HTTPException(
status_code=500, detail="Failed to get JWT key mapping info."
)
@router.post(
"/jwt_client/new",
tags=["JWT Key Mapping"],
response_model=JWTKeyMappingResponse,
)
async def create_jwt_client(
data: CreateJWTClientRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Atomically create a virtual key + JWT mapping in one call.
The key is automatically typed as JWT_CLIENT (LLM API routes only)
and stamped with jwt_bound metadata.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
)
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403, detail="Only proxy admins can create JWT clients"
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
jwt_metadata = {
**(data.metadata or {}),
"jwt_bound": True,
"jwt_claim_name": data.jwt_claim_name,
"jwt_claim_value": data.jwt_claim_value,
}
key_data = await generate_key_helper_fn(
request_type="key",
table_name="key", # skip user-table upsert (user_id may be None for master key)
models=data.models or [],
metadata=jwt_metadata,
key_max_budget=data.max_budget,
key_budget_duration=data.budget_duration,
tpm_limit=data.tpm_limit,
rpm_limit=data.rpm_limit,
team_id=data.team_id,
key_alias=data.key_alias,
duration=data.duration,
allowed_routes=["llm_api_routes"],
created_by=user_api_key_dict.user_id,
updated_by=user_api_key_dict.user_id,
)
cleartext_token: str = key_data["token"]
hashed_key = hash_token(cleartext_token)
try:
create_data = {
"jwt_claim_name": data.jwt_claim_name,
"jwt_claim_value": data.jwt_claim_value,
"issuer": data.issuer,
"token": hashed_key,
"created_by": user_api_key_dict.user_id,
"updated_by": user_api_key_dict.user_id,
}
if data.description is not None:
create_data["description"] = data.description
new_mapping = await prisma_client.db.litellm_jwtkeymapping.create(
data=create_data
)
except Exception as e:
# Best-effort cleanup of the orphaned key
try:
await prisma_client.db.litellm_verificationtoken.delete(
where={"token": hashed_key}
)
except Exception:
pass
error_str = str(e).lower()
if "unique" in error_str or "p2002" in error_str:
raise HTTPException(
status_code=409,
detail=f"A JWT client for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.",
)
raise HTTPException(status_code=500, detail="Failed to create JWT client.")
cache_key = (
f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}:{data.issuer}"
)
await user_api_key_cache.async_delete_cache(cache_key)
key_row = await prisma_client.db.litellm_verificationtoken.find_first(
where={"token": hashed_key}
)
resp = _to_response(new_mapping, key_row=key_row)
resp.key = cleartext_token # return cleartext key only on creation
return resp
@router.post(
"/jwt_client/update",
tags=["JWT Key Mapping"],
response_model=JWTKeyMappingResponse,
)
async def update_jwt_client(
data: UpdateJWTClientRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Update a JWT client's virtual key configuration and/or mapping metadata.
Uses a request body so callers can explicitly null out fields — e.g.
``{"id": "...", "max_budget": null}`` removes the budget cap.
Fields absent from the body are left unchanged.
"""
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403, detail="Only proxy admins can update JWT clients"
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique(
where={"id": data.id}
)
if mapping is None:
raise HTTPException(status_code=404, detail="JWT client not found")
# Build mapping update from explicitly-set fields only
set_fields = data.model_fields_set
mapping_update: dict = {"updated_by": user_api_key_dict.user_id}
if "description" in set_fields:
mapping_update["description"] = data.description
if "is_active" in set_fields:
mapping_update["is_active"] = data.is_active
updated_mapping = await prisma_client.db.litellm_jwtkeymapping.update(
where={"id": data.id}, data=mapping_update
)
# Build key update — null values are intentional (clear the limit)
key_update: dict = {}
for field in ("models", "max_budget", "budget_duration", "tpm_limit", "rpm_limit"):
if field in set_fields:
key_update[field] = getattr(data, field)
key_row = None
if key_update:
key_row = await prisma_client.db.litellm_verificationtoken.update(
where={"token": mapping.token},
data=key_update,
)
else:
key_row = await prisma_client.db.litellm_verificationtoken.find_first(
where={"token": mapping.token}
)
issuer = getattr(mapping, "issuer", "")
cache_key = (
f"jwt_key_mapping:{mapping.jwt_claim_name}:{mapping.jwt_claim_value}:{issuer}"
)
await user_api_key_cache.async_delete_cache(cache_key)
return _to_response(updated_mapping, key_row=key_row)

62
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -453,6 +453,8 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
data_json["allowed_routes"] = ["management_routes"]
elif key_type == LiteLLMKeyType.READ_ONLY:
data_json["allowed_routes"] = ["info_routes"]
elif key_type == LiteLLMKeyType.JWT_CLIENT:
data_json["allowed_routes"] = ["llm_api_routes"]
return data_json
@ -732,9 +734,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
request_type="key", **data_json, table_name="key"
)
response[
"soft_budget"
] = data.soft_budget # include the user-input soft budget in the response
response["soft_budget"] = (
data.soft_budget
) # include the user-input soft budget in the response
response = GenerateKeyResponse(**response)
@ -2116,6 +2118,19 @@ async def update_key_fn(
prisma_client=prisma_client,
)
# JWT-bound keys can only be modified by proxy admins
_key_metadata = existing_key_row.metadata or {}
if isinstance(_key_metadata, str):
_key_metadata = json.loads(_key_metadata)
if (
_key_metadata.get("jwt_bound")
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
):
raise HTTPException(
status_code=403,
detail="JWT-bound keys can only be modified by proxy admins",
)
await _validate_update_key_data(
data=data,
existing_key_row=existing_key_row,
@ -3103,6 +3118,15 @@ async def can_modify_verification_token(
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return True
# 1b. JWT-bound keys cannot be modified by non-admin sessions
key_metadata = key_info.metadata or {}
if isinstance(key_metadata, str):
import json as _json
key_metadata = _json.loads(key_metadata)
if key_metadata.get("jwt_bound"):
return False
# 2. Internal jobs service account can modify any key (for auto-rotation)
if user_api_key_dict.api_key == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME:
return True
@ -3175,10 +3199,10 @@ async def delete_verification_tokens(
try:
if prisma_client:
tokens = [_hash_token_if_needed(token=key) for key in tokens]
_keys_being_deleted: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
_keys_being_deleted: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
)
)
if len(_keys_being_deleted) == 0:
@ -3378,9 +3402,9 @@ async def _rotate_master_key( # noqa: PLR0915
from litellm.proxy.proxy_server import proxy_config
try:
models: Optional[
List
] = await prisma_client.db.litellm_proxymodeltable.find_many()
models: Optional[List] = (
await prisma_client.db.litellm_proxymodeltable.find_many()
)
except Exception:
models = None
# 2. process model table
@ -4020,11 +4044,11 @@ async def validate_key_list_check(
param="user_id",
code=status.HTTP_403_FORBIDDEN,
)
complete_user_info_db_obj: Optional[
BaseModel
] = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
complete_user_info_db_obj: Optional[BaseModel] = (
await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
)
)
if complete_user_info_db_obj is None:
@ -4107,10 +4131,10 @@ async def _fetch_user_team_objects(
if complete_user_info is None or not complete_user_info.teams:
return []
teams: Optional[
List[BaseModel]
] = await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
teams: Optional[List[BaseModel]] = (
await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
)
)
if teams is None:
return []

5
litellm/proxy/schema.prisma
View file

@ -408,6 +408,7 @@ model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
issuer String @default("") // JWT "iss" claim — differentiates identical claims across IdPs
token String // Hashed virtual key (FK)
description String?
is_active Boolean @default(true)
@ -418,8 +419,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_claim_name, jwt_claim_value, issuer])
@@index([jwt_claim_name, jwt_claim_value, issuer, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

693
tests/proxy_unit_tests/test_jwt_key_mapping.py
View file

@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
from litellm.proxy.auth.user_api_key_auth import (
_auto_register_jwt_client,
_resolve_jwt_to_virtual_key,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
@ -15,6 +16,7 @@ from litellm.proxy._types import (
JWTKeyMappingResponse,
LiteLLM_JWTAuth,
LitellmUserRoles,
UpdateJWTClientRequest,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
@ -22,6 +24,7 @@ from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
create_jwt_key_mapping,
delete_jwt_key_mapping,
info_jwt_key_mapping,
update_jwt_client,
update_jwt_key_mapping,
)
from litellm.caching.caching import DualCache
@ -222,6 +225,7 @@ def test_to_response_excludes_token():
mock_mapping.id = "mapping-1"
mock_mapping.jwt_claim_name = "email"
mock_mapping.jwt_claim_value = "user@example.com"
mock_mapping.issuer = ""
mock_mapping.token = "hashed_secret_value"
mock_mapping.description = "test"
mock_mapping.is_active = True
@ -265,6 +269,10 @@ def _mock_prisma():
prisma.db.litellm_jwtkeymapping.update = AsyncMock()
prisma.db.litellm_jwtkeymapping.delete = AsyncMock()
prisma.db.litellm_jwtkeymapping.count = AsyncMock(return_value=0)
# Needed by create_jwt_key_mapping (stamps jwt_bound metadata) and info endpoint
prisma.db.litellm_verificationtoken.find_first = AsyncMock(return_value=None)
prisma.db.litellm_verificationtoken.update = AsyncMock()
prisma.db.litellm_verificationtoken.delete = AsyncMock()
return prisma
@ -272,12 +280,14 @@ def _mock_mapping(
id="mapping-1",
claim_name="email",
claim_value="user@example.com",
issuer="",
):
now = datetime.now(timezone.utc)
m = MagicMock()
m.id = id
m.jwt_claim_name = claim_name
m.jwt_claim_value = claim_value
m.issuer = issuer
m.token = "hashed_token"
m.description = None
m.is_active = True
@ -398,8 +408,11 @@ async def test_info_returns_404_when_not_found():
"""Getting info for non-existent mapping should return 404."""
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None
mock_cache = AsyncMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
):
with pytest.raises(HTTPException) as exc_info:
await info_jwt_key_mapping(id="nonexistent-id", user_api_key_dict=_make_admin_auth())
assert exc_info.value.status_code == 404
@ -425,3 +438,681 @@ async def test_create_success_returns_response_without_token():
assert isinstance(result, JWTKeyMappingResponse)
assert "token" not in result.model_fields
assert result.jwt_claim_name == "email"
# ──────────────────────────────────────────────
# Block 1: Gap #2 — JWT-bound metadata stamping
# ──────────────────────────────────────────────
def test_handle_key_type_jwt_client():
"""JWT_CLIENT key type should resolve to llm_api_routes."""
from litellm.proxy._types import LiteLLMKeyType
from litellm.proxy.management_endpoints.key_management_endpoints import (
handle_key_type,
)
from litellm.proxy._types import GenerateKeyRequest
data = GenerateKeyRequest(key_type=LiteLLMKeyType.JWT_CLIENT)
data_json = data.model_dump()
result = handle_key_type(data=data, data_json=data_json)
assert result["allowed_routes"] == ["llm_api_routes"]
@pytest.mark.asyncio
async def test_create_mapping_stamps_jwt_bound_metadata():
"""create_jwt_key_mapping should update the key's metadata with jwt_bound=True."""
from litellm.proxy._types import CreateJWTKeyMappingRequest
import json as _json
mock_prisma = _mock_prisma()
mock_mapping = _mock_mapping()
mock_prisma.db.litellm_jwtkeymapping.create.return_value = mock_mapping
# Simulate key with no existing metadata
mock_key_row = MagicMock()
mock_key_row.metadata = "{}"
mock_prisma.db.litellm_verificationtoken.find_first.return_value = mock_key_row
mock_cache = AsyncMock()
data = CreateJWTKeyMappingRequest(
jwt_claim_name="sub",
jwt_claim_value="svc-a",
key="sk-test-key",
)
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
):
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
mock_prisma.db.litellm_verificationtoken.update.assert_called_once()
call_kwargs = mock_prisma.db.litellm_verificationtoken.update.call_args
update_data = call_kwargs.kwargs["data"]
stored_metadata = _json.loads(update_data["metadata"])
assert stored_metadata["jwt_bound"] is True
assert stored_metadata["jwt_claim_name"] == "sub"
assert stored_metadata["jwt_claim_value"] == "svc-a"
assert update_data["allowed_routes"] == ["llm_api_routes"]
@pytest.mark.asyncio
async def test_create_mapping_preserves_existing_metadata():
"""Existing metadata keys should be preserved when stamping jwt_bound."""
from litellm.proxy._types import CreateJWTKeyMappingRequest
import json as _json
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping()
mock_key_row = MagicMock()
mock_key_row.metadata = _json.dumps({"custom_field": "keep_me"})
mock_prisma.db.litellm_verificationtoken.find_first.return_value = mock_key_row
mock_cache = AsyncMock()
data = CreateJWTKeyMappingRequest(
jwt_claim_name="email",
jwt_claim_value="user@example.com",
key="sk-test-key",
)
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
):
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
update_data = mock_prisma.db.litellm_verificationtoken.update.call_args.kwargs["data"]
stored_metadata = _json.loads(update_data["metadata"])
assert stored_metadata["custom_field"] == "keep_me"
assert stored_metadata["jwt_bound"] is True
# ──────────────────────────────────────────────
# Block 2: Gap #1 — CRUD block on jwt-bound keys
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_can_modify_jwt_bound_key_as_admin_returns_true():
"""Proxy admin can always modify JWT-bound keys."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
can_modify_verification_token,
)
from litellm.proxy._types import LiteLLM_VerificationToken
key_info = MagicMock(spec=LiteLLM_VerificationToken)
key_info.metadata = {"jwt_bound": True}
key_info.team_id = None
key_info.user_id = "some-user"
admin_dict = _make_admin_auth()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=MagicMock(),
user_api_key_dict=admin_dict,
prisma_client=MagicMock(),
)
assert result is True
@pytest.mark.asyncio
async def test_can_modify_jwt_bound_key_as_internal_user_returns_false():
"""Non-admin cannot modify a JWT-bound key even if they own it."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
can_modify_verification_token,
)
from litellm.proxy._types import LiteLLM_VerificationToken
key_info = MagicMock(spec=LiteLLM_VerificationToken)
key_info.metadata = {"jwt_bound": True}
key_info.team_id = None
key_info.user_id = "user-123"
non_admin = _make_non_admin_auth()
non_admin.user_id = "user-123" # same user_id as key owner
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=MagicMock(),
user_api_key_dict=non_admin,
prisma_client=MagicMock(),
)
assert result is False
def test_jwt_session_blocked_from_key_management_route():
"""A JWT-authenticated session (jwt_claims set) must not reach key management routes."""
from litellm.proxy.auth.route_checks import RouteChecks
from unittest.mock import MagicMock
user_obj = MagicMock()
user_obj.user_id = "user-123"
valid_token = UserAPIKeyAuth(
token="sk-bound",
user_role=LitellmUserRoles.INTERNAL_USER,
jwt_claims={"sub": "svc-a"}, # marks this as a JWT session
)
with pytest.raises(Exception):
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER,
route="/key/update",
request=MagicMock(),
valid_token=valid_token,
request_data={},
)
def test_jwt_session_allowed_llm_api_route():
"""A JWT-authenticated session must be allowed to call LLM API routes."""
from litellm.proxy.auth.route_checks import RouteChecks
valid_token = UserAPIKeyAuth(
token="sk-bound",
user_role=LitellmUserRoles.INTERNAL_USER,
jwt_claims={"sub": "svc-a"},
)
# Should not raise for an LLM API route
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=MagicMock(),
_user_role=LitellmUserRoles.INTERNAL_USER,
route="/v1/chat/completions",
request=MagicMock(),
valid_token=valid_token,
request_data={},
)
def test_non_jwt_session_internal_user_can_access_key_management():
"""Regular API key session (no jwt_claims) keeps existing internal_user access."""
from litellm.proxy.auth.route_checks import RouteChecks
valid_token = UserAPIKeyAuth(
token="sk-regular",
user_role=LitellmUserRoles.INTERNAL_USER,
jwt_claims=None, # not a JWT session
)
# Should not raise — existing behavior preserved
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=MagicMock(),
_user_role=LitellmUserRoles.INTERNAL_USER,
route="/key/update",
request=MagicMock(),
valid_token=valid_token,
request_data={},
)
# ──────────────────────────────────────────────
# Block 3: Gap #3 — Unified /jwt_client/new
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_jwt_client_new_creates_key_and_mapping():
"""/jwt_client/new should call generate_key_helper_fn and create a mapping row."""
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
create_jwt_client,
)
from litellm.proxy._types import CreateJWTClientRequest
import json as _json
mock_prisma = _mock_prisma()
mock_mapping = _mock_mapping(claim_name="sub", claim_value="svc-a")
mock_prisma.db.litellm_jwtkeymapping.create.return_value = mock_mapping
mock_cache = AsyncMock()
fake_key_data = {"token": "sk-auto-generated-key"}
data = CreateJWTClientRequest(
jwt_claim_name="sub",
jwt_claim_value="svc-a",
models=["gpt-4o"],
max_budget=10.0,
)
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value=fake_key_data,
) as mock_gen:
result = await create_jwt_client(data=data, user_api_key_dict=_make_admin_auth())
mock_gen.assert_called_once()
call_kwargs = mock_gen.call_args.kwargs
assert call_kwargs["allowed_routes"] == ["llm_api_routes"]
assert call_kwargs["metadata"]["jwt_bound"] is True
assert call_kwargs["metadata"]["jwt_claim_name"] == "sub"
mock_prisma.db.litellm_jwtkeymapping.create.assert_called_once()
assert isinstance(result, JWTKeyMappingResponse)
@pytest.mark.asyncio
async def test_jwt_client_new_cache_invalidated():
"""/jwt_client/new must invalidate the cache for the new mapping."""
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
create_jwt_client,
)
from litellm.proxy._types import CreateJWTClientRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping(
claim_name="sub", claim_value="svc-b"
)
mock_cache = MagicMock()
mock_cache.async_delete_cache = AsyncMock()
data = CreateJWTClientRequest(jwt_claim_name="sub", jwt_claim_value="svc-b")
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"token": "sk-xyz"},
):
await create_jwt_client(data=data, user_api_key_dict=_make_admin_auth())
mock_cache.async_delete_cache.assert_called_once()
cache_key_arg = mock_cache.async_delete_cache.call_args.args[0]
assert "sub" in cache_key_arg and "svc-b" in cache_key_arg
@pytest.mark.asyncio
async def test_jwt_client_new_non_admin_rejected():
"""/jwt_client/new must reject non-admin users with 403."""
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
create_jwt_client,
)
from litellm.proxy._types import CreateJWTClientRequest
data = CreateJWTClientRequest(jwt_claim_name="sub", jwt_claim_value="svc-c")
with pytest.raises(HTTPException) as exc_info:
await create_jwt_client(data=data, user_api_key_dict=_make_non_admin_auth())
assert exc_info.value.status_code == 403
# ──────────────────────────────────────────────
# Block 4: Gap #4 — unregistered_jwt_client_behavior
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_reject_behavior_raises_403_for_unknown_jwt():
"""'reject' mode should raise 403 when no mapping exists."""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub",
unregistered_jwt_client_behavior="reject",
)
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
with pytest.raises(HTTPException) as exc_info:
await _resolve_jwt_to_virtual_key(
jwt_claims={"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_info.value.status_code == 403
@pytest.mark.asyncio
async def test_fallback_behavior_returns_none_for_unknown_jwt():
"""Default 'fallback_team_mapping' mode should return None for unknown clients."""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub",
# default: fallback_team_mapping
)
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
with patch(
"litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock
):
result = await _resolve_jwt_to_virtual_key(
jwt_claims={"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 result is None
@pytest.mark.asyncio
async def test_auto_register_creates_key_on_first_request():
"""'auto_register' mode should create a key+mapping on first unknown JWT."""
from litellm.proxy._types import JWTClientAutoRegisterDefaults
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub",
unregistered_jwt_client_behavior="auto_register",
auto_register_defaults=JWTClientAutoRegisterDefaults(
models=["gpt-4o-mini"], max_budget=5.0
),
)
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
mock_key_obj = UserAPIKeyAuth(token="sk-auto", team_id=None)
with patch(
"litellm.proxy.auth.user_api_key_auth.get_key_object",
new_callable=AsyncMock,
return_value=mock_key_obj,
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"token": "sk-auto"},
) as mock_gen:
result = await _resolve_jwt_to_virtual_key(
jwt_claims={"sub": "new-svc"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert result is not None
mock_gen.assert_called_once()
call_kwargs = mock_gen.call_args.kwargs
assert call_kwargs["allowed_routes"] == ["llm_api_routes"]
assert call_kwargs["metadata"]["jwt_bound"] is True
# ──────────────────────────────────────────────
# Block 5: Gap #5 — issuer column
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_get_jwt_key_mapping_object_passes_issuer_to_db():
"""get_jwt_key_mapping_object must include issuer in the DB where clause."""
from litellm.proxy.auth.auth_checks import get_jwt_key_mapping_object
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
await get_jwt_key_mapping_object(
jwt_claim_name="sub",
jwt_claim_value="svc",
prisma_client=prisma_client,
issuer="https://idp1.example.com",
)
call_kwargs = prisma_client.db.litellm_jwtkeymapping.find_first.call_args.kwargs
assert call_kwargs["where"]["issuer"] == "https://idp1.example.com"
@pytest.mark.asyncio
async def test_resolve_extracts_issuer_from_jwt_claims():
"""_resolve_jwt_to_virtual_key must pass iss claim as issuer to DB lookup."""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_claim_field="sub")
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
with patch(
"litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock
):
await _resolve_jwt_to_virtual_key(
jwt_claims={"sub": "svc-a", "iss": "https://idp2.example.com"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
call_kwargs = prisma_client.db.litellm_jwtkeymapping.find_first.call_args.kwargs
assert call_kwargs["where"]["issuer"] == "https://idp2.example.com"
@pytest.mark.asyncio
async def test_no_issuer_in_jwt_defaults_to_empty_string():
"""When JWT has no 'iss' field, issuer should default to empty string."""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_claim_field="sub")
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
with patch(
"litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock
):
await _resolve_jwt_to_virtual_key(
jwt_claims={"sub": "svc-b"}, # no "iss"
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
call_kwargs = prisma_client.db.litellm_jwtkeymapping.find_first.call_args.kwargs
assert call_kwargs["where"]["issuer"] == ""
# ──────────────────────────────────────────────
# Block 6: Gap #6 — endpoints expose virtual key properties
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_info_endpoint_returns_virtual_key_fields():
"""/jwt/key/mapping/info should return virtual key fields when the key row exists."""
now = datetime.now(timezone.utc)
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping()
key_row = MagicMock()
key_row.models = ["gpt-4o"]
key_row.max_budget = 50.0
key_row.budget_duration = "30d"
key_row.tpm_limit = 10000
key_row.rpm_limit = 100
key_row.team_id = "team-1"
key_row.key_alias = "my-client"
key_row.spend = 2.5
key_row.expires = now
mock_prisma.db.litellm_verificationtoken.find_first.return_value = key_row
mock_cache = AsyncMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
):
result = await info_jwt_key_mapping(
id="mapping-1", user_api_key_dict=_make_admin_auth()
)
assert result.models == ["gpt-4o"]
assert result.max_budget == 50.0
assert result.team_id == "team-1"
assert result.key_alias == "my-client"
assert result.spend == 2.5
@pytest.mark.asyncio
async def test_info_endpoint_still_returns_mapping_fields():
"""/jwt/key/mapping/info response must still include mapping metadata."""
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping(
claim_name="email", claim_value="user@corp.com"
)
mock_prisma.db.litellm_verificationtoken.find_first.return_value = None
mock_cache = AsyncMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
):
result = await info_jwt_key_mapping(
id="mapping-1", user_api_key_dict=_make_admin_auth()
)
assert result.jwt_claim_name == "email"
assert result.jwt_claim_value == "user@corp.com"
assert result.is_active is True
# ---------------------------------------------------------------------------
# P1/P2 fix tests
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_auto_register_forwards_issuer_to_db():
"""issuer must be written to the mapping row, not hardcoded to ''."""
from litellm.proxy._types import (
JWTClientAutoRegisterDefaults,
LiteLLM_JWTAuth,
hash_token,
)
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.create.return_value = MagicMock()
mock_cache = AsyncMock()
jwtauth = LiteLLM_JWTAuth(
unregistered_jwt_client_behavior="auto_register",
auto_register_defaults=JWTClientAutoRegisterDefaults(models=["gpt-4o-mini"]),
)
jwt_handler = MagicMock()
jwt_handler.litellm_jwtauth = jwtauth
fake_key_data = {"token": "sk-autotest"}
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new=AsyncMock(return_value=fake_key_data),
), patch(
"litellm.proxy.auth.user_api_key_auth.get_key_object",
new=AsyncMock(return_value=MagicMock()),
):
await _auto_register_jwt_client(
jwt_claim_name="sub",
jwt_claim_value="svc-a",
issuer="https://idp1.example.com",
jwt_handler=jwt_handler,
prisma_client=mock_prisma,
user_api_key_cache=mock_cache,
cache_key="jwt_key_mapping:sub:svc-a:https://idp1.example.com",
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
create_call = mock_prisma.db.litellm_jwtkeymapping.create.call_args
assert create_call.kwargs["data"]["issuer"] == "https://idp1.example.com"
@pytest.mark.asyncio
async def test_auto_register_race_condition_cleans_up_orphan():
"""When the DB insert races and loses (P2002), the orphaned key is deleted
and the winning mapping's token_hash is used for the cache entry."""
from litellm.proxy._types import LiteLLM_JWTAuth, hash_token
winner_hash = hash_token("sk-winner")
mock_prisma = _mock_prisma()
# Simulate the unique-constraint violation on insert
mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception(
"Unique constraint failed (p2002)"
)
mock_prisma.db.litellm_verificationtoken.delete.return_value = None
# get_jwt_key_mapping_object path — the find_first called by get_jwt_key_mapping_object
mock_prisma.db.litellm_jwtkeymapping.find_first.return_value = MagicMock(
token=winner_hash
)
mock_cache = AsyncMock()
jwt_handler = MagicMock()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
fake_key_data = {"token": "sk-loser"}
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new=AsyncMock(return_value=fake_key_data),
), patch(
"litellm.proxy.auth.user_api_key_auth.get_key_object",
new=AsyncMock(return_value=MagicMock()),
), patch(
"litellm.proxy.auth.user_api_key_auth.get_jwt_key_mapping_object",
new=AsyncMock(return_value=winner_hash),
):
await _auto_register_jwt_client(
jwt_claim_name="sub",
jwt_claim_value="svc-race",
issuer="",
jwt_handler=jwt_handler,
prisma_client=mock_prisma,
user_api_key_cache=mock_cache,
cache_key="jwt_key_mapping:sub:svc-race:",
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
# Orphan key was deleted
mock_prisma.db.litellm_verificationtoken.delete.assert_called_once()
# Cache was populated with the winner's hash, not the loser's
set_cache_call = mock_cache.async_set_cache.call_args
assert set_cache_call.kwargs["value"] == winner_hash
@pytest.mark.asyncio
async def test_jwt_client_update_null_clears_budget():
"""Sending max_budget=null in the body should explicitly set it to None
(unlimited), not be silently skipped."""
mock_prisma = _mock_prisma()
from datetime import datetime as _dt
mock_mapping = MagicMock()
mock_mapping.id = "map-1"
mock_mapping.token = "hashed-token"
mock_mapping.jwt_claim_name = "sub"
mock_mapping.jwt_claim_value = "svc"
mock_mapping.issuer = ""
mock_mapping.description = None
mock_mapping.is_active = True
mock_mapping.created_at = _dt.utcnow()
mock_mapping.updated_at = _dt.utcnow()
mock_mapping.created_by = None
mock_mapping.updated_by = None
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = mock_mapping
mock_prisma.db.litellm_jwtkeymapping.update.return_value = mock_mapping
mock_key_row = MagicMock()
mock_key_row.max_budget = None
mock_prisma.db.litellm_verificationtoken.update.return_value = mock_key_row
mock_cache = AsyncMock()
# Explicit null for max_budget — field IS in model_fields_set
req = UpdateJWTClientRequest.model_validate({"id": "map-1", "max_budget": None})
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
):
await update_jwt_client(data=req, user_api_key_dict=_make_admin_auth())
key_update_call = mock_prisma.db.litellm_verificationtoken.update.call_args
assert "max_budget" in key_update_call.kwargs["data"]
assert key_update_call.kwargs["data"]["max_budget"] is None