oino_ai/
stream.rs

1//! Streaming events, the partial-free provider form crossing the
2//! WIT boundary ([`ProviderEvent`]) and request types ([`Context`], [`StreamOptions`]).
3
4use indexmap::IndexMap;
5use serde::{Deserialize, Serialize};
6
7use crate::content::{AssistantContent, TextContent, ThinkingContent, ToolCall};
8use crate::json::{JsonObject, JsonValue, js_f64_opt};
9use crate::message::{AssistantMessage, Message};
10use crate::primitives::{CacheRetention, StopReason, ThinkingLevel};
11
12/// Successful terminal reason (`done.reason`).
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "camelCase")]
15pub enum DoneReason {
16    /// End of turn.
17    Stop,
18    /// Token limit.
19    Length,
20    /// Tool calls requested.
21    ToolUse,
22    /// Deferred response.
23    Deferred,
24}
25
26impl From<DoneReason> for StopReason {
27    fn from(reason: DoneReason) -> Self {
28        match reason {
29            DoneReason::Stop => Self::Stop,
30            DoneReason::Length => Self::Length,
31            DoneReason::ToolUse => Self::ToolUse,
32            DoneReason::Deferred => Self::Deferred,
33        }
34    }
35}
36
37/// Failure terminal reason (`error.reason`).
38#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(rename_all = "camelCase")]
40pub enum ErrorReason {
41    /// Aborted by the caller.
42    Aborted,
43    /// Provider/transport error.
44    Error,
45}
46
47impl From<ErrorReason> for StopReason {
48    fn from(reason: ErrorReason) -> Self {
49        match reason {
50            ErrorReason::Aborted => Self::Aborted,
51            ErrorReason::Error => Self::Error,
52        }
53    }
54}
55
56/// An assistant message event (with the live `partial` message).
57#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
58#[serde(tag = "type", rename_all = "snake_case")]
59pub enum AssistantMessageEvent {
60    /// Stream started.
61    Start {
62        /// Live message.
63        partial: AssistantMessage,
64    },
65    /// Text block opened.
66    #[serde(rename_all = "camelCase")]
67    TextStart {
68        /// Block index.
69        content_index: usize,
70        /// Live message.
71        partial: AssistantMessage,
72    },
73    /// Text appended.
74    #[serde(rename_all = "camelCase")]
75    TextDelta {
76        /// Block index.
77        content_index: usize,
78        /// Appended text.
79        delta: String,
80        /// Live message.
81        partial: AssistantMessage,
82    },
83    /// Text block closed (authoritative content).
84    #[serde(rename_all = "camelCase")]
85    TextEnd {
86        /// Block index.
87        content_index: usize,
88        /// Final text.
89        content: String,
90        /// Live message.
91        partial: AssistantMessage,
92    },
93    /// Thinking block opened.
94    #[serde(rename_all = "camelCase")]
95    ThinkingStart {
96        /// Block index.
97        content_index: usize,
98        /// Live message.
99        partial: AssistantMessage,
100    },
101    /// Thinking appended.
102    #[serde(rename_all = "camelCase")]
103    ThinkingDelta {
104        /// Block index.
105        content_index: usize,
106        /// Appended text.
107        delta: String,
108        /// Live message.
109        partial: AssistantMessage,
110    },
111    /// Thinking block closed.
112    #[serde(rename_all = "camelCase")]
113    ThinkingEnd {
114        /// Block index.
115        content_index: usize,
116        /// Final text.
117        content: String,
118        /// Live message.
119        partial: AssistantMessage,
120    },
121    /// Tool call opened.
122    #[serde(rename = "toolcall_start", rename_all = "camelCase")]
123    ToolcallStart {
124        /// Block index.
125        content_index: usize,
126        /// Live message.
127        partial: AssistantMessage,
128    },
129    /// Tool-call argument JSON appended.
130    #[serde(rename = "toolcall_delta", rename_all = "camelCase")]
131    ToolcallDelta {
132        /// Block index.
133        content_index: usize,
134        /// Appended argument JSON.
135        delta: String,
136        /// Live message.
137        partial: AssistantMessage,
138    },
139    /// Tool call closed (authoritative call).
140    #[serde(rename = "toolcall_end", rename_all = "camelCase")]
141    ToolcallEnd {
142        /// Block index.
143        content_index: usize,
144        /// Final call.
145        tool_call: ToolCall,
146        /// Live message.
147        partial: AssistantMessage,
148    },
149    /// Successful end.
150    Done {
151        /// Reason.
152        reason: DoneReason,
153        /// Final message.
154        message: AssistantMessage,
155    },
156    /// Failed end.
157    Error {
158        /// Reason.
159        reason: ErrorReason,
160        /// Final (error) message.
161        error: AssistantMessage,
162    },
163}
164
165/// Provider-emitted event: [`AssistantMessageEvent`] without `partial` (D5). `toolcall_start`
166/// may carry the call `id` and `toolName` early.
167#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
168#[serde(tag = "type", rename_all = "snake_case")]
169pub enum ProviderEvent {
170    /// Stream started.
171    Start {},
172    /// Text block opened.
173    #[serde(rename_all = "camelCase")]
174    TextStart {
175        /// Block index.
176        content_index: usize,
177    },
178    /// Text appended.
179    #[serde(rename_all = "camelCase")]
180    TextDelta {
181        /// Block index.
182        content_index: usize,
183        /// Appended text.
184        delta: String,
185    },
186    /// Text block closed.
187    #[serde(rename_all = "camelCase")]
188    TextEnd {
189        /// Block index.
190        content_index: usize,
191        /// Final text.
192        content: String,
193    },
194    /// Thinking block opened.
195    #[serde(rename_all = "camelCase")]
196    ThinkingStart {
197        /// Block index.
198        content_index: usize,
199    },
200    /// Thinking appended.
201    #[serde(rename_all = "camelCase")]
202    ThinkingDelta {
203        /// Block index.
204        content_index: usize,
205        /// Appended text.
206        delta: String,
207    },
208    /// Thinking block closed.
209    #[serde(rename_all = "camelCase")]
210    ThinkingEnd {
211        /// Block index.
212        content_index: usize,
213        /// Final text.
214        content: String,
215    },
216    /// Tool call opened.
217    #[serde(rename = "toolcall_start", rename_all = "camelCase")]
218    ToolcallStart {
219        /// Block index.
220        content_index: usize,
221        /// Call id when already known.
222        #[serde(default, skip_serializing_if = "Option::is_none")]
223        id: Option<String>,
224        /// Tool name when already known.
225        #[serde(default, skip_serializing_if = "Option::is_none")]
226        tool_name: Option<String>,
227    },
228    /// Tool-call argument JSON appended.
229    #[serde(rename = "toolcall_delta", rename_all = "camelCase")]
230    ToolcallDelta {
231        /// Block index.
232        content_index: usize,
233        /// Appended argument JSON.
234        delta: String,
235    },
236    /// Tool call closed.
237    #[serde(rename = "toolcall_end", rename_all = "camelCase")]
238    ToolcallEnd {
239        /// Block index.
240        content_index: usize,
241        /// Final call.
242        tool_call: ToolCall,
243    },
244    /// Successful end.
245    Done {
246        /// Reason.
247        reason: DoneReason,
248        /// Final message.
249        message: AssistantMessage,
250    },
251    /// Failed end.
252    Error {
253        /// Reason.
254        reason: ErrorReason,
255        /// Final message.
256        error: AssistantMessage,
257    },
258}
259
260/// Rebuilds `partial` messages from [`ProviderEvent`]s so the host emits exact events.
261#[derive(Clone, Debug)]
262pub struct PartialAssembler {
263    partial: AssistantMessage,
264    tool_args: IndexMap<usize, String>,
265}
266
267impl PartialAssembler {
268    /// Starts from `template` (api/provider/model/timestamp; content is cleared and
269    /// `stopReason` set to `pending`).
270    #[must_use]
271    pub fn new(mut template: AssistantMessage) -> Self {
272        template.content.clear();
273        template.stop_reason = StopReason::Pending;
274        Self {
275            partial: template,
276            tool_args: IndexMap::new(),
277        }
278    }
279
280    /// Current partial message.
281    #[must_use]
282    pub const fn partial(&self) -> &AssistantMessage {
283        &self.partial
284    }
285
286    /// Applies one provider event and returns the matching client event.
287    #[allow(
288        clippy::too_many_lines,
289        reason = "one arm per provider event keeps the 1:1 mapping readable"
290    )]
291    pub fn apply(&mut self, event: ProviderEvent) -> AssistantMessageEvent {
292        match event {
293            ProviderEvent::Start {} => AssistantMessageEvent::Start {
294                partial: self.partial.clone(),
295            },
296            ProviderEvent::TextStart { content_index } => {
297                self.put(
298                    content_index,
299                    AssistantContent::Text(TextContent::default()),
300                );
301                AssistantMessageEvent::TextStart {
302                    content_index,
303                    partial: self.partial.clone(),
304                }
305            }
306            ProviderEvent::TextDelta {
307                content_index,
308                delta,
309            } => {
310                if let Some(AssistantContent::Text(block)) = self.slot(content_index) {
311                    block.text.push_str(&delta);
312                }
313                AssistantMessageEvent::TextDelta {
314                    content_index,
315                    delta,
316                    partial: self.partial.clone(),
317                }
318            }
319            ProviderEvent::TextEnd {
320                content_index,
321                content,
322            } => {
323                if let Some(AssistantContent::Text(block)) = self.slot(content_index) {
324                    block.text.clone_from(&content);
325                }
326                AssistantMessageEvent::TextEnd {
327                    content_index,
328                    content,
329                    partial: self.partial.clone(),
330                }
331            }
332            ProviderEvent::ThinkingStart { content_index } => {
333                self.put(
334                    content_index,
335                    AssistantContent::Thinking(ThinkingContent::default()),
336                );
337                AssistantMessageEvent::ThinkingStart {
338                    content_index,
339                    partial: self.partial.clone(),
340                }
341            }
342            ProviderEvent::ThinkingDelta {
343                content_index,
344                delta,
345            } => {
346                if let Some(AssistantContent::Thinking(block)) = self.slot(content_index) {
347                    block.thinking.push_str(&delta);
348                }
349                AssistantMessageEvent::ThinkingDelta {
350                    content_index,
351                    delta,
352                    partial: self.partial.clone(),
353                }
354            }
355            ProviderEvent::ThinkingEnd {
356                content_index,
357                content,
358            } => {
359                if let Some(AssistantContent::Thinking(block)) = self.slot(content_index) {
360                    block.thinking.clone_from(&content);
361                }
362                AssistantMessageEvent::ThinkingEnd {
363                    content_index,
364                    content,
365                    partial: self.partial.clone(),
366                }
367            }
368            ProviderEvent::ToolcallStart {
369                content_index,
370                id,
371                tool_name,
372            } => {
373                self.put(
374                    content_index,
375                    AssistantContent::ToolCall(ToolCall {
376                        id: id.unwrap_or_default(),
377                        name: tool_name.unwrap_or_default(),
378                        ..ToolCall::default()
379                    }),
380                );
381                self.tool_args.insert(content_index, String::new());
382                AssistantMessageEvent::ToolcallStart {
383                    content_index,
384                    partial: self.partial.clone(),
385                }
386            }
387            ProviderEvent::ToolcallDelta {
388                content_index,
389                delta,
390            } => {
391                let buffer = self.tool_args.entry(content_index).or_default();
392                buffer.push_str(&delta);
393                let parsed = serde_json::from_str::<JsonObject>(buffer).ok();
394                if let (Some(arguments), Some(AssistantContent::ToolCall(call))) =
395                    (parsed, self.slot(content_index))
396                {
397                    call.arguments = arguments;
398                }
399                AssistantMessageEvent::ToolcallDelta {
400                    content_index,
401                    delta,
402                    partial: self.partial.clone(),
403                }
404            }
405            ProviderEvent::ToolcallEnd {
406                content_index,
407                tool_call,
408            } => {
409                self.tool_args.shift_remove(&content_index);
410                self.put(content_index, AssistantContent::ToolCall(tool_call.clone()));
411                AssistantMessageEvent::ToolcallEnd {
412                    content_index,
413                    tool_call,
414                    partial: self.partial.clone(),
415                }
416            }
417            ProviderEvent::Done { reason, message } => {
418                self.partial = message.clone();
419                AssistantMessageEvent::Done { reason, message }
420            }
421            ProviderEvent::Error { reason, error } => {
422                self.partial = error.clone();
423                AssistantMessageEvent::Error { reason, error }
424            }
425        }
426    }
427
428    fn put(&mut self, index: usize, block: AssistantContent) {
429        let content = &mut self.partial.content;
430        if index < content.len() {
431            content[index] = block;
432        } else {
433            while content.len() < index {
434                content.push(AssistantContent::Text(TextContent::default()));
435            }
436            content.push(block);
437        }
438    }
439
440    fn slot(&mut self, index: usize) -> Option<&mut AssistantContent> {
441        self.partial.content.get_mut(index)
442    }
443}
444
445/// Normalised transcript sent to providers: system messages carry prompt and tools.
446#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
447pub struct Context {
448    /// Messages in order.
449    pub messages: Vec<Message>,
450}
451
452/// Provider transport preference.
453#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
454#[serde(rename_all = "kebab-case")]
455pub enum Transport {
456    /// Server-sent events.
457    Sse,
458    /// WebSocket.
459    Websocket,
460    /// Cached WebSocket connection.
461    WebsocketCached,
462    /// Provider default.
463    Auto,
464}
465
466/// Token budgets per thinking level.
467#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
468pub struct ThinkingBudgets {
469    /// `minimal` budget.
470    #[serde(default, skip_serializing_if = "Option::is_none")]
471    pub minimal: Option<u64>,
472    /// `low` budget.
473    #[serde(default, skip_serializing_if = "Option::is_none")]
474    pub low: Option<u64>,
475    /// `medium` budget.
476    #[serde(default, skip_serializing_if = "Option::is_none")]
477    pub medium: Option<u64>,
478    /// `high` budget.
479    #[serde(default, skip_serializing_if = "Option::is_none")]
480    pub high: Option<u64>,
481}
482
483/// Serializable subset of the stream options.
484#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
485#[serde(rename_all = "camelCase")]
486pub struct StreamOptions {
487    /// Requested thinking level.
488    #[serde(default, skip_serializing_if = "Option::is_none")]
489    pub reasoning: Option<ThinkingLevel>,
490    /// Output token cap.
491    #[serde(default, skip_serializing_if = "Option::is_none")]
492    pub max_tokens: Option<u64>,
493    /// Sampling temperature.
494    #[serde(default, skip_serializing_if = "Option::is_none", with = "js_f64_opt")]
495    pub temperature: Option<f64>,
496    /// Extra sampling parameters.
497    #[serde(default, skip_serializing_if = "Option::is_none")]
498    pub sampling_params: Option<JsonObject>,
499    /// Prompt cache retention.
500    #[serde(default, skip_serializing_if = "Option::is_none")]
501    pub cache_retention: Option<CacheRetention>,
502    /// Session id (cache affinity).
503    #[serde(default, skip_serializing_if = "Option::is_none")]
504    pub session_id: Option<String>,
505    /// Transport preference.
506    #[serde(default, skip_serializing_if = "Option::is_none")]
507    pub transport: Option<Transport>,
508    /// Thinking budgets.
509    #[serde(default, skip_serializing_if = "Option::is_none")]
510    pub thinking_budgets: Option<ThinkingBudgets>,
511    /// Tool choice (`"auto"`, `"none"`, provider objects).
512    #[serde(default, skip_serializing_if = "Option::is_none")]
513    pub tool_choice: Option<JsonValue>,
514    /// Request timeout.
515    #[serde(default, skip_serializing_if = "Option::is_none")]
516    pub timeout_ms: Option<u64>,
517    /// Provider-level retries.
518    #[serde(default, skip_serializing_if = "Option::is_none")]
519    pub max_retries: Option<u32>,
520    /// Retry delay cap.
521    #[serde(default, skip_serializing_if = "Option::is_none")]
522    pub max_retry_delay_ms: Option<u64>,
523    /// Request metadata.
524    #[serde(default, skip_serializing_if = "Option::is_none")]
525    pub metadata: Option<JsonObject>,
526    /// Header overrides (`null` removes a header).
527    #[serde(default, skip_serializing_if = "Option::is_none")]
528    pub headers: Option<IndexMap<String, Option<String>>>,
529}
530
531#[cfg(test)]
532mod tests {
533    use super::*;
534    use serde_json::json;
535
536    fn template() -> AssistantMessage {
537        AssistantMessage::empty("faux", "faux", "faux-1", 0)
538    }
539
540    #[test]
541    fn assembler_rebuilds_partials() {
542        let mut assembler = PartialAssembler::new(template());
543        let events: Vec<ProviderEvent> = serde_json::from_value(json!([
544            {"type":"start"},
545            {"type":"thinking_start","contentIndex":0},
546            {"type":"thinking_delta","contentIndex":0,"delta":"pl"},
547            {"type":"thinking_end","contentIndex":0,"content":"plan"},
548            {"type":"text_start","contentIndex":1},
549            {"type":"text_delta","contentIndex":1,"delta":"Hel"},
550            {"type":"text_delta","contentIndex":1,"delta":"lo"},
551            {"type":"toolcall_start","contentIndex":2,"id":"call_1","toolName":"read"},
552            {"type":"toolcall_delta","contentIndex":2,"delta":"{\"path\":"},
553            {"type":"toolcall_delta","contentIndex":2,"delta":"\"a\"}"},
554        ]))
555        .unwrap_or_default();
556        assert_eq!(events.len(), 10);
557        let mut last = None;
558        for event in events {
559            last = Some(assembler.apply(event));
560        }
561        let partial = assembler.partial();
562        assert_eq!(partial.content.len(), 3);
563        assert_eq!(
564            serde_json::to_value(&partial.content).ok(),
565            Some(
566                json!([{"type":"thinking","thinking":"plan"},{"type":"text","text":"Hello"},
567                {"type":"toolCall","id":"call_1","name":"read","arguments":{"path":"a"}}])
568            )
569        );
570        let wire = last
571            .and_then(|e| serde_json::to_value(e).ok())
572            .unwrap_or_default();
573        assert_eq!(wire["type"], json!("toolcall_delta"));
574        assert_eq!(wire["contentIndex"], json!(2));
575        assert_eq!(wire["partial"]["stopReason"], json!("pending"));
576    }
577
578    #[test]
579    fn stream_options_wire_names() {
580        let options = StreamOptions {
581            reasoning: Some(ThinkingLevel::High),
582            transport: Some(Transport::WebsocketCached),
583            temperature: Some(1.0),
584            ..StreamOptions::default()
585        };
586        assert_eq!(
587            serde_json::to_value(options).ok(),
588            Some(json!({"reasoning":"high","temperature":1,"transport":"websocket-cached"}))
589        );
590    }
591}
oino_ai/stream.rs

oino_ai/
stream.rs
