feat(rust): stream Anthropic Messages SSE on POST /v1/messages

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

litellm-rust/crates/ai-gateway/src/messages/handler.rs

@@ -5,6 +5,7 @@ use serde_json::Value;
 use super::client::http_client;
 use super::common_utils::truncate_error_body;
 use super::types::ProviderMessagesRequest;
+use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
 
 pub(super) async fn execute_messages_provider_call(
     request: ProviderMessagesRequest,
@@ -45,3 +46,38 @@ pub(super) async fn execute_messages_provider_call(
         CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
     })
 }
+
+pub(super) async fn execute_messages_provider_stream(
+    request: ProviderMessagesRequest,
+) -> CoreResult<reqwest::Response> {
+    if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
+        return Err(CoreError::InvalidRequest(
+            "streaming messages is not supported for this provider".to_string(),
+        ));
+    }
+
+    let mut request_builder = http_client().post(&request.url).json(&request.body);
+    for (key, value) in &request.upstream_headers {
+        request_builder = request_builder.header(key, value);
+    }
+    if let Some(duration) = request.timeout {
+        request_builder = request_builder.timeout(duration);
+    }
+
+    let response = request_builder
+        .send()
+        .await
+        .map_err(|err| CoreError::Network(err.to_string()))?;
+    let status = response.status();
+    if !status.is_success() {
+        let text = response
+            .text()
+            .await
+            .map_err(|err| CoreError::Network(err.to_string()))?;
+        return Err(CoreError::Http {
+            status: status.as_u16(),
+            body: truncate_error_body(&text),
+        });
+    }
+    Ok(response)
+}

litellm-rust/crates/ai-gateway/src/messages/mod.rs

@@ -9,12 +9,37 @@ mod types;
 
 pub use types::MessagesRequest;
 
-use handler::execute_messages_provider_call;
+use handler::{execute_messages_provider_call, execute_messages_provider_stream};
 use prepare::prepare_messages_call;
 
 pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<Value> {
+    match execute_messages(request, false).await? {
+        MessagesResponse::Json(body) => Ok(body),
+        MessagesResponse::Stream(_) => Err(litellm_core::CoreError::InvalidResponse(
+            "non-streaming messages execution returned a stream".to_string(),
+        )),
+    }
+}
+
+pub(crate) enum MessagesResponse {
+    Json(Value),
+    Stream(reqwest::Response),
+}
+
+pub(crate) async fn execute_messages(
+    request: MessagesRequest<'_>,
+    stream: bool,
+) -> CoreResult<MessagesResponse> {
     let prepared = prepare_messages_call(request)?;
-    execute_messages_provider_call(prepared).await
+    if stream {
+        execute_messages_provider_stream(prepared)
+            .await
+            .map(MessagesResponse::Stream)
+    } else {
+        execute_messages_provider_call(prepared)
+            .await
+            .map(MessagesResponse::Json)
+    }
 }
 
 #[cfg(test)]

litellm-rust/crates/ai-gateway/src/messages/prepare.rs

