diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5b1134650f2..5b7dc3a7c73 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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=( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 01034d0cf58..3da899a5610 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9248576b599..9dc450befc2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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, ) diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index b8a69705b99..eaf37b0bcf1 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 14c6f8b8779..84f8685f5dc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 907c4ac7344..58a85171cc7 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 757991b8ff3..298f8a31b64 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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", [ diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index bc5bb877bd0..6fc44ef7519 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -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() diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py index 0074d7062b8..71671966d1a 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -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() diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1a584423fac..34adb4d2091 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index efd2ccb3e53..580a58885d9 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab8198304bb..2921592acbd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23256 }, "LIT002": { - "limit": 27213 + "limit": 27195 }, "LIT003": { "limit": 269 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 752572f9863..7675c0506c0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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