diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 00f6d44e25a..e8ecf253112 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3605,34 +3605,46 @@ async def _virtual_key_soft_budget_check( Triggers a budget alert if the token is over it's soft budget. """ + if valid_token.soft_budget is None: + return - if valid_token.soft_budget and valid_token.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, - valid_token.soft_budget, - ) - call_info = CallInfo( - token=valid_token.token, - spend=valid_token.spend, - max_budget=valid_token.max_budget, - soft_budget=valid_token.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=user_obj.user_email if user_obj else None, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - ) + from litellm.proxy.proxy_server import get_current_spend - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, - ) + spend = await get_current_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + max_budget=valid_token.soft_budget, + ) + + if spend < valid_token.soft_budget: + return + + verbose_proxy_logger.debug( + "Crossed Soft Budget for token %s, spend %s, soft_budget %s", + valid_token.token, + spend, + valid_token.soft_budget, + ) + call_info = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, ) + ) def _parse_email_list(raw: Any) -> List[str]: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 27f43c4948f..7d9b4dbf4de 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 + (0.0, 0.0, True), # Zero soft budget is enforced, not treated as disabled + (100.0, 0.0, True), # Any spend crosses a zero soft budget ], ) @pytest.mark.asyncio @@ -2113,6 +2115,10 @@ async def test_virtual_key_soft_budget_check_scenarios( spend, soft_budget, expect_alert ): """Test _virtual_key_soft_budget_check with various spend and soft_budget scenarios""" + from litellm.proxy.proxy_server import spend_counter_cache + + spend_counter_cache.in_memory_cache.flush_cache() + alert_triggered = False class MockProxyLogging: @@ -2145,6 +2151,49 @@ 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.mark.asyncio +async def test_virtual_key_soft_budget_check_uses_fresh_cross_pod_spend(): + """Regression: the soft budget check must read the live cross-pod spend + counter, not the stale ``valid_token.spend`` snapshot loaded at auth time. + + The cached token spend (10.0) is below the soft budget, so relying on it + would suppress the alert; the fresh counter (100.0) is over it and must + trigger the alert with the fresh value. + """ + from litellm.proxy.proxy_server import spend_counter_cache + + captured_call_info = None + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal captured_call_info + captured_call_info = user_info + + valid_token = UserAPIKeyAuth( + token="stale-spend-token", + spend=10.0, + soft_budget=50.0, + user_id="test-user", + key_alias="test-key", + ) + + spend_counter_cache.in_memory_cache.flush_cache() + spend_counter_cache.in_memory_cache.set_cache( + key=f"spend:key:{valid_token.token}", value=100.0 + ) + + 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 + assert captured_call_info.spend == 100.0 + + @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"""