From 099bc973209fa43e56aff37e2cdeff41d4e8c784 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 3 Jul 2026 14:59:34 -0700 Subject: [PATCH] fix: prevent duplicate budget alert emails on concurrent threshold crossings (#32011) * 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 --- .../send_emails/base_email.py | 61 +++--- .../send_emails/test_base_email.py | 199 ++++++++++++++---- 2 files changed, 193 insertions(+), 67 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index 9d15f45079f..be80a12c80a 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -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, diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py index db23b712125..c1ccc454305 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -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()