From 2b9ad4d4ebd454aff6216fbd7daa43d432fcdeaa Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Wed, 29 Apr 2026 10:14:00 -0700 Subject: [PATCH 1/3] Fall through to team default when per-member budget max_budget is NULL --- litellm/proxy/auth/auth_checks.py | 4 + .../proxy/auth/test_auth_checks.py | 141 ++++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 840f64cfed..ebf83c2226 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3289,10 +3289,14 @@ async def _check_team_member_budget( # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. + # A per-member row whose max_budget is NULL is *not* an override - + # it can result from cloning a team default that itself had no cap + # at member-add time. Treat it as "no override" and fall through. team_member_budget: Optional[float] = None if ( team_membership is not None and team_membership.litellm_budget_table is not None + and team_membership.litellm_budget_table.max_budget is not None ): team_member_budget = team_membership.litellm_budget_table.max_budget else: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 676a32c202..e7a3e4bcac 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2336,3 +2336,144 @@ async def test_team_member_budget_check_per_member_override_wins_over_team_defau ) assert exc_info.value.current_cost == 250.0 assert exc_info.value.max_budget == 200.0 + + +@pytest.mark.asyncio +async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): + """A per-member budget row with max_budget=NULL is not an explicit + no-cap override - it can be the result of cloning a team default that + had no value at member-add time. Enforcement must fall through to the + team default and apply that cap. + + Pre-fix: NULL on the clone short-circuited the comparison block + (team_member_budget = None -> if not None: skipped) so the user + spent unbounded against an apparent $65/$X 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 exists with NULL max_budget (the cloned-from-incomplete-default case). + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + budget_id="budget-clone", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=None), + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + + fake_default_row = MagicMock() + fake_default_row.max_budget = 65.0 + fake_default_row.dict = MagicMock( + return_value={"budget_id": "budget-default", "max_budget": 65.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 500.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.current_cost == 500.0 + assert exc_info.value.max_budget == 65.0 + prisma_client.db.litellm_budgettable.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_team_member_budget_check_null_clone_with_null_default_skips_enforcement(): + """Sanity check: when both the per-member clone AND the team default + have max_budget=NULL, enforcement still skips (the team genuinely has + no cap configured). Confirms the NULL-fall-through is defensive, not + overzealous.""" + 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", + ) + + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + budget_id="budget-clone", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=None), + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + + fake_default_row = MagicMock() + fake_default_row.max_budget = None + fake_default_row.dict = MagicMock( + return_value={"budget_id": "budget-default", "max_budget": None} + ) + + 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 1000.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: both rows are NULL, so enforcement is correctly skipped. + 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, + ) From 04687ba48e706d4d09a578c941fcd1a72b51c0e9 Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Wed, 29 Apr 2026 11:07:35 -0700 Subject: [PATCH 2/3] Trim verbose comments and docstrings --- litellm/proxy/auth/auth_checks.py | 3 --- tests/test_litellm/proxy/auth/test_auth_checks.py | 14 ++------------ 2 files changed, 2 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index ebf83c2226..4d19fb2e35 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3289,9 +3289,6 @@ async def _check_team_member_budget( # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. - # A per-member row whose max_budget is NULL is *not* an override - - # it can result from cloning a team default that itself had no cap - # at member-add time. Treat it as "no override" and fall through. team_member_budget: Optional[float] = None if ( team_membership 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 e7a3e4bcac..841b45cef8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2340,14 +2340,7 @@ async def test_team_member_budget_check_per_member_override_wins_over_team_defau @pytest.mark.asyncio async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): - """A per-member budget row with max_budget=NULL is not an explicit - no-cap override - it can be the result of cloning a team default that - had no value at member-add time. Enforcement must fall through to the - team default and apply that cap. - - Pre-fix: NULL on the clone short-circuited the comparison block - (team_member_budget = None -> if not None: skipped) so the user - spent unbounded against an apparent $65/$X cap.""" + """Per-member NULL max_budget falls through to the team default cap.""" from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership from litellm.proxy.utils import ProxyLogging @@ -2415,10 +2408,7 @@ async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): @pytest.mark.asyncio async def test_team_member_budget_check_null_clone_with_null_default_skips_enforcement(): - """Sanity check: when both the per-member clone AND the team default - have max_budget=NULL, enforcement still skips (the team genuinely has - no cap configured). Confirms the NULL-fall-through is defensive, not - overzealous.""" + """When per-member and team default are both NULL, enforcement still skips.""" from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership from litellm.proxy.utils import ProxyLogging From b5b07089dd0219a44649e78cfcf84319d1320d63 Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Wed, 29 Apr 2026 14:08:28 -0700 Subject: [PATCH 3/3] Update get_team_member_default_budget docstring for NULL fallback --- litellm/proxy/auth/auth_checks.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4d19fb2e35..3f9cd4cf86 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -915,7 +915,8 @@ async def get_team_member_default_budget( Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"]. This budget is applied to team members whose TeamMembership row has no - linked budget. Results are cached for performance. + linked budget, or whose linked budget has max_budget=NULL. Results are + cached for performance. Args: budget_id: The budget_id pulled from team.metadata["team_member_budget_id"]