From 817d277b41daed13eb7e9fbbf5a87deb01e82757 Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Sat, 5 Sep 2026 10:46:46 -0700 Subject: [PATCH] fix(libsy): preserve provider events when escalation serves a buffered reply The escalation router calls the efficient tier, buffers the reply so the judge can read the completed turn, then serves that same reply when the judge declines. It rebuilt the served stream with AggLlmResponse::into_stream(), which is documented as lossy: it emits synthetic chunks and drops response extensions and preservation metadata. into_agg() consumed only each event's normalized chunks and discarded its preservation payload, so the payload was already gone by the time the response was rebuilt. Every unlatched escalation turn therefore reached the outbound codec without the provider bodies it uses for faithful same-format responses. Escalation is the only route that buffers, so no other route was affected. Add LlmResponse::into_agg_retaining_events(), which aggregates for the judge while optionally retaining the original events, and replay_stream_events(), which serves them verbatim. Escalation now replays the originals when the request is streaming and keeps returning the aggregate otherwise. Signed-off-by: Lin Jia --- crates/libsy/src/algorithms/escalation.rs | 18 +++- crates/protocol/src/stream.rs | 101 +++++++++++++++++++++- 2 files changed, 112 insertions(+), 7 deletions(-) diff --git a/crates/libsy/src/algorithms/escalation.rs b/crates/libsy/src/algorithms/escalation.rs index 593c22925..3c625dbb7 100644 --- a/crates/libsy/src/algorithms/escalation.rs +++ b/crates/libsy/src/algorithms/escalation.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use async_trait::async_trait; use switchyard_protocol::{ AggLlmResponse, LlmClientError, LlmResponse, Message, ModelId, Request, Response, Role, + replay_stream_events, }; use super::util::classifier_contract::ClassifierContractConfig; @@ -116,8 +117,17 @@ impl Classifier for EscalationClassifier { }; // The call resolves when its stream handle arrives; transport can still fail while // buffering. Fall back only for that availability failure and keep other errors typed. - let agg = match efficient_response.llm_response.into_agg().await { - Ok(agg) => agg, + // Retain the provider events: the buffered reply is served downstream when the judge + // declines, and replaying the originals keeps `preservation` intact. Rebuilding from the + // aggregate instead would emit synthetic chunks and drop the provider payloads the + // outbound codec needs, degrading every unlatched turn. + let streaming = request.llm_request.stream; + let (agg, retained_events) = match efficient_response + .llm_response + .into_agg_retaining_events(streaming) + .await + { + Ok(pair) => pair, Err(LlmClientError::Transport { .. }) => { return Ok((decisive(&self.capable), None)); } @@ -132,8 +142,8 @@ impl Classifier for EscalationClassifier { .messages .push(assistant_message(&agg)); let efficient_response = Response { - llm_response: if request.llm_request.stream { - LlmResponse::Stream(agg.into_stream()) + llm_response: if streaming { + LlmResponse::Stream(replay_stream_events(retained_events)) } else { LlmResponse::Agg(agg) }, diff --git a/crates/protocol/src/stream.rs b/crates/protocol/src/stream.rs index 6cf4eb6d9..bd144bc7b 100644 --- a/crates/protocol/src/stream.rs +++ b/crates/protocol/src/stream.rs @@ -126,6 +126,16 @@ impl From for LlmResponseStreamEvent { } } +/// Replays previously buffered events as a stream, preserving their provider payloads. +/// +/// Pair with [`LlmResponse::into_agg_retaining_events`] when a caller had to buffer a +/// response but must still serve the original bytes downstream. Unlike +/// [`AggLlmResponse::into_stream`], this is lossless: `preservation` survives, so the +/// outbound codec can emit a faithful same-format response instead of a synthetic one. +pub fn replay_stream_events(events: Vec) -> LlmResponseStream { + Box::pin(futures::stream::iter(events.into_iter().map(Ok))) +} + /// A model response: either a live [`Stream`](LlmResponse::Stream) of events or a /// terminal buffered [`LlmResponse::Agg`] response. /// @@ -155,16 +165,37 @@ impl LlmResponse { /// in-band [`LlmResponseChunk::DecodeError`] (as `ResponseTranslation`) or /// [`LlmResponseChunk::StreamError`] (as `UpstreamHttp`). pub async fn into_agg(self) -> Result { + let (agg, _) = self.into_agg_retaining_events(false).await?; + Ok(agg) + } + + /// Reduce to the buffered aggregate, optionally retaining the original provider events. + /// + /// [`Self::into_agg`] discards each event's `preservation` payload, so rebuilding a stream + /// from the aggregate via [`AggLlmResponse::into_stream`] can only emit synthetic chunks. + /// A caller that must buffer a response (to judge it) and then still serve it downstream + /// should set `retain_events` and replay the returned events with + /// [`replay_stream_events`], which preserves the provider payloads the outbound codec + /// needs for a faithful same-format response. + pub async fn into_agg_retaining_events( + self, + retain_events: bool, + ) -> Result<(AggLlmResponse, Vec), LlmClientError> { match self { - LlmResponse::Agg(agg) => Ok(agg), + LlmResponse::Agg(agg) => Ok((agg, Vec::new())), LlmResponse::Stream(mut stream) => { let mut accumulator = ResponseAccumulator::new(); + let mut retained = Vec::new(); while let Some(item) = stream.next().await { - for chunk in item?.normalized { + let event = item?; + if retain_events { + retained.push(event.clone()); + } + for chunk in event.normalized { push_checked_chunk(&mut accumulator, chunk)?; } } - Ok(accumulator.finish()) + Ok((accumulator.finish(), retained)) } } } @@ -557,6 +588,70 @@ mod tests { ); } + #[test] + fn retained_events_replay_with_preservation_intact() { + let raw = json!({ + "choices": [{"delta": {"content": "hello"}}], + "system_fingerprint": "fp_exact" + }); + let source_event = LlmResponseStreamEvent::preserved( + crate::WireFormat::OpenAiChat, + raw.clone(), + vec![LlmResponseChunk::TextDelta { + index: 0, + text: "hello".to_string(), + }], + ); + let response = LlmResponse::Stream(Box::pin(stream::iter([Ok(source_event)]))); + + let (aggregate, retained) = block_on(response.into_agg_retaining_events(true)) + .expect("stream should aggregate while retaining events"); + + // The judge still sees the aggregated reply. + assert_eq!( + aggregate.outputs[0].content, + vec![ContentBlock::Text { + text: "hello".to_string() + }] + ); + + // Replaying the retained events keeps the provider payload, which + // `AggLlmResponse::into_stream` would have dropped. + let replayed: Vec<_> = block_on(replay_stream_events(retained).collect::>()) + .into_iter() + .map(|item| item.expect("replayed event")) + .collect(); + assert_eq!(replayed.len(), 1); + assert_eq!(replayed[0].preservation().expect("preservation kept").raw(), &raw); + + // Contrast: the synthetic path loses it. + let synthetic: Vec<_> = block_on( + block_on( + LlmResponse::Stream(Box::pin(stream::iter([Ok( + LlmResponseStreamEvent::preserved( + crate::WireFormat::OpenAiChat, + raw.clone(), + vec![LlmResponseChunk::TextDelta { + index: 0, + text: "hello".to_string(), + }], + ), + )]))) + .into_agg(), + ) + .expect("aggregate") + .into_stream() + .collect::>(), + ) + .into_iter() + .map(|item| item.expect("synthetic event")) + .collect(); + assert!( + synthetic.iter().all(|event| event.preservation().is_none()), + "into_stream is expected to emit synthetic chunks without preservation" + ); + } + #[test] fn replacing_normalized_content_drops_preservation() { let event = LlmResponseStreamEvent::preserved(