From c24927cf2ac363af4c55cf701181732f1849c8f7 Mon Sep 17 00:00:00 2001 From: tin Date: Tue, 4 Aug 2026 20:31:55 +0000 Subject: [PATCH 1/3] fix(proxy): report requested model on Anthropic streaming message_start Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_model_restamp.py | 79 +++++++++++++ litellm/proxy/common_request_processing.py | 25 +++- .../test_streaming_model_restamp.py | 109 ++++++++++++++++++ 3 files changed, 212 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/anthropic_endpoints/streaming_model_restamp.py create mode 100644 tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py new file mode 100644 index 00000000000..857b6abd065 --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -0,0 +1,79 @@ +""" +Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only +stream event carrying a model, so streamed responses report the requested model like +non-streaming ones do. + +Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the +provider passthrough path) or as event dicts (fake-stream and agentic paths). +""" + +import json + +from pydantic import TypeAdapter, ValidationError + +_MESSAGE_START_EVENT = "message_start" +_SSE_DATA_FIELD = "data:" + +_EVENT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) + + +def _restamped_event(event: dict[str, object], requested_model: str) -> dict[str, object] | None: + message = event.get("message") + if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict): + return None + if message.get("model") == requested_model: + return None + return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is + + +def _restamped_data_line(line: str, requested_model: str) -> str | None: + stripped = line.strip() + if not stripped.startswith(_SSE_DATA_FIELD): + return None + payload = stripped[len(_SSE_DATA_FIELD) :].strip() + if not payload or payload == "[DONE]": + return None + try: + event = _EVENT_ADAPTER.validate_json(payload) + except ValidationError: + return None + restamped = _restamped_event(event, requested_model) + if restamped is None: + return None + return f"data: {json.dumps(restamped, separators=(',', ':'))}" + + +def _restamped_frame(frame: str, requested_model: str) -> str | None: + lines = frame.split("\n") + restamped = tuple(_restamped_data_line(line, requested_model) for line in lines) + if all(line is None for line in restamped): + return None + return "\n".join(new if new is not None else old for new, old in zip(restamped, lines)) + + +def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object: + """ + Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``. + + Chunks that carry no model are returned unchanged. + """ + if isinstance(chunk, dict): + try: + event = _EVENT_ADAPTER.validate_python(chunk) + except ValidationError: + return chunk + return _restamped_event(event, requested_model) or chunk + + if isinstance(chunk, (bytes, bytearray)): + if _MESSAGE_START_EVENT.encode() not in chunk: + return chunk + restamped = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model) + return chunk if restamped is None else restamped.encode("utf-8") + + if isinstance(chunk, str): + if _MESSAGE_START_EVENT not in chunk: + return chunk + restamped = _restamped_frame(chunk, requested_model) + return chunk if restamped is None else restamped + + return chunk diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f9cad283166..fd5de0debec 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -70,6 +70,9 @@ if TYPE_CHECKING: ProxyConfig = _ProxyConfig else: ProxyConfig = Any +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + restamp_anthropic_stream_chunk_model, +) from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.types.utils import ( ModelResponse, @@ -1953,6 +1956,9 @@ class ProxyBaseLLMRequestProcessing: request_data=self.data, proxy_logging_obj=proxy_logging_obj, request=request, + restamp_model=( + None if _should_return_raw_model_name(self.data) else requested_model_from_client + ), ) return await create_response( generator=selected_data_generator, @@ -2801,6 +2807,18 @@ class ProxyBaseLLMRequestProcessing: else: return chunk + @staticmethod + def _sse_chunk_serializer(restamp_model: str | None) -> StreamChunkSerializer: + if not restamp_model: + return ProxyBaseLLMRequestProcessing.return_sse_chunk + + def serialize(chunk: object) -> str: + return ProxyBaseLLMRequestProcessing.return_sse_chunk( + restamp_anthropic_stream_chunk_model(chunk, restamp_model) + ) + + return serialize + @staticmethod async def _finalize_streaming_generator_cleanup( request: Request | None, @@ -2990,6 +3008,7 @@ class ProxyBaseLLMRequestProcessing: request_data: dict, proxy_logging_obj: ProxyLogging, request: Request | None = None, + restamp_model: str | None = None, ) -> AsyncGenerator[str, None]: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events. @@ -2998,13 +3017,17 @@ class ProxyBaseLLMRequestProcessing: SSE serializers directly (rather than re-wrapping it in another ``async for: yield`` trampoline), so a streamed chunk traverses one fewer async-generator layer / coroutine resume on the hot path. + + ``restamp_model`` publishes that name on the Anthropic ``message_start`` + event in place of the provider's model, matching what the non-streaming + response reports. """ return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, proxy_logging_obj=proxy_logging_obj, - serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamp_model), serialize_error=lambda proxy_exc: ( f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" ), 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 new file mode 100644 index 00000000000..6e2c3f49445 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -0,0 +1,109 @@ +""" +Tests for restamping the public model on Anthropic Messages streaming chunks. +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + restamp_anthropic_stream_chunk_model, +) +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + +def _message_start_frame(model: str) -> bytes: + payload = { + "type": "message_start", + "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []}, + } + return f"event: message_start\ndata: {json.dumps(payload)}\n\n".encode() + + +def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: + async def _iterator_hook(**_kwargs): + for frame in frames: + yield frame + + proxy_logging_obj = MagicMock() + proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook + proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + return proxy_logging_obj + + +def _model_from_frame(frame: bytes | str) -> str: + text = frame.decode("utf-8") if isinstance(frame, bytes) else frame + data_line = next(line for line in text.split("\n") if line.startswith("data:")) + return json.loads(data_line[len("data:") :])["message"]["model"] + + +def test_restamps_sse_bytes_frame(): + restamped = restamp_anthropic_stream_chunk_model( + _message_start_frame("claude-haiku-4-5-20251001"), "claude-auto-1" + ) + + assert isinstance(restamped, bytes) + assert _model_from_frame(restamped) == "claude-auto-1" + assert b"event: message_start" in restamped + + +def test_restamps_event_dict(): + chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}} + + restamped = restamp_anthropic_stream_chunk_model(chunk, "claude-auto-2") + + assert restamped == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-2"}} + assert chunk["message"]["model"] == "claude-sonnet-4-6" + + +@pytest.mark.parametrize( + "chunk", + [ + b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n', + {"type": "content_block_delta", "delta": {"text": "hi"}}, + {"type": "message_start", "message": "not-a-dict"}, + b"event: message_start\ndata: not-json\n\n", + b"data: [DONE]\n\n", + ], +) +def test_leaves_chunks_without_a_model_untouched(chunk): + assert restamp_anthropic_stream_chunk_model(chunk, "claude-auto-1") == chunk + + +@pytest.mark.asyncio +async def test_sse_generator_publishes_requested_model_on_message_start(): + """The message_start event reports the requested model, not the provider's.""" + delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame]) + + 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", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta_frame + + +@pytest.mark.asyncio +async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")]) + + 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, + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001" 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 2/3] 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" From c21e895fe26b16cc97f0d4ed8e389f6ede9f2688 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:16:27 -0700 Subject: [PATCH 3/3] fix(proxy): handle CRLF and CR SSE frame terminators and flush held tail in anthropic stream restamper --- .../streaming_model_restamp.py | 49 ++++++--- litellm/proxy/common_request_processing.py | 21 ++-- .../test_streaming_model_restamp.py | 100 +++++++++++++++++- 3 files changed, 144 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py index e8d54f03949..7da5e5099fc 100644 --- a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -8,6 +8,7 @@ provider passthrough path) or as event dicts (fake-stream and agentic paths). """ import json +import re from collections.abc import Mapping from typing import Final @@ -16,7 +17,7 @@ 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" +_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n") _MAX_HELD_BYTES: Final = 65536 _PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"') @@ -46,15 +47,16 @@ def _restamped_data_line(line: str, requested_model: str) -> str | None: restamped: Final = _restamped_event(event, requested_model) if restamped is None: return None - return f"data: {json.dumps(restamped, separators=(',', ':'))}" + terminator: Final = line[len(line.rstrip("\r\n")) :] + return f"data: {json.dumps(restamped, separators=(',', ':'))}{terminator}" def _restamped_frame(frame: str, requested_model: str) -> str | None: - lines: Final = frame.split("\n") + lines: Final = frame.splitlines(keepends=True) restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines) if all(line is None for line in restamped): return None - return "\n".join(new if new is not None else old for new, old in zip(restamped, lines)) + return "".join(new if new is not None else old for new, old in zip(restamped, lines)) def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object: @@ -93,10 +95,12 @@ 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 + chunks or coalesced with later frames. Complete frames (``\\n\\n``, + ``\\r\\n\\r\\n``, or ``\\r\\r`` terminated) are emitted as their terminator + closes them and an incomplete tail is held until it completes, so the + restamp never misses a torn frame; ``flush`` returns whatever is still held + when the stream ends so no bytes are swallowed. Once ``message_start`` has + been handled, or the first real event proves the stream carries none, every later chunk passes through untouched. """ @@ -117,17 +121,27 @@ class AnthropicStreamModelRestamper: self._armed = False return restamped + def flush(self) -> bytes: + held: Final = self._held + self._held = b"" + self._armed = False + if not held: + return b"" + restamped: Final = restamp_anthropic_stream_chunk_model(held, self._requested_model) + return restamped if isinstance(restamped, bytes) else held + def _process_encoded(self, data: bytes) -> bytes: combined: Final = self._held + data - if _SSE_FRAME_END not in combined: + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(combined)) + if not boundaries: 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) + emitted: Final = self._restamped_closed_block(combined[: boundaries[-1]]) + tail: Final = combined[boundaries[-1] :] if not self._armed: self._held = b"" return emitted + tail @@ -135,7 +149,8 @@ class AnthropicStreamModelRestamper: return emitted def _restamped_closed_block(self, closed: bytes) -> bytes: - frames: Final = tuple(closed.split(_SSE_FRAME_END)[:-1]) + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(closed)) + frames: Final = tuple(closed[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries)) decider: Final = next( ( index @@ -147,13 +162,13 @@ class AnthropicStreamModelRestamper: 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: + if _MESSAGE_START_MARKER not in frames[decider]: return closed - restamped_text: Final = _restamped_frame(decider_frame.decode("utf-8", errors="ignore"), self._requested_model) + restamped_text: Final = _restamped_frame( + frames[decider].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) + restamped_text.encode("utf-8") if index == decider else frame for index, frame in enumerate(frames) ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 989cc7c18fb..eda7ebfa624 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3410,12 +3410,10 @@ class ProxyBaseLLMRequestProcessing: return chunk @staticmethod - def _sse_chunk_serializer(restamp_model: str | None) -> StreamChunkSerializer: - if not restamp_model: + def _sse_chunk_serializer(restamper: AnthropicStreamModelRestamper | None) -> StreamChunkSerializer: + if restamper is None: return ProxyBaseLLMRequestProcessing.return_sse_chunk - restamper: Final = AnthropicStreamModelRestamper(restamp_model) - def serialize(chunk: object) -> str: return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk)) @@ -3481,11 +3479,16 @@ class ProxyBaseLLMRequestProcessing: serialize_chunk: StreamChunkSerializer, serialize_error: StreamErrorSerializer, request: Request | None = None, + flush_tail: Callable[[], bytes] | None = None, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, cost injection, then yields chunks via serialize_chunk; on exception runs failure hook and yields via serialize_error. Use for SSE or NDJSON. + + ``flush_tail`` runs once after the upstream iterator completes cleanly and + its non-empty result is yielded, so a serializer that buffers bytes across + chunks can emit anything still held at end of stream. """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path @@ -3548,6 +3551,9 @@ class ProxyBaseLLMRequestProcessing: # so it must not suppress that refund. delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) + held_tail: Final = flush_tail() if flush_tail is not None else b"" + if held_tail: + yield serialize_chunk(held_tail) stream_completed = True except (asyncio.CancelledError, GeneratorExit): # Client disconnected mid-stream. CancelledError / GeneratorExit @@ -3558,8 +3564,7 @@ class ProxyBaseLLMRequestProcessing: # billing and release exactly once. This is the outermost generator # Starlette closes on disconnect, so the nested iterator hook (which # only sees GeneratorExit on GC) cannot own the refund. - if not stream_completed: - client_disconnected = True + client_disconnected = not stream_completed if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, @@ -3627,16 +3632,18 @@ class ProxyBaseLLMRequestProcessing: event in place of the provider's model, matching what the non-streaming response reports. """ + restamper: Final = AnthropicStreamModelRestamper(restamp_model) if restamp_model else None return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, proxy_logging_obj=proxy_logging_obj, - serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamp_model), + serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper), serialize_error=lambda proxy_exc: ( f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" ), request=request, + flush_tail=None if restamper is None else restamper.flush, ) @overload 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 385173b24a9..b7bc670c7f8 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 @@ -14,12 +14,12 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -def _message_start_frame(model: str) -> bytes: +def _message_start_frame(model: str, line_end: str = "\n") -> bytes: payload = { "type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []}, } - return f"event: message_start\ndata: {json.dumps(payload)}\n\n".encode() + return f"event: message_start{line_end}data: {json.dumps(payload)}{line_end}{line_end}".encode() def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: @@ -190,3 +190,99 @@ async def test_sse_generator_restamps_message_start_split_across_chunks(): 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" + + +def test_restamps_crlf_terminated_message_start_frame(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + emitted = restamper.process(frame) + + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + assert emitted.endswith(b"\r\n\r\n") + assert restamper.process(delta) == delta + + +def test_restamps_cr_terminated_message_start_frame(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + emitted = restamper.process(frame) + + assert isinstance(emitted, bytes) + assert b'"model":"claude-auto-1"' in emitted + assert emitted.endswith(b"\r\r") + + +def test_restamps_crlf_message_start_split_across_transport_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + 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_flush_returns_restamped_held_tail(): + unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(unterminated) == b"" + flushed = restamper.flush() + + assert b'"model":"claude-auto-1"' in flushed + assert restamper.flush() == b"" + + +def test_flush_disarms_the_restamper(): + restamper = AnthropicStreamModelRestamper("claude-auto-1") + frame = _message_start_frame("claude-haiku-4-5-20251001") + + assert restamper.flush() == b"" + assert restamper.process(frame) == frame + + +@pytest.mark.asyncio +async def test_sse_generator_flushes_held_tail_at_end_of_stream(): + unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] + proxy_logging_obj = _proxy_logging_obj_streaming([unterminated]) + + 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 b'"model":"claude-auto-1"' in joined + + +@pytest.mark.asyncio +async def test_sse_generator_restamps_crlf_stream(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' + proxy_logging_obj = _proxy_logging_obj_streaming([frame, delta]) + + 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", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta