scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 77

fix(proxy): dedupe budget alert emails across replicas without redis

Browse source 
The dedup claim added in #32011 lives in DualCache, which is memory-only
when Redis is not configured, so each replica wins its own claim and sends
its own copy of the same alert. It is also lost on restart and expires with
EMAIL_BUDGET_ALERT_TTL, so a long-crossed threshold keeps re-alerting.

Record the claim in a new LiteLLM_BudgetAlertSent table instead. The unique
constraint makes the insert an atomic cross-replica claim, and the stored
budget window re-arms the alert when the budget period resets or max_budget
changes rather than on a timer. Falls open when the claim cannot be recorded.

Resolves LIT-4172
This commit is contained in:
Yucheng Zhu 2026-08-04 13:27:31 -07:00
parent 8445cf158b
commit 47401f29f8
16 changed files with 864 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -20,6 +20,8 @@ from litellm.caching.caching import DualCache
from litellm.constants import (
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
EMAIL_BUDGET_ALERT_TTL,
MAX_BUDGET_ALERT_TYPE,
MAX_BUDGET_DEFAULT_ALERT_TYPE,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER
@ -44,6 +46,11 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
WebhookEvent,
)
from litellm.proxy.db.budget_alert_claim import (
claim_budget_alert_slot,
get_budget_window,
release_budget_alert_slot,
)
from litellm.secret_managers.main import get_secret_bool
from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL
@ -542,16 +549,18 @@ class BaseEmailLogger(CustomLogger):
_id = user_info.token or user_info.user_id or "default_id"
_cache_key = f"email_budget_alerts:max_budget_alert:{_id}"
send_count = await _cache.async_increment_cache(
key=_cache_key,
value=1,
ttl=EMAIL_BUDGET_ALERT_TTL,
# Calculate percentage
percentage = int(
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100
)
if send_count is None or send_count <= 1:
# Calculate percentage
percentage = int(
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100
)
if await self._claim_max_budget_alert_send(
cache=_cache,
cache_key=_cache_key,
user_info=user_info,
threshold_pct=percentage,
alert_type=MAX_BUDGET_DEFAULT_ALERT_TYPE,
):
# Create WebhookEvent for max budget alert
event_message = f"Max Budget Alert - {percentage}% of Maximum Budget Reached"
@ -572,6 +581,7 @@ class BaseEmailLogger(CustomLogger):
projected_exceeded_date=user_info.projected_exceeded_date,
projected_spend=user_info.projected_spend,
event_group=user_info.event_group,
budget_reset_at=user_info.budget_reset_at,
)
try:
@ -581,7 +591,13 @@ class BaseEmailLogger(CustomLogger):
f"Error sending max budget alert email: {e}",
exc_info=True,
)
await self._release_budget_alert_claim(_cache, _cache_key)
await self._release_max_budget_alert_claim(
cache=_cache,
cache_key=_cache_key,
user_info=user_info,
threshold_pct=percentage,
alert_type=MAX_BUDGET_DEFAULT_ALERT_TYPE,
)
return
async def _handle_multi_threshold_max_budget_alert(
@ -624,12 +640,12 @@ class BaseEmailLogger(CustomLogger):
continue
recipient_emails = list(set(emails))
send_count = await _cache.async_increment_cache(
key=_cache_key,
value=1,
ttl=EMAIL_BUDGET_ALERT_TTL,
)
if send_count is not None and send_count > 1:
if not await self._claim_max_budget_alert_send(
cache=_cache,
cache_key=_cache_key,
user_info=user_info,
threshold_pct=threshold_pct,
):
continue
event_message = f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached"
@ -650,6 +666,7 @@ class BaseEmailLogger(CustomLogger):
projected_exceeded_date=user_info.projected_exceeded_date,
projected_spend=user_info.projected_spend,
event_group=user_info.event_group,
budget_reset_at=user_info.budget_reset_at,
)
try:
@ -663,7 +680,12 @@ class BaseEmailLogger(CustomLogger):
f"Error sending multi-threshold max budget alert email for {threshold_pct}%: {e}",
exc_info=True,
)
await self._release_budget_alert_claim(_cache, _cache_key)
await self._release_max_budget_alert_claim(
cache=_cache,
cache_key=_cache_key,
user_info=user_info,
threshold_pct=threshold_pct,
)
async def _release_budget_alert_claim(self, cache: DualCache, cache_key: str) -> None:
try:
@ -674,6 +696,95 @@ class BaseEmailLogger(CustomLogger):
cache_key,
)
@staticmethod
def _budget_window(user_info: CallInfo) -> str:
return get_budget_window(user_info.budget_reset_at, user_info.max_budget)
@staticmethod
def _windowed_cache_key(cache_key: str, budget_window: str) -> str:
"""
Scope the in-memory claim to the budget window it belongs to. Without this
the claim would still suppress the alert for up to EMAIL_BUDGET_ALERT_TTL
after the budget resets, even though the durable claim has re-armed.
"""
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,
cache_key: str,
user_info: CallInfo,
threshold_pct: int,
alert_type: str = MAX_BUDGET_ALERT_TYPE,
) -> bool:
"""
Decide whether this process should send one max budget alert.
The in-memory claim stops a replica re-sending on every request it serves.
The durable claim is what stops a SECOND replica sending the same alert, and
what keeps the alert from firing again after a restart or a cache TTL expiry:
the in-memory claim is process-local whenever Redis is not configured, so on
its own it cannot dedupe across replicas or across restarts.
"""
budget_window = self._budget_window(user_info)
send_count = await cache.async_increment_cache(
key=self._windowed_cache_key(cache_key, budget_window),
value=1,
ttl=EMAIL_BUDGET_ALERT_TTL,
)
if send_count is not None and send_count > 1:
return False
entity = self._durable_claim_entity(user_info)
if entity is None:
return True
entity_type, entity_id = entity
return await claim_budget_alert_slot(
entity_type=entity_type,
entity_id=entity_id,
alert_type=alert_type,
threshold_pct=threshold_pct,
budget_window=budget_window,
)
async def _release_max_budget_alert_claim(
self,
cache: DualCache,
cache_key: str,
user_info: CallInfo,
threshold_pct: int,
alert_type: str = MAX_BUDGET_ALERT_TYPE,
) -> None:
"""Hand a won claim back after a failed send, so the alert can be retried."""
budget_window = self._budget_window(user_info)
await self._release_budget_alert_claim(
cache, self._windowed_cache_key(cache_key, budget_window)
)
entity = self._durable_claim_entity(user_info)
if entity is None:
return
entity_type, entity_id = entity
await release_budget_alert_slot(
entity_type=entity_type,
entity_id=entity_id,
alert_type=alert_type,
threshold_pct=threshold_pct,
budget_window=budget_window,
)
async def _get_email_params(
self,
email_event: EmailEvent,

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

@ -0,0 +1,15 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetAlertSent" (
"id" TEXT NOT NULL,
"entity_type" TEXT NOT NULL,
"entity_id" TEXT NOT NULL,
"alert_type" TEXT NOT NULL,
"threshold_pct" INTEGER NOT NULL,
"budget_window" TEXT NOT NULL,
"sent_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_BudgetAlertSent_pkey" PRIMARY KEY ("id")
);
-- 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");

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

