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