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