From eee239cd9e6b3b3e4bc3d3a515fa7d5160b51df9 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 9 Sep 2026 15:20:58 -0700 Subject: [PATCH] fix(jwt): cascade-delete JWT key mappings when their virtual key is deleted LiteLLM_JWTKeyMapping_token_fkey was created ON DELETE RESTRICT, so deleting a virtual key that a JWT mapping pointed at failed with a foreign key violation on every deletion path (/key/delete, Admin UI, alias delete, team and user cascades). Declaring onDelete: Cascade on the relation lets the database clean the mapping up uniformly, so the next JWT call from that identity re-registers against the newly created key. Rebase of #33703 onto current staging. Claude-Session: https://claude.ai/code/session_011Tn3657NkV6ojLqewL64Kb --- .../migration.sql | 15 +++++ .../litellm_proxy_extras/schema.prisma | 2 +- litellm/proxy/schema.prisma | 2 +- schema.prisma | 2 +- .../test_litellm_proxy_extras_utils.py | 67 +++++++++++++++++++ 5 files changed, 85 insertions(+), 3 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql new file mode 100644 index 00000000000..e5d48abcb52 --- /dev/null +++ b/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 $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/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]) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/litellm/proxy/schema.prisma +++ b/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]) diff --git a/schema.prisma b/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/schema.prisma +++ b/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]) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 3fab20a28ad..e5826b18668 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/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)" + )