1use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19pub struct WorkflowConfig {
20 pub name: String,
22
23 #[serde(default)]
25 pub description: String,
26
27 #[serde(default = "default_version")]
29 pub version: String,
30
31 pub steps: Vec<StepConfig>,
33
34 #[serde(default)]
36 pub variables: HashMap<String, String>,
37
38 #[serde(default)]
40 pub timeout_secs: Option<u64>,
41
42 #[serde(default)]
44 pub default_retry_policy: RetryPolicy,
45
46 #[serde(default)]
48 pub continue_on_failure: bool,
49}
50
51fn default_version() -> String {
52 "1.0".to_string()
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
57pub struct StepConfig {
58 pub id: String,
60
61 pub description: String,
63
64 #[serde(default)]
66 pub depends_on: Vec<String>,
67
68 #[serde(default)]
70 pub tool: Option<String>,
71
72 #[serde(default)]
74 pub tool_args: Option<serde_json::Value>,
75
76 #[serde(default)]
78 pub prompt: Option<String>,
79
80 #[serde(default)]
82 pub requires_approval: bool,
83
84 #[serde(default)]
86 pub timeout_secs: Option<u64>,
87
88 #[serde(default)]
90 pub retry_policy: Option<RetryPolicy>,
91
92 #[serde(default)]
94 pub metadata: HashMap<String, String>,
95}
96
97#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
99pub struct RetryPolicy {
100 #[serde(default = "default_max_retries")]
102 pub max_retries: usize,
103
104 #[serde(default = "default_retry_delay_secs")]
106 pub retry_delay_secs: u64,
107
108 #[serde(default)]
110 pub retry_on_all_errors: bool,
111}
112
113fn default_max_retries() -> usize {
114 3
115}
116
117fn default_retry_delay_secs() -> u64 {
118 5
119}
120
121impl Default for RetryPolicy {
122 fn default() -> Self {
123 Self {
124 max_retries: default_max_retries(),
125 retry_delay_secs: default_retry_delay_secs(),
126 retry_on_all_errors: false,
127 }
128 }
129}
130
131impl WorkflowConfig {
132 pub fn new(name: impl Into<String>, steps: Vec<StepConfig>) -> Self {
134 Self {
135 name: name.into(),
136 description: String::new(),
137 version: default_version(),
138 steps,
139 variables: HashMap::new(),
140 timeout_secs: None,
141 default_retry_policy: RetryPolicy::default(),
142 continue_on_failure: false,
143 }
144 }
145
146 pub fn from_json(json: &str) -> serde_json::Result<Self> {
148 serde_json::from_str(json)
149 }
150
151 pub fn to_json(&self) -> serde_json::Result<String> {
153 serde_json::to_string_pretty(self)
154 }
155
156 pub fn step_ids(&self) -> Vec<String> {
158 self.steps.iter().map(|s| s.id.clone()).collect()
159 }
160
161 pub fn get_step(&self, id: &str) -> Option<&StepConfig> {
163 self.steps.iter().find(|s| s.id == id)
164 }
165
166 pub fn effective_retry_policy(&self, step_id: &str) -> RetryPolicy {
168 self.get_step(step_id)
169 .and_then(|s| s.retry_policy.clone())
170 .unwrap_or_else(|| self.default_retry_policy.clone())
171 }
172
173 pub fn effective_timeout(&self, step_id: &str) -> Option<u64> {
175 self.get_step(step_id)
176 .and_then(|s| s.timeout_secs)
177 .or(self.timeout_secs)
178 }
179}
180
181#[cfg(test)]
182mod tests {
183 use super::*;
184
185 #[test]
186 fn test_workflow_config_creation() {
187 let config = WorkflowConfig::new("test_workflow", vec![
188 StepConfig {
189 id: "step1".to_string(),
190 description: "First step".to_string(),
191 depends_on: vec![],
192 tool: None,
193 tool_args: None,
194 prompt: Some("Do something".to_string()),
195 requires_approval: false,
196 timeout_secs: None,
197 retry_policy: None,
198 metadata: HashMap::new(),
199 },
200 StepConfig {
201 id: "step2".to_string(),
202 description: "Second step".to_string(),
203 depends_on: vec!["step1".to_string()],
204 tool: Some("calculator".to_string()),
205 tool_args: Some(serde_json::json!({"expression": "2+2"})),
206 prompt: None,
207 requires_approval: false,
208 timeout_secs: None,
209 retry_policy: None,
210 metadata: HashMap::new(),
211 },
212 ]);
213
214 assert_eq!(config.name, "test_workflow");
215 assert_eq!(config.steps.len(), 2);
216 assert_eq!(config.step_ids(), vec!["step1", "step2"]);
217 }
218
219 #[test]
220 fn test_workflow_config_serialization() {
221 let config = WorkflowConfig::new("test_workflow", vec![
222 StepConfig {
223 id: "step1".to_string(),
224 description: "First step".to_string(),
225 depends_on: vec![],
226 tool: None,
227 tool_args: None,
228 prompt: Some("Calculate something".to_string()),
229 requires_approval: false,
230 timeout_secs: Some(30),
231 retry_policy: None,
232 metadata: HashMap::new(),
233 },
234 ]);
235
236 let json = config.to_json().unwrap();
237 let parsed: WorkflowConfig = WorkflowConfig::from_json(&json).unwrap();
238 assert_eq!(parsed, config);
239 }
240
241 #[test]
242 fn test_effective_retry_policy() {
243 let config = WorkflowConfig::new("test", vec![
244 StepConfig {
245 id: "step1".to_string(),
246 description: "Step 1".to_string(),
247 depends_on: vec![],
248 tool: None,
249 tool_args: None,
250 prompt: None,
251 requires_approval: false,
252 timeout_secs: None,
253 retry_policy: Some(RetryPolicy { max_retries: 5, retry_delay_secs: 10, retry_on_all_errors: true }),
254 metadata: HashMap::new(),
255 },
256 StepConfig {
257 id: "step2".to_string(),
258 description: "Step 2".to_string(),
259 depends_on: vec!["step1".to_string()],
260 tool: None,
261 tool_args: None,
262 prompt: None,
263 requires_approval: false,
264 timeout_secs: None,
265 retry_policy: None,
266 metadata: HashMap::new(),
267 },
268 ]);
269
270 let policy1 = config.effective_retry_policy("step1");
271 assert_eq!(policy1.max_retries, 5);
272
273 let policy2 = config.effective_retry_policy("step2");
274 assert_eq!(policy2.max_retries, 3);
275 }
276
277 #[test]
278 fn test_effective_timeout() {
279 let mut config = WorkflowConfig::new("test", vec![
280 StepConfig {
281 id: "step1".to_string(),
282 description: "Step 1".to_string(),
283 depends_on: vec![],
284 tool: None,
285 tool_args: None,
286 prompt: None,
287 requires_approval: false,
288 timeout_secs: Some(60),
289 retry_policy: None,
290 metadata: HashMap::new(),
291 },
292 ]);
293 config.timeout_secs = Some(120);
294
295 assert_eq!(config.effective_timeout("step1"), Some(60));
296 assert_eq!(config.effective_timeout("nonexistent"), Some(120));
297 }
298}