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    /// Clone only the history and artifacts included by the query response.
396    pub(crate) fn clone_for_query(&self, history_length: usize, include_artifacts: bool) -> Self {
397        Self {
398            id: self.id.clone(),
399            context_id: self.context_id.clone(),
400            status: self.status.clone(),
401            artifacts: if include_artifacts {
402                self.artifacts.clone()
403            } else {
404                Vec::new()
405            },
406            history: self
407                .history
408                .iter()
409                .skip(self.history.len().saturating_sub(history_length))
410                .cloned()
411                .collect(),
412            metadata: self.metadata.clone(),
413            kind: self.kind.clone(),
414        }
415    }
416
417    /// Get the current state
418    pub(crate) fn state(&self) -> TaskState {
419        self.status.state
420    }
421
422    /// Check if the task is in a terminal state
423    pub(crate) fn is_terminal(&self) -> bool {
424        self.status.state.is_terminal()
425    }
426
427    /// Check if the task can be canceled
428    pub(crate) fn is_cancelable(&self) -> bool {
429        self.status.state.is_cancelable()
430    }
431
432    /// Update the task status
433    fn update_status(&mut self, state: TaskState, message: Option<Message>) {
434        self.status = match message {
435            Some(msg) => TaskStatus::with_message(state, msg),
436            None => TaskStatus::new(state),
437        };
438    }
439
440    /// Add an artifact to the task
441    pub fn add_artifact(&mut self, artifact: Artifact) {
442        self.artifacts.push(artifact);
443    }
444
445    /// Add a message to the history
446    pub fn add_message(&mut self, message: Message) {
447        self.history.push(message);
448    }
449}
450
451impl Default for Task {
452    fn default() -> Self {
453        Self::new()
454    }
455}
456
457#[cfg(test)]
458mod tests {
459    use super::*;
460
461    #[test]
462    fn test_task_state_terminal() {
463        assert!(!TaskState::Submitted.is_terminal());
464        assert!(!TaskState::Working.is_terminal());
465        assert!(TaskState::Completed.is_terminal());
466        assert!(TaskState::Failed.is_terminal());
467        assert!(TaskState::Canceled.is_terminal());
468    }
469
470    #[test]
471    fn test_task_state_cancelable() {
472        assert!(TaskState::Submitted.is_cancelable());
473        assert!(TaskState::Working.is_cancelable());
474        assert!(TaskState::InputRequired.is_cancelable());
475        assert!(!TaskState::Completed.is_cancelable());
476        assert!(!TaskState::Failed.is_cancelable());
477    }
478
479    #[test]
480    fn test_message_creation() {
481        let msg = Message::agent_text("Hello, world!");
482        assert_eq!(msg.role, MessageRole::Agent);
483        assert_eq!(msg.parts.len(), 1);
484        assert_eq!(msg.parts[0].as_text(), Some("Hello, world!"));
485    }
486
487    #[test]
488    fn test_task_lifecycle() {
489        let mut task = Task::new();
490        assert_eq!(task.state(), TaskState::Submitted);
491        assert!(!task.is_terminal());
492        assert!(task.is_cancelable());
493
494        task.update_status(TaskState::Working, None);
495        assert_eq!(task.state(), TaskState::Working);
496
497        task.update_status(TaskState::Completed, Some(Message::agent_text("Task completed")));
498        assert!(task.is_terminal());
499        assert!(!task.is_cancelable());
500    }
501
502    #[test]
503    fn test_part_serialization() {
504        let text_part = Part::text("Hello");
505        let json = serde_json::to_string(&text_part).expect("serialize");
506        assert!(json.contains("\"type\":\"text\""));
507
508        let data_part = Part::data(serde_json::json!({"key": "value"}));
509        let json = serde_json::to_string(&data_part).expect("serialize");
510        assert!(json.contains("\"type\":\"data\""));
511    }
512
513    #[test]
514    fn test_artifact_creation() {
515        let artifact = Artifact::text("art-1", "Generated content")
516            .with_name("output.txt")
517            .with_description("The generated output file");
518
519        assert_eq!(artifact.id, "art-1");
520        assert_eq!(artifact.name, Some("output.txt".to_string()));
521        assert_eq!(artifact.parts.len(), 1);
522    }
523}