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(
188 "test_workflow",
189 vec![
190 StepConfig {
191 id: "step1".to_string(),
192 description: "First step".to_string(),
193 depends_on: vec![],
194 tool: None,
195 tool_args: None,
196 prompt: Some("Do something".to_string()),
197 requires_approval: false,
198 timeout_secs: None,
199 retry_policy: None,
200 metadata: HashMap::new(),
201 },
202 StepConfig {
203 id: "step2".to_string(),
204 description: "Second step".to_string(),
205 depends_on: vec!["step1".to_string()],
206 tool: Some("calculator".to_string()),
207 tool_args: Some(serde_json::json!({"expression": "2+2"})),
208 prompt: None,
209 requires_approval: false,
210 timeout_secs: None,
211 retry_policy: None,
212 metadata: HashMap::new(),
213 },
214 ],
215 );
216
217 assert_eq!(config.name, "test_workflow");
218 assert_eq!(config.steps.len(), 2);
219 assert_eq!(config.step_ids(), vec!["step1", "step2"]);
220 }
221
222 #[test]
223 fn test_workflow_config_serialization() {
224 let config = WorkflowConfig::new(
225 "test_workflow",
226 vec![StepConfig {
227 id: "step1".to_string(),
228 description: "First step".to_string(),
229 depends_on: vec![],
230 tool: None,
231 tool_args: None,
232 prompt: Some("Calculate something".to_string()),
233 requires_approval: false,
234 timeout_secs: Some(30),
235 retry_policy: None,
236 metadata: HashMap::new(),
237 }],
238 );
239
240 let json = config.to_json().unwrap();
241 let parsed: WorkflowConfig = WorkflowConfig::from_json(&json).unwrap();
242 assert_eq!(parsed, config);
243 }
244
245 #[test]
246 fn test_effective_retry_policy() {
247 let config = WorkflowConfig::new(
248 "test",
249 vec![
250 StepConfig {
251 id: "step1".to_string(),
252 description: "Step 1".to_string(),
253 depends_on: vec![],
254 tool: None,
255 tool_args: None,
256 prompt: None,
257 requires_approval: false,
258 timeout_secs: None,
259 retry_policy: Some(RetryPolicy {
260 max_retries: 5,
261 retry_delay_secs: 10,
262 retry_on_all_errors: true,
263 }),
264 metadata: HashMap::new(),
265 },
266 StepConfig {
267 id: "step2".to_string(),
268 description: "Step 2".to_string(),
269 depends_on: vec!["step1".to_string()],
270 tool: None,
271 tool_args: None,
272 prompt: None,
273 requires_approval: false,
274 timeout_secs: None,
275 retry_policy: None,
276 metadata: HashMap::new(),
277 },
278 ],
279 );
280
281 let policy1 = config.effective_retry_policy("step1");
282 assert_eq!(policy1.max_retries, 5);
283
284 let policy2 = config.effective_retry_policy("step2");
285 assert_eq!(policy2.max_retries, 3);
286 }
287
288 #[test]
289 fn test_effective_timeout() {
290 let mut config = WorkflowConfig::new(
291 "test",
292 vec![StepConfig {
293 id: "step1".to_string(),
294 description: "Step 1".to_string(),
295 depends_on: vec![],
296 tool: None,
297 tool_args: None,
298 prompt: None,
299 requires_approval: false,
300 timeout_secs: Some(60),
301 retry_policy: None,
302 metadata: HashMap::new(),
303 }],
304 );
305 config.timeout_secs = Some(120);
306
307 assert_eq!(config.effective_timeout("step1"), Some(60));
308 assert_eq!(config.effective_timeout("nonexistent"), Some(120));
309 }
310}