Skip to main content

rai_sdk/
wire.rs

1//! Serializable stream events for proxying a generation across a network hop.
2//!
3//! The in-process streaming types ([`ProviderStreamEvent`] and [`StreamEvent`])
4//! are shaped for a consumer that lives in the same process as the SDK. A
5//! *proxy* deployment does not: a server holds the provider credentials, calls
6//! the provider through `rai-sdk`, and re-emits each event to a client over
7//! server-sent events (SSE) or a WebSocket. The client then rebuilds the same
8//! stream semantics — text deltas, tool-call activity, token usage, stop
9//! reason — from what arrived on the wire.
10//!
11//! ```text
12//!   client ──HTTP──▶ your server ──rai-sdk──▶ provider
13//!          ◀──SSE─── WireStreamEvent ◀────────┘
14//! ```
15//!
16//! This module is that wire layer:
17//!
18//! - [`WireStreamEvent`] is a `Serialize`/`Deserialize` event enum with an
19//!   explicit, stable JSON representation.
20//! - [`WireError`] projects the crate's [`enum@Error`] into something that can
21//!   cross a wire, so a mid-stream provider failure arrives as an *event*
22//!   rather than as a dropped connection.
23//! - [`StreamAccumulator`] is the receiving end of
24//!   [`RequestBuilder::stream_accumulated`](crate::RequestBuilder::stream_accumulated):
25//!   feed it the events a client parsed off the wire and it hands back one
26//!   assembled [`Response`].
27//!
28//! Produce the events with
29//! [`RequestBuilder::stream_wire_events`](crate::RequestBuilder::stream_wire_events).
30//!
31//! # Wire format
32//!
33//! [`WireStreamEvent`] is an **internally tagged** serde enum. Every event is a
34//! JSON object carrying a `"type"` discriminant alongside that variant's
35//! fields:
36//!
37//! ```json
38//! {"type":"message_start","protocol_version":1,"model":"gpt-4o-mini","provider":"openai"}
39//! {"type":"text_delta","text":"Hello"}
40//! {"type":"usage","usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}
41//! {"type":"message_stop","finish_reason":"stop"}
42//! ```
43//!
44//! The full set of `"type"` values is:
45//!
46//! | `"type"` | Variant | Meaning |
47//! | --- | --- | --- |
48//! | `message_start` | [`WireStreamEvent::MessageStart`] | First event of every stream; names the protocol version, model, and provider. |
49//! | `text_delta` | [`WireStreamEvent::TextDelta`] | Append this text to the output so far. |
50//! | `tool_call_start` | [`WireStreamEvent::ToolCallStart`] | A tool call has begun; arguments follow. |
51//! | `tool_call_delta` | [`WireStreamEvent::ToolCallDelta`] | A fragment of that call's JSON arguments. |
52//! | `tool_call_end` | [`WireStreamEvent::ToolCallEnd`] | The call is complete, with its arguments assembled. |
53//! | `tool_result` | [`WireStreamEvent::ToolResult`] | The output of executing a tool call. |
54//! | `usage` | [`WireStreamEvent::Usage`] | Token counts for the generation. |
55//! | `message_stop` | [`WireStreamEvent::MessageStop`] | Terminal event of a successful stream. |
56//! | `turn_complete` | [`WireStreamEvent::TurnComplete`] | An assembled [`ConversationTurn`], for history. |
57//! | `error` | [`WireStreamEvent::Error`] | Terminal event of a failed stream. |
58//!
59//! ## The variant names are a compatibility surface
60//!
61//! Those `"type"` strings — and the field names inside each event — are part of
62//! this crate's public API in exactly the way a function signature is. A server
63//! and a client can be built from different `rai-sdk` versions, so renaming a
64//! tag silently breaks every deployed client. Treat them as frozen:
65//!
66//! - **Renaming or removing a `"type"` value, or renaming a field, is a
67//!   breaking change** and will only happen in a major (pre-1.0: minor) release,
68//!   with a changelog entry.
69//! - **Adding a variant, or adding an optional field to an existing variant, is
70//!   additive** and can happen in a patch or minor release. Both
71//!   [`WireStreamEvent`] and [`WireErrorKind`] are `#[non_exhaustive]`, and
72//!   unknown [`WireErrorKind`] values deserialize into
73//!   [`WireErrorKind::Other`], so a client compiled against an older version
74//!   keeps parsing streams from a newer server. Match with a `_ => {}` arm and
75//!   ignore what you do not recognize.
76//!
77//! [`WIRE_PROTOCOL_VERSION`] names the current revision of this contract and is
78//! carried on every [`WireStreamEvent::MessageStart`]. It is bumped only when
79//! the *framing* changes in a way a client must react to — not for additive
80//! variants. A client that sees a `protocol_version` it does not understand
81//! should refuse the stream rather than guess.
82//!
83//! # Cancellation
84//!
85//! Dropping a stream aborts the upstream provider request. The "Cancellation"
86//! section of
87//! [`RequestBuilder::stream_wire_events`](crate::RequestBuilder::stream_wire_events)
88//! has the details and what they mean for a proxy.
89//!
90//! # Examples
91//!
92//! Server side — turn each event into an SSE `data:` payload:
93//!
94//! ```no_run
95//! use futures::StreamExt;
96//! use rai_sdk::{ClientBuilder, Model};
97//!
98//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
99//! # let client = ClientBuilder::new().from_env().model(Model::gpt4o_mini()).build()?;
100//! let mut events = client
101//!     .request()
102//!     .prompt("Explain SSE in one sentence.")
103//!     .stream_wire_events()
104//!     .await?;
105//!
106//! while let Some(event) = events.next().await {
107//!     let payload = serde_json::to_string(&event)?;
108//!     println!("event: {}\ndata: {payload}\n", event.tag());
109//! }
110//! # Ok(())
111//! # }
112//! ```
113//!
114//! Client side — rebuild one [`Response`] from the payloads:
115//!
116//! ```
117//! use rai_sdk::wire::{StreamAccumulator, WireStreamEvent};
118//!
119//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
120//! let payloads = [
121//!     r#"{"type":"message_start","protocol_version":1,"model":"gpt-4o-mini","provider":"openai"}"#,
122//!     r#"{"type":"text_delta","text":"Hello, "}"#,
123//!     r#"{"type":"text_delta","text":"world."}"#,
124//!     r#"{"type":"usage","usage":{"prompt_tokens":9,"completion_tokens":3,"total_tokens":12}}"#,
125//!     r#"{"type":"message_stop","finish_reason":"stop"}"#,
126//! ];
127//!
128//! let mut accumulator = StreamAccumulator::new();
129//! for payload in payloads {
130//!     accumulator.push(serde_json::from_str::<WireStreamEvent>(payload)?)?;
131//! }
132//!
133//! let response = accumulator.finish()?;
134//! assert_eq!(response.text(), "Hello, world.");
135//! assert_eq!(response.usage.unwrap().total_tokens, Some(12));
136//! # Ok(())
137//! # }
138//! ```
139//!
140//! [`ProviderStreamEvent`]: crate::provider::ProviderStreamEvent
141
142use futures::{Stream, StreamExt};
143use serde::{Deserialize, Serialize};
144
145use crate::error::{Error, ProviderKind};
146use crate::message::{ConversationTurn, Message, Response, StreamEvent, ToolCall, Usage};
147
148/// The revision of the [`WireStreamEvent`] framing this build speaks.
149///
150/// Carried on every [`WireStreamEvent::MessageStart`] so a client can check the
151/// contract before consuming a stream. See the
152/// [module documentation](self#the-variant-names-are-a-compatibility-surface)
153/// for when this is bumped.
154pub const WIRE_PROTOCOL_VERSION: u32 = 1;
155
156fn default_protocol_version() -> u32 {
157    WIRE_PROTOCOL_VERSION
158}
159
160/// A stream event in its serializable, over-the-wire form.
161///
162/// See the [module documentation](self#wire-format) for the JSON shape and the
163/// compatibility guarantees attached to the `"type"` tags.
164///
165/// A well-formed stream starts with exactly one
166/// [`MessageStart`](WireStreamEvent::MessageStart) and ends with exactly one
167/// terminal event — [`MessageStop`](WireStreamEvent::MessageStop) on success,
168/// [`Error`](WireStreamEvent::Error) on failure (see
169/// [`is_terminal`](WireStreamEvent::is_terminal)). A stream that simply stops
170/// is a *truncated* stream: the connection died. That distinction is the whole
171/// point of putting errors on the wire, and [`StreamAccumulator::finish`]
172/// enforces it.
173///
174/// This enum is `#[non_exhaustive]`: new event types are additive, so match
175/// with a catch-all arm.
176#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
177#[serde(tag = "type", rename_all = "snake_case")]
178#[non_exhaustive]
179pub enum WireStreamEvent {
180    /// The stream has opened. Always the first event.
181    MessageStart {
182        /// Revision of the wire framing the sender speaks.
183        ///
184        /// Defaults to [`WIRE_PROTOCOL_VERSION`] when absent, so payloads
185        /// written before this field existed still parse.
186        #[serde(default = "default_protocol_version")]
187        protocol_version: u32,
188        /// Model identifier the request was routed to.
189        model: String,
190        /// Provider serving the request.
191        provider: ProviderKind,
192    },
193
194    /// An incremental piece of assistant text to append to the output so far.
195    TextDelta {
196        /// Text to append.
197        text: String,
198    },
199
200    /// A tool call has begun; its arguments arrive in later events.
201    ToolCallStart {
202        /// Provider-assigned call identifier.
203        id: String,
204        /// Name of the tool the model wants to run.
205        name: String,
206    },
207
208    /// A fragment of the JSON argument string for a tool call.
209    ToolCallDelta {
210        /// Identifier of the call these arguments belong to.
211        id: String,
212        /// Partial JSON text to append to previously received fragments.
213        arguments: String,
214    },
215
216    /// A tool call has finished streaming, with its arguments assembled.
217    ToolCallEnd {
218        /// Provider-assigned call identifier.
219        id: String,
220        /// Name of the tool the model wants to run.
221        name: String,
222        /// Complete raw JSON argument string.
223        arguments: String,
224    },
225
226    /// The result of executing a tool call.
227    ///
228    /// Never emitted by
229    /// [`stream_wire_events`](crate::RequestBuilder::stream_wire_events), which
230    /// does not run tools. It exists so a proxy that executes tools itself can
231    /// report them to its client using the same envelope.
232    ToolResult {
233        /// Identifier of the call this result answers.
234        id: String,
235        /// Serialized tool output.
236        result: String,
237    },
238
239    /// Token usage for the generation.
240    ///
241    /// Emitted once, just before the terminal event, when the provider reports
242    /// usage. This is the event a server bills against and a client displays,
243    /// so it deliberately carries the counts rather than a summary.
244    Usage {
245        /// Token counts as reported by the provider.
246        usage: Usage,
247    },
248
249    /// Generation finished cleanly. Terminal.
250    MessageStop {
251        /// Why generation stopped, e.g. `"stop"`, `"length"`, or `"tool_use"`.
252        #[serde(default, skip_serializing_if = "Option::is_none")]
253        finish_reason: Option<String>,
254    },
255
256    /// An assembled conversation turn, ready to be stored as history.
257    ///
258    /// Never emitted by
259    /// [`stream_wire_events`](crate::RequestBuilder::stream_wire_events); it is
260    /// the wire form of [`StreamEvent::TurnComplete`], so events from
261    /// [`generate_stream_events`](crate::RequestBuilder::generate_stream_events)
262    /// can be forwarded without loss.
263    TurnComplete {
264        /// The complete turn.
265        turn: ConversationTurn,
266    },
267
268    /// Generation failed. Terminal.
269    ///
270    /// Receiving this means the provider or the SDK rejected the request
271    /// *after* the stream opened. A client that never receives a terminal event
272    /// lost its connection instead — a different failure with a different
273    /// remedy.
274    Error {
275        /// What went wrong.
276        error: WireError,
277    },
278}
279
280impl WireStreamEvent {
281    /// Build the opening event for a stream, stamped with
282    /// [`WIRE_PROTOCOL_VERSION`].
283    pub fn message_start(model: impl Into<String>, provider: ProviderKind) -> Self {
284        Self::MessageStart {
285            protocol_version: WIRE_PROTOCOL_VERSION,
286            model: model.into(),
287            provider,
288        }
289    }
290
291    /// Build a terminal error event from a crate [`enum@Error`].
292    pub fn error(error: &Error) -> Self {
293        Self::Error {
294            error: WireError::from(error),
295        }
296    }
297
298    /// The `"type"` discriminant this event serializes with.
299    ///
300    /// Useful as the SSE `event:` name, or for logging and metrics labels
301    /// without serializing the whole payload.
302    pub fn tag(&self) -> &'static str {
303        match self {
304            Self::MessageStart { .. } => "message_start",
305            Self::TextDelta { .. } => "text_delta",
306            Self::ToolCallStart { .. } => "tool_call_start",
307            Self::ToolCallDelta { .. } => "tool_call_delta",
308            Self::ToolCallEnd { .. } => "tool_call_end",
309            Self::ToolResult { .. } => "tool_result",
310            Self::Usage { .. } => "usage",
311            Self::MessageStop { .. } => "message_stop",
312            Self::TurnComplete { .. } => "turn_complete",
313            Self::Error { .. } => "error",
314        }
315    }
316
317    /// Whether this event ends the stream: [`MessageStop`] or [`Error`].
318    ///
319    /// [`MessageStop`]: WireStreamEvent::MessageStop
320    /// [`Error`]: WireStreamEvent::Error
321    pub fn is_terminal(&self) -> bool {
322        matches!(self, Self::MessageStop { .. } | Self::Error { .. })
323    }
324}
325
326impl From<StreamEvent> for WireStreamEvent {
327    /// Forward a high-level [`StreamEvent`] onto the wire without loss.
328    ///
329    /// Every [`StreamEvent`] has an exact counterpart here, and
330    /// [`TryFrom<WireStreamEvent>`](WireStreamEvent) converts it back to the
331    /// same value.
332    fn from(event: StreamEvent) -> Self {
333        match event {
334            StreamEvent::TextDelta { text } => Self::TextDelta { text },
335            StreamEvent::ToolCall {
336                id,
337                name,
338                arguments,
339            } => Self::ToolCallEnd {
340                id,
341                name,
342                arguments,
343            },
344            StreamEvent::ToolResult { id, result } => Self::ToolResult { id, result },
345            StreamEvent::TurnComplete { turn } => Self::TurnComplete { turn },
346        }
347    }
348}
349
350/// A [`WireStreamEvent`] with no [`StreamEvent`] counterpart.
351///
352/// The wire enum is a superset: it also models stream framing (`message_start`,
353/// `message_stop`), token usage, partial tool-call progress, and errors, none of
354/// which [`StreamEvent`] represents.
355#[derive(Debug, Clone, PartialEq, Eq)]
356pub struct UnrepresentableWireEvent {
357    /// The `"type"` tag of the event that could not be converted.
358    pub tag: &'static str,
359}
360
361impl std::fmt::Display for UnrepresentableWireEvent {
362    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
363        write!(f, "wire event '{}' has no StreamEvent equivalent", self.tag)
364    }
365}
366
367impl std::error::Error for UnrepresentableWireEvent {}
368
369impl TryFrom<WireStreamEvent> for StreamEvent {
370    type Error = UnrepresentableWireEvent;
371
372    fn try_from(event: WireStreamEvent) -> std::result::Result<Self, UnrepresentableWireEvent> {
373        let tag = event.tag();
374        match event {
375            WireStreamEvent::TextDelta { text } => Ok(Self::TextDelta { text }),
376            WireStreamEvent::ToolCallEnd {
377                id,
378                name,
379                arguments,
380            } => Ok(Self::ToolCall {
381                id,
382                name,
383                arguments,
384            }),
385            WireStreamEvent::ToolResult { id, result } => Ok(Self::ToolResult { id, result }),
386            WireStreamEvent::TurnComplete { turn } => Ok(Self::TurnComplete { turn }),
387            _ => Err(UnrepresentableWireEvent { tag }),
388        }
389    }
390}
391
392// ── Errors on the wire ─────────────────────────────────────────────────────
393
394/// A serializable projection of the crate's [`enum@Error`].
395///
396/// [`enum@Error`] cannot cross a wire: it wraps `reqwest::Error` and
397/// `serde_json::Error`, neither of which is `Serialize`. `WireError` keeps the
398/// parts a remote client can act on — the category, a human-readable message,
399/// the provider, and whether a retry might help — and drops the rest.
400///
401/// The conversion is deliberately one-way. Rebuilding a `reqwest::Error` from
402/// JSON is not possible, so a client matches on [`kind`](WireError::kind)
403/// rather than on the original error variant.
404///
405/// # Examples
406///
407/// ```
408/// use rai_sdk::{Error, ProviderKind};
409/// use rai_sdk::wire::{WireError, WireErrorKind};
410///
411/// let error = Error::RateLimit {
412///     provider: ProviderKind::OpenAI,
413///     message: "slow down".to_string(),
414/// };
415///
416/// let wire = WireError::from(&error);
417/// assert_eq!(wire.kind, WireErrorKind::RateLimit);
418/// assert_eq!(wire.provider, Some(ProviderKind::OpenAI));
419/// assert!(wire.retryable);
420/// ```
421#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
422pub struct WireError {
423    /// Error category, mirroring [`Error::kind_str`].
424    pub kind: WireErrorKind,
425
426    /// Human-readable description, from the original error's `Display`.
427    pub message: String,
428
429    /// Provider the failure came from, when the error names one.
430    #[serde(default, skip_serializing_if = "Option::is_none")]
431    pub provider: Option<ProviderKind>,
432
433    /// Whether retrying the request might succeed, from [`Error::is_retryable`].
434    pub retryable: bool,
435}
436
437impl WireError {
438    /// Build an error of `kind` with `message`, with no provider and not
439    /// retryable.
440    ///
441    /// Use this for failures a proxy raises itself — a truncated upstream, a
442    /// rejected protocol version — that have no [`enum@Error`] behind them.
443    pub fn new(kind: WireErrorKind, message: impl Into<String>) -> Self {
444        Self {
445            kind,
446            message: message.into(),
447            provider: None,
448            retryable: false,
449        }
450    }
451}
452
453impl From<&Error> for WireError {
454    fn from(error: &Error) -> Self {
455        Self {
456            kind: WireErrorKind::from(error.kind_str()),
457            message: error.to_string(),
458            provider: error.provider(),
459            retryable: error.is_retryable(),
460        }
461    }
462}
463
464impl std::fmt::Display for WireError {
465    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
466        write!(f, "{}: {}", self.kind.as_str(), self.message)
467    }
468}
469
470impl std::error::Error for WireError {}
471
472/// The category of a [`WireError`], mirroring [`Error::kind_str`].
473///
474/// Serializes as the same snake_case string [`Error::kind_str`] returns, so the
475/// two stay interchangeable in logs and metrics. A value this build does not
476/// know deserializes into [`Other`](WireErrorKind::Other) rather than failing,
477/// which is what lets an older client keep parsing streams from a newer server.
478///
479/// This enum is `#[non_exhaustive]`: match with a catch-all arm.
480#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
481#[serde(rename_all = "snake_case")]
482#[non_exhaustive]
483pub enum WireErrorKind {
484    /// Authentication failed.
485    Auth,
486    /// The provider request failed for a reason with no more specific variant.
487    Request,
488    /// The provider throttled the request.
489    RateLimit,
490    /// The request itself was invalid.
491    InvalidRequest,
492    /// The requested model is unavailable.
493    ModelNotAvailable,
494    /// The provider has no credentials configured.
495    ProviderNotConfigured,
496    /// The provider's Cargo feature is disabled in the server build.
497    ProviderNotEnabled,
498    /// The endpoint does not implement a capability the request needed.
499    CapabilityUnsupported,
500    /// The provider filtered the content.
501    ContentFiltered,
502    /// The SDK was misconfigured.
503    Config,
504    /// A value failed to serialize or deserialize.
505    Serialization,
506    /// The HTTP transport failed.
507    Http,
508    /// The stream itself failed partway through.
509    Stream,
510    /// The request timed out.
511    Timeout,
512    /// The provider does not support tool calling.
513    ToolProviderUnsupported,
514    /// Tool arguments failed schema validation.
515    ToolArguments,
516    /// A requested tool is not registered.
517    ToolNotFound,
518    /// The tool loop exceeded its configured round limit.
519    ToolLoopLimitExceeded,
520    /// Structured output failed validation.
521    StructuredOutput,
522    /// A category this build does not know about.
523    ///
524    /// Produced when deserializing a stream from a newer `rai-sdk`. Serializes
525    /// back to the original string, so a proxy can forward it untouched.
526    #[serde(untagged)]
527    Other(String),
528}
529
530impl WireErrorKind {
531    /// The snake_case name this kind serializes as.
532    pub fn as_str(&self) -> &str {
533        match self {
534            Self::Auth => "auth",
535            Self::Request => "request",
536            Self::RateLimit => "rate_limit",
537            Self::InvalidRequest => "invalid_request",
538            Self::ModelNotAvailable => "model_not_available",
539            Self::ProviderNotConfigured => "provider_not_configured",
540            Self::ProviderNotEnabled => "provider_not_enabled",
541            Self::CapabilityUnsupported => "capability_unsupported",
542            Self::ContentFiltered => "content_filtered",
543            Self::Config => "config",
544            Self::Serialization => "serialization",
545            Self::Http => "http",
546            Self::Stream => "stream",
547            Self::Timeout => "timeout",
548            Self::ToolProviderUnsupported => "tool_provider_unsupported",
549            Self::ToolArguments => "tool_arguments",
550            Self::ToolNotFound => "tool_not_found",
551            Self::ToolLoopLimitExceeded => "tool_loop_limit_exceeded",
552            Self::StructuredOutput => "structured_output",
553            Self::Other(kind) => kind,
554        }
555    }
556}
557
558impl From<&str> for WireErrorKind {
559    fn from(kind: &str) -> Self {
560        match kind {
561            "auth" => Self::Auth,
562            "request" => Self::Request,
563            "rate_limit" => Self::RateLimit,
564            "invalid_request" => Self::InvalidRequest,
565            "model_not_available" => Self::ModelNotAvailable,
566            "provider_not_configured" => Self::ProviderNotConfigured,
567            "provider_not_enabled" => Self::ProviderNotEnabled,
568            "capability_unsupported" => Self::CapabilityUnsupported,
569            "content_filtered" => Self::ContentFiltered,
570            "config" => Self::Config,
571            "serialization" => Self::Serialization,
572            "http" => Self::Http,
573            "stream" => Self::Stream,
574            "timeout" => Self::Timeout,
575            "tool_provider_unsupported" => Self::ToolProviderUnsupported,
576            "tool_arguments" => Self::ToolArguments,
577            "tool_not_found" => Self::ToolNotFound,
578            "tool_loop_limit_exceeded" => Self::ToolLoopLimitExceeded,
579            "structured_output" => Self::StructuredOutput,
580            other => Self::Other(other.to_string()),
581        }
582    }
583}
584
585impl std::fmt::Display for WireErrorKind {
586    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
587        f.write_str(self.as_str())
588    }
589}
590
591// ── Client-side reassembly ─────────────────────────────────────────────────
592
593/// A tool call whose arguments are still arriving.
594#[derive(Debug, Clone)]
595struct PendingToolCall {
596    id: String,
597    name: String,
598    arguments: String,
599}
600
601impl PendingToolCall {
602    fn finish(self) -> ToolCall {
603        ToolCall {
604            id: self.id,
605            name: self.name,
606            arguments: serde_json::from_str(&self.arguments).unwrap_or(serde_json::Value::Null),
607        }
608    }
609}
610
611/// Rebuilds one [`Response`] from a stream of [`WireStreamEvent`]s.
612///
613/// This is the receiving half of the proxy pattern: the client-side equivalent
614/// of [`RequestBuilder::stream_accumulated`](crate::RequestBuilder::stream_accumulated),
615/// operating on events that arrived over a network rather than on a live
616/// provider connection.
617///
618/// Unlike `stream_accumulated`, this *does* reassemble tool calls: fragments
619/// from [`ToolCallStart`](WireStreamEvent::ToolCallStart) /
620/// [`ToolCallDelta`](WireStreamEvent::ToolCallDelta) /
621/// [`ToolCallEnd`](WireStreamEvent::ToolCallEnd) are concatenated and attached
622/// to the assistant message.
623///
624/// # Truncation is an error
625///
626/// [`finish`](StreamAccumulator::finish) fails unless the stream was
627/// well-formed: it must have opened with
628/// [`MessageStart`](WireStreamEvent::MessageStart) and ended with a terminal
629/// event. A stream that just stops — the client's socket died mid-generation —
630/// yields [`WireErrorKind::Stream`], which is what distinguishes "the network
631/// died" from "the provider refused" (the latter arrives as
632/// [`WireStreamEvent::Error`] and is returned as its own [`WireError`]).
633///
634/// # Examples
635///
636/// ```
637/// use rai_sdk::ProviderKind;
638/// use rai_sdk::wire::{StreamAccumulator, WireStreamEvent};
639///
640/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
641/// let mut accumulator = StreamAccumulator::new();
642/// accumulator.push(WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI))?;
643/// accumulator.push(WireStreamEvent::TextDelta { text: "done".into() })?;
644/// accumulator.push(WireStreamEvent::MessageStop { finish_reason: Some("stop".into()) })?;
645///
646/// let response = accumulator.finish()?;
647/// assert_eq!(response.text(), "done");
648/// assert_eq!(response.finish_reason.as_deref(), Some("stop"));
649/// # Ok(())
650/// # }
651/// ```
652#[derive(Debug, Clone, Default)]
653pub struct StreamAccumulator {
654    protocol_version: Option<u32>,
655    model: Option<String>,
656    provider: Option<ProviderKind>,
657    text: String,
658    tool_calls: Vec<ToolCall>,
659    pending_tool_call: Option<PendingToolCall>,
660    tool_results: Vec<Message>,
661    usage: Option<Usage>,
662    finish_reason: Option<String>,
663    error: Option<WireError>,
664    stopped: bool,
665}
666
667impl StreamAccumulator {
668    /// Start with an empty accumulator.
669    pub fn new() -> Self {
670        Self::default()
671    }
672
673    /// The protocol version the sender advertised, once
674    /// [`MessageStart`](WireStreamEvent::MessageStart) has been seen.
675    pub fn protocol_version(&self) -> Option<u32> {
676        self.protocol_version
677    }
678
679    /// The assistant text accumulated so far.
680    pub fn text(&self) -> &str {
681        &self.text
682    }
683
684    /// Token usage, once a [`Usage`](WireStreamEvent::Usage) event has arrived.
685    pub fn usage(&self) -> Option<&Usage> {
686        self.usage.as_ref()
687    }
688
689    /// Whether a terminal event has been seen.
690    ///
691    /// A stream whose events are exhausted while this is `false` was truncated.
692    pub fn is_complete(&self) -> bool {
693        self.stopped || self.error.is_some()
694    }
695
696    /// Absorb one event.
697    ///
698    /// # Errors
699    ///
700    /// Returns the carried [`WireError`] when `event` is
701    /// [`WireStreamEvent::Error`], so a caller driving this in a loop can stop
702    /// with `?`. The error is remembered, so a later
703    /// [`finish`](StreamAccumulator::finish) reports it too.
704    pub fn push(&mut self, event: WireStreamEvent) -> std::result::Result<(), WireError> {
705        match event {
706            WireStreamEvent::MessageStart {
707                protocol_version,
708                model,
709                provider,
710            } => {
711                self.protocol_version = Some(protocol_version);
712                self.model = Some(model);
713                self.provider = Some(provider);
714            }
715
716            WireStreamEvent::TextDelta { text } => self.text.push_str(&text),
717
718            WireStreamEvent::ToolCallStart { id, name } => {
719                self.flush_pending_tool_call();
720                self.pending_tool_call = Some(PendingToolCall {
721                    id,
722                    name,
723                    arguments: String::new(),
724                });
725            }
726
727            WireStreamEvent::ToolCallDelta { id, arguments } => match &mut self.pending_tool_call {
728                Some(pending) if pending.id == id => pending.arguments.push_str(&arguments),
729                _ => {
730                    self.flush_pending_tool_call();
731                    self.pending_tool_call = Some(PendingToolCall {
732                        id,
733                        name: String::new(),
734                        arguments,
735                    });
736                }
737            },
738
739            WireStreamEvent::ToolCallEnd {
740                id,
741                name,
742                arguments,
743            } => {
744                // The assembled form supersedes whatever was buffered for the
745                // same call, so drop the partial rather than emitting both.
746                if self
747                    .pending_tool_call
748                    .as_ref()
749                    .is_some_and(|pending| pending.id == id)
750                {
751                    self.pending_tool_call = None;
752                } else {
753                    self.flush_pending_tool_call();
754                }
755                self.tool_calls.push(
756                    PendingToolCall {
757                        id,
758                        name,
759                        arguments,
760                    }
761                    .finish(),
762                );
763            }
764
765            WireStreamEvent::ToolResult { id, result } => {
766                self.tool_results.push(Message::tool(result, id));
767            }
768
769            WireStreamEvent::Usage { usage } => self.usage = Some(usage),
770
771            WireStreamEvent::MessageStop { finish_reason } => {
772                self.flush_pending_tool_call();
773                if finish_reason.is_some() {
774                    self.finish_reason = finish_reason;
775                }
776                self.stopped = true;
777            }
778
779            WireStreamEvent::TurnComplete { turn } => {
780                self.flush_pending_tool_call();
781                if self.text.is_empty() {
782                    self.text = turn.assistant_message.text_content();
783                }
784                if self.tool_calls.is_empty() {
785                    self.tool_calls = turn.assistant_message.tool_calls.clone();
786                }
787                self.tool_results.extend(turn.tool_results);
788            }
789
790            WireStreamEvent::Error { error } => {
791                self.error = Some(error.clone());
792                return Err(error);
793            }
794
795            // `WireStreamEvent` is `#[non_exhaustive]` in spirit for consumers,
796            // but this match is inside the defining crate, so every arm above
797            // is checked exhaustively and a new variant is a compile error here.
798            #[allow(unreachable_patterns)]
799            _ => {}
800        }
801
802        Ok(())
803    }
804
805    fn flush_pending_tool_call(&mut self) {
806        if let Some(pending) = self.pending_tool_call.take() {
807            self.tool_calls.push(pending.finish());
808        }
809    }
810
811    /// Assemble everything pushed so far into one [`Response`].
812    ///
813    /// # Errors
814    ///
815    /// - The [`WireError`] from a [`WireStreamEvent::Error`], if one arrived.
816    /// - [`WireErrorKind::Stream`] if no
817    ///   [`MessageStart`](WireStreamEvent::MessageStart) was seen, since the
818    ///   model and provider a [`Response`] requires are unknown.
819    /// - [`WireErrorKind::Stream`] if no terminal event was seen, meaning the
820    ///   stream was cut short rather than finished.
821    pub fn finish(mut self) -> std::result::Result<Response, WireError> {
822        if let Some(error) = self.error {
823            return Err(error);
824        }
825
826        self.flush_pending_tool_call();
827
828        let (Some(model), Some(provider)) = (self.model, self.provider) else {
829            return Err(WireError::new(
830                WireErrorKind::Stream,
831                "stream ended without a message_start event, so the model and provider are unknown",
832            ));
833        };
834
835        if !self.stopped {
836            return Err(WireError::new(
837                WireErrorKind::Stream,
838                "stream ended before its terminal event (message_stop or error); \
839                 the connection was most likely dropped mid-generation",
840            ));
841        }
842
843        let mut assistant = Message::assistant(self.text);
844        assistant.tool_calls = self.tool_calls;
845
846        let mut messages = vec![assistant];
847        messages.extend(self.tool_results);
848
849        Ok(Response {
850            messages,
851            usage: self.usage,
852            model,
853            provider,
854            finish_reason: self.finish_reason,
855        })
856    }
857
858    /// Drain a stream of wire events into one [`Response`].
859    ///
860    /// The convenience form of `new()` + `push()` in a loop + `finish()`, for
861    /// the common client-side case where the whole stream is consumed at once.
862    ///
863    /// # Errors
864    ///
865    /// Same as [`push`](StreamAccumulator::push) and
866    /// [`finish`](StreamAccumulator::finish).
867    ///
868    /// # Examples
869    ///
870    /// ```
871    /// use rai_sdk::ProviderKind;
872    /// use rai_sdk::wire::{StreamAccumulator, WireStreamEvent};
873    ///
874    /// # #[tokio::main]
875    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
876    /// let events = futures::stream::iter(vec![
877    ///     WireStreamEvent::message_start("claude-sonnet-4-6", ProviderKind::Anthropic),
878    ///     WireStreamEvent::TextDelta { text: "hi".into() },
879    ///     WireStreamEvent::MessageStop { finish_reason: Some("stop".into()) },
880    /// ]);
881    ///
882    /// let response = StreamAccumulator::accumulate(events).await?;
883    /// assert_eq!(response.text(), "hi");
884    /// # Ok(())
885    /// # }
886    /// ```
887    pub async fn accumulate<S>(stream: S) -> std::result::Result<Response, WireError>
888    where
889        S: Stream<Item = WireStreamEvent>,
890    {
891        let mut stream = std::pin::pin!(stream);
892        let mut accumulator = Self::new();
893
894        while let Some(event) = stream.next().await {
895            accumulator.push(event)?;
896        }
897
898        accumulator.finish()
899    }
900}
901
902#[cfg(test)]
903mod tests {
904    //! Wire-format tests deliberately live *inside* the crate.
905    //!
906    //! [`WireStreamEvent`] is `#[non_exhaustive]`, so an integration test in
907    //! `tests/` — a separate crate — would be forced to write a `_ => {}` arm
908    //! and adding a variant would not break it. In here the match is checked
909    //! exhaustively, so a new variant fails to compile until it has both a
910    //! round-trip case and a committed JSON fixture.
911
912    use super::*;
913    use crate::message::{ContentBlock, Message};
914
915    /// Directory holding one committed JSON fixture per variant.
916    const FIXTURE_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/wire");
917
918    /// Set this to rewrite the fixtures instead of asserting against them.
919    const UPDATE_ENV: &str = "RAI_SDK_UPDATE_WIRE_FIXTURES";
920
921    fn sample_usage() -> Usage {
922        Usage {
923            prompt_tokens: Some(1_024),
924            completion_tokens: Some(256),
925            total_tokens: Some(1_280),
926        }
927    }
928
929    fn sample_turn() -> ConversationTurn {
930        let mut assistant = Message::assistant("Sunny.");
931        assistant.tool_calls = vec![ToolCall {
932            id: "call_1".to_string(),
933            name: "get_weather".to_string(),
934            arguments: serde_json::json!({ "city": "Paris" }),
935        }];
936
937        ConversationTurn {
938            user_message: Message::user("Weather in Paris?"),
939            assistant_message: assistant,
940            tool_results: vec![Message::tool("{\"c\":21}", "call_1")],
941        }
942    }
943
944    /// One representative value per variant.
945    ///
946    /// The exhaustive `match` in [`fixture_name`] is what forces this list to
947    /// grow with the enum.
948    fn every_variant() -> Vec<WireStreamEvent> {
949        vec![
950            WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI),
951            WireStreamEvent::TextDelta {
952                text: "Hello, world.".to_string(),
953            },
954            WireStreamEvent::ToolCallStart {
955                id: "call_1".to_string(),
956                name: "get_weather".to_string(),
957            },
958            WireStreamEvent::ToolCallDelta {
959                id: "call_1".to_string(),
960                arguments: "{\"city\":".to_string(),
961            },
962            WireStreamEvent::ToolCallEnd {
963                id: "call_1".to_string(),
964                name: "get_weather".to_string(),
965                arguments: "{\"city\":\"Paris\"}".to_string(),
966            },
967            WireStreamEvent::ToolResult {
968                id: "call_1".to_string(),
969                result: "{\"celsius\":21}".to_string(),
970            },
971            WireStreamEvent::Usage {
972                usage: sample_usage(),
973            },
974            WireStreamEvent::MessageStop {
975                finish_reason: Some("stop".to_string()),
976            },
977            WireStreamEvent::TurnComplete {
978                turn: sample_turn(),
979            },
980            WireStreamEvent::error(&Error::RateLimit {
981                provider: ProviderKind::Anthropic,
982                message: "too many requests".to_string(),
983            }),
984        ]
985    }
986
987    /// Fixture file stem for an event.
988    ///
989    /// Exhaustive on purpose: adding a variant to [`WireStreamEvent`] breaks
990    /// this match, which is the signal to add it to [`every_variant`] and
991    /// commit a fixture for it.
992    fn fixture_name(event: &WireStreamEvent) -> &'static str {
993        match event {
994            WireStreamEvent::MessageStart { .. } => "message_start",
995            WireStreamEvent::TextDelta { .. } => "text_delta",
996            WireStreamEvent::ToolCallStart { .. } => "tool_call_start",
997            WireStreamEvent::ToolCallDelta { .. } => "tool_call_delta",
998            WireStreamEvent::ToolCallEnd { .. } => "tool_call_end",
999            WireStreamEvent::ToolResult { .. } => "tool_result",
1000            WireStreamEvent::Usage { .. } => "usage",
1001            WireStreamEvent::MessageStop { .. } => "message_stop",
1002            WireStreamEvent::TurnComplete { .. } => "turn_complete",
1003            WireStreamEvent::Error { .. } => "error",
1004        }
1005    }
1006
1007    #[test]
1008    fn every_variant_is_covered_exactly_once() {
1009        let mut names: Vec<&str> = every_variant().iter().map(fixture_name).collect();
1010        let total = names.len();
1011        names.sort_unstable();
1012        names.dedup();
1013        assert_eq!(
1014            names.len(),
1015            total,
1016            "every_variant() must contain each variant exactly once"
1017        );
1018    }
1019
1020    #[test]
1021    fn round_trip_preserves_every_variant() {
1022        for event in every_variant() {
1023            let json = serde_json::to_string(&event).expect("event should serialize");
1024            let parsed: WireStreamEvent =
1025                serde_json::from_str(&json).expect("event should deserialize");
1026            assert_eq!(
1027                parsed,
1028                event,
1029                "round trip changed the '{}' event: {json}",
1030                fixture_name(&event)
1031            );
1032        }
1033    }
1034
1035    #[test]
1036    fn the_tag_matches_the_serialized_discriminant() {
1037        for event in every_variant() {
1038            let json = serde_json::to_value(&event).expect("event should serialize");
1039            assert_eq!(
1040                json.get("type").and_then(serde_json::Value::as_str),
1041                Some(event.tag()),
1042                "tag() disagrees with the serialized discriminant: {json}"
1043            );
1044            assert_eq!(event.tag(), fixture_name(&event));
1045        }
1046    }
1047
1048    /// Pins the exact JSON each variant produces.
1049    ///
1050    /// The fixtures are the compatibility contract with deployed clients, so an
1051    /// accidental rename shows up here as a diff rather than in production.
1052    #[test]
1053    fn wire_format_matches_the_committed_fixtures() {
1054        let updating = std::env::var_os(UPDATE_ENV).is_some();
1055
1056        for event in every_variant() {
1057            let name = fixture_name(&event);
1058            let path = std::path::Path::new(FIXTURE_DIR).join(format!("{name}.json"));
1059            let actual = serde_json::to_value(&event).expect("event should serialize");
1060
1061            if updating {
1062                std::fs::create_dir_all(FIXTURE_DIR).expect("create the fixture directory");
1063                let pretty =
1064                    serde_json::to_string_pretty(&actual).expect("fixture should serialize");
1065                std::fs::write(&path, format!("{pretty}\n")).expect("write the fixture");
1066                continue;
1067            }
1068
1069            let raw = std::fs::read_to_string(&path).unwrap_or_else(|error| {
1070                panic!(
1071                    "missing wire-format fixture {}: {error}\n\
1072                     If this variant is new, regenerate the fixtures with:\n    \
1073                     {UPDATE_ENV}=1 cargo test --all-features wire_format_matches_the_committed_fixtures",
1074                    path.display()
1075                )
1076            });
1077            let expected: serde_json::Value =
1078                serde_json::from_str(&raw).expect("fixture should be valid JSON");
1079
1080            assert_eq!(
1081                actual,
1082                expected,
1083                "the '{name}' event no longer matches {}.\n\
1084                 Wire-format changes break clients built against an older rai-sdk. \
1085                 If this change is intentional, note it in CHANGELOG.md and regenerate with:\n    \
1086                 {UPDATE_ENV}=1 cargo test --all-features wire_format_matches_the_committed_fixtures",
1087                path.display()
1088            );
1089
1090            // The fixture must also parse back, so a hand-edited file cannot
1091            // drift into a shape the crate can no longer read.
1092            let parsed: WireStreamEvent =
1093                serde_json::from_value(expected).expect("fixture should deserialize");
1094            assert_eq!(parsed, event);
1095        }
1096    }
1097
1098    #[test]
1099    fn stream_event_conversion_round_trips() {
1100        let events = vec![
1101            StreamEvent::TextDelta {
1102                text: "hi".to_string(),
1103            },
1104            StreamEvent::ToolCall {
1105                id: "call_1".to_string(),
1106                name: "get_weather".to_string(),
1107                arguments: "{\"city\":\"Paris\"}".to_string(),
1108            },
1109            StreamEvent::ToolResult {
1110                id: "call_1".to_string(),
1111                result: "{\"celsius\":21}".to_string(),
1112            },
1113            StreamEvent::TurnComplete {
1114                turn: sample_turn(),
1115            },
1116        ];
1117
1118        for event in events {
1119            let wire = WireStreamEvent::from(event.clone());
1120            let back = StreamEvent::try_from(wire).expect("wire event should convert back");
1121            assert_eq!(back, event);
1122        }
1123    }
1124
1125    #[test]
1126    fn framing_events_have_no_stream_event_equivalent() {
1127        let wire = WireStreamEvent::MessageStop {
1128            finish_reason: None,
1129        };
1130        let error = StreamEvent::try_from(wire).expect_err("message_stop is wire-only");
1131        assert_eq!(error.tag, "message_stop");
1132        assert!(error.to_string().contains("message_stop"));
1133    }
1134
1135    #[test]
1136    fn message_start_defaults_the_protocol_version_when_absent() {
1137        let parsed: WireStreamEvent = serde_json::from_str(
1138            r#"{"type":"message_start","model":"gpt-4o-mini","provider":"openai"}"#,
1139        )
1140        .expect("legacy payload should still parse");
1141
1142        assert_eq!(
1143            parsed,
1144            WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI)
1145        );
1146    }
1147
1148    #[test]
1149    fn is_terminal_marks_exactly_the_stream_enders() {
1150        for event in every_variant() {
1151            let expected = matches!(fixture_name(&event), "message_stop" | "error");
1152            assert_eq!(event.is_terminal(), expected, "{}", event.tag());
1153        }
1154    }
1155
1156    /// Every [`enum@Error`] variant must map onto a named [`WireErrorKind`].
1157    ///
1158    /// Exhaustive so a new `Error` variant cannot silently start serializing as
1159    /// [`WireErrorKind::Other`].
1160    #[test]
1161    fn every_error_variant_maps_to_a_named_kind() {
1162        let provider = ProviderKind::OpenAI;
1163        let errors = vec![
1164            Error::Auth {
1165                provider,
1166                message: "bad key".into(),
1167            },
1168            Error::Request {
1169                provider,
1170                message: "boom".into(),
1171            },
1172            Error::RateLimit {
1173                provider,
1174                message: "slow down".into(),
1175            },
1176            Error::InvalidRequest("nope".into()),
1177            Error::ModelNotAvailable {
1178                provider,
1179                model: "gpt-9".into(),
1180            },
1181            Error::ProviderNotConfigured(provider),
1182            Error::ProviderNotEnabled(provider),
1183            Error::ContentFiltered {
1184                provider,
1185                reason: "policy".into(),
1186            },
1187            Error::Config("missing".into()),
1188            Error::Serialization(
1189                serde_json::from_str::<serde_json::Value>("{").expect_err("invalid JSON"),
1190            ),
1191            Error::Stream("truncated".into()),
1192            Error::Timeout { provider },
1193            Error::ToolProviderUnsupported { provider },
1194            Error::ToolArguments {
1195                name: "echo".into(),
1196                message: "bad".into(),
1197                issues: Vec::new(),
1198            },
1199            Error::ToolNotFound {
1200                name: "echo".into(),
1201            },
1202            Error::ToolLoopLimitExceeded { max_rounds: 8 },
1203            Error::StructuredOutput {
1204                provider,
1205                model: "gpt-4o-mini".into(),
1206                message: "invalid".into(),
1207            },
1208        ];
1209
1210        for error in &errors {
1211            let wire = WireError::from(error);
1212            assert!(
1213                !matches!(wire.kind, WireErrorKind::Other(_)),
1214                "{} has no named WireErrorKind",
1215                error.kind_str()
1216            );
1217            assert_eq!(wire.kind.as_str(), error.kind_str());
1218            assert_eq!(wire.message, error.to_string());
1219            assert_eq!(wire.provider, error.provider());
1220            assert_eq!(wire.retryable, error.is_retryable());
1221        }
1222
1223        // `Error::Http` is the one variant that cannot be constructed outside
1224        // `reqwest`, so its mapping is asserted on the kind string directly.
1225        assert_eq!(WireErrorKind::from("http"), WireErrorKind::Http);
1226    }
1227
1228    #[test]
1229    fn unknown_error_kinds_survive_a_round_trip() {
1230        let json = r#"{"kind":"quantum_flux","message":"from the future","retryable":true}"#;
1231        let parsed: WireError = serde_json::from_str(json).expect("unknown kind should parse");
1232
1233        assert_eq!(parsed.kind, WireErrorKind::Other("quantum_flux".into()));
1234        assert_eq!(parsed.kind.as_str(), "quantum_flux");
1235        assert_eq!(parsed.provider, None);
1236
1237        let reserialized = serde_json::to_value(&parsed).expect("should serialize");
1238        assert_eq!(reserialized["kind"], "quantum_flux");
1239    }
1240
1241    #[test]
1242    fn accumulator_rebuilds_text_usage_and_finish_reason() {
1243        let mut accumulator = StreamAccumulator::new();
1244        for event in [
1245            WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI),
1246            WireStreamEvent::TextDelta {
1247                text: "Hello, ".into(),
1248            },
1249            WireStreamEvent::TextDelta {
1250                text: "world.".into(),
1251            },
1252            WireStreamEvent::Usage {
1253                usage: sample_usage(),
1254            },
1255            WireStreamEvent::MessageStop {
1256                finish_reason: Some("stop".into()),
1257            },
1258        ] {
1259            accumulator.push(event).expect("event should be absorbed");
1260        }
1261
1262        assert_eq!(accumulator.protocol_version(), Some(WIRE_PROTOCOL_VERSION));
1263        assert!(accumulator.is_complete());
1264
1265        let response = accumulator.finish().expect("stream was well formed");
1266        assert_eq!(response.text(), "Hello, world.");
1267        assert_eq!(response.model, "gpt-4o-mini");
1268        assert_eq!(response.provider, ProviderKind::OpenAI);
1269        assert_eq!(response.finish_reason.as_deref(), Some("stop"));
1270        assert_eq!(response.usage, Some(sample_usage()));
1271    }
1272
1273    #[test]
1274    fn accumulator_assembles_tool_calls_from_fragments() {
1275        let mut accumulator = StreamAccumulator::new();
1276        for event in [
1277            WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI),
1278            WireStreamEvent::ToolCallStart {
1279                id: "call_1".into(),
1280                name: "get_weather".into(),
1281            },
1282            WireStreamEvent::ToolCallDelta {
1283                id: "call_1".into(),
1284                arguments: "{\"city\":".into(),
1285            },
1286            WireStreamEvent::ToolCallDelta {
1287                id: "call_1".into(),
1288                arguments: "\"Paris\"}".into(),
1289            },
1290            WireStreamEvent::MessageStop {
1291                finish_reason: Some("tool_use".into()),
1292            },
1293        ] {
1294            accumulator.push(event).expect("event should be absorbed");
1295        }
1296
1297        let response = accumulator.finish().expect("stream was well formed");
1298        let tool_calls = &response.messages[0].tool_calls;
1299        assert_eq!(tool_calls.len(), 1);
1300        assert_eq!(tool_calls[0].name, "get_weather");
1301        assert_eq!(tool_calls[0].arguments["city"], "Paris");
1302    }
1303
1304    #[test]
1305    fn a_tool_call_end_supersedes_its_own_fragments() {
1306        let mut accumulator = StreamAccumulator::new();
1307        for event in [
1308            WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI),
1309            WireStreamEvent::ToolCallStart {
1310                id: "call_1".into(),
1311                name: "get_weather".into(),
1312            },
1313            WireStreamEvent::ToolCallDelta {
1314                id: "call_1".into(),
1315                arguments: "{\"city\":".into(),
1316            },
1317            WireStreamEvent::ToolCallEnd {
1318                id: "call_1".into(),
1319                name: "get_weather".into(),
1320                arguments: "{\"city\":\"Paris\"}".into(),
1321            },
1322            WireStreamEvent::MessageStop {
1323                finish_reason: Some("tool_use".into()),
1324            },
1325        ] {
1326            accumulator.push(event).expect("event should be absorbed");
1327        }
1328
1329        let response = accumulator.finish().expect("stream was well formed");
1330        assert_eq!(response.messages[0].tool_calls.len(), 1);
1331    }
1332
1333    #[test]
1334    fn accumulator_keeps_tool_results_as_messages() {
1335        let mut accumulator = StreamAccumulator::new();
1336        for event in [
1337            WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI),
1338            WireStreamEvent::ToolResult {
1339                id: "call_1".into(),
1340                result: "{\"celsius\":21}".into(),
1341            },
1342            WireStreamEvent::MessageStop {
1343                finish_reason: Some("stop".into()),
1344            },
1345        ] {
1346            accumulator.push(event).expect("event should be absorbed");
1347        }
1348
1349        let response = accumulator.finish().expect("stream was well formed");
1350        assert_eq!(response.messages.len(), 2);
1351        assert_eq!(response.messages[1].tool_call_id.as_deref(), Some("call_1"));
1352    }
1353
1354    #[test]
1355    fn an_error_event_stops_the_accumulator() {
1356        let mut accumulator = StreamAccumulator::new();
1357        accumulator
1358            .push(WireStreamEvent::message_start(
1359                "gpt-4o-mini",
1360                ProviderKind::OpenAI,
1361            ))
1362            .expect("event should be absorbed");
1363
1364        let error = accumulator
1365            .push(WireStreamEvent::error(&Error::ContentFiltered {
1366                provider: ProviderKind::OpenAI,
1367                reason: "policy".into(),
1368            }))
1369            .expect_err("an error event should surface as an error");
1370
1371        assert_eq!(error.kind, WireErrorKind::ContentFiltered);
1372        assert!(accumulator.is_complete());
1373        assert_eq!(
1374            accumulator
1375                .finish()
1376                .expect_err("finish reports it too")
1377                .kind,
1378            WireErrorKind::ContentFiltered
1379        );
1380    }
1381
1382    #[test]
1383    fn a_truncated_stream_is_distinguishable_from_a_provider_error() {
1384        let mut accumulator = StreamAccumulator::new();
1385        for event in [
1386            WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI),
1387            WireStreamEvent::TextDelta {
1388                text: "half a sen".into(),
1389            },
1390        ] {
1391            accumulator.push(event).expect("event should be absorbed");
1392        }
1393        assert!(!accumulator.is_complete());
1394
1395        let error = accumulator
1396            .finish()
1397            .expect_err("a truncated stream is an error");
1398        assert_eq!(error.kind, WireErrorKind::Stream);
1399        assert!(
1400            error.message.contains("terminal event"),
1401            "unhelpful truncation message: {}",
1402            error.message
1403        );
1404    }
1405
1406    #[test]
1407    fn a_stream_without_message_start_is_rejected() {
1408        let mut accumulator = StreamAccumulator::new();
1409        accumulator
1410            .push(WireStreamEvent::MessageStop {
1411                finish_reason: Some("stop".into()),
1412            })
1413            .expect("event should be absorbed");
1414
1415        let error = accumulator
1416            .finish()
1417            .expect_err("model and provider are unknown");
1418        assert_eq!(error.kind, WireErrorKind::Stream);
1419        assert!(error.message.contains("message_start"));
1420    }
1421
1422    #[test]
1423    fn turn_complete_backfills_an_otherwise_empty_response() {
1424        let mut accumulator = StreamAccumulator::new();
1425        for event in [
1426            WireStreamEvent::message_start("claude-sonnet-4-6", ProviderKind::Anthropic),
1427            WireStreamEvent::TurnComplete {
1428                turn: sample_turn(),
1429            },
1430            WireStreamEvent::MessageStop {
1431                finish_reason: Some("tool_use".into()),
1432            },
1433        ] {
1434            accumulator.push(event).expect("event should be absorbed");
1435        }
1436
1437        let response = accumulator.finish().expect("stream was well formed");
1438        assert_eq!(response.text(), "Sunny.");
1439        assert_eq!(response.messages[0].tool_calls.len(), 1);
1440    }
1441
1442    #[tokio::test]
1443    async fn accumulate_drains_a_whole_stream() {
1444        let response = StreamAccumulator::accumulate(futures::stream::iter(vec![
1445            WireStreamEvent::message_start("gpt-4o-mini", ProviderKind::OpenAI),
1446            WireStreamEvent::TextDelta { text: "ok".into() },
1447            WireStreamEvent::MessageStop {
1448                finish_reason: Some("stop".into()),
1449            },
1450        ]))
1451        .await
1452        .expect("stream was well formed");
1453
1454        assert_eq!(response.text(), "ok");
1455    }
1456
1457    #[test]
1458    fn multimodal_turns_survive_the_wire() {
1459        let turn = ConversationTurn {
1460            user_message: Message::user_multimodal(vec![
1461                ContentBlock::text("What is this?"),
1462                ContentBlock::image_url("https://example.com/cat.png"),
1463            ]),
1464            assistant_message: Message::assistant("A cat."),
1465            tool_results: Vec::new(),
1466        };
1467
1468        let event = WireStreamEvent::TurnComplete { turn };
1469        let json = serde_json::to_string(&event).expect("event should serialize");
1470        let parsed: WireStreamEvent = serde_json::from_str(&json).expect("should deserialize");
1471        assert_eq!(parsed, event);
1472    }
1473}