From 8bcc1a3bd4fb352b89c114f8e7d2c8be7318d10f Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 4 Aug 2026 14:10:53 -0700 Subject: [PATCH] fix(proxy): cascade budget alert claims from the key instead of sweeping call sites Key deletion is not the only path that removes verification tokens; team deletion and the expired-session cleanup delete them directly, so an application-level sweep on /key/delete left claim rows orphaned. Key the claim table on the token with an ON DELETE CASCADE foreign key so every deletion path clears the rows with no application code. Verified live: deleting a key, and deleting a team that owns the key, both drop the claim rows. --- .../send_emails/base_email.py | 26 ++----- .../migration.sql | 9 ++- .../litellm_proxy_extras/schema.prisma | 8 ++- litellm/proxy/db/budget_alert_claim.py | 67 +++++-------------- .../key_management_endpoints.py | 6 -- litellm/proxy/schema.prisma | 8 ++- schema.prisma | 8 ++- .../send_emails/test_base_email.py | 30 +-------- .../test_key_management_endpoints.py | 8 --- 9 files changed, 47 insertions(+), 123 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index dbcd1a7042f..cd2e30d65fe 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -709,17 +709,6 @@ class BaseEmailLogger(CustomLogger): """ return f"{cache_key}:{budget_window}" - @staticmethod - def _durable_claim_entity(user_info: CallInfo) -> tuple[str, str] | None: - """ - (entity_type, entity_id) to key the durable claim on, or None when the alert - carries no stable identifier and can only be deduped in memory. - """ - entity_id = user_info.token or user_info.user_id - if not entity_id: - return None - return user_info.event_group.value, entity_id - async def _claim_max_budget_alert_send( self, cache: DualCache, @@ -746,14 +735,11 @@ class BaseEmailLogger(CustomLogger): if send_count is not None and send_count > 1: return False - entity = self._durable_claim_entity(user_info) - if entity is None: + if not user_info.token: return True - entity_type, entity_id = entity return await claim_budget_alert_slot( - entity_type=entity_type, - entity_id=entity_id, + token=user_info.token, alert_type=alert_type, threshold_pct=threshold_pct, budget_window=budget_window, @@ -773,13 +759,11 @@ class BaseEmailLogger(CustomLogger): cache, self._windowed_cache_key(cache_key, budget_window) ) - entity = self._durable_claim_entity(user_info) - if entity is None: + if not user_info.token: return - entity_type, entity_id = entity + await release_budget_alert_slot( - entity_type=entity_type, - entity_id=entity_id, + token=user_info.token, alert_type=alert_type, threshold_pct=threshold_pct, budget_window=budget_window, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804000000_add_budget_alert_sent/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804000000_add_budget_alert_sent/migration.sql index 958799cebf0..9bb0ed640f3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804000000_add_budget_alert_sent/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804000000_add_budget_alert_sent/migration.sql @@ -1,8 +1,7 @@ -- CreateTable CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetAlertSent" ( "id" TEXT NOT NULL, - "entity_type" TEXT NOT NULL, - "entity_id" TEXT NOT NULL, + "token" TEXT NOT NULL, "alert_type" TEXT NOT NULL, "threshold_pct" INTEGER NOT NULL, "budget_window" TEXT NOT NULL, @@ -12,4 +11,8 @@ CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetAlertSent" ( ); -- CreateIndex -CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_BudgetAlertSent_entity_type_entity_id_alert_type_th_key" ON "LiteLLM_BudgetAlertSent"("entity_type", "entity_id", "alert_type", "threshold_pct"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_BudgetAlertSent_token_alert_type_threshold_pct_key" ON "LiteLLM_BudgetAlertSent"("token", "alert_type", "threshold_pct"); + +-- AddForeignKey +ALTER TABLE "LiteLLM_BudgetAlertSent" DROP CONSTRAINT IF EXISTS "LiteLLM_BudgetAlertSent_token_fkey"; +ALTER TABLE "LiteLLM_BudgetAlertSent" ADD CONSTRAINT "LiteLLM_BudgetAlertSent_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c5f20e3b633..5077b203940 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -464,6 +464,7 @@ model LiteLLM_VerificationToken { litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) jwt_key_mappings LiteLLM_JWTKeyMapping[] + budget_alerts_sent LiteLLM_BudgetAlertSent[] // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 @@ -1470,12 +1471,13 @@ model LiteLLM_WorkflowMessage { model LiteLLM_BudgetAlertSent { id String @id @default(uuid()) - entity_type String - entity_id String + token String alert_type String threshold_pct Int budget_window String sent_at DateTime @default(now()) - @@unique([entity_type, entity_id, alert_type, threshold_pct]) + verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) + + @@unique([token, alert_type, threshold_pct]) } diff --git a/litellm/proxy/db/budget_alert_claim.py b/litellm/proxy/db/budget_alert_claim.py index 66f209af08a..9f808121669 100644 --- a/litellm/proxy/db/budget_alert_claim.py +++ b/litellm/proxy/db/budget_alert_claim.py @@ -7,14 +7,15 @@ claim and sends its own copy of the same alert, and the claim is lost on restart This module records the claim in the database instead, so it is shared by all replicas and survives restarts. -The claim is scoped to the entity's budget window: the budget period it belongs to -plus the budget it was measured against. The alert therefore re-arms when the budget -period rolls over or the budget itself is changed, rather than on a fixed TTL. +The claim is scoped to the key's budget window: the budget period it belongs to +plus the budget it was measured against. The alert therefore re-arms when the +budget period rolls over or the budget itself is changed, rather than on a fixed +TTL. Rows are removed by the schema's cascade when the key is deleted. """ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from datetime import datetime, timezone from litellm._logging import verbose_proxy_logger @@ -34,24 +35,17 @@ def get_budget_window(budget_reset_at: datetime | None, max_budget: float | None return f"{period}|{max_budget}" -def _claim_identity( - entity_type: str, - entity_id: str, - alert_type: str, - threshold_pct: int, -) -> Mapping[str, str | int]: +def _claim_identity(token: str, alert_type: str, threshold_pct: int) -> Mapping[str, str | int]: """The columns the unique constraint is built on, as a prisma query payload.""" return { # mutable-ok: prisma query payloads must be plain dicts - "entity_type": entity_type, - "entity_id": entity_id, + "token": token, "alert_type": alert_type, "threshold_pct": threshold_pct, } async def claim_budget_alert_slot( - entity_type: str, - entity_id: str, + token: str, alert_type: str, threshold_pct: int, budget_window: str, @@ -76,7 +70,7 @@ async def claim_budget_alert_slot( if prisma_client is None: return True - identity = _claim_identity(entity_type, entity_id, alert_type, threshold_pct) + identity = _claim_identity(token, alert_type, threshold_pct) try: await prisma_client.db.litellm_budgetalertsent.create( @@ -86,9 +80,8 @@ async def claim_budget_alert_slot( except Exception as e: # noqa: BLE001 # best-effort dedup, any failure must fall through to sending if not PrismaDBExceptionHandler.is_unique_constraint_violation(e): verbose_proxy_logger.warning( - "Could not record budget alert claim for %s %s at %d%%, sending anyway: %s", - entity_type, - entity_id, + "Could not record budget alert claim for key %s at %d%%, sending anyway: %s", + token, threshold_pct, e, ) @@ -105,9 +98,8 @@ async def claim_budget_alert_slot( return int(rows_updated) > 0 except Exception as e: # noqa: BLE001 # best-effort dedup, any failure must fall through to sending verbose_proxy_logger.warning( - "Could not roll over budget alert claim for %s %s at %d%%, sending anyway: %s", - entity_type, - entity_id, + "Could not roll over budget alert claim for key %s at %d%%, sending anyway: %s", + token, threshold_pct, e, ) @@ -115,8 +107,7 @@ async def claim_budget_alert_slot( async def release_budget_alert_slot( - entity_type: str, - entity_id: str, + token: str, alert_type: str, threshold_pct: int, budget_window: str, @@ -134,38 +125,16 @@ async def release_budget_alert_slot( if prisma_client is None: return - identity = _claim_identity(entity_type, entity_id, alert_type, threshold_pct) + identity = _claim_identity(token, alert_type, threshold_pct) try: await prisma_client.db.litellm_budgetalertsent.delete_many( where={**identity, "budget_window": budget_window} # mutable-ok: prisma query payload ) except Exception as e: # noqa: BLE001 # releasing the claim is best-effort, it also expires with the window - verbose_proxy_logger.debug( - "Failed to release budget alert claim for %s %s at %d%%: %s", - entity_type, - entity_id, + verbose_proxy_logger.warning( + "Failed to release budget alert claim for key %s at %d%%: %s", + token, threshold_pct, e, ) - - -async def delete_budget_alert_claims(entity_type: str, entity_ids: Sequence[str]) -> None: - """ - Drop the claim rows for entities that no longer exist, so deleting a key does - not leave its alert claims behind. - """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None or not entity_ids: - return - - try: - await prisma_client.db.litellm_budgetalertsent.delete_many( - where={ # mutable-ok: prisma query payload - "entity_type": entity_type, - "entity_id": {"in": tuple(entity_ids)}, # mutable-ok: prisma query payload - } - ) - except Exception as e: # noqa: BLE001 # cleanup is best-effort, it must never fail a deletion - verbose_proxy_logger.warning("Failed to delete budget alert claims for %s %s: %s", entity_type, entity_ids, e) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3cc56eaf480..f37bc6da23e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -68,7 +68,6 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.db.budget_alert_claim import delete_budget_alert_claims from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import ( VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, @@ -4121,11 +4120,6 @@ async def delete_verification_tokens( verbose_proxy_logger.debug(traceback.format_exc()) raise e - await delete_budget_alert_claims( - entity_type=Litellm_EntityType.KEY.value, - entity_ids=tuple(key.token for key in _keys_being_deleted), - ) - for key in tokens: user_api_key_cache.delete_cache(key) # remove hash token from cache diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c5f20e3b633..5077b203940 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -464,6 +464,7 @@ model LiteLLM_VerificationToken { litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) jwt_key_mappings LiteLLM_JWTKeyMapping[] + budget_alerts_sent LiteLLM_BudgetAlertSent[] // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 @@ -1470,12 +1471,13 @@ model LiteLLM_WorkflowMessage { model LiteLLM_BudgetAlertSent { id String @id @default(uuid()) - entity_type String - entity_id String + token String alert_type String threshold_pct Int budget_window String sent_at DateTime @default(now()) - @@unique([entity_type, entity_id, alert_type, threshold_pct]) + verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) + + @@unique([token, alert_type, threshold_pct]) } diff --git a/schema.prisma b/schema.prisma index c5f20e3b633..5077b203940 100644 --- a/schema.prisma +++ b/schema.prisma @@ -464,6 +464,7 @@ model LiteLLM_VerificationToken { litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) jwt_key_mappings LiteLLM_JWTKeyMapping[] + budget_alerts_sent LiteLLM_BudgetAlertSent[] // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 @@ -1470,12 +1471,13 @@ model LiteLLM_WorkflowMessage { model LiteLLM_BudgetAlertSent { id String @id @default(uuid()) - entity_type String - entity_id String + token String alert_type String threshold_pct Int budget_window String sent_at DateTime @default(now()) - @@unique([entity_type, entity_id, alert_type, threshold_pct]) + verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) + + @@unique([token, alert_type, threshold_pct]) } diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 5548b4cecdf..387d3c4100c 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -29,7 +29,6 @@ from litellm.constants import EMAIL_BUDGET_ALERT_TTL from litellm.constants import MAX_BUDGET_ALERT_TYPE from litellm.proxy.db.budget_alert_claim import ( claim_budget_alert_slot, - delete_budget_alert_claims, release_budget_alert_slot, ) @@ -1397,7 +1396,7 @@ class _FakeBudgetAlertTable: one replica is visible to every other replica sharing this instance. """ - _UNIQUE_FIELDS = ("entity_type", "entity_id", "alert_type", "threshold_pct") + _UNIQUE_FIELDS = ("token", "alert_type", "threshold_pct") def __init__(self): self.rows: list[dict] = [] @@ -1414,9 +1413,6 @@ class _FakeBudgetAlertTable: if isinstance(expected, dict) and "not" in expected: if actual == expected["not"]: break - elif isinstance(expected, dict) and "in" in expected: - if actual not in expected["in"]: - break elif actual != expected: break else: @@ -1650,16 +1646,14 @@ async def test_release_does_not_steal_a_claim_from_a_later_window(shared_alert_t """A slow failing send must only release the claim it actually took, or it deletes the row another replica has already taken over for a later window.""" await claim_budget_alert_slot( - entity_type="key", - entity_id="hashed_key_1", + token="hashed_key_1", alert_type=MAX_BUDGET_ALERT_TYPE, threshold_pct=50, budget_window="later-window", ) await release_budget_alert_slot( - entity_type="key", - entity_id="hashed_key_1", + token="hashed_key_1", alert_type=MAX_BUDGET_ALERT_TYPE, threshold_pct=50, budget_window="earlier-window", @@ -1738,24 +1732,6 @@ async def test_max_budget_alert_sends_when_claim_table_is_unavailable(monkeypatc assert len(sends) == 2 -@pytest.mark.asyncio -async def test_deleting_a_key_drops_its_claims_only(shared_alert_table): - """Claim rows outlive the key unless deletion sweeps them, and a stale row for a - recycled id would suppress a real alert.""" - for entity_id in ("hashed_key_1", "hashed_key_2"): - await claim_budget_alert_slot( - entity_type="key", - entity_id=entity_id, - alert_type=MAX_BUDGET_ALERT_TYPE, - threshold_pct=50, - budget_window="|100.0", - ) - - await delete_budget_alert_claims(entity_type="key", entity_ids=("hashed_key_1",)) - - assert [r["entity_id"] for r in shared_alert_table.rows] == ["hashed_key_2"] - - @pytest.mark.asyncio async def test_in_memory_claim_does_not_outlive_its_budget_window(shared_alert_table): """One long-lived replica must re-alert after the window moves. The in-memory diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 15af6be3b59..867ef759fb3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -4923,14 +4923,6 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch): # delete_data returns the list directly, which gets wrapped in {"deleted_keys": ...} assert isinstance(result["deleted_keys"], list) assert set(result["deleted_keys"]) == {"hashed-token-1", "hashed-token-2"} - - # budget alert claim rows are keyed on the token and have no FK, so deleting a key - # has to sweep them or a stale claim silences the alert for a recycled token - claim_delete = mock_prisma_client.db.litellm_budgetalertsent.delete_many - claim_delete.assert_awaited_once() - claim_where = claim_delete.await_args.kwargs["where"] - assert claim_where["entity_type"] == "key" - assert set(claim_where["entity_id"]["in"]) == {"hashed-token-1", "hashed-token-2"} assert len(deleted_keys) == 2