feat(messages): route eligible Azure Anthropic streaming through Rust via buffered fake-stream

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm/llms/custom_httpx/llm_http_handler.py

@@ -153,6 +153,9 @@ if TYPE_CHECKING:
     from aiohttp import ClientSession
 
     from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
+    from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+        AnthropicMessagesStreamingResponse,
+    )
     from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
     from litellm.types.llms.openai_evals import (
         CancelEvalResponse,
@@ -2095,6 +2098,7 @@ class BaseLLMHTTPHandler:
             custom_llm_provider=custom_llm_provider,
             litellm_params=litellm_params,
             stream=stream or False,
+            rust_stream_eligible=bool(stream) and not self._has_agentic_completion_hook(logging_obj),
             model=model,
             api_key=api_key,
             api_base=api_base,
@@ -2107,6 +2111,8 @@ class BaseLLMHTTPHandler:
             ),
         )
         if rust_messages_response is not None:
+            if stream:
+                return self._rust_anthropic_messages_fake_stream(rust_messages_response)
             return await self._finalize_anthropic_messages_response(
                 initial_response=rust_messages_response,
                 model=model,
@@ -2247,6 +2253,7 @@ class BaseLLMHTTPHandler:
         custom_llm_provider: str,
         litellm_params: GenericLiteLLMParams,
         stream: bool,
+        rust_stream_eligible: bool,
         model: str,
         api_key: str | None,
         api_base: str | None,
@@ -2254,14 +2261,17 @@ class BaseLLMHTTPHandler:
         request_body: dict,
         timeout: float | httpx.Timeout | None,
     ) -> AnthropicMessagesResponse | None:
-        if stream or custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True:
+        if custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True:
+            return None
+        if stream and not rust_stream_eligible:
             return None
 
         from litellm.rust_bridge import messages as rust_messages_bridge
 
+        upstream_body = {key: value for key, value in request_body.items() if key != "stream"}
         rust_response = await rust_messages_bridge.amessages(
             model=model,
-            body=request_body,
+            body=upstream_body,
             api_key=api_key,
             api_base=api_base,
             custom_llm_provider=custom_llm_provider,
@@ -2275,6 +2285,25 @@ class BaseLLMHTTPHandler:
         response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}}
         return response_obj
 
+    @staticmethod
+    def _rust_anthropic_messages_fake_stream(
+        rust_response: AnthropicMessagesResponse,
+    ) -> "AnthropicMessagesStreamingResponse":
+        from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
+            FakeAnthropicMessagesStreamIterator,
+        )
+        from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+            AnthropicMessagesStreamHiddenParams,
+            AnthropicMessagesStreamingResponse,
+        )
+
+        completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response))
+        hidden_params = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"})
+        return AnthropicMessagesStreamingResponse(
+            completion_stream=completion_stream,
+            hidden_params=hidden_params,
+        )
+
     def anthropic_messages_handler(
         self,
         model: str,

tests/test_litellm/rust_bridge/test_messages.py

@@ -1,12 +1,16 @@
 """Tests for the optional Rust-backed Anthropic Messages path."""
 
 import importlib
+from typing import cast
 
 import httpx
 import pytest
 
 import litellm
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
+from litellm.types.llms.anthropic_messages.anthropic_response import (
+    AnthropicMessagesResponse,
+)
 from litellm.types.router import GenericLiteLLMParams
 
 rust_messages = importlib.import_module("litellm.rust_bridge.messages")
@@ -207,6 +211,7 @@ def _gate(**overrides):
         "custom_llm_provider": "azure_ai",
         "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True),
         "stream": False,
+        "rust_stream_eligible": False,
         "model": "claude-sonnet-4-5",
         "api_key": "sk-azure",
         "api_base": "https://resource.services.ai.azure.com/anthropic",
@@ -271,16 +276,50 @@ async def test_gate_skips_rust_for_non_azure_provider():
 
 
 @pytest.mark.asyncio
-async def test_gate_skips_rust_when_streaming():
+async def test_gate_skips_rust_when_streaming_but_not_eligible():
     bridge = ExplodingAsyncMessages()
     litellm.use_litellm_rust(True, amessages=bridge)
 
-    response = await _gate(stream=True)
+    response = await _gate(stream=True, rust_stream_eligible=False)
 
     assert response is None
     assert bridge.calls == 0
 
 
+@pytest.mark.asyncio
+async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag():
+    bridge = RecordingAsyncMessages()
+    litellm.use_litellm_rust(True, amessages=bridge)
+
+    streaming_body = {**REQUEST_BODY, "stream": True}
+    response = await _gate(
+        stream=True,
+        rust_stream_eligible=True,
+        request_body=streaming_body,
+    )
+
+    assert response is not None
+    assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
+    assert "stream" not in bridge.calls[0]["body"]
+    assert bridge.calls[0]["body"] == REQUEST_BODY
+
+
+@pytest.mark.asyncio
+async def test_fake_stream_wraps_rust_response_as_anthropic_sse():
+    response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE))
+    stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response)
+
+    assert stream._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
+
+    chunks = [chunk async for chunk in stream]
+    joined = b"".join(chunks)
+
+    assert b"event: message_start" in joined
+    assert b"event: content_block_delta" in joined
+    assert b"hello world" in joined
+    assert b"event: message_stop" in joined
+
+
 @pytest.mark.asyncio
 async def test_gate_falls_back_when_bridge_unavailable(monkeypatch):
     monkeypatch.setattr(