From 773a451125231e69209cadb697666a1f03d8f909 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Wed, 29 Jul 2026 17:47:08 -0600 Subject: [PATCH] feat(translation): preserve raw stream events Signed-off-by: Bryan Bednarski --- crates/protocol/src/stream.rs | 109 +++++++++++++++--- .../src/codecs/anthropic/stream.rs | 9 ++ .../src/codecs/openai_chat/stream.rs | 9 ++ .../src/codecs/responses/stream.rs | 9 ++ .../src/codecs/stream.rs | 15 +++ crates/switchyard-translation/src/engine.rs | 60 ++++++++++ .../tests/stream_translation.rs | 86 +++++++++++++- 7 files changed, 278 insertions(+), 19 deletions(-) diff --git a/crates/protocol/src/stream.rs b/crates/protocol/src/stream.rs index 0db72066..35504034 100644 --- a/crates/protocol/src/stream.rs +++ b/crates/protocol/src/stream.rs @@ -13,6 +13,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::{ + format::FormatId, llm::{AggLlmResponse, ContentBlock, ResponseOutput, Role, StopReason, ToolCall, Usage}, LlmClientError, }; @@ -58,21 +59,7 @@ impl LlmResponse { LlmResponse::Stream(mut stream) => { let mut accumulator = ResponseAccumulator::new(); while let Some(item) = stream.next().await { - match item? { - LlmResponseChunk::DecodeError { message } => { - return Err(LlmClientError::ResponseTranslation(message)); - } - LlmResponseChunk::StreamError { message } => { - // The upstream reported the failure inside the response body, so - // there is no real status line to carry; 502 stands in for "the - // upstream failed" the same way a failed non-streaming call would. - return Err(LlmClientError::UpstreamHttp { - status: MID_STREAM_UPSTREAM_STATUS, - body: message, - }); - } - chunk => accumulator.push(chunk), - } + push_checked_chunk(&mut accumulator, item?)?; } Ok(accumulator.finish()) } @@ -88,11 +75,48 @@ impl LlmResponse { } } -/// One provider-neutral streaming event — the normalized counterpart to -/// [`AggLlmResponse`](crate::AggLlmResponse), sitting between stream decoders and -/// encoders. `switchyard-translation` re-exports it as `ConversationStreamEvent`. +fn push_checked_chunk( + accumulator: &mut ResponseAccumulator, + chunk: LlmResponseChunk, +) -> Result<(), LlmClientError> { + match chunk { + LlmResponseChunk::ProviderEvent { normalized, .. } => { + for chunk in normalized { + push_checked_chunk(accumulator, chunk)?; + } + Ok(()) + } + LlmResponseChunk::DecodeError { message } => { + Err(LlmClientError::ResponseTranslation(message)) + } + LlmResponseChunk::StreamError { message } => Err(LlmClientError::UpstreamHttp { + status: MID_STREAM_UPSTREAM_STATUS, + body: message, + }), + chunk => { + accumulator.push(chunk); + Ok(()) + } + } +} + +/// One streaming event carried between a host and an algorithm. +/// +/// Normalized variants expose provider-neutral meaning. [`ProviderEvent`](Self::ProviderEvent) +/// additionally retains exact source JSON for a lossless same-format round trip while keeping +/// normalized children available to algorithms and cross-format encoders. +/// `switchyard-translation` re-exports this type as `ConversationStreamEvent`. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub enum LlmResponseChunk { + /// One exact provider event paired with the neutral events decoded from it. + ProviderEvent { + /// Provider format that produced `raw`. + source: FormatId, + /// Exact parsed provider event. + raw: Value, + /// Provider-neutral events decoded from `raw`, in source order. + normalized: Vec, + }, MessageStart { id: Option, model: Option, @@ -160,6 +184,11 @@ impl ResponseAccumulator { /// earlier ones; text, reasoning, and tool-call arguments append. pub fn push(&mut self, chunk: LlmResponseChunk) { match chunk { + LlmResponseChunk::ProviderEvent { normalized, .. } => { + for chunk in normalized { + self.push(chunk); + } + } LlmResponseChunk::MessageStart { id, model } => { if id.is_some() { self.id = id; @@ -305,6 +334,50 @@ mod tests { ); } + #[test] + fn folds_normalized_chunks_inside_provider_event() { + let aggregate = fold(vec![LlmResponseChunk::ProviderEvent { + source: crate::WireFormat::OpenAiChat.into(), + raw: json!({ + "choices": [{"delta": {"content": "hello"}}], + "system_fingerprint": "fp_exact" + }), + normalized: vec![LlmResponseChunk::TextDelta { + index: 0, + text: "hello".to_string(), + }], + }]); + + assert_eq!( + aggregate.outputs[0].content, + vec![ContentBlock::Text { + text: "hello".to_string() + }] + ); + } + + #[test] + fn stream_errors_inside_provider_events_remain_typed() { + let response = LlmResponse::Stream(Box::pin(stream::iter([Ok( + LlmResponseChunk::ProviderEvent { + source: crate::WireFormat::OpenAiChat.into(), + raw: json!({"error": {"message": "provider failed"}}), + normalized: vec![LlmResponseChunk::StreamError { + message: "provider failed".to_string(), + }], + }, + )]))); + + let error = block_on(response.into_agg()).err(); + assert!(matches!( + error, + Some(LlmClientError::UpstreamHttp { + status: MID_STREAM_UPSTREAM_STATUS, + .. + }) + )); + } + #[test] fn assembles_tool_calls_by_index() { // id/name arrive once, arguments stream across deltas and parse as JSON. diff --git a/crates/switchyard-translation/src/codecs/anthropic/stream.rs b/crates/switchyard-translation/src/codecs/anthropic/stream.rs index 1a9bfdec..2fd8e42c 100644 --- a/crates/switchyard-translation/src/codecs/anthropic/stream.rs +++ b/crates/switchyard-translation/src/codecs/anthropic/stream.rs @@ -126,6 +126,15 @@ fn encode_anthropic_stream( event: LlmResponseChunk, ) -> Vec { match event { + LlmResponseChunk::ProviderEvent { + source, + raw, + normalized: _, + } if source == WireFormat::AnthropicMessages.into() => vec![raw], + LlmResponseChunk::ProviderEvent { normalized, .. } => normalized + .into_iter() + .flat_map(|event| encode_anthropic_stream(state, event)) + .collect(), LlmResponseChunk::MessageStart { id, model } => { record_source_identity(state, id, model); if state.emitted_message_start { diff --git a/crates/switchyard-translation/src/codecs/openai_chat/stream.rs b/crates/switchyard-translation/src/codecs/openai_chat/stream.rs index b2dbca71..22e4750b 100644 --- a/crates/switchyard-translation/src/codecs/openai_chat/stream.rs +++ b/crates/switchyard-translation/src/codecs/openai_chat/stream.rs @@ -154,6 +154,15 @@ fn encode_openai_chat_stream( event: LlmResponseChunk, ) -> Vec { match event { + LlmResponseChunk::ProviderEvent { + source, + raw, + normalized: _, + } if source == WireFormat::OpenAiChat.into() => vec![raw], + LlmResponseChunk::ProviderEvent { normalized, .. } => normalized + .into_iter() + .flat_map(|event| encode_openai_chat_stream(state, event)) + .collect(), LlmResponseChunk::MessageStart { id, model } => { record_source_identity(state, id, model); if state.emitted_message_start diff --git a/crates/switchyard-translation/src/codecs/responses/stream.rs b/crates/switchyard-translation/src/codecs/responses/stream.rs index 60610e0d..992b3447 100644 --- a/crates/switchyard-translation/src/codecs/responses/stream.rs +++ b/crates/switchyard-translation/src/codecs/responses/stream.rs @@ -154,6 +154,15 @@ fn encode_responses_stream( event: LlmResponseChunk, ) -> Vec { match event { + LlmResponseChunk::ProviderEvent { + source, + raw, + normalized: _, + } if source == WireFormat::OpenAiResponses.into() => vec![raw], + LlmResponseChunk::ProviderEvent { normalized, .. } => normalized + .into_iter() + .flat_map(|event| encode_responses_stream(state, event)) + .collect(), LlmResponseChunk::MessageStart { id, model } => { record_source_identity(state, id, model); ensure_responses_created(state) diff --git a/crates/switchyard-translation/src/codecs/stream.rs b/crates/switchyard-translation/src/codecs/stream.rs index 38cd73b9..4fec5dff 100644 --- a/crates/switchyard-translation/src/codecs/stream.rs +++ b/crates/switchyard-translation/src/codecs/stream.rs @@ -214,6 +214,21 @@ pub fn decode_stream_event( }) } +/// Decodes one provider stream event and retains its exact source JSON. +pub fn decode_stream_event_preserving( + state: &mut StreamTranslationState, + source: impl Into, + event: &Value, +) -> LlmResponseChunk { + let source = source.into(); + state.source = Some(source.clone()); + LlmResponseChunk::ProviderEvent { + source: source.clone(), + raw: event.clone(), + normalized: decode_stream_event(state, source, event), + } +} + /// Encodes one neutral stream event with the built-in codec registry. pub fn encode_stream_event( state: &mut StreamTranslationState, diff --git a/crates/switchyard-translation/src/engine.rs b/crates/switchyard-translation/src/engine.rs index 5303c3ed..d1e88176 100644 --- a/crates/switchyard-translation/src/engine.rs +++ b/crates/switchyard-translation/src/engine.rs @@ -18,6 +18,7 @@ use crate::error::{Result, TranslationError}; use crate::format::FormatId; use crate::llm::{AggLlmResponse, LlmRequest}; use crate::policy::TranslationPolicy; +use crate::LlmResponseChunk; /// Encoded translation result with any diagnostics emitted along the way. #[derive(Debug)] @@ -243,6 +244,45 @@ impl TranslationEngine { .collect()) } + /// Decodes one provider event while retaining the exact source JSON. + pub fn decode_stream_event( + &self, + state: &mut StreamTranslationState, + source: impl Into, + event: &Value, + ) -> Result { + let source = source.into(); + let source_codec = self.stream_registry.codec(source.clone())?; + state.source = Some(source.clone()); + Ok(LlmResponseChunk::ProviderEvent { + source, + raw: event.clone(), + normalized: source_codec.decode_event(state, event), + }) + } + + /// Encodes one neutral or preserved stream event for a target provider. + /// + /// A preserved event is replayed exactly when its source and target formats + /// match. Cross-format encoding intentionally uses only its normalized + /// events. + pub fn encode_stream_event( + &self, + state: &mut StreamTranslationState, + target: impl Into, + event: LlmResponseChunk, + ) -> Result> { + let target = target.into(); + let target_codec = self.stream_registry.codec(target.clone())?; + state.target = Some(target.clone()); + Ok(encode_stream_chunk( + state, + target_codec.as_ref(), + &target, + event, + )) + } + /// Finishes target-provider stream emission after the source stream closes. pub fn finish_stream( &self, @@ -255,6 +295,26 @@ impl TranslationEngine { } } +fn encode_stream_chunk( + state: &mut StreamTranslationState, + target_codec: &dyn crate::codecs::stream::StreamCodec, + target: &FormatId, + event: LlmResponseChunk, +) -> Vec { + match event { + LlmResponseChunk::ProviderEvent { + source, + raw, + normalized: _, + } if &source == target => vec![raw], + LlmResponseChunk::ProviderEvent { normalized, .. } => normalized + .into_iter() + .flat_map(|event| encode_stream_chunk(state, target_codec, target, event)) + .collect(), + event => target_codec.encode_event(state, event), + } +} + // Attaches source and target formats to every diagnostic emitted across both passes. fn with_formats( decoded: Vec, diff --git a/crates/switchyard-translation/tests/stream_translation.rs b/crates/switchyard-translation/tests/stream_translation.rs index 27bd9c41..fa7a8ea1 100644 --- a/crates/switchyard-translation/tests/stream_translation.rs +++ b/crates/switchyard-translation/tests/stream_translation.rs @@ -7,11 +7,95 @@ use pretty_assertions::assert_eq; use serde_json::json; use switchyard_protocol::{ResponseAccumulator, StopReason}; use switchyard_translation::{ - decode_stream_event, LlmResponseChunk, StreamTranslationState, TranslationEngine, WireFormat, + decode_stream_event, decode_stream_event_preserving, encode_stream_event, LlmResponseChunk, + StreamTranslationState, TranslationEngine, WireFormat, }; type TestResult = std::result::Result<(), Box>; +#[test] +fn preserved_same_format_events_replay_unknown_fields_exactly() -> TestResult { + let cases = [ + ( + WireFormat::OpenAiChat, + json!({ + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "model": "gpt-4o", + "system_fingerprint": "fp_provider_specific", + "choices": [{ + "index": 0, + "delta": {"content": "Hi"}, + "finish_reason": null + }] + }), + ), + ( + WireFormat::OpenAiResponses, + json!({ + "type": "response.output_text.delta", + "item_id": "item-1", + "output_index": 0, + "content_index": 0, + "delta": "Hi", + "sequence_number": 2, + "provider_extension": {"exact": true} + }), + ), + ( + WireFormat::AnthropicMessages, + json!({ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hi"}, + "provider_extension": {"exact": true} + }), + ), + ]; + + for (format, event) in cases { + let engine = TranslationEngine::default(); + let mut state = StreamTranslationState::new(format, format); + let preserved = engine.decode_stream_event(&mut state, format, &event)?; + let replayed = engine.encode_stream_event(&mut state, format, preserved)?; + assert_eq!(replayed, vec![event.clone()]); + + let mut state = StreamTranslationState::new(format, format); + let preserved = decode_stream_event_preserving(&mut state, format, &event); + let replayed = encode_stream_event(&mut state, format, preserved); + assert_eq!(replayed, vec![event]); + } + Ok(()) +} + +#[test] +fn preserved_cross_format_event_uses_normalized_content() -> TestResult { + let engine = TranslationEngine::default(); + let mut state = + StreamTranslationState::new(WireFormat::OpenAiChat, WireFormat::AnthropicMessages); + let event = json!({ + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "model": "gpt-4o", + "system_fingerprint": "fp_provider_specific", + "choices": [{ + "index": 0, + "delta": {"content": "Hi"}, + "finish_reason": null + }] + }); + + let preserved = engine.decode_stream_event(&mut state, WireFormat::OpenAiChat, &event)?; + let translated = + engine.encode_stream_event(&mut state, WireFormat::AnthropicMessages, preserved)?; + + assert_eq!(translated[2]["delta"]["text"], "Hi"); + assert!(translated + .iter() + .all(|event| event.get("system_fingerprint").is_none())); + Ok(()) +} + // Verifies an OpenAI text delta opens the expected Anthropic message and content blocks. #[test] fn openai_chat_stream_event_translates_to_anthropic_message_events() -> TestResult {