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("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}