mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
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:
parent
80383a6172
commit
4406bb3867
2 changed files with 37 additions and 1 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue