fix(streaming): assemble partial stream usage off the event loop

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/litellm_core_utils/streaming_handler.py

@@ -2068,7 +2068,7 @@ class CustomStreamWrapper:
             ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT
             traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}"
             if self.logging_obj is not None:
-                self._record_partial_usage_for_failure()
+                await asyncio.to_thread(self._record_partial_usage_for_failure)
                 ## LOGGING
                 asyncio.create_task(
                     self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
@@ -2076,10 +2076,10 @@ class CustomStreamWrapper:
             self._handle_stream_fallback_error(e)
         except (httpx.ReadError, httpx.RemoteProtocolError) as e:
             if self.received_finish_reason is None:
-                self._log_stream_failure_and_raise(e)
+                await self._alog_stream_failure_and_raise(e)
             return await self._finalize_completed_stream(cache_hit=cache_hit)
         except Exception as e:
-            self._log_stream_failure_and_raise(e)
+            await self._alog_stream_failure_and_raise(e)
 
     async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
         if self.sent_last_chunk is True:
@@ -2183,10 +2183,10 @@ class CustomStreamWrapper:
             # relies on aclose() or the best-effort __del__ guard.
             return processed_chunk
 
-    def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn:
+    async def _alog_stream_failure_and_raise(self, e: Exception) -> NoReturn:
         traceback_exception: Final = traceback.format_exc()
         if self.logging_obj is not None:
-            self._record_partial_usage_for_failure()
+            await asyncio.to_thread(self._record_partial_usage_for_failure)
             ## LOGGING
             asyncio.create_task(
                 self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)

litellm/proxy/common_request_processing.py

@@ -183,7 +183,10 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
 
     Awaited directly by the shielded cleanup rather than scheduled with
     create_task: the client is already gone so the extra latency is harmless,
-    and an unrooted task could be garbage-collected before it bills.
+    and an unrooted task could be garbage-collected before it bills. Assembly
+    is offloaded to a thread because a partial stream never carries a usage
+    chunk, so it re-tokenizes the whole delivered response with tiktoken, which
+    on a multi-MB stream blocks the event loop for hundreds of milliseconds.
 
     Returns True when a disconnect-time success event owns the request's
     max_parallel_requests slot release (one was dispatched here, or one had
@@ -211,7 +214,8 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
     )
     messages: Final[object] = getattr(response, "messages", None)
     try:
-        partial_response: Final = litellm.stream_chunk_builder(
+        partial_response: Final = await asyncio.to_thread(
+            litellm.stream_chunk_builder,
             chunks=chunks,
             messages=messages if isinstance(messages, list) else None,
             logging_obj=logging_obj,

tests/test_litellm/litellm_core_utils/test_streaming_handler.py

@@ -3190,6 +3190,63 @@ def test_record_partial_usage_for_failure_stashes_usage_and_cost():
     assert isinstance(logging_obj.model_call_details["response_cost"], float)
 
 
+@pytest.mark.asyncio
+async def test_async_stream_failure_recovers_partial_usage_off_the_event_loop():
+    """Recovering partial usage re-tokenizes everything streamed so far. On a
+    multi-MB stream that costs hundreds of ms of tiktoken, so the async failure
+    path must not run it on the event loop thread and stall the whole worker.
+    """
+    import threading
+
+    class _FailingStream:
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            raise ValueError("upstream died mid-stream")
+
+    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-offload",
+        function_id="1245",
+    )
+    logging_obj.model_call_details["custom_llm_provider"] = "openai"
+
+    wrapper = CustomStreamWrapper(
+        completion_stream=_FailingStream(),
+        model="gpt-4o-mini",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+    )
+    wrapper.chunks = [
+        ModelResponseStream(
+            id="chatcmpl-partial-offload",
+            created=1742056047,
+            model="gpt-4o-mini",
+            object="chat.completion.chunk",
+            choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Rome fell", role="assistant"))],
+        )
+    ]
+
+    builder_threads = []
+    real_builder = litellm.stream_chunk_builder
+
+    def _record_thread(*args, **kwargs):
+        builder_threads.append(threading.get_ident())
+        return real_builder(*args, **kwargs)
+
+    with patch.object(litellm, "stream_chunk_builder", _record_thread):
+        with pytest.raises(Exception):
+            await wrapper.__anext__()
+
+    assert len(builder_threads) == 1
+    assert builder_threads[0] != threading.get_ident()
+
+
 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.

tests/test_litellm/proxy/test_common_request_processing.py

@@ -1,6 +1,7 @@
 import asyncio
 import copy
 import datetime
+import threading
 from types import SimpleNamespace
 from typing import AsyncGenerator, Callable, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -5097,6 +5098,31 @@ class TestStreamingClientDisconnectBilling:
         await stream_iter.__anext__()
         return response
 
+    @pytest.mark.asyncio
+    async def test_disconnect_billing_assembles_off_the_event_loop(self):
+        """
+        A partial stream carries no usage chunk, so assembling it re-tokenizes
+        the whole delivered response. On a multi-MB stream that is hundreds of
+        ms of tiktoken, which stalls every other request on the worker if it
+        runs on the event loop thread.
+        """
+        response = await self._start_partial_stream()
+        builder_threads: list[int] = []
+        real_builder = litellm.stream_chunk_builder
+
+        def _record_thread(*args, **kwargs):
+            builder_threads.append(threading.get_ident())
+            return real_builder(*args, **kwargs)
+
+        with patch.object(litellm, "stream_chunk_builder", _record_thread):
+            billed = await _bill_partial_streamed_spend_on_disconnect(
+                {"litellm_logging_obj": response.logging_obj}, response
+            )
+
+        assert billed is True
+        assert len(builder_threads) == 1
+        assert builder_threads[0] != threading.get_ident()
+
     @pytest.mark.asyncio
     async def test_disconnect_bills_partial_streamed_spend(self):
         recorder = _RecordingSuccessLogger()