fix(anthropic): 限制流式别名重写缓冲

- 将 Anthropic SSE 重写逻辑下沉到提供商模块。

- 首帧缓冲超过上限时切换为原样透传。

- 重写 message_start 后停止缓冲并补充边界测试。

litellm/llms/anthropic/experimental_pass_through/messages/model_alias_rewriter.py

@@ -0,0 +1,106 @@
+"""Bounded Anthropic SSE rewriting for the public model alias."""
+
+import json
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Final
+
+from pydantic import TypeAdapter, ValidationError
+
+from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
+
+MAX_MESSAGE_START_BUFFER_BYTES: Final = 1024 * 1024
+_SSE_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object])
+
+
+def _restamp_message_start_data_line(line: bytes, requested_model: str) -> tuple[bytes, bool]:
+    """Replace the model in one Anthropic message_start data line."""
+    body: Final = line.rstrip(b"\r\n")
+    line_ending: Final = line[len(body) :]
+    if not body.startswith(b"data:"):
+        return line, False
+    try:
+        payload: Final = _SSE_OBJECT_ADAPTER.validate_json(body[len(b"data:") :].strip())
+    except ValidationError:
+        return line, False
+    if payload.get("type") != "message_start":
+        return line, False
+    try:
+        message: Final = _SSE_OBJECT_ADAPTER.validate_python(payload.get("message"))
+    except ValidationError:
+        return line, False
+    if "model" not in message:
+        return line, False
+    restamped_payload: Final = {**payload, "message": {**message, "model": requested_model}}
+    rewritten: Final = b"data: " + json.dumps(restamped_payload, separators=(",", ":")).encode() + line_ending
+    return rewritten, True
+
+
+def _restamp_message_start_frames(frames: bytes, requested_model: str) -> tuple[bytes, bool]:
+    """Restore the alias in complete frames and report when message_start was found."""
+    output: list[bytes] = []
+    restamped = False
+    for line in frames.splitlines(keepends=True):
+        rewritten_line, line_restamped = _restamp_message_start_data_line(line, requested_model)
+        output.append(rewritten_line)
+        restamped = restamped or line_restamped
+    return b"".join(output), restamped
+
+
+class AnthropicMessageStartModelRewriter:
+    """Reassemble only the first Anthropic frames with a bounded pending buffer."""
+
+    __slots__ = ("_finished", "_pending", "_requested_model")
+
+    def __init__(self, requested_model: str) -> None:
+        self._requested_model: Final = requested_model
+        self._pending = b""
+        self._finished = False
+
+    def feed(self, chunk: bytes) -> bytes:
+        """Return rewritten frames, or bounded pass-through for a malformed stream."""
+        if self._finished:
+            return chunk
+
+        buffered: Final = self._pending + chunk
+        if len(buffered) > MAX_MESSAGE_START_BUFFER_BYTES:
+            self._pending = b""
+            self._finished = True
+            return buffered
+
+        complete_frames, self._pending = split_complete_sse_frames(buffered)
+        rewritten_frames, restamped = _restamp_message_start_frames(
+            complete_frames,
+            self._requested_model,
+        )
+        if not restamped:
+            return rewritten_frames
+
+        tail: Final = self._pending
+        self._pending = b""
+        self._finished = True
+        return rewritten_frames + tail
+
+    def flush(self) -> bytes:
+        """Return the final unterminated tail without dropping provider bytes."""
+        tail: Final = self._pending
+        self._pending = b""
+        return tail
+
+
+async def restamp_anthropic_streaming_response_model(
+    response: AsyncIterator[object],
+    requested_model: str,
+) -> AsyncGenerator[object, None]:
+    """Restamp message_start while preserving arbitrary transport chunk boundaries."""
+    rewriter: Final = AnthropicMessageStartModelRewriter(requested_model)
+    async for chunk in response:
+        if isinstance(chunk, (bytes, bytearray)):
+            if rewritten_frames := rewriter.feed(bytes(chunk)):
+                yield rewritten_frames
+            continue
+        if buffered_tail := rewriter.flush():
+            yield buffered_tail
+        yield chunk
+    final_tail: Final = rewriter.flush()
+    if final_tail:
+        yield final_tail

