From e0ac559ae1525ee0dae535f535d6029dcf51ca28 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 12 Jan 2026 15:06:12 -0800 Subject: [PATCH] feat(workflow): add scheduled/delayed steps and conditional branching (#16980) --- sgl-model-gateway/src/workflow/definition.rs | 114 +++- sgl-model-gateway/src/workflow/engine.rs | 299 ++++++++- sgl-model-gateway/tests/workflow_test.rs | 638 +++++++++++++++++++ 3 files changed, 1025 insertions(+), 26 deletions(-) diff --git a/sgl-model-gateway/src/workflow/definition.rs b/sgl-model-gateway/src/workflow/definition.rs index 6c4aa721c..f88e411f5 100644 --- a/sgl-model-gateway/src/workflow/definition.rs +++ b/sgl-model-gateway/src/workflow/definition.rs @@ -7,9 +7,11 @@ use std::{ time::Duration, }; +use chrono::{DateTime, Utc}; + use super::{ executor::StepExecutor, - types::{FailureAction, RetryPolicy, StepId, WorkflowData, WorkflowId}, + types::{FailureAction, RetryPolicy, StepId, WorkflowContext, WorkflowData, WorkflowId}, }; /// Errors that can occur during workflow validation @@ -32,7 +34,17 @@ pub struct StepDefinition { pub retry_policy: Option, pub timeout: Option, pub on_failure: FailureAction, + /// Dependencies that must ALL complete before this step runs pub depends_on: Vec, + /// Dependencies where ANY completing triggers this step (used with depends_on) + pub depends_on_any: Vec, + /// Delay before starting the step (after dependencies satisfied) + pub delay: Option, + /// Run step at or after this time (after dependencies satisfied) + pub scheduled_at: Option>, + /// Condition to evaluate; if false, step is skipped + #[allow(clippy::type_complexity)] + pub run_if: Option) -> bool + Send + Sync>>, } impl fmt::Debug for StepDefinition { @@ -44,6 +56,10 @@ impl fmt::Debug for StepDefinition { .field("timeout", &self.timeout) .field("on_failure", &self.on_failure) .field("depends_on", &self.depends_on) + .field("depends_on_any", &self.depends_on_any) + .field("delay", &self.delay) + .field("scheduled_at", &self.scheduled_at) + .field("run_if", &self.run_if.as_ref().map(|_| "")) .finish_non_exhaustive() } } @@ -62,6 +78,10 @@ impl StepDefinition { timeout: None, on_failure: FailureAction::FailWorkflow, depends_on: Vec::new(), + depends_on_any: Vec::new(), + delay: None, + scheduled_at: None, + run_if: None, } } @@ -81,12 +101,74 @@ impl StepDefinition { } /// Set dependencies for this step. - /// The step will only run after all specified dependencies have completed successfully. + /// The step will only run after ALL specified dependencies have completed successfully. /// Empty slice means no dependencies - step can run immediately in parallel with others. pub fn depends_on(mut self, deps: &[&str]) -> Self { self.depends_on = deps.iter().map(|s| StepId::new(*s)).collect(); self } + + /// Set "any of" dependencies for this step. + /// The step will run when ANY of these dependencies complete (in addition to + /// all `depends_on` dependencies). + /// + /// **Combined semantics** (when both `depends_on` and `depends_on_any` are set): + /// - ALL `depends_on` must complete successfully, AND + /// - AT LEAST ONE `depends_on_any` must complete successfully + /// + /// **Failure handling**: + /// - For `depends_on`: if ANY fails, this step is blocked + /// - For `depends_on_any`: only blocked if ALL fail (since we only need one) + /// + /// **Skipped dependencies**: A skipped dependency (e.g., via `run_if`) counts as + /// "completed" for dependency satisfaction purposes. + pub fn depends_on_any(mut self, deps: &[&str]) -> Self { + self.depends_on_any = deps.iter().map(|s| StepId::new(*s)).collect(); + self + } + + /// Set a delay before starting the step (after dependencies are satisfied). + /// + /// If both `delay` and `scheduled_at` are set, the step will wait for the + /// scheduled time AND THEN wait for the delay duration (they stack). + pub fn with_delay(mut self, delay: Duration) -> Self { + self.delay = Some(delay); + self + } + + /// Schedule the step to run at or after the specified time. + /// The step will wait until this time even if dependencies are satisfied earlier. + /// + /// If the scheduled time is in the past when the step becomes ready, it will + /// proceed immediately (a debug log is emitted). + /// + /// If both `delay` and `scheduled_at` are set, the step will wait for the + /// scheduled time AND THEN wait for the delay duration (they stack). + pub fn scheduled_at(mut self, time: DateTime) -> Self { + self.scheduled_at = Some(time); + self + } + + /// Set a condition for running this step. + /// If the condition returns false, the step is skipped. + /// + /// **Skipped step semantics**: When a step is skipped: + /// - It is marked as `StepStatus::Skipped` + /// - Downstream steps that depend on it (via `depends_on`) will consider it satisfied + /// - This allows conditional branches without blocking the workflow + /// + /// **Error handling**: If the context cannot be retrieved to evaluate the condition, + /// the step will **fail** (not proceed blindly). This is a safety measure. + /// + /// **Note**: The condition closure is not serializable, so workflows with `run_if` + /// cannot be persisted and resumed from external storage. + pub fn run_if(mut self, condition: F) -> Self + where + F: Fn(&WorkflowContext) -> bool + Send + Sync + 'static, + { + self.run_if = Some(Arc::new(condition)); + self + } } /// Complete workflow definition @@ -166,7 +248,7 @@ impl WorkflowDefinition { let steps_map: HashMap<&StepId, &StepDefinition> = self.steps.iter().map(|s| (&s.id, s)).collect(); - // Check all dependencies exist + // Check all dependencies exist (both depends_on and depends_on_any) for step in &self.steps { for dep in &step.depends_on { if !steps_map.contains_key(dep) { @@ -176,9 +258,17 @@ impl WorkflowDefinition { }); } } + for dep in &step.depends_on_any { + if !steps_map.contains_key(dep) { + return Err(ValidationError::MissingDependency { + step: step.id.clone(), + dependency: dep.clone(), + }); + } + } } - // Check for cycles using DFS + // Check for cycles using DFS (considers both dependency types) let mut visited = HashSet::new(); let mut rec_stack = HashSet::new(); @@ -191,6 +281,7 @@ impl WorkflowDefinition { } // Build reverse dependency map: for each step, which steps depend on it? + // Include both depends_on and depends_on_any self.reverse_deps.clear(); for (idx, step) in self.steps.iter().enumerate() { for dep_id in &step.depends_on { @@ -199,14 +290,21 @@ impl WorkflowDefinition { .or_default() .push(idx); } + for dep_id in &step.depends_on_any { + self.reverse_deps + .entry(dep_id.clone()) + .or_default() + .push(idx); + } } // Cache indices of steps with no dependencies (can start immediately) + // A step with only depends_on_any still needs at least one to complete self.initial_step_indices = self .steps .iter() .enumerate() - .filter(|(_, s)| s.depends_on.is_empty()) + .filter(|(_, s)| s.depends_on.is_empty() && s.depends_on_any.is_empty()) .map(|(i, _)| i) .collect(); @@ -231,12 +329,18 @@ impl WorkflowDefinition { rec_stack.insert(step_id); // O(1) lookup instead of linear search + // Check both depends_on and depends_on_any for cycles if let Some(step) = steps_map.get(step_id) { for dep in &step.depends_on { if Self::has_cycle(dep, steps_map, visited, rec_stack) { return true; } } + for dep in &step.depends_on_any { + if Self::has_cycle(dep, steps_map, visited, rec_stack) { + return true; + } + } } rec_stack.remove(step_id); diff --git a/sgl-model-gateway/src/workflow/engine.rs b/sgl-model-gateway/src/workflow/engine.rs index 259615a10..dec72d93f 100644 --- a/sgl-model-gateway/src/workflow/engine.rs +++ b/sgl-model-gateway/src/workflow/engine.rs @@ -35,6 +35,9 @@ struct StepTracker { failed: HashSet, skipped: HashSet, running: HashSet, + /// Steps waiting for delay/scheduled_at: maps step INDEX to ready time + /// Using index instead of StepId for O(w) iteration in the main loop + waiting_until: HashMap, } impl StepTracker { @@ -42,22 +45,59 @@ impl StepTracker { self.completed.len() + self.failed.len() + self.skipped.len() } - fn is_step_processable(&self, step_id: &StepId) -> bool { + fn is_step_processable(&self, step_id: &StepId, step_idx: usize) -> bool { !self.completed.contains(step_id) && !self.failed.contains(step_id) && !self.skipped.contains(step_id) && !self.running.contains(step_id) + && !self.waiting_until.contains_key(&step_idx) } + /// Get indices of waiting steps that are now ready to run - O(w) where w = waiting count + fn get_ready_waiting_indices(&self) -> Vec { + let now = std::time::Instant::now(); + self.waiting_until + .iter() + .filter(|(_, &ready_at)| now >= ready_at) + .map(|(&idx, _)| idx) + .collect() + } + + /// Mark a step as waiting until a specific time (by index) + fn set_waiting(&mut self, step_idx: usize, ready_at: std::time::Instant) { + self.waiting_until.insert(step_idx, ready_at); + } + + /// Clear waiting status for a step (it's now ready to run) + fn clear_waiting(&mut self, step_idx: usize) { + self.waiting_until.remove(&step_idx); + } + + /// Check if ALL dependencies are satisfied (completed or skipped) fn are_dependencies_satisfied(&self, depends_on: &[StepId]) -> bool { depends_on .iter() .all(|dep| self.completed.contains(dep) || self.skipped.contains(dep)) } + /// Check if ANY dependency is satisfied (completed or skipped) + /// Returns true if the list is empty (no "any" dependencies) + fn is_any_dependency_satisfied(&self, depends_on_any: &[StepId]) -> bool { + depends_on_any.is_empty() + || depends_on_any + .iter() + .any(|dep| self.completed.contains(dep) || self.skipped.contains(dep)) + } + + /// Check if ANY dependency has failed (for depends_on - blocks if any fail) fn has_failed_dependency(&self, depends_on: &[StepId]) -> bool { depends_on.iter().any(|dep| self.failed.contains(dep)) } + + /// Check if ALL dependencies have failed (for depends_on_any - blocks only if all fail) + fn have_all_any_deps_failed(&self, depends_on_any: &[StepId]) -> bool { + !depends_on_any.is_empty() && depends_on_any.iter().all(|dep| self.failed.contains(dep)) + } } /// Fixed backoff that returns the same delay every time @@ -264,6 +304,35 @@ impl + 'static> WorkflowEngine { self.active_workflows.fetch_sub(1, Ordering::Release); } + /// Calculate how long a step needs to wait based on delay and/or scheduled_at. + /// Returns None if no waiting is needed. + fn calculate_wait_duration(step: &StepDefinition) -> Option { + let now = Utc::now(); + + // Calculate wait time for scheduled_at + let schedule_wait = step.scheduled_at.and_then(|scheduled_time| { + if now < scheduled_time { + (scheduled_time - now).to_std().ok() + } else { + // Scheduled time is in the past + tracing::debug!( + step_id = %step.id, + scheduled_time = %scheduled_time, + "Step scheduled_at is in the past, proceeding immediately" + ); + None + } + }); + + // Combine delay and schedule_wait (both apply if both are set) + match (step.delay, schedule_wait) { + (Some(delay), Some(schedule)) => Some(delay + schedule), + (Some(delay), None) => Some(delay), + (None, Some(schedule)) => Some(schedule), + (None, None) => None, + } + } + /// Create a guard that decrements active_workflows on drop. /// This ensures the count is decremented even if a task panics. fn active_workflow_guard(&self) -> ActiveWorkflowGuard { @@ -428,22 +497,95 @@ impl + 'static> WorkflowEngine { return Ok(()); } - // Find ready steps from pending_check (not all steps) - let (ready_step_indices, total_processed, running_count) = { + // Phase 0: Drain any pending completion signals to ensure dependents are added + // This prevents race conditions where tasks finish but their signals aren't processed + while let Ok((step_id, result)) = rx.try_recv() { + tracing::debug!( + step_id = %step_id, + result = ?result, + "Step completed (startup drain)" + ); + if matches!(result, StepResult::Success | StepResult::Skip) { + for &dep_idx in definition.get_dependent_indices(&step_id) { + pending_check.push_back(dep_idx); + } + } + } + + // Phase 1: Check waiting steps + deps-ready steps + blocked detection + // Single lock acquisition for all read operations + let ( + newly_ready_from_wait, + deps_ready_indices, + total_processed, + current_running, + current_waiting, + ) = { let t = tracker.read(); - // Only check steps in pending_check, not all steps - let ready: Vec = pending_check + // O(w) iteration over waiting_until keys instead of O(n) over all steps + let wait_ready: Vec = t.get_ready_waiting_indices(); + + // Check pending_check for dependency-satisfied steps + let deps_ready: Vec = pending_check .drain(..) .filter(|&idx| { let step = &definition.steps[idx]; - t.is_step_processable(&step.id) + t.is_step_processable(&step.id, idx) && t.are_dependencies_satisfied(&step.depends_on) + && t.is_any_dependency_satisfied(&step.depends_on_any) && !t.has_failed_dependency(&step.depends_on) + && !t.have_all_any_deps_failed(&step.depends_on_any) }) .collect(); - (ready, t.total_processed(), t.running.len()) + ( + wait_ready, + deps_ready, + t.total_processed(), + t.running.len(), + t.waiting_until.len(), + ) + }; + + // Phase 2: Process waiting and deps-ready steps, update waiting_until + // Single write lock for all mutations + // Returns (ready_to_launch, steps_added_to_waiting) + let (ready_to_launch, steps_added_to_waiting) = { + let now = std::time::Instant::now(); + let mut t = tracker.write(); + let mut added_to_waiting = 0usize; + + // Clear waiting status for ready steps + for &idx in &newly_ready_from_wait { + t.clear_waiting(idx); + } + + // Collect steps ready to launch + let mut ready = newly_ready_from_wait; + + for idx in deps_ready_indices { + let step = &definition.steps[idx]; + let wait_duration = Self::calculate_wait_duration(step); + + if let Some(duration) = wait_duration { + if duration > Duration::ZERO { + // Step needs to wait - add to waiting_until + let ready_at = now + duration; + tracing::debug!( + step_id = %step.id, + wait_ms = duration.as_millis(), + "Step waiting for delay/schedule" + ); + t.set_waiting(idx, ready_at); + added_to_waiting += 1; + continue; + } + } + ready.push(idx); + } + + (ready, added_to_waiting) }; // Check if we're done @@ -451,13 +593,23 @@ impl + 'static> WorkflowEngine { break; } - // Handle blocked workflow (no ready steps, none running, but work remains) - if ready_step_indices.is_empty() && running_count == 0 && pending_check.is_empty() { + // Handle blocked workflow (no ready steps, none running/waiting, but work remains) + // Use current_running/current_waiting from Phase 1, adjusted for Phase 2 changes. + // current_running may be stale-high (tasks completed since), but that's safe - + // we'd just do an extra loop iteration. current_waiting needs adjustment for + // steps we just added to waiting in Phase 2. + let effective_waiting = current_waiting + steps_added_to_waiting; + if ready_to_launch.is_empty() + && current_running == 0 + && effective_waiting == 0 + && pending_check.is_empty() + { let failed_step = tracker.read().failed.iter().next().cloned(); - let error_message = if failed_step.is_some() { - "Workflow failed due to step dependency failure".to_string() + // Use &'static str to avoid allocation in common error paths + let error_message: &'static str = if failed_step.is_some() { + "Workflow failed due to step dependency failure" } else { - "Workflow deadlocked: no steps ready and none running. This may indicate a scheduler bug.".to_string() + "Workflow deadlocked: no steps ready and none running" }; self.state_store @@ -470,19 +622,23 @@ impl + 'static> WorkflowEngine { instance_id, failed_step: failed_step .unwrap_or_else(|| StepId::new("internal_scheduler")), - error: error_message, + error: error_message.to_string(), }) .await; return Ok(()); } // Launch ready steps in parallel - let mut tasks_launched = 0; - for step_idx in ready_step_indices { - let step = &definition.steps[step_idx]; - tracker.write().running.insert(step.id.clone()); - tasks_launched += 1; + let tasks_launched = ready_to_launch.len(); + if tasks_launched > 0 { + let mut t = tracker.write(); + for &idx in &ready_to_launch { + t.running.insert(definition.steps[idx].id.clone()); + } + } + for step_idx in ready_to_launch { + let step = &definition.steps[step_idx]; let engine = self.clone_for_execution(); let def = Arc::clone(&definition); let step_id = step.id.clone(); @@ -491,6 +647,76 @@ impl + 'static> WorkflowEngine { tokio::spawn(async move { let step = &def.steps[step_idx]; + + // Evaluate run_if condition if present + if let Some(ref condition) = step.run_if { + match engine.state_store.get_context(instance_id).await { + Ok(ctx) => { + if !condition(&ctx) { + // Condition is false - skip this step + tracing::debug!( + step_id = %step_id, + "Step skipped due to run_if condition" + ); + + // Update tracker and send skip signal + { + let mut t = tracker.write(); + t.running.remove(&step_id); + t.skipped.insert(step_id.clone()); + } + + // Update state store + let _ = engine + .state_store + .update(instance_id, |s| { + if let Some(step_state) = + s.step_states.get_mut(&step_id) + { + step_state.status = StepStatus::Skipped; + } + }) + .await; + + // Send skip signal + let _ = tx.try_send((step_id, StepResult::Skip)); + return; + } + } + Err(e) => { + // Failed to get context - fail the step rather than proceeding blindly + tracing::error!( + step_id = %step_id, + error = ?e, + "Failed to get context for run_if evaluation, failing step" + ); + + // Update tracker + { + let mut t = tracker.write(); + t.running.remove(&step_id); + t.failed.insert(step_id.clone()); + } + + // Update state store + let _ = engine + .state_store + .update(instance_id, |s| { + if let Some(step_state) = s.step_states.get_mut(&step_id) { + step_state.status = StepStatus::Failed; + step_state.last_error = + Some(format!("run_if context error: {e}")); + } + }) + .await; + + // Send failure signal + let _ = tx.try_send((step_id, StepResult::Failure)); + return; + } + } + } + let result = engine .execute_step_with_retry(instance_id, step, &def) .await; @@ -568,8 +794,17 @@ impl + 'static> WorkflowEngine { }); } - let should_wait = tasks_launched > 0 || !tracker.read().running.is_empty(); - if should_wait { + // Single lock read for both checks + let (has_running, has_waiting) = { + let t = tracker.read(); + ( + tasks_launched > 0 || !t.running.is_empty(), + !t.waiting_until.is_empty(), + ) + }; + + if has_running { + // Wait for a step to complete; Phase 0 will drain any others if let Some((completed_step_id, result)) = rx.recv().await { tracing::debug!( step_id = %completed_step_id, @@ -585,6 +820,28 @@ impl + 'static> WorkflowEngine { } } } + } else if has_waiting { + // No running tasks, but some steps are waiting for their delay/schedule + // Calculate how long to sleep based on the soonest waiting step + let sleep_duration = { + let t = tracker.read(); + let now = std::time::Instant::now(); + t.waiting_until + .values() + .filter_map(|&ready_at| { + if ready_at > now { + Some(ready_at - now) + } else { + None + } + }) + .min() + .unwrap_or(Duration::from_millis(10)) + }; + + // Sleep until the next step is ready (with a cap to allow cancellation checks) + let capped_sleep = sleep_duration.min(Duration::from_millis(100)); + tokio::time::sleep(capped_sleep).await; } } @@ -603,7 +860,7 @@ impl + 'static> WorkflowEngine { .publish(WorkflowEvent::WorkflowFailed { instance_id, failed_step: step.clone(), - error: "One or more steps failed".to_string(), + error: "One or more steps failed".into(), }) .await; } else { diff --git a/sgl-model-gateway/tests/workflow_test.rs b/sgl-model-gateway/tests/workflow_test.rs index c1b6b08ed..a70509f8c 100644 --- a/sgl-model-gateway/tests/workflow_test.rs +++ b/sgl-model-gateway/tests/workflow_test.rs @@ -68,6 +68,18 @@ impl StepExecutor for AlwaysSucceedStep { } } +struct AlwaysFailStep; + +#[async_trait::async_trait] +impl StepExecutor for AlwaysFailStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { + Ok(StepResult::Failure) + } +} + #[tokio::test] async fn test_simple_workflow_execution() { let engine: WorkflowEngine = WorkflowEngine::new(); @@ -703,3 +715,629 @@ fn test_dag_validation_valid_workflow() { let result = workflow.validate(); assert!(result.is_ok()); } + +// ============================================================================ +// Scheduled/Delayed Steps Tests (#24) +// ============================================================================ + +#[tokio::test] +async fn test_step_delay() { + let engine: WorkflowEngine = WorkflowEngine::new(); + + // Create a workflow with a 100ms delay + let workflow = WorkflowDefinition::new("delay_workflow", "Delay Test").add_step( + StepDefinition::new("delayed_step", "Delayed Step", Arc::new(AlwaysSucceedStep)) + .with_delay(Duration::from_millis(100)), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let start = std::time::Instant::now(); + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + // Wait for completion + engine + .wait_for_completion(instance_id, "test", Duration::from_secs(5)) + .await + .unwrap(); + + let duration = start.elapsed(); + + // Verify delay was applied (should take at least 100ms) + // Note: wait_for_completion cleans up state, so we verify via timing + assert!( + duration >= Duration::from_millis(100), + "Step delay not applied, duration: {:?}", + duration + ); +} + +#[tokio::test] +async fn test_step_scheduled_at() { + use chrono::Utc; + + let engine: WorkflowEngine = WorkflowEngine::new(); + + // Schedule step to run 100ms in the future + let scheduled_time = Utc::now() + chrono::Duration::milliseconds(100); + + let workflow = WorkflowDefinition::new("scheduled_workflow", "Scheduled Test").add_step( + StepDefinition::new( + "scheduled_step", + "Scheduled Step", + Arc::new(AlwaysSucceedStep), + ) + .scheduled_at(scheduled_time), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let start = std::time::Instant::now(); + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + // Wait for completion + engine + .wait_for_completion(instance_id, "test", Duration::from_secs(5)) + .await + .unwrap(); + + let duration = start.elapsed(); + + // Verify scheduled time was respected (should take at least 100ms) + // Note: wait_for_completion cleans up state, so we verify via timing + assert!( + duration >= Duration::from_millis(100), + "Scheduled time not respected, duration: {:?}", + duration + ); +} + +// ============================================================================ +// Conditional Branching Tests (#25) +// ============================================================================ + +#[tokio::test] +async fn test_run_if_true() { + let engine: WorkflowEngine = WorkflowEngine::new(); + + let executed = Arc::new(AtomicU32::new(0)); + let executed_clone = Arc::clone(&executed); + + struct TrackingStep { + counter: Arc, + } + + #[async_trait::async_trait] + impl StepExecutor for TrackingStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { + self.counter.fetch_add(1, Ordering::SeqCst); + Ok(StepResult::Success) + } + } + + // Step with run_if that always returns true + let workflow = WorkflowDefinition::new("run_if_true_workflow", "Run If True Test").add_step( + StepDefinition::new( + "conditional_step", + "Conditional Step", + Arc::new(TrackingStep { counter: executed }), + ) + .run_if(|_ctx| true), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + engine + .wait_for_completion(instance_id, "test", Duration::from_secs(5)) + .await + .unwrap(); + + // Step should have executed (condition was true) + // Note: wait_for_completion cleans up state, so we verify via counter + assert_eq!(executed_clone.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_run_if_false() { + use tokio::time::sleep; + + let engine: WorkflowEngine = WorkflowEngine::new(); + + let executed = Arc::new(AtomicU32::new(0)); + let executed_clone = Arc::clone(&executed); + + struct TrackingStep { + counter: Arc, + } + + #[async_trait::async_trait] + impl StepExecutor for TrackingStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { + self.counter.fetch_add(1, Ordering::SeqCst); + Ok(StepResult::Success) + } + } + + // Step with run_if that always returns false + let workflow = WorkflowDefinition::new("run_if_false_workflow", "Run If False Test").add_step( + StepDefinition::new( + "conditional_step", + "Conditional Step", + Arc::new(TrackingStep { counter: executed }), + ) + .run_if(|_ctx| false), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + // Use polling to check status (don't use wait_for_completion which cleans up state) + let mut state = engine.get_status(instance_id).await.unwrap(); + for _ in 0..50 { + if state.status != WorkflowStatus::Running && state.status != WorkflowStatus::Pending { + break; + } + sleep(Duration::from_millis(50)).await; + state = engine.get_status(instance_id).await.unwrap(); + } + + assert_eq!(state.status, WorkflowStatus::Completed); + + // Step should NOT have executed (skipped due to run_if) + assert_eq!(executed_clone.load(Ordering::SeqCst), 0); + + // Verify step was marked as skipped + let step_state = state + .step_states + .get(&StepId::new("conditional_step")) + .unwrap(); + assert_eq!(step_state.status, StepStatus::Skipped); +} + +#[tokio::test] +async fn test_run_if_context_based() { + let engine: WorkflowEngine = WorkflowEngine::new(); + + // Step that sets test_key in context + struct SetKeyStep; + + #[async_trait::async_trait] + impl StepExecutor for SetKeyStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + context.data.test_key = Some("execute_next".to_string()); + Ok(StepResult::Success) + } + } + + let executed = Arc::new(AtomicU32::new(0)); + let executed_clone = Arc::clone(&executed); + + struct TrackingStep { + counter: Arc, + } + + #[async_trait::async_trait] + impl StepExecutor for TrackingStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { + self.counter.fetch_add(1, Ordering::SeqCst); + Ok(StepResult::Success) + } + } + + // Workflow where second step only runs if first step sets the right key + let workflow = WorkflowDefinition::new("context_run_if_workflow", "Context Run If Test") + .add_step(StepDefinition::new( + "set_key_step", + "Set Key", + Arc::new(SetKeyStep), + )) + .add_step( + StepDefinition::new( + "conditional_step", + "Conditional Step", + Arc::new(TrackingStep { counter: executed }), + ) + .depends_on(&["set_key_step"]) + .run_if(|ctx| ctx.data.test_key.as_deref() == Some("execute_next")), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + engine + .wait_for_completion(instance_id, "test", Duration::from_secs(5)) + .await + .unwrap(); + + // Step should have executed because context had the right value + // Note: wait_for_completion cleans up state, so we verify via counter + assert_eq!(executed_clone.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_depends_on_any() { + // DAG: A and B run in parallel, C waits for ANY (not both) + // A ──┐ + // ├──> C (any_of) + // B ──┘ + let engine: WorkflowEngine = WorkflowEngine::new(); + + let start_times: Arc>> = + Arc::new(parking_lot::RwLock::new(Vec::new())); + let end_times: Arc>> = + Arc::new(parking_lot::RwLock::new(Vec::new())); + + // A takes 50ms, B takes 200ms + // C should start after A finishes (not wait for B) + let workflow = WorkflowDefinition::new("depends_on_any_workflow", "Depends On Any Test") + .add_step(StepDefinition::new( + "step_a", + "Step A", + Arc::new(TimingStep { + step_name: "step_a".to_string(), + duration_ms: 50, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step(StepDefinition::new( + "step_b", + "Step B", + Arc::new(TimingStep { + step_name: "step_b".to_string(), + duration_ms: 200, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step( + StepDefinition::new( + "step_c", + "Step C", + Arc::new(TimingStep { + step_name: "step_c".to_string(), + duration_ms: 50, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + ) + .depends_on_any(&["step_a", "step_b"]), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + engine + .wait_for_completion(instance_id, "test", Duration::from_secs(5)) + .await + .unwrap(); + + // Verify step C started after A finished but before B finished + // Note: wait_for_completion cleans up state, so we verify via timing + let starts = start_times.read(); + let ends = end_times.read(); + + let c_start = starts.iter().find(|(n, _)| n == "step_c").unwrap().1; + let a_end = ends.iter().find(|(n, _)| n == "step_a").unwrap().1; + let b_end = ends.iter().find(|(n, _)| n == "step_b").unwrap().1; + + assert!( + c_start >= a_end, + "Step C should start after Step A finishes" + ); + assert!( + c_start < b_end, + "Step C should start before Step B finishes (any_of semantics)" + ); +} + +#[tokio::test] +async fn test_depends_on_any_combined_with_depends_on() { + // DAG: C requires ALL of [A] AND ANY of [B, D] + // A takes 50ms, B takes 100ms, D takes 200ms + // C should start after A AND (B or D) complete + let engine: WorkflowEngine = WorkflowEngine::new(); + + let start_times: Arc>> = + Arc::new(parking_lot::RwLock::new(Vec::new())); + let end_times: Arc>> = + Arc::new(parking_lot::RwLock::new(Vec::new())); + + let workflow = WorkflowDefinition::new("combined_deps_workflow", "Combined Dependencies Test") + .add_step(StepDefinition::new( + "step_a", + "Step A", + Arc::new(TimingStep { + step_name: "step_a".to_string(), + duration_ms: 50, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step(StepDefinition::new( + "step_b", + "Step B", + Arc::new(TimingStep { + step_name: "step_b".to_string(), + duration_ms: 100, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step(StepDefinition::new( + "step_d", + "Step D", + Arc::new(TimingStep { + step_name: "step_d".to_string(), + duration_ms: 200, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step( + StepDefinition::new( + "step_c", + "Step C", + Arc::new(TimingStep { + step_name: "step_c".to_string(), + duration_ms: 50, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + ) + .depends_on(&["step_a"]) // Must wait for A + .depends_on_any(&["step_b", "step_d"]), // AND any of B or D + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + engine + .wait_for_completion(instance_id, "test", Duration::from_secs(5)) + .await + .unwrap(); + + // Verify step C started after both A AND B finished + // (B finishes at 100ms, which is after A at 50ms) + // Note: wait_for_completion cleans up state, so we verify via timing + let starts = start_times.read(); + let ends = end_times.read(); + + let c_start = starts.iter().find(|(n, _)| n == "step_c").unwrap().1; + let a_end = ends.iter().find(|(n, _)| n == "step_a").unwrap().1; + let b_end = ends.iter().find(|(n, _)| n == "step_b").unwrap().1; + let d_end = ends.iter().find(|(n, _)| n == "step_d").unwrap().1; + + assert!( + c_start >= a_end, + "Step C should start after Step A (depends_on)" + ); + assert!( + c_start >= b_end || c_start >= d_end, + "Step C should start after at least one of B or D (depends_on_any)" + ); + // Since B finishes first (100ms) and A finishes before B, C should start around 100ms + assert!( + c_start < d_end, + "Step C should start before D finishes (any_of semantics)" + ); +} + +#[test] +fn test_dag_validation_depends_on_any_missing() { + // Create a workflow with a missing depends_on_any dependency + let mut workflow = WorkflowDefinition::new("missing_any_dep_workflow", "Missing Any Dep Test") + .add_step(StepDefinition::new( + "step_a", + "Step A", + Arc::new(AlwaysSucceedStep), + )) + .add_step( + StepDefinition::new("step_b", "Step B", Arc::new(AlwaysSucceedStep)) + .depends_on_any(&["nonexistent_step"]), + ); + + let result = workflow.validate(); + assert!(result.is_err()); + assert!(matches!( + result.unwrap_err(), + ValidationError::MissingDependency { .. } + )); +} + +#[tokio::test] +async fn test_depends_on_any_all_fail() { + // When ALL depends_on_any dependencies fail, the step should be blocked + // Workflow: A and B both fail, C depends_on_any([A, B]) + // Expected: C should not run, workflow should fail + use tokio::time::sleep; + + let engine: WorkflowEngine = WorkflowEngine::new(); + + let c_executed = Arc::new(AtomicU32::new(0)); + let c_executed_clone = Arc::clone(&c_executed); + + struct TrackingStep { + counter: Arc, + } + + #[async_trait::async_trait] + impl StepExecutor for TrackingStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { + self.counter.fetch_add(1, Ordering::SeqCst); + Ok(StepResult::Success) + } + } + + let workflow = WorkflowDefinition::new("all_any_fail_workflow", "All Any Fail Test") + .add_step(StepDefinition::new( + "step_a", + "Step A (fails)", + Arc::new(AlwaysFailStep), + )) + .add_step(StepDefinition::new( + "step_b", + "Step B (fails)", + Arc::new(AlwaysFailStep), + )) + .add_step( + StepDefinition::new( + "step_c", + "Step C (depends on any of A, B)", + Arc::new(TrackingStep { + counter: c_executed, + }), + ) + .depends_on_any(&["step_a", "step_b"]), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + // Use polling to check status + let mut state = engine.get_status(instance_id).await.unwrap(); + for _ in 0..50 { + if state.status != WorkflowStatus::Running && state.status != WorkflowStatus::Pending { + break; + } + sleep(Duration::from_millis(50)).await; + state = engine.get_status(instance_id).await.unwrap(); + } + + // Workflow should have failed (because all depends_on_any deps failed) + assert_eq!(state.status, WorkflowStatus::Failed); + + // Step C should NOT have executed + assert_eq!(c_executed_clone.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn test_depends_on_any_one_fails_one_succeeds() { + // When only SOME depends_on_any dependencies fail (but at least one succeeds), + // the step should still run + // Workflow: A fails, B succeeds, C depends_on_any([A, B]) + // Expected: C should run (B succeeded) + use tokio::time::sleep; + + let engine: WorkflowEngine = WorkflowEngine::new(); + + let c_executed = Arc::new(AtomicU32::new(0)); + let c_executed_clone = Arc::clone(&c_executed); + + struct TrackingStep { + counter: Arc, + } + + #[async_trait::async_trait] + impl StepExecutor for TrackingStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { + self.counter.fetch_add(1, Ordering::SeqCst); + Ok(StepResult::Success) + } + } + + // Note: Step A uses ContinueNextStep so its failure doesn't fail the workflow. + // This is the correct way to model "any of" semantics where a failing path + // shouldn't fail the entire workflow if another path succeeds. + let workflow = WorkflowDefinition::new("one_any_fail_workflow", "One Any Fail Test") + .add_step( + StepDefinition::new("step_a", "Step A (fails)", Arc::new(AlwaysFailStep)) + .with_failure_action(FailureAction::ContinueNextStep), + ) + .add_step(StepDefinition::new( + "step_b", + "Step B (succeeds)", + Arc::new(AlwaysSucceedStep), + )) + .add_step( + StepDefinition::new( + "step_c", + "Step C (depends on any of A, B)", + Arc::new(TrackingStep { + counter: c_executed, + }), + ) + .depends_on_any(&["step_a", "step_b"]), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, TestWorkflowData::default()) + .await + .unwrap(); + + // Use polling to check status + let mut state = engine.get_status(instance_id).await.unwrap(); + for _ in 0..50 { + if state.status != WorkflowStatus::Running && state.status != WorkflowStatus::Pending { + break; + } + sleep(Duration::from_millis(50)).await; + state = engine.get_status(instance_id).await.unwrap(); + } + + // Workflow should have completed (B succeeded, so C could run) + assert_eq!(state.status, WorkflowStatus::Completed); + + // Step C SHOULD have executed (because B succeeded) + assert_eq!(c_executed_clone.load(Ordering::SeqCst), 1); +}