diff --git a/crates/switchyard-translation/src/helpers.rs b/crates/switchyard-translation/src/helpers.rs index f1fc7c32..b49a57b2 100644 --- a/crates/switchyard-translation/src/helpers.rs +++ b/crates/switchyard-translation/src/helpers.rs @@ -18,8 +18,9 @@ use switchyard_protocol::LlmClientError; use crate::codecs::stream::encode_response_stream_event; use crate::sse; use crate::{ - AggLlmResponse, FormatId, LlmRequest, LlmResponseStream, LlmResponseStreamEvent, Result, - StreamCodecRegistry, StreamTranslationState, TranslationEngine, TranslationPolicy, WireFormat, + AggLlmResponse, FormatId, LlmRequest, LlmResponseChunk, LlmResponseStream, + LlmResponseStreamEvent, Result, StreamCodecRegistry, StreamTranslationState, TranslationEngine, + TranslationPolicy, WireFormat, }; static DEFAULT_TRANSLATION_POLICY: LazyLock = @@ -165,12 +166,16 @@ fn stamp_streamed_response_model( } } -/// Decodes a byte stream of `source`-format SSE frames into neutral IR chunks. +/// Decodes provider SSE bytes into normalized stream events. /// /// Operates on raw bytes, not any HTTP client type: the caller adapts its /// transport's body stream into `Stream, _>>`. Frames are /// buffered across chunks (a partial frame waits for its boundary); the source /// stream codec is resolved once and reused for every frame. +/// +/// The decoder tracks the source format's protocol-specific terminal event. If +/// EOF arrives without that required event, the stream yields a deferred +/// [`LlmClientError::ResponseTranslation`] error after any valid decoded events. pub fn decode_stream( bytes: S, source: WireFormat, @@ -200,6 +205,8 @@ where ..StreamTranslationState::default() }; let mut frame = String::new(); + let mut saw_terminal = false; + let mut saw_error = false; let stream = Box::pin(try_stream! { futures::pin_mut!(lines); while let Some(line) = lines.next().await { @@ -211,9 +218,18 @@ where frame.clear(); match parsed { sse::SseFrame::Empty => {} - sse::SseFrame::Done => break, + sse::SseFrame::Done => { + saw_terminal |= source != WireFormat::AnthropicMessages; + break; + } sse::SseFrame::Data(value) => { + saw_terminal |= sse::is_terminal_event(source, &value); let normalized = codec.decode_event(&mut state, &value); + saw_error |= normalized.iter().any(|chunk| matches!( + chunk, + LlmResponseChunk::DecodeError { .. } + | LlmResponseChunk::StreamError { .. } + )); yield LlmResponseStreamEvent::preserved( source_format.clone(), value, @@ -233,11 +249,29 @@ where if !frame.trim_end().is_empty() { let parsed = sse::parse_json_sse_frame(&frame, marker) .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; - if let sse::SseFrame::Data(value) = parsed { - let normalized = codec.decode_event(&mut state, &value); - yield LlmResponseStreamEvent::preserved(source_format, value, normalized); + match parsed { + sse::SseFrame::Done => { + saw_terminal |= source != WireFormat::AnthropicMessages; + } + sse::SseFrame::Data(value) => { + saw_terminal |= sse::is_terminal_event(source, &value); + let normalized = codec.decode_event(&mut state, &value); + saw_error |= normalized.iter().any(|chunk| matches!( + chunk, + LlmResponseChunk::DecodeError { .. } + | LlmResponseChunk::StreamError { .. } + )); + yield LlmResponseStreamEvent::preserved(source_format, value, normalized); + } + sse::SseFrame::Empty => {} } } + + if !saw_terminal && !saw_error { + Err(LlmClientError::ResponseTranslation(format!( + "{source} stream ended before a terminal event" + )))?; + } }); Ok(stream) } @@ -683,13 +717,32 @@ mod tests { fn decode_stream_decodes_trailing_frame_without_blank_line() -> Result<(), BoxError> { // A non-standard upstream omits the final blank line; the last frame // must still be decoded rather than dropped. - let sse = b"data: {\"choices\":[{\"delta\":{\"content\":\"tail\"}}]}".to_vec(); + let sse = + b"data: {\"choices\":[{\"delta\":{\"content\":\"tail\"}}]}\n\ndata: [DONE]".to_vec(); let bytes = stream::once(async move { Ok::, LlmClientError>(sse) }); let chunks = decode_all(bytes, WireFormat::OpenAiChat)?; assert_eq!(text_of(&chunks), "tail"); Ok(()) } + #[test] + fn decode_stream_rejects_eof_before_a_terminal_event() -> Result<(), BoxError> { + // Preserve valid content before surfacing premature EOF as the terminal stream error. + let sse = b"data: {\"choices\":[{\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n".to_vec(); + let bytes = stream::once(async move { Ok::, LlmClientError>(sse) }); + let results = block_on(decode_stream(bytes, WireFormat::OpenAiChat)?.collect::>()); + + let Some(Ok(first)) = results.first() else { + return Err("expected the partial event".into()); + }; + assert_eq!(text_of(std::slice::from_ref(first)), "partial"); + let Some(Err(LlmClientError::ResponseTranslation(message))) = results.last() else { + panic!("expected incomplete OpenAI stream to fail"); + }; + assert_eq!(message, "openai_chat stream ended before a terminal event"); + Ok(()) + } + #[test] fn decode_stream_decodes_crlf_delimited_frames() -> Result<(), BoxError> { // CRLF framing: blank lines are `\r\n\r\n` and the bare `\r` must not diff --git a/crates/switchyard-translation/src/sse.rs b/crates/switchyard-translation/src/sse.rs index dce5a361..b58e4142 100644 --- a/crates/switchyard-translation/src/sse.rs +++ b/crates/switchyard-translation/src/sse.rs @@ -29,6 +29,33 @@ pub(crate) fn done_marker(_format: WireFormat) -> Option<&'static str> { Some("[DONE]") } +/// Returns whether a provider event explicitly completes its wire-format stream. +pub(crate) fn is_terminal_event(format: WireFormat, event: &Value) -> bool { + match format { + WireFormat::OpenAiChat => event + .get("choices") + .and_then(Value::as_array) + .into_iter() + .flatten() + .any(|choice| { + choice + .get("finish_reason") + .and_then(Value::as_str) + .is_some() + }), + WireFormat::AnthropicMessages => { + event.get("type").and_then(Value::as_str) == Some("message_stop") + } + WireFormat::OpenAiResponses => matches!( + event + .get("type") + .or_else(|| event.get("event")) + .and_then(Value::as_str), + Some("response.completed" | "response.incomplete") + ), + } +} + pub(crate) fn parse_json_sse_frame( frame: &str, done_marker: Option<&str>, @@ -125,4 +152,25 @@ mod tests { fn anthropic_accepts_optional_done_marker() { assert_eq!(done_marker(WireFormat::AnthropicMessages), Some("[DONE]")); } + + #[test] + fn recognizes_provider_terminal_events() { + // Each source format requires its own protocol-specific terminal event. + assert!(is_terminal_event( + WireFormat::OpenAiChat, + &json!({"choices": [{"finish_reason": "stop"}]}) + )); + assert!(is_terminal_event( + WireFormat::AnthropicMessages, + &json!({"type": "message_stop"}) + )); + assert!(is_terminal_event( + WireFormat::OpenAiResponses, + &json!({"type": "response.completed"}) + )); + assert!(!is_terminal_event( + WireFormat::OpenAiChat, + &json!({"choices": [{"finish_reason": null}]}) + )); + } }