From 744868b7b1672230054bbcc160bbfca63cd1df41 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 01:17:51 +0000 Subject: [PATCH] feat(rust): stream Anthropic Messages SSE on POST /v1/messages Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .../crates/ai-gateway/src/messages/handler.rs | 36 ++++ .../crates/ai-gateway/src/messages/mod.rs | 29 ++- .../crates/ai-gateway/src/messages/prepare.rs | 1 + .../crates/ai-gateway/src/messages/types.rs | 1 + .../ai-gateway/src/routes/messages/mod.rs | 197 +++++++++++++++++- .../ai-gateway/src/routes/messages/service.rs | 23 +- 6 files changed, 272 insertions(+), 15 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/messages/handler.rs b/litellm-rust/crates/ai-gateway/src/messages/handler.rs index dd4a2f22aa7..d3b9d3b3fba 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/handler.rs +++ b/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 { + 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) +} diff --git a/litellm-rust/crates/ai-gateway/src/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/messages/mod.rs index 7ed81474c47..883cfb40336 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/mod.rs +++ b/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 { + 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 { 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)] diff --git a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs index 47105b39954..6176f9cb67f 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs +++ b/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, diff --git a/litellm-rust/crates/ai-gateway/src/messages/types.rs b/litellm-rust/crates/ai-gateway/src/messages/types.rs index 6840ff57cc4..848fadb4b02 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/types.rs +++ b/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, diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index 8fd7c4f8d4b..3604eebb8e0 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/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, headers: HeaderMap, Json(body): Json, -) -> Result, MessagesRouteError> { +) -> Result { 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 { + 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>, 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) { + 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::().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::(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::(&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( diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs index 639cb73f77e..7f00123ca39 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/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, body: Value, extra_headers: Option>, -) -> CoreResult { +) -> CoreResult { 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) + } + }) }