fix(proxy): reassemble fragmented SSE frames and inject logging dependency

litellm/proxy/pass_through_endpoints/streaming_handler.py

@@ -1,5 +1,6 @@
+from collections.abc import Coroutine
 from datetime import datetime
-from typing import Final
+from typing import Final, Protocol
 
 import httpx
 
@@ -24,6 +25,21 @@ from .llm_provider_handlers.vertex_passthrough_logging_handler import (
 from .success_handler import PassThroughEndpointLogging
 
 
+class RouteStreamingLogging(Protocol):
+    def __call__(
+        self,
+        *,
+        litellm_logging_obj: LiteLLMLoggingObj,
+        passthrough_success_handler_obj: PassThroughEndpointLogging,
+        url_route: str,
+        request_body: dict,
+        endpoint_type: EndpointType,
+        start_time: datetime,
+        raw_bytes: list[bytes],
+        end_time: datetime,
+    ) -> Coroutine[None, None, None]: ...
+
+
 class PassThroughStreamingHandler:
     @staticmethod
     def _stamp_first_chunk_if_needed(litellm_logging_obj: LiteLLMLoggingObj) -> None:
@@ -39,7 +55,11 @@ class PassThroughStreamingHandler:
         start_time: datetime,
         passthrough_success_handler_obj: PassThroughEndpointLogging,
         url_route: str,
+        route_streaming_logging: RouteStreamingLogging | None = None,
     ):
+        resolved_route_streaming_logging: Final[RouteStreamingLogging] = (
+            route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler
+        )
         raw_bytes: Final[list[bytes]] = []
         logging_scheduled = False
         model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection(
@@ -77,10 +97,19 @@ class PassThroughStreamingHandler:
                 # -> ``str`` for the per-chunk call site.
                 assert model_name is not None
                 resolved_model_name: Final[str] = model_name
+                pending = b""
                 async for chunk in response.aiter_bytes():
                     raw_bytes.append(chunk)
                     PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj)
-                    yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, resolved_model_name)
+                    complete_frames, pending = PassThroughStreamingHandler._split_complete_sse_frames(
+                        pending + chunk
+                    )  # rebind-ok: SSE frame reassembly buffer across transport chunks
+                    if complete_frames:
+                        yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
+                            complete_frames, resolved_model_name
+                        )
+                if pending:
+                    yield pending
         except Exception as e:
             verbose_proxy_logger.error("Error in chunk_processor: %s", e)
             raise
@@ -96,7 +125,7 @@ class PassThroughStreamingHandler:
                 logging_scheduled = True
                 try:
                     GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
-                        async_coroutine=PassThroughStreamingHandler._route_streaming_logging_to_handler(
+                        async_coroutine=resolved_route_streaming_logging(
                             litellm_logging_obj=litellm_logging_obj,
                             passthrough_success_handler_obj=passthrough_success_handler_obj,
                             url_route=url_route,
@@ -110,6 +139,13 @@ class PassThroughStreamingHandler:
                 except Exception as e:
                     verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e)
 
+    @staticmethod
+    def _split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]:
+        frame_boundary: Final = pending.rfind(b"\n\n")
+        if frame_boundary == -1:
+            return b"", pending
+        return pending[: frame_boundary + 2], pending[frame_boundary + 2 :]
+
     @staticmethod
     async def _route_streaming_logging_to_handler(
         litellm_logging_obj: LiteLLMLoggingObj,

tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py

@@ -383,23 +383,19 @@ def _openai_passthrough_stream_chunks():
 
 async def _collect_openai_passthrough_chunks(chunks, endpoint_type):
     response = _make_streaming_response(chunks)
-    with patch.object(
-        PassThroughStreamingHandler,
-        "_route_streaming_logging_to_handler",
-        new=AsyncMock(),
+    received = []
+    async for chunk in PassThroughStreamingHandler.chunk_processor(
+        response=response,
+        request_body={"model": "gpt-4o-mini", "stream": True},
+        litellm_logging_obj=MagicMock(),
+        endpoint_type=endpoint_type,
+        start_time=datetime.now(),
+        passthrough_success_handler_obj=MagicMock(),
+        url_route="/openai/v1/chat/completions",
+        route_streaming_logging=AsyncMock(),
     ):
-        received = []
-        async for chunk in PassThroughStreamingHandler.chunk_processor(
-            response=response,
-            request_body={"model": "gpt-4o-mini", "stream": True},
-            litellm_logging_obj=MagicMock(),
-            endpoint_type=endpoint_type,
-            start_time=datetime.now(),
-            passthrough_success_handler_obj=MagicMock(),
-            url_route="/openai/v1/chat/completions",
-        ):
-            received.append(chunk)
-        await asyncio.sleep(0)
+        received.append(chunk)
+    await asyncio.sleep(0)
     return received
 
 
@@ -426,6 +422,27 @@ async def test_chunk_processor_injects_cost_into_openai_passthrough_usage_frame(
     assert final_payload["usage"]["total_tokens"] == 15
 
 
+@pytest.mark.asyncio
+async def test_chunk_processor_injects_cost_into_usage_frame_fragmented_across_chunks(monkeypatch):
+    """Regression: an SSE usage frame split across transport chunks must still get
+    cost injected once the frame completes, instead of passing through untouched."""
+    monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
+    whole = _openai_passthrough_stream_chunks()
+    usage_frame = whole[2]
+    split_at = len(usage_frame) // 2
+    chunks = [whole[0], whole[1], usage_frame[:split_at], usage_frame[split_at:], whole[3]]
+
+    received = await _collect_openai_passthrough_chunks(chunks, EndpointType.OPENAI)
+
+    reassembled = b"".join(received).decode("utf-8")
+    usage_lines = [ln for ln in reassembled.split("\n") if '"total_tokens"' in ln]
+    assert len(usage_lines) == 1
+    final_payload = json.loads(usage_lines[0].split("data:", 1)[1].strip())
+    assert final_payload["usage"]["cost"] > 0
+    assert final_payload["usage"]["prompt_tokens"] == 11
+    assert reassembled.endswith("data: [DONE]\n\n")
+
+
 @pytest.mark.asyncio
 async def test_chunk_processor_flag_off_leaves_openai_passthrough_stream_byte_identical(monkeypatch):
     monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)