1use std::collections::HashMap;
10use std::fmt;
11use std::sync::Arc;
12use std::time::Instant;
13
14use chrono::{DateTime, TimeDelta, 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::artifact::ArtifactSink;
35use crate::budget::{BudgetConfig, month_start};
36use crate::context::WorkflowContext;
37use crate::error::EngineError;
38use crate::executor::StepResult;
39use crate::handler::{WorkflowHandler, WorkflowInfo};
40use crate::log_sender::LogSender;
41use crate::notify::{Event, EventPublisher, EventSubscriber, WorkflowEventBus};
42use crate::retry_policy::{backoff_for_retry, is_run_retryable};
43use crate::schedule::CronSchedule;
44
45#[derive(Debug, Clone)]
64pub struct WorkflowResult {
65 pub run: Run,
67 pub steps: Vec<StepResult>,
69}
70
71#[derive(Debug, Clone, Default)]
90pub struct EnqueueOptions {
91 pub max_retries: u32,
93 pub labels: HashMap<String, String>,
95 pub scheduled_at: Option<DateTime<Utc>>,
98 pub max_cost_usd: Option<Decimal>,
102 pub created_by: Option<RunActor>,
105 pub idempotency_key: Option<String>,
111}
112
113pub struct Engine {
154 store: Arc<dyn Store>,
155 provider: Arc<dyn AgentProvider>,
156 handlers: HashMap<String, Arc<dyn WorkflowHandler>>,
157 event_publisher: EventPublisher,
158 log_sender: Option<LogSender>,
159 budget: BudgetConfig,
160 artifact_sink: Option<Arc<dyn ArtifactSink>>,
161 event_bus: Option<WorkflowEventBus>,
162}
163
164fn validate_category(handler_name: &str, category: &str) -> Result<(), EngineError> {
174 let reject = |reason: &str| {
175 Err(EngineError::InvalidWorkflow(format!(
176 "handler '{handler_name}' has invalid category '{category}': {reason}"
177 )))
178 };
179
180 if category.is_empty() {
181 return reject("empty category");
182 }
183 if category.starts_with('/') {
184 return reject("leading '/'");
185 }
186 if category.ends_with('/') {
187 return reject("trailing '/'");
188 }
189 for segment in category.split('/') {
190 if segment.is_empty() {
191 return reject("empty segment (double '/')");
192 }
193 if segment.trim().is_empty() {
194 return reject("whitespace-only segment");
195 }
196 }
197 Ok(())
198}
199
200impl Engine {
201 pub fn new(store: Arc<dyn Store>, provider: Arc<dyn AgentProvider>) -> Self {
217 Self {
218 store,
219 provider,
220 handlers: HashMap::new(),
221 event_publisher: EventPublisher::new(),
222 log_sender: None,
223 budget: BudgetConfig::new(),
224 artifact_sink: None,
225 event_bus: None,
226 }
227 }
228
229 pub fn with_budget_config(mut self, budget: BudgetConfig) -> Self {
250 self.budget = budget;
251 self
252 }
253
254 pub fn budget_config(&self) -> &BudgetConfig {
256 &self.budget
257 }
258
259 pub fn set_log_sender(&mut self, sender: LogSender) {
265 self.log_sender = Some(sender);
266 }
267
268 pub fn set_artifact_sink(&mut self, sink: Arc<dyn ArtifactSink>) {
286 self.artifact_sink = Some(sink);
287 }
288
289 pub fn artifact_sink(&self) -> Option<&Arc<dyn ArtifactSink>> {
291 self.artifact_sink.as_ref()
292 }
293
294 pub fn set_event_bus(&mut self, bus: WorkflowEventBus) {
311 self.event_bus = Some(bus);
312 }
313
314 pub fn event_bus(&self) -> Option<&WorkflowEventBus> {
316 self.event_bus.as_ref()
317 }
318
319 pub fn store(&self) -> &Arc<dyn Store> {
321 &self.store
322 }
323
324 pub fn provider(&self) -> &Arc<dyn AgentProvider> {
326 &self.provider
327 }
328
329 fn build_context(&self, run: &Run) -> WorkflowContext {
338 let handlers = self.handlers.clone();
339 let resolver: crate::context::HandlerResolver =
340 Arc::new(move |name: &str| handlers.get(name).cloned());
341 let mut ctx = WorkflowContext::with_handler_resolver(
342 run.id,
343 self.store.clone(),
344 self.provider.clone(),
345 resolver,
346 );
347 ctx.carry_over_run_totals(run.retry_count + 1, run.cost_usd, run.duration_ms);
348 ctx.set_max_cost_usd(run.max_cost_usd);
349 if let Some(ref sender) = self.log_sender {
350 ctx.set_log_sender(sender.clone());
351 }
352 if let Some(ref sink) = self.artifact_sink {
353 ctx.set_artifact_sink(sink.clone());
354 }
355 if let Some(ref bus) = self.event_bus {
356 ctx.set_event_bus(bus.clone());
357 }
358 ctx
359 }
360
361 async fn check_monthly_quota(&self, workflow_name: &str) -> Result<(), EngineError> {
372 let Some(limit) = self.budget.monthly_cost_limit_usd else {
373 return Ok(());
374 };
375
376 let stats = self
377 .store
378 .get_stats(RunFilter {
379 created_after: Some(month_start(Utc::now())),
380 ..RunFilter::default()
381 })
382 .await?;
383
384 if stats.total_cost_usd < limit {
385 return Ok(());
386 }
387
388 warn!(
389 workflow = %workflow_name,
390 limit_usd = %limit,
391 spent_usd = %stats.total_cost_usd,
392 "monthly cost quota exhausted, refusing new run"
393 );
394
395 #[cfg(feature = "prometheus")]
396 counter!(
397 RUN_BUDGET_EXCEEDED_TOTAL,
398 "workflow" => workflow_name.to_string(),
399 "scope" => "monthly",
400 )
401 .increment(1);
402
403 Err(EngineError::MonthlyBudgetExceeded {
404 limit_usd: limit,
405 spent_usd: stats.total_cost_usd,
406 })
407 }
408
409 pub fn register(&mut self, handler: impl WorkflowHandler + 'static) -> Result<(), EngineError> {
453 let name = handler.name().to_string();
454 if self.handlers.contains_key(&name) {
455 return Err(EngineError::InvalidWorkflow(format!(
456 "handler '{}' already registered",
457 name
458 )));
459 }
460 if let Some(category) = handler.category() {
461 validate_category(&name, category)?;
462 }
463 self.handlers.insert(name, Arc::new(handler));
464 Ok(())
465 }
466
467 pub fn register_boxed(&mut self, handler: Box<dyn WorkflowHandler>) -> Result<(), EngineError> {
474 let name = handler.name().to_string();
475 if self.handlers.contains_key(&name) {
476 return Err(EngineError::InvalidWorkflow(format!(
477 "handler '{}' already registered",
478 name
479 )));
480 }
481 if let Some(category) = handler.category() {
482 validate_category(&name, category)?;
483 }
484 self.handlers.insert(name, Arc::from(handler));
485 Ok(())
486 }
487
488 pub fn get_handler(&self, name: &str) -> Option<&Arc<dyn WorkflowHandler>> {
490 self.handlers.get(name)
491 }
492
493 pub fn handler_names(&self) -> Vec<&str> {
495 self.handlers.keys().map(|s| s.as_str()).collect()
496 }
497
498 pub fn handler_info(&self, name: &str) -> Option<WorkflowInfo> {
500 self.handlers.get(name).map(|h| h.describe())
501 }
502
503 pub fn scheduled_handlers(&self) -> Vec<(&str, &CronSchedule)> {
527 self.handlers
528 .iter()
529 .filter_map(|(name, handler)| handler.schedule().map(|sched| (name.as_str(), sched)))
530 .collect()
531 }
532
533 pub fn subscribe(
558 &mut self,
559 subscriber: impl EventSubscriber + 'static,
560 event_types: &[&'static str],
561 ) {
562 self.event_publisher.subscribe(subscriber, event_types);
563 }
564
565 pub fn event_publisher(&self) -> &EventPublisher {
570 &self.event_publisher
571 }
572
573 #[tracing::instrument(name = "engine.run_handler", skip_all, fields(workflow = %handler_name))]
603 pub async fn run_handler(
604 &self,
605 handler_name: &str,
606 trigger: TriggerKind,
607 payload: Value,
608 ) -> Result<WorkflowResult, EngineError> {
609 let handler = self
610 .handlers
611 .get(handler_name)
612 .ok_or_else(|| {
613 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
614 })?
615 .clone();
616
617 self.check_monthly_quota(handler_name).await?;
618
619 let handler_version = handler.version().map(str::to_string);
620 let max_cost_usd = self
621 .budget
622 .resolve_run_cap(None, handler.default_max_cost_usd());
623 let run = self
624 .store
625 .create_run(NewRun {
626 created_by: None,
627 workflow_name: handler_name.to_string(),
628 trigger,
629 payload,
630 max_retries: 0,
631 handler_version,
632 labels: handler.default_labels(),
633 scheduled_at: None,
634 idempotency_key: None,
635 max_cost_usd,
636 })
637 .await?
638 .into_run();
639
640 let run_id = run.id;
641 info!(run_id = %run_id, handler_version = run.handler_version.as_deref().unwrap_or(""), "run created");
642
643 self.store
644 .update_run_status(run_id, RunStatus::Running)
645 .await?;
646
647 #[cfg(feature = "prometheus")]
648 gauge!(RUNS_ACTIVE, "workflow" => handler_name.to_string()).increment(1.0);
649
650 let run_start = Instant::now();
651 let mut ctx = self.build_context(&run);
652
653 let result = handler.execute(&mut ctx).await;
654 self.finalize_run(run_id, handler_name, result, &ctx, run_start, run.labels)
655 .await
656 }
657
658 #[tracing::instrument(name = "engine.enqueue_handler", skip_all, fields(workflow = %handler_name))]
669 pub async fn enqueue_handler(
670 &self,
671 handler_name: &str,
672 trigger: TriggerKind,
673 payload: Value,
674 max_retries: u32,
675 ) -> Result<Run, EngineError> {
676 self.enqueue_handler_with_options(
677 handler_name,
678 trigger,
679 payload,
680 EnqueueOptions {
681 max_retries,
682 ..Default::default()
683 },
684 )
685 .await
686 .map(RunCreation::into_run)
687 }
688
689 #[tracing::instrument(name = "engine.enqueue_handler_with_options", skip_all, fields(workflow = %handler_name))]
734 pub async fn enqueue_handler_with_options(
735 &self,
736 handler_name: &str,
737 trigger: TriggerKind,
738 payload: Value,
739 options: EnqueueOptions,
740 ) -> Result<RunCreation, EngineError> {
741 let EnqueueOptions {
742 max_retries,
743 labels,
744 scheduled_at,
745 max_cost_usd,
746 created_by,
747 idempotency_key,
748 } = options;
749
750 let handler = self.handlers.get(handler_name).ok_or_else(|| {
751 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
752 })?;
753
754 self.check_monthly_quota(handler_name).await?;
755
756 let handler_version = handler.version().map(str::to_string);
757 let mut merged_labels = handler.default_labels();
758 merged_labels.extend(labels);
759 let resolved_cap = self
760 .budget
761 .resolve_run_cap(max_cost_usd, handler.default_max_cost_usd());
762
763 let creation = self
764 .store
765 .create_run(NewRun {
766 workflow_name: handler_name.to_string(),
767 trigger,
768 payload,
769 max_retries,
770 handler_version,
771 labels: merged_labels,
772 scheduled_at,
773 created_by,
774 idempotency_key,
775 max_cost_usd: resolved_cap,
776 })
777 .await?;
778
779 match &creation {
780 RunCreation::Created(run) => info!(
781 run_id = %run.id,
782 workflow = %handler_name,
783 max_cost_usd = ?resolved_cap,
784 "handler run enqueued"
785 ),
786 RunCreation::Existing(run) => info!(
787 run_id = %run.id,
788 workflow = %handler_name,
789 "idempotent replay, nothing enqueued"
790 ),
791 }
792
793 Ok(creation)
794 }
795
796 #[tracing::instrument(name = "engine.execute_handler_run", skip_all, fields(run_id = %run_id))]
805 pub async fn execute_handler_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
806 let run = self
807 .store
808 .get_run(run_id)
809 .await?
810 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
811
812 let handler = self
813 .handlers
814 .get(&run.workflow_name)
815 .ok_or_else(|| {
816 EngineError::InvalidWorkflow(format!(
817 "no handler registered: {}",
818 run.workflow_name
819 ))
820 })?
821 .clone();
822
823 #[cfg(feature = "prometheus")]
824 gauge!(RUNS_ACTIVE, "workflow" => run.workflow_name.clone()).increment(1.0);
825
826 let run_start = Instant::now();
827 let mut ctx = self.build_context(&run);
828
829 if run.retry_count > 0 {
832 ctx.load_replay_steps().await?;
833 }
834
835 let result = handler.execute(&mut ctx).await;
836 self.finalize_run(
837 run_id,
838 &run.workflow_name,
839 result,
840 &ctx,
841 run_start,
842 run.labels,
843 )
844 .await
845 }
846
847 #[tracing::instrument(name = "engine.execute_run", skip_all, fields(run_id = %run_id))]
855 pub async fn execute_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
856 self.execute_handler_run(run_id).await
857 }
858
859 #[tracing::instrument(name = "engine.resume_run", skip_all, fields(run_id = %run_id))]
873 pub async fn resume_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
874 let run = self
875 .store
876 .get_run(run_id)
877 .await?
878 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
879
880 let handler = self
881 .handlers
882 .get(&run.workflow_name)
883 .ok_or_else(|| {
884 EngineError::InvalidWorkflow(format!(
885 "no handler registered: {}",
886 run.workflow_name
887 ))
888 })?
889 .clone();
890
891 info!(run_id = %run_id, workflow = %run.workflow_name, "resuming run after approval");
892
893 let run_start = Instant::now();
894 let mut ctx = self.build_context(&run);
895 ctx.load_replay_steps().await?;
896
897 let result = handler.execute(&mut ctx).await;
898 self.finalize_run(
899 run_id,
900 &run.workflow_name,
901 result,
902 &ctx,
903 run_start,
904 run.labels,
905 )
906 .await
907 }
908
909 pub async fn fail_or_schedule_retry(
954 &self,
955 run_id: Uuid,
956 error: &str,
957 retryable: bool,
958 cost_usd: Option<Decimal>,
959 duration_ms: Option<u64>,
960 ) -> Result<RunStatus, EngineError> {
961 let run = self
962 .store
963 .get_run(run_id)
964 .await?
965 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
966
967 let has_attempts_left = run.retry_count < run.max_retries;
968 let update = if retryable && has_attempts_left {
969 let backoff = backoff_for_retry(run.retry_count);
970 let scheduled_at = Utc::now() + TimeDelta::milliseconds(backoff.as_millis() as i64);
971
972 info!(
973 run_id = %run_id,
974 workflow = %run.workflow_name,
975 attempt = run.retry_count + 1,
976 max_retries = run.max_retries,
977 backoff_secs = backoff.as_secs(),
978 scheduled_at = %scheduled_at,
979 "run failed, scheduling retry"
980 );
981
982 RunUpdate {
983 status: Some(RunStatus::Retrying),
984 error: Some(error.to_string()),
985 increment_retry: true,
986 cost_usd,
987 duration_ms,
988 scheduled_at: Some(scheduled_at),
989 ..RunUpdate::default()
990 }
991 } else {
992 RunUpdate {
993 status: Some(RunStatus::Failed),
994 error: Some(error.to_string()),
995 cost_usd,
996 duration_ms,
997 completed_at: Some(Utc::now()),
998 ..RunUpdate::default()
999 }
1000 };
1001
1002 let status = update.status.unwrap_or(RunStatus::Failed);
1003 self.store.update_run(run_id, update).await?;
1004 self.fail_orphaned_steps(run_id, error).await?;
1005
1006 Ok(status)
1007 }
1008
1009 pub async fn fail_orphaned_steps(
1023 &self,
1024 run_id: Uuid,
1025 error_message: &str,
1026 ) -> Result<(), EngineError> {
1027 let steps = self.store.list_steps(run_id).await?;
1028 let now = Utc::now();
1029
1030 for step in steps {
1031 if step.status.state.is_terminal() {
1032 continue;
1033 }
1034
1035 let (target_status, error) = match step.status.state {
1036 StepStatus::Running | StepStatus::AwaitingApproval => {
1037 let err = if step.error.is_some() {
1038 None
1039 } else {
1040 Some(error_message.to_string())
1041 };
1042 (StepStatus::Failed, err)
1043 }
1044 StepStatus::Pending => (StepStatus::Skipped, None),
1045 _ => continue,
1046 };
1047
1048 if let Err(e) = self
1049 .store
1050 .update_step(
1051 step.id,
1052 StepUpdate {
1053 status: Some(target_status),
1054 error,
1055 completed_at: Some(now),
1056 ..StepUpdate::default()
1057 },
1058 )
1059 .await
1060 {
1061 warn!(
1062 run_id = %run_id,
1063 step_id = %step.id,
1064 step_name = %step.name,
1065 error = %e,
1066 "failed to cleanup orphaned step"
1067 );
1068 } else {
1069 info!(
1070 run_id = %run_id,
1071 step_id = %step.id,
1072 step_name = %step.name,
1073 from = %step.status.state,
1074 to = %target_status,
1075 "cleaned up orphaned step"
1076 );
1077 }
1078 }
1079
1080 Ok(())
1081 }
1082
1083 async fn finalize_run(
1089 &self,
1090 run_id: Uuid,
1091 workflow_name: &str,
1092 result: Result<(), EngineError>,
1093 ctx: &WorkflowContext,
1094 run_start: Instant,
1095 run_labels: HashMap<String, String>,
1096 ) -> Result<WorkflowResult, EngineError> {
1097 let total_duration = ctx.carried_duration_ms() + run_start.elapsed().as_millis() as u64;
1100 let completed_at = Utc::now();
1101
1102 let final_status;
1103 let final_run;
1104
1105 match result {
1106 Ok(()) => {
1107 final_status = if ctx.has_allowed_failure() {
1108 RunStatus::Warning
1109 } else {
1110 RunStatus::Completed
1111 };
1112 final_run = self
1113 .store
1114 .update_run_returning(
1115 run_id,
1116 RunUpdate {
1117 status: Some(final_status),
1118 cost_usd: Some(ctx.total_cost_usd()),
1119 duration_ms: Some(total_duration),
1120 completed_at: Some(completed_at),
1121 ..RunUpdate::default()
1122 },
1123 )
1124 .await?;
1125
1126 info!(
1127 run_id = %run_id,
1128 status = %final_status,
1129 cost_usd = %ctx.total_cost_usd(),
1130 duration_ms = total_duration,
1131 "run completed"
1132 );
1133 }
1134 Err(EngineError::ApprovalRequired {
1135 run_id: approval_run_id,
1136 step_id,
1137 ref message,
1138 }) => {
1139 final_status = RunStatus::AwaitingApproval;
1140 final_run = self
1141 .store
1142 .update_run_returning(
1143 run_id,
1144 RunUpdate {
1145 status: Some(RunStatus::AwaitingApproval),
1146 cost_usd: Some(ctx.total_cost_usd()),
1147 duration_ms: Some(total_duration),
1148 ..RunUpdate::default()
1149 },
1150 )
1151 .await?;
1152
1153 info!(
1154 run_id = %approval_run_id,
1155 step_id = %step_id,
1156 message = %message,
1157 "run awaiting approval"
1158 );
1159 }
1160 Err(err) => {
1161 let budget_exceeded = matches!(err, EngineError::RunBudgetExceeded { .. });
1165
1166 final_status = if budget_exceeded {
1167 if let Err(store_err) = self
1168 .store
1169 .update_run(
1170 run_id,
1171 RunUpdate {
1172 status: Some(RunStatus::Cancelled),
1173 error: Some(err.to_string()),
1174 cost_usd: Some(ctx.total_cost_usd()),
1175 duration_ms: Some(total_duration),
1176 completed_at: Some(completed_at),
1177 ..RunUpdate::default()
1178 },
1179 )
1180 .await
1181 {
1182 error!(run_id = %run_id, store_error = %store_err, "failed to persist run cancellation");
1183 }
1184 if let Err(cleanup_err) = self
1185 .fail_orphaned_steps(run_id, "run stopped: cost cap reached")
1186 .await
1187 {
1188 error!(run_id = %run_id, store_error = %cleanup_err, "failed to cleanup orphaned steps");
1189 }
1190 RunStatus::Cancelled
1191 } else {
1192 self.fail_or_schedule_retry(
1193 run_id,
1194 &err.to_string(),
1195 is_run_retryable(&err),
1196 Some(ctx.total_cost_usd()),
1197 Some(total_duration),
1198 )
1199 .await
1200 .unwrap_or_else(|store_err| {
1201 error!(run_id = %run_id, store_error = %store_err, "failed to persist run failure");
1202 RunStatus::Failed
1203 })
1204 };
1205
1206 if budget_exceeded {
1207 self.on_run_budget_exceeded(workflow_name, run_id, &err);
1208 }
1209
1210 error!(run_id = %run_id, status = %final_status, error = %err, "run stopped");
1211
1212 self.publish_run_status_changed(
1213 workflow_name,
1214 run_id,
1215 final_status,
1216 Some(err.to_string()),
1217 ctx,
1218 total_duration,
1219 run_labels,
1220 );
1221
1222 #[cfg(feature = "prometheus")]
1223 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
1224
1225 return Err(err);
1226 }
1227 }
1228
1229 self.publish_run_status_changed(
1230 workflow_name,
1231 run_id,
1232 final_status,
1233 None,
1234 ctx,
1235 total_duration,
1236 run_labels,
1237 );
1238
1239 #[cfg(feature = "prometheus")]
1240 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
1241
1242 Ok(WorkflowResult {
1243 run: final_run,
1244 steps: ctx.step_results().to_vec(),
1245 })
1246 }
1247
1248 #[cfg(feature = "prometheus")]
1250 fn emit_run_metrics(
1251 &self,
1252 workflow_name: &str,
1253 status: RunStatus,
1254 duration_ms: u64,
1255 ctx: &WorkflowContext,
1256 ) {
1257 let status_str = status.to_string();
1258 let wf = workflow_name.to_string();
1259
1260 counter!(RUNS_TOTAL, "workflow" => wf.clone(), "status" => status_str.clone()).increment(1);
1261 histogram!(RUN_DURATION_SECONDS, "workflow" => wf.clone(), "status" => status_str)
1262 .record(duration_ms as f64 / 1000.0);
1263 histogram!(RUN_COST_USD, "workflow" => wf.clone()).record(
1264 ctx.total_cost_usd()
1265 .to_string()
1266 .parse::<f64>()
1267 .unwrap_or(0.0),
1268 );
1269 gauge!(RUNS_ACTIVE, "workflow" => wf).decrement(1.0);
1270 }
1271
1272 fn on_run_budget_exceeded(&self, workflow_name: &str, run_id: Uuid, err: &EngineError) {
1278 let EngineError::RunBudgetExceeded {
1279 limit_usd,
1280 spent_usd,
1281 step_budget_usd,
1282 ..
1283 } = err
1284 else {
1285 return;
1286 };
1287
1288 #[cfg(feature = "prometheus")]
1289 counter!(
1290 RUN_BUDGET_EXCEEDED_TOTAL,
1291 "workflow" => workflow_name.to_string(),
1292 "scope" => "run",
1293 )
1294 .increment(1);
1295
1296 self.event_publisher.publish(Event::RunBudgetExceeded {
1297 run_id,
1298 workflow_name: workflow_name.to_string(),
1299 limit_usd: *limit_usd,
1300 spent_usd: *spent_usd,
1301 step_budget_usd: *step_budget_usd,
1302 at: Utc::now(),
1303 });
1304 }
1305
1306 #[allow(clippy::too_many_arguments)]
1311 fn publish_run_status_changed(
1312 &self,
1313 workflow_name: &str,
1314 run_id: Uuid,
1315 to: RunStatus,
1316 error: Option<String>,
1317 ctx: &WorkflowContext,
1318 duration_ms: u64,
1319 labels: HashMap<String, String>,
1320 ) {
1321 let now = Utc::now();
1322 let cost_usd = ctx.total_cost_usd();
1323 let wf = workflow_name.to_string();
1324
1325 self.event_publisher.publish(Event::RunStatusChanged {
1326 run_id,
1327 workflow_name: wf.clone(),
1328 from: RunStatus::Running,
1329 to,
1330 error: error.clone(),
1331 cost_usd,
1332 duration_ms,
1333 labels: labels.clone(),
1334 at: now,
1335 });
1336
1337 if to == RunStatus::Failed {
1338 self.event_publisher.publish(Event::RunFailed {
1339 run_id,
1340 workflow_name: wf,
1341 error,
1342 cost_usd,
1343 duration_ms,
1344 labels,
1345 at: now,
1346 });
1347 }
1348 }
1349}
1350
1351impl fmt::Debug for Engine {
1352 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1353 f.debug_struct("Engine")
1354 .field("handlers", &self.handlers.keys().collect::<Vec<_>>())
1355 .finish_non_exhaustive()
1356 }
1357}
1358
1359#[cfg(test)]
1360mod tests {
1361 use super::*;
1362 use crate::config::ShellConfig;
1363 use crate::handler::{HandlerFuture, WorkflowHandler};
1364 use ironflow_core::providers::claude::ClaudeCodeProvider;
1365 use ironflow_core::providers::record_replay::RecordReplayProvider;
1366 use ironflow_store::memory::InMemoryStore;
1367 use ironflow_store::models::StepStatus;
1368 use serde_json::json;
1369
1370 struct EchoWorkflow;
1372
1373 impl WorkflowHandler for EchoWorkflow {
1374 fn name(&self) -> &str {
1375 "echo-workflow"
1376 }
1377
1378 fn describe(&self) -> WorkflowInfo {
1379 WorkflowInfo {
1380 description: "A simple workflow that echoes hello".to_string(),
1381 source_code: None,
1382 sub_workflows: Vec::new(),
1383 category: None,
1384 version: self.version().map(str::to_string),
1385 compatible_versions: Vec::new(),
1386 input_schema: None,
1387 default_labels: HashMap::new(),
1388 schedule: self.schedule().cloned(),
1389 default_max_cost_usd: self.default_max_cost_usd(),
1390 }
1391 }
1392
1393 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1394 Box::pin(async move {
1395 ctx.shell("greet", ShellConfig::new("echo hello")).await?;
1396 Ok(())
1397 })
1398 }
1399 }
1400
1401 struct FailingWorkflow;
1403
1404 impl WorkflowHandler for FailingWorkflow {
1405 fn name(&self) -> &str {
1406 "failing-workflow"
1407 }
1408
1409 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1410 Box::pin(async move {
1411 ctx.shell("fail", ShellConfig::new("exit 1")).await?;
1412 Ok(())
1413 })
1414 }
1415 }
1416
1417 fn create_test_engine() -> Engine {
1418 let store = Arc::new(InMemoryStore::new());
1419 let inner = ClaudeCodeProvider::new();
1420 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
1421 inner,
1422 "/tmp/ironflow-fixtures",
1423 ));
1424 Engine::new(store, provider)
1425 }
1426
1427 #[test]
1428 fn engine_new_creates_instance() {
1429 let engine = create_test_engine();
1430 assert_eq!(engine.handler_names().len(), 0);
1431 }
1432
1433 #[test]
1434 fn engine_register_handler() {
1435 let mut engine = create_test_engine();
1436 let result = engine.register(EchoWorkflow);
1437 assert!(result.is_ok());
1438 assert_eq!(engine.handler_names().len(), 1);
1439 assert!(engine.handler_names().contains(&"echo-workflow"));
1440 }
1441
1442 #[test]
1443 fn engine_register_duplicate_returns_error() {
1444 let mut engine = create_test_engine();
1445 engine.register(EchoWorkflow).unwrap();
1446 let result = engine.register(EchoWorkflow);
1447 assert!(result.is_err());
1448 }
1449
1450 #[test]
1451 fn engine_get_handler_found() {
1452 let mut engine = create_test_engine();
1453 engine.register(EchoWorkflow).unwrap();
1454 let handler = engine.get_handler("echo-workflow");
1455 assert!(handler.is_some());
1456 }
1457
1458 #[test]
1459 fn engine_get_handler_not_found() {
1460 let engine = create_test_engine();
1461 let handler = engine.get_handler("nonexistent");
1462 assert!(handler.is_none());
1463 }
1464
1465 #[test]
1466 fn engine_handler_names_lists_all() {
1467 let mut engine = create_test_engine();
1468 engine.register(EchoWorkflow).unwrap();
1469 engine.register(FailingWorkflow).unwrap();
1470 let names = engine.handler_names();
1471 assert_eq!(names.len(), 2);
1472 assert!(names.contains(&"echo-workflow"));
1473 assert!(names.contains(&"failing-workflow"));
1474 }
1475
1476 #[test]
1477 fn engine_handler_info_returns_description() {
1478 let mut engine = create_test_engine();
1479 engine.register(EchoWorkflow).unwrap();
1480 let info = engine.handler_info("echo-workflow");
1481 assert!(info.is_some());
1482 let info = info.unwrap();
1483 assert_eq!(info.description, "A simple workflow that echoes hello");
1484 }
1485
1486 struct CategorizedWorkflow;
1487
1488 impl WorkflowHandler for CategorizedWorkflow {
1489 fn name(&self) -> &str {
1490 "categorized"
1491 }
1492 fn category(&self) -> Option<&str> {
1493 Some("data/etl")
1494 }
1495 fn execute<'a>(
1496 &'a self,
1497 _ctx: &'a mut WorkflowContext,
1498 ) -> crate::handler::HandlerFuture<'a> {
1499 Box::pin(async move { Ok(()) })
1500 }
1501 }
1502
1503 #[test]
1504 fn engine_default_describe_propagates_category() {
1505 let mut engine = create_test_engine();
1506 engine.register(CategorizedWorkflow).unwrap();
1507 let info = engine.handler_info("categorized").unwrap();
1508 assert_eq!(info.category.as_deref(), Some("data/etl"));
1509 }
1510
1511 #[test]
1512 fn engine_default_describe_without_category() {
1513 let mut engine = create_test_engine();
1514 engine.register(EchoWorkflow).unwrap();
1515 let info = engine.handler_info("echo-workflow").unwrap();
1516 assert!(info.category.is_none());
1517 }
1518
1519 struct ScheduledWorkflow {
1524 schedule: CronSchedule,
1525 }
1526
1527 impl ScheduledWorkflow {
1528 fn new() -> Self {
1529 Self {
1530 schedule: CronSchedule::new("0 0 * * * *").unwrap(),
1531 }
1532 }
1533 }
1534
1535 impl WorkflowHandler for ScheduledWorkflow {
1536 fn name(&self) -> &str {
1537 "scheduled"
1538 }
1539 fn schedule(&self) -> Option<&CronSchedule> {
1540 Some(&self.schedule)
1541 }
1542 fn execute<'a>(
1543 &'a self,
1544 _ctx: &'a mut WorkflowContext,
1545 ) -> crate::handler::HandlerFuture<'a> {
1546 Box::pin(async move { Ok(()) })
1547 }
1548 }
1549
1550 #[test]
1551 fn engine_default_describe_propagates_schedule() {
1552 let mut engine = create_test_engine();
1553 engine.register(ScheduledWorkflow::new()).unwrap();
1554 let info = engine.handler_info("scheduled").unwrap();
1555 assert_eq!(
1556 info.schedule.as_ref().map(|s| s.as_str()),
1557 Some("0 0 * * * *")
1558 );
1559 }
1560
1561 #[test]
1562 fn engine_default_describe_without_schedule() {
1563 let mut engine = create_test_engine();
1564 engine.register(EchoWorkflow).unwrap();
1565 let info = engine.handler_info("echo-workflow").unwrap();
1566 assert!(info.schedule.is_none());
1567 }
1568
1569 #[test]
1570 fn scheduled_handlers_returns_only_scheduled() {
1571 let mut engine = create_test_engine();
1572 engine.register(EchoWorkflow).unwrap();
1573 engine.register(ScheduledWorkflow::new()).unwrap();
1574 engine.register(FailingWorkflow).unwrap();
1575
1576 let scheduled = engine.scheduled_handlers();
1577 assert_eq!(scheduled.len(), 1);
1578 assert_eq!(scheduled[0].0, "scheduled");
1579 assert_eq!(scheduled[0].1.as_str(), "0 0 * * * *");
1580 }
1581
1582 #[test]
1583 fn scheduled_handlers_empty_when_none_scheduled() {
1584 let mut engine = create_test_engine();
1585 engine.register(EchoWorkflow).unwrap();
1586 engine.register(FailingWorkflow).unwrap();
1587
1588 let scheduled = engine.scheduled_handlers();
1589 assert!(scheduled.is_empty());
1590 }
1591
1592 struct BadCategoryWorkflow(&'static str);
1593
1594 impl WorkflowHandler for BadCategoryWorkflow {
1595 fn name(&self) -> &str {
1596 "bad-category"
1597 }
1598 fn category(&self) -> Option<&str> {
1599 Some(self.0)
1600 }
1601 fn execute<'a>(
1602 &'a self,
1603 _ctx: &'a mut WorkflowContext,
1604 ) -> crate::handler::HandlerFuture<'a> {
1605 Box::pin(async move { Ok(()) })
1606 }
1607 }
1608
1609 #[test]
1610 fn engine_register_rejects_empty_category() {
1611 let mut engine = create_test_engine();
1612 let err = engine.register(BadCategoryWorkflow("")).unwrap_err();
1613 match err {
1614 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty category")),
1615 other => panic!("expected InvalidWorkflow, got {other:?}"),
1616 }
1617 }
1618
1619 #[test]
1620 fn engine_register_rejects_leading_slash_category() {
1621 let mut engine = create_test_engine();
1622 let err = engine
1623 .register(BadCategoryWorkflow("/data/etl"))
1624 .unwrap_err();
1625 match err {
1626 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("leading '/'")),
1627 other => panic!("expected InvalidWorkflow, got {other:?}"),
1628 }
1629 }
1630
1631 #[test]
1632 fn engine_register_rejects_trailing_slash_category() {
1633 let mut engine = create_test_engine();
1634 let err = engine
1635 .register(BadCategoryWorkflow("data/etl/"))
1636 .unwrap_err();
1637 match err {
1638 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("trailing '/'")),
1639 other => panic!("expected InvalidWorkflow, got {other:?}"),
1640 }
1641 }
1642
1643 #[test]
1644 fn engine_register_rejects_double_slash_category() {
1645 let mut engine = create_test_engine();
1646 let err = engine
1647 .register(BadCategoryWorkflow("data//etl"))
1648 .unwrap_err();
1649 match err {
1650 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty segment")),
1651 other => panic!("expected InvalidWorkflow, got {other:?}"),
1652 }
1653 }
1654
1655 #[test]
1656 fn engine_register_rejects_whitespace_only_segment_category() {
1657 let mut engine = create_test_engine();
1658 let err = engine
1659 .register(BadCategoryWorkflow("data/ /etl"))
1660 .unwrap_err();
1661 match err {
1662 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("whitespace-only segment")),
1663 other => panic!("expected InvalidWorkflow, got {other:?}"),
1664 }
1665 }
1666
1667 #[test]
1668 fn engine_register_accepts_valid_nested_category() {
1669 let mut engine = create_test_engine();
1670 assert!(engine.register(CategorizedWorkflow).is_ok());
1671 }
1672
1673 #[tokio::test]
1674 async fn engine_unknown_workflow_returns_error() {
1675 let engine = create_test_engine();
1676 let result = engine
1677 .run_handler("unknown", TriggerKind::Manual, json!({}))
1678 .await;
1679 assert!(result.is_err());
1680 match result {
1681 Err(EngineError::InvalidWorkflow(msg)) => {
1682 assert!(msg.contains("no handler registered"));
1683 }
1684 _ => panic!("expected InvalidWorkflow error"),
1685 }
1686 }
1687
1688 #[tokio::test]
1689 async fn engine_enqueue_handler_creates_pending_run() {
1690 let mut engine = create_test_engine();
1691 engine.register(EchoWorkflow).unwrap();
1692
1693 let run = engine
1694 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
1695 .await
1696 .unwrap();
1697 assert_eq!(run.status.state, RunStatus::Pending);
1698 assert_eq!(run.workflow_name, "echo-workflow");
1699 }
1700
1701 #[tokio::test]
1702 async fn enqueue_handler_leaves_the_run_unattributed() {
1703 let mut engine = create_test_engine();
1704 engine.register(EchoWorkflow).unwrap();
1705
1706 let run = engine
1707 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
1708 .await
1709 .unwrap();
1710
1711 assert!(run.created_by.is_none());
1712 }
1713
1714 #[tokio::test]
1715 async fn enqueue_handler_with_options_records_the_author() {
1716 let mut engine = create_test_engine();
1717 engine.register(EchoWorkflow).unwrap();
1718 let actor = RunActor::User {
1719 user_id: Uuid::now_v7(),
1720 };
1721
1722 let run = engine
1723 .enqueue_handler_with_options(
1724 "echo-workflow",
1725 TriggerKind::Api,
1726 json!({}),
1727 EnqueueOptions {
1728 created_by: Some(actor.clone()),
1729 ..Default::default()
1730 },
1731 )
1732 .await
1733 .unwrap()
1734 .into_run();
1735
1736 assert_eq!(run.created_by, Some(actor));
1737 }
1738
1739 #[tokio::test]
1740 async fn enqueue_handler_with_options_accepts_no_author() {
1741 let mut engine = create_test_engine();
1742 engine.register(EchoWorkflow).unwrap();
1743
1744 let run = engine
1745 .enqueue_handler_with_options(
1746 "echo-workflow",
1747 TriggerKind::Cron {
1748 schedule: "0 * * * * *".to_string(),
1749 },
1750 json!({}),
1751 EnqueueOptions::default(),
1752 )
1753 .await
1754 .unwrap()
1755 .into_run();
1756
1757 assert!(run.created_by.is_none());
1758 }
1759
1760 #[tokio::test]
1761 async fn run_handler_leaves_the_run_unattributed() {
1762 let mut engine = create_test_engine();
1763 engine.register(EchoWorkflow).unwrap();
1764
1765 let run = engine
1766 .run_handler("echo-workflow", TriggerKind::Manual, json!({}))
1767 .await
1768 .unwrap()
1769 .run;
1770
1771 assert!(run.created_by.is_none());
1772 }
1773
1774 #[tokio::test]
1775 async fn engine_register_boxed() {
1776 let mut engine = create_test_engine();
1777 let handler: Box<dyn WorkflowHandler> = Box::new(EchoWorkflow);
1778 let result = engine.register_boxed(handler);
1779 assert!(result.is_ok());
1780 assert_eq!(engine.handler_names().len(), 1);
1781 }
1782
1783 #[tokio::test]
1784 async fn engine_store_and_provider_accessors() {
1785 let store = Arc::new(InMemoryStore::new());
1786 let inner = ClaudeCodeProvider::new();
1787 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
1788 inner,
1789 "/tmp/ironflow-fixtures",
1790 ));
1791 let engine = Engine::new(store.clone(), provider.clone());
1792
1793 let _ = engine.store();
1795 let _ = engine.provider();
1796 }
1797
1798 use crate::operation::Operation;
1803 use ironflow_store::models::StepKind;
1804 use std::future::Future;
1805 use std::pin::Pin;
1806
1807 struct FakeGitlabOp {
1808 project_id: u64,
1809 title: String,
1810 }
1811
1812 impl Operation for FakeGitlabOp {
1813 fn kind(&self) -> &str {
1814 "gitlab"
1815 }
1816
1817 fn execute(&self) -> Pin<Box<dyn Future<Output = Result<Value, EngineError>> + Send + '_>> {
1818 Box::pin(async move {
1819 Ok(json!({
1820 "issue_id": 42,
1821 "project_id": self.project_id,
1822 "title": self.title,
1823 }))
1824 })
1825 }
1826
1827 fn input(&self) -> Option<Value> {
1828 Some(json!({
1829 "project_id": self.project_id,
1830 "title": self.title,
1831 }))
1832 }
1833 }
1834
1835 struct FailingOp;
1836
1837 impl Operation for FailingOp {
1838 fn kind(&self) -> &str {
1839 "broken-service"
1840 }
1841
1842 fn execute(&self) -> Pin<Box<dyn Future<Output = Result<Value, EngineError>> + Send + '_>> {
1843 Box::pin(async move { Err(EngineError::StepConfig("service unavailable".to_string())) })
1844 }
1845 }
1846
1847 struct OperationWorkflow;
1848
1849 impl WorkflowHandler for OperationWorkflow {
1850 fn name(&self) -> &str {
1851 "operation-workflow"
1852 }
1853
1854 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1855 Box::pin(async move {
1856 let op = FakeGitlabOp {
1857 project_id: 123,
1858 title: "Bug report".to_string(),
1859 };
1860 ctx.operation("create-issue", &op).await?;
1861 Ok(())
1862 })
1863 }
1864 }
1865
1866 struct FailingOperationWorkflow;
1867
1868 impl WorkflowHandler for FailingOperationWorkflow {
1869 fn name(&self) -> &str {
1870 "failing-operation-workflow"
1871 }
1872
1873 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1874 Box::pin(async move {
1875 ctx.operation("broken-call", &FailingOp).await?;
1876 Ok(())
1877 })
1878 }
1879 }
1880
1881 struct MixedWorkflow;
1882
1883 impl WorkflowHandler for MixedWorkflow {
1884 fn name(&self) -> &str {
1885 "mixed-workflow"
1886 }
1887
1888 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1889 Box::pin(async move {
1890 ctx.shell("build", ShellConfig::new("echo built")).await?;
1891 let op = FakeGitlabOp {
1892 project_id: 456,
1893 title: "Deploy done".to_string(),
1894 };
1895 let result = ctx.operation("notify-gitlab", &op).await?;
1896 assert_eq!(result.output["issue_id"], 42);
1897 Ok(())
1898 })
1899 }
1900 }
1901
1902 #[tokio::test]
1903 async fn operation_step_happy_path() {
1904 let mut engine = create_test_engine();
1905 engine.register(OperationWorkflow).unwrap();
1906
1907 let run = engine
1908 .run_handler("operation-workflow", TriggerKind::Manual, json!({}))
1909 .await
1910 .unwrap()
1911 .run;
1912
1913 assert_eq!(run.status.state, RunStatus::Completed);
1914
1915 let steps = engine.store().list_steps(run.id).await.unwrap();
1916
1917 assert_eq!(steps.len(), 1);
1918 assert_eq!(steps[0].name, "create-issue");
1919 assert_eq!(steps[0].kind, StepKind::Custom("gitlab".to_string()));
1920 assert_eq!(
1921 steps[0].status.state,
1922 ironflow_store::models::StepStatus::Completed
1923 );
1924
1925 let output = steps[0].output.as_ref().unwrap();
1926 assert_eq!(output["issue_id"], 42);
1927 assert_eq!(output["project_id"], 123);
1928
1929 let input = steps[0].input.as_ref().unwrap();
1930 assert_eq!(input["project_id"], 123);
1931 assert_eq!(input["title"], "Bug report");
1932 }
1933
1934 #[tokio::test]
1935 async fn operation_step_failure_marks_run_failed() {
1936 let mut engine = create_test_engine();
1937 engine.register(FailingOperationWorkflow).unwrap();
1938
1939 let result = engine
1940 .run_handler("failing-operation-workflow", TriggerKind::Manual, json!({}))
1941 .await;
1942
1943 assert!(result.is_err());
1944 }
1945
1946 #[tokio::test]
1947 async fn operation_mixed_with_shell_steps() {
1948 let mut engine = create_test_engine();
1949 engine.register(MixedWorkflow).unwrap();
1950
1951 let run = engine
1952 .run_handler("mixed-workflow", TriggerKind::Manual, json!({}))
1953 .await
1954 .unwrap()
1955 .run;
1956
1957 assert_eq!(run.status.state, RunStatus::Completed);
1958
1959 let steps = engine.store().list_steps(run.id).await.unwrap();
1960
1961 assert_eq!(steps.len(), 2);
1962 assert_eq!(steps[0].kind, StepKind::Shell);
1963 assert_eq!(steps[1].kind, StepKind::Custom("gitlab".to_string()));
1964 assert_eq!(steps[0].position, 0);
1965 assert_eq!(steps[1].position, 1);
1966 }
1967
1968 use crate::config::ApprovalConfig;
1973
1974 struct SingleApprovalWorkflow;
1975
1976 impl WorkflowHandler for SingleApprovalWorkflow {
1977 fn name(&self) -> &str {
1978 "single-approval"
1979 }
1980
1981 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1982 Box::pin(async move {
1983 ctx.shell("build", ShellConfig::new("echo built")).await?;
1984 ctx.approval("gate", ApprovalConfig::new("OK?")).await?;
1985 ctx.shell("deploy", ShellConfig::new("echo deployed"))
1986 .await?;
1987 Ok(())
1988 })
1989 }
1990 }
1991
1992 struct DoubleApprovalWorkflow;
1993
1994 impl WorkflowHandler for DoubleApprovalWorkflow {
1995 fn name(&self) -> &str {
1996 "double-approval"
1997 }
1998
1999 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2000 Box::pin(async move {
2001 ctx.shell("build", ShellConfig::new("echo built")).await?;
2002 ctx.approval("staging-gate", ApprovalConfig::new("Deploy staging?"))
2003 .await?;
2004 ctx.shell("deploy-staging", ShellConfig::new("echo staging"))
2005 .await?;
2006 ctx.approval("prod-gate", ApprovalConfig::new("Deploy prod?"))
2007 .await?;
2008 ctx.shell("deploy-prod", ShellConfig::new("echo prod"))
2009 .await?;
2010 Ok(())
2011 })
2012 }
2013 }
2014
2015 #[tokio::test]
2016 async fn approval_pauses_run() {
2017 let mut engine = create_test_engine();
2018 engine.register(SingleApprovalWorkflow).unwrap();
2019
2020 let run = engine
2021 .run_handler("single-approval", TriggerKind::Manual, json!({}))
2022 .await
2023 .unwrap()
2024 .run;
2025
2026 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
2027
2028 let steps = engine.store().list_steps(run.id).await.unwrap();
2029 assert_eq!(steps.len(), 2); assert_eq!(steps[0].kind, StepKind::Shell);
2031 assert_eq!(steps[0].status.state, StepStatus::Completed);
2032 assert_eq!(steps[1].kind, StepKind::Approval);
2033 assert_eq!(steps[1].status.state, StepStatus::AwaitingApproval);
2034 }
2035
2036 #[tokio::test]
2037 async fn approval_resume_completes_run() {
2038 let mut engine = create_test_engine();
2039 engine.register(SingleApprovalWorkflow).unwrap();
2040
2041 let run = engine
2043 .run_handler("single-approval", TriggerKind::Manual, json!({}))
2044 .await
2045 .unwrap()
2046 .run;
2047 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
2048
2049 engine
2051 .store()
2052 .update_run_status(run.id, RunStatus::Running)
2053 .await
2054 .unwrap();
2055
2056 let resumed = engine.resume_run(run.id).await.unwrap().run;
2058 assert_eq!(resumed.status.state, RunStatus::Completed);
2059
2060 let steps = engine.store().list_steps(run.id).await.unwrap();
2061 assert_eq!(steps.len(), 3); assert_eq!(steps[0].name, "build");
2063 assert_eq!(steps[0].status.state, StepStatus::Completed);
2064 assert_eq!(steps[1].name, "gate");
2065 assert_eq!(steps[1].kind, StepKind::Approval);
2066 assert_eq!(steps[1].status.state, StepStatus::Completed);
2067 assert_eq!(steps[2].name, "deploy");
2068 assert_eq!(steps[2].status.state, StepStatus::Completed);
2069 }
2070
2071 #[tokio::test]
2072 async fn double_approval_two_resumes() {
2073 let mut engine = create_test_engine();
2074 engine.register(DoubleApprovalWorkflow).unwrap();
2075
2076 let run = engine
2078 .run_handler("double-approval", TriggerKind::Manual, json!({}))
2079 .await
2080 .unwrap()
2081 .run;
2082 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
2083
2084 let steps = engine.store().list_steps(run.id).await.unwrap();
2085 assert_eq!(steps.len(), 2); engine
2089 .store()
2090 .update_run_status(run.id, RunStatus::Running)
2091 .await
2092 .unwrap();
2093
2094 let resumed = engine.resume_run(run.id).await.unwrap().run;
2095 assert_eq!(resumed.status.state, RunStatus::AwaitingApproval);
2096
2097 let steps = engine.store().list_steps(run.id).await.unwrap();
2098 assert_eq!(steps.len(), 4); engine
2102 .store()
2103 .update_run_status(run.id, RunStatus::Running)
2104 .await
2105 .unwrap();
2106
2107 let final_run = engine.resume_run(run.id).await.unwrap().run;
2108 assert_eq!(final_run.status.state, RunStatus::Completed);
2109
2110 let steps = engine.store().list_steps(run.id).await.unwrap();
2111 assert_eq!(steps.len(), 5);
2112 assert_eq!(steps[0].name, "build");
2113 assert_eq!(steps[1].name, "staging-gate");
2114 assert_eq!(steps[2].name, "deploy-staging");
2115 assert_eq!(steps[3].name, "prod-gate");
2116 assert_eq!(steps[4].name, "deploy-prod");
2117
2118 for step in &steps {
2119 assert_eq!(step.status.state, StepStatus::Completed);
2120 }
2121 }
2122
2123 use ironflow_store::models::{NewStep, StepUpdate, step_trace_id};
2128
2129 async fn create_step_with_status(
2130 store: &Arc<dyn Store>,
2131 run_id: Uuid,
2132 name: &str,
2133 position: u32,
2134 status: StepStatus,
2135 ) -> ironflow_store::models::Step {
2136 let step = store
2137 .create_step(NewStep {
2138 run_id,
2139 trace_id: step_trace_id(run_id, name, position),
2140 name: name.to_string(),
2141 kind: StepKind::Shell,
2142 position,
2143 input: None,
2144 is_error_handler: false,
2145 })
2146 .await
2147 .unwrap();
2148
2149 match status {
2150 StepStatus::Pending => {}
2151 StepStatus::Running => {
2152 store
2153 .update_step(
2154 step.id,
2155 StepUpdate {
2156 status: Some(StepStatus::Running),
2157 ..StepUpdate::default()
2158 },
2159 )
2160 .await
2161 .unwrap();
2162 }
2163 StepStatus::Completed => {
2164 store
2165 .update_step(
2166 step.id,
2167 StepUpdate {
2168 status: Some(StepStatus::Running),
2169 ..StepUpdate::default()
2170 },
2171 )
2172 .await
2173 .unwrap();
2174 store
2175 .update_step(
2176 step.id,
2177 StepUpdate {
2178 status: Some(StepStatus::Completed),
2179 ..StepUpdate::default()
2180 },
2181 )
2182 .await
2183 .unwrap();
2184 }
2185 StepStatus::AwaitingApproval => {
2186 store
2187 .update_step(
2188 step.id,
2189 StepUpdate {
2190 status: Some(StepStatus::Running),
2191 ..StepUpdate::default()
2192 },
2193 )
2194 .await
2195 .unwrap();
2196 store
2197 .update_step(
2198 step.id,
2199 StepUpdate {
2200 status: Some(StepStatus::AwaitingApproval),
2201 ..StepUpdate::default()
2202 },
2203 )
2204 .await
2205 .unwrap();
2206 }
2207 _ => panic!("unsupported status for test helper: {status}"),
2208 }
2209
2210 store.get_step(step.id).await.unwrap().unwrap()
2211 }
2212
2213 #[tokio::test]
2214 async fn fail_orphaned_steps_marks_running_as_failed() {
2215 let engine = create_test_engine();
2216 let run = engine
2217 .store()
2218 .create_run(NewRun {
2219 created_by: None,
2220 workflow_name: "test".to_string(),
2221 trigger: TriggerKind::Manual,
2222 payload: json!({}),
2223 max_retries: 0,
2224 handler_version: None,
2225 labels: HashMap::new(),
2226 scheduled_at: None,
2227 idempotency_key: None,
2228 max_cost_usd: None,
2229 })
2230 .await
2231 .unwrap()
2232 .into_run();
2233
2234 let step = create_step_with_status(
2235 engine.store(),
2236 run.id,
2237 "running-step",
2238 0,
2239 StepStatus::Running,
2240 )
2241 .await;
2242
2243 engine
2244 .fail_orphaned_steps(run.id, "parent run timed out")
2245 .await
2246 .unwrap();
2247
2248 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
2249 assert_eq!(updated.status.state, StepStatus::Failed);
2250 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
2251 assert!(updated.completed_at.is_some());
2252 }
2253
2254 #[tokio::test]
2255 async fn fail_orphaned_steps_marks_pending_as_skipped() {
2256 let engine = create_test_engine();
2257 let run = engine
2258 .store()
2259 .create_run(NewRun {
2260 created_by: None,
2261 workflow_name: "test".to_string(),
2262 trigger: TriggerKind::Manual,
2263 payload: json!({}),
2264 max_retries: 0,
2265 handler_version: None,
2266 labels: HashMap::new(),
2267 scheduled_at: None,
2268 idempotency_key: None,
2269 max_cost_usd: None,
2270 })
2271 .await
2272 .unwrap()
2273 .into_run();
2274
2275 let step = create_step_with_status(
2276 engine.store(),
2277 run.id,
2278 "pending-step",
2279 0,
2280 StepStatus::Pending,
2281 )
2282 .await;
2283
2284 engine
2285 .fail_orphaned_steps(run.id, "parent run timed out")
2286 .await
2287 .unwrap();
2288
2289 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
2290 assert_eq!(updated.status.state, StepStatus::Skipped);
2291 assert!(updated.error.is_none());
2292 assert!(updated.completed_at.is_some());
2293 }
2294
2295 #[tokio::test]
2296 async fn fail_orphaned_steps_marks_awaiting_approval_as_failed() {
2297 let engine = create_test_engine();
2298 let run = engine
2299 .store()
2300 .create_run(NewRun {
2301 created_by: None,
2302 workflow_name: "test".to_string(),
2303 trigger: TriggerKind::Manual,
2304 payload: json!({}),
2305 max_retries: 0,
2306 handler_version: None,
2307 labels: HashMap::new(),
2308 scheduled_at: None,
2309 idempotency_key: None,
2310 max_cost_usd: None,
2311 })
2312 .await
2313 .unwrap()
2314 .into_run();
2315
2316 let step = create_step_with_status(
2317 engine.store(),
2318 run.id,
2319 "approval-step",
2320 0,
2321 StepStatus::AwaitingApproval,
2322 )
2323 .await;
2324
2325 engine
2326 .fail_orphaned_steps(run.id, "parent run timed out")
2327 .await
2328 .unwrap();
2329
2330 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
2331 assert_eq!(updated.status.state, StepStatus::Failed);
2332 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
2333 assert!(updated.completed_at.is_some());
2334 }
2335
2336 #[tokio::test]
2337 async fn fail_orphaned_steps_skips_terminal_steps() {
2338 let engine = create_test_engine();
2339 let run = engine
2340 .store()
2341 .create_run(NewRun {
2342 created_by: None,
2343 workflow_name: "test".to_string(),
2344 trigger: TriggerKind::Manual,
2345 payload: json!({}),
2346 max_retries: 0,
2347 handler_version: None,
2348 labels: HashMap::new(),
2349 scheduled_at: None,
2350 idempotency_key: None,
2351 max_cost_usd: None,
2352 })
2353 .await
2354 .unwrap()
2355 .into_run();
2356
2357 let completed_step =
2358 create_step_with_status(engine.store(), run.id, "done", 0, StepStatus::Completed).await;
2359 let running_step =
2360 create_step_with_status(engine.store(), run.id, "in-flight", 1, StepStatus::Running)
2361 .await;
2362
2363 engine
2364 .fail_orphaned_steps(run.id, "parent run timed out")
2365 .await
2366 .unwrap();
2367
2368 let completed = engine
2369 .store()
2370 .get_step(completed_step.id)
2371 .await
2372 .unwrap()
2373 .unwrap();
2374 assert_eq!(completed.status.state, StepStatus::Completed);
2375
2376 let failed = engine
2377 .store()
2378 .get_step(running_step.id)
2379 .await
2380 .unwrap()
2381 .unwrap();
2382 assert_eq!(failed.status.state, StepStatus::Failed);
2383 }
2384
2385 #[tokio::test]
2386 async fn fail_orphaned_steps_mixed_states() {
2387 let engine = create_test_engine();
2388 let run = engine
2389 .store()
2390 .create_run(NewRun {
2391 created_by: None,
2392 workflow_name: "test".to_string(),
2393 trigger: TriggerKind::Manual,
2394 payload: json!({}),
2395 max_retries: 0,
2396 handler_version: None,
2397 labels: HashMap::new(),
2398 scheduled_at: None,
2399 idempotency_key: None,
2400 max_cost_usd: None,
2401 })
2402 .await
2403 .unwrap()
2404 .into_run();
2405
2406 let s_completed =
2407 create_step_with_status(engine.store(), run.id, "step-1", 0, StepStatus::Completed)
2408 .await;
2409 let s_running =
2410 create_step_with_status(engine.store(), run.id, "step-2", 1, StepStatus::Running).await;
2411 let s_pending =
2412 create_step_with_status(engine.store(), run.id, "step-3", 2, StepStatus::Pending).await;
2413
2414 engine.fail_orphaned_steps(run.id, "timeout").await.unwrap();
2415
2416 let r_completed = engine
2417 .store()
2418 .get_step(s_completed.id)
2419 .await
2420 .unwrap()
2421 .unwrap();
2422 assert_eq!(r_completed.status.state, StepStatus::Completed);
2423
2424 let r_running = engine
2425 .store()
2426 .get_step(s_running.id)
2427 .await
2428 .unwrap()
2429 .unwrap();
2430 assert_eq!(r_running.status.state, StepStatus::Failed);
2431 assert_eq!(r_running.error.as_deref(), Some("timeout"));
2432
2433 let r_pending = engine
2434 .store()
2435 .get_step(s_pending.id)
2436 .await
2437 .unwrap()
2438 .unwrap();
2439 assert_eq!(r_pending.status.state, StepStatus::Skipped);
2440 assert!(r_pending.error.is_none());
2441 }
2442
2443 #[tokio::test]
2444 async fn fail_orphaned_steps_no_steps_is_noop() {
2445 let engine = create_test_engine();
2446 let run = engine
2447 .store()
2448 .create_run(NewRun {
2449 created_by: None,
2450 workflow_name: "test".to_string(),
2451 trigger: TriggerKind::Manual,
2452 payload: json!({}),
2453 max_retries: 0,
2454 handler_version: None,
2455 labels: HashMap::new(),
2456 scheduled_at: None,
2457 idempotency_key: None,
2458 max_cost_usd: None,
2459 })
2460 .await
2461 .unwrap()
2462 .into_run();
2463
2464 let result = engine.fail_orphaned_steps(run.id, "timeout").await;
2465 assert!(result.is_ok());
2466 }
2467
2468 #[tokio::test]
2469 async fn fail_orphaned_steps_preserves_existing_error() {
2470 let engine = create_test_engine();
2471 let run = engine
2472 .store()
2473 .create_run(NewRun {
2474 created_by: None,
2475 workflow_name: "test".to_string(),
2476 trigger: TriggerKind::Manual,
2477 payload: json!({}),
2478 max_retries: 0,
2479 handler_version: None,
2480 labels: HashMap::new(),
2481 scheduled_at: None,
2482 idempotency_key: None,
2483 max_cost_usd: None,
2484 })
2485 .await
2486 .unwrap()
2487 .into_run();
2488
2489 let step_with_error = create_step_with_status(
2490 engine.store(),
2491 run.id,
2492 "already-errored",
2493 0,
2494 StepStatus::Running,
2495 )
2496 .await;
2497
2498 engine
2499 .store()
2500 .update_step(
2501 step_with_error.id,
2502 StepUpdate {
2503 error: Some("real error from provider".to_string()),
2504 ..StepUpdate::default()
2505 },
2506 )
2507 .await
2508 .unwrap();
2509
2510 let step_no_error = create_step_with_status(
2511 engine.store(),
2512 run.id,
2513 "no-error-yet",
2514 1,
2515 StepStatus::Running,
2516 )
2517 .await;
2518
2519 engine
2520 .fail_orphaned_steps(run.id, "parent run failed")
2521 .await
2522 .unwrap();
2523
2524 let updated_with = engine
2525 .store()
2526 .get_step(step_with_error.id)
2527 .await
2528 .unwrap()
2529 .unwrap();
2530 assert_eq!(updated_with.status.state, StepStatus::Failed);
2531 assert_eq!(
2532 updated_with.error.as_deref(),
2533 Some("real error from provider"),
2534 );
2535
2536 let updated_without = engine
2537 .store()
2538 .get_step(step_no_error.id)
2539 .await
2540 .unwrap()
2541 .unwrap();
2542 assert_eq!(updated_without.status.state, StepStatus::Failed);
2543 assert_eq!(updated_without.error.as_deref(), Some("parent run failed"),);
2544 }
2545}