fix(proxy): release budget reservation when a request is cancelled mid-flight

litellm/proxy/common_request_processing.py

@@ -2531,6 +2531,7 @@ async def async_streaming_data_generator(
         debug_enabled = verbose_proxy_logger.isEnabledFor(logging.DEBUG)
         stream_completed = False
         client_disconnected = False
+        delivered_chunk = False
         try:
             str_so_far = ""
             async for (
@@ -2547,36 +2548,38 @@ async def async_streaming_data_generator(
                         "async_data_generator: received streaming chunk - %s", chunk
                     )
 
-                if fast_path:
-                    yield serialize_chunk(chunk)
-                    continue
-
-                chunk = await proxy_logging_obj.async_post_call_streaming_hook(
-                    user_api_key_dict=user_api_key_dict,
-                    response=chunk,
-                    data=request_data,
-                    str_so_far=str_so_far,
-                )
-
-                if isinstance(chunk, (ModelResponse, ModelResponseStream)):
-                    response_str = litellm.get_response_string(response_obj=chunk)
-                    str_so_far += response_str
-                elif hasattr(chunk, "model_dump"):
-                    try:
-                        d = chunk.model_dump(mode="json", exclude_none=True)
-                        if isinstance(d, dict):
-                            str_so_far += str(d.get("content", ""))
-                    except Exception:
-                        pass
-                elif isinstance(chunk, dict):
-                    str_so_far += str(chunk.get("content", ""))
+                if not fast_path:
+                    chunk = await proxy_logging_obj.async_post_call_streaming_hook(
+                        user_api_key_dict=user_api_key_dict,
+                        response=chunk,
+                        data=request_data,
+                        str_so_far=str_so_far,
+                    )
 
-                model_name = request_data.get("model", "")
-                chunk = (
-                    ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
+                    if isinstance(chunk, (ModelResponse, ModelResponseStream)):
+                        response_str = litellm.get_response_string(response_obj=chunk)
+                        str_so_far += response_str
+                    elif hasattr(chunk, "model_dump"):
+                        try:
+                            d = chunk.model_dump(mode="json", exclude_none=True)
+                            if isinstance(d, dict):
+                                str_so_far += str(d.get("content", ""))
+                        except Exception:
+                            pass
+                    elif isinstance(chunk, dict):
+                        str_so_far += str(chunk.get("content", ""))
+
+                    model_name = request_data.get("model", "")
+                    chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
                         chunk, model_name
                     )
-                )
+
+                # Set before the yield: an async generator suspends at the yield,
+                # so a GeneratorExit on client disconnect is raised there and any
+                # statement after the yield never runs. The slow-path hook is
+                # awaited above, so a cancellation during it still leaves this
+                # False and refunds.
+                delivered_chunk = True
                 yield serialize_chunk(chunk)
             stream_completed = True
         except (asyncio.CancelledError, GeneratorExit):
@@ -2591,6 +2594,14 @@ async def async_streaming_data_generator(
                     user_api_key_dict
                 )
                 client_disconnected = True
+            if not delivered_chunk:
+                from litellm.proxy.spend_tracking.budget_reservation import (
+                    release_budget_reservation_on_cancel,
+                )
+
+                await release_budget_reservation_on_cancel(
+                    getattr(user_api_key_dict, "budget_reservation", None)
+                )
             raise
         except Exception as e:
             verbose_proxy_logger.exception(

litellm/proxy/spend_tracking/budget_reservation.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import json
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
@@ -162,10 +163,14 @@ async def reserve_budget_for_request(
     if not applied_entries:
         return None
 
+    input_cost = estimate_request_input_cost(
+        request_body=request_body, route=route, llm_router=llm_router
+    )
     return {
         "reserved_cost": reservation_cost,
         "entries": applied_entries,
         "finalized": False,
+        "input_cost": min(float(input_cost or 0.0), reservation_cost),
     }
 
 
@@ -195,6 +200,41 @@ async def release_budget_reservation(budget_reservation: Optional[dict]) -> None
     )
 
 
+async def release_budget_reservation_on_cancel(
+    budget_reservation: dict | None,
+) -> None:
+    """Reconcile a still-open reservation when the request is cancelled mid-flight.
+
+    A client disconnect or timeout cancels the request task, which surfaces as
+    CancelledError / GeneratorExit rather than a normal exception, so neither the
+    success cost callback nor the failure hook runs and the pre-call reservation
+    is never reconciled. Left alone it pins the spend counter above real spend
+    and 429s subsequent requests until the counter's TTL expires.
+
+    Reconcile to the request's input-token cost rather than refunding to zero:
+    by the time a request is cancelled in-flight the provider call was already
+    dispatched, so the input tokens were billed even if no chunk reached the
+    client. Refunding to zero would let a caller abort pre-token to dodge that
+    charge; the worst-case output portion of the reservation is still released.
+
+    asyncio.shield keeps the reconcile running to completion even though the
+    surrounding task is being cancelled. The `finalized` guard makes this a no-op
+    when success/failure handling already reconciled, so calling it on every
+    cancellation path is safe.
+    """
+    if not budget_reservation or budget_reservation.get("finalized") is True:
+        return
+    incurred_cost = float(budget_reservation.get("input_cost") or 0.0)
+    try:
+        await asyncio.shield(
+            reconcile_budget_reservation(
+                budget_reservation=budget_reservation, actual_cost=incurred_cost
+            )
+        )
+    except (asyncio.CancelledError, Exception):
+        pass
+
+
 async def invalidate_budget_reservation_counters(
     budget_reservation: Optional[dict],
 ) -> None:
@@ -817,6 +857,61 @@ def estimate_request_max_cost(
     return max(cast(List[float], estimates))
 
 
+def estimate_request_input_cost(
+    request_body: dict,
+    route: str,
+    llm_router: Router | None,
+) -> float | None:
+    """Cost of the request's input tokens alone.
+
+    Once the provider request is dispatched the input tokens are billed even if
+    the client disconnects before the first chunk, so this is the cost floor a
+    cancelled in-flight request has already incurred. A cancelled reservation is
+    reconciled to this instead of being refunded to zero.
+    """
+    model = get_model_from_request(request_body, route, llm_router=llm_router)
+    if model is None:
+        return None
+
+    models = [model] if isinstance(model, str) else model
+    estimates = [
+        _estimate_request_input_cost_for_model(
+            request_body=request_body,
+            route=route,
+            model=model_name,
+            llm_router=llm_router,
+        )
+        for model_name in models
+    ]
+    estimates = [estimate for estimate in estimates if estimate is not None]
+    if not estimates:
+        return None
+    return max(cast("list[float]", estimates))
+
+
+def _estimate_request_input_cost_for_model(
+    request_body: dict,
+    route: str,
+    model: str,
+    llm_router: Router | None,
+) -> float | None:
+    model_info = _get_model_cost_info(model=model, llm_router=llm_router)
+    if model_info is None:
+        return None
+    input_cost_per_token = _to_float(model_info.get("input_cost_per_token"))
+    if input_cost_per_token is None:
+        return None
+    input_tokens = _estimate_input_tokens(
+        request_body=request_body,
+        route=route,
+        model=model,
+        model_info=model_info,
+    )
+    if input_tokens is None:
+        return None
+    return input_tokens * input_cost_per_token
+
+
 def _estimate_request_max_cost_for_model(
     request_body: dict,
     route: str,