From 8b92b365733cfdbbbb3dc4c82c44d38bde116942 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 30 Jul 2026 13:19:49 -0700 Subject: [PATCH] revert(proxy)!: stop enforcing user budget on team keys (#35271) Reverts #32005. Team-scoped keys are governed by the team and team-member budgets only; the key owner personal max_budget no longer applies to them, restoring the hierarchy that existed before that PR. The skip_user_budget_on_team_key opt-out existed solely to turn the new behavior back off, so it is removed along with the behavior: the ConfigGeneralSettings field, the /config/list allowed_args entry that surfaced it as an Admin UI toggle, and the argument threaded through reserve_budget_for_request and _get_budget_counters. Regression tests cover both enforcement points in the restored direction: test_common_checks_personal_user_budget_skipped_for_team_key for the read-time check and test_should_not_reserve_user_budget_counter_for_team_key for the optimistic reservation path. (cherry picked from commit 6f1625d23bb1649394ad59c6a98c31633a1ba924) --- litellm/proxy/_types.py | 10 ---- litellm/proxy/auth/auth_checks.py | 41 ++++++------- litellm/proxy/auth/user_api_key_auth.py | 1 - litellm/proxy/proxy_server.py | 1 - .../spend_tracking/budget_reservation.py | 6 +- .../test_user_api_key_auth.py | 11 ++-- .../proxy/auth/test_auth_checks.py | 58 +++---------------- .../proxy/test_budget_reservation.py | 49 ++-------------- tests/test_litellm/proxy/test_proxy_server.py | 32 ---------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 -- 10 files changed, 42 insertions(+), 172 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 20d0d535b87..c19092dd49f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2455,16 +2455,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "is active as a reminder that hard enforcement is relaxed." ), ) - skip_user_budget_on_team_key: bool | None = Field( - None, - description=( - "If True, restores the legacy behavior where a user's personal " - "max_budget is NOT enforced when their key belongs to a team; only " - "the team (and team-member) budgets apply. Defaults to False, meaning " - "the user's personal max_budget is always enforced regardless of " - "whether the key belongs to a team (see GitHub issue #12905)." - ), - ) user_url_validation: Optional[bool] = Field( None, description=( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 07dfdc4fb43..6889acedc55 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -632,31 +632,28 @@ async def common_checks( ) async def _user_max_budget_check() -> None: - if user_object is None or user_object.max_budget is None: - return - skip_for_team = ( - general_settings.get("skip_user_budget_on_team_key") is True - and team_object is not None - and team_object.team_id is not None - ) - if skip_for_team: - return - from litellm.proxy.proxy_server import get_current_spend + # 4.1 personal budget, if personal key + if ( + (team_object is None or team_object.team_id is None) + and user_object is not None + and user_object.max_budget is not None + ): + from litellm.proxy.proxy_server import get_current_spend - user_budget = user_object.max_budget - user_spend = await get_current_spend( - counter_key=f"spend:user:{user_object.user_id}", - fallback_spend=user_object.spend or 0.0, - max_budget=user_budget, - ) - if math.isfinite(user_budget) and user_spend >= user_budget: - raise litellm.BudgetExceededError( - current_cost=user_spend, + user_budget = user_object.max_budget + user_spend = await get_current_spend( + counter_key=f"spend:user:{user_object.user_id}", + fallback_spend=user_object.spend or 0.0, max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", - entity_type=Litellm_EntityType.USER.value, - entity_id=user_object.user_id, ) + if math.isfinite(user_budget) and user_spend >= user_budget: + raise litellm.BudgetExceededError( + current_cost=user_spend, + max_budget=user_budget, + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", + entity_type=Litellm_EntityType.USER.value, + entity_id=user_object.user_id, + ) # Each scope reads a distinct counter key with no cross-scope ordering # dependency, so the per-scope Redis-first reads run concurrently instead diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e7c12448435..4d55f1404c6 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2461,7 +2461,6 @@ async def _reserve_budget_after_common_checks( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, - skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True, fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e0365b7d621..a70114437d9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15169,7 +15169,6 @@ async def get_config_list( "forward_client_headers_to_llm_api": {"type": "Boolean"}, "mcp_required_fields": {"type": "List"}, "cancel_on_disconnect": {"type": "Boolean"}, - "skip_user_budget_on_team_key": {"type": "Boolean"}, "disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"}, } diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 013873179c6..12ecd2ec300 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -155,7 +155,6 @@ async def reserve_budget_for_request( proxy_logging_obj: ProxyLogging, end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, - skip_user_budget_on_team_key: bool = False, fail_closed_budget_enforcement: bool = False, ) -> Optional[dict]: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): @@ -175,7 +174,6 @@ async def reserve_budget_for_request( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, - skip_user_budget_on_team_key=skip_user_budget_on_team_key, ) if not counters: return None @@ -333,7 +331,6 @@ async def _get_budget_counters( proxy_logging_obj: ProxyLogging, end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, - skip_user_budget_on_team_key: bool = False, ) -> List[_BudgetCounter]: counters: List[_BudgetCounter] = [] @@ -382,9 +379,8 @@ async def _get_budget_counters( ) ) - is_team_key = team_object is not None and team_object.team_id is not None if ( - not (is_team_key and skip_user_budget_on_team_key) + (team_object is None or team_object.team_id is None) and user_object is not None and user_object.user_id is not None and user_object.max_budget is not None diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 59c4caefa33..01dbb65a648 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -219,8 +219,8 @@ async def test_aaauser_personal_budgets(key_ownership): """ Set a personal budget on a user - User budget is enforced regardless of key ownership (personal or team). - Both cases should raise BudgetExceededError when the user is over budget. + - have it only apply when key belongs to user -> raises BudgetExceededError + - if key belongs to team, have key respect team budget -> allows call to go through """ import asyncio import time @@ -278,9 +278,12 @@ async def test_aaauser_personal_budgets(key_ownership): == valid_token ) - with pytest.raises(ProxyException) as exc_info: + if key_ownership == "user_key": + with pytest.raises(ProxyException) as exc_info: + await user_api_key_auth(request=request, api_key="Bearer " + user_key) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + else: await user_api_key_auth(request=request, api_key="Bearer " + user_key) - assert exc_info.value.type == ProxyErrorTypes.budget_exceeded @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index ccb20976df9..b56340ff1fb 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -4683,66 +4683,26 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): @pytest.mark.asyncio -async def test_user_budget_enforced_on_team_key(): - """User budget must be enforced even when the key belongs to a team. +async def test_common_checks_personal_user_budget_skipped_for_team_key(): + """A user's personal max_budget does not apply to a team-scoped key. - Previously _user_max_budget_check skipped enforcement for team keys, - letting a user with a $100 personal budget spend unlimited through a - team key. This regression test ensures that is no longer the case. + Team keys are governed by the team (and team-member) budgets only; the key + owner's personal budget is deliberately out of scope. This asserts the read + path lets a team key through even when the user is far over their personal + budget, and fails if personal enforcement is reintroduced for team keys. """ from fastapi import Request from litellm.proxy.auth.auth_checks import common_checks user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) - team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0) + team = LiteLLM_TeamTable(team_id="t1", spend=0.0, max_budget=1000.0) token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1") async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 - async def _no_membership(*a, **kw): - return None - - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter - ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): - with pytest.raises(litellm.BudgetExceededError) as over: - await common_checks( - request_body={"messages": [{"role": "user", "content": "hi"}]}, - team_object=team, - user_object=user, - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/chat/completions", - llm_router=None, - proxy_logging_obj=MagicMock(), - valid_token=token, - request=MagicMock(spec=Request), - ) - assert "User=u1" in str(over.value) - - -@pytest.mark.asyncio -async def test_skip_user_budget_on_team_key_flag_restores_old_behavior(): - """Setting skip_user_budget_on_team_key=True skips user budget for team keys. - - This is the opt-in escape hatch that restores the legacy behavior where - user budgets were not enforced when the key belonged to a team. - """ - from fastapi import Request - - from litellm.proxy.auth.auth_checks import common_checks - - user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) - team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0) - token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1") - - async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): - return 999.0 if counter_key == "spend:user:u1" else 0.0 - - async def _no_membership(*a, **kw): + async def _no_membership(*args, **kwargs): return None with patch("litellm.proxy.proxy_server.prisma_client", None), patch( @@ -4754,7 +4714,7 @@ async def test_skip_user_budget_on_team_key_flag_restores_old_behavior(): user_object=user, end_user_object=None, global_proxy_spend=None, - general_settings={"skip_user_budget_on_team_key": True}, + general_settings={}, route="/chat/completions", llm_router=None, proxy_logging_obj=MagicMock(), diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index e6ccbc579b5..1a584423fac 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -611,12 +611,12 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_ @pytest.mark.asyncio -async def test_should_reserve_user_budget_counter_for_team_key(spend_counter_state): - """A user's personal budget must be reserved even when the key belongs to a team. +async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter_state): + """The reservation path mirrors the read path: no personal user counter for a team key. - Regression for GitHub issue #12905: previously the reservation path skipped the - user spend counter whenever the key had a team, so a team key could overshoot the - user's personal max_budget under concurrency. + A team-scoped key reserves against the key and team counters only, so the key + owner's personal max_budget never gates a team request. Fails if the user + counter is reserved for team keys again. """ counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) @@ -645,44 +645,7 @@ async def test_should_reserve_user_budget_counter_for_team_key(spend_counter_sta proxy_logging_obj=proxy_logging_obj, ) - assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team") == pytest.approx(0.3) - - await release_budget_reservation(reservation) - - -@pytest.mark.asyncio -async def test_should_skip_user_budget_counter_for_team_key_when_flag_set(spend_counter_state): - """skip_user_budget_on_team_key=True restores the legacy behavior where a user's - personal budget is not reserved for a team key.""" - counter_cache, key_cache = spend_counter_state - proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) - valid_token = UserAPIKeyAuth( - token="key-user-on-team-skip", - spend=0.0, - user_id="user-on-team-skip", - team_id="team-no-budget-skip", - ) - team_object = LiteLLM_TeamTable(team_id="team-no-budget-skip", spend=0.0, max_budget=None) - user_object = LiteLLM_UserTable(user_id="user-on-team-skip", spend=0.0, max_budget=5.0) - - with patch( - "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", - return_value=0.3, - ): - reservation = await reserve_budget_for_request( - request_body=_request_body(), - route="/chat/completions", - llm_router=None, - valid_token=valid_token, - team_object=team_object, - user_object=user_object, - prisma_client=None, - user_api_key_cache=key_cache, - proxy_logging_obj=proxy_logging_obj, - skip_user_budget_on_team_key=True, - ) - - assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-skip") is None + assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team") is None await release_budget_reservation(reservation) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 536d24d4b4e..8c630211966 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9201,38 +9201,6 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): app.dependency_overrides.clear() -def test_get_config_list_includes_skip_user_budget_on_team_key(monkeypatch): - """Related to #12905: the opt-out flag must be discoverable via /config/list so - it renders as a Boolean toggle on the Admin UI General Settings table. This - requires both the ConfigGeneralSettings field and the allowed_args entry.""" - import types - from unittest.mock import AsyncMock, MagicMock - - from fastapi.testclient import TestClient - - import litellm.proxy.proxy_server as ps - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.proxy_server import app - - mock_prisma = MagicMock() - mock_config_table = MagicMock() - mock_config_table.find_first = AsyncMock(return_value=None) - mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) - monkeypatch.setattr(ps, "prisma_client", mock_prisma) - app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) - try: - client = TestClient(app) - resp = client.get("/config/list", params={"config_type": "general_settings"}) - assert resp.status_code == 200, resp.text - fields = {item["field_name"]: item for item in resp.json()} - assert "skip_user_budget_on_team_key" in fields - assert fields["skip_user_budget_on_team_key"]["field_type"] == "Boolean" - finally: - app.dependency_overrides.clear() - - def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch): """The throttle fraction is a litellm_settings scalar surfaced on the General Settings table as a Float field so it sits with the other global limits; it diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fc08864f19f..2b6ee759970 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22793,11 +22793,6 @@ export interface components { * @description When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata. */ reject_clientside_metadata_tags?: boolean | null; - /** - * Skip User Budget On Team Key - * @description If True, restores the legacy behavior where a user's personal max_budget is NOT enforced when their key belongs to a team; only the team (and team-member) budgets apply. Defaults to False, meaning the user's personal max_budget is always enforced regardless of whether the key belongs to a team (see GitHub issue #12905). - */ - skip_user_budget_on_team_key?: boolean | null; /** * Store Model In Db * @description If True, models and config are stored in and loaded from the database. Default is False.