fix(rust): expose transcoded messages as SSE

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

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<String> {
     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"
+        );
+    }
+}

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::<CoreResult<bytes::Bytes>>()),
+        })
+        .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");