fix(jwt): evict jwt_key_mapping cache when a virtual key is deleted

The FK cascade drops the LiteLLM_JWTKeyMapping row, but the cached
jwt_key_mapping:{claim}:{value} entry still resolved to the deleted token
hash, so every JWT call from that identity failed until
virtual_key_mapping_cache_ttl expired instead of auto-registering against a
recreated key. delete_verification_tokens now snapshots the mapping cache
keys before the delete and evicts them across replicas afterwards, the same
way /key/regenerate already does.

Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb
This commit is contained in:
ryan-crabbe-berri 2026-09-09 16:40:25 -07:00
parent eee239cd9e
commit 969d152c4b
2 changed files with 117 additions and 0 deletions

View file

@ -4559,6 +4559,23 @@ async def delete_verification_tokens(
litellm_changed_by=litellm_changed_by,
)
# Snapshot before the delete: the FK cascade drops the mapping rows, but their
# cached jwt_key_mapping entries still resolve to the now-dead token (LIT-5380).
jwt_mapping_cache_keys: Final[tuple[str, ...]] = tuple(
cache_key
for keys_for_token in await asyncio.gather(
*(
get_jwt_key_mapping_cache_keys_for_token(
hashed_token=key.token,
prisma_client=prisma_client,
)
for key in authorized_keys
if key.token is not None
)
)
for cache_key in keys_for_token
)
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
deleted_tokens = await prisma_client.delete_data(tokens=tokens)
if deleted_tokens is not None and len(deleted_tokens) != len(tokens):
@ -4571,6 +4588,8 @@ async def delete_verification_tokens(
if len(deleted_tokens) != len(tokens):
failed_tokens = [token for token in tokens if token not in deleted_tokens]
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
else:
raise Exception("DB not connected. prisma_client is None")
except Exception as e:

View file

@ -5085,6 +5085,104 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
assert len(deleted_keys) == 2
class _JWTMappingRow:
def __init__(self, token, jwt_claim_name, jwt_claim_value):
self.token = token
self.jwt_claim_name = jwt_claim_name
self.jwt_claim_value = jwt_claim_value
class _CascadingJWTMappingTable:
"""Mapping rows that LiteLLM_JWTKeyMapping_token_fkey drops when their key is deleted."""
def __init__(self, rows):
self.rows = rows
async def find_many(self, where, **kwargs):
return [row for row in self.rows if row.token == where["token"]]
def cascade(self, deleted_tokens):
self.rows = [row for row in self.rows if row.token not in deleted_tokens]
class _RecordingEvict:
def __init__(self):
self.cache_keys = ()
async def __call__(self, cache_keys, user_api_key_cache):
self.cache_keys = tuple(cache_keys)
@pytest.mark.asyncio
async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypatch):
"""Deleting a key must evict its jwt_key_mapping cache entries (LIT-5380).
The FK cascade removes the mapping rows, so a surviving cache entry would keep
resolving the deleted token hash and 401 every JWT call from that identity until
virtual_key_mapping_cache_ttl expires, instead of auto-registering again.
"""
jwt_table = _CascadingJWTMappingTable(
[_JWTMappingRow("hashed-token-1", "email", "user@example.com")]
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id=None,
key_alias="jwt-mapped-key",
spend=0.0,
max_budget=None,
models=[],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[key1]
)
mock_prisma_client.db.litellm_jwtkeymapping = jwt_table
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
async def cascading_delete_data(tokens):
jwt_table.cascade(tokens)
return list(tokens)
mock_prisma_client.delete_data = AsyncMock(side_effect=cascading_delete_data)
recording_evict = _RecordingEvict()
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.evict_and_broadcast",
recording_evict,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
lambda token: token,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
await delete_verification_tokens(
tokens=["hashed-token-1"],
user_api_key_cache=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
),
)
assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",)
@pytest.mark.asyncio
async def test_delete_key_fn_persists_deleted_keys(monkeypatch):
from litellm.proxy._types import KeyRequest