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