From 07f7fc224e8db662c4668f1f6f00ba035923e64a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 23 Jul 2026 16:57:37 -0700 Subject: [PATCH] fix(proxy): reject failed atomic budget reservations under fail_closed_budget_enforcement (#34429) * fix(proxy): reject request when budget reservation write fails under fail_closed_budget_enforcement With general_settings.fail_closed_budget_enforcement set to true, the read-time spend check already returns 503 when spend cannot be verified, but the atomic pre-call reservation still failed open: reserve_budget_for_request swallowed _CounterReservationUnavailable per counter and degraded to read-time-only enforcement, so concurrent requests could all pass the same under-budget read during a Redis outage and overspend past the configured budget. Now the strict flag is threaded into reserve_budget_for_request and a failed reservation write raises 503, releasing any counters that already reserved. Default behavior with the flag absent or false is unchanged. Fixes #33923 * fix(proxy): pass 503 budget-enforcement detail as plain string --- litellm/proxy/auth/user_api_key_auth.py | 1 + .../spend_tracking/budget_reservation.py | 23 ++++- .../proxy/auth/test_user_api_key_auth.py | 39 ++++++++ .../proxy/test_budget_reservation.py | 95 +++++++++++++++++++ 4 files changed, 157 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d776a626251..709df5e64df 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2452,6 +2452,7 @@ async def _reserve_budget_after_common_checks( end_user_id=end_user_id, end_user_object=end_user_object, skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True, + fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 80fd8a1594e..013873179c6 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -4,7 +4,9 @@ import asyncio import json from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Mapping, Optional, Sequence, cast +from typing import Any, Dict, List, Mapping, NoReturn, Optional, Sequence, cast + +from fastapi import HTTPException, status import litellm from litellm._logging import verbose_proxy_logger @@ -59,6 +61,22 @@ class _CounterReservationUnavailable(Exception): super().__init__("Counter reservation unavailable") +def _raise_reservation_unavailable(counter_key: str) -> NoReturn: + verbose_proxy_logger.warning( + "fail_closed_budget_enforcement: rejecting request — budget reservation for %s could not be written", + counter_key, + ) + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + "Budget enforcement unavailable: the budget reservation could not " + "be written to the spend counter backend, and " + "fail_closed_budget_enforcement is enabled, so the request was " + "rejected to avoid exceeding the configured budget. Retry shortly." + ), + ) + + def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set: if not budget_reservation: return set() @@ -138,6 +156,7 @@ async def reserve_budget_for_request( end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, skip_user_budget_on_team_key: bool = False, + fail_closed_budget_enforcement: bool = False, ) -> Optional[dict]: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None @@ -193,6 +212,8 @@ async def reserve_budget_for_request( default_reserved_cost=reservation_cost, ) applied_entries.remove(entry) + if fail_closed_budget_enforcement: + _raise_reservation_unavailable(counter_key=counter.counter_key) continue if reserved_value is not None: 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 a2445ecb975..0359d974d19 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 @@ -179,6 +179,45 @@ async def test_budget_reservation_runs_when_not_disabled(): assert user_api_key_auth_obj.budget_reservation == reservation +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings,expected_flag", + [ + ({"fail_closed_budget_enforcement": True}, True), + ({}, False), + ], +) +async def test_fail_closed_budget_enforcement_reaches_reservation( + general_settings, expected_flag +): + """#33923: the strict flag must be threaded into reserve_budget_for_request so a + failed reservation write can reject instead of failing open.""" + 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["fail_closed_budget_enforcement"] + 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/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1db76aed61d..e6ccbc579b5 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -3,6 +3,7 @@ from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache @@ -1562,6 +1563,100 @@ async def test_should_skip_reservation_when_counter_increment_fails( ) +@pytest.mark.asyncio +async def test_should_raise_503_when_counter_increment_fails_and_fail_closed( + spend_counter_state, + monkeypatch, +): + """#33923: with fail_closed_budget_enforcement on, a failed reservation write + must reject instead of silently degrading to read-time-only enforcement.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-reserve-fail-closed", + spend=0.0, + max_budget=1.0, + ) + + async def fail_increment_cache(*args, **kwargs): + raise RuntimeError("counter unavailable") + + monkeypatch.setattr(counter_cache, "async_increment_cache", fail_increment_cache) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(HTTPException) as exc_info: + 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, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.status_code == 503 + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-reserve-fail-closed" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_fail_closed_releases_earlier_counters_before_503( + spend_counter_state, +): + """#33923: when a later counter's reservation write fails in strict mode, the + counters that already reserved must be released before the 503 propagates.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-fail-closed-release", + spend=0.0, + max_budget=1.0, + budget_limits=[ + { + "budget_duration": "1h", + "max_budget": 1.0, + } + ], + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(HTTPException) as exc_info: + 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, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.status_code == 503 + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-fail-closed-release" + ) + == 0.0 + ) + + @pytest.mark.asyncio async def test_should_skip_reservation_when_counter_initialization_fails( spend_counter_state,