fix(proxy): bound the pass-through stream buffer kept for logging

chunk_processor appended every chunk of a config-defined pass-through
stream to a list that lived until end-of-stream logging, so a stream of
large base64 images grew without bound and the proxy was OOM-killed.

Keep the first and last halves of
litellm.pass_through_max_logged_stream_bytes (default 16 MiB, None for
unbounded) and drop the middle. Both ends are cut back to whole SSE
events, so the first and final usage events survive. The client still
receives every byte, and over the limit one warning logs the sizes.

Fixes #48010

litellm/__init__.py

@@ -247,6 +247,10 @@ use_chat_completions_url_for_anthropic_messages: bool = bool(
 strip_anthropic_total_tokens: bool = False
 anthropic_sse_ping_interval_seconds: float = 15.0
 sse_keepalive_ping_interval_seconds: float | None = None
+# Most bytes of one streamed pass-through response kept in memory for end-of-stream logging and cost
+# tracking. Past it the middle is dropped and the first and last halves are kept, so the final usage
+# event survives. None keeps everything. Set via `litellm_settings.pass_through_max_logged_stream_bytes`.
+pass_through_max_logged_stream_bytes: int | None = 16 * 1024 * 1024
 route_all_chat_openai_to_responses: bool = (
     os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true"
 )  # When True, routes all OpenAI /chat/completions requests through the Responses API bridge

litellm/proxy/pass_through_endpoints/bounded_chunk_buffer.py

@@ -0,0 +1,120 @@
+from collections import deque
+from collections.abc import Iterator, Sequence
+from typing import Final, overload
+
+_FRAME_SEPARATORS: Final = (b"\r\n\r\n", b"\n\n")
+
+
+def _end_of_first_frame(data: bytes) -> int:
+    """Index just past the first complete frame in ``data``; 0 when it holds none."""
+    ends: Final = [i + len(sep) for sep in _FRAME_SEPARATORS if (i := data.find(sep)) != -1]
+    if ends:
+        return min(ends)
+    newline: Final = data.find(b"\n")
+    return 0 if newline == -1 else newline + 1
+
+
+def _end_of_last_frame(data: bytes) -> int:
+    """Index just past the last complete frame in ``data``; 0 when it holds none."""
+    ends: Final = [i + len(sep) for sep in _FRAME_SEPARATORS if (i := data.rfind(sep)) != -1]
+    if ends:
+        return max(ends)
+    newline: Final = data.rfind(b"\n")
+    return 0 if newline == -1 else newline + 1
+
+
+class BoundedChunkBuffer(Sequence[bytes]):
+    """
+    Collects the chunks of a streamed response for end-of-stream logging without letting
+    an unbounded stream exhaust memory.
+
+    Half of ``max_bytes`` keeps the start of the stream and half keeps the end, so the
+    first event (input usage, model) and the terminal event (final usage) both survive.
+    When the stream fits, every chunk is kept untouched. When it does not, the bytes in
+    between are dropped and both ends are trimmed back to a frame boundary, so the
+    retained bytes never contain a half-cut event.
+
+    ``max_bytes=None`` keeps everything.
+    """
+
+    def __init__(self, max_bytes: int | None = None) -> None:
+        self._head_limit: Final = None if max_bytes is None else max_bytes // 2
+        self._tail_limit: Final = None if max_bytes is None else max_bytes - (max_bytes // 2)
+        self._head: Final[list[bytes]] = []
+        self._tail: Final[deque[bytes]] = deque()
+        self._head_size = 0  # rebind-ok: bytes currently held in the head
+        self._tail_size = 0  # rebind-ok: bytes currently held in the tail
+        self.total_bytes = 0  # rebind-ok: every byte seen, kept or not
+        self._head_closed = False  # rebind-ok: set once the head budget is spent
+
+    @property
+    def truncated(self) -> bool:
+        return self.total_bytes > self._head_size + self._tail_size
+
+    @property
+    def retained_bytes(self) -> int:
+        return sum(len(chunk) for chunk in self)
+
+    def append(self, chunk: bytes) -> None:
+        self.total_bytes += len(chunk)
+        if self._head_limit is None or self._tail_limit is None:
+            self._head.append(chunk)
+            self._head_size += len(chunk)
+            return
+        if not self._head_closed:
+            room = self._head_limit - self._head_size
+            if len(chunk) <= room:
+                self._head.append(chunk)
+                self._head_size += len(chunk)
+                return
+            if room > 0:
+                self._head.append(chunk[:room])
+                self._head_size += room
+            self._head_closed = True  # rebind-ok: the head never reopens
+            chunk = chunk[room:]
+        self._tail.append(chunk[-self._tail_limit :] if self._tail_limit else b"")
+        self._tail_size += len(self._tail[-1])
+        while self._tail_size > self._tail_limit:
+            overflow = self._tail_size - self._tail_limit
+            oldest = self._tail.popleft()
+            if len(oldest) > overflow:
+                self._tail.appendleft(oldest[overflow:])
+                self._tail_size -= overflow
+                break
+            self._tail_size -= len(oldest)
+
+    def _frames(self) -> list[bytes]:
+        if not self.truncated:
+            return [*self._head, *self._tail]
+        head: Final = b"".join(self._head)
+        tail: Final = b"".join(self._tail)
+        head_end: Final = _end_of_last_frame(head)
+        tail_start: Final = 0 if not self._tail_size else self._tail_start(tail)
+        return [part for part in (head[:head_end], tail[tail_start:]) if part]
+
+    def _tail_start(self, tail: bytes) -> int:
+        # the tail begins wherever the eviction cut it, which is mid-frame unless nothing before it was dropped
+        return _end_of_first_frame(tail) if self._head_closed else 0
+
+    def __iter__(self) -> Iterator[bytes]:
+        return iter(self._frames())
+
+    def __len__(self) -> int:
+        return len(self._frames())
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, Sequence) and list(self) == list(other)
+
+    __hash__ = None  # type: ignore[assignment]
+
+    def __bool__(self) -> bool:
+        return self.total_bytes > 0
+
+    @overload
+    def __getitem__(self, index: int) -> bytes: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> Sequence[bytes]: ...
+
+    def __getitem__(self, index: int | slice) -> bytes | Sequence[bytes]:
+        return self._frames()[index]

litellm/proxy/pass_through_endpoints/streaming_handler.py

@@ -19,6 +19,7 @@ from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
 from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
 from litellm.types.utils import StandardPassThroughResponseObject, Usage
 
+from .bounded_chunk_buffer import BoundedChunkBuffer
 from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
     AnthropicPassthroughLoggingHandler,
 )
@@ -49,7 +50,7 @@ class RouteStreamingLogging(Protocol):
         request_body: dict,
         endpoint_type: EndpointType,
         start_time: datetime,
-        raw_bytes: list[bytes],
+        raw_bytes: Sequence[bytes],
         end_time: datetime,
     ) -> Coroutine[None, None, None]: ...
 
@@ -170,10 +171,11 @@ class PassThroughStreamingHandler:
         resolved_route_streaming_logging: Final[RouteStreamingLogging] = (
             route_streaming_logging or PassThroughStreamingHandler.route_streaming_logging_to_handler
         )
-        raw_bytes: Final[list[bytes]] = []
+        raw_bytes: Final = BoundedChunkBuffer(litellm.pass_through_max_logged_stream_bytes)
         resolved_request_body: Final[dict[str, object]] = request_body or {}
 
         def _build_logging_coroutine() -> Coroutine[None, None, None]:
+            PassThroughStreamingHandler._warn_if_truncated(raw_bytes, url_route)
             return resolved_route_streaming_logging(
                 litellm_logging_obj=litellm_logging_obj,
                 passthrough_success_handler_obj=passthrough_success_handler_obj,
@@ -295,6 +297,7 @@ class PassThroughStreamingHandler:
                     if stream_ended
                     else PassThroughStreamingHandler._error_frame_closed_at(endpoint_type, raw_bytes)
                 )
+                PassThroughStreamingHandler._warn_if_truncated(raw_bytes, url_route)
                 try:
                     await PassThroughStreamingHandler._schedule_stream_end_logging(
                         litellm_logging_obj=litellm_logging_obj,
@@ -309,6 +312,17 @@ class PassThroughStreamingHandler:
                 else:
                     bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params)
 
