From 7866e06deeecddc7f7c2340eab7595d8e54992c0 Mon Sep 17 00:00:00 2001 From: Ali Tariq Date: Wed, 15 Jul 2026 15:13:14 +0500 Subject: [PATCH] fix(proxy): key soft budget reads shared spend, honors zero threshold _virtual_key_soft_budget_check gated on `valid_token.soft_budget` being truthy, so a soft_budget of 0 (accepted by key/budget validation) could never alert. It also compared against the cached key object's spend while the hard key-budget check had already moved to the Redis-first counter, so in multi-replica deployments the alert stayed silent after shared spend crossed the threshold. Gate on `is not None` and read the same `spend:key:{token}` counter, with the cached spend as fallback. max_budget is left unset on that read: this path only alerts, so it must not raise on an unverifiable counter. Fixes #33330 --- litellm/proxy/auth/auth_checks.py | 19 ++++- .../proxy/auth/test_auth_checks.py | 78 +++++++++++++++++++ 2 files changed, 94 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 00f6d44e25a..ee9271c3a10 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3604,18 +3604,31 @@ async def _virtual_key_soft_budget_check( """ Triggers a budget alert if the token is over it's soft budget. + Reads the same cross-pod spend counter as `_virtual_key_max_budget_check` so the alert + fires off shared spend rather than this pod's cached key object. `max_budget` is left + unset on that read: this path only alerts, so it must never raise on an unverifiable + counter the way an enforcement read may. """ + if valid_token.soft_budget is None: + return - if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: + from litellm.proxy.proxy_server import get_current_spend + + spend = await get_current_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + ) + + if spend >= valid_token.soft_budget: verbose_proxy_logger.debug( "Crossed Soft Budget for token %s, spend %s, soft_budget %s", valid_token.token, - valid_token.spend, + spend, valid_token.soft_budget, ) call_info = CallInfo( token=valid_token.token, - spend=valid_token.spend, + spend=spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 27f43c4948f..2a5bdf01dbc 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2106,6 +2106,8 @@ async def test_virtual_key_soft_budget_check_without_user_obj(): (50.0, 50.0, True), # At soft budget (25.0, 50.0, False), # Under soft budget (100.0, None, False), # No soft budget set + (1.0, 0.0, True), # Zero is a valid threshold, not "unset" + (0.0, 0.0, True), # At a zero soft budget ], ) @pytest.mark.asyncio @@ -2145,6 +2147,82 @@ async def test_virtual_key_soft_budget_check_scenarios( ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" +@pytest.fixture() +def spend_counter_cache(): + import litellm.proxy.proxy_server as ps + from litellm.caching.dual_cache import DualCache + + original_counter_cache = ps.spend_counter_cache + original_prisma_client = ps.prisma_client + counter_cache = DualCache() + ps.spend_counter_cache = counter_cache + ps.prisma_client = None + try: + yield counter_cache + finally: + ps.spend_counter_cache = original_counter_cache + ps.prisma_client = original_prisma_client + + +@pytest.mark.asyncio +async def test_virtual_key_soft_budget_check_uses_shared_spend_counter( + spend_counter_cache, +): + """The alert must fire off the cross-pod spend counter, even when this pod's cached + key object is stale-low and still under the soft budget.""" + captured_call_info = None + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal captured_call_info + captured_call_info = user_info + + await spend_counter_cache.async_increment_cache( + key="spend:key:stale-key", value=12.0 + ) + + valid_token = UserAPIKeyAuth( + token="stale-key", + spend=4.0, + soft_budget=10.0, + key_alias="test-key", + ) + + await _virtual_key_soft_budget_check( + valid_token=valid_token, + proxy_logging_obj=MockProxyLogging(), + user_obj=None, + ) + + await asyncio.sleep(0.1) + + assert captured_call_info is not None, "expected soft budget alert off shared spend 12.0" + assert captured_call_info.spend == 12.0 + + +@pytest.mark.asyncio +async def test_virtual_key_soft_budget_check_does_not_raise_on_unverifiable_spend( + spend_counter_cache, monkeypatch +): + """Soft budgets only alert. Under fail-closed enforcement an unverifiable counter must + still not block the request the way a max_budget read may.""" + import litellm.proxy.proxy_server as ps + + monkeypatch.setitem(ps.general_settings, "fail_closed_budget_enforcement", True) + + valid_token = UserAPIKeyAuth( + token="unverifiable-key", + spend=1.0, + soft_budget=10.0, + ) + + await _virtual_key_soft_budget_check( + valid_token=valid_token, + proxy_logging_obj=MagicMock(budget_alerts=AsyncMock()), + user_obj=None, + ) + + @pytest.mark.asyncio async def test_virtual_key_max_budget_alert_check_with_user_obj(): """Test _virtual_key_max_budget_alert_check includes user_email when user_obj is provided"""