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: