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