From 96db81326265cc4d6f4ccb491c20833df118c320 Mon Sep 17 00:00:00 2001 From: Keval Mahajan Date: Thu, 27 Aug 2026 20:30:35 +0530 Subject: [PATCH 1/2] fix(proxy): apply the zero-cost bypass in the personal budget hook --- litellm/proxy/hooks/max_budget_limiter.py | 20 +++++++ .../proxy/hooks/test_max_budget_limiter.py | 55 +++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index eaf37b0bcf1..dd1769d58e9 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -40,6 +40,26 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): ): return + # Zero-cost models bypass budget enforcement everywhere else - both + # common_checks and the reservation path skip them via + # _is_model_cost_zero. Without the same exemption here, a user at + # their personal max_budget also loses access to models that cannot + # incur cost, so a free-model fallback stops working under a + # personal cap while it keeps working under a team cap. + # Imported lazily to avoid a circular import via proxy.utils. + from litellm.proxy.auth.auth_checks import _is_model_cost_zero + from litellm.proxy.proxy_server import llm_router + + if _is_model_cost_zero( + model=data.get("model") if data else None, + llm_router=llm_router, + ): + verbose_proxy_logger.debug( + "MaxBudgetLimiter: skipping personal budget check for zero-cost model: %s", + data.get("model") if data else None, + ) + return + # The reservation path admits at the strict-`<` boundary and # atomically pre-fills the same counter we'd read here. Re-checking # with `>=` would reject a request the reservation already admitted diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py index 71671966d1a..37a377b50fc 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -235,3 +235,58 @@ async def test_no_max_budget_passes(): assert result is None mock_get_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_skips_for_zero_cost_model_when_over_budget(): + """ + Zero-cost models bypass budget enforcement in common_checks and in the + reservation path (both via `_is_model_cost_zero`). This hook must apply the + same exemption, otherwise a user at their personal `max_budget` also loses + access to models that cannot incur cost - while the same models stay usable + when a *team* budget is the one exhausted, since this hook only ever reads + `user_max_budget`. + """ + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) + + with patch( + "litellm.proxy.auth.auth_checks._is_model_cost_zero", + return_value=True, + ), patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ) as mock_get_spend: + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "free-model"}, + call_type="completion", + ) + + assert result is None + mock_get_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_still_rejects_priced_model_when_over_budget(): + """The zero-cost exemption must not weaken enforcement for priced models.""" + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) + + with patch( + "litellm.proxy.auth.auth_checks._is_model_cost_zero", + return_value=False, + ), patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ), pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "paid-model"}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + assert "Max budget limit reached." in exc_info.value.detail From 03c1348153599db35e7cf0ff11a1bbe53dc306f4 Mon Sep 17 00:00:00 2001 From: Keval Mahajan Date: Thu, 27 Aug 2026 20:52:28 +0530 Subject: [PATCH 2/2] test(proxy): explain why the budget hook tests patch module globals --- tests/test_litellm/proxy/hooks/test_max_budget_limiter.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py index 37a377b50fc..97a4e0c9477 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -250,10 +250,10 @@ async def test_skips_for_zero_cost_model_when_over_budget(): handler = _PROXY_MaxBudgetLimiter() user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) - with patch( + with patch( # test-quality-ok: hook resolves both collaborators as module globals, no injection seam "litellm.proxy.auth.auth_checks._is_model_cost_zero", return_value=True, - ), patch( + ), patch( # test-quality-ok: proxy module global, no injection seam "litellm.proxy.proxy_server.get_current_spend", new=AsyncMock(return_value=10.0), ) as mock_get_spend: @@ -274,10 +274,10 @@ async def test_still_rejects_priced_model_when_over_budget(): handler = _PROXY_MaxBudgetLimiter() user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) - with patch( + with patch( # test-quality-ok: hook resolves both collaborators as module globals, no injection seam "litellm.proxy.auth.auth_checks._is_model_cost_zero", return_value=False, - ), patch( + ), patch( # test-quality-ok: proxy module global, no injection seam "litellm.proxy.proxy_server.get_current_spend", new=AsyncMock(return_value=10.0), ), pytest.raises(HTTPException) as exc_info: