scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100

revert(proxy)!: stop enforcing user budget on team keys (#35271)

Browse source 
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 6f1625d23b)
This commit is contained in:
yuneng-jiang 2026-07-30 13:19:49 -07:00 • committed by Yuneng Jiang
parent 38f2e023f1
commit 87c1fc3f41
No known key found for this signature in database
10 changed files with 42 additions and 172 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/_types.py
View file

@ -2433,16 +2433,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=(

41
litellm/proxy/auth/auth_checks.py
View file

@ -631,31 +631,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

1
litellm/proxy/auth/user_api_key_auth.py
View file

@ -2439,7 +2439,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,
)

1
litellm/proxy/proxy_server.py
View file

@ -15016,7 +15016,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"},
}

6
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -137,7 +137,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,
) -> Optional[dict]:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
return None
@ -156,7 +155,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
@ -312,7 +310,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] = []
@ -361,9 +358,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

11
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -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

58
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -4583,66 +4583,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(
@ -4654,7 +4614,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(),

49
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -610,12 +610,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)
@ -644,44 +644,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)

32
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -8875,38 +8875,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

5
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -22666,11 +22666,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.