Merge pull request #39808 from BerriAI/litellm_lit_5379_jwt_mapping_cache_invalidation

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

litellm/proxy/auth/auth_checks.py

@@ -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,

litellm/proxy/auth/user_api_key_auth.py

@@ -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:

litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py

@@ -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,20 @@ 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)
-
         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)
+        # 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)
+        new_cache_key: Final = jwt_key_mapping_cache_key(
+            updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
+        )
+        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)
 
         return _to_response(updated_mapping)
     except HTTPException:
@@ -219,10 +223,12 @@ 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)
-
         await _mapping_table(prisma_client).delete(where={"id": data.id})
+
+        # 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)
+        await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
         return {"status": "success"}
     except HTTPException:
         raise

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -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,

tests/proxy_unit_tests/test_jwt_key_mapping.py

@@ -1333,3 +1333,86 @@ def test_jwt_client_id_field_does_not_raise_on_duplicate():
         virtual_key_claim_field="new_field",
     )
     assert auth.virtual_key_claim_field == "new_field"
+
+
+# ──────────────────────────────────────────────
+# Tests: cache eviction must happen AFTER the DB write commits
+# ──────────────────────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_delete_evicts_cache_after_row_is_gone():
+    """A JWT request racing the delete must not keep the removed mapping authorized.
+
+    The DB delete simulates a concurrent request re-caching the mapping mid-write.
+    If the endpoint evicts before the delete commits, that repopulated entry
+    survives until TTL and the deleted mapping stays usable.
+    """
+    from litellm.proxy._types import DeleteJWTKeyMappingRequest
+    from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
+
+    cache_key = jwt_key_mapping_cache_key("email", "user@example.com")
+    user_api_key_cache = DualCache()
+    await user_api_key_cache.async_set_cache(key=cache_key, value="hashed_token")
+
+    mock_prisma = _mock_prisma()
+    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping()
+
+    async def concurrent_reader_repopulates(**kwargs):
+        await user_api_key_cache.async_set_cache(key=cache_key, value="hashed_token")
+        return _mock_mapping()
+
+    mock_prisma.db.litellm_jwtkeymapping.delete.side_effect = concurrent_reader_repopulates
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+    ):
+        result = await delete_jwt_key_mapping(
+            data=DeleteJWTKeyMappingRequest(id="mapping-1"),
+            user_api_key_dict=_make_admin_auth(),
+        )
+
+    assert result == {"status": "success"}
+    assert await user_api_key_cache.async_get_cache(cache_key) is None
+
+
+@pytest.mark.asyncio
+async def test_update_evicts_old_and_new_cache_keys_after_write():
+    """Renaming a mapping's claim must leave neither claim serving stale cache.
+
+    The DB update simulates a concurrent request re-caching the OLD mapping
+    mid-write. Both the old claim's entry (would restore the pre-rename token)
+    and the new claim's __NO_MAPPING__ sentinel (would 403 the renamed claim)
+    must be gone once the endpoint returns.
+    """
+    from litellm.proxy._types import UpdateJWTKeyMappingRequest
+    from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
+
+    old_cache_key = jwt_key_mapping_cache_key("email", "user@example.com")
+    new_cache_key = jwt_key_mapping_cache_key("email", "renamed@example.com")
+    user_api_key_cache = DualCache()
+    await user_api_key_cache.async_set_cache(key=old_cache_key, value="hashed_token")
+    await user_api_key_cache.async_set_cache(key=new_cache_key, value="__NO_MAPPING__")
+
+    mock_prisma = _mock_prisma()
+    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping()
+
+    async def concurrent_reader_repopulates(**kwargs):
+        await user_api_key_cache.async_set_cache(key=old_cache_key, value="hashed_token")
+        return _mock_mapping(claim_value="renamed@example.com")
+
+    mock_prisma.db.litellm_jwtkeymapping.update.side_effect = concurrent_reader_repopulates
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+    ):
+        result = await update_jwt_key_mapping(
+            data=UpdateJWTKeyMappingRequest(id="mapping-1", jwt_claim_value="renamed@example.com"),
+            user_api_key_dict=_make_admin_auth(),
+        )
+
+    assert result.jwt_claim_value == "renamed@example.com"
+    assert await user_api_key_cache.async_get_cache(old_cache_key) is None
+    assert await user_api_key_cache.async_get_cache(new_cache_key) is None

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -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(  # test-quality-ok: deterministic token; same pattern as sibling regenerate tests
+            "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
+            new_callable=AsyncMock,
+            return_value="sk-newtoken1234ab12",
+        ),
+        patch(  # test-quality-ok: grace-period path not under test; same pattern as sibling regenerate tests
+            "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
+            new_callable=AsyncMock,
+        ),
+        patch(  # test-quality-ok: key-object eviction is separate from the mapping eviction under test
+            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+            new_callable=AsyncMock,
+        ),
+        patch(  # test-quality-ok: background rotation hook is irrelevant to cache eviction
+            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
+            new_callable=AsyncMock,
+        ),
+        patch(  # test-quality-ok: captures the cross-worker broadcast without a redis instance
+            "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(  # test-quality-ok: DB-backed resolve; fake asserts it receives the rotated hash
+        "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."""