Merge pull request #27133 from BerriAI/litellm_zeroBudgetTreatedAsNoCap

[Fix] Treat 0 team_member_budget as no cap
This commit is contained in:
Michael-RZ-Berri
2026-05-04 15:05:18 -07:00
committed by GitHub
2 changed files with 133 additions and 1 deletions
+7 -1
View File
@@ -3547,7 +3547,13 @@ async def _check_team_member_budget(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
if default_budget is not None:
# Treat 0 on the team default as "no cap".
# Per-member rows still respect 0 as an explicit admin disable.
if (
default_budget is not None
and default_budget.max_budget is not None
and default_budget.max_budget > 0
):
team_member_budget = default_budget.max_budget
if team_member_budget is not None:
@@ -2562,3 +2562,129 @@ async def test_team_member_budget_check_null_clone_with_null_default_skips_enfor
user_api_key_cache=DualCache(),
proxy_logging_obj=proxy_logging_obj,
)
@pytest.mark.asyncio
async def test_team_member_budget_check_zero_team_default_treated_as_no_cap():
"""A team default budget with max_budget=0.0 (likely a stale/accidental
write) must not block every member. The fallback path treats 0 as
"no cap"; per-member rows still respect 0 as an explicit disable."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LiteLLM_TeamMembership
from litellm.proxy.utils import ProxyLogging
team_object = LiteLLM_TeamTable(
team_id="test-team",
metadata={"team_member_budget_id": "budget-default"},
)
user_object = LiteLLM_UserTable(user_id="test-user")
valid_token = UserAPIKeyAuth(
token="test-token",
user_id="test-user",
team_id="test-team",
)
# No per-member row -> falls through to team default.
team_membership = LiteLLM_TeamMembership(
user_id="test-user",
team_id="test-team",
spend=0.0,
budget_id=None,
litellm_budget_table=None,
)
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
# Team default budget row with max_budget=0.0 (the regression trigger).
fake_default_row = MagicMock()
fake_default_row.max_budget = 0.0
fake_default_row.dict = MagicMock(
return_value={"budget_id": "budget-default", "max_budget": 0.0}
)
prisma_client = MagicMock()
prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
return_value=fake_default_row
)
async def mock_get_current_spend(counter_key, fallback_spend):
if counter_key == "spend:team_member:test-user:test-team":
return 0.0
return fallback_spend
with (
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend),
patch(
"litellm.proxy.auth.auth_checks.get_team_membership",
new_callable=AsyncMock,
return_value=team_membership,
),
):
# No raise: 0.0 cap is treated as "no cap configured".
await _check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
proxy_logging_obj=proxy_logging_obj,
)
@pytest.mark.asyncio
async def test_team_member_budget_check_zero_per_member_row_still_blocks():
"""A per-member row with max_budget=0.0 is treated as an explicit admin
disable - enforcement still blocks. Only the team-default fallback
path treats 0 as no cap."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership
from litellm.proxy.utils import ProxyLogging
team_object = LiteLLM_TeamTable(
team_id="test-team",
metadata={"team_member_budget_id": "budget-default"},
)
user_object = LiteLLM_UserTable(user_id="test-user")
valid_token = UserAPIKeyAuth(
token="test-token",
user_id="test-user",
team_id="test-team",
)
# Per-member row with max_budget=0.0 - admin intent: disable this user.
team_membership = LiteLLM_TeamMembership(
user_id="test-user",
team_id="test-team",
spend=0.0,
budget_id="budget-disable",
litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.0),
)
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
prisma_client = MagicMock()
prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None)
async def mock_get_current_spend(counter_key, fallback_spend):
if counter_key == "spend:team_member:test-user:test-team":
return 0.0
return fallback_spend
with (
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend),
patch(
"litellm.proxy.auth.auth_checks.get_team_membership",
new_callable=AsyncMock,
return_value=team_membership,
),
):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
proxy_logging_obj=proxy_logging_obj,
)
assert exc_info.value.max_budget == 0.0