tower_a2a/protocol/
message.rs

1//! A2A message types
2
3use std::collections::HashMap;
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7
8/// A message in the A2A protocol
9///
10/// Messages are the primary unit of communication between agents.
11/// Each message has a role (user or assistant), one or more parts (text, file, or data),
12/// and optional metadata and extensions.
13#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
14pub struct Message {
15    /// Role of the message sender
16    pub role: Role,
17
18    /// Message content parts (at least one required)
19    pub parts: Vec<MessagePart>,
20
21    /// Optional message identifier
22    #[serde(rename = "messageId", skip_serializing_if = "Option::is_none")]
23    pub message_id: Option<String>,
24
25    /// Optional task identifier (for associating message with a task)
26    #[serde(rename = "taskId", skip_serializing_if = "Option::is_none")]
27    pub task_id: Option<String>,
28
29    /// Optional context identifier (for multi-turn conversations)
30    #[serde(rename = "contextId", skip_serializing_if = "Option::is_none")]
31    pub context_id: Option<String>,
32
33    /// Optional metadata for the message
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub metadata: Option<HashMap<String, Value>>,
36
37    /// Optional extensions indicating additional protocol features
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub extensions: Option<HashMap<String, Value>>,
40}
41
42impl Message {
43    /// Create a new message with text content
44    pub fn new(role: Role, text: impl Into<String>) -> Self {
45        Self {
46            role,
47            parts: vec![MessagePart::Text { text: text.into() }],
48            message_id: None,
49            task_id: None,
50            context_id: None,
51            metadata: None,
52            extensions: None,
53        }
54    }
55
56    /// Create a user message with text content
57    pub fn user(text: impl Into<String>) -> Self {
58        Self::new(Role::User, text)
59    }
60
61    /// Create an assistant message with text content
62    pub fn assistant(text: impl Into<String>) -> Self {
63        Self::new(Role::Assistant, text)
64    }
65
66    /// Create a new message builder
67    pub fn builder() -> MessageBuilder {
68        MessageBuilder::new()
69    }
70
71    /// Add a metadata field to the message
72    pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
73        self.metadata
74            .get_or_insert_with(HashMap::new)
75            .insert(key.into(), value);
76        self
77    }
78
79    /// Add an extension to the message
80    pub fn with_extension(mut self, key: impl Into<String>, value: Value) -> Self {
81        self.extensions
82            .get_or_insert_with(HashMap::new)
83            .insert(key.into(), value);
84        self
85    }
86
87    /// Add a message part
88    pub fn with_part(mut self, part: MessagePart) -> Self {
89        self.parts.push(part);
90        self
91    }
92}
93
94/// Builder for constructing Message instances
95#[derive(Debug, Default)]
96pub struct MessageBuilder {
97    role: Option<Role>,
98    parts: Vec<MessagePart>,
99    message_id: Option<String>,
100    task_id: Option<String>,
101    context_id: Option<String>,
102    metadata: Option<HashMap<String, Value>>,
103    extensions: Option<HashMap<String, Value>>,
104}
105
106impl MessageBuilder {
107    /// Create a new message builder
108    pub fn new() -> Self {
109        Self::default()
110    }
111
112    /// Set the role of the message
113    pub fn role(mut self, role: Role) -> Self {
114        self.role = Some(role);
115        self
116    }
117
118    /// Set the message parts
119    pub fn parts(mut self, parts: Vec<MessagePart>) -> Self {
120        self.parts = parts;
121        self
122    }
123
124    /// Add a single part to the message
125    pub fn part(mut self, part: MessagePart) -> Self {
126        self.parts.push(part);
127        self
128    }
129
130    /// Set the message ID
131    pub fn message_id(mut self, id: impl Into<String>) -> Self {
132        self.message_id = Some(id.into());
133        self
134    }
135
136    /// Set the task ID
137    pub fn task_id(mut self, id: impl Into<String>) -> Self {
138        self.task_id = Some(id.into());
139        self
140    }
141
142    /// Set the context ID
143    pub fn context_id(mut self, id: impl Into<String>) -> Self {
144        self.context_id = Some(id.into());
145        self
146    }
147
148    /// Add a metadata field
149    pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
150        self.metadata
151            .get_or_insert_with(HashMap::new)
152            .insert(key.into(), value);
153        self
154    }
155
156    /// Add an extension
157    pub fn extension(mut self, key: impl Into<String>, value: Value) -> Self {
158        self.extensions
159            .get_or_insert_with(HashMap::new)
160            .insert(key.into(), value);
161        self
162    }
163
164    /// Build the message
165    ///
166    /// # Panics
167    ///
168    /// Panics if role is not set or if parts are empty
169    pub fn build(self) -> Message {
170        let role = self.role.expect("Message role is required");
171        assert!(!self.parts.is_empty(), "Message must have at least one part");
172
173        Message {
174            role,
175            parts: self.parts,
176            message_id: self.message_id,
177            task_id: self.task_id,
178            context_id: self.context_id,
179            metadata: self.metadata,
180            extensions: self.extensions,
181        }
182    }
183}
184
185/// Role of a message sender
186#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
187#[serde(rename_all = "lowercase")]
188pub enum Role {
189    /// Message from a user
190    User,
191
192    /// Message from an AI assistant/agent
193    Assistant,
194
195    /// System message (optional, for context)
196    System,
197}
198
199/// A part of a message
200///
201/// According to the A2A spec: "A Part MUST contain exactly one of the following: text, file, data"
202#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
203#[serde(untagged)]
204pub enum MessagePart {
205    /// Text content
206    Text {
207        /// The text content
208        text: String,
209    },
210
211    /// File reference
212    File {
213        /// URI of the file
214        #[serde(rename = "fileUri")]
215        file_uri: String,
216
217        /// MIME type of the file
218        #[serde(rename = "mimeType", skip_serializing_if = "Option::is_none")]
219        mime_type: Option<String>,
220    },
221
222    /// Structured data
223    Data {
224        /// The structured data
225        data: Value,
226
227        /// MIME type of the data
228        #[serde(rename = "mimeType", skip_serializing_if = "Option::is_none")]
229        mime_type: Option<String>,
230    },
231}
232
233impl MessagePart {
234    /// Create a text part
235    pub fn text(text: impl Into<String>) -> Self {
236        Self::Text { text: text.into() }
237    }
238
239    /// Create a file part
240    pub fn file(file_uri: impl Into<String>) -> Self {
241        Self::File {
242            file_uri: file_uri.into(),
243            mime_type: None,
244        }
245    }
246
247    /// Create a file part with MIME type
248    pub fn file_with_type(file_uri: impl Into<String>, mime_type: impl Into<String>) -> Self {
249        Self::File {
250            file_uri: file_uri.into(),
251            mime_type: Some(mime_type.into()),
252        }
253    }
254
255    /// Create a data part
256    pub fn data(data: Value) -> Self {
257        Self::Data {
258            data,
259            mime_type: None,
260        }
261    }
262
263    /// Create a data part with MIME type
264    pub fn data_with_type(data: Value, mime_type: impl Into<String>) -> Self {
265        Self::Data {
266            data,
267            mime_type: Some(mime_type.into()),
268        }
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use serde_json::json;
275
276    use super::*;
277
278    #[test]
279    fn test_message_creation() {
280        let msg = Message::user("Hello, agent!");
281        assert_eq!(msg.role, Role::User);
282        assert_eq!(msg.parts.len(), 1);
283
284        match &msg.parts[0] {
285            MessagePart::Text { text } => assert_eq!(text, "Hello, agent!"),
286            _ => panic!("Expected text part"),
287        }
288    }
289
290    #[test]
291    fn test_message_with_metadata() {
292        let msg = Message::user("Test")
293            .with_metadata("key", json!("value"))
294            .with_extension("ext", json!({"enabled": true}));
295
296        assert!(msg.metadata.is_some());
297        assert!(msg.extensions.is_some());
298    }
299
300    #[test]
301    fn test_message_serialization() {
302        let msg = Message::user("Test message");
303        let json = serde_json::to_string(&msg).unwrap();
304        assert!(json.contains("\"role\":\"user\""));
305        assert!(json.contains("\"text\":\"Test message\""));
306
307        let deserialized: Message = serde_json::from_str(&json).unwrap();
308        assert_eq!(msg, deserialized);
309    }
310
311    #[test]
312    fn test_message_part_types() {
313        let text = MessagePart::text("Hello");
314        let file = MessagePart::file("file://path/to/file");
315        let data = MessagePart::data(json!({"key": "value"}));
316
317        assert!(matches!(text, MessagePart::Text { .. }));
318        assert!(matches!(file, MessagePart::File { .. }));
319        assert!(matches!(data, MessagePart::Data { .. }));
320    }
321
322    #[test]
323    fn test_message_builder() {
324        let msg = Message::builder()
325            .role(Role::User)
326            .parts(vec![MessagePart::text("Hello")])
327            .message_id("msg-123")
328            .task_id("task-456")
329            .context_id("ctx-789")
330            .build();
331
332        assert_eq!(msg.role, Role::User);
333        assert_eq!(msg.parts.len(), 1);
334        assert_eq!(msg.message_id, Some("msg-123".to_string()));
335        assert_eq!(msg.task_id, Some("task-456".to_string()));
336        assert_eq!(msg.context_id, Some("ctx-789".to_string()));
337    }
338
339    #[test]
340    fn test_message_builder_with_part() {
341        let msg = Message::builder()
342            .role(Role::Assistant)
343            .part(MessagePart::text("First"))
344            .part(MessagePart::text("Second"))
345            .build();
346
347        assert_eq!(msg.parts.len(), 2);
348    }
349
350    #[test]
351    #[should_panic(expected = "Message role is required")]
352    fn test_message_builder_missing_role() {
353        Message::builder()
354            .parts(vec![MessagePart::text("Hello")])
355            .build();
356    }
357
358    #[test]
359    #[should_panic(expected = "Message must have at least one part")]
360    fn test_message_builder_no_parts() {
361        Message::builder().role(Role::User).build();
362    }
363
364    #[test]
365    fn test_message_serialization_with_ids() {
366        let msg = Message::builder()
367            .role(Role::User)
368            .parts(vec![MessagePart::text("Test")])
369            .message_id("msg-123")
370            .task_id("task-456")
371            .build();
372
373        let json = serde_json::to_string(&msg).unwrap();
374        assert!(json.contains("\"messageId\":\"msg-123\""));
375        assert!(json.contains("\"taskId\":\"task-456\""));
376
377        let deserialized: Message = serde_json::from_str(&json).unwrap();
378        assert_eq!(msg, deserialized);
379    }
380}