litellm/proxy/common_request_processing.py

@@ -4,7 +4,7 @@ import json
 import logging
 import math
 import traceback
-from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
+from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
@@ -15,7 +15,6 @@ import httpx
 import orjson
 from fastapi import HTTPException, Request, status
 from fastapi.responses import JSONResponse, Response, StreamingResponse
-from pydantic import TypeAdapter, ValidationError
 from starlette.types import Receive, Scope, Send
 
 import litellm
@@ -49,6 +48,9 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.streaming_handler import (
     backfill_missing_cache_usage_fields,
 )
+from litellm.llms.anthropic.experimental_pass_through.messages.model_alias_rewriter import (
+    restamp_anthropic_streaming_response_model,
+)
 from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
 from litellm.proxy.auth.auth_utils import check_response_size_is_safe
@@ -60,7 +62,6 @@ from litellm.proxy.common_utils.sse_keepalive import (
     SSE_COMMENT_PING_BYTES,
     coerce_keepalive_interval,
     resolve_ttft_keepalive_interval,
-    split_complete_sse_frames,
     wrap_sse_stream_with_keepalive_pings,
 )
 from litellm.proxy.dd_span_tagger import DDSpanTagger
@@ -1268,78 +1269,6 @@ def _override_openai_response_model(
         )
 
 
-_ANTHROPIC_SSE_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object])
-
-
-def _restamp_anthropic_message_start_data_line(line: bytes, requested_model: str) -> bytes:
-    """Replace the model in one Anthropic message_start data line."""
-    body: Final = line.rstrip(b"\r\n")
-    line_ending: Final = line[len(body) :]
-    if not body.startswith(b"data:"):
-        return line
-    try:
-        payload: Final = _ANTHROPIC_SSE_OBJECT_ADAPTER.validate_json(body[len(b"data:") :].strip())
-    except ValidationError:
-        return line
-    if payload.get("type") != "message_start":
-        return line
-    try:
-        message: Final = _ANTHROPIC_SSE_OBJECT_ADAPTER.validate_python(payload.get("message"))
-    except ValidationError:
-        return line
-    if "model" not in message:
-        return line
-    restamped_payload: Final = {**payload, "message": {**message, "model": requested_model}}
-    return b"data: " + json.dumps(restamped_payload, separators=(",", ":")).encode() + line_ending
-
-
-def _restamp_anthropic_message_start_frames(frames: bytes, requested_model: str) -> bytes:
-    """Restore the client-requested model in complete Anthropic SSE frames."""
-    return b"".join(
-        _restamp_anthropic_message_start_data_line(line, requested_model) for line in frames.splitlines(keepends=True)
-    )
-
-
-class _AnthropicMessageStartModelRewriter:
-    """Reassemble transport chunks before rewriting Anthropic SSE frames."""
-
-    __slots__ = ("_pending", "_requested_model")
-
-    def __init__(self, requested_model: str) -> None:
-        self._requested_model: Final = requested_model
-        self._pending = b""
-
-    def feed(self, chunk: bytes) -> bytes:
-        """Buffer a transport chunk and return rewritten complete frames."""
-        complete_frames, self._pending = split_complete_sse_frames(self._pending + chunk)
-        return _restamp_anthropic_message_start_frames(complete_frames, self._requested_model)
-
-    def flush(self) -> bytes:
-        """Return the final unterminated tail without dropping provider bytes."""
-        tail: Final = self._pending
-        self._pending = b""
-        return tail
-
-
-async def _restamp_anthropic_streaming_response_model(
-    response: AsyncIterator[object],
-    requested_model: str,
-) -> AsyncGenerator[object, None]:
-    """Restamp message_start while preserving arbitrary transport chunk boundaries."""
-    rewriter: Final = _AnthropicMessageStartModelRewriter(requested_model)
-    async for chunk in response:
-        if isinstance(chunk, (bytes, bytearray)):
-            if rewritten_frames := rewriter.feed(bytes(chunk)):
-                yield rewritten_frames
-            continue
-        if buffered_tail := rewriter.flush():
-            yield buffered_tail
-        yield chunk
-    final_tail: Final = rewriter.flush()
-    if final_tail:
-        yield final_tail
-
-
 def _anthropic_client_requested_model(request_data: Mapping[str, object]) -> str | None:
     """Return the public model name only for native Anthropic Messages streams."""
     logging_obj: Final = request_data.get("litellm_logging_obj")
