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

switchyard_protocol/
stream.rs
