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;
10use std::fmt;
11use std::sync::Arc;
12use std::time::Instant;
13
14use chrono::{DateTime, TimeDelta, Utc};
15use rust_decimal::Decimal;
16use serde_json::Value;
17use tracing::{error, info, warn};
18use uuid::Uuid;
19
20#[cfg(feature = "prometheus")]
21use ironflow_core::metric_names::{
22    RUN_BUDGET_EXCEEDED_TOTAL, RUN_COST_USD, RUN_DURATION_SECONDS, RUNS_ACTIVE, RUNS_TOTAL,
23};
24use ironflow_core::provider::AgentProvider;
25use ironflow_store::error::StoreError;
26use ironflow_store::models::{
27    NewRun, Run, RunActor, RunCreation, RunFilter, RunStatus, RunUpdate, StepStatus, StepUpdate,
28    TriggerKind,
29};
30use ironflow_store::store::Store;
31#[cfg(feature = "prometheus")]
32use metrics::{counter, gauge, histogram};
33
34use crate::artifact::ArtifactSink;
35use crate::budget::{BudgetConfig, month_start};
36use crate::context::WorkflowContext;
37use crate::error::EngineError;
38use crate::executor::StepResult;
39use crate::guard::{WorkflowGuardConfig, new_shared_guard_state};
40use crate::handler::{WorkflowHandler, WorkflowInfo};
41use crate::log_sender::LogSender;
42use crate::notify::{Event, EventPublisher, EventSubscriber, WorkflowEventBus};
43use crate::retry_policy::{backoff_for_retry, is_run_retryable};
44use crate::schedule::CronSchedule;
45use ironflow_core::decision::DecisionProvider;
46
47/// Result of a workflow execution, carrying the final [`Run`] and per-step
48/// metrics collected during execution.
49///
50/// Returned by [`Engine::run_handler`], [`Engine::execute_handler_run`],
51/// [`Engine::execute_run`], and [`Engine::resume_run`].
52///
53/// # Examples
54///
55/// ```no_run
56/// use ironflow_engine::engine::WorkflowResult;
57///
58/// # fn example(result: WorkflowResult) {
59/// println!("run {} finished with {} steps", result.run.id, result.steps.len());
60/// for step in &result.steps {
61///     println!("  {} ({:?}): {}ms", step.name, step.status, step.duration_ms);
62/// }
63/// # }
64/// ```
65#[derive(Debug, Clone)]
66pub struct WorkflowResult {
67    /// The finalized run record.
68    pub run: Run,
69    /// Per-step results in execution order.
70    pub steps: Vec<StepResult>,
71}
72
73/// Optional settings for [`Engine::enqueue_handler_with_options`].
74///
75/// All fields fall back to handler or server defaults when left at their
76/// [`Default`] value.
77///
78/// # Examples
79///
80/// ```
81/// use ironflow_engine::engine::EnqueueOptions;
82/// use rust_decimal::Decimal;
83///
84/// let options = EnqueueOptions {
85///     max_retries: 3,
86///     max_cost_usd: Some(Decimal::new(50, 2)),
87///     ..Default::default()
88/// };
89/// assert_eq!(options.max_retries, 3);
90/// ```
91#[derive(Debug, Clone, Default)]
92pub struct EnqueueOptions {
93    /// Number of automatic retries granted to the run.
94    pub max_retries: u32,
95    /// Labels merged on top of the handler's default labels.
96    pub labels: HashMap<String, String>,
97    /// Defer execution until this instant instead of running as soon as a
98    /// worker picks the run up.
99    pub scheduled_at: Option<DateTime<Utc>>,
100    /// Cost cap for the run. Overrides both the handler default and the server
101    /// default. `None` falls back to
102    /// [`BudgetConfig::resolve_run_cap`](crate::budget::BudgetConfig::resolve_run_cap).
103    pub max_cost_usd: Option<Decimal>,
104    /// Authenticated principal that triggered the run. `None` for cron,
105    /// webhook, and programmatic triggers.
106    pub created_by: Option<RunActor>,
107    /// Idempotency key binding this enqueue to a single run.
108    ///
109    /// When set and already bound to a run created within
110    /// [`IDEMPOTENCY_WINDOW`](ironflow_store::entities::IDEMPOTENCY_WINDOW),
111    /// nothing is enqueued and the original run is replayed.
112    pub idempotency_key: Option<String>,
113}
114
115/// The workflow orchestration engine.
116///
117/// Holds references to the store, agent provider, and a registry of
118/// [`WorkflowHandler`]s.
119///
120/// # Examples
121///
122/// ```no_run
123/// use std::sync::Arc;
124/// use ironflow_engine::engine::Engine;
125/// use ironflow_engine::config::ShellConfig;
126/// use ironflow_engine::handler::{WorkflowHandler, HandlerFuture, WorkflowInfo};
127/// use ironflow_engine::context::WorkflowContext;
128/// use ironflow_store::memory::InMemoryStore;
129/// use ironflow_store::models::TriggerKind;
130/// use ironflow_core::providers::claude::ClaudeCodeProvider;
131/// use serde_json::json;
132///
133/// struct CiWorkflow;
134/// impl WorkflowHandler for CiWorkflow {
135///     fn name(&self) -> &str { "ci" }
136///     fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
137///         Box::pin(async move {
138///             ctx.shell("test", ShellConfig::new("cargo test")).await?;
139///             Ok(())
140///         })
141///     }
142/// }
143///
144/// # async fn example() -> Result<(), ironflow_engine::error::EngineError> {
145/// let store = Arc::new(InMemoryStore::new());
146/// let provider = Arc::new(ClaudeCodeProvider::new());
147/// let mut engine = Engine::new(store, provider);
148/// engine.register(CiWorkflow)?;
149///
150/// let result = engine.run_handler("ci", TriggerKind::Manual, json!({})).await?;
151/// tracing::info!(run_id = %result.run.id, status = ?result.run.status, steps = result.steps.len(), "run completed");
152/// # Ok(())
153/// # }
154/// ```
155pub struct Engine {
156    store: Arc<dyn Store>,
157    provider: Arc<dyn AgentProvider>,
158    handlers: HashMap<String, Arc<dyn WorkflowHandler>>,
159    event_publisher: EventPublisher,
160    log_sender: Option<LogSender>,
161    budget: BudgetConfig,
162    artifact_sink: Option<Arc<dyn ArtifactSink>>,
163    guard_config: Option<WorkflowGuardConfig>,
164    event_bus: Option<WorkflowEventBus>,
165    decision_provider: Option<Arc<dyn DecisionProvider>>,
166}
167
168/// Validate a workflow category path.
169///
170/// A category is a `/`-separated list of non-empty segments. This function
171/// rejects empty paths, leading or trailing `/`, consecutive `/`, and
172/// segments containing only whitespace.
173///
174/// # Errors
175///
176/// Returns [`EngineError::InvalidWorkflow`] when the category is malformed.
177fn validate_category(handler_name: &str, category: &str) -> Result<(), EngineError> {
178    let reject = |reason: &str| {
179        Err(EngineError::InvalidWorkflow(format!(
180            "handler '{handler_name}' has invalid category '{category}': {reason}"
181        )))
182    };
183
184    if category.is_empty() {
185        return reject("empty category");
186    }
187    if category.starts_with('/') {
188        return reject("leading '/'");
189    }
190    if category.ends_with('/') {
191        return reject("trailing '/'");
192    }
193    for segment in category.split('/') {
194        if segment.is_empty() {
195            return reject("empty segment (double '/')");
196        }
197        if segment.trim().is_empty() {
198            return reject("whitespace-only segment");
199        }
200    }
201    Ok(())
202}
203
204impl Engine {
205    /// Create a new engine with the given store and agent provider.
206    ///
207    /// # Examples
208    ///
209    /// ```no_run
210    /// use std::sync::Arc;
211    /// use ironflow_engine::engine::Engine;
212    /// use ironflow_store::memory::InMemoryStore;
213    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
214    ///
215    /// let engine = Engine::new(
216    ///     Arc::new(InMemoryStore::new()),
217    ///     Arc::new(ClaudeCodeProvider::new()),
218    /// );
219    /// ```
220    pub fn new(store: Arc<dyn Store>, provider: Arc<dyn AgentProvider>) -> Self {
221        Self {
222            store,
223            provider,
224            handlers: HashMap::new(),
225            event_publisher: EventPublisher::new(),
226            log_sender: None,
227            budget: BudgetConfig::new(),
228            artifact_sink: None,
229            guard_config: None,
230            event_bus: None,
231            decision_provider: None,
232        }
233    }
234
235    /// Wire a [`DecisionProvider`] backend for `ctx.decision(...)` steps.
236    ///
237    /// Without this, a workflow that reaches a decision step fails with
238    /// [`EngineError::NoDecisionProvider`].
239    ///
240    /// # Examples
241    ///
242    /// ```no_run
243    /// use std::sync::Arc;
244    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
245    /// use ironflow_core::providers::record_replay_decision::RecordReplayDecisionProvider;
246    /// use ironflow_engine::engine::Engine;
247    /// use ironflow_store::memory::InMemoryStore;
248    ///
249    /// let engine = Engine::new(
250    ///     Arc::new(InMemoryStore::new()),
251    ///     Arc::new(ClaudeCodeProvider::new()),
252    /// )
253    /// .with_decision_provider(Arc::new(RecordReplayDecisionProvider::replay("tests/fixtures")));
254    /// ```
255    pub fn with_decision_provider(mut self, provider: Arc<dyn DecisionProvider>) -> Self {
256        self.decision_provider = Some(provider);
257        self
258    }
259
260    /// Apply cost guardrails to this engine.
261    ///
262    /// Without this, both the per-run cap default and the monthly quota are
263    /// disabled and the engine behaves exactly as before.
264    ///
265    /// # Examples
266    ///
267    /// ```no_run
268    /// use std::sync::Arc;
269    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
270    /// use ironflow_engine::budget::BudgetConfig;
271    /// use ironflow_engine::engine::Engine;
272    /// use ironflow_store::memory::InMemoryStore;
273    ///
274    /// let engine = Engine::new(
275    ///     Arc::new(InMemoryStore::new()),
276    ///     Arc::new(ClaudeCodeProvider::new()),
277    /// )
278    /// .with_budget_config(BudgetConfig::from_env());
279    /// ```
280    pub fn with_budget_config(mut self, budget: BudgetConfig) -> Self {
281        self.budget = budget;
282        self
283    }
284
285    /// Returns the cost guardrails applied by this engine.
286    pub fn budget_config(&self) -> &BudgetConfig {
287        &self.budget
288    }
289
290    /// Apply workflow guard configuration to this engine.
291    ///
292    /// When set, every workflow run created by this engine is protected
293    /// by the guard. Handlers can override this via
294    /// [`WorkflowHandler::guard_config`].
295    ///
296    /// # Examples
297    ///
298    /// ```no_run
299    /// use std::sync::Arc;
300    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
301    /// use ironflow_engine::engine::Engine;
302    /// use ironflow_engine::guard::WorkflowGuardConfig;
303    /// use ironflow_store::memory::InMemoryStore;
304    ///
305    /// let engine = Engine::new(
306    ///     Arc::new(InMemoryStore::new()),
307    ///     Arc::new(ClaudeCodeProvider::new()),
308    /// )
309    /// .with_guard_config(WorkflowGuardConfig::default());
310    /// ```
311    pub fn with_guard_config(mut self, config: WorkflowGuardConfig) -> Self {
312        self.guard_config = Some(config);
313        self
314    }
315
316    /// Returns the workflow guard configuration, if any.
317    pub fn guard_config(&self) -> Option<&WorkflowGuardConfig> {
318        self.guard_config.as_ref()
319    }
320
321    /// Attach a log sender for real-time step output streaming.
322    ///
323    /// When set, all workflow contexts created by this engine will forward
324    /// step output (shell stdout/stderr, agent system messages) to the
325    /// given sender.
326    pub fn set_log_sender(&mut self, sender: LogSender) {
327        self.log_sender = Some(sender);
328    }
329
330    /// Attach the backend that stores and serves artifact bytes.
331    ///
332    /// Every context this engine builds inherits it. Without one, steps that
333    /// declare artifacts fail explicitly and every other step is unaffected.
334    ///
335    /// # Examples
336    ///
337    /// ```no_run
338    /// use std::sync::Arc;
339    ///
340    /// use ironflow_engine::artifact::ArtifactSink;
341    /// use ironflow_engine::engine::Engine;
342    ///
343    /// # fn example(engine: &mut Engine, sink: Arc<dyn ArtifactSink>) {
344    /// engine.set_artifact_sink(sink);
345    /// # }
346    /// ```
347    pub fn set_artifact_sink(&mut self, sink: Arc<dyn ArtifactSink>) {
348        self.artifact_sink = Some(sink);
349    }
350
351    /// The artifact backend attached to this engine, if any.
352    pub fn artifact_sink(&self) -> Option<&Arc<dyn ArtifactSink>> {
353        self.artifact_sink.as_ref()
354    }
355
356    /// Attach a [`WorkflowEventBus`] for per-run real-time monitoring.
357    ///
358    /// When set, all workflow contexts created by this engine will publish
359    /// step-level events (`StepStarted`, `StepCompleted`, `StepFailed`)
360    /// to the bus.
361    ///
362    /// # Examples
363    ///
364    /// ```no_run
365    /// use ironflow_engine::engine::Engine;
366    /// use ironflow_engine::notify::WorkflowEventBus;
367    ///
368    /// # fn example(engine: &mut Engine) {
369    /// engine.set_event_bus(WorkflowEventBus::new());
370    /// # }
371    /// ```
372    pub fn set_event_bus(&mut self, bus: WorkflowEventBus) {
373        self.event_bus = Some(bus);
374    }
375
376    /// The workflow event bus attached to this engine, if any.
377    pub fn event_bus(&self) -> Option<&WorkflowEventBus> {
378        self.event_bus.as_ref()
379    }
380
381    /// Returns a reference to the backing store.
382    pub fn store(&self) -> &Arc<dyn Store> {
383        &self.store
384    }
385
386    /// Returns a reference to the agent provider.
387    pub fn provider(&self) -> &Arc<dyn AgentProvider> {
388        &self.provider
389    }
390
391    /// Build a [`WorkflowContext`] with access to the handler registry.
392    ///
393    /// The context is seeded with the run's attempt number and the cost and
394    /// duration already accumulated by previous attempts, so a retried run
395    /// reports the total it really consumed rather than only its last attempt.
396    ///
397    /// The run's persisted `max_cost_usd` becomes the context cost cap; `None`
398    /// disables the per-run budget check for that context.
399    fn build_context(&self, run: &Run) -> WorkflowContext {
400        let handlers = self.handlers.clone();
401        let resolver: crate::context::HandlerResolver =
402            Arc::new(move |name: &str| handlers.get(name).cloned());
403        let mut ctx = WorkflowContext::with_handler_resolver(
404            run.id,
405            run.workflow_name.clone(),
406            self.store.clone(),
407            self.provider.clone(),
408            resolver,
409        );
410        ctx.carry_over_run_totals(run.retry_count + 1, run.cost_usd, run.duration_ms);
411        ctx.set_max_cost_usd(run.max_cost_usd);
412        if let Some(ref sender) = self.log_sender {
413            ctx.set_log_sender(sender.clone());
414        }
415        if let Some(ref sink) = self.artifact_sink {
416            ctx.set_artifact_sink(sink.clone());
417        }
418        if let Some(ref bus) = self.event_bus {
419            ctx.set_event_bus(bus.clone());
420        }
421        if let Some(ref provider) = self.decision_provider {
422            ctx.set_decision_provider(provider.clone());
423        }
424        ctx
425    }
426
427    /// Build a context with the guard attached.
428    ///
429    /// Uses the handler's `guard_config()` if present, otherwise falls back
430    /// to the engine's global configuration. Creates a fresh shared state
431    /// for each top-level run.
432    fn build_context_with_guard(
433        &self,
434        run: &Run,
435        handler: &dyn WorkflowHandler,
436    ) -> WorkflowContext {
437        let mut ctx = self.build_context(run);
438        let guard_config = handler.guard_config().or_else(|| self.guard_config.clone());
439        if let Some(config) = guard_config {
440            ctx.set_guard(config, new_shared_guard_state());
441        }
442        ctx
443    }
444
445    /// Reject the creation of a new run when the monthly quota is exhausted.
446    ///
447    /// The window is the current calendar month in UTC. Runs already in flight
448    /// are never interrupted -- only creation is refused.
449    ///
450    /// # Errors
451    ///
452    /// Returns [`EngineError::MonthlyBudgetExceeded`] when the accumulated cost
453    /// of the month has reached the configured quota. Returns
454    /// [`EngineError::Store`] when the aggregate query fails.
455    async fn check_monthly_quota(&self, workflow_name: &str) -> Result<(), EngineError> {
456        let Some(limit) = self.budget.monthly_cost_limit_usd else {
457            return Ok(());
458        };
459
460        let stats = self
461            .store
462            .get_stats(RunFilter {
463                created_after: Some(month_start(Utc::now())),
464                ..RunFilter::default()
465            })
466            .await?;
467
468        if stats.total_cost_usd < limit {
469            return Ok(());
470        }
471
472        warn!(
473            workflow = %workflow_name,
474            limit_usd = %limit,
475            spent_usd = %stats.total_cost_usd,
476            "monthly cost quota exhausted, refusing new run"
477        );
478
479        #[cfg(feature = "prometheus")]
480        counter!(
481            RUN_BUDGET_EXCEEDED_TOTAL,
482            "workflow" => workflow_name.to_string(),
483            "scope" => "monthly",
484        )
485        .increment(1);
486
487        Err(EngineError::MonthlyBudgetExceeded {
488            limit_usd: limit,
489            spent_usd: stats.total_cost_usd,
490        })
491    }
492
493    // -----------------------------------------------------------------------
494    // Handler registration
495    // -----------------------------------------------------------------------
496
497    /// Register a [`WorkflowHandler`] for dynamic workflow execution.
498    ///
499    /// The handler is looked up by [`WorkflowHandler::name`] when executing
500    /// or enqueuing.
501    ///
502    /// # Errors
503    ///
504    /// Returns [`EngineError::InvalidWorkflow`] if a handler with the same
505    /// name is already registered.
506    ///
507    /// # Examples
508    ///
509    /// ```no_run
510    /// use std::sync::Arc;
511    /// use ironflow_engine::engine::Engine;
512    /// use ironflow_engine::handler::{WorkflowHandler, HandlerFuture};
513    /// use ironflow_engine::context::WorkflowContext;
514    /// use ironflow_engine::config::ShellConfig;
515    /// use ironflow_store::memory::InMemoryStore;
516    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
517    ///
518    /// struct MyWorkflow;
519    /// impl WorkflowHandler for MyWorkflow {
520    ///     fn name(&self) -> &str { "my-workflow" }
521    ///     fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
522    ///         Box::pin(async move {
523    ///             ctx.shell("step1", ShellConfig::new("echo done")).await?;
524    ///             Ok(())
525    ///         })
526    ///     }
527    /// }
528    ///
529    /// let mut engine = Engine::new(
530    ///     Arc::new(InMemoryStore::new()),
531    ///     Arc::new(ClaudeCodeProvider::new()),
532    /// );
533    /// engine.register(MyWorkflow)?;
534    /// # Ok::<(), ironflow_engine::error::EngineError>(())
535    /// ```
536    pub fn register(&mut self, handler: impl WorkflowHandler + 'static) -> Result<(), EngineError> {
537        let name = handler.name().to_string();
538        if self.handlers.contains_key(&name) {
539            return Err(EngineError::InvalidWorkflow(format!(
540                "handler '{}' already registered",
541                name
542            )));
543        }
544        if let Some(category) = handler.category() {
545            validate_category(&name, category)?;
546        }
547        self.handlers.insert(name, Arc::new(handler));
548        Ok(())
549    }
550
551    /// Register a pre-boxed workflow handler.
552    ///
553    /// # Errors
554    ///
555    /// Returns [`EngineError::InvalidWorkflow`] if a handler with the same
556    /// name is already registered or if its category is invalid.
557    pub fn register_boxed(&mut self, handler: Box<dyn WorkflowHandler>) -> Result<(), EngineError> {
558        let name = handler.name().to_string();
559        if self.handlers.contains_key(&name) {
560            return Err(EngineError::InvalidWorkflow(format!(
561                "handler '{}' already registered",
562                name
563            )));
564        }
565        if let Some(category) = handler.category() {
566            validate_category(&name, category)?;
567        }
568        self.handlers.insert(name, Arc::from(handler));
569        Ok(())
570    }
571
572    /// Get a registered handler by name.
573    pub fn get_handler(&self, name: &str) -> Option<&Arc<dyn WorkflowHandler>> {
574        self.handlers.get(name)
575    }
576
577    /// List registered handler names.
578    pub fn handler_names(&self) -> Vec<&str> {
579        self.handlers.keys().map(|s| s.as_str()).collect()
580    }
581
582    /// Get detailed info about a registered workflow handler.
583    pub fn handler_info(&self, name: &str) -> Option<WorkflowInfo> {
584        self.handlers.get(name).map(|h| h.describe())
585    }
586
587    /// List handlers that have a cron schedule configured.
588    ///
589    /// Returns pairs of `(workflow_name, cron_expression)` for all handlers
590    /// where [`WorkflowHandler::schedule`] returns `Some`.
591    ///
592    /// Use this to wire scheduled handlers into a scheduler
593    /// (e.g. the schedule ticker in `ironflow-api`).
594    ///
595    /// # Examples
596    ///
597    /// ```no_run
598    /// # use std::sync::Arc;
599    /// # use ironflow_engine::engine::Engine;
600    /// # use ironflow_store::memory::InMemoryStore;
601    /// # use ironflow_core::providers::claude::ClaudeCodeProvider;
602    /// let engine = Engine::new(
603    ///     Arc::new(InMemoryStore::new()),
604    ///     Arc::new(ClaudeCodeProvider::new()),
605    /// );
606    /// for (name, schedule) in engine.scheduled_handlers() {
607    ///     tracing::info!("{name} runs on schedule: {schedule}");
608    /// }
609    /// ```
610    pub fn scheduled_handlers(&self) -> Vec<(&str, &CronSchedule)> {
611        self.handlers
612            .iter()
613            .filter_map(|(name, handler)| handler.schedule().map(|sched| (name.as_str(), sched)))
614            .collect()
615    }
616
617    /// Register an event subscriber for domain events.
618    ///
619    /// The subscriber is called only for events whose type is in
620    /// `event_types`. Pass [`Event::ALL`] to receive every event.
621    ///
622    /// # Examples
623    ///
624    /// ```no_run
625    /// use ironflow_engine::engine::Engine;
626    /// use ironflow_engine::notify::{Event, WebhookSubscriber};
627    /// use ironflow_store::memory::InMemoryStore;
628    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
629    /// use std::sync::Arc;
630    ///
631    /// let mut engine = Engine::new(
632    ///     Arc::new(InMemoryStore::new()),
633    ///     Arc::new(ClaudeCodeProvider::new()),
634    /// );
635    ///
636    /// engine.subscribe(
637    ///     WebhookSubscriber::new("https://hooks.example.com/events"),
638    ///     &[Event::RUN_STATUS_CHANGED, Event::STEP_FAILED],
639    /// );
640    /// ```
641    pub fn subscribe(
642        &mut self,
643        subscriber: impl EventSubscriber + 'static,
644        event_types: &[&'static str],
645    ) {
646        self.event_publisher.subscribe(subscriber, event_types);
647    }
648
649    /// Returns a reference to the event publisher.
650    ///
651    /// Useful for publishing events from outside the engine (e.g. auth
652    /// routes in the API layer).
653    pub fn event_publisher(&self) -> &EventPublisher {
654        &self.event_publisher
655    }
656
657    // -----------------------------------------------------------------------
658    // Dynamic workflow execution (WorkflowHandler)
659    // -----------------------------------------------------------------------
660
661    /// Execute a registered handler inline.
662    ///
663    /// Creates a run, builds a [`WorkflowContext`], calls the handler's
664    /// [`execute`](WorkflowHandler::execute), and finalizes the run.
665    ///
666    /// # Errors
667    ///
668    /// Returns [`EngineError::InvalidWorkflow`] if no handler is registered
669    /// with that name. Returns [`EngineError`] if execution fails.
670    ///
671    /// # Examples
672    ///
673    /// ```no_run
674    /// use std::sync::Arc;
675    /// use ironflow_engine::engine::Engine;
676    /// use ironflow_store::memory::InMemoryStore;
677    /// use ironflow_store::models::TriggerKind;
678    /// use ironflow_core::providers::claude::ClaudeCodeProvider;
679    /// use serde_json::json;
680    ///
681    /// # async fn example(engine: &Engine) -> Result<(), ironflow_engine::error::EngineError> {
682    /// let run = engine.run_handler("deploy", TriggerKind::Manual, json!({})).await?;
683    /// # Ok(())
684    /// # }
685    /// ```
686    #[tracing::instrument(name = "engine.run_handler", skip_all, fields(workflow = %handler_name))]
687    pub async fn run_handler(
688        &self,
689        handler_name: &str,
690        trigger: TriggerKind,
691        payload: Value,
692    ) -> Result<WorkflowResult, EngineError> {
693        let handler = self
694            .handlers
695            .get(handler_name)
696            .ok_or_else(|| {
697                EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
698            })?
699            .clone();
700
701        self.check_monthly_quota(handler_name).await?;
702
703        let handler_version = handler.version().map(str::to_string);
704        let max_cost_usd = self
705            .budget
706            .resolve_run_cap(None, handler.default_max_cost_usd());
707        let run = self
708            .store
709            .create_run(NewRun {
710                created_by: None,
711                workflow_name: handler_name.to_string(),
712                trigger,
713                payload,
714                max_retries: 0,
715                handler_version,
716                labels: handler.default_labels(),
717                scheduled_at: None,
718                idempotency_key: None,
719                max_cost_usd,
720            })
721            .await?
722            .into_run();
723
724        let run_id = run.id;
725        info!(run_id = %run_id, handler_version = run.handler_version.as_deref().unwrap_or(""), "run created");
726
727        self.store
728            .update_run_status(run_id, RunStatus::Running)
729            .await?;
730
731        #[cfg(feature = "prometheus")]
732        gauge!(RUNS_ACTIVE, "workflow" => handler_name.to_string()).increment(1.0);
733
734        let run_start = Instant::now();
735        let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
736
737        let result = handler.execute(&mut ctx).await;
738        self.finalize_run(run_id, handler_name, result, &ctx, run_start, run.labels)
739            .await
740    }
741
742    /// Enqueue a handler-based workflow for worker execution.
743    ///
744    /// The workflow name is stored in the run. The worker looks up the
745    /// handler by name when executing.
746    ///
747    /// # Errors
748    ///
749    /// Returns [`EngineError::InvalidWorkflow`] if no handler is registered.
750    /// Returns [`EngineError::MonthlyBudgetExceeded`] if the monthly cost quota
751    /// is exhausted.
752    #[tracing::instrument(name = "engine.enqueue_handler", skip_all, fields(workflow = %handler_name))]
753    pub async fn enqueue_handler(
754        &self,
755        handler_name: &str,
756        trigger: TriggerKind,
757        payload: Value,
758        max_retries: u32,
759    ) -> Result<Run, EngineError> {
760        self.enqueue_handler_with_options(
761            handler_name,
762            trigger,
763            payload,
764            EnqueueOptions {
765                max_retries,
766                ..Default::default()
767            },
768        )
769        .await
770        .map(RunCreation::into_run)
771    }
772
773    /// Enqueue a handler-based workflow with labels, deferred scheduling, an
774    /// optional cost cap, an optional author, and an optional idempotency key.
775    ///
776    /// See [`EnqueueOptions`] for the individual settings.
777    ///
778    /// When [`EnqueueOptions::idempotency_key`] is set and already bound to a run
779    /// created within
780    /// [`IDEMPOTENCY_WINDOW`](ironflow_store::entities::IDEMPOTENCY_WINDOW), nothing
781    /// is enqueued and the original run is returned as [`RunCreation::Existing`].
782    ///
783    /// # Errors
784    ///
785    /// Returns [`EngineError::InvalidWorkflow`] if no handler is registered.
786    /// Returns [`EngineError::MonthlyBudgetExceeded`] if the monthly cost quota
787    /// is exhausted. Returns [`EngineError::Store`] if the run cannot be
788    /// persisted.
789    ///
790    /// # Examples
791    ///
792    /// ```no_run
793    /// use ironflow_engine::engine::{Engine, EnqueueOptions};
794    /// use ironflow_store::models::TriggerKind;
795    /// use serde_json::json;
796    ///
797    /// # async fn example(engine: &Engine) -> Result<(), ironflow_engine::error::EngineError> {
798    /// let creation = engine
799    ///     .enqueue_handler_with_options(
800    ///         "deploy",
801    ///         TriggerKind::Api,
802    ///         json!({"env": "prod"}),
803    ///         EnqueueOptions {
804    ///             max_retries: 3,
805    ///             idempotency_key: Some("github:abc-123".to_string()),
806    ///             ..Default::default()
807    ///         },
808    ///     )
809    ///     .await?;
810    ///
811    /// if creation.is_created() {
812    ///     println!("enqueued {}", creation.run().id);
813    /// }
814    /// # Ok(())
815    /// # }
816    /// ```
817    #[tracing::instrument(name = "engine.enqueue_handler_with_options", skip_all, fields(workflow = %handler_name))]
818    pub async fn enqueue_handler_with_options(
819        &self,
820        handler_name: &str,
821        trigger: TriggerKind,
822        payload: Value,
823        options: EnqueueOptions,
824    ) -> Result<RunCreation, EngineError> {
825        let EnqueueOptions {
826            max_retries,
827            labels,
828            scheduled_at,
829            max_cost_usd,
830            created_by,
831            idempotency_key,
832        } = options;
833
834        let handler = self.handlers.get(handler_name).ok_or_else(|| {
835            EngineError::InvalidWorkflow(format!("no handler registered: {handler_name}"))
836        })?;
837
838        self.check_monthly_quota(handler_name).await?;
839
840        let handler_version = handler.version().map(str::to_string);
841        let mut merged_labels = handler.default_labels();
842        merged_labels.extend(labels);
843        let resolved_cap = self
844            .budget
845            .resolve_run_cap(max_cost_usd, handler.default_max_cost_usd());
846
847        let creation = self
848            .store
849            .create_run(NewRun {
850                workflow_name: handler_name.to_string(),
851                trigger,
852                payload,
853                max_retries,
854                handler_version,
855                labels: merged_labels,
856                scheduled_at,
857                created_by,
858                idempotency_key,
859                max_cost_usd: resolved_cap,
860            })
861            .await?;
862
863        match &creation {
864            RunCreation::Created(run) => info!(
865                run_id = %run.id,
866                workflow = %handler_name,
867                max_cost_usd = ?resolved_cap,
868                "handler run enqueued"
869            ),
870            RunCreation::Existing(run) => info!(
871                run_id = %run.id,
872                workflow = %handler_name,
873                "idempotent replay, nothing enqueued"
874            ),
875        }
876
877        Ok(creation)
878    }
879
880    /// Execute a handler-based run (used by the worker after pick_next_pending).
881    ///
882    /// Looks up the handler by the run's `workflow_name` and executes it
883    /// with a fresh [`WorkflowContext`].
884    ///
885    /// # Errors
886    ///
887    /// Returns [`EngineError::InvalidWorkflow`] if no handler matches.
888    #[tracing::instrument(name = "engine.execute_handler_run", skip_all, fields(run_id = %run_id))]
889    pub async fn execute_handler_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
890        let run = self
891            .store
892            .get_run(run_id)
893            .await?
894            .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
895
896        let handler = self
897            .handlers
898            .get(&run.workflow_name)
899            .ok_or_else(|| {
900                EngineError::InvalidWorkflow(format!(
901                    "no handler registered: {}",
902                    run.workflow_name
903                ))
904            })?
905            .clone();
906
907        #[cfg(feature = "prometheus")]
908        gauge!(RUNS_ACTIVE, "workflow" => run.workflow_name.clone()).increment(1.0);
909
910        let run_start = Instant::now();
911        let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
912
913        // On a retry the whole run is replayed from the start, but an approval a
914        // human already granted must not be asked again.
915        if run.retry_count > 0 {
916            ctx.load_replay_steps().await?;
917        }
918
919        let result = handler.execute(&mut ctx).await;
920        self.finalize_run(
921            run_id,
922            &run.workflow_name,
923            result,
924            &ctx,
925            run_start,
926            run.labels,
927        )
928        .await
929    }
930
931    /// Execute a run by its ID (used by the worker after pick_next_pending).
932    ///
933    /// Delegates to [`execute_handler_run`](Self::execute_handler_run).
934    ///
935    /// # Errors
936    ///
937    /// Returns [`EngineError`] if the run is not found or execution fails.
938    #[tracing::instrument(name = "engine.execute_run", skip_all, fields(run_id = %run_id))]
939    pub async fn execute_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
940        self.execute_handler_run(run_id).await
941    }
942
943    /// Resume a run after human approval.
944    ///
945    /// Re-executes the handler with step replay: completed steps return
946    /// cached output, approved approval steps are skipped, and execution
947    /// continues from the first unexecuted step.
948    ///
949    /// Supports multiple approval gates -- each resume replays all prior
950    /// steps and stops at the next approval (or completes the run).
951    ///
952    /// # Errors
953    ///
954    /// Returns [`EngineError::InvalidWorkflow`] if no handler matches.
955    /// Returns [`EngineError`] if execution fails or hits another approval.
956    #[tracing::instrument(name = "engine.resume_run", skip_all, fields(run_id = %run_id))]
957    pub async fn resume_run(&self, run_id: Uuid) -> Result<WorkflowResult, EngineError> {
958        let run = self
959            .store
960            .get_run(run_id)
961            .await?
962            .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
963
964        let handler = self
965            .handlers
966            .get(&run.workflow_name)
967            .ok_or_else(|| {
968                EngineError::InvalidWorkflow(format!(
969                    "no handler registered: {}",
970                    run.workflow_name
971                ))
972            })?
973            .clone();
974
975        info!(run_id = %run_id, workflow = %run.workflow_name, "resuming run after approval");
976
977        let run_start = Instant::now();
978        let mut ctx = self.build_context_with_guard(&run, handler.as_ref());
979        ctx.load_replay_steps().await?;
980
981        let result = handler.execute(&mut ctx).await;
982        self.finalize_run(
983            run_id,
984            &run.workflow_name,
985            result,
986            &ctx,
987            run_start,
988            run.labels,
989        )
990        .await
991    }
992
993    /// Record a run failure, replaying the run later when retries remain.
994    ///
995    /// This is the single place where a failed run's fate is decided. When
996    /// `retryable` is true and the run has not exhausted `max_retries`, the run
997    /// moves to [`RunStatus::Retrying`] with `scheduled_at` set to
998    /// `now + backoff` -- [`pick_next_pending`](ironflow_store::store::RunStore::pick_next_pending)
999    /// picks it up again once that time has passed. Otherwise the run moves to
1000    /// [`RunStatus::Failed`].
1001    ///
1002    /// Either way, steps left non-terminal by the failed attempt are closed via
1003    /// [`fail_orphaned_steps`](Self::fail_orphaned_steps) so they are never
1004    /// confused with the next attempt's steps.
1005    ///
1006    /// Callers pass `retryable` explicitly rather than an error value, because
1007    /// the worker classifies failures it observes from the outside (a timeout, a
1008    /// panicked task) that never produce an [`EngineError`]. Use
1009    /// [`is_run_retryable`] to classify an
1010    /// [`EngineError`].
1011    ///
1012    /// `cost_usd` and `duration_ms` are `None` when the caller does not know the
1013    /// totals (a timeout or a panic observed from outside the handler); the
1014    /// values already stored on the run are then left untouched.
1015    ///
1016    /// Returns the status the run was moved to.
1017    ///
1018    /// # Errors
1019    ///
1020    /// Returns [`EngineError::Store`] if the run does not exist or the update
1021    /// cannot be persisted.
1022    ///
1023    /// # Examples
1024    ///
1025    /// ```no_run
1026    /// use ironflow_engine::engine::Engine;
1027    /// use ironflow_engine::error::EngineError;
1028    /// use uuid::Uuid;
1029    ///
1030    /// # async fn example(engine: &Engine, run_id: Uuid) -> Result<(), EngineError> {
1031    /// let status = engine
1032    ///     .fail_or_schedule_retry(run_id, "run timed out after 300s", true, None, None)
1033    ///     .await?;
1034    /// # Ok(())
1035    /// # }
1036    /// ```
1037    pub async fn fail_or_schedule_retry(
1038        &self,
1039        run_id: Uuid,
1040        error: &str,
1041        retryable: bool,
1042        cost_usd: Option<Decimal>,
1043        duration_ms: Option<u64>,
1044    ) -> Result<RunStatus, EngineError> {
1045        let run = self
1046            .store
1047            .get_run(run_id)
1048            .await?
1049            .ok_or(EngineError::Store(StoreError::RunNotFound(run_id)))?;
1050
1051        let has_attempts_left = run.retry_count < run.max_retries;
1052        let update = if retryable && has_attempts_left {
1053            let backoff = backoff_for_retry(run.retry_count);
1054            let scheduled_at = Utc::now() + TimeDelta::milliseconds(backoff.as_millis() as i64);
1055
1056            info!(
1057                run_id = %run_id,
1058                workflow = %run.workflow_name,
1059                attempt = run.retry_count + 1,
1060                max_retries = run.max_retries,
1061                backoff_secs = backoff.as_secs(),
1062                scheduled_at = %scheduled_at,
1063                "run failed, scheduling retry"
1064            );
1065
1066            RunUpdate {
1067                status: Some(RunStatus::Retrying),
1068                error: Some(error.to_string()),
1069                increment_retry: true,
1070                cost_usd,
1071                duration_ms,
1072                scheduled_at: Some(scheduled_at),
1073                ..RunUpdate::default()
1074            }
1075        } else {
1076            RunUpdate {
1077                status: Some(RunStatus::Failed),
1078                error: Some(error.to_string()),
1079                cost_usd,
1080                duration_ms,
1081                completed_at: Some(Utc::now()),
1082                ..RunUpdate::default()
1083            }
1084        };
1085
1086        let status = update.status.unwrap_or(RunStatus::Failed);
1087        self.store.update_run(run_id, update).await?;
1088        self.fail_orphaned_steps(run_id, error).await?;
1089
1090        Ok(status)
1091    }
1092
1093    /// Fail all non-terminal steps for a run.
1094    ///
1095    /// Called after a run is marked as failed (timeout, error, panic) to clean up
1096    /// orphaned steps that are still in `Running`, `Pending`, or `AwaitingApproval`.
1097    ///
1098    /// - `Running` / `AwaitingApproval` steps are marked `Failed`.
1099    /// - `Pending` steps are marked `Skipped` (FSM does not allow Pending -> Failed).
1100    ///
1101    /// Errors from individual step updates are logged but do not abort the cleanup.
1102    ///
1103    /// # Errors
1104    ///
1105    /// Returns [`EngineError`] if listing steps fails.
1106    pub async fn fail_orphaned_steps(
1107        &self,
1108        run_id: Uuid,
1109        error_message: &str,
1110    ) -> Result<(), EngineError> {
1111        let steps = self.store.list_steps(run_id).await?;
1112        let now = Utc::now();
1113
1114        for step in steps {
1115            if step.status.state.is_terminal() {
1116                continue;
1117            }
1118
1119            let (target_status, error) = match step.status.state {
1120                StepStatus::Running | StepStatus::AwaitingApproval => {
1121                    let err = if step.error.is_some() {
1122                        None
1123                    } else {
1124                        Some(error_message.to_string())
1125                    };
1126                    (StepStatus::Failed, err)
1127                }
1128                StepStatus::Pending => (StepStatus::Skipped, None),
1129                _ => continue,
1130            };
1131
1132            if let Err(e) = self
1133                .store
1134                .update_step(
1135                    step.id,
1136                    StepUpdate {
1137                        status: Some(target_status),
1138                        error,
1139                        completed_at: Some(now),
1140                        ..StepUpdate::default()
1141                    },
1142                )
1143                .await
1144            {
1145                warn!(
1146                    run_id = %run_id,
1147                    step_id = %step.id,
1148                    step_name = %step.name,
1149                    error = %e,
1150                    "failed to cleanup orphaned step"
1151                );
1152            } else {
1153                info!(
1154                    run_id = %run_id,
1155                    step_id = %step.id,
1156                    step_name = %step.name,
1157                    from = %step.status.state,
1158                    to = %target_status,
1159                    "cleaned up orphaned step"
1160                );
1161            }
1162        }
1163
1164        Ok(())
1165    }
1166
1167    /// Finalize a run with the given result and context.
1168    ///
1169    /// On success: updates run to Completed with cost, duration, and completed_at.
1170    /// On failure: updates run to Failed with error, cost, duration, and completed_at.
1171    /// Always: fetches and returns the final Run.
1172    async fn finalize_run(
1173        &self,
1174        run_id: Uuid,
1175        workflow_name: &str,
1176        result: Result<(), EngineError>,
1177        ctx: &WorkflowContext,
1178        run_start: Instant,
1179        run_labels: HashMap<String, String>,
1180    ) -> Result<WorkflowResult, EngineError> {
1181        // Covers the whole run: previous attempts plus this one, so a retried
1182        // run reports the time it really consumed.
1183        let total_duration = ctx.carried_duration_ms() + run_start.elapsed().as_millis() as u64;
1184        let completed_at = Utc::now();
1185
1186        let final_status;
1187        let final_run;
1188
1189        match result {
1190            Ok(()) => {
1191                final_status = if ctx.has_allowed_failure() {
1192                    RunStatus::Warning
1193                } else {
1194                    RunStatus::Completed
1195                };
1196                final_run = self
1197                    .store
1198                    .update_run_returning(
1199                        run_id,
1200                        RunUpdate {
1201                            status: Some(final_status),
1202                            cost_usd: Some(ctx.total_cost_usd()),
1203                            duration_ms: Some(total_duration),
1204                            completed_at: Some(completed_at),
1205                            ..RunUpdate::default()
1206                        },
1207                    )
1208                    .await?;
1209
1210                info!(
1211                    run_id = %run_id,
1212                    status = %final_status,
1213                    cost_usd = %ctx.total_cost_usd(),
1214                    duration_ms = total_duration,
1215                    "run completed"
1216                );
1217            }
1218            Err(EngineError::ApprovalRequired {
1219                run_id: approval_run_id,
1220                step_id,
1221                ref message,
1222            }) => {
1223                final_status = RunStatus::AwaitingApproval;
1224                final_run = self
1225                    .store
1226                    .update_run_returning(
1227                        run_id,
1228                        RunUpdate {
1229                            status: Some(RunStatus::AwaitingApproval),
1230                            cost_usd: Some(ctx.total_cost_usd()),
1231                            duration_ms: Some(total_duration),
1232                            ..RunUpdate::default()
1233                        },
1234                    )
1235                    .await?;
1236
1237                info!(
1238                    run_id = %approval_run_id,
1239                    step_id = %step_id,
1240                    message = %message,
1241                    "run awaiting approval"
1242                );
1243            }
1244            Err(EngineError::DelaySleeping {
1245                run_id: delay_run_id,
1246                step_id,
1247                wake_at,
1248            }) => {
1249                final_status = RunStatus::Sleeping;
1250                final_run = self
1251                    .store
1252                    .update_run_returning(
1253                        run_id,
1254                        RunUpdate {
1255                            status: Some(RunStatus::Sleeping),
1256                            cost_usd: Some(ctx.total_cost_usd()),
1257                            duration_ms: Some(total_duration),
1258                            scheduled_at: Some(wake_at),
1259                            ..RunUpdate::default()
1260                        },
1261                    )
1262                    .await?;
1263
1264                info!(
1265                    run_id = %delay_run_id,
1266                    step_id = %step_id,
1267                    wake_at = %wake_at,
1268                    "run sleeping until delay elapses"
1269                );
1270            }
1271            Err(err) => {
1272                // A guardrail stop (budget or workflow guard) is deliberate,
1273                // not a breakage: the run is cancelled, never failed and
1274                // never replayed.
1275                let guardrail_stop = matches!(
1276                    err,
1277                    EngineError::RunBudgetExceeded { .. } | EngineError::WorkflowGuardRejected(_)
1278                );
1279
1280                final_status = if guardrail_stop {
1281                    if let Err(store_err) = self
1282                        .store
1283                        .update_run(
1284                            run_id,
1285                            RunUpdate {
1286                                status: Some(RunStatus::Cancelled),
1287                                error: Some(err.to_string()),
1288                                cost_usd: Some(ctx.total_cost_usd()),
1289                                duration_ms: Some(total_duration),
1290                                completed_at: Some(completed_at),
1291                                ..RunUpdate::default()
1292                            },
1293                        )
1294                        .await
1295                    {
1296                        error!(run_id = %run_id, store_error = %store_err, "failed to persist run cancellation");
1297                    }
1298                    if let Err(cleanup_err) = self
1299                        .fail_orphaned_steps(run_id, "run stopped: guardrail limit reached")
1300                        .await
1301                    {
1302                        error!(run_id = %run_id, store_error = %cleanup_err, "failed to cleanup orphaned steps");
1303                    }
1304                    RunStatus::Cancelled
1305                } else {
1306                    self.fail_or_schedule_retry(
1307                        run_id,
1308                        &err.to_string(),
1309                        is_run_retryable(&err),
1310                        Some(ctx.total_cost_usd()),
1311                        Some(total_duration),
1312                    )
1313                    .await
1314                    .unwrap_or_else(|store_err| {
1315                        error!(run_id = %run_id, store_error = %store_err, "failed to persist run failure");
1316                        RunStatus::Failed
1317                    })
1318                };
1319
1320                if matches!(err, EngineError::RunBudgetExceeded { .. }) {
1321                    self.on_run_budget_exceeded(workflow_name, run_id, &err);
1322                }
1323
1324                error!(run_id = %run_id, status = %final_status, error = %err, "run stopped");
1325
1326                self.publish_run_status_changed(
1327                    workflow_name,
1328                    run_id,
1329                    final_status,
1330                    Some(err.to_string()),
1331                    ctx,
1332                    total_duration,
1333                    run_labels,
1334                );
1335
1336                #[cfg(feature = "prometheus")]
1337                self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
1338
1339                return Err(err);
1340            }
1341        }
1342
1343        self.publish_run_status_changed(
1344            workflow_name,
1345            run_id,
1346            final_status,
1347            None,
1348            ctx,
1349            total_duration,
1350            run_labels,
1351        );
1352
1353        #[cfg(feature = "prometheus")]
1354        self.emit_run_metrics(workflow_name, final_status, total_duration, ctx);
1355
1356        Ok(WorkflowResult {
1357            run: final_run,
1358            steps: ctx.step_results().to_vec(),
1359        })
1360    }
1361
1362    /// Emit Prometheus metrics for a completed run.
1363    #[cfg(feature = "prometheus")]
1364    fn emit_run_metrics(
1365        &self,
1366        workflow_name: &str,
1367        status: RunStatus,
1368        duration_ms: u64,
1369        ctx: &WorkflowContext,
1370    ) {
1371        let status_str = status.to_string();
1372        let wf = workflow_name.to_string();
1373
1374        counter!(RUNS_TOTAL, "workflow" => wf.clone(), "status" => status_str.clone()).increment(1);
1375        histogram!(RUN_DURATION_SECONDS, "workflow" => wf.clone(), "status" => status_str)
1376            .record(duration_ms as f64 / 1000.0);
1377        histogram!(RUN_COST_USD, "workflow" => wf.clone()).record(
1378            ctx.total_cost_usd()
1379                .to_string()
1380                .parse::<f64>()
1381                .unwrap_or(0.0),
1382        );
1383        gauge!(RUNS_ACTIVE, "workflow" => wf).decrement(1.0);
1384    }
1385
1386    /// Record the metric and publish the audit event for a run that hit its
1387    /// cost cap.
1388    ///
1389    /// A non-[`RunBudgetExceeded`](EngineError::RunBudgetExceeded) error is
1390    /// ignored, so callers can pass the error unconditionally.
1391    fn on_run_budget_exceeded(&self, workflow_name: &str, run_id: Uuid, err: &EngineError) {
1392        let EngineError::RunBudgetExceeded {
1393            limit_usd,
1394            spent_usd,
1395            step_budget_usd,
1396            ..
1397        } = err
1398        else {
1399            return;
1400        };
1401
1402        #[cfg(feature = "prometheus")]
1403        counter!(
1404            RUN_BUDGET_EXCEEDED_TOTAL,
1405            "workflow" => workflow_name.to_string(),
1406            "scope" => "run",
1407        )
1408        .increment(1);
1409
1410        self.event_publisher.publish(Event::RunBudgetExceeded {
1411            run_id,
1412            workflow_name: workflow_name.to_string(),
1413            limit_usd: *limit_usd,
1414            spent_usd: *spent_usd,
1415            step_budget_usd: *step_budget_usd,
1416            at: Utc::now(),
1417        });
1418    }
1419
1420    /// Publish a run status changed event to all registered subscribers.
1421    ///
1422    /// `from` is always `Running` because `finalize_run` is only called
1423    /// from a running state.
1424    #[allow(clippy::too_many_arguments)]
1425    fn publish_run_status_changed(
1426        &self,
1427        workflow_name: &str,
1428        run_id: Uuid,
1429        to: RunStatus,
1430        error: Option<String>,
1431        ctx: &WorkflowContext,
1432        duration_ms: u64,
1433        labels: HashMap<String, String>,
1434    ) {
1435        let now = Utc::now();
1436        let cost_usd = ctx.total_cost_usd();
1437        let wf = workflow_name.to_string();
1438
1439        self.event_publisher.publish(Event::RunStatusChanged {
1440            run_id,
1441            workflow_name: wf.clone(),
1442            from: RunStatus::Running,
1443            to,
1444            error: error.clone(),
1445            cost_usd,
1446            duration_ms,
1447            labels: labels.clone(),
1448            at: now,
1449        });
1450
1451        if to == RunStatus::Failed {
1452            self.event_publisher.publish(Event::RunFailed {
1453                run_id,
1454                workflow_name: wf,
1455                error,
1456                cost_usd,
1457                duration_ms,
1458                labels,
1459                at: now,
1460            });
1461        }
1462    }
1463}
1464
1465impl fmt::Debug for Engine {
1466    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1467        f.debug_struct("Engine")
1468            .field("handlers", &self.handlers.keys().collect::<Vec<_>>())
1469            .finish_non_exhaustive()
1470    }
1471}
1472
1473#[cfg(test)]
1474mod tests {
1475    use super::*;
1476    use crate::config::ShellConfig;
1477    use crate::handler::{HandlerFuture, WorkflowHandler};
1478    use ironflow_core::providers::claude::ClaudeCodeProvider;
1479    use ironflow_core::providers::record_replay::RecordReplayProvider;
1480    use ironflow_store::memory::InMemoryStore;
1481    use ironflow_store::models::StepStatus;
1482    use serde_json::json;
1483
1484    // Test handler that echoes a message via shell
1485    struct EchoWorkflow;
1486
1487    impl WorkflowHandler for EchoWorkflow {
1488        fn name(&self) -> &str {
1489            "echo-workflow"
1490        }
1491
1492        fn describe(&self) -> WorkflowInfo {
1493            WorkflowInfo {
1494                description: "A simple workflow that echoes hello".to_string(),
1495                source_code: None,
1496                sub_workflows: Vec::new(),
1497                category: None,
1498                version: self.version().map(str::to_string),
1499                compatible_versions: Vec::new(),
1500                input_schema: None,
1501                default_labels: HashMap::new(),
1502                schedule: self.schedule().cloned(),
1503                default_max_cost_usd: self.default_max_cost_usd(),
1504            }
1505        }
1506
1507        fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1508            Box::pin(async move {
1509                ctx.shell("greet", ShellConfig::new("echo hello")).await?;
1510                Ok(())
1511            })
1512        }
1513    }
1514
1515    // Test handler that fails
1516    struct FailingWorkflow;
1517
1518    impl WorkflowHandler for FailingWorkflow {
1519        fn name(&self) -> &str {
1520            "failing-workflow"
1521        }
1522
1523        fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1524            Box::pin(async move {
1525                ctx.shell("fail", ShellConfig::new("exit 1")).await?;
1526                Ok(())
1527            })
1528        }
1529    }
1530
1531    fn create_test_engine() -> Engine {
1532        let store = Arc::new(InMemoryStore::new());
1533        let inner = ClaudeCodeProvider::new();
1534        let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
1535            inner,
1536            "/tmp/ironflow-fixtures",
1537        ));
1538        Engine::new(store, provider)
1539    }
1540
1541    #[test]
1542    fn engine_new_creates_instance() {
1543        let engine = create_test_engine();
1544        assert_eq!(engine.handler_names().len(), 0);
1545    }
1546
1547    #[test]
1548    fn engine_register_handler() {
1549        let mut engine = create_test_engine();
1550        let result = engine.register(EchoWorkflow);
1551        assert!(result.is_ok());
1552        assert_eq!(engine.handler_names().len(), 1);
1553        assert!(engine.handler_names().contains(&"echo-workflow"));
1554    }
1555
1556    #[test]
1557    fn engine_register_duplicate_returns_error() {
1558        let mut engine = create_test_engine();
1559        engine.register(EchoWorkflow).unwrap();
1560        let result = engine.register(EchoWorkflow);
1561        assert!(result.is_err());
1562    }
1563
1564    #[test]
1565    fn engine_get_handler_found() {
1566        let mut engine = create_test_engine();
1567        engine.register(EchoWorkflow).unwrap();
1568        let handler = engine.get_handler("echo-workflow");
1569        assert!(handler.is_some());
1570    }
1571
1572    #[test]
1573    fn engine_get_handler_not_found() {
1574        let engine = create_test_engine();
1575        let handler = engine.get_handler("nonexistent");
1576        assert!(handler.is_none());
1577    }
1578
1579    #[test]
1580    fn engine_handler_names_lists_all() {
1581        let mut engine = create_test_engine();
1582        engine.register(EchoWorkflow).unwrap();
1583        engine.register(FailingWorkflow).unwrap();
1584        let names = engine.handler_names();
1585        assert_eq!(names.len(), 2);
1586        assert!(names.contains(&"echo-workflow"));
1587        assert!(names.contains(&"failing-workflow"));
1588    }
1589
1590    #[test]
1591    fn engine_handler_info_returns_description() {
1592        let mut engine = create_test_engine();
1593        engine.register(EchoWorkflow).unwrap();
1594        let info = engine.handler_info("echo-workflow");
1595        assert!(info.is_some());
1596        let info = info.unwrap();
1597        assert_eq!(info.description, "A simple workflow that echoes hello");
1598    }
1599
1600    struct CategorizedWorkflow;
1601
1602    impl WorkflowHandler for CategorizedWorkflow {
1603        fn name(&self) -> &str {
1604            "categorized"
1605        }
1606        fn category(&self) -> Option<&str> {
1607            Some("data/etl")
1608        }
1609        fn execute<'a>(
1610            &'a self,
1611            _ctx: &'a mut WorkflowContext,
1612        ) -> crate::handler::HandlerFuture<'a> {
1613            Box::pin(async move { Ok(()) })
1614        }
1615    }
1616
1617    #[test]
1618    fn engine_default_describe_propagates_category() {
1619        let mut engine = create_test_engine();
1620        engine.register(CategorizedWorkflow).unwrap();
1621        let info = engine.handler_info("categorized").unwrap();
1622        assert_eq!(info.category.as_deref(), Some("data/etl"));
1623    }
1624
1625    #[test]
1626    fn engine_default_describe_without_category() {
1627        let mut engine = create_test_engine();
1628        engine.register(EchoWorkflow).unwrap();
1629        let info = engine.handler_info("echo-workflow").unwrap();
1630        assert!(info.category.is_none());
1631    }
1632
1633    // -----------------------------------------------------------------------
1634    // Schedule tests
1635    // -----------------------------------------------------------------------
1636
1637    struct ScheduledWorkflow {
1638        schedule: CronSchedule,
1639    }
1640
1641    impl ScheduledWorkflow {
1642        fn new() -> Self {
1643            Self {
1644                schedule: CronSchedule::new("0 0 * * * *").unwrap(),
1645            }
1646        }
1647    }
1648
1649    impl WorkflowHandler for ScheduledWorkflow {
1650        fn name(&self) -> &str {
1651            "scheduled"
1652        }
1653        fn schedule(&self) -> Option<&CronSchedule> {
1654            Some(&self.schedule)
1655        }
1656        fn execute<'a>(
1657            &'a self,
1658            _ctx: &'a mut WorkflowContext,
1659        ) -> crate::handler::HandlerFuture<'a> {
1660            Box::pin(async move { Ok(()) })
1661        }
1662    }
1663
1664    #[test]
1665    fn engine_default_describe_propagates_schedule() {
1666        let mut engine = create_test_engine();
1667        engine.register(ScheduledWorkflow::new()).unwrap();
1668        let info = engine.handler_info("scheduled").unwrap();
1669        assert_eq!(
1670            info.schedule.as_ref().map(|s| s.as_str()),
1671            Some("0 0 * * * *")
1672        );
1673    }
1674
1675    #[test]
1676    fn engine_default_describe_without_schedule() {
1677        let mut engine = create_test_engine();
1678        engine.register(EchoWorkflow).unwrap();
1679        let info = engine.handler_info("echo-workflow").unwrap();
1680        assert!(info.schedule.is_none());
1681    }
1682
1683    #[test]
1684    fn scheduled_handlers_returns_only_scheduled() {
1685        let mut engine = create_test_engine();
1686        engine.register(EchoWorkflow).unwrap();
1687        engine.register(ScheduledWorkflow::new()).unwrap();
1688        engine.register(FailingWorkflow).unwrap();
1689
1690        let scheduled = engine.scheduled_handlers();
1691        assert_eq!(scheduled.len(), 1);
1692        assert_eq!(scheduled[0].0, "scheduled");
1693        assert_eq!(scheduled[0].1.as_str(), "0 0 * * * *");
1694    }
1695
1696    #[test]
1697    fn scheduled_handlers_empty_when_none_scheduled() {
1698        let mut engine = create_test_engine();
1699        engine.register(EchoWorkflow).unwrap();
1700        engine.register(FailingWorkflow).unwrap();
1701
1702        let scheduled = engine.scheduled_handlers();
1703        assert!(scheduled.is_empty());
1704    }
1705
1706    struct BadCategoryWorkflow(&'static str);
1707
1708    impl WorkflowHandler for BadCategoryWorkflow {
1709        fn name(&self) -> &str {
1710            "bad-category"
1711        }
1712        fn category(&self) -> Option<&str> {
1713            Some(self.0)
1714        }
1715        fn execute<'a>(
1716            &'a self,
1717            _ctx: &'a mut WorkflowContext,
1718        ) -> crate::handler::HandlerFuture<'a> {
1719            Box::pin(async move { Ok(()) })
1720        }
1721    }
1722
1723    #[test]
1724    fn engine_register_rejects_empty_category() {
1725        let mut engine = create_test_engine();
1726        let err = engine.register(BadCategoryWorkflow("")).unwrap_err();
1727        match err {
1728            EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty category")),
1729            other => panic!("expected InvalidWorkflow, got {other:?}"),
1730        }
1731    }
1732
1733    #[test]
1734    fn engine_register_rejects_leading_slash_category() {
1735        let mut engine = create_test_engine();
1736        let err = engine
1737            .register(BadCategoryWorkflow("/data/etl"))
1738            .unwrap_err();
1739        match err {
1740            EngineError::InvalidWorkflow(msg) => assert!(msg.contains("leading '/'")),
1741            other => panic!("expected InvalidWorkflow, got {other:?}"),
1742        }
1743    }
1744
1745    #[test]
1746    fn engine_register_rejects_trailing_slash_category() {
1747        let mut engine = create_test_engine();
1748        let err = engine
1749            .register(BadCategoryWorkflow("data/etl/"))
1750            .unwrap_err();
1751        match err {
1752            EngineError::InvalidWorkflow(msg) => assert!(msg.contains("trailing '/'")),
1753            other => panic!("expected InvalidWorkflow, got {other:?}"),
1754        }
1755    }
1756
1757    #[test]
1758    fn engine_register_rejects_double_slash_category() {
1759        let mut engine = create_test_engine();
1760        let err = engine
1761            .register(BadCategoryWorkflow("data//etl"))
1762            .unwrap_err();
1763        match err {
1764            EngineError::InvalidWorkflow(msg) => assert!(msg.contains("empty segment")),
1765            other => panic!("expected InvalidWorkflow, got {other:?}"),
1766        }
1767    }
1768
1769    #[test]
1770    fn engine_register_rejects_whitespace_only_segment_category() {
1771        let mut engine = create_test_engine();
1772        let err = engine
1773            .register(BadCategoryWorkflow("data/ /etl"))
1774            .unwrap_err();
1775        match err {
1776            EngineError::InvalidWorkflow(msg) => assert!(msg.contains("whitespace-only segment")),
1777            other => panic!("expected InvalidWorkflow, got {other:?}"),
1778        }
1779    }
1780
1781    #[test]
1782    fn engine_register_accepts_valid_nested_category() {
1783        let mut engine = create_test_engine();
1784        assert!(engine.register(CategorizedWorkflow).is_ok());
1785    }
1786
1787    #[tokio::test]
1788    async fn engine_unknown_workflow_returns_error() {
1789        let engine = create_test_engine();
1790        let result = engine
1791            .run_handler("unknown", TriggerKind::Manual, json!({}))
1792            .await;
1793        assert!(result.is_err());
1794        match result {
1795            Err(EngineError::InvalidWorkflow(msg)) => {
1796                assert!(msg.contains("no handler registered"));
1797            }
1798            _ => panic!("expected InvalidWorkflow error"),
1799        }
1800    }
1801
1802    #[tokio::test]
1803    async fn engine_enqueue_handler_creates_pending_run() {
1804        let mut engine = create_test_engine();
1805        engine.register(EchoWorkflow).unwrap();
1806
1807        let run = engine
1808            .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
1809            .await
1810            .unwrap();
1811        assert_eq!(run.status.state, RunStatus::Pending);
1812        assert_eq!(run.workflow_name, "echo-workflow");
1813    }
1814
1815    #[tokio::test]
1816    async fn enqueue_handler_leaves_the_run_unattributed() {
1817        let mut engine = create_test_engine();
1818        engine.register(EchoWorkflow).unwrap();
1819
1820        let run = engine
1821            .enqueue_handler("echo-workflow", TriggerKind::Manual, json!({}), 0)
1822            .await
1823            .unwrap();
1824
1825        assert!(run.created_by.is_none());
1826    }
1827
1828    #[tokio::test]
1829    async fn enqueue_handler_with_options_records_the_author() {
1830        let mut engine = create_test_engine();
1831        engine.register(EchoWorkflow).unwrap();
1832        let actor = RunActor::User {
1833            user_id: Uuid::now_v7(),
1834        };
1835
1836        let run = engine
1837            .enqueue_handler_with_options(
1838                "echo-workflow",
1839                TriggerKind::Api,
1840                json!({}),
1841                EnqueueOptions {
1842                    created_by: Some(actor.clone()),
1843                    ..Default::default()
1844                },
1845            )
1846            .await
1847            .unwrap()
1848            .into_run();
1849
1850        assert_eq!(run.created_by, Some(actor));
1851    }
1852
1853    #[tokio::test]
1854    async fn enqueue_handler_with_options_accepts_no_author() {
1855        let mut engine = create_test_engine();
1856        engine.register(EchoWorkflow).unwrap();
1857
1858        let run = engine
1859            .enqueue_handler_with_options(
1860                "echo-workflow",
1861                TriggerKind::Cron {
1862                    schedule: "0 * * * * *".to_string(),
1863                },
1864                json!({}),
1865                EnqueueOptions::default(),
1866            )
1867            .await
1868            .unwrap()
1869            .into_run();
1870
1871        assert!(run.created_by.is_none());
1872    }
1873
1874    #[tokio::test]
1875    async fn run_handler_leaves_the_run_unattributed() {
1876        let mut engine = create_test_engine();
1877        engine.register(EchoWorkflow).unwrap();
1878
1879        let run = engine
1880            .run_handler("echo-workflow", TriggerKind::Manual, json!({}))
1881            .await
1882            .unwrap()
1883            .run;
1884
1885        assert!(run.created_by.is_none());
1886    }
1887
1888    #[tokio::test]
1889    async fn engine_register_boxed() {
1890        let mut engine = create_test_engine();
1891        let handler: Box<dyn WorkflowHandler> = Box::new(EchoWorkflow);
1892        let result = engine.register_boxed(handler);
1893        assert!(result.is_ok());
1894        assert_eq!(engine.handler_names().len(), 1);
1895    }
1896
1897    #[tokio::test]
1898    async fn engine_store_and_provider_accessors() {
1899        let store = Arc::new(InMemoryStore::new());
1900        let inner = ClaudeCodeProvider::new();
1901        let provider: Arc<dyn AgentProvider> = Arc::new(RecordReplayProvider::replay(
1902            inner,
1903            "/tmp/ironflow-fixtures",
1904        ));
1905        let engine = Engine::new(store.clone(), provider.clone());
1906
1907        // Verify accessors return references
1908        let _ = engine.store();
1909        let _ = engine.provider();
1910    }
1911
1912    // -----------------------------------------------------------------------
1913    // Operation trait tests
1914    // -----------------------------------------------------------------------
1915
1916    use crate::operation::{Operation, OperationContext};
1917    use async_trait::async_trait;
1918    use ironflow_core::error::OperationError;
1919    use ironflow_store::models::StepKind;
1920
1921    struct FakeGitlabOp {
1922        project_id: u64,
1923        title: String,
1924    }
1925
1926    #[async_trait]
1927    impl Operation for FakeGitlabOp {
1928        fn kind(&self) -> &str {
1929            "gitlab"
1930        }
1931
1932        async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
1933            Ok(json!({
1934                "issue_id": 42,
1935                "project_id": self.project_id,
1936                "title": self.title,
1937            }))
1938        }
1939
1940        fn input(&self) -> Option<Value> {
1941            Some(json!({
1942                "project_id": self.project_id,
1943                "title": self.title,
1944            }))
1945        }
1946    }
1947
1948    struct FailingOp;
1949
1950    #[async_trait]
1951    impl Operation for FailingOp {
1952        fn kind(&self) -> &str {
1953            "broken-service"
1954        }
1955
1956        async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
1957            Err(OperationError::Http {
1958                status: None,
1959                message: "service unavailable".to_string(),
1960            })
1961        }
1962    }
1963
1964    struct OperationWorkflow;
1965
1966    impl WorkflowHandler for OperationWorkflow {
1967        fn name(&self) -> &str {
1968            "operation-workflow"
1969        }
1970
1971        fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1972            Box::pin(async move {
1973                let op = FakeGitlabOp {
1974                    project_id: 123,
1975                    title: "Bug report".to_string(),
1976                };
1977                ctx.operation("create-issue", &op).await?;
1978                Ok(())
1979            })
1980        }
1981    }
1982
1983    struct FailingOperationWorkflow;
1984
1985    impl WorkflowHandler for FailingOperationWorkflow {
1986        fn name(&self) -> &str {
1987            "failing-operation-workflow"
1988        }
1989
1990        fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
1991            Box::pin(async move {
1992                ctx.operation("broken-call", &FailingOp).await?;
1993                Ok(())
1994            })
1995        }
1996    }
1997
1998    struct MixedWorkflow;
1999
2000    impl WorkflowHandler for MixedWorkflow {
2001        fn name(&self) -> &str {
2002            "mixed-workflow"
2003        }
2004
2005        fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2006            Box::pin(async move {
2007                ctx.shell("build", ShellConfig::new("echo built")).await?;
2008                let op = FakeGitlabOp {
2009                    project_id: 456,
2010                    title: "Deploy done".to_string(),
2011                };
2012                let result = ctx.operation("notify-gitlab", &op).await?;
2013                assert_eq!(result.output["issue_id"], 42);
2014                Ok(())
2015            })
2016        }
2017    }
2018
2019    #[tokio::test]
2020    async fn operation_step_happy_path() {
2021        let mut engine = create_test_engine();
2022        engine.register(OperationWorkflow).unwrap();
2023
2024        let run = engine
2025            .run_handler("operation-workflow", TriggerKind::Manual, json!({}))
2026            .await
2027            .unwrap()
2028            .run;
2029
2030        assert_eq!(run.status.state, RunStatus::Completed);
2031
2032        let steps = engine.store().list_steps(run.id).await.unwrap();
2033
2034        assert_eq!(steps.len(), 1);
2035        assert_eq!(steps[0].name, "create-issue");
2036        assert_eq!(steps[0].kind, StepKind::Custom("gitlab".to_string()));
2037        assert_eq!(
2038            steps[0].status.state,
2039            ironflow_store::models::StepStatus::Completed
2040        );
2041
2042        let output = steps[0].output.as_ref().unwrap();
2043        assert_eq!(output["issue_id"], 42);
2044        assert_eq!(output["project_id"], 123);
2045
2046        let input = steps[0].input.as_ref().unwrap();
2047        assert_eq!(input["project_id"], 123);
2048        assert_eq!(input["title"], "Bug report");
2049    }
2050
2051    #[tokio::test]
2052    async fn operation_step_failure_marks_run_failed() {
2053        let mut engine = create_test_engine();
2054        engine.register(FailingOperationWorkflow).unwrap();
2055
2056        let result = engine
2057            .run_handler("failing-operation-workflow", TriggerKind::Manual, json!({}))
2058            .await;
2059
2060        assert!(result.is_err());
2061    }
2062
2063    #[tokio::test]
2064    async fn operation_mixed_with_shell_steps() {
2065        let mut engine = create_test_engine();
2066        engine.register(MixedWorkflow).unwrap();
2067
2068        let run = engine
2069            .run_handler("mixed-workflow", TriggerKind::Manual, json!({}))
2070            .await
2071            .unwrap()
2072            .run;
2073
2074        assert_eq!(run.status.state, RunStatus::Completed);
2075
2076        let steps = engine.store().list_steps(run.id).await.unwrap();
2077
2078        assert_eq!(steps.len(), 2);
2079        assert_eq!(steps[0].kind, StepKind::Shell);
2080        assert_eq!(steps[1].kind, StepKind::Custom("gitlab".to_string()));
2081        assert_eq!(steps[0].position, 0);
2082        assert_eq!(steps[1].position, 1);
2083    }
2084
2085    // -----------------------------------------------------------------------
2086    // Approval + resume tests
2087    // -----------------------------------------------------------------------
2088
2089    use crate::config::ApprovalConfig;
2090
2091    struct SingleApprovalWorkflow;
2092
2093    impl WorkflowHandler for SingleApprovalWorkflow {
2094        fn name(&self) -> &str {
2095            "single-approval"
2096        }
2097
2098        fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2099            Box::pin(async move {
2100                ctx.shell("build", ShellConfig::new("echo built")).await?;
2101                ctx.approval("gate", ApprovalConfig::new("OK?")).await?;
2102                ctx.shell("deploy", ShellConfig::new("echo deployed"))
2103                    .await?;
2104                Ok(())
2105            })
2106        }
2107    }
2108
2109    struct DoubleApprovalWorkflow;
2110
2111    impl WorkflowHandler for DoubleApprovalWorkflow {
2112        fn name(&self) -> &str {
2113            "double-approval"
2114        }
2115
2116        fn execute<'a>(&'a self, ctx: &'a mut WorkflowContext) -> HandlerFuture<'a> {
2117            Box::pin(async move {
2118                ctx.shell("build", ShellConfig::new("echo built")).await?;
2119                ctx.approval("staging-gate", ApprovalConfig::new("Deploy staging?"))
2120                    .await?;
2121                ctx.shell("deploy-staging", ShellConfig::new("echo staging"))
2122                    .await?;
2123                ctx.approval("prod-gate", ApprovalConfig::new("Deploy prod?"))
2124                    .await?;
2125                ctx.shell("deploy-prod", ShellConfig::new("echo prod"))
2126                    .await?;
2127                Ok(())
2128            })
2129        }
2130    }
2131
2132    #[tokio::test]
2133    async fn approval_pauses_run() {
2134        let mut engine = create_test_engine();
2135        engine.register(SingleApprovalWorkflow).unwrap();
2136
2137        let run = engine
2138            .run_handler("single-approval", TriggerKind::Manual, json!({}))
2139            .await
2140            .unwrap()
2141            .run;
2142
2143        assert_eq!(run.status.state, RunStatus::AwaitingApproval);
2144
2145        let steps = engine.store().list_steps(run.id).await.unwrap();
2146        assert_eq!(steps.len(), 2); // build + approval gate
2147        assert_eq!(steps[0].kind, StepKind::Shell);
2148        assert_eq!(steps[0].status.state, StepStatus::Completed);
2149        assert_eq!(steps[1].kind, StepKind::Approval);
2150        assert_eq!(steps[1].status.state, StepStatus::AwaitingApproval);
2151    }
2152
2153    #[tokio::test]
2154    async fn approval_resume_completes_run() {
2155        let mut engine = create_test_engine();
2156        engine.register(SingleApprovalWorkflow).unwrap();
2157
2158        // First execution: pauses at approval
2159        let run = engine
2160            .run_handler("single-approval", TriggerKind::Manual, json!({}))
2161            .await
2162            .unwrap()
2163            .run;
2164        assert_eq!(run.status.state, RunStatus::AwaitingApproval);
2165
2166        // Simulate approval: transition to Running
2167        engine
2168            .store()
2169            .update_run_status(run.id, RunStatus::Running)
2170            .await
2171            .unwrap();
2172
2173        // Resume: replays build, skips approval, executes deploy
2174        let resumed = engine.resume_run(run.id).await.unwrap().run;
2175        assert_eq!(resumed.status.state, RunStatus::Completed);
2176
2177        let steps = engine.store().list_steps(run.id).await.unwrap();
2178        assert_eq!(steps.len(), 3); // build + approval + deploy
2179        assert_eq!(steps[0].name, "build");
2180        assert_eq!(steps[0].status.state, StepStatus::Completed);
2181        assert_eq!(steps[1].name, "gate");
2182        assert_eq!(steps[1].kind, StepKind::Approval);
2183        assert_eq!(steps[1].status.state, StepStatus::Completed);
2184        assert_eq!(steps[2].name, "deploy");
2185        assert_eq!(steps[2].status.state, StepStatus::Completed);
2186    }
2187
2188    #[tokio::test]
2189    async fn double_approval_two_resumes() {
2190        let mut engine = create_test_engine();
2191        engine.register(DoubleApprovalWorkflow).unwrap();
2192
2193        // First execution: pauses at staging-gate
2194        let run = engine
2195            .run_handler("double-approval", TriggerKind::Manual, json!({}))
2196            .await
2197            .unwrap()
2198            .run;
2199        assert_eq!(run.status.state, RunStatus::AwaitingApproval);
2200
2201        let steps = engine.store().list_steps(run.id).await.unwrap();
2202        assert_eq!(steps.len(), 2); // build + staging-gate
2203
2204        // First approval
2205        engine
2206            .store()
2207            .update_run_status(run.id, RunStatus::Running)
2208            .await
2209            .unwrap();
2210
2211        let resumed = engine.resume_run(run.id).await.unwrap().run;
2212        assert_eq!(resumed.status.state, RunStatus::AwaitingApproval);
2213
2214        let steps = engine.store().list_steps(run.id).await.unwrap();
2215        assert_eq!(steps.len(), 4); // build + staging-gate + deploy-staging + prod-gate
2216
2217        // Second approval
2218        engine
2219            .store()
2220            .update_run_status(run.id, RunStatus::Running)
2221            .await
2222            .unwrap();
2223
2224        let final_run = engine.resume_run(run.id).await.unwrap().run;
2225        assert_eq!(final_run.status.state, RunStatus::Completed);
2226
2227        let steps = engine.store().list_steps(run.id).await.unwrap();
2228        assert_eq!(steps.len(), 5);
2229        assert_eq!(steps[0].name, "build");
2230        assert_eq!(steps[1].name, "staging-gate");
2231        assert_eq!(steps[2].name, "deploy-staging");
2232        assert_eq!(steps[3].name, "prod-gate");
2233        assert_eq!(steps[4].name, "deploy-prod");
2234
2235        for step in &steps {
2236            assert_eq!(step.status.state, StepStatus::Completed);
2237        }
2238    }
2239
2240    // -----------------------------------------------------------------------
2241    // fail_orphaned_steps tests
2242    // -----------------------------------------------------------------------
2243
2244    use ironflow_store::models::{NewStep, StepUpdate, step_trace_id};
2245
2246    async fn create_step_with_status(
2247        store: &Arc<dyn Store>,
2248        run_id: Uuid,
2249        name: &str,
2250        position: u32,
2251        status: StepStatus,
2252    ) -> ironflow_store::models::Step {
2253        let step = store
2254            .create_step(NewStep {
2255                run_id,
2256                trace_id: step_trace_id(run_id, name, position),
2257                name: name.to_string(),
2258                kind: StepKind::Shell,
2259                position,
2260                input: None,
2261                is_error_handler: false,
2262            })
2263            .await
2264            .unwrap();
2265
2266        match status {
2267            StepStatus::Pending => {}
2268            StepStatus::Running => {
2269                store
2270                    .update_step(
2271                        step.id,
2272                        StepUpdate {
2273                            status: Some(StepStatus::Running),
2274                            ..StepUpdate::default()
2275                        },
2276                    )
2277                    .await
2278                    .unwrap();
2279            }
2280            StepStatus::Completed => {
2281                store
2282                    .update_step(
2283                        step.id,
2284                        StepUpdate {
2285                            status: Some(StepStatus::Running),
2286                            ..StepUpdate::default()
2287                        },
2288                    )
2289                    .await
2290                    .unwrap();
2291                store
2292                    .update_step(
2293                        step.id,
2294                        StepUpdate {
2295                            status: Some(StepStatus::Completed),
2296                            ..StepUpdate::default()
2297                        },
2298                    )
2299                    .await
2300                    .unwrap();
2301            }
2302            StepStatus::AwaitingApproval => {
2303                store
2304                    .update_step(
2305                        step.id,
2306                        StepUpdate {
2307                            status: Some(StepStatus::Running),
2308                            ..StepUpdate::default()
2309                        },
2310                    )
2311                    .await
2312                    .unwrap();
2313                store
2314                    .update_step(
2315                        step.id,
2316                        StepUpdate {
2317                            status: Some(StepStatus::AwaitingApproval),
2318                            ..StepUpdate::default()
2319                        },
2320                    )
2321                    .await
2322                    .unwrap();
2323            }
2324            _ => panic!("unsupported status for test helper: {status}"),
2325        }
2326
2327        store.get_step(step.id).await.unwrap().unwrap()
2328    }
2329
2330    #[tokio::test]
2331    async fn fail_orphaned_steps_marks_running_as_failed() {
2332        let engine = create_test_engine();
2333        let run = engine
2334            .store()
2335            .create_run(NewRun {
2336                created_by: None,
2337                workflow_name: "test".to_string(),
2338                trigger: TriggerKind::Manual,
2339                payload: json!({}),
2340                max_retries: 0,
2341                handler_version: None,
2342                labels: HashMap::new(),
2343                scheduled_at: None,
2344                idempotency_key: None,
2345                max_cost_usd: None,
2346            })
2347            .await
2348            .unwrap()
2349            .into_run();
2350
2351        let step = create_step_with_status(
2352            engine.store(),
2353            run.id,
2354            "running-step",
2355            0,
2356            StepStatus::Running,
2357        )
2358        .await;
2359
2360        engine
2361            .fail_orphaned_steps(run.id, "parent run timed out")
2362            .await
2363            .unwrap();
2364
2365        let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
2366        assert_eq!(updated.status.state, StepStatus::Failed);
2367        assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
2368        assert!(updated.completed_at.is_some());
2369    }
2370
2371    #[tokio::test]
2372    async fn fail_orphaned_steps_marks_pending_as_skipped() {
2373        let engine = create_test_engine();
2374        let run = engine
2375            .store()
2376            .create_run(NewRun {
2377                created_by: None,
2378                workflow_name: "test".to_string(),
2379                trigger: TriggerKind::Manual,
2380                payload: json!({}),
2381                max_retries: 0,
2382                handler_version: None,
2383                labels: HashMap::new(),
2384                scheduled_at: None,
2385                idempotency_key: None,
2386                max_cost_usd: None,
2387            })
2388            .await
2389            .unwrap()
2390            .into_run();
2391
2392        let step = create_step_with_status(
2393            engine.store(),
2394            run.id,
2395            "pending-step",
2396            0,
2397            StepStatus::Pending,
2398        )
2399        .await;
2400
2401        engine
2402            .fail_orphaned_steps(run.id, "parent run timed out")
2403            .await
2404            .unwrap();
2405
2406        let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
2407        assert_eq!(updated.status.state, StepStatus::Skipped);
2408        assert!(updated.error.is_none());
2409        assert!(updated.completed_at.is_some());
2410    }
2411
2412    #[tokio::test]
2413    async fn fail_orphaned_steps_marks_awaiting_approval_as_failed() {
2414        let engine = create_test_engine();
2415        let run = engine
2416            .store()
2417            .create_run(NewRun {
2418                created_by: None,
2419                workflow_name: "test".to_string(),
2420                trigger: TriggerKind::Manual,
2421                payload: json!({}),
2422                max_retries: 0,
2423                handler_version: None,
2424                labels: HashMap::new(),
2425                scheduled_at: None,
2426                idempotency_key: None,
2427                max_cost_usd: None,
2428            })
2429            .await
2430            .unwrap()
2431            .into_run();
2432
2433        let step = create_step_with_status(
2434            engine.store(),
2435            run.id,
2436            "approval-step",
2437            0,
2438            StepStatus::AwaitingApproval,
2439        )
2440        .await;
2441
2442        engine
2443            .fail_orphaned_steps(run.id, "parent run timed out")
2444            .await
2445            .unwrap();
2446
2447        let updated = engine.store().get_step(step.id).await.unwrap().unwrap();
2448        assert_eq!(updated.status.state, StepStatus::Failed);
2449        assert_eq!(updated.error.as_deref(), Some("parent run timed out"));
2450        assert!(updated.completed_at.is_some());
2451    }
2452
2453    #[tokio::test]
2454    async fn fail_orphaned_steps_skips_terminal_steps() {
2455        let engine = create_test_engine();
2456        let run = engine
2457            .store()
2458            .create_run(NewRun {
2459                created_by: None,
2460                workflow_name: "test".to_string(),
2461                trigger: TriggerKind::Manual,
2462                payload: json!({}),
2463                max_retries: 0,
2464                handler_version: None,
2465                labels: HashMap::new(),
2466                scheduled_at: None,
2467                idempotency_key: None,
2468                max_cost_usd: None,
2469            })
2470            .await
2471            .unwrap()
2472            .into_run();
2473
2474        let completed_step =
2475            create_step_with_status(engine.store(), run.id, "done", 0, StepStatus::Completed).await;
2476        let running_step =
2477            create_step_with_status(engine.store(), run.id, "in-flight", 1, StepStatus::Running)
2478                .await;
2479
2480        engine
2481            .fail_orphaned_steps(run.id, "parent run timed out")
2482            .await
2483            .unwrap();
2484
2485        let completed = engine
2486            .store()
2487            .get_step(completed_step.id)
2488            .await
2489            .unwrap()
2490            .unwrap();
2491        assert_eq!(completed.status.state, StepStatus::Completed);
2492
2493        let failed = engine
2494            .store()
2495            .get_step(running_step.id)
2496            .await
2497            .unwrap()
2498            .unwrap();
2499        assert_eq!(failed.status.state, StepStatus::Failed);
2500    }
2501
2502    #[tokio::test]
2503    async fn fail_orphaned_steps_mixed_states() {
2504        let engine = create_test_engine();
2505        let run = engine
2506            .store()
2507            .create_run(NewRun {
2508                created_by: None,
2509                workflow_name: "test".to_string(),
2510                trigger: TriggerKind::Manual,
2511                payload: json!({}),
2512                max_retries: 0,
2513                handler_version: None,
2514                labels: HashMap::new(),
2515                scheduled_at: None,
2516                idempotency_key: None,
2517                max_cost_usd: None,
2518            })
2519            .await
2520            .unwrap()
2521            .into_run();
2522
2523        let s_completed =
2524            create_step_with_status(engine.store(), run.id, "step-1", 0, StepStatus::Completed)
2525                .await;
2526        let s_running =
2527            create_step_with_status(engine.store(), run.id, "step-2", 1, StepStatus::Running).await;
2528        let s_pending =
2529            create_step_with_status(engine.store(), run.id, "step-3", 2, StepStatus::Pending).await;
2530
2531        engine.fail_orphaned_steps(run.id, "timeout").await.unwrap();
2532
2533        let r_completed = engine
2534            .store()
2535            .get_step(s_completed.id)
2536            .await
2537            .unwrap()
2538            .unwrap();
2539        assert_eq!(r_completed.status.state, StepStatus::Completed);
2540
2541        let r_running = engine
2542            .store()
2543            .get_step(s_running.id)
2544            .await
2545            .unwrap()
2546            .unwrap();
2547        assert_eq!(r_running.status.state, StepStatus::Failed);
2548        assert_eq!(r_running.error.as_deref(), Some("timeout"));
2549
2550        let r_pending = engine
2551            .store()
2552            .get_step(s_pending.id)
2553            .await
2554            .unwrap()
2555            .unwrap();
2556        assert_eq!(r_pending.status.state, StepStatus::Skipped);
2557        assert!(r_pending.error.is_none());
2558    }
2559
2560    #[tokio::test]
2561    async fn fail_orphaned_steps_no_steps_is_noop() {
2562        let engine = create_test_engine();
2563        let run = engine
2564            .store()
2565            .create_run(NewRun {
2566                created_by: None,
2567                workflow_name: "test".to_string(),
2568                trigger: TriggerKind::Manual,
2569                payload: json!({}),
2570                max_retries: 0,
2571                handler_version: None,
2572                labels: HashMap::new(),
2573                scheduled_at: None,
2574                idempotency_key: None,
2575                max_cost_usd: None,
2576            })
2577            .await
2578            .unwrap()
2579            .into_run();
2580
2581        let result = engine.fail_orphaned_steps(run.id, "timeout").await;
2582        assert!(result.is_ok());
2583    }
2584
2585    #[tokio::test]
2586    async fn fail_orphaned_steps_preserves_existing_error() {
2587        let engine = create_test_engine();
2588        let run = engine
2589            .store()
2590            .create_run(NewRun {
2591                created_by: None,
2592                workflow_name: "test".to_string(),
2593                trigger: TriggerKind::Manual,
2594                payload: json!({}),
2595                max_retries: 0,
2596                handler_version: None,
2597                labels: HashMap::new(),
2598                scheduled_at: None,
2599                idempotency_key: None,
2600                max_cost_usd: None,
2601            })
2602            .await
2603            .unwrap()
2604            .into_run();
2605
2606        let step_with_error = create_step_with_status(
2607            engine.store(),
2608            run.id,
2609            "already-errored",
2610            0,
2611            StepStatus::Running,
2612        )
2613        .await;
2614
2615        engine
2616            .store()
2617            .update_step(
2618                step_with_error.id,
2619                StepUpdate {
2620                    error: Some("real error from provider".to_string()),
2621                    ..StepUpdate::default()
2622                },
2623            )
2624            .await
2625            .unwrap();
2626
2627        let step_no_error = create_step_with_status(
2628            engine.store(),
2629            run.id,
2630            "no-error-yet",
2631            1,
2632            StepStatus::Running,
2633        )
2634        .await;
2635
2636        engine
2637            .fail_orphaned_steps(run.id, "parent run failed")
2638            .await
2639            .unwrap();
2640
2641        let updated_with = engine
2642            .store()
2643            .get_step(step_with_error.id)
2644            .await
2645            .unwrap()
2646            .unwrap();
2647        assert_eq!(updated_with.status.state, StepStatus::Failed);
2648        assert_eq!(
2649            updated_with.error.as_deref(),
2650            Some("real error from provider"),
2651        );
2652
2653        let updated_without = engine
2654            .store()
2655            .get_step(step_no_error.id)
2656            .await
2657            .unwrap()
2658            .unwrap();
2659        assert_eq!(updated_without.status.state, StepStatus::Failed);
2660        assert_eq!(updated_without.error.as_deref(), Some("parent run failed"),);
2661    }
2662}