Merge remote-tracking branch 'origin/main' into litellm_/release-version-bump-940c92

This commit is contained in:
Yuneng Jiang 2026-09-15 15:56:18 -07:00
commit 818fd7bbb1
No known key found for this signature in database
13 changed files with 556 additions and 43 deletions

View file

@ -0,0 +1,18 @@
-- DropIndex
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx";
-- DropIndex
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key";
-- AlterTable
-- NOT NULL DEFAULT '' (not nullable): Postgres unique constraints treat every
-- NULL as distinct, so a nullable column would let multiple unscoped mappings
-- collide on the same claim without a constraint violation. The constant
-- default is a fast, metadata-only backfill for existing rows, not a rewrite.
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD COLUMN IF NOT EXISTS "jwt_issuer" TEXT NOT NULL DEFAULT '';
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_idx" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value", "is_active");
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_key" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value");

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -15226,6 +15226,17 @@
"title": "Jwt Claim Value",
"type": "string"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"key": {
"title": "Key",
"type": "string"
@ -15310,6 +15321,17 @@
"title": "Jwt Claim Value",
"type": "string"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"updated_at": {
"format": "date-time",
"title": "Updated At",
@ -15366,6 +15388,17 @@
],
"title": "Is Active"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"key": {
"anyOf": [
{

View file

@ -4486,12 +4486,14 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
jwt_claim_name: str
jwt_claim_value: str
key: str
jwt_issuer: str | None = None
description: str | None = None
class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
id: str
key: str | None = None
jwt_issuer: str | None = None
description: str | None = None
is_active: bool | None = None
@ -4502,6 +4504,7 @@ class DeleteJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
class JWTKeyMappingResponse(LiteLLMPydanticObjectBase):
id: str
jwt_issuer: str | None = None
jwt_claim_name: str
jwt_claim_value: str
description: str | None = None

View file

@ -148,6 +148,7 @@ class _PrismaDictableRow(Protocol):
class _PrismaJWTKeyMappingRow(Protocol):
token: str
jwt_issuer: str
jwt_claim_name: str
jwt_claim_value: str
@ -3601,9 +3602,18 @@ async def _fetch_key_object_from_db_with_reconnect(
raise
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str:
"""Cache key under which ``_resolve_jwt_to_virtual_key`` stores a JWT-claim-to-key mapping."""
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str, jwt_issuer: str | None = None) -> str:
"""Cache key under which a JWT-claim-to-key mapping is stored, scoped to one
issuer (or the issuer-agnostic/global scope when ``jwt_issuer`` is falsy).
Scoped by issuer (when one is configured) so a cached hit or ``__NO_MAPPING__`` miss
for one issuer's claim value can never be served to a different issuer whose claim
value happens to collide. Unchanged for the global scope, keeping the single-issuer
(no ``litellm_jwtauth.issuers`` configured) cache key format stable across this fix.
"""
if not jwt_issuer:
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
return f"jwt_key_mapping:{jwt_issuer}:{jwt_claim_name}:{jwt_claim_value}"
@log_db_metrics
@ -3615,7 +3625,7 @@ async def get_jwt_key_mapping_cache_keys_for_token(
mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many(
where={"token": hashed_token}
)
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value) for m in mappings)
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
@log_db_metrics
@ -3623,9 +3633,14 @@ async def get_jwt_key_mapping_object(
jwt_claim_name: str,
jwt_claim_value: str,
prisma_client: PrismaClient,
jwt_issuer: str | None = None,
) -> str | None:
"""
Lookup a JWT-to-virtual-key mapping from the database.
Lookup a JWT-to-virtual-key mapping from the database for one exact scope:
``jwt_issuer`` (or the global/issuer-agnostic scope when falsy). Does not fall
back to the global scope itself -- a caller that wants "issuer-scoped mapping,
else the global one" queries both scopes itself, so each result can be cached
under its own scope's key (see ``_resolve_jwt_to_virtual_key``).
Returns the hashed token (str) if a matching active mapping is found, else None.
"""
@ -3633,6 +3648,7 @@ async def get_jwt_key_mapping_object(
where={
"jwt_claim_name": jwt_claim_name,
"jwt_claim_value": jwt_claim_value,
"jwt_issuer": jwt_issuer or "",
"is_active": True,
}
)

View file

@ -269,6 +269,17 @@ class _TokenTeamModels(Protocol):
def team_models(self) -> list[str]: ...
class _RawCacheRead(Protocol):
async def async_get_cache(self, *, key: str) -> object: ...
def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead:
"""View an untyped cache object's ``async_get_cache`` as returning ``object``
instead of ``Any``, so a caller can ``isinstance``-narrow it without paying
the ``reportAny`` cost of the underlying (unannotated) cache implementation."""
return cache
def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
return valid_token.team_models
@ -842,6 +853,7 @@ class _PendingAutoRegister(NamedTuple):
claim_field: str
claim_value: str
cache_key: str
jwt_issuer: str | None = None
async def _auto_register_jwt_mapping(
@ -853,6 +865,7 @@ async def _auto_register_jwt_mapping(
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
cache_key: str,
jwt_issuer: str | None = None,
team_id: str | None = None,
user_id: str | None = None,
org_id: str | None = None,
@ -905,6 +918,7 @@ async def _auto_register_jwt_mapping(
try:
await prisma_client.db.litellm_jwtkeymapping.create(
data={
"jwt_issuer": jwt_issuer or "",
"jwt_claim_name": virtual_key_claim_field,
"jwt_claim_value": claim_value,
"token": token_hash,
@ -939,6 +953,7 @@ async def _auto_register_jwt_mapping(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=jwt_issuer,
)
if token_hash is None:
# The winner's mapping vanished between the unique-constraint
@ -983,6 +998,43 @@ async def _auto_register_jwt_mapping(
return auto_registered_key
async def _lookup_jwt_mapping_token_hash(
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
virtual_key_claim_field: str,
claim_value: str,
normalized_issuer: str | None,
cache_key: str,
ttl: float,
) -> str | None:
issuer_scoped: Final = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=normalized_issuer,
)
if issuer_scoped is not None:
await user_api_key_cache.async_set_cache(key=cache_key, value=issuer_scoped, ttl=ttl)
return issuer_scoped
if normalized_issuer is None:
return None
# Another issuer may have already resolved (and cached) this same
# global mapping -- check its cache entry before re-querying the DB.
global_cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, claim_value)
cached_global: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=global_cache_key)
if isinstance(cached_global, str) and cached_global != "__NO_MAPPING__":
return cached_global
global_row: Final = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=None,
)
if global_row is not None:
await user_api_key_cache.async_set_cache(key=global_cache_key, value=global_row, ttl=ttl)
return global_row
async def _resolve_jwt_to_virtual_key(
jwt_claims: dict,
jwt_handler: JWTHandler,
@ -1041,7 +1093,7 @@ async def _resolve_jwt_to_virtual_key(
)
return None
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value), normalized_issuer)
raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER
cached_mapping: Final = (
@ -1081,6 +1133,7 @@ async def _resolve_jwt_to_virtual_key(
claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
cache_key=cache_key,
jwt_issuer=normalized_issuer,
)
return None
elif cached_mapping is not None:
@ -1094,21 +1147,30 @@ async def _resolve_jwt_to_virtual_key(
)
# Resolve the mapping from DB, or treat prisma_client=None as a definitive
# miss (no DB → no mapping can exist → apply no-match policy below).
token_hash: str | None = None
if prisma_client is not None:
token_hash = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=str(claim_value),
# miss (no DB → no mapping can exist → apply no-match policy below). An
# issuer-scoped row wins; falling back to the global (no-issuer) row keeps
# mappings created before issuer scoping existed working for every issuer.
# Each tier is cached under ITS OWN key (the global tier under the
# issuer-less cache key, not under `cache_key`/this issuer's key) so that
# updating or deleting either row invalidates exactly the cache entries it
# can affect. Caching a global-row hit under the requesting issuer's key
# would leave every OTHER issuer that had fallen back to that same global
# mapping serving its stale token until TTL after the row changes.
token_hash: Final = (
await _lookup_jwt_mapping_token_hash(
prisma_client=prisma_client,
)
if token_hash is not None:
await user_api_key_cache.async_set_cache(
key=cache_key,
value=token_hash,
user_api_key_cache=user_api_key_cache,
virtual_key_claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
normalized_issuer=normalized_issuer,
cache_key=cache_key,
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
)
if prisma_client is not None
else None
)
if token_hash is not None:
return IdentityStore.key_from_principal(
await IdentityStore(
prisma_client,
@ -1149,6 +1211,7 @@ async def _resolve_jwt_to_virtual_key(
claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
cache_key=cache_key,
jwt_issuer=normalized_issuer,
)
# FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the
@ -1641,6 +1704,7 @@ async def _user_api_key_auth_builder(
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
cache_key=pending_auto_register.cache_key,
jwt_issuer=pending_auto_register.jwt_issuer,
team_id=team_id,
user_id=user_id,
org_id=org_id,

View file

@ -28,6 +28,9 @@ class _JWTKeyMappingRecord(Protocol):
@property
def id(self) -> str: ...
@property
def jwt_issuer(self) -> str: ...
@property
def jwt_claim_name(self) -> str: ...
@ -78,6 +81,7 @@ def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse:
"""Convert a Prisma mapping object to a safe response (no hashed token)."""
return JWTKeyMappingResponse(
id=mapping.id,
jwt_issuer=mapping.jwt_issuer or None,
jwt_claim_name=mapping.jwt_claim_name,
jwt_claim_value=mapping.jwt_claim_value,
description=mapping.description,
@ -109,6 +113,7 @@ async def create_jwt_key_mapping(
try:
hashed_key: Final = hash_token(data.key)
create_data: Final = {
"jwt_issuer": data.jwt_issuer or "",
"jwt_claim_name": data.jwt_claim_name,
"jwt_claim_value": data.jwt_claim_value,
"token": hashed_key,
@ -120,7 +125,7 @@ async def create_jwt_key_mapping(
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value)
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value, data.jwt_issuer)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return _to_response(new_mapping)
@ -131,7 +136,10 @@ async def create_jwt_key_mapping(
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.",
detail=(
f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' "
f"already exists for issuer '{data.jwt_issuer}'."
),
)
if "foreign" in error_str or "p2003" in error_str:
raise HTTPException(
@ -161,6 +169,9 @@ async def update_jwt_key_mapping(
update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"})
if data.key is not None:
update_data["token"] = hash_token(data.key)
if "jwt_issuer" in update_data:
# DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel.
update_data["jwt_issuer"] = update_data["jwt_issuer"] or ""
update_data["updated_by"] = user_api_key_dict.user_id
try:
@ -178,9 +189,11 @@ async def update_jwt_key_mapping(
# Evict only after the write commits: a concurrent request between an
# early eviction and the commit would re-cache the old mapping and keep
# it authorized until TTL.
old_cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
old_cache_key: Final = jwt_key_mapping_cache_key(
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
)
new_cache_key: Final = jwt_key_mapping_cache_key(
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value, updated_mapping.jwt_issuer
)
cache_keys: Final = (old_cache_key,) if old_cache_key == new_cache_key else (old_cache_key, new_cache_key)
await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache)
@ -227,7 +240,9 @@ async def delete_jwt_key_mapping(
# Evict only after the row is gone, else a concurrent request can
# re-cache the deleted mapping and keep it authorized until TTL.
cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
cache_key: Final = jwt_key_mapping_cache_key(
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return {"status": "success"}
except HTTPException:

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -91,6 +91,154 @@ async def test_jwt_to_virtual_key_mapping_resolution():
prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
@pytest.mark.asyncio
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
DIFFERENT issuer whose claim value happens to collide, even though both issuers
map the same claim field (``sub``) to a virtual key."""
issuer_a = "https://issuer-a.example.com"
issuer_b = "https://issuer-b.example.com"
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=3600
)
rows = [
{
"jwt_issuer": issuer_b,
"jwt_claim_name": "sub",
"jwt_claim_value": "dev-alice",
"token": "hashed-issuer-b-key",
"is_active": True,
}
]
async def fake_find_first(where):
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return MagicMock(**row)
return None
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(side_effect=fake_find_first)
# Dependency-inject the resolved key via the cache (IdentityStore._resolve_key
# reads it from here) instead of monkeypatching IdentityStore itself.
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-issuer-b-key",
value=UserAPIKeyAuth(token="hashed-issuer-b-key", team_id="issuer-b-team"),
)
# The rightful owner: issuer-b's own claim resolves to its mapping.
owner_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "dev-alice"},
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 isinstance(owner_result, UserAPIKeyAuth)
assert owner_result.token == "hashed-issuer-b-key"
# A validly-signed token from issuer-a carrying the SAME claim value must not
# inherit issuer-b's mapping. Default behavior is fallback_team_mapping, so a
# correctly-scoped miss returns None instead of resolving to issuer-b's key.
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "dev-alice"},
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 colliding_result is None
@pytest.mark.asyncio
async def test_global_mapping_resolution_is_cached_under_the_global_key_not_the_requesting_issuer():
"""LIT-7417: caching a global (unscoped) mapping's hit under the REQUESTING
issuer's key would leave every issuer that falls back to it holding its own
stale copy after the row is updated/deleted -- CRUD only evicts the cache key
computed from the row's own scope (global), so a copy cached under some other
issuer's key would keep resolving to the old token until TTL. Caching it under
the global key instead means every issuer shares (and CRUD correctly evicts)
the exact same entry."""
issuer_a = "https://issuer-a.example.com"
issuer_b = "https://issuer-b.example.com"
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
issuers=[
{
"issuer": issuer_a,
"jwks_url": f"{issuer_a}/jwks",
"virtual_key_claim_field": "sub",
"disable_audience_validation": True,
},
{
"issuer": issuer_b,
"jwks_url": f"{issuer_b}/jwks",
"virtual_key_claim_field": "sub",
"disable_audience_validation": True,
},
]
)
rows = [
{
"jwt_issuer": "",
"jwt_claim_name": "sub",
"jwt_claim_value": "legacy-user",
"token": "hashed-legacy-key",
"is_active": True,
}
]
async def fake_find_first(where):
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return MagicMock(**row)
return None
prisma_client = MagicMock()
find_first = AsyncMock(side_effect=fake_find_first)
prisma_client.db.litellm_jwtkeymapping.find_first = find_first
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-legacy-key",
value=UserAPIKeyAuth(token="hashed-legacy-key", team_id="legacy-team"),
)
resolved_a = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "legacy-user"},
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 isinstance(resolved_a, UserAPIKeyAuth)
assert find_first.await_count == 2 # issuer-a-scoped miss, then global hit
# issuer-b resolving the SAME global mapping must hit the cache issuer-a's
# resolution populated, not issue a fresh DB query for the global row again.
resolved_b = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "legacy-user"},
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 isinstance(resolved_b, UserAPIKeyAuth)
assert resolved_b.token == "hashed-legacy-key"
assert find_first.await_count == 3 # +1 for issuer-b's own issuer-scoped miss; global tier served from cache
@pytest.mark.asyncio
async def test_jwt_to_virtual_key_mapping_no_mapping():
"""
@ -223,6 +371,7 @@ def test_to_response_excludes_token():
now = datetime.now(timezone.utc)
mock_mapping = MagicMock()
mock_mapping.id = "mapping-1"
mock_mapping.jwt_issuer = None
mock_mapping.jwt_claim_name = "email"
mock_mapping.jwt_claim_value = "user@example.com"
mock_mapping.token = "hashed_secret_value"
@ -275,10 +424,12 @@ def _mock_mapping(
id="mapping-1",
claim_name="email",
claim_value="user@example.com",
issuer=None,
):
now = datetime.now(timezone.utc)
m = MagicMock()
m.id = id
m.jwt_issuer = issuer
m.jwt_claim_name = claim_name
m.jwt_claim_value = claim_value
m.token = "hashed_token"
@ -485,6 +636,35 @@ async def test_create_success_returns_response_without_token():
assert result.jwt_claim_name == "email"
@pytest.mark.asyncio
async def test_create_without_issuer_stores_empty_string_not_null():
"""LIT-7417: the DB column is NOT NULL (see schema.prisma). Storing a real NULL
for an unscoped mapping would let Postgres accept unlimited duplicate unscoped
rows for the same claim (NULL is never equal to NULL in a unique constraint),
so two mappings for the same claim value could point at two different keys with
no conflict, and resolution would pick whichever one Postgres returns first."""
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="sub", jwt_claim_value="dev-alice", key="sk-test-key")
with (
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
"litellm.proxy.proxy_server.prisma_client", mock_prisma
),
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
),
):
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
sent_data = mock_prisma.db.litellm_jwtkeymapping.create.call_args.kwargs["data"]
assert sent_data["jwt_issuer"] == ""
# ──────────────────────────────────────────────
# Tests: unregistered_jwt_client_behavior
# ──────────────────────────────────────────────

View file

@ -32,7 +32,13 @@ from litellm.proxy._types import (
JWTRoutingOverride,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object
from litellm.proxy.auth.auth_checks import (
TeamNotFoundError,
UserNotFoundError,
get_key_object,
_cache_key_object,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import (
_check_key_model_budget_with_fallback,
@ -7948,13 +7954,38 @@ def _per_issuer_virtual_key_jwt_handler(
def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]:
"""Every ``find_first`` call (issuer-scoped or global fallback) resolves the same way."""
find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token))
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
return prisma_client, find_first
def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]:
return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True}
def _fake_prisma_jwt_key_mapping_table(rows: list[dict[str, object]]) -> tuple[SimpleNamespace, AsyncMock]:
"""A ``find_first`` whose result depends on the ``where`` clause, like a real table.
Matches a row when every key present in ``where`` equals that key on the row --
a key ``get_jwt_key_mapping_object`` omits (e.g. old, issuer-blind code never
sending ``jwt_issuer``) does not constrain the match, exactly like Prisma.
"""
async def _find_first(where: dict[str, object]) -> SimpleNamespace | None:
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return SimpleNamespace(**row)
return None
find_first = AsyncMock(side_effect=_find_first)
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
return prisma_client, find_first
def _mapping_where(claim_name: str, claim_value: str, jwt_issuer: str | None) -> dict[str, str | bool]:
return {
"jwt_claim_name": claim_name,
"jwt_claim_value": claim_value,
"jwt_issuer": jwt_issuer or "",
"is_active": True,
}
@pytest.mark.asyncio
@ -7978,11 +8009,13 @@ async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for
proxy_logging_obj=MagicMock(),
)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7"))
# Issuer-scoped lookup hits on the first query, so no global fallback query runs.
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7", ISSUER_TWO))
assert isinstance(resolved, UserAPIKeyAuth)
assert resolved.token == "hashed-mapped-key"
assert resolved.team_id == "svc-team"
assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key"
cache_key = jwt_key_mapping_cache_key("sub", "svc-account-7", ISSUER_TWO)
assert await user_api_key_cache.async_get_cache(cache_key) == "hashed-mapped-key"
@pytest.mark.asyncio
@ -8015,7 +8048,11 @@ async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer():
assert exc.value.status_code == 403
assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc"))
# REJECT checks the issuer-scoped row first, then falls back to a global (NULL-issuer) row.
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("sub", "unknown-svc", ISSUER_TWO),
_mapping_where("sub", "unknown-svc", None),
]
@pytest.mark.asyncio
@ -8025,7 +8062,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register")
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL)
# Sentinel cached under issuer-one's own key -- must never answer issuer-two's lookup.
await user_api_key_cache.async_set_cache(
key=jwt_key_mapping_cache_key("sub", "admin-7", ISSUER_ONE), value=_JWT_PROXY_ADMIN_SENTINEL
)
auto_register_issuer_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"},
@ -8050,7 +8090,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
assert exc.value.status_code == 403
assert "No registered mapping for sub='admin-7'" in str(exc.value.detail)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7"))
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("sub", "admin-7", ISSUER_TWO),
_mapping_where("sub", "admin-7", None),
]
@pytest.mark.asyncio
@ -8079,7 +8122,125 @@ async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_f
assert with_claim is None
assert without_claim is None
find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9"))
# without_claim has no claim value and returns before ever reaching the DB.
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("client_id", "app-9", ISSUER_ONE),
_mapping_where("client_id", "app-9", None),
]
@pytest.mark.asyncio
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
DIFFERENT issuer whose claim value happens to collide, even though both issuers
use the same claim field (``sub``) for their virtual-key mapping."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table(
[
{
"jwt_issuer": ISSUER_TWO,
"jwt_claim_name": "sub",
"jwt_claim_value": "dev-alice",
"token": "hashed-issuer-b-key",
"is_active": True,
}
]
)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-issuer-b-key",
value=UserAPIKeyAuth(token="hashed-issuer-b-key", api_key="hashed-issuer-b-key", team_id="issuer-b-team"),
)
owner_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert isinstance(owner_result, UserAPIKeyAuth)
assert owner_result.token == "hashed-issuer-b-key"
# issuer-one's behavior is fallback_team_mapping: a correctly-scoped miss must
# return None (fall through to team-based JWT auth), never issuer-two's key.
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert colliding_result is None
assert find_first.await_count == 3 # owner hit (1 call) + colliding miss (issuer-scoped + global fallback)
@pytest.mark.asyncio
async def test_cached_resolution_for_one_issuer_does_not_leak_to_a_colliding_issuer():
"""A cached positive resolution must be keyed by issuer too, or a colliding
claim value from another issuer could be served straight from cache without
ever reaching the (correctly issuer-scoped) DB lookup."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table([])
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key=jwt_key_mapping_cache_key("sub", "dev-alice", ISSUER_TWO), value="hashed-issuer-b-key"
)
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert colliding_result is None
# Must have gone to the DB rather than serving issuer-two's cached token.
assert find_first.await_count == 2
@pytest.mark.asyncio
async def test_issuer_agnostic_mapping_matches_every_issuer():
"""A mapping created before issuer scoping existed (``jwt_issuer`` is NULL) keeps
matching any issuer, so existing global mappings are not broken by this fix."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, _find_first = _fake_prisma_jwt_key_mapping_table(
[
{
"jwt_issuer": "",
"jwt_claim_name": "sub",
"jwt_claim_value": "legacy-user",
"token": "hashed-legacy-key",
"is_active": True,
}
]
)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-legacy-key",
value=UserAPIKeyAuth(token="hashed-legacy-key", api_key="hashed-legacy-key", team_id="legacy-team"),
)
for issuer in (ISSUER_ONE, ISSUER_TWO):
resolved = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer, "sub": "legacy-user"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert isinstance(resolved, UserAPIKeyAuth)
assert resolved.token == "hashed-legacy-key"
@pytest.mark.asyncio

View file

@ -36,7 +36,11 @@ from litellm.proxy._types import (
UpdateKeyRequest,
)
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key
from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
_project_cache_key,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
@ -5132,10 +5136,11 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
class _JWTMappingRow:
def __init__(self, token, jwt_claim_name, jwt_claim_value):
def __init__(self, token, jwt_claim_name, jwt_claim_value, jwt_issuer=None):
self.token = token
self.jwt_claim_name = jwt_claim_name
self.jwt_claim_value = jwt_claim_value
self.jwt_issuer = jwt_issuer
class _CascadingJWTMappingTable:
@ -5226,7 +5231,7 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat
),
)
assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",)
assert recording_evict.cache_keys == (jwt_key_mapping_cache_key("email", "user@example.com", None),)
@pytest.mark.asyncio
@ -13131,11 +13136,11 @@ async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new
_execute_virtual_key_regeneration,
)
stale_cache_key = "jwt_key_mapping:sub:user1"
stale_cache_key = jwt_key_mapping_cache_key("sub", "user1", None)
existing_key = _make_regenerate_existing_key()
mock_prisma_client = _make_regenerate_mock_prisma()
mock_prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1")]
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1", jwt_issuer=None)]
)
mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
return_value=MagicMock(token="new-hashed-token")

View file

@ -27297,6 +27297,8 @@ export interface components {
jwt_claim_name: string;
/** Jwt Claim Value */
jwt_claim_value: string;
/** Jwt Issuer */
jwt_issuer?: string | null;
/** Key */
key: string;
};
@ -28910,6 +28912,8 @@ export interface components {
jwt_claim_name: string;
/** Jwt Claim Value */
jwt_claim_value: string;
/** Jwt Issuer */
jwt_issuer?: string | null;
/**
* Updated At
* Format: date-time
@ -38643,6 +38647,8 @@ export interface components {
id: string;
/** Is Active */
is_active?: boolean | null;
/** Jwt Issuer */
jwt_issuer?: string | null;
/** Key */
key?: string | null;
};