switchyard_protocol/
stream.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Streaming half of the neutral IR: incremental response chunks ([`LlmResponseChunk`]),
5//! their stream envelope ([`LlmResponseStreamEvent`]), and the streamed response
6//! ([`LlmResponse`]) that carries either a live stream or the terminal [`AggLlmResponse`].
7
8use std::collections::BTreeMap;
9use std::pin::Pin;
10
11use futures::{Stream, StreamExt};
12use http::StatusCode;
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use thiserror::Error;
16
17use crate::{
18    LlmClientError,
19    format::FormatId,
20    llm::{AggLlmResponse, ContentBlock, ResponseOutput, Role, StopReason, ToolCall, Usage},
21};
22
23/// Status reported for an upstream error delivered inside a streaming body. The
24/// upstream already sent a success status line before failing, so there is no real
25/// code to propagate; 502 matches how a failed upstream call surfaces elsewhere.
26const MID_STREAM_UPSTREAM_STATUS: StatusCode = StatusCode::BAD_GATEWAY;
27
28/// Why a translated event stream stopped early.
29#[derive(Debug, Error)]
30pub enum LlmStreamError {
31    /// An upstream sse error
32    #[error("upstream stream error: {0}")]
33    Upstream(Value),
34
35    /// The stream itself failed
36    #[error(transparent)]
37    Client(#[from] LlmClientError),
38}
39
40/// A boxed, `Send` stream of response events. Each item may fail independently mid-stream.
41pub type LlmResponseStream =
42    Pin<Box<dyn Stream<Item = Result<LlmResponseStreamEvent, LlmClientError>> + Send>>;
43
44/// Parsed provider event retained for same-format replay.
45#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
46pub struct ProviderStreamEvent {
47    source: FormatId,
48    raw: Value,
49}
50
51impl ProviderStreamEvent {
52    /// Source wire format of the retained event.
53    pub fn source(&self) -> &FormatId {
54        &self.source
55    }
56
57    /// Parsed provider JSON retained for replay.
58    pub fn raw(&self) -> &Value {
59        &self.raw
60    }
61
62    /// Consumes the preservation value into its source format and parsed JSON.
63    pub fn into_parts(self) -> (FormatId, Value) {
64        (self.source, self.raw)
65    }
66}
67
68/// One streaming item crossing the host/algorithm boundary.
69///
70/// `normalized` contains only provider-neutral chunks. `preservation` is opaque to
71/// algorithms and is interpreted by `switchyard-translation` for same-format replay.
72#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
73pub struct LlmResponseStreamEvent {
74    preservation: Option<ProviderStreamEvent>,
75    normalized: Vec<LlmResponseChunk>,
76}
77
78impl LlmResponseStreamEvent {
79    /// Creates an event containing only normalized chunks.
80    pub fn new(normalized: Vec<LlmResponseChunk>) -> Self {
81        Self {
82            preservation: None,
83            normalized,
84        }
85    }
86
87    /// Creates an event with parsed provider JSON retained for replay.
88    pub fn preserved(
89        source: impl Into<FormatId>,
90        raw: Value,
91        normalized: Vec<LlmResponseChunk>,
92    ) -> Self {
93        Self {
94            preservation: Some(ProviderStreamEvent {
95                source: source.into(),
96                raw,
97            }),
98            normalized,
99        }
100    }
101
102    /// Retained provider event, when this event came directly from a provider.
103    pub fn preservation(&self) -> Option<&ProviderStreamEvent> {
104        self.preservation.as_ref()
105    }
106
107    /// Provider-neutral chunks carried by this event.
108    pub fn normalized(&self) -> &[LlmResponseChunk] {
109        &self.normalized
110    }
111
112    /// Consumes the event into its preservation and normalized content.
113    pub fn into_parts(self) -> (Option<ProviderStreamEvent>, Vec<LlmResponseChunk>) {
114        (self.preservation, self.normalized)
115    }
116
117    /// Replaces semantic content and drops raw replay data that no longer describes it.
118    pub fn replace_normalized(self, normalized: Vec<LlmResponseChunk>) -> Self {
119        Self::new(normalized)
120    }
121}
122
123impl From<LlmResponseChunk> for LlmResponseStreamEvent {
124    fn from(chunk: LlmResponseChunk) -> Self {
125        Self::new(vec![chunk])
126    }
127}
128
129/// Replays previously buffered events as a stream, preserving their provider payloads.
130///
131/// Pair with [`LlmResponse::into_agg_retaining_events`] when a caller had to buffer a
132/// response but must still serve the original bytes downstream. Unlike
133/// [`AggLlmResponse::into_stream`], this is lossless: `preservation` survives, so the
134/// outbound codec can emit a faithful same-format response instead of a synthetic one.
135pub fn replay_stream_events(events: Vec<LlmResponseStreamEvent>) -> LlmResponseStream {
136    Box::pin(futures::stream::iter(events.into_iter().map(Ok)))
137}
138
139/// A model response: either a live [`Stream`](LlmResponse::Stream) of events or a
140/// terminal buffered [`LlmResponse::Agg`] response.
141///
142/// Not `Clone` — the `Stream` variant owns a single-consumption stream. A buffered
143/// backend returns `Agg` directly; a streaming one returns `Stream` and the consumer
144/// drives it, folding to an [`AggLlmResponse`] when it needs the whole response.
145#[allow(clippy::large_enum_variant)]
146pub enum LlmResponse {
147    /// Live, single-consumption response stream.
148    Stream(LlmResponseStream),
149    /// Fully buffered terminal response.
150    Agg(AggLlmResponse),
151}
152
153impl LlmResponse {
154    /// Borrow the aggregate; `None` while this is still a stream.
155    pub fn as_agg(&self) -> Option<&AggLlmResponse> {
156        match self {
157            LlmResponse::Agg(agg) => Some(agg),
158            LlmResponse::Stream(_) => None,
159        }
160    }
161
162    /// Reduce to the buffered aggregate: return an `Agg` unchanged, or drive a `Stream`
163    /// to completion, folding its normalized chunks into an [`AggLlmResponse`] via
164    /// [`ResponseAccumulator`]. A stream item error aborts with `Err`, as does an
165    /// in-band [`LlmResponseChunk::DecodeError`] (as `ResponseTranslation`) or
166    /// [`LlmResponseChunk::StreamError`] (as `UpstreamHttp`).
167    pub async fn into_agg(self) -> Result<AggLlmResponse, LlmClientError> {
168        let (agg, _) = self.into_agg_retaining_events(false).await?;
169        Ok(agg)
170    }
171
172    /// Reduce to the buffered aggregate, optionally retaining the original provider events.
173    ///
174    /// [`Self::into_agg`] discards each event's `preservation` payload, so rebuilding a stream
175    /// from the aggregate via [`AggLlmResponse::into_stream`] can only emit synthetic chunks.
176    /// A caller that must buffer a response (to judge it) and then still serve it downstream
177    /// should set `retain_events` and replay the returned events with
178    /// [`replay_stream_events`], which preserves the provider payloads the outbound codec
179    /// needs for a faithful same-format response.
180    pub async fn into_agg_retaining_events(
181        self,
182        retain_events: bool,
183    ) -> Result<(AggLlmResponse, Vec<LlmResponseStreamEvent>), LlmClientError> {
184        match self {
185            LlmResponse::Agg(agg) => Ok((agg, Vec::new())),
186            LlmResponse::Stream(mut stream) => {
187                let mut accumulator = ResponseAccumulator::new();
188                let mut retained = Vec::new();
189                while let Some(item) = stream.next().await {
190                    let event = item?;
191                    if retain_events {
192                        retained.push(event.clone());
193                    }
194                    for chunk in event.normalized {
195                        push_checked_chunk(&mut accumulator, chunk)?;
196                    }
197                }
198                Ok((accumulator.finish(), retained))
199            }
200        }
201    }
202
203    /// Returns the model recorded by a buffered response.
204    ///
205    /// Returns `None` for a live stream because reading its model-bearing
206    /// [`LlmResponseChunk::MessageStart`] event would consume the stream.
207    pub fn selected_model(&self) -> Option<&str> {
208        match self {
209            LlmResponse::Agg(agg) => agg.model.as_deref(),
210            // TODO: How do we get the model name on a stream?
211            LlmResponse::Stream(_) => None,
212        }
213    }
214}
215
216impl AggLlmResponse {
217    /// Converts a fully-buffered response into a synthetic chunk stream.
218    ///
219    /// Useful when a caller had to buffer the response (e.g. to judge it) but the
220    /// downstream expects `LlmResponse::Stream` — for instance when `stream: true` was
221    /// requested and the algorithm had to aggregate before it could return.
222    ///
223    /// This conversion is lossy: only text, reasoning, and tool-call content has
224    /// a synthetic chunk representation. Refusals, tool results, media, files,
225    /// unknown blocks, response extensions, and preservation metadata are omitted.
226    pub fn into_stream(self) -> LlmResponseStream {
227        let mut chunks: Vec<LlmResponseChunk> = Vec::new();
228        chunks.push(LlmResponseChunk::MessageStart {
229            id: self.id,
230            model: self.model,
231        });
232        let mut tool_call_index = 0usize;
233        for (output_index, output) in self.outputs.into_iter().enumerate() {
234            for block in output.content {
235                match block {
236                    ContentBlock::Text { text } => {
237                        chunks.push(LlmResponseChunk::TextDelta {
238                            index: output_index,
239                            text,
240                        });
241                    }
242                    ContentBlock::Reasoning { text, details, .. } => {
243                        if !details.is_empty() {
244                            chunks.push(LlmResponseChunk::ReasoningDetailsDelta {
245                                index: output_index,
246                                details,
247                                text,
248                            });
249                        } else {
250                            chunks.push(LlmResponseChunk::ReasoningDelta {
251                                index: output_index,
252                                text,
253                            });
254                        }
255                    }
256                    ContentBlock::ToolCall(tool) => {
257                        let args = serde_json::to_string(&tool.arguments).unwrap_or_default();
258                        chunks.push(LlmResponseChunk::ToolCallDelta {
259                            index: tool_call_index,
260                            id: Some(tool.id),
261                            name: Some(tool.name),
262                            arguments_delta: Some(args),
263                        });
264                        tool_call_index += 1;
265                    }
266                    // Other variants (Image, ToolResult, Refusal, etc.) don't have
267                    // a streaming chunk representation and don't appear in assistant outputs.
268                    _ => {}
269                }
270            }
271            chunks.push(LlmResponseChunk::MessageStop {
272                reason: output.stop_reason.and_then(|r| {
273                    serde_json::to_value(r)
274                        .ok()
275                        .and_then(|v| v.as_str().map(String::from))
276                }),
277            });
278        }
279        chunks.push(LlmResponseChunk::Usage(self.usage));
280        Box::pin(futures::stream::iter(
281            chunks.into_iter().map(|chunk| Ok(chunk.into())),
282        ))
283    }
284}
285
286fn push_checked_chunk(
287    accumulator: &mut ResponseAccumulator,
288    chunk: LlmResponseChunk,
289) -> Result<(), LlmClientError> {
290    match chunk {
291        LlmResponseChunk::DecodeError { message } => {
292            Err(LlmClientError::ResponseTranslation(message))
293        }
294        LlmResponseChunk::StreamError { message } => Err(LlmClientError::UpstreamHttp {
295            status: MID_STREAM_UPSTREAM_STATUS,
296            body: message,
297        }),
298        chunk => {
299            accumulator.push(chunk);
300            Ok(())
301        }
302    }
303}
304
305/// One provider-neutral streaming response chunk.
306#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
307pub enum LlmResponseChunk {
308    /// Starts a response message.
309    MessageStart {
310        /// Provider response identifier.
311        id: Option<String>,
312        /// Model reported by the provider.
313        model: Option<String>,
314    },
315    /// Adds text to one output index.
316    TextDelta {
317        /// Provider output index.
318        index: usize,
319        /// Text fragment.
320        text: String,
321    },
322    /// Adds reasoning text to one output index.
323    ReasoningDelta {
324        /// Provider output index.
325        index: usize,
326        /// Reasoning fragment.
327        text: String,
328    },
329    /// Adds structured reasoning details to one output index.
330    ReasoningDetailsDelta {
331        /// Provider output index.
332        index: usize,
333        /// Reasoning detail objects in provider order.
334        details: Vec<Value>,
335        /// Normalized reasoning text represented by or accompanying the details.
336        text: String,
337    },
338    /// Adds or updates a tool call at one index.
339    ToolCallDelta {
340        /// Tool-call index within the response.
341        index: usize,
342        /// Tool-call identifier, normally supplied by the first delta.
343        id: Option<String>,
344        /// Tool name, normally supplied by the first delta.
345        name: Option<String>,
346        /// Fragment of the serialized tool arguments.
347        arguments_delta: Option<String>,
348    },
349    /// Reports token usage, normally near the end of the stream.
350    Usage(Usage),
351    /// Ends a response message.
352    MessageStop {
353        /// Provider stop-reason string before normalization.
354        reason: Option<String>,
355    },
356    /// Reports that an inbound stream event could not be decoded.
357    DecodeError {
358        /// Human-readable decoding failure.
359        message: String,
360    },
361    /// Reports an upstream failure delivered inside an otherwise successful stream.
362    StreamError {
363        /// Human-readable upstream failure.
364        message: String,
365    },
366}
367
368/// Folds a sequence of [`LlmResponseChunk`]s into the terminal [`AggLlmResponse`].
369///
370/// Text and reasoning deltas concatenate; tool-call deltas assemble by index (name,
371/// id, and a growing arguments string parsed as JSON at the end); `MessageStart`,
372/// `Usage`, and `MessageStop` set the corresponding fields.
373/// [`DecodeError`](LlmResponseChunk::DecodeError) and
374/// [`StreamError`](LlmResponseChunk::StreamError) chunks are ignored here; use
375/// [`LlmResponse::into_agg`] when they must become errors.
376///
377/// Folding is lossy for multiple outputs: text and reasoning indices are ignored,
378/// and [`finish`](Self::finish) produces one assistant output. Prefer
379/// [`LlmResponse::into_agg`] when stream errors must be surfaced.
380///
381/// Drive it by `push`-ing each chunk in order, then call [`finish`](Self::finish).
382#[derive(Default)]
383pub struct ResponseAccumulator {
384    id: Option<String>,
385    model: Option<String>,
386    text: String,
387    reasoning: Option<String>,
388    reasoning_details: Vec<Value>,
389    tool_calls: BTreeMap<usize, PartialToolCall>,
390    usage: Usage,
391    stop_reason: Option<StopReason>,
392}
393
394/// A tool call being assembled from streamed [`LlmResponseChunk::ToolCallDelta`]s.
395#[derive(Default)]
396struct PartialToolCall {
397    id: Option<String>,
398    name: Option<String>,
399    arguments: String,
400}
401
402impl ResponseAccumulator {
403    /// A fresh accumulator with no chunks applied.
404    pub fn new() -> Self {
405        Self::default()
406    }
407
408    /// Apply one chunk. Later `MessageStart`/`Usage`/`MessageStop` fields overwrite
409    /// earlier ones; text, reasoning, and tool-call arguments append.
410    pub fn push(&mut self, chunk: LlmResponseChunk) {
411        match chunk {
412            LlmResponseChunk::MessageStart { id, model } => {
413                if id.is_some() {
414                    self.id = id;
415                }
416                if model.is_some() {
417                    self.model = model;
418                }
419            }
420            LlmResponseChunk::TextDelta { text, .. } => self.text.push_str(&text),
421            LlmResponseChunk::ReasoningDelta { text, .. } => {
422                self.reasoning
423                    .get_or_insert_with(String::new)
424                    .push_str(&text);
425            }
426            LlmResponseChunk::ReasoningDetailsDelta { details, text, .. } => {
427                self.reasoning_details.extend(details);
428                if !text.is_empty() {
429                    self.reasoning
430                        .get_or_insert_with(String::new)
431                        .push_str(&text);
432                }
433            }
434            LlmResponseChunk::ToolCallDelta {
435                index,
436                id,
437                name,
438                arguments_delta,
439            } => {
440                let call = self.tool_calls.entry(index).or_default();
441                if id.is_some() {
442                    call.id = id;
443                }
444                if name.is_some() {
445                    call.name = name;
446                }
447                if let Some(delta) = arguments_delta {
448                    call.arguments.push_str(&delta);
449                }
450            }
451            LlmResponseChunk::Usage(usage) => self.usage = usage,
452            LlmResponseChunk::MessageStop { reason } => {
453                self.stop_reason = Some(stop_reason_from_str(reason.as_deref()));
454            }
455            LlmResponseChunk::DecodeError { .. } | LlmResponseChunk::StreamError { .. } => {}
456        }
457    }
458
459    /// Build the buffered response. Content is ordered reasoning, then text, then
460    /// tool calls (by ascending delta index) — a single assistant output.
461    pub fn finish(self) -> AggLlmResponse {
462        let mut content = Vec::new();
463        if self.reasoning.is_some() || !self.reasoning_details.is_empty() {
464            content.push(ContentBlock::Reasoning {
465                text: self.reasoning.unwrap_or_default(),
466                signature: None,
467                details: self.reasoning_details,
468            });
469        }
470        if !self.text.is_empty() {
471            content.push(ContentBlock::Text { text: self.text });
472        }
473        for call in self.tool_calls.into_values() {
474            content.push(ContentBlock::ToolCall(ToolCall {
475                id: call.id.unwrap_or_default(),
476                name: call.name.unwrap_or_default(),
477                arguments: parse_tool_arguments(&call.arguments),
478            }));
479        }
480        AggLlmResponse {
481            id: self.id,
482            model: self.model,
483            outputs: vec![ResponseOutput {
484                role: Role::Assistant,
485                content,
486                stop_reason: self.stop_reason,
487            }],
488            usage: self.usage,
489            ..AggLlmResponse::default()
490        }
491    }
492}
493
494/// Parse an assembled tool-call arguments string as JSON, falling back to a JSON
495/// string when it is not valid JSON and to an empty object when it is empty.
496fn parse_tool_arguments(arguments: &str) -> Value {
497    if arguments.is_empty() {
498        return Value::Object(serde_json::Map::new());
499    }
500    serde_json::from_str(arguments).unwrap_or_else(|_| Value::String(arguments.to_string()))
501}
502
503/// Map a provider stop-reason string (as carried by [`LlmResponseChunk::MessageStop`])
504/// to a normalized [`StopReason`], covering the common OpenAI and Anthropic spellings.
505fn stop_reason_from_str(reason: Option<&str>) -> StopReason {
506    match reason {
507        Some("length" | "max_tokens") => StopReason::MaxTokens,
508        Some("tool_calls" | "function_call" | "tool_use") => StopReason::ToolUse,
509        Some("content_filter") => StopReason::ContentFilter,
510        Some("stop" | "end_turn" | "stop_sequence") | None => StopReason::EndTurn,
511        Some(_) => StopReason::Unknown,
512    }
513}
514
515#[cfg(test)]
516mod tests {
517    use futures::executor::block_on;
518    use futures::stream;
519
520    use super::*;
521    use serde_json::json;
522
523    fn fold(chunks: Vec<LlmResponseChunk>) -> AggLlmResponse {
524        let mut accumulator = ResponseAccumulator::new();
525        for chunk in chunks {
526            accumulator.push(chunk);
527        }
528        accumulator.finish()
529    }
530
531    #[test]
532    fn folds_text_usage_and_stop_reason() {
533        let agg = fold(vec![
534            LlmResponseChunk::MessageStart {
535                id: Some("id1".to_string()),
536                model: Some("m".to_string()),
537            },
538            LlmResponseChunk::TextDelta {
539                index: 0,
540                text: "Hel".to_string(),
541            },
542            LlmResponseChunk::TextDelta {
543                index: 0,
544                text: "lo".to_string(),
545            },
546            LlmResponseChunk::Usage(Usage {
547                output_tokens: Some(2),
548                ..Usage::default()
549            }),
550            LlmResponseChunk::MessageStop {
551                reason: Some("length".to_string()),
552            },
553        ]);
554        assert_eq!(agg.id.as_deref(), Some("id1"));
555        assert_eq!(agg.model.as_deref(), Some("m"));
556        assert_eq!(agg.usage.output_tokens, Some(2));
557        assert_eq!(agg.outputs[0].stop_reason, Some(StopReason::MaxTokens));
558        assert_eq!(
559            agg.outputs[0].content,
560            vec![ContentBlock::Text {
561                text: "Hello".to_string()
562            }]
563        );
564    }
565
566    #[test]
567    fn aggregates_normalized_chunks_inside_stream_event() {
568        let response = LlmResponse::Stream(Box::pin(stream::iter([Ok(
569            LlmResponseStreamEvent::preserved(
570                crate::WireFormat::OpenAiChat,
571                json!({
572                "choices": [{"delta": {"content": "hello"}}],
573                "system_fingerprint": "fp_exact"
574                }),
575                vec![LlmResponseChunk::TextDelta {
576                    index: 0,
577                    text: "hello".to_string(),
578                }],
579            ),
580        )])));
581        let aggregate = block_on(response.into_agg()).expect("stream event should aggregate");
582
583        assert_eq!(
584            aggregate.outputs[0].content,
585            vec![ContentBlock::Text {
586                text: "hello".to_string()
587            }]
588        );
589    }
590
591    #[test]
592    fn retained_events_replay_with_preservation_intact() {
593        let raw = json!({
594            "choices": [{"delta": {"content": "hello"}}],
595            "system_fingerprint": "fp_exact"
596        });
597        let source_event = LlmResponseStreamEvent::preserved(
598            crate::WireFormat::OpenAiChat,
599            raw.clone(),
600            vec![LlmResponseChunk::TextDelta {
601                index: 0,
602                text: "hello".to_string(),
603            }],
604        );
605        let response = LlmResponse::Stream(Box::pin(stream::iter([Ok(source_event)])));
606
607        let (aggregate, retained) = block_on(response.into_agg_retaining_events(true))
608            .expect("stream should aggregate while retaining events");
609
610        // The judge still sees the aggregated reply.
611        assert_eq!(
612            aggregate.outputs[0].content,
613            vec![ContentBlock::Text {
614                text: "hello".to_string()
615            }]
616        );
617
618        // Replaying the retained events keeps the provider payload, which
619        // `AggLlmResponse::into_stream` would have dropped.
620        let replayed: Vec<_> = block_on(replay_stream_events(retained).collect::<Vec<_>>())
621            .into_iter()
622            .map(|item| item.expect("replayed event"))
623            .collect();
624        assert_eq!(replayed.len(), 1);
625        assert_eq!(replayed[0].preservation().expect("preservation kept").raw(), &raw);
626
627        // Contrast: the synthetic path loses it.
628        let synthetic: Vec<_> = block_on(
629            block_on(
630                LlmResponse::Stream(Box::pin(stream::iter([Ok(
631                    LlmResponseStreamEvent::preserved(
632                        crate::WireFormat::OpenAiChat,
633                        raw.clone(),
634                        vec![LlmResponseChunk::TextDelta {
635                            index: 0,
636                            text: "hello".to_string(),
637                        }],
638                    ),
639                )])))
640                .into_agg(),
641            )
642            .expect("aggregate")
643            .into_stream()
644            .collect::<Vec<_>>(),
645        )
646        .into_iter()
647        .map(|item| item.expect("synthetic event"))
648        .collect();
649        assert!(
650            synthetic.iter().all(|event| event.preservation().is_none()),
651            "into_stream is expected to emit synthetic chunks without preservation"
652        );
653    }
654
655    #[test]
656    fn replacing_normalized_content_drops_preservation() {
657        let event = LlmResponseStreamEvent::preserved(
658            crate::WireFormat::OpenAiChat,
659            json!({"choices": [{"delta": {"content": "old"}}]}),
660            vec![LlmResponseChunk::TextDelta {
661                index: 0,
662                text: "old".to_string(),
663            }],
664        )
665        .replace_normalized(vec![LlmResponseChunk::TextDelta {
666            index: 0,
667            text: "new".to_string(),
668        }]);
669
670        assert!(event.preservation().is_none());
671        assert_eq!(
672            event.normalized(),
673            &[LlmResponseChunk::TextDelta {
674                index: 0,
675                text: "new".to_string(),
676            }]
677        );
678    }
679
680    #[test]
681    fn stream_errors_inside_preserved_events_remain_typed() {
682        let response = LlmResponse::Stream(Box::pin(stream::iter([Ok(
683            LlmResponseStreamEvent::preserved(
684                crate::WireFormat::OpenAiChat,
685                json!({"error": {"message": "provider failed"}}),
686                vec![LlmResponseChunk::StreamError {
687                    message: "provider failed".to_string(),
688                }],
689            ),
690        )])));
691
692        let error = block_on(response.into_agg()).err();
693        assert!(matches!(
694            error,
695            Some(LlmClientError::UpstreamHttp {
696                status: MID_STREAM_UPSTREAM_STATUS,
697                ..
698            })
699        ));
700    }
701
702    #[test]
703    fn assembles_tool_calls_by_index() {
704        // id/name arrive once, arguments stream across deltas and parse as JSON.
705        let agg = fold(vec![
706            LlmResponseChunk::ToolCallDelta {
707                index: 0,
708                id: Some("call_1".to_string()),
709                name: Some("lookup".to_string()),
710                arguments_delta: Some("{\"q\":".to_string()),
711            },
712            LlmResponseChunk::ToolCallDelta {
713                index: 0,
714                id: None,
715                name: None,
716                arguments_delta: Some("\"rust\"}".to_string()),
717            },
718            LlmResponseChunk::MessageStop {
719                reason: Some("tool_calls".to_string()),
720            },
721        ]);
722        assert_eq!(agg.outputs[0].stop_reason, Some(StopReason::ToolUse));
723        assert_eq!(
724            agg.outputs[0].content,
725            vec![ContentBlock::ToolCall(ToolCall {
726                id: "call_1".to_string(),
727                name: "lookup".to_string(),
728                arguments: json!({"q": "rust"}),
729            })]
730        );
731    }
732
733    #[test]
734    fn reasoning_precedes_text_in_content() {
735        let agg = fold(vec![
736            LlmResponseChunk::ReasoningDelta {
737                index: 0,
738                text: "think".to_string(),
739            },
740            LlmResponseChunk::TextDelta {
741                index: 0,
742                text: "answer".to_string(),
743            },
744        ]);
745        assert_eq!(
746            agg.outputs[0].content,
747            vec![
748                ContentBlock::Reasoning {
749                    text: "think".to_string(),
750                    signature: None,
751                    details: Vec::new(),
752                },
753                ContentBlock::Text {
754                    text: "answer".to_string(),
755                },
756            ]
757        );
758    }
759
760    #[test]
761    fn into_stream_round_trips_through_into_agg() {
762        let original = AggLlmResponse {
763            id: Some("id1".to_string()),
764            model: Some("m".to_string()),
765            outputs: vec![ResponseOutput {
766                role: Role::Assistant,
767                content: vec![ContentBlock::Text {
768                    text: "hello".to_string(),
769                }],
770                stop_reason: Some(StopReason::EndTurn),
771            }],
772            usage: Usage {
773                output_tokens: Some(3),
774                ..Usage::default()
775            },
776            ..AggLlmResponse::default()
777        };
778        let stream = LlmResponse::Stream(original.clone().into_stream());
779        let recovered = block_on(stream.into_agg()).expect("into_agg failed");
780        assert_eq!(recovered.id, original.id);
781        assert_eq!(recovered.model, original.model);
782        assert_eq!(recovered.usage.output_tokens, original.usage.output_tokens);
783        assert_eq!(
784            recovered.outputs[0].stop_reason,
785            original.outputs[0].stop_reason
786        );
787        assert_eq!(recovered.outputs[0].content, original.outputs[0].content);
788    }
789
790    #[test]
791    fn into_stream_retains_encrypted_reasoning_and_text() {
792        let details = vec![json!({
793            "type": "reasoning.encrypted",
794            "data": "opaque-encrypted-reasoning"
795        })];
796        let original = AggLlmResponse {
797            outputs: vec![ResponseOutput {
798                role: Role::Assistant,
799                content: vec![ContentBlock::Reasoning {
800                    text: "fallback reasoning".to_string(),
801                    signature: None,
802                    details: details.clone(),
803                }],
804                stop_reason: Some(StopReason::EndTurn),
805            }],
806            ..AggLlmResponse::default()
807        };
808
809        let recovered = block_on(LlmResponse::Stream(original.into_stream()).into_agg())
810            .expect("into_agg failed");
811        let ContentBlock::Reasoning {
812            text,
813            details: recovered_details,
814            ..
815        } = &recovered.outputs[0].content[0]
816        else {
817            panic!("expected reasoning block");
818        };
819        assert_eq!(text, "fallback reasoning");
820        assert_eq!(recovered_details, &details);
821    }
822
823    #[test]
824    fn into_agg_preserves_stream_item_error() {
825        let response = LlmResponse::Stream(Box::pin(stream::once(async {
826            Err(LlmClientError::Timeout {
827                source: Box::new(std::io::Error::other("timed out")),
828            })
829        })));
830
831        let Err(error) = block_on(response.into_agg()) else {
832            panic!("expected stream aggregation to fail");
833        };
834        assert!(matches!(error, LlmClientError::Timeout { .. }));
835    }
836}
switchyard_protocol/stream.rs

switchyard_protocol/
stream.rs
