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/// A model response: either a live [`Stream`](LlmResponse::Stream) of events or a
130/// terminal buffered [`LlmResponse::Agg`] response.
131///
132/// Not `Clone` — the `Stream` variant owns a single-consumption stream. A buffered
133/// backend returns `Agg` directly; a streaming one returns `Stream` and the consumer
134/// drives it, folding to an [`AggLlmResponse`] when it needs the whole response.
135#[allow(clippy::large_enum_variant)]
136pub enum LlmResponse {
137    /// Live, single-consumption response stream.
138    Stream(LlmResponseStream),
139    /// Fully buffered terminal response.
140    Agg(AggLlmResponse),
141}
142
143impl LlmResponse {
144    /// Borrow the aggregate; `None` while this is still a stream.
145    pub fn as_agg(&self) -> Option<&AggLlmResponse> {
146        match self {
147            LlmResponse::Agg(agg) => Some(agg),
148            LlmResponse::Stream(_) => None,
149        }
150    }
151
152    /// Reduce to the buffered aggregate: return an `Agg` unchanged, or drive a `Stream`
153    /// to completion, folding its normalized chunks into an [`AggLlmResponse`] via
154    /// [`ResponseAccumulator`]. A stream item error aborts with `Err`, as does an
155    /// in-band [`LlmResponseChunk::DecodeError`] (as `ResponseTranslation`) or
156    /// [`LlmResponseChunk::StreamError`] (as `UpstreamHttp`).
157    pub async fn into_agg(self) -> Result<AggLlmResponse, LlmClientError> {
158        match self {
159            LlmResponse::Agg(agg) => Ok(agg),
160            LlmResponse::Stream(mut stream) => {
161                let mut accumulator = ResponseAccumulator::new();
162                while let Some(item) = stream.next().await {
163                    for chunk in item?.normalized {
164                        push_checked_chunk(&mut accumulator, chunk)?;
165                    }
166                }
167                Ok(accumulator.finish())
168            }
169        }
170    }
171
172    /// Returns the model recorded by a buffered response.
173    ///
174    /// Returns `None` for a live stream because reading its model-bearing
175    /// [`LlmResponseChunk::MessageStart`] event would consume the stream.
176    pub fn selected_model(&self) -> Option<&str> {
177        match self {
178            LlmResponse::Agg(agg) => agg.model.as_deref(),
179            // TODO: How do we get the model name on a stream?
180            LlmResponse::Stream(_) => None,
181        }
182    }
183}
184
185/// An encrypted reasoning detail that names its provider item id but carries no payload yet.
186fn is_reasoning_id_announcement(detail: &Value) -> bool {
187    detail.get("type").and_then(Value::as_str) == Some("reasoning.encrypted")
188        && detail.get("data").is_none()
189}
190
191/// Appends a reasoning detail to an accumulated block. A Responses stream decoder announces a
192/// reasoning item's provider id (`{"type": "reasoning.encrypted", "id"}`) before the payload
193/// arrives; the payload detail then replaces that announcement so history holds one detail.
194fn push_reasoning_detail(details: &mut Vec<Value>, detail: Value) {
195    if detail.get("type").and_then(Value::as_str) == Some("reasoning.encrypted")
196        && let Some(id) = detail.get("id").and_then(Value::as_str)
197        && let Some(announcement) = details.iter_mut().find(|existing| {
198            is_reasoning_id_announcement(existing)
199                && existing.get("id").and_then(Value::as_str) == Some(id)
200        })
201    {
202        *announcement = detail;
203        return;
204    }
205    details.push(detail);
206}
207
208impl AggLlmResponse {
209    /// Converts a fully-buffered response into a synthetic chunk stream.
210    ///
211    /// Useful when a caller had to buffer the response (e.g. to judge it) but the
212    /// downstream expects `LlmResponse::Stream` — for instance when `stream: true` was
213    /// requested and the algorithm had to aggregate before it could return.
214    ///
215    /// This conversion is lossy: only text, reasoning, and tool-call content has
216    /// a synthetic chunk representation. Refusals, tool results, media, files,
217    /// MCP activity, unknown blocks, response extensions, and preservation metadata are omitted.
218    pub fn into_stream(self) -> LlmResponseStream {
219        let mut chunks: Vec<LlmResponseChunk> = Vec::new();
220        chunks.push(LlmResponseChunk::MessageStart {
221            id: self.id,
222            model: self.model,
223        });
224        let mut tool_call_index = 0usize;
225        for (output_index, output) in self.outputs.into_iter().enumerate() {
226            for block in output.content {
227                match block {
228                    ContentBlock::Text { text, .. } => {
229                        chunks.push(LlmResponseChunk::TextDelta {
230                            index: output_index,
231                            text,
232                        });
233                    }
234                    ContentBlock::Reasoning { text, details, .. } => {
235                        if !details.is_empty() {
236                            chunks.push(LlmResponseChunk::ReasoningDetailsDelta {
237                                index: output_index,
238                                details,
239                                text,
240                            });
241                        } else {
242                            chunks.push(LlmResponseChunk::ReasoningDelta {
243                                index: output_index,
244                                text,
245                            });
246                        }
247                    }
248                    ContentBlock::ToolCall(tool) => {
249                        let args = serde_json::to_string(&tool.arguments).unwrap_or_default();
250                        chunks.push(LlmResponseChunk::ToolCallDelta {
251                            index: tool_call_index,
252                            id: Some(tool.id),
253                            name: Some(tool.name),
254                            arguments_delta: Some(args),
255                        });
256                        tool_call_index += 1;
257                    }
258                    // Other variants, including backend MCP activity, have no synthetic chunks.
259                    _ => {}
260                }
261            }
262            chunks.push(LlmResponseChunk::MessageStop {
263                reason: output.stop_reason.and_then(|r| {
264                    serde_json::to_value(r)
265                        .ok()
266                        .and_then(|v| v.as_str().map(String::from))
267                }),
268            });
269        }
270        chunks.push(LlmResponseChunk::Usage(self.usage));
271        Box::pin(futures::stream::iter(
272            chunks.into_iter().map(|chunk| Ok(chunk.into())),
273        ))
274    }
275}
276
277fn push_checked_chunk(
278    accumulator: &mut ResponseAccumulator,
279    chunk: LlmResponseChunk,
280) -> Result<(), LlmClientError> {
281    match chunk {
282        LlmResponseChunk::DecodeError { message } => {
283            Err(LlmClientError::ResponseTranslation(message))
284        }
285        LlmResponseChunk::StreamError { message } => Err(LlmClientError::UpstreamHttp {
286            status: MID_STREAM_UPSTREAM_STATUS,
287            body: message,
288        }),
289        chunk => {
290            accumulator.push(chunk);
291            Ok(())
292        }
293    }
294}
295
296/// One provider-neutral streaming response chunk.
297#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
298pub enum LlmResponseChunk {
299    /// Starts a response message, or fills in identity first supplied by a later frame.
300    MessageStart {
301        /// Provider response identifier.
302        id: Option<String>,
303        /// Model reported by the provider.
304        model: Option<String>,
305    },
306    /// Adds text to one output index.
307    TextDelta {
308        /// Provider output index.
309        index: usize,
310        /// Text fragment.
311        text: String,
312    },
313    /// Adds reasoning text to one output index.
314    ReasoningDelta {
315        /// Provider output index.
316        index: usize,
317        /// Reasoning fragment.
318        text: String,
319    },
320    /// Adds structured reasoning details to one output index.
321    ReasoningDetailsDelta {
322        /// Provider output index.
323        index: usize,
324        /// Reasoning detail objects in provider order.
325        details: Vec<Value>,
326        /// Normalized reasoning text represented by or accompanying the details.
327        text: String,
328    },
329    /// Adds or updates a tool call at one index.
330    ToolCallDelta {
331        /// Tool-call index within the response.
332        index: usize,
333        /// Tool-call identifier, normally supplied by the first delta.
334        id: Option<String>,
335        /// Tool name, normally supplied by the first delta.
336        name: Option<String>,
337        /// Fragment of the serialized tool arguments.
338        arguments_delta: Option<String>,
339    },
340    /// Reports token usage, normally near the end of the stream.
341    Usage(Usage),
342    /// Ends a response message.
343    MessageStop {
344        /// Provider stop-reason string before normalization.
345        reason: Option<String>,
346    },
347    /// Reports that an inbound stream event could not be decoded.
348    DecodeError {
349        /// Human-readable decoding failure.
350        message: String,
351    },
352    /// Reports an upstream failure delivered inside an otherwise successful stream.
353    StreamError {
354        /// Human-readable upstream failure.
355        message: String,
356    },
357}
358
359/// Folds a sequence of [`LlmResponseChunk`]s into the terminal [`AggLlmResponse`].
360///
361/// Text and reasoning deltas concatenate; tool-call deltas assemble by index (name,
362/// id, and a growing arguments string parsed as JSON at the end); `MessageStart`,
363/// `Usage`, and `MessageStop` set the corresponding fields.
364/// [`DecodeError`](LlmResponseChunk::DecodeError) and
365/// [`StreamError`](LlmResponseChunk::StreamError) chunks are ignored here; use
366/// [`LlmResponse::into_agg`] when they must become errors.
367///
368/// Folding is lossy for multiple outputs: text and reasoning indices are ignored,
369/// and [`finish`](Self::finish) produces one assistant output. Prefer
370/// [`LlmResponse::into_agg`] when stream errors must be surfaced.
371///
372/// Drive it by `push`-ing each chunk in order, then call [`finish`](Self::finish).
373#[derive(Default)]
374pub struct ResponseAccumulator {
375    id: Option<String>,
376    model: Option<String>,
377    text: String,
378    reasoning: Option<String>,
379    reasoning_details: Vec<Value>,
380    tool_calls: BTreeMap<usize, PartialToolCall>,
381    usage: Usage,
382    stop_reason: Option<StopReason>,
383}
384
385/// A tool call being assembled from streamed [`LlmResponseChunk::ToolCallDelta`]s.
386#[derive(Default)]
387struct PartialToolCall {
388    id: Option<String>,
389    name: Option<String>,
390    arguments: String,
391}
392
393impl ResponseAccumulator {
394    /// A fresh accumulator with no chunks applied.
395    pub fn new() -> Self {
396        Self::default()
397    }
398
399    /// Apply one chunk. Later `MessageStart`/`Usage`/`MessageStop` fields overwrite
400    /// earlier ones; text, reasoning, and tool-call arguments append.
401    pub fn push(&mut self, chunk: LlmResponseChunk) {
402        match chunk {
403            LlmResponseChunk::MessageStart { id, model } => {
404                if id.is_some() {
405                    self.id = id;
406                }
407                if model.is_some() {
408                    self.model = model;
409                }
410            }
411            LlmResponseChunk::TextDelta { text, .. } => self.text.push_str(&text),
412            LlmResponseChunk::ReasoningDelta { text, .. } => {
413                self.reasoning
414                    .get_or_insert_with(String::new)
415                    .push_str(&text);
416            }
417            LlmResponseChunk::ReasoningDetailsDelta { details, text, .. } => {
418                for detail in details {
419                    push_reasoning_detail(&mut self.reasoning_details, detail);
420                }
421                if !text.is_empty() {
422                    self.reasoning
423                        .get_or_insert_with(String::new)
424                        .push_str(&text);
425                }
426            }
427            LlmResponseChunk::ToolCallDelta {
428                index,
429                id,
430                name,
431                arguments_delta,
432            } => {
433                let call = self.tool_calls.entry(index).or_default();
434                if id.is_some() {
435                    call.id = id;
436                }
437                if name.is_some() {
438                    call.name = name;
439                }
440                if let Some(delta) = arguments_delta {
441                    call.arguments.push_str(&delta);
442                }
443            }
444            LlmResponseChunk::Usage(usage) => self.usage = usage,
445            LlmResponseChunk::MessageStop { reason } => {
446                self.stop_reason = Some(stop_reason_from_str(reason.as_deref()));
447            }
448            LlmResponseChunk::DecodeError { .. } | LlmResponseChunk::StreamError { .. } => {}
449        }
450    }
451
452    /// Build the buffered response. Content is ordered reasoning, then text, then
453    /// tool calls (by ascending delta index) — a single assistant output.
454    pub fn finish(self) -> AggLlmResponse {
455        let mut content = Vec::new();
456        if self.reasoning.is_some() || !self.reasoning_details.is_empty() {
457            content.push(ContentBlock::Reasoning {
458                text: self.reasoning.unwrap_or_default(),
459                signature: None,
460                // An announcement whose payload never arrived is a stream-level hint only.
461                details: self
462                    .reasoning_details
463                    .into_iter()
464                    .filter(|detail| !is_reasoning_id_announcement(detail))
465                    .collect(),
466            });
467        }
468        if !self.text.is_empty() {
469            content.push(ContentBlock::Text {
470                responses: None,
471                text: self.text,
472            });
473        }
474        for call in self.tool_calls.into_values() {
475            content.push(ContentBlock::ToolCall(ToolCall {
476                id: call.id.unwrap_or_default(),
477                name: call.name.unwrap_or_default(),
478                arguments: parse_tool_arguments(&call.arguments),
479            }));
480        }
481        AggLlmResponse {
482            id: self.id,
483            model: self.model,
484            outputs: vec![ResponseOutput {
485                role: Role::Assistant,
486                content,
487                url_citations: Vec::new(),
488                stop_reason: self.stop_reason,
489            }],
490            usage: self.usage,
491            ..AggLlmResponse::default()
492        }
493    }
494}
495
496/// Parse an assembled tool-call arguments string as JSON, falling back to a JSON
497/// string when it is not valid JSON and to an empty object when it is empty.
498fn parse_tool_arguments(arguments: &str) -> Value {
499    if arguments.is_empty() {
500        return Value::Object(serde_json::Map::new());
501    }
502    serde_json::from_str(arguments).unwrap_or_else(|_| Value::String(arguments.to_string()))
503}
504
505/// Map a provider stop-reason string (as carried by [`LlmResponseChunk::MessageStop`])
506/// to a normalized [`StopReason`], covering the common OpenAI and Anthropic spellings.
507fn stop_reason_from_str(reason: Option<&str>) -> StopReason {
508    match reason {
509        Some("length" | "max_tokens") => StopReason::MaxTokens,
510        Some("tool_calls" | "function_call" | "tool_use") => StopReason::ToolUse,
511        // Anthropic spells a content-policy stop `refusal`; the buffered codec
512        // already maps it (anthropic/buffered.rs, map_anthropic_stop_reason).
513        Some("content_filter" | "refusal") => StopReason::ContentFilter,
514        Some("stop" | "end_turn" | "stop_sequence") | None => StopReason::EndTurn,
515        Some(_) => StopReason::Unknown,
516    }
517}
518
519#[cfg(test)]
520mod tests {
521    use futures::executor::block_on;
522    use futures::stream;
523
524    use super::*;
525    use serde_json::json;
526
527    fn fold(chunks: Vec<LlmResponseChunk>) -> AggLlmResponse {
528        let mut accumulator = ResponseAccumulator::new();
529        for chunk in chunks {
530            accumulator.push(chunk);
531        }
532        accumulator.finish()
533    }
534
535    #[test]
536    fn folds_text_usage_and_stop_reason() {
537        let agg = fold(vec![
538            LlmResponseChunk::MessageStart {
539                id: Some("id1".to_string()),
540                model: Some("m".to_string()),
541            },
542            LlmResponseChunk::TextDelta {
543                index: 0,
544                text: "Hel".to_string(),
545            },
546            LlmResponseChunk::TextDelta {
547                index: 0,
548                text: "lo".to_string(),
549            },
550            LlmResponseChunk::Usage(Usage {
551                output_tokens: Some(2),
552                ..Usage::default()
553            }),
554            LlmResponseChunk::MessageStop {
555                reason: Some("length".to_string()),
556            },
557        ]);
558        assert_eq!(agg.id.as_deref(), Some("id1"));
559        assert_eq!(agg.model.as_deref(), Some("m"));
560        assert_eq!(agg.usage.output_tokens, Some(2));
561        assert_eq!(agg.outputs[0].stop_reason, Some(StopReason::MaxTokens));
562        assert_eq!(
563            agg.outputs[0].content,
564            vec![ContentBlock::Text {
565                responses: None,
566                text: "Hello".to_string()
567            }]
568        );
569    }
570
571    #[test]
572    fn stop_reason_spellings_normalize_the_same_streamed_as_buffered() {
573        // A stop reason has to mean the same thing whether the response arrived
574        // buffered or was folded out of a stream. The expectations here are the
575        // ones the codecs apply to a buffered body.
576        for (reason, expected) in [
577            ("length", StopReason::MaxTokens),
578            ("max_tokens", StopReason::MaxTokens),
579            ("tool_calls", StopReason::ToolUse),
580            ("tool_use", StopReason::ToolUse),
581            ("content_filter", StopReason::ContentFilter),
582            ("refusal", StopReason::ContentFilter),
583            ("stop", StopReason::EndTurn),
584            ("end_turn", StopReason::EndTurn),
585        ] {
586            let agg = fold(vec![LlmResponseChunk::MessageStop {
587                reason: Some(reason.to_string()),
588            }]);
589
590            assert_eq!(
591                agg.outputs.first().and_then(|output| output.stop_reason),
592                Some(expected),
593                "stop reason {reason:?}"
594            );
595        }
596    }
597
598    #[test]
599    fn aggregates_normalized_chunks_inside_stream_event() {
600        let response = LlmResponse::Stream(Box::pin(stream::iter([Ok(
601            LlmResponseStreamEvent::preserved(
602                crate::WireFormat::OpenAiChat,
603                json!({
604                "choices": [{"delta": {"content": "hello"}}],
605                "system_fingerprint": "fp_exact"
606                }),
607                vec![LlmResponseChunk::TextDelta {
608                    index: 0,
609                    text: "hello".to_string(),
610                }],
611            ),
612        )])));
613        let aggregate = block_on(response.into_agg()).expect("stream event should aggregate");
614
615        assert_eq!(
616            aggregate.outputs[0].content,
617            vec![ContentBlock::Text {
618                responses: None,
619                text: "hello".to_string()
620            }]
621        );
622    }
623
624    #[test]
625    fn replacing_normalized_content_drops_preservation() {
626        let event = LlmResponseStreamEvent::preserved(
627            crate::WireFormat::OpenAiChat,
628            json!({"choices": [{"delta": {"content": "old"}}]}),
629            vec![LlmResponseChunk::TextDelta {
630                index: 0,
631                text: "old".to_string(),
632            }],
633        )
634        .replace_normalized(vec![LlmResponseChunk::TextDelta {
635            index: 0,
636            text: "new".to_string(),
637        }]);
638
639        assert!(event.preservation().is_none());
640        assert_eq!(
641            event.normalized(),
642            &[LlmResponseChunk::TextDelta {
643                index: 0,
644                text: "new".to_string(),
645            }]
646        );
647    }
648
649    #[test]
650    fn stream_errors_inside_preserved_events_remain_typed() {
651        let response = LlmResponse::Stream(Box::pin(stream::iter([Ok(
652            LlmResponseStreamEvent::preserved(
653                crate::WireFormat::OpenAiChat,
654                json!({"error": {"message": "provider failed"}}),
655                vec![LlmResponseChunk::StreamError {
656                    message: "provider failed".to_string(),
657                }],
658            ),
659        )])));
660
661        let error = block_on(response.into_agg()).err();
662        assert!(matches!(
663            error,
664            Some(LlmClientError::UpstreamHttp {
665                status: MID_STREAM_UPSTREAM_STATUS,
666                ..
667            })
668        ));
669    }
670
671    #[test]
672    fn assembles_tool_calls_by_index() {
673        // id/name arrive once, arguments stream across deltas and parse as JSON.
674        let agg = fold(vec![
675            LlmResponseChunk::ToolCallDelta {
676                index: 0,
677                id: Some("call_1".to_string()),
678                name: Some("lookup".to_string()),
679                arguments_delta: Some("{\"q\":".to_string()),
680            },
681            LlmResponseChunk::ToolCallDelta {
682                index: 0,
683                id: None,
684                name: None,
685                arguments_delta: Some("\"rust\"}".to_string()),
686            },
687            LlmResponseChunk::MessageStop {
688                reason: Some("tool_calls".to_string()),
689            },
690        ]);
691        assert_eq!(agg.outputs[0].stop_reason, Some(StopReason::ToolUse));
692        assert_eq!(
693            agg.outputs[0].content,
694            vec![ContentBlock::ToolCall(ToolCall {
695                id: "call_1".to_string(),
696                name: "lookup".to_string(),
697                arguments: json!({"q": "rust"}),
698            })]
699        );
700    }
701
702    #[test]
703    fn reasoning_precedes_text_in_content() {
704        let agg = fold(vec![
705            LlmResponseChunk::ReasoningDelta {
706                index: 0,
707                text: "think".to_string(),
708            },
709            LlmResponseChunk::TextDelta {
710                index: 0,
711                text: "answer".to_string(),
712            },
713        ]);
714        assert_eq!(
715            agg.outputs[0].content,
716            vec![
717                ContentBlock::Reasoning {
718                    text: "think".to_string(),
719                    signature: None,
720                    details: Vec::new(),
721                },
722                ContentBlock::Text {
723                    responses: None,
724                    text: "answer".to_string(),
725                },
726            ]
727        );
728    }
729
730    #[test]
731    fn into_stream_round_trips_through_into_agg() {
732        let original = AggLlmResponse {
733            id: Some("id1".to_string()),
734            model: Some("m".to_string()),
735            outputs: vec![ResponseOutput {
736                role: Role::Assistant,
737                content: vec![ContentBlock::Text {
738                    responses: None,
739                    text: "hello".to_string(),
740                }],
741                url_citations: Vec::new(),
742                stop_reason: Some(StopReason::EndTurn),
743            }],
744            usage: Usage {
745                output_tokens: Some(3),
746                ..Usage::default()
747            },
748            ..AggLlmResponse::default()
749        };
750        let stream = LlmResponse::Stream(original.clone().into_stream());
751        let recovered = block_on(stream.into_agg()).expect("into_agg failed");
752        assert_eq!(recovered.id, original.id);
753        assert_eq!(recovered.model, original.model);
754        assert_eq!(recovered.usage.output_tokens, original.usage.output_tokens);
755        assert_eq!(
756            recovered.outputs[0].stop_reason,
757            original.outputs[0].stop_reason
758        );
759        assert_eq!(recovered.outputs[0].content, original.outputs[0].content);
760    }
761
762    #[test]
763    fn into_stream_retains_encrypted_reasoning_and_text() {
764        let details = vec![json!({
765            "type": "reasoning.encrypted",
766            "data": "opaque-encrypted-reasoning"
767        })];
768        let original = AggLlmResponse {
769            outputs: vec![ResponseOutput {
770                role: Role::Assistant,
771                content: vec![ContentBlock::Reasoning {
772                    text: "fallback reasoning".to_string(),
773                    signature: None,
774                    details: details.clone(),
775                }],
776                url_citations: Vec::new(),
777                stop_reason: Some(StopReason::EndTurn),
778            }],
779            ..AggLlmResponse::default()
780        };
781
782        let recovered = block_on(LlmResponse::Stream(original.into_stream()).into_agg())
783            .expect("into_agg failed");
784        let ContentBlock::Reasoning {
785            text,
786            details: recovered_details,
787            ..
788        } = &recovered.outputs[0].content[0]
789        else {
790            panic!("expected reasoning block");
791        };
792        assert_eq!(text, "fallback reasoning");
793        assert_eq!(recovered_details, &details);
794    }
795
796    #[test]
797    fn into_agg_preserves_stream_item_error() {
798        let response = LlmResponse::Stream(Box::pin(stream::once(async {
799            Err(LlmClientError::Timeout {
800                source: Box::new(std::io::Error::other("timed out")),
801            })
802        })));
803
804        let Err(error) = block_on(response.into_agg()) else {
805            panic!("expected stream aggregation to fail");
806        };
807        assert!(matches!(error, LlmClientError::Timeout { .. }));
808    }
809}
switchyard_protocol/stream.rs

switchyard_protocol/
stream.rs
