fix(proxy): report requested model on Anthropic streaming message_start

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

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"
             ),

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"