runifold_model/
stream.rs

1use std::collections::BTreeMap;
2
3use base64::{Engine as _, engine::general_purpose::STANDARD};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7use crate::{
8    ContentPart, DEFAULT_MAX_ARTIFACT_BYTES, FinishReason, MediaSource, ModelError, ModelErrorKind,
9    ModelRef, ModelResponse, ModelUsage, ModelWarning, ProviderData, ReasoningPart, ToolCall,
10};
11
12/// The type and initial metadata of a streamed content block.
13#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
14#[serde(tag = "type", rename_all = "snake_case")]
15#[non_exhaustive]
16pub enum ContentBlockKind {
17    /// Text output.
18    Text,
19    /// Reasoning or thinking output.
20    Reasoning {
21        /// Initial signature or continuation token.
22        signature: Option<String>,
23        /// Whether the provider redacted the reasoning body.
24        redacted: bool,
25    },
26    /// A streamed tool call.
27    ToolCall {
28        /// Provider- or runtime-assigned call identity.
29        id: String,
30        /// Tool name.
31        name: String,
32    },
33    /// A streamed refusal.
34    Refusal,
35    /// Streamed image output.
36    Image {
37        /// MIME type of the completed image.
38        media_type: String,
39    },
40    /// Streamed audio output.
41    Audio {
42        /// MIME type of the completed audio.
43        media_type: String,
44    },
45    /// Streamed document output.
46    Document {
47        /// MIME type of the completed document.
48        media_type: String,
49        /// Optional neutral display name.
50        name: Option<String>,
51    },
52}
53
54/// A provider event retained without normalization.
55#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
56pub struct ProviderEvent {
57    /// Provider namespace.
58    pub provider: String,
59    /// Provider event name.
60    pub name: String,
61    /// Original structured payload.
62    pub payload: Value,
63}
64
65/// Canonical events emitted by a streaming model call.
66#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
67#[serde(tag = "type", rename_all = "snake_case")]
68#[non_exhaustive]
69pub enum ModelStreamEvent {
70    /// The provider accepted the request and started a response.
71    ResponseStarted {
72        /// Provider response identity.
73        id: Option<String>,
74        /// Actual model serving the request.
75        model: ModelRef,
76    },
77    /// A delta-capable content block started.
78    ContentBlockStarted {
79        /// Stable output ordering index.
80        index: u32,
81        /// Block type and initial metadata.
82        kind: ContentBlockKind,
83    },
84    /// A text delta.
85    TextDelta {
86        /// Target block index.
87        index: u32,
88        /// Appended text.
89        text: String,
90    },
91    /// A reasoning-text delta.
92    ReasoningDelta {
93        /// Target block index.
94        index: u32,
95        /// Appended reasoning text.
96        text: String,
97    },
98    /// A reasoning-signature delta.
99    ReasoningSignatureDelta {
100        /// Target block index.
101        index: u32,
102        /// Appended signature data.
103        signature: String,
104    },
105    /// A raw JSON fragment for tool arguments.
106    ToolArgumentsDelta {
107        /// Target block index.
108        index: u32,
109        /// Appended raw JSON text.
110        json: String,
111    },
112    /// A refusal-text delta.
113    RefusalDelta {
114        /// Target block index.
115        index: u32,
116        /// Appended refusal text.
117        text: String,
118    },
119    /// One independently base64-encoded binary media chunk.
120    BinaryDelta {
121        /// Target image, audio, or document block index.
122        index: u32,
123        /// Base64-encoded chunk bytes.
124        data: String,
125    },
126    /// A delta-capable block completed.
127    ContentBlockCompleted {
128        /// Completed block index.
129        index: u32,
130    },
131    /// A complete non-delta content part arrived.
132    ContentPartCompleted {
133        /// Stable output ordering index.
134        index: u32,
135        /// Completed content.
136        part: ContentPart,
137    },
138    /// A cumulative usage snapshot.
139    UsageUpdated {
140        /// Latest cumulative model usage.
141        usage: ModelUsage,
142    },
143    /// A translation or feature-degradation warning.
144    Warning {
145        /// Visible warning.
146        warning: ModelWarning,
147    },
148    /// A provider heartbeat without model content.
149    Heartbeat,
150    /// An unknown or provider-specific event.
151    Provider {
152        /// Retained provider event.
153        event: ProviderEvent,
154    },
155    /// The response completed.
156    ResponseCompleted {
157        /// Normalized terminal reason.
158        finish_reason: FinishReason,
159        /// Namespaced terminal provider metadata.
160        provider_metadata: BTreeMap<String, Value>,
161    },
162    /// The provider's authoritative complete JSON argument text.
163    ToolArgumentsCompleted {
164        /// Target tool-call block index.
165        index: u32,
166        /// Complete raw JSON text, replacing any accumulated deltas.
167        json: String,
168    },
169    /// Namespaced metadata for an open or completed content block.
170    ///
171    /// Providers may learn final item metadata only after the delta-capable
172    /// block itself has completed. Accumulators therefore apply these updates
173    /// to either state without reopening the block.
174    ContentBlockMetadata {
175        /// Target block index.
176        index: u32,
177        /// Namespaced metadata to merge into the canonical content part.
178        metadata: BTreeMap<String, Value>,
179    },
180}
181
182#[derive(Debug)]
183enum PartialBlock {
184    Text(String),
185    Reasoning {
186        text: String,
187        signature: Option<String>,
188        redacted: bool,
189    },
190    ToolCall {
191        id: String,
192        name: String,
193        arguments: String,
194        metadata: BTreeMap<String, Value>,
195    },
196    Refusal(String),
197    Media {
198        kind: PartialMediaKind,
199        bytes: Vec<u8>,
200    },
201}
202
203#[derive(Debug)]
204enum PartialMediaKind {
205    Image {
206        media_type: String,
207    },
208    Audio {
209        media_type: String,
210    },
211    Document {
212        media_type: String,
213        name: Option<String>,
214    },
215}
216
217impl PartialBlock {
218    fn from_kind(kind: ContentBlockKind) -> Self {
219        match kind {
220            ContentBlockKind::Text => Self::Text(String::new()),
221            ContentBlockKind::Reasoning {
222                signature,
223                redacted,
224            } => Self::Reasoning {
225                text: String::new(),
226                signature,
227                redacted,
228            },
229            ContentBlockKind::ToolCall { id, name } => Self::ToolCall {
230                id,
231                name,
232                arguments: String::new(),
233                metadata: BTreeMap::new(),
234            },
235            ContentBlockKind::Refusal => Self::Refusal(String::new()),
236            ContentBlockKind::Image { media_type } => Self::Media {
237                kind: PartialMediaKind::Image { media_type },
238                bytes: Vec::new(),
239            },
240            ContentBlockKind::Audio { media_type } => Self::Media {
241                kind: PartialMediaKind::Audio { media_type },
242                bytes: Vec::new(),
243            },
244            ContentBlockKind::Document { media_type, name } => Self::Media {
245                kind: PartialMediaKind::Document { media_type, name },
246                bytes: Vec::new(),
247            },
248        }
249    }
250
251    fn complete(self) -> Result<ContentPart, ModelError> {
252        match self {
253            Self::Text(text) => Ok(ContentPart::Text { text }),
254            Self::Reasoning {
255                text,
256                signature,
257                redacted,
258            } => Ok(ContentPart::Reasoning(ReasoningPart {
259                text: (!text.is_empty()).then_some(text),
260                signature,
261                redacted,
262                provider_data: Vec::new(),
263            })),
264            Self::ToolCall {
265                id,
266                name,
267                arguments,
268                metadata,
269            } => {
270                let parsed = if arguments.trim().is_empty() {
271                    serde_json::json!({})
272                } else {
273                    serde_json::from_str(&arguments).map_err(|error| {
274                        ModelError::local(
275                            ModelErrorKind::MalformedToolArguments,
276                            format!("tool call {id} returned invalid JSON arguments: {error}"),
277                        )
278                    })?
279                };
280                Ok(ContentPart::ToolCall(ToolCall {
281                    id,
282                    name,
283                    arguments: parsed,
284                    raw_arguments: Some(arguments),
285                    metadata,
286                }))
287            }
288            Self::Refusal(text) => Ok(ContentPart::Refusal { text }),
289            Self::Media { kind, bytes } => {
290                let data = STANDARD.encode(bytes);
291                Ok(match kind {
292                    PartialMediaKind::Image { media_type } => ContentPart::Image {
293                        source: MediaSource::Base64 { media_type, data },
294                    },
295                    PartialMediaKind::Audio { media_type } => ContentPart::Audio {
296                        source: MediaSource::Base64 { media_type, data },
297                    },
298                    PartialMediaKind::Document { media_type, name } => ContentPart::Document {
299                        source: MediaSource::Base64 { media_type, data },
300                        name,
301                    },
302                })
303            }
304        }
305    }
306}
307
308/// Strictly reconstructs a canonical response from model stream events.
309#[derive(Debug, Default)]
310pub struct ModelStreamAccumulator {
311    started: bool,
312    completed: bool,
313    id: Option<String>,
314    model: Option<ModelRef>,
315    open_blocks: BTreeMap<u32, PartialBlock>,
316    content: BTreeMap<u32, ContentPart>,
317    usage: ModelUsage,
318    warnings: Vec<ModelWarning>,
319    provider_events: Vec<ProviderData>,
320}
321
322impl ModelStreamAccumulator {
323    /// Creates an empty accumulator.
324    pub fn new() -> Self {
325        Self::default()
326    }
327
328    /// Applies one event and returns the response when the terminal event arrives.
329    ///
330    /// # Errors
331    ///
332    /// Returns [`ModelError`] when event order is invalid, a delta targets the
333    /// wrong block type, indices collide, a response completes with open
334    /// blocks, or tool arguments contain malformed JSON.
335    pub fn push(&mut self, event: ModelStreamEvent) -> Result<Option<ModelResponse>, ModelError> {
336        if self.completed {
337            return Err(state_error("received an event after response completion"));
338        }
339
340        match event {
341            ModelStreamEvent::ResponseStarted { id, model } => self.start(id, model),
342            ModelStreamEvent::ContentBlockStarted { index, kind } => self.start_block(index, kind),
343            ModelStreamEvent::TextDelta { index, text } => {
344                match self.open_block_mut(index)? {
345                    PartialBlock::Text(current) => current.push_str(&text),
346                    _ => return Err(wrong_delta(index, "text")),
347                }
348                Ok(None)
349            }
350            ModelStreamEvent::ReasoningDelta { index, text } => {
351                match self.open_block_mut(index)? {
352                    PartialBlock::Reasoning { text: current, .. } => current.push_str(&text),
353                    _ => return Err(wrong_delta(index, "reasoning")),
354                }
355                Ok(None)
356            }
357            ModelStreamEvent::ReasoningSignatureDelta { index, signature } => {
358                match self.open_block_mut(index)? {
359                    PartialBlock::Reasoning {
360                        signature: current, ..
361                    } => current.get_or_insert_with(String::new).push_str(&signature),
362                    _ => return Err(wrong_delta(index, "reasoning signature")),
363                }
364                Ok(None)
365            }
366            ModelStreamEvent::ToolArgumentsDelta { index, json } => {
367                self.update_tool_arguments(index, json, false)?;
368                Ok(None)
369            }
370            ModelStreamEvent::ToolArgumentsCompleted { index, json } => {
371                self.update_tool_arguments(index, json, true)?;
372                Ok(None)
373            }
374            ModelStreamEvent::ContentBlockMetadata { index, metadata } => {
375                self.merge_block_metadata(index, metadata)
376            }
377            ModelStreamEvent::RefusalDelta { index, text } => {
378                match self.open_block_mut(index)? {
379                    PartialBlock::Refusal(current) => current.push_str(&text),
380                    _ => return Err(wrong_delta(index, "refusal")),
381                }
382                Ok(None)
383            }
384            ModelStreamEvent::BinaryDelta { index, data } => {
385                let decoded = STANDARD.decode(data).map_err(|error| {
386                    state_error(format!("binary delta {index} is invalid base64: {error}"))
387                })?;
388                match self.open_block_mut(index)? {
389                    PartialBlock::Media { bytes, .. } => {
390                        let next = bytes.len().checked_add(decoded.len()).ok_or_else(|| {
391                            state_error(format!("binary block {index} size overflow"))
392                        })?;
393                        if next > DEFAULT_MAX_ARTIFACT_BYTES {
394                            return Err(state_error(format!(
395                                "binary block {index} exceeds the {DEFAULT_MAX_ARTIFACT_BYTES}-byte limit"
396                            )));
397                        }
398                        bytes.extend_from_slice(&decoded);
399                    }
400                    _ => return Err(wrong_delta(index, "binary media")),
401                }
402                Ok(None)
403            }
404            ModelStreamEvent::ContentBlockCompleted { index } => self.complete_block(index),
405            ModelStreamEvent::ContentPartCompleted { index, part } => {
406                self.complete_part(index, part)
407            }
408            ModelStreamEvent::UsageUpdated { usage } => {
409                self.require_started()?;
410                self.usage = usage;
411                Ok(None)
412            }
413            ModelStreamEvent::Warning { warning } => {
414                self.require_started()?;
415                self.warnings.push(warning);
416                Ok(None)
417            }
418            ModelStreamEvent::Heartbeat => {
419                self.require_started()?;
420                Ok(None)
421            }
422            ModelStreamEvent::Provider { event } => {
423                self.require_started()?;
424                self.provider_events.push(ProviderData {
425                    provider: event.provider,
426                    kind: event.name,
427                    value: event.payload,
428                });
429                Ok(None)
430            }
431            ModelStreamEvent::ResponseCompleted {
432                finish_reason,
433                provider_metadata,
434            } => self.complete(finish_reason, provider_metadata),
435        }
436    }
437
438    fn start(
439        &mut self,
440        id: Option<String>,
441        model: ModelRef,
442    ) -> Result<Option<ModelResponse>, ModelError> {
443        if self.started {
444            return Err(state_error("received more than one response-start event"));
445        }
446        self.started = true;
447        self.id = id;
448        self.model = Some(model);
449        Ok(None)
450    }
451
452    fn update_tool_arguments(
453        &mut self,
454        index: u32,
455        json: String,
456        complete: bool,
457    ) -> Result<(), ModelError> {
458        match self.open_block_mut(index)? {
459            PartialBlock::ToolCall { arguments, .. } if complete => *arguments = json,
460            PartialBlock::ToolCall { arguments, .. } => arguments.push_str(&json),
461            _ if complete => return Err(wrong_delta(index, "completed tool arguments")),
462            _ => return Err(wrong_delta(index, "tool arguments")),
463        }
464        Ok(())
465    }
466
467    fn start_block(
468        &mut self,
469        index: u32,
470        kind: ContentBlockKind,
471    ) -> Result<Option<ModelResponse>, ModelError> {
472        self.require_started()?;
473        self.require_unused_index(index)?;
474        self.open_blocks
475            .insert(index, PartialBlock::from_kind(kind));
476        Ok(None)
477    }
478
479    fn complete_block(&mut self, index: u32) -> Result<Option<ModelResponse>, ModelError> {
480        self.require_started()?;
481        let block = self
482            .open_blocks
483            .remove(&index)
484            .ok_or_else(|| state_error(format!("content block {index} is not open")))?;
485        self.content.insert(index, block.complete()?);
486        Ok(None)
487    }
488
489    fn complete_part(
490        &mut self,
491        index: u32,
492        part: ContentPart,
493    ) -> Result<Option<ModelResponse>, ModelError> {
494        self.require_started()?;
495        self.require_unused_index(index)?;
496        self.content.insert(index, part);
497        Ok(None)
498    }
499
500    fn merge_block_metadata(
501        &mut self,
502        index: u32,
503        metadata: BTreeMap<String, Value>,
504    ) -> Result<Option<ModelResponse>, ModelError> {
505        self.require_started()?;
506        if let Some(block) = self.open_blocks.get_mut(&index) {
507            return match block {
508                PartialBlock::ToolCall {
509                    metadata: current, ..
510                } => {
511                    current.extend(metadata);
512                    Ok(None)
513                }
514                _ => Err(wrong_delta(index, "content metadata")),
515            };
516        }
517        match self.content.get_mut(&index) {
518            Some(ContentPart::ToolCall(call)) => {
519                call.metadata.extend(metadata);
520                Ok(None)
521            }
522            Some(_) => Err(wrong_delta(index, "content metadata")),
523            None => Err(state_error(format!(
524                "content metadata targeted unknown block {index}"
525            ))),
526        }
527    }
528
529    fn complete(
530        &mut self,
531        mut finish_reason: FinishReason,
532        provider_metadata: BTreeMap<String, Value>,
533    ) -> Result<Option<ModelResponse>, ModelError> {
534        self.require_started()?;
535        if !self.open_blocks.is_empty() {
536            let open = self
537                .open_blocks
538                .keys()
539                .map(u32::to_string)
540                .collect::<Vec<_>>()
541                .join(", ");
542            return Err(state_error(format!(
543                "response completed with open content blocks: {open}"
544            )));
545        }
546        self.completed = true;
547        let has_tool_calls = self
548            .content
549            .values()
550            .any(|part| matches!(part, ContentPart::ToolCall(_)));
551        if has_tool_calls && matches!(finish_reason, FinishReason::Stop) {
552            finish_reason = FinishReason::ToolCalls;
553        }
554        let model = self
555            .model
556            .clone()
557            .ok_or_else(|| state_error("response model is missing"))?;
558        Ok(Some(ModelResponse {
559            id: self.id.clone(),
560            model,
561            content: std::mem::take(&mut self.content).into_values().collect(),
562            finish_reason,
563            usage: self.usage,
564            warnings: std::mem::take(&mut self.warnings),
565            provider_metadata,
566            provider_events: std::mem::take(&mut self.provider_events),
567        }))
568    }
569
570    fn require_started(&self) -> Result<(), ModelError> {
571        if self.started {
572            Ok(())
573        } else {
574            Err(state_error("received content before response start"))
575        }
576    }
577
578    fn require_unused_index(&self, index: u32) -> Result<(), ModelError> {
579        if self.open_blocks.contains_key(&index) || self.content.contains_key(&index) {
580            Err(state_error(format!(
581                "content block index {index} was already used"
582            )))
583        } else {
584            Ok(())
585        }
586    }
587
588    fn open_block_mut(&mut self, index: u32) -> Result<&mut PartialBlock, ModelError> {
589        self.require_started()?;
590        self.open_blocks
591            .get_mut(&index)
592            .ok_or_else(|| state_error(format!("content block {index} is not open")))
593    }
594}
595
596fn wrong_delta(index: u32, delta: &str) -> ModelError {
597    state_error(format!(
598        "{delta} delta does not match content block {index}"
599    ))
600}
601
602fn state_error(message: impl Into<String>) -> ModelError {
603    ModelError::local(ModelErrorKind::StreamState, message)
604}
605
606#[cfg(test)]
607mod tests {
608    use std::collections::BTreeMap;
609
610    use super::{ContentBlockKind, ModelStreamAccumulator, ModelStreamEvent, ProviderEvent};
611    use crate::{
612        ContentPart, FinishReason, MediaSource, ModelErrorKind, ModelRef, ModelUsage, ModelWarning,
613        ToolCall,
614    };
615
616    fn started() -> ModelStreamEvent {
617        ModelStreamEvent::ResponseStarted {
618            id: Some("response-1".into()),
619            model: ModelRef::new("test", "model"),
620        }
621    }
622
623    fn completed() -> ModelStreamEvent {
624        ModelStreamEvent::ResponseCompleted {
625            finish_reason: FinishReason::Stop,
626            provider_metadata: BTreeMap::new(),
627        }
628    }
629
630    #[test]
631    fn accumulates_ordered_text_and_tool_calls() {
632        let mut accumulator = ModelStreamAccumulator::new();
633        let events = [
634            started(),
635            ModelStreamEvent::ContentBlockStarted {
636                index: 1,
637                kind: ContentBlockKind::ToolCall {
638                    id: "call-1".into(),
639                    name: "search".into(),
640                },
641            },
642            ModelStreamEvent::ToolArgumentsDelta {
643                index: 1,
644                json: "{\"query\":".into(),
645            },
646            ModelStreamEvent::ContentBlockStarted {
647                index: 0,
648                kind: ContentBlockKind::Text,
649            },
650            ModelStreamEvent::TextDelta {
651                index: 0,
652                text: "I will search.".into(),
653            },
654            ModelStreamEvent::ToolArgumentsDelta {
655                index: 1,
656                json: "\"rust\"}".into(),
657            },
658            ModelStreamEvent::ContentBlockMetadata {
659                index: 1,
660                metadata: BTreeMap::from([("test.status".into(), serde_json::json!("completed"))]),
661            },
662            ModelStreamEvent::ContentBlockCompleted { index: 0 },
663            ModelStreamEvent::ContentBlockCompleted { index: 1 },
664            ModelStreamEvent::UsageUpdated {
665                usage: ModelUsage {
666                    input_tokens: 5,
667                    output_tokens: 3,
668                    ..ModelUsage::default()
669                },
670            },
671            completed(),
672        ];
673
674        let response = events
675            .into_iter()
676            .find_map(|event| accumulator.push(event).unwrap())
677            .unwrap();
678
679        assert_eq!(response.content[0], ContentPart::text("I will search."));
680        assert_eq!(
681            response.content[1],
682            ContentPart::ToolCall(ToolCall {
683                id: "call-1".into(),
684                name: "search".into(),
685                arguments: serde_json::json!({"query": "rust"}),
686                raw_arguments: Some("{\"query\":\"rust\"}".into()),
687                metadata: BTreeMap::from([("test.status".into(), serde_json::json!("completed"),)]),
688            })
689        );
690        assert_eq!(response.usage.input_tokens, 5);
691    }
692
693    #[test]
694    fn preserves_provider_events_and_warnings() {
695        let mut accumulator = ModelStreamAccumulator::new();
696        accumulator.push(started()).unwrap();
697        accumulator
698            .push(ModelStreamEvent::Provider {
699                event: ProviderEvent {
700                    provider: "test".into(),
701                    name: "ping".into(),
702                    payload: serde_json::json!({"alive": true}),
703                },
704            })
705            .unwrap();
706        accumulator
707            .push(ModelStreamEvent::Warning {
708                warning: ModelWarning {
709                    code: "emulated".into(),
710                    message: "structured output was emulated".into(),
711                    metadata: BTreeMap::new(),
712                },
713            })
714            .unwrap();
715        let response = accumulator.push(completed()).unwrap().unwrap();
716
717        assert_eq!(response.provider_events.len(), 1);
718        assert_eq!(response.provider_events[0].kind, "ping");
719        assert_eq!(response.warnings.len(), 1);
720    }
721
722    #[test]
723    fn rejects_delta_without_matching_open_block() {
724        let mut accumulator = ModelStreamAccumulator::new();
725        accumulator.push(started()).unwrap();
726
727        let error = accumulator
728            .push(ModelStreamEvent::TextDelta {
729                index: 4,
730                text: "orphan".into(),
731            })
732            .unwrap_err();
733
734        assert_eq!(error.kind, ModelErrorKind::StreamState);
735    }
736
737    #[test]
738    fn rejects_completion_with_open_blocks() {
739        let mut accumulator = ModelStreamAccumulator::new();
740        accumulator.push(started()).unwrap();
741        accumulator
742            .push(ModelStreamEvent::ContentBlockStarted {
743                index: 0,
744                kind: ContentBlockKind::Text,
745            })
746            .unwrap();
747
748        let error = accumulator.push(completed()).unwrap_err();
749
750        assert_eq!(error.kind, ModelErrorKind::StreamState);
751    }
752
753    #[test]
754    fn rejects_malformed_tool_arguments() {
755        let mut accumulator = ModelStreamAccumulator::new();
756        accumulator.push(started()).unwrap();
757        accumulator
758            .push(ModelStreamEvent::ContentBlockStarted {
759                index: 0,
760                kind: ContentBlockKind::ToolCall {
761                    id: "bad".into(),
762                    name: "tool".into(),
763                },
764            })
765            .unwrap();
766        accumulator
767            .push(ModelStreamEvent::ToolArgumentsDelta {
768                index: 0,
769                json: "{invalid".into(),
770            })
771            .unwrap();
772
773        let error = accumulator
774            .push(ModelStreamEvent::ContentBlockCompleted { index: 0 })
775            .unwrap_err();
776
777        assert_eq!(error.kind, ModelErrorKind::MalformedToolArguments);
778    }
779
780    #[test]
781    fn completed_arguments_replace_partial_deltas() {
782        let mut accumulator = ModelStreamAccumulator::new();
783        accumulator.push(started()).unwrap();
784        accumulator
785            .push(ModelStreamEvent::ContentBlockStarted {
786                index: 0,
787                kind: ContentBlockKind::ToolCall {
788                    id: "call".into(),
789                    name: "tool".into(),
790                },
791            })
792            .unwrap();
793        accumulator
794            .push(ModelStreamEvent::ToolArgumentsDelta {
795                index: 0,
796                json: "{\"stale\":".into(),
797            })
798            .unwrap();
799        accumulator
800            .push(ModelStreamEvent::ToolArgumentsCompleted {
801                index: 0,
802                json: "{\"final\":true}".into(),
803            })
804            .unwrap();
805        accumulator
806            .push(ModelStreamEvent::ContentBlockCompleted { index: 0 })
807            .unwrap();
808        let response = accumulator.push(completed()).unwrap().unwrap();
809
810        let ContentPart::ToolCall(call) = &response.content[0] else {
811            panic!("fixture must produce a tool call");
812        };
813        assert_eq!(call.arguments, serde_json::json!({"final": true}));
814        assert_eq!(call.raw_arguments.as_deref(), Some("{\"final\":true}"));
815    }
816
817    #[test]
818    fn tool_content_normalizes_stop_reason() {
819        let mut accumulator = ModelStreamAccumulator::new();
820        accumulator.push(started()).unwrap();
821        accumulator
822            .push(ModelStreamEvent::ContentPartCompleted {
823                index: 0,
824                part: ContentPart::ToolCall(ToolCall {
825                    id: "call".into(),
826                    name: "tool".into(),
827                    arguments: serde_json::json!({}),
828                    raw_arguments: Some("{}".into()),
829                    metadata: BTreeMap::new(),
830                }),
831            })
832            .unwrap();
833        let response = accumulator.push(completed()).unwrap().unwrap();
834
835        assert_eq!(response.finish_reason, FinishReason::ToolCalls);
836    }
837
838    #[test]
839    fn tool_content_does_not_hide_unknown_finish_reason() {
840        let mut accumulator = ModelStreamAccumulator::new();
841        accumulator
842            .push(ModelStreamEvent::ResponseStarted {
843                id: None,
844                model: ModelRef::new("test", "model"),
845            })
846            .unwrap();
847        accumulator
848            .push(ModelStreamEvent::ContentPartCompleted {
849                index: 0,
850                part: ContentPart::ToolCall(ToolCall {
851                    id: "call-1".into(),
852                    name: "lookup".into(),
853                    arguments: serde_json::json!({}),
854                    raw_arguments: Some("{}".into()),
855                    metadata: BTreeMap::new(),
856                }),
857            })
858            .unwrap();
859        let response = accumulator
860            .push(ModelStreamEvent::ResponseCompleted {
861                finish_reason: FinishReason::Unknown,
862                provider_metadata: BTreeMap::new(),
863            })
864            .unwrap()
865            .unwrap();
866
867        assert_eq!(response.finish_reason, FinishReason::Unknown);
868    }
869
870    #[test]
871    fn explicit_failure_reason_is_not_hidden_by_tool_content() {
872        let mut accumulator = ModelStreamAccumulator::new();
873        accumulator.push(started()).unwrap();
874        accumulator
875            .push(ModelStreamEvent::ContentPartCompleted {
876                index: 0,
877                part: ContentPart::ToolCall(ToolCall {
878                    id: "partial".into(),
879                    name: "tool".into(),
880                    arguments: serde_json::json!({}),
881                    raw_arguments: Some("{}".into()),
882                    metadata: BTreeMap::new(),
883                }),
884            })
885            .unwrap();
886        let response = accumulator
887            .push(ModelStreamEvent::ResponseCompleted {
888                finish_reason: FinishReason::Length,
889                provider_metadata: BTreeMap::new(),
890            })
891            .unwrap()
892            .unwrap();
893
894        assert_eq!(response.finish_reason, FinishReason::Length);
895    }
896
897    #[test]
898    fn accumulates_bounded_binary_media_chunks() {
899        let mut accumulator = ModelStreamAccumulator::new();
900        accumulator.push(started()).unwrap();
901        accumulator
902            .push(ModelStreamEvent::ContentBlockStarted {
903                index: 0,
904                kind: ContentBlockKind::Image {
905                    media_type: "image/png".into(),
906                },
907            })
908            .unwrap();
909        accumulator
910            .push(ModelStreamEvent::BinaryDelta {
911                index: 0,
912                data: "cG5n".into(),
913            })
914            .unwrap();
915        accumulator
916            .push(ModelStreamEvent::ContentBlockCompleted { index: 0 })
917            .unwrap();
918        let response = accumulator.push(completed()).unwrap().unwrap();
919
920        assert!(matches!(
921            &response.content[0],
922            ContentPart::Image {
923                source: MediaSource::Base64 { media_type, data }
924            } if media_type == "image/png" && data == "cG5n"
925        ));
926    }
927}
runifold_model/stream.rs

runifold_model/
stream.rs
