1use std::collections::{HashMap, HashSet};
10use std::fmt;
11use std::sync::{Arc, Mutex};
12use std::time::{Duration, Instant};
13
14use chrono::{DateTime, TimeDelta, Utc};
15use rust_decimal::Decimal;
16use serde_json::{Value, to_value};
17use tokio::spawn;
18use tokio::sync::{Mutex as AsyncMutex, OwnedMutexGuard};
19use tokio::time::timeout;
20use tracing::{error, info, warn};
21use uuid::Uuid;
22
23use ironflow_core::error::OperationError;
24#[cfg(feature = "prometheus")]
25use ironflow_core::metric_names::{
26 RUN_BUDGET_EXCEEDED_TOTAL, RUN_COST_USD, RUN_DURATION_SECONDS, RUNS_ACTIVE, RUNS_TOTAL,
27};
28use ironflow_core::provider::{AgentProvider, LABEL_ROOT_RUN_ID};
29use ironflow_store::error::StoreError;
30use ironflow_store::models::{
31 ConcurrencyLimit, LeaseUpdate, NewRun, NewSignal, ProviderKind, Run, RunActor, RunCreation,
32 RunFilter, RunStatus, RunUpdate, SignalInsert, SignalStepResolution, StepKind, StepStatus,
33 StepUpdate, TriggerKind, normalize_worker_tags, validate_concurrency_limits, validate_priority,
34 validate_worker_tags,
35};
36use ironflow_store::store::Store;
37#[cfg(feature = "prometheus")]
38use metrics::{counter, gauge, histogram};
39
40use crate::artifact::ArtifactSink;
41use crate::budget::{BudgetConfig, month_start};
42use crate::chain::ChainClaim;
43use crate::context::{PARENT_RUN_ID_LABEL, WorkflowContext, interrupt_running_steps};
44use crate::error::EngineError;
45use crate::executor::{StepInterceptor, StepOutput, StepResult};
46use crate::guard::{WorkflowGuardConfig, WorkflowGuardState};
47use crate::handler::{WorkflowHandler, WorkflowInfo, clamp_priority};
48use crate::log_sender::LogSender;
49use crate::notify::{
50 ApprovalRequestedEvent, Event, EventBusReleaser, EventPublisher, EventSubscriber,
51 RunBudgetExceededEvent, RunFailedEvent, RunStatusChangedEvent, SignalAwaitedEvent,
52 SignalReceivedEvent, WorkflowEventBus,
53};
54use crate::plan::{
55 ExecutionPlan, PlanOptions, PlanRecorder, SharedPlanRecorder, estimate_durations, lock_plan,
56};
57use crate::retry_policy::{backoff_for_retry, is_run_retryable};
58use crate::schedule::CronSchedule;
59use crate::signal::{
60 Signal, SignalDelivery, SignalOrigin, SignalRejected, SignalResumed, SignalSendFuture,
61 SignalSender, received_output, validate_signal, validate_step_payload,
62};
63use ironflow_core::decision::DecisionProvider;
64
65const LIST_RUNS_PAGE_SIZE: u32 = 100;
67
68#[derive(Debug, Clone)]
87pub struct WorkflowResult {
88 pub run: Run,
90 pub steps: Vec<StepResult>,
92}
93
94#[derive(Debug, Clone, Default)]
113pub struct EnqueueOptions {
114 pub max_retries: u32,
116 pub labels: HashMap<String, String>,
118 pub scheduled_at: Option<DateTime<Utc>>,
121 pub max_cost_usd: Option<Decimal>,
125 pub created_by: Option<RunActor>,
128 pub idempotency_key: Option<String>,
134 pub concurrency_key: Option<String>,
140 pub concurrency_limits: Vec<ConcurrencyLimit>,
148 pub priority: Option<i16>,
157 pub worker_tags: Vec<String>,
163}
164
165#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
176pub enum ExecutionMode {
177 #[default]
181 Local,
182 Workers,
185}
186
187pub struct Engine {
228 store: Arc<dyn Store>,
229 provider: Arc<dyn AgentProvider>,
230 handlers: HashMap<String, Arc<dyn WorkflowHandler>>,
231 event_publisher: EventPublisher,
232 log_sender: Option<LogSender>,
233 budget: BudgetConfig,
234 artifact_sink: Option<Arc<dyn ArtifactSink>>,
235 guard_config: Option<WorkflowGuardConfig>,
236 event_bus: Option<WorkflowEventBus>,
237 decision_provider: Option<Arc<dyn DecisionProvider>>,
238 step_interceptor: Option<Arc<dyn StepInterceptor>>,
239 execution_mode: ExecutionMode,
240 worker_tags: Option<Arc<Vec<String>>>,
241 signal_sender: Option<Arc<dyn SignalSender>>,
243 active_runs: ActiveRuns,
244 local_resumes: ActiveRuns,
245}
246
247type ActiveRuns = Arc<Mutex<HashMap<Uuid, Arc<AsyncMutex<()>>>>>;
249
250struct ActiveRunGuard {
256 runs: ActiveRuns,
257 run_id: Uuid,
258 entry: Arc<AsyncMutex<()>>,
259 lock: Option<OwnedMutexGuard<()>>,
260}
261
262impl Drop for ActiveRunGuard {
263 fn drop(&mut self) {
264 self.lock.take();
265 prune_active_run(&self.runs, self.run_id, &self.entry);
266 }
267}
268
269async fn lock_run(runs: &ActiveRuns, run_id: Uuid) -> ActiveRunGuard {
271 let entry = {
272 let mut runs = runs.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
273 Arc::clone(runs.entry(run_id).or_default())
274 };
275 let lock = Arc::clone(&entry).lock_owned().await;
276 ActiveRunGuard {
277 runs: Arc::clone(runs),
278 run_id,
279 entry,
280 lock: Some(lock),
281 }
282}
283
284fn prune_active_run(runs: &ActiveRuns, run_id: Uuid, entry: &Arc<AsyncMutex<()>>) {
289 let mut runs = runs.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
290 let unused = runs
291 .get(&run_id)
292 .is_some_and(|current| Arc::ptr_eq(current, entry))
293 && Arc::strong_count(entry) == 2;
294 if unused {
295 runs.remove(&run_id);
296 }
297}
298
299struct InProcessSignalSender(Arc<Engine>);
301
302impl SignalSender for InProcessSignalSender {
303 fn send(&self, signal: NewSignal, origin: SignalOrigin) -> SignalSendFuture<'_> {
304 Box::pin(self.0.deliver_signal_from(signal, Some(origin)))
305 }
306}
307
308fn validate_category(handler_name: &str, category: &str) -> Result<(), EngineError> {
318 let reject = |reason: &str| {
319 Err(EngineError::InvalidWorkflow(format!(
320 "handler '{handler_name}' has invalid category '{category}': {reason}"
321 )))
322 };
323
324 if category.is_empty() {
325 return reject("empty category");
326 }
327 if category.starts_with('/') {
328 return reject("leading '/'");
329 }
330 if category.ends_with('/') {
331 return reject("trailing '/'");
332 }
333 for segment in category.split('/') {
334 if segment.is_empty() {
335 return reject("empty segment (double '/')");
336 }
337 if segment.trim().is_empty() {
338 return reject("whitespace-only segment");
339 }
340 }
341 Ok(())
342}
343
344fn chain_label(run: &Run, key: &str) -> Option<Uuid> {
349 if !matches!(run.trigger, TriggerKind::Workflow) {
350 return None;
351 }
352 let id = Uuid::parse_str(run.labels.get(key)?).ok()?;
353 (id != run.id).then_some(id)
354}
355
356pub fn chain_root(run: &Run) -> Option<Uuid> {
376 chain_label(run, LABEL_ROOT_RUN_ID)
377}
378
379pub(crate) fn chain_parent(run: &Run) -> Option<Uuid> {
381 chain_label(run, PARENT_RUN_ID_LABEL)
382}
383
384impl Engine {
385 pub fn new(store: Arc<dyn Store>, provider: Arc<dyn AgentProvider>) -> Self {
401 Self {
402 store,
403 provider,
404 handlers: HashMap::new(),
405 event_publisher: EventPublisher::new(),
406 log_sender: None,
407 budget: BudgetConfig::new(),
408 artifact_sink: None,
409 guard_config: None,
410 event_bus: None,
411 decision_provider: None,
412 step_interceptor: None,
413 execution_mode: ExecutionMode::default(),
414 worker_tags: None,
415 signal_sender: None,
416 active_runs: Arc::new(Mutex::new(HashMap::new())),
417 local_resumes: Arc::new(Mutex::new(HashMap::new())),
418 }
419 }
420
421 async fn track_execution(&self, run_id: Uuid) -> ActiveRunGuard {
424 lock_run(&self.active_runs, run_id).await
425 }
426
427 async fn wait_until_idle(&self, run_id: Uuid) {
429 let entry = {
430 let runs = self
431 .active_runs
432 .lock()
433 .unwrap_or_else(|poisoned| poisoned.into_inner());
434 runs.get(&run_id).map(Arc::clone)
435 };
436 if let Some(entry) = entry {
437 drop(Arc::clone(&entry).lock_owned().await);
438 prune_active_run(&self.active_runs, run_id, &entry);
439 }
440 }
441
442 pub(crate) fn is_executing(&self, run_id: Uuid) -> bool {
444 let runs = self
445 .active_runs
446 .lock()
447 .unwrap_or_else(|poisoned| poisoned.into_inner());
448 runs.get(&run_id)
449 .is_some_and(|entry| entry.try_lock().is_err())
450 }
451
452 pub fn with_signal_sender(mut self, sender: Arc<dyn SignalSender>) -> Self {
471 self.signal_sender = Some(sender);
472 self
473 }
474
475 pub fn with_decision_provider(mut self, provider: Arc<dyn DecisionProvider>) -> Self {
496 self.decision_provider = Some(provider);
497 self
498 }
499
500 pub fn with_step_interceptor(mut self, interceptor: Arc<dyn StepInterceptor>) -> Self {
524 self.step_interceptor = Some(interceptor);
525 self
526 }
527
528 pub fn step_interceptor(&self) -> Option<&Arc<dyn StepInterceptor>> {
530 self.step_interceptor.as_ref()
531 }
532
533 pub fn with_budget_config(mut self, budget: BudgetConfig) -> Self {
554 self.budget = budget;
555 self
556 }
557
558 pub fn budget_config(&self) -> &BudgetConfig {
560 &self.budget
561 }
562
563 pub fn with_guard_config(mut self, config: WorkflowGuardConfig) -> Self {
585 self.guard_config = Some(config);
586 self
587 }
588
589 pub fn guard_config(&self) -> Option<&WorkflowGuardConfig> {
591 self.guard_config.as_ref()
592 }
593
594 pub fn with_execution_mode(mut self, mode: ExecutionMode) -> Self {
615 self.execution_mode = mode;
616 self
617 }
618
619 pub fn execution_mode(&self) -> ExecutionMode {
621 self.execution_mode
622 }
623
624 pub fn set_worker_tags(&mut self, tags: Vec<String>) {
647 self.worker_tags = Some(Arc::new(tags));
648 }
649
650 pub fn worker_tags(&self) -> Option<&[String]> {
668 self.worker_tags.as_deref().map(Vec::as_slice)
669 }
670
671 pub fn set_log_sender(&mut self, sender: LogSender) {
677 self.log_sender = Some(sender);
678 }
679
680 pub fn set_artifact_sink(&mut self, sink: Arc<dyn ArtifactSink>) {
698 self.artifact_sink = Some(sink);
699 }
700
701 pub fn artifact_sink(&self) -> Option<&Arc<dyn ArtifactSink>> {
703 self.artifact_sink.as_ref()
704 }
705
706 pub fn set_event_bus(&mut self, bus: WorkflowEventBus) {
727 self.event_bus = Some(bus.clone());
728 self.event_publisher
729 .subscribe(EventBusReleaser::new(bus), &[Event::RUN_STATUS_CHANGED]);
730 }
731
732 pub fn event_bus(&self) -> Option<&WorkflowEventBus> {
734 self.event_bus.as_ref()
735 }
736
737 pub fn store(&self) -> &Arc<dyn Store> {
739 &self.store
740 }
741
742 pub fn provider(&self) -> &Arc<dyn AgentProvider> {
744 &self.provider
745 }
746
747 fn build_context(&self, run: &Run) -> WorkflowContext {
756 let handlers = self.handlers.clone();
757 let resolver: crate::context::HandlerResolver =
758 Arc::new(move |name: &str| handlers.get(name).cloned());
759 let mut ctx = WorkflowContext::with_handler_resolver(
760 run.id,
761 run.workflow_name.clone(),
762 self.store.clone(),
763 self.provider.clone(),
764 resolver,
765 );
766 ctx.carry_over_run_totals(run.retry_count + 1, run.cost_usd, run.duration_ms);
767 ctx.set_max_cost_usd(run.max_cost_usd);
768 ctx.set_run_created_at(run.created_at);
769 if let Some(ref sender) = self.log_sender {
770 ctx.set_log_sender(sender.clone());
771 }
772 if let Some(ref sink) = self.artifact_sink {
773 ctx.set_artifact_sink(sink.clone());
774 }
775 if let Some(ref bus) = self.event_bus {
776 ctx.set_event_bus(bus.clone());
777 }
778 if let Some(ref provider) = self.decision_provider {
779 ctx.set_decision_provider(provider.clone());
780 }
781 if let Some(ref interceptor) = self.step_interceptor {
782 ctx.set_step_interceptor(interceptor.clone());
783 }
784 if let Some(ref tags) = self.worker_tags {
785 ctx.set_worker_tags(tags.clone());
786 }
787 let sender = self
790 .signal_sender
791 .clone()
792 .unwrap_or_else(|| Arc::new(InProcessSignalSender(Arc::new(self.shared_copy()))));
793 ctx.set_signal_sender(sender);
794 ctx
795 }
796
797 fn shared_copy(&self) -> Engine {
800 Engine {
801 store: self.store.clone(),
802 provider: self.provider.clone(),
803 handlers: self.handlers.clone(),
804 event_publisher: self.event_publisher.clone(),
805 log_sender: self.log_sender.clone(),
806 budget: self.budget.clone(),
807 artifact_sink: self.artifact_sink.clone(),
808 guard_config: self.guard_config.clone(),
809 event_bus: self.event_bus.clone(),
810 decision_provider: self.decision_provider.clone(),
811 step_interceptor: self.step_interceptor.clone(),
812 execution_mode: self.execution_mode,
813 worker_tags: self.worker_tags.clone(),
814 signal_sender: self.signal_sender.clone(),
815 active_runs: self.active_runs.clone(),
816 local_resumes: self.local_resumes.clone(),
817 }
818 }
819
820 pub(crate) fn effective_guard_config(
823 &self,
824 handler: &dyn WorkflowHandler,
825 ) -> Option<WorkflowGuardConfig> {
826 handler.guard_config().or_else(|| self.guard_config.clone())
827 }
828
829 async fn build_context_with_guard(
848 &self,
849 run: &Run,
850 handler: &dyn WorkflowHandler,
851 ) -> Result<WorkflowContext, EngineError> {
852 let mut ctx = self.build_context(run);
853 if let Some(config) = self.effective_guard_config(handler) {
854 let mut state = WorkflowGuardState::new()
855 .with_tokens_used(self.tokens_used_by_run_tree(run).await?)
856 .with_elapsed(Duration::from_millis(run.duration_ms));
857 if let Some(at) = config.deadline_for(run.created_at) {
858 state = state.with_deadline_at(at);
859 }
860 ctx.set_guard(config, Arc::new(Mutex::new(state)));
861 }
862 Ok(ctx)
863 }
864
865 async fn fail_before_execution(&self, run_id: Uuid, err: EngineError) -> EngineError {
871 let retryable = is_run_retryable(&err);
872 if let Err(store_err) = self
873 .fail_or_schedule_retry(run_id, &err.to_string(), retryable, None, None)
874 .await
875 {
876 error!(run_id = %run_id, store_error = %store_err, "failed to persist run failure");
877 }
878 err
879 }
880
881 pub(crate) async fn list_all_runs(&self, filter: RunFilter) -> Result<Vec<Run>, EngineError> {
890 let mut runs = Vec::new();
891 let mut page = 1;
892 loop {
893 let result = self
894 .store
895 .list_runs(filter.clone(), page, LIST_RUNS_PAGE_SIZE)
896 .await?;
897 if result.items.is_empty() {
898 break;
899 }
900 runs.extend(result.items);
901 if runs.len() as u64 >= result.total {
902 break;
903 }
904 page += 1;
905 }
906 Ok(runs)
907 }
908
909 async fn tokens_used_by_run_tree(&self, root: &Run) -> Result<u64, EngineError> {
918 let tree_label = HashMap::from([(LABEL_ROOT_RUN_ID.to_string(), root.id.to_string())]);
919 let filter = RunFilter {
920 labels: Some(tree_label),
921 ..RunFilter::default()
922 };
923 let mut run_ids = vec![root.id];
924 run_ids.extend(
925 self.list_all_runs(filter)
926 .await?
927 .into_iter()
928 .map(|run| run.id),
929 );
930
931 let mut tokens: u64 = 0;
932 for run_id in run_ids {
933 for step in self.store.list_steps(run_id).await? {
934 if step.kind == StepKind::Agent {
935 let step_tokens = StepOutput::from(&step).total_tokens();
936 tokens = tokens.saturating_add(step_tokens);
937 }
938 }
939 }
940 Ok(tokens)
941 }
942
943 pub async fn check_monthly_quota(&self, workflow_name: &str) -> Result<(), EngineError> {
970 let Some(limit) = self.budget.monthly_cost_limit_usd else {
971 return Ok(());
972 };
973
974 let stats = self
975 .store
976 .get_stats(RunFilter {
977 created_after: Some(month_start(Utc::now())),
978 ..RunFilter::default()
979 })
980 .await?;
981
982 if stats.total_cost_usd < limit {
983 return Ok(());
984 }
985
986 warn!(
987 workflow = %workflow_name,
988 limit_usd = %limit,
989 spent_usd = %stats.total_cost_usd,
990 "monthly cost quota exhausted, refusing new run"
991 );
992
993 #[cfg(feature = "prometheus")]
994 counter!(
995 RUN_BUDGET_EXCEEDED_TOTAL,
996 "workflow" => workflow_name.to_string(),
997 "scope" => "monthly",
998 )
999 .increment(1);
1000
1001 Err(EngineError::MonthlyBudgetExceeded {
1002 limit_usd: limit,
1003 spent_usd: stats.total_cost_usd,
1004 })
1005 }
1006
1007 pub fn register(&mut self, handler: impl WorkflowHandler + 'static) -> Result<(), EngineError> {
1051 let name = handler.name().to_string();
1052 if self.handlers.contains_key(&name) {
1053 return Err(EngineError::InvalidWorkflow(format!(
1054 "handler '{}' already registered",
1055 name
1056 )));
1057 }
1058 if let Some(category) = handler.category() {
1059 validate_category(&name, category)?;
1060 }
1061 self.handlers.insert(name, Arc::new(handler));
1062 Ok(())
1063 }
1064
1065 pub fn register_boxed(&mut self, handler: Box<dyn WorkflowHandler>) -> Result<(), EngineError> {
1072 let name = handler.name().to_string();
1073 if self.handlers.contains_key(&name) {
1074 return Err(EngineError::InvalidWorkflow(format!(
1075 "handler '{}' already registered",
1076 name
1077 )));
1078 }
1079 if let Some(category) = handler.category() {
1080 validate_category(&name, category)?;
1081 }
1082 self.handlers.insert(name, Arc::from(handler));
1083 Ok(())
1084 }
1085
1086 pub fn get_handler(&self, name: &str) -> Option<&Arc<dyn WorkflowHandler>> {
1088 self.handlers.get(name)
1089 }
1090
1091 pub fn handler_names(&self) -> Vec<&str> {
1093 self.handlers.keys().map(|s| s.as_str()).collect()
1094 }
1095
1096 pub fn handler_info(&self, name: &str) -> Option<WorkflowInfo> {
1098 self.handlers.get(name).map(|h| h.describe())
1099 }
1100
1101 pub fn scheduled_handlers(&self) -> Vec<(&str, &CronSchedule)> {
1125 self.handlers
1126 .iter()
1127 .filter_map(|(name, handler)| handler.schedule().map(|sched| (name.as_str(), sched)))
1128 .collect()
1129 }
1130
1131 pub fn subscribe(
1156 &mut self,
1157 subscriber: impl EventSubscriber + 'static,
1158 event_types: &[&'static str],
1159 ) {
1160 self.event_publisher.subscribe(subscriber, event_types);
1161 }
1162
1163 pub fn event_publisher(&self) -> &EventPublisher {
1168 &self.event_publisher
1169 }
1170
1171 #[tracing::instrument(name = "engine.run_handler", skip_all, fields(workflow = %handler_name))]
1207 pub async fn run_handler(
1208 &self,
1209 handler_name: &str,
1210 trigger: TriggerKind,
1211 payload: Value,
1212 ) -> Result<WorkflowResult, EngineError> {
1213 let handler = self
1214 .handlers
1215 .get(handler_name)
1216 .ok_or_else(|| {
1217 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1218 })?
1219 .clone();
1220
1221 self.check_monthly_quota(handler_name).await?;
1222
1223 let handler_version = handler.version().map(str::to_string);
1224 let max_cost_usd = self
1225 .budget
1226 .resolve_run_cap(None, handler.default_max_cost_usd());
1227 let run = self
1228 .store
1229 .create_run(NewRun {
1230 created_by: None,
1231 workflow_name: handler_name.to_string(),
1232 trigger,
1233 payload,
1234 max_retries: 0,
1235 handler_version,
1236 labels: handler.default_labels(),
1237 scheduled_at: None,
1238 idempotency_key: None,
1239 concurrency_key: None,
1240 priority: clamp_priority(handler.priority()),
1241 concurrency_limits: Vec::new(),
1242 max_cost_usd,
1243 worker_tags: normalize_worker_tags(handler.required_worker_tags()),
1244 })
1245 .await?
1246 .into_run();
1247
1248 let run_id = run.id;
1249 info!(run_id = %run_id, handler_version = run.handler_version.as_deref().unwrap_or(""), "run created");
1250
1251 if self.is_workflow_paused(handler_name).await? {
1252 info!(run_id = %run_id, workflow = %handler_name, "workflow paused, run left pending");
1253 return Ok(WorkflowResult {
1254 run,
1255 steps: Vec::new(),
1256 });
1257 }
1258
1259 let mut ctx = self
1263 .build_context_with_guard(&run, handler.as_ref())
1264 .await?;
1265
1266 self.store
1267 .update_run_status(run_id, RunStatus::Running)
1268 .await?;
1269
1270 #[cfg(feature = "prometheus")]
1271 gauge!(RUNS_ACTIVE, "workflow" => handler_name.to_string()).increment(1.0);
1272
1273 let run_start = Instant::now();
1274 let result = handler.execute(&mut ctx).await;
1275 self.finalize_run(run_id, handler_name, result, &ctx, run_start, run.labels)
1276 .await
1277 }
1278
1279 #[tracing::instrument(name = "engine.plan_handler", skip_all, fields(workflow = %handler_name))]
1328 pub async fn plan_handler(
1329 &self,
1330 handler_name: &str,
1331 payload: Value,
1332 options: PlanOptions,
1333 ) -> Result<ExecutionPlan, EngineError> {
1334 if options.max_depth == 0 {
1335 return Err(EngineError::InvalidWorkflow(
1336 "max_depth must be at least 1".to_string(),
1337 ));
1338 }
1339
1340 let handler = self
1341 .handlers
1342 .get(handler_name)
1343 .ok_or_else(|| {
1344 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1345 })?
1346 .clone();
1347
1348 let estimates = if options.estimate_durations {
1349 estimate_durations(&self.store, handler_name, options.sample_runs).await?
1350 } else {
1351 HashMap::new()
1352 };
1353
1354 let shared: SharedPlanRecorder = Arc::new(Mutex::new(PlanRecorder::new(
1355 handler_name.to_string(),
1356 payload,
1357 options.max_depth,
1358 estimates,
1359 )));
1360
1361 let handlers = self.handlers.clone();
1364 let resolver: crate::context::HandlerResolver =
1365 Arc::new(move |name: &str| handlers.get(name).cloned());
1366 let mut ctx = WorkflowContext::with_handler_resolver(
1367 Uuid::now_v7(),
1368 handler_name.to_string(),
1369 self.store.clone(),
1370 self.provider.clone(),
1371 resolver,
1372 );
1373 ctx.set_plan(shared.clone());
1374
1375 match timeout(options.timeout, handler.execute(&mut ctx)).await {
1376 Ok(Ok(())) => {}
1377 Ok(Err(err)) => lock_plan(&shared).fail(err.to_string()),
1378 Err(_) => lock_plan(&shared).fail(format!(
1379 "plan timed out after {} s",
1380 options.timeout.as_secs_f64()
1381 )),
1382 }
1383 drop(ctx);
1384
1385 let plan = match Arc::try_unwrap(shared) {
1386 Ok(mutex) => mutex
1387 .into_inner()
1388 .unwrap_or_else(|poisoned| poisoned.into_inner())
1389 .into_plan(),
1390 Err(shared) => lock_plan(&shared).snapshot(),
1391 };
1392
1393 info!(
1394 workflow = %handler_name,
1395 steps = plan.steps.len(),
1396 truncated = plan.truncated,
1397 "execution plan built"
1398 );
1399
1400 Ok(plan)
1401 }
1402
1403 #[tracing::instrument(name = "engine.enqueue_handler", skip_all, fields(workflow = %handler_name))]
1414 pub async fn enqueue_handler(
1415 &self,
1416 handler_name: &str,
1417 trigger: TriggerKind,
1418 payload: Value,
1419 max_retries: u32,
1420 ) -> Result<Run, EngineError> {
1421 self.enqueue_handler_with_options(
1422 handler_name,
1423 trigger,
1424 payload,
1425 EnqueueOptions {
1426 max_retries,
1427 ..Default::default()
1428 },
1429 )
1430 .await
1431 .map(RunCreation::into_run)
1432 }
1433
1434 #[tracing::instrument(name = "engine.enqueue_handler_with_options", skip_all, fields(workflow = %handler_name))]
1486 pub async fn enqueue_handler_with_options(
1487 &self,
1488 handler_name: &str,
1489 trigger: TriggerKind,
1490 payload: Value,
1491 options: EnqueueOptions,
1492 ) -> Result<RunCreation, EngineError> {
1493 let EnqueueOptions {
1494 max_retries,
1495 labels,
1496 scheduled_at,
1497 max_cost_usd,
1498 created_by,
1499 idempotency_key,
1500 concurrency_key,
1501 concurrency_limits,
1502 priority,
1503 worker_tags,
1504 } = options;
1505
1506 validate_concurrency_limits(&concurrency_limits)
1509 .map_err(EngineError::InvalidConcurrencyLimit)?;
1510 if let Some(priority) = priority {
1511 validate_priority(priority).map_err(EngineError::InvalidPriority)?;
1512 }
1513 validate_worker_tags(&worker_tags).map_err(EngineError::InvalidWorkerTag)?;
1514
1515 let handler = self.handlers.get(handler_name).ok_or_else(|| {
1516 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1517 })?;
1518
1519 self.check_monthly_quota(handler_name).await?;
1520
1521 let handler_version = handler.version().map(str::to_string);
1522 let mut merged_labels = handler.default_labels();
1523 merged_labels.extend(labels);
1524 let resolved_cap = self
1525 .budget
1526 .resolve_run_cap(max_cost_usd, handler.default_max_cost_usd());
1527 let priority = priority.unwrap_or_else(|| clamp_priority(handler.priority()));
1528 let required_tags = normalize_worker_tags(
1529 handler
1530 .required_worker_tags()
1531 .into_iter()
1532 .chain(worker_tags),
1533 );
1534
1535 let creation = self
1536 .store
1537 .create_run(NewRun {
1538 workflow_name: handler_name.to_string(),
1539 trigger,
1540 payload,
1541 max_retries,
1542 handler_version,
1543 labels: merged_labels,
1544 scheduled_at,
1545 created_by,
1546 idempotency_key,
1547 concurrency_key,
1548 priority,
1549 concurrency_limits,
1550 max_cost_usd: resolved_cap,
1551 worker_tags: required_tags,
1552 })
1553 .await?;
1554
1555 match &creation {
1556 RunCreation::Created(run) => info!(
1557 run_id = %run.id,
1558 workflow = %handler_name,
1559 max_cost_usd = ?resolved_cap,
1560 "handler run enqueued"
1561 ),
1562 RunCreation::Existing(run) => info!(
1563 run_id = %run.id,
1564 workflow = %handler_name,
1565 "idempotent replay, nothing enqueued"
1566 ),
1567 }
1568
1569 Ok(creation)
1570 }
1571
1572 #[tracing::instrument(name = "engine.execute_handler_run", skip_all, fields(run_id = %run_id))]
1592 pub async fn execute_handler_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1593 let run = self
1594 .store
1595 .get_run(run_id)
1596 .await?
1597 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1598
1599 if let Some(root_run_id) = chain_root(&run) {
1600 return self.resume_chain(run, root_run_id).await;
1601 }
1602
1603 let _active = self.track_execution(run_id).await;
1604
1605 let handler = self
1606 .handlers
1607 .get(&run.workflow_name)
1608 .ok_or_else(|| {
1609 EngineError::InvalidWorkflow(format!(
1610 "no handler registered: {}",
1611 run.workflow_name
1612 ))
1613 })?
1614 .clone();
1615
1616 #[cfg(feature = "prometheus")]
1617 gauge!(RUNS_ACTIVE, "workflow" => run.workflow_name.clone()).increment(1.0);
1618
1619 let run_start = Instant::now();
1620 let mut ctx = match self.build_context_with_guard(&run, handler.as_ref()).await {
1621 Ok(ctx) => ctx,
1622 Err(err) => return Err(self.fail_before_execution(run_id, err).await),
1623 };
1624
1625 let result = if handler.is_version_compatible(run.handler_version.as_deref()) {
1638 if let Err(err) = ctx.load_replay_steps().await {
1639 return Err(self.fail_before_execution(run_id, err).await);
1640 }
1641 self.release_then_execute(run_id, handler.as_ref(), &mut ctx)
1642 .await
1643 } else {
1644 Err(EngineError::HandlerVersionMismatch {
1645 run_id,
1646 workflow_name: run.workflow_name.clone(),
1647 run_version: run
1648 .handler_version
1649 .clone()
1650 .unwrap_or_else(|| "unknown".to_string()),
1651 current_version: handler
1652 .version()
1653 .map(str::to_string)
1654 .unwrap_or_else(|| "unknown".to_string()),
1655 })
1656 };
1657
1658 self.finalize_run(
1659 run_id,
1660 &run.workflow_name,
1661 result,
1662 &ctx,
1663 run_start,
1664 run.labels,
1665 )
1666 .await
1667 }
1668
1669 #[tracing::instrument(name = "engine.execute_run", skip_all, fields(run_id = %run_id))]
1677 pub async fn execute_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1678 self.execute_handler_run(run_id).await
1679 }
1680
1681 #[tracing::instrument(name = "engine.resume_run", skip_all, fields(run_id = %run_id))]
1710 pub async fn resume_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1711 let run = self
1712 .store
1713 .get_run(run_id)
1714 .await?
1715 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1716
1717 if let Some(root_run_id) = chain_root(&run) {
1718 return self.resume_chain(run, root_run_id).await;
1719 }
1720
1721 self.resume_loaded_run(run).await
1722 }
1723
1724 async fn resume_chain(
1741 &self,
1742 child: Run,
1743 root_run_id: Uuid,
1744 ) -> Result<WorkflowResult, EngineError> {
1745 let child_run_id = child.id;
1746 let lease = child.worker_id.zip(child.lease_expires_at);
1747 match self
1748 .claim_chain_root(child_run_id, root_run_id, lease.as_ref())
1749 .await?
1750 {
1751 ChainClaim::Claimed => {}
1752 ChainClaim::Settled(root) => {
1753 return Ok(WorkflowResult {
1754 run: *root,
1755 steps: Vec::new(),
1756 });
1757 }
1758 }
1759
1760 if lease.is_some() {
1761 self.store
1762 .update_run(
1763 child_run_id,
1764 RunUpdate {
1765 lease: Some(LeaseUpdate::Release),
1766 ..RunUpdate::default()
1767 },
1768 )
1769 .await?;
1770 }
1771
1772 info!(
1773 run_id = %child_run_id,
1774 root_run_id = %root_run_id,
1775 lease_transferred = lease.is_some(),
1776 "child run resumed through its root run"
1777 );
1778
1779 let root = self
1780 .store
1781 .get_run(root_run_id)
1782 .await?
1783 .ok_or(EngineError::Store(StoreError::RunNotFound(root_run_id)))?;
1784 self.resume_loaded_run(root).await
1785 }
1786
1787 pub(crate) async fn move_root_to_running(
1793 &self,
1794 root_run_id: Uuid,
1795 lease: Option<&(String, DateTime<Utc>)>,
1796 ) -> Result<(), EngineError> {
1797 match lease {
1798 Some((worker_id, expires_at)) => {
1799 self.store
1800 .update_run(
1801 root_run_id,
1802 RunUpdate {
1803 status: Some(RunStatus::Running),
1804 lease: Some(LeaseUpdate::Set {
1805 worker_id: worker_id.clone(),
1806 expires_at: *expires_at,
1807 }),
1808 ..RunUpdate::default()
1809 },
1810 )
1811 .await?;
1812 }
1813 None => {
1814 self.store
1815 .update_run_status(root_run_id, RunStatus::Running)
1816 .await?;
1817 }
1818 }
1819 Ok(())
1820 }
1821
1822 async fn resume_loaded_run(&self, run: Run) -> Result<WorkflowResult, EngineError> {
1824 let run_id = run.id;
1825 let _active = self.track_execution(run_id).await;
1826 let handler = self
1827 .handlers
1828 .get(&run.workflow_name)
1829 .ok_or_else(|| {
1830 EngineError::InvalidWorkflow(format!(
1831 "no handler registered: {}",
1832 run.workflow_name
1833 ))
1834 })?
1835 .clone();
1836
1837 info!(run_id = %run_id, workflow = %run.workflow_name, "resuming run after approval");
1838
1839 let run_start = Instant::now();
1840 let mut ctx = match self.build_context_with_guard(&run, handler.as_ref()).await {
1841 Ok(ctx) => ctx,
1842 Err(err) => return Err(self.fail_before_execution(run_id, err).await),
1843 };
1844
1845 let result = if handler.is_version_compatible(run.handler_version.as_deref()) {
1846 if let Err(err) = ctx.load_replay_steps().await {
1847 return Err(self.fail_before_execution(run_id, err).await);
1848 }
1849 self.release_then_execute(run_id, handler.as_ref(), &mut ctx)
1850 .await
1851 } else {
1852 Err(EngineError::HandlerVersionMismatch {
1853 run_id,
1854 workflow_name: run.workflow_name.clone(),
1855 run_version: run
1856 .handler_version
1857 .clone()
1858 .unwrap_or_else(|| "unknown".to_string()),
1859 current_version: handler
1860 .version()
1861 .map(str::to_string)
1862 .unwrap_or_else(|| "unknown".to_string()),
1863 })
1864 };
1865
1866 self.finalize_run(
1867 run_id,
1868 &run.workflow_name,
1869 result,
1870 &ctx,
1871 run_start,
1872 run.labels,
1873 )
1874 .await
1875 }
1876
1877 pub async fn deliver_signal(
1925 self: &Arc<Self>,
1926 signal: NewSignal,
1927 ) -> Result<SignalDelivery, EngineError> {
1928 self.deliver_signal_from(signal, None).await
1929 }
1930
1931 pub async fn deliver_signal_from(
1961 self: &Arc<Self>,
1962 signal: NewSignal,
1963 origin: Option<SignalOrigin>,
1964 ) -> Result<SignalDelivery, EngineError> {
1965 validate_signal(&signal)?;
1966
1967 let stored = match self.store.insert_signal(signal).await? {
1968 SignalInsert::Created(stored) => stored,
1969 SignalInsert::Duplicate(existing) => {
1970 info!(
1971 signal_id = %existing.id,
1972 signal = %existing.name,
1973 key = %existing.key,
1974 "duplicate signal ignored"
1975 );
1976 return Ok(SignalDelivery {
1977 signal_id: existing.id,
1978 duplicate: true,
1979 resumed: Vec::new(),
1980 rejected: Vec::new(),
1981 });
1982 }
1983 };
1984
1985 let waiters = self
1986 .store
1987 .list_signal_waiters(&stored.name, &stored.key)
1988 .await?;
1989 let mut resumed = Vec::new();
1990 let mut rejected = Vec::new();
1991
1992 for step in waiters {
1993 if let Err(error) = validate_step_payload(step.input.as_ref(), &stored.payload) {
1994 rejected.push(SignalRejected {
1995 run_id: step.run_id,
1996 step_id: step.id,
1997 error,
1998 });
1999 continue;
2000 }
2001
2002 match self
2003 .store
2004 .resolve_signal_step(step.id, received_output(&stored))
2005 .await
2006 {
2007 Ok(SignalStepResolution::Resolved {
2008 run_id,
2009 run_resumed,
2010 }) => {
2011 resumed.push(SignalResumed {
2012 run_id,
2013 step_id: step.id,
2014 });
2015 if run_resumed && self.execution_mode == ExecutionMode::Local {
2016 self.spawn_local_resume(run_id);
2017 }
2018 }
2019 Ok(SignalStepResolution::NotWaiting { .. }) => {}
2021 Err(err) => {
2022 error!(
2023 run_id = %step.run_id,
2024 step_id = %step.id,
2025 error = %err,
2026 "failed to resolve a waiting signal step"
2027 );
2028 rejected.push(SignalRejected {
2029 run_id: step.run_id,
2030 step_id: step.id,
2031 error: err.to_string(),
2032 });
2033 }
2034 }
2035 }
2036
2037 info!(
2038 signal_id = %stored.id,
2039 signal = %stored.name,
2040 key = %stored.key,
2041 resumed = resumed.len(),
2042 rejected = rejected.len(),
2043 "signal received"
2044 );
2045 self.event_publisher
2046 .publish(Event::SignalReceived(SignalReceivedEvent {
2047 signal_id: stored.id,
2048 name: stored.name.clone(),
2049 key: stored.key.clone(),
2050 resumed_runs: resumed.iter().map(|r| r.run_id).collect(),
2051 sender_run_id: origin.map(|o| o.run_id),
2052 sender_step_id: origin.map(|o| o.step_id),
2053 at: stored.received_at,
2054 }));
2055
2056 Ok(SignalDelivery {
2057 signal_id: stored.id,
2058 duplicate: false,
2059 resumed,
2060 rejected,
2061 })
2062 }
2063
2064 pub async fn send_signal<S: Signal>(
2102 self: &Arc<Self>,
2103 signal: &S,
2104 key: &str,
2105 idempotency_id: Option<&str>,
2106 ) -> Result<SignalDelivery, EngineError> {
2107 let payload = to_value(signal)?;
2108 self.deliver_signal(NewSignal {
2109 name: S::NAME.to_string(),
2110 key: key.to_string(),
2111 payload,
2112 idempotency_id: idempotency_id.map(str::to_string),
2113 })
2114 .await
2115 }
2116
2117 pub(crate) fn spawn_local_resume(self: &Arc<Self>, run_id: Uuid) {
2142 let engine = Arc::clone(self);
2143 spawn(async move {
2144 let _resuming = lock_run(&engine.local_resumes, run_id).await;
2148 engine.wait_until_idle(run_id).await;
2149 let run = match engine.store.get_run(run_id).await {
2150 Ok(Some(run)) => run,
2151 Ok(None) => {
2152 error!(run_id = %run_id, "run to restart not found");
2153 return;
2154 }
2155 Err(err) => {
2156 error!(run_id = %run_id, error = %err, "failed to load a run to restart");
2157 return;
2158 }
2159 };
2160 match run.status.state {
2161 RunStatus::Pending => {
2162 match engine
2163 .store
2164 .update_run_status(run_id, RunStatus::Running)
2165 .await
2166 {
2167 Ok(()) => {}
2168 Err(StoreError::InvalidTransition { from, .. }) => {
2169 info!(
2170 run_id = %run_id,
2171 status = %from,
2172 "run claimed elsewhere, left to its worker"
2173 );
2174 return;
2175 }
2176 Err(err) => {
2177 error!(run_id = %run_id, error = %err, "failed to restart a woken run");
2178 return;
2179 }
2180 }
2181 }
2182 RunStatus::Running if run.worker_id.is_some() => {
2183 info!(
2184 run_id = %run_id,
2185 worker_id = ?run.worker_id,
2186 "run owned by a worker lease, local resume skipped"
2187 );
2188 return;
2189 }
2190 RunStatus::Running => {
2191 if engine.is_executing(run_id) {
2192 info!(run_id = %run_id, "run already executing in this process, resume skipped");
2193 return;
2194 }
2195 }
2196 status => {
2197 info!(
2198 run_id = %run_id,
2199 status = %status,
2200 "run no longer waiting to restart, resume skipped"
2201 );
2202 return;
2203 }
2204 }
2205 if let Err(err) = engine.resume_run(run_id).await {
2206 error!(run_id = %run_id, error = %err, "failed to resume a woken run");
2207 }
2208 });
2209 }
2210
2211 pub async fn fail_or_schedule_retry(
2262 &self,
2263 run_id: Uuid,
2264 error: &str,
2265 retryable: bool,
2266 cost_usd: Option<Decimal>,
2267 duration_ms: Option<u64>,
2268 ) -> Result<RunStatus, EngineError> {
2269 let run = self
2270 .store
2271 .get_run(run_id)
2272 .await?
2273 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2274
2275 if run.status.state == RunStatus::Paused {
2278 info!(run_id = %run_id, error = %error, "run paused, failure not recorded");
2279 return Ok(RunStatus::Paused);
2280 }
2281
2282 let has_attempts_left = run.retry_count < run.max_retries;
2283 let update = if retryable && has_attempts_left {
2284 let backoff = backoff_for_retry(run.retry_count);
2285 let scheduled_at = Utc::now() + TimeDelta::milliseconds(backoff.as_millis() as i64);
2286
2287 info!(
2288 run_id = %run_id,
2289 workflow = %run.workflow_name,
2290 attempt = run.retry_count + 1,
2291 max_retries = run.max_retries,
2292 backoff_secs = backoff.as_secs(),
2293 scheduled_at = %scheduled_at,
2294 "run failed, scheduling retry"
2295 );
2296
2297 RunUpdate {
2298 status: Some(RunStatus::Retrying),
2299 error: Some(error.to_string()),
2300 increment_retry: true,
2301 cost_usd,
2302 duration_ms,
2303 scheduled_at: Some(scheduled_at),
2304 ..RunUpdate::default()
2305 }
2306 } else {
2307 RunUpdate {
2308 status: Some(RunStatus::Failed),
2309 error: Some(error.to_string()),
2310 cost_usd,
2311 duration_ms,
2312 completed_at: Some(Utc::now()),
2313 ..RunUpdate::default()
2314 }
2315 };
2316
2317 let status = update.status.unwrap_or(RunStatus::Failed);
2318 self.store.update_run(run_id, update).await?;
2319 self.fail_orphaned_steps(run_id, error).await?;
2320 self.cancel_descendants_of_stopped_run(run_id, error).await;
2323
2324 Ok(status)
2325 }
2326
2327 pub async fn interrupt_running_steps(&self, run_id: Uuid) -> Result<(), EngineError> {
2357 interrupt_running_steps(self.store.as_ref(), run_id).await
2358 }
2359
2360 pub async fn fail_orphaned_steps(
2374 &self,
2375 run_id: Uuid,
2376 error_message: &str,
2377 ) -> Result<(), EngineError> {
2378 let steps = self.store.list_steps(run_id).await?;
2379 let now = Utc::now();
2380
2381 for step in steps {
2382 if step.status.state.is_terminal() {
2383 continue;
2384 }
2385
2386 let (target_status, error) = match step.status.state {
2387 StepStatus::Running | StepStatus::AwaitingApproval => {
2388 let err = if step.error.is_some() {
2389 None
2390 } else {
2391 Some(error_message.to_string())
2392 };
2393 (StepStatus::Failed, err)
2394 }
2395 StepStatus::Pending => (StepStatus::Skipped, None),
2396 _ => continue,
2397 };
2398
2399 if let Err(e) = self
2400 .store
2401 .update_step(
2402 step.id,
2403 StepUpdate {
2404 status: Some(target_status),
2405 error,
2406 completed_at: Some(now),
2407 ..StepUpdate::default()
2408 },
2409 )
2410 .await
2411 {
2412 warn!(
2413 run_id = %run_id,
2414 step_id = %step.id,
2415 step_name = %step.name,
2416 error = %e,
2417 "failed to cleanup orphaned step"
2418 );
2419 } else {
2420 info!(
2421 run_id = %run_id,
2422 step_id = %step.id,
2423 step_name = %step.name,
2424 from = %step.status.state,
2425 to = %target_status,
2426 "cleaned up orphaned step"
2427 );
2428 }
2429 }
2430
2431 Ok(())
2432 }
2433
2434 async fn release_then_execute(
2439 &self,
2440 run_id: Uuid,
2441 handler: &dyn WorkflowHandler,
2442 ctx: &mut WorkflowContext,
2443 ) -> Result<(), EngineError> {
2444 match self.provider.release_run(&run_id.to_string()).await {
2445 Ok(()) => handler.execute(ctx).await,
2446 Err(e) => Err(EngineError::Operation(OperationError::Agent(e))),
2447 }
2448 }
2449
2450 async fn finalize_run(
2456 &self,
2457 run_id: Uuid,
2458 workflow_name: &str,
2459 result: Result<(), EngineError>,
2460 ctx: &WorkflowContext,
2461 run_start: Instant,
2462 run_labels: HashMap<String, String>,
2463 ) -> Result<WorkflowResult, EngineError> {
2464 let total_duration = ctx.carried_duration_ms() + run_start.elapsed().as_millis() as u64;
2467 let completed_at = Utc::now();
2468
2469 if let Some(run) = self.store.get_run(run_id).await?
2473 && run.status.state == RunStatus::Paused
2474 {
2475 let run = self
2476 .store
2477 .update_run_returning(
2478 run_id,
2479 RunUpdate {
2480 cost_usd: Some(ctx.total_cost_usd()),
2481 duration_ms: Some(total_duration),
2482 ..RunUpdate::default()
2483 },
2484 )
2485 .await?;
2486 info!(
2487 run_id = %run_id,
2488 outcome = ?result.err().map(|err| err.to_string()),
2489 "run paused, execution stopped"
2490 );
2491 return Ok(WorkflowResult {
2492 run,
2493 steps: ctx.step_results().to_vec(),
2494 });
2495 }
2496
2497 if matches!(result, Err(EngineError::RunPaused { .. })) {
2500 let run = self
2501 .store
2502 .get_run(run_id)
2503 .await?
2504 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2505 info!(
2506 run_id = %run_id,
2507 status = %run.status.state,
2508 "run resumed after the pause, execution stopped"
2509 );
2510 return Ok(WorkflowResult {
2511 run,
2512 steps: ctx.step_results().to_vec(),
2513 });
2514 }
2515
2516 let final_status;
2517 let final_run;
2518
2519 match result {
2520 Ok(()) => {
2521 final_status = if ctx.has_allowed_failure() {
2522 RunStatus::Warning
2523 } else {
2524 RunStatus::Completed
2525 };
2526 final_run = self
2527 .store
2528 .update_run_returning(
2529 run_id,
2530 RunUpdate {
2531 status: Some(final_status),
2532 cost_usd: Some(ctx.total_cost_usd()),
2533 duration_ms: Some(total_duration),
2534 completed_at: Some(completed_at),
2535 output: ctx.output().cloned(),
2536 ..RunUpdate::default()
2537 },
2538 )
2539 .await?;
2540
2541 info!(
2542 run_id = %run_id,
2543 status = %final_status,
2544 cost_usd = %ctx.total_cost_usd(),
2545 duration_ms = total_duration,
2546 "run completed"
2547 );
2548 }
2549 Err(EngineError::ApprovalRequired {
2550 run_id: approval_run_id,
2551 step_id,
2552 ref message,
2553 }) => {
2554 final_status = RunStatus::AwaitingApproval;
2555 final_run = self
2556 .store
2557 .update_run_returning(
2558 run_id,
2559 RunUpdate {
2560 status: Some(RunStatus::AwaitingApproval),
2561 cost_usd: Some(ctx.total_cost_usd()),
2562 duration_ms: Some(total_duration),
2563 ..RunUpdate::default()
2564 },
2565 )
2566 .await?;
2567
2568 info!(
2569 run_id = %approval_run_id,
2570 step_id = %step_id,
2571 message = %message,
2572 "run awaiting approval"
2573 );
2574
2575 self.publish_approval_requested(approval_run_id, step_id, message)
2576 .await?;
2577 }
2578 Err(EngineError::ChildSuspended {
2579 run_id: child_run_id,
2580 ref cause,
2581 }) => {
2582 final_status = cause.suspension_status();
2583 final_run = self
2587 .store
2588 .update_run_returning(
2589 run_id,
2590 RunUpdate {
2591 status: Some(final_status),
2592 cost_usd: Some(ctx.total_cost_usd()),
2593 duration_ms: Some(total_duration),
2594 ..RunUpdate::default()
2595 },
2596 )
2597 .await?;
2598
2599 let leaf = cause.suspension_leaf();
2600 info!(
2601 run_id = %run_id,
2602 child_run_id = %child_run_id,
2603 status = %final_status,
2604 cause = %leaf,
2605 "run suspended with its child run"
2606 );
2607
2608 self.publish_suspension(leaf).await?;
2609 for sibling in ctx.sibling_suspensions() {
2611 self.publish_suspension(sibling.suspension_leaf()).await?;
2612 }
2613 }
2614 Err(EngineError::HumanInputRequired {
2615 run_id: input_run_id,
2616 step_id,
2617 ref message,
2618 }) => {
2619 final_status = RunStatus::AwaitingApproval;
2620 final_run = self
2621 .store
2622 .update_run_returning(
2623 run_id,
2624 RunUpdate {
2625 status: Some(RunStatus::AwaitingApproval),
2626 cost_usd: Some(ctx.total_cost_usd()),
2627 duration_ms: Some(total_duration),
2628 ..RunUpdate::default()
2629 },
2630 )
2631 .await?;
2632
2633 info!(
2635 run_id = %input_run_id,
2636 step_id = %step_id,
2637 message = %message,
2638 "run awaiting human input"
2639 );
2640 }
2641 Err(EngineError::DelaySleeping {
2642 run_id: delay_run_id,
2643 step_id,
2644 wake_at,
2645 }) => {
2646 final_status = RunStatus::Sleeping;
2647 final_run = self
2648 .store
2649 .update_run_returning(
2650 run_id,
2651 RunUpdate {
2652 status: Some(RunStatus::Sleeping),
2653 cost_usd: Some(ctx.total_cost_usd()),
2654 duration_ms: Some(total_duration),
2655 scheduled_at: Some(wake_at),
2656 ..RunUpdate::default()
2657 },
2658 )
2659 .await?;
2660
2661 info!(
2662 run_id = %delay_run_id,
2663 step_id = %step_id,
2664 wake_at = %wake_at,
2665 "run sleeping until delay elapses"
2666 );
2667 }
2668 Err(EngineError::CapacitySleeping {
2669 run_id: capacity_run_id,
2670 step_id,
2671 ref kind,
2672 wake_at,
2673 }) => {
2674 final_status = RunStatus::Sleeping;
2675 final_run = self
2676 .store
2677 .update_run_returning(
2678 run_id,
2679 RunUpdate {
2680 status: Some(RunStatus::Sleeping),
2681 cost_usd: Some(ctx.total_cost_usd()),
2682 duration_ms: Some(total_duration),
2683 scheduled_at: Some(wake_at),
2684 capacity_wait_kind: Some(ProviderKind::new(kind.as_str())),
2685 ..RunUpdate::default()
2686 },
2687 )
2688 .await?;
2689
2690 info!(
2691 run_id = %capacity_run_id,
2692 step_id = %step_id,
2693 kind = %kind,
2694 wake_at = %wake_at,
2695 "run sleeping until provider capacity returns"
2696 );
2697 }
2698 Err(EngineError::SignalWaiting {
2699 run_id: wait_run_id,
2700 step_id,
2701 ref step_name,
2702 ref name,
2703 ref key,
2704 deadline_at,
2705 }) => {
2706 final_status = RunStatus::Sleeping;
2707 let waiting = self
2711 .store
2712 .suspend_run_on_signal(run_id, step_id, deadline_at)
2713 .await?;
2714 final_run = self
2715 .store
2716 .update_run_returning(
2717 run_id,
2718 RunUpdate {
2719 cost_usd: Some(ctx.total_cost_usd()),
2720 duration_ms: Some(total_duration),
2721 ..RunUpdate::default()
2722 },
2723 )
2724 .await?;
2725
2726 if waiting {
2727 self.event_publisher
2728 .publish(Event::SignalAwaited(SignalAwaitedEvent {
2729 run_id: wait_run_id,
2730 step_id,
2731 step_name: step_name.clone(),
2732 name: name.clone(),
2733 key: key.clone(),
2734 deadline_at,
2735 at: Utc::now(),
2736 }));
2737 }
2738
2739 info!(
2740 run_id = %wait_run_id,
2741 step_id = %step_id,
2742 signal = %name,
2743 key = %key,
2744 deadline_at = %deadline_at,
2745 waiting,
2746 "run sleeping until a signal arrives"
2747 );
2748 }
2749 Err(err) => {
2750 let guardrail_stop = matches!(
2754 err,
2755 EngineError::RunBudgetExceeded { .. } | EngineError::WorkflowGuardRejected(_)
2756 );
2757
2758 final_status = if guardrail_stop {
2759 if let Err(store_err) = self
2760 .store
2761 .update_run(
2762 run_id,
2763 RunUpdate {
2764 status: Some(RunStatus::Cancelled),
2765 error: Some(err.to_string()),
2766 cost_usd: Some(ctx.total_cost_usd()),
2767 duration_ms: Some(total_duration),
2768 completed_at: Some(completed_at),
2769 output: ctx.output().cloned(),
2770 ..RunUpdate::default()
2771 },
2772 )
2773 .await
2774 {
2775 error!(run_id = %run_id, store_error = %store_err, "failed to persist run cancellation");
2776 }
2777 if let Err(cleanup_err) = self
2778 .fail_orphaned_steps(run_id, "run stopped: guardrail limit reached")
2779 .await
2780 {
2781 error!(run_id = %run_id, store_error = %cleanup_err, "failed to cleanup orphaned steps");
2782 }
2783 RunStatus::Cancelled
2784 } else {
2785 if let Some(output) = ctx.output()
2788 && let Err(store_err) = self
2789 .store
2790 .update_run(
2791 run_id,
2792 RunUpdate {
2793 output: Some(output.clone()),
2794 ..RunUpdate::default()
2795 },
2796 )
2797 .await
2798 {
2799 error!(run_id = %run_id, store_error = %store_err, "failed to persist run output");
2800 }
2801 self.fail_or_schedule_retry(
2802 run_id,
2803 &err.to_string(),
2804 is_run_retryable(&err),
2805 Some(ctx.total_cost_usd()),
2806 Some(total_duration),
2807 )
2808 .await
2809 .unwrap_or_else(|store_err| {
2810 error!(run_id = %run_id, store_error = %store_err, "failed to persist run failure");
2811 RunStatus::Failed
2812 })
2813 };
2814
2815 if matches!(err, EngineError::RunBudgetExceeded { .. }) {
2816 self.on_run_budget_exceeded(workflow_name, run_id, &err);
2817 }
2818
2819 error!(run_id = %run_id, status = %final_status, error = %err, "run stopped");
2820
2821 self.publish_run_status_changed(
2822 workflow_name,
2823 run_id,
2824 final_status,
2825 Some(err.to_string()),
2826 ctx,
2827 total_duration,
2828 run_labels,
2829 );
2830
2831 #[cfg(feature = "prometheus")]
2832 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
2833
2834 return Err(err);
2835 }
2836 }
2837
2838 self.publish_run_status_changed(
2839 workflow_name,
2840 run_id,
2841 final_status,
2842 None,
2843 ctx,
2844 total_duration,
2845 run_labels,
2846 );
2847
2848 #[cfg(feature = "prometheus")]
2849 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
2850
2851 Ok(WorkflowResult {
2852 run: final_run,
2853 steps: ctx.step_results().to_vec(),
2854 })
2855 }
2856
2857 pub(crate) async fn publish_approval_requested(
2860 &self,
2861 run_id: Uuid,
2862 step_id: Uuid,
2863 message: &str,
2864 ) -> Result<(), EngineError> {
2865 let requirement = self
2866 .store
2867 .get_step(step_id)
2868 .await?
2869 .and_then(|s| s.approval_requirement);
2870 self.event_publisher
2871 .publish(Event::ApprovalRequested(ApprovalRequestedEvent {
2872 run_id,
2873 step_id,
2874 message: message.to_string(),
2875 requirement,
2876 at: Utc::now(),
2877 }));
2878 Ok(())
2879 }
2880
2881 pub async fn fail_ancestors(&self, run_id: Uuid, reason: &str) -> Result<(), EngineError> {
2914 let mut current = self
2915 .store
2916 .get_run(run_id)
2917 .await?
2918 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2919 let mut visited = HashSet::from([run_id]);
2921
2922 while let Some(parent_id) = chain_parent(¤t) {
2923 if !visited.insert(parent_id) {
2924 break;
2925 }
2926 let status = self
2927 .fail_or_schedule_retry(parent_id, reason, false, None, None)
2928 .await?;
2929 if status == RunStatus::Paused {
2932 self.requeue_paused_root(¤t).await?;
2933 break;
2934 }
2935 info!(
2936 run_id = %run_id,
2937 ancestor_run_id = %parent_id,
2938 status = %status,
2939 "ancestor run failed with its child"
2940 );
2941 current = self
2942 .store
2943 .get_run(parent_id)
2944 .await?
2945 .ok_or(EngineError::Store(StoreError::RunNotFound(parent_id)))?;
2946 }
2947
2948 Ok(())
2949 }
2950
2951 #[cfg(feature = "prometheus")]
2953 fn emit_run_metrics(
2954 &self,
2955 workflow_name: &str,
2956 status: RunStatus,
2957 duration_ms: u64,
2958 ctx: &WorkflowContext,
2959 ) {
2960 let status_str = status.to_string();
2961 let wf = workflow_name.to_string();
2962
2963 counter!(RUNS_TOTAL, "workflow" => wf.clone(), "status" => status_str.clone()).increment(1);
2964 histogram!(RUN_DURATION_SECONDS, "workflow" => wf.clone(), "status" => status_str)
2965 .record(duration_ms as f64 / 1000.0);
2966 histogram!(RUN_COST_USD, "workflow" => wf.clone()).record(
2967 ctx.total_cost_usd()
2968 .to_string()
2969 .parse::<f64>()
2970 .unwrap_or(0.0),
2971 );
2972 gauge!(RUNS_ACTIVE, "workflow" => wf).decrement(1.0);
2973 }
2974
2975 fn on_run_budget_exceeded(&self, workflow_name: &str, run_id: Uuid, err: &EngineError) {
2981 let EngineError::RunBudgetExceeded {
2982 limit_usd,
2983 spent_usd,
2984 step_budget_usd,
2985 ..
2986 } = err
2987 else {
2988 return;
2989 };
2990
2991 #[cfg(feature = "prometheus")]
2992 counter!(
2993 RUN_BUDGET_EXCEEDED_TOTAL,
2994 "workflow" => workflow_name.to_string(),
2995 "scope" => "run",
2996 )
2997 .increment(1);
2998
2999 self.event_publisher
3000 .publish(Event::RunBudgetExceeded(RunBudgetExceededEvent {
3001 run_id,
3002 workflow_name: workflow_name.to_string(),
3003 limit_usd: *limit_usd,
3004 spent_usd: *spent_usd,
3005 step_budget_usd: *step_budget_usd,
3006 at: Utc::now(),
3007 }));
3008 }
3009
3010 #[allow(clippy::too_many_arguments)]
3015 fn publish_run_status_changed(
3016 &self,
3017 workflow_name: &str,
3018 run_id: Uuid,
3019 to: RunStatus,
3020 error: Option<String>,
3021 ctx: &WorkflowContext,
3022 duration_ms: u64,
3023 labels: HashMap<String, String>,
3024 ) {
3025 let now = Utc::now();
3026 let cost_usd = ctx.total_cost_usd();
3027 let wf = workflow_name.to_string();
3028
3029 self.event_publisher
3030 .publish(Event::RunStatusChanged(RunStatusChangedEvent {
3031 run_id,
3032 workflow_name: wf.clone(),
3033 from: RunStatus::Running,
3034 to,
3035 error: error.clone(),
3036 cost_usd,
3037 duration_ms,
3038 labels: labels.clone(),
3039 at: now,
3040 }));
3041
3042 if to == RunStatus::Failed {
3043 self.event_publisher
3044 .publish(Event::RunFailed(RunFailedEvent {
3045 run_id,
3046 workflow_name: wf,
3047 error,
3048 cost_usd,
3049 duration_ms,
3050 labels,
3051 at: now,
3052 }));
3053 }
3054 }
3055}
3056
3057impl fmt::Debug for Engine {
3058 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3059 f.debug_struct("Engine")
3060 .field("handlers", &self.handlers.keys().collect::<Vec<_>>())
3061 .finish_non_exhaustive()
3062 }
3063}
3064
3065#[cfg(test)]
3066mod tests {
3067 use super::*;
3068 use crate::config::ShellConfig;
3069 use crate::handler::{HandlerFuture, WorkflowHandler};
3070 use ironflow_core::providers::claude::ClaudeCodeProvider;
3071 use ironflow_core::providers::record_replay::RecordReplayProvider;
3072 use ironflow_store::memory::InMemoryStore;
3073 use ironflow_store::models::{MAX_PRIORITY, MIN_PRIORITY, StepApproval, StepStatus};
3074 use serde_json::json;
3075
3076 struct EchoWorkflow;
3078
3079 impl WorkflowHandler for EchoWorkflow {
3080 fn name(&self) -> &str {
3081 "echo-workflow"
3082 }
3083
3084 fn describe(&self) -> WorkflowInfo {
3085 WorkflowInfo {
3086 description: "A simple workflow that echoes hello".to_string(),
3087 source_code: None,
3088 sub_workflows: Vec::new(),
3089 category: None,
3090 version: self.version().map(str::to_string),
3091 compatible_versions: Vec::new(),
3092 input_schema: None,
3093 default_labels: HashMap::new(),
3094 schedule: self.schedule().cloned(),
3095 default_max_cost_usd: self.default_max_cost_usd(),
3096 priority: self.priority(),
3097 }
3098 }
3099
3100 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3101 Box::pin(async move {
3102 ctx.shell("greet", ShellConfig::new("echo hello")).await?;
3103 Ok(())
3104 })
3105 }
3106 }
3107
3108 struct FailingWorkflow;
3110
3111 impl WorkflowHandler for FailingWorkflow {
3112 fn name(&self) -> &str {
3113 "failing-workflow"
3114 }
3115
3116 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3117 Box::pin(async move {
3118 ctx.shell("fail", ShellConfig::new("exit 1")).await?;
3119 Ok(())
3120 })
3121 }
3122 }
3123
3124 fn create_test_engine() -> Engine {
3125 let store = Arc::new(InMemoryStore::new());
3126 let inner = ClaudeCodeProvider::new();
3127 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
3128 inner,
3129 "/tmp/ironflow-fixtures",
3130 ));
3131 Engine::new(store, provider)
3132 }
3133
3134 #[test]
3135 fn engine_new_creates_instance() {
3136 let engine = create_test_engine();
3137 assert_eq!(engine.handler_names().len(), 0);
3138 }
3139
3140 #[test]
3141 fn execution_mode_defaults_to_local() {
3142 let engine = create_test_engine();
3143 assert_eq!(engine.execution_mode(), ExecutionMode::Local);
3144 }
3145
3146 #[test]
3147 fn with_execution_mode_overrides_the_default() {
3148 let engine = create_test_engine().with_execution_mode(ExecutionMode::Workers);
3149 assert_eq!(engine.execution_mode(), ExecutionMode::Workers);
3150 }
3151
3152 #[test]
3153 fn engine_register_handler() {
3154 let mut engine = create_test_engine();
3155 let result = engine.register(EchoWorkflow);
3156 assert!(result.is_ok());
3157 assert_eq!(engine.handler_names().len(), 1);
3158 assert!(engine.handler_names().contains(&"echo-workflow"));
3159 }
3160
3161 #[test]
3162 fn engine_register_duplicate_returns_error() {
3163 let mut engine = create_test_engine();
3164 engine.register(EchoWorkflow).unwrap();
3165 let result = engine.register(EchoWorkflow);
3166 assert!(result.is_err());
3167 }
3168
3169 #[test]
3170 fn engine_get_handler_found() {
3171 let mut engine = create_test_engine();
3172 engine.register(EchoWorkflow).unwrap();
3173 let handler = engine.get_handler("echo-workflow");
3174 assert!(handler.is_some());
3175 }
3176
3177 #[test]
3178 fn engine_get_handler_not_found() {
3179 let engine = create_test_engine();
3180 let handler = engine.get_handler("nonexistent");
3181 assert!(handler.is_none());
3182 }
3183
3184 #[test]
3185 fn engine_handler_names_lists_all() {
3186 let mut engine = create_test_engine();
3187 engine.register(EchoWorkflow).unwrap();
3188 engine.register(FailingWorkflow).unwrap();
3189 let names = engine.handler_names();
3190 assert_eq!(names.len(), 2);
3191 assert!(names.contains(&"echo-workflow"));
3192 assert!(names.contains(&"failing-workflow"));
3193 }
3194
3195 #[test]
3196 fn engine_handler_info_returns_description() {
3197 let mut engine = create_test_engine();
3198 engine.register(EchoWorkflow).unwrap();
3199 let info = engine.handler_info("echo-workflow");
3200 assert!(info.is_some());
3201 let info = info.unwrap();
3202 assert_eq!(info.description, "A simple workflow that echoes hello");
3203 }
3204
3205 struct CategorizedWorkflow;
3206
3207 impl WorkflowHandler for CategorizedWorkflow {
3208 fn name(&self) -> &str {
3209 "categorized"
3210 }
3211 fn category(&self) -> Option<&str> {
3212 Some("data/etl")
3213 }
3214 fn execute<'a>(
3215 &'a self,
3216 _ctx: &'a mut WorkflowContext,
3217 ) -> crate::handler::HandlerFuture<'a> {
3218 Box::pin(async move { Ok(()) })
3219 }
3220 }
3221
3222 #[test]
3223 fn engine_default_describe_propagates_category() {
3224 let mut engine = create_test_engine();
3225 engine.register(CategorizedWorkflow).unwrap();
3226 let info = engine.handler_info("categorized").unwrap();
3227 assert_eq!(info.category.as_deref(), Some("data/etl"));
3228 }
3229
3230 #[test]
3231 fn engine_default_describe_without_category() {
3232 let mut engine = create_test_engine();
3233 engine.register(EchoWorkflow).unwrap();
3234 let info = engine.handler_info("echo-workflow").unwrap();
3235 assert!(info.category.is_none());
3236 }
3237
3238 struct ScheduledWorkflow {
3243 schedule: CronSchedule,
3244 }
3245
3246 impl ScheduledWorkflow {
3247 fn new() -> Self {
3248 Self {
3249 schedule: CronSchedule::new("0 0 * * * *").unwrap(),
3250 }
3251 }
3252 }
3253
3254 impl WorkflowHandler for ScheduledWorkflow {
3255 fn name(&self) -> &str {
3256 "scheduled"
3257 }
3258 fn schedule(&self) -> Option<&CronSchedule> {
3259 Some(&self.schedule)
3260 }
3261 fn execute<'a>(
3262 &'a self,
3263 _ctx: &'a mut WorkflowContext,
3264 ) -> crate::handler::HandlerFuture<'a> {
3265 Box::pin(async move { Ok(()) })
3266 }
3267 }
3268
3269 #[test]
3270 fn engine_default_describe_propagates_schedule() {
3271 let mut engine = create_test_engine();
3272 engine.register(ScheduledWorkflow::new()).unwrap();
3273 let info = engine.handler_info("scheduled").unwrap();
3274 assert_eq!(
3275 info.schedule.as_ref().map(|s| s.as_str()),
3276 Some("0 0 * * * *")
3277 );
3278 }
3279
3280 #[test]
3281 fn engine_default_describe_without_schedule() {
3282 let mut engine = create_test_engine();
3283 engine.register(EchoWorkflow).unwrap();
3284 let info = engine.handler_info("echo-workflow").unwrap();
3285 assert!(info.schedule.is_none());
3286 }
3287
3288 #[test]
3289 fn scheduled_handlers_returns_only_scheduled() {
3290 let mut engine = create_test_engine();
3291 engine.register(EchoWorkflow).unwrap();
3292 engine.register(ScheduledWorkflow::new()).unwrap();
3293 engine.register(FailingWorkflow).unwrap();
3294
3295 let scheduled = engine.scheduled_handlers();
3296 assert_eq!(scheduled.len(), 1);
3297 assert_eq!(scheduled[0].0, "scheduled");
3298 assert_eq!(scheduled[0].1.as_str(), "0 0 * * * *");
3299 }
3300
3301 #[test]
3302 fn scheduled_handlers_empty_when_none_scheduled() {
3303 let mut engine = create_test_engine();
3304 engine.register(EchoWorkflow).unwrap();
3305 engine.register(FailingWorkflow).unwrap();
3306
3307 let scheduled = engine.scheduled_handlers();
3308 assert!(scheduled.is_empty());
3309 }
3310
3311 struct BadCategoryWorkflow(&'static str);
3312
3313 impl WorkflowHandler for BadCategoryWorkflow {
3314 fn name(&self) -> &str {
3315 "bad-category"
3316 }
3317 fn category(&self) -> Option<&str> {
3318 Some(self.0)
3319 }
3320 fn execute<'a>(
3321 &'a self,
3322 _ctx: &'a mut WorkflowContext,
3323 ) -> crate::handler::HandlerFuture<'a> {
3324 Box::pin(async move { Ok(()) })
3325 }
3326 }
3327
3328 #[test]
3329 fn engine_register_rejects_empty_category() {
3330 let mut engine = create_test_engine();
3331 let err = engine.register(BadCategoryWorkflow("")).unwrap_err();
3332 match err {
3333 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty category")),
3334 other => panic!("expected InvalidWorkflow, got {other:?}"),
3335 }
3336 }
3337
3338 #[test]
3339 fn engine_register_rejects_leading_slash_category() {
3340 let mut engine = create_test_engine();
3341 let err = engine
3342 .register(BadCategoryWorkflow("/data/etl"))
3343 .unwrap_err();
3344 match err {
3345 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("leading '/'")),
3346 other => panic!("expected InvalidWorkflow, got {other:?}"),
3347 }
3348 }
3349
3350 #[test]
3351 fn engine_register_rejects_trailing_slash_category() {
3352 let mut engine = create_test_engine();
3353 let err = engine
3354 .register(BadCategoryWorkflow("data/etl/"))
3355 .unwrap_err();
3356 match err {
3357 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("trailing '/'")),
3358 other => panic!("expected InvalidWorkflow, got {other:?}"),
3359 }
3360 }
3361
3362 #[test]
3363 fn engine_register_rejects_double_slash_category() {
3364 let mut engine = create_test_engine();
3365 let err = engine
3366 .register(BadCategoryWorkflow("data//etl"))
3367 .unwrap_err();
3368 match err {
3369 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty segment")),
3370 other => panic!("expected InvalidWorkflow, got {other:?}"),
3371 }
3372 }
3373
3374 #[test]
3375 fn engine_register_rejects_whitespace_only_segment_category() {
3376 let mut engine = create_test_engine();
3377 let err = engine
3378 .register(BadCategoryWorkflow("data/ /etl"))
3379 .unwrap_err();
3380 match err {
3381 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("whitespace-only segment")),
3382 other => panic!("expected InvalidWorkflow, got {other:?}"),
3383 }
3384 }
3385
3386 #[test]
3387 fn engine_register_accepts_valid_nested_category() {
3388 let mut engine = create_test_engine();
3389 assert!(engine.register(CategorizedWorkflow).is_ok());
3390 }
3391
3392 #[tokio::test]
3393 async fn engine_unknown_workflow_returns_error() {
3394 let engine = create_test_engine();
3395 let result = engine
3396 .run_handler("unknown", TriggerKind::Manual, json!({}))
3397 .await;
3398 assert!(result.is_err());
3399 match result {
3400 Err(EngineError::InvalidWorkflow(msg)) => {
3401 assert!(msg.contains("no handler registered"));
3402 }
3403 _ => panic!("expected InvalidWorkflow error"),
3404 }
3405 }
3406
3407 #[tokio::test]
3408 async fn engine_enqueue_handler_creates_pending_run() {
3409 let mut engine = create_test_engine();
3410 engine.register(EchoWorkflow).unwrap();
3411
3412 let run = engine
3413 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
3414 .await
3415 .unwrap();
3416 assert_eq!(run.status.state, RunStatus::Pending);
3417 assert_eq!(run.workflow_name, "echo-workflow");
3418 }
3419
3420 #[tokio::test]
3421 async fn enqueue_handler_leaves_the_run_unattributed() {
3422 let mut engine = create_test_engine();
3423 engine.register(EchoWorkflow).unwrap();
3424
3425 let run = engine
3426 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
3427 .await
3428 .unwrap();
3429
3430 assert!(run.created_by.is_none());
3431 }
3432
3433 #[tokio::test]
3434 async fn enqueue_handler_with_options_records_the_author() {
3435 let mut engine = create_test_engine();
3436 engine.register(EchoWorkflow).unwrap();
3437 let actor = RunActor::User {
3438 user_id: Uuid::now_v7(),
3439 };
3440
3441 let run = engine
3442 .enqueue_handler_with_options(
3443 "echo-workflow",
3444 TriggerKind::Api,
3445 json!({}),
3446 EnqueueOptions {
3447 created_by: Some(actor.clone()),
3448 ..Default::default()
3449 },
3450 )
3451 .await
3452 .unwrap()
3453 .into_run();
3454
3455 assert_eq!(run.created_by, Some(actor));
3456 }
3457
3458 #[tokio::test]
3459 async fn enqueue_handler_with_options_accepts_no_author() {
3460 let mut engine = create_test_engine();
3461 engine.register(EchoWorkflow).unwrap();
3462
3463 let run = engine
3464 .enqueue_handler_with_options(
3465 "echo-workflow",
3466 TriggerKind::Cron {
3467 schedule: "0 * * * * *".to_string(),
3468 schedule_id: None,
3469 scheduled_for: None,
3470 },
3471 json!({}),
3472 EnqueueOptions::default(),
3473 )
3474 .await
3475 .unwrap()
3476 .into_run();
3477
3478 assert!(run.created_by.is_none());
3479 }
3480
3481 #[tokio::test]
3482 async fn enqueue_handler_with_options_stores_concurrency_limits() {
3483 let mut engine = create_test_engine();
3484 engine.register(EchoWorkflow).unwrap();
3485 let limits = vec![
3486 ConcurrencyLimit::new("repo:acme", 2),
3487 ConcurrencyLimit::new("tenant:42", 5),
3488 ];
3489
3490 let run = engine
3491 .enqueue_handler_with_options(
3492 "echo-workflow",
3493 TriggerKind::Api,
3494 json!({}),
3495 EnqueueOptions {
3496 concurrency_limits: limits.clone(),
3497 ..Default::default()
3498 },
3499 )
3500 .await
3501 .unwrap()
3502 .into_run();
3503
3504 assert_eq!(run.concurrency_limits, limits);
3505 }
3506
3507 #[tokio::test]
3508 async fn enqueue_rejects_invalid_concurrency_limits() {
3509 let mut engine = create_test_engine();
3510 engine.register(EchoWorkflow).unwrap();
3511
3512 let invalid = [
3513 vec![ConcurrencyLimit::new("repo:acme", 0)],
3514 vec![ConcurrencyLimit::new("", 1)],
3515 vec![
3516 ConcurrencyLimit::new("repo:acme", 1),
3517 ConcurrencyLimit::new("repo:acme", 2),
3518 ],
3519 ];
3520 for concurrency_limits in invalid {
3521 let err = engine
3522 .enqueue_handler_with_options(
3523 "echo-workflow",
3524 TriggerKind::Api,
3525 json!({}),
3526 EnqueueOptions {
3527 concurrency_limits,
3528 ..Default::default()
3529 },
3530 )
3531 .await
3532 .unwrap_err();
3533 assert!(
3534 matches!(err, EngineError::InvalidConcurrencyLimit(_)),
3535 "{err:?}"
3536 );
3537 }
3538
3539 let err = engine
3541 .enqueue_handler_with_options(
3542 "not-registered",
3543 TriggerKind::Api,
3544 json!({}),
3545 EnqueueOptions {
3546 concurrency_limits: vec![ConcurrencyLimit::new("repo:acme", 0)],
3547 ..Default::default()
3548 },
3549 )
3550 .await
3551 .unwrap_err();
3552 assert!(
3553 matches!(err, EngineError::InvalidConcurrencyLimit(_)),
3554 "{err:?}"
3555 );
3556
3557 let page = engine
3558 .store()
3559 .list_runs(RunFilter::default(), 1, 10)
3560 .await
3561 .unwrap();
3562 assert_eq!(page.total, 0, "no run may be created");
3563 }
3564
3565 struct UrgentWorkflow;
3566
3567 impl WorkflowHandler for UrgentWorkflow {
3568 fn name(&self) -> &str {
3569 "urgent-workflow"
3570 }
3571
3572 fn priority(&self) -> i16 {
3573 60
3574 }
3575
3576 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3577 Box::pin(async { Ok(()) })
3578 }
3579 }
3580
3581 #[tokio::test]
3582 async fn enqueue_priority_defaults_to_the_handler_priority() {
3583 let mut engine = create_test_engine();
3584 engine.register(EchoWorkflow).unwrap();
3585 engine.register(UrgentWorkflow).unwrap();
3586
3587 let echo = engine
3588 .enqueue_handler_with_options(
3589 "echo-workflow",
3590 TriggerKind::Api,
3591 json!({}),
3592 EnqueueOptions::default(),
3593 )
3594 .await
3595 .unwrap()
3596 .into_run();
3597 assert_eq!(echo.priority, 0);
3598
3599 let urgent = engine
3600 .enqueue_handler_with_options(
3601 "urgent-workflow",
3602 TriggerKind::Api,
3603 json!({}),
3604 EnqueueOptions::default(),
3605 )
3606 .await
3607 .unwrap()
3608 .into_run();
3609 assert_eq!(urgent.priority, 60);
3610 }
3611
3612 #[tokio::test]
3613 async fn enqueue_priority_explicit_value_overrides_the_handler() {
3614 let mut engine = create_test_engine();
3615 engine.register(UrgentWorkflow).unwrap();
3616
3617 let run = engine
3618 .enqueue_handler_with_options(
3619 "urgent-workflow",
3620 TriggerKind::Api,
3621 json!({}),
3622 EnqueueOptions {
3623 priority: Some(-20),
3624 ..Default::default()
3625 },
3626 )
3627 .await
3628 .unwrap()
3629 .into_run();
3630 assert_eq!(run.priority, -20);
3631
3632 let stored = engine.store().get_run(run.id).await.unwrap().unwrap();
3633 assert_eq!(stored.priority, -20);
3634 }
3635
3636 #[tokio::test]
3637 async fn enqueue_priority_out_of_range_is_rejected() {
3638 let mut engine = create_test_engine();
3639 engine.register(EchoWorkflow).unwrap();
3640
3641 for priority in [MAX_PRIORITY + 1, MIN_PRIORITY - 1] {
3642 let err = engine
3643 .enqueue_handler_with_options(
3644 "echo-workflow",
3645 TriggerKind::Api,
3646 json!({}),
3647 EnqueueOptions {
3648 priority: Some(priority),
3649 ..Default::default()
3650 },
3651 )
3652 .await
3653 .unwrap_err();
3654 assert!(matches!(err, EngineError::InvalidPriority(_)), "{err:?}");
3655 }
3656
3657 let err = engine
3659 .enqueue_handler_with_options(
3660 "not-registered",
3661 TriggerKind::Api,
3662 json!({}),
3663 EnqueueOptions {
3664 priority: Some(MAX_PRIORITY + 1),
3665 ..Default::default()
3666 },
3667 )
3668 .await
3669 .unwrap_err();
3670 assert!(matches!(err, EngineError::InvalidPriority(_)), "{err:?}");
3671
3672 let page = engine
3673 .store()
3674 .list_runs(RunFilter::default(), 1, 10)
3675 .await
3676 .unwrap();
3677 assert_eq!(page.total, 0, "no run may be created");
3678 }
3679
3680 #[tokio::test]
3681 async fn enqueue_priority_bounds_are_accepted() {
3682 let mut engine = create_test_engine();
3683 engine.register(EchoWorkflow).unwrap();
3684
3685 for priority in [MIN_PRIORITY, MAX_PRIORITY] {
3686 let run = engine
3687 .enqueue_handler_with_options(
3688 "echo-workflow",
3689 TriggerKind::Api,
3690 json!({}),
3691 EnqueueOptions {
3692 priority: Some(priority),
3693 ..Default::default()
3694 },
3695 )
3696 .await
3697 .unwrap()
3698 .into_run();
3699 assert_eq!(run.priority, priority);
3700 }
3701 }
3702
3703 struct GpuWorkflow;
3704
3705 impl WorkflowHandler for GpuWorkflow {
3706 fn name(&self) -> &str {
3707 "gpu-workflow"
3708 }
3709
3710 fn required_worker_tags(&self) -> Vec<String> {
3711 vec!["gpu".to_string()]
3712 }
3713
3714 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3715 Box::pin(async move { Ok(()) })
3716 }
3717 }
3718
3719 #[tokio::test]
3720 async fn enqueue_merges_handler_and_request_worker_tags() {
3721 let mut engine = create_test_engine();
3722 engine.register(GpuWorkflow).unwrap();
3723
3724 let run = engine
3725 .enqueue_handler_with_options(
3726 "gpu-workflow",
3727 TriggerKind::Api,
3728 json!({}),
3729 EnqueueOptions {
3730 worker_tags: vec!["region:eu".to_string(), "gpu".to_string()],
3731 ..Default::default()
3732 },
3733 )
3734 .await
3735 .unwrap()
3736 .into_run();
3737
3738 assert_eq!(
3739 run.worker_tags,
3740 vec!["gpu".to_string(), "region:eu".to_string()]
3741 );
3742 }
3743
3744 #[tokio::test]
3745 async fn enqueue_without_worker_tags_keeps_handler_tags() {
3746 let mut engine = create_test_engine();
3747 engine.register(GpuWorkflow).unwrap();
3748 engine.register(EchoWorkflow).unwrap();
3749
3750 let gpu = engine
3751 .enqueue_handler("gpu-workflow", TriggerKind::Api, json!({}), 0)
3752 .await
3753 .unwrap();
3754 assert_eq!(gpu.worker_tags, vec!["gpu".to_string()]);
3755
3756 let echo = engine
3757 .enqueue_handler("echo-workflow", TriggerKind::Api, json!({}), 0)
3758 .await
3759 .unwrap();
3760 assert!(echo.worker_tags.is_empty());
3761 }
3762
3763 #[tokio::test]
3764 async fn enqueue_rejects_invalid_worker_tags() {
3765 let mut engine = create_test_engine();
3766 engine.register(EchoWorkflow).unwrap();
3767
3768 for worker_tags in [
3769 vec!["bad,tag".to_string()],
3770 vec![" ".to_string()],
3771 vec!["x".repeat(65)],
3772 ] {
3773 let err = engine
3774 .enqueue_handler_with_options(
3775 "echo-workflow",
3776 TriggerKind::Api,
3777 json!({}),
3778 EnqueueOptions {
3779 worker_tags,
3780 ..Default::default()
3781 },
3782 )
3783 .await
3784 .unwrap_err();
3785 assert!(matches!(err, EngineError::InvalidWorkerTag(_)), "{err:?}");
3786 }
3787
3788 let err = engine
3790 .enqueue_handler_with_options(
3791 "not-registered",
3792 TriggerKind::Api,
3793 json!({}),
3794 EnqueueOptions {
3795 worker_tags: vec!["bad,tag".to_string()],
3796 ..Default::default()
3797 },
3798 )
3799 .await
3800 .unwrap_err();
3801 assert!(matches!(err, EngineError::InvalidWorkerTag(_)), "{err:?}");
3802 }
3803
3804 #[test]
3805 fn worker_tags_are_unset_by_default() {
3806 let engine = create_test_engine();
3807 assert!(engine.worker_tags().is_none());
3808 }
3809
3810 #[test]
3811 fn set_worker_tags_stores_the_tags() {
3812 let mut engine = create_test_engine();
3813 engine.set_worker_tags(vec!["gpu".to_string()]);
3814 assert_eq!(engine.worker_tags(), Some(&["gpu".to_string()][..]));
3815
3816 engine.set_worker_tags(Vec::new());
3817 assert_eq!(engine.worker_tags(), Some(&[][..]));
3818 }
3819
3820 #[tokio::test]
3821 async fn run_handler_records_handler_worker_tags() {
3822 let mut engine = create_test_engine();
3823 engine.register(GpuWorkflow).unwrap();
3824
3825 let result = engine
3826 .run_handler("gpu-workflow", TriggerKind::Manual, json!({}))
3827 .await
3828 .unwrap();
3829 assert_eq!(result.run.worker_tags, vec!["gpu".to_string()]);
3830 }
3831
3832 #[tokio::test]
3833 async fn run_handler_leaves_the_run_unattributed() {
3834 let mut engine = create_test_engine();
3835 engine.register(EchoWorkflow).unwrap();
3836
3837 let run = engine
3838 .run_handler("echo-workflow", TriggerKind::Manual, json!({}))
3839 .await
3840 .unwrap()
3841 .run;
3842
3843 assert!(run.created_by.is_none());
3844 }
3845
3846 #[tokio::test]
3847 async fn run_handler_priority_comes_from_the_handler() {
3848 let mut engine = create_test_engine();
3849 engine.register(UrgentWorkflow).unwrap();
3850
3851 let run = engine
3852 .run_handler("urgent-workflow", TriggerKind::Manual, json!({}))
3853 .await
3854 .unwrap()
3855 .run;
3856
3857 assert_eq!(run.priority, 60);
3858 }
3859
3860 #[tokio::test]
3861 async fn engine_register_boxed() {
3862 let mut engine = create_test_engine();
3863 let handler: Box<dyn WorkflowHandler> = Box::new(EchoWorkflow);
3864 let result = engine.register_boxed(handler);
3865 assert!(result.is_ok());
3866 assert_eq!(engine.handler_names().len(), 1);
3867 }
3868
3869 #[tokio::test]
3870 async fn engine_store_and_provider_accessors() {
3871 let store = Arc::new(InMemoryStore::new());
3872 let inner = ClaudeCodeProvider::new();
3873 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
3874 inner,
3875 "/tmp/ironflow-fixtures",
3876 ));
3877 let engine = Engine::new(store.clone(), provider.clone());
3878
3879 let _ = engine.store();
3881 let _ = engine.provider();
3882 }
3883
3884 use crate::operation::{Operation, OperationContext};
3889 use async_trait::async_trait;
3890 use ironflow_core::error::OperationError;
3891 use ironflow_store::models::StepKind;
3892
3893 struct FakeGitlabOp {
3894 project_id: u64,
3895 title: String,
3896 }
3897
3898 #[async_trait]
3899 impl Operation for FakeGitlabOp {
3900 fn kind(&self) -> &str {
3901 "gitlab"
3902 }
3903
3904 async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
3905 Ok(json!({
3906 "issue_id": 42,
3907 "project_id": self.project_id,
3908 "title": self.title,
3909 }))
3910 }
3911
3912 fn input(&self) -> Option<Value> {
3913 Some(json!({
3914 "project_id": self.project_id,
3915 "title": self.title,
3916 }))
3917 }
3918 }
3919
3920 struct FailingOp;
3921
3922 #[async_trait]
3923 impl Operation for FailingOp {
3924 fn kind(&self) -> &str {
3925 "broken-service"
3926 }
3927
3928 async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
3929 Err(OperationError::Http {
3930 status: None,
3931 message: "service unavailable".to_string(),
3932 })
3933 }
3934 }
3935
3936 struct OperationWorkflow;
3937
3938 impl WorkflowHandler for OperationWorkflow {
3939 fn name(&self) -> &str {
3940 "operation-workflow"
3941 }
3942
3943 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3944 Box::pin(async move {
3945 let op = FakeGitlabOp {
3946 project_id: 123,
3947 title: "Bug report".to_string(),
3948 };
3949 ctx.operation("create-issue", &op).await?;
3950 Ok(())
3951 })
3952 }
3953 }
3954
3955 struct FailingOperationWorkflow;
3956
3957 impl WorkflowHandler for FailingOperationWorkflow {
3958 fn name(&self) -> &str {
3959 "failing-operation-workflow"
3960 }
3961
3962 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3963 Box::pin(async move {
3964 ctx.operation("broken-call", &FailingOp).await?;
3965 Ok(())
3966 })
3967 }
3968 }
3969
3970 struct MixedWorkflow;
3971
3972 impl WorkflowHandler for MixedWorkflow {
3973 fn name(&self) -> &str {
3974 "mixed-workflow"
3975 }
3976
3977 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3978 Box::pin(async move {
3979 ctx.shell("build", ShellConfig::new("echo built")).await?;
3980 let op = FakeGitlabOp {
3981 project_id: 456,
3982 title: "Deploy done".to_string(),
3983 };
3984 let result = ctx.operation("notify-gitlab", &op).await?;
3985 assert_eq!(result.output["issue_id"], 42);
3986 Ok(())
3987 })
3988 }
3989 }
3990
3991 #[tokio::test]
3992 async fn operation_step_happy_path() {
3993 let mut engine = create_test_engine();
3994 engine.register(OperationWorkflow).unwrap();
3995
3996 let run = engine
3997 .run_handler("operation-workflow", TriggerKind::Manual, json!({}))
3998 .await
3999 .unwrap()
4000 .run;
4001
4002 assert_eq!(run.status.state, RunStatus::Completed);
4003
4004 let steps = engine.store().list_steps(run.id).await.unwrap();
4005
4006 assert_eq!(steps.len(), 1);
4007 assert_eq!(steps[0].name, "create-issue");
4008 assert_eq!(steps[0].kind, StepKind::Custom("gitlab".to_string()));
4009 assert_eq!(
4010 steps[0].status.state,
4011 ironflow_store::models::StepStatus::Completed
4012 );
4013
4014 let output = steps[0].output.as_ref().unwrap();
4015 assert_eq!(output["issue_id"], 42);
4016 assert_eq!(output["project_id"], 123);
4017
4018 let input = steps[0].input.as_ref().unwrap();
4019 assert_eq!(input["project_id"], 123);
4020 assert_eq!(input["title"], "Bug report");
4021 }
4022
4023 #[tokio::test]
4024 async fn operation_step_failure_marks_run_failed() {
4025 let mut engine = create_test_engine();
4026 engine.register(FailingOperationWorkflow).unwrap();
4027
4028 let result = engine
4029 .run_handler("failing-operation-workflow", TriggerKind::Manual, json!({}))
4030 .await;
4031
4032 assert!(result.is_err());
4033 }
4034
4035 #[tokio::test]
4036 async fn operation_mixed_with_shell_steps() {
4037 let mut engine = create_test_engine();
4038 engine.register(MixedWorkflow).unwrap();
4039
4040 let run = engine
4041 .run_handler("mixed-workflow", TriggerKind::Manual, json!({}))
4042 .await
4043 .unwrap()
4044 .run;
4045
4046 assert_eq!(run.status.state, RunStatus::Completed);
4047
4048 let steps = engine.store().list_steps(run.id).await.unwrap();
4049
4050 assert_eq!(steps.len(), 2);
4051 assert_eq!(steps[0].kind, StepKind::Shell);
4052 assert_eq!(steps[1].kind, StepKind::Custom("gitlab".to_string()));
4053 assert_eq!(steps[0].position, 0);
4054 assert_eq!(steps[1].position, 1);
4055 }
4056
4057 use crate::config::ApprovalConfig;
4062
4063 struct SingleApprovalWorkflow;
4064
4065 impl WorkflowHandler for SingleApprovalWorkflow {
4066 fn name(&self) -> &str {
4067 "single-approval"
4068 }
4069
4070 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
4071 Box::pin(async move {
4072 ctx.shell("build", ShellConfig::new("echo built")).await?;
4073 ctx.approval("gate", ApprovalConfig::new("OK?")).await?;
4074 ctx.shell("deploy", ShellConfig::new("echo deployed"))
4075 .await?;
4076 Ok(())
4077 })
4078 }
4079 }
4080
4081 struct DoubleApprovalWorkflow;
4082
4083 impl WorkflowHandler for DoubleApprovalWorkflow {
4084 fn name(&self) -> &str {
4085 "double-approval"
4086 }
4087
4088 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
4089 Box::pin(async move {
4090 ctx.shell("build", ShellConfig::new("echo built")).await?;
4091 ctx.approval("staging-gate", ApprovalConfig::new("Deploy staging?"))
4092 .await?;
4093 ctx.shell("deploy-staging", ShellConfig::new("echo staging"))
4094 .await?;
4095 ctx.approval("prod-gate", ApprovalConfig::new("Deploy prod?"))
4096 .await?;
4097 ctx.shell("deploy-prod", ShellConfig::new("echo prod"))
4098 .await?;
4099 Ok(())
4100 })
4101 }
4102 }
4103
4104 #[tokio::test]
4105 async fn approval_pauses_run() {
4106 let mut engine = create_test_engine();
4107 engine.register(SingleApprovalWorkflow).unwrap();
4108
4109 let run = engine
4110 .run_handler("single-approval", TriggerKind::Manual, json!({}))
4111 .await
4112 .unwrap()
4113 .run;
4114
4115 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
4116
4117 let steps = engine.store().list_steps(run.id).await.unwrap();
4118 assert_eq!(steps.len(), 2); assert_eq!(steps[0].kind, StepKind::Shell);
4120 assert_eq!(steps[0].status.state, StepStatus::Completed);
4121 assert_eq!(steps[1].kind, StepKind::Approval);
4122 assert_eq!(steps[1].status.state, StepStatus::AwaitingApproval);
4123 }
4124
4125 async fn approve_open_gate(store: &Arc<dyn Store>, run_id: Uuid) {
4128 let gate = store
4129 .list_steps(run_id)
4130 .await
4131 .unwrap()
4132 .into_iter()
4133 .find(|step| step.status.state == StepStatus::AwaitingApproval)
4134 .expect("an open approval gate");
4135 store
4136 .record_step_approval(
4137 gate.id,
4138 StepApproval {
4139 user_id: Uuid::now_v7(),
4140 approved_by: "tester".to_string(),
4141 at: Utc::now(),
4142 },
4143 )
4144 .await
4145 .unwrap();
4146 store
4147 .update_run_status(run_id, RunStatus::Running)
4148 .await
4149 .unwrap();
4150 }
4151
4152 #[tokio::test]
4153 async fn approval_resume_completes_run() {
4154 let mut engine = create_test_engine();
4155 engine.register(SingleApprovalWorkflow).unwrap();
4156
4157 let run = engine
4159 .run_handler("single-approval", TriggerKind::Manual, json!({}))
4160 .await
4161 .unwrap()
4162 .run;
4163 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
4164
4165 approve_open_gate(engine.store(), run.id).await;
4167
4168 let resumed = engine.resume_run(run.id).await.unwrap().run;
4170 assert_eq!(resumed.status.state, RunStatus::Completed);
4171
4172 let steps = engine.store().list_steps(run.id).await.unwrap();
4173 assert_eq!(steps.len(), 3); assert_eq!(steps[0].name, "build");
4175 assert_eq!(steps[0].status.state, StepStatus::Completed);
4176 assert_eq!(steps[1].name, "gate");
4177 assert_eq!(steps[1].kind, StepKind::Approval);
4178 assert_eq!(steps[1].status.state, StepStatus::Completed);
4179 assert_eq!(steps[2].name, "deploy");
4180 assert_eq!(steps[2].status.state, StepStatus::Completed);
4181 }
4182
4183 #[tokio::test]
4184 async fn double_approval_two_resumes() {
4185 let mut engine = create_test_engine();
4186 engine.register(DoubleApprovalWorkflow).unwrap();
4187
4188 let run = engine
4190 .run_handler("double-approval", TriggerKind::Manual, json!({}))
4191 .await
4192 .unwrap()
4193 .run;
4194 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
4195
4196 let steps = engine.store().list_steps(run.id).await.unwrap();
4197 assert_eq!(steps.len(), 2); approve_open_gate(engine.store(), run.id).await;
4201
4202 let resumed = engine.resume_run(run.id).await.unwrap().run;
4203 assert_eq!(resumed.status.state, RunStatus::AwaitingApproval);
4204
4205 let steps = engine.store().list_steps(run.id).await.unwrap();
4206 assert_eq!(steps.len(), 4); approve_open_gate(engine.store(), run.id).await;
4210
4211 let final_run = engine.resume_run(run.id).await.unwrap().run;
4212 assert_eq!(final_run.status.state, RunStatus::Completed);
4213
4214 let steps = engine.store().list_steps(run.id).await.unwrap();
4215 assert_eq!(steps.len(), 5);
4216 assert_eq!(steps[0].name, "build");
4217 assert_eq!(steps[1].name, "staging-gate");
4218 assert_eq!(steps[2].name, "deploy-staging");
4219 assert_eq!(steps[3].name, "prod-gate");
4220 assert_eq!(steps[4].name, "deploy-prod");
4221
4222 for step in &steps {
4223 assert_eq!(step.status.state, StepStatus::Completed);
4224 }
4225 }
4226
4227 use ironflow_store::models::{NewStep, StepUpdate, step_trace_id};
4232
4233 async fn create_step_with_status(
4234 store: &Arc<dyn Store>,
4235 run_id: Uuid,
4236 name: &str,
4237 position: u32,
4238 status: StepStatus,
4239 ) -> ironflow_store::models::Step {
4240 let step = store
4241 .create_step(NewStep {
4242 run_id,
4243 trace_id: step_trace_id(run_id, name, position),
4244 name: name.to_string(),
4245 kind: StepKind::Shell,
4246 position,
4247 input: None,
4248 is_error_handler: false,
4249 })
4250 .await
4251 .unwrap();
4252
4253 match status {
4254 StepStatus::Pending => {}
4255 StepStatus::Running => {
4256 store
4257 .update_step(
4258 step.id,
4259 StepUpdate {
4260 status: Some(StepStatus::Running),
4261 ..StepUpdate::default()
4262 },
4263 )
4264 .await
4265 .unwrap();
4266 }
4267 StepStatus::Completed => {
4268 store
4269 .update_step(
4270 step.id,
4271 StepUpdate {
4272 status: Some(StepStatus::Running),
4273 ..StepUpdate::default()
4274 },
4275 )
4276 .await
4277 .unwrap();
4278 store
4279 .update_step(
4280 step.id,
4281 StepUpdate {
4282 status: Some(StepStatus::Completed),
4283 ..StepUpdate::default()
4284 },
4285 )
4286 .await
4287 .unwrap();
4288 }
4289 StepStatus::AwaitingApproval => {
4290 store
4291 .update_step(
4292 step.id,
4293 StepUpdate {
4294 status: Some(StepStatus::Running),
4295 ..StepUpdate::default()
4296 },
4297 )
4298 .await
4299 .unwrap();
4300 store
4301 .update_step(
4302 step.id,
4303 StepUpdate {
4304 status: Some(StepStatus::AwaitingApproval),
4305 ..StepUpdate::default()
4306 },
4307 )
4308 .await
4309 .unwrap();
4310 }
4311 _ => panic!("unsupported status for test helper: {status}"),
4312 }
4313
4314 store.get_step(step.id).await.unwrap().unwrap()
4315 }
4316
4317 #[tokio::test]
4318 async fn fail_orphaned_steps_marks_running_as_failed() {
4319 let engine = create_test_engine();
4320 let run = engine
4321 .store()
4322 .create_run(NewRun {
4323 created_by: None,
4324 workflow_name: "test".to_string(),
4325 trigger: TriggerKind::Manual,
4326 payload: json!({}),
4327 max_retries: 0,
4328 handler_version: None,
4329 labels: HashMap::new(),
4330 scheduled_at: None,
4331 idempotency_key: None,
4332 concurrency_key: None,
4333 priority: 0,
4334 concurrency_limits: Vec::new(),
4335 max_cost_usd: None,
4336 worker_tags: Vec::new(),
4337 })
4338 .await
4339 .unwrap()
4340 .into_run();
4341
4342 let step = create_step_with_status(
4343 engine.store(),
4344 run.id,
4345 "running-step",
4346 0,
4347 StepStatus::Running,
4348 )
4349 .await;
4350
4351 engine
4352 .fail_orphaned_steps(run.id, "parent run timed out")
4353 .await
4354 .unwrap();
4355
4356 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4357 assert_eq!(updated.status.state, StepStatus::Failed);
4358 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
4359 assert!(updated.completed_at.is_some());
4360 }
4361
4362 #[tokio::test]
4363 async fn fail_orphaned_steps_marks_pending_as_skipped() {
4364 let engine = create_test_engine();
4365 let run = engine
4366 .store()
4367 .create_run(NewRun {
4368 created_by: None,
4369 workflow_name: "test".to_string(),
4370 trigger: TriggerKind::Manual,
4371 payload: json!({}),
4372 max_retries: 0,
4373 handler_version: None,
4374 labels: HashMap::new(),
4375 scheduled_at: None,
4376 idempotency_key: None,
4377 concurrency_key: None,
4378 priority: 0,
4379 concurrency_limits: Vec::new(),
4380 max_cost_usd: None,
4381 worker_tags: Vec::new(),
4382 })
4383 .await
4384 .unwrap()
4385 .into_run();
4386
4387 let step = create_step_with_status(
4388 engine.store(),
4389 run.id,
4390 "pending-step",
4391 0,
4392 StepStatus::Pending,
4393 )
4394 .await;
4395
4396 engine
4397 .fail_orphaned_steps(run.id, "parent run timed out")
4398 .await
4399 .unwrap();
4400
4401 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4402 assert_eq!(updated.status.state, StepStatus::Skipped);
4403 assert!(updated.error.is_none());
4404 assert!(updated.completed_at.is_some());
4405 }
4406
4407 #[tokio::test]
4408 async fn fail_orphaned_steps_marks_awaiting_approval_as_failed() {
4409 let engine = create_test_engine();
4410 let run = engine
4411 .store()
4412 .create_run(NewRun {
4413 created_by: None,
4414 workflow_name: "test".to_string(),
4415 trigger: TriggerKind::Manual,
4416 payload: json!({}),
4417 max_retries: 0,
4418 handler_version: None,
4419 labels: HashMap::new(),
4420 scheduled_at: None,
4421 idempotency_key: None,
4422 concurrency_key: None,
4423 priority: 0,
4424 concurrency_limits: Vec::new(),
4425 max_cost_usd: None,
4426 worker_tags: Vec::new(),
4427 })
4428 .await
4429 .unwrap()
4430 .into_run();
4431
4432 let step = create_step_with_status(
4433 engine.store(),
4434 run.id,
4435 "approval-step",
4436 0,
4437 StepStatus::AwaitingApproval,
4438 )
4439 .await;
4440
4441 engine
4442 .fail_orphaned_steps(run.id, "parent run timed out")
4443 .await
4444 .unwrap();
4445
4446 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4447 assert_eq!(updated.status.state, StepStatus::Failed);
4448 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
4449 assert!(updated.completed_at.is_some());
4450 }
4451
4452 #[tokio::test]
4453 async fn fail_orphaned_steps_skips_terminal_steps() {
4454 let engine = create_test_engine();
4455 let run = engine
4456 .store()
4457 .create_run(NewRun {
4458 created_by: None,
4459 workflow_name: "test".to_string(),
4460 trigger: TriggerKind::Manual,
4461 payload: json!({}),
4462 max_retries: 0,
4463 handler_version: None,
4464 labels: HashMap::new(),
4465 scheduled_at: None,
4466 idempotency_key: None,
4467 concurrency_key: None,
4468 priority: 0,
4469 concurrency_limits: Vec::new(),
4470 max_cost_usd: None,
4471 worker_tags: Vec::new(),
4472 })
4473 .await
4474 .unwrap()
4475 .into_run();
4476
4477 let completed_step =
4478 create_step_with_status(engine.store(), run.id, "done", 0, StepStatus::Completed).await;
4479 let running_step =
4480 create_step_with_status(engine.store(), run.id, "in-flight", 1, StepStatus::Running)
4481 .await;
4482
4483 engine
4484 .fail_orphaned_steps(run.id, "parent run timed out")
4485 .await
4486 .unwrap();
4487
4488 let completed = engine
4489 .store()
4490 .get_step(completed_step.id)
4491 .await
4492 .unwrap()
4493 .unwrap();
4494 assert_eq!(completed.status.state, StepStatus::Completed);
4495
4496 let failed = engine
4497 .store()
4498 .get_step(running_step.id)
4499 .await
4500 .unwrap()
4501 .unwrap();
4502 assert_eq!(failed.status.state, StepStatus::Failed);
4503 }
4504
4505 #[tokio::test]
4506 async fn fail_orphaned_steps_mixed_states() {
4507 let engine = create_test_engine();
4508 let run = engine
4509 .store()
4510 .create_run(NewRun {
4511 created_by: None,
4512 workflow_name: "test".to_string(),
4513 trigger: TriggerKind::Manual,
4514 payload: json!({}),
4515 max_retries: 0,
4516 handler_version: None,
4517 labels: HashMap::new(),
4518 scheduled_at: None,
4519 idempotency_key: None,
4520 concurrency_key: None,
4521 priority: 0,
4522 concurrency_limits: Vec::new(),
4523 max_cost_usd: None,
4524 worker_tags: Vec::new(),
4525 })
4526 .await
4527 .unwrap()
4528 .into_run();
4529
4530 let s_completed =
4531 create_step_with_status(engine.store(), run.id, "step-1", 0, StepStatus::Completed)
4532 .await;
4533 let s_running =
4534 create_step_with_status(engine.store(), run.id, "step-2", 1, StepStatus::Running).await;
4535 let s_pending =
4536 create_step_with_status(engine.store(), run.id, "step-3", 2, StepStatus::Pending).await;
4537
4538 engine.fail_orphaned_steps(run.id, "timeout").await.unwrap();
4539
4540 let r_completed = engine
4541 .store()
4542 .get_step(s_completed.id)
4543 .await
4544 .unwrap()
4545 .unwrap();
4546 assert_eq!(r_completed.status.state, StepStatus::Completed);
4547
4548 let r_running = engine
4549 .store()
4550 .get_step(s_running.id)
4551 .await
4552 .unwrap()
4553 .unwrap();
4554 assert_eq!(r_running.status.state, StepStatus::Failed);
4555 assert_eq!(r_running.error.as_deref(), Some("timeout"));
4556
4557 let r_pending = engine
4558 .store()
4559 .get_step(s_pending.id)
4560 .await
4561 .unwrap()
4562 .unwrap();
4563 assert_eq!(r_pending.status.state, StepStatus::Skipped);
4564 assert!(r_pending.error.is_none());
4565 }
4566
4567 #[tokio::test]
4568 async fn fail_orphaned_steps_no_steps_is_noop() {
4569 let engine = create_test_engine();
4570 let run = engine
4571 .store()
4572 .create_run(NewRun {
4573 created_by: None,
4574 workflow_name: "test".to_string(),
4575 trigger: TriggerKind::Manual,
4576 payload: json!({}),
4577 max_retries: 0,
4578 handler_version: None,
4579 labels: HashMap::new(),
4580 scheduled_at: None,
4581 idempotency_key: None,
4582 concurrency_key: None,
4583 priority: 0,
4584 concurrency_limits: Vec::new(),
4585 max_cost_usd: None,
4586 worker_tags: Vec::new(),
4587 })
4588 .await
4589 .unwrap()
4590 .into_run();
4591
4592 let result = engine.fail_orphaned_steps(run.id, "timeout").await;
4593 assert!(result.is_ok());
4594 }
4595
4596 #[tokio::test]
4597 async fn fail_orphaned_steps_preserves_existing_error() {
4598 let engine = create_test_engine();
4599 let run = engine
4600 .store()
4601 .create_run(NewRun {
4602 created_by: None,
4603 workflow_name: "test".to_string(),
4604 trigger: TriggerKind::Manual,
4605 payload: json!({}),
4606 max_retries: 0,
4607 handler_version: None,
4608 labels: HashMap::new(),
4609 scheduled_at: None,
4610 idempotency_key: None,
4611 concurrency_key: None,
4612 priority: 0,
4613 concurrency_limits: Vec::new(),
4614 max_cost_usd: None,
4615 worker_tags: Vec::new(),
4616 })
4617 .await
4618 .unwrap()
4619 .into_run();
4620
4621 let step_with_error = create_step_with_status(
4622 engine.store(),
4623 run.id,
4624 "already-errored",
4625 0,
4626 StepStatus::Running,
4627 )
4628 .await;
4629
4630 engine
4631 .store()
4632 .update_step(
4633 step_with_error.id,
4634 StepUpdate {
4635 error: Some("real error from provider".to_string()),
4636 ..StepUpdate::default()
4637 },
4638 )
4639 .await
4640 .unwrap();
4641
4642 let step_no_error = create_step_with_status(
4643 engine.store(),
4644 run.id,
4645 "no-error-yet",
4646 1,
4647 StepStatus::Running,
4648 )
4649 .await;
4650
4651 engine
4652 .fail_orphaned_steps(run.id, "parent run failed")
4653 .await
4654 .unwrap();
4655
4656 let updated_with = engine
4657 .store()
4658 .get_step(step_with_error.id)
4659 .await
4660 .unwrap()
4661 .unwrap();
4662 assert_eq!(updated_with.status.state, StepStatus::Failed);
4663 assert_eq!(
4664 updated_with.error.as_deref(),
4665 Some("real error from provider"),
4666 );
4667
4668 let updated_without = engine
4669 .store()
4670 .get_step(step_no_error.id)
4671 .await
4672 .unwrap()
4673 .unwrap();
4674 assert_eq!(updated_without.status.state, StepStatus::Failed);
4675 assert_eq!(updated_without.error.as_deref(), Some("parent run failed"),);
4676 }
4677
4678 use std::sync::atomic::{AtomicUsize, Ordering};
4683 use std::time::Duration;
4684
4685 use ironflow_store::models::LeaseRequest;
4686 use tokio::time::sleep;
4687
4688 struct CountingWorkflow {
4690 executions: Arc<AtomicUsize>,
4691 }
4692
4693 impl WorkflowHandler for CountingWorkflow {
4694 fn name(&self) -> &str {
4695 "local-resume-counter"
4696 }
4697
4698 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
4699 Box::pin(async move {
4700 self.executions.fetch_add(1, Ordering::SeqCst);
4701 Ok(())
4702 })
4703 }
4704 }
4705
4706 async fn engine_with_pending_counter() -> (Arc<Engine>, Arc<AtomicUsize>, Uuid) {
4708 let executions = Arc::new(AtomicUsize::new(0));
4709 let mut engine = create_test_engine();
4710 engine
4711 .register(CountingWorkflow {
4712 executions: executions.clone(),
4713 })
4714 .unwrap();
4715 let run = engine
4716 .enqueue_handler("local-resume-counter", TriggerKind::Manual, json!({}), 0)
4717 .await
4718 .unwrap();
4719 (Arc::new(engine), executions, run.id)
4720 }
4721
4722 #[tokio::test]
4723 async fn local_resume_leaves_a_run_claimed_by_a_worker_alone() {
4724 let (engine, executions, run_id) = engine_with_pending_counter().await;
4725 let lease = LeaseRequest {
4726 worker_id: "worker-1".to_string(),
4727 ttl: Duration::from_secs(30),
4728 };
4729 let picked = engine
4730 .store()
4731 .pick_next_pending_for(Some(lease), None)
4732 .await
4733 .unwrap()
4734 .expect("the run is due");
4735 assert_eq!(picked.id, run_id);
4736
4737 engine.spawn_local_resume(run_id);
4738 sleep(Duration::from_millis(200)).await;
4739
4740 assert_eq!(executions.load(Ordering::SeqCst), 0);
4741 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4742 assert_eq!(run.status.state, RunStatus::Running);
4743 assert_eq!(run.worker_id.as_deref(), Some("worker-1"));
4744 }
4745
4746 #[tokio::test]
4747 async fn local_resume_claims_a_pending_run_once() {
4748 let (engine, executions, run_id) = engine_with_pending_counter().await;
4749
4750 engine.spawn_local_resume(run_id);
4751 engine.spawn_local_resume(run_id);
4752 sleep(Duration::from_millis(200)).await;
4753
4754 assert_eq!(executions.load(Ordering::SeqCst), 1);
4755 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4756 assert_eq!(run.status.state, RunStatus::Completed);
4757 }
4758
4759 #[tokio::test]
4760 async fn local_resume_of_a_running_run_executes_it_once() {
4761 let (engine, executions, run_id) = engine_with_pending_counter().await;
4762 engine
4763 .store()
4764 .update_run_status(run_id, RunStatus::Running)
4765 .await
4766 .unwrap();
4767
4768 engine.spawn_local_resume(run_id);
4769 engine.spawn_local_resume(run_id);
4770 sleep(Duration::from_millis(200)).await;
4771
4772 assert_eq!(executions.load(Ordering::SeqCst), 1);
4773 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4774 assert_eq!(run.status.state, RunStatus::Completed);
4775 }
4776}