fix(proxy): record partial spend on the failure row for interrupted streams (#30788)

A streaming request that breaks mid-flight, for example on a mid-stream read
timeout, still bills the provider for the chunks already delivered, yet the proxy
recorded that interrupted request as a zero-spend failure. An earlier revision
logged the recovered partial usage through the success path, which mislabeled a
failed request as a success and produced a misleading spend row

This recovers the partial usage where the failure is actually logged. The
streaming handler assembles the usage from the chunks seen so far and stashes it,
with its cost, on the logging object before firing the failure handlers. The
proxy failure hook lifts that usage and cost onto request_data before the
non-serialisable logging object is popped, and the spend-log writer records the
real partial spend on the failure row instead of a hardcoded zero;
get_logging_payload honors the recovered usage for the token columns and
_failure_handler_helper_fn preserves the recovered cost so the non-DB failure
loggers stay consistent

A request that recovers via a successful fallback is unaffected: the failure hook
only fires when the whole request fails, so the fallback's combined-usage success
row stays the single source of truth and there is no double counting

Resolves LIT-3825

Co-authored-by: veria-ai[bot] <224490171+veria-ai[bot]@users.noreply.github.com>
(cherry picked from commit 4847fa5dd5)

litellm/litellm_core_utils/litellm_logging.py

@@ -2951,7 +2951,12 @@ class Logging(LiteLLMLoggingBaseClass):
         )
         self.model_call_details["end_time"] = end_time
         self.model_call_details.setdefault("original_response", None)
-        self.model_call_details["response_cost"] = 0
+        # A stream interrupted mid-flight still billed the provider for the
+        # chunks already delivered; the router stashes that recovered usage as
+        # ``combined_usage_object`` and pre-computes its cost, so preserve it
+        # here instead of zeroing the spend on an otherwise-failed request.
+        if self.model_call_details.get("combined_usage_object") is None:
+            self.model_call_details["response_cost"] = 0
 
         if hasattr(exception, "headers") and isinstance(exception.headers, dict):
             self.model_call_details.setdefault("litellm_params", {})

litellm/litellm_core_utils/streaming_handler.py

@@ -2233,6 +2233,7 @@ class CustomStreamWrapper:
                 litellm.request_timeout
             )
             if self.logging_obj is not None:
