diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6440bf0ed81..6b51df709f7 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 d3b028f63f9..6ca7b290a07 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 9680a92ae67..c5d91d3699b 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