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}