Skip to main content

systemprompt_models/a2a/
task_metadata.rs

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