+                self._record_partial_usage_for_failure()
                 ## LOGGING
                 threading.Thread(
                     target=self.logging_obj.failure_handler,
@@ -2246,6 +2247,7 @@ class CustomStreamWrapper:
         except Exception as e:
             traceback_exception = traceback.format_exc()
             if self.logging_obj is not None:
+                self._record_partial_usage_for_failure()
                 ## LOGGING
                 threading.Thread(
                     target=self.logging_obj.failure_handler,
@@ -2257,6 +2259,33 @@ class CustomStreamWrapper:
                 )
             self._handle_stream_fallback_error(e)
 
+    def _record_partial_usage_for_failure(self) -> None:
+        """
+        A stream that breaks mid-flight still billed the provider for the chunks
+        already delivered. Recover that partial usage from the chunks seen so
+        far and stash it, with its cost, on the logging object so the failure
+        handler records the real partial spend instead of zero. A request that
+        later recovers via a router fallback overwrites this with the combined
+        success log on the same request id, so this never double counts.
+        """
+        if self.logging_obj is None or not self.chunks:
+            return
+        try:
+            partial_response = litellm.stream_chunk_builder(chunks=self.chunks)
+            usage = cast(Optional[Usage], getattr(partial_response, "usage", None))
+            if usage is None:
+                return
+            self.logging_obj.model_call_details["combined_usage_object"] = usage
+            self.logging_obj.model_call_details["response_cost"] = (
+                self.logging_obj._response_cost_calculator(result=partial_response)
+                or 0.0
+            )
+        except Exception as recover_error:
+            verbose_logger.debug(
+                "could not recover partial usage for interrupted stream: %s",
+                recover_error,
+            )
+
     def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn":
         """
         Common error handling for both __next__ and __anext__.

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -40,7 +40,7 @@ class _ProxyDBLogger(CustomLogger):
             kwargs, response_obj, start_time, end_time
         )
 
-    async def async_post_call_failure_hook(
+    async def async_post_call_failure_hook(  # noqa: PLR0915
         self,
         request_data: dict,
         original_exception: Exception,
@@ -162,9 +162,20 @@ class _ProxyDBLogger(CustomLogger):
             if obj_start is not None:
                 actual_start_time = obj_start
 
+        # A stream that broke mid-flight still billed the provider for the
+        # chunks already delivered. ``post_call_failure_hook`` lifts that
+        # recovered cost onto request_data (the usage rides along in
+        # ``combined_usage_object`` for the token columns), so attribute the
+        # real partial spend to this failure row instead of zero.
+        recovered_response_cost = 0.0
+        if isinstance(request_data.get("combined_usage_object"), litellm.Usage):
+            recovered_response_cost = max(
+                float(request_data.get("response_cost") or 0.0), 0.0
+            )
+
         await proxy_logging_obj.db_spend_update_writer.update_database(
             token=user_api_key_dict.api_key,
-            response_cost=0.0,
+            response_cost=recovered_response_cost,
             user_id=user_api_key_dict.user_id,
             end_user_id=user_api_key_dict.end_user_id,
             team_id=user_api_key_dict.team_id,

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -265,6 +265,13 @@ def get_logging_payload(  # noqa: PLR0915
         elif isinstance(_usage, dict):
             usage = _usage
 
+    # A request that failed mid-stream has no usable response_obj usage, but the
+    # streaming handler may have recovered the usage from the chunks already
+    # delivered. Honor that override so the partial usage lands in spend tracking.
+    _combined_usage = kwargs.get("combined_usage_object")
+    if not usage and isinstance(_combined_usage, litellm.Usage):
+        usage = _combined_usage.model_dump()
+
     id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs)
     standard_logging_payload = cast(
         Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None)

litellm/proxy/utils.py

@@ -2006,12 +2006,21 @@ class ProxyLogging:
         # compute preprocessing latency after the logging object is popped.
         _logging_obj = request_data.get("litellm_logging_obj")
         if _logging_obj is not None:
-            _first_handoff = getattr(_logging_obj, "model_call_details", {}).get(
-                "first_api_call_start_time"
-            )
+            _model_call_details = getattr(_logging_obj, "model_call_details", {})
+            _first_handoff = _model_call_details.get("first_api_call_start_time")
             if _first_handoff is not None:
                 request_data["first_api_call_start_time"] = _first_handoff
 
+            # A stream that broke mid-flight still billed the provider for the
+            # chunks already delivered; the streaming handler stashes that
+            # recovered usage and cost here. Lift them onto request_data so the
+            # failure-path spend callbacks (which run after the logging object
+            # is popped) record the real partial spend instead of zero.
+            _recovered_usage = _model_call_details.get("combined_usage_object")
+            if _recovered_usage is not None:
+                request_data["combined_usage_object"] = _recovered_usage
+                request_data["response_cost"] = _model_call_details.get("response_cost")
+
         # Remove before callbacks iterate — not serialisable
         request_data.pop("litellm_logging_obj", None)
 

tests/test_litellm/litellm_core_utils/test_litellm_logging.py

@@ -3078,3 +3078,46 @@ class TestFirstApiCallStartTimeSetOnce:
         assert obj.model_call_details["api_call_start_time"] > first
         assert obj.model_call_details["first_api_call_start_time"] == first
         assert user_meta == {}
+
+
+def test_failure_handler_records_recovered_partial_spend(logging_obj):
+    """A stream interrupted mid-flight still billed the provider for the chunks
+    already delivered. When the router stashes that recovered usage as
+    ``combined_usage_object`` and pre-computes ``response_cost``, the failure
+    handler must preserve them so the failure row carries the real partial
+    spend instead of zero.
+    """
+    from litellm.types.utils import Usage
+
+    logging_obj.model_call_details["combined_usage_object"] = Usage(
+        prompt_tokens=17, completion_tokens=9, total_tokens=26
+    )
+    logging_obj.model_call_details["response_cost"] = 0.00012
+
+    logging_obj._failure_handler_helper_fn(
+        exception=Exception("Connection lost"),
+        traceback_exception="Traceback ...",
+    )
+
+    payload = logging_obj.model_call_details["standard_logging_object"]
+    assert payload["status"] == "failure"
+    assert payload["response_cost"] == 0.00012
+    assert payload["prompt_tokens"] == 17
+    assert payload["completion_tokens"] == 9
+    assert payload["total_tokens"] == 26
+
+
+def test_failure_handler_zeroes_spend_without_recovered_usage(logging_obj):
+    """A failure with no recovered partial usage keeps the existing behavior of
+    recording zero spend, so the partial-spend preservation does not leak into
+    ordinary failures.
+    """
+    logging_obj._failure_handler_helper_fn(
+        exception=Exception("boom"),
+        traceback_exception="Traceback ...",
+    )
+
+    payload = logging_obj.model_call_details["standard_logging_object"]
+    assert payload["status"] == "failure"
+    assert payload["response_cost"] == 0
+    assert payload["total_tokens"] == 0

tests/test_litellm/litellm_core_utils/test_streaming_handler.py

@@ -2118,3 +2118,79 @@ def test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum():
         f"Expected 'tool_calls' but got {final.choices[0].finish_reason!r}. "
         "STOP enum was not normalised through map_finish_reason()."
     )
+
+
+def test_record_partial_usage_for_failure_stashes_usage_and_cost():
+    """A stream that breaks mid-flight must surface the usage assembled from the
+    chunks already delivered, plus its cost, on the logging object so the
+    failure handler records the real partial spend instead of zero.
+    """
+    logging_obj = Logging(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": "Hey"}],
+        stream=True,
+        call_type="completion",
+        start_time=time.time(),
+        litellm_call_id="partial-usage-1",
+        function_id="1245",
+    )
+    logging_obj.model_call_details["custom_llm_provider"] = "openai"
+
+    wrapper = CustomStreamWrapper(
+        completion_stream=None,
+        model="gpt-4o-mini",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+    )
+    wrapper.chunks = [
+        ModelResponseStream(
+            id="chatcmpl-partial-1",
+            created=1742056047,
+            model="gpt-4o-mini",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason=None,
+                    index=0,
+                    delta=Delta(
+                        content="The Roman Empire began when", role="assistant"
+                    ),
+                )
+            ],
+            usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31),
+        )
+    ]
+
+    wrapper._record_partial_usage_for_failure()
+
+    stashed = logging_obj.model_call_details["combined_usage_object"]
+    assert stashed.prompt_tokens == 30
+    assert stashed.completion_tokens == 1
+    assert stashed.total_tokens == 31
+    assert isinstance(logging_obj.model_call_details["response_cost"], float)
+
+
+def test_record_partial_usage_for_failure_noop_without_chunks():
+    """With no chunks delivered there is nothing billed to recover, so the
+    failure stash must stay absent and not force a zero-usage row.
+    """
+    logging_obj = Logging(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": "Hey"}],
+        stream=True,
+        call_type="completion",
+        start_time=time.time(),
+        litellm_call_id="partial-usage-2",
+        function_id="1245",
+    )
+    wrapper = CustomStreamWrapper(
+        completion_stream=None,
+        model="gpt-4o-mini",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+    )
+    wrapper.chunks = []
+
+    wrapper._record_partial_usage_for_failure()
+
+    assert "combined_usage_object" not in logging_obj.model_call_details

tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py

@@ -1067,3 +1067,40 @@ async def test_failure_hook_drops_error_information_traceback_when_env_set(
     assert "traceback" not in error_information
     assert error_information["error_class"] == "RuntimeError"
     assert error_information["error_message"] == "boom-with-traceback"
+
+
+@pytest.mark.asyncio
+async def test_async_post_call_failure_hook_records_recovered_partial_spend():
+    """A stream that broke mid-flight still billed the provider. The failure
+    hook lifts the recovered cost onto request_data as ``response_cost``; this
+    hook must pass it through to update_database so the failure row records the
+    real partial spend instead of the hardcoded zero.
+    """
+    from litellm.types.utils import Usage
+
+    logger = _ProxyDBLogger()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key", user_id="u", team_id="t")
+
+    request_data = {
+        "model": "anthropic/claude-haiku-4-5",
+        "messages": [{"role": "user", "content": "Hello"}],
+        "metadata": {},
+        "proxy_server_request": {"request_id": "rid"},
+        "response_cost": 3.5e-05,
+        "combined_usage_object": Usage(
+            prompt_tokens=30, completion_tokens=1, total_tokens=31
+        ),
+    }
+
+    with patch(
+        "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
+        new_callable=AsyncMock,
+    ) as mock_update_database:
+        await logger.async_post_call_failure_hook(
+            request_data=request_data,
+            original_exception=Exception("MidStreamFallbackError: read timeout"),
+            user_api_key_dict=user_api_key_dict,
+        )
+
+        mock_update_database.assert_called_once()
+        assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05

tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py

@@ -2009,3 +2009,50 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form(
     assert sanitized is not None
     assert "leaked-via-pydantic-msg" not in sanitized["error_message"]
     assert REDACTED_BY_LITELM_STRING in sanitized["error_message"]
+
+
+def test_get_logging_payload_uses_recovered_combined_usage_on_failure():
+    """A request that fails mid-stream has no usable response_obj usage, but the
+    streaming handler recovers the usage from the chunks already delivered and
+    the failure hook surfaces it as ``combined_usage_object``. The spend-log
+    payload must record those token counts instead of zero.
+    """
+    from litellm.types.utils import Usage
+
+    kwargs = {
+        "model": "anthropic/claude-haiku-4-5",
+        "call_type": "acompletion",
+        "litellm_params": {"metadata": {"user_api_key": "sk-test"}},
+        "combined_usage_object": Usage(
+            prompt_tokens=30, completion_tokens=1, total_tokens=31
+        ),
+    }
+    response_obj = Exception("MidStreamFallbackError: read timeout")
+    now = datetime.datetime.now(timezone.utc)
+
+    payload = get_logging_payload(
+        kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now
+    )
+
+    assert payload["prompt_tokens"] == 30
+    assert payload["completion_tokens"] == 1
+    assert payload["total_tokens"] == 31
+
+
+def test_get_logging_payload_failure_without_recovered_usage_is_zero():
+    """A failure with no recovered usage keeps zero token counts, so the
+    combined-usage override never invents tokens for ordinary failures.
+    """
+    kwargs = {
+        "model": "anthropic/claude-haiku-4-5",
+        "call_type": "acompletion",
+        "litellm_params": {"metadata": {"user_api_key": "sk-test"}},
+    }
+    response_obj = Exception("BadRequestError")
+    now = datetime.datetime.now(timezone.utc)
+
+    payload = get_logging_payload(
+        kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now
+    )
+
+    assert payload["total_tokens"] == 0

tests/test_litellm/proxy/test_proxy_utils.py

@@ -435,3 +435,52 @@ class TestPostCallFailureHookProxyExceptionLogging:
             )
             is False
         )
+class TestPostCallFailureHookLiftsRecoveredPartialSpend:
+    """A stream that broke mid-flight still billed the provider for the chunks
+    already delivered. The streaming handler stashes that recovered usage and
+    cost on the logging object; post_call_failure_hook must lift them onto
+    request_data before the logging object is popped, so the failure-path spend
+    callbacks (which run after the pop) record the real partial spend.
+    """
+
+    async def _run(self, request_data):
+        from unittest.mock import AsyncMock, patch
+
+        from litellm.proxy._types import UserAPIKeyAuth
+
+        proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
+        proxy_logging_obj.alert_types = []
+        with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()):
+            await proxy_logging_obj.post_call_failure_hook(
+                request_data=request_data,
+                original_exception=Exception("boom"),
+                user_api_key_dict=UserAPIKeyAuth(),
+            )
+
+    @pytest.mark.asyncio
+    async def test_lifts_recovered_usage_and_cost(self):
+        from litellm.types.utils import Usage
+
+        recovered_usage = Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31)
+        logging_obj = MagicMock()
+        logging_obj.model_call_details = {
+            "combined_usage_object": recovered_usage,
+            "response_cost": 3.5e-05,
+        }
+        request_data = {"litellm_logging_obj": logging_obj, "metadata": {}}
+        await self._run(request_data)
+
+        assert request_data["combined_usage_object"] is recovered_usage
+        assert request_data["response_cost"] == 3.5e-05
+        assert "litellm_logging_obj" not in request_data
+
+    @pytest.mark.asyncio
+    async def test_no_recovered_usage_is_noop(self):
+        logging_obj = MagicMock()
+        logging_obj.model_call_details = {}
+        request_data = {"litellm_logging_obj": logging_obj, "metadata": {}}
+        await self._run(request_data)
+        assert "combined_usage_object" not in request_data
+        assert "response_cost" not in request_data
+
+

tests/test_litellm/test_router.py

@@ -3340,6 +3340,151 @@ def test_combine_fallback_usage():
     assert chunk.usage.total_tokens == 15
 
 
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure():
+    """A mid-stream failure with no successful fallback raises and is logged as
+    a failure, so the router must never dispatch it as a success. Partial-spend
+    recovery for the failure row happens in the streaming handler, not here, so
+    this guards only against reintroducing a success log for a failed stream.
+    """
+    from litellm.exceptions import MidStreamFallbackError
+    from litellm.types.utils import Delta, StreamingChoices, Usage
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"},
+            },
+        ],
+        set_verbose=True,
+    )
+
+    error = MidStreamFallbackError(
+        message="Connection lost",
+        model="gpt-4",
+        llm_provider="openai",
+        generated_content="The Roman Empire began when",
+    )
+
+    def _make_interrupted_model_response():
+        partial_chunk = litellm.ModelResponseStream(
+            id="chatcmpl-partial-1",
+            created=1742056047,
+            model="gpt-4",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason=None,
+                    index=0,
+                    delta=Delta(
+                        content="The Roman Empire began when", role="assistant"
+                    ),
+                )
+            ],
+            usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26),
+        )
+
+        class _RaisingStream:
+            def __init__(self):
+                self.index = 0
+                self.chunks = [partial_chunk]
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                if self.index == 0:
+                    self.index += 1
+                    return partial_chunk
+                raise error
+
+        stream = _RaisingStream()
+        logging_obj = MagicMock()
+        logging_obj.dispatch_success_handlers = AsyncMock()
+        logging_obj.model_call_details = {}
+        setattr(stream, "model", "gpt-4")
+        setattr(stream, "custom_llm_provider", "openai")
+        setattr(stream, "logging_obj", logging_obj)
+        return stream, logging_obj
+
+    messages = [{"role": "user", "content": "Hello"}]
+    initial_kwargs = {"model": "gpt-4", "stream": True}
+
+    # Terminal path: no successful fallback -> the error propagates and the
+    # router never dispatches a success for the failed stream.
+    model_response, logging_obj = _make_interrupted_model_response()
+    with patch.object(
+        router,
+        "async_function_with_fallbacks_common_utils",
+        new=AsyncMock(side_effect=error),
+    ):
+        result = await router._acompletion_streaming_iterator(
+            model_response=model_response,
+            messages=messages,
+            initial_kwargs=dict(initial_kwargs),
+        )
+        collected = []
+        with pytest.raises(MidStreamFallbackError):
+            async for chunk in result:
+                collected.append(chunk)
+
+    assert len(collected) == 1
+    logging_obj.dispatch_success_handlers.assert_not_called()
+
+    # Fallback success: the fallback stream owns success accounting via
+    # _combine_fallback_usage, so this iterator must not dispatch its own.
+    model_response, logging_obj = _make_interrupted_model_response()
+
+    class _FallbackStream:
+        def __init__(self, items):
+            self.items = items
+            self.index = 0
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            if self.index >= len(self.items):
+                raise StopAsyncIteration
+            item = self.items[self.index]
+            self.index += 1
+            return item
+
+    fallback_stream = _FallbackStream(
+        [
+            litellm.ModelResponseStream(
+                id="chatcmpl-fallback-1",
+                model="gpt-3.5-turbo",
+                object="chat.completion.chunk",
+                choices=[
+                    StreamingChoices(
+                        finish_reason=None,
+                        index=0,
+                        delta=Delta(content=" continued", role="assistant"),
+                    )
+                ],
+            )
+        ]
+    )
+    with patch.object(
+        router,
+        "async_function_with_fallbacks_common_utils",
+        new=AsyncMock(return_value=fallback_stream),
+    ):
+        result = await router._acompletion_streaming_iterator(
+            model_response=model_response,
+            messages=messages,
+            initial_kwargs=dict(initial_kwargs),
+        )
+        collected = []
+        async for chunk in result:
+            collected.append(chunk)
+
+    assert len(collected) == 2
+    logging_obj.dispatch_success_handlers.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_team_scoped_model_fallback():
     """