tirea_protocol_ai_sdk_v6/
message.rs

1use serde::de::Error as DeError;
2use serde::{Deserialize, Deserializer, Serialize};
3use serde_json::Value;
4
5/// Streaming state for UI parts.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(rename_all = "lowercase")]
8pub enum StreamState {
9    /// Content is still streaming.
10    Streaming,
11    /// Content streaming is complete.
12    Done,
13}
14
15/// Tool execution state in UI.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(rename_all = "kebab-case")]
18pub enum ToolState {
19    /// Tool input is being streamed.
20    InputStreaming,
21    /// Tool input is complete, ready for execution.
22    InputAvailable,
23    /// User approval has been requested for this tool call.
24    ApprovalRequested,
25    /// User has responded to the approval request.
26    ApprovalResponded,
27    /// Tool execution completed with output.
28    OutputAvailable,
29    /// Tool execution resulted in error.
30    OutputError,
31    /// Tool execution was denied.
32    OutputDenied,
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
36enum TextPartType {
37    #[serde(rename = "text")]
38    Text,
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
42enum ReasoningPartType {
43    #[serde(rename = "reasoning")]
44    Reasoning,
45}
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
48enum SourceUrlPartType {
49    #[serde(rename = "source-url")]
50    SourceUrl,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
54enum SourceDocumentPartType {
55    #[serde(rename = "source-document")]
56    SourceDocument,
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
60enum FilePartType {
61    #[serde(rename = "file")]
62    File,
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
66enum StepStartPartType {
67    #[serde(rename = "step-start")]
68    StepStart,
69}
70
71/// User approval details attached to a tool invocation part.
72#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
73pub struct ToolApproval {
74    /// Approval request ID.
75    pub id: String,
76    /// Optional final decision.
77    #[serde(skip_serializing_if = "Option::is_none")]
78    pub approved: Option<bool>,
79    /// Optional reason from the approver.
80    #[serde(skip_serializing_if = "Option::is_none")]
81    pub reason: Option<String>,
82    /// Optional request to remember this decision.
83    #[serde(skip_serializing_if = "Option::is_none")]
84    pub remember: Option<bool>,
85}
86
87/// Text part.
88#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
89pub struct TextUIPart {
90    #[serde(rename = "type")]
91    part_type: TextPartType,
92    /// The text content.
93    pub text: String,
94    /// Optional streaming state.
95    #[serde(skip_serializing_if = "Option::is_none")]
96    pub state: Option<StreamState>,
97    /// Optional provider metadata.
98    #[serde(rename = "providerMetadata", skip_serializing_if = "Option::is_none")]
99    pub provider_metadata: Option<Value>,
100}
101
102impl TextUIPart {
103    pub fn new(text: impl Into<String>, state: Option<StreamState>) -> Self {
104        Self {
105            part_type: TextPartType::Text,
106            text: text.into(),
107            state,
108            provider_metadata: None,
109        }
110    }
111}
112
113/// Reasoning part.
114#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
115pub struct ReasoningUIPart {
116    #[serde(rename = "type")]
117    part_type: ReasoningPartType,
118    /// The reasoning text.
119    pub text: String,
120    /// Optional streaming state.
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub state: Option<StreamState>,
123    /// Optional provider metadata.
124    #[serde(rename = "providerMetadata", skip_serializing_if = "Option::is_none")]
125    pub provider_metadata: Option<Value>,
126}
127
128impl ReasoningUIPart {
129    pub fn new(text: impl Into<String>, state: Option<StreamState>) -> Self {
130        Self {
131            part_type: ReasoningPartType::Reasoning,
132            text: text.into(),
133            state,
134            provider_metadata: None,
135        }
136    }
137}
138
139/// Tool invocation part.
140#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
141pub struct ToolUIPart {
142    /// Part type. Must be `dynamic-tool` or `tool-*`.
143    #[serde(rename = "type", deserialize_with = "deserialize_tool_part_type")]
144    pub part_type: String,
145    /// Tool call identifier.
146    #[serde(rename = "toolCallId")]
147    pub tool_call_id: String,
148    /// Tool name (required for `dynamic-tool`).
149    #[serde(rename = "toolName", skip_serializing_if = "Option::is_none")]
150    pub tool_name: Option<String>,
151    /// Optional display title.
152    #[serde(skip_serializing_if = "Option::is_none")]
153    pub title: Option<String>,
154    /// Whether the provider executed this tool directly.
155    #[serde(rename = "providerExecuted", skip_serializing_if = "Option::is_none")]
156    pub provider_executed: Option<bool>,
157    /// Tool execution state.
158    pub state: ToolState,
159    /// Tool input payload.
160    #[serde(skip_serializing_if = "Option::is_none")]
161    pub input: Option<Value>,
162    /// Tool output payload.
163    #[serde(skip_serializing_if = "Option::is_none")]
164    pub output: Option<Value>,
165    /// Tool error text.
166    #[serde(rename = "errorText", skip_serializing_if = "Option::is_none")]
167    pub error_text: Option<String>,
168    /// Provider metadata associated with the tool input call.
169    #[serde(
170        rename = "callProviderMetadata",
171        skip_serializing_if = "Option::is_none"
172    )]
173    pub call_provider_metadata: Option<Value>,
174    /// Raw input when parsing fails (AI SDK compatibility for output-error).
175    #[serde(rename = "rawInput", skip_serializing_if = "Option::is_none")]
176    pub raw_input: Option<Value>,
177    /// Marks provisional tool outputs.
178    #[serde(skip_serializing_if = "Option::is_none")]
179    pub preliminary: Option<bool>,
180    /// Optional approval state payload.
181    #[serde(skip_serializing_if = "Option::is_none")]
182    pub approval: Option<ToolApproval>,
183}
184
185impl ToolUIPart {
186    /// Create a static tool part (`type = tool-{name}`).
187    pub fn static_tool(
188        tool_name: impl Into<String>,
189        tool_call_id: impl Into<String>,
190        state: ToolState,
191    ) -> Self {
192        let tool_name = tool_name.into();
193        Self {
194            part_type: format!("tool-{tool_name}"),
195            tool_call_id: tool_call_id.into(),
196            tool_name: None,
197            title: None,
198            provider_executed: None,
199            state,
200            input: None,
201            output: None,
202            error_text: None,
203            call_provider_metadata: None,
204            raw_input: None,
205            preliminary: None,
206            approval: None,
207        }
208    }
209
210    /// Create a dynamic tool part (`type = dynamic-tool`).
211    pub fn dynamic_tool(
212        tool_name: impl Into<String>,
213        tool_call_id: impl Into<String>,
214        state: ToolState,
215    ) -> Self {
216        Self {
217            part_type: "dynamic-tool".to_string(),
218            tool_call_id: tool_call_id.into(),
219            tool_name: Some(tool_name.into()),
220            title: None,
221            provider_executed: None,
222            state,
223            input: None,
224            output: None,
225            error_text: None,
226            call_provider_metadata: None,
227            raw_input: None,
228            preliminary: None,
229            approval: None,
230        }
231    }
232}
233
234/// URL source reference.
235#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
236pub struct SourceUrlUIPart {
237    #[serde(rename = "type")]
238    part_type: SourceUrlPartType,
239    /// Source identifier.
240    #[serde(rename = "sourceId")]
241    pub source_id: String,
242    /// The URL.
243    pub url: String,
244    /// Optional title.
245    #[serde(skip_serializing_if = "Option::is_none")]
246    pub title: Option<String>,
247    /// Optional provider metadata.
248    #[serde(rename = "providerMetadata", skip_serializing_if = "Option::is_none")]
249    pub provider_metadata: Option<Value>,
250}
251
252/// Document source reference.
253#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
254pub struct SourceDocumentUIPart {
255    #[serde(rename = "type")]
256    part_type: SourceDocumentPartType,
257    /// Source identifier.
258    #[serde(rename = "sourceId")]
259    pub source_id: String,
260    /// IANA media type.
261    #[serde(rename = "mediaType")]
262    pub media_type: String,
263    /// Document title.
264    pub title: String,
265    /// Optional filename.
266    #[serde(skip_serializing_if = "Option::is_none")]
267    pub filename: Option<String>,
268    /// Optional provider metadata.
269    #[serde(rename = "providerMetadata", skip_serializing_if = "Option::is_none")]
270    pub provider_metadata: Option<Value>,
271}
272
273/// File attachment.
274#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
275pub struct FileUIPart {
276    #[serde(rename = "type")]
277    part_type: FilePartType,
278    /// File URL.
279    pub url: String,
280    /// IANA media type.
281    #[serde(rename = "mediaType")]
282    pub media_type: String,
283    /// Optional filename.
284    #[serde(skip_serializing_if = "Option::is_none")]
285    pub filename: Option<String>,
286    /// Optional provider metadata.
287    #[serde(rename = "providerMetadata", skip_serializing_if = "Option::is_none")]
288    pub provider_metadata: Option<Value>,
289}
290
291/// Step start marker.
292#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
293pub struct StepStartUIPart {
294    #[serde(rename = "type")]
295    part_type: StepStartPartType,
296}
297
298/// Custom data part.
299#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
300pub struct DataUIPart {
301    /// Custom type, must start with `data-`.
302    #[serde(rename = "type", deserialize_with = "deserialize_data_part_type")]
303    pub data_type: String,
304    /// Optional stable data part ID.
305    #[serde(skip_serializing_if = "Option::is_none")]
306    pub id: Option<String>,
307    /// Data payload.
308    pub data: Value,
309}
310
311/// A part of a UI message.
312#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
313#[serde(untagged)]
314pub enum UIMessagePart {
315    /// Text content part.
316    Text(TextUIPart),
317    /// Reasoning content part.
318    Reasoning(ReasoningUIPart),
319    /// Tool invocation part.
320    Tool(ToolUIPart),
321    /// URL source reference.
322    SourceUrl(SourceUrlUIPart),
323    /// Document source reference.
324    SourceDocument(SourceDocumentUIPart),
325    /// File attachment.
326    File(FileUIPart),
327    /// Step start marker.
328    StepStart(StepStartUIPart),
329    /// Custom data part.
330    Data(DataUIPart),
331}
332
333/// Role in the conversation.
334#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
335#[serde(rename_all = "lowercase")]
336pub enum UIRole {
337    /// System message.
338    System,
339    /// User message.
340    User,
341    /// Assistant message.
342    Assistant,
343}
344
345/// A UI message with rich parts.
346///
347/// This is the source of truth for application state, representing the complete
348/// message including metadata, data parts, and all contextual information.
349#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
350pub struct UIMessage {
351    /// Unique identifier for this message.
352    pub id: String,
353    /// Role of the message sender.
354    pub role: UIRole,
355    /// Optional metadata.
356    #[serde(skip_serializing_if = "Option::is_none")]
357    pub metadata: Option<Value>,
358    /// Message parts.
359    pub parts: Vec<UIMessagePart>,
360}
361
362impl UIMessage {
363    /// Create a new UI message.
364    pub fn new(id: impl Into<String>, role: UIRole) -> Self {
365        Self {
366            id: id.into(),
367            role,
368            metadata: None,
369            parts: Vec::new(),
370        }
371    }
372
373    /// Create a user message.
374    pub fn user(id: impl Into<String>, text: impl Into<String>) -> Self {
375        Self {
376            id: id.into(),
377            role: UIRole::User,
378            metadata: None,
379            parts: vec![UIMessagePart::Text(TextUIPart::new(
380                text,
381                Some(StreamState::Done),
382            ))],
383        }
384    }
385
386    /// Create an assistant message.
387    pub fn assistant(id: impl Into<String>) -> Self {
388        Self {
389            id: id.into(),
390            role: UIRole::Assistant,
391            metadata: None,
392            parts: Vec::new(),
393        }
394    }
395
396    /// Add metadata.
397    #[must_use]
398    pub fn with_metadata(mut self, metadata: Value) -> Self {
399        self.metadata = Some(metadata);
400        self
401    }
402
403    /// Add a part.
404    #[must_use]
405    pub fn with_part(mut self, part: UIMessagePart) -> Self {
406        self.parts.push(part);
407        self
408    }
409
410    /// Add multiple parts.
411    #[must_use]
412    pub fn with_parts(mut self, parts: impl IntoIterator<Item = UIMessagePart>) -> Self {
413        self.parts.extend(parts);
414        self
415    }
416
417    /// Add a text part.
418    #[must_use]
419    pub fn with_text(self, text: impl Into<String>) -> Self {
420        self.with_part(UIMessagePart::Text(TextUIPart::new(
421            text,
422            Some(StreamState::Done),
423        )))
424    }
425
426    /// Get all text content concatenated.
427    pub fn text_content(&self) -> String {
428        self.parts
429            .iter()
430            .filter_map(|p| match p {
431                UIMessagePart::Text(part) => Some(part.text.as_str()),
432                _ => None,
433            })
434            .collect::<Vec<_>>()
435            .join("")
436    }
437}
438
439fn deserialize_tool_part_type<'de, D>(deserializer: D) -> Result<String, D::Error>
440where
441    D: Deserializer<'de>,
442{
443    let value = String::deserialize(deserializer)?;
444    if value == "dynamic-tool" || value.starts_with("tool-") {
445        Ok(value)
446    } else {
447        Err(D::Error::custom(format!(
448            "invalid tool part type '{value}', expected 'dynamic-tool' or 'tool-*'"
449        )))
450    }
451}
452
453fn deserialize_data_part_type<'de, D>(deserializer: D) -> Result<String, D::Error>
454where
455    D: Deserializer<'de>,
456{
457    let value = String::deserialize(deserializer)?;
458    if value.starts_with("data-") {
459        Ok(value)
460    } else {
461        Err(D::Error::custom(format!(
462            "invalid data part type '{value}', expected 'data-*'"
463        )))
464    }
465}
466
467#[cfg(test)]
468mod tests {
469    use super::*;
470    use serde_json::json;
471
472    #[test]
473    fn static_tool_part_serializes_with_tool_prefix_type() {
474        let mut part = ToolUIPart::static_tool("search", "call_1", ToolState::InputAvailable);
475        part.input = Some(json!({ "q": "rust" }));
476        let value = serde_json::to_value(UIMessagePart::Tool(part)).expect("serialize tool part");
477
478        assert_eq!(value["type"], "tool-search");
479        assert_eq!(value["toolCallId"], "call_1");
480        assert!(value.get("toolName").is_none());
481    }
482
483    #[test]
484    fn dynamic_tool_part_serializes_with_tool_name() {
485        let part = ToolUIPart::dynamic_tool("search", "call_2", ToolState::InputStreaming);
486        let value = serde_json::to_value(UIMessagePart::Tool(part)).expect("serialize tool part");
487
488        assert_eq!(value["type"], "dynamic-tool");
489        assert_eq!(value["toolCallId"], "call_2");
490        assert_eq!(value["toolName"], "search");
491    }
492
493    #[test]
494    fn tool_part_rejects_invalid_type() {
495        let err = serde_json::from_value::<ToolUIPart>(json!({
496            "type": "tool",
497            "toolCallId": "call_1",
498            "state": "input-available"
499        }))
500        .expect_err("invalid tool type must be rejected");
501
502        assert!(err.to_string().contains("dynamic-tool"));
503    }
504
505    #[test]
506    fn data_part_rejects_invalid_type_prefix() {
507        let err = serde_json::from_value::<DataUIPart>(json!({
508            "type": "reasoning-encrypted",
509            "data": { "v": 1 }
510        }))
511        .expect_err("invalid data type must be rejected");
512
513        assert!(err.to_string().contains("data-*"));
514    }
515
516    #[test]
517    fn file_part_roundtrip_preserves_filename_and_provider_metadata() {
518        let part = UIMessagePart::File(FileUIPart {
519            part_type: FilePartType::File,
520            url: "https://example.com/a.png".to_string(),
521            media_type: "image/png".to_string(),
522            filename: Some("a.png".to_string()),
523            provider_metadata: Some(json!({ "source": "upload" })),
524        });
525
526        let raw = serde_json::to_string(&part).expect("serialize file part");
527        let restored: UIMessagePart = serde_json::from_str(&raw).expect("deserialize file part");
528
529        assert!(matches!(
530            restored,
531            UIMessagePart::File(FileUIPart {
532                filename: Some(filename),
533                provider_metadata: Some(provider_metadata),
534                ..
535            }) if filename == "a.png" && provider_metadata["source"] == "upload"
536        ));
537    }
538}