Skip to main content

allama/chat/
stream.rs

1//! Chat streams: events, collection, transcripts, and commit-point helpers.
2//!
3//! The types every stream consumer names ([`StreamEvent`],
4//! [`ChatStream`], [`ChatStreamExt`], [`CollectedEvent`]) are also
5//! exported at the crate root. Other operations that stream (transcription,
6//! speech, …) get their own event types beside their own modules; these
7//! are the chat ones.
8//!
9//! # The commit point
10//!
11//! A stream's commit point is its first event that carries model output
12//! ([`StreamEvent::is_output`](crate::StreamEvent::is_output): text,
13//! reasoning, a refusal, a tool call, an image), after which the caller may
14//! have seen the answer. A stream that ends without output commits at its
15//! end.
16//!
17//! - **Before it**, nothing is exposed, so an error can still be retried,
18//!   sent to another provider, or answered with an HTTP status code.
19//!   [`RetryingProvider`](crate::wrappers::RetryingProvider) and
20//!   [`FallbackProvider`](crate::wrappers::FallbackProvider) hold a stream
21//!   back until its commit point, and [`until_committed`] gives an
22//!   application the same wait, replaying the held events so nothing is
23//!   lost and the time to the first token is unchanged.
24//! - **After it**, an error can only end the stream: retry and fallback
25//!   pass it through, and a server that already sent `200 OK` reports it as
26//!   an event.
27//!
28//! The commit point also decides what a timeout outside retry and fallback
29//! bounds: its `stream_setup` is the wait for the first output (see
30//! [attempt timeout vs deadline](crate::wrappers::TimeoutPolicy#attempt-timeout-vs-deadline)).
31
32use std::borrow::Cow;
33use std::collections::VecDeque;
34use std::fmt;
35use std::pin::Pin;
36use std::task::{Context, Poll};
37
38use futures_core::Stream;
39use serde::{Deserialize, Serialize};
40
41use crate::chat::content::is_valid_json;
42pub use crate::chat::transcript::StreamTranscript;
43use crate::{
44    ChatResponse, Citation, ContentBlock, Error, ExtraMap, FinishReason, Image, MediaSource,
45    OtherContent, ProviderData, ProviderToolCall, ProviderToolResult, Reasoning, Refusal, Result,
46    StreamErrorKind, Text, ToolCall, Usage, Warning,
47};
48
49/// A boxed, type-erased chat stream.
50///
51/// # A chat stream ends after its first error
52///
53/// An `Err` item is the last item of a chat stream: nothing follows it, and
54/// polling again yields `None`. Consumers rely on it: collecting a stream
55/// ([`collect_response`](ChatStreamExt::collect_response)) and
56/// [`until_committed`] stop at the first
57/// error and treat it as the stream's terminal error.
58///
59/// The streams allama builds keep the rule however the stream under them
60/// behaves: [`Prepared::finish_stream`](crate::adapter::chat::Prepared::finish_stream),
61/// [`with_collector`](ChatStreamExt::with_collector) and the extraction
62/// stream stop polling what they wrap after its first error. A stream built
63/// by hand (an adapter without `finish_stream`, a wrapper, a test double)
64/// must honour it too: end after yielding an error, and do not poll what
65/// it wraps again.
66pub type ChatStream = Pin<Box<dyn Stream<Item = Result<StreamEvent>> + Send + 'static>>;
67
68/// How to interpret the `usage` of a [`StreamEvent::ResponseMetadata`] event.
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
70#[serde(rename_all = "snake_case")]
71pub enum UsageMetadataMode {
72    /// Incremental usage, added to prior usage.
73    Delta,
74    /// Cumulative usage, replacing prior usage.
75    #[default]
76    Snapshot,
77}
78
79fn is_snapshot(mode: &UsageMetadataMode) -> bool {
80    *mode == UsageMetadataMode::Snapshot
81}
82
83/// Portable type of a streamed content block.
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
85#[serde(rename_all = "snake_case")]
86#[non_exhaustive]
87pub enum StreamBlockType {
88    /// A text block.
89    Text,
90    /// An image block (complete at start).
91    Image,
92    /// A reasoning block.
93    Reasoning,
94    /// A refusal block; its text streams as `RefusalDelta`.
95    Refusal,
96    /// A client tool call; its arguments stream as `ToolCallDelta`. Start
97    /// it with [`StreamEvent::tool_call_start`], which carries the call's
98    /// id and name: a tool call without them is dropped as incomplete.
99    ToolCall,
100    /// A provider-executed tool call; its arguments stream as
101    /// `ToolCallDelta`. Start it with
102    /// [`StreamEvent::provider_tool_call_start`], which carries its id and
103    /// name.
104    ProviderToolCall,
105    /// A provider tool result (complete at start: `id` is the call id,
106    /// `data` holds `output` and `is_error`).
107    ProviderToolResult,
108    /// A block the portable model does not represent (complete at start).
109    Other,
110}
111
112impl StreamBlockType {
113    /// Every block type, in declaration order (for audits such as the
114    /// conformance honesty contract). Grows when block types are added.
115    pub const ALL: &'static [Self] = &[
116        Self::Text,
117        Self::Image,
118        Self::Reasoning,
119        Self::Refusal,
120        Self::ToolCall,
121        Self::ProviderToolCall,
122        Self::ProviderToolResult,
123        Self::Other,
124    ];
125
126    /// The name used in protocol error messages (the serialized form).
127    fn label(self) -> &'static str {
128        match self {
129            Self::Text => "text",
130            Self::Image => "image",
131            Self::Reasoning => "reasoning",
132            Self::Refusal => "refusal",
133            Self::ToolCall => "tool_call",
134            Self::ProviderToolCall => "provider_tool_call",
135            Self::ProviderToolResult => "provider_tool_result",
136            Self::Other => "other",
137        }
138    }
139}
140
141/// One event of a normalized chat stream.
142///
143/// A stream is `ResponseStart`, then for each block `BlockStart`, deltas and
144/// `BlockStop`, then `ResponseMetadata` (possibly several), then
145/// `ResponseStop`. `Warning` and `ResponseProviderData` events may appear at
146/// any point. Blocks may overlap: a block can stay open while later blocks
147/// start and stream (Gemini keeps a text block open until its candidate
148/// ends, to attach grounding citations), so route deltas by their `index`,
149/// not by the last `BlockStart`. [`StreamCollector`] is lenient about the
150/// order (a delta before its `BlockStart` opens the block); adapters keep
151/// it anyway, for consumers of the raw events. Wrappers keep this order when they add events: for example
152/// [`FallbackProvider`](crate::wrappers::FallbackProvider) reports its
153/// failed attempts in a `ResponseMetadata` event right after the fallback's
154/// `ResponseStart`.
155///
156/// The enum and its struct variants are `#[non_exhaustive]`, so new events and
157/// new fields are not breaking changes. Build events with the constructors
158/// ([`StreamEvent::text_delta`], [`StreamEvent::metadata`], ...) or, for
159/// `ResponseMetadata`, a [`MetadataBuilder`];
160/// set further fields through a mutable match, and match variants with `..`:
161///
162/// ```
163/// use allama::chat::stream::MetadataBuilder;
164/// use allama::{FinishReason, StreamBlockType, StreamEvent};
165///
166/// let metadata = MetadataBuilder::new()
167///     .with_finish_reason(FinishReason::Stop)
168///     .build();
169/// assert!(matches!(metadata, StreamEvent::ResponseMetadata { .. }));
170///
171/// let mut start = StreamEvent::block_start(0, StreamBlockType::Other);
172/// if let StreamEvent::BlockStart { type_name, .. } = &mut start {
173///     *type_name = Some("vendor_block".into());
174/// }
175///
176/// match StreamEvent::text_delta(0, "Hello") {
177///     StreamEvent::TextDelta { text, .. } => assert_eq!(text, "Hello"),
178///     _ => unreachable!(),
179/// }
180/// ```
181///
182/// A tool call starts as a `BlockStart` with the `ToolCall` block type,
183/// which names the tool before its arguments stream in:
184///
185/// ```
186/// use allama::{StreamBlockType, StreamEvent};
187///
188/// fn announce(event: &StreamEvent) {
189///     if let StreamEvent::BlockStart {
190///         block_type: StreamBlockType::ToolCall,
191///         name: Some(name),
192///         ..
193///     } = event
194///     {
195///         println!("calling {name}…");
196///     }
197/// }
198/// # announce(&StreamEvent::block_start(0, StreamBlockType::ToolCall));
199/// ```
200#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
201#[serde(tag = "event", rename_all = "snake_case")]
202#[non_exhaustive]
203pub enum StreamEvent {
204    /// The response started.
205    #[non_exhaustive]
206    ResponseStart {
207        /// Response identifier, when known.
208        #[serde(default, skip_serializing_if = "Option::is_none")]
209        id: Option<String>,
210        /// Model identifier, when known.
211        #[serde(default, skip_serializing_if = "Option::is_none")]
212        model: Option<String>,
213    },
214    /// A content block started. A tool call's start names the tool (see
215    /// the [enum docs](Self)).
216    #[doc(alias = "ToolCallStart")]
217    #[non_exhaustive]
218    BlockStart {
219        /// Block index, stable across the block's events.
220        index: usize,
221        /// Portable block type.
222        block_type: StreamBlockType,
223        /// Tool call id (tool-call blocks).
224        #[serde(default, skip_serializing_if = "Option::is_none")]
225        id: Option<String>,
226        /// Tool name (tool-call blocks).
227        #[serde(default, skip_serializing_if = "Option::is_none")]
228        name: Option<String>,
229        /// Wire type name (`Other` blocks).
230        #[serde(default, skip_serializing_if = "Option::is_none")]
231        type_name: Option<String>,
232        /// Block payload for blocks that are complete at start (image, other).
233        #[serde(default, skip_serializing_if = "Option::is_none")]
234        data: Option<ExtraMap>,
235        /// Origin namespace (`Other` blocks).
236        #[serde(default, skip_serializing_if = "Option::is_none")]
237        origin: Option<String>,
238        /// Provider data known at block start.
239        #[serde(default, skip_serializing_if = "ProviderData::is_empty")]
240        provider_data: ProviderData,
241    },
242    /// Incremental text.
243    #[non_exhaustive]
244    TextDelta {
245        /// Block index.
246        index: usize,
247        /// Text chunk.
248        text: String,
249    },
250    /// Incremental reasoning text.
251    #[non_exhaustive]
252    ReasoningDelta {
253        /// Block index.
254        index: usize,
255        /// Reasoning chunk.
256        text: String,
257    },
258    /// Incremental refusal text: the model declining to answer (see
259    /// [`Refusal`]).
260    ///
261    /// It is never a `TextDelta`, so a loop that prints only text deltas
262    /// never shows a refusal as the answer. It is model output, so it is a
263    /// commit point ([`is_output`](Self::is_output)).
264    #[non_exhaustive]
265    RefusalDelta {
266        /// Block index.
267        index: usize,
268        /// Refusal chunk.
269        text: String,
270    },
271    /// Incremental tool-call arguments (a JSON fragment).
272    #[non_exhaustive]
273    ToolCallDelta {
274        /// Block index.
275        index: usize,
276        /// Arguments chunk.
277        arguments: String,
278    },
279    /// A citation for a text block, appended to the block's
280    /// [`citations`](crate::Text::citations).
281    #[non_exhaustive]
282    Citation {
283        /// Block index.
284        index: usize,
285        /// The citation.
286        citation: Citation,
287    },
288    /// Provider data for a block (e.g. a signature), merged into the block.
289    #[non_exhaustive]
290    BlockProviderData {
291        /// Block index.
292        index: usize,
293        /// Data to merge (per key, later values win).
294        provider_data: ProviderData,
295    },
296    /// Provider data for the response as a whole, merged into
297    /// [`ChatResponse::provider_data`](crate::ChatResponse::provider_data)
298    /// (per key, later values win).
299    ///
300    /// The response-level counterpart of `BlockProviderData`, for data that
301    /// arrives on its own rather than with `ResponseMetadata`: for example
302    /// vendor stream frames the adapter does not understand, surfaced
303    /// verbatim under the provider's namespace.
304    #[non_exhaustive]
305    ResponseProviderData {
306        /// Data to merge (per key, later values win).
307        provider_data: ProviderData,
308    },
309    /// A provider-specific incremental update that the portable events do
310    /// not represent: hosted-tool progress, delta types an adapter does not
311    /// map, and so on.
312    ///
313    /// Observation only: collectors do not fold it into the response, so it
314    /// never changes what [`collect_response`](ChatStreamExt::collect_response)
315    /// returns. Adapters emit it only when no portable event fits; it is not
316    /// raw wire passthrough.
317    #[non_exhaustive]
318    ProviderDelta {
319        /// Index of the block the update belongs to.
320        index: usize,
321        /// Provider event or delta type, e.g. `"web_search_call.searching"`.
322        type_name: String,
323        /// Payload, verbatim.
324        #[serde(default, skip_serializing_if = "ExtraMap::is_empty")]
325        data: ExtraMap,
326        /// Namespace of the provider that produced the update.
327        #[serde(default, skip_serializing_if = "Option::is_none")]
328        origin: Option<String>,
329    },
330    /// A content block finished.
331    #[non_exhaustive]
332    BlockStop {
333        /// Block index.
334        index: usize,
335    },
336    /// A non-fatal warning.
337    #[non_exhaustive]
338    Warning {
339        /// The warning.
340        warning: Warning,
341    },
342    /// Response-level metadata. May arrive several times; later values win.
343    #[non_exhaustive]
344    ResponseMetadata {
345        /// Normalized finish reason.
346        #[serde(default, skip_serializing_if = "Option::is_none")]
347        finish_reason: Option<FinishReason>,
348        /// Raw provider finish reason.
349        #[serde(default, skip_serializing_if = "Option::is_none")]
350        finish_reason_raw: Option<String>,
351        /// Usage, interpreted per `usage_mode`.
352        #[serde(default, skip_serializing_if = "Option::is_none")]
353        usage: Option<Usage>,
354        /// How to interpret `usage`.
355        #[serde(default, skip_serializing_if = "is_snapshot")]
356        usage_mode: UsageMetadataMode,
357        /// Response identifier.
358        #[serde(default, skip_serializing_if = "Option::is_none")]
359        id: Option<String>,
360        /// Model identifier.
361        #[serde(default, skip_serializing_if = "Option::is_none")]
362        model: Option<String>,
363        /// Provider request identifier.
364        #[serde(default, skip_serializing_if = "Option::is_none")]
365        request_id: Option<String>,
366        /// Provider that served the request.
367        #[serde(default, skip_serializing_if = "Option::is_none")]
368        provider: Option<Cow<'static, str>>,
369        /// Response-level provider data (merged).
370        #[serde(default, skip_serializing_if = "ProviderData::is_empty")]
371        provider_data: ProviderData,
372    },
373    /// The response ended.
374    ///
375    /// A unit variant, so it stays constructible and matchable as
376    /// `StreamEvent::ResponseStop`. Response-level data belongs in
377    /// `ResponseMetadata`.
378    ResponseStop,
379}
380
381impl StreamEvent {
382    /// A `ResponseStart` event.
383    #[must_use]
384    pub fn response_start(id: Option<String>, model: Option<String>) -> Self {
385        Self::ResponseStart { id, model }
386    }
387
388    /// A `TextDelta` event.
389    #[must_use]
390    pub fn text_delta(index: usize, text: impl Into<String>) -> Self {
391        Self::TextDelta {
392            index,
393            text: text.into(),
394        }
395    }
396
397    /// A `ReasoningDelta` event.
398    #[must_use]
399    pub fn reasoning_delta(index: usize, text: impl Into<String>) -> Self {
400        Self::ReasoningDelta {
401            index,
402            text: text.into(),
403        }
404    }
405
406    /// A `RefusalDelta` event.
407    #[must_use]
408    pub fn refusal_delta(index: usize, text: impl Into<String>) -> Self {
409        Self::RefusalDelta {
410            index,
411            text: text.into(),
412        }
413    }
414
415    /// A `ToolCallDelta` event carrying a JSON fragment of the arguments.
416    #[must_use]
417    pub fn tool_call_delta(index: usize, arguments: impl Into<String>) -> Self {
418        Self::ToolCallDelta {
419            index,
420            arguments: arguments.into(),
421        }
422    }
423
424    /// A `Citation` event for the text block at `index`.
425    #[must_use]
426    pub fn citation(index: usize, citation: Citation) -> Self {
427        Self::Citation { index, citation }
428    }
429
430    /// A `BlockProviderData` event, merged into the block at `index`.
431    #[must_use]
432    pub fn block_provider_data(index: usize, provider_data: ProviderData) -> Self {
433        Self::BlockProviderData {
434            index,
435            provider_data,
436        }
437    }
438
439    /// A `ResponseProviderData` event, merged into the response's provider
440    /// data.
441    #[must_use]
442    pub fn response_provider_data(provider_data: ProviderData) -> Self {
443        Self::ResponseProviderData { provider_data }
444    }
445
446    /// A `BlockStop` event.
447    #[must_use]
448    pub fn block_stop(index: usize) -> Self {
449        Self::BlockStop { index }
450    }
451
452    /// A `Warning` event.
453    #[must_use]
454    pub fn warning(warning: Warning) -> Self {
455        Self::Warning { warning }
456    }
457
458    /// A `ResponseMetadata` event with every field empty.
459    #[must_use]
460    pub fn metadata() -> Self {
461        Self::ResponseMetadata {
462            finish_reason: None,
463            finish_reason_raw: None,
464            usage: None,
465            usage_mode: UsageMetadataMode::Snapshot,
466            id: None,
467            model: None,
468            request_id: None,
469            provider: None,
470            provider_data: ProviderData::new(),
471        }
472    }
473
474    /// A `BlockStart` event for a block of `block_type` with no payload:
475    /// text, reasoning, or a refusal.
476    ///
477    /// A tool call needs its id and name, so start one with
478    /// [`tool_call_start`](Self::tool_call_start) or
479    /// [`provider_tool_call_start`](Self::provider_tool_call_start); a
480    /// `ToolCall` started here has neither and is dropped as incomplete. A
481    /// block that is complete at start needs its payload:
482    /// [`image_start`](Self::image_start),
483    /// [`provider_tool_result_start`](Self::provider_tool_result_start),
484    /// [`other_start`](Self::other_start).
485    #[must_use]
486    pub fn block_start(index: usize, block_type: StreamBlockType) -> Self {
487        Self::BlockStart {
488            index,
489            block_type,
490            id: None,
491            name: None,
492            type_name: None,
493            data: None,
494            origin: None,
495            provider_data: ProviderData::new(),
496        }
497    }
498
499    /// A `BlockStart` event for a provider-executed tool call; its arguments
500    /// follow as `ToolCallDelta` events.
501    #[must_use]
502    pub fn provider_tool_call_start(
503        index: usize,
504        id: impl Into<String>,
505        name: impl Into<String>,
506        origin: Option<String>,
507    ) -> Self {
508        Self::BlockStart {
509            index,
510            block_type: StreamBlockType::ProviderToolCall,
511            id: Some(id.into()),
512            name: Some(name.into()),
513            type_name: None,
514            data: None,
515            origin,
516            provider_data: ProviderData::new(),
517        }
518    }
519
520    /// The `BlockStart` event of a provider tool result, which is complete at
521    /// start.
522    #[must_use]
523    pub fn provider_tool_result_start(index: usize, result: ProviderToolResult) -> Self {
524        let mut data = ExtraMap::new();
525        data.insert("output".into(), result.output);
526        if result.is_error {
527            data.insert("is_error".into(), true.into());
528        }
529        Self::BlockStart {
530            index,
531            block_type: StreamBlockType::ProviderToolResult,
532            id: Some(result.call_id),
533            name: None,
534            type_name: None,
535            data: Some(data),
536            origin: result.origin,
537            provider_data: result.provider_data,
538        }
539    }
540
541    /// The `BlockStart` event of an image block, which is complete at start.
542    #[must_use]
543    pub fn image_start(index: usize, image: Image) -> Self {
544        Self::BlockStart {
545            index,
546            block_type: StreamBlockType::Image,
547            id: None,
548            name: None,
549            type_name: None,
550            data: Some(image_data(&image.source, image.media_type.as_deref())),
551            origin: None,
552            provider_data: image.provider_data,
553        }
554    }
555
556    /// The `BlockStart` event of a block the portable model does not
557    /// represent, which is complete at start.
558    #[must_use]
559    pub fn other_start(index: usize, other: OtherContent) -> Self {
560        Self::BlockStart {
561            index,
562            block_type: StreamBlockType::Other,
563            id: None,
564            name: None,
565            type_name: Some(other.type_name),
566            data: Some(other.data),
567            origin: other.origin,
568            provider_data: other.provider_data,
569        }
570    }
571
572    /// A `ProviderDelta` event.
573    #[must_use]
574    pub fn provider_delta(
575        index: usize,
576        type_name: impl Into<String>,
577        data: ExtraMap,
578        origin: Option<String>,
579    ) -> Self {
580        Self::ProviderDelta {
581            index,
582            type_name: type_name.into(),
583            data,
584            origin,
585        }
586    }
587
588    /// A `BlockStart` event for a tool call.
589    #[must_use]
590    pub fn tool_call_start(index: usize, id: impl Into<String>, name: impl Into<String>) -> Self {
591        Self::BlockStart {
592            index,
593            block_type: StreamBlockType::ToolCall,
594            id: Some(id.into()),
595            name: Some(name.into()),
596            type_name: None,
597            data: None,
598            origin: None,
599            provider_data: ProviderData::new(),
600        }
601    }
602
603    /// Returns `true` for events that carry model output. Fallback and retry
604    /// treat the first such event as the commit point.
605    ///
606    /// A refusal is output: once the caller may have seen the model decline,
607    /// a later error must not retry the call or send it to another provider.
608    #[must_use]
609    pub fn is_output(&self) -> bool {
610        match self {
611            Self::TextDelta { .. }
612            | Self::ReasoningDelta { .. }
613            | Self::RefusalDelta { .. }
614            | Self::ToolCallDelta { .. } => true,
615            Self::BlockStart { block_type, .. } => {
616                matches!(
617                    block_type,
618                    StreamBlockType::Image
619                        | StreamBlockType::Refusal
620                        | StreamBlockType::ProviderToolResult
621                        | StreamBlockType::Other
622                )
623            }
624            _ => false,
625        }
626    }
627
628    /// Returns `true` for events that address a content block by index.
629    fn is_block_event(&self) -> bool {
630        matches!(
631            self,
632            Self::BlockStart { .. }
633                | Self::TextDelta { .. }
634                | Self::ReasoningDelta { .. }
635                | Self::RefusalDelta { .. }
636                | Self::ToolCallDelta { .. }
637                | Self::Citation { .. }
638                | Self::BlockProviderData { .. }
639                | Self::ProviderDelta { .. }
640                | Self::BlockStop { .. }
641        )
642    }
643
644    pub(crate) fn name(&self) -> &'static str {
645        match self {
646            Self::ResponseStart { .. } => "response_start",
647            Self::BlockStart { .. } => "block_start",
648            Self::TextDelta { .. } => "text_delta",
649            Self::ReasoningDelta { .. } => "reasoning_delta",
650            Self::RefusalDelta { .. } => "refusal_delta",
651            Self::ToolCallDelta { .. } => "tool_call_delta",
652            Self::Citation { .. } => "citation",
653            Self::BlockProviderData { .. } => "block_provider_data",
654            Self::ResponseProviderData { .. } => "response_provider_data",
655            Self::ProviderDelta { .. } => "provider_delta",
656            Self::BlockStop { .. } => "block_stop",
657            Self::Warning { .. } => "warning",
658            Self::ResponseMetadata { .. } => "response_metadata",
659            Self::ResponseStop => "response_stop",
660        }
661    }
662}
663
664/// Builds a [`StreamEvent::ResponseMetadata`] event field by field.
665///
666/// An alternative to setting the fields of [`StreamEvent::metadata`]
667/// through a mutable match. Unset fields stay empty, and `usage_mode`
668/// defaults to [`UsageMetadataMode::Snapshot`].
669///
670/// The setters follow the rule of the [`Error`] builders:
671/// **`with_*` takes the value**, and each field a wire response may lack
672/// also has a **`with_*_opt`** twin that takes an `Option` and keeps the
673/// current value on `None`. Text twins take an `Option<&str>`, so a
674/// decoder passes an owned field with `.as_deref()` and a bare `None`
675/// compiles.
676///
677/// ```
678/// use allama::chat::stream::{MetadataBuilder, UsageMetadataMode};
679/// use allama::{FinishReason, StreamEvent, Usage};
680///
681/// // Whatever the wire had.
682/// let request_id: Option<String> = Some("req_1".into());
683/// let model: Option<String> = None;
684/// let event = MetadataBuilder::new()
685///     .with_finish_reason(FinishReason::Stop)
686///     .with_finish_reason_raw("end_turn")
687///     .with_usage(Usage::new().input_tokens(12).output_tokens(3))
688///     .with_request_id_opt(request_id.as_deref())
689///     .with_model_opt(model.as_deref())
690///     .with_provider("anthropic")
691///     .build();
692///
693/// let StreamEvent::ResponseMetadata {
694///     finish_reason,
695///     usage_mode,
696///     request_id,
697///     model,
698///     ..
699/// } = &event
700/// else {
701///     unreachable!()
702/// };
703/// assert_eq!(*finish_reason, Some(FinishReason::Stop));
704/// assert_eq!(*usage_mode, UsageMetadataMode::Snapshot);
705/// assert_eq!(request_id.as_deref(), Some("req_1"));
706/// assert_eq!(*model, None);
707/// ```
708#[derive(Debug, Clone, PartialEq, Eq)]
709#[must_use]
710pub struct MetadataBuilder {
711    /// Always a `ResponseMetadata` event.
712    event: StreamEvent,
713}
714
715impl Default for MetadataBuilder {
716    fn default() -> Self {
717        Self::new()
718    }
719}
720
721impl MetadataBuilder {
722    /// A builder with every field empty (as [`StreamEvent::metadata`]).
723    pub fn new() -> Self {
724        Self {
725            event: StreamEvent::metadata(),
726        }
727    }
728
729    /// Set the normalized finish reason.
730    pub fn with_finish_reason(self, value: FinishReason) -> Self {
731        self.with_finish_reason_opt(Some(value))
732    }
733
734    /// Set the normalized finish reason when there is one; `None` keeps
735    /// the current value. The `_opt` form of
736    /// [`with_finish_reason`](Self::with_finish_reason).
737    pub fn with_finish_reason_opt(mut self, value: Option<FinishReason>) -> Self {
738        if let StreamEvent::ResponseMetadata { finish_reason, .. } = &mut self.event {
739            merge(finish_reason, value);
740        }
741        self
742    }
743
744    /// Set the raw provider finish reason.
745    pub fn with_finish_reason_raw(mut self, value: impl Into<String>) -> Self {
746        if let StreamEvent::ResponseMetadata {
747            finish_reason_raw, ..
748        } = &mut self.event
749        {
750            *finish_reason_raw = Some(value.into());
751        }
752        self
753    }
754
755    /// Set the raw provider finish reason when there is one; `None` keeps
756    /// the current value. The `_opt` form of
757    /// [`with_finish_reason_raw`](Self::with_finish_reason_raw).
758    pub fn with_finish_reason_raw_opt(self, value: Option<&str>) -> Self {
759        match value {
760            Some(value) => self.with_finish_reason_raw(value),
761            None => self,
762        }
763    }
764
765    /// Set the usage, interpreted per [`with_usage_mode`](Self::with_usage_mode).
766    pub fn with_usage(self, value: Usage) -> Self {
767        self.with_usage_opt(Some(value))
768    }
769
770    /// Set the usage when the wire reported one; `None` keeps the current
771    /// value. The `_opt` form of [`with_usage`](Self::with_usage).
772    pub fn with_usage_opt(mut self, value: Option<Usage>) -> Self {
773        if let StreamEvent::ResponseMetadata { usage, .. } = &mut self.event {
774            merge(usage, value);
775        }
776        self
777    }
778
779    /// Set how to interpret the usage (default: a snapshot).
780    pub fn with_usage_mode(mut self, value: UsageMetadataMode) -> Self {
781        if let StreamEvent::ResponseMetadata { usage_mode, .. } = &mut self.event {
782            *usage_mode = value;
783        }
784        self
785    }
786
787    /// Set the response identifier.
788    pub fn with_id(mut self, value: impl Into<String>) -> Self {
789        if let StreamEvent::ResponseMetadata { id, .. } = &mut self.event {
790            *id = Some(value.into());
791        }
792        self
793    }
794
795    /// Set the response identifier when there is one; `None` keeps the
796    /// current value. The `_opt` form of [`with_id`](Self::with_id).
797    pub fn with_id_opt(self, value: Option<&str>) -> Self {
798        match value {
799            Some(value) => self.with_id(value),
800            None => self,
801        }
802    }
803
804    /// Set the model identifier.
805    pub fn with_model(mut self, value: impl Into<String>) -> Self {
806        if let StreamEvent::ResponseMetadata { model, .. } = &mut self.event {
807            *model = Some(value.into());
808        }
809        self
810    }
811
812    /// Set the model identifier when there is one; `None` keeps the
813    /// current value. The `_opt` form of [`with_model`](Self::with_model).
814    pub fn with_model_opt(self, value: Option<&str>) -> Self {
815        match value {
816            Some(value) => self.with_model(value),
817            None => self,
818        }
819    }
820
821    /// Set the provider request identifier.
822    pub fn with_request_id(mut self, value: impl Into<String>) -> Self {
823        if let StreamEvent::ResponseMetadata { request_id, .. } = &mut self.event {
824            *request_id = Some(value.into());
825        }
826        self
827    }
828
829    /// Set the provider request identifier when there is one; `None` keeps
830    /// the current value. The `_opt` form of
831    /// [`with_request_id`](Self::with_request_id).
832    pub fn with_request_id_opt(self, value: Option<&str>) -> Self {
833        match value {
834            Some(value) => self.with_request_id(value),
835            None => self,
836        }
837    }
838
839    /// Set the provider that served the request, such as an adapter's
840    /// `PROVIDER_NAME` constant.
841    ///
842    /// An adapter built on [`prepare`](crate::adapter::chat::prepare) need
843    /// not set it: [`finish_stream`](crate::adapter::chat::Prepared::finish_stream)
844    /// fills it in when the stream leaves it empty.
845    pub fn with_provider(self, value: impl Into<Cow<'static, str>>) -> Self {
846        self.with_provider_opt(Some(value.into()))
847    }
848
849    /// Set the provider that served the request when there is one; `None`
850    /// keeps the current value. The `_opt` form of
851    /// [`with_provider`](Self::with_provider); it takes a `Cow` so a
852    /// `&'static str` name is not copied.
853    pub fn with_provider_opt(mut self, value: Option<Cow<'static, str>>) -> Self {
854        if let StreamEvent::ResponseMetadata { provider, .. } = &mut self.event {
855            merge(provider, value);
856        }
857        self
858    }
859
860    /// Set the response-level provider data.
861    pub fn with_provider_data(mut self, value: ProviderData) -> Self {
862        if let StreamEvent::ResponseMetadata { provider_data, .. } = &mut self.event {
863            *provider_data = value;
864        }
865        self
866    }
867
868    /// The `ResponseMetadata` event.
869    #[must_use]
870    pub fn build(self) -> StreamEvent {
871        self.event
872    }
873}
874
875impl From<MetadataBuilder> for StreamEvent {
876    fn from(builder: MetadataBuilder) -> Self {
877        builder.build()
878    }
879}
880
881/// Completeness flags of a collected stream.
882#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
883#[non_exhaustive]
884pub struct StreamCompleteness {
885    /// A `ResponseStart` was seen.
886    pub saw_response_start: bool,
887    /// A `ResponseStop` was seen.
888    pub saw_response_stop: bool,
889    /// Some blocks never received `BlockStop`.
890    pub has_open_blocks: bool,
891    /// Incomplete tool calls were dropped from the partial response.
892    pub dropped_incomplete_tool_calls: bool,
893}
894
895impl StreamCompleteness {
896    /// Returns `true` when the stream was fully delivered.
897    #[must_use]
898    pub fn is_complete(&self) -> bool {
899        self.saw_response_start
900            && self.saw_response_stop
901            && !self.has_open_blocks
902            && !self.dropped_incomplete_tool_calls
903    }
904}
905
906/// A best-effort collected response plus completeness diagnostics.
907#[derive(Debug, Clone)]
908#[non_exhaustive]
909pub struct CollectedResponse {
910    /// The response assembled so far.
911    pub response: ChatResponse,
912    /// Completeness flags.
913    pub completeness: StreamCompleteness,
914    /// The error that ended the stream, if any.
915    pub terminal_error: Option<Error>,
916}
917
918/// Accumulates [`StreamEvent`]s into a [`ChatResponse`].
919///
920/// A block is final once its `BlockStop` arrives: [`push`](Self::push)
921/// returns it then, so callers can act on complete blocks (a tool call, a
922/// paragraph) while the stream continues. Events for a stopped block are
923/// protocol errors.
924///
925/// Block indices should be dense (`0, 1, 2, ...`). A decoder may skip the
926/// indices of blocks it does not model, but a stream that skips more than
927/// 1024 indices in total is a protocol error: the collector keeps a slot per
928/// index, so its memory stays linear in the events pushed.
929///
930/// Application code that reads a stream usually wants
931/// [`ChatStreamExt::with_collector`], which drives a collector for it. Push
932/// events into a collector directly when you write a wrapper or a stream
933/// adapter that sees the events one at a time.
934#[derive(Default)]
935pub struct StreamCollector {
936    blocks: Vec<Option<Slot>>,
937    /// Indices left empty by gaps so far (see [`reserve`](Self::reserve)).
938    skipped_indices: usize,
939    response: ChatResponse,
940    saw_response_start: bool,
941    saw_response_stop: bool,
942    terminal_error: Option<Error>,
943    /// Skip block events; see [`metadata_only`](Self::metadata_only).
944    metadata_only: bool,
945}
946
947enum Slot {
948    Open(BlockAccumulator),
949    Closed {
950        block_type: StreamBlockType,
951        block: Finished,
952    },
953}
954
955/// What a block turned into.
956enum Finished {
957    Block(ContentBlock),
958    /// Text or reasoning with nothing in it.
959    Empty,
960    /// A closed tool call whose arguments are not valid JSON.
961    InvalidToolCall(ToolCall),
962    /// A tool call without an id or a name, or one that never closed.
963    IncompleteToolCall,
964}
965
966struct BlockAccumulator {
967    kind: BlockKind,
968    provider_data: ProviderData,
969    started: bool,
970}
971
972impl BlockAccumulator {
973    fn finish(self, closed: bool) -> Finished {
974        let provider_data = self.provider_data;
975        match self.kind {
976            BlockKind::Text { text, citations }
977                if text.is_empty() && citations.is_empty() && provider_data.is_empty() =>
978            {
979                Finished::Empty
980            }
981            BlockKind::Text { text, citations } => Finished::Block(ContentBlock::Text(
982                Text::new(text)
983                    .with_citations(citations)
984                    .with_provider_data(provider_data),
985            )),
986            BlockKind::Image { source, media_type } => {
987                let mut image = Image::new(source);
988                image.media_type = media_type;
989                image.provider_data = provider_data;
990                Finished::Block(ContentBlock::Image(image))
991            }
992            BlockKind::Reasoning(text) if text.is_empty() && provider_data.is_empty() => {
993                Finished::Empty
994            }
995            BlockKind::Reasoning(text) => Finished::Block(ContentBlock::Reasoning(
996                Reasoning::new(text).with_provider_data(provider_data),
997            )),
998            BlockKind::Refusal(text) if text.is_empty() && provider_data.is_empty() => {
999                Finished::Empty
1000            }
1001            BlockKind::Refusal(text) => Finished::Block(ContentBlock::Refusal(
1002                Refusal::new(text).with_provider_data(provider_data),
1003            )),
1004            BlockKind::ToolCall {
1005                id: Some(id),
1006                name: Some(name),
1007                arguments,
1008            } if closed => {
1009                let arguments = if arguments.trim().is_empty() {
1010                    "{}".to_owned()
1011                } else {
1012                    arguments
1013                };
1014                let valid = is_valid_json(&arguments);
1015                let call = ToolCall::new(id, name, arguments).with_provider_data(provider_data);
1016                if valid {
1017                    Finished::Block(ContentBlock::ToolCall(call))
1018                } else {
1019                    Finished::InvalidToolCall(call)
1020                }
1021            }
1022            BlockKind::ToolCall { .. } => Finished::IncompleteToolCall,
1023            BlockKind::ProviderToolCall {
1024                id: Some(id),
1025                name: Some(name),
1026                arguments,
1027                origin,
1028            } if closed => {
1029                let arguments = if arguments.trim().is_empty() {
1030                    "{}".to_owned()
1031                } else {
1032                    arguments
1033                };
1034                let mut call =
1035                    ProviderToolCall::new(id, name, arguments).with_provider_data(provider_data);
1036                call.origin = origin;
1037                Finished::Block(ContentBlock::ProviderToolCall(call))
1038            }
1039            BlockKind::ProviderToolCall { .. } => Finished::IncompleteToolCall,
1040            BlockKind::ProviderToolResult {
1041                call_id,
1042                output,
1043                is_error,
1044                origin,
1045            } => {
1046                let mut result = ProviderToolResult::new(call_id, output)
1047                    .with_error(is_error)
1048                    .with_provider_data(provider_data);
1049                result.origin = origin;
1050                Finished::Block(ContentBlock::ProviderToolResult(result))
1051            }
1052            BlockKind::Other {
1053                type_name,
1054                data,
1055                origin,
1056            } => {
1057                let mut other =
1058                    OtherContent::new(type_name, data).with_provider_data(provider_data);
1059                other.origin = origin;
1060                Finished::Block(ContentBlock::Other(other))
1061            }
1062        }
1063    }
1064}
1065
1066enum BlockKind {
1067    Text {
1068        text: String,
1069        citations: Vec<Citation>,
1070    },
1071    Image {
1072        source: MediaSource,
1073        /// The declared media type of a URL image.
1074        media_type: Option<String>,
1075    },
1076    Reasoning(String),
1077    Refusal(String),
1078    ToolCall {
1079        id: Option<String>,
1080        name: Option<String>,
1081        arguments: String,
1082    },
1083    ProviderToolCall {
1084        id: Option<String>,
1085        name: Option<String>,
1086        arguments: String,
1087        origin: Option<String>,
1088    },
1089    ProviderToolResult {
1090        call_id: String,
1091        output: serde_json::Value,
1092        is_error: bool,
1093        origin: Option<String>,
1094    },
1095    Other {
1096        type_name: String,
1097        data: ExtraMap,
1098        origin: Option<String>,
1099    },
1100}
1101
1102impl BlockKind {
1103    fn block_type(&self) -> StreamBlockType {
1104        match self {
1105            Self::Text { .. } => StreamBlockType::Text,
1106            Self::Image { .. } => StreamBlockType::Image,
1107            Self::Reasoning(_) => StreamBlockType::Reasoning,
1108            Self::Refusal(_) => StreamBlockType::Refusal,
1109            Self::ToolCall { .. } => StreamBlockType::ToolCall,
1110            Self::ProviderToolCall { .. } => StreamBlockType::ProviderToolCall,
1111            Self::ProviderToolResult { .. } => StreamBlockType::ProviderToolResult,
1112            Self::Other { .. } => StreamBlockType::Other,
1113        }
1114    }
1115
1116    fn label(&self) -> &'static str {
1117        self.block_type().label()
1118    }
1119
1120    fn text() -> Self {
1121        Self::Text {
1122            text: String::new(),
1123            citations: Vec::new(),
1124        }
1125    }
1126}
1127
1128impl StreamCollector {
1129    /// Create an empty collector.
1130    #[must_use]
1131    pub fn new() -> Self {
1132        Self::default()
1133    }
1134
1135    /// Create a collector that merges response metadata and ignores
1136    /// content blocks.
1137    ///
1138    /// Response-level events (`ResponseStart`, `ResponseMetadata`,
1139    /// `ResponseProviderData`, `Warning`, `ResponseStop`) are merged exactly
1140    /// as [`new`](Self::new) merges them. Usage is folded by each event's
1141    /// [`UsageMetadataMode`]: a `Snapshot` replaces the usage so far, a
1142    /// `Delta` adds to it. Block events are dropped without buffering or
1143    /// protocol checks, so the collected response has no content and never
1144    /// reports open blocks or dropped tool calls.
1145    ///
1146    /// For observers that need usage, finish reason, id and model but not
1147    /// the output itself, such as a budget guard that passes a stream on
1148    /// and counts its tokens. Whether a stream reports usage at all is the
1149    /// provider's [`ChatCapability::StreamUsage`](crate::ChatCapability::StreamUsage)
1150    /// answer for the model.
1151    ///
1152    /// ```
1153    /// use allama::chat::stream::{MetadataBuilder, StreamCollector, UsageMetadataMode};
1154    /// use allama::{StreamEvent, Usage};
1155    ///
1156    /// let delta = |usage: Usage| {
1157    ///     MetadataBuilder::new()
1158    ///         .with_usage(usage)
1159    ///         .with_usage_mode(UsageMetadataMode::Delta)
1160    ///         .build()
1161    /// };
1162    /// let mut budget = StreamCollector::metadata_only();
1163    /// for event in [
1164    ///     StreamEvent::response_start(Some("resp_1".into()), None),
1165    ///     delta(Usage::new().input_tokens(12).output_tokens(1)),
1166    ///     StreamEvent::text_delta(0, "Hello"),
1167    ///     delta(Usage::new().output_tokens(4)),
1168    ///     StreamEvent::ResponseStop,
1169    /// ] {
1170    ///     budget.push_ref(&event)?; // and pass `event` on
1171    /// }
1172    ///
1173    /// let collected = budget.finish_partial()?;
1174    /// let usage = collected.response.usage.expect("the stream reported usage");
1175    /// assert_eq!((usage.input_tokens, usage.output_tokens), (Some(12), Some(5)));
1176    /// assert!(collected.response.content.is_empty());
1177    /// # Ok::<(), allama::Error>(())
1178    /// ```
1179    #[must_use]
1180    pub fn metadata_only() -> Self {
1181        Self {
1182            metadata_only: true,
1183            ..Self::default()
1184        }
1185    }
1186
1187    /// Feed one event.
1188    ///
1189    /// # Errors
1190    ///
1191    /// Returns a protocol [`Error::Stream`] when events violate the stream
1192    /// protocol (events after `ResponseStop`, conflicting blocks, …).
1193    pub fn push(&mut self, event: StreamEvent) -> Result<Option<&ContentBlock>> {
1194        if self.saw_response_stop {
1195            return Err(Error::stream_protocol(match event {
1196                StreamEvent::ResponseStop => "duplicate response_stop".to_owned(),
1197                other => format!("{} received after response_stop", other.name()),
1198            }));
1199        }
1200
1201        if self.metadata_only && event.is_block_event() {
1202            return Ok(None);
1203        }
1204
1205        match event {
1206            StreamEvent::ResponseStart { id, model } => {
1207                self.saw_response_start = true;
1208                merge(&mut self.response.id, id);
1209                merge(&mut self.response.model, model);
1210            }
1211            StreamEvent::BlockStart {
1212                index,
1213                block_type,
1214                id,
1215                name,
1216                type_name,
1217                data,
1218                origin,
1219                provider_data,
1220            } => self.start_block(
1221                index,
1222                block_type,
1223                id,
1224                name,
1225                type_name,
1226                data,
1227                origin,
1228                provider_data,
1229            )?,
1230            StreamEvent::TextDelta { index, text } => {
1231                match &mut self.block(index, BlockKind::text)?.kind {
1232                    BlockKind::Text { text: buf, .. } => buf.push_str(&text),
1233                    other => return Err(mismatch("text_delta", other, index)),
1234                }
1235            }
1236            StreamEvent::ReasoningDelta { index, text } => {
1237                match &mut self
1238                    .block(index, || BlockKind::Reasoning(String::new()))?
1239                    .kind
1240                {
1241                    BlockKind::Reasoning(buf) => buf.push_str(&text),
1242                    other => return Err(mismatch("reasoning_delta", other, index)),
1243                }
1244            }
1245            StreamEvent::RefusalDelta { index, text } => {
1246                match &mut self
1247                    .block(index, || BlockKind::Refusal(String::new()))?
1248                    .kind
1249                {
1250                    BlockKind::Refusal(buf) => buf.push_str(&text),
1251                    other => return Err(mismatch("refusal_delta", other, index)),
1252                }
1253            }
1254            StreamEvent::ToolCallDelta { index, arguments } => {
1255                let block = self.block(index, || BlockKind::ToolCall {
1256                    id: None,
1257                    name: None,
1258                    arguments: String::new(),
1259                })?;
1260                match &mut block.kind {
1261                    BlockKind::ToolCall { arguments: buf, .. }
1262                    | BlockKind::ProviderToolCall { arguments: buf, .. } => {
1263                        buf.push_str(&arguments)
1264                    }
1265                    other => return Err(mismatch("tool_call_delta", other, index)),
1266                }
1267            }
1268            StreamEvent::Citation { index, citation } => {
1269                match &mut self.block(index, BlockKind::text)?.kind {
1270                    BlockKind::Text { citations, .. } => citations.push(citation),
1271                    other => return Err(mismatch("citation", other, index)),
1272                }
1273            }
1274            StreamEvent::BlockProviderData {
1275                index,
1276                provider_data,
1277            } => {
1278                let block = self.open_block(index, "block_provider_data")?;
1279                block.provider_data.merge(provider_data);
1280            }
1281            StreamEvent::ResponseProviderData { provider_data } => {
1282                self.response.provider_data.merge(provider_data);
1283            }
1284            StreamEvent::ProviderDelta { .. } => {}
1285            StreamEvent::BlockStop { index } => return self.stop_block(index),
1286            StreamEvent::Warning { warning } => self.response.warnings.push(warning),
1287            StreamEvent::ResponseMetadata {
1288                finish_reason,
1289                finish_reason_raw,
1290                usage,
1291                usage_mode,
1292                id,
1293                model,
1294                request_id,
1295                provider,
1296                provider_data,
1297            } => {
1298                merge(&mut self.response.finish_reason, finish_reason);
1299                merge(&mut self.response.finish_reason_raw, finish_reason_raw);
1300                merge(&mut self.response.id, id);
1301                merge(&mut self.response.model, model);
1302                merge(&mut self.response.request_id, request_id);
1303                merge(&mut self.response.provider, provider);
1304                self.response.provider_data.merge(provider_data);
1305                if let Some(usage) = usage {
1306                    match (usage_mode, &mut self.response.usage) {
1307                        (UsageMetadataMode::Delta, Some(existing)) => *existing += usage,
1308                        _ => self.response.usage = Some(usage),
1309                    }
1310                }
1311            }
1312            StreamEvent::ResponseStop => self.saw_response_stop = true,
1313        }
1314        Ok(None)
1315    }
1316
1317    /// Feed one event by reference.
1318    ///
1319    /// # Errors
1320    ///
1321    /// See [`push`](Self::push).
1322    pub fn push_ref(&mut self, event: &StreamEvent) -> Result<Option<&ContentBlock>> {
1323        // A skipped block event needs no copy.
1324        if self.metadata_only && event.is_block_event() && !self.saw_response_stop {
1325            return Ok(None);
1326        }
1327        self.push(event.clone())
1328    }
1329
1330    /// Record the error that ended the stream.
1331    pub fn set_terminal_error(&mut self, error: Error) {
1332        self.terminal_error = Some(error);
1333    }
1334
1335    /// The response-level metadata merged so far (usage, warnings, id,
1336    /// model, finish reason), without content: blocks are assembled when
1337    /// the collector finishes.
1338    #[cfg_attr(not(feature = "extract"), allow(dead_code))]
1339    pub(crate) fn response_metadata(&self) -> &ChatResponse {
1340        &self.response
1341    }
1342
1343    /// The type of the block at `index`, open or closed, or `None` when no
1344    /// block has that index.
1345    ///
1346    /// Lets a UI route a delta without tracking `BlockStart`s itself: a
1347    /// `ToolCallDelta` belongs to a client call
1348    /// ([`StreamBlockType::ToolCall`]) or a provider-executed one
1349    /// ([`StreamBlockType::ProviderToolCall`]). A block that began with a
1350    /// delta has the type that delta implies.
1351    #[must_use]
1352    pub fn block_type(&self, index: usize) -> Option<StreamBlockType> {
1353        match self.blocks.get(index)?.as_ref()? {
1354            Slot::Open(block) => Some(block.kind.block_type()),
1355            Slot::Closed { block_type, .. } => Some(*block_type),
1356        }
1357    }
1358
1359    /// Finish collection, requiring a complete stream.
1360    ///
1361    /// A complete stream collects to `Ok` whatever its finish reason, as a
1362    /// unary response does: a tool call whose arguments are not valid JSON
1363    /// (cut off at the output cap, or malformed) is kept with the arguments
1364    /// as they arrived ([`ToolCall::has_valid_arguments`] is `false`).
1365    /// Errors are reserved for streams that ended early or broke the
1366    /// protocol.
1367    ///
1368    /// # Errors
1369    ///
1370    /// Returns an [`Error::Stream`] of kind
1371    /// [`Incomplete`](StreamErrorKind::Incomplete) (retryable) when the
1372    /// stream ended early, or the terminal error when one was recorded.
1373    pub fn finish(self) -> Result<ChatResponse> {
1374        let unidentified = self.has_unidentified_tool_call();
1375        let collected = self.build()?;
1376        if let Some(error) = collected.terminal_error {
1377            return Err(error);
1378        }
1379        let c = &collected.completeness;
1380        if !c.is_complete() {
1381            // An adapter that opened a tool call with `block_start` gave it
1382            // no id or name: say which constructor carries them.
1383            let hint = if unidentified {
1384                "; a tool call has no id or name: start it with \
1385                 `StreamEvent::tool_call_start` (`provider_tool_call_start` for a \
1386                 provider-executed one), not `block_start`"
1387            } else {
1388                ""
1389            };
1390            return Err(Error::stream(
1391                StreamErrorKind::Incomplete,
1392                format!(
1393                    "stream incomplete: start={}, stop={}, open_blocks={}, dropped_incomplete_tool_calls={}{hint}",
1394                    c.saw_response_start,
1395                    c.saw_response_stop,
1396                    c.has_open_blocks,
1397                    c.dropped_incomplete_tool_calls
1398                ),
1399            ));
1400        }
1401        Ok(collected.response)
1402    }
1403
1404    /// Finish collection, returning whatever was assembled.
1405    ///
1406    /// # Errors
1407    ///
1408    /// Never fails for incomplete streams; reserved for future validation.
1409    pub fn finish_partial(self) -> Result<CollectedResponse> {
1410        self.build()
1411    }
1412
1413    /// The open block at `index`, created from `default` when the stream
1414    /// sends deltas without a `BlockStart`.
1415    fn block(
1416        &mut self,
1417        index: usize,
1418        default: impl FnOnce() -> BlockKind,
1419    ) -> Result<&mut BlockAccumulator> {
1420        self.reserve(index)?;
1421        let slot = self.blocks[index].get_or_insert_with(|| {
1422            Slot::Open(BlockAccumulator {
1423                kind: default(),
1424                provider_data: ProviderData::new(),
1425                started: false,
1426            })
1427        });
1428        match slot {
1429            Slot::Open(block) => Ok(block),
1430            Slot::Closed { block_type, .. } => Err(Error::stream_protocol(format!(
1431                "delta received for closed {} block at index {index}",
1432                block_type.label()
1433            ))),
1434        }
1435    }
1436
1437    /// Make room for a block at `index`, refusing to skip more than
1438    /// [`MAX_BLOCK_INDEX_GAP`] indices over the whole stream: indices come
1439    /// from the wire, and neither one absurd index nor many small gaps may
1440    /// allocate memory for blocks that never arrive (or overflow). Memory
1441    /// stays linear in the events pushed.
1442    fn reserve(&mut self, index: usize) -> Result<()> {
1443        if index >= self.blocks.len() {
1444            let gap = index - self.blocks.len();
1445            let skipped = self.skipped_indices.saturating_add(gap);
1446            if skipped > MAX_BLOCK_INDEX_GAP {
1447                return Err(Error::stream_protocol(format!(
1448                    "block index {index} leaves {skipped} block indices unused in this \
1449                     stream (at most {MAX_BLOCK_INDEX_GAP}); after {} blocks",
1450                    self.blocks.len()
1451                )));
1452            }
1453            self.skipped_indices = skipped;
1454            self.blocks.resize_with(index + 1, || None);
1455        }
1456        Ok(())
1457    }
1458
1459    /// The open block at `index`, which must exist.
1460    fn open_block(&mut self, index: usize, event: &str) -> Result<&mut BlockAccumulator> {
1461        match self.blocks.get_mut(index).and_then(Option::as_mut) {
1462            Some(Slot::Open(block)) => Ok(block),
1463            Some(Slot::Closed { block_type, .. }) => Err(Error::stream_protocol(format!(
1464                "{event} for closed {} block at index {index}",
1465                block_type.label()
1466            ))),
1467            None => Err(Error::stream_protocol(format!(
1468                "{event} for unknown index {index}"
1469            ))),
1470        }
1471    }
1472
1473    fn stop_block(&mut self, index: usize) -> Result<Option<&ContentBlock>> {
1474        let slot = self
1475            .blocks
1476            .get_mut(index)
1477            .and_then(Option::take)
1478            .ok_or_else(|| {
1479                Error::stream_protocol(format!("block_stop for unknown index {index}"))
1480            })?;
1481        let closed = match slot {
1482            Slot::Open(block) => Slot::Closed {
1483                block_type: block.kind.block_type(),
1484                block: block.finish(true),
1485            },
1486            Slot::Closed { block_type, block } => {
1487                self.blocks[index] = Some(Slot::Closed { block_type, block });
1488                return Err(Error::stream_protocol(format!(
1489                    "duplicate block_stop for {} block at index {index}",
1490                    block_type.label()
1491                )));
1492            }
1493        };
1494        let slot = self.blocks[index].insert(closed);
1495        Ok(match slot {
1496            Slot::Closed {
1497                block: Finished::Block(block),
1498                ..
1499            } => Some(block),
1500            _ => None,
1501        })
1502    }
1503
1504    #[allow(clippy::too_many_arguments)]
1505    fn start_block(
1506        &mut self,
1507        index: usize,
1508        block_type: StreamBlockType,
1509        id: Option<String>,
1510        name: Option<String>,
1511        type_name: Option<String>,
1512        data: Option<ExtraMap>,
1513        origin: Option<String>,
1514        provider_data: ProviderData,
1515    ) -> Result<()> {
1516        let fresh = match block_type {
1517            StreamBlockType::Text => BlockKind::text(),
1518            StreamBlockType::Reasoning => BlockKind::Reasoning(String::new()),
1519            StreamBlockType::Refusal => BlockKind::Refusal(String::new()),
1520            StreamBlockType::ToolCall => BlockKind::ToolCall {
1521                id: id.clone(),
1522                name: name.clone(),
1523                arguments: String::new(),
1524            },
1525            StreamBlockType::Image => {
1526                let (source, media_type) = image_from_data(data.clone())?;
1527                BlockKind::Image { source, media_type }
1528            }
1529            StreamBlockType::ProviderToolCall => BlockKind::ProviderToolCall {
1530                id: id.clone(),
1531                name: name.clone(),
1532                arguments: String::new(),
1533                origin,
1534            },
1535            StreamBlockType::ProviderToolResult => {
1536                let call_id = id.clone().ok_or_else(|| {
1537                    Error::stream_protocol(format!(
1538                        "provider_tool_result block at index {index} has no call id"
1539                    ))
1540                })?;
1541                let mut data = data.clone().unwrap_or_default();
1542                BlockKind::ProviderToolResult {
1543                    call_id,
1544                    output: data.remove("output").unwrap_or_default(),
1545                    is_error: data
1546                        .remove("is_error")
1547                        .and_then(|v| v.as_bool())
1548                        .unwrap_or(false),
1549                    origin,
1550                }
1551            }
1552            StreamBlockType::Other => BlockKind::Other {
1553                type_name: type_name.unwrap_or_else(|| "other".into()),
1554                data: data.unwrap_or_default(),
1555                origin,
1556            },
1557        };
1558
1559        self.reserve(index)?;
1560        let slot = &mut self.blocks[index];
1561        match slot {
1562            None => {
1563                *slot = Some(Slot::Open(BlockAccumulator {
1564                    kind: fresh,
1565                    provider_data,
1566                    started: true,
1567                }));
1568                Ok(())
1569            }
1570            Some(Slot::Closed { block_type, .. }) => Err(Error::stream_protocol(format!(
1571                "block_start for closed {} block at index {index}",
1572                block_type.label()
1573            ))),
1574            Some(Slot::Open(block)) => {
1575                let label = block.kind.label();
1576                if block.started {
1577                    return Err(Error::stream_protocol(format!(
1578                        "duplicate block_start for {label} block at index {index}"
1579                    )));
1580                }
1581                if label != fresh.label() {
1582                    return Err(Error::stream_protocol(format!(
1583                        "block_start for {} conflicted with existing {label} block at index {index}",
1584                        fresh.label()
1585                    )));
1586                }
1587                block.started = true;
1588                block.provider_data.merge(provider_data);
1589                if let BlockKind::ToolCall {
1590                    id: existing_id,
1591                    name: existing_name,
1592                    ..
1593                } = &mut block.kind
1594                {
1595                    merge_strict(existing_id, id, index, "id")?;
1596                    merge_strict(existing_name, name, index, "name")?;
1597                }
1598                Ok(())
1599            }
1600        }
1601    }
1602
1603    /// Whether a tool call ended up without its id or name.
1604    fn has_unidentified_tool_call(&self) -> bool {
1605        self.blocks.iter().flatten().any(|slot| match slot {
1606            Slot::Open(block) => matches!(
1607                &block.kind,
1608                BlockKind::ToolCall { id, name, .. }
1609                    | BlockKind::ProviderToolCall { id, name, .. }
1610                    if id.is_none() || name.is_none()
1611            ),
1612            // A closed call is incomplete only when it lacks either.
1613            Slot::Closed { block, .. } => matches!(block, Finished::IncompleteToolCall),
1614        })
1615    }
1616
1617    fn build(self) -> Result<CollectedResponse> {
1618        let mut content = Vec::with_capacity(self.blocks.len());
1619        let mut has_open_blocks = false;
1620        let mut dropped_incomplete_tool_calls = false;
1621
1622        for slot in self.blocks.into_iter().flatten() {
1623            let finished = match slot {
1624                Slot::Open(block) => {
1625                    has_open_blocks = true;
1626                    block.finish(false)
1627                }
1628                Slot::Closed { block, .. } => block,
1629            };
1630            match finished {
1631                Finished::Block(block) => content.push(block),
1632                Finished::Empty => {}
1633                // A closed call keeps its arguments as they arrived, as a
1634                // unary response does: `ToolCall::parse_arguments` reports
1635                // them, and the finish reason says why (an output cap cut
1636                // them off, for example).
1637                Finished::InvalidToolCall(call) => content.push(ContentBlock::ToolCall(call)),
1638                Finished::IncompleteToolCall => {
1639                    dropped_incomplete_tool_calls = true;
1640                }
1641            }
1642        }
1643
1644        let mut response = self.response;
1645        response.content = content;
1646        Ok(CollectedResponse {
1647            response,
1648            completeness: StreamCompleteness {
1649                saw_response_start: self.saw_response_start,
1650                saw_response_stop: self.saw_response_stop,
1651                has_open_blocks,
1652                dropped_incomplete_tool_calls,
1653            },
1654            terminal_error: self.terminal_error,
1655        })
1656    }
1657}
1658
1659/// Counts and state, never content.
1660impl fmt::Debug for StreamCollector {
1661    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1662        f.debug_struct("StreamCollector")
1663            .field("blocks", &self.blocks.len())
1664            .field("saw_response_start", &self.saw_response_start)
1665            .field("saw_response_stop", &self.saw_response_stop)
1666            .field(
1667                "terminal_error",
1668                &self.terminal_error.as_ref().map(Error::kind_name),
1669            )
1670            .field("metadata_only", &self.metadata_only)
1671            .finish_non_exhaustive()
1672    }
1673}
1674
1675/// How many block indices a stream may skip in total. Decoders number
1676/// blocks densely, so more gaps are a protocol violation, not a stream to
1677/// allocate for.
1678const MAX_BLOCK_INDEX_GAP: usize = 1024;
1679
1680fn merge<T>(slot: &mut Option<T>, value: Option<T>) {
1681    if value.is_some() {
1682        *slot = value;
1683    }
1684}
1685
1686fn merge_strict(
1687    existing: &mut Option<String>,
1688    incoming: Option<String>,
1689    index: usize,
1690    field: &str,
1691) -> Result<()> {
1692    match (existing.as_ref(), incoming) {
1693        (Some(current), Some(new)) if current != &new => Err(Error::stream_protocol(format!(
1694            "conflicting {field} for tool call block at index {index}"
1695        ))),
1696        (None, Some(new)) => {
1697            *existing = Some(new);
1698            Ok(())
1699        }
1700        _ => Ok(()),
1701    }
1702}
1703
1704fn mismatch(event: &str, kind: &BlockKind, index: usize) -> Error {
1705    Error::stream_protocol(format!(
1706        "{event} received for {} block at index {index}",
1707        kind.label()
1708    ))
1709}
1710
1711/// An image block's source and declared media type, from its `BlockStart`
1712/// data.
1713fn image_from_data(data: Option<ExtraMap>) -> Result<(MediaSource, Option<String>)> {
1714    let mut data = data.ok_or_else(|| Error::stream_protocol("image block missing source"))?;
1715    let source = data
1716        .remove("source")
1717        .ok_or_else(|| Error::stream_protocol("image block missing source"))?;
1718    let source = serde_json::from_value(source)
1719        .map_err(|e| Error::stream_decode(format!("image block source was invalid: {e}")))?;
1720    let media_type = match data.remove("media_type") {
1721        Some(serde_json::Value::String(media_type)) => Some(media_type),
1722        _ => None,
1723    };
1724    Ok((source, media_type))
1725}
1726
1727fn image_data(source: &MediaSource, media_type: Option<&str>) -> ExtraMap {
1728    let mut data = ExtraMap::new();
1729    data.insert(
1730        "source".into(),
1731        serde_json::to_value(source).expect("media source serializes"),
1732    );
1733    if let Some(media_type) = media_type {
1734        data.insert(
1735            "media_type".into(),
1736            serde_json::Value::String(media_type.to_owned()),
1737        );
1738    }
1739    data
1740}
1741
1742/// A stream that replays a complete [`ChatResponse`] as events.
1743///
1744/// Providers without native streaming return this from `chat_stream`.
1745pub struct SingleResponseStream {
1746    events: std::vec::IntoIter<StreamEvent>,
1747}
1748
1749impl SingleResponseStream {
1750    /// Create a stream from a response.
1751    #[must_use]
1752    pub fn new(response: ChatResponse) -> Self {
1753        Self {
1754            events: response
1755                .into_stream_events()
1756                .collect::<Vec<_>>()
1757                .into_iter(),
1758        }
1759    }
1760}
1761
1762impl fmt::Debug for SingleResponseStream {
1763    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1764        f.debug_struct("SingleResponseStream")
1765            .field("remaining_events", &self.events.len())
1766            .finish()
1767    }
1768}
1769
1770impl From<ChatResponse> for SingleResponseStream {
1771    fn from(response: ChatResponse) -> Self {
1772        Self::new(response)
1773    }
1774}
1775
1776impl Stream for SingleResponseStream {
1777    type Item = Result<StreamEvent>;
1778
1779    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1780        Poll::Ready(self.events.next().map(Ok))
1781    }
1782}
1783
1784impl ChatResponse {
1785    /// Replay this response as a stream.
1786    #[must_use]
1787    pub fn into_stream(self) -> SingleResponseStream {
1788        SingleResponseStream::new(self)
1789    }
1790
1791    /// Replay this response as stream events.
1792    pub fn into_stream_events(self) -> impl Iterator<Item = StreamEvent> {
1793        let ChatResponse {
1794            content,
1795            finish_reason,
1796            finish_reason_raw,
1797            usage,
1798            model,
1799            id,
1800            request_id,
1801            provider,
1802            warnings,
1803            provider_data,
1804            finish_reported: _,
1805        } = self;
1806
1807        let mut events = Vec::with_capacity(content.len() * 3 + warnings.len() + 3);
1808        events.push(StreamEvent::ResponseStart {
1809            id: id.clone(),
1810            model: model.clone(),
1811        });
1812        events.extend(
1813            warnings
1814                .into_iter()
1815                .map(|warning| StreamEvent::Warning { warning }),
1816        );
1817        for (index, block) in content.into_iter().enumerate() {
1818            block_events(index, block, &mut events);
1819        }
1820        events.push(StreamEvent::ResponseMetadata {
1821            finish_reason,
1822            finish_reason_raw,
1823            usage,
1824            usage_mode: UsageMetadataMode::Snapshot,
1825            id,
1826            model,
1827            request_id,
1828            provider,
1829            provider_data,
1830        });
1831        events.push(StreamEvent::ResponseStop);
1832        events.into_iter()
1833    }
1834}
1835
1836fn block_events(index: usize, block: ContentBlock, events: &mut Vec<StreamEvent>) {
1837    let mut start = |block_type, id, name, type_name, data, origin, provider_data| {
1838        events.push(StreamEvent::BlockStart {
1839            index,
1840            block_type,
1841            id,
1842            name,
1843            type_name,
1844            data,
1845            origin,
1846            provider_data,
1847        });
1848    };
1849    match block {
1850        ContentBlock::Text(text) => {
1851            start(
1852                StreamBlockType::Text,
1853                None,
1854                None,
1855                None,
1856                None,
1857                None,
1858                text.provider_data,
1859            );
1860            if !text.text.is_empty() {
1861                events.push(StreamEvent::TextDelta {
1862                    index,
1863                    text: text.text,
1864                });
1865            }
1866            events.extend(
1867                text.citations
1868                    .into_iter()
1869                    .map(|citation| StreamEvent::Citation { index, citation }),
1870            );
1871        }
1872        ContentBlock::Image(image) => events.push(StreamEvent::image_start(index, image)),
1873        ContentBlock::Reasoning(reasoning) => {
1874            start(
1875                StreamBlockType::Reasoning,
1876                None,
1877                None,
1878                None,
1879                None,
1880                None,
1881                reasoning.provider_data,
1882            );
1883            if !reasoning.text.is_empty() {
1884                events.push(StreamEvent::ReasoningDelta {
1885                    index,
1886                    text: reasoning.text,
1887                });
1888            }
1889        }
1890        ContentBlock::Refusal(refusal) => {
1891            start(
1892                StreamBlockType::Refusal,
1893                None,
1894                None,
1895                None,
1896                None,
1897                None,
1898                refusal.provider_data,
1899            );
1900            if !refusal.text.is_empty() {
1901                events.push(StreamEvent::RefusalDelta {
1902                    index,
1903                    text: refusal.text,
1904                });
1905            }
1906        }
1907        ContentBlock::ToolCall(call) => {
1908            start(
1909                StreamBlockType::ToolCall,
1910                Some(call.id),
1911                Some(call.name),
1912                None,
1913                None,
1914                None,
1915                call.provider_data,
1916            );
1917            events.push(StreamEvent::ToolCallDelta {
1918                index,
1919                arguments: call.arguments,
1920            });
1921        }
1922        ContentBlock::ProviderToolCall(call) => {
1923            start(
1924                StreamBlockType::ProviderToolCall,
1925                Some(call.id),
1926                Some(call.name),
1927                None,
1928                None,
1929                call.origin,
1930                call.provider_data,
1931            );
1932            events.push(StreamEvent::ToolCallDelta {
1933                index,
1934                arguments: call.arguments,
1935            });
1936        }
1937        ContentBlock::ProviderToolResult(result) => {
1938            events.push(StreamEvent::provider_tool_result_start(index, result));
1939        }
1940        ContentBlock::Other(other) => events.push(StreamEvent::other_start(index, other)),
1941    }
1942    events.push(StreamEvent::BlockStop { index });
1943}
1944
1945/// Extension methods for chat streams.
1946pub trait ChatStreamExt: Stream<Item = Result<StreamEvent>> + Send {
1947    /// Collect a complete response.
1948    ///
1949    /// # Errors
1950    ///
1951    /// Fails on the first stream error, on protocol violations, and with a
1952    /// retryable [`StreamErrorKind::Incomplete`] error when the stream ends
1953    /// early.
1954    fn collect_response(self) -> impl std::future::Future<Output = Result<ChatResponse>> + Send;
1955
1956    /// The next event, or `None` at the end of the stream.
1957    ///
1958    /// The same as `StreamExt::next` from `futures-util`, without importing
1959    /// a stream extension trait: a loop over a chat stream needs nothing
1960    /// beyond [`allama::prelude`](crate::prelude), and a `tokio_stream` or
1961    /// `futures` import of your own never makes the call ambiguous.
1962    ///
1963    /// ```
1964    /// # #[cfg(feature = "mock")]
1965    /// # async fn example() -> allama::Result<()> {
1966    /// use allama::prelude::*;
1967    /// use allama::mock::MockProvider;
1968    ///
1969    /// let provider = MockProvider::with_text("Hello");
1970    /// let mut stream = provider.chat_stream(&ChatRequest::new("m").user("hi")).await?;
1971    /// while let Some(event) = stream.next_event().await {
1972    ///     if let StreamEvent::TextDelta { text, .. } = event? {
1973    ///         print!("{text}");
1974    ///     }
1975    /// }
1976    /// # Ok(()) }
1977    /// ```
1978    fn next_event(
1979        &mut self,
1980    ) -> impl std::future::Future<Output = Option<Result<StreamEvent>>> + Send
1981    where
1982        Self: Unpin,
1983    {
1984        std::future::poll_fn(move |cx| Pin::new(&mut *self).poll_next(cx))
1985    }
1986
1987    /// Collect the response while streaming; each event comes with the block
1988    /// it completes.
1989    ///
1990    /// Use it to show events as they arrive and still get the
1991    /// [`ChatResponse`] at the end; use [`collect_response`](Self::collect_response)
1992    /// when only the response matters.
1993    ///
1994    /// Each item is a [`CollectedEvent`]: the [`event`](CollectedEvent::event)
1995    /// and, for a `BlockStop`, the block it [`completed`](CollectedEvent::completed),
1996    /// so a caller can act on a complete tool call or text block while the
1997    /// rest of the response is still streaming. Read the response with
1998    /// [`WithCollector::finish`]. [`WithCollector::next_event`] reads the
1999    /// next item without a stream extension trait in scope:
2000    ///
2001    /// ```
2002    /// # #[cfg(feature = "mock")]
2003    /// # async fn example() -> allama::Result<()> {
2004    /// use allama::prelude::*;
2005    /// use allama::mock::MockProvider;
2006    ///
2007    /// let provider = MockProvider::with_text("Hello");
2008    /// let mut stream = provider
2009    ///     .chat_stream(&ChatRequest::new("m").user("hi"))
2010    ///     .await?
2011    ///     .with_collector();
2012    /// while let Some(item) = stream.next_event().await {
2013    ///     let item = item?;
2014    ///     if let StreamEvent::TextDelta { text, .. } = &item.event {
2015    ///         print!("{text}");
2016    ///     }
2017    ///     if let Some(ContentBlock::ToolCall(call)) = &item.completed {
2018    ///         println!("call {}", call.name);
2019    ///     }
2020    /// }
2021    /// let response = stream.finish()?;
2022    /// # let _ = response; Ok(()) }
2023    /// ```
2024    fn with_collector(self) -> WithCollector<Self>
2025    where
2026        Self: Sized,
2027    {
2028        WithCollector {
2029            inner: Box::pin(self),
2030            collector: StreamCollector::new(),
2031            done: false,
2032        }
2033    }
2034}
2035
2036/// One item of a [`WithCollector`] stream: an event, and the block it
2037/// completed.
2038///
2039/// The struct is `#[non_exhaustive]`, so it can grow (a block index, a
2040/// usage snapshot, …) without breaking the loops that read it; read its
2041/// fields by name, or take it apart with [`into_parts`](Self::into_parts).
2042#[derive(Debug, Clone, PartialEq, Eq)]
2043#[non_exhaustive]
2044pub struct CollectedEvent {
2045    /// The event, as the stream yielded it.
2046    pub event: StreamEvent,
2047    /// The block this event completed: `Some` for the `BlockStop` that
2048    /// closed a block the collector assembled, `None` for every other
2049    /// event.
2050    ///
2051    /// A tool call's `BlockStop` has `None` for one of two reasons:
2052    ///
2053    /// - its arguments are not valid JSON. The call is still in the
2054    ///   collected response, with the arguments as they arrived
2055    ///   ([`ToolCall::has_valid_arguments`] is `false`); answer it with
2056    ///   [`ChatRequest::push_tool_error`](crate::ChatRequest::push_tool_error).
2057    /// - it is incomplete: it never got an id or a name. The call is dropped,
2058    ///   and the stream does not collect as complete
2059    ///   ([`StreamCompleteness::dropped_incomplete_tool_calls`]).
2060    pub completed: Option<ContentBlock>,
2061}
2062
2063impl CollectedEvent {
2064    /// An event and the block it completed (for test doubles and
2065    /// adapters of `WithCollector`-shaped streams).
2066    #[must_use]
2067    pub fn new(event: StreamEvent, completed: Option<ContentBlock>) -> Self {
2068        Self { event, completed }
2069    }
2070
2071    /// The event and the block it completed.
2072    #[must_use]
2073    pub fn into_parts(self) -> (StreamEvent, Option<ContentBlock>) {
2074        (self.event, self.completed)
2075    }
2076}
2077
2078impl From<CollectedEvent> for (StreamEvent, Option<ContentBlock>) {
2079    fn from(item: CollectedEvent) -> Self {
2080        item.into_parts()
2081    }
2082}
2083
2084/// A stream that collects the response while yielding each event with the
2085/// block it completes ([`CollectedEvent`]). Created by
2086/// [`ChatStreamExt::with_collector`].
2087///
2088/// A stream error is passed through, recorded as the terminal error, and
2089/// ends the stream, as does a protocol violation (events that cannot be
2090/// assembled), which is yielded as an error. The wrapped stream is not
2091/// polled after either.
2092pub struct WithCollector<S> {
2093    inner: Pin<Box<S>>,
2094    collector: StreamCollector,
2095    done: bool,
2096}
2097
2098impl<S> WithCollector<S>
2099where
2100    S: Stream<Item = Result<StreamEvent>>,
2101{
2102    /// The next item, or `None` at the end of the stream.
2103    ///
2104    /// The same as `StreamExt::next`, without importing a stream extension
2105    /// trait (see [`ChatStreamExt::next_event`]).
2106    pub async fn next_event(&mut self) -> Option<Result<CollectedEvent>> {
2107        std::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
2108    }
2109}
2110
2111impl<S> WithCollector<S> {
2112    /// Borrow the collector that has seen every event so far.
2113    #[must_use]
2114    pub fn collector(&self) -> &StreamCollector {
2115        &self.collector
2116    }
2117
2118    /// The collected response; see [`StreamCollector::finish`].
2119    ///
2120    /// # Errors
2121    ///
2122    /// Fails when the stream was incomplete or ended with an error.
2123    pub fn finish(self) -> Result<ChatResponse> {
2124        self.collector.finish()
2125    }
2126
2127    /// Whatever was collected; see [`StreamCollector::finish_partial`].
2128    ///
2129    /// # Errors
2130    ///
2131    /// Reserved for future validation.
2132    pub fn finish_partial(self) -> Result<CollectedResponse> {
2133        self.collector.finish_partial()
2134    }
2135}
2136
2137impl<S> fmt::Debug for WithCollector<S> {
2138    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2139        f.debug_struct("WithCollector")
2140            .field("collector", &self.collector)
2141            .field("done", &self.done)
2142            .finish_non_exhaustive()
2143    }
2144}
2145
2146impl<S> Stream for WithCollector<S>
2147where
2148    S: Stream<Item = Result<StreamEvent>>,
2149{
2150    type Item = Result<CollectedEvent>;
2151
2152    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2153        let this = &mut *self;
2154        if this.done {
2155            return Poll::Ready(None);
2156        }
2157        match this.inner.as_mut().poll_next(cx) {
2158            Poll::Ready(Some(Ok(event))) => match this.collector.push_ref(&event) {
2159                Ok(completed) => {
2160                    let completed = completed.cloned();
2161                    Poll::Ready(Some(Ok(CollectedEvent { event, completed })))
2162                }
2163                Err(error) => {
2164                    this.done = true;
2165                    this.collector.set_terminal_error(error.clone());
2166                    Poll::Ready(Some(Err(error)))
2167                }
2168            },
2169            Poll::Ready(Some(Err(error))) => {
2170                this.done = true;
2171                this.collector.set_terminal_error(error.clone());
2172                Poll::Ready(Some(Err(error)))
2173            }
2174            Poll::Ready(None) => {
2175                this.done = true;
2176                Poll::Ready(None)
2177            }
2178            Poll::Pending => Poll::Pending,
2179        }
2180    }
2181}
2182
2183impl<S: Stream<Item = Result<StreamEvent>> + Send> ChatStreamExt for S {
2184    async fn collect_response(self) -> Result<ChatResponse> {
2185        use futures_util::StreamExt;
2186        let mut collector = StreamCollector::new();
2187        let mut stream = std::pin::pin!(self);
2188        while let Some(event) = stream.next().await {
2189            collector.push(event?)?;
2190        }
2191        collector.finish()
2192    }
2193}
2194
2195/// A stream that first yields buffered events, then the rest of `inner`.
2196///
2197/// It is the committed stream [`until_committed`] returns: an ordinary
2198/// `Stream` of the same items, to use as you would the original.
2199///
2200/// [`until_committed`] returns one holding the events up to the commit
2201/// point. A wrapper that adds to that prefix before handing the stream on,
2202/// as [`FallbackProvider`](crate::wrappers::FallbackProvider) does with its
2203/// failed attempts, takes it apart with [`into_parts`](Self::into_parts)
2204/// and rebuilds it with [`new`](Self::new).
2205#[doc(alias = "CommittedStream")]
2206pub struct Prefixed<S> {
2207    buffer: VecDeque<StreamEvent>,
2208    inner: Option<S>,
2209}
2210
2211impl<S> Prefixed<S> {
2212    /// Create a prefixed stream.
2213    #[must_use]
2214    pub fn new(buffer: impl IntoIterator<Item = StreamEvent>, inner: Option<S>) -> Self {
2215        Self {
2216            buffer: buffer.into_iter().collect(),
2217            inner,
2218        }
2219    }
2220
2221    /// The buffered events not yet yielded, and the rest of the stream
2222    /// (`None` when it already ended).
2223    ///
2224    /// Pass both back to [`new`](Self::new) to resume the stream where it
2225    /// was, with any events added to the prefix:
2226    ///
2227    /// ```
2228    /// # tokio::runtime::Builder::new_current_thread().build().unwrap().block_on(async {
2229    /// use allama::StreamEvent;
2230    /// use allama::chat::stream::{Prefixed, until_committed};
2231    /// use futures_util::{StreamExt, stream};
2232    ///
2233    /// let events = stream::iter([
2234    ///     Ok(StreamEvent::response_start(None, None)),
2235    ///     Ok(StreamEvent::text_delta(0, "Hi")),
2236    /// ]);
2237    /// let committed = until_committed(events).await.unwrap();
2238    ///
2239    /// let (mut buffered, rest) = committed.into_parts();
2240    /// buffered.insert(1, StreamEvent::metadata());
2241    /// let resumed: Vec<_> = Prefixed::new(buffered, rest).collect().await;
2242    /// assert_eq!(resumed.len(), 3);
2243    /// # });
2244    /// ```
2245    #[must_use]
2246    pub fn into_parts(self) -> (Vec<StreamEvent>, Option<S>) {
2247        (self.buffer.into(), self.inner)
2248    }
2249}
2250
2251impl<S> fmt::Debug for Prefixed<S> {
2252    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2253        f.debug_struct("Prefixed")
2254            .field("buffered_events", &self.buffer.len())
2255            .field("ended", &self.inner.is_none())
2256            .finish()
2257    }
2258}
2259
2260impl<S: Stream<Item = Result<StreamEvent>> + Unpin> Stream for Prefixed<S> {
2261    type Item = Result<StreamEvent>;
2262
2263    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2264        if let Some(event) = self.buffer.pop_front() {
2265            return Poll::Ready(Some(Ok(event)));
2266        }
2267        match self.inner.as_mut() {
2268            Some(inner) => Pin::new(inner).poll_next(cx),
2269            None => Poll::Ready(None),
2270        }
2271    }
2272}
2273
2274/// Read `stream` until its [commit point](crate::chat::stream#the-commit-point):
2275/// the first event carrying model output (see [`StreamEvent::is_output`]),
2276/// or the end of the stream. Before that point an error can still be
2277/// retried, failed over, or answered with an HTTP status; after it, an
2278/// error can only end the stream.
2279///
2280/// Returns a [`Prefixed`] stream (an ordinary `Stream` of the same items)
2281/// that replays the buffered events followed by the rest, so the caller
2282/// observes every event and time-to-first-token is unchanged.
2283/// When an error arrives before the commit point, nothing has been exposed
2284/// yet, so the caller may safely retry or fail over.
2285///
2286/// The events received before such an error carry no model output, so they
2287/// are dropped with it.
2288///
2289/// ```
2290/// # tokio::runtime::Builder::new_current_thread().build().unwrap().block_on(async {
2291/// use allama::StreamEvent;
2292/// use allama::chat::stream::until_committed;
2293/// use futures_util::stream;
2294///
2295/// async fn respond() -> allama::Result<usize> {
2296///     let events = stream::iter([
2297///         Ok(StreamEvent::response_start(None, None)),
2298///         Ok(StreamEvent::text_delta(0, "Hi")),
2299///     ]);
2300///     let committed = until_committed(events).await?;
2301///     Ok(committed.into_parts().0.len())
2302/// }
2303///
2304/// async fn handler() -> anyhow::Result<usize> {
2305///     let failing = stream::iter([Err(allama::Error::overloaded("busy"))]);
2306///     let committed = until_committed(failing).await?;
2307///     Ok(committed.into_parts().0.len())
2308/// }
2309///
2310/// assert_eq!(respond().await.unwrap(), 2);
2311/// assert!(handler().await.is_err());
2312/// # });
2313/// ```
2314///
2315/// # Errors
2316///
2317/// Returns the stream's error when it fails before committing.
2318pub async fn until_committed<S>(mut stream: S) -> Result<Prefixed<S>>
2319where
2320    S: Stream<Item = Result<StreamEvent>> + Unpin,
2321{
2322    use futures_util::StreamExt;
2323    let mut buffered = Vec::new();
2324    loop {
2325        match stream.next().await {
2326            Some(Ok(event)) => {
2327                let committed = event.is_output();
2328                buffered.push(event);
2329                if committed {
2330                    return Ok(Prefixed::new(buffered, Some(stream)));
2331                }
2332            }
2333            Some(Err(error)) => return Err(error),
2334            None => return Ok(Prefixed::new(buffered, None)),
2335        }
2336    }
2337}
2338
2339#[cfg(test)]
2340mod tests {
2341    use super::*;
2342    use crate::{CitationSource, RequestField, WarningKind};
2343    use futures_util::StreamExt;
2344    use serde_json::json;
2345
2346    /// Application state holding a collecting stream can derive `Debug`.
2347    /// The output shows counts and state, never the streamed text.
2348    #[tokio::test]
2349    async fn stream_types_debug_without_content() {
2350        let stream: ChatStream =
2351            Box::pin(ChatResponse::new(vec![ContentBlock::text("secret answer")]).into_stream());
2352        let single = SingleResponseStream::from(ChatResponse::new(vec![ContentBlock::text(
2353            "secret answer",
2354        )]));
2355        assert!(format!("{single:?}").starts_with("SingleResponseStream { remaining_events: "));
2356
2357        let mut collecting = stream.with_collector();
2358        while collecting.next_event().await.is_some() {}
2359        let debug = format!("{collecting:?}");
2360        assert!(
2361            debug.starts_with("WithCollector { collector: StreamCollector { blocks: 1"),
2362            "{debug}"
2363        );
2364        assert!(!debug.contains("secret"), "{debug}");
2365
2366        let prefixed = Prefixed::<ChatStream>::new([StreamEvent::text_delta(0, "secret")], None);
2367        assert_eq!(
2368            format!("{prefixed:?}"),
2369            "Prefixed { buffered_events: 1, ended: true }"
2370        );
2371    }
2372
2373    #[test]
2374    fn fieldless_enums_are_hashable() {
2375        use std::collections::HashSet;
2376
2377        use crate::{CapabilitySupport, ChatCapability, StreamErrorKind};
2378        let capabilities: HashSet<(ChatCapability, CapabilitySupport)> =
2379            [(ChatCapability::ToolCalls, CapabilitySupport::Supported)].into();
2380        assert!(capabilities.contains(&(ChatCapability::ToolCalls, CapabilitySupport::Supported)));
2381        let _: HashSet<StreamErrorKind> = HashSet::new();
2382        let _: HashSet<UsageMetadataMode> = HashSet::new();
2383        #[cfg(feature = "extract")]
2384        {
2385            let _: HashSet<crate::extract::RefusalSource> = HashSet::new();
2386            let _: HashSet<crate::extract::ExtractionMode> = HashSet::new();
2387        }
2388    }
2389
2390    fn sample_response() -> ChatResponse {
2391        let mut call = ToolCall::new("c1", "lookup", r#"{"q":1}"#);
2392        call.provider_data
2393            .insert_raw("google-ai", "thought_signature", json!("sig"));
2394        ChatResponse::new(vec![
2395            ContentBlock::Reasoning(Reasoning::new("think").with_provider_data(
2396                ProviderData::new().with_raw("anthropic", "signature", json!("s")),
2397            )),
2398            ContentBlock::Text(Text::new("hello").with_citations(vec![
2399                Citation::new(CitationSource::url("https://x", Some("X".into()))).with_span(0..5),
2400            ])),
2401            ContentBlock::Refusal(Refusal::new("I can't help with that.")),
2402            ContentBlock::ToolCall(call),
2403            ContentBlock::Image(Image::url("https://x/y.png")),
2404            ContentBlock::Other(
2405                OtherContent::new("citation", ExtraMap::new())
2406                    .with_origin("anthropic")
2407                    .with_provider_data(ProviderData::new().with_raw(
2408                        "anthropic",
2409                        "signature",
2410                        json!("other-sig"),
2411                    )),
2412            ),
2413        ])
2414        .finish_reason(FinishReason::ToolCalls)
2415        .finish_reason_raw("tool_use")
2416        .usage(Usage::new().input_tokens(1).output_tokens(2))
2417        .id("r1")
2418        .model("m")
2419        .request_id("req")
2420        .provider("demo")
2421        .warning(Warning::unsupported(RequestField::Seed, "ignored"))
2422    }
2423
2424    #[tokio::test]
2425    async fn response_round_trips_through_stream() {
2426        let response = sample_response();
2427        let collected = response
2428            .clone()
2429            .into_stream()
2430            .collect_response()
2431            .await
2432            .unwrap();
2433        assert_eq!(collected, response);
2434    }
2435
2436    /// A refusal streams as its own block and delta, never as text, and
2437    /// collects back into a `Refusal` block that `text()` skips.
2438    #[tokio::test]
2439    async fn refusals_stream_apart_from_text() {
2440        let response = ChatResponse::new(vec![
2441            ContentBlock::text("Sure."),
2442            ContentBlock::Refusal(Refusal::new("I can't help with that.")),
2443        ]);
2444        let events: Vec<StreamEvent> = response.clone().into_stream_events().collect();
2445        assert!(events.iter().any(|event| matches!(
2446            event,
2447            StreamEvent::BlockStart {
2448                index: 1,
2449                block_type: StreamBlockType::Refusal,
2450                ..
2451            }
2452        )));
2453        assert!(events.contains(&StreamEvent::refusal_delta(1, "I can't help with that.")));
2454        assert!(!events.iter().any(|event| matches!(
2455            event,
2456            StreamEvent::TextDelta { text, .. } if text.contains("can't")
2457        )));
2458
2459        let collected = response
2460            .clone()
2461            .into_stream()
2462            .collect_response()
2463            .await
2464            .unwrap();
2465        assert_eq!(collected, response);
2466        assert_eq!(collected.first_refusal(), Some("I can't help with that."));
2467        assert_eq!(collected.text().as_deref(), Some("Sure."));
2468
2469        // A delta without a `BlockStart` opens a refusal block.
2470        let mut collector = StreamCollector::new();
2471        collector
2472            .push(StreamEvent::refusal_delta(0, "No."))
2473            .unwrap();
2474        assert_eq!(collector.block_type(0), Some(StreamBlockType::Refusal));
2475        let error = collector.push(StreamEvent::text_delta(0, "x")).unwrap_err();
2476        assert!(error.to_string().contains("refusal block"), "{error}");
2477    }
2478
2479    #[test]
2480    fn all_block_types_listed() {
2481        // Exhaustive on purpose: a new variant fails to compile here, as a
2482        // reminder to add it to `StreamBlockType::ALL` and bump the count.
2483        fn _exhaustive(block_type: StreamBlockType) {
2484            match block_type {
2485                StreamBlockType::Text
2486                | StreamBlockType::Image
2487                | StreamBlockType::Reasoning
2488                | StreamBlockType::Refusal
2489                | StreamBlockType::ToolCall
2490                | StreamBlockType::ProviderToolCall
2491                | StreamBlockType::ProviderToolResult
2492                | StreamBlockType::Other => {}
2493            }
2494        }
2495        let unique: std::collections::HashSet<_> = StreamBlockType::ALL.iter().collect();
2496        assert_eq!(unique.len(), StreamBlockType::ALL.len());
2497        assert_eq!(StreamBlockType::ALL.len(), 8);
2498    }
2499
2500    #[test]
2501    fn refusals_are_output() {
2502        assert!(StreamEvent::refusal_delta(0, "No.").is_output());
2503        assert!(StreamEvent::block_start(0, StreamBlockType::Refusal).is_output());
2504        assert!(StreamEvent::refusal_delta(0, "No.").is_block_event());
2505    }
2506
2507    #[test]
2508    fn events_serialize_round_trip() {
2509        let events: Vec<StreamEvent> = sample_response().into_stream_events().collect();
2510        let value = serde_json::to_value(&events).unwrap();
2511        let back: Vec<StreamEvent> = serde_json::from_value(value).unwrap();
2512        assert_eq!(back, events);
2513    }
2514
2515    #[test]
2516    fn block_provider_data_merges() {
2517        let mut c = StreamCollector::new();
2518        c.push(StreamEvent::ResponseStart {
2519            id: None,
2520            model: None,
2521        })
2522        .unwrap();
2523        c.push(StreamEvent::block_start(0, StreamBlockType::Reasoning))
2524            .unwrap();
2525        c.push(StreamEvent::ReasoningDelta {
2526            index: 0,
2527            text: "a".into(),
2528        })
2529        .unwrap();
2530        c.push(StreamEvent::BlockProviderData {
2531            index: 0,
2532            provider_data: ProviderData::new().with_raw("anthropic", "signature", json!("s")),
2533        })
2534        .unwrap();
2535        c.push(StreamEvent::BlockStop { index: 0 }).unwrap();
2536        c.push(StreamEvent::Warning {
2537            warning: Warning::new(WarningKind::Dropped, crate::Subject::Response, "x"),
2538        })
2539        .unwrap();
2540        c.push(StreamEvent::ResponseStop).unwrap();
2541        let response = c.finish().unwrap();
2542        let reasoning = response.content[0].as_reasoning().unwrap();
2543        assert_eq!(
2544            reasoning.provider_data.get_raw("anthropic", "signature"),
2545            Some(&json!("s"))
2546        );
2547        assert_eq!(response.warnings.len(), 1);
2548    }
2549
2550    /// Regression: provider data aimed at an `Other` block was accepted and
2551    /// then dropped when the block finished.
2552    #[tokio::test]
2553    async fn other_block_keeps_its_provider_data() {
2554        let other = OtherContent::new("compaction", ExtraMap::new())
2555            .with_origin("anthropic")
2556            .with_provider_data(ProviderData::new().with_raw("anthropic", "start", json!(1)));
2557        let stream = futures_util::stream::iter(vec![
2558            Ok(StreamEvent::response_start(None, None)),
2559            Ok(StreamEvent::other_start(0, other)),
2560            Ok(StreamEvent::block_provider_data(
2561                0,
2562                ProviderData::new().with_raw("anthropic", "signature", json!("other-sig")),
2563            )),
2564            Ok(StreamEvent::block_stop(0)),
2565            Ok(StreamEvent::ResponseStop),
2566        ]);
2567        let response = stream.collect_response().await.unwrap();
2568        assert_eq!(
2569            response.content[0].provider_data(),
2570            &ProviderData::new()
2571                .with_raw("anthropic", "start", json!(1))
2572                .with_raw("anthropic", "signature", json!("other-sig"))
2573        );
2574    }
2575
2576    #[test]
2577    fn response_provider_data_merges_into_the_response() {
2578        let event = StreamEvent::response_provider_data(ProviderData::new().with_raw(
2579            "openai-compat:lumen",
2580            "search_results",
2581            json!([{"url": "https://a"}]),
2582        ));
2583        let value = serde_json::to_value(&event).unwrap();
2584        assert_eq!(
2585            value,
2586            json!({"event": "response_provider_data", "provider_data": {"openai-compat:lumen": {"search_results": [{"url": "https://a"}]}}})
2587        );
2588        assert_eq!(serde_json::from_value::<StreamEvent>(value).unwrap(), event);
2589        assert!(!event.is_output());
2590
2591        let mut c = StreamCollector::new();
2592        for item in sample_response().into_stream_events() {
2593            if matches!(item, StreamEvent::BlockStart { index: 0, .. }) {
2594                c.push(event.clone()).unwrap();
2595            }
2596            c.push(item).unwrap();
2597        }
2598        let response = c.finish().unwrap();
2599        assert_eq!(
2600            response
2601                .provider_data
2602                .get_raw("openai-compat:lumen", "search_results"),
2603            Some(&json!([{"url": "https://a"}]))
2604        );
2605        let mut expected = sample_response();
2606        expected.provider_data.insert_raw(
2607            "openai-compat:lumen",
2608            "search_results",
2609            json!([{"url": "https://a"}]),
2610        );
2611        assert_eq!(response, expected);
2612
2613        // After `ResponseStop` it is a protocol error, like every event.
2614        let mut c = StreamCollector::new();
2615        c.push(StreamEvent::ResponseStop).unwrap();
2616        let err = c
2617            .push(StreamEvent::response_provider_data(ProviderData::new()))
2618            .unwrap_err();
2619        assert!(err.to_string().contains("response_provider_data"), "{err}");
2620    }
2621
2622    #[test]
2623    fn provider_deltas_are_observed_not_collected() {
2624        let mut data = ExtraMap::new();
2625        data.insert("query".into(), json!("rust"));
2626        let delta = StreamEvent::provider_delta(
2627            3,
2628            "web_search_call.searching",
2629            data,
2630            Some("openai".into()),
2631        );
2632        let value = serde_json::to_value(&delta).unwrap();
2633        assert_eq!(
2634            value,
2635            json!({"event": "provider_delta", "index": 3, "type_name": "web_search_call.searching", "data": {"query": "rust"}, "origin": "openai"})
2636        );
2637        assert_eq!(serde_json::from_value::<StreamEvent>(value).unwrap(), delta);
2638        assert!(!delta.is_output());
2639
2640        let mut c = StreamCollector::new();
2641        for event in sample_response().into_stream_events() {
2642            if matches!(event, StreamEvent::ResponseStop) {
2643                c.push(delta.clone()).unwrap();
2644            }
2645            c.push(event).unwrap();
2646        }
2647        assert_eq!(c.finish().unwrap(), sample_response());
2648    }
2649
2650    #[test]
2651    fn push_returns_blocks_as_they_complete() {
2652        let mut c = StreamCollector::new();
2653        let mut completed = Vec::new();
2654        for event in sample_response().into_stream_events() {
2655            if let Some(block) = c.push(event).unwrap() {
2656                completed.push(block.clone());
2657            }
2658        }
2659        assert_eq!(completed, sample_response().content);
2660
2661        // Empty text blocks and incomplete tool calls complete as nothing.
2662        let mut c = StreamCollector::new();
2663        c.push(StreamEvent::block_start(0, StreamBlockType::Text))
2664            .unwrap();
2665        assert!(
2666            c.push(StreamEvent::BlockStop { index: 0 })
2667                .unwrap()
2668                .is_none()
2669        );
2670        c.push(StreamEvent::block_start(1, StreamBlockType::ToolCall))
2671            .unwrap();
2672        assert!(
2673            c.push(StreamEvent::BlockStop { index: 1 })
2674                .unwrap()
2675                .is_none()
2676        );
2677        c.push(StreamEvent::tool_call_start(2, "c", "f")).unwrap();
2678        c.push(StreamEvent::ToolCallDelta {
2679            index: 2,
2680            arguments: "{bad".into(),
2681        })
2682        .unwrap();
2683        assert!(
2684            c.push(StreamEvent::BlockStop { index: 2 })
2685                .unwrap()
2686                .is_none()
2687        );
2688        // The call with invalid arguments is kept; the one with no id or
2689        // name is dropped.
2690        let partial = c.finish_partial().unwrap();
2691        let calls: Vec<_> = partial.response.tool_calls().collect();
2692        assert_eq!(calls, [&ToolCall::new("c", "f", "{bad")]);
2693        assert!(!calls[0].has_valid_arguments());
2694        assert!(partial.completeness.dropped_incomplete_tool_calls);
2695    }
2696
2697    #[test]
2698    fn a_tool_call_opened_without_an_id_names_its_constructor() {
2699        let mut c = StreamCollector::new();
2700        c.push(StreamEvent::response_start(Some("r".into()), None))
2701            .unwrap();
2702        c.push(StreamEvent::block_start(0, StreamBlockType::ToolCall))
2703            .unwrap();
2704        c.push(StreamEvent::BlockStop { index: 0 }).unwrap();
2705        c.push(StreamEvent::ResponseStop).unwrap();
2706        let err = c.finish().unwrap_err();
2707        assert!(
2708            err.message().contains("`StreamEvent::tool_call_start`"),
2709            "{err}"
2710        );
2711
2712        // A call that has its id and name but never closed gets no hint.
2713        let mut c = StreamCollector::new();
2714        c.push(StreamEvent::response_start(Some("r".into()), None))
2715            .unwrap();
2716        c.push(StreamEvent::tool_call_start(0, "c", "f")).unwrap();
2717        let err = c.finish().unwrap_err();
2718        assert!(!err.message().contains("tool_call_start"), "{err}");
2719    }
2720
2721    #[test]
2722    fn unfinished_provider_tool_calls_are_dropped_from_partial_responses() {
2723        let mut c = StreamCollector::new();
2724        c.push(StreamEvent::provider_tool_call_start(
2725            0,
2726            "srv_1",
2727            "web_search",
2728            Some("anthropic".into()),
2729        ))
2730        .unwrap();
2731        c.push(StreamEvent::ToolCallDelta {
2732            index: 0,
2733            arguments: r#"{"query":"ru"#.into(),
2734        })
2735        .unwrap();
2736        let partial = c.finish_partial().unwrap();
2737        assert!(partial.response.content.is_empty());
2738        assert!(partial.completeness.dropped_incomplete_tool_calls);
2739    }
2740
2741    #[test]
2742    fn open_tool_calls_are_dropped_from_partial_responses() {
2743        let mut c = StreamCollector::new();
2744        c.push(StreamEvent::tool_call_start(0, "c1", "get_weather"))
2745            .unwrap();
2746        c.push(StreamEvent::tool_call_delta(0, r#"{"city":"Pa"#))
2747            .unwrap();
2748        let partial = c.finish_partial().unwrap();
2749        assert_eq!(partial.response.tool_calls().count(), 0);
2750        assert!(partial.completeness.dropped_incomplete_tool_calls);
2751    }
2752
2753    #[test]
2754    fn open_tool_call_without_id_or_name_has_a_block_type() {
2755        let mut c = StreamCollector::new();
2756        // A delta with no `BlockStart` opens a call with no id or name.
2757        c.push(StreamEvent::tool_call_delta(0, "{")).unwrap();
2758        assert_eq!(c.block_type(0), Some(StreamBlockType::ToolCall));
2759    }
2760
2761    #[test]
2762    fn closed_tool_calls_complete_and_are_collected() {
2763        let mut c = StreamCollector::new();
2764        c.push(StreamEvent::tool_call_start(0, "c1", "get_weather"))
2765            .unwrap();
2766        c.push(StreamEvent::tool_call_delta(0, r#"{"city":"Paris"}"#))
2767            .unwrap();
2768        let first = c.push(StreamEvent::block_stop(0)).unwrap().cloned();
2769        c.push(StreamEvent::tool_call_start(1, "c2", "now"))
2770            .unwrap();
2771        let second = c.push(StreamEvent::block_stop(1)).unwrap().cloned();
2772        let expected = [
2773            ToolCall::new("c1", "get_weather", r#"{"city":"Paris"}"#),
2774            ToolCall::new("c2", "now", "{}"),
2775        ];
2776        assert_eq!(
2777            [first, second],
2778            expected
2779                .clone()
2780                .map(|call| Some(ContentBlock::ToolCall(call)))
2781        );
2782        let partial = c.finish_partial().unwrap();
2783        assert!(partial.response.tool_calls().eq(&expected));
2784    }
2785
2786    #[test]
2787    fn invalid_arguments_complete_as_nothing_but_are_collected() {
2788        let mut c = StreamCollector::new();
2789        c.push(StreamEvent::tool_call_start(0, "c1", "get_weather"))
2790            .unwrap();
2791        c.push(StreamEvent::tool_call_delta(0, r#"{"city":"#))
2792            .unwrap();
2793        assert!(c.push(StreamEvent::block_stop(0)).unwrap().is_none());
2794        let partial = c.finish_partial().unwrap();
2795        let call = partial.response.first_tool_call().unwrap();
2796        assert_eq!(
2797            (call.id.as_str(), call.arguments.as_str()),
2798            ("c1", r#"{"city":"#)
2799        );
2800        assert!(!call.has_valid_arguments());
2801    }
2802
2803    #[test]
2804    fn provider_tool_calls_never_appear_as_tool_calls() {
2805        let mut c = StreamCollector::new();
2806        c.push(StreamEvent::provider_tool_call_start(
2807            0,
2808            "srv_1",
2809            "web_search",
2810            Some("anthropic".into()),
2811        ))
2812        .unwrap();
2813        c.push(StreamEvent::tool_call_delta(0, r#"{"query":"rust"}"#))
2814            .unwrap();
2815        c.push(StreamEvent::block_stop(0)).unwrap();
2816
2817        c.push(StreamEvent::tool_call_start(1, "c1", "get_weather"))
2818            .unwrap();
2819        c.push(StreamEvent::block_stop(1)).unwrap();
2820        let partial = c.finish_partial().unwrap();
2821        let ids: Vec<_> = partial
2822            .response
2823            .tool_calls()
2824            .map(|call| call.id.as_str())
2825            .collect();
2826        assert_eq!(ids, ["c1"]);
2827        assert_eq!(partial.response.content.len(), 2);
2828    }
2829
2830    #[test]
2831    fn block_type_reports_open_and_closed_blocks() {
2832        let mut c = StreamCollector::new();
2833        assert_eq!(c.block_type(0), None);
2834        c.push(StreamEvent::text_delta(0, "Hi")).unwrap();
2835        c.push(StreamEvent::provider_tool_call_start(
2836            1,
2837            "srv_1",
2838            "web_search",
2839            None,
2840        ))
2841        .unwrap();
2842        c.push(StreamEvent::tool_call_start(3, "c1", "f")).unwrap();
2843        assert_eq!(c.block_type(0), Some(StreamBlockType::Text));
2844        assert_eq!(c.block_type(1), Some(StreamBlockType::ProviderToolCall));
2845        assert_eq!(c.block_type(2), None, "gap in the indexes");
2846        assert_eq!(c.block_type(3), Some(StreamBlockType::ToolCall));
2847        assert_eq!(c.block_type(4), None);
2848
2849        for index in [0, 1, 3] {
2850            c.push(StreamEvent::block_stop(index)).unwrap();
2851        }
2852        assert_eq!(c.block_type(0), Some(StreamBlockType::Text));
2853        assert_eq!(c.block_type(1), Some(StreamBlockType::ProviderToolCall));
2854        assert_eq!(c.block_type(2), None);
2855        assert_eq!(c.block_type(3), Some(StreamBlockType::ToolCall));
2856
2857        // An empty text block closes as nothing but keeps its type.
2858        c.push(StreamEvent::block_start(5, StreamBlockType::Text))
2859            .unwrap();
2860        c.push(StreamEvent::block_stop(5)).unwrap();
2861        assert_eq!(c.block_type(5), Some(StreamBlockType::Text));
2862    }
2863
2864    #[test]
2865    fn stopped_blocks_reject_further_events() {
2866        let mut c = StreamCollector::new();
2867        c.push(StreamEvent::block_start(0, StreamBlockType::Text))
2868            .unwrap();
2869        c.push(StreamEvent::BlockStop { index: 0 }).unwrap();
2870        for event in [
2871            StreamEvent::TextDelta {
2872                index: 0,
2873                text: "late".into(),
2874            },
2875            StreamEvent::BlockProviderData {
2876                index: 0,
2877                provider_data: ProviderData::new().with_raw("x", "k", json!(1)),
2878            },
2879            StreamEvent::block_start(0, StreamBlockType::Text),
2880            StreamEvent::BlockStop { index: 0 },
2881        ] {
2882            let err = c.push(event).unwrap_err();
2883            assert!(err.to_string().contains("text block"), "{err}");
2884        }
2885    }
2886
2887    #[tokio::test]
2888    async fn with_collector_pairs_stops_with_blocks() {
2889        let mut stream = sample_response().into_stream().with_collector();
2890        let mut completed = Vec::new();
2891        while let Some(item) = stream.next_event().await {
2892            let (event, block) = item.unwrap().into_parts();
2893            assert_eq!(
2894                block.is_some(),
2895                matches!(event, StreamEvent::BlockStop { .. })
2896            );
2897            completed.extend(block);
2898        }
2899        assert_eq!(completed, sample_response().content);
2900        assert_eq!(stream.finish().unwrap(), sample_response());
2901    }
2902
2903    #[tokio::test]
2904    async fn with_collector_records_stream_errors() {
2905        let events = futures_util::stream::iter(vec![
2906            Ok(StreamEvent::ResponseStart {
2907                id: None,
2908                model: None,
2909            }),
2910            Ok(StreamEvent::TextDelta {
2911                index: 0,
2912                text: "partial".into(),
2913            }),
2914            Err(Error::overloaded("busy")),
2915        ]);
2916        let after = futures_util::stream::iter([Ok(StreamEvent::text_delta(0, "after"))]);
2917        let mut stream = events.chain(after).with_collector();
2918        let items: Vec<_> = stream.by_ref().collect().await;
2919        // The stream ends after its first error.
2920        assert_eq!(items.len(), 3);
2921        assert!(items[2].is_err());
2922        let partial = stream.finish_partial().unwrap();
2923        assert_eq!(partial.response.text().as_deref(), Some("partial"));
2924        assert!(matches!(
2925            partial.terminal_error,
2926            Some(Error::Overloaded { .. })
2927        ));
2928    }
2929
2930    #[test]
2931    fn incomplete_stream_is_retryable() {
2932        let mut c = StreamCollector::new();
2933        c.push(StreamEvent::ResponseStart {
2934            id: None,
2935            model: None,
2936        })
2937        .unwrap();
2938        c.push(StreamEvent::TextDelta {
2939            index: 0,
2940            text: "partial".into(),
2941        })
2942        .unwrap();
2943        let err = c.finish().unwrap_err();
2944        assert!(matches!(
2945            err,
2946            Error::Stream {
2947                kind: StreamErrorKind::Incomplete,
2948                ..
2949            }
2950        ));
2951        assert!(err.is_retryable());
2952    }
2953
2954    #[test]
2955    fn protocol_violations_are_rejected() {
2956        let mut c = StreamCollector::new();
2957        c.push(StreamEvent::ResponseStop).unwrap();
2958        assert!(c.push(StreamEvent::ResponseStop).is_err());
2959
2960        let mut c = StreamCollector::new();
2961        c.push(StreamEvent::block_start(0, StreamBlockType::Text))
2962            .unwrap();
2963        assert!(
2964            c.push(StreamEvent::ToolCallDelta {
2965                index: 0,
2966                arguments: "{}".into()
2967            })
2968            .is_err()
2969        );
2970
2971        let mut c = StreamCollector::new();
2972        c.push(StreamEvent::tool_call_start(0, "a", "f")).unwrap();
2973        assert!(c.push(StreamEvent::tool_call_start(0, "b", "f")).is_err());
2974    }
2975
2976    #[test]
2977    fn invalid_tool_arguments_are_kept_like_unary() {
2978        let events = || {
2979            vec![
2980                StreamEvent::ResponseStart {
2981                    id: None,
2982                    model: None,
2983                },
2984                StreamEvent::tool_call_start(0, "c", "f"),
2985                StreamEvent::ToolCallDelta {
2986                    index: 0,
2987                    arguments: "{bad".into(),
2988                },
2989                StreamEvent::BlockStop { index: 0 },
2990                StreamEvent::ResponseStop,
2991            ]
2992        };
2993        // The same response is `Ok` unary and
2994        // streamed, with the call's raw arguments and the finish reason.
2995        let mut strict = StreamCollector::new();
2996        for e in events() {
2997            strict.push(e).unwrap();
2998        }
2999        let response = strict.finish().unwrap();
3000        let call = response.first_tool_call().unwrap();
3001        assert_eq!(call.arguments, "{bad");
3002        assert!(!call.has_valid_arguments());
3003
3004        let mut partial = StreamCollector::new();
3005        for e in events() {
3006            partial.push(e).unwrap();
3007        }
3008        let collected = partial.finish_partial().unwrap();
3009        assert!(!collected.completeness.dropped_incomplete_tool_calls);
3010        assert_eq!(collected.response.content.len(), 1);
3011    }
3012
3013    /// One absurd block index from the wire is a protocol error, not a
3014    /// panic or an allocation for every block before it.
3015    #[test]
3016    fn absurd_block_indices_are_protocol_errors() {
3017        for index in [usize::MAX, usize::MAX - 1, 10_000_000_000, 50_000_000] {
3018            let mut c = StreamCollector::new();
3019            let err = c.push(StreamEvent::text_delta(index, "x")).unwrap_err();
3020            assert!(
3021                matches!(
3022                    err,
3023                    Error::Stream {
3024                        kind: StreamErrorKind::Protocol,
3025                        ..
3026                    }
3027                ),
3028                "{err:?}"
3029            );
3030            let mut c = StreamCollector::new();
3031            assert!(
3032                c.push(StreamEvent::block_start(index, StreamBlockType::Text))
3033                    .is_err()
3034            );
3035        }
3036        // A modest gap is fine: decoders may skip blocks they do not model.
3037        let mut c = StreamCollector::new();
3038        c.push(StreamEvent::text_delta(MAX_BLOCK_INDEX_GAP, "x"))
3039            .unwrap();
3040        assert_eq!(
3041            c.block_type(MAX_BLOCK_INDEX_GAP),
3042            Some(StreamBlockType::Text)
3043        );
3044    }
3045
3046    /// Gaps are bounded over the whole stream, not per event: an index
3047    /// that climbs by the per-event limit on every small event must not
3048    /// grow memory a thousand times faster than the input.
3049    #[test]
3050    fn stepped_block_indices_are_bounded_in_total() {
3051        let mut c = StreamCollector::new();
3052        let mut failed = None;
3053        for step in 1..=4 {
3054            let index = step * (MAX_BLOCK_INDEX_GAP / 2 + 1);
3055            if let Err(err) = c.push(StreamEvent::text_delta(index, "x")) {
3056                failed = Some((step, err));
3057                break;
3058            }
3059        }
3060        let (step, err) = failed.expect("stepped indices are refused");
3061        assert_eq!(step, 2, "the second gap passes the stream's total");
3062        assert!(
3063            matches!(
3064                err,
3065                Error::Stream {
3066                    kind: StreamErrorKind::Protocol,
3067                    ..
3068                }
3069            ),
3070            "{err:?}"
3071        );
3072        assert!(
3073            c.blocks.len() <= MAX_BLOCK_INDEX_GAP + 2,
3074            "{}",
3075            c.blocks.len()
3076        );
3077
3078        // Dense blocks are never limited, however many there are.
3079        let mut c = StreamCollector::new();
3080        for index in 0..4 * MAX_BLOCK_INDEX_GAP {
3081            c.push(StreamEvent::text_delta(index, "x")).unwrap();
3082        }
3083    }
3084
3085    #[test]
3086    fn delta_usage_accumulates() {
3087        let mut c = StreamCollector::new();
3088        for usage in [Usage::new().output_tokens(2), Usage::new().output_tokens(3)] {
3089            c.push(StreamEvent::ResponseMetadata {
3090                usage: Some(usage),
3091                usage_mode: UsageMetadataMode::Delta,
3092                finish_reason: None,
3093                finish_reason_raw: None,
3094                id: None,
3095                model: None,
3096                request_id: None,
3097                provider: None,
3098                provider_data: ProviderData::new(),
3099            })
3100            .unwrap();
3101        }
3102        let collected = c.finish_partial().unwrap();
3103        assert_eq!(collected.response.usage.unwrap().output_tokens, Some(5));
3104    }
3105
3106    /// Every recorded `*_events.json` stream fixture in the workspace, or
3107    /// `None` outside the workspace (a packaged `allama` has no sibling
3108    /// crates). In the workspace there are at least ten.
3109    fn event_fixtures() -> Option<Vec<(String, Vec<StreamEvent>)>> {
3110        fn visit(dir: &std::path::Path, out: &mut Vec<std::path::PathBuf>) {
3111            let Ok(entries) = std::fs::read_dir(dir) else {
3112                return;
3113            };
3114            for path in entries.map(|entry| entry.unwrap().path()) {
3115                if path.is_dir() {
3116                    visit(&path, out);
3117                } else if path.to_string_lossy().ends_with("_events.json") {
3118                    out.push(path);
3119                }
3120            }
3121        }
3122        let crates = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("..");
3123        let workspace = std::fs::read_to_string(crates.join("../Cargo.toml"))
3124            .is_ok_and(|manifest| manifest.contains("[workspace]"));
3125        if !workspace {
3126            eprintln!("skipped: the workspace stream fixtures are not available");
3127            return None;
3128        }
3129        let mut paths = Vec::new();
3130        for entry in std::fs::read_dir(&crates).unwrap() {
3131            visit(&entry.unwrap().path().join("fixtures"), &mut paths);
3132        }
3133        paths.sort();
3134        assert!(paths.len() >= 10, "found {} fixtures", paths.len());
3135        let fixtures = paths
3136            .into_iter()
3137            .map(|path| {
3138                let text = std::fs::read_to_string(&path).unwrap();
3139                let items: Vec<serde_json::Value> = serde_json::from_str(&text).unwrap();
3140                // A recorded stream error ends the events.
3141                let events = items
3142                    .into_iter()
3143                    .take_while(|item| item.get("error").is_none())
3144                    .map(|item| {
3145                        serde_json::from_value(item)
3146                            .unwrap_or_else(|err| panic!("{}: {err}", path.display()))
3147                    })
3148                    .collect();
3149                (path.display().to_string(), events)
3150            })
3151            .collect();
3152        Some(fixtures)
3153    }
3154
3155    /// The response `events` collect to, in full and metadata-only mode.
3156    fn collect_both(events: &[StreamEvent]) -> (CollectedResponse, CollectedResponse) {
3157        let mut full = StreamCollector::new();
3158        let mut metadata = StreamCollector::metadata_only();
3159        for event in events {
3160            let _ = full.push_ref(event);
3161            metadata.push_ref(event).unwrap();
3162        }
3163        (
3164            full.finish_partial().unwrap(),
3165            metadata.finish_partial().unwrap(),
3166        )
3167    }
3168
3169    /// Metadata-only collection matches full collection on everything but
3170    /// content.
3171    fn assert_same_metadata(label: &str, events: &[StreamEvent]) {
3172        let (full, metadata) = collect_both(events);
3173        assert!(metadata.response.content.is_empty(), "{label}");
3174        assert_eq!(
3175            metadata.response.usage, full.response.usage,
3176            "{label}: usage"
3177        );
3178        assert_eq!(
3179            metadata.response.finish_reason, full.response.finish_reason,
3180            "{label}: finish reason"
3181        );
3182        assert_eq!(metadata.response.id, full.response.id, "{label}: id");
3183        assert_eq!(
3184            metadata.response.model, full.response.model,
3185            "{label}: model"
3186        );
3187        let mut expected = full.response;
3188        expected.content.clear();
3189        assert_eq!(metadata.response, expected, "{label}");
3190    }
3191
3192    #[test]
3193    fn metadata_only_collection_matches_full_collection_on_fixtures() {
3194        let Some(fixtures) = event_fixtures() else {
3195            return;
3196        };
3197        let mut with_usage = 0;
3198        for (label, events) in &fixtures {
3199            assert_same_metadata(label, events);
3200            if collect_both(events).1.response.usage.is_some() {
3201                with_usage += 1;
3202            }
3203        }
3204        assert!(with_usage > 0);
3205    }
3206
3207    #[test]
3208    fn metadata_only_collection_matches_full_collection_in_delta_usage_mode() {
3209        // The recorded fixtures report cumulative usage; replay each with
3210        // its usage as deltas, so later reports add to earlier ones.
3211        for (label, mut events) in event_fixtures().unwrap_or_default() {
3212            for event in &mut events {
3213                if let StreamEvent::ResponseMetadata { usage_mode, .. } = event {
3214                    *usage_mode = UsageMetadataMode::Delta;
3215                }
3216            }
3217            assert_same_metadata(&format!("{label} (delta)"), &events);
3218        }
3219    }
3220
3221    #[test]
3222    fn metadata_only_collection_merges_metadata_and_skips_blocks() {
3223        let delta = |usage: Usage| {
3224            let mut event = StreamEvent::metadata();
3225            if let StreamEvent::ResponseMetadata {
3226                usage: slot,
3227                usage_mode,
3228                ..
3229            } = &mut event
3230            {
3231                *slot = Some(usage);
3232                *usage_mode = UsageMetadataMode::Delta;
3233            }
3234            event
3235        };
3236        let mut finish = StreamEvent::metadata();
3237        if let StreamEvent::ResponseMetadata {
3238            finish_reason,
3239            model,
3240            provider,
3241            ..
3242        } = &mut finish
3243        {
3244            *finish_reason = Some(FinishReason::Length);
3245            *model = Some("m-2".into());
3246            *provider = Some("demo".into());
3247        }
3248        let mut events = vec![
3249            StreamEvent::response_start(Some("r1".into()), Some("m".into())),
3250            delta(Usage::new().input_tokens(4).output_tokens(1)),
3251        ];
3252        events.extend(sample_response().into_stream_events().filter(|event| {
3253            !matches!(
3254                event,
3255                StreamEvent::ResponseStart { .. }
3256                    | StreamEvent::ResponseMetadata { .. }
3257                    | StreamEvent::ResponseStop
3258            )
3259        }));
3260        // An unclosed block and an out-of-protocol delta are not checked.
3261        events.push(StreamEvent::text_delta(9, "dangling"));
3262        events.push(StreamEvent::block_stop(42));
3263        events.push(delta(Usage::new().output_tokens(2)));
3264        events.push(finish);
3265        events.push(StreamEvent::ResponseStop);
3266
3267        let mut collector = StreamCollector::metadata_only();
3268        for event in &events {
3269            assert!(collector.push_ref(event).unwrap().is_none());
3270        }
3271        assert!(
3272            collector
3273                .push_ref(&StreamEvent::text_delta(0, "late"))
3274                .is_err()
3275        );
3276        let collected = collector.finish_partial().unwrap();
3277        let response = &collected.response;
3278        assert!(response.content.is_empty());
3279        assert_eq!(
3280            response.usage,
3281            Some(Usage::new().input_tokens(4).output_tokens(3))
3282        );
3283        assert_eq!(response.finish_reason, Some(FinishReason::Length));
3284        assert_eq!(response.id.as_deref(), Some("r1"));
3285        assert_eq!(response.model.as_deref(), Some("m-2"));
3286        assert_eq!(response.provider.as_deref(), Some("demo"));
3287        assert_eq!(response.warnings, sample_response().warnings);
3288        assert!(collected.completeness.is_complete());
3289    }
3290
3291    #[tokio::test]
3292    async fn until_committed_buffers_and_replays() {
3293        let events = sample_response().into_stream();
3294        let prefixed = until_committed(events).await.unwrap();
3295        let all: Vec<_> = prefixed.map(|e| e.unwrap()).collect().await;
3296        assert_eq!(
3297            all,
3298            sample_response().into_stream_events().collect::<Vec<_>>()
3299        );
3300    }
3301
3302    #[tokio::test]
3303    async fn until_committed_reports_pre_output_errors() {
3304        let stream = futures_util::stream::iter(vec![
3305            Ok(StreamEvent::ResponseStart {
3306                id: None,
3307                model: None,
3308            }),
3309            Err(Error::overloaded("busy")),
3310        ]);
3311        let err = until_committed(stream).await.err().unwrap();
3312        assert!(matches!(err, Error::Overloaded { .. }));
3313    }
3314
3315    #[tokio::test]
3316    async fn until_committed_passes_post_output_errors_through() {
3317        let stream = futures_util::stream::iter(vec![
3318            Ok(StreamEvent::ResponseStart {
3319                id: None,
3320                model: None,
3321            }),
3322            Ok(StreamEvent::TextDelta {
3323                index: 0,
3324                text: "x".into(),
3325            }),
3326            Err(Error::overloaded("busy")),
3327        ]);
3328        let prefixed = until_committed(stream).await.unwrap();
3329        let items: Vec<_> = prefixed.collect().await;
3330        assert_eq!(items.len(), 3);
3331        assert!(items[2].is_err());
3332    }
3333
3334    #[tokio::test]
3335    async fn into_parts_returns_the_unyielded_prefix_and_the_rest() {
3336        let start = StreamEvent::response_start(None, None);
3337        let first = StreamEvent::text_delta(0, "a");
3338        let second = StreamEvent::text_delta(0, "b");
3339        let rest = futures_util::stream::iter(vec![Ok(second.clone())]);
3340        let mut prefixed = Prefixed::new([start.clone(), first.clone()], Some(rest));
3341
3342        // One buffered event is already yielded; it is not handed back.
3343        assert_eq!(prefixed.next().await.unwrap().unwrap(), start);
3344        let (buffered, rest) = prefixed.into_parts();
3345        assert_eq!(buffered, std::slice::from_ref(&first));
3346
3347        let resumed: Vec<_> = Prefixed::new(buffered, rest)
3348            .map(Result::unwrap)
3349            .collect()
3350            .await;
3351        assert_eq!(resumed, [first, second]);
3352    }
3353
3354    #[tokio::test]
3355    async fn into_parts_of_an_ended_stream_has_no_rest() {
3356        let events = futures_util::stream::iter(vec![Ok(StreamEvent::response_start(None, None))]);
3357        let (buffered, rest) = until_committed(events).await.unwrap().into_parts();
3358        assert_eq!(buffered.len(), 1);
3359        assert!(rest.is_none());
3360    }
3361
3362    #[test]
3363    fn metadata_builder_sets_every_field_of_the_event() {
3364        let provider_data = ProviderData::new().with_raw("demo", "k", serde_json::json!(1));
3365        let usage = Usage::new().input_tokens(4).output_tokens(2);
3366        let built: StreamEvent = MetadataBuilder::new()
3367            .with_finish_reason(FinishReason::Length)
3368            .with_finish_reason_raw("max_tokens")
3369            .with_usage(usage.clone())
3370            .with_usage_mode(UsageMetadataMode::Delta)
3371            .with_id("r1".to_owned())
3372            .with_model("m")
3373            .with_request_id_opt(None)
3374            .with_provider("demo")
3375            .with_provider_data(provider_data.clone())
3376            .into();
3377
3378        let mut expected = StreamEvent::metadata();
3379        if let StreamEvent::ResponseMetadata {
3380            finish_reason,
3381            finish_reason_raw,
3382            usage: slot_usage,
3383            usage_mode,
3384            id,
3385            model,
3386            provider,
3387            provider_data: slot_provider_data,
3388            ..
3389        } = &mut expected
3390        {
3391            *finish_reason = Some(FinishReason::Length);
3392            *finish_reason_raw = Some("max_tokens".into());
3393            *slot_usage = Some(usage);
3394            *usage_mode = UsageMetadataMode::Delta;
3395            *id = Some("r1".into());
3396            *model = Some("m".into());
3397            *provider = Some("demo".into());
3398            *slot_provider_data = provider_data;
3399        }
3400        assert_eq!(built, expected);
3401        assert_eq!(MetadataBuilder::default().build(), StreamEvent::metadata());
3402    }
3403
3404    /// Values, and `_opt` twins that keep the current
3405    /// value on `None` (the `Error` builder rule).
3406    #[test]
3407    fn metadata_builder_opt_twins_keep_the_current_value_on_none() {
3408        let set = MetadataBuilder::new()
3409            .with_finish_reason(FinishReason::Stop)
3410            .with_finish_reason_raw("end_turn")
3411            .with_usage(Usage::new().output_tokens(1))
3412            .with_id("r1")
3413            .with_model("m")
3414            .with_request_id("req_1")
3415            .with_provider("demo");
3416        let kept = set
3417            .clone()
3418            .with_finish_reason_opt(None)
3419            .with_finish_reason_raw_opt(None)
3420            .with_usage_opt(None)
3421            .with_id_opt(None)
3422            .with_model_opt(None)
3423            .with_request_id_opt(None)
3424            .with_provider_opt(None);
3425        assert_eq!(kept, set);
3426
3427        let wire_id: Option<String> = Some("r2".to_owned());
3428        let StreamEvent::ResponseMetadata { id, .. } = set.with_id_opt(wire_id.as_deref()).build()
3429        else {
3430            unreachable!()
3431        };
3432        assert_eq!(id.as_deref(), Some("r2"));
3433
3434        let served_by = Some(Cow::Owned("compat:groq".to_owned()));
3435        let StreamEvent::ResponseMetadata { provider, .. } =
3436            MetadataBuilder::new().with_provider_opt(served_by).build()
3437        else {
3438            unreachable!()
3439        };
3440        assert_eq!(provider.as_deref(), Some("compat:groq"));
3441    }
3442
3443    /// A URL image keeps its declared media type through the stream events,
3444    /// as on a unary response.
3445    #[test]
3446    fn image_media_type_survives_the_stream_events() {
3447        let image =
3448            Image::new(MediaSource::url("https://example.com/a")).with_media_type("image/jpeg");
3449        let response = ChatResponse::new(vec![ContentBlock::Image(image.clone())]);
3450        let mut collector = StreamCollector::new();
3451        for event in response.into_stream_events() {
3452            collector.push(event).unwrap();
3453        }
3454        let collected = collector.finish().unwrap();
3455        assert_eq!(collected.content, vec![ContentBlock::Image(image)]);
3456    }
3457}