diff --git a/crates/switchyard-server/src/response.rs b/crates/switchyard-server/src/response.rs index 2149f7036..98515aab3 100644 --- a/crates/switchyard-server/src/response.rs +++ b/crates/switchyard-server/src/response.rs @@ -8,7 +8,7 @@ use std::error::Error; use axum::Json; use axum::response::{IntoResponse, Response as HttpResponse}; use switchyard_protocol::{LlmResponse, Response as AlgorithmResponse}; -use switchyard_translation::{WireFormat, encode_aggregated_response, encode_stream}; +use switchyard_translation::{WireFormat, encode_aggregated_response, encode_stream_with_outcome}; use crate::sse::frame_stream; @@ -29,10 +29,10 @@ pub(crate) fn into_http_response( served_model.as_deref(), )?) .into_response()), - LlmResponse::Stream(stream) => Ok(frame_stream( - encode_stream(stream, target_format, served_model)?, - target_format, - ) - .into_response()), + LlmResponse::Stream(stream) => { + let (events, outcome) = + encode_stream_with_outcome(stream, target_format, served_model)?; + Ok(frame_stream(events, target_format, outcome).into_response()) + } } } diff --git a/crates/switchyard-server/src/sse.rs b/crates/switchyard-server/src/sse.rs index 6354945ec..2038e70f2 100644 --- a/crates/switchyard-server/src/sse.rs +++ b/crates/switchyard-server/src/sse.rs @@ -8,16 +8,21 @@ use std::convert::Infallible; use axum::response::sse::{Event, Sse}; use futures_util::Stream; use serde_json::{Value, json}; -use switchyard_translation::{RawEventStream, WireFormat}; +use switchyard_translation::{RawEventStream, StreamOutcome, WireFormat}; /// Boxed stream type accepted by Axum's SSE response wrapper. pub(crate) type SseFrameStream = std::pin::Pin> + Send>>; /// Converts translated JSON events into endpoint-specific SSE frames. +/// +/// `outcome` reports whether the translated stream ended on an in-band upstream +/// error. A failed stream must not be closed with a success sentinel, so the +/// OpenAI Chat `[DONE]` marker is emitted only for a clean finish. pub(crate) fn frame_stream( stream: RawEventStream, target_format: WireFormat, + outcome: StreamOutcome, ) -> Sse { let framed = async_stream::stream! { let mut stream = stream; @@ -43,7 +48,10 @@ pub(crate) fn frame_stream( } } - if !failed && target_format == WireFormat::OpenAiChat { + // An upstream in-band error terminates the stream inside the translation + // layer, which reports it through `outcome` rather than a framing error. + let upstream_errored = outcome.load(std::sync::atomic::Ordering::Acquire); + if !failed && !upstream_errored && target_format == WireFormat::OpenAiChat { yield Ok(Event::default().data("[DONE]")); } }; @@ -111,7 +119,12 @@ mod tests { Ok(json!({"id": "after"})), ])); - let response = frame_stream(stream, WireFormat::OpenAiChat).into_response(); + let response = frame_stream( + stream, + WireFormat::OpenAiChat, + std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), + ) + .into_response(); let body = String::from_utf8(to_bytes(response.into_body(), usize::MAX).await?.to_vec())?; // A stream error is terminal: later chunks and success markers must not be emitted. @@ -121,4 +134,46 @@ mod tests { assert!(!body.contains("[DONE]")); Ok(()) } + + // An upstream in-band error ends the translated stream without a framing error, + // so the outcome flag is the only signal that the turn failed. `[DONE]` would + // otherwise tell the client a failed generation completed successfully. + #[tokio::test] + async fn upstream_in_band_error_suppresses_the_done_marker() -> TestResult { + let stream: RawEventStream = Box::pin(stream::iter(vec![ + Ok(json!({"choices": [{"delta": {"content": "partial"}}]})), + Ok(json!({"error": {"message": "upstream failed after stream start"}})), + ])); + let outcome = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)); + + let response = frame_stream(stream, WireFormat::OpenAiChat, outcome).into_response(); + let body = String::from_utf8(to_bytes(response.into_body(), usize::MAX).await?.to_vec())?; + + assert!(body.contains("partial")); + assert!(body.contains("upstream failed after stream start")); + assert!( + !body.contains("[DONE]"), + "failed stream must not be closed with [DONE]:\n{body}" + ); + Ok(()) + } + + // A clean stream still ends with the sentinel the OpenAI Chat contract requires. + #[tokio::test] + async fn clean_stream_still_emits_the_done_marker() -> TestResult { + let stream: RawEventStream = Box::pin(stream::iter(vec![Ok( + json!({"choices": [{"delta": {"content": "hello"}}]}), + )])); + let outcome = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + + let response = frame_stream(stream, WireFormat::OpenAiChat, outcome).into_response(); + let body = String::from_utf8(to_bytes(response.into_body(), usize::MAX).await?.to_vec())?; + + assert!(body.contains("hello")); + assert!( + body.contains("[DONE]"), + "clean stream must end with [DONE]:\n{body}" + ); + Ok(()) + } } diff --git a/crates/switchyard-translation/src/helpers.rs b/crates/switchyard-translation/src/helpers.rs index f1fc7c321..fd4db20a6 100644 --- a/crates/switchyard-translation/src/helpers.rs +++ b/crates/switchyard-translation/src/helpers.rs @@ -75,6 +75,16 @@ pub type RawEventStream = Pin< >, >; +/// Shared flag reporting whether a [`RawEventStream`] ended on an in-band error. +/// +/// An in-band error is terminal: the encoder emits the error event and stops +/// without calling `finish`, so no terminal event follows. The serving layer +/// cannot tell that apart from a clean end-of-stream, yet it must not append a +/// success sentinel (OpenAI Chat's `[DONE]`) to a failed stream. This flag +/// carries that outcome across the layer boundary; `false` until the encoder +/// observes an error. +pub type StreamOutcome = std::sync::Arc; + /// Encodes a stream of IR chunks into a stream of target-format wire events. /// /// `served_model` is exposed as the response model (via the stream state's @@ -86,6 +96,16 @@ pub fn encode_stream( target: WireFormat, served_model: Option, ) -> std::result::Result { + encode_stream_with_outcome(chunks, target, served_model).map(|(stream, _)| stream) +} + +/// Same as [`encode_stream`], plus a [`StreamOutcome`] the caller can read once +/// the stream ends to tell a clean finish from an in-band error termination. +pub fn encode_stream_with_outcome( + chunks: LlmResponseStream, + target: WireFormat, + served_model: Option, +) -> std::result::Result<(RawEventStream, StreamOutcome), LlmClientError> { let target_format: FormatId = target.into(); // The target is always a built-in wire format, so this lookup cannot fail; a // failure returns as an `Err` rather than a panic. @@ -103,6 +123,8 @@ pub fn encode_stream( ..Default::default() }; let mut chunks = chunks; + let outcome: StreamOutcome = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let outcome_for_stream = std::sync::Arc::clone(&outcome); let events = try_stream! { while let Some(item) = chunks.next().await { @@ -118,6 +140,9 @@ pub fn encode_stream( yield value; } if state.errored { + // Terminal: report the outcome so the serving layer skips any + // success sentinel, and stop without calling `finish`. + outcome_for_stream.store(true, std::sync::atomic::Ordering::Release); return; } } @@ -131,7 +156,7 @@ pub fn encode_stream( } }; - Ok(Box::pin(events)) + Ok((Box::pin(events), outcome)) } // The raw-response helper promises that the caller sees the model that served the