tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-10 22:24:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: jwt mapping vkeyv

Browse Source
This commit is contained in:
Harshit28j
2026-02-28 13:29:46 +05:30
parent 50ecfbc614
commit 9dc085694c
7 changed files with 553 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+45
View File
@@ -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:
litellm/proxy/auth/handle_jwt.py
+2
View File
@@ -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,
)
litellm/proxy/auth/user_api_key_auth.py
+219 -129
View File
@@ -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
litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py
+152
View File
@@ -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))
litellm/proxy/proxy_server.py
+4
View File
@@ -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)
litellm/proxy/schema.prisma
+19
View File
@@ -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())
tests/proxy_unit_tests/test_jwt_key_mapping.py
+112
View File
@@ -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()