scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

feat(proxy): add apply_user_budget_to_team_keys opt-in (#36102)

Browse source 
* feat(proxy): add apply_user_budget_to_team_keys opt-in

PR #32005 made a user's personal max_budget apply to their team-scoped keys
too, and PR #35271 reverted the whole thing (behavior plus the
skip_user_budget_on_team_key opt-out) because that flipped the default for
everyone. This brings the behavior back the other way round: default is
unchanged, and general_settings.apply_user_budget_to_team_keys opts a
deployment into charging the key owner's personal budget on team keys.

The flag reaches all three personal-budget gates so an opted-in deployment
enforces consistently: the read-time check in common_checks, the optimistic
reservation counter in _get_budget_counters, and the _PROXY_MaxBudgetLimiter
pre-call hook. It is also in the /config/list allowed args and, unlike the
reverted flag, in the _update_general_settings propagation allowlist, so the
Admin UI General Settings toggle actually takes effect at runtime; an explicit
YAML value still wins over the DB value on reload.

get_config_list's allowed_args moves to a module-level frozen mapping of
field name to type string, dropping 18 LIT002 violations and rebuilding one
less dict per request.

* style(proxy): drop explanatory comments from the budget flag paths
This commit is contained in:
ryan-crabbe-berri 2026-08-07 08:40:13 -07:00 • committed by GitHub
parent 83ab6e08da
commit 527dc0a8bb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 334 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/_types.py
View file

