From a12b4249bd65b4712682cd73bf96055e3389f7a2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 1 May 2026 15:56:07 -0700 Subject: [PATCH] [Fix] Proxy: Skip Personal Budget Hook When Reservation Covers Counter The reservation path (PR #26845) atomically pre-fills `spend:user:{user_id}` and admits at the strict-`<` boundary. The legacy `_PROXY_MaxBudgetLimiter` pre-call hook re-reads the same counter with `>=`, so a reservation that fills the counter to exactly `max_budget` (e.g. a request without a `max_tokens` cap that falls back to reserving the smallest remaining headroom) is rejected by the hook even though the reservation already admitted it. Skip the hook when the request's active `budget_reservation` covers `spend:user:{user_id}`. The reservation is the source of truth for that counter cross-pod; the legacy `>=` path remains in place for requests without a reservation (e.g. paths that bypass the reservation entirely). Reproduces as `tests/otel_tests/test_prometheus.py::test_user_budget_metrics` on a fresh user with `max_budget=10` calling `fake-openai-endpoint` without `max_tokens`. Adds focused unit coverage in `tests/test_litellm/proxy/hooks/test_max_budget_limiter.py`. --- litellm/proxy/hooks/max_budget_limiter.py | 17 +- .../proxy/hooks/test_max_budget_limiter.py | 208 ++++++++++++++++++ 2 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/hooks/test_max_budget_limiter.py diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 7789fa6a34..9a7e511794 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -32,10 +32,25 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): if user_api_key_dict.team_id is not 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 + # when the reservation fills the counter to exactly max_budget. + # Imported lazily to avoid a circular import via proxy.utils. + from litellm.proxy.spend_tracking.budget_reservation import ( + get_reserved_counter_keys, + ) + + user_counter_key = f"spend:user:{user_id}" + if user_counter_key in get_reserved_counter_keys( + user_api_key_dict.budget_reservation + ): + return + from litellm.proxy.proxy_server import get_current_spend curr_spend = await get_current_spend( - counter_key=f"spend:user:{user_id}", + counter_key=user_counter_key, fallback_spend=user_api_key_dict.user_spend or 0.0, ) diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py new file mode 100644 index 0000000000..0074d7062b --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -0,0 +1,208 @@ +""" +Unit tests for the personal-budget pre-call hook. + +The reservation path (added in PR #26845) atomically pre-fills the same +`spend:user:{user_id}` counter this hook reads, admitting at a strict-`<` +boundary. Re-checking with `>=` after reservation would reject requests the +reservation already admitted when the reservation fills the counter to +exactly `max_budget` (e.g. requests with no `max_tokens` cap fall back to +reserving the smallest remaining headroom). + +These tests pin the skip-when-reserved behavior and guard against drift. +""" + +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter + + +def _make_user_api_key_auth( + user_id: str = "user-1", + user_max_budget: float = 10.0, + user_spend: float = 0.0, + team_id=None, + budget_reservation=None, +) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id=user_id, + user_max_budget=user_max_budget, + user_spend=user_spend, + team_id=team_id, + budget_reservation=budget_reservation, + ) + + +@pytest.mark.asyncio +async def test_under_budget_passes(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=3.0), + ): + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_over_budget_rejects_without_reservation(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + assert "Max budget limit reached." in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_skips_when_user_counter_is_reserved(): + """ + Reservation atomically pre-fills `spend:user:{user_id}` and admits the + request. The legacy `>=` check must not double-enforce on the same + counter — that's what produced the boundary regression where a fresh + user with no `max_tokens` cap got 429'd on their first request. + """ + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth( + user_id="user-1", + user_max_budget=10.0, + budget_reservation={ + "reserved_cost": 10.0, + "entries": [ + { + "counter_key": "spend:user:user-1", + "entity_type": "User", + "entity_id": "user-1", + "reserved_cost": 10.0, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + }, + ) + + # `get_current_spend` would return 10.0 here (counter pre-filled by the + # reservation). The hook must skip without reading it. + with 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={}, + call_type="completion", + ) + + assert result is None + mock_get_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_does_not_skip_when_reservation_covers_a_different_counter(): + """ + A reservation that only covers e.g. `spend:team:{team_id}` (not the user + counter) must not exempt the user-budget check. + """ + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth( + user_id="user-1", + user_max_budget=10.0, + budget_reservation={ + "reserved_cost": 5.0, + "entries": [ + { + "counter_key": "spend:team:team-x", + "entity_type": "Team", + "entity_id": "team-x", + "reserved_cost": 5.0, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + }, + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_team_keys_skip_personal_budget(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth( + user_max_budget=10.0, + team_id="team-1", + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=999.0), + ) as mock_get_spend: + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert result is None + mock_get_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_no_max_budget_passes(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", + user_id="user-1", + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=999.0), + ) as mock_get_spend: + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert result is None + mock_get_spend.assert_not_awaited()