From 46fb1cd514cd2b21db62146fae06c704f8193e63 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:07:57 -0700 Subject: [PATCH] fix(proxy): reassemble fragmented SSE frames and inject logging dependency --- .../streaming_handler.py | 42 ++++++++++++++-- .../test_streaming_handler_interrupt.py | 49 +++++++++++++------ 2 files changed, 72 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 192600ba150..da4a8eceb89 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/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, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 6d47897f710..7b649db43ed 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/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)