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