diff --git a/litellm/__init__.py b/litellm/__init__.py index 2ec0830d622..6e2a03b7c7c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -379,6 +379,7 @@ budget_duration: Optional[str] = ( None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). ) default_soft_budget: float = DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 +budget_exceeded_throttle_percentage: Optional[float] = None forward_traceparent_to_llm_provider: bool = False diff --git a/litellm/constants.py b/litellm/constants.py index 1300668cc70..7423d9b2211 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1504,6 +1504,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "public_model_groups_links", "cost_discount_config", "cost_margin_config", + "budget_exceeded_throttle_percentage", ] SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e1d657f293b..3c16c2c3ed7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1070,6 +1070,7 @@ class KeyRequestBase(GenerateRequestBase): budget_id: Optional[str] = None tags: Optional[List[str]] = None disable_global_guardrails: Optional[bool] = None + throttle_on_budget_exceeded: Optional[bool] = None enforced_params: Optional[List[str]] = None allowed_routes: Optional[list] = [] allowed_passthrough_routes: Optional[list] = None @@ -2469,6 +2470,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob request_route: Optional[str] = None is_session_token: bool = False budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True) + budget_throttle_pct: Optional[float] = Field(default=None, exclude=True) user: Optional[Any] = None # Expanded user object when expand=user is used created_by_user: Optional[Any] = None # Expanded created_by user when expand=user is used end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None @@ -3859,6 +3861,7 @@ LiteLLM_ManagementEndpoint_MetadataFields = [ "allowed_vector_store_indexes", "enforced_batch_output_expires_after", "enforced_file_expires_after", + "throttle_on_budget_exceeded", ] LiteLLM_ManagementEndpoint_MetadataFields_Premium = [ diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b749e9fbe0d..ee548ba0a43 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -60,6 +60,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.budget_throttle import ( + budget_throttle_percentage, + should_throttle_budget_exceeded, +) from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( @@ -2496,6 +2500,7 @@ def _copy_user_api_key_auth_for_cache( ) -> UserAPIKeyAuth: copied_key_obj = user_api_key_obj.model_copy() copied_key_obj.budget_reservation = None + copied_key_obj.budget_throttle_pct = None copied_key_obj.parent_otel_span = None copied_key_obj.request_route = None return copied_key_obj @@ -3428,6 +3433,24 @@ async def is_valid_fallback_model( return True +def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: + """ + Throttle an over-budget key instead of blocking it, when the key opted in + via `throttle_on_budget_exceeded` and a global percentage is configured. + + Records the percentage on the request-scoped `budget_throttle_pct` so the + rate limiter scales the key's TPM/RPM down to it; the persistent limits are + left untouched so the throttle never compounds across requests. Returns True + when the key was throttled (caller skips raising), False when it should still + be hard-blocked. + """ + pct = budget_throttle_percentage() + if pct is None or not should_throttle_budget_exceeded(valid_token): + return False + valid_token.budget_throttle_pct = pct + return True + + async def _virtual_key_max_budget_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, @@ -3488,6 +3511,8 @@ async def _virtual_key_max_budget_check( # so a NaN max_budget would silently disable enforcement. Treat a # non-finite max_budget as "no configured limit" rather than as a bypass. if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: + if _apply_budget_exceeded_throttle(valid_token): + return # name the key in the error so operators don't have to reverse-map # spend back to a key; key_name is the masked form (last 4 chars) key_label = valid_token.key_alias or "key" diff --git a/litellm/proxy/auth/budget_throttle.py b/litellm/proxy/auth/budget_throttle.py new file mode 100644 index 00000000000..19dffee462b --- /dev/null +++ b/litellm/proxy/auth/budget_throttle.py @@ -0,0 +1,56 @@ +""" +Throttle a key after it exceeds its own ``max_budget`` instead of blocking it. + +When a key opts in via ``throttle_on_budget_exceeded`` and a global +``budget_exceeded_throttle_percentage`` is configured, an over-budget key keeps +serving requests but at a reduced TPM/RPM (the configured percentage of its +configured limits). The decision (over budget + opted in) is made once during +auth; the scaling is recomputed from the key's original limits on every request +so it never compounds across requests. +""" + +import math +from typing import Optional + +import litellm +from litellm.proxy._types import UserAPIKeyAuth + + +def budget_throttle_percentage() -> Optional[float]: + """ + The global throttle percentage, or None when throttling is disabled / + misconfigured (in which case an over-budget key is hard-blocked, the safe + default). + """ + pct = litellm.budget_exceeded_throttle_percentage + if not isinstance(pct, (int, float)) or isinstance(pct, bool): + return None + if not 0 < pct <= 1: + return None + return float(pct) + + +def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: + """ + True when a key that exceeded its own ``max_budget`` should be throttled + rather than blocked: it opted in, a valid global percentage is set, and the + key has a TPM or RPM limit to scale down. A key with neither limit has + nothing to throttle, so it stays hard-blocked (the safe default) rather than + serving unlimited requests past its budget. + """ + if (valid_token.metadata or {}).get("throttle_on_budget_exceeded") is not True: + return False + if valid_token.tpm_limit is None and valid_token.rpm_limit is None: + return False + return budget_throttle_percentage() is not None + + +def throttled_limit(limit: Optional[int], pct: Optional[float]) -> Optional[int]: + """ + Scale a TPM/RPM limit to ``pct`` of its value, keeping a trickle of at least + 1 so a throttled key is slowed rather than fully locked out. An unset limit + or unset percentage leaves the limit unchanged. + """ + if limit is None or pct is None: + return limit + return max(1, math.floor(limit * pct)) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 7944bb54d67..2613510bd0c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2004,6 +2004,11 @@ async def _user_api_key_auth_builder( valid_token_dict = valid_token.model_dump(exclude_none=True) valid_token_dict.pop("token", None) + # budget_throttle_pct is excluded from model_dump (it must not leak + # into serialized responses), so carry the request-scoped decision + # forward by hand to the auth object the rate limiter receives. + if valid_token.budget_throttle_pct is not None: + valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct if _end_user_object is not None: valid_token_dict.update(end_user_params) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 1ed76d5b1e3..ee6abb13d6b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -17,6 +17,7 @@ from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, get_key_model_tpm_limit, ) +from litellm.proxy.auth.budget_throttle import throttled_limit from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit @@ -248,10 +249,11 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if data is None: data = {} global_max_parallel_requests = data.get("metadata", {}).get("global_max_parallel_requests", None) - tpm_limit = getattr(user_api_key_dict, "tpm_limit", sys.maxsize) + throttle_pct = getattr(user_api_key_dict, "budget_throttle_pct", None) + tpm_limit = throttled_limit(getattr(user_api_key_dict, "tpm_limit", sys.maxsize), throttle_pct) if tpm_limit is None: tpm_limit = sys.maxsize - rpm_limit = getattr(user_api_key_dict, "rpm_limit", sys.maxsize) + rpm_limit = throttled_limit(getattr(user_api_key_dict, "rpm_limit", sys.maxsize), throttle_pct) if rpm_limit is None: rpm_limit = sys.maxsize diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 78c43715dad..ee0a0e1789d 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -32,6 +32,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import get_model_rate_limit_from_metadata +from litellm.proxy.auth.budget_throttle import throttled_limit from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, map_v3_rate_limit_type, @@ -1549,18 +1550,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): or user_api_key_dict.tpm_limit is not None or user_api_key_dict.max_parallel_requests is not None ): + throttle_pct = user_api_key_dict.budget_throttle_pct descriptors.append( RateLimitDescriptor( key="api_key", value=user_api_key_dict.api_key, rate_limit={ "requests_per_unit": self._get_enforced_limit( - limit_value=user_api_key_dict.rpm_limit, + limit_value=throttled_limit(user_api_key_dict.rpm_limit, throttle_pct), limit_type=rpm_limit_type, model_has_failures=model_has_failures, ), "tokens_per_unit": self._get_enforced_limit( - limit_value=user_api_key_dict.tpm_limit, + limit_value=throttled_limit(user_api_key_dict.tpm_limit, throttle_pct), limit_type=tpm_limit_type, model_has_failures=model_has_failures, ), diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d92aea57063..71cf2db3dfb 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -758,6 +758,12 @@ async def _common_key_generation_helper( premium_user=premium_user, ) + if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + raise HTTPException( + status_code=403, + detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."}, + ) + if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None: await validate_team_id_used_in_service_account_request( team_id=data.team_id, @@ -1483,6 +1489,7 @@ async def generate_key_fn( - guardrails: Optional[List[str]] - List of active guardrails for the key - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. + - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. @@ -2317,6 +2324,16 @@ async def _validate_update_key_data( or "budget_limits" in data.model_fields_set ) + _existing_metadata = getattr(existing_key_row, "metadata", None) + _existing_throttle = ( + _existing_metadata.get("throttle_on_budget_exceeded") if isinstance(_existing_metadata, dict) else None + ) + if data.throttle_on_budget_exceeded is True and _existing_throttle is not True and not _is_proxy_admin: + raise HTTPException( + status_code=403, + detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."}, + ) + # Personal-key bypass: the caller both created the key AND still owns it # (user_id == caller). Checking only created_by would let a demoted admin # who originally created a key for another user continue editing it without @@ -2507,6 +2524,7 @@ async def update_key_fn( - guardrails: Optional[List[str]] - List of active guardrails for the key - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. + - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. - blocked: Optional[bool] - Whether the key is blocked - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 810e94cdc27..38788d140e9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14285,6 +14285,9 @@ async def update_config_general_settings( detail={"error": CommonProxyErrors.not_allowed_access.value}, ) + if data.field_name in _GENERAL_SETTINGS_UI_LITELLM_FIELDS: + return await _persist_general_settings_ui_litellm_field(data.field_name, data.field_value, user_api_key_dict) + if data.field_name not in ConfigGeneralSettings.model_fields: raise HTTPException( status_code=400, @@ -14550,6 +14553,55 @@ async def get_config_general_settings( ) +_GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, dict[str, str]] = { + "budget_exceeded_throttle_percentage": { + "type": "Float", + "description": ( + "Fraction (0, 1] of a key's configured TPM/RPM that an over-budget key with " + "'Throttle on budget exceeded' enabled keeps serving at. Leave empty to hard-block " + "over-budget keys." + ), + }, +} + + +def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> Optional[float]: + if value is None or value == "": + return None + if isinstance(value, bool) or not isinstance(value, (int, float)) or not (0 < float(value) <= 1): + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be a number in (0, 1] or empty"}, + ) + return float(value) + + +async def _persist_general_settings_ui_litellm_field( + field_name: str, value: Any, user_api_key_dict: UserAPIKeyAuth +) -> dict: + validated = _validate_general_settings_ui_litellm_value(field_name, value) + config = await proxy_config.get_config() + before_value = config.get("litellm_settings", {}).get(field_name) + setattr(litellm, field_name, validated) + if "litellm_settings" not in config: + config["litellm_settings"] = {} + config["litellm_settings"][field_name] = validated + await proxy_config.save_config(new_config=config) + asyncio.create_task(create_config_audit_log(field_name, "updated", before_value, validated, user_api_key_dict)) + return {"message": f"Field {field_name} updated", "status": "success"} + + +async def _reset_general_settings_ui_litellm_field(field_name: str, user_api_key_dict: UserAPIKeyAuth) -> dict: + config = await proxy_config.get_config() + before_value = config.get("litellm_settings", {}).get(field_name) + setattr(litellm, field_name, None) + if "litellm_settings" in config: + config["litellm_settings"].pop(field_name, None) + await proxy_config.save_config(new_config=config) + asyncio.create_task(create_config_audit_log(field_name, "deleted", before_value, None, user_api_key_dict)) + return {"message": f"Field {field_name} reset", "status": "success"} + + @router.get( "/config/list", tags=["config.yaml"], @@ -14703,6 +14755,35 @@ async def get_config_list( ) return_val.append(_response_obj) + db_litellm_settings_row = await ConfigRepository(prisma_client).table.find_first( + where={"param_name": "litellm_settings"} + ) + db_litellm_settings: dict = ( + dict(db_litellm_settings_row.param_value) + if db_litellm_settings_row is not None and db_litellm_settings_row.param_value is not None + else {} + ) + for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items(): + current_value: Optional[float] = getattr(litellm, litellm_field_name, None) + stored_in_db_litellm: Optional[bool] + if litellm_field_name in db_litellm_settings: + stored_in_db_litellm = True + elif current_value is not None: + stored_in_db_litellm = False + else: + stored_in_db_litellm = None + return_val.append( + ConfigList( + field_name=litellm_field_name, + field_type=spec["type"], + field_description=spec["description"], + field_value=current_value, + stored_in_db=stored_in_db_litellm, + field_default_value=None, + nested_fields=None, + ) + ) + return return_val @@ -14743,6 +14824,9 @@ async def delete_config_general_settings( }, ) + if data.field_name in _GENERAL_SETTINGS_UI_LITELLM_FIELDS: + return await _reset_general_settings_ui_litellm_field(data.field_name, user_api_key_dict) + if data.field_name not in ConfigGeneralSettings.model_fields: raise HTTPException( status_code=400, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ca6c2e86789..b577513fc0e 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -17,6 +17,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_utils import get_model_from_request +from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -54,6 +55,61 @@ def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set: } +def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: Optional[UserAPIKeyAuth]) -> bool: + """ + Whether an over-budget key's own ``max_budget`` reservation should be + released rather than blocked, because the key opted into throttling: the + rate limiter slows it instead. Only the key's own ``max_budget`` counter is + exempt; team/user/window counters still enforce normally, and under-budget + requests never reach this branch so their concurrent-overspend protection is + untouched. + """ + if valid_token is None: + return False + return counter_key == f"spend:key:{valid_token.token}" and should_throttle_budget_exceeded(valid_token) + + +async def _apply_over_budget_reservation_policy( + counter: _BudgetCounter, + valid_token: Optional[UserAPIKeyAuth], + entry: dict[str, Any], + applied_entries: list[dict[str, Any]], + reservation_cost: float, + current_spend: float, +) -> float: + """ + Decide what to do when a counter is over budget, and return the reservation + cost to carry into the next counter. Three outcomes: an over-budget key that + opted into throttling releases its own reservation (the rate limiter slows + it) and keeps the cost; a partially-remaining budget resizes the reservation + down to what is left; anything else hard-blocks by raising. + """ + if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token): + await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) + applied_entries.remove(entry) + return reservation_cost + + remaining_before_reservation = counter.max_budget - (current_spend - reservation_cost) + if remaining_before_reservation > 1e-12: + await _resize_applied_reservation( + entries=applied_entries, + current_reserved_cost=reservation_cost, + new_reserved_cost=remaining_before_reservation, + ) + return remaining_before_reservation + + raise litellm.BudgetExceededError( + current_cost=current_spend, + max_budget=counter.max_budget, + message=( + "Budget has been exceeded! " + f"{counter.entity_type}={counter.entity_id} " + f"Current cost: {current_spend}, " + f"Max budget: {counter.max_budget}" + ), + ) + + async def reserve_budget_for_request( request_body: dict, route: str, @@ -130,25 +186,15 @@ async def reserve_budget_for_request( cached_spend = await _get_current_counter_value(counter=counter) current_spend = cached_spend + reservation_cost if current_spend > counter.max_budget: - remaining_before_reservation = counter.max_budget - (current_spend - reservation_cost) - if remaining_before_reservation > 1e-12: - await _resize_applied_reservation( - entries=applied_entries, - current_reserved_cost=reservation_cost, - new_reserved_cost=remaining_before_reservation, - ) - reservation_cost = remaining_before_reservation - continue - raise litellm.BudgetExceededError( - current_cost=current_spend, - max_budget=counter.max_budget, - message=( - "Budget has been exceeded! " - f"{counter.entity_type}={counter.entity_id} " - f"Current cost: {current_spend}, " - f"Max budget: {counter.max_budget}" - ), + reservation_cost = await _apply_over_budget_reservation_policy( + counter=counter, + valid_token=valid_token, + entry=entry, + applied_entries=applied_entries, + reservation_cost=reservation_cost, + current_spend=current_spend, ) + continue except Exception: await _release_applied_entries_best_effort( entries=applied_entries, diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index bc2f41c8cb4..d12ff20ee5b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2635,6 +2635,170 @@ async def test_virtual_key_budget_check_fallback_no_counter(): assert exc_info.value.current_cost == 15.0 +# ===================================================================== +# Throttle-on-budget-exceeded tests (LIT-3894): an over-budget key that +# opted in is throttled to a global % of its TPM/RPM instead of blocked. +# ===================================================================== + + +def _over_budget_token(**overrides) -> UserAPIKeyAuth: + base = dict( + token="throttle-token", + spend=20.0, + max_budget=10.0, + user_id="test-user", + ) + base.update(overrides) + return UserAPIKeyAuth(**base) + + +def _patched_spend(value: float): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): + return value + + return patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend) + + +def _budget_logging_obj(): + from litellm.proxy.utils import ProxyLogging + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.budget_alerts = AsyncMock() + return proxy_logging_obj + + +@pytest.mark.parametrize( + "limit, pct, expected", + [ + (1000, 0.1, 100), + (100, 0.1, 10), + (1, 0.1, 1), # floor would be 0; trickle of 1 keeps the key alive + (None, 0.1, None), + (50, 0.5, 25), + (1000, None, 1000), # no percentage -> limit unchanged + ], +) +def test_throttled_limit(limit, pct, expected): + from litellm.proxy.auth.budget_throttle import throttled_limit + + assert throttled_limit(limit, pct) == expected + + +@pytest.mark.asyncio +async def test_budget_exceeded_throttles_instead_of_blocking(monkeypatch): + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + valid_token = _over_budget_token( + tpm_limit=1000, + rpm_limit=100, + metadata={"throttle_on_budget_exceeded": True}, + ) + + with _patched_spend(20.0): + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=_budget_logging_obj(), + ) + + # persistent limits are untouched (so the throttle never compounds); the + # request-scoped percentage is what the rate limiter scales by + assert valid_token.budget_throttle_pct == 0.1 + assert valid_token.tpm_limit == 1000 + assert valid_token.rpm_limit == 100 + # the request-scoped decision must not leak into serialized responses + assert "budget_throttle_pct" not in valid_token.model_dump() + + +@pytest.mark.asyncio +async def test_budget_throttle_decision_cleared_before_caching(): + """The request-scoped throttle decision must not persist into the key cache, + otherwise it would re-apply (and compound) on every subsequent request.""" + from litellm.proxy.auth.auth_checks import _copy_user_api_key_auth_for_cache + + valid_token = _over_budget_token( + tpm_limit=1000, rpm_limit=100, metadata={"throttle_on_budget_exceeded": True} + ) + valid_token.budget_throttle_pct = 0.1 + + cached = _copy_user_api_key_auth_for_cache(user_api_key_obj=valid_token) + + assert cached.budget_throttle_pct is None + assert cached.tpm_limit == 1000 + assert cached.rpm_limit == 100 + + +@pytest.mark.asyncio +async def test_budget_exceeded_throttle_no_configured_limits(monkeypatch): + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + valid_token = _over_budget_token(metadata={"throttle_on_budget_exceeded": True}) + assert valid_token.tpm_limit is None + assert valid_token.rpm_limit is None + + with _patched_spend(20.0): + with pytest.raises(litellm.BudgetExceededError): + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=_budget_logging_obj(), + ) + + assert valid_token.budget_throttle_pct is None + + +@pytest.mark.asyncio +async def test_budget_exceeded_not_opted_in_still_blocks(monkeypatch): + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + valid_token = _over_budget_token(tpm_limit=1000, rpm_limit=100) + + with _patched_spend(20.0): + with pytest.raises(litellm.BudgetExceededError): + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=_budget_logging_obj(), + ) + + assert valid_token.budget_throttle_pct is None + + +@pytest.mark.parametrize("pct", [None, 0, 1.5, -0.1, True]) +@pytest.mark.asyncio +async def test_budget_exceeded_invalid_percentage_blocks(monkeypatch, pct): + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", pct) + valid_token = _over_budget_token( + tpm_limit=1000, + rpm_limit=100, + metadata={"throttle_on_budget_exceeded": True}, + ) + + with _patched_spend(20.0): + with pytest.raises(litellm.BudgetExceededError): + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=_budget_logging_obj(), + ) + + assert valid_token.budget_throttle_pct is None + + +@pytest.mark.asyncio +async def test_under_budget_does_not_throttle(monkeypatch): + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + valid_token = _over_budget_token( + max_budget=100.0, + tpm_limit=1000, + rpm_limit=100, + metadata={"throttle_on_budget_exceeded": True}, + ) + + with _patched_spend(5.0): + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=_budget_logging_obj(), + ) + + assert valid_token.budget_throttle_pct is None + + @pytest.mark.asyncio async def test_team_budget_check_reads_from_spend_counter(): """Team budget check should use get_current_spend when counter exists.""" diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 50f471721b1..12f0a64a179 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -48,6 +48,42 @@ def time_controller(monkeypatch): return controller +@pytest.mark.parametrize( + "throttle_pct, expected_rpm, expected_tpm", + [ + (None, 100, 1000), # no throttle -> configured limits + (0.1, 10, 100), # 10% of configured + (0.5, 50, 500), + ], +) +def test_api_key_descriptor_applies_budget_throttle( + throttle_pct, expected_rpm, expected_tpm +): + """The api_key rate-limit descriptor scales the key's configured TPM/RPM by + the request-scoped budget_throttle_pct, leaving the configured limits intact.""" + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-throttle"), + rpm_limit=100, + tpm_limit=1000, + budget_throttle_pct=throttle_pct, + ) + + descriptors = handler._create_rate_limit_descriptors( + user_api_key_dict=user_api_key_dict, + data={}, + rpm_limit_type=None, + tpm_limit_type=None, + model_has_failures=False, + ) + + api_key_descriptor = next(d for d in descriptors if d["key"] == "api_key") + assert api_key_descriptor["rate_limit"]["requests_per_unit"] == expected_rpm + assert api_key_descriptor["rate_limit"]["tokens_per_unit"] == expected_tpm + + @pytest.mark.flaky(reruns=3) @pytest.mark.asyncio async def test_sliding_window_rate_limit_v3(monkeypatch, time_controller): diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 7b97ae60443..4fb3df52cf6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1495,6 +1495,57 @@ async def test_generate_service_account_works_with_team_id(): ) +@pytest.mark.asyncio +async def test_generate_key_throttle_rejected_for_non_admin(): + """Security regression: a non-admin creating a key must not be able to set + throttle_on_budget_exceeded=true, which would let the new key keep spending + past an admin-imposed per-key budget ceiling instead of hard-blocking. The + /key/update gate does not cover generate, so generate needs its own admin + check. Only the enable value is gated, so this must 403.""" + mock_prisma_client = AsyncMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(throttle_on_budget_exceeded=True), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "Only proxy admins can enable throttle_on_budget_exceeded" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_generate_key_throttle_allowed_for_admin(): + """A proxy admin may create a key with throttle_on_budget_exceeded=true; the + generate admin gate must let the admin through to key creation.""" + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", False), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = { + "key": "sk-test-key", + "expires": None, + "user_id": "admin", + "team_id": None, + } + await _common_key_generation_helper( + data=GenerateKeyRequest(throttle_on_budget_exceeded=True), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"), + litellm_changed_by=None, + team_table=None, + ) + assert mock_generate_key.called + + @pytest.mark.asyncio async def test_update_service_account_requires_team_id(): data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"}) @@ -9577,6 +9628,165 @@ async def test_update_key_non_budget_fields_allowed_for_internal_user(monkeypatc assert result is not None +@pytest.mark.asyncio +async def test_update_key_throttle_on_budget_exceeded_rejected_for_internal_user( + monkeypatch, +): + """Security regression: throttle_on_budget_exceeded turns an admin-imposed + hard budget block into a soft throttle that keeps spending past max_budget, + so it is a budget-enforcement change. A non-admin key owner (same setup that + is allowed to change non-budget fields via the caller_is_creator shortcut) + must NOT be able to self-opt-in to it; it has to route through the admin-only + _check_key_admin_access and return 403. Without treating the flag as a budget + change this update would succeed, letting the owner bypass their own cap.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = AsyncMock() + mock_proxy_logging_obj = MagicMock() + + test_hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + + # Owner of the key (created_by == user_id) so caller_is_creator is True. + # This is exactly the setup that is allowed to change non-budget fields; + # the throttle flag must still be rejected. + mock_existing_key = MagicMock() + mock_existing_key.token = test_hashed_token + mock_existing_key.user_id = "internal_user" + mock_existing_key.created_by = "internal_user" + mock_existing_key.team_id = None + mock_existing_key.project_id = None + mock_existing_key.max_budget = 10.0 + mock_existing_key.key_alias = None + mock_existing_key.models = [] + mock_existing_key.metadata = {} + mock_existing_key.model_dump.return_value = { + "token": test_hashed_token, + "user_id": "internal_user", + "team_id": None, + "max_budget": 10.0, + } + + mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=mock_existing_key) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + mock_request = MagicMock() + mock_request.query_params = {} + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ) + + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest( + key=test_hashed_token, + throttle_on_budget_exceeded=True, + ), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can enable throttle_on_budget_exceeded" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_key_throttle_unchanged_allows_non_budget_edit_for_internal_user( + monkeypatch, +): + """A non-admin owner editing a non-budget field must not be blocked just + because the UI resends throttle_on_budget_exceeded unchanged (the edit form + always includes it). Only the transition to enabled is admin-gated, so an + unchanged False here leaves the key owner's non-budget edit working.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = AsyncMock() + mock_proxy_logging_obj = MagicMock() + + test_hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + + mock_existing_key = MagicMock() + mock_existing_key.token = test_hashed_token + mock_existing_key.user_id = "internal_user" + mock_existing_key.created_by = "internal_user" + mock_existing_key.team_id = None + mock_existing_key.project_id = None + mock_existing_key.max_budget = 10.0 + mock_existing_key.key_alias = None + mock_existing_key.models = [] + mock_existing_key.metadata = {"throttle_on_budget_exceeded": False} + mock_existing_key.model_dump.return_value = { + "token": test_hashed_token, + "user_id": "internal_user", + "team_id": None, + "max_budget": 10.0, + } + + mock_updated_key = MagicMock() + mock_updated_key.token = test_hashed_token + mock_updated_key.key_alias = "my-alias" + + mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key) + mock_prisma_client.update_data = AsyncMock(return_value=mock_updated_key) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=mock_existing_key) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + + monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", lambda token: test_hashed_token) + + async def _noop(**kwargs): + pass + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + _noop, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias", + _noop, + ) + + mock_request = MagicMock() + mock_request.query_params = {} + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ) + + result = await update_key_fn( + request=mock_request, + data=UpdateKeyRequest( + key=test_hashed_token, + key_alias="my-alias", + throttle_on_budget_exceeded=False, + ), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + assert result is not None + + @pytest.mark.asyncio async def test_update_key_non_budget_rejects_cross_user_modification(monkeypatch): """Regression: previously _check_key_admin_access was gated on diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index d940f592a83..75242af81f4 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -58,6 +58,101 @@ def _request_body() -> dict: } +async def _reserve(valid_token, cost, key_cache, proxy_logging_obj): + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=cost, + ): + return await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +@pytest.mark.asyncio +async def test_reservation_still_protects_under_budget_throttled_key( + spend_counter_state, monkeypatch +): + """An opted-in key that is still under budget keeps its reservation counter, + so concurrent requests can't collectively overshoot max_budget.""" + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-throttle-under", + spend=0.0, + max_budget=1.0, + metadata={"throttle_on_budget_exceeded": True}, + ) + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + + assert reservation is not None + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:key:key-throttle-under") + == 0.6 + ) + + +@pytest.mark.asyncio +async def test_reservation_does_not_block_over_budget_throttled_key( + spend_counter_state, monkeypatch +): + """Once an opted-in key is over budget the reservation path must not raise; + the rate limiter throttles it instead.""" + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-throttle-over", + spend=0.0, + max_budget=1.0, + tpm_limit=1000, + rpm_limit=100, + metadata={"throttle_on_budget_exceeded": True}, + ) + + # first reservation lands under budget (counter -> 0.6) + await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + + # any further request is over budget (0.6 + 0.6 > 1.0): the opted-in key is + # released and allowed through (None), not blocked, and its over-budget + # increment is released so the counter is not permanently inflated + result = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert result is None + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:key:key-throttle-over") + == 0.6 + ) + + +@pytest.mark.asyncio +async def test_reservation_blocks_over_budget_non_throttled_key( + spend_counter_state, monkeypatch +): + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-no-optin-over", + spend=0.0, + max_budget=1.0, + ) + + await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) # counter -> 1.0 + + with pytest.raises(litellm.BudgetExceededError): + await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + + def test_should_not_serialize_budget_reservation_on_user_api_key_auth(): auth = UserAPIKeyAuth( token="key-budget-runtime-state", diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2dc67c827e3..d06a1c16ab9 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8644,6 +8644,149 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): 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 + must appear in /config/list reading its live litellm. value.""" + 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) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.15) + 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 "budget_exceeded_throttle_percentage" in fields + assert fields["budget_exceeded_throttle_percentage"]["field_type"] == "Float" + assert fields["budget_exceeded_throttle_percentage"]["field_value"] == 0.15 + finally: + app.dependency_overrides.clear() + + +@pytest.mark.asyncio +async def test_update_config_field_throttle_persists_to_litellm_settings(monkeypatch): + """Editing the throttle Float row on the General Settings table routes to + litellm_settings (not general_settings): it sets litellm. live and + persists under litellm_settings so the runtime read is unchanged.""" + from unittest.mock import MagicMock + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + saved: dict = {} + + async def fake_get_config(): + return {"litellm_settings": {}} + + async def fake_save_config(new_config=None): + saved.update(new_config or {}) + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps.proxy_config, "save_config", fake_save_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "store_audit_logs", False) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", None) + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="budget_exceeded_throttle_percentage", + field_value=0.1, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + + assert litellm.budget_exceeded_throttle_percentage == 0.1 + assert saved["litellm_settings"]["budget_exceeded_throttle_percentage"] == 0.1 + + +@pytest.mark.parametrize("bad_value", [0, -0.1, 1.5, True]) +@pytest.mark.asyncio +async def test_update_config_field_throttle_rejects_invalid(monkeypatch, bad_value): + from unittest.mock import MagicMock + + from fastapi import HTTPException + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + async def fake_get_config(): + return {"litellm_settings": {}} + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", None) + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as exc: + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="budget_exceeded_throttle_percentage", + field_value=bad_value, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + assert exc.value.status_code == 400 + assert litellm.budget_exceeded_throttle_percentage is None + + +@pytest.mark.asyncio +async def test_update_config_field_throttle_rejected_for_non_admin(monkeypatch): + from unittest.mock import MagicMock + + from fastapi import HTTPException + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", None) + + non_admin = UserAPIKeyAuth(api_key="k", user_id="u", user_role=LitellmUserRoles.INTERNAL_USER) + with pytest.raises(HTTPException): + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="budget_exceeded_throttle_percentage", + field_value=0.1, + config_type="general_settings", + ), + user_api_key_dict=non_admin, + ) + assert litellm.budget_exceeded_throttle_percentage is None + + def test_preserve_redacted_plugin_keys_keeps_stored_credential(): """A redacted or blank plugin_key on update must not overwrite the real key.""" from litellm.proxy.proxy_server import _preserve_redacted_plugin_keys diff --git a/ui/litellm-dashboard/src/components/general_settings.tsx b/ui/litellm-dashboard/src/components/general_settings.tsx index a7a8af2691e..5b8dec39505 100644 --- a/ui/litellm-dashboard/src/components/general_settings.tsx +++ b/ui/litellm-dashboard/src/components/general_settings.tsx @@ -161,6 +161,14 @@ const GeneralSettings: React.FC = ({ accessToken, user checked={value.field_value === true || value.field_value === "true"} onChange={(checked) => handleInputChange(value.field_name, checked)} /> + ) : value.field_type == "Float" ? ( + handleInputChange(value.field_name, newValue)} + /> ) : null} diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index ca5766a3682..bf0f0cc3fae 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1163,6 +1163,21 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp form={form} showDetailedDescriptions={true} /> + + Throttle on budget exceeded{" "} + + + + + } + name="throttle_on_budget_exceeded" + valuePropName="checked" + > + + diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 40c82c51031..ba68124beee 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -335,6 +335,36 @@ describe("KeyEditView", () => { }); }); + it("should initialize and submit throttle_on_budget_exceeded from key metadata", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + const keyDataWithThrottle = { + ...MOCK_KEY_DATA, + metadata: { ...MOCK_KEY_DATA.metadata, throttle_on_budget_exceeded: true }, + }; + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken={"test-token"} + userID={"test-user"} + userRole={"admin"} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Throttle on budget exceeded")).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalledWith(expect.objectContaining({ throttle_on_budget_exceeded: true })); + }); + }); + it("should disable models field when management routes are selected", async () => { const keyDataWithManagementRoutes = { ...MOCK_KEY_DATA, diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 7dda555daa4..4821ea86b87 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -181,6 +181,7 @@ export function KeyEditView({ metadata: formatMetadataForDisplay(stripTagsFromMetadata(keyData.metadata)), guardrails: keyData.metadata?.guardrails, disable_global_guardrails: keyData.metadata?.disable_global_guardrails || false, + throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false, prompts: keyData.metadata?.prompts, tags: keyData.metadata?.tags, vector_stores: keyData.object_permission?.vector_stores || [], @@ -222,6 +223,7 @@ export function KeyEditView({ accessGroups: keyData.object_permission?.mcp_access_groups || [], }, mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {}, + throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false, logging_settings: extractLoggingSettings(keyData.metadata), disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks) ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) @@ -512,6 +514,21 @@ export function KeyEditView({ + + Throttle on budget exceeded{" "} + + + + + } + name="throttle_on_budget_exceeded" + valuePropName="checked" + > + + + diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 197f87b6dfe..1c0d916f0f3 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -542,6 +542,9 @@ export default function KeyInfoView({
TPM: {currentKeyData.tpm_limit !== null ? currentKeyData.tpm_limit : "Unlimited"} RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"} + {Boolean(currentKeyData.metadata?.throttle_on_budget_exceeded) && ( + Throttle on budget exceeded: Yes + )}
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 71568299529..d9bba85bf4b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -6502,6 +6502,7 @@ export interface paths { * - guardrails: Optional[List[str]] - List of active guardrails for the key * - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. + * - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. * - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. @@ -6905,6 +6906,7 @@ export interface paths { * - guardrails: Optional[List[str]] - List of active guardrails for the key * - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. + * - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. * - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. * - blocked: Optional[bool] - Whether the key is blocked * - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases) @@ -23707,6 +23709,8 @@ export interface components { tags?: string[] | null; /** Team Id */ team_id?: string | null; + /** Throttle On Budget Exceeded */ + throttle_on_budget_exceeded?: boolean | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -23847,6 +23851,8 @@ export interface components { tags?: string[] | null; /** Team Id */ team_id?: string | null; + /** Throttle On Budget Exceeded */ + throttle_on_budget_exceeded?: boolean | null; /** Token */ token?: string | null; /** Token Id */ @@ -28186,6 +28192,8 @@ export interface components { team_id?: string | null; /** Teams */ teams?: unknown[] | null; + /** Throttle On Budget Exceeded */ + throttle_on_budget_exceeded?: boolean | null; /** Token */ token?: string | null; /** Token Id */ @@ -29813,6 +29821,8 @@ export interface components { tags?: string[] | null; /** Team Id */ team_id?: string | null; + /** Throttle On Budget Exceeded */ + throttle_on_budget_exceeded?: boolean | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -31760,6 +31770,8 @@ export interface components { temp_budget_expiry?: string | null; /** Temp Budget Increase */ temp_budget_increase?: number | null; + /** Throttle On Budget Exceeded */ + throttle_on_budget_exceeded?: boolean | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */