scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

Merge e01a05472e into eb0e3f8c18

Browse source
This commit is contained in:
yucheng-berri 2026-08-27 19:19:28 -05:00 • committed by GitHub
commit 44f3bc657f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 789 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

129
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
@ -543,16 +550,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"
@ -573,6 +582,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:
@ -582,7 +592,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(
@ -625,12 +641,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"
@ -651,6 +667,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:
@ -664,7 +681,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:
@ -675,6 +697,79 @@ 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}"
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
if not user_info.token:
return True
return await claim_budget_alert_slot(
token=user_info.token,
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)
)
if not user_info.token:
return
await release_budget_alert_slot(
token=user_info.token,
alert_type=alert_type,
threshold_pct=threshold_pct,
budget_window=budget_window,
)
async def _get_email_params(
self,
email_event: EmailEvent,

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

@ -0,0 +1,18 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetAlertSent" (
"id" TEXT NOT NULL,
"token" 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_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;

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

@ -465,6 +465,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
@ -1609,3 +1610,16 @@ model LiteLLM_WorkflowMessage {
@@unique([run_id, sequence_number])
@@index([run_id])
}
model LiteLLM_BudgetAlertSent {
id String @id @default(uuid())
token String
alert_type String
threshold_pct Int
budget_window String
sent_at DateTime @default(now())
verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([token, alert_type, threshold_pct])
}

2
litellm/constants.py
View file

@ -503,6 +503,8 @@ EMAIL_BUDGET_ALERT_TTL: Final = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = float(
os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)
) # 80% of max budget
MAX_BUDGET_ALERT_TYPE: Final = "max_budget_alert"
MAX_BUDGET_DEFAULT_ALERT_TYPE: Final = "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

@ -659,6 +659,7 @@ class SlackAlerting(CustomBatchLogger):
"""
_all_fields_as_dict: Final = 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

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

@ -4605,6 +4605,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(
@ -4636,6 +4637,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(

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

@ -0,0 +1,141 @@
"""
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 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
from datetime import datetime, timezone
from typing import Final
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: Final = budget_reset_at.isoformat() if budget_reset_at is not None else ""
return f"{period}|{max_budget}"
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
"token": token,
"alert_type": alert_type,
"threshold_pct": threshold_pct,
}
async def claim_budget_alert_slot(
token: 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: Final = _claim_identity(token, 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 key %s at %d%%, sending anyway: %s",
token,
threshold_pct,
e,
)
return True
try:
rows_updated: Final = 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 key %s at %d%%, sending anyway: %s",
token,
threshold_pct,
e,
)
return True
async def release_budget_alert_slot(
token: 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: Final = _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.warning(
"Failed to release budget alert claim for key %s at %d%%: %s",
token,
threshold_pct,
e,
)

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

@ -123,6 +123,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:
"""

14
litellm/proxy/schema.prisma
View file

@ -465,6 +465,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
@ -1609,3 +1610,16 @@ model LiteLLM_WorkflowMessage {
@@unique([run_id, sequence_number])
@@index([run_id])
}
model LiteLLM_BudgetAlertSent {
id String @id @default(uuid())
token String
alert_type String
threshold_pct Int
budget_window String
sent_at DateTime @default(now())
verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([token, alert_type, threshold_pct])
}

14
schema.prisma
View file

@ -465,6 +465,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
@ -1609,3 +1610,16 @@ model LiteLLM_WorkflowMessage {
@@unique([run_id, sequence_number])
@@index([run_id])
}
model LiteLLM_BudgetAlertSent {
id String @id @default(uuid())
token String
alert_type String
threshold_pct Int
budget_window String
sent_at DateTime @default(now())
verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([token, alert_type, threshold_pct])
}

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

@ -2,6 +2,9 @@ import asyncio
import json
import os
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
@ -21,6 +24,11 @@ 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,
release_budget_alert_slot,
)
@pytest.fixture(autouse=True)
@ -912,7 +920,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
@ -953,8 +962,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
@ -1117,7 +1126,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>"
@ -1376,3 +1385,416 @@ 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 = ("token", "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 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(
token="hashed_key_1",
alert_type=MAX_BUDGET_ALERT_TYPE,
threshold_pct=50,
budget_window="later-window",
)
await release_budget_alert_slot(
token="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_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

@ -3051,6 +3051,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(
@ -3082,6 +3083,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
@ -3194,6 +3196,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",
@ -3213,6 +3216,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

@ -428,6 +428,29 @@ def test_handle_db_exception_with_non_db_error():
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
def _permanent_prisma_faults():
"""Every prisma error class that is not a transient outage and not a
data-layer error, built by enumeration so the list cannot drift out of sync