diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dfc2ba59d9..6440bf0ed8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -539,6 +539,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 = [ @@ -3664,6 +3669,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 + token: str + key_alias: Optional[str] = None + description: Optional[str] = None + is_active: bool + created_at: datetime + updated_at: datetime + + class SpecialHeaders(enum.Enum): """Used by user_api_key_auth.py to get litellm key""" @@ -3834,6 +3869,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): @@ -3977,6 +4013,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: diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9921b74b56..210996a0a8 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -857,6 +857,7 @@ class JWTAuthManager: end_user_id=None, org_id=org_id, team_membership=None, + jwt_claims={}, ) @staticmethod @@ -1479,4 +1480,5 @@ class JWTAuthManager: end_user_object=end_user_object, token=api_key, team_membership=team_membership_object, + jwt_claims=jwt_valid_token, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 8ad3b83c04..d453b72164 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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, @@ -438,6 +439,78 @@ 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 + + mapping = await prisma_client.db.litellm_jwtkeymapping.find_first( + where={ + "jwt_claim_name": virtual_key_claim_field, + "jwt_claim_value": str(claim_value), + "is_active": True, + } + ) + + if mapping: + token_hash = mapping.token + 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, @@ -602,132 +675,151 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_headers=_safe_get_request_headers(request), ) - 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 - ) + # JWT-to-Virtual-Key Mapping lookup + do_standard_jwt_auth = True + if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None: + valid_token = await _resolve_jwt_to_virtual_key( + jwt_claims=result["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: + 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 + ) - if is_proxy_admin: - return UserAPIKeyAuth( + 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 +922,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 +1079,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 diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py new file mode 100644 index 0000000000..056db954a5 --- /dev/null +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -0,0 +1,152 @@ +import asyncio +from typing import List, Optional, Union +from fastapi import APIRouter, Depends, HTTPException, Request +import litellm +from litellm.proxy._types import * +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.proxy.auth.auth_checks import _delete_cache_key_object + +router = APIRouter() + +@router.post("/jwt/key/mapping/new", tags=["JWT Key Mapping"]) +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: + new_mapping = await prisma_client.db.litellm_jwtkeymapping.create( + data={ + "jwt_claim_name": data.jwt_claim_name, + "jwt_claim_value": data.jwt_claim_value, + "token": data.token, + "is_active": data.is_active, + } + ) + + # 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 new_mapping + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + +@router.post("/jwt/key/mapping/update", tags=["JWT Key Mapping"]) +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={"mapping_id"}) + + try: + # Get old mapping for cache invalidation + old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + where={"mapping_id": data.mapping_id} + ) + + if old_mapping: + 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={"mapping_id": data.mapping_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 updated_mapping + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + +@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={"mapping_id": data.mapping_id} + ) + + if old_mapping: + 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={"mapping_id": data.mapping_id} + ) + return {"status": "success"} + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + +@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), +): + 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: + mappings = await prisma_client.db.litellm_jwtkeymapping.find_many() + return mappings + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + +@router.get("/jwt/key/mapping/info", tags=["JWT Key Mapping"]) +async def info_jwt_key_mapping( + 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={"mapping_id": mapping_id} + ) + if mapping is None: + raise HTTPException(status_code=404, detail="Mapping not found") + return mapping + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index be76c2ac5f..4c613a4dcb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -376,6 +376,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, ) @@ -12929,6 +12932,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) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a5b0d930f5..5b5ca8abf8 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -351,6 +351,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 @@ -363,6 +364,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()) diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py new file mode 100644 index 0000000000..e44365897a --- /dev/null +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -0,0 +1,112 @@ +import pytest +import sys +import os +from unittest.mock import AsyncMock, MagicMock, patch +from fastapi import Request +from starlette.datastructures import URL +import litellm + +# 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 user_api_key_auth, _resolve_jwt_to_virtual_key +from litellm.proxy.auth.handle_jwt import JWTHandler, JWTAuthManager +from litellm.proxy._types import LiteLLM_JWTAuth, UserAPIKeyAuth +from litellm.caching.caching import DualCache + +@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) as mock_get_key: + 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()