1use serde::de::Error as DeError;
2use serde::{Deserialize, Deserializer, Serialize};
3use serde_json::Value;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(rename_all = "lowercase")]
8pub enum StreamState {
9 Streaming,
11 Done,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(rename_all = "kebab-case")]
18pub enum ToolState {
19 InputStreaming,
21 InputAvailable,
23 ApprovalRequested,
25 ApprovalResponded,
27 OutputAvailable,
29 OutputError,
31 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
73pub struct ToolApproval {
74 pub id: String,
76 #[serde(skip_serializing_if = "Option::is_none")]
78 pub approved: Option<bool>,
79 #[serde(skip_serializing_if = "Option::is_none")]
81 pub reason: Option<String>,
82 #[serde(skip_serializing_if = "Option::is_none")]
84 pub remember: Option<bool>,
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
89pub struct TextUIPart {
90 #[serde(rename = "type")]
91 part_type: TextPartType,
92 pub text: String,
94 #[serde(skip_serializing_if = "Option::is_none")]
96 pub state: Option<StreamState>,
97 #[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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
115pub struct ReasoningUIPart {
116 #[serde(rename = "type")]
117 part_type: ReasoningPartType,
118 pub text: String,
120 #[serde(skip_serializing_if = "Option::is_none")]
122 pub state: Option<StreamState>,
123 #[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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
141pub struct ToolUIPart {
142 #[serde(rename = "type", deserialize_with = "deserialize_tool_part_type")]
144 pub part_type: String,
145 #[serde(rename = "toolCallId")]
147 pub tool_call_id: String,
148 #[serde(rename = "toolName", skip_serializing_if = "Option::is_none")]
150 pub tool_name: Option<String>,
151 #[serde(skip_serializing_if = "Option::is_none")]
153 pub title: Option<String>,
154 #[serde(rename = "providerExecuted", skip_serializing_if = "Option::is_none")]
156 pub provider_executed: Option<bool>,
157 pub state: ToolState,
159 #[serde(skip_serializing_if = "Option::is_none")]
161 pub input: Option<Value>,
162 #[serde(skip_serializing_if = "Option::is_none")]
164 pub output: Option<Value>,
165 #[serde(rename = "errorText", skip_serializing_if = "Option::is_none")]
167 pub error_text: Option<String>,
168 #[serde(
170 rename = "callProviderMetadata",
171 skip_serializing_if = "Option::is_none"
172 )]
173 pub call_provider_metadata: Option<Value>,
174 #[serde(rename = "rawInput", skip_serializing_if = "Option::is_none")]
176 pub raw_input: Option<Value>,
177 #[serde(skip_serializing_if = "Option::is_none")]
179 pub preliminary: Option<bool>,
180 #[serde(skip_serializing_if = "Option::is_none")]
182 pub approval: Option<ToolApproval>,
183}
184
185impl ToolUIPart {
186 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 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
236pub struct SourceUrlUIPart {
237 #[serde(rename = "type")]
238 part_type: SourceUrlPartType,
239 #[serde(rename = "sourceId")]
241 pub source_id: String,
242 pub url: String,
244 #[serde(skip_serializing_if = "Option::is_none")]
246 pub title: Option<String>,
247 #[serde(rename = "providerMetadata", skip_serializing_if = "Option::is_none")]
249 pub provider_metadata: Option<Value>,
250}
251
252#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
254pub struct SourceDocumentUIPart {
255 #[serde(rename = "type")]
256 part_type: SourceDocumentPartType,
257 #[serde(rename = "sourceId")]
259 pub source_id: String,
260 #[serde(rename = "mediaType")]
262 pub media_type: String,
263 pub title: String,
265 #[serde(skip_serializing_if = "Option::is_none")]
267 pub filename: Option<String>,
268 #[serde(rename = "providerMetadata", skip_serializing_if = "Option::is_none")]
270 pub provider_metadata: Option<Value>,
271}
272
273#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
275pub struct FileUIPart {
276 #[serde(rename = "type")]
277 part_type: FilePartType,
278 pub url: String,
280 #[serde(rename = "mediaType")]
282 pub media_type: String,
283 #[serde(skip_serializing_if = "Option::is_none")]
285 pub filename: Option<String>,
286 #[serde(rename = "providerMetadata", skip_serializing_if = "Option::is_none")]
288 pub provider_metadata: Option<Value>,
289}
290
291#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
293pub struct StepStartUIPart {
294 #[serde(rename = "type")]
295 part_type: StepStartPartType,
296}
297
298#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
300pub struct DataUIPart {
301 #[serde(rename = "type", deserialize_with = "deserialize_data_part_type")]
303 pub data_type: String,
304 #[serde(skip_serializing_if = "Option::is_none")]
306 pub id: Option<String>,
307 pub data: Value,
309}
310
311#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
313#[serde(untagged)]
314pub enum UIMessagePart {
315 Text(TextUIPart),
317 Reasoning(ReasoningUIPart),
319 Tool(ToolUIPart),
321 SourceUrl(SourceUrlUIPart),
323 SourceDocument(SourceDocumentUIPart),
325 File(FileUIPart),
327 StepStart(StepStartUIPart),
329 Data(DataUIPart),
331}
332
333#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
335#[serde(rename_all = "lowercase")]
336pub enum UIRole {
337 System,
339 User,
341 Assistant,
343}
344
345#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
350pub struct UIMessage {
351 pub id: String,
353 pub role: UIRole,
355 #[serde(skip_serializing_if = "Option::is_none")]
357 pub metadata: Option<Value>,
358 pub parts: Vec<UIMessagePart>,
360}
361
362impl UIMessage {
363 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 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 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 #[must_use]
398 pub fn with_metadata(mut self, metadata: Value) -> Self {
399 self.metadata = Some(metadata);
400 self
401 }
402
403 #[must_use]
405 pub fn with_part(mut self, part: UIMessagePart) -> Self {
406 self.parts.push(part);
407 self
408 }
409
410 #[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 #[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 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}