fix(server): do not append [DONE] after an upstream in-band stream error

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())
+        }
     }
 }

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<Box<dyn Stream<Item = Result<Event, Infallible>> + 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<SseFrameStream> {
     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(())
+    }
 }

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<std::sync::atomic::AtomicBool>;
+
 /// 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<String>,
 ) -> std::result::Result<RawEventStream, LlmClientError> {
+    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<String>,
+) -> 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