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

fix: prevent duplicate budget alert emails on concurrent threshold crossings (#32011)

Browse source 
* fix: prevent duplicate budget alert emails on concurrent threshold crossings

Budget alert emails were sent more than once for a single threshold crossing. The email dedup guard read the "already sent" marker, awaited the send, then wrote the marker, so concurrent requests crossing the same threshold within the send window all saw no marker and each sent. This affected the multi-threshold path (default_key_max_budget_alert_emails), the legacy single-threshold path (EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE), and the soft budget path, all in EmailBaseCallback.budget_alerts

All three branches now claim the send slot atomically before sending via async_increment_cache, which is atomic per event loop for the in-memory cache and across workers via Redis INCR; only the caller that observes a count of 1 sends. On send failure the marker is released with async_delete_cache so a transient failure does not suppress the alert for the full 24h TTL

* fix: harden budget alert claim release and skip-path event allocation

Addresses review feedback on the claim-before-send change. The claim release in each send-failure handler now logs the send error first and releases the claim best-effort through a shared helper, so a transient cache error during async_delete_cache cannot propagate out of the fire-and-forget budget_alerts task, drop the send-failure log, and leave the claim stuck for the full 24h TTL. In the multi-threshold branch the increment claim now runs before the WebhookEvent is built, so skipped concurrent crossings no longer construct and discard the event, matching the single-threshold and soft budget branches
This commit is contained in:
yucheng-berri 2026-07-03 14:59:34 -07:00 • committed by GitHub
parent 8e4e37d662
commit 099bc97320
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 193 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -477,9 +477,12 @@ class BaseEmailLogger(CustomLogger):
_id = user_info.token or user_info.user_id or "default_id"
_cache_key = f"email_budget_alerts:soft_budget_crossed:{_id}"
# Check if we've already sent this alert
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
send_count = await _cache.async_increment_cache(
key=_cache_key,
value=1,
ttl=EMAIL_BUDGET_ALERT_TTL,
)
if send_count is None or send_count <= 1:
# Create WebhookEvent for soft budget alert
event_message = f"Soft Budget Crossed - Total Soft Budget: ${user_info.soft_budget}"
webhook_event = WebhookEvent(
@ -508,18 +511,12 @@ class BaseEmailLogger(CustomLogger):
await self.send_team_soft_budget_alert_email(webhook_event)
else:
await self.send_soft_budget_alert_email(webhook_event)
# Cache the alert to prevent duplicate sends
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=EMAIL_BUDGET_ALERT_TTL,
)
except Exception as e:
verbose_proxy_logger.error(
f"Error sending soft budget alert email: {e}",
exc_info=True,
)
await self._release_budget_alert_claim(_cache, _cache_key)
return
# For max_budget_alert, check if we've already sent an alert
@ -545,9 +542,12 @@ 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}"
# Check if we've already sent this alert
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
send_count = await _cache.async_increment_cache(
key=_cache_key,
value=1,
ttl=EMAIL_BUDGET_ALERT_TTL,
)
if send_count is None or send_count <= 1:
# Calculate percentage
percentage = int(
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100
@ -576,18 +576,12 @@ class BaseEmailLogger(CustomLogger):
try:
await self.send_max_budget_alert_email(webhook_event)
# Cache the alert to prevent duplicate sends
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=EMAIL_BUDGET_ALERT_TTL,
)
except Exception as e:
verbose_proxy_logger.error(
f"Error sending max budget alert email: {e}",
exc_info=True,
)
await self._release_budget_alert_claim(_cache, _cache_key)
return
async def _handle_multi_threshold_max_budget_alert(
@ -617,10 +611,6 @@ class BaseEmailLogger(CustomLogger):
f"email_budget_alerts:max_budget_alert:{threshold_pct}:{_id}"
)
result = await _cache.async_get_cache(key=_cache_key)
if result is not None:
continue
# Parse emails + auto-include owner
emails = _parse_email_list(raw_emails)
if user_info.user_email:
@ -634,6 +624,14 @@ 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:
continue
event_message = f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached"
webhook_event = WebhookEvent(
event="max_budget_alert",
@ -660,16 +658,21 @@ class BaseEmailLogger(CustomLogger):
threshold_pct=threshold_pct,
recipient_emails=recipient_emails,
)
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=EMAIL_BUDGET_ALERT_TTL,
)
except Exception as e:
verbose_proxy_logger.error(
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)
async def _release_budget_alert_claim(self, cache: DualCache, cache_key: str) -> None:
try:
await cache.async_delete_cache(key=cache_key)
except Exception:
verbose_proxy_logger.debug(
"Failed to release budget alert claim for %s; it expires with the TTL",
cache_key,
)
async def _get_email_params(
self,

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

@ -1,3 +1,4 @@
import asyncio
import json
import os
import sys
@ -7,6 +8,7 @@ from unittest.mock import patch
import pytest
from fastapi.testclient import TestClient
from litellm.caching.caching import DualCache
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
BaseEmailLogger,
)
@ -707,10 +709,9 @@ async def test_budget_alerts_soft_budget_crossed(base_email_logger, mock_send_em
event_group=Litellm_EntityType.USER,
)
# Mock the cache to return None (no previous alert sent)
# Mock the cache so the claim is won (increment returns 1)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(
@ -726,14 +727,14 @@ async def test_budget_alerts_soft_budget_crossed(base_email_logger, mock_send_em
call_args = mock_send_email.call_args[1]
assert call_args["to_email"] == ["test@example.com"]
# Verify cache was set to prevent duplicate alerts
mock_cache.async_set_cache.assert_called_once()
cache_call_args = mock_cache.async_set_cache.call_args[1]
# Verify the send slot was claimed to prevent duplicate alerts
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:soft_budget_crossed:test_user"
)
assert cache_call_args["value"] == "SENT"
assert cache_call_args["value"] == 1
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
@ -774,9 +775,9 @@ async def test_budget_alerts_soft_budget_duplicate_prevention(
event_group=Litellm_EntityType.USER,
)
# Mock the cache to return "SENT" (previous alert already sent)
# Mock the cache so the slot is already claimed (increment returns > 1)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value="SENT")
mock_cache.async_increment_cache = mock.AsyncMock(return_value=2)
base_email_logger.internal_usage_cache = mock_cache
await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
@ -818,10 +819,9 @@ async def test_budget_alerts_uses_token_for_cache_key(
event_group=Litellm_EntityType.KEY,
)
# Mock the cache to return None (no previous alert sent)
# Mock the cache so the claim is won (increment returns 1)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(
@ -833,8 +833,8 @@ async def test_budget_alerts_uses_token_for_cache_key(
await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
# Verify cache key uses token instead of user_id
mock_cache.async_set_cache.assert_called_once()
cache_call_args = mock_cache.async_set_cache.call_args[1]
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:soft_budget_crossed:hashed_token_123"
@ -880,8 +880,7 @@ async def test_budget_alerts_max_budget_alert_crossed(
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(
@ -899,12 +898,12 @@ async def test_budget_alerts_max_budget_alert_crossed(
assert call_args["to_email"] == ["test@example.com"]
assert "Max Budget Alert" in call_args["subject"]
mock_cache.async_set_cache.assert_called_once()
cache_call_args = mock_cache.async_set_cache.call_args[1]
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"
)
assert cache_call_args["value"] == "SENT"
assert cache_call_args["value"] == 1
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
@ -928,8 +927,7 @@ async def test_multi_threshold_sends_crossed_thresholds(
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
@ -941,7 +939,9 @@ async def test_multi_threshold_sends_crossed_thresholds(
assert mock_send_email.call_count == 2
# Check cache keys include threshold percentage
cache_keys = [c[1]["key"] for c in mock_cache.async_set_cache.call_args_list]
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
@ -964,15 +964,14 @@ async def test_multi_threshold_dedup_cache_prevents_resend(
},
)
# Simulate 50% already sent (cached), 75% not yet sent
async def cache_get(key):
# Simulate 50% already claimed (increment returns >1), 75% first send (returns 1)
async def cache_increment(key, value, ttl=None):
if "50:" in key:
return "SENT"
return None
return 2
return 1
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(side_effect=cache_get)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(side_effect=cache_increment)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
@ -982,7 +981,7 @@ async def test_multi_threshold_dedup_cache_prevents_resend(
# Only 75% should fire
assert mock_send_email.call_count == 1
cache_key = mock_cache.async_set_cache.call_args[1]["key"]
cache_key = mock_cache.async_increment_cache.call_args[1]["key"]
assert "75:" in cache_key
@ -1004,8 +1003,7 @@ async def test_multi_threshold_owner_email_auto_included(
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
@ -1038,8 +1036,7 @@ async def test_multi_threshold_malformed_keys_skipped(
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
@ -1069,8 +1066,7 @@ async def test_multi_threshold_empty_emails_only_owner(
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
@ -1097,8 +1093,7 @@ async def test_no_map_preserves_old_single_threshold(
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
@ -1110,7 +1105,7 @@ async def test_no_map_preserves_old_single_threshold(
call_args = mock_send_email.call_args[1]
assert call_args["to_email"] == ["test@example.com"]
# Old path cache key has no threshold percentage
cache_key = mock_cache.async_set_cache.call_args[1]["key"]
cache_key = mock_cache.async_increment_cache.call_args[1]["key"]
assert cache_key == "email_budget_alerts:max_budget_alert:test_user"
@ -1242,3 +1237,131 @@ async def test_send_soft_budget_alert_email_default_footer_when_no_signature(
html_body = mock_send_email.call_args[1]["html_body"]
assert EMAIL_FOOTER in html_body
_BUDGET_ALERT_BRANCHES = [
(
"multi_threshold",
"max_budget_alert",
"send_max_budget_alert_email",
dict(max_budget=100.0, spend=80.0, max_budget_alert_emails={"50": ["finance@co.com"]}),
),
(
"single_threshold",
"max_budget_alert",
"send_max_budget_alert_email",
dict(max_budget=100.0, spend=85.0),
),
(
"soft_budget",
"soft_budget",
"send_soft_budget_alert_email",
dict(soft_budget=50.0, spend=60.0),
),
]
def _budget_alert_user_info(extra: dict) -> CallInfo:
return CallInfo(
token="hashed_key_1",
user_id="test_user",
user_email="owner@co.com",
event_group=Litellm_EntityType.KEY,
**extra,
)
@pytest.mark.parametrize(
"branch, alert_type, send_method, ci_kwargs",
_BUDGET_ALERT_BRANCHES,
ids=[b[0] for b in _BUDGET_ALERT_BRANCHES],
)
@pytest.mark.asyncio
async def test_budget_alert_no_duplicate_on_concurrent_crossing(
base_email_logger, branch, alert_type, send_method, ci_kwargs
):
"""Regression for LIT-4172: two requests crossing the same threshold at the
same time must send exactly one email. The old code wrote the dedup marker
only after the send finished awaiting, so both concurrent tasks passed the
'already sent' check and both sent. Covers all three send branches."""
base_email_logger.internal_usage_cache = DualCache()
sends = []
async def slow_send(*args, **kwargs):
sends.append(1)
await asyncio.sleep(0.05)
with mock.patch.object(base_email_logger, send_method, side_effect=slow_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await asyncio.gather(
base_email_logger.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
),
base_email_logger.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
),
)
assert len(sends) == 1
@pytest.mark.parametrize(
"branch, alert_type, send_method, ci_kwargs",
_BUDGET_ALERT_BRANCHES,
ids=[b[0] for b in _BUDGET_ALERT_BRANCHES],
)
@pytest.mark.asyncio
async def test_budget_alert_failed_send_releases_claim_for_retry(
base_email_logger, branch, alert_type, send_method, ci_kwargs
):
"""Claiming the send slot before sending must not swallow the alert forever
if the send fails; the claim is released so a later request retries. Covers
all three send branches."""
base_email_logger.internal_usage_cache = DualCache()
attempts = []
async def flaky_send(*args, **kwargs):
attempts.append(1)
if len(attempts) == 1:
raise ValueError("transient email backend failure")
with mock.patch.object(base_email_logger, send_method, side_effect=flaky_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
)
await base_email_logger.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
)
assert len(attempts) == 2
@pytest.mark.asyncio
async def test_budget_alert_release_failure_does_not_propagate(base_email_logger):
"""If the send fails and releasing the claim also fails (transient cache
error), budget_alerts must swallow it and still log the send failure rather
than letting the exception escape the fire-and-forget task."""
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
mock_cache.async_delete_cache = mock.AsyncMock(
side_effect=RuntimeError("cache backend unavailable")
)
base_email_logger.internal_usage_cache = mock_cache
async def failing_send(*args, **kwargs):
raise ValueError("smtp backend down")
with mock.patch.object(
base_email_logger, "send_max_budget_alert_email", side_effect=failing_send
):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
# Must not raise even though both the send and the release fail.
await base_email_logger.budget_alerts(
type="max_budget_alert",
user_info=_budget_alert_user_info(dict(max_budget=100.0, spend=85.0)),
)
mock_cache.async_delete_cache.assert_awaited_once()