mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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:
parent
eee239cd9e
commit
969d152c4b
2 changed files with 117 additions and 0 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue