fix(proxy): reassemble split SSE frames before restamping anthropic message_start

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)
+        )

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
 

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"