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) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4f72e30ac6..9acc80b0dd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -543,6 +543,11 @@ class LiteLLMRoutes(enum.Enum): "/model/update", "/model/delete", "/model/info", + "/jwt/key/mapping/new", + "/jwt/key/mapping/update", + "/jwt/key/mapping/delete", + "/jwt/key/mapping/list", + "/jwt/key/mapping/info", ] + key_management_routes spend_tracking_routes = [ @@ -3685,6 +3690,36 @@ class KeyHealthResponse(TypedDict, total=False): logging_callbacks: Optional[LoggingCallbackStatus] +class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): + jwt_claim_name: str + jwt_claim_value: str + key: str + description: Optional[str] = None + + +class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): + id: str + key: Optional[str] = None + description: Optional[str] = None + is_active: Optional[bool] = None + + +class DeleteJWTKeyMappingRequest(LiteLLMPydanticObjectBase): + id: str + + +class JWTKeyMappingResponse(LiteLLMPydanticObjectBase): + id: str + jwt_claim_name: str + jwt_claim_value: str + description: Optional[str] = None + is_active: bool + created_at: datetime + updated_at: datetime + created_by: Optional[str] = None + updated_by: Optional[str] = None + + class SpecialHeaders(enum.Enum): """Used by user_api_key_auth.py to get litellm key""" @@ -3857,6 +3892,7 @@ class JWTAuthBuilderResult(TypedDict): end_user_id: Optional[str] org_id: Optional[str] team_membership: Optional[LiteLLM_TeamMembership] + jwt_claims: dict # Decoded JWT token claims (avoids re-decoding) class ClientSideFallbackModel(TypedDict, total=False): @@ -4000,6 +4036,15 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=300, description="TTL (in seconds) for caching UserInfo responses. Default: 300s (5 minutes).", ) + # JWT-to-Virtual-Key Mapping + virtual_key_claim_field: Optional[str] = Field( + default=None, + description="JWT claim field for virtual key mapping lookup (e.g. 'sub', 'email'). Supports dot notation.", + ) + virtual_key_mapping_cache_ttl: float = Field( + default=300, + description="TTL (seconds) for caching JWT-to-virtual-key mapping lookups.", + ) ######################################################### def __init__(self, **kwargs: Any) -> None: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 41b0a0bc38..af51b9fe44 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2081,6 +2081,29 @@ async def _fetch_key_object_from_db_with_reconnect( raise +@log_db_metrics +async def get_jwt_key_mapping_object( + jwt_claim_name: str, + jwt_claim_value: str, + prisma_client: PrismaClient, +) -> Optional[str]: + """ + Lookup a JWT-to-virtual-key mapping from the database. + + Returns the hashed token (str) if a matching active mapping is found, else None. + """ + mapping = await prisma_client.db.litellm_jwtkeymapping.find_first( + where={ + "jwt_claim_name": jwt_claim_name, + "jwt_claim_value": jwt_claim_value, + "is_active": True, + } + ) + if mapping is not None: + return mapping.token + return None + + @log_db_metrics async def get_key_object( hashed_token: str, diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 553ba4d6c4..bfad9f0c3c 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -166,7 +166,6 @@ class JWTHandler: return False def get_team_ids_from_jwt(self, token: dict) -> List[str]: - if self.litellm_jwtauth.team_ids_jwt_field is not None: team_ids: Optional[List[str]] = get_nested_value( data=token, @@ -256,7 +255,9 @@ class JWTHandler: team_id = default_value return team_id - def get_team_alias(self, token: dict, default_value: Optional[str]) -> Optional[str]: + def get_team_alias( + self, token: dict, default_value: Optional[str] + ) -> Optional[str]: """ Extract team name/alias from JWT token using the configured team_alias_jwt_field. @@ -596,17 +597,17 @@ class JWTHandler: async def get_oidc_userinfo(self, token: str) -> dict: """ Fetch user information from OIDC UserInfo endpoint. - + This follows the OpenID Connect protocol where an access token is sent to the identity provider's UserInfo endpoint to retrieve user identity information. - + Args: token: The access token to use for authentication - + Returns: dict: User information from the UserInfo endpoint - + Raises: Exception: If UserInfo endpoint is not configured or request fails """ @@ -614,19 +615,21 @@ class JWTHandler: raise Exception( "OIDC UserInfo endpoint not configured. Set 'oidc_userinfo_endpoint' in JWT auth config." ) - + # Check cache first - cache_key = f"oidc_userinfo_{token[:20]}" # Use first 20 chars of token as cache key + cache_key = ( + f"oidc_userinfo_{token[:20]}" # Use first 20 chars of token as cache key + ) cached_userinfo = await self.user_api_key_cache.async_get_cache(cache_key) - + if cached_userinfo is not None: verbose_proxy_logger.debug("Returning cached OIDC UserInfo") return cached_userinfo - + verbose_proxy_logger.debug( f"Calling OIDC UserInfo endpoint: {self.litellm_jwtauth.oidc_userinfo_endpoint}" ) - + try: # Call the UserInfo endpoint with the access token response = await self.http_handler.get( @@ -636,24 +639,24 @@ class JWTHandler: "Accept": "application/json", }, ) - + if response.status_code != 200: raise Exception( f"OIDC UserInfo endpoint returned status {response.status_code}: {response.text}" ) - + userinfo = response.json() verbose_proxy_logger.debug(f"Received OIDC UserInfo: {userinfo}") - + # Cache the userinfo response await self.user_api_key_cache.async_set_cache( key=cache_key, value=userinfo, ttl=self.litellm_jwtauth.oidc_userinfo_cache_ttl, ) - + return userinfo - + except Exception as e: verbose_proxy_logger.error(f"Error fetching OIDC UserInfo: {str(e)}") raise Exception(f"Failed to fetch OIDC UserInfo: {str(e)}") @@ -886,6 +889,7 @@ class JWTAuthManager: user_id: Optional[str], org_id: Optional[str], api_key: str, + jwt_valid_token: Optional[dict] = None, ) -> Optional[JWTAuthBuilderResult]: """Check admin status and route access permissions""" if not jwt_handler.is_admin(scopes=scopes): @@ -915,6 +919,7 @@ class JWTAuthManager: end_user_id=None, org_id=org_id, team_membership=None, + jwt_claims=jwt_valid_token or {}, ) @staticmethod @@ -1111,11 +1116,11 @@ class JWTAuthManager: ) -> Tuple[ Optional[LiteLLM_UserTable], Optional[LiteLLM_OrganizationTable], - Optional[LiteLLM_EndUserTable], + Optional[LiteLLM_EndUserTable], Optional[LiteLLM_TeamMembership], ]: """Get user, org, and end user objects. Also resolves org aliases to IDs if configured.""" - + # Get org object - first try by ID, then by alias org_object: Optional[LiteLLM_OrganizationTable] = None if org_id: @@ -1444,7 +1449,7 @@ class JWTAuthManager: # Check admin access admin_result = await JWTAuthManager.check_admin_access( - jwt_handler, scopes, route, user_id, org_id, api_key + jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token ) if admin_result: return admin_result @@ -1452,7 +1457,9 @@ class JWTAuthManager: # Get team with model access ## Check if team_id is specified via x-litellm-team-id header all_team_ids = JWTAuthManager.get_all_team_ids(jwt_handler, jwt_valid_token) - specific_team_id = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None) + specific_team_id = jwt_handler.get_team_id( + token=jwt_valid_token, default_value=None + ) if specific_team_id: all_team_ids.add(specific_team_id) @@ -1500,22 +1507,25 @@ class JWTAuthManager: org_alias = jwt_handler.get_org_alias(token=jwt_valid_token, default_value=None) # Get other objects - user_object, org_object, end_user_object, team_membership_object = ( - await JWTAuthManager.get_objects( - user_id=user_id, - user_email=user_email, - org_id=org_id, - end_user_id=end_user_id, - team_id=team_id, - valid_user_email=valid_user_email, - jwt_handler=jwt_handler, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - org_alias=org_alias, - ) + ( + user_object, + org_object, + end_user_object, + team_membership_object, + ) = await JWTAuthManager.get_objects( + user_id=user_id, + user_email=user_email, + org_id=org_id, + end_user_id=end_user_id, + team_id=team_id, + valid_user_email=valid_user_email, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + org_alias=org_alias, ) # Derive org_id from org_object if resolved by alias @@ -1559,4 +1569,5 @@ class JWTAuthManager: end_user_object=end_user_object, token=api_key, team_membership=team_membership_object, + jwt_claims=jwt_valid_token, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 7b52f6bb96..ab162bcfdc 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, @@ -35,6 +36,7 @@ from litellm.proxy.auth.auth_checks import ( can_key_call_model, common_checks, get_end_user_object, + get_jwt_key_mapping_object, get_key_object, get_project_object, get_team_object, @@ -438,6 +440,75 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( return api_key +async def _resolve_jwt_to_virtual_key( + jwt_claims: dict, + jwt_handler: JWTHandler, + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + parent_otel_span: Optional[Span], + proxy_logging_obj: ProxyLogging, +) -> Optional[UserAPIKeyAuth]: + virtual_key_claim_field = jwt_handler.litellm_jwtauth.virtual_key_claim_field + if virtual_key_claim_field is None: + return None + + claim_value = get_nested_value( + data=jwt_claims, + key_path=virtual_key_claim_field, + default=None, + ) + + if claim_value is None: + verbose_proxy_logger.debug( + f"JWT Key Mapping: Claim field '{virtual_key_claim_field}' not found in JWT claims." + ) + return None + + cache_key = f"jwt_key_mapping:{virtual_key_claim_field}:{claim_value}" + cached_mapping = await user_api_key_cache.async_get_cache(cache_key) + + if cached_mapping == "__NO_MAPPING__": + return None + elif cached_mapping is not None: + return await get_key_object( + hashed_token=cached_mapping, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + + if prisma_client is None: + return None + + token_hash = await get_jwt_key_mapping_object( + jwt_claim_name=virtual_key_claim_field, + jwt_claim_value=str(claim_value), + prisma_client=prisma_client, + ) + + if token_hash is not None: + await user_api_key_cache.async_set_cache( + key=cache_key, + value=token_hash, + ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, + ) + return await get_key_object( + hashed_token=token_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + else: + await user_api_key_cache.async_set_cache( + key=cache_key, + value="__NO_MAPPING__", + ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, + ) + return None + + async def _user_api_key_auth_builder( # noqa: PLR0915 request: Request, api_key: str, @@ -589,145 +660,174 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 is_jwt = jwt_handler.is_jwt(token=api_key) verbose_proxy_logger.debug("is_jwt: %s", is_jwt) if is_jwt: - result = await JWTAuthManager.auth_builder( - request_data=request_data, - general_settings=general_settings, - api_key=api_key, - jwt_handler=jwt_handler, - route=route, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - parent_otel_span=parent_otel_span, - request_headers=_safe_get_request_headers(request), - ) + # Try JWT-to-Virtual-Key mapping first to avoid + # unnecessary DB queries in auth_builder + do_standard_jwt_auth = True + if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None: + # Decode JWT to get claims without running full auth_builder + if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled: + jwt_claims = await jwt_handler.get_oidc_userinfo(token=api_key) + else: + jwt_claims = await jwt_handler.auth_jwt(token=api_key) - is_proxy_admin = result["is_proxy_admin"] - team_id = result["team_id"] - team_object = result["team_object"] - user_id = result["user_id"] - user_object = result["user_object"] - end_user_id = result["end_user_id"] - end_user_object = result["end_user_object"] - org_id = result["org_id"] - token = result["token"] - team_membership: Optional[LiteLLM_TeamMembership] = result.get( - "team_membership", None - ) + valid_token = await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if valid_token is not None: + api_key = valid_token.token or "" + do_standard_jwt_auth = False + # Fall through to virtual key checks - global_proxy_spend = await get_global_proxy_spend( - litellm_proxy_admin_name=litellm_proxy_admin_name, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - token=token, - proxy_logging_obj=proxy_logging_obj, - ) + if do_standard_jwt_auth: + result = await JWTAuthManager.auth_builder( + request_data=request_data, + general_settings=general_settings, + api_key=api_key, + jwt_handler=jwt_handler, + route=route, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + parent_otel_span=parent_otel_span, + request_headers=_safe_get_request_headers(request), + ) - if is_proxy_admin: - return UserAPIKeyAuth( + is_proxy_admin = result["is_proxy_admin"] + team_id = result["team_id"] + team_object = result["team_object"] + user_id = result["user_id"] + user_object = result["user_object"] + end_user_id = result["end_user_id"] + end_user_object = result["end_user_object"] + org_id = result["org_id"] + token = result["token"] + team_membership: Optional[LiteLLM_TeamMembership] = result.get( + "team_membership", None + ) + + global_proxy_spend = await get_global_proxy_spend( + litellm_proxy_admin_name=litellm_proxy_admin_name, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + token=token, + proxy_logging_obj=proxy_logging_obj, + ) + + if is_proxy_admin: + return UserAPIKeyAuth( + api_key=None, + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id=user_id, + team_id=team_id, + team_alias=( + team_object.team_alias + if team_object is not None + else None + ), + team_metadata=team_object.metadata + if team_object is not None + else None, + org_id=org_id, + end_user_id=end_user_id, + parent_otel_span=parent_otel_span, + ) + + valid_token = UserAPIKeyAuth( api_key=None, - user_role=LitellmUserRoles.PROXY_ADMIN, - user_id=user_id, team_id=team_id, team_alias=( team_object.team_alias if team_object is not None else None ), + team_tpm_limit=( + team_object.tpm_limit if team_object is not None else None + ), + team_rpm_limit=( + team_object.rpm_limit if team_object is not None else None + ), + team_models=team_object.models + if team_object is not None + else [], + user_role=( + LitellmUserRoles(user_object.user_role) + if user_object is not None + and user_object.user_role is not None + else LitellmUserRoles.INTERNAL_USER + ), + user_id=user_id, + org_id=org_id, + parent_otel_span=parent_otel_span, + end_user_id=end_user_id, + user_tpm_limit=( + user_object.tpm_limit if user_object is not None else None + ), + user_rpm_limit=( + user_object.rpm_limit if user_object is not None else None + ), + team_member_rpm_limit=( + team_membership.safe_get_team_member_rpm_limit() + if team_membership is not None + else None + ), + team_member_tpm_limit=( + team_membership.safe_get_team_member_tpm_limit() + if team_membership is not None + else None + ), team_metadata=team_object.metadata if team_object is not None else None, - org_id=org_id, - end_user_id=end_user_id, - parent_otel_span=parent_otel_span, ) - valid_token = UserAPIKeyAuth( - api_key=None, - team_id=team_id, - team_alias=( - team_object.team_alias if team_object is not None else None - ), - team_tpm_limit=( - team_object.tpm_limit if team_object is not None else None - ), - team_rpm_limit=( - team_object.rpm_limit if team_object is not None else None - ), - team_models=team_object.models if team_object is not None else [], - user_role=( - LitellmUserRoles(user_object.user_role) - if user_object is not None and user_object.user_role is not None - else LitellmUserRoles.INTERNAL_USER - ), - user_id=user_id, - org_id=org_id, - parent_otel_span=parent_otel_span, - end_user_id=end_user_id, - user_tpm_limit=( - user_object.tpm_limit if user_object is not None else None - ), - user_rpm_limit=( - user_object.rpm_limit if user_object is not None else None - ), - team_member_rpm_limit=( - team_membership.safe_get_team_member_rpm_limit() - if team_membership is not None - else None - ), - team_member_tpm_limit=( - team_membership.safe_get_team_member_tpm_limit() - if team_membership is not None - else None - ), - team_metadata=team_object.metadata - if team_object is not None - else None, - ) + # Check if model has zero cost - if so, skip all budget checks + model = get_model_from_request(request_data, route) + skip_budget_checks = False + if model is not None and llm_router is not None: + from litellm.proxy.auth.auth_checks import _is_model_cost_zero - # Check if model has zero cost - if so, skip all budget checks - model = get_model_from_request(request_data, route) - skip_budget_checks = False - if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero - - skip_budget_checks = _is_model_cost_zero( - model=model, llm_router=llm_router - ) - if skip_budget_checks: - verbose_proxy_logger.info( - f"Skipping all budget checks for zero-cost model: {model}" + skip_budget_checks = _is_model_cost_zero( + model=model, llm_router=llm_router ) + if skip_budget_checks: + verbose_proxy_logger.info( + f"Skipping all budget checks for zero-cost model: {model}" + ) - # Fetch project object for JWT path if project_id is set - _jwt_project_obj = None - if valid_token.project_id is not None: - _jwt_project_obj = await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + # Fetch project object for JWT path if project_id is set + _jwt_project_obj = None + if valid_token.project_id is not None: + _jwt_project_obj = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if _jwt_project_obj is not None: + valid_token.project_metadata = _jwt_project_obj.metadata + + # run through common checks + _ = await common_checks( + request=request, + request_body=request_data, + team_object=team_object, + user_object=user_object, + end_user_object=end_user_object, + general_settings=general_settings, + global_proxy_spend=global_proxy_spend, + route=route, + llm_router=llm_router, proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + skip_budget_checks=skip_budget_checks, + project_object=_jwt_project_obj, ) - if _jwt_project_obj is not None: - valid_token.project_metadata = _jwt_project_obj.metadata - # run through common checks - _ = await common_checks( - request=request, - request_body=request_data, - team_object=team_object, - user_object=user_object, - end_user_object=end_user_object, - general_settings=general_settings, - global_proxy_spend=global_proxy_spend, - route=route, - llm_router=llm_router, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - skip_budget_checks=skip_budget_checks, - project_object=_jwt_project_obj, - ) - - # return UserAPIKeyAuth object - return cast(UserAPIKeyAuth, valid_token) + # return UserAPIKeyAuth object + return cast(UserAPIKeyAuth, valid_token) #### ELSE #### ## CHECK PASS-THROUGH ENDPOINTS ## @@ -830,25 +930,26 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead ### CHECK IF ADMIN ### # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead - ## Check CACHE - try: - valid_token = await get_key_object( - hashed_token=hash_token(api_key), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, - ) - except Exception: - verbose_logger.debug("api key not found in cache.") - valid_token = None + if valid_token is None: + ## Check CACHE + try: + valid_token = await get_key_object( + hashed_token=hash_token(api_key), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=True, + ) + except Exception: + verbose_logger.debug("api key not found in cache.") + valid_token = None - ## Check UI Hash Key - if valid_token is None and get_secret_bool("EXPERIMENTAL_UI_LOGIN"): - valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key( - api_key - ) + ## Check UI Hash Key + if valid_token is None and get_secret_bool("EXPERIMENTAL_UI_LOGIN"): + valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key( + api_key + ) if ( valid_token is not None @@ -986,9 +1087,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 param=None, ) - ## check for cache hit (In-Memory Cache) - _user_role = None - if valid_token is None: if isinstance( api_key, str 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..a2a38cad14 --- /dev/null +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -0,0 +1,257 @@ + +from fastapi import APIRouter, Depends, HTTPException, Query + +from litellm.proxy._types import ( + CreateJWTKeyMappingRequest, + DeleteJWTKeyMappingRequest, + JWTKeyMappingResponse, + LitellmUserRoles, + UpdateJWTKeyMappingRequest, + UserAPIKeyAuth, + hash_token, +) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + + +def _to_response(mapping) -> JWTKeyMappingResponse: + """Convert a Prisma mapping object to a safe response (no hashed token).""" + return JWTKeyMappingResponse( + id=mapping.id, + jwt_claim_name=mapping.jwt_claim_name, + jwt_claim_value=mapping.jwt_claim_value, + description=mapping.description, + is_active=mapping.is_active, + created_at=mapping.created_at, + updated_at=mapping.updated_at, + created_by=mapping.created_by, + updated_by=mapping.updated_by, + ) + + +@router.post( + "/jwt/key/mapping/new", + tags=["JWT Key Mapping"], + response_model=JWTKeyMappingResponse, +) +async def create_jwt_key_mapping( + data: CreateJWTKeyMappingRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, detail="Only proxy admins can create JWT key mappings" + ) + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + hashed_key = hash_token(data.key) + create_data = { + "jwt_claim_name": data.jwt_claim_name, + "jwt_claim_value": data.jwt_claim_value, + "token": hashed_key, + "created_by": user_api_key_dict.user_id, + "updated_by": user_api_key_dict.user_id, + } + if data.description is not None: + create_data["description"] = data.description + + new_mapping = await prisma_client.db.litellm_jwtkeymapping.create( + data=create_data + ) + + # Invalidate cache + cache_key = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}" + await user_api_key_cache.async_delete_cache(cache_key) + + return _to_response(new_mapping) + except HTTPException: + raise + except Exception as e: + error_str = str(e).lower() + if "unique" in error_str or "p2002" in error_str: + raise HTTPException( + status_code=409, + detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.", + ) + if "foreign" in error_str or "p2003" in error_str: + raise HTTPException( + status_code=400, + detail="The provided key does not match an existing virtual key.", + ) + raise HTTPException(status_code=500, detail="Failed to create JWT key mapping.") + + +@router.post( + "/jwt/key/mapping/update", + tags=["JWT Key Mapping"], + response_model=JWTKeyMappingResponse, +) +async def update_jwt_key_mapping( + data: UpdateJWTKeyMappingRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, detail="Only proxy admins can update JWT key mappings" + ) + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + update_data = data.model_dump(exclude_unset=True, exclude={"id", "key"}) + if data.key is not None: + update_data["token"] = hash_token(data.key) + update_data["updated_by"] = user_api_key_dict.user_id + + try: + # Get old mapping for cache invalidation + old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + where={"id": data.id} + ) + + if old_mapping is None: + raise HTTPException(status_code=404, detail="Mapping not found") + + cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" + await user_api_key_cache.async_delete_cache(cache_key) + + updated_mapping = await prisma_client.db.litellm_jwtkeymapping.update( + where={"id": data.id}, data=update_data + ) + + # Invalidate new cache key if claim fields changed + cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}" + await user_api_key_cache.async_delete_cache(cache_key) + + return _to_response(updated_mapping) + except HTTPException: + raise + except Exception as e: + error_str = str(e).lower() + if "unique" in error_str or "p2002" in error_str: + raise HTTPException( + status_code=409, + detail="A mapping with those claim values already exists.", + ) + if "foreign" in error_str or "p2003" in error_str: + raise HTTPException( + status_code=400, + detail="The provided key does not match an existing virtual key.", + ) + raise HTTPException(status_code=500, detail="Failed to update JWT key mapping.") + + +@router.post("/jwt/key/mapping/delete", tags=["JWT Key Mapping"]) +async def delete_jwt_key_mapping( + data: DeleteJWTKeyMappingRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, detail="Only proxy admins can delete JWT key mappings" + ) + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + # Get old mapping for cache invalidation + old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + where={"id": data.id} + ) + + if old_mapping is None: + raise HTTPException(status_code=404, detail="Mapping not found") + + cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" + await user_api_key_cache.async_delete_cache(cache_key) + + await prisma_client.db.litellm_jwtkeymapping.delete(where={"id": data.id}) + return {"status": "success"} + except HTTPException: + raise + except Exception: + raise HTTPException(status_code=500, detail="Failed to delete JWT key mapping.") + + +@router.get( + "/jwt/key/mapping/list", + tags=["JWT Key Mapping"], +) +async def list_jwt_key_mappings( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + page: int = Query(1, description="Page number", ge=1), + size: int = Query(50, description="Page size", ge=1, le=100), +): + from litellm.proxy.proxy_server import prisma_client + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, detail="Only proxy admins can list JWT key mappings" + ) + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + skip = (page - 1) * size + mappings = await prisma_client.db.litellm_jwtkeymapping.find_many( + skip=skip, + take=size, + order={"created_at": "desc"}, + ) + total_count = await prisma_client.db.litellm_jwtkeymapping.count() + return { + "mappings": [_to_response(m) for m in mappings], + "total_count": total_count, + "current_page": page, + "total_pages": -(-total_count // size), # ceiling division + } + except HTTPException: + raise + except Exception: + raise HTTPException(status_code=500, detail="Failed to list JWT key mappings.") + + +@router.get( + "/jwt/key/mapping/info", + tags=["JWT Key Mapping"], + response_model=JWTKeyMappingResponse, +) +async def info_jwt_key_mapping( + id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + from litellm.proxy.proxy_server import prisma_client + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, detail="Only proxy admins can get JWT key mapping info" + ) + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + where={"id": id} + ) + if mapping is None: + raise HTTPException(status_code=404, detail="Mapping not found") + return _to_response(mapping) + except HTTPException: + raise + except Exception: + raise HTTPException( + status_code=500, detail="Failed to get JWT key mapping info." + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6a2b0accb0..ad31ff3380 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -377,6 +377,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.key_management_endpoints import ( router as key_management_router, ) +from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( + router as jwt_key_mapping_router, +) from litellm.proxy.management_endpoints.mcp_management_endpoints import ( router as mcp_management_router, ) @@ -13004,6 +13007,7 @@ app.include_router(debugging_endpoints_router) app.include_router(ui_crud_endpoints_router) app.include_router(openai_files_router) app.include_router(team_callback_router) +app.include_router(jwt_key_mapping_router) app.include_router(budget_management_router) app.include_router(model_management_router) app.include_router(model_access_group_management_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f7c0711241..9e847993ef 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -353,6 +353,7 @@ model LiteLLM_VerificationToken { litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + jwt_key_mappings LiteLLM_JWTKeyMapping[] // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 @@ -365,6 +366,24 @@ model LiteLLM_VerificationToken { @@index([budget_reset_at, expires]) } +model LiteLLM_JWTKeyMapping { + id String @id @default(uuid()) + jwt_claim_name String // e.g. "sub", "email" + jwt_claim_value String // The claim value to match + token String // Hashed virtual key (FK) + description String? + is_active Boolean @default(true) + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? + + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + + @@unique([jwt_claim_name, jwt_claim_value]) + @@index([jwt_claim_name, jwt_claim_value, is_active]) +} + // Deprecated keys during grace period - allows old key to work until revoke_at model LiteLLM_DeprecatedVerificationToken { id String @id @default(uuid()) 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..b67dd2792f --- /dev/null +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -0,0 +1,352 @@ +import pytest +import sys +import os +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +# Add project root to sys.path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) + +from litellm.proxy.auth.user_api_key_auth import ( + _resolve_jwt_to_virtual_key, +) +from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy._types import ( + JWTKeyMappingResponse, + LiteLLM_JWTAuth, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( + _to_response, + create_jwt_key_mapping, + delete_jwt_key_mapping, + info_jwt_key_mapping, + update_jwt_key_mapping, +) +from litellm.caching.caching import DualCache +from fastapi import HTTPException + + +# ────────────────────────────────────────────── +# Tests: _resolve_jwt_to_virtual_key +# ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_jwt_to_virtual_key_mapping_resolution(): + """ + Test that a JWT claim is correctly resolved to a virtual key token. + """ + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="email", virtual_key_mapping_cache_ttl=3600 + ) + + jwt_claims = {"email": "user@example.com", "sub": "123"} + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock() + + # Mock finding a mapping + mock_mapping = MagicMock() + mock_mapping.token = "sk-1234" + mock_mapping.is_active = True + prisma_client.db.litellm_jwtkeymapping.find_first.return_value = mock_mapping + + # Mock getting the key object + mock_key_obj = UserAPIKeyAuth(token="sk-1234", team_id="team1") + + user_api_key_cache = DualCache() + + # Use patch to mock get_key_object in the module where it's used + with patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + ) as mock_get_key: + mock_get_key.return_value = mock_key_obj + + result = await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert result == mock_key_obj + prisma_client.db.litellm_jwtkeymapping.find_first.assert_called_once() + + # Test Cache hit + prisma_client.db.litellm_jwtkeymapping.find_first.reset_mock() + result_cached = await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert result_cached == mock_key_obj + prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called() + + +@pytest.mark.asyncio +async def test_jwt_to_virtual_key_mapping_no_mapping(): + """ + Test that when no mapping exists, resolve returns None. + """ + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_claim_field="email") + jwt_claims = {"email": "unknown@example.com"} + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock() + prisma_client.db.litellm_jwtkeymapping.find_first.return_value = None + + # Mock get_key_object just in case + with patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + ): + user_api_key_cache = DualCache() + + result = await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert result is None + + # Test Negative Cache hit + prisma_client.db.litellm_jwtkeymapping.find_first.reset_mock() + result_cached = await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert result_cached is None + prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called() + + +# ────────────────────────────────────────────── +# Tests: _to_response redacts hashed token +# ────────────────────────────────────────────── + + +def test_to_response_excludes_token(): + """_to_response should not expose the hashed token field.""" + now = datetime.now(timezone.utc) + mock_mapping = MagicMock() + mock_mapping.id = "mapping-1" + mock_mapping.jwt_claim_name = "email" + mock_mapping.jwt_claim_value = "user@example.com" + mock_mapping.token = "hashed_secret_value" + mock_mapping.description = "test" + mock_mapping.is_active = True + mock_mapping.created_at = now + mock_mapping.updated_at = now + mock_mapping.created_by = "admin" + mock_mapping.updated_by = "admin" + + resp = _to_response(mock_mapping) + + assert isinstance(resp, JWTKeyMappingResponse) + assert resp.id == "mapping-1" + assert resp.jwt_claim_name == "email" + assert "token" not in resp.model_fields + + +# ────────────────────────────────────────────── +# Helpers +# ────────────────────────────────────────────── + + +def _make_admin_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + +def _make_non_admin_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="sk-user", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + +def _mock_prisma(): + prisma = MagicMock() + prisma.db.litellm_jwtkeymapping.create = AsyncMock() + prisma.db.litellm_jwtkeymapping.find_unique = AsyncMock() + prisma.db.litellm_jwtkeymapping.find_many = AsyncMock() + prisma.db.litellm_jwtkeymapping.update = AsyncMock() + prisma.db.litellm_jwtkeymapping.delete = AsyncMock() + prisma.db.litellm_jwtkeymapping.count = AsyncMock(return_value=0) + return prisma + + +def _mock_mapping( + id="mapping-1", + claim_name="email", + claim_value="user@example.com", +): + now = datetime.now(timezone.utc) + m = MagicMock() + m.id = id + m.jwt_claim_name = claim_name + m.jwt_claim_value = claim_value + m.token = "hashed_token" + m.description = None + m.is_active = True + m.created_at = now + m.updated_at = now + m.created_by = "admin" + m.updated_by = "admin" + return m + + +# ────────────────────────────────────────────── +# Tests: CRUD endpoint error handling +# ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_create_returns_409_on_unique_violation(): + """Duplicate mapping should return 409, not 500.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception( + "Unique constraint failed (P2002)" + ) + mock_cache = AsyncMock() + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + with pytest.raises(HTTPException) as exc_info: + await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 409 + assert "already exists" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_create_returns_400_on_foreign_key_violation(): + """Non-existent key should return 400, not 500.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception( + "Foreign key constraint failed on field: `token` (P2003)" + ) + mock_cache = AsyncMock() + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="sub", jwt_claim_value="user-999", key="sk-nonexistent", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + with pytest.raises(HTTPException) as exc_info: + await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 400 + assert "does not match" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_create_non_admin_returns_403(): + """Non-admin users should get 403.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test", + ) + + with pytest.raises(HTTPException) as exc_info: + await create_jwt_key_mapping(data=data, user_api_key_dict=_make_non_admin_auth()) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_delete_returns_404_when_not_found(): + """Deleting non-existent mapping should return 404.""" + from litellm.proxy._types import DeleteJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None + mock_cache = AsyncMock() + + data = DeleteJWTKeyMappingRequest(id="nonexistent-id") + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + with pytest.raises(HTTPException) as exc_info: + await delete_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_update_returns_404_when_not_found(): + """Updating non-existent mapping should return 404.""" + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None + mock_cache = AsyncMock() + + data = UpdateJWTKeyMappingRequest(id="nonexistent-id", description="test") + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + with pytest.raises(HTTPException) as exc_info: + await update_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_info_returns_404_when_not_found(): + """Getting info for non-existent mapping should return 404.""" + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with pytest.raises(HTTPException) as exc_info: + await info_jwt_key_mapping(id="nonexistent-id", user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_create_success_returns_response_without_token(): + """Successful create should return JWTKeyMappingResponse without hashed token.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping() + mock_cache = AsyncMock() + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + result = await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert isinstance(result, JWTKeyMappingResponse) + assert "token" not in result.model_fields + assert result.jwt_claim_name == "email"