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, SignalRejected, SignalResumed, received_output, validate_step_payload,
61};
62use ironflow_core::decision::DecisionProvider;
63
64const LIST_RUNS_PAGE_SIZE: u32 = 100;
66
67#[derive(Debug, Clone)]
86pub struct WorkflowResult {
87 pub run: Run,
89 pub steps: Vec<StepResult>,
91}
92
93#[derive(Debug, Clone, Default)]
112pub struct EnqueueOptions {
113 pub max_retries: u32,
115 pub labels: HashMap<String, String>,
117 pub scheduled_at: Option<DateTime<Utc>>,
120 pub max_cost_usd: Option<Decimal>,
124 pub created_by: Option<RunActor>,
127 pub idempotency_key: Option<String>,
133 pub concurrency_key: Option<String>,
139 pub concurrency_limits: Vec<ConcurrencyLimit>,
147 pub priority: Option<i16>,
156 pub worker_tags: Vec<String>,
162}
163
164#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
175pub enum ExecutionMode {
176 #[default]
180 Local,
181 Workers,
184}
185
186pub struct Engine {
227 store: Arc<dyn Store>,
228 provider: Arc<dyn AgentProvider>,
229 handlers: HashMap<String, Arc<dyn WorkflowHandler>>,
230 event_publisher: EventPublisher,
231 log_sender: Option<LogSender>,
232 budget: BudgetConfig,
233 artifact_sink: Option<Arc<dyn ArtifactSink>>,
234 guard_config: Option<WorkflowGuardConfig>,
235 event_bus: Option<WorkflowEventBus>,
236 decision_provider: Option<Arc<dyn DecisionProvider>>,
237 step_interceptor: Option<Arc<dyn StepInterceptor>>,
238 execution_mode: ExecutionMode,
239 worker_tags: Option<Arc<Vec<String>>>,
240 active_runs: ActiveRuns,
241 local_resumes: ActiveRuns,
242}
243
244type ActiveRuns = Arc<Mutex<HashMap<Uuid, Arc<AsyncMutex<()>>>>>;
246
247struct ActiveRunGuard {
253 runs: ActiveRuns,
254 run_id: Uuid,
255 entry: Arc<AsyncMutex<()>>,
256 lock: Option<OwnedMutexGuard<()>>,
257}
258
259impl Drop for ActiveRunGuard {
260 fn drop(&mut self) {
261 self.lock.take();
262 prune_active_run(&self.runs, self.run_id, &self.entry);
263 }
264}
265
266async fn lock_run(runs: &ActiveRuns, run_id: Uuid) -> ActiveRunGuard {
268 let entry = {
269 let mut runs = runs.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
270 Arc::clone(runs.entry(run_id).or_default())
271 };
272 let lock = Arc::clone(&entry).lock_owned().await;
273 ActiveRunGuard {
274 runs: Arc::clone(runs),
275 run_id,
276 entry,
277 lock: Some(lock),
278 }
279}
280
281fn prune_active_run(runs: &ActiveRuns, run_id: Uuid, entry: &Arc<AsyncMutex<()>>) {
286 let mut runs = runs.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
287 let unused = runs
288 .get(&run_id)
289 .is_some_and(|current| Arc::ptr_eq(current, entry))
290 && Arc::strong_count(entry) == 2;
291 if unused {
292 runs.remove(&run_id);
293 }
294}
295
296fn validate_category(handler_name: &str, category: &str) -> Result<(), EngineError> {
306 let reject = |reason: &str| {
307 Err(EngineError::InvalidWorkflow(format!(
308 "handler '{handler_name}' has invalid category '{category}': {reason}"
309 )))
310 };
311
312 if category.is_empty() {
313 return reject("empty category");
314 }
315 if category.starts_with('/') {
316 return reject("leading '/'");
317 }
318 if category.ends_with('/') {
319 return reject("trailing '/'");
320 }
321 for segment in category.split('/') {
322 if segment.is_empty() {
323 return reject("empty segment (double '/')");
324 }
325 if segment.trim().is_empty() {
326 return reject("whitespace-only segment");
327 }
328 }
329 Ok(())
330}
331
332fn chain_label(run: &Run, key: &str) -> Option<Uuid> {
337 if !matches!(run.trigger, TriggerKind::Workflow) {
338 return None;
339 }
340 let id = Uuid::parse_str(run.labels.get(key)?).ok()?;
341 (id != run.id).then_some(id)
342}
343
344pub fn chain_root(run: &Run) -> Option<Uuid> {
364 chain_label(run, LABEL_ROOT_RUN_ID)
365}
366
367pub(crate) fn chain_parent(run: &Run) -> Option<Uuid> {
369 chain_label(run, PARENT_RUN_ID_LABEL)
370}
371
372impl Engine {
373 pub fn new(store: Arc<dyn Store>, provider: Arc<dyn AgentProvider>) -> Self {
389 Self {
390 store,
391 provider,
392 handlers: HashMap::new(),
393 event_publisher: EventPublisher::new(),
394 log_sender: None,
395 budget: BudgetConfig::new(),
396 artifact_sink: None,
397 guard_config: None,
398 event_bus: None,
399 decision_provider: None,
400 step_interceptor: None,
401 execution_mode: ExecutionMode::default(),
402 worker_tags: None,
403 active_runs: Arc::new(Mutex::new(HashMap::new())),
404 local_resumes: Arc::new(Mutex::new(HashMap::new())),
405 }
406 }
407
408 async fn track_execution(&self, run_id: Uuid) -> ActiveRunGuard {
411 lock_run(&self.active_runs, run_id).await
412 }
413
414 async fn wait_until_idle(&self, run_id: Uuid) {
416 let entry = {
417 let runs = self
418 .active_runs
419 .lock()
420 .unwrap_or_else(|poisoned| poisoned.into_inner());
421 runs.get(&run_id).map(Arc::clone)
422 };
423 if let Some(entry) = entry {
424 drop(Arc::clone(&entry).lock_owned().await);
425 prune_active_run(&self.active_runs, run_id, &entry);
426 }
427 }
428
429 pub(crate) fn is_executing(&self, run_id: Uuid) -> bool {
431 let runs = self
432 .active_runs
433 .lock()
434 .unwrap_or_else(|poisoned| poisoned.into_inner());
435 runs.get(&run_id)
436 .is_some_and(|entry| entry.try_lock().is_err())
437 }
438
439 pub fn with_decision_provider(mut self, provider: Arc<dyn DecisionProvider>) -> Self {
460 self.decision_provider = Some(provider);
461 self
462 }
463
464 pub fn with_step_interceptor(mut self, interceptor: Arc<dyn StepInterceptor>) -> Self {
488 self.step_interceptor = Some(interceptor);
489 self
490 }
491
492 pub fn step_interceptor(&self) -> Option<&Arc<dyn StepInterceptor>> {
494 self.step_interceptor.as_ref()
495 }
496
497 pub fn with_budget_config(mut self, budget: BudgetConfig) -> Self {
518 self.budget = budget;
519 self
520 }
521
522 pub fn budget_config(&self) -> &BudgetConfig {
524 &self.budget
525 }
526
527 pub fn with_guard_config(mut self, config: WorkflowGuardConfig) -> Self {
549 self.guard_config = Some(config);
550 self
551 }
552
553 pub fn guard_config(&self) -> Option<&WorkflowGuardConfig> {
555 self.guard_config.as_ref()
556 }
557
558 pub fn with_execution_mode(mut self, mode: ExecutionMode) -> Self {
579 self.execution_mode = mode;
580 self
581 }
582
583 pub fn execution_mode(&self) -> ExecutionMode {
585 self.execution_mode
586 }
587
588 pub fn set_worker_tags(&mut self, tags: Vec<String>) {
611 self.worker_tags = Some(Arc::new(tags));
612 }
613
614 pub fn worker_tags(&self) -> Option<&[String]> {
632 self.worker_tags.as_deref().map(Vec::as_slice)
633 }
634
635 pub fn set_log_sender(&mut self, sender: LogSender) {
641 self.log_sender = Some(sender);
642 }
643
644 pub fn set_artifact_sink(&mut self, sink: Arc<dyn ArtifactSink>) {
662 self.artifact_sink = Some(sink);
663 }
664
665 pub fn artifact_sink(&self) -> Option<&Arc<dyn ArtifactSink>> {
667 self.artifact_sink.as_ref()
668 }
669
670 pub fn set_event_bus(&mut self, bus: WorkflowEventBus) {
691 self.event_bus = Some(bus.clone());
692 self.event_publisher
693 .subscribe(EventBusReleaser::new(bus), &[Event::RUN_STATUS_CHANGED]);
694 }
695
696 pub fn event_bus(&self) -> Option<&WorkflowEventBus> {
698 self.event_bus.as_ref()
699 }
700
701 pub fn store(&self) -> &Arc<dyn Store> {
703 &self.store
704 }
705
706 pub fn provider(&self) -> &Arc<dyn AgentProvider> {
708 &self.provider
709 }
710
711 fn build_context(&self, run: &Run) -> WorkflowContext {
720 let handlers = self.handlers.clone();
721 let resolver: crate::context::HandlerResolver =
722 Arc::new(move |name: &str| handlers.get(name).cloned());
723 let mut ctx = WorkflowContext::with_handler_resolver(
724 run.id,
725 run.workflow_name.clone(),
726 self.store.clone(),
727 self.provider.clone(),
728 resolver,
729 );
730 ctx.carry_over_run_totals(run.retry_count + 1, run.cost_usd, run.duration_ms);
731 ctx.set_max_cost_usd(run.max_cost_usd);
732 ctx.set_run_created_at(run.created_at);
733 if let Some(ref sender) = self.log_sender {
734 ctx.set_log_sender(sender.clone());
735 }
736 if let Some(ref sink) = self.artifact_sink {
737 ctx.set_artifact_sink(sink.clone());
738 }
739 if let Some(ref bus) = self.event_bus {
740 ctx.set_event_bus(bus.clone());
741 }
742 if let Some(ref provider) = self.decision_provider {
743 ctx.set_decision_provider(provider.clone());
744 }
745 if let Some(ref interceptor) = self.step_interceptor {
746 ctx.set_step_interceptor(interceptor.clone());
747 }
748 if let Some(ref tags) = self.worker_tags {
749 ctx.set_worker_tags(tags.clone());
750 }
751 ctx
752 }
753
754 pub(crate) fn effective_guard_config(
757 &self,
758 handler: &dyn WorkflowHandler,
759 ) -> Option<WorkflowGuardConfig> {
760 handler.guard_config().or_else(|| self.guard_config.clone())
761 }
762
763 async fn build_context_with_guard(
782 &self,
783 run: &Run,
784 handler: &dyn WorkflowHandler,
785 ) -> Result<WorkflowContext, EngineError> {
786 let mut ctx = self.build_context(run);
787 if let Some(config) = self.effective_guard_config(handler) {
788 let mut state = WorkflowGuardState::new()
789 .with_tokens_used(self.tokens_used_by_run_tree(run).await?)
790 .with_elapsed(Duration::from_millis(run.duration_ms));
791 if let Some(at) = config.deadline_for(run.created_at) {
792 state = state.with_deadline_at(at);
793 }
794 ctx.set_guard(config, Arc::new(Mutex::new(state)));
795 }
796 Ok(ctx)
797 }
798
799 pub(crate) async fn list_all_runs(&self, filter: RunFilter) -> Result<Vec<Run>, EngineError> {
808 let mut runs = Vec::new();
809 let mut page = 1;
810 loop {
811 let result = self
812 .store
813 .list_runs(filter.clone(), page, LIST_RUNS_PAGE_SIZE)
814 .await?;
815 if result.items.is_empty() {
816 break;
817 }
818 runs.extend(result.items);
819 if runs.len() as u64 >= result.total {
820 break;
821 }
822 page += 1;
823 }
824 Ok(runs)
825 }
826
827 async fn tokens_used_by_run_tree(&self, root: &Run) -> Result<u64, EngineError> {
836 let tree_label = HashMap::from([(LABEL_ROOT_RUN_ID.to_string(), root.id.to_string())]);
837 let filter = RunFilter {
838 labels: Some(tree_label),
839 ..RunFilter::default()
840 };
841 let mut run_ids = vec![root.id];
842 run_ids.extend(
843 self.list_all_runs(filter)
844 .await?
845 .into_iter()
846 .map(|run| run.id),
847 );
848
849 let mut tokens: u64 = 0;
850 for run_id in run_ids {
851 for step in self.store.list_steps(run_id).await? {
852 if step.kind == StepKind::Agent {
853 let step_tokens = StepOutput::from(&step).total_tokens();
854 tokens = tokens.saturating_add(step_tokens);
855 }
856 }
857 }
858 Ok(tokens)
859 }
860
861 pub async fn check_monthly_quota(&self, workflow_name: &str) -> Result<(), EngineError> {
888 let Some(limit) = self.budget.monthly_cost_limit_usd else {
889 return Ok(());
890 };
891
892 let stats = self
893 .store
894 .get_stats(RunFilter {
895 created_after: Some(month_start(Utc::now())),
896 ..RunFilter::default()
897 })
898 .await?;
899
900 if stats.total_cost_usd < limit {
901 return Ok(());
902 }
903
904 warn!(
905 workflow = %workflow_name,
906 limit_usd = %limit,
907 spent_usd = %stats.total_cost_usd,
908 "monthly cost quota exhausted, refusing new run"
909 );
910
911 #[cfg(feature = "prometheus")]
912 counter!(
913 RUN_BUDGET_EXCEEDED_TOTAL,
914 "workflow" => workflow_name.to_string(),
915 "scope" => "monthly",
916 )
917 .increment(1);
918
919 Err(EngineError::MonthlyBudgetExceeded {
920 limit_usd: limit,
921 spent_usd: stats.total_cost_usd,
922 })
923 }
924
925 pub fn register(&mut self, handler: impl WorkflowHandler + 'static) -> Result<(), EngineError> {
969 let name = handler.name().to_string();
970 if self.handlers.contains_key(&name) {
971 return Err(EngineError::InvalidWorkflow(format!(
972 "handler '{}' already registered",
973 name
974 )));
975 }
976 if let Some(category) = handler.category() {
977 validate_category(&name, category)?;
978 }
979 self.handlers.insert(name, Arc::new(handler));
980 Ok(())
981 }
982
983 pub fn register_boxed(&mut self, handler: Box<dyn WorkflowHandler>) -> Result<(), EngineError> {
990 let name = handler.name().to_string();
991 if self.handlers.contains_key(&name) {
992 return Err(EngineError::InvalidWorkflow(format!(
993 "handler '{}' already registered",
994 name
995 )));
996 }
997 if let Some(category) = handler.category() {
998 validate_category(&name, category)?;
999 }
1000 self.handlers.insert(name, Arc::from(handler));
1001 Ok(())
1002 }
1003
1004 pub fn get_handler(&self, name: &str) -> Option<&Arc<dyn WorkflowHandler>> {
1006 self.handlers.get(name)
1007 }
1008
1009 pub fn handler_names(&self) -> Vec<&str> {
1011 self.handlers.keys().map(|s| s.as_str()).collect()
1012 }
1013
1014 pub fn handler_info(&self, name: &str) -> Option<WorkflowInfo> {
1016 self.handlers.get(name).map(|h| h.describe())
1017 }
1018
1019 pub fn scheduled_handlers(&self) -> Vec<(&str, &CronSchedule)> {
1043 self.handlers
1044 .iter()
1045 .filter_map(|(name, handler)| handler.schedule().map(|sched| (name.as_str(), sched)))
1046 .collect()
1047 }
1048
1049 pub fn subscribe(
1074 &mut self,
1075 subscriber: impl EventSubscriber + 'static,
1076 event_types: &[&'static str],
1077 ) {
1078 self.event_publisher.subscribe(subscriber, event_types);
1079 }
1080
1081 pub fn event_publisher(&self) -> &EventPublisher {
1086 &self.event_publisher
1087 }
1088
1089 #[tracing::instrument(name = "engine.run_handler", skip_all, fields(workflow = %handler_name))]
1125 pub async fn run_handler(
1126 &self,
1127 handler_name: &str,
1128 trigger: TriggerKind,
1129 payload: Value,
1130 ) -> Result<WorkflowResult, EngineError> {
1131 let handler = self
1132 .handlers
1133 .get(handler_name)
1134 .ok_or_else(|| {
1135 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1136 })?
1137 .clone();
1138
1139 self.check_monthly_quota(handler_name).await?;
1140
1141 let handler_version = handler.version().map(str::to_string);
1142 let max_cost_usd = self
1143 .budget
1144 .resolve_run_cap(None, handler.default_max_cost_usd());
1145 let run = self
1146 .store
1147 .create_run(NewRun {
1148 created_by: None,
1149 workflow_name: handler_name.to_string(),
1150 trigger,
1151 payload,
1152 max_retries: 0,
1153 handler_version,
1154 labels: handler.default_labels(),
1155 scheduled_at: None,
1156 idempotency_key: None,
1157 concurrency_key: None,
1158 priority: clamp_priority(handler.priority()),
1159 concurrency_limits: Vec::new(),
1160 max_cost_usd,
1161 worker_tags: normalize_worker_tags(handler.required_worker_tags()),
1162 })
1163 .await?
1164 .into_run();
1165
1166 let run_id = run.id;
1167 info!(run_id = %run_id, handler_version = run.handler_version.as_deref().unwrap_or(""), "run created");
1168
1169 if self.is_workflow_paused(handler_name).await? {
1170 info!(run_id = %run_id, workflow = %handler_name, "workflow paused, run left pending");
1171 return Ok(WorkflowResult {
1172 run,
1173 steps: Vec::new(),
1174 });
1175 }
1176
1177 let mut ctx = self
1181 .build_context_with_guard(&run, handler.as_ref())
1182 .await?;
1183
1184 self.store
1185 .update_run_status(run_id, RunStatus::Running)
1186 .await?;
1187
1188 #[cfg(feature = "prometheus")]
1189 gauge!(RUNS_ACTIVE, "workflow" => handler_name.to_string()).increment(1.0);
1190
1191 let run_start = Instant::now();
1192 let result = handler.execute(&mut ctx).await;
1193 self.finalize_run(run_id, handler_name, result, &ctx, run_start, run.labels)
1194 .await
1195 }
1196
1197 #[tracing::instrument(name = "engine.plan_handler", skip_all, fields(workflow = %handler_name))]
1246 pub async fn plan_handler(
1247 &self,
1248 handler_name: &str,
1249 payload: Value,
1250 options: PlanOptions,
1251 ) -> Result<ExecutionPlan, EngineError> {
1252 if options.max_depth == 0 {
1253 return Err(EngineError::InvalidWorkflow(
1254 "max_depth must be at least 1".to_string(),
1255 ));
1256 }
1257
1258 let handler = self
1259 .handlers
1260 .get(handler_name)
1261 .ok_or_else(|| {
1262 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1263 })?
1264 .clone();
1265
1266 let estimates = if options.estimate_durations {
1267 estimate_durations(&self.store, handler_name, options.sample_runs).await?
1268 } else {
1269 HashMap::new()
1270 };
1271
1272 let shared: SharedPlanRecorder = Arc::new(Mutex::new(PlanRecorder::new(
1273 handler_name.to_string(),
1274 payload,
1275 options.max_depth,
1276 estimates,
1277 )));
1278
1279 let handlers = self.handlers.clone();
1282 let resolver: crate::context::HandlerResolver =
1283 Arc::new(move |name: &str| handlers.get(name).cloned());
1284 let mut ctx = WorkflowContext::with_handler_resolver(
1285 Uuid::now_v7(),
1286 handler_name.to_string(),
1287 self.store.clone(),
1288 self.provider.clone(),
1289 resolver,
1290 );
1291 ctx.set_plan(shared.clone());
1292
1293 match timeout(options.timeout, handler.execute(&mut ctx)).await {
1294 Ok(Ok(())) => {}
1295 Ok(Err(err)) => lock_plan(&shared).fail(err.to_string()),
1296 Err(_) => lock_plan(&shared).fail(format!(
1297 "plan timed out after {} s",
1298 options.timeout.as_secs_f64()
1299 )),
1300 }
1301 drop(ctx);
1302
1303 let plan = match Arc::try_unwrap(shared) {
1304 Ok(mutex) => mutex
1305 .into_inner()
1306 .unwrap_or_else(|poisoned| poisoned.into_inner())
1307 .into_plan(),
1308 Err(shared) => lock_plan(&shared).snapshot(),
1309 };
1310
1311 info!(
1312 workflow = %handler_name,
1313 steps = plan.steps.len(),
1314 truncated = plan.truncated,
1315 "execution plan built"
1316 );
1317
1318 Ok(plan)
1319 }
1320
1321 #[tracing::instrument(name = "engine.enqueue_handler", skip_all, fields(workflow = %handler_name))]
1332 pub async fn enqueue_handler(
1333 &self,
1334 handler_name: &str,
1335 trigger: TriggerKind,
1336 payload: Value,
1337 max_retries: u32,
1338 ) -> Result<Run, EngineError> {
1339 self.enqueue_handler_with_options(
1340 handler_name,
1341 trigger,
1342 payload,
1343 EnqueueOptions {
1344 max_retries,
1345 ..Default::default()
1346 },
1347 )
1348 .await
1349 .map(RunCreation::into_run)
1350 }
1351
1352 #[tracing::instrument(name = "engine.enqueue_handler_with_options", skip_all, fields(workflow = %handler_name))]
1404 pub async fn enqueue_handler_with_options(
1405 &self,
1406 handler_name: &str,
1407 trigger: TriggerKind,
1408 payload: Value,
1409 options: EnqueueOptions,
1410 ) -> Result<RunCreation, EngineError> {
1411 let EnqueueOptions {
1412 max_retries,
1413 labels,
1414 scheduled_at,
1415 max_cost_usd,
1416 created_by,
1417 idempotency_key,
1418 concurrency_key,
1419 concurrency_limits,
1420 priority,
1421 worker_tags,
1422 } = options;
1423
1424 validate_concurrency_limits(&concurrency_limits)
1427 .map_err(EngineError::InvalidConcurrencyLimit)?;
1428 if let Some(priority) = priority {
1429 validate_priority(priority).map_err(EngineError::InvalidPriority)?;
1430 }
1431 validate_worker_tags(&worker_tags).map_err(EngineError::InvalidWorkerTag)?;
1432
1433 let handler = self.handlers.get(handler_name).ok_or_else(|| {
1434 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1435 })?;
1436
1437 self.check_monthly_quota(handler_name).await?;
1438
1439 let handler_version = handler.version().map(str::to_string);
1440 let mut merged_labels = handler.default_labels();
1441 merged_labels.extend(labels);
1442 let resolved_cap = self
1443 .budget
1444 .resolve_run_cap(max_cost_usd, handler.default_max_cost_usd());
1445 let priority = priority.unwrap_or_else(|| clamp_priority(handler.priority()));
1446 let required_tags = normalize_worker_tags(
1447 handler
1448 .required_worker_tags()
1449 .into_iter()
1450 .chain(worker_tags),
1451 );
1452
1453 let creation = self
1454 .store
1455 .create_run(NewRun {
1456 workflow_name: handler_name.to_string(),
1457 trigger,
1458 payload,
1459 max_retries,
1460 handler_version,
1461 labels: merged_labels,
1462 scheduled_at,
1463 created_by,
1464 idempotency_key,
1465 concurrency_key,
1466 priority,
1467 concurrency_limits,
1468 max_cost_usd: resolved_cap,
1469 worker_tags: required_tags,
1470 })
1471 .await?;
1472
1473 match &creation {
1474 RunCreation::Created(run) => info!(
1475 run_id = %run.id,
1476 workflow = %handler_name,
1477 max_cost_usd = ?resolved_cap,
1478 "handler run enqueued"
1479 ),
1480 RunCreation::Existing(run) => info!(
1481 run_id = %run.id,
1482 workflow = %handler_name,
1483 "idempotent replay, nothing enqueued"
1484 ),
1485 }
1486
1487 Ok(creation)
1488 }
1489
1490 #[tracing::instrument(name = "engine.execute_handler_run", skip_all, fields(run_id = %run_id))]
1510 pub async fn execute_handler_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1511 let run = self
1512 .store
1513 .get_run(run_id)
1514 .await?
1515 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1516
1517 if let Some(root_run_id) = chain_root(&run) {
1518 return self.resume_chain(run, root_run_id).await;
1519 }
1520
1521 let _active = self.track_execution(run_id).await;
1522
1523 let handler = self
1524 .handlers
1525 .get(&run.workflow_name)
1526 .ok_or_else(|| {
1527 EngineError::InvalidWorkflow(format!(
1528 "no handler registered: {}",
1529 run.workflow_name
1530 ))
1531 })?
1532 .clone();
1533
1534 #[cfg(feature = "prometheus")]
1535 gauge!(RUNS_ACTIVE, "workflow" => run.workflow_name.clone()).increment(1.0);
1536
1537 let run_start = Instant::now();
1538 let mut ctx = self
1539 .build_context_with_guard(&run, handler.as_ref())
1540 .await?;
1541
1542 let result = if handler.is_version_compatible(run.handler_version.as_deref()) {
1555 ctx.load_replay_steps().await?;
1556 self.release_then_execute(run_id, handler.as_ref(), &mut ctx)
1557 .await
1558 } else {
1559 Err(EngineError::HandlerVersionMismatch {
1560 run_id,
1561 workflow_name: run.workflow_name.clone(),
1562 run_version: run
1563 .handler_version
1564 .clone()
1565 .unwrap_or_else(|| "unknown".to_string()),
1566 current_version: handler
1567 .version()
1568 .map(str::to_string)
1569 .unwrap_or_else(|| "unknown".to_string()),
1570 })
1571 };
1572
1573 self.finalize_run(
1574 run_id,
1575 &run.workflow_name,
1576 result,
1577 &ctx,
1578 run_start,
1579 run.labels,
1580 )
1581 .await
1582 }
1583
1584 #[tracing::instrument(name = "engine.execute_run", skip_all, fields(run_id = %run_id))]
1592 pub async fn execute_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1593 self.execute_handler_run(run_id).await
1594 }
1595
1596 #[tracing::instrument(name = "engine.resume_run", skip_all, fields(run_id = %run_id))]
1625 pub async fn resume_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1626 let run = self
1627 .store
1628 .get_run(run_id)
1629 .await?
1630 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1631
1632 if let Some(root_run_id) = chain_root(&run) {
1633 return self.resume_chain(run, root_run_id).await;
1634 }
1635
1636 self.resume_loaded_run(run).await
1637 }
1638
1639 async fn resume_chain(
1656 &self,
1657 child: Run,
1658 root_run_id: Uuid,
1659 ) -> Result<WorkflowResult, EngineError> {
1660 let child_run_id = child.id;
1661 let lease = child.worker_id.zip(child.lease_expires_at);
1662 match self
1663 .claim_chain_root(child_run_id, root_run_id, lease.as_ref())
1664 .await?
1665 {
1666 ChainClaim::Claimed => {}
1667 ChainClaim::Settled(root) => {
1668 return Ok(WorkflowResult {
1669 run: *root,
1670 steps: Vec::new(),
1671 });
1672 }
1673 }
1674
1675 if lease.is_some() {
1676 self.store
1677 .update_run(
1678 child_run_id,
1679 RunUpdate {
1680 lease: Some(LeaseUpdate::Release),
1681 ..RunUpdate::default()
1682 },
1683 )
1684 .await?;
1685 }
1686
1687 info!(
1688 run_id = %child_run_id,
1689 root_run_id = %root_run_id,
1690 lease_transferred = lease.is_some(),
1691 "child run resumed through its root run"
1692 );
1693
1694 let root = self
1695 .store
1696 .get_run(root_run_id)
1697 .await?
1698 .ok_or(EngineError::Store(StoreError::RunNotFound(root_run_id)))?;
1699 self.resume_loaded_run(root).await
1700 }
1701
1702 pub(crate) async fn move_root_to_running(
1708 &self,
1709 root_run_id: Uuid,
1710 lease: Option<&(String, DateTime<Utc>)>,
1711 ) -> Result<(), EngineError> {
1712 match lease {
1713 Some((worker_id, expires_at)) => {
1714 self.store
1715 .update_run(
1716 root_run_id,
1717 RunUpdate {
1718 status: Some(RunStatus::Running),
1719 lease: Some(LeaseUpdate::Set {
1720 worker_id: worker_id.clone(),
1721 expires_at: *expires_at,
1722 }),
1723 ..RunUpdate::default()
1724 },
1725 )
1726 .await?;
1727 }
1728 None => {
1729 self.store
1730 .update_run_status(root_run_id, RunStatus::Running)
1731 .await?;
1732 }
1733 }
1734 Ok(())
1735 }
1736
1737 async fn resume_loaded_run(&self, run: Run) -> Result<WorkflowResult, EngineError> {
1739 let run_id = run.id;
1740 let _active = self.track_execution(run_id).await;
1741 let handler = self
1742 .handlers
1743 .get(&run.workflow_name)
1744 .ok_or_else(|| {
1745 EngineError::InvalidWorkflow(format!(
1746 "no handler registered: {}",
1747 run.workflow_name
1748 ))
1749 })?
1750 .clone();
1751
1752 info!(run_id = %run_id, workflow = %run.workflow_name, "resuming run after approval");
1753
1754 let run_start = Instant::now();
1755 let mut ctx = self
1756 .build_context_with_guard(&run, handler.as_ref())
1757 .await?;
1758
1759 let result = if handler.is_version_compatible(run.handler_version.as_deref()) {
1760 ctx.load_replay_steps().await?;
1761 self.release_then_execute(run_id, handler.as_ref(), &mut ctx)
1762 .await
1763 } else {
1764 Err(EngineError::HandlerVersionMismatch {
1765 run_id,
1766 workflow_name: run.workflow_name.clone(),
1767 run_version: run
1768 .handler_version
1769 .clone()
1770 .unwrap_or_else(|| "unknown".to_string()),
1771 current_version: handler
1772 .version()
1773 .map(str::to_string)
1774 .unwrap_or_else(|| "unknown".to_string()),
1775 })
1776 };
1777
1778 self.finalize_run(
1779 run_id,
1780 &run.workflow_name,
1781 result,
1782 &ctx,
1783 run_start,
1784 run.labels,
1785 )
1786 .await
1787 }
1788
1789 pub async fn deliver_signal(
1835 self: &Arc<Self>,
1836 signal: NewSignal,
1837 ) -> Result<SignalDelivery, EngineError> {
1838 if signal.name.trim().is_empty() {
1839 return Err(EngineError::InvalidSignal(
1840 "signal name must not be empty".to_string(),
1841 ));
1842 }
1843 if signal.key.trim().is_empty() {
1844 return Err(EngineError::InvalidSignal(
1845 "signal key must not be empty".to_string(),
1846 ));
1847 }
1848
1849 let stored = match self.store.insert_signal(signal).await? {
1850 SignalInsert::Created(stored) => stored,
1851 SignalInsert::Duplicate(existing) => {
1852 info!(
1853 signal_id = %existing.id,
1854 signal = %existing.name,
1855 key = %existing.key,
1856 "duplicate signal ignored"
1857 );
1858 return Ok(SignalDelivery {
1859 signal_id: existing.id,
1860 duplicate: true,
1861 resumed: Vec::new(),
1862 rejected: Vec::new(),
1863 });
1864 }
1865 };
1866
1867 let waiters = self
1868 .store
1869 .list_signal_waiters(&stored.name, &stored.key)
1870 .await?;
1871 let mut resumed = Vec::new();
1872 let mut rejected = Vec::new();
1873
1874 for step in waiters {
1875 if let Err(error) = validate_step_payload(step.input.as_ref(), &stored.payload) {
1876 rejected.push(SignalRejected {
1877 run_id: step.run_id,
1878 step_id: step.id,
1879 error,
1880 });
1881 continue;
1882 }
1883
1884 match self
1885 .store
1886 .resolve_signal_step(step.id, received_output(&stored))
1887 .await
1888 {
1889 Ok(SignalStepResolution::Resolved {
1890 run_id,
1891 run_resumed,
1892 }) => {
1893 resumed.push(SignalResumed {
1894 run_id,
1895 step_id: step.id,
1896 });
1897 if run_resumed && self.execution_mode == ExecutionMode::Local {
1898 self.spawn_local_resume(run_id);
1899 }
1900 }
1901 Ok(SignalStepResolution::NotWaiting { .. }) => {}
1903 Err(err) => {
1904 error!(
1905 run_id = %step.run_id,
1906 step_id = %step.id,
1907 error = %err,
1908 "failed to resolve a waiting signal step"
1909 );
1910 rejected.push(SignalRejected {
1911 run_id: step.run_id,
1912 step_id: step.id,
1913 error: err.to_string(),
1914 });
1915 }
1916 }
1917 }
1918
1919 info!(
1920 signal_id = %stored.id,
1921 signal = %stored.name,
1922 key = %stored.key,
1923 resumed = resumed.len(),
1924 rejected = rejected.len(),
1925 "signal received"
1926 );
1927 self.event_publisher
1928 .publish(Event::SignalReceived(SignalReceivedEvent {
1929 signal_id: stored.id,
1930 name: stored.name.clone(),
1931 key: stored.key.clone(),
1932 resumed_runs: resumed.iter().map(|r| r.run_id).collect(),
1933 at: stored.received_at,
1934 }));
1935
1936 Ok(SignalDelivery {
1937 signal_id: stored.id,
1938 duplicate: false,
1939 resumed,
1940 rejected,
1941 })
1942 }
1943
1944 pub async fn send_signal<S: Signal>(
1982 self: &Arc<Self>,
1983 signal: &S,
1984 key: &str,
1985 idempotency_id: Option<&str>,
1986 ) -> Result<SignalDelivery, EngineError> {
1987 let payload = to_value(signal)?;
1988 self.deliver_signal(NewSignal {
1989 name: S::NAME.to_string(),
1990 key: key.to_string(),
1991 payload,
1992 idempotency_id: idempotency_id.map(str::to_string),
1993 })
1994 .await
1995 }
1996
1997 pub(crate) fn spawn_local_resume(self: &Arc<Self>, run_id: Uuid) {
2022 let engine = Arc::clone(self);
2023 spawn(async move {
2024 let _resuming = lock_run(&engine.local_resumes, run_id).await;
2028 engine.wait_until_idle(run_id).await;
2029 let run = match engine.store.get_run(run_id).await {
2030 Ok(Some(run)) => run,
2031 Ok(None) => {
2032 error!(run_id = %run_id, "run to restart not found");
2033 return;
2034 }
2035 Err(err) => {
2036 error!(run_id = %run_id, error = %err, "failed to load a run to restart");
2037 return;
2038 }
2039 };
2040 match run.status.state {
2041 RunStatus::Pending => {
2042 match engine
2043 .store
2044 .update_run_status(run_id, RunStatus::Running)
2045 .await
2046 {
2047 Ok(()) => {}
2048 Err(StoreError::InvalidTransition { from, .. }) => {
2049 info!(
2050 run_id = %run_id,
2051 status = %from,
2052 "run claimed elsewhere, left to its worker"
2053 );
2054 return;
2055 }
2056 Err(err) => {
2057 error!(run_id = %run_id, error = %err, "failed to restart a woken run");
2058 return;
2059 }
2060 }
2061 }
2062 RunStatus::Running if run.worker_id.is_some() => {
2063 info!(
2064 run_id = %run_id,
2065 worker_id = ?run.worker_id,
2066 "run owned by a worker lease, local resume skipped"
2067 );
2068 return;
2069 }
2070 RunStatus::Running => {
2071 if engine.is_executing(run_id) {
2072 info!(run_id = %run_id, "run already executing in this process, resume skipped");
2073 return;
2074 }
2075 }
2076 status => {
2077 info!(
2078 run_id = %run_id,
2079 status = %status,
2080 "run no longer waiting to restart, resume skipped"
2081 );
2082 return;
2083 }
2084 }
2085 if let Err(err) = engine.resume_run(run_id).await {
2086 error!(run_id = %run_id, error = %err, "failed to resume a woken run");
2087 }
2088 });
2089 }
2090
2091 pub async fn fail_or_schedule_retry(
2142 &self,
2143 run_id: Uuid,
2144 error: &str,
2145 retryable: bool,
2146 cost_usd: Option<Decimal>,
2147 duration_ms: Option<u64>,
2148 ) -> Result<RunStatus, EngineError> {
2149 let run = self
2150 .store
2151 .get_run(run_id)
2152 .await?
2153 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2154
2155 if run.status.state == RunStatus::Paused {
2158 info!(run_id = %run_id, error = %error, "run paused, failure not recorded");
2159 return Ok(RunStatus::Paused);
2160 }
2161
2162 let has_attempts_left = run.retry_count < run.max_retries;
2163 let update = if retryable && has_attempts_left {
2164 let backoff = backoff_for_retry(run.retry_count);
2165 let scheduled_at = Utc::now() + TimeDelta::milliseconds(backoff.as_millis() as i64);
2166
2167 info!(
2168 run_id = %run_id,
2169 workflow = %run.workflow_name,
2170 attempt = run.retry_count + 1,
2171 max_retries = run.max_retries,
2172 backoff_secs = backoff.as_secs(),
2173 scheduled_at = %scheduled_at,
2174 "run failed, scheduling retry"
2175 );
2176
2177 RunUpdate {
2178 status: Some(RunStatus::Retrying),
2179 error: Some(error.to_string()),
2180 increment_retry: true,
2181 cost_usd,
2182 duration_ms,
2183 scheduled_at: Some(scheduled_at),
2184 ..RunUpdate::default()
2185 }
2186 } else {
2187 RunUpdate {
2188 status: Some(RunStatus::Failed),
2189 error: Some(error.to_string()),
2190 cost_usd,
2191 duration_ms,
2192 completed_at: Some(Utc::now()),
2193 ..RunUpdate::default()
2194 }
2195 };
2196
2197 let status = update.status.unwrap_or(RunStatus::Failed);
2198 self.store.update_run(run_id, update).await?;
2199 self.fail_orphaned_steps(run_id, error).await?;
2200 self.cancel_descendants_of_stopped_run(run_id, error).await;
2203
2204 Ok(status)
2205 }
2206
2207 pub async fn interrupt_running_steps(&self, run_id: Uuid) -> Result<(), EngineError> {
2237 interrupt_running_steps(self.store.as_ref(), run_id).await
2238 }
2239
2240 pub async fn fail_orphaned_steps(
2254 &self,
2255 run_id: Uuid,
2256 error_message: &str,
2257 ) -> Result<(), EngineError> {
2258 let steps = self.store.list_steps(run_id).await?;
2259 let now = Utc::now();
2260
2261 for step in steps {
2262 if step.status.state.is_terminal() {
2263 continue;
2264 }
2265
2266 let (target_status, error) = match step.status.state {
2267 StepStatus::Running | StepStatus::AwaitingApproval => {
2268 let err = if step.error.is_some() {
2269 None
2270 } else {
2271 Some(error_message.to_string())
2272 };
2273 (StepStatus::Failed, err)
2274 }
2275 StepStatus::Pending => (StepStatus::Skipped, None),
2276 _ => continue,
2277 };
2278
2279 if let Err(e) = self
2280 .store
2281 .update_step(
2282 step.id,
2283 StepUpdate {
2284 status: Some(target_status),
2285 error,
2286 completed_at: Some(now),
2287 ..StepUpdate::default()
2288 },
2289 )
2290 .await
2291 {
2292 warn!(
2293 run_id = %run_id,
2294 step_id = %step.id,
2295 step_name = %step.name,
2296 error = %e,
2297 "failed to cleanup orphaned step"
2298 );
2299 } else {
2300 info!(
2301 run_id = %run_id,
2302 step_id = %step.id,
2303 step_name = %step.name,
2304 from = %step.status.state,
2305 to = %target_status,
2306 "cleaned up orphaned step"
2307 );
2308 }
2309 }
2310
2311 Ok(())
2312 }
2313
2314 async fn release_then_execute(
2319 &self,
2320 run_id: Uuid,
2321 handler: &dyn WorkflowHandler,
2322 ctx: &mut WorkflowContext,
2323 ) -> Result<(), EngineError> {
2324 match self.provider.release_run(&run_id.to_string()).await {
2325 Ok(()) => handler.execute(ctx).await,
2326 Err(e) => Err(EngineError::Operation(OperationError::Agent(e))),
2327 }
2328 }
2329
2330 async fn finalize_run(
2336 &self,
2337 run_id: Uuid,
2338 workflow_name: &str,
2339 result: Result<(), EngineError>,
2340 ctx: &WorkflowContext,
2341 run_start: Instant,
2342 run_labels: HashMap<String, String>,
2343 ) -> Result<WorkflowResult, EngineError> {
2344 let total_duration = ctx.carried_duration_ms() + run_start.elapsed().as_millis() as u64;
2347 let completed_at = Utc::now();
2348
2349 if let Some(run) = self.store.get_run(run_id).await?
2353 && run.status.state == RunStatus::Paused
2354 {
2355 let run = self
2356 .store
2357 .update_run_returning(
2358 run_id,
2359 RunUpdate {
2360 cost_usd: Some(ctx.total_cost_usd()),
2361 duration_ms: Some(total_duration),
2362 ..RunUpdate::default()
2363 },
2364 )
2365 .await?;
2366 info!(
2367 run_id = %run_id,
2368 outcome = ?result.err().map(|err| err.to_string()),
2369 "run paused, execution stopped"
2370 );
2371 return Ok(WorkflowResult {
2372 run,
2373 steps: ctx.step_results().to_vec(),
2374 });
2375 }
2376
2377 if matches!(result, Err(EngineError::RunPaused { .. })) {
2380 let run = self
2381 .store
2382 .get_run(run_id)
2383 .await?
2384 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2385 info!(
2386 run_id = %run_id,
2387 status = %run.status.state,
2388 "run resumed after the pause, execution stopped"
2389 );
2390 return Ok(WorkflowResult {
2391 run,
2392 steps: ctx.step_results().to_vec(),
2393 });
2394 }
2395
2396 let final_status;
2397 let final_run;
2398
2399 match result {
2400 Ok(()) => {
2401 final_status = if ctx.has_allowed_failure() {
2402 RunStatus::Warning
2403 } else {
2404 RunStatus::Completed
2405 };
2406 final_run = self
2407 .store
2408 .update_run_returning(
2409 run_id,
2410 RunUpdate {
2411 status: Some(final_status),
2412 cost_usd: Some(ctx.total_cost_usd()),
2413 duration_ms: Some(total_duration),
2414 completed_at: Some(completed_at),
2415 output: ctx.output().cloned(),
2416 ..RunUpdate::default()
2417 },
2418 )
2419 .await?;
2420
2421 info!(
2422 run_id = %run_id,
2423 status = %final_status,
2424 cost_usd = %ctx.total_cost_usd(),
2425 duration_ms = total_duration,
2426 "run completed"
2427 );
2428 }
2429 Err(EngineError::ApprovalRequired {
2430 run_id: approval_run_id,
2431 step_id,
2432 ref message,
2433 }) => {
2434 final_status = RunStatus::AwaitingApproval;
2435 final_run = self
2436 .store
2437 .update_run_returning(
2438 run_id,
2439 RunUpdate {
2440 status: Some(RunStatus::AwaitingApproval),
2441 cost_usd: Some(ctx.total_cost_usd()),
2442 duration_ms: Some(total_duration),
2443 ..RunUpdate::default()
2444 },
2445 )
2446 .await?;
2447
2448 info!(
2449 run_id = %approval_run_id,
2450 step_id = %step_id,
2451 message = %message,
2452 "run awaiting approval"
2453 );
2454
2455 self.publish_approval_requested(approval_run_id, step_id, message)
2456 .await?;
2457 }
2458 Err(EngineError::ChildSuspended {
2459 run_id: child_run_id,
2460 ref cause,
2461 }) => {
2462 final_status = cause.suspension_status();
2463 final_run = self
2467 .store
2468 .update_run_returning(
2469 run_id,
2470 RunUpdate {
2471 status: Some(final_status),
2472 cost_usd: Some(ctx.total_cost_usd()),
2473 duration_ms: Some(total_duration),
2474 ..RunUpdate::default()
2475 },
2476 )
2477 .await?;
2478
2479 let leaf = cause.suspension_leaf();
2480 info!(
2481 run_id = %run_id,
2482 child_run_id = %child_run_id,
2483 status = %final_status,
2484 cause = %leaf,
2485 "run suspended with its child run"
2486 );
2487
2488 self.publish_suspension(leaf).await?;
2489 for sibling in ctx.sibling_suspensions() {
2491 self.publish_suspension(sibling.suspension_leaf()).await?;
2492 }
2493 }
2494 Err(EngineError::HumanInputRequired {
2495 run_id: input_run_id,
2496 step_id,
2497 ref message,
2498 }) => {
2499 final_status = RunStatus::AwaitingApproval;
2500 final_run = self
2501 .store
2502 .update_run_returning(
2503 run_id,
2504 RunUpdate {
2505 status: Some(RunStatus::AwaitingApproval),
2506 cost_usd: Some(ctx.total_cost_usd()),
2507 duration_ms: Some(total_duration),
2508 ..RunUpdate::default()
2509 },
2510 )
2511 .await?;
2512
2513 info!(
2515 run_id = %input_run_id,
2516 step_id = %step_id,
2517 message = %message,
2518 "run awaiting human input"
2519 );
2520 }
2521 Err(EngineError::DelaySleeping {
2522 run_id: delay_run_id,
2523 step_id,
2524 wake_at,
2525 }) => {
2526 final_status = RunStatus::Sleeping;
2527 final_run = self
2528 .store
2529 .update_run_returning(
2530 run_id,
2531 RunUpdate {
2532 status: Some(RunStatus::Sleeping),
2533 cost_usd: Some(ctx.total_cost_usd()),
2534 duration_ms: Some(total_duration),
2535 scheduled_at: Some(wake_at),
2536 ..RunUpdate::default()
2537 },
2538 )
2539 .await?;
2540
2541 info!(
2542 run_id = %delay_run_id,
2543 step_id = %step_id,
2544 wake_at = %wake_at,
2545 "run sleeping until delay elapses"
2546 );
2547 }
2548 Err(EngineError::CapacitySleeping {
2549 run_id: capacity_run_id,
2550 step_id,
2551 ref kind,
2552 wake_at,
2553 }) => {
2554 final_status = RunStatus::Sleeping;
2555 final_run = self
2556 .store
2557 .update_run_returning(
2558 run_id,
2559 RunUpdate {
2560 status: Some(RunStatus::Sleeping),
2561 cost_usd: Some(ctx.total_cost_usd()),
2562 duration_ms: Some(total_duration),
2563 scheduled_at: Some(wake_at),
2564 capacity_wait_kind: Some(ProviderKind::new(kind.as_str())),
2565 ..RunUpdate::default()
2566 },
2567 )
2568 .await?;
2569
2570 info!(
2571 run_id = %capacity_run_id,
2572 step_id = %step_id,
2573 kind = %kind,
2574 wake_at = %wake_at,
2575 "run sleeping until provider capacity returns"
2576 );
2577 }
2578 Err(EngineError::SignalWaiting {
2579 run_id: wait_run_id,
2580 step_id,
2581 ref step_name,
2582 ref name,
2583 ref key,
2584 deadline_at,
2585 }) => {
2586 final_status = RunStatus::Sleeping;
2587 let waiting = self
2591 .store
2592 .suspend_run_on_signal(run_id, step_id, deadline_at)
2593 .await?;
2594 final_run = self
2595 .store
2596 .update_run_returning(
2597 run_id,
2598 RunUpdate {
2599 cost_usd: Some(ctx.total_cost_usd()),
2600 duration_ms: Some(total_duration),
2601 ..RunUpdate::default()
2602 },
2603 )
2604 .await?;
2605
2606 if waiting {
2607 self.event_publisher
2608 .publish(Event::SignalAwaited(SignalAwaitedEvent {
2609 run_id: wait_run_id,
2610 step_id,
2611 step_name: step_name.clone(),
2612 name: name.clone(),
2613 key: key.clone(),
2614 deadline_at,
2615 at: Utc::now(),
2616 }));
2617 }
2618
2619 info!(
2620 run_id = %wait_run_id,
2621 step_id = %step_id,
2622 signal = %name,
2623 key = %key,
2624 deadline_at = %deadline_at,
2625 waiting,
2626 "run sleeping until a signal arrives"
2627 );
2628 }
2629 Err(err) => {
2630 let guardrail_stop = matches!(
2634 err,
2635 EngineError::RunBudgetExceeded { .. } | EngineError::WorkflowGuardRejected(_)
2636 );
2637
2638 final_status = if guardrail_stop {
2639 if let Err(store_err) = self
2640 .store
2641 .update_run(
2642 run_id,
2643 RunUpdate {
2644 status: Some(RunStatus::Cancelled),
2645 error: Some(err.to_string()),
2646 cost_usd: Some(ctx.total_cost_usd()),
2647 duration_ms: Some(total_duration),
2648 completed_at: Some(completed_at),
2649 output: ctx.output().cloned(),
2650 ..RunUpdate::default()
2651 },
2652 )
2653 .await
2654 {
2655 error!(run_id = %run_id, store_error = %store_err, "failed to persist run cancellation");
2656 }
2657 if let Err(cleanup_err) = self
2658 .fail_orphaned_steps(run_id, "run stopped: guardrail limit reached")
2659 .await
2660 {
2661 error!(run_id = %run_id, store_error = %cleanup_err, "failed to cleanup orphaned steps");
2662 }
2663 RunStatus::Cancelled
2664 } else {
2665 if let Some(output) = ctx.output()
2668 && let Err(store_err) = self
2669 .store
2670 .update_run(
2671 run_id,
2672 RunUpdate {
2673 output: Some(output.clone()),
2674 ..RunUpdate::default()
2675 },
2676 )
2677 .await
2678 {
2679 error!(run_id = %run_id, store_error = %store_err, "failed to persist run output");
2680 }
2681 self.fail_or_schedule_retry(
2682 run_id,
2683 &err.to_string(),
2684 is_run_retryable(&err),
2685 Some(ctx.total_cost_usd()),
2686 Some(total_duration),
2687 )
2688 .await
2689 .unwrap_or_else(|store_err| {
2690 error!(run_id = %run_id, store_error = %store_err, "failed to persist run failure");
2691 RunStatus::Failed
2692 })
2693 };
2694
2695 if matches!(err, EngineError::RunBudgetExceeded { .. }) {
2696 self.on_run_budget_exceeded(workflow_name, run_id, &err);
2697 }
2698
2699 error!(run_id = %run_id, status = %final_status, error = %err, "run stopped");
2700
2701 self.publish_run_status_changed(
2702 workflow_name,
2703 run_id,
2704 final_status,
2705 Some(err.to_string()),
2706 ctx,
2707 total_duration,
2708 run_labels,
2709 );
2710
2711 #[cfg(feature = "prometheus")]
2712 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
2713
2714 return Err(err);
2715 }
2716 }
2717
2718 self.publish_run_status_changed(
2719 workflow_name,
2720 run_id,
2721 final_status,
2722 None,
2723 ctx,
2724 total_duration,
2725 run_labels,
2726 );
2727
2728 #[cfg(feature = "prometheus")]
2729 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
2730
2731 Ok(WorkflowResult {
2732 run: final_run,
2733 steps: ctx.step_results().to_vec(),
2734 })
2735 }
2736
2737 pub(crate) async fn publish_approval_requested(
2740 &self,
2741 run_id: Uuid,
2742 step_id: Uuid,
2743 message: &str,
2744 ) -> Result<(), EngineError> {
2745 let requirement = self
2746 .store
2747 .get_step(step_id)
2748 .await?
2749 .and_then(|s| s.approval_requirement);
2750 self.event_publisher
2751 .publish(Event::ApprovalRequested(ApprovalRequestedEvent {
2752 run_id,
2753 step_id,
2754 message: message.to_string(),
2755 requirement,
2756 at: Utc::now(),
2757 }));
2758 Ok(())
2759 }
2760
2761 pub async fn fail_ancestors(&self, run_id: Uuid, reason: &str) -> Result<(), EngineError> {
2794 let mut current = self
2795 .store
2796 .get_run(run_id)
2797 .await?
2798 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2799 let mut visited = HashSet::from([run_id]);
2801
2802 while let Some(parent_id) = chain_parent(¤t) {
2803 if !visited.insert(parent_id) {
2804 break;
2805 }
2806 let status = self
2807 .fail_or_schedule_retry(parent_id, reason, false, None, None)
2808 .await?;
2809 if status == RunStatus::Paused {
2812 self.requeue_paused_root(¤t).await?;
2813 break;
2814 }
2815 info!(
2816 run_id = %run_id,
2817 ancestor_run_id = %parent_id,
2818 status = %status,
2819 "ancestor run failed with its child"
2820 );
2821 current = self
2822 .store
2823 .get_run(parent_id)
2824 .await?
2825 .ok_or(EngineError::Store(StoreError::RunNotFound(parent_id)))?;
2826 }
2827
2828 Ok(())
2829 }
2830
2831 #[cfg(feature = "prometheus")]
2833 fn emit_run_metrics(
2834 &self,
2835 workflow_name: &str,
2836 status: RunStatus,
2837 duration_ms: u64,
2838 ctx: &WorkflowContext,
2839 ) {
2840 let status_str = status.to_string();
2841 let wf = workflow_name.to_string();
2842
2843 counter!(RUNS_TOTAL, "workflow" => wf.clone(), "status" => status_str.clone()).increment(1);
2844 histogram!(RUN_DURATION_SECONDS, "workflow" => wf.clone(), "status" => status_str)
2845 .record(duration_ms as f64 / 1000.0);
2846 histogram!(RUN_COST_USD, "workflow" => wf.clone()).record(
2847 ctx.total_cost_usd()
2848 .to_string()
2849 .parse::<f64>()
2850 .unwrap_or(0.0),
2851 );
2852 gauge!(RUNS_ACTIVE, "workflow" => wf).decrement(1.0);
2853 }
2854
2855 fn on_run_budget_exceeded(&self, workflow_name: &str, run_id: Uuid, err: &EngineError) {
2861 let EngineError::RunBudgetExceeded {
2862 limit_usd,
2863 spent_usd,
2864 step_budget_usd,
2865 ..
2866 } = err
2867 else {
2868 return;
2869 };
2870
2871 #[cfg(feature = "prometheus")]
2872 counter!(
2873 RUN_BUDGET_EXCEEDED_TOTAL,
2874 "workflow" => workflow_name.to_string(),
2875 "scope" => "run",
2876 )
2877 .increment(1);
2878
2879 self.event_publisher
2880 .publish(Event::RunBudgetExceeded(RunBudgetExceededEvent {
2881 run_id,
2882 workflow_name: workflow_name.to_string(),
2883 limit_usd: *limit_usd,
2884 spent_usd: *spent_usd,
2885 step_budget_usd: *step_budget_usd,
2886 at: Utc::now(),
2887 }));
2888 }
2889
2890 #[allow(clippy::too_many_arguments)]
2895 fn publish_run_status_changed(
2896 &self,
2897 workflow_name: &str,
2898 run_id: Uuid,
2899 to: RunStatus,
2900 error: Option<String>,
2901 ctx: &WorkflowContext,
2902 duration_ms: u64,
2903 labels: HashMap<String, String>,
2904 ) {
2905 let now = Utc::now();
2906 let cost_usd = ctx.total_cost_usd();
2907 let wf = workflow_name.to_string();
2908
2909 self.event_publisher
2910 .publish(Event::RunStatusChanged(RunStatusChangedEvent {
2911 run_id,
2912 workflow_name: wf.clone(),
2913 from: RunStatus::Running,
2914 to,
2915 error: error.clone(),
2916 cost_usd,
2917 duration_ms,
2918 labels: labels.clone(),
2919 at: now,
2920 }));
2921
2922 if to == RunStatus::Failed {
2923 self.event_publisher
2924 .publish(Event::RunFailed(RunFailedEvent {
2925 run_id,
2926 workflow_name: wf,
2927 error,
2928 cost_usd,
2929 duration_ms,
2930 labels,
2931 at: now,
2932 }));
2933 }
2934 }
2935}
2936
2937impl fmt::Debug for Engine {
2938 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2939 f.debug_struct("Engine")
2940 .field("handlers", &self.handlers.keys().collect::<Vec<_>>())
2941 .finish_non_exhaustive()
2942 }
2943}
2944
2945#[cfg(test)]
2946mod tests {
2947 use super::*;
2948 use crate::config::ShellConfig;
2949 use crate::handler::{HandlerFuture, WorkflowHandler};
2950 use ironflow_core::providers::claude::ClaudeCodeProvider;
2951 use ironflow_core::providers::record_replay::RecordReplayProvider;
2952 use ironflow_store::memory::InMemoryStore;
2953 use ironflow_store::models::{MAX_PRIORITY, MIN_PRIORITY, StepApproval, StepStatus};
2954 use serde_json::json;
2955
2956 struct EchoWorkflow;
2958
2959 impl WorkflowHandler for EchoWorkflow {
2960 fn name(&self) -> &str {
2961 "echo-workflow"
2962 }
2963
2964 fn describe(&self) -> WorkflowInfo {
2965 WorkflowInfo {
2966 description: "A simple workflow that echoes hello".to_string(),
2967 source_code: None,
2968 sub_workflows: Vec::new(),
2969 category: None,
2970 version: self.version().map(str::to_string),
2971 compatible_versions: Vec::new(),
2972 input_schema: None,
2973 default_labels: HashMap::new(),
2974 schedule: self.schedule().cloned(),
2975 default_max_cost_usd: self.default_max_cost_usd(),
2976 priority: self.priority(),
2977 }
2978 }
2979
2980 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2981 Box::pin(async move {
2982 ctx.shell("greet", ShellConfig::new("echo hello")).await?;
2983 Ok(())
2984 })
2985 }
2986 }
2987
2988 struct FailingWorkflow;
2990
2991 impl WorkflowHandler for FailingWorkflow {
2992 fn name(&self) -> &str {
2993 "failing-workflow"
2994 }
2995
2996 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2997 Box::pin(async move {
2998 ctx.shell("fail", ShellConfig::new("exit 1")).await?;
2999 Ok(())
3000 })
3001 }
3002 }
3003
3004 fn create_test_engine() -> Engine {
3005 let store = Arc::new(InMemoryStore::new());
3006 let inner = ClaudeCodeProvider::new();
3007 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
3008 inner,
3009 "/tmp/ironflow-fixtures",
3010 ));
3011 Engine::new(store, provider)
3012 }
3013
3014 #[test]
3015 fn engine_new_creates_instance() {
3016 let engine = create_test_engine();
3017 assert_eq!(engine.handler_names().len(), 0);
3018 }
3019
3020 #[test]
3021 fn execution_mode_defaults_to_local() {
3022 let engine = create_test_engine();
3023 assert_eq!(engine.execution_mode(), ExecutionMode::Local);
3024 }
3025
3026 #[test]
3027 fn with_execution_mode_overrides_the_default() {
3028 let engine = create_test_engine().with_execution_mode(ExecutionMode::Workers);
3029 assert_eq!(engine.execution_mode(), ExecutionMode::Workers);
3030 }
3031
3032 #[test]
3033 fn engine_register_handler() {
3034 let mut engine = create_test_engine();
3035 let result = engine.register(EchoWorkflow);
3036 assert!(result.is_ok());
3037 assert_eq!(engine.handler_names().len(), 1);
3038 assert!(engine.handler_names().contains(&"echo-workflow"));
3039 }
3040
3041 #[test]
3042 fn engine_register_duplicate_returns_error() {
3043 let mut engine = create_test_engine();
3044 engine.register(EchoWorkflow).unwrap();
3045 let result = engine.register(EchoWorkflow);
3046 assert!(result.is_err());
3047 }
3048
3049 #[test]
3050 fn engine_get_handler_found() {
3051 let mut engine = create_test_engine();
3052 engine.register(EchoWorkflow).unwrap();
3053 let handler = engine.get_handler("echo-workflow");
3054 assert!(handler.is_some());
3055 }
3056
3057 #[test]
3058 fn engine_get_handler_not_found() {
3059 let engine = create_test_engine();
3060 let handler = engine.get_handler("nonexistent");
3061 assert!(handler.is_none());
3062 }
3063
3064 #[test]
3065 fn engine_handler_names_lists_all() {
3066 let mut engine = create_test_engine();
3067 engine.register(EchoWorkflow).unwrap();
3068 engine.register(FailingWorkflow).unwrap();
3069 let names = engine.handler_names();
3070 assert_eq!(names.len(), 2);
3071 assert!(names.contains(&"echo-workflow"));
3072 assert!(names.contains(&"failing-workflow"));
3073 }
3074
3075 #[test]
3076 fn engine_handler_info_returns_description() {
3077 let mut engine = create_test_engine();
3078 engine.register(EchoWorkflow).unwrap();
3079 let info = engine.handler_info("echo-workflow");
3080 assert!(info.is_some());
3081 let info = info.unwrap();
3082 assert_eq!(info.description, "A simple workflow that echoes hello");
3083 }
3084
3085 struct CategorizedWorkflow;
3086
3087 impl WorkflowHandler for CategorizedWorkflow {
3088 fn name(&self) -> &str {
3089 "categorized"
3090 }
3091 fn category(&self) -> Option<&str> {
3092 Some("data/etl")
3093 }
3094 fn execute<'a>(
3095 &'a self,
3096 _ctx: &'a mut WorkflowContext,
3097 ) -> crate::handler::HandlerFuture<'a> {
3098 Box::pin(async move { Ok(()) })
3099 }
3100 }
3101
3102 #[test]
3103 fn engine_default_describe_propagates_category() {
3104 let mut engine = create_test_engine();
3105 engine.register(CategorizedWorkflow).unwrap();
3106 let info = engine.handler_info("categorized").unwrap();
3107 assert_eq!(info.category.as_deref(), Some("data/etl"));
3108 }
3109
3110 #[test]
3111 fn engine_default_describe_without_category() {
3112 let mut engine = create_test_engine();
3113 engine.register(EchoWorkflow).unwrap();
3114 let info = engine.handler_info("echo-workflow").unwrap();
3115 assert!(info.category.is_none());
3116 }
3117
3118 struct ScheduledWorkflow {
3123 schedule: CronSchedule,
3124 }
3125
3126 impl ScheduledWorkflow {
3127 fn new() -> Self {
3128 Self {
3129 schedule: CronSchedule::new("0 0 * * * *").unwrap(),
3130 }
3131 }
3132 }
3133
3134 impl WorkflowHandler for ScheduledWorkflow {
3135 fn name(&self) -> &str {
3136 "scheduled"
3137 }
3138 fn schedule(&self) -> Option<&CronSchedule> {
3139 Some(&self.schedule)
3140 }
3141 fn execute<'a>(
3142 &'a self,
3143 _ctx: &'a mut WorkflowContext,
3144 ) -> crate::handler::HandlerFuture<'a> {
3145 Box::pin(async move { Ok(()) })
3146 }
3147 }
3148
3149 #[test]
3150 fn engine_default_describe_propagates_schedule() {
3151 let mut engine = create_test_engine();
3152 engine.register(ScheduledWorkflow::new()).unwrap();
3153 let info = engine.handler_info("scheduled").unwrap();
3154 assert_eq!(
3155 info.schedule.as_ref().map(|s| s.as_str()),
3156 Some("0 0 * * * *")
3157 );
3158 }
3159
3160 #[test]
3161 fn engine_default_describe_without_schedule() {
3162 let mut engine = create_test_engine();
3163 engine.register(EchoWorkflow).unwrap();
3164 let info = engine.handler_info("echo-workflow").unwrap();
3165 assert!(info.schedule.is_none());
3166 }
3167
3168 #[test]
3169 fn scheduled_handlers_returns_only_scheduled() {
3170 let mut engine = create_test_engine();
3171 engine.register(EchoWorkflow).unwrap();
3172 engine.register(ScheduledWorkflow::new()).unwrap();
3173 engine.register(FailingWorkflow).unwrap();
3174
3175 let scheduled = engine.scheduled_handlers();
3176 assert_eq!(scheduled.len(), 1);
3177 assert_eq!(scheduled[0].0, "scheduled");
3178 assert_eq!(scheduled[0].1.as_str(), "0 0 * * * *");
3179 }
3180
3181 #[test]
3182 fn scheduled_handlers_empty_when_none_scheduled() {
3183 let mut engine = create_test_engine();
3184 engine.register(EchoWorkflow).unwrap();
3185 engine.register(FailingWorkflow).unwrap();
3186
3187 let scheduled = engine.scheduled_handlers();
3188 assert!(scheduled.is_empty());
3189 }
3190
3191 struct BadCategoryWorkflow(&'static str);
3192
3193 impl WorkflowHandler for BadCategoryWorkflow {
3194 fn name(&self) -> &str {
3195 "bad-category"
3196 }
3197 fn category(&self) -> Option<&str> {
3198 Some(self.0)
3199 }
3200 fn execute<'a>(
3201 &'a self,
3202 _ctx: &'a mut WorkflowContext,
3203 ) -> crate::handler::HandlerFuture<'a> {
3204 Box::pin(async move { Ok(()) })
3205 }
3206 }
3207
3208 #[test]
3209 fn engine_register_rejects_empty_category() {
3210 let mut engine = create_test_engine();
3211 let err = engine.register(BadCategoryWorkflow("")).unwrap_err();
3212 match err {
3213 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty category")),
3214 other => panic!("expected InvalidWorkflow, got {other:?}"),
3215 }
3216 }
3217
3218 #[test]
3219 fn engine_register_rejects_leading_slash_category() {
3220 let mut engine = create_test_engine();
3221 let err = engine
3222 .register(BadCategoryWorkflow("/data/etl"))
3223 .unwrap_err();
3224 match err {
3225 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("leading '/'")),
3226 other => panic!("expected InvalidWorkflow, got {other:?}"),
3227 }
3228 }
3229
3230 #[test]
3231 fn engine_register_rejects_trailing_slash_category() {
3232 let mut engine = create_test_engine();
3233 let err = engine
3234 .register(BadCategoryWorkflow("data/etl/"))
3235 .unwrap_err();
3236 match err {
3237 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("trailing '/'")),
3238 other => panic!("expected InvalidWorkflow, got {other:?}"),
3239 }
3240 }
3241
3242 #[test]
3243 fn engine_register_rejects_double_slash_category() {
3244 let mut engine = create_test_engine();
3245 let err = engine
3246 .register(BadCategoryWorkflow("data//etl"))
3247 .unwrap_err();
3248 match err {
3249 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty segment")),
3250 other => panic!("expected InvalidWorkflow, got {other:?}"),
3251 }
3252 }
3253
3254 #[test]
3255 fn engine_register_rejects_whitespace_only_segment_category() {
3256 let mut engine = create_test_engine();
3257 let err = engine
3258 .register(BadCategoryWorkflow("data/ /etl"))
3259 .unwrap_err();
3260 match err {
3261 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("whitespace-only segment")),
3262 other => panic!("expected InvalidWorkflow, got {other:?}"),
3263 }
3264 }
3265
3266 #[test]
3267 fn engine_register_accepts_valid_nested_category() {
3268 let mut engine = create_test_engine();
3269 assert!(engine.register(CategorizedWorkflow).is_ok());
3270 }
3271
3272 #[tokio::test]
3273 async fn engine_unknown_workflow_returns_error() {
3274 let engine = create_test_engine();
3275 let result = engine
3276 .run_handler("unknown", TriggerKind::Manual, json!({}))
3277 .await;
3278 assert!(result.is_err());
3279 match result {
3280 Err(EngineError::InvalidWorkflow(msg)) => {
3281 assert!(msg.contains("no handler registered"));
3282 }
3283 _ => panic!("expected InvalidWorkflow error"),
3284 }
3285 }
3286
3287 #[tokio::test]
3288 async fn engine_enqueue_handler_creates_pending_run() {
3289 let mut engine = create_test_engine();
3290 engine.register(EchoWorkflow).unwrap();
3291
3292 let run = engine
3293 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
3294 .await
3295 .unwrap();
3296 assert_eq!(run.status.state, RunStatus::Pending);
3297 assert_eq!(run.workflow_name, "echo-workflow");
3298 }
3299
3300 #[tokio::test]
3301 async fn enqueue_handler_leaves_the_run_unattributed() {
3302 let mut engine = create_test_engine();
3303 engine.register(EchoWorkflow).unwrap();
3304
3305 let run = engine
3306 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
3307 .await
3308 .unwrap();
3309
3310 assert!(run.created_by.is_none());
3311 }
3312
3313 #[tokio::test]
3314 async fn enqueue_handler_with_options_records_the_author() {
3315 let mut engine = create_test_engine();
3316 engine.register(EchoWorkflow).unwrap();
3317 let actor = RunActor::User {
3318 user_id: Uuid::now_v7(),
3319 };
3320
3321 let run = engine
3322 .enqueue_handler_with_options(
3323 "echo-workflow",
3324 TriggerKind::Api,
3325 json!({}),
3326 EnqueueOptions {
3327 created_by: Some(actor.clone()),
3328 ..Default::default()
3329 },
3330 )
3331 .await
3332 .unwrap()
3333 .into_run();
3334
3335 assert_eq!(run.created_by, Some(actor));
3336 }
3337
3338 #[tokio::test]
3339 async fn enqueue_handler_with_options_accepts_no_author() {
3340 let mut engine = create_test_engine();
3341 engine.register(EchoWorkflow).unwrap();
3342
3343 let run = engine
3344 .enqueue_handler_with_options(
3345 "echo-workflow",
3346 TriggerKind::Cron {
3347 schedule: "0 * * * * *".to_string(),
3348 schedule_id: None,
3349 scheduled_for: None,
3350 },
3351 json!({}),
3352 EnqueueOptions::default(),
3353 )
3354 .await
3355 .unwrap()
3356 .into_run();
3357
3358 assert!(run.created_by.is_none());
3359 }
3360
3361 #[tokio::test]
3362 async fn enqueue_handler_with_options_stores_concurrency_limits() {
3363 let mut engine = create_test_engine();
3364 engine.register(EchoWorkflow).unwrap();
3365 let limits = vec![
3366 ConcurrencyLimit::new("repo:acme", 2),
3367 ConcurrencyLimit::new("tenant:42", 5),
3368 ];
3369
3370 let run = engine
3371 .enqueue_handler_with_options(
3372 "echo-workflow",
3373 TriggerKind::Api,
3374 json!({}),
3375 EnqueueOptions {
3376 concurrency_limits: limits.clone(),
3377 ..Default::default()
3378 },
3379 )
3380 .await
3381 .unwrap()
3382 .into_run();
3383
3384 assert_eq!(run.concurrency_limits, limits);
3385 }
3386
3387 #[tokio::test]
3388 async fn enqueue_rejects_invalid_concurrency_limits() {
3389 let mut engine = create_test_engine();
3390 engine.register(EchoWorkflow).unwrap();
3391
3392 let invalid = [
3393 vec![ConcurrencyLimit::new("repo:acme", 0)],
3394 vec![ConcurrencyLimit::new("", 1)],
3395 vec![
3396 ConcurrencyLimit::new("repo:acme", 1),
3397 ConcurrencyLimit::new("repo:acme", 2),
3398 ],
3399 ];
3400 for concurrency_limits in invalid {
3401 let err = engine
3402 .enqueue_handler_with_options(
3403 "echo-workflow",
3404 TriggerKind::Api,
3405 json!({}),
3406 EnqueueOptions {
3407 concurrency_limits,
3408 ..Default::default()
3409 },
3410 )
3411 .await
3412 .unwrap_err();
3413 assert!(
3414 matches!(err, EngineError::InvalidConcurrencyLimit(_)),
3415 "{err:?}"
3416 );
3417 }
3418
3419 let err = engine
3421 .enqueue_handler_with_options(
3422 "not-registered",
3423 TriggerKind::Api,
3424 json!({}),
3425 EnqueueOptions {
3426 concurrency_limits: vec![ConcurrencyLimit::new("repo:acme", 0)],
3427 ..Default::default()
3428 },
3429 )
3430 .await
3431 .unwrap_err();
3432 assert!(
3433 matches!(err, EngineError::InvalidConcurrencyLimit(_)),
3434 "{err:?}"
3435 );
3436
3437 let page = engine
3438 .store()
3439 .list_runs(RunFilter::default(), 1, 10)
3440 .await
3441 .unwrap();
3442 assert_eq!(page.total, 0, "no run may be created");
3443 }
3444
3445 struct UrgentWorkflow;
3446
3447 impl WorkflowHandler for UrgentWorkflow {
3448 fn name(&self) -> &str {
3449 "urgent-workflow"
3450 }
3451
3452 fn priority(&self) -> i16 {
3453 60
3454 }
3455
3456 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3457 Box::pin(async { Ok(()) })
3458 }
3459 }
3460
3461 #[tokio::test]
3462 async fn enqueue_priority_defaults_to_the_handler_priority() {
3463 let mut engine = create_test_engine();
3464 engine.register(EchoWorkflow).unwrap();
3465 engine.register(UrgentWorkflow).unwrap();
3466
3467 let echo = engine
3468 .enqueue_handler_with_options(
3469 "echo-workflow",
3470 TriggerKind::Api,
3471 json!({}),
3472 EnqueueOptions::default(),
3473 )
3474 .await
3475 .unwrap()
3476 .into_run();
3477 assert_eq!(echo.priority, 0);
3478
3479 let urgent = engine
3480 .enqueue_handler_with_options(
3481 "urgent-workflow",
3482 TriggerKind::Api,
3483 json!({}),
3484 EnqueueOptions::default(),
3485 )
3486 .await
3487 .unwrap()
3488 .into_run();
3489 assert_eq!(urgent.priority, 60);
3490 }
3491
3492 #[tokio::test]
3493 async fn enqueue_priority_explicit_value_overrides_the_handler() {
3494 let mut engine = create_test_engine();
3495 engine.register(UrgentWorkflow).unwrap();
3496
3497 let run = engine
3498 .enqueue_handler_with_options(
3499 "urgent-workflow",
3500 TriggerKind::Api,
3501 json!({}),
3502 EnqueueOptions {
3503 priority: Some(-20),
3504 ..Default::default()
3505 },
3506 )
3507 .await
3508 .unwrap()
3509 .into_run();
3510 assert_eq!(run.priority, -20);
3511
3512 let stored = engine.store().get_run(run.id).await.unwrap().unwrap();
3513 assert_eq!(stored.priority, -20);
3514 }
3515
3516 #[tokio::test]
3517 async fn enqueue_priority_out_of_range_is_rejected() {
3518 let mut engine = create_test_engine();
3519 engine.register(EchoWorkflow).unwrap();
3520
3521 for priority in [MAX_PRIORITY + 1, MIN_PRIORITY - 1] {
3522 let err = engine
3523 .enqueue_handler_with_options(
3524 "echo-workflow",
3525 TriggerKind::Api,
3526 json!({}),
3527 EnqueueOptions {
3528 priority: Some(priority),
3529 ..Default::default()
3530 },
3531 )
3532 .await
3533 .unwrap_err();
3534 assert!(matches!(err, EngineError::InvalidPriority(_)), "{err:?}");
3535 }
3536
3537 let err = engine
3539 .enqueue_handler_with_options(
3540 "not-registered",
3541 TriggerKind::Api,
3542 json!({}),
3543 EnqueueOptions {
3544 priority: Some(MAX_PRIORITY + 1),
3545 ..Default::default()
3546 },
3547 )
3548 .await
3549 .unwrap_err();
3550 assert!(matches!(err, EngineError::InvalidPriority(_)), "{err:?}");
3551
3552 let page = engine
3553 .store()
3554 .list_runs(RunFilter::default(), 1, 10)
3555 .await
3556 .unwrap();
3557 assert_eq!(page.total, 0, "no run may be created");
3558 }
3559
3560 #[tokio::test]
3561 async fn enqueue_priority_bounds_are_accepted() {
3562 let mut engine = create_test_engine();
3563 engine.register(EchoWorkflow).unwrap();
3564
3565 for priority in [MIN_PRIORITY, MAX_PRIORITY] {
3566 let run = engine
3567 .enqueue_handler_with_options(
3568 "echo-workflow",
3569 TriggerKind::Api,
3570 json!({}),
3571 EnqueueOptions {
3572 priority: Some(priority),
3573 ..Default::default()
3574 },
3575 )
3576 .await
3577 .unwrap()
3578 .into_run();
3579 assert_eq!(run.priority, priority);
3580 }
3581 }
3582
3583 struct GpuWorkflow;
3584
3585 impl WorkflowHandler for GpuWorkflow {
3586 fn name(&self) -> &str {
3587 "gpu-workflow"
3588 }
3589
3590 fn required_worker_tags(&self) -> Vec<String> {
3591 vec!["gpu".to_string()]
3592 }
3593
3594 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3595 Box::pin(async move { Ok(()) })
3596 }
3597 }
3598
3599 #[tokio::test]
3600 async fn enqueue_merges_handler_and_request_worker_tags() {
3601 let mut engine = create_test_engine();
3602 engine.register(GpuWorkflow).unwrap();
3603
3604 let run = engine
3605 .enqueue_handler_with_options(
3606 "gpu-workflow",
3607 TriggerKind::Api,
3608 json!({}),
3609 EnqueueOptions {
3610 worker_tags: vec!["region:eu".to_string(), "gpu".to_string()],
3611 ..Default::default()
3612 },
3613 )
3614 .await
3615 .unwrap()
3616 .into_run();
3617
3618 assert_eq!(
3619 run.worker_tags,
3620 vec!["gpu".to_string(), "region:eu".to_string()]
3621 );
3622 }
3623
3624 #[tokio::test]
3625 async fn enqueue_without_worker_tags_keeps_handler_tags() {
3626 let mut engine = create_test_engine();
3627 engine.register(GpuWorkflow).unwrap();
3628 engine.register(EchoWorkflow).unwrap();
3629
3630 let gpu = engine
3631 .enqueue_handler("gpu-workflow", TriggerKind::Api, json!({}), 0)
3632 .await
3633 .unwrap();
3634 assert_eq!(gpu.worker_tags, vec!["gpu".to_string()]);
3635
3636 let echo = engine
3637 .enqueue_handler("echo-workflow", TriggerKind::Api, json!({}), 0)
3638 .await
3639 .unwrap();
3640 assert!(echo.worker_tags.is_empty());
3641 }
3642
3643 #[tokio::test]
3644 async fn enqueue_rejects_invalid_worker_tags() {
3645 let mut engine = create_test_engine();
3646 engine.register(EchoWorkflow).unwrap();
3647
3648 for worker_tags in [
3649 vec!["bad,tag".to_string()],
3650 vec![" ".to_string()],
3651 vec!["x".repeat(65)],
3652 ] {
3653 let err = engine
3654 .enqueue_handler_with_options(
3655 "echo-workflow",
3656 TriggerKind::Api,
3657 json!({}),
3658 EnqueueOptions {
3659 worker_tags,
3660 ..Default::default()
3661 },
3662 )
3663 .await
3664 .unwrap_err();
3665 assert!(matches!(err, EngineError::InvalidWorkerTag(_)), "{err:?}");
3666 }
3667
3668 let err = engine
3670 .enqueue_handler_with_options(
3671 "not-registered",
3672 TriggerKind::Api,
3673 json!({}),
3674 EnqueueOptions {
3675 worker_tags: vec!["bad,tag".to_string()],
3676 ..Default::default()
3677 },
3678 )
3679 .await
3680 .unwrap_err();
3681 assert!(matches!(err, EngineError::InvalidWorkerTag(_)), "{err:?}");
3682 }
3683
3684 #[test]
3685 fn worker_tags_are_unset_by_default() {
3686 let engine = create_test_engine();
3687 assert!(engine.worker_tags().is_none());
3688 }
3689
3690 #[test]
3691 fn set_worker_tags_stores_the_tags() {
3692 let mut engine = create_test_engine();
3693 engine.set_worker_tags(vec!["gpu".to_string()]);
3694 assert_eq!(engine.worker_tags(), Some(&["gpu".to_string()][..]));
3695
3696 engine.set_worker_tags(Vec::new());
3697 assert_eq!(engine.worker_tags(), Some(&[][..]));
3698 }
3699
3700 #[tokio::test]
3701 async fn run_handler_records_handler_worker_tags() {
3702 let mut engine = create_test_engine();
3703 engine.register(GpuWorkflow).unwrap();
3704
3705 let result = engine
3706 .run_handler("gpu-workflow", TriggerKind::Manual, json!({}))
3707 .await
3708 .unwrap();
3709 assert_eq!(result.run.worker_tags, vec!["gpu".to_string()]);
3710 }
3711
3712 #[tokio::test]
3713 async fn run_handler_leaves_the_run_unattributed() {
3714 let mut engine = create_test_engine();
3715 engine.register(EchoWorkflow).unwrap();
3716
3717 let run = engine
3718 .run_handler("echo-workflow", TriggerKind::Manual, json!({}))
3719 .await
3720 .unwrap()
3721 .run;
3722
3723 assert!(run.created_by.is_none());
3724 }
3725
3726 #[tokio::test]
3727 async fn run_handler_priority_comes_from_the_handler() {
3728 let mut engine = create_test_engine();
3729 engine.register(UrgentWorkflow).unwrap();
3730
3731 let run = engine
3732 .run_handler("urgent-workflow", TriggerKind::Manual, json!({}))
3733 .await
3734 .unwrap()
3735 .run;
3736
3737 assert_eq!(run.priority, 60);
3738 }
3739
3740 #[tokio::test]
3741 async fn engine_register_boxed() {
3742 let mut engine = create_test_engine();
3743 let handler: Box<dyn WorkflowHandler> = Box::new(EchoWorkflow);
3744 let result = engine.register_boxed(handler);
3745 assert!(result.is_ok());
3746 assert_eq!(engine.handler_names().len(), 1);
3747 }
3748
3749 #[tokio::test]
3750 async fn engine_store_and_provider_accessors() {
3751 let store = Arc::new(InMemoryStore::new());
3752 let inner = ClaudeCodeProvider::new();
3753 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
3754 inner,
3755 "/tmp/ironflow-fixtures",
3756 ));
3757 let engine = Engine::new(store.clone(), provider.clone());
3758
3759 let _ = engine.store();
3761 let _ = engine.provider();
3762 }
3763
3764 use crate::operation::{Operation, OperationContext};
3769 use async_trait::async_trait;
3770 use ironflow_core::error::OperationError;
3771 use ironflow_store::models::StepKind;
3772
3773 struct FakeGitlabOp {
3774 project_id: u64,
3775 title: String,
3776 }
3777
3778 #[async_trait]
3779 impl Operation for FakeGitlabOp {
3780 fn kind(&self) -> &str {
3781 "gitlab"
3782 }
3783
3784 async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
3785 Ok(json!({
3786 "issue_id": 42,
3787 "project_id": self.project_id,
3788 "title": self.title,
3789 }))
3790 }
3791
3792 fn input(&self) -> Option<Value> {
3793 Some(json!({
3794 "project_id": self.project_id,
3795 "title": self.title,
3796 }))
3797 }
3798 }
3799
3800 struct FailingOp;
3801
3802 #[async_trait]
3803 impl Operation for FailingOp {
3804 fn kind(&self) -> &str {
3805 "broken-service"
3806 }
3807
3808 async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
3809 Err(OperationError::Http {
3810 status: None,
3811 message: "service unavailable".to_string(),
3812 })
3813 }
3814 }
3815
3816 struct OperationWorkflow;
3817
3818 impl WorkflowHandler for OperationWorkflow {
3819 fn name(&self) -> &str {
3820 "operation-workflow"
3821 }
3822
3823 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3824 Box::pin(async move {
3825 let op = FakeGitlabOp {
3826 project_id: 123,
3827 title: "Bug report".to_string(),
3828 };
3829 ctx.operation("create-issue", &op).await?;
3830 Ok(())
3831 })
3832 }
3833 }
3834
3835 struct FailingOperationWorkflow;
3836
3837 impl WorkflowHandler for FailingOperationWorkflow {
3838 fn name(&self) -> &str {
3839 "failing-operation-workflow"
3840 }
3841
3842 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3843 Box::pin(async move {
3844 ctx.operation("broken-call", &FailingOp).await?;
3845 Ok(())
3846 })
3847 }
3848 }
3849
3850 struct MixedWorkflow;
3851
3852 impl WorkflowHandler for MixedWorkflow {
3853 fn name(&self) -> &str {
3854 "mixed-workflow"
3855 }
3856
3857 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3858 Box::pin(async move {
3859 ctx.shell("build", ShellConfig::new("echo built")).await?;
3860 let op = FakeGitlabOp {
3861 project_id: 456,
3862 title: "Deploy done".to_string(),
3863 };
3864 let result = ctx.operation("notify-gitlab", &op).await?;
3865 assert_eq!(result.output["issue_id"], 42);
3866 Ok(())
3867 })
3868 }
3869 }
3870
3871 #[tokio::test]
3872 async fn operation_step_happy_path() {
3873 let mut engine = create_test_engine();
3874 engine.register(OperationWorkflow).unwrap();
3875
3876 let run = engine
3877 .run_handler("operation-workflow", TriggerKind::Manual, json!({}))
3878 .await
3879 .unwrap()
3880 .run;
3881
3882 assert_eq!(run.status.state, RunStatus::Completed);
3883
3884 let steps = engine.store().list_steps(run.id).await.unwrap();
3885
3886 assert_eq!(steps.len(), 1);
3887 assert_eq!(steps[0].name, "create-issue");
3888 assert_eq!(steps[0].kind, StepKind::Custom("gitlab".to_string()));
3889 assert_eq!(
3890 steps[0].status.state,
3891 ironflow_store::models::StepStatus::Completed
3892 );
3893
3894 let output = steps[0].output.as_ref().unwrap();
3895 assert_eq!(output["issue_id"], 42);
3896 assert_eq!(output["project_id"], 123);
3897
3898 let input = steps[0].input.as_ref().unwrap();
3899 assert_eq!(input["project_id"], 123);
3900 assert_eq!(input["title"], "Bug report");
3901 }
3902
3903 #[tokio::test]
3904 async fn operation_step_failure_marks_run_failed() {
3905 let mut engine = create_test_engine();
3906 engine.register(FailingOperationWorkflow).unwrap();
3907
3908 let result = engine
3909 .run_handler("failing-operation-workflow", TriggerKind::Manual, json!({}))
3910 .await;
3911
3912 assert!(result.is_err());
3913 }
3914
3915 #[tokio::test]
3916 async fn operation_mixed_with_shell_steps() {
3917 let mut engine = create_test_engine();
3918 engine.register(MixedWorkflow).unwrap();
3919
3920 let run = engine
3921 .run_handler("mixed-workflow", TriggerKind::Manual, json!({}))
3922 .await
3923 .unwrap()
3924 .run;
3925
3926 assert_eq!(run.status.state, RunStatus::Completed);
3927
3928 let steps = engine.store().list_steps(run.id).await.unwrap();
3929
3930 assert_eq!(steps.len(), 2);
3931 assert_eq!(steps[0].kind, StepKind::Shell);
3932 assert_eq!(steps[1].kind, StepKind::Custom("gitlab".to_string()));
3933 assert_eq!(steps[0].position, 0);
3934 assert_eq!(steps[1].position, 1);
3935 }
3936
3937 use crate::config::ApprovalConfig;
3942
3943 struct SingleApprovalWorkflow;
3944
3945 impl WorkflowHandler for SingleApprovalWorkflow {
3946 fn name(&self) -> &str {
3947 "single-approval"
3948 }
3949
3950 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3951 Box::pin(async move {
3952 ctx.shell("build", ShellConfig::new("echo built")).await?;
3953 ctx.approval("gate", ApprovalConfig::new("OK?")).await?;
3954 ctx.shell("deploy", ShellConfig::new("echo deployed"))
3955 .await?;
3956 Ok(())
3957 })
3958 }
3959 }
3960
3961 struct DoubleApprovalWorkflow;
3962
3963 impl WorkflowHandler for DoubleApprovalWorkflow {
3964 fn name(&self) -> &str {
3965 "double-approval"
3966 }
3967
3968 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3969 Box::pin(async move {
3970 ctx.shell("build", ShellConfig::new("echo built")).await?;
3971 ctx.approval("staging-gate", ApprovalConfig::new("Deploy staging?"))
3972 .await?;
3973 ctx.shell("deploy-staging", ShellConfig::new("echo staging"))
3974 .await?;
3975 ctx.approval("prod-gate", ApprovalConfig::new("Deploy prod?"))
3976 .await?;
3977 ctx.shell("deploy-prod", ShellConfig::new("echo prod"))
3978 .await?;
3979 Ok(())
3980 })
3981 }
3982 }
3983
3984 #[tokio::test]
3985 async fn approval_pauses_run() {
3986 let mut engine = create_test_engine();
3987 engine.register(SingleApprovalWorkflow).unwrap();
3988
3989 let run = engine
3990 .run_handler("single-approval", TriggerKind::Manual, json!({}))
3991 .await
3992 .unwrap()
3993 .run;
3994
3995 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
3996
3997 let steps = engine.store().list_steps(run.id).await.unwrap();
3998 assert_eq!(steps.len(), 2); assert_eq!(steps[0].kind, StepKind::Shell);
4000 assert_eq!(steps[0].status.state, StepStatus::Completed);
4001 assert_eq!(steps[1].kind, StepKind::Approval);
4002 assert_eq!(steps[1].status.state, StepStatus::AwaitingApproval);
4003 }
4004
4005 async fn approve_open_gate(store: &Arc<dyn Store>, run_id: Uuid) {
4008 let gate = store
4009 .list_steps(run_id)
4010 .await
4011 .unwrap()
4012 .into_iter()
4013 .find(|step| step.status.state == StepStatus::AwaitingApproval)
4014 .expect("an open approval gate");
4015 store
4016 .record_step_approval(
4017 gate.id,
4018 StepApproval {
4019 user_id: Uuid::now_v7(),
4020 approved_by: "tester".to_string(),
4021 at: Utc::now(),
4022 },
4023 )
4024 .await
4025 .unwrap();
4026 store
4027 .update_run_status(run_id, RunStatus::Running)
4028 .await
4029 .unwrap();
4030 }
4031
4032 #[tokio::test]
4033 async fn approval_resume_completes_run() {
4034 let mut engine = create_test_engine();
4035 engine.register(SingleApprovalWorkflow).unwrap();
4036
4037 let run = engine
4039 .run_handler("single-approval", TriggerKind::Manual, json!({}))
4040 .await
4041 .unwrap()
4042 .run;
4043 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
4044
4045 approve_open_gate(engine.store(), run.id).await;
4047
4048 let resumed = engine.resume_run(run.id).await.unwrap().run;
4050 assert_eq!(resumed.status.state, RunStatus::Completed);
4051
4052 let steps = engine.store().list_steps(run.id).await.unwrap();
4053 assert_eq!(steps.len(), 3); assert_eq!(steps[0].name, "build");
4055 assert_eq!(steps[0].status.state, StepStatus::Completed);
4056 assert_eq!(steps[1].name, "gate");
4057 assert_eq!(steps[1].kind, StepKind::Approval);
4058 assert_eq!(steps[1].status.state, StepStatus::Completed);
4059 assert_eq!(steps[2].name, "deploy");
4060 assert_eq!(steps[2].status.state, StepStatus::Completed);
4061 }
4062
4063 #[tokio::test]
4064 async fn double_approval_two_resumes() {
4065 let mut engine = create_test_engine();
4066 engine.register(DoubleApprovalWorkflow).unwrap();
4067
4068 let run = engine
4070 .run_handler("double-approval", TriggerKind::Manual, json!({}))
4071 .await
4072 .unwrap()
4073 .run;
4074 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
4075
4076 let steps = engine.store().list_steps(run.id).await.unwrap();
4077 assert_eq!(steps.len(), 2); approve_open_gate(engine.store(), run.id).await;
4081
4082 let resumed = engine.resume_run(run.id).await.unwrap().run;
4083 assert_eq!(resumed.status.state, RunStatus::AwaitingApproval);
4084
4085 let steps = engine.store().list_steps(run.id).await.unwrap();
4086 assert_eq!(steps.len(), 4); approve_open_gate(engine.store(), run.id).await;
4090
4091 let final_run = engine.resume_run(run.id).await.unwrap().run;
4092 assert_eq!(final_run.status.state, RunStatus::Completed);
4093
4094 let steps = engine.store().list_steps(run.id).await.unwrap();
4095 assert_eq!(steps.len(), 5);
4096 assert_eq!(steps[0].name, "build");
4097 assert_eq!(steps[1].name, "staging-gate");
4098 assert_eq!(steps[2].name, "deploy-staging");
4099 assert_eq!(steps[3].name, "prod-gate");
4100 assert_eq!(steps[4].name, "deploy-prod");
4101
4102 for step in &steps {
4103 assert_eq!(step.status.state, StepStatus::Completed);
4104 }
4105 }
4106
4107 use ironflow_store::models::{NewStep, StepUpdate, step_trace_id};
4112
4113 async fn create_step_with_status(
4114 store: &Arc<dyn Store>,
4115 run_id: Uuid,
4116 name: &str,
4117 position: u32,
4118 status: StepStatus,
4119 ) -> ironflow_store::models::Step {
4120 let step = store
4121 .create_step(NewStep {
4122 run_id,
4123 trace_id: step_trace_id(run_id, name, position),
4124 name: name.to_string(),
4125 kind: StepKind::Shell,
4126 position,
4127 input: None,
4128 is_error_handler: false,
4129 })
4130 .await
4131 .unwrap();
4132
4133 match status {
4134 StepStatus::Pending => {}
4135 StepStatus::Running => {
4136 store
4137 .update_step(
4138 step.id,
4139 StepUpdate {
4140 status: Some(StepStatus::Running),
4141 ..StepUpdate::default()
4142 },
4143 )
4144 .await
4145 .unwrap();
4146 }
4147 StepStatus::Completed => {
4148 store
4149 .update_step(
4150 step.id,
4151 StepUpdate {
4152 status: Some(StepStatus::Running),
4153 ..StepUpdate::default()
4154 },
4155 )
4156 .await
4157 .unwrap();
4158 store
4159 .update_step(
4160 step.id,
4161 StepUpdate {
4162 status: Some(StepStatus::Completed),
4163 ..StepUpdate::default()
4164 },
4165 )
4166 .await
4167 .unwrap();
4168 }
4169 StepStatus::AwaitingApproval => {
4170 store
4171 .update_step(
4172 step.id,
4173 StepUpdate {
4174 status: Some(StepStatus::Running),
4175 ..StepUpdate::default()
4176 },
4177 )
4178 .await
4179 .unwrap();
4180 store
4181 .update_step(
4182 step.id,
4183 StepUpdate {
4184 status: Some(StepStatus::AwaitingApproval),
4185 ..StepUpdate::default()
4186 },
4187 )
4188 .await
4189 .unwrap();
4190 }
4191 _ => panic!("unsupported status for test helper: {status}"),
4192 }
4193
4194 store.get_step(step.id).await.unwrap().unwrap()
4195 }
4196
4197 #[tokio::test]
4198 async fn fail_orphaned_steps_marks_running_as_failed() {
4199 let engine = create_test_engine();
4200 let run = engine
4201 .store()
4202 .create_run(NewRun {
4203 created_by: None,
4204 workflow_name: "test".to_string(),
4205 trigger: TriggerKind::Manual,
4206 payload: json!({}),
4207 max_retries: 0,
4208 handler_version: None,
4209 labels: HashMap::new(),
4210 scheduled_at: None,
4211 idempotency_key: None,
4212 concurrency_key: None,
4213 priority: 0,
4214 concurrency_limits: Vec::new(),
4215 max_cost_usd: None,
4216 worker_tags: Vec::new(),
4217 })
4218 .await
4219 .unwrap()
4220 .into_run();
4221
4222 let step = create_step_with_status(
4223 engine.store(),
4224 run.id,
4225 "running-step",
4226 0,
4227 StepStatus::Running,
4228 )
4229 .await;
4230
4231 engine
4232 .fail_orphaned_steps(run.id, "parent run timed out")
4233 .await
4234 .unwrap();
4235
4236 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4237 assert_eq!(updated.status.state, StepStatus::Failed);
4238 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
4239 assert!(updated.completed_at.is_some());
4240 }
4241
4242 #[tokio::test]
4243 async fn fail_orphaned_steps_marks_pending_as_skipped() {
4244 let engine = create_test_engine();
4245 let run = engine
4246 .store()
4247 .create_run(NewRun {
4248 created_by: None,
4249 workflow_name: "test".to_string(),
4250 trigger: TriggerKind::Manual,
4251 payload: json!({}),
4252 max_retries: 0,
4253 handler_version: None,
4254 labels: HashMap::new(),
4255 scheduled_at: None,
4256 idempotency_key: None,
4257 concurrency_key: None,
4258 priority: 0,
4259 concurrency_limits: Vec::new(),
4260 max_cost_usd: None,
4261 worker_tags: Vec::new(),
4262 })
4263 .await
4264 .unwrap()
4265 .into_run();
4266
4267 let step = create_step_with_status(
4268 engine.store(),
4269 run.id,
4270 "pending-step",
4271 0,
4272 StepStatus::Pending,
4273 )
4274 .await;
4275
4276 engine
4277 .fail_orphaned_steps(run.id, "parent run timed out")
4278 .await
4279 .unwrap();
4280
4281 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4282 assert_eq!(updated.status.state, StepStatus::Skipped);
4283 assert!(updated.error.is_none());
4284 assert!(updated.completed_at.is_some());
4285 }
4286
4287 #[tokio::test]
4288 async fn fail_orphaned_steps_marks_awaiting_approval_as_failed() {
4289 let engine = create_test_engine();
4290 let run = engine
4291 .store()
4292 .create_run(NewRun {
4293 created_by: None,
4294 workflow_name: "test".to_string(),
4295 trigger: TriggerKind::Manual,
4296 payload: json!({}),
4297 max_retries: 0,
4298 handler_version: None,
4299 labels: HashMap::new(),
4300 scheduled_at: None,
4301 idempotency_key: None,
4302 concurrency_key: None,
4303 priority: 0,
4304 concurrency_limits: Vec::new(),
4305 max_cost_usd: None,
4306 worker_tags: Vec::new(),
4307 })
4308 .await
4309 .unwrap()
4310 .into_run();
4311
4312 let step = create_step_with_status(
4313 engine.store(),
4314 run.id,
4315 "approval-step",
4316 0,
4317 StepStatus::AwaitingApproval,
4318 )
4319 .await;
4320
4321 engine
4322 .fail_orphaned_steps(run.id, "parent run timed out")
4323 .await
4324 .unwrap();
4325
4326 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4327 assert_eq!(updated.status.state, StepStatus::Failed);
4328 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
4329 assert!(updated.completed_at.is_some());
4330 }
4331
4332 #[tokio::test]
4333 async fn fail_orphaned_steps_skips_terminal_steps() {
4334 let engine = create_test_engine();
4335 let run = engine
4336 .store()
4337 .create_run(NewRun {
4338 created_by: None,
4339 workflow_name: "test".to_string(),
4340 trigger: TriggerKind::Manual,
4341 payload: json!({}),
4342 max_retries: 0,
4343 handler_version: None,
4344 labels: HashMap::new(),
4345 scheduled_at: None,
4346 idempotency_key: None,
4347 concurrency_key: None,
4348 priority: 0,
4349 concurrency_limits: Vec::new(),
4350 max_cost_usd: None,
4351 worker_tags: Vec::new(),
4352 })
4353 .await
4354 .unwrap()
4355 .into_run();
4356
4357 let completed_step =
4358 create_step_with_status(engine.store(), run.id, "done", 0, StepStatus::Completed).await;
4359 let running_step =
4360 create_step_with_status(engine.store(), run.id, "in-flight", 1, StepStatus::Running)
4361 .await;
4362
4363 engine
4364 .fail_orphaned_steps(run.id, "parent run timed out")
4365 .await
4366 .unwrap();
4367
4368 let completed = engine
4369 .store()
4370 .get_step(completed_step.id)
4371 .await
4372 .unwrap()
4373 .unwrap();
4374 assert_eq!(completed.status.state, StepStatus::Completed);
4375
4376 let failed = engine
4377 .store()
4378 .get_step(running_step.id)
4379 .await
4380 .unwrap()
4381 .unwrap();
4382 assert_eq!(failed.status.state, StepStatus::Failed);
4383 }
4384
4385 #[tokio::test]
4386 async fn fail_orphaned_steps_mixed_states() {
4387 let engine = create_test_engine();
4388 let run = engine
4389 .store()
4390 .create_run(NewRun {
4391 created_by: None,
4392 workflow_name: "test".to_string(),
4393 trigger: TriggerKind::Manual,
4394 payload: json!({}),
4395 max_retries: 0,
4396 handler_version: None,
4397 labels: HashMap::new(),
4398 scheduled_at: None,
4399 idempotency_key: None,
4400 concurrency_key: None,
4401 priority: 0,
4402 concurrency_limits: Vec::new(),
4403 max_cost_usd: None,
4404 worker_tags: Vec::new(),
4405 })
4406 .await
4407 .unwrap()
4408 .into_run();
4409
4410 let s_completed =
4411 create_step_with_status(engine.store(), run.id, "step-1", 0, StepStatus::Completed)
4412 .await;
4413 let s_running =
4414 create_step_with_status(engine.store(), run.id, "step-2", 1, StepStatus::Running).await;
4415 let s_pending =
4416 create_step_with_status(engine.store(), run.id, "step-3", 2, StepStatus::Pending).await;
4417
4418 engine.fail_orphaned_steps(run.id, "timeout").await.unwrap();
4419
4420 let r_completed = engine
4421 .store()
4422 .get_step(s_completed.id)
4423 .await
4424 .unwrap()
4425 .unwrap();
4426 assert_eq!(r_completed.status.state, StepStatus::Completed);
4427
4428 let r_running = engine
4429 .store()
4430 .get_step(s_running.id)
4431 .await
4432 .unwrap()
4433 .unwrap();
4434 assert_eq!(r_running.status.state, StepStatus::Failed);
4435 assert_eq!(r_running.error.as_deref(), Some("timeout"));
4436
4437 let r_pending = engine
4438 .store()
4439 .get_step(s_pending.id)
4440 .await
4441 .unwrap()
4442 .unwrap();
4443 assert_eq!(r_pending.status.state, StepStatus::Skipped);
4444 assert!(r_pending.error.is_none());
4445 }
4446
4447 #[tokio::test]
4448 async fn fail_orphaned_steps_no_steps_is_noop() {
4449 let engine = create_test_engine();
4450 let run = engine
4451 .store()
4452 .create_run(NewRun {
4453 created_by: None,
4454 workflow_name: "test".to_string(),
4455 trigger: TriggerKind::Manual,
4456 payload: json!({}),
4457 max_retries: 0,
4458 handler_version: None,
4459 labels: HashMap::new(),
4460 scheduled_at: None,
4461 idempotency_key: None,
4462 concurrency_key: None,
4463 priority: 0,
4464 concurrency_limits: Vec::new(),
4465 max_cost_usd: None,
4466 worker_tags: Vec::new(),
4467 })
4468 .await
4469 .unwrap()
4470 .into_run();
4471
4472 let result = engine.fail_orphaned_steps(run.id, "timeout").await;
4473 assert!(result.is_ok());
4474 }
4475
4476 #[tokio::test]
4477 async fn fail_orphaned_steps_preserves_existing_error() {
4478 let engine = create_test_engine();
4479 let run = engine
4480 .store()
4481 .create_run(NewRun {
4482 created_by: None,
4483 workflow_name: "test".to_string(),
4484 trigger: TriggerKind::Manual,
4485 payload: json!({}),
4486 max_retries: 0,
4487 handler_version: None,
4488 labels: HashMap::new(),
4489 scheduled_at: None,
4490 idempotency_key: None,
4491 concurrency_key: None,
4492 priority: 0,
4493 concurrency_limits: Vec::new(),
4494 max_cost_usd: None,
4495 worker_tags: Vec::new(),
4496 })
4497 .await
4498 .unwrap()
4499 .into_run();
4500
4501 let step_with_error = create_step_with_status(
4502 engine.store(),
4503 run.id,
4504 "already-errored",
4505 0,
4506 StepStatus::Running,
4507 )
4508 .await;
4509
4510 engine
4511 .store()
4512 .update_step(
4513 step_with_error.id,
4514 StepUpdate {
4515 error: Some("real error from provider".to_string()),
4516 ..StepUpdate::default()
4517 },
4518 )
4519 .await
4520 .unwrap();
4521
4522 let step_no_error = create_step_with_status(
4523 engine.store(),
4524 run.id,
4525 "no-error-yet",
4526 1,
4527 StepStatus::Running,
4528 )
4529 .await;
4530
4531 engine
4532 .fail_orphaned_steps(run.id, "parent run failed")
4533 .await
4534 .unwrap();
4535
4536 let updated_with = engine
4537 .store()
4538 .get_step(step_with_error.id)
4539 .await
4540 .unwrap()
4541 .unwrap();
4542 assert_eq!(updated_with.status.state, StepStatus::Failed);
4543 assert_eq!(
4544 updated_with.error.as_deref(),
4545 Some("real error from provider"),
4546 );
4547
4548 let updated_without = engine
4549 .store()
4550 .get_step(step_no_error.id)
4551 .await
4552 .unwrap()
4553 .unwrap();
4554 assert_eq!(updated_without.status.state, StepStatus::Failed);
4555 assert_eq!(updated_without.error.as_deref(), Some("parent run failed"),);
4556 }
4557
4558 use std::sync::atomic::{AtomicUsize, Ordering};
4563 use std::time::Duration;
4564
4565 use ironflow_store::models::LeaseRequest;
4566 use tokio::time::sleep;
4567
4568 struct CountingWorkflow {
4570 executions: Arc<AtomicUsize>,
4571 }
4572
4573 impl WorkflowHandler for CountingWorkflow {
4574 fn name(&self) -> &str {
4575 "local-resume-counter"
4576 }
4577
4578 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
4579 Box::pin(async move {
4580 self.executions.fetch_add(1, Ordering::SeqCst);
4581 Ok(())
4582 })
4583 }
4584 }
4585
4586 async fn engine_with_pending_counter() -> (Arc<Engine>, Arc<AtomicUsize>, Uuid) {
4588 let executions = Arc::new(AtomicUsize::new(0));
4589 let mut engine = create_test_engine();
4590 engine
4591 .register(CountingWorkflow {
4592 executions: executions.clone(),
4593 })
4594 .unwrap();
4595 let run = engine
4596 .enqueue_handler("local-resume-counter", TriggerKind::Manual, json!({}), 0)
4597 .await
4598 .unwrap();
4599 (Arc::new(engine), executions, run.id)
4600 }
4601
4602 #[tokio::test]
4603 async fn local_resume_leaves_a_run_claimed_by_a_worker_alone() {
4604 let (engine, executions, run_id) = engine_with_pending_counter().await;
4605 let lease = LeaseRequest {
4606 worker_id: "worker-1".to_string(),
4607 ttl: Duration::from_secs(30),
4608 };
4609 let picked = engine
4610 .store()
4611 .pick_next_pending_for(Some(lease), None)
4612 .await
4613 .unwrap()
4614 .expect("the run is due");
4615 assert_eq!(picked.id, run_id);
4616
4617 engine.spawn_local_resume(run_id);
4618 sleep(Duration::from_millis(200)).await;
4619
4620 assert_eq!(executions.load(Ordering::SeqCst), 0);
4621 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4622 assert_eq!(run.status.state, RunStatus::Running);
4623 assert_eq!(run.worker_id.as_deref(), Some("worker-1"));
4624 }
4625
4626 #[tokio::test]
4627 async fn local_resume_claims_a_pending_run_once() {
4628 let (engine, executions, run_id) = engine_with_pending_counter().await;
4629
4630 engine.spawn_local_resume(run_id);
4631 engine.spawn_local_resume(run_id);
4632 sleep(Duration::from_millis(200)).await;
4633
4634 assert_eq!(executions.load(Ordering::SeqCst), 1);
4635 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4636 assert_eq!(run.status.state, RunStatus::Completed);
4637 }
4638
4639 #[tokio::test]
4640 async fn local_resume_of_a_running_run_executes_it_once() {
4641 let (engine, executions, run_id) = engine_with_pending_counter().await;
4642 engine
4643 .store()
4644 .update_run_status(run_id, RunStatus::Running)
4645 .await
4646 .unwrap();
4647
4648 engine.spawn_local_resume(run_id);
4649 engine.spawn_local_resume(run_id);
4650 sleep(Duration::from_millis(200)).await;
4651
4652 assert_eq!(executions.load(Ordering::SeqCst), 1);
4653 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4654 assert_eq!(run.status.state, RunStatus::Completed);
4655 }
4656}