Skip to main content

systemprompt_models/a2a/
task_metadata.rs

1//! Lifecycle and accounting metadata for A2A tasks.
2//!
3//! [`TaskMetadata`] distinguishes the two [`TaskType`] flavours (MCP tool
4//! execution versus agent message), tracks timing and token usage, and carries
5//! an open-ended `extensions` map flattened into the serialized form. The
6//! `new_validated_*` constructors enforce the required-field contract before a
7//! task is recorded.
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12use chrono::Utc;
13use serde::{Deserialize, Serialize};
14use systemprompt_traits::validation::{
15    MetadataValidation, MetadataValidationError, Validate, ValidationResult,
16};
17
18use crate::execution::ExecutionStep;
19
20pub mod agent_names {
21    pub const SYSTEM: &str = "system";
22}
23
24#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
25#[serde(rename_all = "snake_case")]
26pub enum TaskType {
27    McpExecution,
28    AgentMessage,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
32pub struct TaskMetadata {
33    pub task_type: TaskType,
34    pub agent_name: String,
35    pub created_at: String,
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub updated_at: Option<String>,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub started_at: Option<String>,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub completed_at: Option<String>,
42    #[serde(skip_serializing_if = "Option::is_none")]
43    pub execution_time_ms: Option<i64>,
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub tool_name: Option<String>,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub mcp_server_name: Option<String>,
48    #[serde(skip_serializing_if = "Option::is_none")]
49    pub input_tokens: Option<u32>,
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub output_tokens: Option<u32>,
52    #[serde(skip_serializing_if = "Option::is_none")]
53    pub model: Option<String>,
54    #[serde(rename = "executionSteps", skip_serializing_if = "Option::is_none")]
55    pub execution_steps: Option<Vec<ExecutionStep>>,
56    #[serde(flatten, default)]
57    // JSON: A2A task metadata extension map (vendor-prefixed keys).
58    pub extensions: serde_json::Map<String, serde_json::Value>,
59}
60
61impl TaskMetadata {
62    pub fn new_mcp_execution(
63        agent_name: String,
64        tool_name: String,
65        mcp_server_name: String,
66    ) -> Self {
67        Self {
68            task_type: TaskType::McpExecution,
69            agent_name,
70            tool_name: Some(tool_name),
71            mcp_server_name: Some(mcp_server_name),
72            created_at: Utc::now().to_rfc3339(),
73            updated_at: None,
74            started_at: None,
75            completed_at: None,
76            execution_time_ms: None,
77            input_tokens: None,
78            output_tokens: None,
79            model: None,
80            execution_steps: None,
81            extensions: serde_json::Map::new(),
82        }
83    }
84
85    pub fn new_agent_message(agent_name: String) -> Self {
86        Self {
87            task_type: TaskType::AgentMessage,
88            agent_name,
89            tool_name: None,
90            mcp_server_name: None,
91            created_at: Utc::now().to_rfc3339(),
92            updated_at: None,
93            started_at: None,
94            completed_at: None,
95            execution_time_ms: None,
96            input_tokens: None,
97            output_tokens: None,
98            model: None,
99            execution_steps: None,
100            extensions: serde_json::Map::new(),
101        }
102    }
103
104    pub const fn with_token_usage(mut self, input_tokens: u32, output_tokens: u32) -> Self {
105        self.input_tokens = Some(input_tokens);
106        self.output_tokens = Some(output_tokens);
107        self
108    }
109
110    pub fn with_model(mut self, model: impl Into<String>) -> Self {
111        self.model = Some(model.into());
112        self
113    }
114
115    pub fn with_updated_at(mut self) -> Self {
116        self.updated_at = Some(Utc::now().to_rfc3339());
117        self
118    }
119
120    pub fn with_tool_name(mut self, tool_name: impl Into<String>) -> Self {
121        self.tool_name = Some(tool_name.into());
122        self
123    }
124
125    pub fn with_execution_steps(mut self, steps: Vec<ExecutionStep>) -> Self {
126        self.execution_steps = Some(steps);
127        self
128    }
129
130    // JSON: A2A task metadata extension map (vendor-prefixed keys).
131    pub fn with_extension(mut self, key: String, value: serde_json::Value) -> Self {
132        self.extensions.insert(key, value);
133        self
134    }
135
136    pub fn new_validated_agent_message(agent_name: String) -> ValidationResult<Self> {
137        if agent_name.is_empty() {
138            return Err(MetadataValidationError::new(
139                "agent_name",
140                "Cannot create TaskMetadata: agent_name is empty",
141            )
142            .with_context(format!("agent_name={agent_name:?}")));
143        }
144
145        let metadata = Self::new_agent_message(agent_name);
146        metadata.validate()?;
147        Ok(metadata)
148    }
149
150    pub fn new_validated_mcp_execution(
151        agent_name: String,
152        tool_name: String,
153        mcp_server_name: String,
154    ) -> ValidationResult<Self> {
155        if agent_name.is_empty() {
156            return Err(MetadataValidationError::new(
157                "agent_name",
158                "Cannot create TaskMetadata: agent_name is empty",
159            )
160            .with_context(format!("agent_name={agent_name:?}")));
161        }
162
163        if tool_name.is_empty() {
164            return Err(MetadataValidationError::new(
165                "tool_name",
166                "Cannot create TaskMetadata: tool_name is empty for MCP execution",
167            )
168            .with_context(format!("tool_name={tool_name:?}")));
169        }
170
171        let metadata = Self::new_mcp_execution(agent_name, tool_name, mcp_server_name);
172        metadata.validate()?;
173        Ok(metadata)
174    }
175}
176
177impl Validate for TaskMetadata {
178    fn validate(&self) -> ValidationResult<()> {
179        self.validate_required_fields()?;
180        Ok(())
181    }
182}
183
184impl MetadataValidation for TaskMetadata {
185    fn required_string_fields(&self) -> Vec<(&'static str, &str)> {
186        vec![
187            ("agent_name", &self.agent_name),
188            ("created_at", &self.created_at),
189        ]
190    }
191}