@ -2488,6 +2488,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"is active as a reminder that hard enforcement is relaxed."
),
)
apply_user_budget_to_team_keys: bool | None = Field(
None,
description=(
"If True, a user's personal max_budget is enforced on every request they "
"make, including requests made with a team-scoped key. Defaults to False, "
"where a team-scoped key is governed only by the team and team-member "
"budgets and the key owner's personal max_budget does not apply "
"(see GitHub issue #12905)."
),
)
user_url_validation: bool | None = Field(
None,
description=(

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

@ -648,28 +648,29 @@ async def common_checks(
)
async def _user_max_budget_check() -> None:
# 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
# 4.1 personal budget
if user_object is None or user_object.max_budget is None:
return
is_team_key: Final = team_object is not None and team_object.team_id is not None
if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True:
return
user_budget: Final = user_object.max_budget
user_spend: Final = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
from litellm.proxy.proxy_server import get_current_spend
user_budget: Final = user_object.max_budget
user_spend: Final = 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,
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

@ -2470,6 +2470,7 @@ 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,
apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
)

9
litellm/proxy/hooks/max_budget_limiter.py
View file

@ -32,9 +32,12 @@ class _PROXY_MaxBudgetLimiter(CustomLogger):
if max_budget is None or user_id is None:
return
# Personal budget applies only to non-team requests, matching
# the explicit team-key exemption in common_checks section 4.1.
if user_api_key_dict.team_id is not None:
from litellm.proxy.proxy_server import general_settings
if (
user_api_key_dict.team_id is not None
and general_settings.get("apply_user_budget_to_team_keys") is not True
):
return
# The reservation path admits at the strict-`<` boundary and

60
litellm/proxy/proxy_server.py
View file

@ -17,7 +17,7 @@ import traceback
import warnings
from collections.abc import AsyncGenerator, Callable, Mapping
from datetime import datetime, timedelta, timezone
from types import UnionType
from types import MappingProxyType, UnionType
from typing import (
TYPE_CHECKING,
Any,
@ -6096,6 +6096,15 @@ class ProxyConfig:
else:
general_settings["disable_auto_add_proxy_admin_to_teams"] = value if value is None else bool(value)
if "apply_user_budget_to_team_keys" in _general_settings and (
"apply_user_budget_to_team_keys" not in self._yaml_general_settings_keys
):
db_value: Final = _general_settings["apply_user_budget_to_team_keys"]
if isinstance(db_value, str):
general_settings["apply_user_budget_to_team_keys"] = db_value.lower() == "true"
else:
general_settings["apply_user_budget_to_team_keys"] = db_value if db_value is None else bool(db_value)
## STORE MODEL IN DB ##
if "store_model_in_db" in _general_settings:
value = _general_settings["store_model_in_db"]
@ -14982,6 +14991,29 @@ Keep it more precise, to prevent overwrite other values unintentially
_PLUGIN_KEY_REDACTED: Final = "***"
_GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingProxyType(
{
"max_parallel_requests": "Integer",
"global_max_parallel_requests": "Integer",
"max_request_size_mb": "Integer",
"max_response_size_mb": "Integer",
"proxy_config_reload_interval_seconds": "Integer",
"pass_through_endpoints": "PydanticModel",
"store_model_in_db": "Boolean",
"store_prompts_in_spend_logs": "Boolean",
"maximum_spend_logs_retention_period": "String",
"mcp_internal_ip_ranges": "List",
"mcp_trusted_proxy_ranges": "List",
"mcp_xff_num_trusted_hops": "Integer",
"always_include_stream_usage": "Boolean",
"forward_client_headers_to_llm_api": "Boolean",
"mcp_required_fields": "List",
"cancel_on_disconnect": "Boolean",
"disable_auto_add_proxy_admin_to_teams": "Boolean",
"apply_user_budget_to_team_keys": "Boolean",
}
)
def _preserve_redacted_plugin_keys(incoming: object, existing: object) -> object:
"""Restore real plugin_key values the client never sees.
@ -15480,25 +15512,7 @@ async def get_config_list(
else:
db_general_settings_dict = {}
allowed_args: Final = {
"max_parallel_requests": {"type": "Integer"},
"global_max_parallel_requests": {"type": "Integer"},
"max_request_size_mb": {"type": "Integer"},
"max_response_size_mb": {"type": "Integer"},
"proxy_config_reload_interval_seconds": {"type": "Integer"},
"pass_through_endpoints": {"type": "PydanticModel"},
"store_model_in_db": {"type": "Boolean"},
"store_prompts_in_spend_logs": {"type": "Boolean"},
"maximum_spend_logs_retention_period": {"type": "String"},
"mcp_internal_ip_ranges": {"type": "List"},
"mcp_trusted_proxy_ranges": {"type": "List"},
"mcp_xff_num_trusted_hops": {"type": "Integer"},
"always_include_stream_usage": {"type": "Boolean"},
"forward_client_headers_to_llm_api": {"type": "Boolean"},
"mcp_required_fields": {"type": "List"},
"cancel_on_disconnect": {"type": "Boolean"},
"disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"},
}
allowed_args: Final = _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES
return_val: Final = []
@ -15506,7 +15520,7 @@ async def get_config_list(
if field_name in allowed_args:
## HANDLE TYPED DICT
typed_dict_type = allowed_args[field_name]["type"]
typed_dict_type = allowed_args[field_name]
if typed_dict_type == "PydanticModel":
if field_name == "pass_through_endpoints":
@ -15548,7 +15562,7 @@ async def get_config_list(
_response_obj = ConfigList(
field_name=field_name,
field_type=allowed_args[field_name]["type"],
field_type=allowed_args[field_name],
field_description=field_info.description or "",
field_value=_redact_general_setting_value(
field_name,
@ -15576,7 +15590,7 @@ async def get_config_list(
_response_obj = ConfigList(
field_name=field_name,
field_type=allowed_args[field_name]["type"],
field_type=allowed_args[field_name],
field_description=field_info.description or "",
field_value=_redact_general_setting_value(field_name, _field_value, is_full_admin),
stored_in_db=_stored_in_db,

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

@ -156,6 +156,7 @@ async def reserve_budget_for_request(
proxy_logging_obj: ProxyLogging,
end_user_id: str | None = None,
end_user_object: Any | None = None,
apply_user_budget_to_team_keys: bool = False,
fail_closed_budget_enforcement: bool = False,
) -> dict | None:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
@ -175,6 +176,7 @@ async def reserve_budget_for_request(
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
apply_user_budget_to_team_keys=apply_user_budget_to_team_keys,
)
if not counters:
return None
@ -332,6 +334,7 @@ async def _get_budget_counters(
proxy_logging_obj: ProxyLogging,
end_user_id: str | None = None,
end_user_object: Any | None = None,
apply_user_budget_to_team_keys: bool = False,
) -> list[_BudgetCounter]:
counters: Final[list[_BudgetCounter]] = []
@ -380,8 +383,9 @@ async def _get_budget_counters(
)
)
is_team_key: Final = team_object is not None and team_object.team_id is not None
if (
(team_object is None or team_object.team_id is None)
(not is_team_key or apply_user_budget_to_team_keys)
and user_object is not None
and user_object.user_id is not None
and user_object.max_budget is not None

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

@ -5105,6 +5105,80 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key():
assert result is True
@pytest.mark.asyncio
async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag_enabled():
"""general_settings.apply_user_budget_to_team_keys opts a deployment into
charging the key owner's personal budget on team-scoped keys too.
Same fixture as the default-off test above, so a regression that ignores the
flag lets this call through instead of raising.
"""
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", 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(*args, **kwargs):
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 exc_info:
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={"apply_user_budget_to_team_keys": True},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=MagicMock(spec=Request),
)
assert "ExceededBudget: User=u1" in str(exc_info.value)
@pytest.mark.asyncio
async def test_common_checks_personal_user_budget_still_enforced_on_personal_key_with_flag_enabled():
"""The flag only widens enforcement to team keys; personal keys keep blocking."""
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)
token = UserAPIKeyAuth(token="k1", user_id="u1")
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
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
"litellm.proxy.proxy_server.get_current_spend", _spend_by_counter
):
with pytest.raises(litellm.BudgetExceededError):
await common_checks(
request_body={"messages": [{"role": "user", "content": "hi"}]},
team_object=None,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={"apply_user_budget_to_team_keys": True},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=MagicMock(spec=Request),
)
@pytest.mark.parametrize(
"scope, route, expect_blocked",
[

40
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -218,6 +218,46 @@ async def test_fail_closed_budget_enforcement_reaches_reservation(
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"general_settings,expected_flag",
[
({"apply_user_budget_to_team_keys": True}, True),
({"apply_user_budget_to_team_keys": False}, False),
({}, False),
],
)
async def test_apply_user_budget_to_team_keys_reaches_reservation(
general_settings, expected_flag
):
"""The opt-in lives in general_settings but is consumed inside
_get_budget_counters, so it has to be threaded through reserve_budget_for_request
or the reservation path keeps exempting team keys while the read path enforces."""
user_api_key_auth_obj = UserAPIKeyAuth(token="test_token")
with patch(
"litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request",
new=AsyncMock(return_value=None),
) as mock_reserve:
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request_data={"model": "gpt-4o"},
route="/v1/chat/completions",
llm_router=None,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
skip_budget_checks=False,
general_settings=general_settings,
)
assert (
mock_reserve.await_args.kwargs["apply_user_budget_to_team_keys"] is expected_flag
)
@pytest.mark.asyncio
async def test_should_not_reuse_cached_key_object_for_request_state():
key_cache = DualCache()

29
tests/test_litellm/proxy/hooks/test_max_budget_limiter.py
View file

@ -185,6 +185,35 @@ async def test_team_keys_skip_personal_budget():
mock_get_spend.assert_not_awaited()
@pytest.mark.asyncio
async def test_team_keys_enforce_personal_budget_when_flag_enabled():
"""This hook is the third personal-budget gate alongside common_checks and the
reservation path, so apply_user_budget_to_team_keys has to reach it too or an
opted-in deployment enforces in two places out of three."""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_max_budget=10.0,
team_id="team-1",
)
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"apply_user_budget_to_team_keys": True},
), patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=999.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
@pytest.mark.asyncio
async def test_no_max_budget_passes():
handler = _PROXY_MaxBudgetLimiter()

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

@ -650,6 +650,47 @@ async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter
await release_budget_reservation(reservation)
@pytest.mark.asyncio
async def test_should_reserve_user_budget_counter_for_team_key_when_flag_enabled(spend_counter_state):
"""apply_user_budget_to_team_keys must widen the reservation path too.
Read-time enforcement alone leaks budget under concurrency, so the opt-in has
to reserve against the personal counter as well or a burst of team-key
requests slips past the owner's max_budget.
"""
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-flagged",
spend=0.0,
user_id="user-on-team-flagged",
team_id="team-no-budget",
)
team_object = LiteLLM_TeamTable(team_id="team-no-budget", spend=0.0, max_budget=None)
user_object = LiteLLM_UserTable(user_id="user-on-team-flagged", 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,
apply_user_budget_to_team_keys=True,
)
assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-flagged") == pytest.approx(0.3)
await release_budget_reservation(reservation)
@pytest.mark.asyncio
async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_state):
counter_cache, key_cache = spend_counter_state

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

@ -7043,6 +7043,39 @@ async def test_update_general_settings_store_model_in_db_false():
assert ps.general_settings["store_model_in_db"] is False
@pytest.mark.asyncio
async def test_update_general_settings_propagates_apply_user_budget_to_team_keys():
"""The Admin UI toggle writes to the DB config, so the flag has to be in the
runtime propagation allowlist. The reverted skip_user_budget_on_team_key was
exposed in /config/list but never propagated, so its toggle did nothing."""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
with patch("litellm.proxy.proxy_server.general_settings", {}):
await proxy_config._update_general_settings(db_general_settings={"apply_user_budget_to_team_keys": "true"})
import litellm.proxy.proxy_server as ps
assert ps.general_settings["apply_user_budget_to_team_keys"] is True
@pytest.mark.asyncio
async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins():
"""A DB value must not silently override an explicit YAML setting on reload."""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
proxy_config._yaml_general_settings_keys = {"apply_user_budget_to_team_keys"}
with patch("litellm.proxy.proxy_server.general_settings", {"apply_user_budget_to_team_keys": True}):
await proxy_config._update_general_settings(db_general_settings={"apply_user_budget_to_team_keys": False})
import litellm.proxy.proxy_server as ps
assert ps.general_settings["apply_user_budget_to_team_keys"] is True
@pytest.mark.asyncio
@pytest.mark.parametrize(
"db_value,expected",
@ -9536,6 +9569,38 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch):
app.dependency_overrides.clear()
def test_get_config_list_includes_apply_user_budget_to_team_keys(monkeypatch):
"""Related to #12905: the opt-in must be discoverable via /config/list so it
renders as a Boolean toggle on the Admin UI General Settings table. This needs
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 "apply_user_budget_to_team_keys" in fields
assert fields["apply_user_budget_to_team_keys"]["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

2
type-discipline-budget.json
View file

@ -3,7 +3,7 @@
"limit": 23256
},
"LIT002": {
"limit": 27213
"limit": 27195
},
"LIT003": {
"limit": 269

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

@ -23389,6 +23389,11 @@ export interface components {
* @description Proxy API Endpoints you want users to be able to access
*/
allowed_routes?: unknown[] | null;
/**
* Apply User Budget To Team Keys
* @description If True, a user's personal max_budget is enforced on every request they make, including requests made with a team-scoped key. Defaults to False, where a team-scoped key is governed only by the team and team-member budgets and the key owner's personal max_budget does not apply (see GitHub issue #12905).
*/
apply_user_budget_to_team_keys?: boolean | null;
/**
* Background Health Checks
* @description run health checks in background