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