From ede4c635475fa778b3b3edd3be2e04ec2edf5ce2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 02:11:42 +0000 Subject: [PATCH] fix(proxy/budgets): reject requests at the exact hard-budget limit --- litellm/proxy/auth/auth_checks.py | 6 +- litellm/proxy/auth/user_api_key_auth.py | 2 +- .../proxy/hooks/model_max_budget_limiter.py | 4 +- .../proxy/auth/test_auth_checks.py | 88 +++++++++++++++++++ .../hooks/test_model_max_budget_limiter.py | 84 ++++++++++++++++++ 5 files changed, 178 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 00f6d44e25a..c2e874864db 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1088,7 +1088,7 @@ async def _check_end_user_budget( max_budget=end_user_budget, fallback_authoritative=True, ) - if end_user_spend > end_user_budget: + if end_user_spend >= end_user_budget: raise litellm.BudgetExceededError( current_cost=end_user_spend, max_budget=end_user_budget, @@ -3900,7 +3900,7 @@ async def _team_max_budget_check( max_budget=team_object.max_budget, ) - if math.isfinite(team_object.max_budget) and spend > team_object.max_budget: + if math.isfinite(team_object.max_budget) and spend >= team_object.max_budget: if valid_token: call_info = CallInfo( token=valid_token.token, @@ -4320,7 +4320,7 @@ async def _tag_max_budget_check( max_budget=tag_object.litellm_budget_table.max_budget, fallback_authoritative=True, ) - if tag_spend <= tag_object.litellm_budget_table.max_budget: + if tag_spend < tag_object.litellm_budget_table.max_budget: continue raise litellm.BudgetExceededError( current_cost=tag_spend, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b402212fb2e..c849965d90e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1793,7 +1793,7 @@ async def _user_api_key_auth_builder( fallback_spend=team_member_spend, max_budget=team_member_budget, ) - if team_member_spend > team_member_budget: + if team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 803fe64c193..d07144e8e4f 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -70,7 +70,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): if ( _current_spend is not None and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget + and _current_spend >= _current_model_budget_info.max_budget ): raise litellm.BudgetExceededError( message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}", @@ -134,7 +134,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): if ( _current_spend is not None and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget + and _current_spend >= _current_model_budget_info.max_budget ): raise litellm.BudgetExceededError( message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 27f43c4948f..de1144226e5 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2978,6 +2978,94 @@ async def test_team_member_budget_check_reads_from_spend_counter(): assert exc_info.value.current_cost == 1.5 +@pytest.mark.asyncio +async def test_team_budget_check_rejects_at_exact_limit(): + """Team budget must be rejected when spend equals max_budget (hard limit).""" + from litellm.proxy.utils import ProxyLogging + + team_object = LiteLLM_TeamTable(team_id="test-team", spend=0.0, max_budget=1.0) + valid_token = UserAPIKeyAuth(token="test-token", team_id="test-team") + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.budget_alerts = AsyncMock() + + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + if counter_key == "spend:team:test-team": + return 1.0 + return fallback_spend + + with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _team_max_budget_check( + team_object=team_object, + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.current_cost == 1.0 + assert exc_info.value.max_budget == 1.0 + + +@pytest.mark.asyncio +async def test_end_user_budget_check_rejects_at_exact_limit(): + """End-user budget must be rejected when spend equals max_budget (hard limit).""" + end_user_object = LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + if counter_key == "spend:end_user:customer-1": + return 1.0 + return fallback_spend + + with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_end_user_budget( + end_user_obj=end_user_object, + route="/chat/completions", + ) + assert exc_info.value.current_cost == 1.0 + assert exc_info.value.max_budget == 1.0 + + +@pytest.mark.asyncio +async def test_tag_budget_check_rejects_at_exact_limit(): + """Tag budget must be rejected when spend equals max_budget (hard limit).""" + from litellm.proxy.utils import ProxyLogging + + tag_object = LiteLLM_TagTable( + tag_name="paid-tag", + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + if counter_key == "spend:tag:paid-tag": + return 1.0 + return fallback_spend + + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={"paid-tag": tag_object}, + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _tag_max_budget_check( + request_body={"metadata": {"tags": ["paid-tag"]}}, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token"), + ) + assert exc_info.value.current_cost == 1.0 + assert exc_info.value.max_budget == 1.0 + + class TestGuardrailModificationCheck: """Defense-in-depth: `_guardrail_modification_check` must 403 when the caller's metadata attempts to modify any guardrail-related key and the diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py new file mode 100644 index 00000000000..ba28222decd --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py @@ -0,0 +1,84 @@ +from unittest.mock import AsyncMock + +import pytest + +import litellm +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, +) + + +@pytest.mark.asyncio +async def test_virtual_key_model_budget_rejects_at_exact_limit(): + """Per-model virtual-key budget must reject when spend equals max_budget.""" + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + limiter._get_virtual_key_spend_for_model = AsyncMock(return_value=10.0) + + user_api_key_dict = UserAPIKeyAuth( + token="test-token", + model_max_budget={"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await limiter.is_key_within_model_budget( + user_api_key_dict=user_api_key_dict, model="gpt-4o" + ) + assert exc_info.value.current_cost == 10.0 + assert exc_info.value.max_budget == 10.0 + + +@pytest.mark.asyncio +async def test_virtual_key_model_budget_admits_below_limit(): + """Per-model virtual-key budget must admit when spend is below max_budget.""" + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + limiter._get_virtual_key_spend_for_model = AsyncMock(return_value=9.99) + + user_api_key_dict = UserAPIKeyAuth( + token="test-token", + model_max_budget={"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + assert ( + await limiter.is_key_within_model_budget( + user_api_key_dict=user_api_key_dict, model="gpt-4o" + ) + is True + ) + + +@pytest.mark.asyncio +async def test_end_user_model_budget_rejects_at_exact_limit(): + """Per-model end-user budget must reject when spend equals max_budget.""" + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + limiter._get_end_user_spend_for_model = AsyncMock(return_value=10.0) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await limiter.is_end_user_within_model_budget( + end_user_id="customer-1", + end_user_model_max_budget={ + "gpt-4o": {"budget_limit": 10.0, "time_period": "1d"} + }, + model="gpt-4o", + ) + assert exc_info.value.current_cost == 10.0 + assert exc_info.value.max_budget == 10.0 + + +@pytest.mark.asyncio +async def test_end_user_model_budget_admits_below_limit(): + """Per-model end-user budget must admit when spend is below max_budget.""" + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + limiter._get_end_user_spend_for_model = AsyncMock(return_value=9.99) + + assert ( + await limiter.is_end_user_within_model_budget( + end_user_id="customer-1", + end_user_model_max_budget={ + "gpt-4o": {"budget_limit": 10.0, "time_period": "1d"} + }, + model="gpt-4o", + ) + is True + )