1use 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 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 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}