1use std::collections::{HashMap, HashSet};
10use std::fmt;
11use std::sync::{Arc, Mutex};
12use std::time::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, StepStatus, StepUpdate,
33 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::context::{PARENT_RUN_ID_LABEL, WorkflowContext, interrupt_running_steps};
43use crate::error::EngineError;
44use crate::executor::{StepInterceptor, StepResult};
45use crate::guard::{WorkflowGuardConfig, new_shared_guard_state};
46use crate::handler::{WorkflowHandler, WorkflowInfo, clamp_priority};
47use crate::log_sender::LogSender;
48use crate::notify::{
49 ApprovalRequestedEvent, Event, EventBusReleaser, EventPublisher, EventSubscriber,
50 RunBudgetExceededEvent, RunFailedEvent, RunStatusChangedEvent, SignalAwaitedEvent,
51 SignalReceivedEvent, WorkflowEventBus,
52};
53use crate::plan::{
54 ExecutionPlan, PlanOptions, PlanRecorder, SharedPlanRecorder, estimate_durations, lock_plan,
55};
56use crate::retry_policy::{backoff_for_retry, is_run_retryable};
57use crate::schedule::CronSchedule;
58use crate::signal::{
59 Signal, SignalDelivery, SignalRejected, SignalResumed, received_output, validate_step_payload,
60};
61use ironflow_core::decision::DecisionProvider;
62
63#[derive(Debug, Clone)]
82pub struct WorkflowResult {
83 pub run: Run,
85 pub steps: Vec<StepResult>,
87}
88
89#[derive(Debug, Clone, Default)]
108pub struct EnqueueOptions {
109 pub max_retries: u32,
111 pub labels: HashMap<String, String>,
113 pub scheduled_at: Option<DateTime<Utc>>,
116 pub max_cost_usd: Option<Decimal>,
120 pub created_by: Option<RunActor>,
123 pub idempotency_key: Option<String>,
129 pub concurrency_key: Option<String>,
135 pub concurrency_limits: Vec<ConcurrencyLimit>,
143 pub priority: Option<i16>,
152 pub worker_tags: Vec<String>,
158}
159
160#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
171pub enum ExecutionMode {
172 #[default]
176 Local,
177 Workers,
180}
181
182pub struct Engine {
223 store: Arc<dyn Store>,
224 provider: Arc<dyn AgentProvider>,
225 handlers: HashMap<String, Arc<dyn WorkflowHandler>>,
226 event_publisher: EventPublisher,
227 log_sender: Option<LogSender>,
228 budget: BudgetConfig,
229 artifact_sink: Option<Arc<dyn ArtifactSink>>,
230 guard_config: Option<WorkflowGuardConfig>,
231 event_bus: Option<WorkflowEventBus>,
232 decision_provider: Option<Arc<dyn DecisionProvider>>,
233 step_interceptor: Option<Arc<dyn StepInterceptor>>,
234 execution_mode: ExecutionMode,
235 worker_tags: Option<Arc<Vec<String>>>,
236 active_runs: ActiveRuns,
237 local_resumes: ActiveRuns,
238}
239
240type ActiveRuns = Arc<Mutex<HashMap<Uuid, Arc<AsyncMutex<()>>>>>;
242
243struct ActiveRunGuard {
249 runs: ActiveRuns,
250 run_id: Uuid,
251 entry: Arc<AsyncMutex<()>>,
252 lock: Option<OwnedMutexGuard<()>>,
253}
254
255impl Drop for ActiveRunGuard {
256 fn drop(&mut self) {
257 self.lock.take();
258 prune_active_run(&self.runs, self.run_id, &self.entry);
259 }
260}
261
262async fn lock_run(runs: &ActiveRuns, run_id: Uuid) -> ActiveRunGuard {
264 let entry = {
265 let mut runs = runs.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
266 Arc::clone(runs.entry(run_id).or_default())
267 };
268 let lock = Arc::clone(&entry).lock_owned().await;
269 ActiveRunGuard {
270 runs: Arc::clone(runs),
271 run_id,
272 entry,
273 lock: Some(lock),
274 }
275}
276
277fn prune_active_run(runs: &ActiveRuns, run_id: Uuid, entry: &Arc<AsyncMutex<()>>) {
282 let mut runs = runs.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
283 let unused = runs
284 .get(&run_id)
285 .is_some_and(|current| Arc::ptr_eq(current, entry))
286 && Arc::strong_count(entry) == 2;
287 if unused {
288 runs.remove(&run_id);
289 }
290}
291
292fn validate_category(handler_name: &str, category: &str) -> Result<(), EngineError> {
302 let reject = |reason: &str| {
303 Err(EngineError::InvalidWorkflow(format!(
304 "handler '{handler_name}' has invalid category '{category}': {reason}"
305 )))
306 };
307
308 if category.is_empty() {
309 return reject("empty category");
310 }
311 if category.starts_with('/') {
312 return reject("leading '/'");
313 }
314 if category.ends_with('/') {
315 return reject("trailing '/'");
316 }
317 for segment in category.split('/') {
318 if segment.is_empty() {
319 return reject("empty segment (double '/')");
320 }
321 if segment.trim().is_empty() {
322 return reject("whitespace-only segment");
323 }
324 }
325 Ok(())
326}
327
328fn chain_label(run: &Run, key: &str) -> Option<Uuid> {
333 if !matches!(run.trigger, TriggerKind::Workflow) {
334 return None;
335 }
336 let id = Uuid::parse_str(run.labels.get(key)?).ok()?;
337 (id != run.id).then_some(id)
338}
339
340pub fn chain_root(run: &Run) -> Option<Uuid> {
360 chain_label(run, LABEL_ROOT_RUN_ID)
361}
362
363pub(crate) fn chain_parent(run: &Run) -> Option<Uuid> {
365 chain_label(run, PARENT_RUN_ID_LABEL)
366}
367
368impl Engine {
369 pub fn new(store: Arc<dyn Store>, provider: Arc<dyn AgentProvider>) -> Self {
385 Self {
386 store,
387 provider,
388 handlers: HashMap::new(),
389 event_publisher: EventPublisher::new(),
390 log_sender: None,
391 budget: BudgetConfig::new(),
392 artifact_sink: None,
393 guard_config: None,
394 event_bus: None,
395 decision_provider: None,
396 step_interceptor: None,
397 execution_mode: ExecutionMode::default(),
398 worker_tags: None,
399 active_runs: Arc::new(Mutex::new(HashMap::new())),
400 local_resumes: Arc::new(Mutex::new(HashMap::new())),
401 }
402 }
403
404 async fn track_execution(&self, run_id: Uuid) -> ActiveRunGuard {
407 lock_run(&self.active_runs, run_id).await
408 }
409
410 async fn wait_until_idle(&self, run_id: Uuid) {
412 let entry = {
413 let runs = self
414 .active_runs
415 .lock()
416 .unwrap_or_else(|poisoned| poisoned.into_inner());
417 runs.get(&run_id).map(Arc::clone)
418 };
419 if let Some(entry) = entry {
420 drop(Arc::clone(&entry).lock_owned().await);
421 prune_active_run(&self.active_runs, run_id, &entry);
422 }
423 }
424
425 pub(crate) fn is_executing(&self, run_id: Uuid) -> bool {
427 let runs = self
428 .active_runs
429 .lock()
430 .unwrap_or_else(|poisoned| poisoned.into_inner());
431 runs.get(&run_id)
432 .is_some_and(|entry| entry.try_lock().is_err())
433 }
434
435 pub fn with_decision_provider(mut self, provider: Arc<dyn DecisionProvider>) -> Self {
456 self.decision_provider = Some(provider);
457 self
458 }
459
460 pub fn with_step_interceptor(mut self, interceptor: Arc<dyn StepInterceptor>) -> Self {
484 self.step_interceptor = Some(interceptor);
485 self
486 }
487
488 pub fn step_interceptor(&self) -> Option<&Arc<dyn StepInterceptor>> {
490 self.step_interceptor.as_ref()
491 }
492
493 pub fn with_budget_config(mut self, budget: BudgetConfig) -> Self {
514 self.budget = budget;
515 self
516 }
517
518 pub fn budget_config(&self) -> &BudgetConfig {
520 &self.budget
521 }
522
523 pub fn with_guard_config(mut self, config: WorkflowGuardConfig) -> Self {
545 self.guard_config = Some(config);
546 self
547 }
548
549 pub fn guard_config(&self) -> Option<&WorkflowGuardConfig> {
551 self.guard_config.as_ref()
552 }
553
554 pub fn with_execution_mode(mut self, mode: ExecutionMode) -> Self {
575 self.execution_mode = mode;
576 self
577 }
578
579 pub fn execution_mode(&self) -> ExecutionMode {
581 self.execution_mode
582 }
583
584 pub fn set_worker_tags(&mut self, tags: Vec<String>) {
607 self.worker_tags = Some(Arc::new(tags));
608 }
609
610 pub fn worker_tags(&self) -> Option<&[String]> {
628 self.worker_tags.as_deref().map(Vec::as_slice)
629 }
630
631 pub fn set_log_sender(&mut self, sender: LogSender) {
637 self.log_sender = Some(sender);
638 }
639
640 pub fn set_artifact_sink(&mut self, sink: Arc<dyn ArtifactSink>) {
658 self.artifact_sink = Some(sink);
659 }
660
661 pub fn artifact_sink(&self) -> Option<&Arc<dyn ArtifactSink>> {
663 self.artifact_sink.as_ref()
664 }
665
666 pub fn set_event_bus(&mut self, bus: WorkflowEventBus) {
687 self.event_bus = Some(bus.clone());
688 self.event_publisher
689 .subscribe(EventBusReleaser::new(bus), &[Event::RUN_STATUS_CHANGED]);
690 }
691
692 pub fn event_bus(&self) -> Option<&WorkflowEventBus> {
694 self.event_bus.as_ref()
695 }
696
697 pub fn store(&self) -> &Arc<dyn Store> {
699 &self.store
700 }
701
702 pub fn provider(&self) -> &Arc<dyn AgentProvider> {
704 &self.provider
705 }
706
707 fn build_context(&self, run: &Run) -> WorkflowContext {
716 let handlers = self.handlers.clone();
717 let resolver: crate::context::HandlerResolver =
718 Arc::new(move |name: &str| handlers.get(name).cloned());
719 let mut ctx = WorkflowContext::with_handler_resolver(
720 run.id,
721 run.workflow_name.clone(),
722 self.store.clone(),
723 self.provider.clone(),
724 resolver,
725 );
726 ctx.carry_over_run_totals(run.retry_count + 1, run.cost_usd, run.duration_ms);
727 ctx.set_max_cost_usd(run.max_cost_usd);
728 ctx.set_run_created_at(run.created_at);
729 if let Some(ref sender) = self.log_sender {
730 ctx.set_log_sender(sender.clone());
731 }
732 if let Some(ref sink) = self.artifact_sink {
733 ctx.set_artifact_sink(sink.clone());
734 }
735 if let Some(ref bus) = self.event_bus {
736 ctx.set_event_bus(bus.clone());
737 }
738 if let Some(ref provider) = self.decision_provider {
739 ctx.set_decision_provider(provider.clone());
740 }
741 if let Some(ref interceptor) = self.step_interceptor {
742 ctx.set_step_interceptor(interceptor.clone());
743 }
744 if let Some(ref tags) = self.worker_tags {
745 ctx.set_worker_tags(tags.clone());
746 }
747 ctx
748 }
749
750 fn build_context_with_guard(
756 &self,
757 run: &Run,
758 handler: &dyn WorkflowHandler,
759 ) -> WorkflowContext {
760 let mut ctx = self.build_context(run);
761 let guard_config = handler.guard_config().or_else(|| self.guard_config.clone());
762 if let Some(config) = guard_config {
763 ctx.set_guard(config, new_shared_guard_state());
764 }
765 ctx
766 }
767
768 pub async fn check_monthly_quota(&self, workflow_name: &str) -> Result<(), EngineError> {
795 let Some(limit) = self.budget.monthly_cost_limit_usd else {
796 return Ok(());
797 };
798
799 let stats = self
800 .store
801 .get_stats(RunFilter {
802 created_after: Some(month_start(Utc::now())),
803 ..RunFilter::default()
804 })
805 .await?;
806
807 if stats.total_cost_usd < limit {
808 return Ok(());
809 }
810
811 warn!(
812 workflow = %workflow_name,
813 limit_usd = %limit,
814 spent_usd = %stats.total_cost_usd,
815 "monthly cost quota exhausted, refusing new run"
816 );
817
818 #[cfg(feature = "prometheus")]
819 counter!(
820 RUN_BUDGET_EXCEEDED_TOTAL,
821 "workflow" => workflow_name.to_string(),
822 "scope" => "monthly",
823 )
824 .increment(1);
825
826 Err(EngineError::MonthlyBudgetExceeded {
827 limit_usd: limit,
828 spent_usd: stats.total_cost_usd,
829 })
830 }
831
832 pub fn register(&mut self, handler: impl WorkflowHandler + 'static) -> Result<(), EngineError> {
876 let name = handler.name().to_string();
877 if self.handlers.contains_key(&name) {
878 return Err(EngineError::InvalidWorkflow(format!(
879 "handler '{}' already registered",
880 name
881 )));
882 }
883 if let Some(category) = handler.category() {
884 validate_category(&name, category)?;
885 }
886 self.handlers.insert(name, Arc::new(handler));
887 Ok(())
888 }
889
890 pub fn register_boxed(&mut self, handler: Box<dyn WorkflowHandler>) -> Result<(), EngineError> {
897 let name = handler.name().to_string();
898 if self.handlers.contains_key(&name) {
899 return Err(EngineError::InvalidWorkflow(format!(
900 "handler '{}' already registered",
901 name
902 )));
903 }
904 if let Some(category) = handler.category() {
905 validate_category(&name, category)?;
906 }
907 self.handlers.insert(name, Arc::from(handler));
908 Ok(())
909 }
910
911 pub fn get_handler(&self, name: &str) -> Option<&Arc<dyn WorkflowHandler>> {
913 self.handlers.get(name)
914 }
915
916 pub fn handler_names(&self) -> Vec<&str> {
918 self.handlers.keys().map(|s| s.as_str()).collect()
919 }
920
921 pub fn handler_info(&self, name: &str) -> Option<WorkflowInfo> {
923 self.handlers.get(name).map(|h| h.describe())
924 }
925
926 pub fn scheduled_handlers(&self) -> Vec<(&str, &CronSchedule)> {
950 self.handlers
951 .iter()
952 .filter_map(|(name, handler)| handler.schedule().map(|sched| (name.as_str(), sched)))
953 .collect()
954 }
955
956 pub fn subscribe(
981 &mut self,
982 subscriber: impl EventSubscriber + 'static,
983 event_types: &[&'static str],
984 ) {
985 self.event_publisher.subscribe(subscriber, event_types);
986 }
987
988 pub fn event_publisher(&self) -> &EventPublisher {
993 &self.event_publisher
994 }
995
996 #[tracing::instrument(name = "engine.run_handler", skip_all, fields(workflow = %handler_name))]
1032 pub async fn run_handler(
1033 &self,
1034 handler_name: &str,
1035 trigger: TriggerKind,
1036 payload: Value,
1037 ) -> Result<WorkflowResult, EngineError> {
1038 let handler = self
1039 .handlers
1040 .get(handler_name)
1041 .ok_or_else(|| {
1042 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1043 })?
1044 .clone();
1045
1046 self.check_monthly_quota(handler_name).await?;
1047
1048 let handler_version = handler.version().map(str::to_string);
1049 let max_cost_usd = self
1050 .budget
1051 .resolve_run_cap(None, handler.default_max_cost_usd());
1052 let run = self
1053 .store
1054 .create_run(NewRun {
1055 created_by: None,
1056 workflow_name: handler_name.to_string(),
1057 trigger,
1058 payload,
1059 max_retries: 0,
1060 handler_version,
1061 labels: handler.default_labels(),
1062 scheduled_at: None,
1063 idempotency_key: None,
1064 concurrency_key: None,
1065 priority: clamp_priority(handler.priority()),
1066 concurrency_limits: Vec::new(),
1067 max_cost_usd,
1068 worker_tags: normalize_worker_tags(handler.required_worker_tags()),
1069 })
1070 .await?
1071 .into_run();
1072
1073 let run_id = run.id;
1074 info!(run_id = %run_id, handler_version = run.handler_version.as_deref().unwrap_or(""), "run created");
1075
1076 if self.is_workflow_paused(handler_name).await? {
1077 info!(run_id = %run_id, workflow = %handler_name, "workflow paused, run left pending");
1078 return Ok(WorkflowResult {
1079 run,
1080 steps: Vec::new(),
1081 });
1082 }
1083
1084 self.store
1085 .update_run_status(run_id, RunStatus::Running)
1086 .await?;
1087
1088 #[cfg(feature = "prometheus")]
1089 gauge!(RUNS_ACTIVE, "workflow" => handler_name.to_string()).increment(1.0);
1090
1091 let run_start = Instant::now();
1092 let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
1093
1094 let result = handler.execute(&mut ctx).await;
1095 self.finalize_run(run_id, handler_name, result, &ctx, run_start, run.labels)
1096 .await
1097 }
1098
1099 #[tracing::instrument(name = "engine.plan_handler", skip_all, fields(workflow = %handler_name))]
1148 pub async fn plan_handler(
1149 &self,
1150 handler_name: &str,
1151 payload: Value,
1152 options: PlanOptions,
1153 ) -> Result<ExecutionPlan, EngineError> {
1154 if options.max_depth == 0 {
1155 return Err(EngineError::InvalidWorkflow(
1156 "max_depth must be at least 1".to_string(),
1157 ));
1158 }
1159
1160 let handler = self
1161 .handlers
1162 .get(handler_name)
1163 .ok_or_else(|| {
1164 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1165 })?
1166 .clone();
1167
1168 let estimates = if options.estimate_durations {
1169 estimate_durations(&self.store, handler_name, options.sample_runs).await?
1170 } else {
1171 HashMap::new()
1172 };
1173
1174 let shared: SharedPlanRecorder = Arc::new(Mutex::new(PlanRecorder::new(
1175 handler_name.to_string(),
1176 payload,
1177 options.max_depth,
1178 estimates,
1179 )));
1180
1181 let handlers = self.handlers.clone();
1184 let resolver: crate::context::HandlerResolver =
1185 Arc::new(move |name: &str| handlers.get(name).cloned());
1186 let mut ctx = WorkflowContext::with_handler_resolver(
1187 Uuid::now_v7(),
1188 handler_name.to_string(),
1189 self.store.clone(),
1190 self.provider.clone(),
1191 resolver,
1192 );
1193 ctx.set_plan(shared.clone());
1194
1195 match timeout(options.timeout, handler.execute(&mut ctx)).await {
1196 Ok(Ok(())) => {}
1197 Ok(Err(err)) => lock_plan(&shared).fail(err.to_string()),
1198 Err(_) => lock_plan(&shared).fail(format!(
1199 "plan timed out after {} s",
1200 options.timeout.as_secs_f64()
1201 )),
1202 }
1203 drop(ctx);
1204
1205 let plan = match Arc::try_unwrap(shared) {
1206 Ok(mutex) => mutex
1207 .into_inner()
1208 .unwrap_or_else(|poisoned| poisoned.into_inner())
1209 .into_plan(),
1210 Err(shared) => lock_plan(&shared).snapshot(),
1211 };
1212
1213 info!(
1214 workflow = %handler_name,
1215 steps = plan.steps.len(),
1216 truncated = plan.truncated,
1217 "execution plan built"
1218 );
1219
1220 Ok(plan)
1221 }
1222
1223 #[tracing::instrument(name = "engine.enqueue_handler", skip_all, fields(workflow = %handler_name))]
1234 pub async fn enqueue_handler(
1235 &self,
1236 handler_name: &str,
1237 trigger: TriggerKind,
1238 payload: Value,
1239 max_retries: u32,
1240 ) -> Result<Run, EngineError> {
1241 self.enqueue_handler_with_options(
1242 handler_name,
1243 trigger,
1244 payload,
1245 EnqueueOptions {
1246 max_retries,
1247 ..Default::default()
1248 },
1249 )
1250 .await
1251 .map(RunCreation::into_run)
1252 }
1253
1254 #[tracing::instrument(name = "engine.enqueue_handler_with_options", skip_all, fields(workflow = %handler_name))]
1306 pub async fn enqueue_handler_with_options(
1307 &self,
1308 handler_name: &str,
1309 trigger: TriggerKind,
1310 payload: Value,
1311 options: EnqueueOptions,
1312 ) -> Result<RunCreation, EngineError> {
1313 let EnqueueOptions {
1314 max_retries,
1315 labels,
1316 scheduled_at,
1317 max_cost_usd,
1318 created_by,
1319 idempotency_key,
1320 concurrency_key,
1321 concurrency_limits,
1322 priority,
1323 worker_tags,
1324 } = options;
1325
1326 validate_concurrency_limits(&concurrency_limits)
1329 .map_err(EngineError::InvalidConcurrencyLimit)?;
1330 if let Some(priority) = priority {
1331 validate_priority(priority).map_err(EngineError::InvalidPriority)?;
1332 }
1333 validate_worker_tags(&worker_tags).map_err(EngineError::InvalidWorkerTag)?;
1334
1335 let handler = self.handlers.get(handler_name).ok_or_else(|| {
1336 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1337 })?;
1338
1339 self.check_monthly_quota(handler_name).await?;
1340
1341 let handler_version = handler.version().map(str::to_string);
1342 let mut merged_labels = handler.default_labels();
1343 merged_labels.extend(labels);
1344 let resolved_cap = self
1345 .budget
1346 .resolve_run_cap(max_cost_usd, handler.default_max_cost_usd());
1347 let priority = priority.unwrap_or_else(|| clamp_priority(handler.priority()));
1348 let required_tags = normalize_worker_tags(
1349 handler
1350 .required_worker_tags()
1351 .into_iter()
1352 .chain(worker_tags),
1353 );
1354
1355 let creation = self
1356 .store
1357 .create_run(NewRun {
1358 workflow_name: handler_name.to_string(),
1359 trigger,
1360 payload,
1361 max_retries,
1362 handler_version,
1363 labels: merged_labels,
1364 scheduled_at,
1365 created_by,
1366 idempotency_key,
1367 concurrency_key,
1368 priority,
1369 concurrency_limits,
1370 max_cost_usd: resolved_cap,
1371 worker_tags: required_tags,
1372 })
1373 .await?;
1374
1375 match &creation {
1376 RunCreation::Created(run) => info!(
1377 run_id = %run.id,
1378 workflow = %handler_name,
1379 max_cost_usd = ?resolved_cap,
1380 "handler run enqueued"
1381 ),
1382 RunCreation::Existing(run) => info!(
1383 run_id = %run.id,
1384 workflow = %handler_name,
1385 "idempotent replay, nothing enqueued"
1386 ),
1387 }
1388
1389 Ok(creation)
1390 }
1391
1392 #[tracing::instrument(name = "engine.execute_handler_run", skip_all, fields(run_id = %run_id))]
1412 pub async fn execute_handler_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1413 let run = self
1414 .store
1415 .get_run(run_id)
1416 .await?
1417 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1418
1419 if let Some(root_run_id) = chain_root(&run) {
1420 return self.resume_chain(run, root_run_id).await;
1421 }
1422
1423 let _active = self.track_execution(run_id).await;
1424
1425 let handler = self
1426 .handlers
1427 .get(&run.workflow_name)
1428 .ok_or_else(|| {
1429 EngineError::InvalidWorkflow(format!(
1430 "no handler registered: {}",
1431 run.workflow_name
1432 ))
1433 })?
1434 .clone();
1435
1436 #[cfg(feature = "prometheus")]
1437 gauge!(RUNS_ACTIVE, "workflow" => run.workflow_name.clone()).increment(1.0);
1438
1439 let run_start = Instant::now();
1440 let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
1441
1442 let result = if handler.is_version_compatible(run.handler_version.as_deref()) {
1455 ctx.load_replay_steps().await?;
1456 self.release_then_execute(run_id, handler.as_ref(), &mut ctx)
1457 .await
1458 } else {
1459 Err(EngineError::HandlerVersionMismatch {
1460 run_id,
1461 workflow_name: run.workflow_name.clone(),
1462 run_version: run
1463 .handler_version
1464 .clone()
1465 .unwrap_or_else(|| "unknown".to_string()),
1466 current_version: handler
1467 .version()
1468 .map(str::to_string)
1469 .unwrap_or_else(|| "unknown".to_string()),
1470 })
1471 };
1472
1473 self.finalize_run(
1474 run_id,
1475 &run.workflow_name,
1476 result,
1477 &ctx,
1478 run_start,
1479 run.labels,
1480 )
1481 .await
1482 }
1483
1484 #[tracing::instrument(name = "engine.execute_run", skip_all, fields(run_id = %run_id))]
1492 pub async fn execute_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1493 self.execute_handler_run(run_id).await
1494 }
1495
1496 #[tracing::instrument(name = "engine.resume_run", skip_all, fields(run_id = %run_id))]
1525 pub async fn resume_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1526 let run = self
1527 .store
1528 .get_run(run_id)
1529 .await?
1530 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1531
1532 if let Some(root_run_id) = chain_root(&run) {
1533 return self.resume_chain(run, root_run_id).await;
1534 }
1535
1536 self.resume_loaded_run(run).await
1537 }
1538
1539 async fn resume_chain(
1554 &self,
1555 child: Run,
1556 root_run_id: Uuid,
1557 ) -> Result<WorkflowResult, EngineError> {
1558 let child_run_id = child.id;
1559 let lease = child.worker_id.zip(child.lease_expires_at);
1560 let root = self
1561 .store
1562 .get_run(root_run_id)
1563 .await?
1564 .ok_or(EngineError::Store(StoreError::RunNotFound(root_run_id)))?;
1565
1566 match root.status.state {
1567 RunStatus::AwaitingApproval | RunStatus::Pending => {
1568 self.move_root_to_running(root_run_id, lease.as_ref())
1569 .await?;
1570 }
1571 RunStatus::Sleeping => {
1572 self.store
1573 .update_run_status(root_run_id, RunStatus::Pending)
1574 .await?;
1575 self.move_root_to_running(root_run_id, lease.as_ref())
1576 .await?;
1577 }
1578 other => {
1579 let reason = format!(
1580 "cannot resume child run {child_run_id}: root run {root_run_id} is {other}"
1581 );
1582 if let Err(err) = self
1583 .fail_or_schedule_retry(child_run_id, &reason, false, None, None)
1584 .await
1585 {
1586 error!(
1587 run_id = %child_run_id,
1588 error = %err,
1589 "failed to fail a child run whose root cannot resume"
1590 );
1591 }
1592 return Err(EngineError::InvalidWorkflow(reason));
1593 }
1594 }
1595
1596 if lease.is_some() {
1597 self.store
1598 .update_run(
1599 child_run_id,
1600 RunUpdate {
1601 lease: Some(LeaseUpdate::Release),
1602 ..RunUpdate::default()
1603 },
1604 )
1605 .await?;
1606 }
1607
1608 info!(
1609 run_id = %child_run_id,
1610 root_run_id = %root_run_id,
1611 lease_transferred = lease.is_some(),
1612 "child run resumed through its root run"
1613 );
1614
1615 let root = self
1616 .store
1617 .get_run(root_run_id)
1618 .await?
1619 .ok_or(EngineError::Store(StoreError::RunNotFound(root_run_id)))?;
1620 self.resume_loaded_run(root).await
1621 }
1622
1623 async fn move_root_to_running(
1629 &self,
1630 root_run_id: Uuid,
1631 lease: Option<&(String, DateTime<Utc>)>,
1632 ) -> Result<(), EngineError> {
1633 match lease {
1634 Some((worker_id, expires_at)) => {
1635 self.store
1636 .update_run(
1637 root_run_id,
1638 RunUpdate {
1639 status: Some(RunStatus::Running),
1640 lease: Some(LeaseUpdate::Set {
1641 worker_id: worker_id.clone(),
1642 expires_at: *expires_at,
1643 }),
1644 ..RunUpdate::default()
1645 },
1646 )
1647 .await?;
1648 }
1649 None => {
1650 self.store
1651 .update_run_status(root_run_id, RunStatus::Running)
1652 .await?;
1653 }
1654 }
1655 Ok(())
1656 }
1657
1658 async fn resume_loaded_run(&self, run: Run) -> Result<WorkflowResult, EngineError> {
1660 let run_id = run.id;
1661 let _active = self.track_execution(run_id).await;
1662 let handler = self
1663 .handlers
1664 .get(&run.workflow_name)
1665 .ok_or_else(|| {
1666 EngineError::InvalidWorkflow(format!(
1667 "no handler registered: {}",
1668 run.workflow_name
1669 ))
1670 })?
1671 .clone();
1672
1673 info!(run_id = %run_id, workflow = %run.workflow_name, "resuming run after approval");
1674
1675 let run_start = Instant::now();
1676 let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
1677
1678 let result = if handler.is_version_compatible(run.handler_version.as_deref()) {
1679 ctx.load_replay_steps().await?;
1680 self.release_then_execute(run_id, handler.as_ref(), &mut ctx)
1681 .await
1682 } else {
1683 Err(EngineError::HandlerVersionMismatch {
1684 run_id,
1685 workflow_name: run.workflow_name.clone(),
1686 run_version: run
1687 .handler_version
1688 .clone()
1689 .unwrap_or_else(|| "unknown".to_string()),
1690 current_version: handler
1691 .version()
1692 .map(str::to_string)
1693 .unwrap_or_else(|| "unknown".to_string()),
1694 })
1695 };
1696
1697 self.finalize_run(
1698 run_id,
1699 &run.workflow_name,
1700 result,
1701 &ctx,
1702 run_start,
1703 run.labels,
1704 )
1705 .await
1706 }
1707
1708 pub async fn deliver_signal(
1754 self: &Arc<Self>,
1755 signal: NewSignal,
1756 ) -> Result<SignalDelivery, EngineError> {
1757 if signal.name.trim().is_empty() {
1758 return Err(EngineError::InvalidSignal(
1759 "signal name must not be empty".to_string(),
1760 ));
1761 }
1762 if signal.key.trim().is_empty() {
1763 return Err(EngineError::InvalidSignal(
1764 "signal key must not be empty".to_string(),
1765 ));
1766 }
1767
1768 let stored = match self.store.insert_signal(signal).await? {
1769 SignalInsert::Created(stored) => stored,
1770 SignalInsert::Duplicate(existing) => {
1771 info!(
1772 signal_id = %existing.id,
1773 signal = %existing.name,
1774 key = %existing.key,
1775 "duplicate signal ignored"
1776 );
1777 return Ok(SignalDelivery {
1778 signal_id: existing.id,
1779 duplicate: true,
1780 resumed: Vec::new(),
1781 rejected: Vec::new(),
1782 });
1783 }
1784 };
1785
1786 let waiters = self
1787 .store
1788 .list_signal_waiters(&stored.name, &stored.key)
1789 .await?;
1790 let mut resumed = Vec::new();
1791 let mut rejected = Vec::new();
1792
1793 for step in waiters {
1794 if let Err(error) = validate_step_payload(step.input.as_ref(), &stored.payload) {
1795 rejected.push(SignalRejected {
1796 run_id: step.run_id,
1797 step_id: step.id,
1798 error,
1799 });
1800 continue;
1801 }
1802
1803 match self
1804 .store
1805 .resolve_signal_step(step.id, received_output(&stored))
1806 .await
1807 {
1808 Ok(SignalStepResolution::Resolved {
1809 run_id,
1810 run_resumed,
1811 }) => {
1812 resumed.push(SignalResumed {
1813 run_id,
1814 step_id: step.id,
1815 });
1816 if run_resumed && self.execution_mode == ExecutionMode::Local {
1817 self.spawn_local_resume(run_id);
1818 }
1819 }
1820 Ok(SignalStepResolution::NotWaiting { .. }) => {}
1822 Err(err) => {
1823 error!(
1824 run_id = %step.run_id,
1825 step_id = %step.id,
1826 error = %err,
1827 "failed to resolve a waiting signal step"
1828 );
1829 rejected.push(SignalRejected {
1830 run_id: step.run_id,
1831 step_id: step.id,
1832 error: err.to_string(),
1833 });
1834 }
1835 }
1836 }
1837
1838 info!(
1839 signal_id = %stored.id,
1840 signal = %stored.name,
1841 key = %stored.key,
1842 resumed = resumed.len(),
1843 rejected = rejected.len(),
1844 "signal received"
1845 );
1846 self.event_publisher
1847 .publish(Event::SignalReceived(SignalReceivedEvent {
1848 signal_id: stored.id,
1849 name: stored.name.clone(),
1850 key: stored.key.clone(),
1851 resumed_runs: resumed.iter().map(|r| r.run_id).collect(),
1852 at: stored.received_at,
1853 }));
1854
1855 Ok(SignalDelivery {
1856 signal_id: stored.id,
1857 duplicate: false,
1858 resumed,
1859 rejected,
1860 })
1861 }
1862
1863 pub async fn send_signal<S: Signal>(
1901 self: &Arc<Self>,
1902 signal: &S,
1903 key: &str,
1904 idempotency_id: Option<&str>,
1905 ) -> Result<SignalDelivery, EngineError> {
1906 let payload = to_value(signal)?;
1907 self.deliver_signal(NewSignal {
1908 name: S::NAME.to_string(),
1909 key: key.to_string(),
1910 payload,
1911 idempotency_id: idempotency_id.map(str::to_string),
1912 })
1913 .await
1914 }
1915
1916 pub(crate) fn spawn_local_resume(self: &Arc<Self>, run_id: Uuid) {
1941 let engine = Arc::clone(self);
1942 spawn(async move {
1943 let _resuming = lock_run(&engine.local_resumes, run_id).await;
1947 engine.wait_until_idle(run_id).await;
1948 let run = match engine.store.get_run(run_id).await {
1949 Ok(Some(run)) => run,
1950 Ok(None) => {
1951 error!(run_id = %run_id, "run to restart not found");
1952 return;
1953 }
1954 Err(err) => {
1955 error!(run_id = %run_id, error = %err, "failed to load a run to restart");
1956 return;
1957 }
1958 };
1959 match run.status.state {
1960 RunStatus::Pending => {
1961 match engine
1962 .store
1963 .update_run_status(run_id, RunStatus::Running)
1964 .await
1965 {
1966 Ok(()) => {}
1967 Err(StoreError::InvalidTransition { from, .. }) => {
1968 info!(
1969 run_id = %run_id,
1970 status = %from,
1971 "run claimed elsewhere, left to its worker"
1972 );
1973 return;
1974 }
1975 Err(err) => {
1976 error!(run_id = %run_id, error = %err, "failed to restart a woken run");
1977 return;
1978 }
1979 }
1980 }
1981 RunStatus::Running if run.worker_id.is_some() => {
1982 info!(
1983 run_id = %run_id,
1984 worker_id = ?run.worker_id,
1985 "run owned by a worker lease, local resume skipped"
1986 );
1987 return;
1988 }
1989 RunStatus::Running => {
1990 if engine.is_executing(run_id) {
1991 info!(run_id = %run_id, "run already executing in this process, resume skipped");
1992 return;
1993 }
1994 }
1995 status => {
1996 info!(
1997 run_id = %run_id,
1998 status = %status,
1999 "run no longer waiting to restart, resume skipped"
2000 );
2001 return;
2002 }
2003 }
2004 if let Err(err) = engine.resume_run(run_id).await {
2005 error!(run_id = %run_id, error = %err, "failed to resume a woken run");
2006 }
2007 });
2008 }
2009
2010 pub async fn fail_or_schedule_retry(
2061 &self,
2062 run_id: Uuid,
2063 error: &str,
2064 retryable: bool,
2065 cost_usd: Option<Decimal>,
2066 duration_ms: Option<u64>,
2067 ) -> Result<RunStatus, EngineError> {
2068 let run = self
2069 .store
2070 .get_run(run_id)
2071 .await?
2072 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2073
2074 if run.status.state == RunStatus::Paused {
2077 info!(run_id = %run_id, error = %error, "run paused, failure not recorded");
2078 return Ok(RunStatus::Paused);
2079 }
2080
2081 let has_attempts_left = run.retry_count < run.max_retries;
2082 let update = if retryable && has_attempts_left {
2083 let backoff = backoff_for_retry(run.retry_count);
2084 let scheduled_at = Utc::now() + TimeDelta::milliseconds(backoff.as_millis() as i64);
2085
2086 info!(
2087 run_id = %run_id,
2088 workflow = %run.workflow_name,
2089 attempt = run.retry_count + 1,
2090 max_retries = run.max_retries,
2091 backoff_secs = backoff.as_secs(),
2092 scheduled_at = %scheduled_at,
2093 "run failed, scheduling retry"
2094 );
2095
2096 RunUpdate {
2097 status: Some(RunStatus::Retrying),
2098 error: Some(error.to_string()),
2099 increment_retry: true,
2100 cost_usd,
2101 duration_ms,
2102 scheduled_at: Some(scheduled_at),
2103 ..RunUpdate::default()
2104 }
2105 } else {
2106 RunUpdate {
2107 status: Some(RunStatus::Failed),
2108 error: Some(error.to_string()),
2109 cost_usd,
2110 duration_ms,
2111 completed_at: Some(Utc::now()),
2112 ..RunUpdate::default()
2113 }
2114 };
2115
2116 let status = update.status.unwrap_or(RunStatus::Failed);
2117 self.store.update_run(run_id, update).await?;
2118 self.fail_orphaned_steps(run_id, error).await?;
2119 self.cancel_descendants_of_stopped_run(run_id, error).await;
2122
2123 Ok(status)
2124 }
2125
2126 pub async fn interrupt_running_steps(&self, run_id: Uuid) -> Result<(), EngineError> {
2156 interrupt_running_steps(self.store.as_ref(), run_id).await
2157 }
2158
2159 pub async fn fail_orphaned_steps(
2173 &self,
2174 run_id: Uuid,
2175 error_message: &str,
2176 ) -> Result<(), EngineError> {
2177 let steps = self.store.list_steps(run_id).await?;
2178 let now = Utc::now();
2179
2180 for step in steps {
2181 if step.status.state.is_terminal() {
2182 continue;
2183 }
2184
2185 let (target_status, error) = match step.status.state {
2186 StepStatus::Running | StepStatus::AwaitingApproval => {
2187 let err = if step.error.is_some() {
2188 None
2189 } else {
2190 Some(error_message.to_string())
2191 };
2192 (StepStatus::Failed, err)
2193 }
2194 StepStatus::Pending => (StepStatus::Skipped, None),
2195 _ => continue,
2196 };
2197
2198 if let Err(e) = self
2199 .store
2200 .update_step(
2201 step.id,
2202 StepUpdate {
2203 status: Some(target_status),
2204 error,
2205 completed_at: Some(now),
2206 ..StepUpdate::default()
2207 },
2208 )
2209 .await
2210 {
2211 warn!(
2212 run_id = %run_id,
2213 step_id = %step.id,
2214 step_name = %step.name,
2215 error = %e,
2216 "failed to cleanup orphaned step"
2217 );
2218 } else {
2219 info!(
2220 run_id = %run_id,
2221 step_id = %step.id,
2222 step_name = %step.name,
2223 from = %step.status.state,
2224 to = %target_status,
2225 "cleaned up orphaned step"
2226 );
2227 }
2228 }
2229
2230 Ok(())
2231 }
2232
2233 async fn release_then_execute(
2238 &self,
2239 run_id: Uuid,
2240 handler: &dyn WorkflowHandler,
2241 ctx: &mut WorkflowContext,
2242 ) -> Result<(), EngineError> {
2243 match self.provider.release_run(&run_id.to_string()).await {
2244 Ok(()) => handler.execute(ctx).await,
2245 Err(e) => Err(EngineError::Operation(OperationError::Agent(e))),
2246 }
2247 }
2248
2249 async fn finalize_run(
2255 &self,
2256 run_id: Uuid,
2257 workflow_name: &str,
2258 result: Result<(), EngineError>,
2259 ctx: &WorkflowContext,
2260 run_start: Instant,
2261 run_labels: HashMap<String, String>,
2262 ) -> Result<WorkflowResult, EngineError> {
2263 let total_duration = ctx.carried_duration_ms() + run_start.elapsed().as_millis() as u64;
2266 let completed_at = Utc::now();
2267
2268 if let Some(run) = self.store.get_run(run_id).await?
2272 && run.status.state == RunStatus::Paused
2273 {
2274 let run = self
2275 .store
2276 .update_run_returning(
2277 run_id,
2278 RunUpdate {
2279 cost_usd: Some(ctx.total_cost_usd()),
2280 duration_ms: Some(total_duration),
2281 ..RunUpdate::default()
2282 },
2283 )
2284 .await?;
2285 info!(
2286 run_id = %run_id,
2287 outcome = ?result.err().map(|err| err.to_string()),
2288 "run paused, execution stopped"
2289 );
2290 return Ok(WorkflowResult {
2291 run,
2292 steps: ctx.step_results().to_vec(),
2293 });
2294 }
2295
2296 if matches!(result, Err(EngineError::RunPaused { .. })) {
2299 let run = self
2300 .store
2301 .get_run(run_id)
2302 .await?
2303 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2304 info!(
2305 run_id = %run_id,
2306 status = %run.status.state,
2307 "run resumed after the pause, execution stopped"
2308 );
2309 return Ok(WorkflowResult {
2310 run,
2311 steps: ctx.step_results().to_vec(),
2312 });
2313 }
2314
2315 let final_status;
2316 let final_run;
2317
2318 match result {
2319 Ok(()) => {
2320 final_status = if ctx.has_allowed_failure() {
2321 RunStatus::Warning
2322 } else {
2323 RunStatus::Completed
2324 };
2325 final_run = self
2326 .store
2327 .update_run_returning(
2328 run_id,
2329 RunUpdate {
2330 status: Some(final_status),
2331 cost_usd: Some(ctx.total_cost_usd()),
2332 duration_ms: Some(total_duration),
2333 completed_at: Some(completed_at),
2334 output: ctx.output().cloned(),
2335 ..RunUpdate::default()
2336 },
2337 )
2338 .await?;
2339
2340 info!(
2341 run_id = %run_id,
2342 status = %final_status,
2343 cost_usd = %ctx.total_cost_usd(),
2344 duration_ms = total_duration,
2345 "run completed"
2346 );
2347 }
2348 Err(EngineError::ApprovalRequired {
2349 run_id: approval_run_id,
2350 step_id,
2351 ref message,
2352 }) => {
2353 final_status = RunStatus::AwaitingApproval;
2354 final_run = self
2355 .store
2356 .update_run_returning(
2357 run_id,
2358 RunUpdate {
2359 status: Some(RunStatus::AwaitingApproval),
2360 cost_usd: Some(ctx.total_cost_usd()),
2361 duration_ms: Some(total_duration),
2362 ..RunUpdate::default()
2363 },
2364 )
2365 .await?;
2366
2367 info!(
2368 run_id = %approval_run_id,
2369 step_id = %step_id,
2370 message = %message,
2371 "run awaiting approval"
2372 );
2373
2374 self.publish_approval_requested(approval_run_id, step_id, message)
2375 .await?;
2376 }
2377 Err(EngineError::ChildSuspended {
2378 run_id: child_run_id,
2379 ref cause,
2380 }) => {
2381 final_status = cause.suspension_status();
2382 final_run = self
2386 .store
2387 .update_run_returning(
2388 run_id,
2389 RunUpdate {
2390 status: Some(final_status),
2391 cost_usd: Some(ctx.total_cost_usd()),
2392 duration_ms: Some(total_duration),
2393 ..RunUpdate::default()
2394 },
2395 )
2396 .await?;
2397
2398 let leaf = cause.suspension_leaf();
2399 info!(
2400 run_id = %run_id,
2401 child_run_id = %child_run_id,
2402 status = %final_status,
2403 cause = %leaf,
2404 "run suspended with its child run"
2405 );
2406
2407 match leaf {
2408 EngineError::ApprovalRequired {
2409 run_id: approval_run_id,
2410 step_id,
2411 message,
2412 } => {
2413 self.publish_approval_requested(*approval_run_id, *step_id, message)
2414 .await?;
2415 }
2416 EngineError::SignalWaiting {
2417 run_id: wait_run_id,
2418 step_id,
2419 step_name,
2420 name,
2421 key,
2422 deadline_at,
2423 } => {
2424 self.event_publisher
2425 .publish(Event::SignalAwaited(SignalAwaitedEvent {
2426 run_id: *wait_run_id,
2427 step_id: *step_id,
2428 step_name: step_name.clone(),
2429 name: name.clone(),
2430 key: key.clone(),
2431 deadline_at: *deadline_at,
2432 at: Utc::now(),
2433 }));
2434 }
2435 _ => {}
2438 }
2439 }
2440 Err(EngineError::HumanInputRequired {
2441 run_id: input_run_id,
2442 step_id,
2443 ref message,
2444 }) => {
2445 final_status = RunStatus::AwaitingApproval;
2446 final_run = self
2447 .store
2448 .update_run_returning(
2449 run_id,
2450 RunUpdate {
2451 status: Some(RunStatus::AwaitingApproval),
2452 cost_usd: Some(ctx.total_cost_usd()),
2453 duration_ms: Some(total_duration),
2454 ..RunUpdate::default()
2455 },
2456 )
2457 .await?;
2458
2459 info!(
2461 run_id = %input_run_id,
2462 step_id = %step_id,
2463 message = %message,
2464 "run awaiting human input"
2465 );
2466 }
2467 Err(EngineError::DelaySleeping {
2468 run_id: delay_run_id,
2469 step_id,
2470 wake_at,
2471 }) => {
2472 final_status = RunStatus::Sleeping;
2473 final_run = self
2474 .store
2475 .update_run_returning(
2476 run_id,
2477 RunUpdate {
2478 status: Some(RunStatus::Sleeping),
2479 cost_usd: Some(ctx.total_cost_usd()),
2480 duration_ms: Some(total_duration),
2481 scheduled_at: Some(wake_at),
2482 ..RunUpdate::default()
2483 },
2484 )
2485 .await?;
2486
2487 info!(
2488 run_id = %delay_run_id,
2489 step_id = %step_id,
2490 wake_at = %wake_at,
2491 "run sleeping until delay elapses"
2492 );
2493 }
2494 Err(EngineError::CapacitySleeping {
2495 run_id: capacity_run_id,
2496 step_id,
2497 ref kind,
2498 wake_at,
2499 }) => {
2500 final_status = RunStatus::Sleeping;
2501 final_run = self
2502 .store
2503 .update_run_returning(
2504 run_id,
2505 RunUpdate {
2506 status: Some(RunStatus::Sleeping),
2507 cost_usd: Some(ctx.total_cost_usd()),
2508 duration_ms: Some(total_duration),
2509 scheduled_at: Some(wake_at),
2510 capacity_wait_kind: Some(ProviderKind::new(kind.as_str())),
2511 ..RunUpdate::default()
2512 },
2513 )
2514 .await?;
2515
2516 info!(
2517 run_id = %capacity_run_id,
2518 step_id = %step_id,
2519 kind = %kind,
2520 wake_at = %wake_at,
2521 "run sleeping until provider capacity returns"
2522 );
2523 }
2524 Err(EngineError::SignalWaiting {
2525 run_id: wait_run_id,
2526 step_id,
2527 ref step_name,
2528 ref name,
2529 ref key,
2530 deadline_at,
2531 }) => {
2532 final_status = RunStatus::Sleeping;
2533 let waiting = self
2537 .store
2538 .suspend_run_on_signal(run_id, step_id, deadline_at)
2539 .await?;
2540 final_run = self
2541 .store
2542 .update_run_returning(
2543 run_id,
2544 RunUpdate {
2545 cost_usd: Some(ctx.total_cost_usd()),
2546 duration_ms: Some(total_duration),
2547 ..RunUpdate::default()
2548 },
2549 )
2550 .await?;
2551
2552 if waiting {
2553 self.event_publisher
2554 .publish(Event::SignalAwaited(SignalAwaitedEvent {
2555 run_id: wait_run_id,
2556 step_id,
2557 step_name: step_name.clone(),
2558 name: name.clone(),
2559 key: key.clone(),
2560 deadline_at,
2561 at: Utc::now(),
2562 }));
2563 }
2564
2565 info!(
2566 run_id = %wait_run_id,
2567 step_id = %step_id,
2568 signal = %name,
2569 key = %key,
2570 deadline_at = %deadline_at,
2571 waiting,
2572 "run sleeping until a signal arrives"
2573 );
2574 }
2575 Err(err) => {
2576 let guardrail_stop = matches!(
2580 err,
2581 EngineError::RunBudgetExceeded { .. } | EngineError::WorkflowGuardRejected(_)
2582 );
2583
2584 final_status = if guardrail_stop {
2585 if let Err(store_err) = self
2586 .store
2587 .update_run(
2588 run_id,
2589 RunUpdate {
2590 status: Some(RunStatus::Cancelled),
2591 error: Some(err.to_string()),
2592 cost_usd: Some(ctx.total_cost_usd()),
2593 duration_ms: Some(total_duration),
2594 completed_at: Some(completed_at),
2595 output: ctx.output().cloned(),
2596 ..RunUpdate::default()
2597 },
2598 )
2599 .await
2600 {
2601 error!(run_id = %run_id, store_error = %store_err, "failed to persist run cancellation");
2602 }
2603 if let Err(cleanup_err) = self
2604 .fail_orphaned_steps(run_id, "run stopped: guardrail limit reached")
2605 .await
2606 {
2607 error!(run_id = %run_id, store_error = %cleanup_err, "failed to cleanup orphaned steps");
2608 }
2609 RunStatus::Cancelled
2610 } else {
2611 if let Some(output) = ctx.output()
2614 && let Err(store_err) = self
2615 .store
2616 .update_run(
2617 run_id,
2618 RunUpdate {
2619 output: Some(output.clone()),
2620 ..RunUpdate::default()
2621 },
2622 )
2623 .await
2624 {
2625 error!(run_id = %run_id, store_error = %store_err, "failed to persist run output");
2626 }
2627 self.fail_or_schedule_retry(
2628 run_id,
2629 &err.to_string(),
2630 is_run_retryable(&err),
2631 Some(ctx.total_cost_usd()),
2632 Some(total_duration),
2633 )
2634 .await
2635 .unwrap_or_else(|store_err| {
2636 error!(run_id = %run_id, store_error = %store_err, "failed to persist run failure");
2637 RunStatus::Failed
2638 })
2639 };
2640
2641 if matches!(err, EngineError::RunBudgetExceeded { .. }) {
2642 self.on_run_budget_exceeded(workflow_name, run_id, &err);
2643 }
2644
2645 error!(run_id = %run_id, status = %final_status, error = %err, "run stopped");
2646
2647 self.publish_run_status_changed(
2648 workflow_name,
2649 run_id,
2650 final_status,
2651 Some(err.to_string()),
2652 ctx,
2653 total_duration,
2654 run_labels,
2655 );
2656
2657 #[cfg(feature = "prometheus")]
2658 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
2659
2660 return Err(err);
2661 }
2662 }
2663
2664 self.publish_run_status_changed(
2665 workflow_name,
2666 run_id,
2667 final_status,
2668 None,
2669 ctx,
2670 total_duration,
2671 run_labels,
2672 );
2673
2674 #[cfg(feature = "prometheus")]
2675 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
2676
2677 Ok(WorkflowResult {
2678 run: final_run,
2679 steps: ctx.step_results().to_vec(),
2680 })
2681 }
2682
2683 async fn publish_approval_requested(
2686 &self,
2687 run_id: Uuid,
2688 step_id: Uuid,
2689 message: &str,
2690 ) -> Result<(), EngineError> {
2691 let requirement = self
2692 .store
2693 .get_step(step_id)
2694 .await?
2695 .and_then(|s| s.approval_requirement);
2696 self.event_publisher
2697 .publish(Event::ApprovalRequested(ApprovalRequestedEvent {
2698 run_id,
2699 step_id,
2700 message: message.to_string(),
2701 requirement,
2702 at: Utc::now(),
2703 }));
2704 Ok(())
2705 }
2706
2707 pub async fn fail_ancestors(&self, run_id: Uuid, reason: &str) -> Result<(), EngineError> {
2740 let mut current = self
2741 .store
2742 .get_run(run_id)
2743 .await?
2744 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2745 let mut visited = HashSet::from([run_id]);
2747
2748 while let Some(parent_id) = chain_parent(¤t) {
2749 if !visited.insert(parent_id) {
2750 break;
2751 }
2752 let status = self
2753 .fail_or_schedule_retry(parent_id, reason, false, None, None)
2754 .await?;
2755 if status == RunStatus::Paused {
2758 self.requeue_paused_root(¤t).await?;
2759 break;
2760 }
2761 info!(
2762 run_id = %run_id,
2763 ancestor_run_id = %parent_id,
2764 status = %status,
2765 "ancestor run failed with its child"
2766 );
2767 current = self
2768 .store
2769 .get_run(parent_id)
2770 .await?
2771 .ok_or(EngineError::Store(StoreError::RunNotFound(parent_id)))?;
2772 }
2773
2774 Ok(())
2775 }
2776
2777 #[cfg(feature = "prometheus")]
2779 fn emit_run_metrics(
2780 &self,
2781 workflow_name: &str,
2782 status: RunStatus,
2783 duration_ms: u64,
2784 ctx: &WorkflowContext,
2785 ) {
2786 let status_str = status.to_string();
2787 let wf = workflow_name.to_string();
2788
2789 counter!(RUNS_TOTAL, "workflow" => wf.clone(), "status" => status_str.clone()).increment(1);
2790 histogram!(RUN_DURATION_SECONDS, "workflow" => wf.clone(), "status" => status_str)
2791 .record(duration_ms as f64 / 1000.0);
2792 histogram!(RUN_COST_USD, "workflow" => wf.clone()).record(
2793 ctx.total_cost_usd()
2794 .to_string()
2795 .parse::<f64>()
2796 .unwrap_or(0.0),
2797 );
2798 gauge!(RUNS_ACTIVE, "workflow" => wf).decrement(1.0);
2799 }
2800
2801 fn on_run_budget_exceeded(&self, workflow_name: &str, run_id: Uuid, err: &EngineError) {
2807 let EngineError::RunBudgetExceeded {
2808 limit_usd,
2809 spent_usd,
2810 step_budget_usd,
2811 ..
2812 } = err
2813 else {
2814 return;
2815 };
2816
2817 #[cfg(feature = "prometheus")]
2818 counter!(
2819 RUN_BUDGET_EXCEEDED_TOTAL,
2820 "workflow" => workflow_name.to_string(),
2821 "scope" => "run",
2822 )
2823 .increment(1);
2824
2825 self.event_publisher
2826 .publish(Event::RunBudgetExceeded(RunBudgetExceededEvent {
2827 run_id,
2828 workflow_name: workflow_name.to_string(),
2829 limit_usd: *limit_usd,
2830 spent_usd: *spent_usd,
2831 step_budget_usd: *step_budget_usd,
2832 at: Utc::now(),
2833 }));
2834 }
2835
2836 #[allow(clippy::too_many_arguments)]
2841 fn publish_run_status_changed(
2842 &self,
2843 workflow_name: &str,
2844 run_id: Uuid,
2845 to: RunStatus,
2846 error: Option<String>,
2847 ctx: &WorkflowContext,
2848 duration_ms: u64,
2849 labels: HashMap<String, String>,
2850 ) {
2851 let now = Utc::now();
2852 let cost_usd = ctx.total_cost_usd();
2853 let wf = workflow_name.to_string();
2854
2855 self.event_publisher
2856 .publish(Event::RunStatusChanged(RunStatusChangedEvent {
2857 run_id,
2858 workflow_name: wf.clone(),
2859 from: RunStatus::Running,
2860 to,
2861 error: error.clone(),
2862 cost_usd,
2863 duration_ms,
2864 labels: labels.clone(),
2865 at: now,
2866 }));
2867
2868 if to == RunStatus::Failed {
2869 self.event_publisher
2870 .publish(Event::RunFailed(RunFailedEvent {
2871 run_id,
2872 workflow_name: wf,
2873 error,
2874 cost_usd,
2875 duration_ms,
2876 labels,
2877 at: now,
2878 }));
2879 }
2880 }
2881}
2882
2883impl fmt::Debug for Engine {
2884 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2885 f.debug_struct("Engine")
2886 .field("handlers", &self.handlers.keys().collect::<Vec<_>>())
2887 .finish_non_exhaustive()
2888 }
2889}
2890
2891#[cfg(test)]
2892mod tests {
2893 use super::*;
2894 use crate::config::ShellConfig;
2895 use crate::handler::{HandlerFuture, WorkflowHandler};
2896 use ironflow_core::providers::claude::ClaudeCodeProvider;
2897 use ironflow_core::providers::record_replay::RecordReplayProvider;
2898 use ironflow_store::memory::InMemoryStore;
2899 use ironflow_store::models::{MAX_PRIORITY, MIN_PRIORITY, StepApproval, StepStatus};
2900 use serde_json::json;
2901
2902 struct EchoWorkflow;
2904
2905 impl WorkflowHandler for EchoWorkflow {
2906 fn name(&self) -> &str {
2907 "echo-workflow"
2908 }
2909
2910 fn describe(&self) -> WorkflowInfo {
2911 WorkflowInfo {
2912 description: "A simple workflow that echoes hello".to_string(),
2913 source_code: None,
2914 sub_workflows: Vec::new(),
2915 category: None,
2916 version: self.version().map(str::to_string),
2917 compatible_versions: Vec::new(),
2918 input_schema: None,
2919 default_labels: HashMap::new(),
2920 schedule: self.schedule().cloned(),
2921 default_max_cost_usd: self.default_max_cost_usd(),
2922 priority: self.priority(),
2923 }
2924 }
2925
2926 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2927 Box::pin(async move {
2928 ctx.shell("greet", ShellConfig::new("echo hello")).await?;
2929 Ok(())
2930 })
2931 }
2932 }
2933
2934 struct FailingWorkflow;
2936
2937 impl WorkflowHandler for FailingWorkflow {
2938 fn name(&self) -> &str {
2939 "failing-workflow"
2940 }
2941
2942 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2943 Box::pin(async move {
2944 ctx.shell("fail", ShellConfig::new("exit 1")).await?;
2945 Ok(())
2946 })
2947 }
2948 }
2949
2950 fn create_test_engine() -> Engine {
2951 let store = Arc::new(InMemoryStore::new());
2952 let inner = ClaudeCodeProvider::new();
2953 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
2954 inner,
2955 "/tmp/ironflow-fixtures",
2956 ));
2957 Engine::new(store, provider)
2958 }
2959
2960 #[test]
2961 fn engine_new_creates_instance() {
2962 let engine = create_test_engine();
2963 assert_eq!(engine.handler_names().len(), 0);
2964 }
2965
2966 #[test]
2967 fn execution_mode_defaults_to_local() {
2968 let engine = create_test_engine();
2969 assert_eq!(engine.execution_mode(), ExecutionMode::Local);
2970 }
2971
2972 #[test]
2973 fn with_execution_mode_overrides_the_default() {
2974 let engine = create_test_engine().with_execution_mode(ExecutionMode::Workers);
2975 assert_eq!(engine.execution_mode(), ExecutionMode::Workers);
2976 }
2977
2978 #[test]
2979 fn engine_register_handler() {
2980 let mut engine = create_test_engine();
2981 let result = engine.register(EchoWorkflow);
2982 assert!(result.is_ok());
2983 assert_eq!(engine.handler_names().len(), 1);
2984 assert!(engine.handler_names().contains(&"echo-workflow"));
2985 }
2986
2987 #[test]
2988 fn engine_register_duplicate_returns_error() {
2989 let mut engine = create_test_engine();
2990 engine.register(EchoWorkflow).unwrap();
2991 let result = engine.register(EchoWorkflow);
2992 assert!(result.is_err());
2993 }
2994
2995 #[test]
2996 fn engine_get_handler_found() {
2997 let mut engine = create_test_engine();
2998 engine.register(EchoWorkflow).unwrap();
2999 let handler = engine.get_handler("echo-workflow");
3000 assert!(handler.is_some());
3001 }
3002
3003 #[test]
3004 fn engine_get_handler_not_found() {
3005 let engine = create_test_engine();
3006 let handler = engine.get_handler("nonexistent");
3007 assert!(handler.is_none());
3008 }
3009
3010 #[test]
3011 fn engine_handler_names_lists_all() {
3012 let mut engine = create_test_engine();
3013 engine.register(EchoWorkflow).unwrap();
3014 engine.register(FailingWorkflow).unwrap();
3015 let names = engine.handler_names();
3016 assert_eq!(names.len(), 2);
3017 assert!(names.contains(&"echo-workflow"));
3018 assert!(names.contains(&"failing-workflow"));
3019 }
3020
3021 #[test]
3022 fn engine_handler_info_returns_description() {
3023 let mut engine = create_test_engine();
3024 engine.register(EchoWorkflow).unwrap();
3025 let info = engine.handler_info("echo-workflow");
3026 assert!(info.is_some());
3027 let info = info.unwrap();
3028 assert_eq!(info.description, "A simple workflow that echoes hello");
3029 }
3030
3031 struct CategorizedWorkflow;
3032
3033 impl WorkflowHandler for CategorizedWorkflow {
3034 fn name(&self) -> &str {
3035 "categorized"
3036 }
3037 fn category(&self) -> Option<&str> {
3038 Some("data/etl")
3039 }
3040 fn execute<'a>(
3041 &'a self,
3042 _ctx: &'a mut WorkflowContext,
3043 ) -> crate::handler::HandlerFuture<'a> {
3044 Box::pin(async move { Ok(()) })
3045 }
3046 }
3047
3048 #[test]
3049 fn engine_default_describe_propagates_category() {
3050 let mut engine = create_test_engine();
3051 engine.register(CategorizedWorkflow).unwrap();
3052 let info = engine.handler_info("categorized").unwrap();
3053 assert_eq!(info.category.as_deref(), Some("data/etl"));
3054 }
3055
3056 #[test]
3057 fn engine_default_describe_without_category() {
3058 let mut engine = create_test_engine();
3059 engine.register(EchoWorkflow).unwrap();
3060 let info = engine.handler_info("echo-workflow").unwrap();
3061 assert!(info.category.is_none());
3062 }
3063
3064 struct ScheduledWorkflow {
3069 schedule: CronSchedule,
3070 }
3071
3072 impl ScheduledWorkflow {
3073 fn new() -> Self {
3074 Self {
3075 schedule: CronSchedule::new("0 0 * * * *").unwrap(),
3076 }
3077 }
3078 }
3079
3080 impl WorkflowHandler for ScheduledWorkflow {
3081 fn name(&self) -> &str {
3082 "scheduled"
3083 }
3084 fn schedule(&self) -> Option<&CronSchedule> {
3085 Some(&self.schedule)
3086 }
3087 fn execute<'a>(
3088 &'a self,
3089 _ctx: &'a mut WorkflowContext,
3090 ) -> crate::handler::HandlerFuture<'a> {
3091 Box::pin(async move { Ok(()) })
3092 }
3093 }
3094
3095 #[test]
3096 fn engine_default_describe_propagates_schedule() {
3097 let mut engine = create_test_engine();
3098 engine.register(ScheduledWorkflow::new()).unwrap();
3099 let info = engine.handler_info("scheduled").unwrap();
3100 assert_eq!(
3101 info.schedule.as_ref().map(|s| s.as_str()),
3102 Some("0 0 * * * *")
3103 );
3104 }
3105
3106 #[test]
3107 fn engine_default_describe_without_schedule() {
3108 let mut engine = create_test_engine();
3109 engine.register(EchoWorkflow).unwrap();
3110 let info = engine.handler_info("echo-workflow").unwrap();
3111 assert!(info.schedule.is_none());
3112 }
3113
3114 #[test]
3115 fn scheduled_handlers_returns_only_scheduled() {
3116 let mut engine = create_test_engine();
3117 engine.register(EchoWorkflow).unwrap();
3118 engine.register(ScheduledWorkflow::new()).unwrap();
3119 engine.register(FailingWorkflow).unwrap();
3120
3121 let scheduled = engine.scheduled_handlers();
3122 assert_eq!(scheduled.len(), 1);
3123 assert_eq!(scheduled[0].0, "scheduled");
3124 assert_eq!(scheduled[0].1.as_str(), "0 0 * * * *");
3125 }
3126
3127 #[test]
3128 fn scheduled_handlers_empty_when_none_scheduled() {
3129 let mut engine = create_test_engine();
3130 engine.register(EchoWorkflow).unwrap();
3131 engine.register(FailingWorkflow).unwrap();
3132
3133 let scheduled = engine.scheduled_handlers();
3134 assert!(scheduled.is_empty());
3135 }
3136
3137 struct BadCategoryWorkflow(&'static str);
3138
3139 impl WorkflowHandler for BadCategoryWorkflow {
3140 fn name(&self) -> &str {
3141 "bad-category"
3142 }
3143 fn category(&self) -> Option<&str> {
3144 Some(self.0)
3145 }
3146 fn execute<'a>(
3147 &'a self,
3148 _ctx: &'a mut WorkflowContext,
3149 ) -> crate::handler::HandlerFuture<'a> {
3150 Box::pin(async move { Ok(()) })
3151 }
3152 }
3153
3154 #[test]
3155 fn engine_register_rejects_empty_category() {
3156 let mut engine = create_test_engine();
3157 let err = engine.register(BadCategoryWorkflow("")).unwrap_err();
3158 match err {
3159 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty category")),
3160 other => panic!("expected InvalidWorkflow, got {other:?}"),
3161 }
3162 }
3163
3164 #[test]
3165 fn engine_register_rejects_leading_slash_category() {
3166 let mut engine = create_test_engine();
3167 let err = engine
3168 .register(BadCategoryWorkflow("/data/etl"))
3169 .unwrap_err();
3170 match err {
3171 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("leading '/'")),
3172 other => panic!("expected InvalidWorkflow, got {other:?}"),
3173 }
3174 }
3175
3176 #[test]
3177 fn engine_register_rejects_trailing_slash_category() {
3178 let mut engine = create_test_engine();
3179 let err = engine
3180 .register(BadCategoryWorkflow("data/etl/"))
3181 .unwrap_err();
3182 match err {
3183 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("trailing '/'")),
3184 other => panic!("expected InvalidWorkflow, got {other:?}"),
3185 }
3186 }
3187
3188 #[test]
3189 fn engine_register_rejects_double_slash_category() {
3190 let mut engine = create_test_engine();
3191 let err = engine
3192 .register(BadCategoryWorkflow("data//etl"))
3193 .unwrap_err();
3194 match err {
3195 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty segment")),
3196 other => panic!("expected InvalidWorkflow, got {other:?}"),
3197 }
3198 }
3199
3200 #[test]
3201 fn engine_register_rejects_whitespace_only_segment_category() {
3202 let mut engine = create_test_engine();
3203 let err = engine
3204 .register(BadCategoryWorkflow("data/ /etl"))
3205 .unwrap_err();
3206 match err {
3207 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("whitespace-only segment")),
3208 other => panic!("expected InvalidWorkflow, got {other:?}"),
3209 }
3210 }
3211
3212 #[test]
3213 fn engine_register_accepts_valid_nested_category() {
3214 let mut engine = create_test_engine();
3215 assert!(engine.register(CategorizedWorkflow).is_ok());
3216 }
3217
3218 #[tokio::test]
3219 async fn engine_unknown_workflow_returns_error() {
3220 let engine = create_test_engine();
3221 let result = engine
3222 .run_handler("unknown", TriggerKind::Manual, json!({}))
3223 .await;
3224 assert!(result.is_err());
3225 match result {
3226 Err(EngineError::InvalidWorkflow(msg)) => {
3227 assert!(msg.contains("no handler registered"));
3228 }
3229 _ => panic!("expected InvalidWorkflow error"),
3230 }
3231 }
3232
3233 #[tokio::test]
3234 async fn engine_enqueue_handler_creates_pending_run() {
3235 let mut engine = create_test_engine();
3236 engine.register(EchoWorkflow).unwrap();
3237
3238 let run = engine
3239 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
3240 .await
3241 .unwrap();
3242 assert_eq!(run.status.state, RunStatus::Pending);
3243 assert_eq!(run.workflow_name, "echo-workflow");
3244 }
3245
3246 #[tokio::test]
3247 async fn enqueue_handler_leaves_the_run_unattributed() {
3248 let mut engine = create_test_engine();
3249 engine.register(EchoWorkflow).unwrap();
3250
3251 let run = engine
3252 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
3253 .await
3254 .unwrap();
3255
3256 assert!(run.created_by.is_none());
3257 }
3258
3259 #[tokio::test]
3260 async fn enqueue_handler_with_options_records_the_author() {
3261 let mut engine = create_test_engine();
3262 engine.register(EchoWorkflow).unwrap();
3263 let actor = RunActor::User {
3264 user_id: Uuid::now_v7(),
3265 };
3266
3267 let run = engine
3268 .enqueue_handler_with_options(
3269 "echo-workflow",
3270 TriggerKind::Api,
3271 json!({}),
3272 EnqueueOptions {
3273 created_by: Some(actor.clone()),
3274 ..Default::default()
3275 },
3276 )
3277 .await
3278 .unwrap()
3279 .into_run();
3280
3281 assert_eq!(run.created_by, Some(actor));
3282 }
3283
3284 #[tokio::test]
3285 async fn enqueue_handler_with_options_accepts_no_author() {
3286 let mut engine = create_test_engine();
3287 engine.register(EchoWorkflow).unwrap();
3288
3289 let run = engine
3290 .enqueue_handler_with_options(
3291 "echo-workflow",
3292 TriggerKind::Cron {
3293 schedule: "0 * * * * *".to_string(),
3294 schedule_id: None,
3295 scheduled_for: None,
3296 },
3297 json!({}),
3298 EnqueueOptions::default(),
3299 )
3300 .await
3301 .unwrap()
3302 .into_run();
3303
3304 assert!(run.created_by.is_none());
3305 }
3306
3307 #[tokio::test]
3308 async fn enqueue_handler_with_options_stores_concurrency_limits() {
3309 let mut engine = create_test_engine();
3310 engine.register(EchoWorkflow).unwrap();
3311 let limits = vec![
3312 ConcurrencyLimit::new("repo:acme", 2),
3313 ConcurrencyLimit::new("tenant:42", 5),
3314 ];
3315
3316 let run = engine
3317 .enqueue_handler_with_options(
3318 "echo-workflow",
3319 TriggerKind::Api,
3320 json!({}),
3321 EnqueueOptions {
3322 concurrency_limits: limits.clone(),
3323 ..Default::default()
3324 },
3325 )
3326 .await
3327 .unwrap()
3328 .into_run();
3329
3330 assert_eq!(run.concurrency_limits, limits);
3331 }
3332
3333 #[tokio::test]
3334 async fn enqueue_rejects_invalid_concurrency_limits() {
3335 let mut engine = create_test_engine();
3336 engine.register(EchoWorkflow).unwrap();
3337
3338 let invalid = [
3339 vec![ConcurrencyLimit::new("repo:acme", 0)],
3340 vec![ConcurrencyLimit::new("", 1)],
3341 vec![
3342 ConcurrencyLimit::new("repo:acme", 1),
3343 ConcurrencyLimit::new("repo:acme", 2),
3344 ],
3345 ];
3346 for concurrency_limits in invalid {
3347 let err = engine
3348 .enqueue_handler_with_options(
3349 "echo-workflow",
3350 TriggerKind::Api,
3351 json!({}),
3352 EnqueueOptions {
3353 concurrency_limits,
3354 ..Default::default()
3355 },
3356 )
3357 .await
3358 .unwrap_err();
3359 assert!(
3360 matches!(err, EngineError::InvalidConcurrencyLimit(_)),
3361 "{err:?}"
3362 );
3363 }
3364
3365 let err = engine
3367 .enqueue_handler_with_options(
3368 "not-registered",
3369 TriggerKind::Api,
3370 json!({}),
3371 EnqueueOptions {
3372 concurrency_limits: vec![ConcurrencyLimit::new("repo:acme", 0)],
3373 ..Default::default()
3374 },
3375 )
3376 .await
3377 .unwrap_err();
3378 assert!(
3379 matches!(err, EngineError::InvalidConcurrencyLimit(_)),
3380 "{err:?}"
3381 );
3382
3383 let page = engine
3384 .store()
3385 .list_runs(RunFilter::default(), 1, 10)
3386 .await
3387 .unwrap();
3388 assert_eq!(page.total, 0, "no run may be created");
3389 }
3390
3391 struct UrgentWorkflow;
3392
3393 impl WorkflowHandler for UrgentWorkflow {
3394 fn name(&self) -> &str {
3395 "urgent-workflow"
3396 }
3397
3398 fn priority(&self) -> i16 {
3399 60
3400 }
3401
3402 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3403 Box::pin(async { Ok(()) })
3404 }
3405 }
3406
3407 #[tokio::test]
3408 async fn enqueue_priority_defaults_to_the_handler_priority() {
3409 let mut engine = create_test_engine();
3410 engine.register(EchoWorkflow).unwrap();
3411 engine.register(UrgentWorkflow).unwrap();
3412
3413 let echo = engine
3414 .enqueue_handler_with_options(
3415 "echo-workflow",
3416 TriggerKind::Api,
3417 json!({}),
3418 EnqueueOptions::default(),
3419 )
3420 .await
3421 .unwrap()
3422 .into_run();
3423 assert_eq!(echo.priority, 0);
3424
3425 let urgent = engine
3426 .enqueue_handler_with_options(
3427 "urgent-workflow",
3428 TriggerKind::Api,
3429 json!({}),
3430 EnqueueOptions::default(),
3431 )
3432 .await
3433 .unwrap()
3434 .into_run();
3435 assert_eq!(urgent.priority, 60);
3436 }
3437
3438 #[tokio::test]
3439 async fn enqueue_priority_explicit_value_overrides_the_handler() {
3440 let mut engine = create_test_engine();
3441 engine.register(UrgentWorkflow).unwrap();
3442
3443 let run = engine
3444 .enqueue_handler_with_options(
3445 "urgent-workflow",
3446 TriggerKind::Api,
3447 json!({}),
3448 EnqueueOptions {
3449 priority: Some(-20),
3450 ..Default::default()
3451 },
3452 )
3453 .await
3454 .unwrap()
3455 .into_run();
3456 assert_eq!(run.priority, -20);
3457
3458 let stored = engine.store().get_run(run.id).await.unwrap().unwrap();
3459 assert_eq!(stored.priority, -20);
3460 }
3461
3462 #[tokio::test]
3463 async fn enqueue_priority_out_of_range_is_rejected() {
3464 let mut engine = create_test_engine();
3465 engine.register(EchoWorkflow).unwrap();
3466
3467 for priority in [MAX_PRIORITY + 1, MIN_PRIORITY - 1] {
3468 let err = engine
3469 .enqueue_handler_with_options(
3470 "echo-workflow",
3471 TriggerKind::Api,
3472 json!({}),
3473 EnqueueOptions {
3474 priority: Some(priority),
3475 ..Default::default()
3476 },
3477 )
3478 .await
3479 .unwrap_err();
3480 assert!(matches!(err, EngineError::InvalidPriority(_)), "{err:?}");
3481 }
3482
3483 let err = engine
3485 .enqueue_handler_with_options(
3486 "not-registered",
3487 TriggerKind::Api,
3488 json!({}),
3489 EnqueueOptions {
3490 priority: Some(MAX_PRIORITY + 1),
3491 ..Default::default()
3492 },
3493 )
3494 .await
3495 .unwrap_err();
3496 assert!(matches!(err, EngineError::InvalidPriority(_)), "{err:?}");
3497
3498 let page = engine
3499 .store()
3500 .list_runs(RunFilter::default(), 1, 10)
3501 .await
3502 .unwrap();
3503 assert_eq!(page.total, 0, "no run may be created");
3504 }
3505
3506 #[tokio::test]
3507 async fn enqueue_priority_bounds_are_accepted() {
3508 let mut engine = create_test_engine();
3509 engine.register(EchoWorkflow).unwrap();
3510
3511 for priority in [MIN_PRIORITY, MAX_PRIORITY] {
3512 let run = engine
3513 .enqueue_handler_with_options(
3514 "echo-workflow",
3515 TriggerKind::Api,
3516 json!({}),
3517 EnqueueOptions {
3518 priority: Some(priority),
3519 ..Default::default()
3520 },
3521 )
3522 .await
3523 .unwrap()
3524 .into_run();
3525 assert_eq!(run.priority, priority);
3526 }
3527 }
3528
3529 struct GpuWorkflow;
3530
3531 impl WorkflowHandler for GpuWorkflow {
3532 fn name(&self) -> &str {
3533 "gpu-workflow"
3534 }
3535
3536 fn required_worker_tags(&self) -> Vec<String> {
3537 vec!["gpu".to_string()]
3538 }
3539
3540 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3541 Box::pin(async move { Ok(()) })
3542 }
3543 }
3544
3545 #[tokio::test]
3546 async fn enqueue_merges_handler_and_request_worker_tags() {
3547 let mut engine = create_test_engine();
3548 engine.register(GpuWorkflow).unwrap();
3549
3550 let run = engine
3551 .enqueue_handler_with_options(
3552 "gpu-workflow",
3553 TriggerKind::Api,
3554 json!({}),
3555 EnqueueOptions {
3556 worker_tags: vec!["region:eu".to_string(), "gpu".to_string()],
3557 ..Default::default()
3558 },
3559 )
3560 .await
3561 .unwrap()
3562 .into_run();
3563
3564 assert_eq!(
3565 run.worker_tags,
3566 vec!["gpu".to_string(), "region:eu".to_string()]
3567 );
3568 }
3569
3570 #[tokio::test]
3571 async fn enqueue_without_worker_tags_keeps_handler_tags() {
3572 let mut engine = create_test_engine();
3573 engine.register(GpuWorkflow).unwrap();
3574 engine.register(EchoWorkflow).unwrap();
3575
3576 let gpu = engine
3577 .enqueue_handler("gpu-workflow", TriggerKind::Api, json!({}), 0)
3578 .await
3579 .unwrap();
3580 assert_eq!(gpu.worker_tags, vec!["gpu".to_string()]);
3581
3582 let echo = engine
3583 .enqueue_handler("echo-workflow", TriggerKind::Api, json!({}), 0)
3584 .await
3585 .unwrap();
3586 assert!(echo.worker_tags.is_empty());
3587 }
3588
3589 #[tokio::test]
3590 async fn enqueue_rejects_invalid_worker_tags() {
3591 let mut engine = create_test_engine();
3592 engine.register(EchoWorkflow).unwrap();
3593
3594 for worker_tags in [
3595 vec!["bad,tag".to_string()],
3596 vec![" ".to_string()],
3597 vec!["x".repeat(65)],
3598 ] {
3599 let err = engine
3600 .enqueue_handler_with_options(
3601 "echo-workflow",
3602 TriggerKind::Api,
3603 json!({}),
3604 EnqueueOptions {
3605 worker_tags,
3606 ..Default::default()
3607 },
3608 )
3609 .await
3610 .unwrap_err();
3611 assert!(matches!(err, EngineError::InvalidWorkerTag(_)), "{err:?}");
3612 }
3613
3614 let err = engine
3616 .enqueue_handler_with_options(
3617 "not-registered",
3618 TriggerKind::Api,
3619 json!({}),
3620 EnqueueOptions {
3621 worker_tags: vec!["bad,tag".to_string()],
3622 ..Default::default()
3623 },
3624 )
3625 .await
3626 .unwrap_err();
3627 assert!(matches!(err, EngineError::InvalidWorkerTag(_)), "{err:?}");
3628 }
3629
3630 #[test]
3631 fn worker_tags_are_unset_by_default() {
3632 let engine = create_test_engine();
3633 assert!(engine.worker_tags().is_none());
3634 }
3635
3636 #[test]
3637 fn set_worker_tags_stores_the_tags() {
3638 let mut engine = create_test_engine();
3639 engine.set_worker_tags(vec!["gpu".to_string()]);
3640 assert_eq!(engine.worker_tags(), Some(&["gpu".to_string()][..]));
3641
3642 engine.set_worker_tags(Vec::new());
3643 assert_eq!(engine.worker_tags(), Some(&[][..]));
3644 }
3645
3646 #[tokio::test]
3647 async fn run_handler_records_handler_worker_tags() {
3648 let mut engine = create_test_engine();
3649 engine.register(GpuWorkflow).unwrap();
3650
3651 let result = engine
3652 .run_handler("gpu-workflow", TriggerKind::Manual, json!({}))
3653 .await
3654 .unwrap();
3655 assert_eq!(result.run.worker_tags, vec!["gpu".to_string()]);
3656 }
3657
3658 #[tokio::test]
3659 async fn run_handler_leaves_the_run_unattributed() {
3660 let mut engine = create_test_engine();
3661 engine.register(EchoWorkflow).unwrap();
3662
3663 let run = engine
3664 .run_handler("echo-workflow", TriggerKind::Manual, json!({}))
3665 .await
3666 .unwrap()
3667 .run;
3668
3669 assert!(run.created_by.is_none());
3670 }
3671
3672 #[tokio::test]
3673 async fn run_handler_priority_comes_from_the_handler() {
3674 let mut engine = create_test_engine();
3675 engine.register(UrgentWorkflow).unwrap();
3676
3677 let run = engine
3678 .run_handler("urgent-workflow", TriggerKind::Manual, json!({}))
3679 .await
3680 .unwrap()
3681 .run;
3682
3683 assert_eq!(run.priority, 60);
3684 }
3685
3686 #[tokio::test]
3687 async fn engine_register_boxed() {
3688 let mut engine = create_test_engine();
3689 let handler: Box<dyn WorkflowHandler> = Box::new(EchoWorkflow);
3690 let result = engine.register_boxed(handler);
3691 assert!(result.is_ok());
3692 assert_eq!(engine.handler_names().len(), 1);
3693 }
3694
3695 #[tokio::test]
3696 async fn engine_store_and_provider_accessors() {
3697 let store = Arc::new(InMemoryStore::new());
3698 let inner = ClaudeCodeProvider::new();
3699 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
3700 inner,
3701 "/tmp/ironflow-fixtures",
3702 ));
3703 let engine = Engine::new(store.clone(), provider.clone());
3704
3705 let _ = engine.store();
3707 let _ = engine.provider();
3708 }
3709
3710 use crate::operation::{Operation, OperationContext};
3715 use async_trait::async_trait;
3716 use ironflow_core::error::OperationError;
3717 use ironflow_store::models::StepKind;
3718
3719 struct FakeGitlabOp {
3720 project_id: u64,
3721 title: String,
3722 }
3723
3724 #[async_trait]
3725 impl Operation for FakeGitlabOp {
3726 fn kind(&self) -> &str {
3727 "gitlab"
3728 }
3729
3730 async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
3731 Ok(json!({
3732 "issue_id": 42,
3733 "project_id": self.project_id,
3734 "title": self.title,
3735 }))
3736 }
3737
3738 fn input(&self) -> Option<Value> {
3739 Some(json!({
3740 "project_id": self.project_id,
3741 "title": self.title,
3742 }))
3743 }
3744 }
3745
3746 struct FailingOp;
3747
3748 #[async_trait]
3749 impl Operation for FailingOp {
3750 fn kind(&self) -> &str {
3751 "broken-service"
3752 }
3753
3754 async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
3755 Err(OperationError::Http {
3756 status: None,
3757 message: "service unavailable".to_string(),
3758 })
3759 }
3760 }
3761
3762 struct OperationWorkflow;
3763
3764 impl WorkflowHandler for OperationWorkflow {
3765 fn name(&self) -> &str {
3766 "operation-workflow"
3767 }
3768
3769 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3770 Box::pin(async move {
3771 let op = FakeGitlabOp {
3772 project_id: 123,
3773 title: "Bug report".to_string(),
3774 };
3775 ctx.operation("create-issue", &op).await?;
3776 Ok(())
3777 })
3778 }
3779 }
3780
3781 struct FailingOperationWorkflow;
3782
3783 impl WorkflowHandler for FailingOperationWorkflow {
3784 fn name(&self) -> &str {
3785 "failing-operation-workflow"
3786 }
3787
3788 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3789 Box::pin(async move {
3790 ctx.operation("broken-call", &FailingOp).await?;
3791 Ok(())
3792 })
3793 }
3794 }
3795
3796 struct MixedWorkflow;
3797
3798 impl WorkflowHandler for MixedWorkflow {
3799 fn name(&self) -> &str {
3800 "mixed-workflow"
3801 }
3802
3803 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3804 Box::pin(async move {
3805 ctx.shell("build", ShellConfig::new("echo built")).await?;
3806 let op = FakeGitlabOp {
3807 project_id: 456,
3808 title: "Deploy done".to_string(),
3809 };
3810 let result = ctx.operation("notify-gitlab", &op).await?;
3811 assert_eq!(result.output["issue_id"], 42);
3812 Ok(())
3813 })
3814 }
3815 }
3816
3817 #[tokio::test]
3818 async fn operation_step_happy_path() {
3819 let mut engine = create_test_engine();
3820 engine.register(OperationWorkflow).unwrap();
3821
3822 let run = engine
3823 .run_handler("operation-workflow", TriggerKind::Manual, json!({}))
3824 .await
3825 .unwrap()
3826 .run;
3827
3828 assert_eq!(run.status.state, RunStatus::Completed);
3829
3830 let steps = engine.store().list_steps(run.id).await.unwrap();
3831
3832 assert_eq!(steps.len(), 1);
3833 assert_eq!(steps[0].name, "create-issue");
3834 assert_eq!(steps[0].kind, StepKind::Custom("gitlab".to_string()));
3835 assert_eq!(
3836 steps[0].status.state,
3837 ironflow_store::models::StepStatus::Completed
3838 );
3839
3840 let output = steps[0].output.as_ref().unwrap();
3841 assert_eq!(output["issue_id"], 42);
3842 assert_eq!(output["project_id"], 123);
3843
3844 let input = steps[0].input.as_ref().unwrap();
3845 assert_eq!(input["project_id"], 123);
3846 assert_eq!(input["title"], "Bug report");
3847 }
3848
3849 #[tokio::test]
3850 async fn operation_step_failure_marks_run_failed() {
3851 let mut engine = create_test_engine();
3852 engine.register(FailingOperationWorkflow).unwrap();
3853
3854 let result = engine
3855 .run_handler("failing-operation-workflow", TriggerKind::Manual, json!({}))
3856 .await;
3857
3858 assert!(result.is_err());
3859 }
3860
3861 #[tokio::test]
3862 async fn operation_mixed_with_shell_steps() {
3863 let mut engine = create_test_engine();
3864 engine.register(MixedWorkflow).unwrap();
3865
3866 let run = engine
3867 .run_handler("mixed-workflow", TriggerKind::Manual, json!({}))
3868 .await
3869 .unwrap()
3870 .run;
3871
3872 assert_eq!(run.status.state, RunStatus::Completed);
3873
3874 let steps = engine.store().list_steps(run.id).await.unwrap();
3875
3876 assert_eq!(steps.len(), 2);
3877 assert_eq!(steps[0].kind, StepKind::Shell);
3878 assert_eq!(steps[1].kind, StepKind::Custom("gitlab".to_string()));
3879 assert_eq!(steps[0].position, 0);
3880 assert_eq!(steps[1].position, 1);
3881 }
3882
3883 use crate::config::ApprovalConfig;
3888
3889 struct SingleApprovalWorkflow;
3890
3891 impl WorkflowHandler for SingleApprovalWorkflow {
3892 fn name(&self) -> &str {
3893 "single-approval"
3894 }
3895
3896 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3897 Box::pin(async move {
3898 ctx.shell("build", ShellConfig::new("echo built")).await?;
3899 ctx.approval("gate", ApprovalConfig::new("OK?")).await?;
3900 ctx.shell("deploy", ShellConfig::new("echo deployed"))
3901 .await?;
3902 Ok(())
3903 })
3904 }
3905 }
3906
3907 struct DoubleApprovalWorkflow;
3908
3909 impl WorkflowHandler for DoubleApprovalWorkflow {
3910 fn name(&self) -> &str {
3911 "double-approval"
3912 }
3913
3914 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3915 Box::pin(async move {
3916 ctx.shell("build", ShellConfig::new("echo built")).await?;
3917 ctx.approval("staging-gate", ApprovalConfig::new("Deploy staging?"))
3918 .await?;
3919 ctx.shell("deploy-staging", ShellConfig::new("echo staging"))
3920 .await?;
3921 ctx.approval("prod-gate", ApprovalConfig::new("Deploy prod?"))
3922 .await?;
3923 ctx.shell("deploy-prod", ShellConfig::new("echo prod"))
3924 .await?;
3925 Ok(())
3926 })
3927 }
3928 }
3929
3930 #[tokio::test]
3931 async fn approval_pauses_run() {
3932 let mut engine = create_test_engine();
3933 engine.register(SingleApprovalWorkflow).unwrap();
3934
3935 let run = engine
3936 .run_handler("single-approval", TriggerKind::Manual, json!({}))
3937 .await
3938 .unwrap()
3939 .run;
3940
3941 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
3942
3943 let steps = engine.store().list_steps(run.id).await.unwrap();
3944 assert_eq!(steps.len(), 2); assert_eq!(steps[0].kind, StepKind::Shell);
3946 assert_eq!(steps[0].status.state, StepStatus::Completed);
3947 assert_eq!(steps[1].kind, StepKind::Approval);
3948 assert_eq!(steps[1].status.state, StepStatus::AwaitingApproval);
3949 }
3950
3951 async fn approve_open_gate(store: &Arc<dyn Store>, run_id: Uuid) {
3954 let gate = store
3955 .list_steps(run_id)
3956 .await
3957 .unwrap()
3958 .into_iter()
3959 .find(|step| step.status.state == StepStatus::AwaitingApproval)
3960 .expect("an open approval gate");
3961 store
3962 .record_step_approval(
3963 gate.id,
3964 StepApproval {
3965 user_id: Uuid::now_v7(),
3966 approved_by: "tester".to_string(),
3967 at: Utc::now(),
3968 },
3969 )
3970 .await
3971 .unwrap();
3972 store
3973 .update_run_status(run_id, RunStatus::Running)
3974 .await
3975 .unwrap();
3976 }
3977
3978 #[tokio::test]
3979 async fn approval_resume_completes_run() {
3980 let mut engine = create_test_engine();
3981 engine.register(SingleApprovalWorkflow).unwrap();
3982
3983 let run = engine
3985 .run_handler("single-approval", TriggerKind::Manual, json!({}))
3986 .await
3987 .unwrap()
3988 .run;
3989 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
3990
3991 approve_open_gate(engine.store(), run.id).await;
3993
3994 let resumed = engine.resume_run(run.id).await.unwrap().run;
3996 assert_eq!(resumed.status.state, RunStatus::Completed);
3997
3998 let steps = engine.store().list_steps(run.id).await.unwrap();
3999 assert_eq!(steps.len(), 3); assert_eq!(steps[0].name, "build");
4001 assert_eq!(steps[0].status.state, StepStatus::Completed);
4002 assert_eq!(steps[1].name, "gate");
4003 assert_eq!(steps[1].kind, StepKind::Approval);
4004 assert_eq!(steps[1].status.state, StepStatus::Completed);
4005 assert_eq!(steps[2].name, "deploy");
4006 assert_eq!(steps[2].status.state, StepStatus::Completed);
4007 }
4008
4009 #[tokio::test]
4010 async fn double_approval_two_resumes() {
4011 let mut engine = create_test_engine();
4012 engine.register(DoubleApprovalWorkflow).unwrap();
4013
4014 let run = engine
4016 .run_handler("double-approval", TriggerKind::Manual, json!({}))
4017 .await
4018 .unwrap()
4019 .run;
4020 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
4021
4022 let steps = engine.store().list_steps(run.id).await.unwrap();
4023 assert_eq!(steps.len(), 2); approve_open_gate(engine.store(), run.id).await;
4027
4028 let resumed = engine.resume_run(run.id).await.unwrap().run;
4029 assert_eq!(resumed.status.state, RunStatus::AwaitingApproval);
4030
4031 let steps = engine.store().list_steps(run.id).await.unwrap();
4032 assert_eq!(steps.len(), 4); approve_open_gate(engine.store(), run.id).await;
4036
4037 let final_run = engine.resume_run(run.id).await.unwrap().run;
4038 assert_eq!(final_run.status.state, RunStatus::Completed);
4039
4040 let steps = engine.store().list_steps(run.id).await.unwrap();
4041 assert_eq!(steps.len(), 5);
4042 assert_eq!(steps[0].name, "build");
4043 assert_eq!(steps[1].name, "staging-gate");
4044 assert_eq!(steps[2].name, "deploy-staging");
4045 assert_eq!(steps[3].name, "prod-gate");
4046 assert_eq!(steps[4].name, "deploy-prod");
4047
4048 for step in &steps {
4049 assert_eq!(step.status.state, StepStatus::Completed);
4050 }
4051 }
4052
4053 use ironflow_store::models::{NewStep, StepUpdate, step_trace_id};
4058
4059 async fn create_step_with_status(
4060 store: &Arc<dyn Store>,
4061 run_id: Uuid,
4062 name: &str,
4063 position: u32,
4064 status: StepStatus,
4065 ) -> ironflow_store::models::Step {
4066 let step = store
4067 .create_step(NewStep {
4068 run_id,
4069 trace_id: step_trace_id(run_id, name, position),
4070 name: name.to_string(),
4071 kind: StepKind::Shell,
4072 position,
4073 input: None,
4074 is_error_handler: false,
4075 })
4076 .await
4077 .unwrap();
4078
4079 match status {
4080 StepStatus::Pending => {}
4081 StepStatus::Running => {
4082 store
4083 .update_step(
4084 step.id,
4085 StepUpdate {
4086 status: Some(StepStatus::Running),
4087 ..StepUpdate::default()
4088 },
4089 )
4090 .await
4091 .unwrap();
4092 }
4093 StepStatus::Completed => {
4094 store
4095 .update_step(
4096 step.id,
4097 StepUpdate {
4098 status: Some(StepStatus::Running),
4099 ..StepUpdate::default()
4100 },
4101 )
4102 .await
4103 .unwrap();
4104 store
4105 .update_step(
4106 step.id,
4107 StepUpdate {
4108 status: Some(StepStatus::Completed),
4109 ..StepUpdate::default()
4110 },
4111 )
4112 .await
4113 .unwrap();
4114 }
4115 StepStatus::AwaitingApproval => {
4116 store
4117 .update_step(
4118 step.id,
4119 StepUpdate {
4120 status: Some(StepStatus::Running),
4121 ..StepUpdate::default()
4122 },
4123 )
4124 .await
4125 .unwrap();
4126 store
4127 .update_step(
4128 step.id,
4129 StepUpdate {
4130 status: Some(StepStatus::AwaitingApproval),
4131 ..StepUpdate::default()
4132 },
4133 )
4134 .await
4135 .unwrap();
4136 }
4137 _ => panic!("unsupported status for test helper: {status}"),
4138 }
4139
4140 store.get_step(step.id).await.unwrap().unwrap()
4141 }
4142
4143 #[tokio::test]
4144 async fn fail_orphaned_steps_marks_running_as_failed() {
4145 let engine = create_test_engine();
4146 let run = engine
4147 .store()
4148 .create_run(NewRun {
4149 created_by: None,
4150 workflow_name: "test".to_string(),
4151 trigger: TriggerKind::Manual,
4152 payload: json!({}),
4153 max_retries: 0,
4154 handler_version: None,
4155 labels: HashMap::new(),
4156 scheduled_at: None,
4157 idempotency_key: None,
4158 concurrency_key: None,
4159 priority: 0,
4160 concurrency_limits: Vec::new(),
4161 max_cost_usd: None,
4162 worker_tags: Vec::new(),
4163 })
4164 .await
4165 .unwrap()
4166 .into_run();
4167
4168 let step = create_step_with_status(
4169 engine.store(),
4170 run.id,
4171 "running-step",
4172 0,
4173 StepStatus::Running,
4174 )
4175 .await;
4176
4177 engine
4178 .fail_orphaned_steps(run.id, "parent run timed out")
4179 .await
4180 .unwrap();
4181
4182 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4183 assert_eq!(updated.status.state, StepStatus::Failed);
4184 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
4185 assert!(updated.completed_at.is_some());
4186 }
4187
4188 #[tokio::test]
4189 async fn fail_orphaned_steps_marks_pending_as_skipped() {
4190 let engine = create_test_engine();
4191 let run = engine
4192 .store()
4193 .create_run(NewRun {
4194 created_by: None,
4195 workflow_name: "test".to_string(),
4196 trigger: TriggerKind::Manual,
4197 payload: json!({}),
4198 max_retries: 0,
4199 handler_version: None,
4200 labels: HashMap::new(),
4201 scheduled_at: None,
4202 idempotency_key: None,
4203 concurrency_key: None,
4204 priority: 0,
4205 concurrency_limits: Vec::new(),
4206 max_cost_usd: None,
4207 worker_tags: Vec::new(),
4208 })
4209 .await
4210 .unwrap()
4211 .into_run();
4212
4213 let step = create_step_with_status(
4214 engine.store(),
4215 run.id,
4216 "pending-step",
4217 0,
4218 StepStatus::Pending,
4219 )
4220 .await;
4221
4222 engine
4223 .fail_orphaned_steps(run.id, "parent run timed out")
4224 .await
4225 .unwrap();
4226
4227 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4228 assert_eq!(updated.status.state, StepStatus::Skipped);
4229 assert!(updated.error.is_none());
4230 assert!(updated.completed_at.is_some());
4231 }
4232
4233 #[tokio::test]
4234 async fn fail_orphaned_steps_marks_awaiting_approval_as_failed() {
4235 let engine = create_test_engine();
4236 let run = engine
4237 .store()
4238 .create_run(NewRun {
4239 created_by: None,
4240 workflow_name: "test".to_string(),
4241 trigger: TriggerKind::Manual,
4242 payload: json!({}),
4243 max_retries: 0,
4244 handler_version: None,
4245 labels: HashMap::new(),
4246 scheduled_at: None,
4247 idempotency_key: None,
4248 concurrency_key: None,
4249 priority: 0,
4250 concurrency_limits: Vec::new(),
4251 max_cost_usd: None,
4252 worker_tags: Vec::new(),
4253 })
4254 .await
4255 .unwrap()
4256 .into_run();
4257
4258 let step = create_step_with_status(
4259 engine.store(),
4260 run.id,
4261 "approval-step",
4262 0,
4263 StepStatus::AwaitingApproval,
4264 )
4265 .await;
4266
4267 engine
4268 .fail_orphaned_steps(run.id, "parent run timed out")
4269 .await
4270 .unwrap();
4271
4272 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4273 assert_eq!(updated.status.state, StepStatus::Failed);
4274 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
4275 assert!(updated.completed_at.is_some());
4276 }
4277
4278 #[tokio::test]
4279 async fn fail_orphaned_steps_skips_terminal_steps() {
4280 let engine = create_test_engine();
4281 let run = engine
4282 .store()
4283 .create_run(NewRun {
4284 created_by: None,
4285 workflow_name: "test".to_string(),
4286 trigger: TriggerKind::Manual,
4287 payload: json!({}),
4288 max_retries: 0,
4289 handler_version: None,
4290 labels: HashMap::new(),
4291 scheduled_at: None,
4292 idempotency_key: None,
4293 concurrency_key: None,
4294 priority: 0,
4295 concurrency_limits: Vec::new(),
4296 max_cost_usd: None,
4297 worker_tags: Vec::new(),
4298 })
4299 .await
4300 .unwrap()
4301 .into_run();
4302
4303 let completed_step =
4304 create_step_with_status(engine.store(), run.id, "done", 0, StepStatus::Completed).await;
4305 let running_step =
4306 create_step_with_status(engine.store(), run.id, "in-flight", 1, StepStatus::Running)
4307 .await;
4308
4309 engine
4310 .fail_orphaned_steps(run.id, "parent run timed out")
4311 .await
4312 .unwrap();
4313
4314 let completed = engine
4315 .store()
4316 .get_step(completed_step.id)
4317 .await
4318 .unwrap()
4319 .unwrap();
4320 assert_eq!(completed.status.state, StepStatus::Completed);
4321
4322 let failed = engine
4323 .store()
4324 .get_step(running_step.id)
4325 .await
4326 .unwrap()
4327 .unwrap();
4328 assert_eq!(failed.status.state, StepStatus::Failed);
4329 }
4330
4331 #[tokio::test]
4332 async fn fail_orphaned_steps_mixed_states() {
4333 let engine = create_test_engine();
4334 let run = engine
4335 .store()
4336 .create_run(NewRun {
4337 created_by: None,
4338 workflow_name: "test".to_string(),
4339 trigger: TriggerKind::Manual,
4340 payload: json!({}),
4341 max_retries: 0,
4342 handler_version: None,
4343 labels: HashMap::new(),
4344 scheduled_at: None,
4345 idempotency_key: None,
4346 concurrency_key: None,
4347 priority: 0,
4348 concurrency_limits: Vec::new(),
4349 max_cost_usd: None,
4350 worker_tags: Vec::new(),
4351 })
4352 .await
4353 .unwrap()
4354 .into_run();
4355
4356 let s_completed =
4357 create_step_with_status(engine.store(), run.id, "step-1", 0, StepStatus::Completed)
4358 .await;
4359 let s_running =
4360 create_step_with_status(engine.store(), run.id, "step-2", 1, StepStatus::Running).await;
4361 let s_pending =
4362 create_step_with_status(engine.store(), run.id, "step-3", 2, StepStatus::Pending).await;
4363
4364 engine.fail_orphaned_steps(run.id, "timeout").await.unwrap();
4365
4366 let r_completed = engine
4367 .store()
4368 .get_step(s_completed.id)
4369 .await
4370 .unwrap()
4371 .unwrap();
4372 assert_eq!(r_completed.status.state, StepStatus::Completed);
4373
4374 let r_running = engine
4375 .store()
4376 .get_step(s_running.id)
4377 .await
4378 .unwrap()
4379 .unwrap();
4380 assert_eq!(r_running.status.state, StepStatus::Failed);
4381 assert_eq!(r_running.error.as_deref(), Some("timeout"));
4382
4383 let r_pending = engine
4384 .store()
4385 .get_step(s_pending.id)
4386 .await
4387 .unwrap()
4388 .unwrap();
4389 assert_eq!(r_pending.status.state, StepStatus::Skipped);
4390 assert!(r_pending.error.is_none());
4391 }
4392
4393 #[tokio::test]
4394 async fn fail_orphaned_steps_no_steps_is_noop() {
4395 let engine = create_test_engine();
4396 let run = engine
4397 .store()
4398 .create_run(NewRun {
4399 created_by: None,
4400 workflow_name: "test".to_string(),
4401 trigger: TriggerKind::Manual,
4402 payload: json!({}),
4403 max_retries: 0,
4404 handler_version: None,
4405 labels: HashMap::new(),
4406 scheduled_at: None,
4407 idempotency_key: None,
4408 concurrency_key: None,
4409 priority: 0,
4410 concurrency_limits: Vec::new(),
4411 max_cost_usd: None,
4412 worker_tags: Vec::new(),
4413 })
4414 .await
4415 .unwrap()
4416 .into_run();
4417
4418 let result = engine.fail_orphaned_steps(run.id, "timeout").await;
4419 assert!(result.is_ok());
4420 }
4421
4422 #[tokio::test]
4423 async fn fail_orphaned_steps_preserves_existing_error() {
4424 let engine = create_test_engine();
4425 let run = engine
4426 .store()
4427 .create_run(NewRun {
4428 created_by: None,
4429 workflow_name: "test".to_string(),
4430 trigger: TriggerKind::Manual,
4431 payload: json!({}),
4432 max_retries: 0,
4433 handler_version: None,
4434 labels: HashMap::new(),
4435 scheduled_at: None,
4436 idempotency_key: None,
4437 concurrency_key: None,
4438 priority: 0,
4439 concurrency_limits: Vec::new(),
4440 max_cost_usd: None,
4441 worker_tags: Vec::new(),
4442 })
4443 .await
4444 .unwrap()
4445 .into_run();
4446
4447 let step_with_error = create_step_with_status(
4448 engine.store(),
4449 run.id,
4450 "already-errored",
4451 0,
4452 StepStatus::Running,
4453 )
4454 .await;
4455
4456 engine
4457 .store()
4458 .update_step(
4459 step_with_error.id,
4460 StepUpdate {
4461 error: Some("real error from provider".to_string()),
4462 ..StepUpdate::default()
4463 },
4464 )
4465 .await
4466 .unwrap();
4467
4468 let step_no_error = create_step_with_status(
4469 engine.store(),
4470 run.id,
4471 "no-error-yet",
4472 1,
4473 StepStatus::Running,
4474 )
4475 .await;
4476
4477 engine
4478 .fail_orphaned_steps(run.id, "parent run failed")
4479 .await
4480 .unwrap();
4481
4482 let updated_with = engine
4483 .store()
4484 .get_step(step_with_error.id)
4485 .await
4486 .unwrap()
4487 .unwrap();
4488 assert_eq!(updated_with.status.state, StepStatus::Failed);
4489 assert_eq!(
4490 updated_with.error.as_deref(),
4491 Some("real error from provider"),
4492 );
4493
4494 let updated_without = engine
4495 .store()
4496 .get_step(step_no_error.id)
4497 .await
4498 .unwrap()
4499 .unwrap();
4500 assert_eq!(updated_without.status.state, StepStatus::Failed);
4501 assert_eq!(updated_without.error.as_deref(), Some("parent run failed"),);
4502 }
4503
4504 use std::sync::atomic::{AtomicUsize, Ordering};
4509 use std::time::Duration;
4510
4511 use ironflow_store::models::LeaseRequest;
4512 use tokio::time::sleep;
4513
4514 struct CountingWorkflow {
4516 executions: Arc<AtomicUsize>,
4517 }
4518
4519 impl WorkflowHandler for CountingWorkflow {
4520 fn name(&self) -> &str {
4521 "local-resume-counter"
4522 }
4523
4524 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
4525 Box::pin(async move {
4526 self.executions.fetch_add(1, Ordering::SeqCst);
4527 Ok(())
4528 })
4529 }
4530 }
4531
4532 async fn engine_with_pending_counter() -> (Arc<Engine>, Arc<AtomicUsize>, Uuid) {
4534 let executions = Arc::new(AtomicUsize::new(0));
4535 let mut engine = create_test_engine();
4536 engine
4537 .register(CountingWorkflow {
4538 executions: executions.clone(),
4539 })
4540 .unwrap();
4541 let run = engine
4542 .enqueue_handler("local-resume-counter", TriggerKind::Manual, json!({}), 0)
4543 .await
4544 .unwrap();
4545 (Arc::new(engine), executions, run.id)
4546 }
4547
4548 #[tokio::test]
4549 async fn local_resume_leaves_a_run_claimed_by_a_worker_alone() {
4550 let (engine, executions, run_id) = engine_with_pending_counter().await;
4551 let lease = LeaseRequest {
4552 worker_id: "worker-1".to_string(),
4553 ttl: Duration::from_secs(30),
4554 };
4555 let picked = engine
4556 .store()
4557 .pick_next_pending_for(Some(lease), None)
4558 .await
4559 .unwrap()
4560 .expect("the run is due");
4561 assert_eq!(picked.id, run_id);
4562
4563 engine.spawn_local_resume(run_id);
4564 sleep(Duration::from_millis(200)).await;
4565
4566 assert_eq!(executions.load(Ordering::SeqCst), 0);
4567 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4568 assert_eq!(run.status.state, RunStatus::Running);
4569 assert_eq!(run.worker_id.as_deref(), Some("worker-1"));
4570 }
4571
4572 #[tokio::test]
4573 async fn local_resume_claims_a_pending_run_once() {
4574 let (engine, executions, run_id) = engine_with_pending_counter().await;
4575
4576 engine.spawn_local_resume(run_id);
4577 engine.spawn_local_resume(run_id);
4578 sleep(Duration::from_millis(200)).await;
4579
4580 assert_eq!(executions.load(Ordering::SeqCst), 1);
4581 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4582 assert_eq!(run.status.state, RunStatus::Completed);
4583 }
4584
4585 #[tokio::test]
4586 async fn local_resume_of_a_running_run_executes_it_once() {
4587 let (engine, executions, run_id) = engine_with_pending_counter().await;
4588 engine
4589 .store()
4590 .update_run_status(run_id, RunStatus::Running)
4591 .await
4592 .unwrap();
4593
4594 engine.spawn_local_resume(run_id);
4595 engine.spawn_local_resume(run_id);
4596 sleep(Duration::from_millis(200)).await;
4597
4598 assert_eq!(executions.load(Ordering::SeqCst), 1);
4599 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4600 assert_eq!(run.status.state, RunStatus::Completed);
4601 }
4602}