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}