Merge pull request #33703 from BerriAI/litellm_fix_jwt_key_mapping_cascade_delete

fix(jwt): cascade-delete JWT key mappings when their virtual key is deleted

litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql

@@ -0,0 +1,15 @@
+-- DropForeignKey
+DO $$
+BEGIN
+    IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
+        ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey";
+    END IF;
+END $$;
+
+-- AddForeignKey
+DO $$
+BEGIN
+    IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
+        ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE;
+    END IF;
+END $$;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
     updated_at        DateTime @default(now()) @updatedAt
     updated_by        String?
 
-    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
+    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])

litellm/proxy/management_endpoints/key_management_endpoints.py

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

litellm/proxy/schema.prisma

@@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
     updated_at        DateTime @default(now()) @updatedAt
     updated_by        String?
 
-    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
+    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])

schema.prisma

@@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
     updated_at        DateTime @default(now()) @updatedAt
     updated_by        String?
 
-    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
+    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])

tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py

@@ -2,6 +2,7 @@ import glob
 import os
 import re
 import sys
+from pathlib import Path
 
 import pytest
 
@@ -870,3 +871,69 @@ class TestMigrateDeployAttemptAccounting:
             harness.run()
         assert len(harness.deploy_calls) == 1
         assert harness.resolved == []
+
+
+class TestJWTKeyMappingCascade:
+    """Regression tests for issue #33702.
+
+    A virtual key referenced by a LiteLLM_JWTKeyMapping row could not be deleted
+    because LiteLLM_JWTKeyMapping_token_fkey was created ON DELETE RESTRICT, so
+    deleting the key (Admin UI, /key/delete, team delete, ...) raised a foreign
+    key violation. The mapping must be removed automatically when its key is
+    deleted, which the FK now enforces via ON DELETE CASCADE.
+    """
+
+    _FK_NAME = "LiteLLM_JWTKeyMapping_token_fkey"
+
+    def _effective_on_delete(self):
+        """Replay every migration in order and return the last ON DELETE action
+        declared for the JWT key mapping FK."""
+        action = None
+        for _migration_name, sql in _get_all_migrations():
+            for match in re.finditer(
+                rf'ADD\s+CONSTRAINT\s+"{re.escape(self._FK_NAME)}".*?'
+                r"ON\s+DELETE\s+(CASCADE|RESTRICT|SET\s+NULL|NO\s+ACTION|SET\s+DEFAULT)",
+                sql,
+                re.IGNORECASE | re.DOTALL,
+            ):
+                action = re.sub(r"\s+", " ", match.group(1).upper())
+        return action
+
+    def test_fk_effective_on_delete_is_cascade(self):
+        """The final FK definition across all migrations must cascade deletes."""
+        assert self._effective_on_delete() == "CASCADE", (
+            f"{self._FK_NAME} must end up ON DELETE CASCADE so deleting a "
+            "virtual key removes its JWT key mapping (issue #33702)"
+        )
+
+    def test_schema_declares_cascade_on_relation(self):
+        """schema.prisma must declare onDelete: Cascade on the mapping relation
+        so the generated client and DB agree."""
+        schema_paths = glob.glob(
+            os.path.abspath(
+                os.path.join(
+                    os.path.dirname(__file__), "../../**/schema.prisma"
+                )
+            ),
+            recursive=True,
+        )
+        declaring = tuple(
+            (path, schema)
+            for path, schema in ((p, Path(p).read_text()) for p in schema_paths)
+            if "model LiteLLM_JWTKeyMapping" in schema
+        )
+        assert declaring, "No schema.prisma declaring LiteLLM_JWTKeyMapping found"
+        for path, schema in declaring:
+            match = re.search(
+                r"litellm_verification_token\s+LiteLLM_VerificationToken\s+@relation\(([^)]*)\)",
+                schema,
+            )
+            assert match is not None, (
+                f"{path} declares LiteLLM_JWTKeyMapping but its verification token "
+                "relation could not be parsed, so this test cannot vouch for it "
+                "(issue #33702)"
+            )
+            assert "onDelete: Cascade" in match.group(1), (
+                f"{path} must declare onDelete: Cascade on the JWT key mapping "
+                "relation (issue #33702)"
+            )

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

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