1use std::collections::HashMap;
10use std::fmt;
11use std::sync::Arc;
12use std::time::Instant;
13
14use chrono::{DateTime, Utc};
15use rust_decimal::Decimal;
16use serde_json::Value;
17use tracing::{error, info, warn};
18use uuid::Uuid;
19
20#[cfg(feature = "prometheus")]
21use ironflow_core::metric_names::{
22 RUN_BUDGET_EXCEEDED_TOTAL, RUN_COST_USD, RUN_DURATION_SECONDS, RUNS_ACTIVE, RUNS_TOTAL,
23};
24use ironflow_core::provider::AgentProvider;
25use ironflow_store::error::StoreError;
26use ironflow_store::models::{
27 NewRun, Run, RunCreation, RunFilter, RunStatus, RunUpdate, StepStatus, StepUpdate, TriggerKind,
28};
29use ironflow_store::store::Store;
30#[cfg(feature = "prometheus")]
31use metrics::{counter, gauge, histogram};
32
33use crate::budget::{BudgetConfig, month_start};
34use crate::context::WorkflowContext;
35use crate::error::EngineError;
36use crate::handler::{WorkflowHandler, WorkflowInfo};
37use crate::log_sender::LogSender;
38use crate::notify::{Event, EventPublisher, EventSubscriber};
39use crate::schedule::CronSchedule;
40
41#[derive(Debug, Clone, Default)]
60pub struct EnqueueOptions {
61 pub max_retries: u32,
63 pub labels: HashMap<String, String>,
65 pub scheduled_at: Option<DateTime<Utc>>,
68 pub max_cost_usd: Option<Decimal>,
72 pub idempotency_key: Option<String>,
78}
79
80pub struct Engine {
121 store: Arc<dyn Store>,
122 provider: Arc<dyn AgentProvider>,
123 handlers: HashMap<String, Arc<dyn WorkflowHandler>>,
124 event_publisher: EventPublisher,
125 log_sender: Option<LogSender>,
126 budget: BudgetConfig,
127}
128
129fn validate_category(handler_name: &str, category: &str) -> Result<(), EngineError> {
139 let reject = |reason: &str| {
140 Err(EngineError::InvalidWorkflow(format!(
141 "handler '{handler_name}' has invalid category '{category}': {reason}"
142 )))
143 };
144
145 if category.is_empty() {
146 return reject("empty category");
147 }
148 if category.starts_with('/') {
149 return reject("leading '/'");
150 }
151 if category.ends_with('/') {
152 return reject("trailing '/'");
153 }
154 for segment in category.split('/') {
155 if segment.is_empty() {
156 return reject("empty segment (double '/')");
157 }
158 if segment.trim().is_empty() {
159 return reject("whitespace-only segment");
160 }
161 }
162 Ok(())
163}
164
165impl Engine {
166 pub fn new(store: Arc<dyn Store>, provider: Arc<dyn AgentProvider>) -> Self {
182 Self {
183 store,
184 provider,
185 handlers: HashMap::new(),
186 event_publisher: EventPublisher::new(),
187 log_sender: None,
188 budget: BudgetConfig::new(),
189 }
190 }
191
192 pub fn with_budget_config(mut self, budget: BudgetConfig) -> Self {
213 self.budget = budget;
214 self
215 }
216
217 pub fn budget_config(&self) -> &BudgetConfig {
219 &self.budget
220 }
221
222 pub fn set_log_sender(&mut self, sender: LogSender) {
228 self.log_sender = Some(sender);
229 }
230
231 pub fn store(&self) -> &Arc<dyn Store> {
233 &self.store
234 }
235
236 pub fn provider(&self) -> &Arc<dyn AgentProvider> {
238 &self.provider
239 }
240
241 fn build_context(&self, run_id: Uuid, max_cost_usd: Option<Decimal>) -> WorkflowContext {
246 let handlers = self.handlers.clone();
247 let resolver: crate::context::HandlerResolver =
248 Arc::new(move |name: &str| handlers.get(name).cloned());
249 let mut ctx = WorkflowContext::with_handler_resolver(
250 run_id,
251 self.store.clone(),
252 self.provider.clone(),
253 resolver,
254 );
255 ctx.set_max_cost_usd(max_cost_usd);
256 if let Some(ref sender) = self.log_sender {
257 ctx.set_log_sender(sender.clone());
258 }
259 ctx
260 }
261
262 async fn check_monthly_quota(&self, workflow_name: &str) -> Result<(), EngineError> {
273 let Some(limit) = self.budget.monthly_cost_limit_usd else {
274 return Ok(());
275 };
276
277 let stats = self
278 .store
279 .get_stats(RunFilter {
280 created_after: Some(month_start(Utc::now())),
281 ..RunFilter::default()
282 })
283 .await?;
284
285 if stats.total_cost_usd < limit {
286 return Ok(());
287 }
288
289 warn!(
290 workflow = %workflow_name,
291 limit_usd = %limit,
292 spent_usd = %stats.total_cost_usd,
293 "monthly cost quota exhausted, refusing new run"
294 );
295
296 #[cfg(feature = "prometheus")]
297 counter!(
298 RUN_BUDGET_EXCEEDED_TOTAL,
299 "workflow" => workflow_name.to_string(),
300 "scope" => "monthly",
301 )
302 .increment(1);
303
304 Err(EngineError::MonthlyBudgetExceeded {
305 limit_usd: limit,
306 spent_usd: stats.total_cost_usd,
307 })
308 }
309
310 pub fn register(&mut self, handler: impl WorkflowHandler + 'static) -> Result<(), EngineError> {
354 let name = handler.name().to_string();
355 if self.handlers.contains_key(&name) {
356 return Err(EngineError::InvalidWorkflow(format!(
357 "handler '{}' already registered",
358 name
359 )));
360 }
361 if let Some(category) = handler.category() {
362 validate_category(&name, category)?;
363 }
364 self.handlers.insert(name, Arc::new(handler));
365 Ok(())
366 }
367
368 pub fn register_boxed(&mut self, handler: Box<dyn WorkflowHandler>) -> Result<(), EngineError> {
375 let name = handler.name().to_string();
376 if self.handlers.contains_key(&name) {
377 return Err(EngineError::InvalidWorkflow(format!(
378 "handler '{}' already registered",
379 name
380 )));
381 }
382 if let Some(category) = handler.category() {
383 validate_category(&name, category)?;
384 }
385 self.handlers.insert(name, Arc::from(handler));
386 Ok(())
387 }
388
389 pub fn get_handler(&self, name: &str) -> Option<&Arc<dyn WorkflowHandler>> {
391 self.handlers.get(name)
392 }
393
394 pub fn handler_names(&self) -> Vec<&str> {
396 self.handlers.keys().map(|s| s.as_str()).collect()
397 }
398
399 pub fn handler_info(&self, name: &str) -> Option<WorkflowInfo> {
401 self.handlers.get(name).map(|h| h.describe())
402 }
403
404 pub fn scheduled_handlers(&self) -> Vec<(&str, &CronSchedule)> {
428 self.handlers
429 .iter()
430 .filter_map(|(name, handler)| handler.schedule().map(|sched| (name.as_str(), sched)))
431 .collect()
432 }
433
434 pub fn subscribe(
459 &mut self,
460 subscriber: impl EventSubscriber + 'static,
461 event_types: &[&'static str],
462 ) {
463 self.event_publisher.subscribe(subscriber, event_types);
464 }
465
466 pub fn event_publisher(&self) -> &EventPublisher {
471 &self.event_publisher
472 }
473
474 #[tracing::instrument(name = "engine.run_handler", skip_all, fields(workflow = %handler_name))]
504 pub async fn run_handler(
505 &self,
506 handler_name: &str,
507 trigger: TriggerKind,
508 payload: Value,
509 ) -> Result<Run, EngineError> {
510 let handler = self
511 .handlers
512 .get(handler_name)
513 .ok_or_else(|| {
514 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
515 })?
516 .clone();
517
518 self.check_monthly_quota(handler_name).await?;
519
520 let handler_version = handler.version().map(str::to_string);
521 let max_cost_usd = self
522 .budget
523 .resolve_run_cap(None, handler.default_max_cost_usd());
524 let run = self
525 .store
526 .create_run(NewRun {
527 workflow_name: handler_name.to_string(),
528 trigger,
529 payload,
530 max_retries: 0,
531 handler_version,
532 labels: handler.default_labels(),
533 scheduled_at: None,
534 idempotency_key: None,
535 max_cost_usd,
536 })
537 .await?
538 .into_run();
539
540 let run_id = run.id;
541 info!(run_id = %run_id, handler_version = run.handler_version.as_deref().unwrap_or(""), "run created");
542
543 self.store
544 .update_run_status(run_id, RunStatus::Running)
545 .await?;
546
547 #[cfg(feature = "prometheus")]
548 gauge!(RUNS_ACTIVE, "workflow" => handler_name.to_string()).increment(1.0);
549
550 let run_start = Instant::now();
551 let mut ctx = self.build_context(run_id, run.max_cost_usd);
552
553 let result = handler.execute(&mut ctx).await;
554 self.finalize_run(run_id, handler_name, result, &ctx, run_start)
555 .await
556 }
557
558 #[tracing::instrument(name = "engine.enqueue_handler", skip_all, fields(workflow = %handler_name))]
569 pub async fn enqueue_handler(
570 &self,
571 handler_name: &str,
572 trigger: TriggerKind,
573 payload: Value,
574 max_retries: u32,
575 ) -> Result<Run, EngineError> {
576 self.enqueue_handler_with_options(
577 handler_name,
578 trigger,
579 payload,
580 EnqueueOptions {
581 max_retries,
582 ..Default::default()
583 },
584 )
585 .await
586 .map(RunCreation::into_run)
587 }
588
589 #[tracing::instrument(name = "engine.enqueue_handler_with_options", skip_all, fields(workflow = %handler_name))]
634 pub async fn enqueue_handler_with_options(
635 &self,
636 handler_name: &str,
637 trigger: TriggerKind,
638 payload: Value,
639 options: EnqueueOptions,
640 ) -> Result<RunCreation, EngineError> {
641 let EnqueueOptions {
642 max_retries,
643 labels,
644 scheduled_at,
645 max_cost_usd,
646 idempotency_key,
647 } = options;
648
649 let handler = self.handlers.get(handler_name).ok_or_else(|| {
650 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
651 })?;
652
653 self.check_monthly_quota(handler_name).await?;
654
655 let handler_version = handler.version().map(str::to_string);
656 let mut merged_labels = handler.default_labels();
657 merged_labels.extend(labels);
658 let resolved_cap = self
659 .budget
660 .resolve_run_cap(max_cost_usd, handler.default_max_cost_usd());
661
662 let creation = self
663 .store
664 .create_run(NewRun {
665 workflow_name: handler_name.to_string(),
666 trigger,
667 payload,
668 max_retries,
669 handler_version,
670 labels: merged_labels,
671 scheduled_at,
672 idempotency_key,
673 max_cost_usd: resolved_cap,
674 })
675 .await?;
676
677 match &creation {
678 RunCreation::Created(run) => info!(
679 run_id = %run.id,
680 workflow = %handler_name,
681 max_cost_usd = ?resolved_cap,
682 "handler run enqueued"
683 ),
684 RunCreation::Existing(run) => info!(
685 run_id = %run.id,
686 workflow = %handler_name,
687 "idempotent replay, nothing enqueued"
688 ),
689 }
690
691 Ok(creation)
692 }
693
694 #[tracing::instrument(name = "engine.execute_handler_run", skip_all, fields(run_id = %run_id))]
703 pub async fn execute_handler_run(&self, run_id: Uuid) -> Result<Run, EngineError> {
704 let run = self
705 .store
706 .get_run(run_id)
707 .await?
708 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
709
710 let handler = self
711 .handlers
712 .get(&run.workflow_name)
713 .ok_or_else(|| {
714 EngineError::InvalidWorkflow(format!(
715 "no handler registered: {}",
716 run.workflow_name
717 ))
718 })?
719 .clone();
720
721 #[cfg(feature = "prometheus")]
722 gauge!(RUNS_ACTIVE, "workflow" => run.workflow_name.clone()).increment(1.0);
723
724 let run_start = Instant::now();
725 let mut ctx = self.build_context(run_id, run.max_cost_usd);
726
727 let result = handler.execute(&mut ctx).await;
728 self.finalize_run(run_id, &run.workflow_name, result, &ctx, run_start)
729 .await
730 }
731
732 #[tracing::instrument(name = "engine.execute_run", skip_all, fields(run_id = %run_id))]
740 pub async fn execute_run(&self, run_id: Uuid) -> Result<Run, EngineError> {
741 self.execute_handler_run(run_id).await
742 }
743
744 #[tracing::instrument(name = "engine.resume_run", skip_all, fields(run_id = %run_id))]
758 pub async fn resume_run(&self, run_id: Uuid) -> Result<Run, EngineError> {
759 let run = self
760 .store
761 .get_run(run_id)
762 .await?
763 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
764
765 let handler = self
766 .handlers
767 .get(&run.workflow_name)
768 .ok_or_else(|| {
769 EngineError::InvalidWorkflow(format!(
770 "no handler registered: {}",
771 run.workflow_name
772 ))
773 })?
774 .clone();
775
776 info!(run_id = %run_id, workflow = %run.workflow_name, "resuming run after approval");
777
778 let run_start = Instant::now();
779 let mut ctx = self.build_context(run_id, run.max_cost_usd);
780 ctx.load_replay_steps().await?;
781
782 let result = handler.execute(&mut ctx).await;
783 self.finalize_run(run_id, &run.workflow_name, result, &ctx, run_start)
784 .await
785 }
786
787 pub async fn fail_orphaned_steps(
801 &self,
802 run_id: Uuid,
803 error_message: &str,
804 ) -> Result<(), EngineError> {
805 let steps = self.store.list_steps(run_id).await?;
806 let now = Utc::now();
807
808 for step in steps {
809 if step.status.state.is_terminal() {
810 continue;
811 }
812
813 let (target_status, error) = match step.status.state {
814 StepStatus::Running | StepStatus::AwaitingApproval => {
815 let err = if step.error.is_some() {
816 None
817 } else {
818 Some(error_message.to_string())
819 };
820 (StepStatus::Failed, err)
821 }
822 StepStatus::Pending => (StepStatus::Skipped, None),
823 _ => continue,
824 };
825
826 if let Err(e) = self
827 .store
828 .update_step(
829 step.id,
830 StepUpdate {
831 status: Some(target_status),
832 error,
833 completed_at: Some(now),
834 ..StepUpdate::default()
835 },
836 )
837 .await
838 {
839 warn!(
840 run_id = %run_id,
841 step_id = %step.id,
842 step_name = %step.name,
843 error = %e,
844 "failed to cleanup orphaned step"
845 );
846 } else {
847 info!(
848 run_id = %run_id,
849 step_id = %step.id,
850 step_name = %step.name,
851 from = %step.status.state,
852 to = %target_status,
853 "cleaned up orphaned step"
854 );
855 }
856 }
857
858 Ok(())
859 }
860
861 async fn finalize_run(
867 &self,
868 run_id: Uuid,
869 workflow_name: &str,
870 result: Result<(), EngineError>,
871 ctx: &WorkflowContext,
872 run_start: Instant,
873 ) -> Result<Run, EngineError> {
874 let total_duration = run_start.elapsed().as_millis() as u64;
875 let completed_at = Utc::now();
876
877 let final_status;
878 let final_run;
879
880 match result {
881 Ok(()) => {
882 final_status = RunStatus::Completed;
883 final_run = self
884 .store
885 .update_run_returning(
886 run_id,
887 RunUpdate {
888 status: Some(RunStatus::Completed),
889 cost_usd: Some(ctx.total_cost_usd()),
890 duration_ms: Some(total_duration),
891 completed_at: Some(completed_at),
892 ..RunUpdate::default()
893 },
894 )
895 .await?;
896
897 info!(
898 run_id = %run_id,
899 cost_usd = %ctx.total_cost_usd(),
900 duration_ms = total_duration,
901 "run completed"
902 );
903 }
904 Err(EngineError::ApprovalRequired {
905 run_id: approval_run_id,
906 step_id,
907 ref message,
908 }) => {
909 final_status = RunStatus::AwaitingApproval;
910 final_run = self
911 .store
912 .update_run_returning(
913 run_id,
914 RunUpdate {
915 status: Some(RunStatus::AwaitingApproval),
916 cost_usd: Some(ctx.total_cost_usd()),
917 duration_ms: Some(total_duration),
918 ..RunUpdate::default()
919 },
920 )
921 .await?;
922
923 info!(
924 run_id = %approval_run_id,
925 step_id = %step_id,
926 message = %message,
927 "run awaiting approval"
928 );
929 }
930 Err(err) => {
931 let budget_exceeded = matches!(err, EngineError::RunBudgetExceeded { .. });
934 final_status = if budget_exceeded {
935 RunStatus::Cancelled
936 } else {
937 RunStatus::Failed
938 };
939
940 if let Err(store_err) = self
941 .store
942 .update_run(
943 run_id,
944 RunUpdate {
945 status: Some(final_status),
946 error: Some(err.to_string()),
947 cost_usd: Some(ctx.total_cost_usd()),
948 duration_ms: Some(total_duration),
949 completed_at: Some(completed_at),
950 ..RunUpdate::default()
951 },
952 )
953 .await
954 {
955 error!(run_id = %run_id, store_error = %store_err, "failed to persist run failure");
956 }
957
958 if budget_exceeded {
959 self.on_run_budget_exceeded(workflow_name, run_id, &err);
960 }
961
962 error!(run_id = %run_id, status = %final_status, error = %err, "run stopped");
963
964 self.publish_run_status_changed(
965 workflow_name,
966 run_id,
967 final_status,
968 Some(err.to_string()),
969 ctx,
970 total_duration,
971 );
972
973 #[cfg(feature = "prometheus")]
974 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
975
976 return Err(err);
977 }
978 }
979
980 self.publish_run_status_changed(
981 workflow_name,
982 run_id,
983 final_status,
984 None,
985 ctx,
986 total_duration,
987 );
988
989 #[cfg(feature = "prometheus")]
990 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
991
992 Ok(final_run)
993 }
994
995 #[cfg(feature = "prometheus")]
997 fn emit_run_metrics(
998 &self,
999 workflow_name: &str,
1000 status: RunStatus,
1001 duration_ms: u64,
1002 ctx: &WorkflowContext,
1003 ) {
1004 let status_str = status.to_string();
1005 let wf = workflow_name.to_string();
1006
1007 counter!(RUNS_TOTAL, "workflow" => wf.clone(), "status" => status_str.clone()).increment(1);
1008 histogram!(RUN_DURATION_SECONDS, "workflow" => wf.clone(), "status" => status_str)
1009 .record(duration_ms as f64 / 1000.0);
1010 histogram!(RUN_COST_USD, "workflow" => wf.clone()).record(
1011 ctx.total_cost_usd()
1012 .to_string()
1013 .parse::<f64>()
1014 .unwrap_or(0.0),
1015 );
1016 gauge!(RUNS_ACTIVE, "workflow" => wf).decrement(1.0);
1017 }
1018
1019 fn on_run_budget_exceeded(&self, workflow_name: &str, run_id: Uuid, err: &EngineError) {
1025 let EngineError::RunBudgetExceeded {
1026 limit_usd,
1027 spent_usd,
1028 step_budget_usd,
1029 ..
1030 } = err
1031 else {
1032 return;
1033 };
1034
1035 #[cfg(feature = "prometheus")]
1036 counter!(
1037 RUN_BUDGET_EXCEEDED_TOTAL,
1038 "workflow" => workflow_name.to_string(),
1039 "scope" => "run",
1040 )
1041 .increment(1);
1042
1043 self.event_publisher.publish(Event::RunBudgetExceeded {
1044 run_id,
1045 workflow_name: workflow_name.to_string(),
1046 limit_usd: *limit_usd,
1047 spent_usd: *spent_usd,
1048 step_budget_usd: *step_budget_usd,
1049 at: Utc::now(),
1050 });
1051 }
1052
1053 fn publish_run_status_changed(
1058 &self,
1059 workflow_name: &str,
1060 run_id: Uuid,
1061 to: RunStatus,
1062 error: Option<String>,
1063 ctx: &WorkflowContext,
1064 duration_ms: u64,
1065 ) {
1066 let now = Utc::now();
1067 let cost_usd = ctx.total_cost_usd();
1068 let wf = workflow_name.to_string();
1069
1070 self.event_publisher.publish(Event::RunStatusChanged {
1071 run_id,
1072 workflow_name: wf.clone(),
1073 from: RunStatus::Running,
1074 to,
1075 error: error.clone(),
1076 cost_usd,
1077 duration_ms,
1078 at: now,
1079 });
1080
1081 if to == RunStatus::Failed {
1082 self.event_publisher.publish(Event::RunFailed {
1083 run_id,
1084 workflow_name: wf,
1085 error,
1086 cost_usd,
1087 duration_ms,
1088 at: now,
1089 });
1090 }
1091 }
1092}
1093
1094impl fmt::Debug for Engine {
1095 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1096 f.debug_struct("Engine")
1097 .field("handlers", &self.handlers.keys().collect::<Vec<_>>())
1098 .finish_non_exhaustive()
1099 }
1100}
1101
1102#[cfg(test)]
1103mod tests {
1104 use super::*;
1105 use crate::config::ShellConfig;
1106 use crate::handler::{HandlerFuture, WorkflowHandler};
1107 use ironflow_core::providers::claude::ClaudeCodeProvider;
1108 use ironflow_core::providers::record_replay::RecordReplayProvider;
1109 use ironflow_store::memory::InMemoryStore;
1110 use ironflow_store::models::StepStatus;
1111 use serde_json::json;
1112
1113 struct EchoWorkflow;
1115
1116 impl WorkflowHandler for EchoWorkflow {
1117 fn name(&self) -> &str {
1118 "echo-workflow"
1119 }
1120
1121 fn describe(&self) -> WorkflowInfo {
1122 WorkflowInfo {
1123 description: "A simple workflow that echoes hello".to_string(),
1124 source_code: None,
1125 sub_workflows: Vec::new(),
1126 category: None,
1127 version: self.version().map(str::to_string),
1128 input_schema: None,
1129 default_labels: HashMap::new(),
1130 schedule: self.schedule().cloned(),
1131 default_max_cost_usd: self.default_max_cost_usd(),
1132 }
1133 }
1134
1135 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1136 Box::pin(async move {
1137 ctx.shell("greet", ShellConfig::new("echo hello")).await?;
1138 Ok(())
1139 })
1140 }
1141 }
1142
1143 struct FailingWorkflow;
1145
1146 impl WorkflowHandler for FailingWorkflow {
1147 fn name(&self) -> &str {
1148 "failing-workflow"
1149 }
1150
1151 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1152 Box::pin(async move {
1153 ctx.shell("fail", ShellConfig::new("exit 1")).await?;
1154 Ok(())
1155 })
1156 }
1157 }
1158
1159 fn create_test_engine() -> Engine {
1160 let store = Arc::new(InMemoryStore::new());
1161 let inner = ClaudeCodeProvider::new();
1162 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
1163 inner,
1164 "/tmp/ironflow-fixtures",
1165 ));
1166 Engine::new(store, provider)
1167 }
1168
1169 #[test]
1170 fn engine_new_creates_instance() {
1171 let engine = create_test_engine();
1172 assert_eq!(engine.handler_names().len(), 0);
1173 }
1174
1175 #[test]
1176 fn engine_register_handler() {
1177 let mut engine = create_test_engine();
1178 let result = engine.register(EchoWorkflow);
1179 assert!(result.is_ok());
1180 assert_eq!(engine.handler_names().len(), 1);
1181 assert!(engine.handler_names().contains(&"echo-workflow"));
1182 }
1183
1184 #[test]
1185 fn engine_register_duplicate_returns_error() {
1186 let mut engine = create_test_engine();
1187 engine.register(EchoWorkflow).unwrap();
1188 let result = engine.register(EchoWorkflow);
1189 assert!(result.is_err());
1190 }
1191
1192 #[test]
1193 fn engine_get_handler_found() {
1194 let mut engine = create_test_engine();
1195 engine.register(EchoWorkflow).unwrap();
1196 let handler = engine.get_handler("echo-workflow");
1197 assert!(handler.is_some());
1198 }
1199
1200 #[test]
1201 fn engine_get_handler_not_found() {
1202 let engine = create_test_engine();
1203 let handler = engine.get_handler("nonexistent");
1204 assert!(handler.is_none());
1205 }
1206
1207 #[test]
1208 fn engine_handler_names_lists_all() {
1209 let mut engine = create_test_engine();
1210 engine.register(EchoWorkflow).unwrap();
1211 engine.register(FailingWorkflow).unwrap();
1212 let names = engine.handler_names();
1213 assert_eq!(names.len(), 2);
1214 assert!(names.contains(&"echo-workflow"));
1215 assert!(names.contains(&"failing-workflow"));
1216 }
1217
1218 #[test]
1219 fn engine_handler_info_returns_description() {
1220 let mut engine = create_test_engine();
1221 engine.register(EchoWorkflow).unwrap();
1222 let info = engine.handler_info("echo-workflow");
1223 assert!(info.is_some());
1224 let info = info.unwrap();
1225 assert_eq!(info.description, "A simple workflow that echoes hello");
1226 }
1227
1228 struct CategorizedWorkflow;
1229
1230 impl WorkflowHandler for CategorizedWorkflow {
1231 fn name(&self) -> &str {
1232 "categorized"
1233 }
1234 fn category(&self) -> Option<&str> {
1235 Some("data/etl")
1236 }
1237 fn execute<'a>(
1238 &'a self,
1239 _ctx: &'a mut WorkflowContext,
1240 ) -> crate::handler::HandlerFuture<'a> {
1241 Box::pin(async move { Ok(()) })
1242 }
1243 }
1244
1245 #[test]
1246 fn engine_default_describe_propagates_category() {
1247 let mut engine = create_test_engine();
1248 engine.register(CategorizedWorkflow).unwrap();
1249 let info = engine.handler_info("categorized").unwrap();
1250 assert_eq!(info.category.as_deref(), Some("data/etl"));
1251 }
1252
1253 #[test]
1254 fn engine_default_describe_without_category() {
1255 let mut engine = create_test_engine();
1256 engine.register(EchoWorkflow).unwrap();
1257 let info = engine.handler_info("echo-workflow").unwrap();
1258 assert!(info.category.is_none());
1259 }
1260
1261 struct ScheduledWorkflow {
1266 schedule: CronSchedule,
1267 }
1268
1269 impl ScheduledWorkflow {
1270 fn new() -> Self {
1271 Self {
1272 schedule: CronSchedule::new("0 0 * * * *").unwrap(),
1273 }
1274 }
1275 }
1276
1277 impl WorkflowHandler for ScheduledWorkflow {
1278 fn name(&self) -> &str {
1279 "scheduled"
1280 }
1281 fn schedule(&self) -> Option<&CronSchedule> {
1282 Some(&self.schedule)
1283 }
1284 fn execute<'a>(
1285 &'a self,
1286 _ctx: &'a mut WorkflowContext,
1287 ) -> crate::handler::HandlerFuture<'a> {
1288 Box::pin(async move { Ok(()) })
1289 }
1290 }
1291
1292 #[test]
1293 fn engine_default_describe_propagates_schedule() {
1294 let mut engine = create_test_engine();
1295 engine.register(ScheduledWorkflow::new()).unwrap();
1296 let info = engine.handler_info("scheduled").unwrap();
1297 assert_eq!(
1298 info.schedule.as_ref().map(|s| s.as_str()),
1299 Some("0 0 * * * *")
1300 );
1301 }
1302
1303 #[test]
1304 fn engine_default_describe_without_schedule() {
1305 let mut engine = create_test_engine();
1306 engine.register(EchoWorkflow).unwrap();
1307 let info = engine.handler_info("echo-workflow").unwrap();
1308 assert!(info.schedule.is_none());
1309 }
1310
1311 #[test]
1312 fn scheduled_handlers_returns_only_scheduled() {
1313 let mut engine = create_test_engine();
1314 engine.register(EchoWorkflow).unwrap();
1315 engine.register(ScheduledWorkflow::new()).unwrap();
1316 engine.register(FailingWorkflow).unwrap();
1317
1318 let scheduled = engine.scheduled_handlers();
1319 assert_eq!(scheduled.len(), 1);
1320 assert_eq!(scheduled[0].0, "scheduled");
1321 assert_eq!(scheduled[0].1.as_str(), "0 0 * * * *");
1322 }
1323
1324 #[test]
1325 fn scheduled_handlers_empty_when_none_scheduled() {
1326 let mut engine = create_test_engine();
1327 engine.register(EchoWorkflow).unwrap();
1328 engine.register(FailingWorkflow).unwrap();
1329
1330 let scheduled = engine.scheduled_handlers();
1331 assert!(scheduled.is_empty());
1332 }
1333
1334 struct BadCategoryWorkflow(&'static str);
1335
1336 impl WorkflowHandler for BadCategoryWorkflow {
1337 fn name(&self) -> &str {
1338 "bad-category"
1339 }
1340 fn category(&self) -> Option<&str> {
1341 Some(self.0)
1342 }
1343 fn execute<'a>(
1344 &'a self,
1345 _ctx: &'a mut WorkflowContext,
1346 ) -> crate::handler::HandlerFuture<'a> {
1347 Box::pin(async move { Ok(()) })
1348 }
1349 }
1350
1351 #[test]
1352 fn engine_register_rejects_empty_category() {
1353 let mut engine = create_test_engine();
1354 let err = engine.register(BadCategoryWorkflow("")).unwrap_err();
1355 match err {
1356 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty category")),
1357 other => panic!("expected InvalidWorkflow, got {other:?}"),
1358 }
1359 }
1360
1361 #[test]
1362 fn engine_register_rejects_leading_slash_category() {
1363 let mut engine = create_test_engine();
1364 let err = engine
1365 .register(BadCategoryWorkflow("/data/etl"))
1366 .unwrap_err();
1367 match err {
1368 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("leading '/'")),
1369 other => panic!("expected InvalidWorkflow, got {other:?}"),
1370 }
1371 }
1372
1373 #[test]
1374 fn engine_register_rejects_trailing_slash_category() {
1375 let mut engine = create_test_engine();
1376 let err = engine
1377 .register(BadCategoryWorkflow("data/etl/"))
1378 .unwrap_err();
1379 match err {
1380 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("trailing '/'")),
1381 other => panic!("expected InvalidWorkflow, got {other:?}"),
1382 }
1383 }
1384
1385 #[test]
1386 fn engine_register_rejects_double_slash_category() {
1387 let mut engine = create_test_engine();
1388 let err = engine
1389 .register(BadCategoryWorkflow("data//etl"))
1390 .unwrap_err();
1391 match err {
1392 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty segment")),
1393 other => panic!("expected InvalidWorkflow, got {other:?}"),
1394 }
1395 }
1396
1397 #[test]
1398 fn engine_register_rejects_whitespace_only_segment_category() {
1399 let mut engine = create_test_engine();
1400 let err = engine
1401 .register(BadCategoryWorkflow("data/ /etl"))
1402 .unwrap_err();
1403 match err {
1404 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("whitespace-only segment")),
1405 other => panic!("expected InvalidWorkflow, got {other:?}"),
1406 }
1407 }
1408
1409 #[test]
1410 fn engine_register_accepts_valid_nested_category() {
1411 let mut engine = create_test_engine();
1412 assert!(engine.register(CategorizedWorkflow).is_ok());
1413 }
1414
1415 #[tokio::test]
1416 async fn engine_unknown_workflow_returns_error() {
1417 let engine = create_test_engine();
1418 let result = engine
1419 .run_handler("unknown", TriggerKind::Manual, json!({}))
1420 .await;
1421 assert!(result.is_err());
1422 match result {
1423 Err(EngineError::InvalidWorkflow(msg)) => {
1424 assert!(msg.contains("no handler registered"));
1425 }
1426 _ => panic!("expected InvalidWorkflow error"),
1427 }
1428 }
1429
1430 #[tokio::test]
1431 async fn engine_enqueue_handler_creates_pending_run() {
1432 let mut engine = create_test_engine();
1433 engine.register(EchoWorkflow).unwrap();
1434
1435 let run = engine
1436 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
1437 .await
1438 .unwrap();
1439 assert_eq!(run.status.state, RunStatus::Pending);
1440 assert_eq!(run.workflow_name, "echo-workflow");
1441 }
1442
1443 #[tokio::test]
1444 async fn engine_register_boxed() {
1445 let mut engine = create_test_engine();
1446 let handler: Box<dyn WorkflowHandler> = Box::new(EchoWorkflow);
1447 let result = engine.register_boxed(handler);
1448 assert!(result.is_ok());
1449 assert_eq!(engine.handler_names().len(), 1);
1450 }
1451
1452 #[tokio::test]
1453 async fn engine_store_and_provider_accessors() {
1454 let store = Arc::new(InMemoryStore::new());
1455 let inner = ClaudeCodeProvider::new();
1456 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
1457 inner,
1458 "/tmp/ironflow-fixtures",
1459 ));
1460 let engine = Engine::new(store.clone(), provider.clone());
1461
1462 let _ = engine.store();
1464 let _ = engine.provider();
1465 }
1466
1467 use crate::operation::Operation;
1472 use ironflow_store::models::StepKind;
1473 use std::future::Future;
1474 use std::pin::Pin;
1475
1476 struct FakeGitlabOp {
1477 project_id: u64,
1478 title: String,
1479 }
1480
1481 impl Operation for FakeGitlabOp {
1482 fn kind(&self) -> &str {
1483 "gitlab"
1484 }
1485
1486 fn execute(&self) -> Pin<Box<dyn Future<Output = Result<Value, EngineError>> + Send + '_>> {
1487 Box::pin(async move {
1488 Ok(json!({
1489 "issue_id": 42,
1490 "project_id": self.project_id,
1491 "title": self.title,
1492 }))
1493 })
1494 }
1495
1496 fn input(&self) -> Option<Value> {
1497 Some(json!({
1498 "project_id": self.project_id,
1499 "title": self.title,
1500 }))
1501 }
1502 }
1503
1504 struct FailingOp;
1505
1506 impl Operation for FailingOp {
1507 fn kind(&self) -> &str {
1508 "broken-service"
1509 }
1510
1511 fn execute(&self) -> Pin<Box<dyn Future<Output = Result<Value, EngineError>> + Send + '_>> {
1512 Box::pin(async move { Err(EngineError::StepConfig("service unavailable".to_string())) })
1513 }
1514 }
1515
1516 struct OperationWorkflow;
1517
1518 impl WorkflowHandler for OperationWorkflow {
1519 fn name(&self) -> &str {
1520 "operation-workflow"
1521 }
1522
1523 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1524 Box::pin(async move {
1525 let op = FakeGitlabOp {
1526 project_id: 123,
1527 title: "Bug report".to_string(),
1528 };
1529 ctx.operation("create-issue", &op).await?;
1530 Ok(())
1531 })
1532 }
1533 }
1534
1535 struct FailingOperationWorkflow;
1536
1537 impl WorkflowHandler for FailingOperationWorkflow {
1538 fn name(&self) -> &str {
1539 "failing-operation-workflow"
1540 }
1541
1542 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1543 Box::pin(async move {
1544 ctx.operation("broken-call", &FailingOp).await?;
1545 Ok(())
1546 })
1547 }
1548 }
1549
1550 struct MixedWorkflow;
1551
1552 impl WorkflowHandler for MixedWorkflow {
1553 fn name(&self) -> &str {
1554 "mixed-workflow"
1555 }
1556
1557 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1558 Box::pin(async move {
1559 ctx.shell("build", ShellConfig::new("echo built")).await?;
1560 let op = FakeGitlabOp {
1561 project_id: 456,
1562 title: "Deploy done".to_string(),
1563 };
1564 let result = ctx.operation("notify-gitlab", &op).await?;
1565 assert_eq!(result.output["issue_id"], 42);
1566 Ok(())
1567 })
1568 }
1569 }
1570
1571 #[tokio::test]
1572 async fn operation_step_happy_path() {
1573 let mut engine = create_test_engine();
1574 engine.register(OperationWorkflow).unwrap();
1575
1576 let run = engine
1577 .run_handler("operation-workflow", TriggerKind::Manual, json!({}))
1578 .await
1579 .unwrap();
1580
1581 assert_eq!(run.status.state, RunStatus::Completed);
1582
1583 let steps = engine.store().list_steps(run.id).await.unwrap();
1584
1585 assert_eq!(steps.len(), 1);
1586 assert_eq!(steps[0].name, "create-issue");
1587 assert_eq!(steps[0].kind, StepKind::Custom("gitlab".to_string()));
1588 assert_eq!(
1589 steps[0].status.state,
1590 ironflow_store::models::StepStatus::Completed
1591 );
1592
1593 let output = steps[0].output.as_ref().unwrap();
1594 assert_eq!(output["issue_id"], 42);
1595 assert_eq!(output["project_id"], 123);
1596
1597 let input = steps[0].input.as_ref().unwrap();
1598 assert_eq!(input["project_id"], 123);
1599 assert_eq!(input["title"], "Bug report");
1600 }
1601
1602 #[tokio::test]
1603 async fn operation_step_failure_marks_run_failed() {
1604 let mut engine = create_test_engine();
1605 engine.register(FailingOperationWorkflow).unwrap();
1606
1607 let result = engine
1608 .run_handler("failing-operation-workflow", TriggerKind::Manual, json!({}))
1609 .await;
1610
1611 assert!(result.is_err());
1612 }
1613
1614 #[tokio::test]
1615 async fn operation_mixed_with_shell_steps() {
1616 let mut engine = create_test_engine();
1617 engine.register(MixedWorkflow).unwrap();
1618
1619 let run = engine
1620 .run_handler("mixed-workflow", TriggerKind::Manual, json!({}))
1621 .await
1622 .unwrap();
1623
1624 assert_eq!(run.status.state, RunStatus::Completed);
1625
1626 let steps = engine.store().list_steps(run.id).await.unwrap();
1627
1628 assert_eq!(steps.len(), 2);
1629 assert_eq!(steps[0].kind, StepKind::Shell);
1630 assert_eq!(steps[1].kind, StepKind::Custom("gitlab".to_string()));
1631 assert_eq!(steps[0].position, 0);
1632 assert_eq!(steps[1].position, 1);
1633 }
1634
1635 use crate::config::ApprovalConfig;
1640
1641 struct SingleApprovalWorkflow;
1642
1643 impl WorkflowHandler for SingleApprovalWorkflow {
1644 fn name(&self) -> &str {
1645 "single-approval"
1646 }
1647
1648 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1649 Box::pin(async move {
1650 ctx.shell("build", ShellConfig::new("echo built")).await?;
1651 ctx.approval("gate", ApprovalConfig::new("OK?")).await?;
1652 ctx.shell("deploy", ShellConfig::new("echo deployed"))
1653 .await?;
1654 Ok(())
1655 })
1656 }
1657 }
1658
1659 struct DoubleApprovalWorkflow;
1660
1661 impl WorkflowHandler for DoubleApprovalWorkflow {
1662 fn name(&self) -> &str {
1663 "double-approval"
1664 }
1665
1666 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1667 Box::pin(async move {
1668 ctx.shell("build", ShellConfig::new("echo built")).await?;
1669 ctx.approval("staging-gate", ApprovalConfig::new("Deploy staging?"))
1670 .await?;
1671 ctx.shell("deploy-staging", ShellConfig::new("echo staging"))
1672 .await?;
1673 ctx.approval("prod-gate", ApprovalConfig::new("Deploy prod?"))
1674 .await?;
1675 ctx.shell("deploy-prod", ShellConfig::new("echo prod"))
1676 .await?;
1677 Ok(())
1678 })
1679 }
1680 }
1681
1682 #[tokio::test]
1683 async fn approval_pauses_run() {
1684 let mut engine = create_test_engine();
1685 engine.register(SingleApprovalWorkflow).unwrap();
1686
1687 let run = engine
1688 .run_handler("single-approval", TriggerKind::Manual, json!({}))
1689 .await
1690 .unwrap();
1691
1692 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
1693
1694 let steps = engine.store().list_steps(run.id).await.unwrap();
1695 assert_eq!(steps.len(), 2); assert_eq!(steps[0].kind, StepKind::Shell);
1697 assert_eq!(steps[0].status.state, StepStatus::Completed);
1698 assert_eq!(steps[1].kind, StepKind::Approval);
1699 assert_eq!(steps[1].status.state, StepStatus::AwaitingApproval);
1700 }
1701
1702 #[tokio::test]
1703 async fn approval_resume_completes_run() {
1704 let mut engine = create_test_engine();
1705 engine.register(SingleApprovalWorkflow).unwrap();
1706
1707 let run = engine
1709 .run_handler("single-approval", TriggerKind::Manual, json!({}))
1710 .await
1711 .unwrap();
1712 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
1713
1714 engine
1716 .store()
1717 .update_run_status(run.id, RunStatus::Running)
1718 .await
1719 .unwrap();
1720
1721 let resumed = engine.resume_run(run.id).await.unwrap();
1723 assert_eq!(resumed.status.state, RunStatus::Completed);
1724
1725 let steps = engine.store().list_steps(run.id).await.unwrap();
1726 assert_eq!(steps.len(), 3); assert_eq!(steps[0].name, "build");
1728 assert_eq!(steps[0].status.state, StepStatus::Completed);
1729 assert_eq!(steps[1].name, "gate");
1730 assert_eq!(steps[1].kind, StepKind::Approval);
1731 assert_eq!(steps[1].status.state, StepStatus::Completed);
1732 assert_eq!(steps[2].name, "deploy");
1733 assert_eq!(steps[2].status.state, StepStatus::Completed);
1734 }
1735
1736 #[tokio::test]
1737 async fn double_approval_two_resumes() {
1738 let mut engine = create_test_engine();
1739 engine.register(DoubleApprovalWorkflow).unwrap();
1740
1741 let run = engine
1743 .run_handler("double-approval", TriggerKind::Manual, json!({}))
1744 .await
1745 .unwrap();
1746 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
1747
1748 let steps = engine.store().list_steps(run.id).await.unwrap();
1749 assert_eq!(steps.len(), 2); engine
1753 .store()
1754 .update_run_status(run.id, RunStatus::Running)
1755 .await
1756 .unwrap();
1757
1758 let resumed = engine.resume_run(run.id).await.unwrap();
1759 assert_eq!(resumed.status.state, RunStatus::AwaitingApproval);
1760
1761 let steps = engine.store().list_steps(run.id).await.unwrap();
1762 assert_eq!(steps.len(), 4); engine
1766 .store()
1767 .update_run_status(run.id, RunStatus::Running)
1768 .await
1769 .unwrap();
1770
1771 let final_run = engine.resume_run(run.id).await.unwrap();
1772 assert_eq!(final_run.status.state, RunStatus::Completed);
1773
1774 let steps = engine.store().list_steps(run.id).await.unwrap();
1775 assert_eq!(steps.len(), 5);
1776 assert_eq!(steps[0].name, "build");
1777 assert_eq!(steps[1].name, "staging-gate");
1778 assert_eq!(steps[2].name, "deploy-staging");
1779 assert_eq!(steps[3].name, "prod-gate");
1780 assert_eq!(steps[4].name, "deploy-prod");
1781
1782 for step in &steps {
1783 assert_eq!(step.status.state, StepStatus::Completed);
1784 }
1785 }
1786
1787 use ironflow_store::models::{NewStep, StepUpdate};
1792
1793 async fn create_step_with_status(
1794 store: &Arc<dyn Store>,
1795 run_id: Uuid,
1796 name: &str,
1797 position: u32,
1798 status: StepStatus,
1799 ) -> ironflow_store::models::Step {
1800 let step = store
1801 .create_step(NewStep {
1802 run_id,
1803 name: name.to_string(),
1804 kind: StepKind::Shell,
1805 position,
1806 input: None,
1807 })
1808 .await
1809 .unwrap();
1810
1811 match status {
1812 StepStatus::Pending => {}
1813 StepStatus::Running => {
1814 store
1815 .update_step(
1816 step.id,
1817 StepUpdate {
1818 status: Some(StepStatus::Running),
1819 ..StepUpdate::default()
1820 },
1821 )
1822 .await
1823 .unwrap();
1824 }
1825 StepStatus::Completed => {
1826 store
1827 .update_step(
1828 step.id,
1829 StepUpdate {
1830 status: Some(StepStatus::Running),
1831 ..StepUpdate::default()
1832 },
1833 )
1834 .await
1835 .unwrap();
1836 store
1837 .update_step(
1838 step.id,
1839 StepUpdate {
1840 status: Some(StepStatus::Completed),
1841 ..StepUpdate::default()
1842 },
1843 )
1844 .await
1845 .unwrap();
1846 }
1847 StepStatus::AwaitingApproval => {
1848 store
1849 .update_step(
1850 step.id,
1851 StepUpdate {
1852 status: Some(StepStatus::Running),
1853 ..StepUpdate::default()
1854 },
1855 )
1856 .await
1857 .unwrap();
1858 store
1859 .update_step(
1860 step.id,
1861 StepUpdate {
1862 status: Some(StepStatus::AwaitingApproval),
1863 ..StepUpdate::default()
1864 },
1865 )
1866 .await
1867 .unwrap();
1868 }
1869 _ => panic!("unsupported status for test helper: {status}"),
1870 }
1871
1872 store.get_step(step.id).await.unwrap().unwrap()
1873 }
1874
1875 #[tokio::test]
1876 async fn fail_orphaned_steps_marks_running_as_failed() {
1877 let engine = create_test_engine();
1878 let run = engine
1879 .store()
1880 .create_run(NewRun {
1881 workflow_name: "test".to_string(),
1882 trigger: TriggerKind::Manual,
1883 payload: json!({}),
1884 max_retries: 0,
1885 handler_version: None,
1886 labels: HashMap::new(),
1887 scheduled_at: None,
1888 idempotency_key: None,
1889 max_cost_usd: None,
1890 })
1891 .await
1892 .unwrap()
1893 .into_run();
1894
1895 let step = create_step_with_status(
1896 engine.store(),
1897 run.id,
1898 "running-step",
1899 0,
1900 StepStatus::Running,
1901 )
1902 .await;
1903
1904 engine
1905 .fail_orphaned_steps(run.id, "parent run timed out")
1906 .await
1907 .unwrap();
1908
1909 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
1910 assert_eq!(updated.status.state, StepStatus::Failed);
1911 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
1912 assert!(updated.completed_at.is_some());
1913 }
1914
1915 #[tokio::test]
1916 async fn fail_orphaned_steps_marks_pending_as_skipped() {
1917 let engine = create_test_engine();
1918 let run = engine
1919 .store()
1920 .create_run(NewRun {
1921 workflow_name: "test".to_string(),
1922 trigger: TriggerKind::Manual,
1923 payload: json!({}),
1924 max_retries: 0,
1925 handler_version: None,
1926 labels: HashMap::new(),
1927 scheduled_at: None,
1928 idempotency_key: None,
1929 max_cost_usd: None,
1930 })
1931 .await
1932 .unwrap()
1933 .into_run();
1934
1935 let step = create_step_with_status(
1936 engine.store(),
1937 run.id,
1938 "pending-step",
1939 0,
1940 StepStatus::Pending,
1941 )
1942 .await;
1943
1944 engine
1945 .fail_orphaned_steps(run.id, "parent run timed out")
1946 .await
1947 .unwrap();
1948
1949 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
1950 assert_eq!(updated.status.state, StepStatus::Skipped);
1951 assert!(updated.error.is_none());
1952 assert!(updated.completed_at.is_some());
1953 }
1954
1955 #[tokio::test]
1956 async fn fail_orphaned_steps_marks_awaiting_approval_as_failed() {
1957 let engine = create_test_engine();
1958 let run = engine
1959 .store()
1960 .create_run(NewRun {
1961 workflow_name: "test".to_string(),
1962 trigger: TriggerKind::Manual,
1963 payload: json!({}),
1964 max_retries: 0,
1965 handler_version: None,
1966 labels: HashMap::new(),
1967 scheduled_at: None,
1968 idempotency_key: None,
1969 max_cost_usd: None,
1970 })
1971 .await
1972 .unwrap()
1973 .into_run();
1974
1975 let step = create_step_with_status(
1976 engine.store(),
1977 run.id,
1978 "approval-step",
1979 0,
1980 StepStatus::AwaitingApproval,
1981 )
1982 .await;
1983
1984 engine
1985 .fail_orphaned_steps(run.id, "parent run timed out")
1986 .await
1987 .unwrap();
1988
1989 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
1990 assert_eq!(updated.status.state, StepStatus::Failed);
1991 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
1992 assert!(updated.completed_at.is_some());
1993 }
1994
1995 #[tokio::test]
1996 async fn fail_orphaned_steps_skips_terminal_steps() {
1997 let engine = create_test_engine();
1998 let run = engine
1999 .store()
2000 .create_run(NewRun {
2001 workflow_name: "test".to_string(),
2002 trigger: TriggerKind::Manual,
2003 payload: json!({}),
2004 max_retries: 0,
2005 handler_version: None,
2006 labels: HashMap::new(),
2007 scheduled_at: None,
2008 idempotency_key: None,
2009 max_cost_usd: None,
2010 })
2011 .await
2012 .unwrap()
2013 .into_run();
2014
2015 let completed_step =
2016 create_step_with_status(engine.store(), run.id, "done", 0, StepStatus::Completed).await;
2017 let running_step =
2018 create_step_with_status(engine.store(), run.id, "in-flight", 1, StepStatus::Running)
2019 .await;
2020
2021 engine
2022 .fail_orphaned_steps(run.id, "parent run timed out")
2023 .await
2024 .unwrap();
2025
2026 let completed = engine
2027 .store()
2028 .get_step(completed_step.id)
2029 .await
2030 .unwrap()
2031 .unwrap();
2032 assert_eq!(completed.status.state, StepStatus::Completed);
2033
2034 let failed = engine
2035 .store()
2036 .get_step(running_step.id)
2037 .await
2038 .unwrap()
2039 .unwrap();
2040 assert_eq!(failed.status.state, StepStatus::Failed);
2041 }
2042
2043 #[tokio::test]
2044 async fn fail_orphaned_steps_mixed_states() {
2045 let engine = create_test_engine();
2046 let run = engine
2047 .store()
2048 .create_run(NewRun {
2049 workflow_name: "test".to_string(),
2050 trigger: TriggerKind::Manual,
2051 payload: json!({}),
2052 max_retries: 0,
2053 handler_version: None,
2054 labels: HashMap::new(),
2055 scheduled_at: None,
2056 idempotency_key: None,
2057 max_cost_usd: None,
2058 })
2059 .await
2060 .unwrap()
2061 .into_run();
2062
2063 let s_completed =
2064 create_step_with_status(engine.store(), run.id, "step-1", 0, StepStatus::Completed)
2065 .await;
2066 let s_running =
2067 create_step_with_status(engine.store(), run.id, "step-2", 1, StepStatus::Running).await;
2068 let s_pending =
2069 create_step_with_status(engine.store(), run.id, "step-3", 2, StepStatus::Pending).await;
2070
2071 engine.fail_orphaned_steps(run.id, "timeout").await.unwrap();
2072
2073 let r_completed = engine
2074 .store()
2075 .get_step(s_completed.id)
2076 .await
2077 .unwrap()
2078 .unwrap();
2079 assert_eq!(r_completed.status.state, StepStatus::Completed);
2080
2081 let r_running = engine
2082 .store()
2083 .get_step(s_running.id)
2084 .await
2085 .unwrap()
2086 .unwrap();
2087 assert_eq!(r_running.status.state, StepStatus::Failed);
2088 assert_eq!(r_running.error.as_deref(), Some("timeout"));
2089
2090 let r_pending = engine
2091 .store()
2092 .get_step(s_pending.id)
2093 .await
2094 .unwrap()
2095 .unwrap();
2096 assert_eq!(r_pending.status.state, StepStatus::Skipped);
2097 assert!(r_pending.error.is_none());
2098 }
2099
2100 #[tokio::test]
2101 async fn fail_orphaned_steps_no_steps_is_noop() {
2102 let engine = create_test_engine();
2103 let run = engine
2104 .store()
2105 .create_run(NewRun {
2106 workflow_name: "test".to_string(),
2107 trigger: TriggerKind::Manual,
2108 payload: json!({}),
2109 max_retries: 0,
2110 handler_version: None,
2111 labels: HashMap::new(),
2112 scheduled_at: None,
2113 idempotency_key: None,
2114 max_cost_usd: None,
2115 })
2116 .await
2117 .unwrap()
2118 .into_run();
2119
2120 let result = engine.fail_orphaned_steps(run.id, "timeout").await;
2121 assert!(result.is_ok());
2122 }
2123
2124 #[tokio::test]
2125 async fn fail_orphaned_steps_preserves_existing_error() {
2126 let engine = create_test_engine();
2127 let run = engine
2128 .store()
2129 .create_run(NewRun {
2130 workflow_name: "test".to_string(),
2131 trigger: TriggerKind::Manual,
2132 payload: json!({}),
2133 max_retries: 0,
2134 handler_version: None,
2135 labels: HashMap::new(),
2136 scheduled_at: None,
2137 idempotency_key: None,
2138 max_cost_usd: None,
2139 })
2140 .await
2141 .unwrap()
2142 .into_run();
2143
2144 let step_with_error = create_step_with_status(
2145 engine.store(),
2146 run.id,
2147 "already-errored",
2148 0,
2149 StepStatus::Running,
2150 )
2151 .await;
2152
2153 engine
2154 .store()
2155 .update_step(
2156 step_with_error.id,
2157 StepUpdate {
2158 error: Some("real error from provider".to_string()),
2159 ..StepUpdate::default()
2160 },
2161 )
2162 .await
2163 .unwrap();
2164
2165 let step_no_error = create_step_with_status(
2166 engine.store(),
2167 run.id,
2168 "no-error-yet",
2169 1,
2170 StepStatus::Running,
2171 )
2172 .await;
2173
2174 engine
2175 .fail_orphaned_steps(run.id, "parent run failed")
2176 .await
2177 .unwrap();
2178
2179 let updated_with = engine
2180 .store()
2181 .get_step(step_with_error.id)
2182 .await
2183 .unwrap()
2184 .unwrap();
2185 assert_eq!(updated_with.status.state, StepStatus::Failed);
2186 assert_eq!(
2187 updated_with.error.as_deref(),
2188 Some("real error from provider"),
2189 );
2190
2191 let updated_without = engine
2192 .store()
2193 .get_step(step_no_error.id)
2194 .await
2195 .unwrap()
2196 .unwrap();
2197 assert_eq!(updated_without.status.state, StepStatus::Failed);
2198 assert_eq!(updated_without.error.as_deref(), Some("parent run failed"),);
2199 }
2200}