From c02c81452c932eff275755a43ff9c686cd8055d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:16:48 -0700 Subject: [PATCH] fix(proxy): reassemble split SSE frames before restamping anthropic message_start --- .../streaming_model_restamp.py | 78 +++++++++++++++++ litellm/proxy/common_request_processing.py | 8 +- .../test_streaming_model_restamp.py | 83 +++++++++++++++++++ 3 files changed, 165 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py index bbb7f2ceaa3..e8d54f03949 100644 --- a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -14,7 +14,11 @@ from typing import Final from pydantic import TypeAdapter, ValidationError _MESSAGE_START_EVENT: Final = "message_start" +_MESSAGE_START_MARKER: Final = b"message_start" _SSE_DATA_FIELD: Final = "data:" +_SSE_FRAME_END: Final = b"\n\n" +_MAX_HELD_BYTES: Final = 65536 +_PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"') _EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -79,3 +83,77 @@ def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> return chunk if restamped_text is None else restamped_text return chunk + + +def _is_ping_frame(frame: bytes) -> bool: + return any(marker in frame for marker in _PING_MARKERS) + + +class AnthropicStreamModelRestamper: + """ + Per-stream restamper for the encoded passthrough path, where chunks are raw + transport reads: the ``message_start`` SSE frame can arrive split across + chunks or coalesced with later frames. Complete frames are emitted as their + terminator closes them and an incomplete tail is held until it completes, + so the restamp never misses a torn frame. Once ``message_start`` has been + handled, or the first real event proves the stream carries none, every + later chunk passes through untouched. + """ + + def __init__(self, requested_model: str) -> None: + self._requested_model: Final = requested_model + self._held = b"" + self._armed = True + + def process(self, chunk: object) -> object: + if not self._armed: + return chunk + if isinstance(chunk, (bytes, bytearray)): + return self._process_encoded(bytes(chunk)) + if isinstance(chunk, str): + return self._process_encoded(chunk.encode("utf-8")) + restamped: Final = restamp_anthropic_stream_chunk_model(chunk, self._requested_model) + if isinstance(chunk, dict) and chunk.get("type") not in (None, "ping"): + self._armed = False + return restamped + + def _process_encoded(self, data: bytes) -> bytes: + combined: Final = self._held + data + if _SSE_FRAME_END not in combined: + if len(combined) > _MAX_HELD_BYTES: + self._held = b"" + self._armed = False + return combined + self._held = combined + return b"" + closed, _, tail = combined.rpartition(_SSE_FRAME_END) + emitted: Final = self._restamped_closed_block(closed + _SSE_FRAME_END) + if not self._armed: + self._held = b"" + return emitted + tail + self._held = tail + return emitted + + def _restamped_closed_block(self, closed: bytes) -> bytes: + frames: Final = tuple(closed.split(_SSE_FRAME_END)[:-1]) + decider: Final = next( + ( + index + for index, frame in enumerate(frames) + if _MESSAGE_START_MARKER in frame or (b"data:" in frame and not _is_ping_frame(frame)) + ), + None, + ) + if decider is None: + return closed + self._armed = False + decider_frame: Final = frames[decider] + _SSE_FRAME_END + if _MESSAGE_START_MARKER not in decider_frame: + return closed + restamped_text: Final = _restamped_frame(decider_frame.decode("utf-8", errors="ignore"), self._requested_model) + if restamped_text is None: + return closed + return b"".join( + restamped_text.encode("utf-8") if index == decider else frame + _SSE_FRAME_END + for index, frame in enumerate(frames) + ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 23e887f34a3..989cc7c18fb 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -177,7 +177,7 @@ if TYPE_CHECKING: else: ProxyConfig = Any from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( - restamp_anthropic_stream_chunk_model, + AnthropicStreamModelRestamper, ) from litellm.proxy.litellm_pre_call_utils import ( add_litellm_data_to_request, @@ -3414,10 +3414,10 @@ class ProxyBaseLLMRequestProcessing: if not restamp_model: return ProxyBaseLLMRequestProcessing.return_sse_chunk + restamper: Final = AnthropicStreamModelRestamper(restamp_model) + def serialize(chunk: object) -> str: - return ProxyBaseLLMRequestProcessing.return_sse_chunk( - restamp_anthropic_stream_chunk_model(chunk, restamp_model) - ) + return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk)) return serialize diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py index 6e2c3f49445..385173b24a9 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + AnthropicStreamModelRestamper, restamp_anthropic_stream_chunk_model, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -107,3 +108,85 @@ async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): ] assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001" + + +def test_restamps_message_start_split_across_transport_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + held = restamper.process(frame[:25]) + emitted = restamper.process(frame[25:]) + + assert held == b"" + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + + +def test_emits_coalesced_frames_with_only_message_start_rewritten(): + delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + combined = _message_start_frame("claude-haiku-4-5-20251001") + delta + + emitted = restamper_output = AnthropicStreamModelRestamper("claude-auto-1").process(combined) + + assert isinstance(restamper_output, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + assert emitted.endswith(delta) + + +def test_ping_frames_keep_the_restamper_armed(): + ping = b'event: ping\ndata: {"type": "ping"}\n\n' + frame = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(ping) == ping + reassembled = restamper.process(frame[:10]) + reassembled += restamper.process(frame[10:]) + + assert _model_from_frame(reassembled) == "claude-auto-1" + + +def test_first_non_ping_event_disarms_the_restamper(): + delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + late_message_start = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(delta) == delta + assert restamper.process(late_message_start) == late_message_start + + +def test_oversized_unterminated_chunk_flushes_unmodified(): + blob = b"data: " + b"x" * 70000 + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(blob) == blob + frame = _message_start_frame("claude-haiku-4-5-20251001") + assert restamper.process(frame) == frame + + +def test_dict_message_start_disarms_after_restamp(): + restamper = AnthropicStreamModelRestamper("claude-auto-1") + first = restamper.process({"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}}) + second = {"type": "message_start", "message": {"id": "msg_2", "model": "claude-sonnet-4-6"}} + + assert first == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-1"}} + assert restamper.process(second) == second + + +@pytest.mark.asyncio +async def test_sse_generator_restamps_message_start_split_across_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001") + proxy_logging_obj = _proxy_logging_obj_streaming([frame[:30], frame[30:]]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks) + assert _model_from_frame(joined) == "claude-auto-1"