diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 7cad3f0a022..d4d3a66212b 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -84,6 +84,38 @@ def _raise_reservation_unavailable(counter_key: str) -> NoReturn: ) +def _raise_cost_estimate_unavailable(route: str) -> NoReturn: + verbose_proxy_logger.warning( + "fail_closed_budget_enforcement: rejecting request — request cost for route %s could not be estimated", + route, + ) + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + "Budget enforcement unavailable: this request's cost could not be " + "estimated before dispatch (e.g. an unpriced model or route), and " + "fail_closed_budget_enforcement is enabled, so the request was " + "rejected to avoid exceeding a configured budget." + ), + ) + + +def _handle_missing_cost_estimate(route: str, fail_closed_budget_enforcement: bool) -> None: + """A budget is configured but estimate_request_max_cost() returned no usable + estimate. Reject when the operator opted into fail_closed_budget_enforcement; + otherwise warn so the read-time-only fallback is visible instead of silent. + """ + if fail_closed_budget_enforcement: + _raise_cost_estimate_unavailable(route=route) + verbose_proxy_logger.warning( + "reserve_budget_for_request: could not estimate a cost for route=%s; a budget is " + "configured for this request but no atomic reservation will be made (non-atomic, " + "read-time-only enforcement). Set fail_closed_budget_enforcement=true to reject " + "these requests instead.", + route, + ) + + def get_reserved_counter_keys(budget_reservation: dict | None) -> set: if not budget_reservation: return set() @@ -202,8 +234,10 @@ async def reserve_budget_for_request( ) # estimate_request_max_cost still returns None when the model is unknown # to the cost map (no token-priced cost fields, e.g. image/audio routes). - # In that case we fall back to read-time enforcement only. + # A budget is configured here (counters is non-empty), so this is a real + # atomicity gap, not a benign no-budget request. if reservation_cost is None or reservation_cost <= 0: + _handle_missing_cost_estimate(route=route, fail_closed_budget_enforcement=fail_closed_budget_enforcement) return None applied_entries: Final[list[dict[str, float | str]]] = [] diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 38a346e7fb7..bfe8eb10b77 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1672,6 +1672,98 @@ async def test_fail_closed_releases_earlier_counters_before_503( ) +@pytest.mark.asyncio +async def test_should_raise_503_when_cost_cannot_be_estimated_and_fail_closed( + spend_counter_state, +): + """A budget is configured but the request's cost can't be estimated (e.g. an + unpriced model or an image/audio route). With fail_closed_budget_enforcement + on, this must reject instead of silently skipping the reservation and + falling back to non-atomic 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-no-estimate-fail-closed", + spend=0.0, + max_budget=1.0, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=None, + ): + 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-no-estimate-fail-closed" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_missing_cost_estimate_falls_back_silently_by_default( + spend_counter_state, +): + """Default behavior (fail_closed_budget_enforcement off) is unchanged: an + unestimable cost with a budget configured returns no reservation rather + than raising, but must now warn that enforcement fell back to non-atomic + read-time-only checks instead of doing so silently.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-no-estimate-default", + spend=0.0, + max_budget=1.0, + ) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=None, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.verbose_proxy_logger.warning" + ) as mock_warning, + ): + result = 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=False, + ) + + assert result is None + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-no-estimate-default" + ) + is None + ) + mock_warning.assert_called_once() + assert "could not estimate a cost" in mock_warning.call_args.args[0] + + @pytest.mark.asyncio async def test_should_skip_reservation_when_counter_initialization_fails( spend_counter_state,