From b59ba74733be454e0058bbfbdae4d19d486a88a0 Mon Sep 17 00:00:00 2001 From: Jahanzeb-git Date: Fri, 16 Oct 2026 23:38:10 +0500 Subject: [PATCH 1/6] fix(fallback_budget): enforce team member budget on router fallback targets --- litellm/proxy/auth/fallback_budget.py | 16 ++++++++++++++++ litellm/types/proxy/auth/user_api_key_auth.py | 1 + tests/unit/proxy/auth/test_fallback_budget.py | 10 ++++++++++ 3 files changed, 27 insertions(+) diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py index a126a53821b..397b3cf7301 100644 --- a/litellm/proxy/auth/fallback_budget.py +++ b/litellm/proxy/auth/fallback_budget.py @@ -130,6 +130,22 @@ async def is_token_within_budget_for_model( if is_model_cost_zero(model=model, llm_router=llm_router): return True + # Gating check: team member overall budget within their team. + # Prevents zero-cost models with paid fallbacks from exceeding the member's team budget cap. + team_member_budget: Final = valid_token.team_member_budget + if ( + team_member_budget is not None + and valid_token.user_id is not None + and valid_token.team_id is not None + ): + team_member_spend: Final = await _counter_spend( + counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", + fallback_spend=valid_token.team_member_spend or 0.0, + max_budget=team_member_budget, + ) + if team_member_spend >= team_member_budget: + return False + key_budget: Final = valid_token.max_budget if key_budget is not None and valid_token.token is not None: key_spend: Final = await _counter_spend( diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index c2f7c487b7b..6e9c115a2e6 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -78,6 +78,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_member_spend: float | None = None team_member_tpm_limit: int | None = None team_member_rpm_limit: int | None = None + team_member_budget: float | None = None # End User Params end_user_id: str | None = None diff --git a/tests/unit/proxy/auth/test_fallback_budget.py b/tests/unit/proxy/auth/test_fallback_budget.py index 1117ca43657..4e7574819da 100644 --- a/tests/unit/proxy/auth/test_fallback_budget.py +++ b/tests/unit/proxy/auth/test_fallback_budget.py @@ -115,6 +115,16 @@ async def test_paid_target_refused_when_over_user_budget(): token = _token(user_spend=1900.0, user_max_budget=50.0) assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False +@pytest.mark.asyncio +async def test_paid_target_refused_when_over_team_member_budget(): + token = _token(user_id="u1", team_id="t1", team_member_budget=50.0, team_member_spend=100.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + +@pytest.mark.asyncio +async def test_paid_target_allowed_when_under_team_member_budget(): + token = _token(user_id="u1", team_id="t1", team_member_budget=50.0, team_member_spend=10.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True + @pytest.mark.asyncio async def test_team_member_budget_refuses_an_exhausted_paid_target() -> None: From 2509efb9e28d753b840630ef69998356886603c2 Mon Sep 17 00:00:00 2001 From: Jahanzeb-git Date: Sat, 17 Oct 2026 00:59:00 +0500 Subject: [PATCH 2/6] fix(auth): populate team_member_budget on token for fallback budget checks --- litellm/proxy/auth/auth_checks.py | 2 ++ litellm/proxy/auth/team_grants.py | 6 ++++++ litellm/proxy/auth/user_api_key_auth.py | 1 + litellm/proxy/utils.py | 1 + 4 files changed, 10 insertions(+) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8b32eb450e8..bfaf2c98bb5 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6753,6 +6753,8 @@ async def _check_team_member_budget( ) if team_member_budget is not None: + if valid_token is not None: + valid_token.team_member_budget = team_member_budget team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py index 450a3ce456e..c782b823884 100644 --- a/litellm/proxy/auth/team_grants.py +++ b/litellm/proxy/auth/team_grants.py @@ -72,6 +72,7 @@ class TeamGrants(TypedDict, total=False): team_member_spend: ReadOnly[float | None] team_member_tpm_limit: ReadOnly[int | None] team_member_rpm_limit: ReadOnly[int | None] + team_member_budget: ReadOnly[float | None] def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns: @@ -124,4 +125,9 @@ def team_grants( team_member_rpm_limit=( team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None ), + team_member_budget=( + team_membership.litellm_budget_table.max_budget + if team_membership is not None and team_membership.litellm_budget_table is not None + else None + ), ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a1abed202a2..76bb52c6323 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2956,6 +2956,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget( now=datetime.now(timezone.utc), ) + valid_token.team_member_budget = team_member_budget if team_member_budget is not None and team_member_budget > 0: # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index be6532ba36b..f8e206fc719 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5443,6 +5443,7 @@ class PrismaClient: tm.spend AS team_member_spend, b_tm.tpm_limit AS team_member_tpm_limit, b_tm.rpm_limit AS team_member_rpm_limit, + b_tm.max_budget AS team_member_budget, m.aliases AS team_model_aliases, -- Added comma to separate b.* columns b.max_budget AS litellm_budget_table_max_budget, From f15b27b0a9e8d1ef9230868419c0812eb86549cf Mon Sep 17 00:00:00 2001 From: Jahanzeb-git Date: Sun, 18 Oct 2026 16:35:08 +0500 Subject: [PATCH 3/6] fix(lint): remove trailing whitespace, update team_grants test, and sync UI schema types --- litellm/proxy/auth/fallback_budget.py | 2 +- tests/unit/proxy/auth/test_team_grants.py | 3 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 ++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py index 397b3cf7301..aef6cc44520 100644 --- a/litellm/proxy/auth/fallback_budget.py +++ b/litellm/proxy/auth/fallback_budget.py @@ -130,7 +130,7 @@ async def is_token_within_budget_for_model( if is_model_cost_zero(model=model, llm_router=llm_router): return True - # Gating check: team member overall budget within their team. + # Gating check: team member overall budget within their team. # Prevents zero-cost models with paid fallbacks from exceeding the member's team budget cap. team_member_budget: Final = valid_token.team_member_budget if ( diff --git a/tests/unit/proxy/auth/test_team_grants.py b/tests/unit/proxy/auth/test_team_grants.py index f74531beaf0..39f9e063cde 100644 --- a/tests/unit/proxy/auth/test_team_grants.py +++ b/tests/unit/proxy/auth/test_team_grants.py @@ -50,7 +50,7 @@ def _membership() -> LiteLLM_TeamMembership: user_id=USER_ID, team_id=TEAM_ID, spend=3.25, - litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=500, rpm_limit=5), + litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=500, rpm_limit=5, max_budget=100.0), ) @@ -85,6 +85,7 @@ def test_team_grants_cover_every_team_field_the_key_path_gets(): assert token.team_member_spend == 3.25 assert token.team_member_tpm_limit == 500 assert token.team_member_rpm_limit == 5 + assert token.team_member_budget == 100.0 def test_team_grants_without_team_leave_token_defaults(): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6f6eae1f98e..0067f81fcf6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -49118,6 +49118,8 @@ export interface components { /** Team Max Budget */ team_max_budget?: number | null; team_member?: components["schemas"]["Member"] | null; + /** Team Member Budget */ + team_member_budget?: number | null; /** Team Member Rpm Limit */ team_member_rpm_limit?: number | null; /** Team Member Spend */ From 4f9db2be37254fbcfad934429f3123d4044a1ea4 Mon Sep 17 00:00:00 2001 From: Jahanzeb-git Date: Sun, 18 Oct 2026 16:57:09 +0500 Subject: [PATCH 4/6] style(auth): format fallback_budget condition with ruff --- litellm/proxy/auth/fallback_budget.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py index aef6cc44520..bd85398fd4c 100644 --- a/litellm/proxy/auth/fallback_budget.py +++ b/litellm/proxy/auth/fallback_budget.py @@ -133,11 +133,7 @@ async def is_token_within_budget_for_model( # Gating check: team member overall budget within their team. # Prevents zero-cost models with paid fallbacks from exceeding the member's team budget cap. team_member_budget: Final = valid_token.team_member_budget - if ( - team_member_budget is not None - and valid_token.user_id is not None - and valid_token.team_id is not None - ): + if team_member_budget is not None and valid_token.user_id is not None and valid_token.team_id is not None: team_member_spend: Final = await _counter_spend( counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", fallback_spend=valid_token.team_member_spend or 0.0, From 3cbd0faea9a207bd755212100deb2a34982c591d Mon Sep 17 00:00:00 2001 From: Jahanzeb-git Date: Sun, 18 Oct 2026 17:25:59 +0500 Subject: [PATCH 5/6] chore(types): add rebind-ok suppressions for valid_token attribute assignments --- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index bfaf2c98bb5..a26dee4a5d7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6754,7 +6754,7 @@ async def _check_team_member_budget( if team_member_budget is not None: if valid_token is not None: - valid_token.team_member_budget = team_member_budget + valid_token.team_member_budget = team_member_budget # rebind-ok: carrier for router fallback checks team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 76bb52c6323..4cc05896305 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2956,7 +2956,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget( now=datetime.now(timezone.utc), ) - valid_token.team_member_budget = team_member_budget + valid_token.team_member_budget = team_member_budget # rebind-ok: carrier for router fallback checks if team_member_budget is not None and team_member_budget > 0: # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend From 7b894385b3d8e40e391220f4a54ca804c4e13f48 Mon Sep 17 00:00:00 2001 From: Jahanzeb-git Date: Sun, 18 Oct 2026 17:44:50 +0500 Subject: [PATCH 6/6] fix(types): remove redundant valid_token None check in auth-checks --- litellm/proxy/auth/auth_checks.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a26dee4a5d7..2b63646a25a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6753,8 +6753,7 @@ async def _check_team_member_budget( ) if team_member_budget is not None: - if valid_token is not None: - valid_token.team_member_budget = team_member_budget # rebind-ok: carrier for router fallback checks + valid_token.team_member_budget = team_member_budget # rebind-ok: carrier for router fallback checks team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available