fix(jwt): invalidate JWT key mapping cache on /key/regenerate

/key/regenerate carries the JWT-to-key mapping to the new token via FK
cascade, but the jwt_key_mapping cache entry kept resolving the old
(now invalid) token for up to virtual_key_mapping_cache_ttl. Snapshot
the key's mapping cache keys before the token update and evict them
with evict_and_broadcast so every worker drops the stale entry.

Also share the cache-key format through jwt_key_mapping_cache_key and
upgrade the /jwt/key/mapping CRUD endpoints from local-only deletes to
evict_and_broadcast, closing the same cross-worker staleness there.
This commit is contained in:
ryan-crabbe-berri 2026-09-04 14:11:31 -07:00
parent 300d335255
commit 2f7ee39545
5 changed files with 147 additions and 11 deletions

View file

@ -142,6 +142,8 @@ class _PrismaDictableRow(Protocol):
class _PrismaJWTKeyMappingRow(Protocol):
token: str
jwt_claim_name: str
jwt_claim_value: str
class _PrismaModelDumpRow(Protocol):
@ -3466,6 +3468,23 @@ 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}"
@log_db_metrics
async def get_jwt_key_mapping_cache_keys_for_token(
hashed_token: str,
prisma_client: PrismaClient,
) -> tuple[str, ...]:
"""Cache keys of every JWT claim mapped to the given virtual key."""
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)
@log_db_metrics
async def get_jwt_key_mapping_object(
jwt_claim_name: str,

View file

@ -58,6 +58,7 @@ from litellm.proxy.auth.auth_checks import (
get_team_object,
get_user_object,
is_valid_fallback_model,
jwt_key_mapping_cache_key,
resolve_and_validate_end_user_id,
)
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler
@ -970,7 +971,7 @@ async def _resolve_jwt_to_virtual_key(
)
return None
cache_key: Final = f"jwt_key_mapping:{virtual_key_claim_field}:{claim_value}"
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
if cached_mapping == _JWT_PROXY_ADMIN_SENTINEL:

View file

@ -13,7 +13,9 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
hash_token,
)
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.repositories.table_repositories import JWTKeyMappingRepository
@ -118,9 +120,8 @@ async def create_jwt_key_mapping(
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
# Invalidate cache
cache_key: Final = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return _to_response(new_mapping)
except HTTPException:
@ -169,17 +170,18 @@ async def update_jwt_key_mapping(
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
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)
old_cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
await evict_and_broadcast(cache_keys=(old_cache_key,), user_api_key_cache=user_api_key_cache)
updated_mapping: Final = await _mapping_table(prisma_client).update(where={"id": data.id}, data=update_data)
if updated_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
# 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)
new_cache_key: Final = jwt_key_mapping_cache_key(
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
)
await evict_and_broadcast(cache_keys=(new_cache_key,), user_api_key_cache=user_api_key_cache)
return _to_response(updated_mapping)
except HTTPException:
@ -219,8 +221,8 @@ async def delete_jwt_key_mapping(
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
cache_key: Final = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
await _mapping_table(prisma_client).delete(where={"id": data.id})
return {"status": "success"}

View file

@ -54,6 +54,7 @@ from litellm.proxy._types import Litellm_EntityType, LiteLLM_VerificationToken,
from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
can_team_access_model,
get_jwt_key_mapping_cache_keys_for_token,
get_org_object,
get_project_object,
get_team_object,
@ -65,6 +66,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
evict_and_broadcast,
publish_auth_cache_invalidation,
)
from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error
@ -4975,6 +4977,13 @@ async def _execute_virtual_key_regeneration(
update_data.update(non_default_values)
jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
# Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash,
# but their cached jwt_key_mapping entries still point at the old token (LIT-5379).
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_token(
hashed_token=hashed_api_key,
prisma_client=prisma_client,
)
# If grace period set, insert deprecated key so old key remains valid
await _insert_deprecated_key(
prisma_client=prisma_client,
@ -5000,6 +5009,8 @@ async def _execute_virtual_key_regeneration(
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
# After credential invalidation, so a failure here can never keep the old key alive.
await sync_key_regeneration_access_group_membership(
prisma_client=prisma_client,

View file

@ -11912,6 +11912,109 @@ async def test_execute_virtual_key_regeneration_allows_within_limit_duration(mon
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1
@pytest.mark.asyncio
async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token():
"""
LIT-5379: /key/regenerate rewrites the JWT mapping row to the new token (FK
cascade) but left the jwt_key_mapping cache entry pointing at the old hash,
so JWT calls kept resolving the dead token until the cache TTL expired.
Regenerate must evict the entry locally, broadcast the eviction to other
workers, and the very next JWT resolve must return the rotated token.
"""
from litellm.caching.caching import DualCache
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.auth_method import AuthMethod
from litellm.proxy.auth.resolvers.models import CredentialRef
from litellm.proxy.auth.resolvers.store import IdentityStore
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
from litellm.proxy.management_endpoints.key_management_endpoints import (
_execute_virtual_key_regeneration,
)
stale_cache_key = "jwt_key_mapping:sub:user1"
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")]
)
mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
return_value=MagicMock(token="new-hashed-token")
)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(key=stale_cache_key, value="abc123")
publish_mock = AsyncMock()
with (
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
new_callable=AsyncMock,
return_value="sk-newtoken1234ab12",
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
publish_mock,
),
):
await _execute_virtual_key_regeneration(
prisma_client=mock_prisma_client,
key_in_db=existing_key,
hashed_api_key="abc123",
key="abc123",
data=None,
user_api_key_dict=_make_regenerate_user_api_key_dict(),
litellm_changed_by=None,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=MagicMock(),
)
assert await user_api_key_cache.async_get_cache(stale_cache_key) is None
publish_mock.assert_any_await(cache_key=stale_cache_key)
mock_prisma_client.db.litellm_jwtkeymapping.find_many.assert_awaited_once_with(where={"token": "abc123"})
rotated_key = UserAPIKeyAuth(token="new-hashed-token", user_id="user-1")
rotated_principal = IdentityStore._principal_from_key(
rotated_key,
auth_method=AuthMethod.API_KEY,
credential_ref=CredentialRef(token_id="new-hashed-token"),
)
async def fake_resolve(hashed_token):
assert hashed_token == "new-hashed-token", f"JWT resolved stale token {hashed_token!r} after regenerate"
return rotated_principal
jwt_handler = MagicMock()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=300
)
with patch(
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
new_callable=AsyncMock,
side_effect=fake_resolve,
):
resolved = await _resolve_jwt_to_virtual_key(
jwt_claims={"sub": "user1"},
jwt_handler=jwt_handler,
prisma_client=mock_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 == "new-hashed-token"
@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_rejects_over_limit_max_budget(monkeypatch):
"""Regenerate must reject max_budget exceeding upperbound — proves the fix covers non-duration fields."""