@@ -62,6 +62,7 @@ pub(super) fn prepare_messages_call(
     })?;
 
     Ok(ProviderMessagesRequest {
+        provider: provider.to_string(),
         model,
         config,
         url,

litellm-rust/crates/ai-gateway/src/messages/types.rs

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

litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs

@@ -1,11 +1,10 @@
-//! `POST /v1/messages`, the non-streaming Anthropic Messages HTTP surface.
-//! SSE streaming is a follow-up because the current `messages()` API returns a
-//! single JSON value.
+//! `POST /v1/messages`, the Anthropic Messages HTTP surface.
 
 mod service;
 
+use axum::body::Body;
 use axum::extract::{Json, State};
-use axum::http::header::HeaderMap;
+use axum::http::header::{HeaderMap, HeaderValue, CACHE_CONTROL, CONTENT_TYPE};
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
 use axum::routing::post;
@@ -27,12 +26,42 @@ async fn handle(
     State(state): State<AppState>,
     headers: HeaderMap,
     Json(body): Json<Value>,
-) -> Result<Json<Value>, MessagesRouteError> {
+) -> Result<Response, MessagesRouteError> {
     let extra_headers = forwarded_headers(&headers)?;
-    service::run(&state.router, body, extra_headers)
+    match service::run(&state.router, body, extra_headers)
         .await
-        .map(Json)
-        .map_err(MessagesRouteError::from)
+        .map_err(MessagesRouteError::from)?
+    {
+        service::MessagesResponse::Json(body) => Ok(Json(body).into_response()),
+        service::MessagesResponse::Stream(upstream) => stream_response(upstream),
+    }
+}
+
+fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRouteError> {
+    let content_type = upstream
+        .headers()
+        .get(CONTENT_TYPE)
+        .cloned()
+        .unwrap_or_else(|| HeaderValue::from_static("text/event-stream"));
+    let mut response = Response::builder()
+        .status(
+            StatusCode::from_u16(upstream.status().as_u16()).map_err(|error| {
+                MessagesRouteError(CoreError::InvalidResponse(format!(
+                    "invalid upstream response status: {error}"
+                )))
+            })?,
+        )
+        .header(CONTENT_TYPE, content_type);
+    if let Some(value) = upstream.headers().get(CACHE_CONTROL) {
+        response = response.header(CACHE_CONTROL, value);
+    }
+    response
+        .body(Body::from_stream(upstream.bytes_stream()))
+        .map_err(|error| {
+            MessagesRouteError(CoreError::InvalidResponse(format!(
+                "failed to build streaming response: {error}"
+            )))
+        })
 }
 
 fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, CoreError> {
@@ -96,6 +125,7 @@ mod tests {
     use std::sync::Arc;
 
     use axum::body::Body;
+    use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE};
     use axum::http::Request;
     use axum::http::StatusCode;
     use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
@@ -169,6 +199,53 @@ mod tests {
         (format!("http://{address}"), server)
     }
 
+    async fn streaming_upstream(
+        listener: TcpListener,
+        status: u16,
+        content_type: &'static str,
+        body: &'static str,
+    ) -> (String, tokio::task::JoinHandle<String>) {
+        let address = listener.local_addr().expect("listener has address");
+        let server = tokio::spawn(async move {
+            let (mut socket, _) = listener.accept().await.expect("accepts request");
+            let mut request = Vec::new();
+            let mut buffer = [0_u8; 4096];
+            loop {
+                let read = socket.read(&mut buffer).await.expect("reads request");
+                request.extend_from_slice(&buffer[..read]);
+                if request.windows(4).any(|window| window == b"\r\n\r\n") {
+                    break;
+                }
+            }
+            let request_text = String::from_utf8(request).expect("request is utf8");
+            let content_length = request_text
+                .lines()
+                .find_map(|line| {
+                    let (name, value) = line.split_once(':')?;
+                    name.eq_ignore_ascii_case("content-length")
+                        .then(|| value.trim().parse::<usize>().ok())
+                        .flatten()
+                })
+                .unwrap_or(0);
+            let header_end = request_text.find("\r\n\r\n").expect("request has headers") + 4;
+            let mut full_request = request_text.into_bytes();
+            while full_request.len().saturating_sub(header_end) < content_length {
+                let read = socket.read(&mut buffer).await.expect("reads body");
+                full_request.extend_from_slice(&buffer[..read]);
+            }
+            let response = format!(
+                "HTTP/1.1 {status} OK\r\ncontent-type: {content_type}\r\ncache-control: no-cache\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
+                body.len()
+            );
+            socket
+                .write_all(response.as_bytes())
+                .await
+                .expect("writes response");
+            String::from_utf8(full_request).expect("request is utf8")
+        });
+        (format!("http://{address}"), server)
+    }
+
     #[tokio::test]
     async fn route_constructs_anthropic_upstream_request() {
         let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
@@ -216,6 +293,110 @@ mod tests {
         assert_eq!(body["messages"][0]["content"], "hello");
     }
 
+    #[tokio::test]
+    async fn route_streams_anthropic_events_without_buffering_or_reordering() {
+        let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
+        let events = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
+        let (api_base, server) =
+            streaming_upstream(listener, 200, "text/event-stream", events).await;
+        let app = app(state("claude-test", api_base, Some("master-key")));
+        let response = app
+            .oneshot(
+                Request::builder()
+                    .method("POST")
+                    .uri("/v1/messages")
+                    .header("authorization", "Bearer master-key")
+                    .header("content-type", "application/json")
+                    .body(Body::from(
+                        json!({
+                            "model": "claude-test",
+                            "max_tokens": 16,
+                            "stream": true,
+                            "messages": [{"role": "user", "content": "hello"}]
+                        })
+                        .to_string(),
+                    ))
+                    .expect("request builds"),
+            )
+            .await
+            .expect("route responds");
+        assert_eq!(response.status(), StatusCode::OK);
+        assert_eq!(
+            response
+                .headers()
+                .get(CONTENT_TYPE)
+                .unwrap()
+                .to_str()
+                .unwrap(),
+            "text/event-stream"
+        );
+        assert_eq!(
+            response
+                .headers()
+                .get(CACHE_CONTROL)
+                .unwrap()
+                .to_str()
+                .unwrap(),
+            "no-cache"
+        );
+        let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
+            .await
+            .expect("response body reads");
+        assert_eq!(response_body, events.as_bytes());
+        let upstream_request = server.await.expect("upstream task completes");
+        let (_, upstream_body) = upstream_request
+            .split_once("\r\n\r\n")
+            .expect("upstream request has body");
+        assert_eq!(
+            serde_json::from_str::<serde_json::Value>(upstream_body)
+                .expect("upstream body is json")["stream"],
+            true
+        );
+    }
+
+    #[tokio::test]
+    async fn route_maps_streaming_upstream_errors_before_starting_response() {
+        let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
+        let (api_base, server) = streaming_upstream(
+            listener,
+            429,
+            "application/json",
+            r#"{"error":"rate limited"}"#,
+        )
+        .await;
+        let app = app(state("claude-test", api_base, Some("master-key")));
+        let response = app
+            .oneshot(
+                Request::builder()
+                    .method("POST")
+                    .uri("/v1/messages")
+                    .header("authorization", "Bearer master-key")
+                    .header("content-type", "application/json")
+                    .body(Body::from(
+                        json!({
+                            "model": "claude-test",
+                            "max_tokens": 16,
+                            "stream": true,
+                            "messages": [{"role": "user", "content": "hello"}]
+                        })
+                        .to_string(),
+                    ))
+                    .expect("request builds"),
+            )
+            .await
+            .expect("route responds");
+        assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
+        let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
+            .await
+            .expect("response body reads");
+        assert_eq!(
+            serde_json::from_slice::<serde_json::Value>(&response_body).expect("error is json")
+                ["error"]["message"],
+            "messages provider request failed"
+        );
+        server.await.expect("upstream task completes");
+    }
+
     #[tokio::test]
     async fn route_rejects_missing_master_key() {
         let app = app(state(

litellm-rust/crates/ai-gateway/src/routes/messages/service.rs

@@ -5,13 +5,18 @@ use litellm_core::{CoreError, CoreResult};
 use serde_json::{Map, Value};
 
 use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
-use crate::messages::{messages, MessagesRequest};
+use crate::messages::{execute_messages, MessagesRequest};
+
+pub(crate) enum MessagesResponse {
+    Json(Value),
+    Stream(reqwest::Response),
+}
 
 pub async fn run(
     router: &Arc<Router>,
     body: Value,
     extra_headers: Option<Map<String, Value>>,
-) -> CoreResult<Value> {
+) -> CoreResult<MessagesResponse> {
     let model = body
         .get("model")
         .and_then(Value::as_str)
@@ -38,7 +43,7 @@ pub async fn run(
             Value::String(upstream_model.to_string()),
         );
 
-    messages(MessagesRequest {
+    let request = MessagesRequest {
         model: provider_model,
         body,
         api_key: deployment.litellm_params.api_key.as_deref(),
@@ -46,6 +51,14 @@ pub async fn run(
         custom_llm_provider,
         extra_headers,
         timeout: None,
-    })
-    .await
+    };
+    let stream = request.body.get("stream").and_then(Value::as_bool) == Some(true);
+    execute_messages(request, stream)
+        .await
+        .map(|response| match response {
+            crate::messages::MessagesResponse::Json(body) => MessagesResponse::Json(body),
+            crate::messages::MessagesResponse::Stream(upstream) => {
+                MessagesResponse::Stream(upstream)
+            }
+        })
 }