From 6ea0873649f9e9152aeafe03a0cec3265a638d4b Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Sat, 19 Sep 2026 16:44:45 -0700 Subject: [PATCH] fix(translation): recover Responses terminal streams Signed-off-by: Anthony Casagrande --- .../src/codecs/responses/stream.rs | 141 ++++++----- crates/switchyard-translation/src/helpers.rs | 76 +++++- crates/switchyard-translation/src/sse.rs | 134 +++++++---- .../tests/stream_translation.rs | 227 ++++++++++++++++-- 4 files changed, 446 insertions(+), 132 deletions(-) diff --git a/crates/switchyard-translation/src/codecs/responses/stream.rs b/crates/switchyard-translation/src/codecs/responses/stream.rs index a179d4ee7..64076cd9b 100644 --- a/crates/switchyard-translation/src/codecs/responses/stream.rs +++ b/crates/switchyard-translation/src/codecs/responses/stream.rs @@ -97,6 +97,44 @@ struct ResponsesStreamResponse { usage: Value, } +fn decode_responses_identity( + state: &mut StreamTranslationState, + response: Option<&serde_json::Map>, +) -> Vec { + let id = response + .and_then(|response| response.get("id")) + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map(ToOwned::to_owned); + let model = response + .and_then(|response| response.get("model")) + .and_then(Value::as_str) + .filter(|model| !model.is_empty()) + .map(ToOwned::to_owned); + let identity_changed = id + .as_deref() + .is_some_and(|id| state.message_id.as_deref() != Some(id)) + || model + .as_deref() + .is_some_and(|model| state.model.as_deref() != Some(model)); + + if id.is_some() { + state.message_id = id; + } + if model.is_some() { + state.model = model; + } + if !state.saw_message_start || identity_changed { + state.saw_message_start = true; + vec![LlmResponseChunk::MessageStart { + id: state.message_id.clone(), + model: state.model.clone(), + }] + } else { + Vec::new() + } +} + // Decodes one OpenAI Responses event into neutral streaming events. fn decode_responses_stream( state: &mut StreamTranslationState, @@ -110,25 +148,8 @@ fn decode_responses_stream( .or_else(|| event.get("event")) .and_then(Value::as_str); match event_type { - Some("response.created") => { - state.saw_message_start = true; - let response = event.get("response").and_then(Value::as_object); - if let Some(model) = response - .and_then(|response| response.get("model")) - .and_then(Value::as_str) - { - state.model = Some(model.to_string()); - } - if let Some(id) = response - .and_then(|response| response.get("id")) - .and_then(Value::as_str) - { - state.message_id = Some(id.to_string()); - } - vec![LlmResponseChunk::MessageStart { - id: state.message_id.clone(), - model: state.model.clone(), - }] + Some("response.created" | "response.in_progress") => { + decode_responses_identity(state, event.get("response").and_then(Value::as_object)) } Some("response.output_text.delta") => event .get("delta") @@ -257,48 +278,10 @@ fn decode_responses_stream( }) .unwrap_or_default() } - Some("response.completed") => { - let mut out = Vec::new(); - // Some providers send output only in the final snapshot. Reconcile it with - // decoded deltas before emitting the stop, without repeating streamed content. - if let Some(items) = event - .get("response") - .and_then(|response| response.get("output")) - .and_then(Value::as_array) - { - for (position, item) in items.iter().enumerate() { - if let Some(item) = item.as_object() { - out.extend(decode_responses_completed_item(item, position, state)); - if matches!(out.last(), Some(LlmResponseChunk::StreamError { .. })) { - return out; - } - } - } - } - if let Some(usage) = event - .get("response") - .and_then(Value::as_object) - .and_then(|response| response.get("usage")) - .and_then(Value::as_object) - { - let usage = responses_usage(usage); - state.usage = usage.clone(); - state.saw_backend_usage = true; - out.push(LlmResponseChunk::Usage(usage)); - } - // A completed response that produced a tool call ended the turn to run that - // tool, not because the assistant was done. The buffered decoder reports tool - // use for such output; the stream must too, or stop-reason-driven tool loops - // (Anthropic `tool_use`, Chat `tool_calls`) stop without running the tool. - // Carries the Anthropic spelling because every encoder already maps it. - let reason = state.decoded_tool_call.then(|| "tool_use".to_string()); - out.push(LlmResponseChunk::MessageStop { reason }); - out + Some("response.completed") => decode_responses_terminal_snapshot(state, event, None), + Some("response.incomplete") => { + decode_responses_terminal_snapshot(state, event, Some("max_tokens".to_string())) } - // Carries the Anthropic spelling because every encoder already maps it. - Some("response.incomplete") => vec![LlmResponseChunk::MessageStop { - reason: Some("max_tokens".to_string()), - }], Some("response.failed") => vec![LlmResponseChunk::StreamError { message: event .get("response") @@ -319,6 +302,44 @@ fn decode_responses_stream( } } +fn decode_responses_terminal_snapshot( + state: &mut StreamTranslationState, + event: &Value, + explicit_stop_reason: Option, +) -> Vec { + let response = event.get("response").and_then(Value::as_object); + let mut out = decode_responses_identity(state, response); + + if let Some(items) = response + .and_then(|response| response.get("output")) + .and_then(Value::as_array) + { + for (position, item) in items.iter().enumerate() { + if let Some(item) = item.as_object() { + out.extend(decode_responses_completed_item(item, position, state)); + if matches!(out.last(), Some(LlmResponseChunk::StreamError { .. })) { + return out; + } + } + } + } + + if let Some(usage) = response + .and_then(|response| response.get("usage")) + .and_then(Value::as_object) + { + let usage = responses_usage(usage); + state.usage = usage.clone(); + state.saw_backend_usage = true; + out.push(LlmResponseChunk::Usage(usage)); + } + + let reason = + explicit_stop_reason.or_else(|| state.decoded_tool_call.then(|| "tool_use".to_string())); + out.push(LlmResponseChunk::MessageStop { reason }); + out +} + // Encodes neutral streaming events into OpenAI Responses events. fn encode_responses_stream( state: &mut StreamTranslationState, diff --git a/crates/switchyard-translation/src/helpers.rs b/crates/switchyard-translation/src/helpers.rs index 95ebb75e7..bb03afee0 100644 --- a/crates/switchyard-translation/src/helpers.rs +++ b/crates/switchyard-translation/src/helpers.rs @@ -288,7 +288,8 @@ where break; } sse::SseFrame::Data(value) => { - saw_terminal |= sse::is_terminal_event(source, &value); + let terminal = sse::terminal_behavior(source, &value); + saw_terminal |= terminal.is_terminal(); let normalized = codec.decode_event(&mut state, &value); saw_error |= normalized.iter().any(|chunk| matches!( chunk, @@ -300,6 +301,9 @@ where value, normalized, ); + if terminal.is_stream_ending() { + break; + } } } } else { @@ -319,7 +323,8 @@ where saw_terminal |= source != WireFormat::AnthropicMessages; } sse::SseFrame::Data(value) => { - saw_terminal |= sse::is_terminal_event(source, &value); + let terminal = sse::terminal_behavior(source, &value); + saw_terminal |= terminal.is_terminal(); let normalized = codec.decode_event(&mut state, &value); saw_error |= normalized.iter().any(|chunk| matches!( chunk, @@ -845,6 +850,73 @@ mod tests { Ok(()) } + #[test] + fn decode_stream_stops_polling_after_explicit_responses_terminals() -> Result<(), BoxError> { + let cases = [ + ( + "response.incomplete", + b"data: {\"type\":\"response.incomplete\",\"response\":{\"status\":\"incomplete\"}}\n\n" + .to_vec(), + ), + ( + "error", + b"data: {\"type\":\"error\",\"message\":\"boom\"}\n\n".to_vec(), + ), + ]; + for (event_type, terminal) in cases { + let bytes = + stream::iter([Ok::, LlmClientError>(terminal)]).chain(stream::poll_fn( + move |_| panic!("decode_stream polled upstream after explicit {event_type}"), + )); + + let events = decode_all(bytes, WireFormat::OpenAiResponses)?; + + assert_eq!(events.len(), 1); + match event_type { + "response.incomplete" => assert!(matches!( + events[0].normalized().last(), + Some(LlmResponseChunk::MessageStop { reason: Some(reason) }) + if reason == "max_tokens" + )), + "error" => assert!(matches!( + events[0].normalized().last(), + Some(LlmResponseChunk::StreamError { message }) if message == "boom" + )), + _ => unreachable!(), + } + } + Ok(()) + } + + #[test] + fn decode_stream_keeps_chat_usage_after_finish_reason() -> Result<(), BoxError> { + let sse = concat!( + "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3,\"total_tokens\":5}}\n\n", + "data: [DONE]\n\n", + ) + .as_bytes() + .to_vec(); + let bytes = + stream::iter([Ok::, LlmClientError>(sse)]).chain(stream::poll_fn(|_| { + panic!("decode_stream polled upstream after [DONE]") + })); + + let events = decode_all(bytes, WireFormat::OpenAiChat)?; + let usage = events + .iter() + .flat_map(LlmResponseStreamEvent::normalized) + .find_map(|chunk| match chunk { + LlmResponseChunk::Usage(usage) => Some(usage), + _ => None, + }) + .ok_or("missing trailing Chat usage")?; + assert_eq!(usage.input_tokens, Some(2)); + assert_eq!(usage.output_tokens, Some(3)); + assert_eq!(usage.total_tokens, Some(5)); + Ok(()) + } + #[test] fn responses_failed_before_done_remains_a_stream_error() -> Result<(), BoxError> { let sse = b"data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"upstream failed\"}}}\n\ndata: [DONE]\n\n".to_vec(); diff --git a/crates/switchyard-translation/src/sse.rs b/crates/switchyard-translation/src/sse.rs index bf24aabbb..da52e7f52 100644 --- a/crates/switchyard-translation/src/sse.rs +++ b/crates/switchyard-translation/src/sse.rs @@ -39,30 +39,63 @@ fn data_field_value(line: &str) -> Option { Some(value.strip_prefix(' ').unwrap_or(value).to_string()) } -/// Returns whether a provider event explicitly completes its wire-format stream. -pub(crate) fn is_terminal_event(format: WireFormat, event: &Value) -> bool { +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum TerminalBehavior { + NotTerminal, + TerminalMayHaveTrailingData, + TerminalEndsStream, +} + +impl TerminalBehavior { + pub(crate) fn is_terminal(self) -> bool { + self != Self::NotTerminal + } + + pub(crate) fn is_stream_ending(self) -> bool { + self == Self::TerminalEndsStream + } +} + +pub(crate) fn terminal_behavior(format: WireFormat, event: &Value) -> TerminalBehavior { 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::OpenAiChat => { + if event + .get("choices") + .and_then(Value::as_array) + .into_iter() + .flatten() + .any(|choice| { + choice + .get("finish_reason") + .and_then(Value::as_str) + .is_some() + }) + { + TerminalBehavior::TerminalMayHaveTrailingData + } else { + TerminalBehavior::NotTerminal + } + } WireFormat::AnthropicMessages => { - event.get("type").and_then(Value::as_str) == Some("message_stop") + if event.get("type").and_then(Value::as_str) == Some("message_stop") { + TerminalBehavior::TerminalEndsStream + } else { + TerminalBehavior::NotTerminal + } + } + WireFormat::OpenAiResponses => { + if matches!( + event + .get("type") + .or_else(|| event.get("event")) + .and_then(Value::as_str), + Some("response.completed" | "response.incomplete" | "response.failed" | "error") + ) { + TerminalBehavior::TerminalEndsStream + } else { + TerminalBehavior::NotTerminal + } } - WireFormat::OpenAiResponses => matches!( - event - .get("type") - .or_else(|| event.get("event")) - .and_then(Value::as_str), - Some("response.completed" | "response.incomplete" | "response.failed") - ), } } @@ -193,31 +226,38 @@ mod tests { } #[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::OpenAiResponses, - &json!({"type": "response.incomplete"}) - )); - assert!(is_terminal_event( - WireFormat::OpenAiResponses, - &json!({"type": "response.failed"}) - )); - assert!(!is_terminal_event( - WireFormat::OpenAiChat, - &json!({"choices": [{"finish_reason": null}]}) - )); + fn classifies_provider_terminal_behavior() { + assert_eq!( + terminal_behavior( + WireFormat::OpenAiChat, + &json!({"choices": [{"finish_reason": "stop"}]}) + ), + TerminalBehavior::TerminalMayHaveTrailingData + ); + assert_eq!( + terminal_behavior( + WireFormat::AnthropicMessages, + &json!({"type": "message_stop"}) + ), + TerminalBehavior::TerminalEndsStream + ); + for event_type in [ + "response.completed", + "response.incomplete", + "response.failed", + "error", + ] { + assert_eq!( + terminal_behavior(WireFormat::OpenAiResponses, &json!({"type": event_type})), + TerminalBehavior::TerminalEndsStream + ); + } + assert_eq!( + terminal_behavior( + WireFormat::OpenAiChat, + &json!({"choices": [{"finish_reason": null}]}) + ), + TerminalBehavior::NotTerminal + ); } } diff --git a/crates/switchyard-translation/tests/stream_translation.rs b/crates/switchyard-translation/tests/stream_translation.rs index e22511c59..80db4c353 100644 --- a/crates/switchyard-translation/tests/stream_translation.rs +++ b/crates/switchyard-translation/tests/stream_translation.rs @@ -1613,7 +1613,13 @@ fn responses_function_call_stream_ends_with_tool_use_on_every_wire() -> TestResu decode_stream_event(&mut state, WireFormat::OpenAiResponses, &text); assert_eq!( decode_stream_event(&mut state, WireFormat::OpenAiResponses, &bare_completed), - vec![LlmResponseChunk::MessageStop { reason: None }] + vec![ + LlmResponseChunk::MessageStart { + id: None, + model: None, + }, + LlmResponseChunk::MessageStop { reason: None }, + ] ); Ok(()) } @@ -2599,6 +2605,11 @@ fn responses_stream_decodes_reasoning_text_from_added_done_and_completed_once() .normalized() .iter() .map(|c| match c { + LlmResponseChunk::MessageStart { id, model } => { + assert_eq!(id.as_deref(), Some("resp_1")); + assert_eq!(model.as_deref(), None); + "start" + } LlmResponseChunk::ReasoningDelta { text, .. } => { assert_eq!(text, "from completed"); "reasoning" @@ -2609,7 +2620,7 @@ fn responses_stream_decodes_reasoning_text_from_added_done_and_completed_once() _ => "other", }) .collect(); - assert_eq!(kinds, vec!["id", "reasoning", "usage", "stop"]); + assert_eq!(kinds, vec!["start", "id", "reasoning", "usage", "stop"]); // (d) completed output repeating already-streamed reasoning adds nothing let mut state = StreamTranslationState::new(format, format); @@ -2994,6 +3005,174 @@ fn responses_stream_empty_reasoning_delta_opens_no_summary_part() -> TestResult Ok(()) } +#[test] +fn responses_lifecycle_and_terminal_snapshots_recover_final_state() -> TestResult { + let source = WireFormat::OpenAiResponses; + for (terminal_type, expected_chunk_reason, expected_stop) in [ + ("response.completed", "tool_use", StopReason::ToolUse), + ("response.incomplete", "max_tokens", StopReason::MaxTokens), + ] { + let mut state = StreamTranslationState::new(source, source); + let mut accumulator = ResponseAccumulator::new(); + + let created = decode_stream_event( + &mut state, + source, + &json!({"type": "response.created", "response": { + "id": "resp_created", "model": "model/created" + }}), + ); + assert_eq!( + created, + vec![LlmResponseChunk::MessageStart { + id: Some("resp_created".into()), + model: Some("model/created".into()), + }] + ); + let in_progress = decode_stream_event( + &mut state, + source, + &json!({"type": "response.in_progress", "response": { + "id": "resp_provisional", "model": "model/provisional" + }}), + ); + assert_eq!( + in_progress, + vec![LlmResponseChunk::MessageStart { + id: Some("resp_provisional".into()), + model: Some("model/provisional".into()), + }] + ); + assert_eq!(state.message_id.as_deref(), Some("resp_provisional")); + assert_eq!(state.model.as_deref(), Some("model/provisional")); + for chunk in created.into_iter().chain(in_progress) { + accumulator.push(chunk); + } + + let terminal = json!({ + "type": terminal_type, + "response": { + "id": "resp_final", + "model": "model/final", + "status": if terminal_type == "response.completed" { "completed" } else { "incomplete" }, + "incomplete_details": if terminal_type == "response.incomplete" { + json!({"reason": "max_output_tokens"}) + } else { + Value::Null + }, + "output": [ + {"type": "message", "role": "assistant", "content": [ + {"type": "output_text", "text": "partial answer"} + ]}, + {"type": "function_call", "id": "fc_1", "call_id": "call_1", + "name": "lookup", "arguments": "{\"query\":\"rust\"}"} + ], + "usage": {"input_tokens": 7, "output_tokens": 3, "total_tokens": 10} + } + }); + let terminal_chunks = decode_stream_event(&mut state, source, &terminal); + assert_eq!(terminal_chunks.len(), 5); + assert!(matches!( + &terminal_chunks[0], + LlmResponseChunk::MessageStart { id: Some(id), model: Some(model) } + if id == "resp_final" && model == "model/final" + )); + assert!(matches!( + &terminal_chunks[1], + LlmResponseChunk::TextDelta { index: 0, text } if text == "partial answer" + )); + assert!(matches!( + &terminal_chunks[2], + LlmResponseChunk::ToolCallDelta { + index: 1, + id: Some(id), + name: Some(name), + arguments_delta: Some(arguments), + } if id == "call_1" + && name == "lookup" + && arguments == "{\"query\":\"rust\"}" + )); + assert!(matches!( + &terminal_chunks[3], + LlmResponseChunk::Usage(usage) + if usage.input_tokens == Some(7) + && usage.output_tokens == Some(3) + && usage.total_tokens == Some(10) + )); + assert!(matches!( + &terminal_chunks[4], + LlmResponseChunk::MessageStop { reason: Some(reason) } + if reason == expected_chunk_reason + )); + for chunk in terminal_chunks { + accumulator.push(chunk); + } + + let aggregate = accumulator.finish(); + assert_eq!(aggregate.id.as_deref(), Some("resp_final")); + assert_eq!(aggregate.model.as_deref(), Some("model/final")); + assert_eq!( + switchyard_protocol::completion_text(&aggregate), + "partial answer" + ); + assert!(matches!( + &aggregate.outputs[0].content[1], + switchyard_protocol::ContentBlock::ToolCall(call) + if call.id == "call_1" + && call.name == "lookup" + && call.arguments == json!({"query": "rust"}) + )); + assert_eq!(aggregate.usage.input_tokens, Some(7)); + assert_eq!(aggregate.usage.output_tokens, Some(3)); + assert_eq!(aggregate.usage.total_tokens, Some(10)); + assert_eq!(aggregate.outputs[0].stop_reason, Some(expected_stop)); + } + Ok(()) +} + +#[test] +fn responses_identity_free_terminal_snapshots_start_before_output() { + let source = WireFormat::OpenAiResponses; + for (terminal_type, stop_reason) in [ + ("response.completed", None), + ("response.incomplete", Some("max_tokens")), + ] { + let mut state = StreamTranslationState::new(source, source); + let decoded = decode_stream_event( + &mut state, + source, + &json!({ + "type": terminal_type, + "response": { + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "terminal text"}] + }], + "usage": {"input_tokens": 2, "output_tokens": 1, "total_tokens": 3} + } + }), + ); + + assert!(matches!( + decoded[0], + LlmResponseChunk::MessageStart { + id: None, + model: None, + } + )); + assert!(matches!( + &decoded[1], + LlmResponseChunk::TextDelta { index: 0, text } if text == "terminal text" + )); + assert!(matches!(decoded[2], LlmResponseChunk::Usage(_))); + assert!(matches!( + &decoded[3], + LlmResponseChunk::MessageStop { reason } if reason.as_deref() == stop_reason + )); + } +} + #[test] fn responses_terminal_snapshots_recover_missing_output_once() -> TestResult { let engine = TranslationEngine::default(); @@ -3139,27 +3318,29 @@ fn responses_terminal_snapshots_recover_missing_output_once() -> TestResult { ); } } - for delta in [ - json!({"type": "response.output_text.delta", "output_index": 0, "delta": "Different"}), - json!({"type": "response.function_call_arguments.delta", "output_index": 1, "delta": "["}), - ] { - let mut state = StreamTranslationState::new(source, WireFormat::OpenAiChat); - engine.decode_stream_event(&mut state, source, delta)?; - let decoded = engine.decode_stream_event( - &mut state, - source, - json!({"type": "response.completed", "response": response}), - )?; - assert!(matches!( - decoded.normalized().last(), - Some(LlmResponseChunk::StreamError { .. }) - )); - assert!( - !decoded - .normalized() - .iter() - .any(|chunk| matches!(chunk, LlmResponseChunk::MessageStop { .. })) - ); + for terminal_type in ["response.completed", "response.incomplete"] { + for delta in [ + json!({"type": "response.output_text.delta", "output_index": 0, "delta": "Different"}), + json!({"type": "response.function_call_arguments.delta", "output_index": 1, "delta": "["}), + ] { + let mut state = StreamTranslationState::new(source, WireFormat::OpenAiChat); + engine.decode_stream_event(&mut state, source, delta)?; + let decoded = engine.decode_stream_event( + &mut state, + source, + json!({"type": terminal_type, "response": response}), + )?; + assert!(matches!( + decoded.normalized().last(), + Some(LlmResponseChunk::StreamError { .. }) + )); + assert!( + !decoded + .normalized() + .iter() + .any(|chunk| matches!(chunk, LlmResponseChunk::MessageStop { .. })) + ); + } } Ok(()) }