fix(proxy): use >= in _team_max_budget_check to match other budget checks

Fixes #28020

_team_max_budget_check used `>` while every other budget enforcement
path (key, organization, budget windows) uses `>=`. A team whose spend
hit exactly `max_budget` was still allowed through.

Change `>` to `>=` and add a regression test that asserts
BudgetExceededError is raised when spend == max_budget.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Haofei Sun 2026-05-16 17:53:45 +08:00 committed by unknown
parent 80383a6172
commit 4406bb3867
2 changed files with 37 additions and 1 deletions

View file

@ -3694,7 +3694,7 @@ async def _team_max_budget_check(
fallback_spend=team_object.spend or 0.0,
)
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,

View file

@ -2297,6 +2297,42 @@ async def test_team_budget_check_reads_from_spend_counter():
assert exc_info.value.current_cost == 1.5
@pytest.mark.asyncio
async def test_team_budget_check_raises_when_spend_equals_max_budget():
"""Team budget check should raise BudgetExceededError when spend == max_budget,
matching the behavior of key / org / budget-window checks which all use >=.
Regression test for the inconsistency where _team_max_budget_check used `>`
while every other budget enforcement path used `>=`.
"""
from litellm.proxy.utils import ProxyLogging
team_object = LiteLLM_TeamTable(
team_id="test-team-equal",
spend=0.0,
max_budget=10.0,
)
valid_token = UserAPIKeyAuth(token="test-token", team_id="test-team-equal")
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):
if counter_key == "spend:team:test-team-equal":
return 10.0 # exactly at the cap
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 == 10.0
assert exc_info.value.max_budget == 10.0
@pytest.mark.asyncio
async def test_end_user_budget_check_reads_from_spend_counter():
"""End-user budget check should use get_current_spend when counter exists."""