Skip to main content

ironflow_engine/
engine.rs

1//! The core [`Engine`] -- orchestrates workflow execution and persistence.
2//!
3//! The engine ties together a `RunStore` for persistence, an [`AgentProvider`]
4//! for AI operations, and a registry of [`WorkflowHandler`]s.
5//!
6//! Handlers are Rust-native: steps can reference previous outputs, use native
7//! `if`/`else`/`match` for conditional branching, and execute in parallel.
8
9use 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/// Result of a workflow execution, carrying the final [`Run`] and per-step
64/// metrics collected during execution.
65///
66/// Returned by [`Engine::run_handler`], [`Engine::execute_handler_run`],
67/// [`Engine::execute_run`], and [`Engine::resume_run`].
68///
69/// # Examples
70///
71/// ```no_run
72/// use ironflow_engine::engine::WorkflowResult;
73///
74/// # fn example(result: WorkflowResult) {
75/// println!("run {} finished with {} steps", result.run.id, result.steps.len());
76/// for step in &result.steps {
77///     println!("  {} ({:?}): {}ms", step.name, step.status, step.duration_ms);
78/// }
79/// # }
80/// ```
81#[derive(Debug, Clone)]
82pub struct WorkflowResult {
83    /// The finalized run record.
84    pub run: Run,
85    /// Per-step results in execution order.
86    pub steps: Vec<StepResult>,
87}
88
89/// Optional settings for [`Engine::enqueue_handler_with_options`].
90///
91/// All fields fall back to handler or server defaults when left at their
92/// [`Default`] value.
93///
94/// # Examples
95///
96/// ```
97/// use ironflow_engine::engine::EnqueueOptions;
98/// use rust_decimal::Decimal;
99///
100/// let options = EnqueueOptions {
101///     max_retries: 3,
102///     max_cost_usd: Some(Decimal::new(50, 2)),
103///     ..Default::default()
104/// };
105/// assert_eq!(options.max_retries, 3);
106/// ```
107#[derive(Debug, Clone, Default)]
108pub struct EnqueueOptions {
109    /// Number of automatic retries granted to the run.
110    pub max_retries: u32,
111    /// Labels merged on top of the handler's default labels.
112    pub labels: HashMap<String, String>,
113    /// Defer execution until this instant instead of running as soon as a
114    /// worker picks the run up.
115    pub scheduled_at: Option<DateTime<Utc>>,
116    /// Cost cap for the run. Overrides both the handler default and the server
117    /// default. `None` falls back to
118    /// [`BudgetConfig::resolve_run_cap`](crate::budget::BudgetConfig::resolve_run_cap).
119    pub max_cost_usd: Option<Decimal>,
120    /// Authenticated principal that triggered the run. `None` for cron,
121    /// webhook, and programmatic triggers.
122    pub created_by: Option<RunActor>,
123    /// Idempotency key binding this enqueue to a single run.
124    ///
125    /// When set and already bound to a run created within
126    /// [`IDEMPOTENCY_WINDOW`](ironflow_store::entities::IDEMPOTENCY_WINDOW),
127    /// nothing is enqueued and the original run is replayed.
128    pub idempotency_key: Option<String>,
129    /// Concurrency key making the run exclusive.
130    ///
131    /// When set, the run is refused with [`EngineError::ConcurrencyConflict`]
132    /// while another non-terminal run holds the same key. The key is released
133    /// once the holder reaches a terminal state.
134    pub concurrency_key: Option<String>,
135    /// Concurrency groups the run belongs to, each with the maximum number of
136    /// root runs of that group allowed to execute at once.
137    ///
138    /// Unlike [`concurrency_key`](Self::concurrency_key), the run is always
139    /// created: it stays pending until every group is under its limit. Empty
140    /// means no limit. Invalid limits are refused with
141    /// [`EngineError::InvalidConcurrencyLimit`].
142    pub concurrency_limits: Vec<ConcurrencyLimit>,
143    /// Queue priority of the run, overriding
144    /// [`WorkflowHandler::priority`](crate::handler::WorkflowHandler::priority).
145    ///
146    /// Workers pick the pending run with the highest priority first, then the
147    /// oldest among equal priorities. `None` falls back to the handler's
148    /// priority, clamped to the accepted range. An explicit value outside
149    /// [`MIN_PRIORITY`](ironflow_store::entities::MIN_PRIORITY)`..=`[`MAX_PRIORITY`](ironflow_store::entities::MAX_PRIORITY)
150    /// is refused with [`EngineError::InvalidPriority`].
151    pub priority: Option<i16>,
152    /// Worker tags the run requires, merged with the handler's
153    /// [`required_worker_tags`](WorkflowHandler::required_worker_tags).
154    ///
155    /// Only a worker carrying every tag picks the run. Invalid tags are
156    /// refused with [`EngineError::InvalidWorkerTag`].
157    pub worker_tags: Vec<String>,
158}
159
160/// Where a run resumes once an approval, a human input or an escalation
161/// resolves the gate it was suspended on.
162///
163/// # Examples
164///
165/// ```
166/// use ironflow_engine::engine::ExecutionMode;
167///
168/// assert_eq!(ExecutionMode::default(), ExecutionMode::Local);
169/// ```
170#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
171pub enum ExecutionMode {
172    /// Resume in-process via [`Engine::resume_run`]. Used by
173    /// [`crate::testing::TestEngine`] and single-process deployments where
174    /// the API also registers the handlers.
175    #[default]
176    Local,
177    /// Requeue the run to `Pending` instead. A worker's `pick_next_pending`
178    /// claims it and finishes it via [`Engine::execute_handler_run`].
179    Workers,
180}
181
182/// The workflow orchestration engine.
183///
184/// Holds references to the store, agent provider, and a registry of
185/// [`WorkflowHandler`]s.
186///
187/// # Examples
188///
189/// ```no_run
190/// use std::sync::Arc;
191/// use ironflow_engine::engine::Engine;
192/// use ironflow_engine::config::ShellConfig;
193/// use ironflow_engine::handler::{WorkflowHandler, HandlerFuture, WorkflowInfo};
194/// use ironflow_engine::context::WorkflowContext;
195/// use ironflow_store::memory::InMemoryStore;
196/// use ironflow_store::models::TriggerKind;
197/// use ironflow_core::providers::claude::ClaudeCodeProvider;
198/// use serde_json::json;
199///
200/// struct CiWorkflow;
201/// impl WorkflowHandler for CiWorkflow {
202///     fn name(&self) -> &str { "ci" }
203///     fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
204///         Box::pin(async move {
205///             ctx.shell("test", ShellConfig::new("cargo test")).await?;
206///             Ok(())
207///         })
208///     }
209/// }
210///
211/// # async fn example() -> Result<(), ironflow_engine::error::EngineError> {
212/// let store = Arc::new(InMemoryStore::new());
213/// let provider = Arc::new(ClaudeCodeProvider::new());
214/// let mut engine = Engine::new(store, provider);
215/// engine.register(CiWorkflow)?;
216///
217/// let result = engine.run_handler("ci", TriggerKind::Manual, json!({})).await?;
218/// tracing::info!(run_id = %result.run.id, status = ?result.run.status, steps = result.steps.len(), "run completed");
219/// # Ok(())
220/// # }
221/// ```
222pub 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
240/// Per-run locks of this process, keyed by run id.
241type ActiveRuns = Arc<Mutex<HashMap<Uuid, Arc<AsyncMutex<()>>>>>;
242
243/// Marks a run as executing in this process, or as being resumed locally,
244/// until it is dropped.
245///
246/// Holds the run's lock, so a resume of the same run waits for the execution
247/// to end ([`Engine::wait_until_idle`]) instead of replaying it in parallel.
248struct 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
262/// Take the lock of `run_id` in `runs`, waiting for its current holder.
263async 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
277/// Remove the lock of `run_id` from `runs` once nobody holds or awaits it.
278///
279/// `entry` is the caller's own clone: the map plus that clone make two
280/// references, any more is a waiter that still needs the lock in the map.
281fn 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
292/// Validate a workflow category path.
293///
294/// A category is a `/`-separated list of non-empty segments. This function
295/// rejects empty paths, leading or trailing `/`, consecutive `/`, and
296/// segments containing only whitespace.
297///
298/// # Errors
299///
300/// Returns [`EngineError::InvalidWorkflow`] when the category is malformed.
301fn 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
328/// Read a run id from the label `key` of a sub-workflow child run.
329///
330/// `None` when the run was not started by a `Workflow` step, when the label
331/// is missing or invalid, or when it points back at the run itself.
332fn 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
340/// The root run of the chain a sub-workflow child run belongs to.
341///
342/// Resuming a child resumes this root instead: the root replays its steps
343/// and re-enters the same child run through its open `Workflow` step. The
344/// worker uses it to follow its lease from a child to the root it resumed.
345///
346/// `None` when `run` is not a sub-workflow child run.
347///
348/// # Examples
349///
350/// ```no_run
351/// use ironflow_engine::engine::chain_root;
352/// use ironflow_store::entities::Run;
353/// use uuid::Uuid;
354///
355/// fn lease_target(run: &Run) -> Uuid {
356///     chain_root(run).unwrap_or(run.id)
357/// }
358/// ```
359pub fn chain_root(run: &Run) -> Option<Uuid> {
360    chain_label(run, LABEL_ROOT_RUN_ID)
361}
362
363/// The run whose `Workflow` step started this sub-workflow child run.
364pub(crate) fn chain_parent(run: &Run) -> Option<Uuid> {
365    chain_label(run, PARENT_RUN_ID_LABEL)
366}
367
368impl Engine {
369    /// Create a new engine with the given store and agent provider.
370    ///
371    /// # Examples
372    ///
373    /// ```no_run
374    /// use std::sync::Arc;
375    /// use ironflow_engine::engine::Engine;
376    /// use ironflow_store::memory::InMemoryStore;
377    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
378    ///
379    /// let engine = Engine::new(
380    ///     Arc::new(InMemoryStore::new()),
381    ///     Arc::new(ClaudeCodeProvider::new()),
382    /// );
383    /// ```
384    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    /// Mark `run_id` as executing in this process, waiting for a previous
405    /// execution of the same run to end first.
406    async fn track_execution(&self, run_id: Uuid) -> ActiveRunGuard {
407        lock_run(&self.active_runs, run_id).await
408    }
409
410    /// Wait until no execution of `run_id` is running in this process.
411    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    /// Whether an execution of `run_id` is running in this process.
426    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    /// Wire a [`DecisionProvider`] backend for `ctx.decision(...)` steps.
436    ///
437    /// Without this, a workflow that reaches a decision step fails with
438    /// [`EngineError::NoDecisionProvider`].
439    ///
440    /// # Examples
441    ///
442    /// ```no_run
443    /// use std::sync::Arc;
444    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
445    /// use ironflow_core::providers::record_replay_decision::RecordReplayDecisionProvider;
446    /// use ironflow_engine::engine::Engine;
447    /// use ironflow_store::memory::InMemoryStore;
448    ///
449    /// let engine = Engine::new(
450    ///     Arc::new(InMemoryStore::new()),
451    ///     Arc::new(ClaudeCodeProvider::new()),
452    /// )
453    /// .with_decision_provider(Arc::new(RecordReplayDecisionProvider::replay("tests/fixtures")));
454    /// ```
455    pub fn with_decision_provider(mut self, provider: Arc<dyn DecisionProvider>) -> Self {
456        self.decision_provider = Some(provider);
457        self
458    }
459
460    /// Wire a [`StepInterceptor`] that resolves steps without executing them.
461    ///
462    /// Used by [`crate::testing::TestEngine`] to mock shell, HTTP and approval
463    /// steps. `None` (the default) executes every step for real.
464    ///
465    /// # Examples
466    ///
467    /// ```no_run
468    /// use std::sync::Arc;
469    ///
470    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
471    /// use ironflow_engine::engine::Engine;
472    /// use ironflow_engine::executor::StepInterceptor;
473    /// use ironflow_store::memory::InMemoryStore;
474    ///
475    /// # fn example(interceptor: Arc<dyn StepInterceptor>) {
476    /// let engine = Engine::new(
477    ///     Arc::new(InMemoryStore::new()),
478    ///     Arc::new(ClaudeCodeProvider::new()),
479    /// )
480    /// .with_step_interceptor(interceptor);
481    /// # }
482    /// ```
483    pub fn with_step_interceptor(mut self, interceptor: Arc<dyn StepInterceptor>) -> Self {
484        self.step_interceptor = Some(interceptor);
485        self
486    }
487
488    /// The step interceptor wired into this engine, if any.
489    pub fn step_interceptor(&self) -> Option<&Arc<dyn StepInterceptor>> {
490        self.step_interceptor.as_ref()
491    }
492
493    /// Apply cost guardrails to this engine.
494    ///
495    /// Without this, both the per-run cap default and the monthly quota are
496    /// disabled and the engine behaves exactly as before.
497    ///
498    /// # Examples
499    ///
500    /// ```no_run
501    /// use std::sync::Arc;
502    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
503    /// use ironflow_engine::budget::BudgetConfig;
504    /// use ironflow_engine::engine::Engine;
505    /// use ironflow_store::memory::InMemoryStore;
506    ///
507    /// let engine = Engine::new(
508    ///     Arc::new(InMemoryStore::new()),
509    ///     Arc::new(ClaudeCodeProvider::new()),
510    /// )
511    /// .with_budget_config(BudgetConfig::from_env());
512    /// ```
513    pub fn with_budget_config(mut self, budget: BudgetConfig) -> Self {
514        self.budget = budget;
515        self
516    }
517
518    /// Returns the cost guardrails applied by this engine.
519    pub fn budget_config(&self) -> &BudgetConfig {
520        &self.budget
521    }
522
523    /// Apply workflow guard configuration to this engine.
524    ///
525    /// When set, every workflow run created by this engine is protected
526    /// by the guard. Handlers can override this via
527    /// [`WorkflowHandler::guard_config`].
528    ///
529    /// # Examples
530    ///
531    /// ```no_run
532    /// use std::sync::Arc;
533    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
534    /// use ironflow_engine::engine::Engine;
535    /// use ironflow_engine::guard::WorkflowGuardConfig;
536    /// use ironflow_store::memory::InMemoryStore;
537    ///
538    /// let engine = Engine::new(
539    ///     Arc::new(InMemoryStore::new()),
540    ///     Arc::new(ClaudeCodeProvider::new()),
541    /// )
542    /// .with_guard_config(WorkflowGuardConfig::default());
543    /// ```
544    pub fn with_guard_config(mut self, config: WorkflowGuardConfig) -> Self {
545        self.guard_config = Some(config);
546        self
547    }
548
549    /// Returns the workflow guard configuration, if any.
550    pub fn guard_config(&self) -> Option<&WorkflowGuardConfig> {
551        self.guard_config.as_ref()
552    }
553
554    /// Choose where a run resumes once its approval, human input or
555    /// escalation is resolved.
556    ///
557    /// Defaults to [`ExecutionMode::Local`]. Set [`ExecutionMode::Workers`]
558    /// on an API process that delegates execution to workers.
559    ///
560    /// # Examples
561    ///
562    /// ```no_run
563    /// use std::sync::Arc;
564    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
565    /// use ironflow_engine::engine::{Engine, ExecutionMode};
566    /// use ironflow_store::memory::InMemoryStore;
567    ///
568    /// let engine = Engine::new(
569    ///     Arc::new(InMemoryStore::new()),
570    ///     Arc::new(ClaudeCodeProvider::new()),
571    /// )
572    /// .with_execution_mode(ExecutionMode::Workers);
573    /// ```
574    pub fn with_execution_mode(mut self, mode: ExecutionMode) -> Self {
575        self.execution_mode = mode;
576        self
577    }
578
579    /// Returns where a run resumes once its gate is resolved.
580    pub fn execution_mode(&self) -> ExecutionMode {
581        self.execution_mode
582    }
583
584    /// Declare the worker tags carried by the process running this engine.
585    ///
586    /// Set by a worker at build time. A `ctx.workflow(..)` step then refuses
587    /// a child whose [`required_worker_tags`](WorkflowHandler::required_worker_tags)
588    /// are not all carried here, instead of running it on the wrong host.
589    /// Unset by default (API or local mode): no check is made.
590    ///
591    /// # Examples
592    ///
593    /// ```no_run
594    /// use std::sync::Arc;
595    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
596    /// use ironflow_engine::engine::Engine;
597    /// use ironflow_store::memory::InMemoryStore;
598    ///
599    /// let mut engine = Engine::new(
600    ///     Arc::new(InMemoryStore::new()),
601    ///     Arc::new(ClaudeCodeProvider::new()),
602    /// );
603    /// engine.set_worker_tags(vec!["gpu".to_string()]);
604    /// assert_eq!(engine.worker_tags(), Some(&["gpu".to_string()][..]));
605    /// ```
606    pub fn set_worker_tags(&mut self, tags: Vec<String>) {
607        self.worker_tags = Some(Arc::new(tags));
608    }
609
610    /// Returns the worker tags set by [`set_worker_tags`](Self::set_worker_tags),
611    /// or `None` when the engine does not run inside a tagged worker.
612    ///
613    /// # Examples
614    ///
615    /// ```no_run
616    /// use std::sync::Arc;
617    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
618    /// use ironflow_engine::engine::Engine;
619    /// use ironflow_store::memory::InMemoryStore;
620    ///
621    /// let engine = Engine::new(
622    ///     Arc::new(InMemoryStore::new()),
623    ///     Arc::new(ClaudeCodeProvider::new()),
624    /// );
625    /// assert!(engine.worker_tags().is_none());
626    /// ```
627    pub fn worker_tags(&self) -> Option<&[String]> {
628        self.worker_tags.as_deref().map(Vec::as_slice)
629    }
630
631    /// Attach a log sender for real-time step output streaming.
632    ///
633    /// When set, all workflow contexts created by this engine will forward
634    /// step output (shell stdout/stderr, agent system messages) to the
635    /// given sender.
636    pub fn set_log_sender(&mut self, sender: LogSender) {
637        self.log_sender = Some(sender);
638    }
639
640    /// Attach the backend that stores and serves artifact bytes.
641    ///
642    /// Every context this engine builds inherits it. Without one, steps that
643    /// declare artifacts fail explicitly and every other step is unaffected.
644    ///
645    /// # Examples
646    ///
647    /// ```no_run
648    /// use std::sync::Arc;
649    ///
650    /// use ironflow_engine::artifact::ArtifactSink;
651    /// use ironflow_engine::engine::Engine;
652    ///
653    /// # fn example(engine: &mut Engine, sink: Arc<dyn ArtifactSink>) {
654    /// engine.set_artifact_sink(sink);
655    /// # }
656    /// ```
657    pub fn set_artifact_sink(&mut self, sink: Arc<dyn ArtifactSink>) {
658        self.artifact_sink = Some(sink);
659    }
660
661    /// The artifact backend attached to this engine, if any.
662    pub fn artifact_sink(&self) -> Option<&Arc<dyn ArtifactSink>> {
663        self.artifact_sink.as_ref()
664    }
665
666    /// Attach a [`WorkflowEventBus`] for per-run real-time monitoring.
667    ///
668    /// When set, all workflow contexts created by this engine will publish
669    /// step-level events (`StepStarted`, `StepCompleted`, `StepFailed`)
670    /// to the bus.
671    ///
672    /// # Examples
673    ///
674    /// ```no_run
675    /// use ironflow_engine::engine::Engine;
676    /// use ironflow_engine::notify::WorkflowEventBus;
677    ///
678    /// # fn example(engine: &mut Engine) {
679    /// engine.set_event_bus(WorkflowEventBus::new());
680    /// # }
681    /// ```
682    pub fn set_event_bus(&mut self, bus: WorkflowEventBus) {
683        self.event_bus = Some(bus);
684    }
685
686    /// The workflow event bus attached to this engine, if any.
687    pub fn event_bus(&self) -> Option<&WorkflowEventBus> {
688        self.event_bus.as_ref()
689    }
690
691    /// Returns a reference to the backing store.
692    pub fn store(&self) -> &Arc<dyn Store> {
693        &self.store
694    }
695
696    /// Returns a reference to the agent provider.
697    pub fn provider(&self) -> &Arc<dyn AgentProvider> {
698        &self.provider
699    }
700
701    /// Build a [`WorkflowContext`] with access to the handler registry.
702    ///
703    /// The context is seeded with the run's attempt number and the cost and
704    /// duration already accumulated by previous attempts, so a retried run
705    /// reports the total it really consumed rather than only its last attempt.
706    ///
707    /// The run's persisted `max_cost_usd` becomes the context cost cap; `None`
708    /// disables the per-run budget check for that context.
709    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    /// Build a context with the guard attached.
745    ///
746    /// Uses the handler's `guard_config()` if present, otherwise falls back
747    /// to the engine's global configuration. Creates a fresh shared state
748    /// for each top-level run.
749    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    /// Reject the creation of a new run when the monthly quota is exhausted.
763    ///
764    /// The window is the current calendar month in UTC. Runs already in flight
765    /// are never interrupted -- only creation is refused.
766    ///
767    /// # Errors
768    ///
769    /// Returns [`EngineError::MonthlyBudgetExceeded`] when the accumulated cost
770    /// of the month has reached the configured quota. Returns
771    /// [`EngineError::Store`] when the aggregate query fails.
772    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    // -----------------------------------------------------------------------
811    // Handler registration
812    // -----------------------------------------------------------------------
813
814    /// Register a [`WorkflowHandler`] for dynamic workflow execution.
815    ///
816    /// The handler is looked up by [`WorkflowHandler::name`] when executing
817    /// or enqueuing.
818    ///
819    /// # Errors
820    ///
821    /// Returns [`EngineError::InvalidWorkflow`] if a handler with the same
822    /// name is already registered.
823    ///
824    /// # Examples
825    ///
826    /// ```no_run
827    /// use std::sync::Arc;
828    /// use ironflow_engine::engine::Engine;
829    /// use ironflow_engine::handler::{WorkflowHandler, HandlerFuture};
830    /// use ironflow_engine::context::WorkflowContext;
831    /// use ironflow_engine::config::ShellConfig;
832    /// use ironflow_store::memory::InMemoryStore;
833    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
834    ///
835    /// struct MyWorkflow;
836    /// impl WorkflowHandler for MyWorkflow {
837    ///     fn name(&self) -> &str { "my-workflow" }
838    ///     fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
839    ///         Box::pin(async move {
840    ///             ctx.shell("step1", ShellConfig::new("echo done")).await?;
841    ///             Ok(())
842    ///         })
843    ///     }
844    /// }
845    ///
846    /// let mut engine = Engine::new(
847    ///     Arc::new(InMemoryStore::new()),
848    ///     Arc::new(ClaudeCodeProvider::new()),
849    /// );
850    /// engine.register(MyWorkflow)?;
851    /// # Ok::<(), ironflow_engine::error::EngineError>(())
852    /// ```
853    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    /// Register a pre-boxed workflow handler.
869    ///
870    /// # Errors
871    ///
872    /// Returns [`EngineError::InvalidWorkflow`] if a handler with the same
873    /// name is already registered or if its category is invalid.
874    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    /// Get a registered handler by name.
890    pub fn get_handler(&self, name: &str) -> Option<&Arc<dyn WorkflowHandler>> {
891        self.handlers.get(name)
892    }
893
894    /// List registered handler names.
895    pub fn handler_names(&self) -> Vec<&str> {
896        self.handlers.keys().map(|s| s.as_str()).collect()
897    }
898
899    /// Get detailed info about a registered workflow handler.
900    pub fn handler_info(&self, name: &str) -> Option<WorkflowInfo> {
901        self.handlers.get(name).map(|h| h.describe())
902    }
903
904    /// List handlers that have a cron schedule configured.
905    ///
906    /// Returns pairs of `(workflow_name, cron_expression)` for all handlers
907    /// where [`WorkflowHandler::schedule`] returns `Some`.
908    ///
909    /// Use this to wire scheduled handlers into a scheduler
910    /// (e.g. the schedule ticker in `ironflow-api`).
911    ///
912    /// # Examples
913    ///
914    /// ```no_run
915    /// # use std::sync::Arc;
916    /// # use ironflow_engine::engine::Engine;
917    /// # use ironflow_store::memory::InMemoryStore;
918    /// # use ironflow_core::providers::claude::ClaudeCodeProvider;
919    /// let engine = Engine::new(
920    ///     Arc::new(InMemoryStore::new()),
921    ///     Arc::new(ClaudeCodeProvider::new()),
922    /// );
923    /// for (name, schedule) in engine.scheduled_handlers() {
924    ///     tracing::info!("{name} runs on schedule: {schedule}");
925    /// }
926    /// ```
927    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    /// Register an event subscriber for domain events.
935    ///
936    /// The subscriber is called only for events whose type is in
937    /// `event_types`. Pass [`Event::ALL`] to receive every event.
938    ///
939    /// # Examples
940    ///
941    /// ```no_run
942    /// use ironflow_engine::engine::Engine;
943    /// use ironflow_engine::notify::{Event, WebhookSubscriber};
944    /// use ironflow_store::memory::InMemoryStore;
945    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
946    /// use std::sync::Arc;
947    ///
948    /// let mut engine = Engine::new(
949    ///     Arc::new(InMemoryStore::new()),
950    ///     Arc::new(ClaudeCodeProvider::new()),
951    /// );
952    ///
953    /// engine.subscribe(
954    ///     WebhookSubscriber::new("https://hooks.example.com/events"),
955    ///     &[Event::RUN_STATUS_CHANGED, Event::STEP_FAILED],
956    /// );
957    /// ```
958    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    /// Returns a reference to the event publisher.
967    ///
968    /// Useful for publishing events from outside the engine (e.g. auth
969    /// routes in the API layer).
970    pub fn event_publisher(&self) -> &EventPublisher {
971        &self.event_publisher
972    }
973
974    // -----------------------------------------------------------------------
975    // Dynamic workflow execution (WorkflowHandler)
976    // -----------------------------------------------------------------------
977
978    /// Execute a registered handler inline.
979    ///
980    /// Creates a run, builds a [`WorkflowContext`], calls the handler's
981    /// [`execute`](WorkflowHandler::execute), and finalizes the run.
982    ///
983    /// # Errors
984    ///
985    /// Returns [`EngineError::InvalidWorkflow`] if no handler is registered
986    /// with that name. Returns [`EngineError`] if execution fails.
987    ///
988    /// When the workflow is paused (see [`Engine::pause_workflow`]) the run is
989    /// created and left `Pending`, with no steps: the returned
990    /// `WorkflowResult.run.status.state` is `Pending`. The hold applies in
991    /// every execution mode; [`Engine::resume_workflow`] runs the held run
992    /// under [`ExecutionMode::Local`].
993    ///
994    /// # Examples
995    ///
996    /// ```no_run
997    /// use std::sync::Arc;
998    /// use ironflow_engine::engine::Engine;
999    /// use ironflow_store::memory::InMemoryStore;
1000    /// use ironflow_store::models::TriggerKind;
1001    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
1002    /// use serde_json::json;
1003    ///
1004    /// # async fn example(engine: &Engine) -> Result<(), ironflow_engine::error::EngineError> {
1005    /// let run = engine.run_handler("deploy", TriggerKind::Manual, json!({})).await?;
1006    /// # Ok(())
1007    /// # }
1008    /// ```
1009    #[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    /// Build the execution plan for a registered handler without running it.
1078    ///
1079    /// Executes the handler with every step method in recording mode: no
1080    /// command is spawned, no HTTP request is sent, no agent is called,
1081    /// nothing is persisted. Conditions declared with
1082    /// [`WorkflowContext::when`](crate::context::WorkflowContext::when) are
1083    /// evaluated against `payload`; those declared with
1084    /// [`WorkflowContext::when_dynamic`](crate::context::WorkflowContext::when_dynamic)
1085    /// are reported as unevaluable.
1086    ///
1087    /// Step outputs are synthetic and success-shaped, so the plan follows the
1088    /// nominal branch. A handler that unwraps a decision answer, or that
1089    /// deserializes `ctx.input::<T>()` against a payload it does not match,
1090    /// aborts the plan: the partial plan is returned with
1091    /// [`ExecutionPlan::incomplete_reason`] set.
1092    ///
1093    /// The handler is bounded twice: by
1094    /// [`MAX_PLANNED_STEPS`](crate::plan::MAX_PLANNED_STEPS) on every step
1095    /// kind, and by [`PlanOptions::timeout`]. Either one
1096    /// returns the partial plan, with [`ExecutionPlan::truncated`] and
1097    /// [`ExecutionPlan::incomplete_reason`] set, not an error. `ctx.store()`
1098    /// and `ctx.secrets()` return an error while planning.
1099    ///
1100    /// # Errors
1101    ///
1102    /// Returns [`EngineError::InvalidWorkflow`] when no handler is registered
1103    /// under `handler_name` or when `options.max_depth` is zero. Returns
1104    /// [`EngineError::Store`] when the duration-history query fails. A handler
1105    /// that errors mid-plan does **not** fail this call.
1106    ///
1107    /// # Examples
1108    ///
1109    /// ```no_run
1110    /// use ironflow_engine::engine::Engine;
1111    /// use ironflow_engine::error::EngineError;
1112    /// use ironflow_engine::plan::PlanOptions;
1113    /// use serde_json::json;
1114    ///
1115    /// # async fn example(engine: &Engine) -> Result<(), EngineError> {
1116    /// let plan = engine
1117    ///     .plan_handler("deploy", json!({"env": "prod"}), PlanOptions::default())
1118    ///     .await?;
1119    /// for step in &plan.steps {
1120    ///     println!("{} ({:?})", step.name, step.kind);
1121    /// }
1122    /// # Ok(())
1123    /// # }
1124    /// ```
1125    #[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        // Deliberately bare: no guard, no event bus, no log sender, no artifact
1160        // sink and no budget. Planning produces no side effect to report.
1161        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    /// Enqueue a handler-based workflow for worker execution.
1202    ///
1203    /// The workflow name is stored in the run. The worker looks up the
1204    /// handler by name when executing.
1205    ///
1206    /// # Errors
1207    ///
1208    /// Returns [`EngineError::InvalidWorkflow`] if no handler is registered.
1209    /// Returns [`EngineError::MonthlyBudgetExceeded`] if the monthly cost quota
1210    /// is exhausted.
1211    #[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    /// Enqueue a handler-based workflow with labels, deferred scheduling, an
1233    /// optional cost cap, an optional author, and an optional idempotency key.
1234    ///
1235    /// See [`EnqueueOptions`] for the individual settings.
1236    ///
1237    /// When [`EnqueueOptions::idempotency_key`] is set and already bound to a run
1238    /// created within
1239    /// [`IDEMPOTENCY_WINDOW`](ironflow_store::entities::IDEMPOTENCY_WINDOW), nothing
1240    /// is enqueued and the original run is returned as [`RunCreation::Existing`].
1241    ///
1242    /// # Errors
1243    ///
1244    /// Returns [`EngineError::InvalidWorkflow`] if no handler is registered.
1245    /// Returns [`EngineError::MonthlyBudgetExceeded`] if the monthly cost quota
1246    /// is exhausted. Returns [`EngineError::ConcurrencyConflict`] if
1247    /// [`EnqueueOptions::concurrency_key`] is held by another non-terminal run.
1248    /// Returns [`EngineError::InvalidConcurrencyLimit`] if
1249    /// [`EnqueueOptions::concurrency_limits`] is invalid, before any other check.
1250    /// Returns [`EngineError::InvalidPriority`] if [`EnqueueOptions::priority`]
1251    /// is out of range, before any other check.
1252    /// Returns [`EngineError::InvalidWorkerTag`] if [`EnqueueOptions::worker_tags`]
1253    /// holds an invalid tag, checked right after the concurrency limits.
1254    /// Returns [`EngineError::Store`] if the run cannot be persisted.
1255    ///
1256    /// # Examples
1257    ///
1258    /// ```no_run
1259    /// use ironflow_engine::engine::{Engine, EnqueueOptions};
1260    /// use ironflow_store::models::TriggerKind;
1261    /// use serde_json::json;
1262    ///
1263    /// # async fn example(engine: &Engine) -> Result<(), ironflow_engine::error::EngineError> {
1264    /// let creation = engine
1265    ///     .enqueue_handler_with_options(
1266    ///         "deploy",
1267    ///         TriggerKind::Api,
1268    ///         json!({"env": "prod"}),
1269    ///         EnqueueOptions {
1270    ///             max_retries: 3,
1271    ///             idempotency_key: Some("github:abc-123".to_string()),
1272    ///             ..Default::default()
1273    ///         },
1274    ///     )
1275    ///     .await?;
1276    ///
1277    /// if creation.is_created() {
1278    ///     println!("enqueued {}", creation.run().id);
1279    /// }
1280    /// # Ok(())
1281    /// # }
1282    /// ```
1283    #[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        // Checked first: a malformed request is the caller's error, whatever
1305        // the handler or the quota.
1306        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    /// Execute a handler-based run (used by the worker after pick_next_pending).
1371    ///
1372    /// Looks up the handler by the run's `workflow_name` and executes it
1373    /// with a fresh [`WorkflowContext`], after
1374    /// [`AgentProvider::release_run`] has stopped whatever a previous
1375    /// execution of the run left running.
1376    ///
1377    /// # Errors
1378    ///
1379    /// Returns [`EngineError::InvalidWorkflow`] if no handler matches. A
1380    /// failed release fails the execution with [`EngineError::Operation`],
1381    /// replayed while the run has retries left. Returns
1382    /// [`EngineError::HandlerVersionMismatch`] when the handler's current
1383    /// version is incompatible with the run's `handler_version` -- checked
1384    /// before any step is replayed.
1385    ///
1386    /// A suspended child run of a sub-workflow is never executed on its own:
1387    /// its root run is resumed instead, and re-enters the child (see
1388    /// [`resume_run`](Self::resume_run)).
1389    #[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        // Replay the steps already persisted for this run: on a retry, an approval
1421        // a human already granted must not be asked again, and a run requeued to
1422        // `Pending` after its approval, human input or escalation resolved
1423        // (`ExecutionMode::Workers`, `retry_count` unchanged) must not re-run
1424        // completed steps. Neither must a run requeued by the reaper after its
1425        // worker lost the lease, which stays in the same attempt too. A
1426        // brand-new run has no steps, so this is a no-op.
1427        //
1428        // The handler version is checked first: replaying an incompatible
1429        // handler's steps risks serving one step's cached output to another
1430        // (`EngineError::ReplayDivergence`), so no step is replayed at all
1431        // when the handler changed incompatibly since the run was created.
1432        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    /// Execute a run by its ID (used by the worker after pick_next_pending).
1463    ///
1464    /// Delegates to [`execute_handler_run`](Self::execute_handler_run).
1465    ///
1466    /// # Errors
1467    ///
1468    /// Returns [`EngineError`] if the run is not found or execution fails.
1469    #[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    /// Resume a run after human approval.
1475    ///
1476    /// Re-executes the handler with step replay: completed steps return
1477    /// cached output, approved approval steps are skipped, and execution
1478    /// continues from the first unexecuted step.
1479    ///
1480    /// Supports multiple approval gates -- each resume replays all prior
1481    /// steps and stops at the next approval (or completes the run).
1482    ///
1483    /// When `run_id` is a child run of a sub-workflow, the root run of its
1484    /// chain is moved back to `Running` and resumed instead: it replays, and
1485    /// its open `Workflow` step re-enters the same child run. The returned
1486    /// result is the root run's.
1487    ///
1488    /// Like [`execute_handler_run`](Self::execute_handler_run), the handler
1489    /// only starts once [`AgentProvider::release_run`] has stopped whatever
1490    /// a previous execution of the run left running.
1491    ///
1492    /// # Errors
1493    ///
1494    /// Returns [`EngineError::InvalidWorkflow`] if no handler matches, or when
1495    /// the root run of a child cannot be resumed (it is running or finished:
1496    /// the child is then failed).
1497    /// Returns [`EngineError`] if execution fails or hits another approval.
1498    /// A failed release fails the execution with [`EngineError::Operation`].
1499    /// Returns [`EngineError::HandlerVersionMismatch`] when the handler's
1500    /// current version is incompatible with the run's `handler_version` --
1501    /// checked before any step is replayed.
1502    #[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    /// Resume the root run of a suspended child run.
1518    ///
1519    /// The root waits without `scheduled_at` while its child is suspended, so
1520    /// nothing but this path ever wakes it. It is moved to `Running` and
1521    /// resumed; its replay re-enters the child. A root that is not suspended
1522    /// (already running, or finished) cannot take the child back: the child
1523    /// is failed.
1524    ///
1525    /// When the child holds a worker lease (it was picked by a worker), the
1526    /// lease is transferred to the root in the same update that moves it to
1527    /// `Running`, then released on the child: the worker keeps renewing the
1528    /// root it now executes, and the reaper recovers the root if that worker
1529    /// dies. A child without a lease (inline execution, API-side resume)
1530    /// leaves the root without one, as before.
1531    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    /// Move the suspended root run of a chain to `Running`.
1602    ///
1603    /// With `lease`, the root takes the worker lease of the child that
1604    /// resumes it, atomically with the transition, so it is never `Running`
1605    /// without an owner.
1606    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    /// Resume `run`, already loaded and already `Running`.
1637    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    /// Store a signal and resolve every step waiting for its `(name, key)`.
1687    ///
1688    /// Each waiting step validates the payload against the JSON schema it
1689    /// stored when it opened. A matching step is completed with the payload
1690    /// and its run, if `Sleeping`, goes back to `Pending`: under
1691    /// [`ExecutionMode::Local`] it resumes in a background task, under
1692    /// [`ExecutionMode::Workers`] a worker picks it up. A step whose schema
1693    /// the payload does not match keeps waiting and is listed in
1694    /// [`SignalDelivery::rejected`].
1695    ///
1696    /// The signal is stored even when nobody waits for it: a run opening its
1697    /// wait step later still finds it. A signal whose `idempotency_id` was
1698    /// already used is not stored nor delivered again, and comes back with
1699    /// [`SignalDelivery::duplicate`] set.
1700    ///
1701    /// Publishes [`Event::SignalReceived`] for every stored signal.
1702    ///
1703    /// # Errors
1704    ///
1705    /// Returns [`EngineError::InvalidSignal`] when `name` or `key` is empty,
1706    /// and [`EngineError::Store`] when the signal cannot be stored or its
1707    /// waiters cannot be listed.
1708    ///
1709    /// # Examples
1710    ///
1711    /// ```no_run
1712    /// use std::sync::Arc;
1713    /// use ironflow_engine::engine::Engine;
1714    /// use ironflow_engine::error::EngineError;
1715    /// use ironflow_store::entities::NewSignal;
1716    /// use serde_json::json;
1717    ///
1718    /// # async fn example(engine: Arc<Engine>) -> Result<(), EngineError> {
1719    /// let delivery = engine
1720    ///     .deliver_signal(NewSignal {
1721    ///         name: "ci.pipeline_finished".to_string(),
1722    ///         key: "4f2a9c1".to_string(),
1723    ///         payload: json!({"status": "success"}),
1724    ///         idempotency_id: Some("delivery-42".to_string()),
1725    ///     })
1726    ///     .await?;
1727    /// println!("{} runs resumed", delivery.resumed.len());
1728    /// # Ok(())
1729    /// # }
1730    /// ```
1731    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                // A concurrent delivery or the timeout resolved it first.
1799                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    /// Send a typed signal: shorthand for [`deliver_signal`](Self::deliver_signal)
1842    /// with `S::NAME` as the name and `signal` as the payload.
1843    ///
1844    /// `key` identifies the occurrence (a commit SHA, an order ID).
1845    /// `idempotency_id`, when set, makes a redelivery of the same event (a
1846    /// webhook retried by its sender) a no-op.
1847    ///
1848    /// # Errors
1849    ///
1850    /// Returns [`EngineError::Serialization`] when `signal` cannot be
1851    /// serialized, and every error of [`deliver_signal`](Self::deliver_signal).
1852    ///
1853    /// # Examples
1854    ///
1855    /// ```no_run
1856    /// use std::sync::Arc;
1857    /// use ironflow_engine::engine::Engine;
1858    /// use ironflow_engine::error::EngineError;
1859    /// use ironflow_engine::signal::Signal;
1860    /// use schemars::JsonSchema;
1861    /// use serde::{Deserialize, Serialize};
1862    ///
1863    /// #[derive(Serialize, Deserialize, JsonSchema)]
1864    /// struct PipelineFinished {
1865    ///     status: String,
1866    /// }
1867    ///
1868    /// impl Signal for PipelineFinished {
1869    ///     const NAME: &'static str = "ci.pipeline_finished";
1870    /// }
1871    ///
1872    /// # async fn example(engine: Arc<Engine>) -> Result<(), EngineError> {
1873    /// let finished = PipelineFinished { status: "success".to_string() };
1874    /// engine.send_signal(&finished, "4f2a9c1", Some("delivery-42")).await?;
1875    /// # Ok(())
1876    /// # }
1877    /// ```
1878    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    /// Resume a run requeued to `Pending` in a background task.
1895    ///
1896    /// Used under [`ExecutionMode::Local`], where no worker would pick the
1897    /// run up. The state change already happened, so a failed resume is
1898    /// logged, not rolled back.
1899    ///
1900    /// The local resumes of a run are serialized: a task reads the run only
1901    /// once the previous local resume of it has returned, so two tasks never
1902    /// both act on the same read. The task then waits for the execution of
1903    /// the run still in flight in this process, if any: a run paused then
1904    /// resumed while a step runs would otherwise be replayed next to the
1905    /// execution that has not noticed the pause yet. Once it is gone, the run
1906    /// is read again:
1907    ///
1908    /// - `Pending` is claimed with the `Pending -> Running` transition, a
1909    ///   compare-and-set in every store, the one a worker performs when it
1910    ///   picks a run. Only a successful claim resumes it here: a run another
1911    ///   executor claimed first is left to it.
1912    /// - `Running` with a worker lease is owned by that worker and left alone.
1913    /// - `Running` without a lease was left unfinished by the previous
1914    ///   execution in this process and is resumed as is, unless an execution
1915    ///   of it is still running here.
1916    /// - any other status means the previous execution finished the run or it
1917    ///   was taken elsewhere, and nothing is done.
1918    pub(crate) fn spawn_local_resume(self: &Arc<Self>, run_id: Uuid) {
1919        let engine = Arc::clone(self);
1920        spawn(async move {
1921            // One local resume of a run at a time, held until its execution
1922            // returns: a sibling task reads the run only once this one is done,
1923            // so it never sees the claim of a resume that has not started yet.
1924            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    /// Record a run failure, replaying the run later when retries remain.
1989    ///
1990    /// This is the single place where a failed run's fate is decided. When
1991    /// `retryable` is true and the run has not exhausted `max_retries`, the run
1992    /// moves to [`RunStatus::Retrying`] with `scheduled_at` set to
1993    /// `now + backoff` -- [`pick_next_pending`](ironflow_store::store::RunStore::pick_next_pending)
1994    /// picks it up again once that time has passed. Otherwise the run moves to
1995    /// [`RunStatus::Failed`].
1996    ///
1997    /// Either way, steps left non-terminal by the failed attempt are closed via
1998    /// [`fail_orphaned_steps`](Self::fail_orphaned_steps) so they are never
1999    /// confused with the next attempt's steps, and the sub-workflow runs it
2000    /// left non-terminal are cancelled by
2001    /// [`cancel_descendants`](Self::cancel_descendants) (a failure there is
2002    /// logged, not returned).
2003    ///
2004    /// Callers pass `retryable` explicitly rather than an error value, because
2005    /// the worker classifies failures it observes from the outside (a timeout, a
2006    /// panicked task) that never produce an [`EngineError`]. Use
2007    /// [`is_run_retryable`] to classify an
2008    /// [`EngineError`].
2009    ///
2010    /// `cost_usd` and `duration_ms` are `None` when the caller does not know the
2011    /// totals (a timeout or a panic observed from outside the handler); the
2012    /// values already stored on the run are then left untouched.
2013    ///
2014    /// A run an operator paused meanwhile is left untouched and
2015    /// [`RunStatus::Paused`] is returned: the resume replays the attempt.
2016    ///
2017    /// Returns the status the run was moved to.
2018    ///
2019    /// # Errors
2020    ///
2021    /// Returns [`EngineError::Store`] if the run does not exist or the update
2022    /// cannot be persisted.
2023    ///
2024    /// # Examples
2025    ///
2026    /// ```no_run
2027    /// use ironflow_engine::engine::Engine;
2028    /// use ironflow_engine::error::EngineError;
2029    /// use uuid::Uuid;
2030    ///
2031    /// # async fn example(engine: &Engine, run_id: Uuid) -> Result<(), EngineError> {
2032    /// let status = engine
2033    ///     .fail_or_schedule_retry(run_id, "run timed out after 300s", true, None, None)
2034    ///     .await?;
2035    /// # Ok(())
2036    /// # }
2037    /// ```
2038    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        // An operator paused the run: its fate is decided by the resume or
2053        // the cancellation, not by the attempt that stopped.
2054        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        // The attempt is over: a retry starts new children, and nothing drives
2098        // those this attempt left running.
2099        self.cancel_descendants_of_stopped_run(run_id, error).await;
2100
2101        Ok(status)
2102    }
2103
2104    /// Mark the steps of a run requeued after a lost worker lease as interrupted.
2105    ///
2106    /// Every `Running` step is marked `Failed` with
2107    /// [`STEP_INTERRUPTED_ERROR`](ironflow_store::store::STEP_INTERRUPTED_ERROR).
2108    /// The run keeps its attempt number, so when it is picked up again its
2109    /// finished steps are replayed and each interrupted step is executed again
2110    /// at the same position; an interrupted `Workflow` step re-enters the same
2111    /// child run. `Pending` and `AwaitingApproval` steps are left as they are,
2112    /// unlike [`fail_orphaned_steps`](Self::fail_orphaned_steps), which ends
2113    /// the run's steps for good.
2114    ///
2115    /// Errors from individual step updates are logged but do not abort the cleanup.
2116    ///
2117    /// # Errors
2118    ///
2119    /// Returns [`EngineError`] if listing steps fails.
2120    ///
2121    /// # Examples
2122    ///
2123    /// ```no_run
2124    /// use ironflow_engine::engine::Engine;
2125    /// use ironflow_engine::error::EngineError;
2126    /// use uuid::Uuid;
2127    ///
2128    /// # async fn example(engine: &Engine, run_id: Uuid) -> Result<(), EngineError> {
2129    /// engine.interrupt_running_steps(run_id).await?;
2130    /// # Ok(())
2131    /// # }
2132    /// ```
2133    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    /// Fail all non-terminal steps for a run.
2138    ///
2139    /// Called after a run is marked as failed (timeout, error, panic) to clean up
2140    /// orphaned steps that are still in `Running`, `Pending`, or `AwaitingApproval`.
2141    ///
2142    /// - `Running` / `AwaitingApproval` steps are marked `Failed`.
2143    /// - `Pending` steps are marked `Skipped` (FSM does not allow Pending -> Failed).
2144    ///
2145    /// Errors from individual step updates are logged but do not abort the cleanup.
2146    ///
2147    /// # Errors
2148    ///
2149    /// Returns [`EngineError`] if listing steps fails.
2150    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    /// Execute `handler` once [`AgentProvider::release_run`] has stopped
2212    /// whatever a previous execution of the run left running (an agent pod
2213    /// writing to a shared worktree). A failed release fails the execution
2214    /// before its first step, with [`EngineError::Operation`].
2215    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    /// Finalize a run with the given result and context.
2228    ///
2229    /// On success: updates run to Completed with cost, duration, and completed_at.
2230    /// On failure: updates run to Failed with error, cost, duration, and completed_at.
2231    /// Always: fetches and returns the final Run.
2232    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        // Covers the whole run: previous attempts plus this one, so a retried
2242        // run reports the time it really consumed.
2243        let total_duration = ctx.carried_duration_ms() + run_start.elapsed().as_millis() as u64;
2244        let completed_at = Utc::now();
2245
2246        // Paused while it ran: whatever the handler returned, the run stays
2247        // paused for the operator. Only what this execution spent is kept;
2248        // the resume replays from the first step that did not complete.
2249        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        // The pause stopped this execution but the run was resumed since: the
2275        // resume continues or restarts it, so nothing is recorded here.
2276        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                // No `scheduled_at`: the suspended descendant owns the wake-up
2361                // and resumes this run through `resume_chain`. A wake-up armed
2362                // here too would resume the chain twice.
2363                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                    // A human input or a delay publishes no suspension event,
2414                    // like on a top-level run.
2415                    _ => {}
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                // No `ApprovalRequested` event: a human input is not an approval.
2438                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                // Atomic with the step lock: a signal delivered since the step
2512                // opened leaves the run due right away instead of until the
2513                // deadline.
2514                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                // A guardrail stop (budget or workflow guard) is deliberate,
2555                // not a breakage: the run is cancelled, never failed and
2556                // never replayed.
2557                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                    // Written before the failure so a failed run keeps the
2590                    // verdict the handler set before returning its error.
2591                    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    /// Publish [`Event::ApprovalRequested`] for an approval gate that opened
2662    /// on `run_id`, with the requirement recorded when the gate opened.
2663    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    /// Fail every ancestor of a child run, closest first.
2686    ///
2687    /// Used when a gate inside a sub-workflow is rejected: the child run
2688    /// fails, and the runs suspended with it (its parent, up to the root)
2689    /// must not stay suspended on a child that will never resume. Each
2690    /// ancestor goes through [`fail_or_schedule_retry`](Self::fail_or_schedule_retry)
2691    /// without a retry, which also fails its open `Workflow` step. A run that
2692    /// is not a child of a sub-workflow has no ancestor: nothing happens.
2693    ///
2694    /// A paused ancestor is not failed: the root of the chain is set to resume
2695    /// to `Pending`, so its replay observes the failed child once an operator
2696    /// resumes it.
2697    ///
2698    /// # Errors
2699    ///
2700    /// Returns [`EngineError::Store`] if a run of the chain does not exist or
2701    /// its failure cannot be persisted.
2702    ///
2703    /// # Examples
2704    ///
2705    /// ```no_run
2706    /// use ironflow_engine::engine::Engine;
2707    /// use ironflow_engine::error::EngineError;
2708    /// use uuid::Uuid;
2709    ///
2710    /// # async fn example(engine: &Engine, child_run_id: Uuid) -> Result<(), EngineError> {
2711    /// engine
2712    ///     .fail_ancestors(child_run_id, "approval rejected in a sub-workflow")
2713    ///     .await?;
2714    /// # Ok(())
2715    /// # }
2716    /// ```
2717    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        // Labels are data: a chain that loops back on itself stops there.
2724        let mut visited = HashSet::from([run_id]);
2725
2726        while let Some(parent_id) = chain_parent(&current) {
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            // A paused chain waits for its operator: the root replays on
2734            // resume and observes the failed child, like a cancelled one.
2735            if status == RunStatus::Paused {
2736                self.requeue_paused_root(&current).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    /// Emit Prometheus metrics for a completed run.
2756    #[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    /// Record the metric and publish the audit event for a run that hit its
2780    /// cost cap.
2781    ///
2782    /// A non-[`RunBudgetExceeded`](EngineError::RunBudgetExceeded) error is
2783    /// ignored, so callers can pass the error unconditionally.
2784    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    /// Publish a run status changed event to all registered subscribers.
2815    ///
2816    /// `from` is always `Running` because `finalize_run` is only called
2817    /// from a running state.
2818    #[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    // Test handler that echoes a message via shell
2881    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    // Test handler that fails
2913    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    // -----------------------------------------------------------------------
3043    // Schedule tests
3044    // -----------------------------------------------------------------------
3045
3046    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        // Validated before the handler lookup.
3344        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        // Validated before the handler lookup.
3462        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        // Validated before the handler lookup.
3593        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        // Verify accessors return references
3684        let _ = engine.store();
3685        let _ = engine.provider();
3686    }
3687
3688    // -----------------------------------------------------------------------
3689    // Operation trait tests
3690    // -----------------------------------------------------------------------
3691
3692    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    // -----------------------------------------------------------------------
3862    // Approval + resume tests
3863    // -----------------------------------------------------------------------
3864
3865    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); // build + approval gate
3923        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    /// Record one vote on the open approval gate of `run_id`, then move the
3930    /// run to `Running`, the way the approve route does.
3931    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        // First execution: pauses at approval
3962        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        // Simulate approval: one vote, then transition to Running
3970        approve_open_gate(engine.store(), run.id).await;
3971
3972        // Resume: replays build, skips approval, executes deploy
3973        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); // build + approval + deploy
3978        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        // First execution: pauses at staging-gate
3993        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); // build + staging-gate
4002
4003        // First approval
4004        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); // build + staging-gate + deploy-staging + prod-gate
4011
4012        // Second approval
4013        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    // -----------------------------------------------------------------------
4032    // fail_orphaned_steps tests
4033    // -----------------------------------------------------------------------
4034
4035    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    // -----------------------------------------------------------------------
4483    // spawn_local_resume tests
4484    // -----------------------------------------------------------------------
4485
4486    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    /// Counts its executions, then completes.
4493    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    /// An engine running [`CountingWorkflow`], with one `Pending` run of it.
4511    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}