@@ -3493,7 +3422,7 @@ class ProxyBaseLLMRequestProcessing:
         delivered_chunk = False
         anthropic_requested_model: Final = _anthropic_client_requested_model(request_data)
         stream_response: Final = (
-            _restamp_anthropic_streaming_response_model(response, anthropic_requested_model)
+            restamp_anthropic_streaming_response_model(response, anthropic_requested_model)
             if anthropic_requested_model is not None
             else response
         )

tests/test_litellm/proxy/test_response_model_sanitization.py

@@ -1,6 +1,6 @@
 import asyncio
 import json
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, AsyncIterator
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -338,22 +338,22 @@ async def test_proxy_streaming_chunks_use_client_requested_model_before_alias_ma
 
 @pytest.mark.asyncio
 async def test_anthropic_streaming_message_start_uses_client_requested_model_across_chunks(
-    monkeypatch,
+    monkeypatch: pytest.MonkeyPatch,
 ):
     """Keep the public model alias when an Anthropic message_start spans transport chunks."""
     from litellm.proxy import proxy_server
     from litellm.proxy._types import UserAPIKeyAuth
     from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 
-    async def upstream_stream():
+    async def upstream_stream() -> AsyncGenerator[bytes, None]:
         yield b'event: message_start\ndata: {"type":"message_start","message":{"model":"backend'
         yield b'-model","id":"msg_1"}}\n\nevent: message_stop\ndata: {"type":"message_stop"}\n\n'
 
     async def passthrough_iterator_hook(
         user_api_key_dict: UserAPIKeyAuth,
-        response: AsyncGenerator,
-        request_data: dict,
-    ):
+        response: AsyncIterator[object],
+        request_data: dict[str, object],
+    ) -> AsyncGenerator[object, None]:
         async for chunk in response:
             yield chunk
 
@@ -385,6 +385,37 @@ async def test_anthropic_streaming_message_start_uses_client_requested_model_acr
     assert b'event: message_stop\ndata: {"type":"message_stop"}\n\n' in streamed
 
 
+def test_anthropic_message_start_rewriter_bounds_unterminated_streams():
+    """Fall back to byte-for-byte pass-through when the first SSE frame is oversized."""
+    from litellm.llms.anthropic.experimental_pass_through.messages.model_alias_rewriter import (
+        MAX_MESSAGE_START_BUFFER_BYTES,
+        AnthropicMessageStartModelRewriter,
+    )
+
+    rewriter = AnthropicMessageStartModelRewriter("my-alias")
+    pending = b"x" * MAX_MESSAGE_START_BUFFER_BYTES
+
+    assert rewriter.feed(pending) == b""
+    assert rewriter.feed(b"y") == pending + b"y"
+    assert rewriter.feed(b"z") == b"z"
+
+
+def test_anthropic_message_start_rewriter_stops_buffering_after_rewrite():
+    """Pass later transport chunks through directly after the alias has been restored."""
+    from litellm.llms.anthropic.experimental_pass_through.messages.model_alias_rewriter import (
+        MAX_MESSAGE_START_BUFFER_BYTES,
+        AnthropicMessageStartModelRewriter,
+    )
+
+    rewriter = AnthropicMessageStartModelRewriter("my-alias")
+    message_start = b'event: message_start\ndata: {"type":"message_start","message":{"model":"backend-model"}}\n\n'
+    rewritten = rewriter.feed(message_start)
+    later_chunk = b"z" * (MAX_MESSAGE_START_BUFFER_BYTES + 1)
+
+    assert b'"model":"my-alias"' in rewritten
+    assert rewriter.feed(later_chunk) == later_chunk
+
+
 @pytest.mark.asyncio
 async def test_proxy_streaming_azure_model_router_preserves_actual_model(monkeypatch):
     """