From 9dc085694c7dcfaedc13df3a873483df226f61dc Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 28 Feb 2026 13:29:46 +0530 Subject: [PATCH 01/12] feat: jwt mapping vkeyv --- litellm/proxy/_types.py | 45 +++ litellm/proxy/auth/handle_jwt.py | 2 + litellm/proxy/auth/user_api_key_auth.py | 348 +++++++++++------- .../jwt_key_mapping_endpoints.py | 152 ++++++++ litellm/proxy/proxy_server.py | 4 + litellm/proxy/schema.prisma | 19 + .../proxy_unit_tests/test_jwt_key_mapping.py | 112 ++++++ 7 files changed, 553 insertions(+), 129 deletions(-) create mode 100644 litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py create mode 100644 tests/proxy_unit_tests/test_jwt_key_mapping.py 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() From 465adce8721249af4ccab094daf307c71ffc955d Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 28 Feb 2026 13:40:53 +0530 Subject: [PATCH 02/12] feat reaq changes --- litellm/proxy/auth/handle_jwt.py | 78 ++++++++++--------- litellm/proxy/auth/user_api_key_auth.py | 7 +- .../jwt_key_mapping_endpoints.py | 45 ++++++----- litellm/proxy/proxy_server.py | 24 +++--- .../proxy_unit_tests/test_jwt_key_mapping.py | 58 +++++++------- 5 files changed, 119 insertions(+), 93 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 210996a0a8..d3b028f63f 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -165,7 +165,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, @@ -245,7 +244,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. @@ -538,17 +539,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 """ @@ -556,19 +557,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( @@ -578,24 +581,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)}") @@ -1032,11 +1035,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: @@ -1373,7 +1376,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) @@ -1421,22 +1426,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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d453b72164..5c529bc69d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -744,10 +744,13 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 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 [], + 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 + if user_object is not None + and user_object.user_role is not None else LitellmUserRoles.INTERNAL_USER ), user_id=user_id, diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 056db954a5..06a526e13a 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,14 +1,10 @@ -import asyncio -from typing import List, Optional, Union -from fastapi import APIRouter, Depends, HTTPException, Request -import litellm +from fastapi import APIRouter, Depends, HTTPException 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, @@ -17,7 +13,9 @@ async def create_jwt_key_mapping( 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") + 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") @@ -31,7 +29,7 @@ async def create_jwt_key_mapping( "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) @@ -40,6 +38,7 @@ async def create_jwt_key_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, @@ -48,28 +47,29 @@ async def update_jwt_key_mapping( 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") + 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 + 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) @@ -78,6 +78,7 @@ async def update_jwt_key_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, @@ -86,7 +87,9 @@ async def delete_jwt_key_mapping( 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") + 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") @@ -96,7 +99,7 @@ async def delete_jwt_key_mapping( 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) @@ -108,6 +111,7 @@ async def delete_jwt_key_mapping( 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), @@ -115,7 +119,9 @@ async def list_jwt_key_mappings( 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") + 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") @@ -126,6 +132,7 @@ async def list_jwt_key_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, @@ -134,7 +141,9 @@ async def info_jwt_key_mapping( 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") + 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") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4c613a4dcb..d80cb6be57 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2165,22 +2165,24 @@ async def _run_background_health_check(): "Error in shared health check, falling back to direct health check: %s", str(e), ) - healthy_endpoints, unhealthy_endpoints = ( - await _run_direct_health_check_with_instrumentation( - _llm_model_list, - health_check_details, - health_check_concurrency, - instrumentation_context, - ) - ) - else: - healthy_endpoints, unhealthy_endpoints = ( - await _run_direct_health_check_with_instrumentation( + ( + healthy_endpoints, + unhealthy_endpoints, + ) = await _run_direct_health_check_with_instrumentation( _llm_model_list, health_check_details, health_check_concurrency, instrumentation_context, ) + else: + ( + healthy_endpoints, + unhealthy_endpoints, + ) = await _run_direct_health_check_with_instrumentation( + _llm_model_list, + health_check_details, + health_check_concurrency, + instrumentation_context, ) # Update the global variable with the health check results diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index e44365897a..7d3e7371b1 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -2,18 +2,18 @@ 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.auth.user_api_key_auth import ( + _resolve_jwt_to_virtual_key, +) +from litellm.proxy.auth.handle_jwt import JWTHandler 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(): """ @@ -21,42 +21,43 @@ async def test_jwt_to_virtual_key_mapping_resolution(): """ jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - virtual_key_claim_field="email", - virtual_key_mapping_cache_ttl=3600 + 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: + 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 + 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( @@ -65,11 +66,12 @@ async def test_jwt_to_virtual_key_mapping_resolution(): prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=None, - proxy_logging_obj=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(): """ @@ -78,26 +80,28 @@ async def test_jwt_to_virtual_key_mapping_no_mapping(): 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: + 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 + 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( @@ -106,7 +110,7 @@ async def test_jwt_to_virtual_key_mapping_no_mapping(): prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=None, - proxy_logging_obj=None + proxy_logging_obj=None, ) assert result_cached is None prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called() From 941129c9e0bcf2206a439b9d1994d23496e15710 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sun, 1 Mar 2026 05:46:02 +0530 Subject: [PATCH 03/12] fix: resolve field mismatches and direct DB query in jwt key mapping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix data.token → hash_token(data.key) and remove non-existent data.is_active in create endpoint - Fix mapping_id → id in update, delete, and info endpoints to match Prisma schema - Extract direct DB query into get_jwt_key_mapping_object helper in auth_checks.py - Add hash_token import for proper key hashing before storage Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/auth/auth_checks.py | 22 +++++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 14 +++++------- .../jwt_key_mapping_endpoints.py | 19 ++++++++-------- 3 files changed, 38 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 500a39d945..a7867fa08c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2028,6 +2028,28 @@ async def _fetch_key_object_from_db_with_reconnect( raise +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, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5c529bc69d..7098b5e18d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -36,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, @@ -480,16 +481,13 @@ async def _resolve_jwt_to_virtual_key( 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, - } + 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 mapping: - token_hash = mapping.token + if token_hash is not None: await user_api_key_cache.async_set_cache( key=cache_key, value=token_hash, diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 06a526e13a..08ca00e66d 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,5 +1,6 @@ from fastapi import APIRouter, Depends, HTTPException from litellm.proxy._types import * +from litellm.proxy._types import hash_token from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() @@ -21,12 +22,12 @@ async def create_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: + hashed_key = hash_token(data.key) 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, + "token": hashed_key, } ) @@ -54,12 +55,12 @@ async def update_jwt_key_mapping( 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"}) + update_data = data.model_dump(exclude_unset=True, exclude={"id"}) try: # Get old mapping for cache invalidation old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( - where={"mapping_id": data.mapping_id} + where={"id": data.id} ) if old_mapping: @@ -67,7 +68,7 @@ async def update_jwt_key_mapping( 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 + where={"id": data.id}, data=update_data ) # Invalidate new cache key if claim fields changed @@ -97,7 +98,7 @@ async def delete_jwt_key_mapping( try: # Get old mapping for cache invalidation old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( - where={"mapping_id": data.mapping_id} + where={"id": data.id} ) if old_mapping: @@ -105,7 +106,7 @@ async def delete_jwt_key_mapping( await user_api_key_cache.async_delete_cache(cache_key) await prisma_client.db.litellm_jwtkeymapping.delete( - where={"mapping_id": data.mapping_id} + where={"id": data.id} ) return {"status": "success"} except Exception as e: @@ -135,7 +136,7 @@ async def list_jwt_key_mappings( @router.get("/jwt/key/mapping/info", tags=["JWT Key Mapping"]) async def info_jwt_key_mapping( - mapping_id: str, + id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): from litellm.proxy.proxy_server import prisma_client @@ -150,7 +151,7 @@ async def info_jwt_key_mapping( try: mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( - where={"mapping_id": mapping_id} + where={"id": id} ) if mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") From 963390928d3be222170178848b20703d8f861391 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 1 Mar 2026 13:08:40 +0530 Subject: [PATCH 04/12] Update litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/jwt_key_mapping_endpoints.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 08ca00e66d..6c19f8c5dd 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -55,7 +55,9 @@ async def update_jwt_key_mapping( if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") - update_data = data.model_dump(exclude_unset=True, exclude={"id"}) + update_data = data.model_dump(exclude_unset=True, exclude={"id", "key"}) + if data.key is not None: + update_data["token"] = hash_token(data.key) try: # Get old mapping for cache invalidation From 0f9d3808748b74a35539c4ad1bf6dcfb1bc395cf Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sun, 1 Mar 2026 13:28:06 +0530 Subject: [PATCH 05/12] fix: add pagination to jwt key mapping list endpoint Add page/size query params with take/skip to prevent unbounded queries. Returns paginated response with total_count, current_page, total_pages. Co-Authored-By: Claude Opus 4.6 --- .../jwt_key_mapping_endpoints.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 6c19f8c5dd..770b9a48dd 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,4 +1,4 @@ -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Query from litellm.proxy._types import * from litellm.proxy._types import hash_token from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -118,6 +118,8 @@ async def 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 @@ -130,8 +132,19 @@ async def list_jwt_key_mappings( raise HTTPException(status_code=500, detail="Database not connected") try: - mappings = await prisma_client.db.litellm_jwtkeymapping.find_many() - return mappings + 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": mappings, + "total_count": total_count, + "current_page": page, + "total_pages": -(-total_count // size), # ceiling division + } except Exception as e: raise HTTPException(status_code=500, detail=str(e)) From 0e2dd4aac1a71623a72af4b2199731c27f5eaccc Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 1 Mar 2026 13:43:06 +0530 Subject: [PATCH 06/12] Update litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/jwt_key_mapping_endpoints.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 770b9a48dd..9680a92ae6 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -22,12 +22,13 @@ async def create_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - hashed_key = hash_token(data.key) + 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": hashed_key, + "token": data.key, + "is_active": True, } ) From 911ba14e45509a250b19ca7c480dad188ac1ccca Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sun, 1 Mar 2026 14:28:50 +0530 Subject: [PATCH 07/12] fix: address remaining greptile feedback for jwt key mapping - Persist description field on create (was silently dropped) - Remove phantom key_alias from JWTKeyMappingResponse (not in schema) - Populate created_by/updated_by audit fields from authenticated user - Pass actual jwt_valid_token in admin path instead of empty dict - Restore hash_token on create and fix duplicate try block Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_types.py | 3 ++- litellm/proxy/auth/handle_jwt.py | 5 +++-- .../jwt_key_mapping_endpoints.py | 20 ++++++++++++------- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6440bf0ed8..6b51df709f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3692,11 +3692,12 @@ class JWTKeyMappingResponse(LiteLLMPydanticObjectBase): 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 + created_by: Optional[str] = None + updated_by: Optional[str] = None class SpecialHeaders(enum.Enum): diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index d3b028f63f..6ca7b290a0 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -831,6 +831,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): @@ -860,7 +861,7 @@ class JWTAuthManager: end_user_id=None, org_id=org_id, team_membership=None, - jwt_claims={}, + jwt_claims=jwt_valid_token or {}, ) @staticmethod @@ -1368,7 +1369,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 diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 9680a92ae6..c5d91d3699 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -22,14 +22,19 @@ async def create_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - 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={ - "jwt_claim_name": data.jwt_claim_name, - "jwt_claim_value": data.jwt_claim_value, - "token": data.key, - "is_active": True, - } + data=create_data ) # Invalidate cache @@ -59,6 +64,7 @@ async def update_jwt_key_mapping( 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 From 28a48acce645deaee4b53b485f1ad9eb022cfc70 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sun, 1 Mar 2026 16:46:49 +0530 Subject: [PATCH 08/12] fix: add @log_db_metrics and move jwt mapping before auth_builder - Add @log_db_metrics decorator to get_jwt_key_mapping_object for consistent DB latency/error tracking with other helpers - Move virtual key mapping lookup before auth_builder() to avoid unnecessary team/user/org DB queries when mapping resolves - JWT is decoded early; auth_builder only runs when no mapping found Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/auth/auth_checks.py | 1 + litellm/proxy/auth/user_api_key_auth.py | 37 +++++++++++++++---------- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a7867fa08c..e3776f2bfb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2028,6 +2028,7 @@ 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, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 7098b5e18d..d6dad5167c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -660,24 +660,18 @@ 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), - ) - - # JWT-to-Virtual-Key Mapping lookup + # 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) + valid_token = await _resolve_jwt_to_virtual_key( - jwt_claims=result["jwt_claims"], + jwt_claims=jwt_claims, jwt_handler=jwt_handler, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, @@ -690,6 +684,19 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Fall through to virtual key checks 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), + ) + is_proxy_admin = result["is_proxy_admin"] team_id = result["team_id"] team_object = result["team_object"] From 2f15686ea2cca8faeb2f52ad64b82a1c83314dde Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Thu, 5 Mar 2026 03:46:03 +0530 Subject: [PATCH 09/12] fix: address greptile feedback - redact hashed tokens, proper error codes, add tests - Remove token field from JWTKeyMappingResponse to prevent hashed key exposure - Use _to_response() helper on all CRUD endpoints to control returned fields - Return 409 for unique constraint violations, 400 for FK violations, 404 for not found - Add response_model to endpoint decorators - Add 8 new unit tests covering error handling and token redaction Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_types.py | 1 - .../jwt_key_mapping_endpoints.py | 123 +++++++-- .../proxy_unit_tests/test_jwt_key_mapping.py | 238 +++++++++++++++++- 3 files changed, 333 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6b51df709f..e7aba56f00 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3691,7 +3691,6 @@ class JWTKeyMappingResponse(LiteLLMPydanticObjectBase): id: str jwt_claim_name: str jwt_claim_value: str - token: str description: Optional[str] = None is_active: bool created_at: datetime diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index c5d91d3699..779700caf5 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,12 +1,41 @@ +from typing import List + from fastapi import APIRouter, Depends, HTTPException, Query -from litellm.proxy._types import * -from litellm.proxy._types import hash_token + +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() -@router.post("/jwt/key/mapping/new", tags=["JWT Key Mapping"]) +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), @@ -41,12 +70,29 @@ async def create_jwt_key_mapping( 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 + return _to_response(new_mapping) + except HTTPException: + raise except Exception as e: - raise HTTPException(status_code=500, detail=str(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"]) +@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), @@ -72,9 +118,11 @@ async def update_jwt_key_mapping( where={"id": data.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) + 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 @@ -84,9 +132,17 @@ async def update_jwt_key_mapping( 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 + return _to_response(updated_mapping) + except HTTPException: + raise except Exception as e: - raise HTTPException(status_code=500, detail=str(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.", + ) + raise HTTPException(status_code=500, detail="Failed to update JWT key mapping.") @router.post("/jwt/key/mapping/delete", tags=["JWT Key Mapping"]) @@ -110,19 +166,24 @@ async def delete_jwt_key_mapping( where={"id": data.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) + if old_mapping is None: + raise HTTPException(status_code=404, detail="Mapping not found") - await prisma_client.db.litellm_jwtkeymapping.delete( - where={"id": data.id} - ) + 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 Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + 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"]) +@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), @@ -147,16 +208,22 @@ async def list_jwt_key_mappings( ) total_count = await prisma_client.db.litellm_jwtkeymapping.count() return { - "mappings": mappings, + "mappings": [_to_response(m) for m in mappings], "total_count": total_count, "current_page": page, "total_pages": -(-total_count // size), # ceiling division } - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + 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"]) +@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), @@ -177,8 +244,10 @@ async def info_jwt_key_mapping( ) if mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") - return mapping + return _to_response(mapping) except HTTPException: raise - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + except Exception: + raise HTTPException( + status_code=500, detail="Failed to get JWT key mapping info." + ) diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 7d3e7371b1..b67dd2792f 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -1,6 +1,7 @@ import pytest import sys import os +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch # Add project root to sys.path @@ -10,8 +11,26 @@ 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 LiteLLM_JWTAuth, UserAPIKeyAuth +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 @@ -114,3 +133,220 @@ async def test_jwt_to_virtual_key_mapping_no_mapping(): ) 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" From 63459d6777245266b835e801bcde5078b29e9bef Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Thu, 5 Mar 2026 03:59:59 +0530 Subject: [PATCH 10/12] docs: add JWT-to-Virtual-Key mapping documentation Co-Authored-By: Claude Opus 4.6 --- docs/my-website/docs/proxy/token_auth.md | 89 ++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/docs/my-website/docs/proxy/token_auth.md b/docs/my-website/docs/proxy/token_auth.md index e8634f0faf..7364ae0fb5 100644 --- a/docs/my-website/docs/proxy/token_auth.md +++ b/docs/my-website/docs/proxy/token_auth.md @@ -1054,6 +1054,95 @@ curl -X GET 'http://0.0.0.0:4000/user/info?user_id=user-123' \ -H 'Authorization: Bearer ' ``` +## [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 `). + +**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=" \ + -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": "", + "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": ""}' +``` + +### 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) From 36e63bd1eebbacbf8a5034536167a8174b7af09a Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Thu, 5 Mar 2026 04:12:54 +0530 Subject: [PATCH 11/12] Update litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/jwt_key_mapping_endpoints.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 779700caf5..ab3ad82e3c 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -142,6 +142,11 @@ async def update_jwt_key_mapping( 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.") From 063a1a437a8afd9c392c386e414accf42500a628 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Thu, 5 Mar 2026 04:43:37 +0530 Subject: [PATCH 12/12] Update litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index ab3ad82e3c..a2a38cad14 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,4 +1,3 @@ -from typing import List from fastapi import APIRouter, Depends, HTTPException, Query