Merge pull request #22372 from BerriAI/litellm_jwt_vkey_map

Litellm jwt vkey map
This commit is contained in:
Harshit Jain 2026-03-05 06:24:49 +05:30 committed by GitHub
commit 07cb6d5bec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 1075 additions and 177 deletions

View File

@ -1054,6 +1054,95 @@ curl -X GET 'http://0.0.0.0:4000/user/info?user_id=user-123' \
-H 'Authorization: Bearer <PROXY_MASTER_KEY>'
```
## [BETA] JWT-to-Virtual-Key Mapping
Map JWT identities to LiteLLM virtual keys so that JWT-authenticated users get per-user budgets, rate limits, model access controls, and spend tracking.
When a JWT comes in, LiteLLM looks up a configured claim (e.g. `email`, `sub`) in a mapping table. If a mapping exists, the request is treated as if it arrived with the corresponding virtual key — all virtual key features apply.
### Setup
Add `virtual_key_claim_field` to your JWT auth config:
```yaml
general_settings:
enable_jwt_auth: True
litellm_jwtauth:
virtual_key_claim_field: "email" # JWT claim to look up (supports dot notation)
virtual_key_mapping_cache_ttl: 300 # Cache TTL in seconds (default: 300)
```
### Managing Mappings
All endpoints require admin auth (`Authorization: Bearer <master_key>`).
**Create a mapping** — link a JWT claim value to an existing virtual key:
```bash
curl -X POST http://localhost:4000/jwt/key/mapping/new \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"jwt_claim_name": "email",
"jwt_claim_value": "user@example.com",
"key": "sk-virtual-key-from-key-generate"
}'
```
**List mappings** (paginated):
```bash
curl http://localhost:4000/jwt/key/mapping/list?page=1&size=50 \
-H "Authorization: Bearer sk-1234"
```
**Get a specific mapping:**
```bash
curl "http://localhost:4000/jwt/key/mapping/info?id=<mapping-id>" \
-H "Authorization: Bearer sk-1234"
```
**Update a mapping:**
```bash
curl -X POST http://localhost:4000/jwt/key/mapping/update \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"id": "<mapping-id>",
"description": "Updated description",
"is_active": true
}'
```
**Delete a mapping:**
```bash
curl -X POST http://localhost:4000/jwt/key/mapping/delete \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{"id": "<mapping-id>"}'
```
### How It Works
1. A request arrives with a JWT bearer token
2. LiteLLM validates the JWT signature
3. Extracts the configured claim (e.g. `email``user@example.com`)
4. Looks up the claim value in the `LiteLLM_JWTKeyMapping` table
5. If a mapping exists, the request proceeds as if the mapped virtual key was used — budgets, rate limits, model access, and spend tracking all apply
6. If no mapping exists, falls back to standard JWT auth (team-level controls)
### Error Codes
| Code | Meaning |
|------|---------|
| 409 | Duplicate mapping — a mapping for that claim name + value already exists |
| 400 | The provided key does not match an existing virtual key |
| 404 | Mapping not found (for update/delete/info) |
| 403 | Non-admin user attempted a mapping operation |
## All JWT Params
[**See Code**](https://github.com/BerriAI/litellm/blob/b204f0c01c703317d812a1553363ab0cb989d5b6/litellm/proxy/_types.py#L95)

View File

@ -543,6 +543,11 @@ class LiteLLMRoutes(enum.Enum):
"/model/update",
"/model/delete",
"/model/info",
"/jwt/key/mapping/new",
"/jwt/key/mapping/update",
"/jwt/key/mapping/delete",
"/jwt/key/mapping/list",
"/jwt/key/mapping/info",
] + key_management_routes
spend_tracking_routes = [
@ -3685,6 +3690,36 @@ class KeyHealthResponse(TypedDict, total=False):
logging_callbacks: Optional[LoggingCallbackStatus]
class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
jwt_claim_name: str
jwt_claim_value: str
key: str
description: Optional[str] = None
class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
id: str
key: Optional[str] = None
description: Optional[str] = None
is_active: Optional[bool] = None
class DeleteJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
id: str
class JWTKeyMappingResponse(LiteLLMPydanticObjectBase):
id: str
jwt_claim_name: str
jwt_claim_value: str
description: Optional[str] = None
is_active: bool
created_at: datetime
updated_at: datetime
created_by: Optional[str] = None
updated_by: Optional[str] = None
class SpecialHeaders(enum.Enum):
"""Used by user_api_key_auth.py to get litellm key"""
@ -3857,6 +3892,7 @@ class JWTAuthBuilderResult(TypedDict):
end_user_id: Optional[str]
org_id: Optional[str]
team_membership: Optional[LiteLLM_TeamMembership]
jwt_claims: dict # Decoded JWT token claims (avoids re-decoding)
class ClientSideFallbackModel(TypedDict, total=False):
@ -4000,6 +4036,15 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
default=300,
description="TTL (in seconds) for caching UserInfo responses. Default: 300s (5 minutes).",
)
# JWT-to-Virtual-Key Mapping
virtual_key_claim_field: Optional[str] = Field(
default=None,
description="JWT claim field for virtual key mapping lookup (e.g. 'sub', 'email'). Supports dot notation.",
)
virtual_key_mapping_cache_ttl: float = Field(
default=300,
description="TTL (seconds) for caching JWT-to-virtual-key mapping lookups.",
)
#########################################################
def __init__(self, **kwargs: Any) -> None:

View File

@ -2081,6 +2081,29 @@ async def _fetch_key_object_from_db_with_reconnect(
raise
@log_db_metrics
async def get_jwt_key_mapping_object(
jwt_claim_name: str,
jwt_claim_value: str,
prisma_client: PrismaClient,
) -> 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.
"""
mapping = await prisma_client.db.litellm_jwtkeymapping.find_first(
where={
"jwt_claim_name": jwt_claim_name,
"jwt_claim_value": jwt_claim_value,
"is_active": True,
}
)
if mapping is not None:
return mapping.token
return None
@log_db_metrics
async def get_key_object(
hashed_token: str,

View File

@ -166,7 +166,6 @@ class JWTHandler:
return False
def get_team_ids_from_jwt(self, token: dict) -> List[str]:
if self.litellm_jwtauth.team_ids_jwt_field is not None:
team_ids: Optional[List[str]] = get_nested_value(
data=token,
@ -256,7 +255,9 @@ class JWTHandler:
team_id = default_value
return team_id
def get_team_alias(self, token: dict, default_value: Optional[str]) -> Optional[str]:
def get_team_alias(
self, token: dict, default_value: Optional[str]
) -> Optional[str]:
"""
Extract team name/alias from JWT token using the configured team_alias_jwt_field.
@ -596,17 +597,17 @@ class JWTHandler:
async def get_oidc_userinfo(self, token: str) -> dict:
"""
Fetch user information from OIDC UserInfo endpoint.
This follows the OpenID Connect protocol where an access token
is sent to the identity provider's UserInfo endpoint to retrieve
user identity information.
Args:
token: The access token to use for authentication
Returns:
dict: User information from the UserInfo endpoint
Raises:
Exception: If UserInfo endpoint is not configured or request fails
"""
@ -614,19 +615,21 @@ class JWTHandler:
raise Exception(
"OIDC UserInfo endpoint not configured. Set 'oidc_userinfo_endpoint' in JWT auth config."
)
# Check cache first
cache_key = f"oidc_userinfo_{token[:20]}" # Use first 20 chars of token as cache key
cache_key = (
f"oidc_userinfo_{token[:20]}" # Use first 20 chars of token as cache key
)
cached_userinfo = await self.user_api_key_cache.async_get_cache(cache_key)
if cached_userinfo is not None:
verbose_proxy_logger.debug("Returning cached OIDC UserInfo")
return cached_userinfo
verbose_proxy_logger.debug(
f"Calling OIDC UserInfo endpoint: {self.litellm_jwtauth.oidc_userinfo_endpoint}"
)
try:
# Call the UserInfo endpoint with the access token
response = await self.http_handler.get(
@ -636,24 +639,24 @@ class JWTHandler:
"Accept": "application/json",
},
)
if response.status_code != 200:
raise Exception(
f"OIDC UserInfo endpoint returned status {response.status_code}: {response.text}"
)
userinfo = response.json()
verbose_proxy_logger.debug(f"Received OIDC UserInfo: {userinfo}")
# Cache the userinfo response
await self.user_api_key_cache.async_set_cache(
key=cache_key,
value=userinfo,
ttl=self.litellm_jwtauth.oidc_userinfo_cache_ttl,
)
return userinfo
except Exception as e:
verbose_proxy_logger.error(f"Error fetching OIDC UserInfo: {str(e)}")
raise Exception(f"Failed to fetch OIDC UserInfo: {str(e)}")
@ -886,6 +889,7 @@ class JWTAuthManager:
user_id: Optional[str],
org_id: Optional[str],
api_key: str,
jwt_valid_token: Optional[dict] = None,
) -> Optional[JWTAuthBuilderResult]:
"""Check admin status and route access permissions"""
if not jwt_handler.is_admin(scopes=scopes):
@ -915,6 +919,7 @@ class JWTAuthManager:
end_user_id=None,
org_id=org_id,
team_membership=None,
jwt_claims=jwt_valid_token or {},
)
@staticmethod
@ -1111,11 +1116,11 @@ class JWTAuthManager:
) -> Tuple[
Optional[LiteLLM_UserTable],
Optional[LiteLLM_OrganizationTable],
Optional[LiteLLM_EndUserTable],
Optional[LiteLLM_EndUserTable],
Optional[LiteLLM_TeamMembership],
]:
"""Get user, org, and end user objects. Also resolves org aliases to IDs if configured."""
# Get org object - first try by ID, then by alias
org_object: Optional[LiteLLM_OrganizationTable] = None
if org_id:
@ -1444,7 +1449,7 @@ class JWTAuthManager:
# Check admin access
admin_result = await JWTAuthManager.check_admin_access(
jwt_handler, scopes, route, user_id, org_id, api_key
jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token
)
if admin_result:
return admin_result
@ -1452,7 +1457,9 @@ class JWTAuthManager:
# Get team with model access
## Check if team_id is specified via x-litellm-team-id header
all_team_ids = JWTAuthManager.get_all_team_ids(jwt_handler, jwt_valid_token)
specific_team_id = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None)
specific_team_id = jwt_handler.get_team_id(
token=jwt_valid_token, default_value=None
)
if specific_team_id:
all_team_ids.add(specific_team_id)
@ -1500,22 +1507,25 @@ class JWTAuthManager:
org_alias = jwt_handler.get_org_alias(token=jwt_valid_token, default_value=None)
# Get other objects
user_object, org_object, end_user_object, team_membership_object = (
await JWTAuthManager.get_objects(
user_id=user_id,
user_email=user_email,
org_id=org_id,
end_user_id=end_user_id,
team_id=team_id,
valid_user_email=valid_user_email,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
org_alias=org_alias,
)
(
user_object,
org_object,
end_user_object,
team_membership_object,
) = await JWTAuthManager.get_objects(
user_id=user_id,
user_email=user_email,
org_id=org_id,
end_user_id=end_user_id,
team_id=team_id,
valid_user_email=valid_user_email,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
org_alias=org_alias,
)
# Derive org_id from org_object if resolved by alias
@ -1559,4 +1569,5 @@ class JWTAuthManager:
end_user_object=end_user_object,
token=api_key,
team_membership=team_membership_object,
jwt_claims=jwt_valid_token,
)

View File

@ -22,6 +22,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.caching import DualCache
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
@ -35,6 +36,7 @@ from litellm.proxy.auth.auth_checks import (
can_key_call_model,
common_checks,
get_end_user_object,
get_jwt_key_mapping_object,
get_key_object,
get_project_object,
get_team_object,
@ -438,6 +440,75 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints(
return api_key
async def _resolve_jwt_to_virtual_key(
jwt_claims: dict,
jwt_handler: JWTHandler,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
) -> Optional[UserAPIKeyAuth]:
virtual_key_claim_field = jwt_handler.litellm_jwtauth.virtual_key_claim_field
if virtual_key_claim_field is None:
return None
claim_value = get_nested_value(
data=jwt_claims,
key_path=virtual_key_claim_field,
default=None,
)
if claim_value is None:
verbose_proxy_logger.debug(
f"JWT Key Mapping: Claim field '{virtual_key_claim_field}' not found in JWT claims."
)
return None
cache_key = f"jwt_key_mapping:{virtual_key_claim_field}:{claim_value}"
cached_mapping = await user_api_key_cache.async_get_cache(cache_key)
if cached_mapping == "__NO_MAPPING__":
return None
elif cached_mapping is not None:
return await get_key_object(
hashed_token=cached_mapping,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if prisma_client is None:
return None
token_hash = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=str(claim_value),
prisma_client=prisma_client,
)
if token_hash is not None:
await user_api_key_cache.async_set_cache(
key=cache_key,
value=token_hash,
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
)
return await get_key_object(
hashed_token=token_hash,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
else:
await user_api_key_cache.async_set_cache(
key=cache_key,
value="__NO_MAPPING__",
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
)
return None
async def _user_api_key_auth_builder( # noqa: PLR0915
request: Request,
api_key: str,
@ -589,145 +660,174 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
is_jwt = jwt_handler.is_jwt(token=api_key)
verbose_proxy_logger.debug("is_jwt: %s", is_jwt)
if is_jwt:
result = await JWTAuthManager.auth_builder(
request_data=request_data,
general_settings=general_settings,
api_key=api_key,
jwt_handler=jwt_handler,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
parent_otel_span=parent_otel_span,
request_headers=_safe_get_request_headers(request),
)
# Try JWT-to-Virtual-Key mapping first to avoid
# unnecessary DB queries in auth_builder
do_standard_jwt_auth = True
if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None:
# Decode JWT to get claims without running full auth_builder
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
jwt_claims = await jwt_handler.get_oidc_userinfo(token=api_key)
else:
jwt_claims = await jwt_handler.auth_jwt(token=api_key)
is_proxy_admin = result["is_proxy_admin"]
team_id = result["team_id"]
team_object = result["team_object"]
user_id = result["user_id"]
user_object = result["user_object"]
end_user_id = result["end_user_id"]
end_user_object = result["end_user_object"]
org_id = result["org_id"]
token = result["token"]
team_membership: Optional[LiteLLM_TeamMembership] = result.get(
"team_membership", None
)
valid_token = await _resolve_jwt_to_virtual_key(
jwt_claims=jwt_claims,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if valid_token is not None:
api_key = valid_token.token or ""
do_standard_jwt_auth = False
# Fall through to virtual key checks
global_proxy_spend = await get_global_proxy_spend(
litellm_proxy_admin_name=litellm_proxy_admin_name,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
token=token,
proxy_logging_obj=proxy_logging_obj,
)
if do_standard_jwt_auth:
result = await JWTAuthManager.auth_builder(
request_data=request_data,
general_settings=general_settings,
api_key=api_key,
jwt_handler=jwt_handler,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
parent_otel_span=parent_otel_span,
request_headers=_safe_get_request_headers(request),
)
if is_proxy_admin:
return UserAPIKeyAuth(
is_proxy_admin = result["is_proxy_admin"]
team_id = result["team_id"]
team_object = result["team_object"]
user_id = result["user_id"]
user_object = result["user_object"]
end_user_id = result["end_user_id"]
end_user_object = result["end_user_object"]
org_id = result["org_id"]
token = result["token"]
team_membership: Optional[LiteLLM_TeamMembership] = result.get(
"team_membership", None
)
global_proxy_spend = await get_global_proxy_spend(
litellm_proxy_admin_name=litellm_proxy_admin_name,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
token=token,
proxy_logging_obj=proxy_logging_obj,
)
if is_proxy_admin:
return UserAPIKeyAuth(
api_key=None,
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id=user_id,
team_id=team_id,
team_alias=(
team_object.team_alias
if team_object is not None
else None
),
team_metadata=team_object.metadata
if team_object is not None
else None,
org_id=org_id,
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
)
valid_token = UserAPIKeyAuth(
api_key=None,
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id=user_id,
team_id=team_id,
team_alias=(
team_object.team_alias if team_object is not None else None
),
team_tpm_limit=(
team_object.tpm_limit if team_object is not None else None
),
team_rpm_limit=(
team_object.rpm_limit if team_object is not None else None
),
team_models=team_object.models
if team_object is not None
else [],
user_role=(
LitellmUserRoles(user_object.user_role)
if user_object is not None
and user_object.user_role is not None
else LitellmUserRoles.INTERNAL_USER
),
user_id=user_id,
org_id=org_id,
parent_otel_span=parent_otel_span,
end_user_id=end_user_id,
user_tpm_limit=(
user_object.tpm_limit if user_object is not None else None
),
user_rpm_limit=(
user_object.rpm_limit if user_object is not None else None
),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit()
if team_membership is not None
else None
),
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit()
if team_membership is not None
else None
),
team_metadata=team_object.metadata
if team_object is not None
else None,
org_id=org_id,
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
)
valid_token = UserAPIKeyAuth(
api_key=None,
team_id=team_id,
team_alias=(
team_object.team_alias if team_object is not None else None
),
team_tpm_limit=(
team_object.tpm_limit if team_object is not None else None
),
team_rpm_limit=(
team_object.rpm_limit if team_object is not None else None
),
team_models=team_object.models if team_object is not None else [],
user_role=(
LitellmUserRoles(user_object.user_role)
if user_object is not None and user_object.user_role is not None
else LitellmUserRoles.INTERNAL_USER
),
user_id=user_id,
org_id=org_id,
parent_otel_span=parent_otel_span,
end_user_id=end_user_id,
user_tpm_limit=(
user_object.tpm_limit if user_object is not None else None
),
user_rpm_limit=(
user_object.rpm_limit if user_object is not None else None
),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit()
if team_membership is not None
else None
),
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit()
if team_membership is not None
else None
),
team_metadata=team_object.metadata
if team_object is not None
else None,
)
# Check if model has zero cost - if so, skip all budget checks
model = get_model_from_request(request_data, route)
skip_budget_checks = False
if model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
# Check if model has zero cost - if so, skip all budget checks
model = get_model_from_request(request_data, route)
skip_budget_checks = False
if model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
skip_budget_checks = _is_model_cost_zero(
model=model, llm_router=llm_router
)
if skip_budget_checks:
verbose_proxy_logger.info(
f"Skipping all budget checks for zero-cost model: {model}"
skip_budget_checks = _is_model_cost_zero(
model=model, llm_router=llm_router
)
if skip_budget_checks:
verbose_proxy_logger.info(
f"Skipping all budget checks for zero-cost model: {model}"
)
# Fetch project object for JWT path if project_id is set
_jwt_project_obj = None
if valid_token.project_id is not None:
_jwt_project_obj = await get_project_object(
project_id=valid_token.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
# Fetch project object for JWT path if project_id is set
_jwt_project_obj = None
if valid_token.project_id is not None:
_jwt_project_obj = await get_project_object(
project_id=valid_token.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if _jwt_project_obj is not None:
valid_token.project_metadata = _jwt_project_obj.metadata
# run through common checks
_ = await common_checks(
request=request,
request_body=request_data,
team_object=team_object,
user_object=user_object,
end_user_object=end_user_object,
general_settings=general_settings,
global_proxy_spend=global_proxy_spend,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=skip_budget_checks,
project_object=_jwt_project_obj,
)
if _jwt_project_obj is not None:
valid_token.project_metadata = _jwt_project_obj.metadata
# run through common checks
_ = await common_checks(
request=request,
request_body=request_data,
team_object=team_object,
user_object=user_object,
end_user_object=end_user_object,
general_settings=general_settings,
global_proxy_spend=global_proxy_spend,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=skip_budget_checks,
project_object=_jwt_project_obj,
)
# return UserAPIKeyAuth object
return cast(UserAPIKeyAuth, valid_token)
# return UserAPIKeyAuth object
return cast(UserAPIKeyAuth, valid_token)
#### ELSE ####
## CHECK PASS-THROUGH ENDPOINTS ##
@ -830,25 +930,26 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
# note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead
### CHECK IF ADMIN ###
# note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead
## Check CACHE
try:
valid_token = await get_key_object(
hashed_token=hash_token(api_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,
check_cache_only=True,
)
except Exception:
verbose_logger.debug("api key not found in cache.")
valid_token = None
if valid_token is None:
## Check CACHE
try:
valid_token = await get_key_object(
hashed_token=hash_token(api_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,
check_cache_only=True,
)
except Exception:
verbose_logger.debug("api key not found in cache.")
valid_token = None
## Check UI Hash Key
if valid_token is None and get_secret_bool("EXPERIMENTAL_UI_LOGIN"):
valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(
api_key
)
## Check UI Hash Key
if valid_token is None and get_secret_bool("EXPERIMENTAL_UI_LOGIN"):
valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(
api_key
)
if (
valid_token is not None
@ -986,9 +1087,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
param=None,
)
## check for cache hit (In-Memory Cache)
_user_role = None
if valid_token is None:
if isinstance(
api_key, str

View File

@ -0,0 +1,257 @@
from fastapi import APIRouter, Depends, HTTPException, Query
from litellm.proxy._types import (
CreateJWTKeyMappingRequest,
DeleteJWTKeyMappingRequest,
JWTKeyMappingResponse,
LitellmUserRoles,
UpdateJWTKeyMappingRequest,
UserAPIKeyAuth,
hash_token,
)
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(
id=mapping.id,
jwt_claim_name=mapping.jwt_claim_name,
jwt_claim_value=mapping.jwt_claim_value,
description=mapping.description,
is_active=mapping.is_active,
created_at=mapping.created_at,
updated_at=mapping.updated_at,
created_by=mapping.created_by,
updated_by=mapping.updated_by,
)
@router.post(
"/jwt/key/mapping/new",
tags=["JWT Key Mapping"],
response_model=JWTKeyMappingResponse,
)
async def create_jwt_key_mapping(
data: CreateJWTKeyMappingRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
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 key mappings"
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
try:
hashed_key = hash_token(data.key)
create_data = {
"jwt_claim_name": data.jwt_claim_name,
"jwt_claim_value": data.jwt_claim_value,
"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
)
# Invalidate cache
cache_key = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
return _to_response(new_mapping)
except HTTPException:
raise
except Exception as e:
error_str = str(e).lower()
if "unique" in error_str or "p2002" in error_str:
raise HTTPException(
status_code=409,
detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.",
)
if "foreign" in error_str or "p2003" in error_str:
raise HTTPException(
status_code=400,
detail="The provided key does not match an existing virtual key.",
)
raise HTTPException(status_code=500, detail="Failed to create JWT key mapping.")
@router.post(
"/jwt/key/mapping/update",
tags=["JWT Key Mapping"],
response_model=JWTKeyMappingResponse,
)
async def update_jwt_key_mapping(
data: UpdateJWTKeyMappingRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
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 key mappings"
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
update_data = data.model_dump(exclude_unset=True, exclude={"id", "key"})
if data.key is not None:
update_data["token"] = hash_token(data.key)
update_data["updated_by"] = user_api_key_dict.user_id
try:
# Get old mapping for cache invalidation
old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique(
where={"id": data.id}
)
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}"
await user_api_key_cache.async_delete_cache(cache_key)
updated_mapping = await prisma_client.db.litellm_jwtkeymapping.update(
where={"id": data.id}, data=update_data
)
# Invalidate new cache key if claim fields changed
cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
return _to_response(updated_mapping)
except HTTPException:
raise
except Exception as e:
error_str = str(e).lower()
if "unique" in error_str or "p2002" in error_str:
raise HTTPException(
status_code=409,
detail="A mapping with those claim values already exists.",
)
if "foreign" in error_str or "p2003" in error_str:
raise HTTPException(
status_code=400,
detail="The provided key does not match an existing virtual key.",
)
raise HTTPException(status_code=500, detail="Failed to update JWT key mapping.")
@router.post("/jwt/key/mapping/delete", tags=["JWT Key Mapping"])
async def delete_jwt_key_mapping(
data: DeleteJWTKeyMappingRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
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 delete JWT key mappings"
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
try:
# Get old mapping for cache invalidation
old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique(
where={"id": data.id}
)
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}"
await user_api_key_cache.async_delete_cache(cache_key)
await prisma_client.db.litellm_jwtkeymapping.delete(where={"id": data.id})
return {"status": "success"}
except HTTPException:
raise
except Exception:
raise HTTPException(status_code=500, detail="Failed to delete JWT key mapping.")
@router.get(
"/jwt/key/mapping/list",
tags=["JWT Key Mapping"],
)
async def list_jwt_key_mappings(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
page: int = Query(1, description="Page number", ge=1),
size: int = Query(50, description="Page size", ge=1, le=100),
):
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403, detail="Only proxy admins can list JWT key mappings"
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
try:
skip = (page - 1) * size
mappings = await prisma_client.db.litellm_jwtkeymapping.find_many(
skip=skip,
take=size,
order={"created_at": "desc"},
)
total_count = await prisma_client.db.litellm_jwtkeymapping.count()
return {
"mappings": [_to_response(m) for m in mappings],
"total_count": total_count,
"current_page": page,
"total_pages": -(-total_count // size), # ceiling division
}
except HTTPException:
raise
except Exception:
raise HTTPException(status_code=500, detail="Failed to list JWT key mappings.")
@router.get(
"/jwt/key/mapping/info",
tags=["JWT Key Mapping"],
response_model=JWTKeyMappingResponse,
)
async def info_jwt_key_mapping(
id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403, detail="Only proxy admins can get JWT key mapping info"
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
try:
mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique(
where={"id": id}
)
if mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
return _to_response(mapping)
except HTTPException:
raise
except Exception:
raise HTTPException(
status_code=500, detail="Failed to get JWT key mapping info."
)

View File

@ -377,6 +377,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
from litellm.proxy.management_endpoints.key_management_endpoints import (
router as key_management_router,
)
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
router as jwt_key_mapping_router,
)
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
router as mcp_management_router,
)
@ -13004,6 +13007,7 @@ app.include_router(debugging_endpoints_router)
app.include_router(ui_crud_endpoints_router)
app.include_router(openai_files_router)
app.include_router(team_callback_router)
app.include_router(jwt_key_mapping_router)
app.include_router(budget_management_router)
app.include_router(model_management_router)
app.include_router(model_access_group_management_router)

View File

@ -353,6 +353,7 @@ model LiteLLM_VerificationToken {
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
jwt_key_mappings LiteLLM_JWTKeyMapping[]
// SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub"
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2
@ -365,6 +366,24 @@ model LiteLLM_VerificationToken {
@@index([budget_reset_at, expires])
}
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
token String // Hashed virtual key (FK)
description String?
is_active Boolean @default(true)
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
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])
}
// Deprecated keys during grace period - allows old key to work until revoke_at
model LiteLLM_DeprecatedVerificationToken {
id String @id @default(uuid())

View File

@ -0,0 +1,352 @@
import pytest
import sys
import os
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
# Add project root to sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
from litellm.proxy.auth.user_api_key_auth import (
_resolve_jwt_to_virtual_key,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy._types import (
JWTKeyMappingResponse,
LiteLLM_JWTAuth,
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
_to_response,
create_jwt_key_mapping,
delete_jwt_key_mapping,
info_jwt_key_mapping,
update_jwt_key_mapping,
)
from litellm.caching.caching import DualCache
from fastapi import HTTPException
# ──────────────────────────────────────────────
# Tests: _resolve_jwt_to_virtual_key
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_jwt_to_virtual_key_mapping_resolution():
"""
Test that a JWT claim is correctly resolved to a virtual key token.
"""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="email", virtual_key_mapping_cache_ttl=3600
)
jwt_claims = {"email": "user@example.com", "sub": "123"}
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock()
# Mock finding a mapping
mock_mapping = MagicMock()
mock_mapping.token = "sk-1234"
mock_mapping.is_active = True
prisma_client.db.litellm_jwtkeymapping.find_first.return_value = mock_mapping
# Mock getting the key object
mock_key_obj = UserAPIKeyAuth(token="sk-1234", team_id="team1")
user_api_key_cache = DualCache()
# Use patch to mock get_key_object in the module where it's used
with patch(
"litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock
) as mock_get_key:
mock_get_key.return_value = mock_key_obj
result = await _resolve_jwt_to_virtual_key(
jwt_claims=jwt_claims,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result == mock_key_obj
prisma_client.db.litellm_jwtkeymapping.find_first.assert_called_once()
# Test Cache hit
prisma_client.db.litellm_jwtkeymapping.find_first.reset_mock()
result_cached = await _resolve_jwt_to_virtual_key(
jwt_claims=jwt_claims,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result_cached == mock_key_obj
prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
@pytest.mark.asyncio
async def test_jwt_to_virtual_key_mapping_no_mapping():
"""
Test that when no mapping exists, resolve returns None.
"""
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_claim_field="email")
jwt_claims = {"email": "unknown@example.com"}
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock()
prisma_client.db.litellm_jwtkeymapping.find_first.return_value = None
# Mock get_key_object just in case
with patch(
"litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock
):
user_api_key_cache = DualCache()
result = await _resolve_jwt_to_virtual_key(
jwt_claims=jwt_claims,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result is None
# Test Negative Cache hit
prisma_client.db.litellm_jwtkeymapping.find_first.reset_mock()
result_cached = await _resolve_jwt_to_virtual_key(
jwt_claims=jwt_claims,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result_cached is None
prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
# ──────────────────────────────────────────────
# Tests: _to_response redacts hashed token
# ──────────────────────────────────────────────
def test_to_response_excludes_token():
"""_to_response should not expose the hashed token field."""
now = datetime.now(timezone.utc)
mock_mapping = MagicMock()
mock_mapping.id = "mapping-1"
mock_mapping.jwt_claim_name = "email"
mock_mapping.jwt_claim_value = "user@example.com"
mock_mapping.token = "hashed_secret_value"
mock_mapping.description = "test"
mock_mapping.is_active = True
mock_mapping.created_at = now
mock_mapping.updated_at = now
mock_mapping.created_by = "admin"
mock_mapping.updated_by = "admin"
resp = _to_response(mock_mapping)
assert isinstance(resp, JWTKeyMappingResponse)
assert resp.id == "mapping-1"
assert resp.jwt_claim_name == "email"
assert "token" not in resp.model_fields
# ──────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────
def _make_admin_auth() -> UserAPIKeyAuth:
return UserAPIKeyAuth(
token="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
def _make_non_admin_auth() -> UserAPIKeyAuth:
return UserAPIKeyAuth(
token="sk-user",
user_role=LitellmUserRoles.INTERNAL_USER,
)
def _mock_prisma():
prisma = MagicMock()
prisma.db.litellm_jwtkeymapping.create = AsyncMock()
prisma.db.litellm_jwtkeymapping.find_unique = AsyncMock()
prisma.db.litellm_jwtkeymapping.find_many = AsyncMock()
prisma.db.litellm_jwtkeymapping.update = AsyncMock()
prisma.db.litellm_jwtkeymapping.delete = AsyncMock()
prisma.db.litellm_jwtkeymapping.count = AsyncMock(return_value=0)
return prisma
def _mock_mapping(
id="mapping-1",
claim_name="email",
claim_value="user@example.com",
):
now = datetime.now(timezone.utc)
m = MagicMock()
m.id = id
m.jwt_claim_name = claim_name
m.jwt_claim_value = claim_value
m.token = "hashed_token"
m.description = None
m.is_active = True
m.created_at = now
m.updated_at = now
m.created_by = "admin"
m.updated_by = "admin"
return m
# ──────────────────────────────────────────────
# Tests: CRUD endpoint error handling
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_create_returns_409_on_unique_violation():
"""Duplicate mapping should return 409, not 500."""
from litellm.proxy._types import CreateJWTKeyMappingRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception(
"Unique constraint failed (P2002)"
)
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
):
with pytest.raises(HTTPException) as exc_info:
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
assert exc_info.value.status_code == 409
assert "already exists" in exc_info.value.detail
@pytest.mark.asyncio
async def test_create_returns_400_on_foreign_key_violation():
"""Non-existent key should return 400, not 500."""
from litellm.proxy._types import CreateJWTKeyMappingRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception(
"Foreign key constraint failed on field: `token` (P2003)"
)
mock_cache = AsyncMock()
data = CreateJWTKeyMappingRequest(
jwt_claim_name="sub", jwt_claim_value="user-999", key="sk-nonexistent",
)
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 create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
assert exc_info.value.status_code == 400
assert "does not match" in exc_info.value.detail
@pytest.mark.asyncio
async def test_create_non_admin_returns_403():
"""Non-admin users should get 403."""
from litellm.proxy._types import CreateJWTKeyMappingRequest
data = CreateJWTKeyMappingRequest(
jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test",
)
with pytest.raises(HTTPException) as exc_info:
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_non_admin_auth())
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_delete_returns_404_when_not_found():
"""Deleting non-existent mapping should return 404."""
from litellm.proxy._types import DeleteJWTKeyMappingRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None
mock_cache = AsyncMock()
data = DeleteJWTKeyMappingRequest(id="nonexistent-id")
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 delete_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
assert exc_info.value.status_code == 404
@pytest.mark.asyncio
async def test_update_returns_404_when_not_found():
"""Updating non-existent mapping should return 404."""
from litellm.proxy._types import UpdateJWTKeyMappingRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None
mock_cache = AsyncMock()
data = UpdateJWTKeyMappingRequest(id="nonexistent-id", description="test")
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 update_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
assert exc_info.value.status_code == 404
@pytest.mark.asyncio
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
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
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
@pytest.mark.asyncio
async def test_create_success_returns_response_without_token():
"""Successful create should return JWTKeyMappingResponse without hashed token."""
from litellm.proxy._types import CreateJWTKeyMappingRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping()
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
):
result = await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
assert isinstance(result, JWTKeyMappingResponse)
assert "token" not in result.model_fields
assert result.jwt_claim_name == "email"