Skip to main content

oneai_workflow/
config.rs

1//! Workflow configuration DSL — defines workflow steps, dependencies, and tool bindings.
2//!
3//! A workflow is defined as a set of steps with dependencies between them.
4//! The configuration specifies:
5//! - Steps: each step has an ID, description, and optional tool bindings
6//! - Dependencies: which steps must complete before others can start
7//! - Variables: shared context variables for passing data between steps
8//! - Retry policy: how to handle step failures
9//! - Timeout: maximum time for the entire workflow or individual steps
10
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13
14/// A workflow configuration — the user-facing definition of a workflow.
15///
16/// Workflows are defined in YAML/JSON and compiled into a DAG for execution.
17/// Each step specifies what it does and what it depends on.
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19pub struct WorkflowConfig {
20    /// Unique workflow name/identifier.
21    pub name: String,
22
23    /// Human-readable description.
24    #[serde(default)]
25    pub description: String,
26
27    /// Version of this workflow configuration.
28    #[serde(default = "default_version")]
29    pub version: String,
30
31    /// The steps in this workflow.
32    pub steps: Vec<StepConfig>,
33
34    /// Shared context variables (initial values).
35    #[serde(default)]
36    pub variables: HashMap<String, String>,
37
38    /// Workflow-level timeout in seconds.
39    #[serde(default)]
40    pub timeout_secs: Option<u64>,
41
42    /// Default retry policy for steps (can be overridden per-step).
43    #[serde(default)]
44    pub default_retry_policy: RetryPolicy,
45
46    /// Whether to continue on step failure.
47    #[serde(default)]
48    pub continue_on_failure: bool,
49}
50
51fn default_version() -> String {
52    "1.0".to_string()
53}
54
55/// Configuration for a single workflow step.
56#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
57pub struct StepConfig {
58    /// Unique step identifier within the workflow.
59    pub id: String,
60
61    /// Human-readable description of what this step does.
62    pub description: String,
63
64    /// IDs of steps that must complete before this step can start.
65    #[serde(default)]
66    pub depends_on: Vec<String>,
67
68    /// The tool to use for this step (if applicable).
69    #[serde(default)]
70    pub tool: Option<String>,
71
72    /// Arguments for the tool (JSON or template expressions).
73    #[serde(default)]
74    pub tool_args: Option<serde_json::Value>,
75
76    /// The prompt template for LLM-based steps.
77    #[serde(default)]
78    pub prompt: Option<String>,
79
80    /// Whether this step is an approval checkpoint (requires human approval).
81    #[serde(default)]
82    pub requires_approval: bool,
83
84    /// Step-level timeout in seconds (overrides workflow timeout).
85    #[serde(default)]
86    pub timeout_secs: Option<u64>,
87
88    /// Step-level retry policy (overrides default).
89    #[serde(default)]
90    pub retry_policy: Option<RetryPolicy>,
91
92    /// Step-specific metadata.
93    #[serde(default)]
94    pub metadata: HashMap<String, String>,
95}
96
97/// Retry policy for workflow steps.
98#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
99pub struct RetryPolicy {
100    /// Maximum number of retries.
101    #[serde(default = "default_max_retries")]
102    pub max_retries: usize,
103
104    /// Delay between retries in seconds.
105    #[serde(default = "default_retry_delay_secs")]
106    pub retry_delay_secs: u64,
107
108    /// Whether to retry on all errors or only specific ones.
109    #[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    /// Create a new workflow config with name and steps.
133    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    /// Parse a workflow config from JSON.
147    pub fn from_json(json: &str) -> serde_json::Result<Self> {
148        serde_json::from_str(json)
149    }
150
151    /// Serialize to JSON.
152    pub fn to_json(&self) -> serde_json::Result<String> {
153        serde_json::to_string_pretty(self)
154    }
155
156    /// Get all step IDs.
157    pub fn step_ids(&self) -> Vec<String> {
158        self.steps.iter().map(|s| s.id.clone()).collect()
159    }
160
161    /// Get a step by ID.
162    pub fn get_step(&self, id: &str) -> Option<&StepConfig> {
163        self.steps.iter().find(|s| s.id == id)
164    }
165
166    /// Get the effective retry policy for a step.
167    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    /// Get the effective timeout for a step.
174    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}