Skip to main content

vtcode_a2a/
types.rs

1//! A2A Protocol core data types
2//!
3//! Implements the core data structures as defined in the A2A specification:
4//! - Task and TaskStatus for task lifecycle management
5//! - Message and Part types for communication
6//! - Artifact for task outputs
7
8use chrono::{DateTime, Utc};
9use hashbrown::HashMap;
10use serde::{Deserialize, Serialize};
11
12// ============================================================================
13// Task State & Status
14// ============================================================================
15
16/// Task lifecycle states as defined in the A2A specification
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
18#[serde(rename_all = "kebab-case")]
19pub enum TaskState {
20    /// Task has been submitted but not yet started
21    #[default]
22    Submitted,
23    /// Task is actively being processed
24    Working,
25    /// Task requires additional input from the user
26    InputRequired,
27    /// Task completed successfully
28    Completed,
29    /// Task failed with an error
30    Failed,
31    /// Task was canceled by request
32    Canceled,
33    /// Task was rejected by the agent
34    Rejected,
35    /// Task requires authentication
36    AuthRequired,
37    /// Task state is unknown or unrecognized from a remote A2A agent.
38    #[serde(other)]
39    Unknown,
40}
41
42impl TaskState {
43    /// Check if this is a terminal state
44    fn is_terminal(&self) -> bool {
45        matches!(self, TaskState::Completed | TaskState::Failed | TaskState::Canceled | TaskState::Rejected)
46    }
47
48    /// Check if the task can be canceled from this state
49    fn is_cancelable(&self) -> bool {
50        matches!(self, TaskState::Submitted | TaskState::Working | TaskState::InputRequired)
51    }
52}
53
54/// Task status with state, optional message, and timestamp
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct TaskStatus {
57    /// Current lifecycle state
58    pub state: TaskState,
59    /// Optional message from the agent providing status details
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub message: Option<Message>,
62    /// ISO-8601 timestamp of the status update
63    pub(crate) timestamp: DateTime<Utc>,
64}
65
66impl TaskStatus {
67    /// Create a new task status
68    pub(crate) fn new(state: TaskState) -> Self {
69        Self { state, message: None, timestamp: Utc::now() }
70    }
71
72    /// Create a new task status with a message
73    pub(crate) fn with_message(state: TaskState, message: Message) -> Self {
74        Self {
75            state,
76            message: Some(message),
77            timestamp: Utc::now(),
78        }
79    }
80}
81
82impl Default for TaskStatus {
83    fn default() -> Self {
84        Self::new(TaskState::Submitted)
85    }
86}
87
88// ============================================================================
89// Message & Parts
90// ============================================================================
91
92/// Role of the message sender
93#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
94#[serde(rename_all = "lowercase")]
95pub enum MessageRole {
96    /// Message from the user/client
97    User,
98    /// Message from the agent
99    Agent,
100}
101
102/// A single unit of communication
103#[derive(Debug, Clone, Serialize, Deserialize)]
104#[serde(rename_all = "camelCase")]
105pub struct Message {
106    /// Sender's role
107    pub role: MessageRole,
108    /// Content of the message
109    pub parts: Vec<Part>,
110    /// Unique message identifier
111    #[serde(skip_serializing_if = "Option::is_none")]
112    message_id: Option<String>,
113    /// Associated task ID
114    #[serde(skip_serializing_if = "Option::is_none")]
115    task_id: Option<String>,
116    /// Conversation context ID
117    #[serde(skip_serializing_if = "Option::is_none")]
118    context_id: Option<String>,
119    /// List of prior task IDs for context
120    #[serde(skip_serializing_if = "Vec::is_empty", default)]
121    reference_task_ids: Vec<String>,
122    /// Custom metadata
123    #[serde(skip_serializing_if = "Option::is_none")]
124    metadata: Option<HashMap<String, serde_json::Value>>,
125}
126
127impl Message {
128    /// Create a new message
129    fn new(role: MessageRole, parts: Vec<Part>) -> Self {
130        Self {
131            role,
132            parts,
133            message_id: None,
134            task_id: None,
135            context_id: None,
136            reference_task_ids: Vec::new(),
137            metadata: None,
138        }
139    }
140
141    /// Create a text message from the agent
142    pub(crate) fn agent_text(text: impl Into<String>) -> Self {
143        Self::new(MessageRole::Agent, vec![Part::text(text)])
144    }
145
146    /// Create a text message from the user
147    pub fn user_text(text: impl Into<String>) -> Self {
148        Self::new(MessageRole::User, vec![Part::text(text)])
149    }
150
151    /// Set the message ID
152    pub fn with_id(mut self, id: impl Into<String>) -> Self {
153        self.message_id = Some(id.into());
154        self
155    }
156
157    /// Set the task ID
158    pub fn with_task_id(mut self, task_id: impl Into<String>) -> Self {
159        self.task_id = Some(task_id.into());
160        self
161    }
162
163    /// Set the context ID
164    pub fn with_context_id(mut self, context_id: impl Into<String>) -> Self {
165        self.context_id = Some(context_id.into());
166        self
167    }
168}
169
170/// Content part types: text, file, or structured data
171#[derive(Debug, Clone, Serialize, Deserialize)]
172#[serde(tag = "type", rename_all = "camelCase")]
173pub enum Part {
174    /// Plain text content
175    #[serde(rename = "text")]
176    Text {
177        /// The text content
178        text: String,
179    },
180    /// File content (URI or inline bytes)
181    #[serde(rename = "file")]
182    File {
183        /// File content details
184        file: FileContent,
185    },
186    /// Structured JSON data
187    #[serde(rename = "data")]
188    Data {
189        /// The structured data
190        data: serde_json::Value,
191    },
192    /// Catch-all for unknown part types added by the A2A spec.
193    #[serde(other)]
194    Unknown,
195}
196
197impl Part {
198    /// Create a text part
199    fn text(text: impl Into<String>) -> Self {
200        Part::Text { text: text.into() }
201    }
202
203    /// Create a file part from URI
204    pub fn file_uri(uri: impl Into<String>, mime_type: Option<String>) -> Self {
205        Part::File {
206            file: FileContent::Uri { uri: uri.into(), mime_type },
207        }
208    }
209
210    /// Create a file part from inline bytes
211    pub fn file_bytes(bytes: Vec<u8>, mime_type: Option<String>, name: Option<String>) -> Self {
212        Part::File {
213            file: FileContent::Bytes {
214                bytes: base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &bytes),
215                mime_type,
216                name,
217            },
218        }
219    }
220
221    /// Create a data part
222    fn data(data: serde_json::Value) -> Self {
223        Part::Data { data }
224    }
225
226    /// Get the text content if this is a text part
227    pub fn as_text(&self) -> Option<&str> {
228        match self {
229            Part::Text { text } => Some(text),
230            _ => None,
231        }
232    }
233
234    /// Returns true if this is an unknown/unsupported part type.
235    pub fn is_unknown(&self) -> bool {
236        matches!(self, Part::Unknown)
237    }
238}
239
240/// File content representation
241#[derive(Debug, Clone, Serialize, Deserialize)]
242#[serde(untagged)]
243pub enum FileContent {
244    /// File referenced by URI
245    Uri {
246        /// The file URI
247        uri: String,
248        /// Optional MIME type
249        #[serde(skip_serializing_if = "Option::is_none")]
250        mime_type: Option<String>,
251    },
252    /// File with inline base64-encoded bytes
253    Bytes {
254        /// Base64-encoded file content
255        bytes: String,
256        /// Optional MIME type
257        #[serde(skip_serializing_if = "Option::is_none")]
258        mime_type: Option<String>,
259        /// Optional file name
260        #[serde(skip_serializing_if = "Option::is_none")]
261        name: Option<String>,
262    },
263}
264
265// ============================================================================
266// Artifact
267// ============================================================================
268
269/// A tangible output generated by a task
270#[derive(Debug, Clone, Serialize, Deserialize)]
271#[serde(rename_all = "camelCase")]
272pub struct Artifact {
273    /// Unique artifact identifier
274    pub id: String,
275    /// Human-readable name
276    #[serde(skip_serializing_if = "Option::is_none")]
277    name: Option<String>,
278    /// Description of the artifact
279    #[serde(skip_serializing_if = "Option::is_none")]
280    description: Option<String>,
281    /// Content parts
282    pub parts: Vec<Part>,
283    /// Index for ordering multiple artifacts
284    #[serde(skip_serializing_if = "Option::is_none")]
285    index: Option<u32>,
286    /// Custom metadata
287    #[serde(skip_serializing_if = "Option::is_none")]
288    metadata: Option<HashMap<String, serde_json::Value>>,
289}
290
291impl Artifact {
292    /// Create a new artifact with text content
293    pub(crate) fn text(id: impl Into<String>, text: impl Into<String>) -> Self {
294        Self {
295            id: id.into(),
296            name: None,
297            description: None,
298            parts: vec![Part::text(text)],
299            index: None,
300            metadata: None,
301        }
302    }
303
304    /// Create a new artifact with file content
305    pub fn file(id: impl Into<String>, file: FileContent) -> Self {
306        Self {
307            id: id.into(),
308            name: None,
309            description: None,
310            parts: vec![Part::File { file }],
311            index: None,
312            metadata: None,
313        }
314    }
315
316    /// Add a name to the artifact
317    fn with_name(mut self, name: impl Into<String>) -> Self {
318        self.name = Some(name.into());
319        self
320    }
321
322    /// Add a description to the artifact
323    fn with_description(mut self, description: impl Into<String>) -> Self {
324        self.description = Some(description.into());
325        self
326    }
327}
328
329// ============================================================================
330// Task
331// ============================================================================
332
333/// Represents a stateful unit of work
334#[derive(Debug, Clone, Serialize, Deserialize)]
335#[serde(rename_all = "camelCase")]
336pub struct Task {
337    /// Unique task identifier
338    pub id: String,
339    /// Conversational context identifier
340    #[serde(skip_serializing_if = "Option::is_none")]
341    pub context_id: Option<String>,
342    /// Current status
343    pub status: TaskStatus,
344    /// Outputs produced by the agent
345    #[serde(skip_serializing_if = "Vec::is_empty", default)]
346    pub artifacts: Vec<Artifact>,
347    /// Conversation history (if enabled)
348    #[serde(skip_serializing_if = "Vec::is_empty", default)]
349    pub history: Vec<Message>,
350    /// Custom metadata
351    #[serde(skip_serializing_if = "Option::is_none")]
352    metadata: Option<HashMap<String, serde_json::Value>>,
353    /// Type discriminator
354    #[serde(default = "default_task_kind")]
355    kind: String,
356}
357
358fn default_task_kind() -> String {
359    "task".to_string()
360}
361
362impl Task {
363    /// Create a new task with a generated ID
364    pub(crate) fn new() -> Self {
365        Self {
366            id: uuid::Uuid::new_v4().to_string(),
367            context_id: None,
368            status: TaskStatus::default(),
369            artifacts: Vec::new(),
370            history: Vec::new(),
371            metadata: None,
372            kind: "task".to_string(),
373        }
374    }
375
376    /// Create a new task with a specific ID
377    pub fn with_id(id: impl Into<String>) -> Self {
378        Self {
379            id: id.into(),
380            context_id: None,
381            status: TaskStatus::default(),
382            artifacts: Vec::new(),
383            history: Vec::new(),
384            metadata: None,
385            kind: "task".to_string(),
386        }
387    }
388
389    /// Set the context ID
390    pub(crate) fn with_context_id(mut self, context_id: impl Into<String>) -> Self {
391        self.context_id = Some(context_id.into());
392        self
393    }
394
395    /// Get the current state
396    pub(crate) fn state(&self) -> TaskState {
397        self.status.state
398    }
399
400    /// Check if the task is in a terminal state
401    pub(crate) fn is_terminal(&self) -> bool {
402        self.status.state.is_terminal()
403    }
404
405    /// Check if the task can be canceled
406    pub(crate) fn is_cancelable(&self) -> bool {
407        self.status.state.is_cancelable()
408    }
409
410    /// Update the task status
411    fn update_status(&mut self, state: TaskState, message: Option<Message>) {
412        self.status = match message {
413            Some(msg) => TaskStatus::with_message(state, msg),
414            None => TaskStatus::new(state),
415        };
416    }
417
418    /// Add an artifact to the task
419    pub fn add_artifact(&mut self, artifact: Artifact) {
420        self.artifacts.push(artifact);
421    }
422
423    /// Add a message to the history
424    pub fn add_message(&mut self, message: Message) {
425        self.history.push(message);
426    }
427}
428
429impl Default for Task {
430    fn default() -> Self {
431        Self::new()
432    }
433}
434
435#[cfg(test)]
436mod tests {
437    use super::*;
438
439    #[test]
440    fn test_task_state_terminal() {
441        assert!(!TaskState::Submitted.is_terminal());
442        assert!(!TaskState::Working.is_terminal());
443        assert!(TaskState::Completed.is_terminal());
444        assert!(TaskState::Failed.is_terminal());
445        assert!(TaskState::Canceled.is_terminal());
446    }
447
448    #[test]
449    fn test_task_state_cancelable() {
450        assert!(TaskState::Submitted.is_cancelable());
451        assert!(TaskState::Working.is_cancelable());
452        assert!(TaskState::InputRequired.is_cancelable());
453        assert!(!TaskState::Completed.is_cancelable());
454        assert!(!TaskState::Failed.is_cancelable());
455    }
456
457    #[test]
458    fn test_message_creation() {
459        let msg = Message::agent_text("Hello, world!");
460        assert_eq!(msg.role, MessageRole::Agent);
461        assert_eq!(msg.parts.len(), 1);
462        assert_eq!(msg.parts[0].as_text(), Some("Hello, world!"));
463    }
464
465    #[test]
466    fn test_task_lifecycle() {
467        let mut task = Task::new();
468        assert_eq!(task.state(), TaskState::Submitted);
469        assert!(!task.is_terminal());
470        assert!(task.is_cancelable());
471
472        task.update_status(TaskState::Working, None);
473        assert_eq!(task.state(), TaskState::Working);
474
475        task.update_status(TaskState::Completed, Some(Message::agent_text("Task completed")));
476        assert!(task.is_terminal());
477        assert!(!task.is_cancelable());
478    }
479
480    #[test]
481    fn test_part_serialization() {
482        let text_part = Part::text("Hello");
483        let json = serde_json::to_string(&text_part).expect("serialize");
484        assert!(json.contains("\"type\":\"text\""));
485
486        let data_part = Part::data(serde_json::json!({"key": "value"}));
487        let json = serde_json::to_string(&data_part).expect("serialize");
488        assert!(json.contains("\"type\":\"data\""));
489    }
490
491    #[test]
492    fn test_artifact_creation() {
493        let artifact = Artifact::text("art-1", "Generated content")
494            .with_name("output.txt")
495            .with_description("The generated output file");
496
497        assert_eq!(artifact.id, "art-1");
498        assert_eq!(artifact.name, Some("output.txt".to_string()));
499        assert_eq!(artifact.parts.len(), 1);
500    }
501}