+    @staticmethod
+    def _warn_if_truncated(raw_bytes: BoundedChunkBuffer, url_route: str) -> None:
+        if raw_bytes.truncated:
+            verbose_proxy_logger.warning(
+                "Pass-through stream for %s was %s bytes; only the first and last %s bytes are kept for logging "
+                "(litellm.pass_through_max_logged_stream_bytes)",
+                url_route,
+                raw_bytes.total_bytes,
+                raw_bytes.retained_bytes,
+            )
+
     @staticmethod
     def _error_frame_closed_at(endpoint_type: EndpointType, raw_bytes: Sequence[bytes]) -> tuple[str, str, int] | None:
         from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (

tests/unit/proxy/pass_through_endpoints/test_bounded_chunk_buffer.py

@@ -0,0 +1,134 @@
+import json
+from datetime import datetime
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+import pytest
+
+import litellm
+from litellm.proxy.pass_through_endpoints.bounded_chunk_buffer import BoundedChunkBuffer
+from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
+from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
+
+
+def _frame(name: str, payload: dict) -> bytes:
+    return f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode()
+
+
+def test_buffer_under_the_limit_keeps_every_chunk_untouched():
+    buffer = BoundedChunkBuffer(max_bytes=1024)
+    chunks = [b"event: a\ndata: 1\n\n", b"event: b\ndata: 2\n\n"]
+
+    for chunk in chunks:
+        buffer.append(chunk)
+
+    assert list(buffer) == chunks, f"truncated={buffer.truncated}"
+    assert not buffer.truncated
+
+
+def test_buffer_without_a_limit_keeps_everything():
+    buffer = BoundedChunkBuffer(max_bytes=None)
+    chunks = [b"x" * 1_000_000 for _ in range(5)]
+
+    for chunk in chunks:
+        buffer.append(chunk)
+
+    assert list(buffer) == chunks
+    assert buffer.total_bytes == 5_000_000
+
+
+def test_buffer_over_the_limit_keeps_the_first_and_last_events_whole():
+    first = _frame("message_start", {"usage": {"input_tokens": 29}})
+    middle = [_frame("content_block_delta", {"text": "x" * 40, "n": i}) for i in range(200)]
+    last = _frame("message_delta", {"usage": {"output_tokens": 77}})
+    buffer = BoundedChunkBuffer(max_bytes=600)
+
+    for chunk in [first, *middle, last]:
+        buffer.append(chunk)
+
+    kept = b"".join(buffer)
+    assert kept.startswith(first), kept
+    assert kept.endswith(last), kept
+    assert buffer.truncated
+    assert buffer.retained_bytes <= 600
+    assert buffer.total_bytes == len(first) + sum(map(len, middle)) + len(last)
+
+
+def test_buffer_never_keeps_a_half_cut_event_when_chunks_split_frames():
+    stream = b"".join(_frame("delta", {"i": i, "pad": "y" * 30}) for i in range(100))
+    buffer = BoundedChunkBuffer(max_bytes=500)
+
+    for start in range(0, len(stream), 7):
+        buffer.append(stream[start : start + 7])
+
+    for part in b"".join(buffer).split(b"\n\n"):
+        if part:
+            assert part.startswith(b"event: delta\ndata: {"), part
+            assert part.endswith(b"}"), part
+
+
+def test_buffer_holds_one_chunk_larger_than_the_limit_to_its_head_and_tail():
+    buffer = BoundedChunkBuffer(max_bytes=100)
+
+    buffer.append(b"event: big\ndata: " + b"z" * 10_000 + b"\n\n")
+
+    assert buffer.truncated
+    assert buffer.retained_bytes <= 100
+    assert buffer.total_bytes == 10_000 + len(b"event: big\ndata: \n\n")
+
+
+def test_buffer_memory_stays_bounded_while_a_huge_stream_passes():
+    buffer = BoundedChunkBuffer(max_bytes=1000)
+
+    for _ in range(10_000):
+        buffer.append(b"data: " + b"q" * 1000 + b"\n\n")
+
+    assert buffer.total_bytes > 10_000_000
+    assert buffer.retained_bytes <= 1000
+
+
+def _response(chunks: list[bytes]) -> MagicMock:
+    response = MagicMock(spec=httpx.Response)
+    response.status_code = 200
+
+    async def _aiter_bytes():
+        for chunk in chunks:
+            yield chunk
+
+    response.aiter_bytes = _aiter_bytes
+    return response
+
+
+@pytest.mark.asyncio
+async def test_chunk_processor_relays_every_byte_but_logs_only_the_bounded_ends(
+    monkeypatch: pytest.MonkeyPatch, caplog
+):
+    monkeypatch.setattr(litellm, "pass_through_max_logged_stream_bytes", 400)
+    first = _frame("message_start", {"usage": {"input_tokens": 29}})
+    middle = [_frame("content_block_delta", {"text": "x" * 40, "n": i}) for i in range(100)]
+    last = _frame("message_delta", {"usage": {"output_tokens": 77}})
+    chunks = [first, *middle, last]
+    logging_obj = MagicMock()
+    logging_obj._on_deferred_stream_complete = None
+
+    with patch.object(PassThroughStreamingHandler, "route_streaming_logging_to_handler", new=AsyncMock()) as route:
+        relayed = [
+            chunk
+            async for chunk in PassThroughStreamingHandler.chunk_processor(
+                response=_response(chunks),
+                request_body={"model": "m"},
+                litellm_logging_obj=logging_obj,
+                endpoint_type=EndpointType.GENERIC,
+                start_time=datetime.now(),
+                passthrough_success_handler_obj=MagicMock(),
+                url_route="/codex/responses",
+            )
+        ]
+        import asyncio
+
+        await asyncio.sleep(0)
+
+    assert relayed == chunks
+    logged = b"".join(route.call_args.kwargs["raw_bytes"])
+    assert logged.startswith(first) and logged.endswith(last), logged
+    assert len(logged) <= 400