Merge pull request #22372 from BerriAI/litellm_jwt_vkey_map
Litellm jwt vkey map
This commit is contained in:
commit
07cb6d5bec
@ -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)
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
@ -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
|
||||
|
||||
257
litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py
Normal file
257
litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py
Normal 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."
|
||||
)
|
||||
@ -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)
|
||||
|
||||
@ -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())
|
||||
|
||||
352
tests/proxy_unit_tests/test_jwt_key_mapping.py
Normal file
352
tests/proxy_unit_tests/test_jwt_key_mapping.py
Normal 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"
|
||||
Loading…
Reference in New Issue
Block a user