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