Skip to main content

starweaver_model/
stream.rs

1//! Canonical model stream events.
2
3use serde::{Deserialize, Serialize};
4use serde_json::{Map, Value};
5
6use crate::message::ModelResponse;
7
8/// Stream event emitted by model adapters.
9#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
10#[serde(tag = "kind", rename_all = "snake_case")]
11pub enum ModelResponseStreamEvent {
12    /// A response part started.
13    PartStart(PartStart),
14    /// A response part delta arrived.
15    PartDelta(PartDelta),
16    /// A response part ended.
17    PartEnd(PartEnd),
18    /// Provider or transport diagnostic sideband event.
19    Diagnostic(StreamDiagnostic),
20    /// Final response is available.
21    FinalResult(Box<ModelResponse>),
22}
23
24/// Diagnostic sideband event emitted during model streaming.
25#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
26pub struct StreamDiagnostic {
27    /// Stable diagnostic event kind.
28    pub kind: String,
29    /// Structured diagnostic payload. Must not contain secrets or raw request bodies.
30    #[serde(default)]
31    pub payload: Value,
32    /// Diagnostic metadata.
33    #[serde(default, skip_serializing_if = "Map::is_empty")]
34    pub metadata: Map<String, Value>,
35}
36
37impl StreamDiagnostic {
38    /// Build a stream diagnostic event.
39    #[must_use]
40    pub fn new(kind: impl Into<String>, payload: Value) -> Self {
41        Self {
42            kind: kind.into(),
43            payload,
44            metadata: Map::new(),
45        }
46    }
47}
48
49/// Part start event.
50#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
51pub struct PartStart {
52    /// Part index in response.
53    pub index: usize,
54    /// Part kind.
55    pub part_kind: String,
56}
57
58/// Part delta event.
59#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
60pub struct PartDelta {
61    /// Part index in response.
62    pub index: usize,
63    /// Typed delta payload.
64    #[serde(flatten)]
65    pub delta: StreamDelta,
66}
67
68impl PartDelta {
69    /// Build a text delta.
70    #[must_use]
71    pub fn text(index: usize, text: impl Into<String>) -> Self {
72        Self {
73            index,
74            delta: StreamDelta::Text { text: text.into() },
75        }
76    }
77
78    /// Build a thinking delta.
79    #[must_use]
80    pub fn thinking(index: usize, text: impl Into<String>) -> Self {
81        Self {
82            index,
83            delta: StreamDelta::Thinking { text: text.into() },
84        }
85    }
86
87    /// Return a text-only display projection.
88    #[must_use]
89    pub fn as_text(&self) -> String {
90        match &self.delta {
91            StreamDelta::Text { text }
92            | StreamDelta::Thinking { text }
93            | StreamDelta::ToolCallName { name: text }
94            | StreamDelta::ToolCallArguments {
95                arguments_delta: text,
96            } => text.clone(),
97            StreamDelta::NativePayload { payload } | StreamDelta::FileMetadata { payload } => {
98                payload.to_string()
99            }
100        }
101    }
102}
103
104/// Typed model stream delta payload.
105#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
106#[serde(tag = "delta_kind", rename_all = "snake_case")]
107pub enum StreamDelta {
108    /// Text output delta.
109    Text {
110        /// Text fragment.
111        text: String,
112    },
113    /// Thinking output delta.
114    Thinking {
115        /// Thinking fragment.
116        text: String,
117    },
118    /// Tool call name delta.
119    ToolCallName {
120        /// Tool name fragment or final value.
121        name: String,
122    },
123    /// Tool call argument delta.
124    ToolCallArguments {
125        /// Argument fragment.
126        arguments_delta: String,
127    },
128    /// Provider-native payload delta.
129    NativePayload {
130        /// Native payload fragment.
131        payload: Value,
132    },
133    /// File metadata delta.
134    FileMetadata {
135        /// File metadata fragment.
136        payload: Value,
137    },
138}
139
140/// Part end event.
141#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
142pub struct PartEnd {
143    /// Part index in response.
144    pub index: usize,
145    /// Part kind when the provider includes it or the parser can infer it.
146    #[serde(default, skip_serializing_if = "Option::is_none")]
147    pub part_kind: Option<String>,
148}
149
150impl PartEnd {
151    /// Build a part end event without a known part kind.
152    #[must_use]
153    pub const fn new(index: usize) -> Self {
154        Self {
155            index,
156            part_kind: None,
157        }
158    }
159
160    /// Build a part end event with an explicit part kind.
161    #[must_use]
162    pub fn with_kind(index: usize, part_kind: impl Into<String>) -> Self {
163        Self {
164            index,
165            part_kind: Some(part_kind.into()),
166        }
167    }
168}
169
170/// Lifecycle state of a streamed model response.
171#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
172#[serde(rename_all = "snake_case")]
173pub enum StreamLifecycle {
174    /// Events are still arriving.
175    #[default]
176    Incomplete,
177    /// Final response has been assembled.
178    Complete,
179    /// Stream ended by explicit cancellation or transport interruption.
180    Interrupted,
181}
182
183/// Lightweight stream assembly state for replay and lifecycle assertions.
184#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
185pub struct ModelStreamState {
186    /// Current lifecycle.
187    pub lifecycle: StreamLifecycle,
188    /// Number of started parts.
189    pub started_parts: usize,
190    /// Number of ended parts.
191    pub ended_parts: usize,
192    /// Final response when available.
193    #[serde(default, skip_serializing_if = "Option::is_none")]
194    pub final_response: Option<Box<ModelResponse>>,
195}
196
197impl ModelStreamState {
198    /// Apply one stream event.
199    pub fn apply(&mut self, event: &ModelResponseStreamEvent) {
200        match event {
201            ModelResponseStreamEvent::PartStart(_) => {
202                self.started_parts += 1;
203            }
204            ModelResponseStreamEvent::PartDelta(_) | ModelResponseStreamEvent::Diagnostic(_) => {}
205            ModelResponseStreamEvent::PartEnd(_) => {
206                self.ended_parts += 1;
207            }
208            ModelResponseStreamEvent::FinalResult(response) => {
209                self.lifecycle = StreamLifecycle::Complete;
210                self.final_response = Some(response.clone());
211            }
212        }
213    }
214
215    /// Mark stream as interrupted.
216    pub const fn interrupt(&mut self) {
217        self.lifecycle = StreamLifecycle::Interrupted;
218    }
219}