From a564b600f25ca59244467282a5829b3ce446c1c9 Mon Sep 17 00:00:00 2001 From: Jayanth <191611727+Jayanth-reflex@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:04:06 +0000 Subject: [PATCH] fix(alerting): send one budget alert when concurrent checks hit the same budget event --- .../SlackAlerting/slack_alerting.py | 30 ++++--- .../SlackAlerting/test_slack_alerting.py | 81 +++++++++++++++++++ 2 files changed, 99 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 7e9a8fe877c..9639e1c7923 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -592,6 +592,13 @@ class SlackAlerting(CustomBatchLogger): _cache_key: Final = f"budget_alerts:{event}:{_id}" result: Final = await _cache.async_get_cache(key=_cache_key) if result is None: + claim_count: Final = await _cache.async_increment_cache( + key=_cache_key, + value=1, + ttl=self.alerting_args.budget_alert_ttl, + ) + if claim_count is not None and claim_count > 1: + return webhook_event = WebhookEvent( event=event, event_message=event_message, @@ -612,18 +619,17 @@ class SlackAlerting(CustomBatchLogger): alert_emails=user_info.alert_emails, max_budget_alert_emails=user_info.max_budget_alert_emails, ) - await self.send_alert( - message=event_message + "\n\n" + user_info_str, - level="High", - alert_type=AlertType.budget_alerts, - user_info=webhook_event, - alerting_metadata={}, - ) - await _cache.async_set_cache( - key=_cache_key, - value="SENT", - ttl=self.alerting_args.budget_alert_ttl, - ) + try: + await self.send_alert( + message=event_message + "\n\n" + user_info_str, + level="High", + alert_type=AlertType.budget_alerts, + user_info=webhook_event, + alerting_metadata={}, + ) + except Exception: + await _cache.async_delete_cache(key=_cache_key) + raise return return diff --git a/tests/unit/integrations/SlackAlerting/test_slack_alerting.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py index b38d7f745b5..385a41e22ed 100644 --- a/tests/unit/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py @@ -8,6 +8,7 @@ import unittest from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import fakeredis import httpx import pytest from openai import APIError @@ -17,6 +18,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._internal_context import current_service_target from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type from litellm.integrations.SlackAlerting.slack_alerting import DeploymentMetrics, SlackAlerting from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -771,6 +773,85 @@ async def test_budget_alerts_crossed_again(slack_alerting): ) mock_send_alert.assert_not_awaited() + +BUDGET_WEBHOOK_URL: Final = "https://budget-webhook.example/alerts" + + +def _accepted_webhook_response() -> MagicMock: + response: Final = MagicMock(spec=httpx.Response) + response.status_code = 200 + return response + + +def _webhook_post_answering_after_network_wait() -> AsyncMock: + response: Final = _accepted_webhook_response() + + async def _answer(**_: object) -> MagicMock: + await asyncio.sleep(0) + return response + + return AsyncMock(side_effect=_answer) + + +def _budget_webhook_alerting(webhook_post: AsyncMock, cache: DualCache | None = None) -> SlackAlerting: + http_handler: Final = MagicMock(spec=AsyncHTTPHandler, post=webhook_post) + slack_alerting: Final = SlackAlerting( + alerting=["webhook"], async_http_handler=http_handler, internal_usage_cache=cache + ) + slack_alerting.periodic_started = True + return slack_alerting + + +def _cache_sharing_redis(server: fakeredis.FakeServer, monkeypatch: pytest.MonkeyPatch) -> DualCache: + redis_cache: Final = RedisCache(host="localhost", port=6379) + redis_client: Final = fakeredis.FakeAsyncRedis(server=server) + monkeypatch.setattr(redis_cache, "init_async_client", lambda: redis_client) + return DualCache(redis_cache=redis_cache) + + +@pytest.mark.asyncio +async def test_concurrent_budget_alerts_deliver_one_webhook_alert(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("WEBHOOK_URL", BUDGET_WEBHOOK_URL) + webhook_post: Final = _webhook_post_answering_after_network_wait() + slack_alerting: Final = _budget_webhook_alerting(webhook_post) + + await asyncio.gather(*(slack_alerting.budget_alerts(type="soft_budget", user_info=team_info) for _ in range(10))) + await slack_alerting.budget_alerts(type="soft_budget", user_info=team_info) + + assert webhook_post.await_count == 1 + + +@pytest.mark.asyncio +async def test_budget_alert_workers_sharing_redis_deliver_one_webhook_alert( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("WEBHOOK_URL", BUDGET_WEBHOOK_URL) + webhook_post: Final = _webhook_post_answering_after_network_wait() + server: Final = fakeredis.FakeServer() + workers: Final = tuple( + _budget_webhook_alerting(webhook_post, _cache_sharing_redis(server, monkeypatch)) for _ in range(2) + ) + + await asyncio.gather(*(worker.budget_alerts(type="soft_budget", user_info=team_info) for worker in workers * 5)) + + assert webhook_post.await_count == 1 + + +@pytest.mark.asyncio +async def test_failed_budget_alert_delivery_lets_the_next_check_retry(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("WEBHOOK_URL", BUDGET_WEBHOOK_URL) + webhook_post: Final = AsyncMock( + side_effect=[httpx.ConnectError("webhook unreachable"), _accepted_webhook_response()] + ) + slack_alerting: Final = _budget_webhook_alerting(webhook_post) + + with pytest.raises(httpx.ConnectError): + await slack_alerting.budget_alerts(type="soft_budget", user_info=team_info) + await slack_alerting.budget_alerts(type="soft_budget", user_info=team_info) + + assert webhook_post.await_count == 2 + + @pytest.mark.asyncio async def test_daily_reports_unit_test(slack_alerting): with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert: