Merge 5f4a19aed4 into eb0e3f8c18

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]]] = []

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,