scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

fix(proxy): cascade budget alert claims from the key instead of sweeping call sites

Browse source 
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.
This commit is contained in:
Yucheng Zhu 2026-08-04 14:10:53 -07:00
parent 47401f29f8
commit 8bcc1a3bd4
9 changed files with 47 additions and 123 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
View file

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

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20260804000000_add_budget_alert_sent/migration.sql
View file

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

8
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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])
}

67
litellm/proxy/db/budget_alert_claim.py
View file

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

6
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

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

8
litellm/proxy/schema.prisma
View file

@ -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])
}

8
schema.prisma
View file

@ -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])
}

30
tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py
View file

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

8
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

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