diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 57a7d860baa..fba968dd95d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2312,6 +2312,24 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", ) + disable_budget_reservation: Optional[bool] = Field( + None, + description=( + "If True, disables the optimistic per-request budget reservation " + "introduced in v1.84.0. " + "WARNING: This weakens hard budget enforcement. Without the reservation, " + "a burst of concurrent requests from a single key can each pass the " + "read-time spend check before any of them is charged, allowing a " + "configured budget to be exceeded under high concurrency. " + "Budgets are still evaluated on every request at read time, so " + "an already-exhausted budget is still rejected. " + "Enable only if your deployment is experiencing phantom " + "BudgetExceededError responses caused by leaked reservations " + "(see GitHub issue #27639). " + "A proxy-level WARNING is logged on every request while this flag " + "is active as a reminder that hard enforcement is relaxed." + ), + ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a970e0ddee8..666c01562b5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2425,6 +2425,7 @@ async def _run_centralized_common_checks( # noqa: PLR0915 user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, skip_budget_checks=skip_budget_checks, + general_settings=general_settings, ) @@ -2445,12 +2446,23 @@ async def _reserve_budget_after_common_checks( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, skip_budget_checks: bool, + general_settings: dict, end_user_id: Optional[str] = None, end_user_object: Optional[LiteLLM_EndUserTable] = None, ) -> None: user_api_key_auth_obj.budget_reservation = None if skip_budget_checks: return + if general_settings.get("disable_budget_reservation") is True: + verbose_proxy_logger.warning( + "disable_budget_reservation is enabled: skipping optimistic budget " + "reservation. Budget enforcement is read-time only — concurrent " + "requests can each pass the spend check before their cost is recorded, " + "so a configured budget may be briefly exceeded under high concurrency. " + "Set disable_budget_reservation to False or remove it to restore " + "hard per-request budget enforcement." + ) + return from litellm.proxy.spend_tracking.budget_reservation import ( reserve_budget_for_request, 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 fa6cc8bed1b..0236646c796 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 @@ -112,11 +112,71 @@ async def test_should_clear_stale_budget_reservation_when_budget_checks_skip(): user_api_key_cache=MagicMock(), proxy_logging_obj=MagicMock(), skip_budget_checks=True, + general_settings={}, ) assert user_api_key_auth_obj.budget_reservation is None +@pytest.mark.asyncio +async def test_disable_budget_reservation_skips_reservation(): + """#27639: general_settings.disable_budget_reservation turns off the optimistic Redis + reservation so operators hit by phantom BudgetExceededError can opt out of it.""" + 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={"reserved_cost": 0.5, "entries": []}), + ) 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={"disable_budget_reservation": True}, + ) + + mock_reserve.assert_not_called() + assert user_api_key_auth_obj.budget_reservation is None + + +@pytest.mark.asyncio +async def test_budget_reservation_runs_when_not_disabled(): + """Control for #27639: with the flag absent, the reservation still runs and is stored.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_token"}], + } + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=reservation), + ) 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={}, + ) + + mock_reserve.assert_awaited_once() + assert user_api_key_auth_obj.budget_reservation == reservation + + @pytest.mark.asyncio async def test_should_not_reuse_cached_key_object_for_request_state(): key_cache = DualCache()