@ -1467,3 +1467,15 @@ model LiteLLM_WorkflowMessage {
@@unique([run_id, sequence_number])
@@index([run_id])
}
model LiteLLM_BudgetAlertSent {
id String @id @default(uuid())
entity_type String
entity_id String
alert_type String
threshold_pct Int
budget_window String
sent_at DateTime @default(now())
@@unique([entity_type, entity_id, alert_type, threshold_pct])
}

2
litellm/constants.py
View file

@ -464,6 +464,8 @@ EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60))
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(
os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)
) # 80% of max budget
MAX_BUDGET_ALERT_TYPE = "max_budget_alert"
MAX_BUDGET_DEFAULT_ALERT_TYPE = "max_budget_alert_default_threshold"
############### LLM Provider Constants ###############
### ANTHROPIC CONSTANTS ###
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv("ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01")

1
litellm/integrations/SlackAlerting/slack_alerting.py
View file

@ -631,6 +631,7 @@ class SlackAlerting(CustomBatchLogger):
"""
_all_fields_as_dict = user_info.model_dump(exclude_none=True)
_all_fields_as_dict.pop("token")
_all_fields_as_dict.pop("budget_reset_at", None)
msg = ""
for k, v in _all_fields_as_dict.items():
if isinstance(v, Litellm_EntityType):

4
litellm/proxy/_types.py
View file

@ -3092,6 +3092,10 @@ class CallInfo(LiteLLMPydanticObjectBase):
default=None,
description="Map of threshold percentage to email recipients (e.g., {'50': ['a@co.com'], '75': ['a@co.com', 'b@co.com']})",
)
budget_reset_at: datetime | None = Field(
default=None,
description="End of the budget period this alert belongs to. Used to re-arm the alert once the budget rolls over.",
)
class WebhookEvent(CallInfo):

2
litellm/proxy/auth/auth_checks.py
View file

@ -3846,6 +3846,7 @@ async def _virtual_key_max_budget_alert_check(
key_alias=valid_token.key_alias,
event_group=Litellm_EntityType.KEY,
max_budget_alert_emails=alert_email_config,
budget_reset_at=valid_token.budget_reset_at,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
@ -3877,6 +3878,7 @@ async def _virtual_key_max_budget_alert_check(
user_email=owner_email,
key_alias=valid_token.key_alias,
event_group=Litellm_EntityType.KEY,
budget_reset_at=valid_token.budget_reset_at,
)
asyncio.create_task(

171
litellm/proxy/db/budget_alert_claim.py Normal file
View file

@ -0,0 +1,171 @@
"""
Durable, cross-replica claim for "this budget alert has already been sent".
The in-process claim used by the email alerting path lives in ``DualCache``, which
is memory-only when Redis is not configured. Every replica therefore wins its own
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.
"""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
def get_budget_window(budget_reset_at: datetime | None, max_budget: float | None) -> str:
"""
Stable identity for the budget an alert was raised against.
``budget_reset_at`` alone is not enough: a key with no ``budget_duration`` never
rolls over, so a claim keyed on it alone would silence that threshold for the
lifetime of the key. Including ``max_budget`` means raising or lowering the budget
moves every threshold and re-arms the alerts, which is what an operator expects.
"""
period = budget_reset_at.isoformat() if budget_reset_at is not None else ""
return f"{period}|{max_budget}"
def _claim_identity(
entity_type: str,
entity_id: 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,
"alert_type": alert_type,
"threshold_pct": threshold_pct,
}
async def claim_budget_alert_slot(
entity_type: str,
entity_id: str,
alert_type: str,
threshold_pct: int,
budget_window: str,
) -> bool:
"""
Try to become the single sender of one budget alert.
Returns True iff this caller won the claim and must send the alert, and False
iff another replica, or an earlier request in this budget window, already sent
it.
Inserting the claim row is the atomic part: the unique constraint lets exactly
one replica succeed. A row that already exists is taken over only when it
belongs to an older budget window, which is the conditional update below.
Fails open. If there is no database, or the database rejects the claim for any
reason other than the alert already being claimed, this returns True, because a
duplicate alert is a better failure than a missed budget alert.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return True
identity = _claim_identity(entity_type, entity_id, alert_type, threshold_pct)
try:
await prisma_client.db.litellm_budgetalertsent.create(
data={**identity, "budget_window": budget_window} # mutable-ok: prisma query payload
)
return True
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,
threshold_pct,
e,
)
return True
try:
rows_updated = await prisma_client.db.litellm_budgetalertsent.update_many(
where={**identity, "budget_window": {"not": budget_window}}, # mutable-ok: prisma query payload
data={ # mutable-ok: prisma query payload
"budget_window": budget_window,
"sent_at": datetime.now(timezone.utc),
},
)
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,
threshold_pct,
e,
)
return True
async def release_budget_alert_slot(
entity_type: str,
entity_id: str,
alert_type: str,
threshold_pct: int,
budget_window: str,
) -> None:
"""
Give a won claim back after the alert failed to send, so the next request can
retry it. Without this, one failed send would silence the whole fleet for the
rest of the budget window.
Scoped to ``budget_window`` so a slow failing send cannot delete a claim that
another replica has already taken over for a later window.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return
identity = _claim_identity(entity_type, entity_id, 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,
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)

14
litellm/proxy/db/exception_handler.py
View file

@ -92,6 +92,20 @@ class PrismaDBExceptionHandler:
return type(e) is prisma.errors.DataError
@staticmethod
def is_unique_constraint_violation(e: Exception) -> bool:
"""True iff ``e`` is a prisma unique-constraint violation, i.e. the row the
caller tried to insert already exists.
Lets callers use an insert as an atomic cross-replica claim without
importing prisma themselves: prisma is a proxy-only dependency, and a
module reachable from a base ``import litellm`` cannot import it at the top
level.
"""
import prisma
return isinstance(e, prisma.errors.UniqueViolationError)
@staticmethod
def is_database_transport_error(e: Exception) -> bool:
"""

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

@ -68,6 +68,7 @@ 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,
@ -4120,6 +4121,11 @@ 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

12
litellm/proxy/schema.prisma
View file

@ -1467,3 +1467,15 @@ model LiteLLM_WorkflowMessage {
@@unique([run_id, sequence_number])
@@index([run_id])
}
model LiteLLM_BudgetAlertSent {
id String @id @default(uuid())
entity_type String
entity_id String
alert_type String
threshold_pct Int
budget_window String
sent_at DateTime @default(now())
@@unique([entity_type, entity_id, alert_type, threshold_pct])
}

12
schema.prisma
View file

@ -1467,3 +1467,15 @@ model LiteLLM_WorkflowMessage {
@@unique([run_id, sequence_number])
@@index([run_id])
}
model LiteLLM_BudgetAlertSent {
id String @id @default(uuid())
entity_type String
entity_id String
alert_type String
threshold_pct Int
budget_window String
sent_at DateTime @default(now())
@@unique([entity_type, entity_id, alert_type, threshold_pct])
}

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

@ -3,6 +3,9 @@ import json
import os
import sys
import unittest.mock as mock
from collections.abc import Mapping
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import patch
import pytest
@ -23,6 +26,12 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import (
from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER
from litellm.proxy._types import CallInfo, Litellm_EntityType, WebhookEvent
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,
)
@pytest.fixture(autouse=True)
@ -914,7 +923,8 @@ async def test_budget_alerts_max_budget_alert_crossed(
mock_cache.async_increment_cache.assert_called_once()
cache_call_args = mock_cache.async_increment_cache.call_args[1]
assert (
cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user"
cache_call_args["key"]
== "email_budget_alerts:max_budget_alert:test_user:|200.0"
)
assert cache_call_args["value"] == 1
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
@ -955,8 +965,8 @@ async def test_multi_threshold_sends_crossed_thresholds(
cache_keys = [
c[1]["key"] for c in mock_cache.async_increment_cache.call_args_list
]
assert "email_budget_alerts:max_budget_alert:50:hashed_key_1" in cache_keys
assert "email_budget_alerts:max_budget_alert:75:hashed_key_1" in cache_keys
assert "email_budget_alerts:max_budget_alert:50:hashed_key_1:|100.0" in cache_keys
assert "email_budget_alerts:max_budget_alert:75:hashed_key_1:|100.0" in cache_keys
@pytest.mark.asyncio
@ -1119,7 +1129,7 @@ async def test_no_map_preserves_old_single_threshold(
assert call_args["to_email"] == ["test@example.com"]
# Old path cache key has no threshold percentage
cache_key = mock_cache.async_increment_cache.call_args[1]["key"]
assert cache_key == "email_budget_alerts:max_budget_alert:test_user"
assert cache_key == "email_budget_alerts:max_budget_alert:test_user:|200.0"
CUSTOM_SIGNATURE = "<div>Best,<br/>The Acme Platform Team</div>"
@ -1378,3 +1388,439 @@ async def test_budget_alert_release_failure_does_not_propagate(base_email_logger
)
mock_cache.async_delete_cache.assert_awaited_once()
class _FakeBudgetAlertTable:
"""In-memory stand-in for the LiteLLM_BudgetAlertSent prisma table.
Enforces the same unique constraint the real table does, so a claim taken by
one replica is visible to every other replica sharing this instance.
"""
_UNIQUE_FIELDS = ("entity_type", "entity_id", "alert_type", "threshold_pct")
def __init__(self):
self.rows: list[dict] = []
@staticmethod
def _identity(row: Mapping[str, object]) -> tuple:
return tuple(row[f] for f in _FakeBudgetAlertTable._UNIQUE_FIELDS)
def _matching(self, where: Mapping[str, object]) -> list[dict]:
matched = []
for row in self.rows:
for field, expected in where.items():
actual = row.get(field)
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:
matched.append(row)
return matched
async def create(self, data: Mapping[str, object]) -> dict:
import prisma
if any(self._identity(r) == self._identity(data) for r in self.rows):
raise prisma.errors.UniqueViolationError({"user_facing_error": {"meta": {}}})
row = dict(data)
self.rows.append(row)
return row
async def update_many(
self, where: Mapping[str, object], data: Mapping[str, object]
) -> int:
matched = self._matching(where)
for row in matched:
row.update(data)
return len(matched)
async def delete_many(self, where: Mapping[str, object]) -> int:
matched = self._matching(where)
self.rows = [r for r in self.rows if r not in matched]
return len(matched)
@pytest.fixture
def shared_alert_table(monkeypatch):
"""A single durable claim store shared by every simulated replica."""
import litellm.proxy.proxy_server as proxy_server
table = _FakeBudgetAlertTable()
fake_client = SimpleNamespace(db=SimpleNamespace(litellm_budgetalertsent=table))
monkeypatch.setattr(proxy_server, "prisma_client", fake_client, raising=False)
return table
def _replica_logger() -> BaseEmailLogger:
"""A logger with its OWN DualCache, i.e. a separate proxy process with no Redis."""
logger = BaseEmailLogger()
logger.internal_usage_cache = DualCache()
return logger
@pytest.mark.parametrize(
"branch, alert_type, send_method, ci_kwargs",
[b for b in _BUDGET_ALERT_BRANCHES if b[1] == "max_budget_alert"],
ids=[b[0] for b in _BUDGET_ALERT_BRANCHES if b[1] == "max_budget_alert"],
)
@pytest.mark.asyncio
async def test_max_budget_alert_deduped_across_replicas_without_redis(
shared_alert_table, branch, alert_type, send_method, ci_kwargs
):
"""Regression for LIT-4172 follow-up: two replicas with no Redis must send
exactly one email for one threshold crossing.
The in-memory claim added by the first fix is process-local, so each replica
won its own claim and sent its own copy. The durable claim is what makes the
two replicas agree.
"""
sends = []
async def record_send(*args, **kwargs):
sends.append(kwargs.get("threshold_pct"))
for _ in range(2):
replica = _replica_logger()
with mock.patch.object(replica, send_method, side_effect=record_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
)
assert len(sends) == 1
assert len(shared_alert_table.rows) == 1
@pytest.mark.asyncio
async def test_max_budget_alert_not_resent_after_replica_restart(shared_alert_table):
"""A restart wipes the in-memory claim. The alert must not fire again, which
is what made a long-crossed threshold re-alert every restart and every TTL."""
sends = []
async def record_send(*args, **kwargs):
sends.append(1)
user_info = _budget_alert_user_info(
dict(max_budget=100.0, spend=80.0, max_budget_alert_emails={"50": []})
)
for _ in range(3):
replica = _replica_logger()
with mock.patch.object(
replica, "send_max_budget_alert_email", side_effect=record_send
):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(
type="max_budget_alert", user_info=user_info
)
assert len(sends) == 1
@pytest.mark.parametrize(
"branch, alert_type, send_method, ci_kwargs",
[b for b in _BUDGET_ALERT_BRANCHES if b[1] == "max_budget_alert"],
ids=[b[0] for b in _BUDGET_ALERT_BRANCHES if b[1] == "max_budget_alert"],
)
@pytest.mark.asyncio
async def test_max_budget_alert_rearms_when_budget_window_rolls_over(
shared_alert_table, branch, alert_type, send_method, ci_kwargs
):
"""A durable claim must not silence the alert forever: once the key's budget
period resets, the same threshold has to alert again."""
sends = []
async def record_send(*args, **kwargs):
sends.append(1)
def user_info_for(reset_at):
return _budget_alert_user_info({**ci_kwargs, "budget_reset_at": reset_at})
same_window = datetime(2026, 9, 1, tzinfo=timezone.utc)
next_window = datetime(2026, 10, 1, tzinfo=timezone.utc)
for window in (same_window, same_window, next_window):
replica = _replica_logger()
with mock.patch.object(replica, send_method, side_effect=record_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(
type=alert_type, user_info=user_info_for(window)
)
assert len(sends) == 2
assert len(shared_alert_table.rows) == 1
@pytest.mark.parametrize(
"branch, alert_type, send_method, ci_kwargs",
[b for b in _BUDGET_ALERT_BRANCHES if b[1] == "max_budget_alert"],
ids=[b[0] for b in _BUDGET_ALERT_BRANCHES if b[1] == "max_budget_alert"],
)
@pytest.mark.asyncio
async def test_max_budget_alert_failed_send_releases_durable_claim(
shared_alert_table, branch, alert_type, send_method, ci_kwargs
):
"""A send that fails must give the durable claim back, or the whole fleet is
permanently silenced for that threshold."""
attempts = []
async def flaky_send(*args, **kwargs):
attempts.append(1)
if len(attempts) == 1:
raise ValueError("transient email backend failure")
user_info = _budget_alert_user_info(ci_kwargs)
for _ in range(2):
replica = _replica_logger()
with mock.patch.object(replica, send_method, side_effect=flaky_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(type=alert_type, user_info=user_info)
assert len(attempts) == 2
assert len(shared_alert_table.rows) == 1
@pytest.mark.asyncio
async def test_max_budget_alert_sends_when_no_database_is_configured(monkeypatch):
"""With no prisma client the durable claim cannot be taken. The alert must
still be sent rather than silently dropped."""
import litellm.proxy.proxy_server as proxy_server
monkeypatch.setattr(proxy_server, "prisma_client", None, raising=False)
sends = []
async def record_send(*args, **kwargs):
sends.append(1)
replica = _replica_logger()
with mock.patch.object(
replica, "send_max_budget_alert_email", side_effect=record_send
):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(
type="max_budget_alert",
user_info=_budget_alert_user_info(
dict(max_budget=100.0, spend=80.0, max_budget_alert_emails={"50": []})
),
)
assert len(sends) == 1
@pytest.mark.parametrize(
"branch, alert_type, send_method, ci_kwargs",
[b for b in _BUDGET_ALERT_BRANCHES if b[1] == "max_budget_alert"],
ids=[b[0] for b in _BUDGET_ALERT_BRANCHES if b[1] == "max_budget_alert"],
)
@pytest.mark.asyncio
async def test_max_budget_alert_rearms_when_max_budget_changes(
shared_alert_table, branch, alert_type, send_method, ci_kwargs
):
"""A key with no budget_duration never rolls over, so a claim keyed only on the
reset date would silence the threshold for the life of the key. Changing the
budget moves every threshold amount and has to re-arm the alerts."""
sends = []
async def record_send(*args, **kwargs):
sends.append(1)
def user_info_for(max_budget: float):
return _budget_alert_user_info(
{**ci_kwargs, "max_budget": max_budget, "budget_reset_at": None}
)
for max_budget in (100.0, 100.0, 90.0):
replica = _replica_logger()
with mock.patch.object(replica, send_method, side_effect=record_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(
type=alert_type, user_info=user_info_for(max_budget)
)
assert len(sends) == 2
assert len(shared_alert_table.rows) == 1
@pytest.mark.asyncio
async def test_release_does_not_steal_a_claim_from_a_later_window(shared_alert_table):
"""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",
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",
alert_type=MAX_BUDGET_ALERT_TYPE,
threshold_pct=50,
budget_window="earlier-window",
)
assert len(shared_alert_table.rows) == 1
@pytest.mark.asyncio
async def test_each_configured_threshold_claims_its_own_slot(shared_alert_table):
"""The claim is per threshold. A key crossing two thresholds at once must send
both emails; collapsing them onto one claim silently drops the higher one."""
sends = []
async def record_send(*args, **kwargs):
sends.append(kwargs.get("threshold_pct"))
user_info = _budget_alert_user_info(
dict(
max_budget=100.0,
spend=80.0,
max_budget_alert_emails={"50": ["finance@co.com"], "75": ["finance@co.com"]},
)
)
replica = _replica_logger()
with mock.patch.object(
replica, "send_max_budget_alert_email", side_effect=record_send
):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(type="max_budget_alert", user_info=user_info)
assert sorted(sends) == [50, 75]
assert len(shared_alert_table.rows) == 2
@pytest.mark.asyncio
async def test_max_budget_alert_sends_when_claim_table_is_unavailable(monkeypatch):
"""Day one of an upgrade the migration may not have run yet. A database error
that is not "already claimed" must fall through to sending, never silence."""
import prisma
import litellm.proxy.proxy_server as proxy_server
class _BrokenTable:
async def create(self, data):
raise prisma.errors.TableNotFoundError({"user_facing_error": {"meta": {}}})
monkeypatch.setattr(
proxy_server,
"prisma_client",
SimpleNamespace(db=SimpleNamespace(litellm_budgetalertsent=_BrokenTable())),
raising=False,
)
sends = []
async def record_send(*args, **kwargs):
sends.append(1)
for _ in range(2):
replica = _replica_logger()
with mock.patch.object(
replica, "send_max_budget_alert_email", side_effect=record_send
):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(
type="max_budget_alert",
user_info=_budget_alert_user_info(
dict(
max_budget=100.0,
spend=80.0,
max_budget_alert_emails={"50": []},
)
),
)
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
claim outlives the window by up to EMAIL_BUDGET_ALERT_TTL unless it is scoped
to the window too."""
sends = []
async def record_send(*args, **kwargs):
sends.append(1)
replica = _replica_logger()
with mock.patch.object(
replica, "send_max_budget_alert_email", side_effect=record_send
):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
for reset_at in (
datetime(2026, 9, 1, tzinfo=timezone.utc),
datetime(2026, 9, 1, tzinfo=timezone.utc),
datetime(2026, 10, 1, tzinfo=timezone.utc),
):
await replica.budget_alerts(
type="max_budget_alert",
user_info=_budget_alert_user_info(
dict(
max_budget=100.0,
spend=80.0,
max_budget_alert_emails={"50": []},
budget_reset_at=reset_at,
)
),
)
assert len(sends) == 2
@pytest.mark.asyncio
async def test_default_threshold_and_configured_threshold_do_not_share_a_claim(
shared_alert_table,
):
"""The 80% default alert and a configured "80" threshold are different alerts
with different recipients; one must not consume the other's claim."""
sends = []
async def record_send(*args, **kwargs):
sends.append(kwargs.get("threshold_pct"))
replica = _replica_logger()
with mock.patch.object(
replica, "send_max_budget_alert_email", side_effect=record_send
):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await replica.budget_alerts(
type="max_budget_alert",
user_info=_budget_alert_user_info(dict(max_budget=100.0, spend=85.0)),
)
await replica.budget_alerts(
type="max_budget_alert",
user_info=_budget_alert_user_info(
dict(
max_budget=100.0,
spend=85.0,
max_budget_alert_emails={"80": ["finance@co.com"]},
)
),
)
assert len(sends) == 2
assert len(shared_alert_table.rows) == 2

4
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -2624,6 +2624,7 @@ async def test_virtual_key_max_budget_alert_check_with_user_obj():
org_id="test-org",
key_alias="test-key",
soft_budget=50.0,
budget_reset_at=datetime(2026, 9, 1, tzinfo=timezone.utc),
)
user_obj = LiteLLM_UserTable(
@ -2655,6 +2656,7 @@ async def test_virtual_key_max_budget_alert_check_with_user_obj():
assert captured_call_info.organization_id == "test-org"
assert captured_call_info.key_alias == "test-key"
assert captured_call_info.event_group == Litellm_EntityType.KEY
assert captured_call_info.budget_reset_at == valid_token.budget_reset_at
@pytest.mark.asyncio
@ -2767,6 +2769,7 @@ async def test_virtual_key_max_budget_alert_check_with_multi_threshold_map():
user_id="test-user",
key_alias="test-key",
metadata={"max_budget_alert_emails": alert_config},
budget_reset_at=datetime(2026, 9, 1, tzinfo=timezone.utc),
)
user_obj = LiteLLM_UserTable(
user_id="test-user",
@ -2786,6 +2789,7 @@ async def test_virtual_key_max_budget_alert_check_with_multi_threshold_map():
assert captured_call_info.max_budget_alert_emails == alert_config
assert captured_call_info.user_email == "owner@co.com"
assert captured_call_info.event_group == Litellm_EntityType.KEY
assert captured_call_info.budget_reset_at == valid_token.budget_reset_at
@pytest.mark.asyncio

23
tests/test_litellm/proxy/db/test_exception_handler.py
View file

@ -426,3 +426,26 @@ def test_handle_db_exception_with_non_db_error():
)
with pytest.raises(litellm.BudgetExceededError):
PrismaDBExceptionHandler.handle_db_exception(regular_error)
def test_is_unique_constraint_violation():
"""
Callers use an insert as an atomic cross-replica claim, so this predicate is what
separates "someone else already claimed it" from a database problem that has to
fail open. Table-not-found in particular must not read as a claim.
"""
import prisma
assert (
PrismaDBExceptionHandler.is_unique_constraint_violation(
prisma.errors.UniqueViolationError({"user_facing_error": {"meta": {}}})
)
is True
)
for other in (
prisma.errors.TableNotFoundError({"user_facing_error": {"meta": {}}}),
prisma.errors.DataError({"user_facing_error": {"meta": {}}}),
ClientNotConnectedError(),
ValueError("not a prisma error"),
):
assert PrismaDBExceptionHandler.is_unique_constraint_violation(other) is False

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

@ -4923,6 +4923,14 @@ 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