From 28bf4647ef26b6dabc01961b55a135109bc37faa Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Mon, 4 May 2026 14:44:27 -0700 Subject: [PATCH] Treat 0 team_member_budget as no cap --- litellm/proxy/auth/auth_checks.py | 8 +- .../proxy/auth/test_auth_checks.py | 126 ++++++++++++++++++ 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 754488367e..1b6bd3aff2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5dedf05215..3433e3e6d8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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