diff --git a/litellm-rust/crates/ai-gateway/src/messages/handler.rs b/litellm-rust/crates/ai-gateway/src/messages/handler.rs index 88faacff473..24244cf33f1 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/handler.rs @@ -93,12 +93,7 @@ pub(super) async fn execute_messages_provider_stream( body: truncate_error_body(&text), }); } - let content_type = response - .headers() - .get(reqwest::header::CONTENT_TYPE) - .and_then(|value| value.to_str().ok()) - .unwrap_or("text/event-stream") - .to_string(); + let content_type = response_content_type(request.streaming, response.headers()); let cache_control = response .headers() .get(reqwest::header::CACHE_CONTROL) @@ -118,6 +113,20 @@ pub(super) async fn execute_messages_provider_stream( }) } +fn response_content_type( + streaming: MessagesStreaming, + headers: &reqwest::header::HeaderMap, +) -> String { + if matches!(streaming, MessagesStreaming::BedrockEventStream) { + return "text/event-stream".to_string(); + } + headers + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or("text/event-stream") + .to_string() +} + #[cfg_attr(not(feature = "server"), allow(dead_code))] pub(crate) struct MessagesStream { pub(crate) content_type: String, @@ -212,3 +221,33 @@ async fn signed_headers( fn environment_lookup(key: &str) -> Option { std::env::var(key).ok() } + +#[cfg(test)] +mod tests { + use super::response_content_type; + use litellm_core::messages::transformation::MessagesStreaming; + use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue}; + + #[test] + fn bedrock_transcoding_uses_anthropic_sse_content_type() { + let mut headers = HeaderMap::new(); + headers.insert( + CONTENT_TYPE, + HeaderValue::from_static("application/vnd.amazon.eventstream"), + ); + assert_eq!( + response_content_type(MessagesStreaming::BedrockEventStream, &headers), + "text/event-stream" + ); + } + + #[test] + fn passthrough_stream_keeps_upstream_content_type() { + let mut headers = HeaderMap::new(); + headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); + assert_eq!( + response_content_type(MessagesStreaming::SsePassthrough, &headers), + "text/event-stream" + ); + } +} 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 04e0ec1bb47..9b63078ca13 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -125,6 +125,8 @@ mod tests { use axum::http::Request; use axum::http::StatusCode; use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE}; + use futures_util::stream; + use litellm_core::CoreResult; use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter}; use serde_json::json; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -132,7 +134,9 @@ mod tests { use tower::ServiceExt; use super::super::app; + use super::stream_response; use crate::io::realtime_pool::RealtimePool; + use crate::messages::MessagesStream; use crate::state::AppState; fn state(model: &str, api_base: String, master_key: Option<&str>) -> AppState { @@ -400,6 +404,23 @@ mod tests { ); } + #[test] + fn transcoded_stream_response_does_not_forward_upstream_framing_headers() { + let response = stream_response(MessagesStream { + content_type: "text/event-stream".to_string(), + cache_control: Some("no-cache".to_string()), + body: Box::pin(stream::empty::>()), + }) + .expect("response builds"); + assert_eq!( + response.headers().get(CONTENT_TYPE).unwrap(), + "text/event-stream" + ); + assert_eq!(response.headers().get(CACHE_CONTROL).unwrap(), "no-cache"); + assert!(response.headers().get("content-length").is_none()); + assert!(response.headers().get("content-encoding").is_none()); + } + #[tokio::test] async fn route_maps_streaming_upstream_errors_before_starting_response() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");