feat(rust-bridge): extend native messages streaming beyond Anthropic

litellm-rust/crates/core/src/messages/handler.rs

@@ -1,10 +1,10 @@
-use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
 use crate::error::Error;
 use crate::http_utils::http_request;
 
 use super::client::http_client;
 use super::common_utils::truncate_error_body;
 use super::prepare::prepare_provider_request;
+use super::transformation::AnthropicMessagesProviderConfig;
 use super::types::{AnthropicMessagesResponse, MessagesRequest};
 
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
@@ -42,15 +42,25 @@ pub(super) async fn execute_messages_provider_call(
     request.config.transform_response(&request.model, response)
 }
 
+/// The streaming-capability gate. A provider whose config does not opt into
+/// streaming declines ([`Error::Unsupported`]) before any request goes out, so
+/// hosts treat it as "fall back", not "fail".
+pub(super) fn ensure_streaming_supported(
+    config: &dyn AnthropicMessagesProviderConfig,
+) -> Result<(), Error> {
+    if config.supports_streaming() {
+        return Ok(());
+    }
+    Err(Error::Unsupported(
+        "streaming messages is not supported for this provider",
+    ))
+}
+
 pub(super) async fn execute_messages_provider_stream(
     request: MessagesRequest<'_>,
 ) -> Result<reqwest::Response, Error> {
     let mut request = prepare_provider_request(request)?;
-    if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
-        return Err(Error::InvalidRequest(
-            "streaming messages is not supported for this provider".to_string(),
-        ));
-    }
+    ensure_streaming_supported(request.config)?;
     // The streaming entrypoints only make sense for `stream: true`; force it so
     // a host cannot accidentally ask for SSE and receive a buffered JSON body.
     if let Some(body) = request.body.as_object_mut() {

litellm-rust/crates/core/src/messages/prepare.rs

@@ -46,7 +46,6 @@ pub(super) fn prepare_provider_request(
     let url = config.complete_url(request.api_base, &model, &env_lookup)?;
 
     Ok(ProviderMessagesRequest {
-        provider: provider.to_string(),
         model,
         config,
         url,

litellm-rust/crates/core/src/messages/tests.rs

@@ -9,8 +9,12 @@ use crate::error::Error;
 use super::common_utils::{
     has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body,
 };
+use super::handler::ensure_streaming_supported;
 use super::messages;
+use super::transformation::AnthropicMessagesProviderConfig;
 use super::types::MessagesRequest;
+use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
+use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
 
 async fn read_http_request(socket: &mut TcpStream) -> String {
     let mut request = Vec::new();
@@ -60,6 +64,37 @@ fn provider_config_resolves_anthropic_and_azure_ai() {
     assert!(messages_provider_config("openai").is_none());
 }
 
+#[test]
+fn streaming_gate_accepts_opted_in_providers_and_declines_the_rest() {
+    struct NonStreamingConfig;
+
+    impl AnthropicMessagesProviderConfig for NonStreamingConfig {
+        fn complete_url(
+            &self,
+            _api_base: Option<&str>,
+            _model: &str,
+            _env_lookup: &dyn Fn(&str) -> Option<String>,
+        ) -> Result<String, crate::error::Error> {
+            unreachable!("the streaming gate must not touch the config's URLs")
+        }
+
+        fn resolve_api_key(
+            &self,
+            _api_key: Option<&str>,
+            _env_lookup: &dyn Fn(&str) -> Option<String>,
+        ) -> Result<String, crate::error::Error> {
+            unreachable!("the streaming gate must not resolve credentials")
+        }
+    }
+
+    assert!(ensure_streaming_supported(&ANTHROPIC_MESSAGES_CONFIG).is_ok());
+    assert!(ensure_streaming_supported(&AZURE_ANTHROPIC_MESSAGES_CONFIG).is_ok());
+    assert!(matches!(
+        ensure_streaming_supported(&NonStreamingConfig).expect_err("default config declines"),
+        Error::Unsupported(_)
+    ));
+}
+
 #[test]
 fn truncate_error_body_caps_long_payloads() {
     let body = "x".repeat(400);
@@ -510,6 +545,97 @@ async fn messages_stream_frames_yields_complete_frames_and_forces_stream_flag()
     assert_eq!(sent_body["stream"], Value::Bool(true));
 }
 
+#[tokio::test]
+async fn messages_stream_frames_streams_azure_ai_through_the_anthropic_endpoint() {
+    use futures_util::StreamExt;
+
+    use super::messages_stream_frames;
+
+    let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
+    let addr = listener.local_addr().expect("addr");
+
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.expect("accepts request");
+        let request = read_http_request(&mut socket).await;
+        let body = concat!(
+            "event: message_start\ndata: {\"type\":\"message_start\"}\n\n",
+            "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
+        );
+        let response = format!(
+            "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n",
+            body.len()
+        );
+        socket
+            .write_all(response.as_bytes())
+            .await
+            .expect("writes response head");
+        socket
+            .write_all(body.as_bytes())
+            .await
+            .expect("writes body");
+        request
+    });
+
+    let mut stream = messages_stream_frames(MessagesRequest {
+        model: "claude-sonnet-4-5",
+        body: json!({
+            "model": "claude-sonnet-4-5",
+            "max_tokens": 1024,
+            "messages": [{"role": "user", "content": "hi"}]
+        }),
+        api_key: Some("sk-azure"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("azure_ai"),
+        extra_headers: None,
+        timeout: Some(Duration::from_secs(5)),
+    })
+    .await
+    .expect("azure_ai stream request succeeds");
+
+    let mut frames = Vec::new();
+    while let Some(frame) = stream.next().await {
+        frames.push(String::from_utf8(frame.expect("frame decodes")).expect("frame is utf8"));
+    }
+    assert_eq!(
+        frames,
+        vec![
+            "event: message_start\ndata: {\"type\":\"message_start\"}\n\n",
+            "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
+        ]
+    );
+
+    let request = server.await.expect("server task completes");
+    let (head, body) = request.split_once("\r\n\r\n").expect("has body");
+    assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}");
+    let head_lower = head.to_ascii_lowercase();
+    assert!(head_lower.contains("x-api-key: sk-azure"), "{head}");
+    let sent_body: Value = serde_json::from_str(body).expect("body is json");
+    assert_eq!(sent_body["stream"], Value::Bool(true));
+}
+
+#[tokio::test]
+async fn messages_stream_declines_unsupported_provider_before_the_call() {
+    use super::messages_stream_frames;
+
+    let err = match messages_stream_frames(MessagesRequest {
+        model: "gpt-4o",
+        body: json!({"model": "gpt-4o", "max_tokens": 8, "messages": []}),
+        api_key: Some("sk"),
+        // Nothing listens here: a decline must return before any connection.
+        api_base: Some("http://127.0.0.1:1"),
+        custom_llm_provider: Some("openai"),
+        extra_headers: None,
+        timeout: Some(Duration::from_millis(50)),
+    })
+    .await
+    {
+        Err(err) => err,
+        Ok(_) => panic!("unsupported provider must decline before the stream"),
+    };
+
+    assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai"));
+}
+
 #[tokio::test]
 async fn messages_stream_frames_maps_upstream_error_status() {
     let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");

litellm-rust/crates/core/src/messages/transformation.rs

@@ -45,6 +45,13 @@ pub trait AnthropicMessagesProviderConfig: Sync {
         ]
     }
 
+    /// Whether this provider serves the SSE streaming variant of the route.
+    /// Opt-in: a config that does not override this declines streaming so the
+    /// host falls back to its own implementation instead of failing.
+    fn supports_streaming(&self) -> bool {
+        false
+    }
+
     #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_request(
         &self,

litellm-rust/crates/core/src/messages/types.rs

@@ -16,7 +16,6 @@ pub struct MessagesRequest<'a> {
 }
 
 pub(super) struct ProviderMessagesRequest {
-    pub(super) provider: String,
     pub(super) model: String,
     pub(super) config: &'static dyn AnthropicMessagesProviderConfig,
     pub(super) url: String,

litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs

@@ -68,6 +68,10 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
     fn auth_strategy(&self) -> MessagesAuthStrategy {
         MessagesAuthStrategy::Header("x-api-key")
     }
+
+    fn supports_streaming(&self) -> bool {
+        true
+    }
 }
 
 #[cfg(test)]

litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs

@@ -164,6 +164,10 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
         self.anthropic.auth_strategy()
     }
 
+    fn supports_streaming(&self) -> bool {
+        self.anthropic.supports_streaming()
+    }
+
     fn accepts_bearer_auth(&self) -> bool {
         true
     }

litellm-rust/crates/python-bridge/src/errors.rs

@@ -33,7 +33,7 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr {
 /// Everything raised before the request goes out is safe for the host to retry
 /// on its own path; anything after it is not, because the provider has already
 /// done the work and billed for it.
-pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
+pub(crate) fn fallback_error_to_pyerr(err: Error) -> PyErr {
     match err {
         Error::Unsupported(_)
         | Error::Auth(_)

litellm-rust/crates/python-bridge/src/routes/chat_completions.rs

@@ -8,7 +8,7 @@ use litellm_core::chat_completions::{
 use pyo3::prelude::*;
 use serde_json::Value;
 
-use crate::errors::chat_completions_error_to_pyerr;
+use crate::errors::fallback_error_to_pyerr;
 use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value};
 
 fn prepare_chat_completions(
@@ -86,6 +86,6 @@ bridge_route! {
         timeout_seconds: Option<f64>,
     },
     prepare = prepare_chat_completions,
-    errors = chat_completions_error_to_pyerr,
+    errors = fallback_error_to_pyerr,
     extra = [chat_completions_decline],
 }

litellm-rust/crates/python-bridge/src/routes/messages_stream.rs

@@ -17,7 +17,7 @@ use pyo3::types::{PyAny, PyBytes};
 use serde_json::Value;
 use tokio::sync::Mutex;
 
-use crate::errors::core_error_to_pyerr;
+use crate::errors::{core_error_to_pyerr, fallback_error_to_pyerr};
 use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value};
 
 /// One complete SSE frame, converted to `bytes` on the attached thread that
@@ -115,6 +115,10 @@ fn amessages_stream<'py>(
             extra_headers,
             timeout,
         } = options;
+        // Only the request start maps through the fallback contract: declines
+        // (unsupported provider, bad request) happen before the provider is
+        // called, so the host may retry on its own path. Failures surfaced by
+        // `__anext__` are mid-stream and keep `core_error_to_pyerr`.
         let frames = messages_stream_frames(MessagesRequest {
             model: &model,
             body,
@@ -125,7 +129,7 @@ fn amessages_stream<'py>(
             timeout,
         })
         .await
-        .map_err(core_error_to_pyerr)?;
+        .map_err(fallback_error_to_pyerr)?;
         Ok(MessagesStream {
             frames: Arc::new(Mutex::new(Some(frames))),
         })
@@ -187,12 +191,12 @@ mod tests {
     }
 
     /// An SSE upstream that sends `body` in chunks separated by short delays
-    /// so frames span network chunk boundaries.
+    /// so frames span network chunk boundaries. Returns the captured request.
     async fn streaming_upstream(
         listener: TcpListener,
         status_line: &'static str,
         body: &'static str,
-    ) {
+    ) -> String {
         let (mut socket, _) = listener.accept().await.expect("accepts request");
         let request = read_http_request(&mut socket).await;
         assert!(
@@ -217,6 +221,7 @@ mod tests {
             tokio::time::sleep(Duration::from_millis(10)).await;
             socket.write_all(chunk).await.expect("writes body chunk");
         }
+        request
     }
 
     fn register_stream_module(py: Python<'_>) -> Bound<'_, PyModule> {
@@ -233,6 +238,18 @@ mod tests {
         locals
             .set_item("url", url)
             .expect("URL should enter Python locals");
+        locals
+            .set_item(
+                "Declined",
+                py.get_type::<crate::errors::RustBridgeDeclined>(),
+            )
+            .expect("declined exception should enter Python locals");
+        locals
+            .set_item(
+                "UpstreamError",
+                py.get_type::<crate::errors::RustUpstreamError>(),
+            )
+            .expect("upstream exception should enter Python locals");
         let source = CString::new(code).expect("Python source should not contain null bytes");
         py.run(&source, Some(&locals), Some(&locals))
             .expect("Python stream exercise should pass");
@@ -324,8 +341,9 @@ async def exercise():
             api_base=url,
             custom_llm_provider="anthropic",
         )
-    except RuntimeError as error:
-        assert "429" in str(error), str(error)
+    except UpstreamError as error:
+        assert not isinstance(error, Declined), "a provider 429 is not a decline"
+        assert error.args[0] == 429, error.args
     else:
         raise AssertionError("upstream error should surface from the awaitable")
 
@@ -340,6 +358,100 @@ asyncio.run(exercise())
             .expect("server task should not panic");
     }
 
+    #[test]
+    fn stream_route_streams_azure_ai_through_the_anthropic_endpoint() {
+        Python::initialize();
+        let runtime = pyo3_async_runtimes::tokio::get_runtime();
+        let listener = runtime
+            .block_on(TcpListener::bind("127.0.0.1:0"))
+            .expect("listener should bind");
+        let address = listener
+            .local_addr()
+            .expect("listener should have an address");
+        let body = concat!(
+            "event: message_start\ndata: {\"type\":\"message_start\"}\n\n",
+            "event: content_block_delta\ndata: {\"delta\":\"hi\"}\n\n",
+            "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
+        );
+        let server = runtime.spawn(streaming_upstream(listener, "HTTP/1.1 200 OK", body));
+
+        Python::attach(|py| {
+            let module = register_stream_module(py);
+            run_async_code(
+                py,
+                &module,
+                &format!("http://{address}"),
+                r#"
+import asyncio
+
+async def exercise():
+    stream = await runtime.amessages_stream(
+        model="claude-sonnet-4-5",
+        body={"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []},
+        api_key="sk-azure",
+        api_base=url,
+        custom_llm_provider="azure_ai",
+    )
+    frames = [chunk async for chunk in stream]
+    assert frames == [
+        b'event: message_start\ndata: {"type":"message_start"}\n\n',
+        b'event: content_block_delta\ndata: {"delta":"hi"}\n\n',
+        b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
+    ], frames
+
+asyncio.run(exercise())
+"#,
+            );
+        });
+
+        let request = runtime
+            .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await })
+            .expect("server should finish")
+            .expect("server task should not panic");
+        let head = request
+            .split_once("\r\n\r\n")
+            .expect("request should have a body")
+            .0
+            .to_ascii_lowercase();
+        assert!(
+            head.starts_with("post /anthropic/v1/messages "),
+            "azure_ai must stream through the anthropic endpoint: {head}"
+        );
+        assert!(head.contains("x-api-key: sk-azure"), "{head}");
+    }
+
+    #[test]
+    fn stream_route_declines_unsupported_provider_before_any_request() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = register_stream_module(py);
+            run_async_code(
+                py,
+                &module,
+                "http://127.0.0.1:1",
+                r#"
+import asyncio
+
+async def exercise():
+    try:
+        await runtime.amessages_stream(
+            model="gpt-4o",
+            body={"model": "gpt-4o", "max_tokens": 8, "messages": []},
+            api_key="sk",
+            api_base=url,
+            custom_llm_provider="openai",
+        )
+    except Declined as error:
+        assert not isinstance(error, UpstreamError), str(error)
+    else:
+        raise AssertionError("unsupported provider must decline from the awaitable")
+
+asyncio.run(exercise())
+"#,
+            );
+        });
+    }
+
     #[test]
     fn stream_route_ends_iteration_after_a_mid_stream_failure() {
         Python::initialize();

tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py

@@ -479,6 +479,29 @@ async def test_stream_gate_returns_sse_frames_and_keeps_stream_flag():
     assert call["timeout_seconds"] == 30.0
 
 
+@pytest.mark.asyncio
+async def test_stream_gate_routes_azure_ai_to_the_native_stream():
+    bridge = RecordingAsyncMessagesStream()
+    rust_messages_stream.set_rust_amessages_stream(amessages_stream=bridge)
+
+    response = await _stream_gate(
+        custom_llm_provider="azure_ai",
+        litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=True),
+        api_key="sk-azure",
+        api_base="https://resource.services.ai.azure.com/anthropic",
+        headers={"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"},
+    )
+
+    assert response is not None
+    assert response._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
+    chunks = [chunk async for chunk in response]
+    assert chunks == list(FAKE_SSE_FRAMES)
+    call = bridge.calls[0]
+    assert call["custom_llm_provider"] == "azure_ai"
+    assert call["api_key"] == "sk-azure"
+    assert call["body"]["stream"] is True
+
+
 @pytest.mark.asyncio
 async def test_stream_gate_falls_back_to_python_when_bridge_raises():
     bridge = RaisingAsyncMessagesStream()