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, EventPublisher, EventSubscriber, RunBudgetExceededEvent,
50 RunFailedEvent, RunStatusChangedEvent, SignalAwaitedEvent, SignalReceivedEvent,
51 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) {
683 self.event_bus = Some(bus);
684 }
685
686 pub fn event_bus(&self) -> Option<&WorkflowEventBus> {
688 self.event_bus.as_ref()
689 }
690
691 pub fn store(&self) -> &Arc<dyn Store> {
693 &self.store
694 }
695
696 pub fn provider(&self) -> &Arc<dyn AgentProvider> {
698 &self.provider
699 }
700
701 fn build_context(&self, run: &Run) -> WorkflowContext {
710 let handlers = self.handlers.clone();
711 let resolver: crate::context::HandlerResolver =
712 Arc::new(move |name: &str| handlers.get(name).cloned());
713 let mut ctx = WorkflowContext::with_handler_resolver(
714 run.id,
715 run.workflow_name.clone(),
716 self.store.clone(),
717 self.provider.clone(),
718 resolver,
719 );
720 ctx.carry_over_run_totals(run.retry_count + 1, run.cost_usd, run.duration_ms);
721 ctx.set_max_cost_usd(run.max_cost_usd);
722 ctx.set_run_created_at(run.created_at);
723 if let Some(ref sender) = self.log_sender {
724 ctx.set_log_sender(sender.clone());
725 }
726 if let Some(ref sink) = self.artifact_sink {
727 ctx.set_artifact_sink(sink.clone());
728 }
729 if let Some(ref bus) = self.event_bus {
730 ctx.set_event_bus(bus.clone());
731 }
732 if let Some(ref provider) = self.decision_provider {
733 ctx.set_decision_provider(provider.clone());
734 }
735 if let Some(ref interceptor) = self.step_interceptor {
736 ctx.set_step_interceptor(interceptor.clone());
737 }
738 if let Some(ref tags) = self.worker_tags {
739 ctx.set_worker_tags(tags.clone());
740 }
741 ctx
742 }
743
744 fn build_context_with_guard(
750 &self,
751 run: &Run,
752 handler: &dyn WorkflowHandler,
753 ) -> WorkflowContext {
754 let mut ctx = self.build_context(run);
755 let guard_config = handler.guard_config().or_else(|| self.guard_config.clone());
756 if let Some(config) = guard_config {
757 ctx.set_guard(config, new_shared_guard_state());
758 }
759 ctx
760 }
761
762 async fn check_monthly_quota(&self, workflow_name: &str) -> Result<(), EngineError> {
773 let Some(limit) = self.budget.monthly_cost_limit_usd else {
774 return Ok(());
775 };
776
777 let stats = self
778 .store
779 .get_stats(RunFilter {
780 created_after: Some(month_start(Utc::now())),
781 ..RunFilter::default()
782 })
783 .await?;
784
785 if stats.total_cost_usd < limit {
786 return Ok(());
787 }
788
789 warn!(
790 workflow = %workflow_name,
791 limit_usd = %limit,
792 spent_usd = %stats.total_cost_usd,
793 "monthly cost quota exhausted, refusing new run"
794 );
795
796 #[cfg(feature = "prometheus")]
797 counter!(
798 RUN_BUDGET_EXCEEDED_TOTAL,
799 "workflow" => workflow_name.to_string(),
800 "scope" => "monthly",
801 )
802 .increment(1);
803
804 Err(EngineError::MonthlyBudgetExceeded {
805 limit_usd: limit,
806 spent_usd: stats.total_cost_usd,
807 })
808 }
809
810 pub fn register(&mut self, handler: impl WorkflowHandler + 'static) -> Result<(), EngineError> {
854 let name = handler.name().to_string();
855 if self.handlers.contains_key(&name) {
856 return Err(EngineError::InvalidWorkflow(format!(
857 "handler '{}' already registered",
858 name
859 )));
860 }
861 if let Some(category) = handler.category() {
862 validate_category(&name, category)?;
863 }
864 self.handlers.insert(name, Arc::new(handler));
865 Ok(())
866 }
867
868 pub fn register_boxed(&mut self, handler: Box<dyn WorkflowHandler>) -> Result<(), EngineError> {
875 let name = handler.name().to_string();
876 if self.handlers.contains_key(&name) {
877 return Err(EngineError::InvalidWorkflow(format!(
878 "handler '{}' already registered",
879 name
880 )));
881 }
882 if let Some(category) = handler.category() {
883 validate_category(&name, category)?;
884 }
885 self.handlers.insert(name, Arc::from(handler));
886 Ok(())
887 }
888
889 pub fn get_handler(&self, name: &str) -> Option<&Arc<dyn WorkflowHandler>> {
891 self.handlers.get(name)
892 }
893
894 pub fn handler_names(&self) -> Vec<&str> {
896 self.handlers.keys().map(|s| s.as_str()).collect()
897 }
898
899 pub fn handler_info(&self, name: &str) -> Option<WorkflowInfo> {
901 self.handlers.get(name).map(|h| h.describe())
902 }
903
904 pub fn scheduled_handlers(&self) -> Vec<(&str, &CronSchedule)> {
928 self.handlers
929 .iter()
930 .filter_map(|(name, handler)| handler.schedule().map(|sched| (name.as_str(), sched)))
931 .collect()
932 }
933
934 pub fn subscribe(
959 &mut self,
960 subscriber: impl EventSubscriber + 'static,
961 event_types: &[&'static str],
962 ) {
963 self.event_publisher.subscribe(subscriber, event_types);
964 }
965
966 pub fn event_publisher(&self) -> &EventPublisher {
971 &self.event_publisher
972 }
973
974 #[tracing::instrument(name = "engine.run_handler", skip_all, fields(workflow = %handler_name))]
1010 pub async fn run_handler(
1011 &self,
1012 handler_name: &str,
1013 trigger: TriggerKind,
1014 payload: Value,
1015 ) -> Result<WorkflowResult, EngineError> {
1016 let handler = self
1017 .handlers
1018 .get(handler_name)
1019 .ok_or_else(|| {
1020 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1021 })?
1022 .clone();
1023
1024 self.check_monthly_quota(handler_name).await?;
1025
1026 let handler_version = handler.version().map(str::to_string);
1027 let max_cost_usd = self
1028 .budget
1029 .resolve_run_cap(None, handler.default_max_cost_usd());
1030 let run = self
1031 .store
1032 .create_run(NewRun {
1033 created_by: None,
1034 workflow_name: handler_name.to_string(),
1035 trigger,
1036 payload,
1037 max_retries: 0,
1038 handler_version,
1039 labels: handler.default_labels(),
1040 scheduled_at: None,
1041 idempotency_key: None,
1042 concurrency_key: None,
1043 priority: clamp_priority(handler.priority()),
1044 concurrency_limits: Vec::new(),
1045 max_cost_usd,
1046 worker_tags: normalize_worker_tags(handler.required_worker_tags()),
1047 })
1048 .await?
1049 .into_run();
1050
1051 let run_id = run.id;
1052 info!(run_id = %run_id, handler_version = run.handler_version.as_deref().unwrap_or(""), "run created");
1053
1054 if self.is_workflow_paused(handler_name).await? {
1055 info!(run_id = %run_id, workflow = %handler_name, "workflow paused, run left pending");
1056 return Ok(WorkflowResult {
1057 run,
1058 steps: Vec::new(),
1059 });
1060 }
1061
1062 self.store
1063 .update_run_status(run_id, RunStatus::Running)
1064 .await?;
1065
1066 #[cfg(feature = "prometheus")]
1067 gauge!(RUNS_ACTIVE, "workflow" => handler_name.to_string()).increment(1.0);
1068
1069 let run_start = Instant::now();
1070 let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
1071
1072 let result = handler.execute(&mut ctx).await;
1073 self.finalize_run(run_id, handler_name, result, &ctx, run_start, run.labels)
1074 .await
1075 }
1076
1077 #[tracing::instrument(name = "engine.plan_handler", skip_all, fields(workflow = %handler_name))]
1126 pub async fn plan_handler(
1127 &self,
1128 handler_name: &str,
1129 payload: Value,
1130 options: PlanOptions,
1131 ) -> Result<ExecutionPlan, EngineError> {
1132 if options.max_depth == 0 {
1133 return Err(EngineError::InvalidWorkflow(
1134 "max_depth must be at least 1".to_string(),
1135 ));
1136 }
1137
1138 let handler = self
1139 .handlers
1140 .get(handler_name)
1141 .ok_or_else(|| {
1142 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1143 })?
1144 .clone();
1145
1146 let estimates = if options.estimate_durations {
1147 estimate_durations(&self.store, handler_name, options.sample_runs).await?
1148 } else {
1149 HashMap::new()
1150 };
1151
1152 let shared: SharedPlanRecorder = Arc::new(Mutex::new(PlanRecorder::new(
1153 handler_name.to_string(),
1154 payload,
1155 options.max_depth,
1156 estimates,
1157 )));
1158
1159 let handlers = self.handlers.clone();
1162 let resolver: crate::context::HandlerResolver =
1163 Arc::new(move |name: &str| handlers.get(name).cloned());
1164 let mut ctx = WorkflowContext::with_handler_resolver(
1165 Uuid::now_v7(),
1166 handler_name.to_string(),
1167 self.store.clone(),
1168 self.provider.clone(),
1169 resolver,
1170 );
1171 ctx.set_plan(shared.clone());
1172
1173 match timeout(options.timeout, handler.execute(&mut ctx)).await {
1174 Ok(Ok(())) => {}
1175 Ok(Err(err)) => lock_plan(&shared).fail(err.to_string()),
1176 Err(_) => lock_plan(&shared).fail(format!(
1177 "plan timed out after {} s",
1178 options.timeout.as_secs_f64()
1179 )),
1180 }
1181 drop(ctx);
1182
1183 let plan = match Arc::try_unwrap(shared) {
1184 Ok(mutex) => mutex
1185 .into_inner()
1186 .unwrap_or_else(|poisoned| poisoned.into_inner())
1187 .into_plan(),
1188 Err(shared) => lock_plan(&shared).snapshot(),
1189 };
1190
1191 info!(
1192 workflow = %handler_name,
1193 steps = plan.steps.len(),
1194 truncated = plan.truncated,
1195 "execution plan built"
1196 );
1197
1198 Ok(plan)
1199 }
1200
1201 #[tracing::instrument(name = "engine.enqueue_handler", skip_all, fields(workflow = %handler_name))]
1212 pub async fn enqueue_handler(
1213 &self,
1214 handler_name: &str,
1215 trigger: TriggerKind,
1216 payload: Value,
1217 max_retries: u32,
1218 ) -> Result<Run, EngineError> {
1219 self.enqueue_handler_with_options(
1220 handler_name,
1221 trigger,
1222 payload,
1223 EnqueueOptions {
1224 max_retries,
1225 ..Default::default()
1226 },
1227 )
1228 .await
1229 .map(RunCreation::into_run)
1230 }
1231
1232 #[tracing::instrument(name = "engine.enqueue_handler_with_options", skip_all, fields(workflow = %handler_name))]
1284 pub async fn enqueue_handler_with_options(
1285 &self,
1286 handler_name: &str,
1287 trigger: TriggerKind,
1288 payload: Value,
1289 options: EnqueueOptions,
1290 ) -> Result<RunCreation, EngineError> {
1291 let EnqueueOptions {
1292 max_retries,
1293 labels,
1294 scheduled_at,
1295 max_cost_usd,
1296 created_by,
1297 idempotency_key,
1298 concurrency_key,
1299 concurrency_limits,
1300 priority,
1301 worker_tags,
1302 } = options;
1303
1304 validate_concurrency_limits(&concurrency_limits)
1307 .map_err(EngineError::InvalidConcurrencyLimit)?;
1308 if let Some(priority) = priority {
1309 validate_priority(priority).map_err(EngineError::InvalidPriority)?;
1310 }
1311 validate_worker_tags(&worker_tags).map_err(EngineError::InvalidWorkerTag)?;
1312
1313 let handler = self.handlers.get(handler_name).ok_or_else(|| {
1314 EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
1315 })?;
1316
1317 self.check_monthly_quota(handler_name).await?;
1318
1319 let handler_version = handler.version().map(str::to_string);
1320 let mut merged_labels = handler.default_labels();
1321 merged_labels.extend(labels);
1322 let resolved_cap = self
1323 .budget
1324 .resolve_run_cap(max_cost_usd, handler.default_max_cost_usd());
1325 let priority = priority.unwrap_or_else(|| clamp_priority(handler.priority()));
1326 let required_tags = normalize_worker_tags(
1327 handler
1328 .required_worker_tags()
1329 .into_iter()
1330 .chain(worker_tags),
1331 );
1332
1333 let creation = self
1334 .store
1335 .create_run(NewRun {
1336 workflow_name: handler_name.to_string(),
1337 trigger,
1338 payload,
1339 max_retries,
1340 handler_version,
1341 labels: merged_labels,
1342 scheduled_at,
1343 created_by,
1344 idempotency_key,
1345 concurrency_key,
1346 priority,
1347 concurrency_limits,
1348 max_cost_usd: resolved_cap,
1349 worker_tags: required_tags,
1350 })
1351 .await?;
1352
1353 match &creation {
1354 RunCreation::Created(run) => info!(
1355 run_id = %run.id,
1356 workflow = %handler_name,
1357 max_cost_usd = ?resolved_cap,
1358 "handler run enqueued"
1359 ),
1360 RunCreation::Existing(run) => info!(
1361 run_id = %run.id,
1362 workflow = %handler_name,
1363 "idempotent replay, nothing enqueued"
1364 ),
1365 }
1366
1367 Ok(creation)
1368 }
1369
1370 #[tracing::instrument(name = "engine.execute_handler_run", skip_all, fields(run_id = %run_id))]
1390 pub async fn execute_handler_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1391 let run = self
1392 .store
1393 .get_run(run_id)
1394 .await?
1395 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1396
1397 if let Some(root_run_id) = chain_root(&run) {
1398 return self.resume_chain(run, root_run_id).await;
1399 }
1400
1401 let _active = self.track_execution(run_id).await;
1402
1403 let handler = self
1404 .handlers
1405 .get(&run.workflow_name)
1406 .ok_or_else(|| {
1407 EngineError::InvalidWorkflow(format!(
1408 "no handler registered: {}",
1409 run.workflow_name
1410 ))
1411 })?
1412 .clone();
1413
1414 #[cfg(feature = "prometheus")]
1415 gauge!(RUNS_ACTIVE, "workflow" => run.workflow_name.clone()).increment(1.0);
1416
1417 let run_start = Instant::now();
1418 let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
1419
1420 let result = if handler.is_version_compatible(run.handler_version.as_deref()) {
1433 ctx.load_replay_steps().await?;
1434 self.release_then_execute(run_id, handler.as_ref(), &mut ctx)
1435 .await
1436 } else {
1437 Err(EngineError::HandlerVersionMismatch {
1438 run_id,
1439 workflow_name: run.workflow_name.clone(),
1440 run_version: run
1441 .handler_version
1442 .clone()
1443 .unwrap_or_else(|| "unknown".to_string()),
1444 current_version: handler
1445 .version()
1446 .map(str::to_string)
1447 .unwrap_or_else(|| "unknown".to_string()),
1448 })
1449 };
1450
1451 self.finalize_run(
1452 run_id,
1453 &run.workflow_name,
1454 result,
1455 &ctx,
1456 run_start,
1457 run.labels,
1458 )
1459 .await
1460 }
1461
1462 #[tracing::instrument(name = "engine.execute_run", skip_all, fields(run_id = %run_id))]
1470 pub async fn execute_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1471 self.execute_handler_run(run_id).await
1472 }
1473
1474 #[tracing::instrument(name = "engine.resume_run", skip_all, fields(run_id = %run_id))]
1503 pub async fn resume_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
1504 let run = self
1505 .store
1506 .get_run(run_id)
1507 .await?
1508 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1509
1510 if let Some(root_run_id) = chain_root(&run) {
1511 return self.resume_chain(run, root_run_id).await;
1512 }
1513
1514 self.resume_loaded_run(run).await
1515 }
1516
1517 async fn resume_chain(
1532 &self,
1533 child: Run,
1534 root_run_id: Uuid,
1535 ) -> Result<WorkflowResult, EngineError> {
1536 let child_run_id = child.id;
1537 let lease = child.worker_id.zip(child.lease_expires_at);
1538 let root = self
1539 .store
1540 .get_run(root_run_id)
1541 .await?
1542 .ok_or(EngineError::Store(StoreError::RunNotFound(root_run_id)))?;
1543
1544 match root.status.state {
1545 RunStatus::AwaitingApproval | RunStatus::Pending => {
1546 self.move_root_to_running(root_run_id, lease.as_ref())
1547 .await?;
1548 }
1549 RunStatus::Sleeping => {
1550 self.store
1551 .update_run_status(root_run_id, RunStatus::Pending)
1552 .await?;
1553 self.move_root_to_running(root_run_id, lease.as_ref())
1554 .await?;
1555 }
1556 other => {
1557 let reason = format!(
1558 "cannot resume child run {child_run_id}: root run {root_run_id} is {other}"
1559 );
1560 if let Err(err) = self
1561 .fail_or_schedule_retry(child_run_id, &reason, false, None, None)
1562 .await
1563 {
1564 error!(
1565 run_id = %child_run_id,
1566 error = %err,
1567 "failed to fail a child run whose root cannot resume"
1568 );
1569 }
1570 return Err(EngineError::InvalidWorkflow(reason));
1571 }
1572 }
1573
1574 if lease.is_some() {
1575 self.store
1576 .update_run(
1577 child_run_id,
1578 RunUpdate {
1579 lease: Some(LeaseUpdate::Release),
1580 ..RunUpdate::default()
1581 },
1582 )
1583 .await?;
1584 }
1585
1586 info!(
1587 run_id = %child_run_id,
1588 root_run_id = %root_run_id,
1589 lease_transferred = lease.is_some(),
1590 "child run resumed through its root run"
1591 );
1592
1593 let root = self
1594 .store
1595 .get_run(root_run_id)
1596 .await?
1597 .ok_or(EngineError::Store(StoreError::RunNotFound(root_run_id)))?;
1598 self.resume_loaded_run(root).await
1599 }
1600
1601 async fn move_root_to_running(
1607 &self,
1608 root_run_id: Uuid,
1609 lease: Option<&(String, DateTime<Utc>)>,
1610 ) -> Result<(), EngineError> {
1611 match lease {
1612 Some((worker_id, expires_at)) => {
1613 self.store
1614 .update_run(
1615 root_run_id,
1616 RunUpdate {
1617 status: Some(RunStatus::Running),
1618 lease: Some(LeaseUpdate::Set {
1619 worker_id: worker_id.clone(),
1620 expires_at: *expires_at,
1621 }),
1622 ..RunUpdate::default()
1623 },
1624 )
1625 .await?;
1626 }
1627 None => {
1628 self.store
1629 .update_run_status(root_run_id, RunStatus::Running)
1630 .await?;
1631 }
1632 }
1633 Ok(())
1634 }
1635
1636 async fn resume_loaded_run(&self, run: Run) -> Result<WorkflowResult, EngineError> {
1638 let run_id = run.id;
1639 let _active = self.track_execution(run_id).await;
1640 let handler = self
1641 .handlers
1642 .get(&run.workflow_name)
1643 .ok_or_else(|| {
1644 EngineError::InvalidWorkflow(format!(
1645 "no handler registered: {}",
1646 run.workflow_name
1647 ))
1648 })?
1649 .clone();
1650
1651 info!(run_id = %run_id, workflow = %run.workflow_name, "resuming run after approval");
1652
1653 let run_start = Instant::now();
1654 let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
1655
1656 let result = if handler.is_version_compatible(run.handler_version.as_deref()) {
1657 ctx.load_replay_steps().await?;
1658 self.release_then_execute(run_id, handler.as_ref(), &mut ctx)
1659 .await
1660 } else {
1661 Err(EngineError::HandlerVersionMismatch {
1662 run_id,
1663 workflow_name: run.workflow_name.clone(),
1664 run_version: run
1665 .handler_version
1666 .clone()
1667 .unwrap_or_else(|| "unknown".to_string()),
1668 current_version: handler
1669 .version()
1670 .map(str::to_string)
1671 .unwrap_or_else(|| "unknown".to_string()),
1672 })
1673 };
1674
1675 self.finalize_run(
1676 run_id,
1677 &run.workflow_name,
1678 result,
1679 &ctx,
1680 run_start,
1681 run.labels,
1682 )
1683 .await
1684 }
1685
1686 pub async fn deliver_signal(
1732 self: &Arc<Self>,
1733 signal: NewSignal,
1734 ) -> Result<SignalDelivery, EngineError> {
1735 if signal.name.trim().is_empty() {
1736 return Err(EngineError::InvalidSignal(
1737 "signal name must not be empty".to_string(),
1738 ));
1739 }
1740 if signal.key.trim().is_empty() {
1741 return Err(EngineError::InvalidSignal(
1742 "signal key must not be empty".to_string(),
1743 ));
1744 }
1745
1746 let stored = match self.store.insert_signal(signal).await? {
1747 SignalInsert::Created(stored) => stored,
1748 SignalInsert::Duplicate(existing) => {
1749 info!(
1750 signal_id = %existing.id,
1751 signal = %existing.name,
1752 key = %existing.key,
1753 "duplicate signal ignored"
1754 );
1755 return Ok(SignalDelivery {
1756 signal_id: existing.id,
1757 duplicate: true,
1758 resumed: Vec::new(),
1759 rejected: Vec::new(),
1760 });
1761 }
1762 };
1763
1764 let waiters = self
1765 .store
1766 .list_signal_waiters(&stored.name, &stored.key)
1767 .await?;
1768 let mut resumed = Vec::new();
1769 let mut rejected = Vec::new();
1770
1771 for step in waiters {
1772 if let Err(error) = validate_step_payload(step.input.as_ref(), &stored.payload) {
1773 rejected.push(SignalRejected {
1774 run_id: step.run_id,
1775 step_id: step.id,
1776 error,
1777 });
1778 continue;
1779 }
1780
1781 match self
1782 .store
1783 .resolve_signal_step(step.id, received_output(&stored))
1784 .await
1785 {
1786 Ok(SignalStepResolution::Resolved {
1787 run_id,
1788 run_resumed,
1789 }) => {
1790 resumed.push(SignalResumed {
1791 run_id,
1792 step_id: step.id,
1793 });
1794 if run_resumed && self.execution_mode == ExecutionMode::Local {
1795 self.spawn_local_resume(run_id);
1796 }
1797 }
1798 Ok(SignalStepResolution::NotWaiting { .. }) => {}
1800 Err(err) => {
1801 error!(
1802 run_id = %step.run_id,
1803 step_id = %step.id,
1804 error = %err,
1805 "failed to resolve a waiting signal step"
1806 );
1807 rejected.push(SignalRejected {
1808 run_id: step.run_id,
1809 step_id: step.id,
1810 error: err.to_string(),
1811 });
1812 }
1813 }
1814 }
1815
1816 info!(
1817 signal_id = %stored.id,
1818 signal = %stored.name,
1819 key = %stored.key,
1820 resumed = resumed.len(),
1821 rejected = rejected.len(),
1822 "signal received"
1823 );
1824 self.event_publisher
1825 .publish(Event::SignalReceived(SignalReceivedEvent {
1826 signal_id: stored.id,
1827 name: stored.name.clone(),
1828 key: stored.key.clone(),
1829 resumed_runs: resumed.iter().map(|r| r.run_id).collect(),
1830 at: stored.received_at,
1831 }));
1832
1833 Ok(SignalDelivery {
1834 signal_id: stored.id,
1835 duplicate: false,
1836 resumed,
1837 rejected,
1838 })
1839 }
1840
1841 pub async fn send_signal<S: Signal>(
1879 self: &Arc<Self>,
1880 signal: &S,
1881 key: &str,
1882 idempotency_id: Option<&str>,
1883 ) -> Result<SignalDelivery, EngineError> {
1884 let payload = to_value(signal)?;
1885 self.deliver_signal(NewSignal {
1886 name: S::NAME.to_string(),
1887 key: key.to_string(),
1888 payload,
1889 idempotency_id: idempotency_id.map(str::to_string),
1890 })
1891 .await
1892 }
1893
1894 pub(crate) fn spawn_local_resume(self: &Arc<Self>, run_id: Uuid) {
1919 let engine = Arc::clone(self);
1920 spawn(async move {
1921 let _resuming = lock_run(&engine.local_resumes, run_id).await;
1925 engine.wait_until_idle(run_id).await;
1926 let run = match engine.store.get_run(run_id).await {
1927 Ok(Some(run)) => run,
1928 Ok(None) => {
1929 error!(run_id = %run_id, "run to restart not found");
1930 return;
1931 }
1932 Err(err) => {
1933 error!(run_id = %run_id, error = %err, "failed to load a run to restart");
1934 return;
1935 }
1936 };
1937 match run.status.state {
1938 RunStatus::Pending => {
1939 match engine
1940 .store
1941 .update_run_status(run_id, RunStatus::Running)
1942 .await
1943 {
1944 Ok(()) => {}
1945 Err(StoreError::InvalidTransition { from, .. }) => {
1946 info!(
1947 run_id = %run_id,
1948 status = %from,
1949 "run claimed elsewhere, left to its worker"
1950 );
1951 return;
1952 }
1953 Err(err) => {
1954 error!(run_id = %run_id, error = %err, "failed to restart a woken run");
1955 return;
1956 }
1957 }
1958 }
1959 RunStatus::Running if run.worker_id.is_some() => {
1960 info!(
1961 run_id = %run_id,
1962 worker_id = ?run.worker_id,
1963 "run owned by a worker lease, local resume skipped"
1964 );
1965 return;
1966 }
1967 RunStatus::Running => {
1968 if engine.is_executing(run_id) {
1969 info!(run_id = %run_id, "run already executing in this process, resume skipped");
1970 return;
1971 }
1972 }
1973 status => {
1974 info!(
1975 run_id = %run_id,
1976 status = %status,
1977 "run no longer waiting to restart, resume skipped"
1978 );
1979 return;
1980 }
1981 }
1982 if let Err(err) = engine.resume_run(run_id).await {
1983 error!(run_id = %run_id, error = %err, "failed to resume a woken run");
1984 }
1985 });
1986 }
1987
1988 pub async fn fail_or_schedule_retry(
2039 &self,
2040 run_id: Uuid,
2041 error: &str,
2042 retryable: bool,
2043 cost_usd: Option<Decimal>,
2044 duration_ms: Option<u64>,
2045 ) -> Result<RunStatus, EngineError> {
2046 let run = self
2047 .store
2048 .get_run(run_id)
2049 .await?
2050 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2051
2052 if run.status.state == RunStatus::Paused {
2055 info!(run_id = %run_id, error = %error, "run paused, failure not recorded");
2056 return Ok(RunStatus::Paused);
2057 }
2058
2059 let has_attempts_left = run.retry_count < run.max_retries;
2060 let update = if retryable && has_attempts_left {
2061 let backoff = backoff_for_retry(run.retry_count);
2062 let scheduled_at = Utc::now() + TimeDelta::milliseconds(backoff.as_millis() as i64);
2063
2064 info!(
2065 run_id = %run_id,
2066 workflow = %run.workflow_name,
2067 attempt = run.retry_count + 1,
2068 max_retries = run.max_retries,
2069 backoff_secs = backoff.as_secs(),
2070 scheduled_at = %scheduled_at,
2071 "run failed, scheduling retry"
2072 );
2073
2074 RunUpdate {
2075 status: Some(RunStatus::Retrying),
2076 error: Some(error.to_string()),
2077 increment_retry: true,
2078 cost_usd,
2079 duration_ms,
2080 scheduled_at: Some(scheduled_at),
2081 ..RunUpdate::default()
2082 }
2083 } else {
2084 RunUpdate {
2085 status: Some(RunStatus::Failed),
2086 error: Some(error.to_string()),
2087 cost_usd,
2088 duration_ms,
2089 completed_at: Some(Utc::now()),
2090 ..RunUpdate::default()
2091 }
2092 };
2093
2094 let status = update.status.unwrap_or(RunStatus::Failed);
2095 self.store.update_run(run_id, update).await?;
2096 self.fail_orphaned_steps(run_id, error).await?;
2097 self.cancel_descendants_of_stopped_run(run_id, error).await;
2100
2101 Ok(status)
2102 }
2103
2104 pub async fn interrupt_running_steps(&self, run_id: Uuid) -> Result<(), EngineError> {
2134 interrupt_running_steps(self.store.as_ref(), run_id).await
2135 }
2136
2137 pub async fn fail_orphaned_steps(
2151 &self,
2152 run_id: Uuid,
2153 error_message: &str,
2154 ) -> Result<(), EngineError> {
2155 let steps = self.store.list_steps(run_id).await?;
2156 let now = Utc::now();
2157
2158 for step in steps {
2159 if step.status.state.is_terminal() {
2160 continue;
2161 }
2162
2163 let (target_status, error) = match step.status.state {
2164 StepStatus::Running | StepStatus::AwaitingApproval => {
2165 let err = if step.error.is_some() {
2166 None
2167 } else {
2168 Some(error_message.to_string())
2169 };
2170 (StepStatus::Failed, err)
2171 }
2172 StepStatus::Pending => (StepStatus::Skipped, None),
2173 _ => continue,
2174 };
2175
2176 if let Err(e) = self
2177 .store
2178 .update_step(
2179 step.id,
2180 StepUpdate {
2181 status: Some(target_status),
2182 error,
2183 completed_at: Some(now),
2184 ..StepUpdate::default()
2185 },
2186 )
2187 .await
2188 {
2189 warn!(
2190 run_id = %run_id,
2191 step_id = %step.id,
2192 step_name = %step.name,
2193 error = %e,
2194 "failed to cleanup orphaned step"
2195 );
2196 } else {
2197 info!(
2198 run_id = %run_id,
2199 step_id = %step.id,
2200 step_name = %step.name,
2201 from = %step.status.state,
2202 to = %target_status,
2203 "cleaned up orphaned step"
2204 );
2205 }
2206 }
2207
2208 Ok(())
2209 }
2210
2211 async fn release_then_execute(
2216 &self,
2217 run_id: Uuid,
2218 handler: &dyn WorkflowHandler,
2219 ctx: &mut WorkflowContext,
2220 ) -> Result<(), EngineError> {
2221 match self.provider.release_run(&run_id.to_string()).await {
2222 Ok(()) => handler.execute(ctx).await,
2223 Err(e) => Err(EngineError::Operation(OperationError::Agent(e))),
2224 }
2225 }
2226
2227 async fn finalize_run(
2233 &self,
2234 run_id: Uuid,
2235 workflow_name: &str,
2236 result: Result<(), EngineError>,
2237 ctx: &WorkflowContext,
2238 run_start: Instant,
2239 run_labels: HashMap<String, String>,
2240 ) -> Result<WorkflowResult, EngineError> {
2241 let total_duration = ctx.carried_duration_ms() + run_start.elapsed().as_millis() as u64;
2244 let completed_at = Utc::now();
2245
2246 if let Some(run) = self.store.get_run(run_id).await?
2250 && run.status.state == RunStatus::Paused
2251 {
2252 let run = self
2253 .store
2254 .update_run_returning(
2255 run_id,
2256 RunUpdate {
2257 cost_usd: Some(ctx.total_cost_usd()),
2258 duration_ms: Some(total_duration),
2259 ..RunUpdate::default()
2260 },
2261 )
2262 .await?;
2263 info!(
2264 run_id = %run_id,
2265 outcome = ?result.err().map(|err| err.to_string()),
2266 "run paused, execution stopped"
2267 );
2268 return Ok(WorkflowResult {
2269 run,
2270 steps: ctx.step_results().to_vec(),
2271 });
2272 }
2273
2274 if matches!(result, Err(EngineError::RunPaused { .. })) {
2277 let run = self
2278 .store
2279 .get_run(run_id)
2280 .await?
2281 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2282 info!(
2283 run_id = %run_id,
2284 status = %run.status.state,
2285 "run resumed after the pause, execution stopped"
2286 );
2287 return Ok(WorkflowResult {
2288 run,
2289 steps: ctx.step_results().to_vec(),
2290 });
2291 }
2292
2293 let final_status;
2294 let final_run;
2295
2296 match result {
2297 Ok(()) => {
2298 final_status = if ctx.has_allowed_failure() {
2299 RunStatus::Warning
2300 } else {
2301 RunStatus::Completed
2302 };
2303 final_run = self
2304 .store
2305 .update_run_returning(
2306 run_id,
2307 RunUpdate {
2308 status: Some(final_status),
2309 cost_usd: Some(ctx.total_cost_usd()),
2310 duration_ms: Some(total_duration),
2311 completed_at: Some(completed_at),
2312 output: ctx.output().cloned(),
2313 ..RunUpdate::default()
2314 },
2315 )
2316 .await?;
2317
2318 info!(
2319 run_id = %run_id,
2320 status = %final_status,
2321 cost_usd = %ctx.total_cost_usd(),
2322 duration_ms = total_duration,
2323 "run completed"
2324 );
2325 }
2326 Err(EngineError::ApprovalRequired {
2327 run_id: approval_run_id,
2328 step_id,
2329 ref message,
2330 }) => {
2331 final_status = RunStatus::AwaitingApproval;
2332 final_run = self
2333 .store
2334 .update_run_returning(
2335 run_id,
2336 RunUpdate {
2337 status: Some(RunStatus::AwaitingApproval),
2338 cost_usd: Some(ctx.total_cost_usd()),
2339 duration_ms: Some(total_duration),
2340 ..RunUpdate::default()
2341 },
2342 )
2343 .await?;
2344
2345 info!(
2346 run_id = %approval_run_id,
2347 step_id = %step_id,
2348 message = %message,
2349 "run awaiting approval"
2350 );
2351
2352 self.publish_approval_requested(approval_run_id, step_id, message)
2353 .await?;
2354 }
2355 Err(EngineError::ChildSuspended {
2356 run_id: child_run_id,
2357 ref cause,
2358 }) => {
2359 final_status = cause.suspension_status();
2360 final_run = self
2364 .store
2365 .update_run_returning(
2366 run_id,
2367 RunUpdate {
2368 status: Some(final_status),
2369 cost_usd: Some(ctx.total_cost_usd()),
2370 duration_ms: Some(total_duration),
2371 ..RunUpdate::default()
2372 },
2373 )
2374 .await?;
2375
2376 let leaf = cause.suspension_leaf();
2377 info!(
2378 run_id = %run_id,
2379 child_run_id = %child_run_id,
2380 status = %final_status,
2381 cause = %leaf,
2382 "run suspended with its child run"
2383 );
2384
2385 match leaf {
2386 EngineError::ApprovalRequired {
2387 run_id: approval_run_id,
2388 step_id,
2389 message,
2390 } => {
2391 self.publish_approval_requested(*approval_run_id, *step_id, message)
2392 .await?;
2393 }
2394 EngineError::SignalWaiting {
2395 run_id: wait_run_id,
2396 step_id,
2397 step_name,
2398 name,
2399 key,
2400 deadline_at,
2401 } => {
2402 self.event_publisher
2403 .publish(Event::SignalAwaited(SignalAwaitedEvent {
2404 run_id: *wait_run_id,
2405 step_id: *step_id,
2406 step_name: step_name.clone(),
2407 name: name.clone(),
2408 key: key.clone(),
2409 deadline_at: *deadline_at,
2410 at: Utc::now(),
2411 }));
2412 }
2413 _ => {}
2416 }
2417 }
2418 Err(EngineError::HumanInputRequired {
2419 run_id: input_run_id,
2420 step_id,
2421 ref message,
2422 }) => {
2423 final_status = RunStatus::AwaitingApproval;
2424 final_run = self
2425 .store
2426 .update_run_returning(
2427 run_id,
2428 RunUpdate {
2429 status: Some(RunStatus::AwaitingApproval),
2430 cost_usd: Some(ctx.total_cost_usd()),
2431 duration_ms: Some(total_duration),
2432 ..RunUpdate::default()
2433 },
2434 )
2435 .await?;
2436
2437 info!(
2439 run_id = %input_run_id,
2440 step_id = %step_id,
2441 message = %message,
2442 "run awaiting human input"
2443 );
2444 }
2445 Err(EngineError::DelaySleeping {
2446 run_id: delay_run_id,
2447 step_id,
2448 wake_at,
2449 }) => {
2450 final_status = RunStatus::Sleeping;
2451 final_run = self
2452 .store
2453 .update_run_returning(
2454 run_id,
2455 RunUpdate {
2456 status: Some(RunStatus::Sleeping),
2457 cost_usd: Some(ctx.total_cost_usd()),
2458 duration_ms: Some(total_duration),
2459 scheduled_at: Some(wake_at),
2460 ..RunUpdate::default()
2461 },
2462 )
2463 .await?;
2464
2465 info!(
2466 run_id = %delay_run_id,
2467 step_id = %step_id,
2468 wake_at = %wake_at,
2469 "run sleeping until delay elapses"
2470 );
2471 }
2472 Err(EngineError::CapacitySleeping {
2473 run_id: capacity_run_id,
2474 step_id,
2475 ref kind,
2476 wake_at,
2477 }) => {
2478 final_status = RunStatus::Sleeping;
2479 final_run = self
2480 .store
2481 .update_run_returning(
2482 run_id,
2483 RunUpdate {
2484 status: Some(RunStatus::Sleeping),
2485 cost_usd: Some(ctx.total_cost_usd()),
2486 duration_ms: Some(total_duration),
2487 scheduled_at: Some(wake_at),
2488 capacity_wait_kind: Some(ProviderKind::new(kind.as_str())),
2489 ..RunUpdate::default()
2490 },
2491 )
2492 .await?;
2493
2494 info!(
2495 run_id = %capacity_run_id,
2496 step_id = %step_id,
2497 kind = %kind,
2498 wake_at = %wake_at,
2499 "run sleeping until provider capacity returns"
2500 );
2501 }
2502 Err(EngineError::SignalWaiting {
2503 run_id: wait_run_id,
2504 step_id,
2505 ref step_name,
2506 ref name,
2507 ref key,
2508 deadline_at,
2509 }) => {
2510 final_status = RunStatus::Sleeping;
2511 let waiting = self
2515 .store
2516 .suspend_run_on_signal(run_id, step_id, deadline_at)
2517 .await?;
2518 final_run = self
2519 .store
2520 .update_run_returning(
2521 run_id,
2522 RunUpdate {
2523 cost_usd: Some(ctx.total_cost_usd()),
2524 duration_ms: Some(total_duration),
2525 ..RunUpdate::default()
2526 },
2527 )
2528 .await?;
2529
2530 if waiting {
2531 self.event_publisher
2532 .publish(Event::SignalAwaited(SignalAwaitedEvent {
2533 run_id: wait_run_id,
2534 step_id,
2535 step_name: step_name.clone(),
2536 name: name.clone(),
2537 key: key.clone(),
2538 deadline_at,
2539 at: Utc::now(),
2540 }));
2541 }
2542
2543 info!(
2544 run_id = %wait_run_id,
2545 step_id = %step_id,
2546 signal = %name,
2547 key = %key,
2548 deadline_at = %deadline_at,
2549 waiting,
2550 "run sleeping until a signal arrives"
2551 );
2552 }
2553 Err(err) => {
2554 let guardrail_stop = matches!(
2558 err,
2559 EngineError::RunBudgetExceeded { .. } | EngineError::WorkflowGuardRejected(_)
2560 );
2561
2562 final_status = if guardrail_stop {
2563 if let Err(store_err) = self
2564 .store
2565 .update_run(
2566 run_id,
2567 RunUpdate {
2568 status: Some(RunStatus::Cancelled),
2569 error: Some(err.to_string()),
2570 cost_usd: Some(ctx.total_cost_usd()),
2571 duration_ms: Some(total_duration),
2572 completed_at: Some(completed_at),
2573 output: ctx.output().cloned(),
2574 ..RunUpdate::default()
2575 },
2576 )
2577 .await
2578 {
2579 error!(run_id = %run_id, store_error = %store_err, "failed to persist run cancellation");
2580 }
2581 if let Err(cleanup_err) = self
2582 .fail_orphaned_steps(run_id, "run stopped: guardrail limit reached")
2583 .await
2584 {
2585 error!(run_id = %run_id, store_error = %cleanup_err, "failed to cleanup orphaned steps");
2586 }
2587 RunStatus::Cancelled
2588 } else {
2589 if let Some(output) = ctx.output()
2592 && let Err(store_err) = self
2593 .store
2594 .update_run(
2595 run_id,
2596 RunUpdate {
2597 output: Some(output.clone()),
2598 ..RunUpdate::default()
2599 },
2600 )
2601 .await
2602 {
2603 error!(run_id = %run_id, store_error = %store_err, "failed to persist run output");
2604 }
2605 self.fail_or_schedule_retry(
2606 run_id,
2607 &err.to_string(),
2608 is_run_retryable(&err),
2609 Some(ctx.total_cost_usd()),
2610 Some(total_duration),
2611 )
2612 .await
2613 .unwrap_or_else(|store_err| {
2614 error!(run_id = %run_id, store_error = %store_err, "failed to persist run failure");
2615 RunStatus::Failed
2616 })
2617 };
2618
2619 if matches!(err, EngineError::RunBudgetExceeded { .. }) {
2620 self.on_run_budget_exceeded(workflow_name, run_id, &err);
2621 }
2622
2623 error!(run_id = %run_id, status = %final_status, error = %err, "run stopped");
2624
2625 self.publish_run_status_changed(
2626 workflow_name,
2627 run_id,
2628 final_status,
2629 Some(err.to_string()),
2630 ctx,
2631 total_duration,
2632 run_labels,
2633 );
2634
2635 #[cfg(feature = "prometheus")]
2636 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
2637
2638 return Err(err);
2639 }
2640 }
2641
2642 self.publish_run_status_changed(
2643 workflow_name,
2644 run_id,
2645 final_status,
2646 None,
2647 ctx,
2648 total_duration,
2649 run_labels,
2650 );
2651
2652 #[cfg(feature = "prometheus")]
2653 self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
2654
2655 Ok(WorkflowResult {
2656 run: final_run,
2657 steps: ctx.step_results().to_vec(),
2658 })
2659 }
2660
2661 async fn publish_approval_requested(
2664 &self,
2665 run_id: Uuid,
2666 step_id: Uuid,
2667 message: &str,
2668 ) -> Result<(), EngineError> {
2669 let requirement = self
2670 .store
2671 .get_step(step_id)
2672 .await?
2673 .and_then(|s| s.approval_requirement);
2674 self.event_publisher
2675 .publish(Event::ApprovalRequested(ApprovalRequestedEvent {
2676 run_id,
2677 step_id,
2678 message: message.to_string(),
2679 requirement,
2680 at: Utc::now(),
2681 }));
2682 Ok(())
2683 }
2684
2685 pub async fn fail_ancestors(&self, run_id: Uuid, reason: &str) -> Result<(), EngineError> {
2718 let mut current = self
2719 .store
2720 .get_run(run_id)
2721 .await?
2722 .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
2723 let mut visited = HashSet::from([run_id]);
2725
2726 while let Some(parent_id) = chain_parent(¤t) {
2727 if !visited.insert(parent_id) {
2728 break;
2729 }
2730 let status = self
2731 .fail_or_schedule_retry(parent_id, reason, false, None, None)
2732 .await?;
2733 if status == RunStatus::Paused {
2736 self.requeue_paused_root(¤t).await?;
2737 break;
2738 }
2739 info!(
2740 run_id = %run_id,
2741 ancestor_run_id = %parent_id,
2742 status = %status,
2743 "ancestor run failed with its child"
2744 );
2745 current = self
2746 .store
2747 .get_run(parent_id)
2748 .await?
2749 .ok_or(EngineError::Store(StoreError::RunNotFound(parent_id)))?;
2750 }
2751
2752 Ok(())
2753 }
2754
2755 #[cfg(feature = "prometheus")]
2757 fn emit_run_metrics(
2758 &self,
2759 workflow_name: &str,
2760 status: RunStatus,
2761 duration_ms: u64,
2762 ctx: &WorkflowContext,
2763 ) {
2764 let status_str = status.to_string();
2765 let wf = workflow_name.to_string();
2766
2767 counter!(RUNS_TOTAL, "workflow" => wf.clone(), "status" => status_str.clone()).increment(1);
2768 histogram!(RUN_DURATION_SECONDS, "workflow" => wf.clone(), "status" => status_str)
2769 .record(duration_ms as f64 / 1000.0);
2770 histogram!(RUN_COST_USD, "workflow" => wf.clone()).record(
2771 ctx.total_cost_usd()
2772 .to_string()
2773 .parse::<f64>()
2774 .unwrap_or(0.0),
2775 );
2776 gauge!(RUNS_ACTIVE, "workflow" => wf).decrement(1.0);
2777 }
2778
2779 fn on_run_budget_exceeded(&self, workflow_name: &str, run_id: Uuid, err: &EngineError) {
2785 let EngineError::RunBudgetExceeded {
2786 limit_usd,
2787 spent_usd,
2788 step_budget_usd,
2789 ..
2790 } = err
2791 else {
2792 return;
2793 };
2794
2795 #[cfg(feature = "prometheus")]
2796 counter!(
2797 RUN_BUDGET_EXCEEDED_TOTAL,
2798 "workflow" => workflow_name.to_string(),
2799 "scope" => "run",
2800 )
2801 .increment(1);
2802
2803 self.event_publisher
2804 .publish(Event::RunBudgetExceeded(RunBudgetExceededEvent {
2805 run_id,
2806 workflow_name: workflow_name.to_string(),
2807 limit_usd: *limit_usd,
2808 spent_usd: *spent_usd,
2809 step_budget_usd: *step_budget_usd,
2810 at: Utc::now(),
2811 }));
2812 }
2813
2814 #[allow(clippy::too_many_arguments)]
2819 fn publish_run_status_changed(
2820 &self,
2821 workflow_name: &str,
2822 run_id: Uuid,
2823 to: RunStatus,
2824 error: Option<String>,
2825 ctx: &WorkflowContext,
2826 duration_ms: u64,
2827 labels: HashMap<String, String>,
2828 ) {
2829 let now = Utc::now();
2830 let cost_usd = ctx.total_cost_usd();
2831 let wf = workflow_name.to_string();
2832
2833 self.event_publisher
2834 .publish(Event::RunStatusChanged(RunStatusChangedEvent {
2835 run_id,
2836 workflow_name: wf.clone(),
2837 from: RunStatus::Running,
2838 to,
2839 error: error.clone(),
2840 cost_usd,
2841 duration_ms,
2842 labels: labels.clone(),
2843 at: now,
2844 }));
2845
2846 if to == RunStatus::Failed {
2847 self.event_publisher
2848 .publish(Event::RunFailed(RunFailedEvent {
2849 run_id,
2850 workflow_name: wf,
2851 error,
2852 cost_usd,
2853 duration_ms,
2854 labels,
2855 at: now,
2856 }));
2857 }
2858 }
2859}
2860
2861impl fmt::Debug for Engine {
2862 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2863 f.debug_struct("Engine")
2864 .field("handlers", &self.handlers.keys().collect::<Vec<_>>())
2865 .finish_non_exhaustive()
2866 }
2867}
2868
2869#[cfg(test)]
2870mod tests {
2871 use super::*;
2872 use crate::config::ShellConfig;
2873 use crate::handler::{HandlerFuture, WorkflowHandler};
2874 use ironflow_core::providers::claude::ClaudeCodeProvider;
2875 use ironflow_core::providers::record_replay::RecordReplayProvider;
2876 use ironflow_store::memory::InMemoryStore;
2877 use ironflow_store::models::{MAX_PRIORITY, MIN_PRIORITY, StepApproval, StepStatus};
2878 use serde_json::json;
2879
2880 struct EchoWorkflow;
2882
2883 impl WorkflowHandler for EchoWorkflow {
2884 fn name(&self) -> &str {
2885 "echo-workflow"
2886 }
2887
2888 fn describe(&self) -> WorkflowInfo {
2889 WorkflowInfo {
2890 description: "A simple workflow that echoes hello".to_string(),
2891 source_code: None,
2892 sub_workflows: Vec::new(),
2893 category: None,
2894 version: self.version().map(str::to_string),
2895 compatible_versions: Vec::new(),
2896 input_schema: None,
2897 default_labels: HashMap::new(),
2898 schedule: self.schedule().cloned(),
2899 default_max_cost_usd: self.default_max_cost_usd(),
2900 priority: self.priority(),
2901 }
2902 }
2903
2904 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2905 Box::pin(async move {
2906 ctx.shell("greet", ShellConfig::new("echo hello")).await?;
2907 Ok(())
2908 })
2909 }
2910 }
2911
2912 struct FailingWorkflow;
2914
2915 impl WorkflowHandler for FailingWorkflow {
2916 fn name(&self) -> &str {
2917 "failing-workflow"
2918 }
2919
2920 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2921 Box::pin(async move {
2922 ctx.shell("fail", ShellConfig::new("exit 1")).await?;
2923 Ok(())
2924 })
2925 }
2926 }
2927
2928 fn create_test_engine() -> Engine {
2929 let store = Arc::new(InMemoryStore::new());
2930 let inner = ClaudeCodeProvider::new();
2931 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
2932 inner,
2933 "/tmp/ironflow-fixtures",
2934 ));
2935 Engine::new(store, provider)
2936 }
2937
2938 #[test]
2939 fn engine_new_creates_instance() {
2940 let engine = create_test_engine();
2941 assert_eq!(engine.handler_names().len(), 0);
2942 }
2943
2944 #[test]
2945 fn execution_mode_defaults_to_local() {
2946 let engine = create_test_engine();
2947 assert_eq!(engine.execution_mode(), ExecutionMode::Local);
2948 }
2949
2950 #[test]
2951 fn with_execution_mode_overrides_the_default() {
2952 let engine = create_test_engine().with_execution_mode(ExecutionMode::Workers);
2953 assert_eq!(engine.execution_mode(), ExecutionMode::Workers);
2954 }
2955
2956 #[test]
2957 fn engine_register_handler() {
2958 let mut engine = create_test_engine();
2959 let result = engine.register(EchoWorkflow);
2960 assert!(result.is_ok());
2961 assert_eq!(engine.handler_names().len(), 1);
2962 assert!(engine.handler_names().contains(&"echo-workflow"));
2963 }
2964
2965 #[test]
2966 fn engine_register_duplicate_returns_error() {
2967 let mut engine = create_test_engine();
2968 engine.register(EchoWorkflow).unwrap();
2969 let result = engine.register(EchoWorkflow);
2970 assert!(result.is_err());
2971 }
2972
2973 #[test]
2974 fn engine_get_handler_found() {
2975 let mut engine = create_test_engine();
2976 engine.register(EchoWorkflow).unwrap();
2977 let handler = engine.get_handler("echo-workflow");
2978 assert!(handler.is_some());
2979 }
2980
2981 #[test]
2982 fn engine_get_handler_not_found() {
2983 let engine = create_test_engine();
2984 let handler = engine.get_handler("nonexistent");
2985 assert!(handler.is_none());
2986 }
2987
2988 #[test]
2989 fn engine_handler_names_lists_all() {
2990 let mut engine = create_test_engine();
2991 engine.register(EchoWorkflow).unwrap();
2992 engine.register(FailingWorkflow).unwrap();
2993 let names = engine.handler_names();
2994 assert_eq!(names.len(), 2);
2995 assert!(names.contains(&"echo-workflow"));
2996 assert!(names.contains(&"failing-workflow"));
2997 }
2998
2999 #[test]
3000 fn engine_handler_info_returns_description() {
3001 let mut engine = create_test_engine();
3002 engine.register(EchoWorkflow).unwrap();
3003 let info = engine.handler_info("echo-workflow");
3004 assert!(info.is_some());
3005 let info = info.unwrap();
3006 assert_eq!(info.description, "A simple workflow that echoes hello");
3007 }
3008
3009 struct CategorizedWorkflow;
3010
3011 impl WorkflowHandler for CategorizedWorkflow {
3012 fn name(&self) -> &str {
3013 "categorized"
3014 }
3015 fn category(&self) -> Option<&str> {
3016 Some("data/etl")
3017 }
3018 fn execute<'a>(
3019 &'a self,
3020 _ctx: &'a mut WorkflowContext,
3021 ) -> crate::handler::HandlerFuture<'a> {
3022 Box::pin(async move { Ok(()) })
3023 }
3024 }
3025
3026 #[test]
3027 fn engine_default_describe_propagates_category() {
3028 let mut engine = create_test_engine();
3029 engine.register(CategorizedWorkflow).unwrap();
3030 let info = engine.handler_info("categorized").unwrap();
3031 assert_eq!(info.category.as_deref(), Some("data/etl"));
3032 }
3033
3034 #[test]
3035 fn engine_default_describe_without_category() {
3036 let mut engine = create_test_engine();
3037 engine.register(EchoWorkflow).unwrap();
3038 let info = engine.handler_info("echo-workflow").unwrap();
3039 assert!(info.category.is_none());
3040 }
3041
3042 struct ScheduledWorkflow {
3047 schedule: CronSchedule,
3048 }
3049
3050 impl ScheduledWorkflow {
3051 fn new() -> Self {
3052 Self {
3053 schedule: CronSchedule::new("0 0 * * * *").unwrap(),
3054 }
3055 }
3056 }
3057
3058 impl WorkflowHandler for ScheduledWorkflow {
3059 fn name(&self) -> &str {
3060 "scheduled"
3061 }
3062 fn schedule(&self) -> Option<&CronSchedule> {
3063 Some(&self.schedule)
3064 }
3065 fn execute<'a>(
3066 &'a self,
3067 _ctx: &'a mut WorkflowContext,
3068 ) -> crate::handler::HandlerFuture<'a> {
3069 Box::pin(async move { Ok(()) })
3070 }
3071 }
3072
3073 #[test]
3074 fn engine_default_describe_propagates_schedule() {
3075 let mut engine = create_test_engine();
3076 engine.register(ScheduledWorkflow::new()).unwrap();
3077 let info = engine.handler_info("scheduled").unwrap();
3078 assert_eq!(
3079 info.schedule.as_ref().map(|s| s.as_str()),
3080 Some("0 0 * * * *")
3081 );
3082 }
3083
3084 #[test]
3085 fn engine_default_describe_without_schedule() {
3086 let mut engine = create_test_engine();
3087 engine.register(EchoWorkflow).unwrap();
3088 let info = engine.handler_info("echo-workflow").unwrap();
3089 assert!(info.schedule.is_none());
3090 }
3091
3092 #[test]
3093 fn scheduled_handlers_returns_only_scheduled() {
3094 let mut engine = create_test_engine();
3095 engine.register(EchoWorkflow).unwrap();
3096 engine.register(ScheduledWorkflow::new()).unwrap();
3097 engine.register(FailingWorkflow).unwrap();
3098
3099 let scheduled = engine.scheduled_handlers();
3100 assert_eq!(scheduled.len(), 1);
3101 assert_eq!(scheduled[0].0, "scheduled");
3102 assert_eq!(scheduled[0].1.as_str(), "0 0 * * * *");
3103 }
3104
3105 #[test]
3106 fn scheduled_handlers_empty_when_none_scheduled() {
3107 let mut engine = create_test_engine();
3108 engine.register(EchoWorkflow).unwrap();
3109 engine.register(FailingWorkflow).unwrap();
3110
3111 let scheduled = engine.scheduled_handlers();
3112 assert!(scheduled.is_empty());
3113 }
3114
3115 struct BadCategoryWorkflow(&'static str);
3116
3117 impl WorkflowHandler for BadCategoryWorkflow {
3118 fn name(&self) -> &str {
3119 "bad-category"
3120 }
3121 fn category(&self) -> Option<&str> {
3122 Some(self.0)
3123 }
3124 fn execute<'a>(
3125 &'a self,
3126 _ctx: &'a mut WorkflowContext,
3127 ) -> crate::handler::HandlerFuture<'a> {
3128 Box::pin(async move { Ok(()) })
3129 }
3130 }
3131
3132 #[test]
3133 fn engine_register_rejects_empty_category() {
3134 let mut engine = create_test_engine();
3135 let err = engine.register(BadCategoryWorkflow("")).unwrap_err();
3136 match err {
3137 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty category")),
3138 other => panic!("expected InvalidWorkflow, got {other:?}"),
3139 }
3140 }
3141
3142 #[test]
3143 fn engine_register_rejects_leading_slash_category() {
3144 let mut engine = create_test_engine();
3145 let err = engine
3146 .register(BadCategoryWorkflow("/data/etl"))
3147 .unwrap_err();
3148 match err {
3149 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("leading '/'")),
3150 other => panic!("expected InvalidWorkflow, got {other:?}"),
3151 }
3152 }
3153
3154 #[test]
3155 fn engine_register_rejects_trailing_slash_category() {
3156 let mut engine = create_test_engine();
3157 let err = engine
3158 .register(BadCategoryWorkflow("data/etl/"))
3159 .unwrap_err();
3160 match err {
3161 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("trailing '/'")),
3162 other => panic!("expected InvalidWorkflow, got {other:?}"),
3163 }
3164 }
3165
3166 #[test]
3167 fn engine_register_rejects_double_slash_category() {
3168 let mut engine = create_test_engine();
3169 let err = engine
3170 .register(BadCategoryWorkflow("data//etl"))
3171 .unwrap_err();
3172 match err {
3173 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty segment")),
3174 other => panic!("expected InvalidWorkflow, got {other:?}"),
3175 }
3176 }
3177
3178 #[test]
3179 fn engine_register_rejects_whitespace_only_segment_category() {
3180 let mut engine = create_test_engine();
3181 let err = engine
3182 .register(BadCategoryWorkflow("data/ /etl"))
3183 .unwrap_err();
3184 match err {
3185 EngineError::InvalidWorkflow(msg) => assert!(msg.contains("whitespace-only segment")),
3186 other => panic!("expected InvalidWorkflow, got {other:?}"),
3187 }
3188 }
3189
3190 #[test]
3191 fn engine_register_accepts_valid_nested_category() {
3192 let mut engine = create_test_engine();
3193 assert!(engine.register(CategorizedWorkflow).is_ok());
3194 }
3195
3196 #[tokio::test]
3197 async fn engine_unknown_workflow_returns_error() {
3198 let engine = create_test_engine();
3199 let result = engine
3200 .run_handler("unknown", TriggerKind::Manual, json!({}))
3201 .await;
3202 assert!(result.is_err());
3203 match result {
3204 Err(EngineError::InvalidWorkflow(msg)) => {
3205 assert!(msg.contains("no handler registered"));
3206 }
3207 _ => panic!("expected InvalidWorkflow error"),
3208 }
3209 }
3210
3211 #[tokio::test]
3212 async fn engine_enqueue_handler_creates_pending_run() {
3213 let mut engine = create_test_engine();
3214 engine.register(EchoWorkflow).unwrap();
3215
3216 let run = engine
3217 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
3218 .await
3219 .unwrap();
3220 assert_eq!(run.status.state, RunStatus::Pending);
3221 assert_eq!(run.workflow_name, "echo-workflow");
3222 }
3223
3224 #[tokio::test]
3225 async fn enqueue_handler_leaves_the_run_unattributed() {
3226 let mut engine = create_test_engine();
3227 engine.register(EchoWorkflow).unwrap();
3228
3229 let run = engine
3230 .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
3231 .await
3232 .unwrap();
3233
3234 assert!(run.created_by.is_none());
3235 }
3236
3237 #[tokio::test]
3238 async fn enqueue_handler_with_options_records_the_author() {
3239 let mut engine = create_test_engine();
3240 engine.register(EchoWorkflow).unwrap();
3241 let actor = RunActor::User {
3242 user_id: Uuid::now_v7(),
3243 };
3244
3245 let run = engine
3246 .enqueue_handler_with_options(
3247 "echo-workflow",
3248 TriggerKind::Api,
3249 json!({}),
3250 EnqueueOptions {
3251 created_by: Some(actor.clone()),
3252 ..Default::default()
3253 },
3254 )
3255 .await
3256 .unwrap()
3257 .into_run();
3258
3259 assert_eq!(run.created_by, Some(actor));
3260 }
3261
3262 #[tokio::test]
3263 async fn enqueue_handler_with_options_accepts_no_author() {
3264 let mut engine = create_test_engine();
3265 engine.register(EchoWorkflow).unwrap();
3266
3267 let run = engine
3268 .enqueue_handler_with_options(
3269 "echo-workflow",
3270 TriggerKind::Cron {
3271 schedule: "0 * * * * *".to_string(),
3272 schedule_id: None,
3273 scheduled_for: None,
3274 },
3275 json!({}),
3276 EnqueueOptions::default(),
3277 )
3278 .await
3279 .unwrap()
3280 .into_run();
3281
3282 assert!(run.created_by.is_none());
3283 }
3284
3285 #[tokio::test]
3286 async fn enqueue_handler_with_options_stores_concurrency_limits() {
3287 let mut engine = create_test_engine();
3288 engine.register(EchoWorkflow).unwrap();
3289 let limits = vec![
3290 ConcurrencyLimit::new("repo:acme", 2),
3291 ConcurrencyLimit::new("tenant:42", 5),
3292 ];
3293
3294 let run = engine
3295 .enqueue_handler_with_options(
3296 "echo-workflow",
3297 TriggerKind::Api,
3298 json!({}),
3299 EnqueueOptions {
3300 concurrency_limits: limits.clone(),
3301 ..Default::default()
3302 },
3303 )
3304 .await
3305 .unwrap()
3306 .into_run();
3307
3308 assert_eq!(run.concurrency_limits, limits);
3309 }
3310
3311 #[tokio::test]
3312 async fn enqueue_rejects_invalid_concurrency_limits() {
3313 let mut engine = create_test_engine();
3314 engine.register(EchoWorkflow).unwrap();
3315
3316 let invalid = [
3317 vec![ConcurrencyLimit::new("repo:acme", 0)],
3318 vec![ConcurrencyLimit::new("", 1)],
3319 vec![
3320 ConcurrencyLimit::new("repo:acme", 1),
3321 ConcurrencyLimit::new("repo:acme", 2),
3322 ],
3323 ];
3324 for concurrency_limits in invalid {
3325 let err = engine
3326 .enqueue_handler_with_options(
3327 "echo-workflow",
3328 TriggerKind::Api,
3329 json!({}),
3330 EnqueueOptions {
3331 concurrency_limits,
3332 ..Default::default()
3333 },
3334 )
3335 .await
3336 .unwrap_err();
3337 assert!(
3338 matches!(err, EngineError::InvalidConcurrencyLimit(_)),
3339 "{err:?}"
3340 );
3341 }
3342
3343 let err = engine
3345 .enqueue_handler_with_options(
3346 "not-registered",
3347 TriggerKind::Api,
3348 json!({}),
3349 EnqueueOptions {
3350 concurrency_limits: vec![ConcurrencyLimit::new("repo:acme", 0)],
3351 ..Default::default()
3352 },
3353 )
3354 .await
3355 .unwrap_err();
3356 assert!(
3357 matches!(err, EngineError::InvalidConcurrencyLimit(_)),
3358 "{err:?}"
3359 );
3360
3361 let page = engine
3362 .store()
3363 .list_runs(RunFilter::default(), 1, 10)
3364 .await
3365 .unwrap();
3366 assert_eq!(page.total, 0, "no run may be created");
3367 }
3368
3369 struct UrgentWorkflow;
3370
3371 impl WorkflowHandler for UrgentWorkflow {
3372 fn name(&self) -> &str {
3373 "urgent-workflow"
3374 }
3375
3376 fn priority(&self) -> i16 {
3377 60
3378 }
3379
3380 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3381 Box::pin(async { Ok(()) })
3382 }
3383 }
3384
3385 #[tokio::test]
3386 async fn enqueue_priority_defaults_to_the_handler_priority() {
3387 let mut engine = create_test_engine();
3388 engine.register(EchoWorkflow).unwrap();
3389 engine.register(UrgentWorkflow).unwrap();
3390
3391 let echo = engine
3392 .enqueue_handler_with_options(
3393 "echo-workflow",
3394 TriggerKind::Api,
3395 json!({}),
3396 EnqueueOptions::default(),
3397 )
3398 .await
3399 .unwrap()
3400 .into_run();
3401 assert_eq!(echo.priority, 0);
3402
3403 let urgent = engine
3404 .enqueue_handler_with_options(
3405 "urgent-workflow",
3406 TriggerKind::Api,
3407 json!({}),
3408 EnqueueOptions::default(),
3409 )
3410 .await
3411 .unwrap()
3412 .into_run();
3413 assert_eq!(urgent.priority, 60);
3414 }
3415
3416 #[tokio::test]
3417 async fn enqueue_priority_explicit_value_overrides_the_handler() {
3418 let mut engine = create_test_engine();
3419 engine.register(UrgentWorkflow).unwrap();
3420
3421 let run = engine
3422 .enqueue_handler_with_options(
3423 "urgent-workflow",
3424 TriggerKind::Api,
3425 json!({}),
3426 EnqueueOptions {
3427 priority: Some(-20),
3428 ..Default::default()
3429 },
3430 )
3431 .await
3432 .unwrap()
3433 .into_run();
3434 assert_eq!(run.priority, -20);
3435
3436 let stored = engine.store().get_run(run.id).await.unwrap().unwrap();
3437 assert_eq!(stored.priority, -20);
3438 }
3439
3440 #[tokio::test]
3441 async fn enqueue_priority_out_of_range_is_rejected() {
3442 let mut engine = create_test_engine();
3443 engine.register(EchoWorkflow).unwrap();
3444
3445 for priority in [MAX_PRIORITY + 1, MIN_PRIORITY - 1] {
3446 let err = engine
3447 .enqueue_handler_with_options(
3448 "echo-workflow",
3449 TriggerKind::Api,
3450 json!({}),
3451 EnqueueOptions {
3452 priority: Some(priority),
3453 ..Default::default()
3454 },
3455 )
3456 .await
3457 .unwrap_err();
3458 assert!(matches!(err, EngineError::InvalidPriority(_)), "{err:?}");
3459 }
3460
3461 let err = engine
3463 .enqueue_handler_with_options(
3464 "not-registered",
3465 TriggerKind::Api,
3466 json!({}),
3467 EnqueueOptions {
3468 priority: Some(MAX_PRIORITY + 1),
3469 ..Default::default()
3470 },
3471 )
3472 .await
3473 .unwrap_err();
3474 assert!(matches!(err, EngineError::InvalidPriority(_)), "{err:?}");
3475
3476 let page = engine
3477 .store()
3478 .list_runs(RunFilter::default(), 1, 10)
3479 .await
3480 .unwrap();
3481 assert_eq!(page.total, 0, "no run may be created");
3482 }
3483
3484 #[tokio::test]
3485 async fn enqueue_priority_bounds_are_accepted() {
3486 let mut engine = create_test_engine();
3487 engine.register(EchoWorkflow).unwrap();
3488
3489 for priority in [MIN_PRIORITY, MAX_PRIORITY] {
3490 let run = engine
3491 .enqueue_handler_with_options(
3492 "echo-workflow",
3493 TriggerKind::Api,
3494 json!({}),
3495 EnqueueOptions {
3496 priority: Some(priority),
3497 ..Default::default()
3498 },
3499 )
3500 .await
3501 .unwrap()
3502 .into_run();
3503 assert_eq!(run.priority, priority);
3504 }
3505 }
3506
3507 struct GpuWorkflow;
3508
3509 impl WorkflowHandler for GpuWorkflow {
3510 fn name(&self) -> &str {
3511 "gpu-workflow"
3512 }
3513
3514 fn required_worker_tags(&self) -> Vec<String> {
3515 vec!["gpu".to_string()]
3516 }
3517
3518 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3519 Box::pin(async move { Ok(()) })
3520 }
3521 }
3522
3523 #[tokio::test]
3524 async fn enqueue_merges_handler_and_request_worker_tags() {
3525 let mut engine = create_test_engine();
3526 engine.register(GpuWorkflow).unwrap();
3527
3528 let run = engine
3529 .enqueue_handler_with_options(
3530 "gpu-workflow",
3531 TriggerKind::Api,
3532 json!({}),
3533 EnqueueOptions {
3534 worker_tags: vec!["region:eu".to_string(), "gpu".to_string()],
3535 ..Default::default()
3536 },
3537 )
3538 .await
3539 .unwrap()
3540 .into_run();
3541
3542 assert_eq!(
3543 run.worker_tags,
3544 vec!["gpu".to_string(), "region:eu".to_string()]
3545 );
3546 }
3547
3548 #[tokio::test]
3549 async fn enqueue_without_worker_tags_keeps_handler_tags() {
3550 let mut engine = create_test_engine();
3551 engine.register(GpuWorkflow).unwrap();
3552 engine.register(EchoWorkflow).unwrap();
3553
3554 let gpu = engine
3555 .enqueue_handler("gpu-workflow", TriggerKind::Api, json!({}), 0)
3556 .await
3557 .unwrap();
3558 assert_eq!(gpu.worker_tags, vec!["gpu".to_string()]);
3559
3560 let echo = engine
3561 .enqueue_handler("echo-workflow", TriggerKind::Api, json!({}), 0)
3562 .await
3563 .unwrap();
3564 assert!(echo.worker_tags.is_empty());
3565 }
3566
3567 #[tokio::test]
3568 async fn enqueue_rejects_invalid_worker_tags() {
3569 let mut engine = create_test_engine();
3570 engine.register(EchoWorkflow).unwrap();
3571
3572 for worker_tags in [
3573 vec!["bad,tag".to_string()],
3574 vec![" ".to_string()],
3575 vec!["x".repeat(65)],
3576 ] {
3577 let err = engine
3578 .enqueue_handler_with_options(
3579 "echo-workflow",
3580 TriggerKind::Api,
3581 json!({}),
3582 EnqueueOptions {
3583 worker_tags,
3584 ..Default::default()
3585 },
3586 )
3587 .await
3588 .unwrap_err();
3589 assert!(matches!(err, EngineError::InvalidWorkerTag(_)), "{err:?}");
3590 }
3591
3592 let err = engine
3594 .enqueue_handler_with_options(
3595 "not-registered",
3596 TriggerKind::Api,
3597 json!({}),
3598 EnqueueOptions {
3599 worker_tags: vec!["bad,tag".to_string()],
3600 ..Default::default()
3601 },
3602 )
3603 .await
3604 .unwrap_err();
3605 assert!(matches!(err, EngineError::InvalidWorkerTag(_)), "{err:?}");
3606 }
3607
3608 #[test]
3609 fn worker_tags_are_unset_by_default() {
3610 let engine = create_test_engine();
3611 assert!(engine.worker_tags().is_none());
3612 }
3613
3614 #[test]
3615 fn set_worker_tags_stores_the_tags() {
3616 let mut engine = create_test_engine();
3617 engine.set_worker_tags(vec!["gpu".to_string()]);
3618 assert_eq!(engine.worker_tags(), Some(&["gpu".to_string()][..]));
3619
3620 engine.set_worker_tags(Vec::new());
3621 assert_eq!(engine.worker_tags(), Some(&[][..]));
3622 }
3623
3624 #[tokio::test]
3625 async fn run_handler_records_handler_worker_tags() {
3626 let mut engine = create_test_engine();
3627 engine.register(GpuWorkflow).unwrap();
3628
3629 let result = engine
3630 .run_handler("gpu-workflow", TriggerKind::Manual, json!({}))
3631 .await
3632 .unwrap();
3633 assert_eq!(result.run.worker_tags, vec!["gpu".to_string()]);
3634 }
3635
3636 #[tokio::test]
3637 async fn run_handler_leaves_the_run_unattributed() {
3638 let mut engine = create_test_engine();
3639 engine.register(EchoWorkflow).unwrap();
3640
3641 let run = engine
3642 .run_handler("echo-workflow", TriggerKind::Manual, json!({}))
3643 .await
3644 .unwrap()
3645 .run;
3646
3647 assert!(run.created_by.is_none());
3648 }
3649
3650 #[tokio::test]
3651 async fn run_handler_priority_comes_from_the_handler() {
3652 let mut engine = create_test_engine();
3653 engine.register(UrgentWorkflow).unwrap();
3654
3655 let run = engine
3656 .run_handler("urgent-workflow", TriggerKind::Manual, json!({}))
3657 .await
3658 .unwrap()
3659 .run;
3660
3661 assert_eq!(run.priority, 60);
3662 }
3663
3664 #[tokio::test]
3665 async fn engine_register_boxed() {
3666 let mut engine = create_test_engine();
3667 let handler: Box<dyn WorkflowHandler> = Box::new(EchoWorkflow);
3668 let result = engine.register_boxed(handler);
3669 assert!(result.is_ok());
3670 assert_eq!(engine.handler_names().len(), 1);
3671 }
3672
3673 #[tokio::test]
3674 async fn engine_store_and_provider_accessors() {
3675 let store = Arc::new(InMemoryStore::new());
3676 let inner = ClaudeCodeProvider::new();
3677 let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
3678 inner,
3679 "/tmp/ironflow-fixtures",
3680 ));
3681 let engine = Engine::new(store.clone(), provider.clone());
3682
3683 let _ = engine.store();
3685 let _ = engine.provider();
3686 }
3687
3688 use crate::operation::{Operation, OperationContext};
3693 use async_trait::async_trait;
3694 use ironflow_core::error::OperationError;
3695 use ironflow_store::models::StepKind;
3696
3697 struct FakeGitlabOp {
3698 project_id: u64,
3699 title: String,
3700 }
3701
3702 #[async_trait]
3703 impl Operation for FakeGitlabOp {
3704 fn kind(&self) -> &str {
3705 "gitlab"
3706 }
3707
3708 async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
3709 Ok(json!({
3710 "issue_id": 42,
3711 "project_id": self.project_id,
3712 "title": self.title,
3713 }))
3714 }
3715
3716 fn input(&self) -> Option<Value> {
3717 Some(json!({
3718 "project_id": self.project_id,
3719 "title": self.title,
3720 }))
3721 }
3722 }
3723
3724 struct FailingOp;
3725
3726 #[async_trait]
3727 impl Operation for FailingOp {
3728 fn kind(&self) -> &str {
3729 "broken-service"
3730 }
3731
3732 async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
3733 Err(OperationError::Http {
3734 status: None,
3735 message: "service unavailable".to_string(),
3736 })
3737 }
3738 }
3739
3740 struct OperationWorkflow;
3741
3742 impl WorkflowHandler for OperationWorkflow {
3743 fn name(&self) -> &str {
3744 "operation-workflow"
3745 }
3746
3747 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3748 Box::pin(async move {
3749 let op = FakeGitlabOp {
3750 project_id: 123,
3751 title: "Bug report".to_string(),
3752 };
3753 ctx.operation("create-issue", &op).await?;
3754 Ok(())
3755 })
3756 }
3757 }
3758
3759 struct FailingOperationWorkflow;
3760
3761 impl WorkflowHandler for FailingOperationWorkflow {
3762 fn name(&self) -> &str {
3763 "failing-operation-workflow"
3764 }
3765
3766 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3767 Box::pin(async move {
3768 ctx.operation("broken-call", &FailingOp).await?;
3769 Ok(())
3770 })
3771 }
3772 }
3773
3774 struct MixedWorkflow;
3775
3776 impl WorkflowHandler for MixedWorkflow {
3777 fn name(&self) -> &str {
3778 "mixed-workflow"
3779 }
3780
3781 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3782 Box::pin(async move {
3783 ctx.shell("build", ShellConfig::new("echo built")).await?;
3784 let op = FakeGitlabOp {
3785 project_id: 456,
3786 title: "Deploy done".to_string(),
3787 };
3788 let result = ctx.operation("notify-gitlab", &op).await?;
3789 assert_eq!(result.output["issue_id"], 42);
3790 Ok(())
3791 })
3792 }
3793 }
3794
3795 #[tokio::test]
3796 async fn operation_step_happy_path() {
3797 let mut engine = create_test_engine();
3798 engine.register(OperationWorkflow).unwrap();
3799
3800 let run = engine
3801 .run_handler("operation-workflow", TriggerKind::Manual, json!({}))
3802 .await
3803 .unwrap()
3804 .run;
3805
3806 assert_eq!(run.status.state, RunStatus::Completed);
3807
3808 let steps = engine.store().list_steps(run.id).await.unwrap();
3809
3810 assert_eq!(steps.len(), 1);
3811 assert_eq!(steps[0].name, "create-issue");
3812 assert_eq!(steps[0].kind, StepKind::Custom("gitlab".to_string()));
3813 assert_eq!(
3814 steps[0].status.state,
3815 ironflow_store::models::StepStatus::Completed
3816 );
3817
3818 let output = steps[0].output.as_ref().unwrap();
3819 assert_eq!(output["issue_id"], 42);
3820 assert_eq!(output["project_id"], 123);
3821
3822 let input = steps[0].input.as_ref().unwrap();
3823 assert_eq!(input["project_id"], 123);
3824 assert_eq!(input["title"], "Bug report");
3825 }
3826
3827 #[tokio::test]
3828 async fn operation_step_failure_marks_run_failed() {
3829 let mut engine = create_test_engine();
3830 engine.register(FailingOperationWorkflow).unwrap();
3831
3832 let result = engine
3833 .run_handler("failing-operation-workflow", TriggerKind::Manual, json!({}))
3834 .await;
3835
3836 assert!(result.is_err());
3837 }
3838
3839 #[tokio::test]
3840 async fn operation_mixed_with_shell_steps() {
3841 let mut engine = create_test_engine();
3842 engine.register(MixedWorkflow).unwrap();
3843
3844 let run = engine
3845 .run_handler("mixed-workflow", TriggerKind::Manual, json!({}))
3846 .await
3847 .unwrap()
3848 .run;
3849
3850 assert_eq!(run.status.state, RunStatus::Completed);
3851
3852 let steps = engine.store().list_steps(run.id).await.unwrap();
3853
3854 assert_eq!(steps.len(), 2);
3855 assert_eq!(steps[0].kind, StepKind::Shell);
3856 assert_eq!(steps[1].kind, StepKind::Custom("gitlab".to_string()));
3857 assert_eq!(steps[0].position, 0);
3858 assert_eq!(steps[1].position, 1);
3859 }
3860
3861 use crate::config::ApprovalConfig;
3866
3867 struct SingleApprovalWorkflow;
3868
3869 impl WorkflowHandler for SingleApprovalWorkflow {
3870 fn name(&self) -> &str {
3871 "single-approval"
3872 }
3873
3874 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3875 Box::pin(async move {
3876 ctx.shell("build", ShellConfig::new("echo built")).await?;
3877 ctx.approval("gate", ApprovalConfig::new("OK?")).await?;
3878 ctx.shell("deploy", ShellConfig::new("echo deployed"))
3879 .await?;
3880 Ok(())
3881 })
3882 }
3883 }
3884
3885 struct DoubleApprovalWorkflow;
3886
3887 impl WorkflowHandler for DoubleApprovalWorkflow {
3888 fn name(&self) -> &str {
3889 "double-approval"
3890 }
3891
3892 fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
3893 Box::pin(async move {
3894 ctx.shell("build", ShellConfig::new("echo built")).await?;
3895 ctx.approval("staging-gate", ApprovalConfig::new("Deploy staging?"))
3896 .await?;
3897 ctx.shell("deploy-staging", ShellConfig::new("echo staging"))
3898 .await?;
3899 ctx.approval("prod-gate", ApprovalConfig::new("Deploy prod?"))
3900 .await?;
3901 ctx.shell("deploy-prod", ShellConfig::new("echo prod"))
3902 .await?;
3903 Ok(())
3904 })
3905 }
3906 }
3907
3908 #[tokio::test]
3909 async fn approval_pauses_run() {
3910 let mut engine = create_test_engine();
3911 engine.register(SingleApprovalWorkflow).unwrap();
3912
3913 let run = engine
3914 .run_handler("single-approval", TriggerKind::Manual, json!({}))
3915 .await
3916 .unwrap()
3917 .run;
3918
3919 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
3920
3921 let steps = engine.store().list_steps(run.id).await.unwrap();
3922 assert_eq!(steps.len(), 2); assert_eq!(steps[0].kind, StepKind::Shell);
3924 assert_eq!(steps[0].status.state, StepStatus::Completed);
3925 assert_eq!(steps[1].kind, StepKind::Approval);
3926 assert_eq!(steps[1].status.state, StepStatus::AwaitingApproval);
3927 }
3928
3929 async fn approve_open_gate(store: &Arc<dyn Store>, run_id: Uuid) {
3932 let gate = store
3933 .list_steps(run_id)
3934 .await
3935 .unwrap()
3936 .into_iter()
3937 .find(|step| step.status.state == StepStatus::AwaitingApproval)
3938 .expect("an open approval gate");
3939 store
3940 .record_step_approval(
3941 gate.id,
3942 StepApproval {
3943 user_id: Uuid::now_v7(),
3944 approved_by: "tester".to_string(),
3945 at: Utc::now(),
3946 },
3947 )
3948 .await
3949 .unwrap();
3950 store
3951 .update_run_status(run_id, RunStatus::Running)
3952 .await
3953 .unwrap();
3954 }
3955
3956 #[tokio::test]
3957 async fn approval_resume_completes_run() {
3958 let mut engine = create_test_engine();
3959 engine.register(SingleApprovalWorkflow).unwrap();
3960
3961 let run = engine
3963 .run_handler("single-approval", TriggerKind::Manual, json!({}))
3964 .await
3965 .unwrap()
3966 .run;
3967 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
3968
3969 approve_open_gate(engine.store(), run.id).await;
3971
3972 let resumed = engine.resume_run(run.id).await.unwrap().run;
3974 assert_eq!(resumed.status.state, RunStatus::Completed);
3975
3976 let steps = engine.store().list_steps(run.id).await.unwrap();
3977 assert_eq!(steps.len(), 3); assert_eq!(steps[0].name, "build");
3979 assert_eq!(steps[0].status.state, StepStatus::Completed);
3980 assert_eq!(steps[1].name, "gate");
3981 assert_eq!(steps[1].kind, StepKind::Approval);
3982 assert_eq!(steps[1].status.state, StepStatus::Completed);
3983 assert_eq!(steps[2].name, "deploy");
3984 assert_eq!(steps[2].status.state, StepStatus::Completed);
3985 }
3986
3987 #[tokio::test]
3988 async fn double_approval_two_resumes() {
3989 let mut engine = create_test_engine();
3990 engine.register(DoubleApprovalWorkflow).unwrap();
3991
3992 let run = engine
3994 .run_handler("double-approval", TriggerKind::Manual, json!({}))
3995 .await
3996 .unwrap()
3997 .run;
3998 assert_eq!(run.status.state, RunStatus::AwaitingApproval);
3999
4000 let steps = engine.store().list_steps(run.id).await.unwrap();
4001 assert_eq!(steps.len(), 2); approve_open_gate(engine.store(), run.id).await;
4005
4006 let resumed = engine.resume_run(run.id).await.unwrap().run;
4007 assert_eq!(resumed.status.state, RunStatus::AwaitingApproval);
4008
4009 let steps = engine.store().list_steps(run.id).await.unwrap();
4010 assert_eq!(steps.len(), 4); approve_open_gate(engine.store(), run.id).await;
4014
4015 let final_run = engine.resume_run(run.id).await.unwrap().run;
4016 assert_eq!(final_run.status.state, RunStatus::Completed);
4017
4018 let steps = engine.store().list_steps(run.id).await.unwrap();
4019 assert_eq!(steps.len(), 5);
4020 assert_eq!(steps[0].name, "build");
4021 assert_eq!(steps[1].name, "staging-gate");
4022 assert_eq!(steps[2].name, "deploy-staging");
4023 assert_eq!(steps[3].name, "prod-gate");
4024 assert_eq!(steps[4].name, "deploy-prod");
4025
4026 for step in &steps {
4027 assert_eq!(step.status.state, StepStatus::Completed);
4028 }
4029 }
4030
4031 use ironflow_store::models::{NewStep, StepUpdate, step_trace_id};
4036
4037 async fn create_step_with_status(
4038 store: &Arc<dyn Store>,
4039 run_id: Uuid,
4040 name: &str,
4041 position: u32,
4042 status: StepStatus,
4043 ) -> ironflow_store::models::Step {
4044 let step = store
4045 .create_step(NewStep {
4046 run_id,
4047 trace_id: step_trace_id(run_id, name, position),
4048 name: name.to_string(),
4049 kind: StepKind::Shell,
4050 position,
4051 input: None,
4052 is_error_handler: false,
4053 })
4054 .await
4055 .unwrap();
4056
4057 match status {
4058 StepStatus::Pending => {}
4059 StepStatus::Running => {
4060 store
4061 .update_step(
4062 step.id,
4063 StepUpdate {
4064 status: Some(StepStatus::Running),
4065 ..StepUpdate::default()
4066 },
4067 )
4068 .await
4069 .unwrap();
4070 }
4071 StepStatus::Completed => {
4072 store
4073 .update_step(
4074 step.id,
4075 StepUpdate {
4076 status: Some(StepStatus::Running),
4077 ..StepUpdate::default()
4078 },
4079 )
4080 .await
4081 .unwrap();
4082 store
4083 .update_step(
4084 step.id,
4085 StepUpdate {
4086 status: Some(StepStatus::Completed),
4087 ..StepUpdate::default()
4088 },
4089 )
4090 .await
4091 .unwrap();
4092 }
4093 StepStatus::AwaitingApproval => {
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::AwaitingApproval),
4109 ..StepUpdate::default()
4110 },
4111 )
4112 .await
4113 .unwrap();
4114 }
4115 _ => panic!("unsupported status for test helper: {status}"),
4116 }
4117
4118 store.get_step(step.id).await.unwrap().unwrap()
4119 }
4120
4121 #[tokio::test]
4122 async fn fail_orphaned_steps_marks_running_as_failed() {
4123 let engine = create_test_engine();
4124 let run = engine
4125 .store()
4126 .create_run(NewRun {
4127 created_by: None,
4128 workflow_name: "test".to_string(),
4129 trigger: TriggerKind::Manual,
4130 payload: json!({}),
4131 max_retries: 0,
4132 handler_version: None,
4133 labels: HashMap::new(),
4134 scheduled_at: None,
4135 idempotency_key: None,
4136 concurrency_key: None,
4137 priority: 0,
4138 concurrency_limits: Vec::new(),
4139 max_cost_usd: None,
4140 worker_tags: Vec::new(),
4141 })
4142 .await
4143 .unwrap()
4144 .into_run();
4145
4146 let step = create_step_with_status(
4147 engine.store(),
4148 run.id,
4149 "running-step",
4150 0,
4151 StepStatus::Running,
4152 )
4153 .await;
4154
4155 engine
4156 .fail_orphaned_steps(run.id, "parent run timed out")
4157 .await
4158 .unwrap();
4159
4160 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4161 assert_eq!(updated.status.state, StepStatus::Failed);
4162 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
4163 assert!(updated.completed_at.is_some());
4164 }
4165
4166 #[tokio::test]
4167 async fn fail_orphaned_steps_marks_pending_as_skipped() {
4168 let engine = create_test_engine();
4169 let run = engine
4170 .store()
4171 .create_run(NewRun {
4172 created_by: None,
4173 workflow_name: "test".to_string(),
4174 trigger: TriggerKind::Manual,
4175 payload: json!({}),
4176 max_retries: 0,
4177 handler_version: None,
4178 labels: HashMap::new(),
4179 scheduled_at: None,
4180 idempotency_key: None,
4181 concurrency_key: None,
4182 priority: 0,
4183 concurrency_limits: Vec::new(),
4184 max_cost_usd: None,
4185 worker_tags: Vec::new(),
4186 })
4187 .await
4188 .unwrap()
4189 .into_run();
4190
4191 let step = create_step_with_status(
4192 engine.store(),
4193 run.id,
4194 "pending-step",
4195 0,
4196 StepStatus::Pending,
4197 )
4198 .await;
4199
4200 engine
4201 .fail_orphaned_steps(run.id, "parent run timed out")
4202 .await
4203 .unwrap();
4204
4205 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4206 assert_eq!(updated.status.state, StepStatus::Skipped);
4207 assert!(updated.error.is_none());
4208 assert!(updated.completed_at.is_some());
4209 }
4210
4211 #[tokio::test]
4212 async fn fail_orphaned_steps_marks_awaiting_approval_as_failed() {
4213 let engine = create_test_engine();
4214 let run = engine
4215 .store()
4216 .create_run(NewRun {
4217 created_by: None,
4218 workflow_name: "test".to_string(),
4219 trigger: TriggerKind::Manual,
4220 payload: json!({}),
4221 max_retries: 0,
4222 handler_version: None,
4223 labels: HashMap::new(),
4224 scheduled_at: None,
4225 idempotency_key: None,
4226 concurrency_key: None,
4227 priority: 0,
4228 concurrency_limits: Vec::new(),
4229 max_cost_usd: None,
4230 worker_tags: Vec::new(),
4231 })
4232 .await
4233 .unwrap()
4234 .into_run();
4235
4236 let step = create_step_with_status(
4237 engine.store(),
4238 run.id,
4239 "approval-step",
4240 0,
4241 StepStatus::AwaitingApproval,
4242 )
4243 .await;
4244
4245 engine
4246 .fail_orphaned_steps(run.id, "parent run timed out")
4247 .await
4248 .unwrap();
4249
4250 let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
4251 assert_eq!(updated.status.state, StepStatus::Failed);
4252 assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
4253 assert!(updated.completed_at.is_some());
4254 }
4255
4256 #[tokio::test]
4257 async fn fail_orphaned_steps_skips_terminal_steps() {
4258 let engine = create_test_engine();
4259 let run = engine
4260 .store()
4261 .create_run(NewRun {
4262 created_by: None,
4263 workflow_name: "test".to_string(),
4264 trigger: TriggerKind::Manual,
4265 payload: json!({}),
4266 max_retries: 0,
4267 handler_version: None,
4268 labels: HashMap::new(),
4269 scheduled_at: None,
4270 idempotency_key: None,
4271 concurrency_key: None,
4272 priority: 0,
4273 concurrency_limits: Vec::new(),
4274 max_cost_usd: None,
4275 worker_tags: Vec::new(),
4276 })
4277 .await
4278 .unwrap()
4279 .into_run();
4280
4281 let completed_step =
4282 create_step_with_status(engine.store(), run.id, "done", 0, StepStatus::Completed).await;
4283 let running_step =
4284 create_step_with_status(engine.store(), run.id, "in-flight", 1, StepStatus::Running)
4285 .await;
4286
4287 engine
4288 .fail_orphaned_steps(run.id, "parent run timed out")
4289 .await
4290 .unwrap();
4291
4292 let completed = engine
4293 .store()
4294 .get_step(completed_step.id)
4295 .await
4296 .unwrap()
4297 .unwrap();
4298 assert_eq!(completed.status.state, StepStatus::Completed);
4299
4300 let failed = engine
4301 .store()
4302 .get_step(running_step.id)
4303 .await
4304 .unwrap()
4305 .unwrap();
4306 assert_eq!(failed.status.state, StepStatus::Failed);
4307 }
4308
4309 #[tokio::test]
4310 async fn fail_orphaned_steps_mixed_states() {
4311 let engine = create_test_engine();
4312 let run = engine
4313 .store()
4314 .create_run(NewRun {
4315 created_by: None,
4316 workflow_name: "test".to_string(),
4317 trigger: TriggerKind::Manual,
4318 payload: json!({}),
4319 max_retries: 0,
4320 handler_version: None,
4321 labels: HashMap::new(),
4322 scheduled_at: None,
4323 idempotency_key: None,
4324 concurrency_key: None,
4325 priority: 0,
4326 concurrency_limits: Vec::new(),
4327 max_cost_usd: None,
4328 worker_tags: Vec::new(),
4329 })
4330 .await
4331 .unwrap()
4332 .into_run();
4333
4334 let s_completed =
4335 create_step_with_status(engine.store(), run.id, "step-1", 0, StepStatus::Completed)
4336 .await;
4337 let s_running =
4338 create_step_with_status(engine.store(), run.id, "step-2", 1, StepStatus::Running).await;
4339 let s_pending =
4340 create_step_with_status(engine.store(), run.id, "step-3", 2, StepStatus::Pending).await;
4341
4342 engine.fail_orphaned_steps(run.id, "timeout").await.unwrap();
4343
4344 let r_completed = engine
4345 .store()
4346 .get_step(s_completed.id)
4347 .await
4348 .unwrap()
4349 .unwrap();
4350 assert_eq!(r_completed.status.state, StepStatus::Completed);
4351
4352 let r_running = engine
4353 .store()
4354 .get_step(s_running.id)
4355 .await
4356 .unwrap()
4357 .unwrap();
4358 assert_eq!(r_running.status.state, StepStatus::Failed);
4359 assert_eq!(r_running.error.as_deref(), Some("timeout"));
4360
4361 let r_pending = engine
4362 .store()
4363 .get_step(s_pending.id)
4364 .await
4365 .unwrap()
4366 .unwrap();
4367 assert_eq!(r_pending.status.state, StepStatus::Skipped);
4368 assert!(r_pending.error.is_none());
4369 }
4370
4371 #[tokio::test]
4372 async fn fail_orphaned_steps_no_steps_is_noop() {
4373 let engine = create_test_engine();
4374 let run = engine
4375 .store()
4376 .create_run(NewRun {
4377 created_by: None,
4378 workflow_name: "test".to_string(),
4379 trigger: TriggerKind::Manual,
4380 payload: json!({}),
4381 max_retries: 0,
4382 handler_version: None,
4383 labels: HashMap::new(),
4384 scheduled_at: None,
4385 idempotency_key: None,
4386 concurrency_key: None,
4387 priority: 0,
4388 concurrency_limits: Vec::new(),
4389 max_cost_usd: None,
4390 worker_tags: Vec::new(),
4391 })
4392 .await
4393 .unwrap()
4394 .into_run();
4395
4396 let result = engine.fail_orphaned_steps(run.id, "timeout").await;
4397 assert!(result.is_ok());
4398 }
4399
4400 #[tokio::test]
4401 async fn fail_orphaned_steps_preserves_existing_error() {
4402 let engine = create_test_engine();
4403 let run = engine
4404 .store()
4405 .create_run(NewRun {
4406 created_by: None,
4407 workflow_name: "test".to_string(),
4408 trigger: TriggerKind::Manual,
4409 payload: json!({}),
4410 max_retries: 0,
4411 handler_version: None,
4412 labels: HashMap::new(),
4413 scheduled_at: None,
4414 idempotency_key: None,
4415 concurrency_key: None,
4416 priority: 0,
4417 concurrency_limits: Vec::new(),
4418 max_cost_usd: None,
4419 worker_tags: Vec::new(),
4420 })
4421 .await
4422 .unwrap()
4423 .into_run();
4424
4425 let step_with_error = create_step_with_status(
4426 engine.store(),
4427 run.id,
4428 "already-errored",
4429 0,
4430 StepStatus::Running,
4431 )
4432 .await;
4433
4434 engine
4435 .store()
4436 .update_step(
4437 step_with_error.id,
4438 StepUpdate {
4439 error: Some("real error from provider".to_string()),
4440 ..StepUpdate::default()
4441 },
4442 )
4443 .await
4444 .unwrap();
4445
4446 let step_no_error = create_step_with_status(
4447 engine.store(),
4448 run.id,
4449 "no-error-yet",
4450 1,
4451 StepStatus::Running,
4452 )
4453 .await;
4454
4455 engine
4456 .fail_orphaned_steps(run.id, "parent run failed")
4457 .await
4458 .unwrap();
4459
4460 let updated_with = engine
4461 .store()
4462 .get_step(step_with_error.id)
4463 .await
4464 .unwrap()
4465 .unwrap();
4466 assert_eq!(updated_with.status.state, StepStatus::Failed);
4467 assert_eq!(
4468 updated_with.error.as_deref(),
4469 Some("real error from provider"),
4470 );
4471
4472 let updated_without = engine
4473 .store()
4474 .get_step(step_no_error.id)
4475 .await
4476 .unwrap()
4477 .unwrap();
4478 assert_eq!(updated_without.status.state, StepStatus::Failed);
4479 assert_eq!(updated_without.error.as_deref(), Some("parent run failed"),);
4480 }
4481
4482 use std::sync::atomic::{AtomicUsize, Ordering};
4487 use std::time::Duration;
4488
4489 use ironflow_store::models::LeaseRequest;
4490 use tokio::time::sleep;
4491
4492 struct CountingWorkflow {
4494 executions: Arc<AtomicUsize>,
4495 }
4496
4497 impl WorkflowHandler for CountingWorkflow {
4498 fn name(&self) -> &str {
4499 "local-resume-counter"
4500 }
4501
4502 fn execute<'a>(&'a self, _ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
4503 Box::pin(async move {
4504 self.executions.fetch_add(1, Ordering::SeqCst);
4505 Ok(())
4506 })
4507 }
4508 }
4509
4510 async fn engine_with_pending_counter() -> (Arc<Engine>, Arc<AtomicUsize>, Uuid) {
4512 let executions = Arc::new(AtomicUsize::new(0));
4513 let mut engine = create_test_engine();
4514 engine
4515 .register(CountingWorkflow {
4516 executions: executions.clone(),
4517 })
4518 .unwrap();
4519 let run = engine
4520 .enqueue_handler("local-resume-counter", TriggerKind::Manual, json!({}), 0)
4521 .await
4522 .unwrap();
4523 (Arc::new(engine), executions, run.id)
4524 }
4525
4526 #[tokio::test]
4527 async fn local_resume_leaves_a_run_claimed_by_a_worker_alone() {
4528 let (engine, executions, run_id) = engine_with_pending_counter().await;
4529 let lease = LeaseRequest {
4530 worker_id: "worker-1".to_string(),
4531 ttl: Duration::from_secs(30),
4532 };
4533 let picked = engine
4534 .store()
4535 .pick_next_pending_for(Some(lease), None)
4536 .await
4537 .unwrap()
4538 .expect("the run is due");
4539 assert_eq!(picked.id, run_id);
4540
4541 engine.spawn_local_resume(run_id);
4542 sleep(Duration::from_millis(200)).await;
4543
4544 assert_eq!(executions.load(Ordering::SeqCst), 0);
4545 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4546 assert_eq!(run.status.state, RunStatus::Running);
4547 assert_eq!(run.worker_id.as_deref(), Some("worker-1"));
4548 }
4549
4550 #[tokio::test]
4551 async fn local_resume_claims_a_pending_run_once() {
4552 let (engine, executions, run_id) = engine_with_pending_counter().await;
4553
4554 engine.spawn_local_resume(run_id);
4555 engine.spawn_local_resume(run_id);
4556 sleep(Duration::from_millis(200)).await;
4557
4558 assert_eq!(executions.load(Ordering::SeqCst), 1);
4559 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4560 assert_eq!(run.status.state, RunStatus::Completed);
4561 }
4562
4563 #[tokio::test]
4564 async fn local_resume_of_a_running_run_executes_it_once() {
4565 let (engine, executions, run_id) = engine_with_pending_counter().await;
4566 engine
4567 .store()
4568 .update_run_status(run_id, RunStatus::Running)
4569 .await
4570 .unwrap();
4571
4572 engine.spawn_local_resume(run_id);
4573 engine.spawn_local_resume(run_id);
4574 sleep(Duration::from_millis(200)).await;
4575
4576 assert_eq!(executions.load(Ordering::SeqCst), 1);
4577 let run = engine.store().get_run(run_id).await.unwrap().unwrap();
4578 assert_eq!(run.status.state, RunStatus::Completed);
4579 }
4580}