ironflow_engine/
context.rs

1//! [`WorkflowContext`] — execution context for dynamic workflows.
2//!
3//! Provides step execution methods that automatically persist results to the
4//! store. Each call to [`shell`](WorkflowContext::shell),
5//! [`http`](WorkflowContext::http), [`agent`](WorkflowContext::agent), or
6//! [`workflow`](WorkflowContext::workflow) creates a step record, executes the
7//! operation, captures the output, and returns a [`StepOutput`] that the next
8//! step can reference.
9//!
10//! # Examples
11//!
12//! ```no_run
13//! use ironflow_engine::context::WorkflowContext;
14//! use ironflow_engine::config::{ShellConfig, AgentStepConfig};
15//! use ironflow_engine::error::EngineError;
16//!
17//! # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
18//! let build = ctx.shell("build", ShellConfig::new("cargo build")).await?;
19//! let review = ctx.agent("review", AgentStepConfig::new(
20//!     &format!("Build output:\n{}", build.output["stdout"])
21//! )).await?;
22//! # Ok(())
23//! # }
24//! ```
25
26use std::collections::HashMap;
27use std::fmt;
28use std::sync::Arc;
29use std::time::Instant;
30
31use chrono::{DateTime, Utc};
32use rust_decimal::Decimal;
33use serde_json::Value;
34use tokio::task::{Id, JoinSet};
35use tracing::{error, info};
36use uuid::Uuid;
37
38use ironflow_core::error::{AgentError, OperationError};
39use ironflow_core::provider::AgentProvider;
40use ironflow_store::models::{
41    NewRun, NewStep, NewStepDependency, RunStatus, RunUpdate, Step, StepKind, StepStatus,
42    StepUpdate, TriggerKind,
43};
44use ironflow_store::store::Store;
45
46use crate::budget::step_budget_usd;
47use crate::config::{
48    AgentStepConfig, ApprovalConfig, HttpConfig, ShellConfig, StepConfig, WorkflowStepConfig,
49};
50use crate::error::EngineError;
51use crate::executor::{ParallelStepResult, StepOutput, execute_step_config};
52use crate::handler::WorkflowHandler;
53use crate::log_sender::{LogSender, StepLogSender};
54use crate::operation::Operation;
55
56/// Callback type for resolving workflow handlers by name.
57pub(crate) type HandlerResolver =
58    Arc<dyn Fn(&str) -> Option<Arc<dyn WorkflowHandler>> + Send + Sync>;
59
60/// Execution context for a single workflow run.
61///
62/// Tracks the current step position and provides convenience methods
63/// for executing operations with automatic persistence.
64///
65/// # Examples
66///
67/// ```no_run
68/// use ironflow_engine::context::WorkflowContext;
69/// use ironflow_engine::config::ShellConfig;
70/// use ironflow_engine::error::EngineError;
71///
72/// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
73/// let result = ctx.shell("greet", ShellConfig::new("echo hello")).await?;
74/// assert!(result.output["stdout"].as_str().unwrap().contains("hello"));
75/// # Ok(())
76/// # }
77/// ```
78pub struct WorkflowContext {
79    run_id: Uuid,
80    store: Arc<dyn Store>,
81    provider: Arc<dyn AgentProvider>,
82    handler_resolver: Option<HandlerResolver>,
83    position: u32,
84    /// IDs of the last executed step(s) -- used to record DAG dependencies.
85    last_step_ids: Vec<Uuid>,
86    /// Accumulated cost across all steps in this run.
87    total_cost_usd: Decimal,
88    /// Accumulated duration across all steps.
89    total_duration_ms: u64,
90    /// Cumulative cost cap for this run, resolved at creation. `None` = no cap.
91    max_cost_usd: Option<Decimal>,
92    /// Cost already spent by ancestor runs when this context belongs to a
93    /// sub-workflow. Zero for a top-level run.
94    inherited_cost_usd: Decimal,
95    /// Steps from a previous execution, keyed by position.
96    /// Used when resuming after approval to replay completed steps.
97    replay_steps: HashMap<u32, Step>,
98    /// Optional sender for real-time log streaming.
99    log_sender: Option<LogSender>,
100}
101
102impl WorkflowContext {
103    /// Create a new context for a run.
104    ///
105    /// Not typically called directly — the [`Engine`](crate::engine::Engine)
106    /// creates this when executing a [`WorkflowHandler`].
107    pub fn new(run_id: Uuid, store: Arc<dyn Store>, provider: Arc<dyn AgentProvider>) -> Self {
108        Self {
109            run_id,
110            store,
111            provider,
112            handler_resolver: None,
113            position: 0,
114            last_step_ids: Vec::new(),
115            total_cost_usd: Decimal::ZERO,
116            total_duration_ms: 0,
117            max_cost_usd: None,
118            inherited_cost_usd: Decimal::ZERO,
119            replay_steps: HashMap::new(),
120            log_sender: None,
121        }
122    }
123
124    /// Create a new context with a handler resolver for sub-workflow support.
125    ///
126    /// The resolver is called when [`workflow`](Self::workflow) is invoked to
127    /// look up registered handlers by name.
128    pub(crate) fn with_handler_resolver(
129        run_id: Uuid,
130        store: Arc<dyn Store>,
131        provider: Arc<dyn AgentProvider>,
132        resolver: HandlerResolver,
133    ) -> Self {
134        Self {
135            run_id,
136            store,
137            provider,
138            handler_resolver: Some(resolver),
139            position: 0,
140            last_step_ids: Vec::new(),
141            total_cost_usd: Decimal::ZERO,
142            total_duration_ms: 0,
143            max_cost_usd: None,
144            inherited_cost_usd: Decimal::ZERO,
145            replay_steps: HashMap::new(),
146            log_sender: None,
147        }
148    }
149
150    /// Attach a log sender for real-time step output streaming.
151    pub fn set_log_sender(&mut self, sender: LogSender) {
152        self.log_sender = Some(sender);
153    }
154
155    /// Set the cumulative cost cap enforced before every agent step.
156    ///
157    /// Called by the [`Engine`](crate::engine::Engine) with the run's persisted
158    /// `max_cost_usd`. `None` disables the check.
159    ///
160    /// # Examples
161    ///
162    /// ```no_run
163    /// use ironflow_engine::context::WorkflowContext;
164    /// use rust_decimal::Decimal;
165    ///
166    /// # fn example(ctx: &mut WorkflowContext) {
167    /// ctx.set_max_cost_usd(Some(Decimal::new(200, 2))); // $2.00
168    /// # }
169    /// ```
170    pub fn set_max_cost_usd(&mut self, cap: Option<Decimal>) {
171        self.max_cost_usd = cap;
172    }
173
174    /// The cumulative cost cap of this run, if any.
175    pub fn max_cost_usd(&self) -> Option<Decimal> {
176        self.max_cost_usd
177    }
178
179    /// Total cost charged against the cap: this run plus every ancestor run.
180    ///
181    /// For a top-level run this equals [`total_cost_usd`](Self::total_cost_usd).
182    /// For a sub-workflow it also includes what the parent chain already spent.
183    pub fn charged_cost_usd(&self) -> Decimal {
184        self.inherited_cost_usd + self.total_cost_usd
185    }
186
187    /// Reject the upcoming agent work when it would cross the run's cost cap.
188    ///
189    /// `step_budget` is the declared budget of the step (or the sum of budgets
190    /// for a parallel wave). Called *before* any step record is created so a
191    /// refused run never launches the work it could not afford.
192    ///
193    /// # Errors
194    ///
195    /// Returns [`EngineError::RunBudgetExceeded`] when
196    /// `charged_cost + step_budget` exceeds the cap.
197    fn check_run_budget(&self, step_budget: Decimal) -> Result<(), EngineError> {
198        let Some(limit) = self.max_cost_usd else {
199            return Ok(());
200        };
201
202        let spent = self.charged_cost_usd();
203        if spent + step_budget <= limit {
204            return Ok(());
205        }
206
207        error!(
208            run_id = %self.run_id,
209            limit_usd = %limit,
210            spent_usd = %spent,
211            step_budget_usd = %step_budget,
212            "run cost cap reached, refusing agent step"
213        );
214
215        Err(EngineError::RunBudgetExceeded {
216            run_id: self.run_id,
217            limit_usd: limit,
218            spent_usd: spent,
219            step_budget_usd: step_budget,
220        })
221    }
222
223    /// Load existing steps from the store for replay after approval.
224    ///
225    /// Called by the engine when resuming a run. All completed steps
226    /// and the approved approval step are indexed by position so that
227    /// `execute_step` and `approval` can skip them.
228    pub(crate) async fn load_replay_steps(&mut self) -> Result<(), EngineError> {
229        let steps = self.store.list_steps(self.run_id).await?;
230        for step in steps {
231            let dominated = matches!(
232                step.status.state,
233                StepStatus::Completed | StepStatus::Running | StepStatus::AwaitingApproval
234            );
235            if dominated {
236                self.replay_steps.insert(step.position, step);
237            }
238        }
239        Ok(())
240    }
241
242    /// The run ID this context is executing for.
243    pub fn run_id(&self) -> Uuid {
244        self.run_id
245    }
246
247    /// Accumulated cost across all executed steps so far.
248    pub fn total_cost_usd(&self) -> Decimal {
249        self.total_cost_usd
250    }
251
252    /// Accumulated duration across all executed steps so far.
253    pub fn total_duration_ms(&self) -> u64 {
254        self.total_duration_ms
255    }
256
257    /// Execute multiple steps concurrently (wait-all model).
258    ///
259    /// All steps in the batch execute in parallel via `tokio::JoinSet`.
260    /// Each step is recorded with the same `position` (execution wave).
261    /// Dependencies on previous steps are recorded automatically.
262    ///
263    /// When `fail_fast` is true, remaining steps are aborted on the first
264    /// failure. When false, all steps run to completion and the first
265    /// error is returned.
266    ///
267    /// # Errors
268    ///
269    /// Returns [`EngineError`] if any step fails.
270    ///
271    /// # Examples
272    ///
273    /// ```no_run
274    /// use ironflow_engine::context::WorkflowContext;
275    /// use ironflow_engine::config::{StepConfig, ShellConfig};
276    /// use ironflow_engine::error::EngineError;
277    ///
278    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
279    /// let results = ctx.parallel(
280    ///     vec![
281    ///         ("test-unit", StepConfig::Shell(ShellConfig::new("cargo test --lib"))),
282    ///         ("lint", StepConfig::Shell(ShellConfig::new("cargo clippy"))),
283    ///     ],
284    ///     true,
285    /// ).await?;
286    ///
287    /// for r in &results {
288    ///     println!("{}: {:?}", r.name, r.output.output);
289    /// }
290    /// # Ok(())
291    /// # }
292    /// ```
293    pub async fn parallel(
294        &mut self,
295        steps: Vec<(&str, StepConfig)>,
296        fail_fast: bool,
297    ) -> Result<Vec<ParallelStepResult>, EngineError> {
298        if steps.is_empty() {
299            return Ok(Vec::new());
300        }
301
302        // Cost cap: the whole wave is charged at once. Refused before any step
303        // record is created, so nothing in the wave starts.
304        let wave_budget: Decimal = steps
305            .iter()
306            .filter_map(|(_, config)| match config {
307                StepConfig::Agent(agent_config) => Some(agent_config.max_budget_usd),
308                _ => None,
309            })
310            .map(step_budget_usd)
311            .sum();
312        self.check_run_budget(wave_budget)?;
313
314        let wave_position = self.position;
315        self.position += 1;
316
317        let now = Utc::now();
318        let mut step_records: Vec<(Uuid, String, StepConfig)> = Vec::with_capacity(steps.len());
319
320        for (name, config) in &steps {
321            let kind = config.kind();
322            let step = self
323                .store
324                .create_step(NewStep {
325                    run_id: self.run_id,
326                    name: name.to_string(),
327                    kind,
328                    position: wave_position,
329                    input: Some(serde_json::to_value(config)?),
330                })
331                .await?;
332
333            self.start_step(step.id, now).await?;
334
335            step_records.push((step.id, name.to_string(), config.clone()));
336        }
337
338        let mut join_set = JoinSet::new();
339        let mut task_index: HashMap<Id, usize> = HashMap::new();
340        for (idx, (step_id, step_name, config)) in step_records.iter().enumerate() {
341            let provider = self.provider.clone();
342            let config = config.clone();
343            let step_log_sender = self
344                .log_sender
345                .as_ref()
346                .map(|s| StepLogSender::new(s.clone(), self.run_id, *step_id, step_name.clone()));
347            let handle = join_set.spawn(async move {
348                (
349                    idx,
350                    execute_step_config(&config, &provider, step_log_sender).await,
351                )
352            });
353            task_index.insert(handle.id(), idx);
354        }
355
356        // JoinSet returns in completion order; indexed_results restores input order.
357        let mut indexed_results: Vec<Option<Result<StepOutput, String>>> =
358            vec![None; step_records.len()];
359        let mut first_error: Option<EngineError> = None;
360
361        while let Some(join_result) = join_set.join_next().await {
362            let (idx, step_result) = match join_result {
363                Ok(r) => r,
364                Err(e) => {
365                    let error_msg = format!("join error: {e}");
366                    if let Some(&idx) = task_index.get(&e.id()) {
367                        let (step_id, step_name, _) = &step_records[idx];
368                        let completed_at = Utc::now();
369                        error!(
370                            run_id = %self.run_id,
371                            step = %step_name,
372                            error = %error_msg,
373                            "parallel step panicked or was cancelled"
374                        );
375                        if let Err(store_err) = self
376                            .store
377                            .update_step(
378                                *step_id,
379                                StepUpdate {
380                                    status: Some(StepStatus::Failed),
381                                    error: Some(error_msg.clone()),
382                                    completed_at: Some(completed_at),
383                                    ..StepUpdate::default()
384                                },
385                            )
386                            .await
387                        {
388                            error!(
389                                run_id = %self.run_id,
390                                step_id = %step_id,
391                                error = %store_err,
392                                "failed to persist JoinError for step"
393                            );
394                        }
395                        indexed_results[idx] = Some(Err(error_msg.clone()));
396                    }
397                    if first_error.is_none() {
398                        first_error = Some(EngineError::StepConfig(error_msg));
399                    }
400                    if fail_fast {
401                        join_set.abort_all();
402                    }
403                    continue;
404                }
405            };
406
407            let (step_id, step_name, _) = &step_records[idx];
408            let completed_at = Utc::now();
409
410            match step_result {
411                Ok(output) => {
412                    self.total_cost_usd += output.cost_usd;
413                    self.total_duration_ms += output.duration_ms;
414
415                    let debug_messages_json = output.debug_messages_json();
416
417                    self.store
418                        .update_step(
419                            *step_id,
420                            StepUpdate {
421                                status: Some(StepStatus::Completed),
422                                output: Some(output.output.clone()),
423                                duration_ms: Some(output.duration_ms),
424                                cost_usd: Some(output.cost_usd),
425                                input_tokens: output.input_tokens,
426                                output_tokens: output.output_tokens,
427                                completed_at: Some(completed_at),
428                                debug_messages: debug_messages_json,
429                                ..StepUpdate::default()
430                            },
431                        )
432                        .await?;
433
434                    info!(
435                        run_id = %self.run_id,
436                        step = %step_name,
437                        duration_ms = output.duration_ms,
438                        "parallel step completed"
439                    );
440
441                    indexed_results[idx] = Some(Ok(output));
442                }
443                Err(err) => {
444                    let err_msg = err.to_string();
445                    let debug_messages_json = extract_debug_messages_from_error(&err);
446                    let partial = extract_partial_usage_from_error(&err);
447                    let raw_response_output = extract_raw_response_from_error(&err);
448
449                    if let Some(ref usage) = partial {
450                        if let Some(cost) = usage.cost_usd {
451                            self.total_cost_usd += cost;
452                        }
453                        if let Some(dur) = usage.duration_ms {
454                            self.total_duration_ms += dur;
455                        }
456                    }
457
458                    if let Err(store_err) = self
459                        .store
460                        .update_step(
461                            *step_id,
462                            StepUpdate {
463                                status: Some(StepStatus::Failed),
464                                error: Some(err_msg.clone()),
465                                output: raw_response_output,
466                                completed_at: Some(completed_at),
467                                debug_messages: debug_messages_json,
468                                duration_ms: partial.as_ref().and_then(|p| p.duration_ms),
469                                cost_usd: partial.as_ref().and_then(|p| p.cost_usd),
470                                input_tokens: partial.as_ref().and_then(|p| p.input_tokens),
471                                output_tokens: partial.as_ref().and_then(|p| p.output_tokens),
472                                ..StepUpdate::default()
473                            },
474                        )
475                        .await
476                    {
477                        tracing::error!(
478                            step_id = %step_id,
479                            error = %store_err,
480                            "failed to persist parallel step failure"
481                        );
482                    }
483
484                    indexed_results[idx] = Some(Err(err_msg.clone()));
485
486                    if first_error.is_none() {
487                        first_error = Some(err);
488                    }
489
490                    if fail_fast {
491                        join_set.abort_all();
492                    }
493                }
494            }
495        }
496
497        if let Some(err) = first_error {
498            return Err(err);
499        }
500
501        self.last_step_ids = step_records.iter().map(|(id, _, _)| *id).collect();
502
503        // Build results in original order.
504        let results: Vec<ParallelStepResult> = step_records
505            .iter()
506            .enumerate()
507            .map(|(idx, (step_id, name, _))| {
508                let output = match indexed_results[idx].take() {
509                    Some(Ok(o)) => o,
510                    _ => unreachable!("all steps succeeded if no error returned"),
511                };
512                ParallelStepResult {
513                    name: name.clone(),
514                    output,
515                    step_id: *step_id,
516                }
517            })
518            .collect();
519
520        Ok(results)
521    }
522
523    /// Execute a shell step.
524    ///
525    /// Creates the step record, runs the command, persists the result,
526    /// and returns the output for use in subsequent steps.
527    ///
528    /// # Errors
529    ///
530    /// Returns [`EngineError`] if the command fails or the store errors.
531    ///
532    /// # Examples
533    ///
534    /// ```no_run
535    /// use ironflow_engine::context::WorkflowContext;
536    /// use ironflow_engine::config::ShellConfig;
537    /// use ironflow_engine::error::EngineError;
538    ///
539    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
540    /// let files = ctx.shell("list", ShellConfig::new("ls -la")).await?;
541    /// println!("stdout: {}", files.output["stdout"]);
542    /// # Ok(())
543    /// # }
544    /// ```
545    pub async fn shell(
546        &mut self,
547        name: &str,
548        config: ShellConfig,
549    ) -> Result<StepOutput, EngineError> {
550        self.execute_step(name, StepKind::Shell, StepConfig::Shell(config))
551            .await
552    }
553
554    /// Execute an HTTP step.
555    ///
556    /// # Errors
557    ///
558    /// Returns [`EngineError`] if the request fails or the store errors.
559    ///
560    /// # Examples
561    ///
562    /// ```no_run
563    /// use ironflow_engine::context::WorkflowContext;
564    /// use ironflow_engine::config::HttpConfig;
565    /// use ironflow_engine::error::EngineError;
566    ///
567    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
568    /// let resp = ctx.http("health", HttpConfig::get("https://api.example.com/health")).await?;
569    /// println!("status: {}", resp.output["status"]);
570    /// # Ok(())
571    /// # }
572    /// ```
573    pub async fn http(
574        &mut self,
575        name: &str,
576        config: HttpConfig,
577    ) -> Result<StepOutput, EngineError> {
578        self.execute_step(name, StepKind::Http, StepConfig::Http(config))
579            .await
580    }
581
582    /// Execute an agent step.
583    ///
584    /// # Errors
585    ///
586    /// Returns [`EngineError`] if the agent invocation fails or the store errors.
587    ///
588    /// # Examples
589    ///
590    /// ```no_run
591    /// use ironflow_engine::context::WorkflowContext;
592    /// use ironflow_engine::config::AgentStepConfig;
593    /// use ironflow_engine::error::EngineError;
594    ///
595    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
596    /// let review = ctx.agent("review", AgentStepConfig::new("Review the code")).await?;
597    /// println!("review: {}", review.output);
598    /// # Ok(())
599    /// # }
600    /// ```
601    pub async fn agent(
602        &mut self,
603        name: &str,
604        config: impl Into<AgentStepConfig>,
605    ) -> Result<StepOutput, EngineError> {
606        self.execute_step(name, StepKind::Agent, StepConfig::Agent(config.into()))
607            .await
608    }
609
610    /// Create a human approval gate.
611    ///
612    /// On first execution, records an approval step and returns
613    /// [`EngineError::ApprovalRequired`] to suspend the run. The engine
614    /// transitions the run to `AwaitingApproval`.
615    ///
616    /// On resume (after a human approved via the API), the approval step
617    /// is replayed: it is marked as `Completed` and execution continues
618    /// past it. Multiple approval gates in the same handler work -- each
619    /// one pauses and resumes independently.
620    ///
621    /// # Errors
622    ///
623    /// Returns [`EngineError::ApprovalRequired`] to pause the run on
624    /// first execution. Returns other [`EngineError`] variants on store
625    /// failures.
626    ///
627    /// # Examples
628    ///
629    /// ```no_run
630    /// use ironflow_engine::context::WorkflowContext;
631    /// use ironflow_engine::config::ApprovalConfig;
632    /// use ironflow_engine::error::EngineError;
633    ///
634    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
635    /// ctx.approval("deploy-gate", ApprovalConfig::new("Approve deployment?")).await?;
636    /// // Execution continues here after approval
637    /// # Ok(())
638    /// # }
639    /// ```
640    pub async fn approval(
641        &mut self,
642        name: &str,
643        config: ApprovalConfig,
644    ) -> Result<(), EngineError> {
645        let position = self.position;
646        self.position += 1;
647
648        // Replay: if this approval step exists from a prior execution,
649        // the run was approved -- mark it completed (if not already) and continue.
650        if let Some(existing) = self.replay_steps.get(&position)
651            && existing.kind == StepKind::Approval
652        {
653            if existing.status.state == StepStatus::AwaitingApproval {
654                self.store
655                    .update_step(
656                        existing.id,
657                        StepUpdate {
658                            status: Some(StepStatus::Completed),
659                            completed_at: Some(Utc::now()),
660                            ..StepUpdate::default()
661                        },
662                    )
663                    .await?;
664            }
665
666            self.last_step_ids = vec![existing.id];
667            info!(
668                run_id = %self.run_id,
669                step = %name,
670                position,
671                "approval step replayed (approved)"
672            );
673            return Ok(());
674        }
675
676        // First execution: create the approval step and suspend.
677        let step = self
678            .store
679            .create_step(NewStep {
680                run_id: self.run_id,
681                name: name.to_string(),
682                kind: StepKind::Approval,
683                position,
684                input: Some(serde_json::to_value(&config)?),
685            })
686            .await?;
687
688        self.start_step(step.id, Utc::now()).await?;
689
690        // Transition the step to AwaitingApproval so it reflects
691        // the suspended state on the dashboard.
692        self.store
693            .update_step(
694                step.id,
695                StepUpdate {
696                    status: Some(StepStatus::AwaitingApproval),
697                    ..StepUpdate::default()
698                },
699            )
700            .await?;
701
702        self.last_step_ids = vec![step.id];
703
704        Err(EngineError::ApprovalRequired {
705            run_id: self.run_id,
706            step_id: step.id,
707            message: config.message().to_string(),
708        })
709    }
710
711    /// Record a step as explicitly skipped.
712    ///
713    /// Use this inside an `if`/`else` branch when a step should not execute
714    /// but must still appear in the DAG and timeline with its reason.
715    ///
716    /// The step is created directly in [`StepStatus::Skipped`] state and the
717    /// reason is stored in the output as `{"reason": "..."}`.
718    ///
719    /// # Errors
720    ///
721    /// Returns [`EngineError`] if the store fails.
722    ///
723    /// # Examples
724    ///
725    /// ```no_run
726    /// use ironflow_engine::context::WorkflowContext;
727    /// use ironflow_engine::error::EngineError;
728    ///
729    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
730    /// let tests_passed = false;
731    /// if tests_passed {
732    ///     // ctx.shell("deploy", ...).await?;
733    /// } else {
734    ///     ctx.skip("deploy", "tests failed").await?;
735    /// }
736    /// # Ok(())
737    /// # }
738    /// ```
739    pub async fn skip(&mut self, name: &str, reason: &str) -> Result<(), EngineError> {
740        let position = self.position;
741        self.position += 1;
742
743        let step = self
744            .store
745            .create_step(NewStep {
746                run_id: self.run_id,
747                name: name.to_string(),
748                kind: StepKind::Custom("skip".to_string()),
749                position,
750                input: None,
751            })
752            .await?;
753
754        if !self.last_step_ids.is_empty() {
755            let deps: Vec<NewStepDependency> = self
756                .last_step_ids
757                .iter()
758                .map(|&depends_on| NewStepDependency {
759                    step_id: step.id,
760                    depends_on,
761                })
762                .collect();
763            self.store.create_step_dependencies(deps).await?;
764        }
765
766        let now = Utc::now();
767        self.store
768            .update_step(
769                step.id,
770                StepUpdate {
771                    status: Some(StepStatus::Skipped),
772                    output: Some(serde_json::json!({"reason": reason})),
773                    completed_at: Some(now),
774                    ..StepUpdate::default()
775                },
776            )
777            .await?;
778
779        self.last_step_ids = vec![step.id];
780
781        info!(
782            run_id = %self.run_id,
783            step = %name,
784            reason,
785            "step skipped"
786        );
787
788        Ok(())
789    }
790
791    /// Execute a custom operation step.
792    ///
793    /// Runs a user-defined [`Operation`] with full step lifecycle management:
794    /// creates the step record, transitions to Running, executes the operation,
795    /// persists the output and duration, and marks the step Completed or Failed.
796    ///
797    /// The operation's [`kind()`](Operation::kind) is stored as
798    /// [`StepKind::Custom`].
799    ///
800    /// # Errors
801    ///
802    /// Returns [`EngineError`] if the operation fails or the store errors.
803    ///
804    /// # Examples
805    ///
806    /// ```no_run
807    /// use ironflow_engine::context::WorkflowContext;
808    /// use ironflow_engine::operation::Operation;
809    /// use ironflow_engine::error::EngineError;
810    /// use serde_json::{Value, json};
811    /// use std::pin::Pin;
812    /// use std::future::Future;
813    ///
814    /// struct MyOp;
815    /// impl Operation for MyOp {
816    ///     fn kind(&self) -> &str { "my-service" }
817    ///     fn execute(&self) -> Pin<Box<dyn Future<Output = Result<Value, EngineError>> + Send + '_>> {
818    ///         Box::pin(async { Ok(json!({"ok": true})) })
819    ///     }
820    /// }
821    ///
822    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
823    /// let result = ctx.operation("call-service", &MyOp).await?;
824    /// println!("output: {}", result.output);
825    /// # Ok(())
826    /// # }
827    /// ```
828    pub async fn operation(
829        &mut self,
830        name: &str,
831        op: &dyn Operation,
832    ) -> Result<StepOutput, EngineError> {
833        let kind = StepKind::Custom(op.kind().to_string());
834        let position = self.position;
835        self.position += 1;
836
837        let step = self
838            .store
839            .create_step(NewStep {
840                run_id: self.run_id,
841                name: name.to_string(),
842                kind,
843                position,
844                input: op.input(),
845            })
846            .await?;
847
848        self.start_step(step.id, Utc::now()).await?;
849
850        let start = Instant::now();
851
852        match op.execute().await {
853            Ok(output_value) => {
854                let duration_ms = start.elapsed().as_millis() as u64;
855                self.total_duration_ms += duration_ms;
856
857                let completed_at = Utc::now();
858                self.store
859                    .update_step(
860                        step.id,
861                        StepUpdate {
862                            status: Some(StepStatus::Completed),
863                            output: Some(output_value.clone()),
864                            duration_ms: Some(duration_ms),
865                            cost_usd: Some(Decimal::ZERO),
866                            completed_at: Some(completed_at),
867                            ..StepUpdate::default()
868                        },
869                    )
870                    .await?;
871
872                info!(
873                    run_id = %self.run_id,
874                    step = %name,
875                    kind = op.kind(),
876                    duration_ms,
877                    "operation step completed"
878                );
879
880                self.last_step_ids = vec![step.id];
881
882                Ok(StepOutput {
883                    output: output_value,
884                    duration_ms,
885                    cost_usd: Decimal::ZERO,
886                    input_tokens: None,
887                    output_tokens: None,
888                    model: None,
889                    debug_messages: None,
890                })
891            }
892            Err(err) => {
893                let completed_at = Utc::now();
894                if let Err(store_err) = self
895                    .store
896                    .update_step(
897                        step.id,
898                        StepUpdate {
899                            status: Some(StepStatus::Failed),
900                            error: Some(err.to_string()),
901                            completed_at: Some(completed_at),
902                            ..StepUpdate::default()
903                        },
904                    )
905                    .await
906                {
907                    error!(step_id = %step.id, error = %store_err, "failed to persist step failure");
908                }
909
910                Err(err)
911            }
912        }
913    }
914
915    /// Execute a sub-workflow step.
916    ///
917    /// Creates a child run for the named workflow handler, executes it with
918    /// its own steps and lifecycle, and returns a [`StepOutput`] containing
919    /// the child run ID and aggregated metrics.
920    ///
921    /// Requires the context to be created with
922    /// `with_handler_resolver`.
923    ///
924    /// # Errors
925    ///
926    /// Returns [`EngineError::InvalidWorkflow`] if no handler is registered
927    /// with the given name, or if no handler resolver is available.
928    ///
929    /// # Examples
930    ///
931    /// ```no_run
932    /// use ironflow_engine::context::WorkflowContext;
933    /// use ironflow_engine::error::EngineError;
934    /// use serde_json::json;
935    ///
936    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
937    /// // let result = ctx.workflow(&MySubWorkflow, json!({})).await?;
938    /// # Ok(())
939    /// # }
940    /// ```
941    pub async fn workflow(
942        &mut self,
943        handler: &dyn WorkflowHandler,
944        payload: Value,
945    ) -> Result<StepOutput, EngineError> {
946        let config = WorkflowStepConfig::new(handler.name(), payload);
947        let position = self.position;
948        self.position += 1;
949
950        let step = self
951            .store
952            .create_step(NewStep {
953                run_id: self.run_id,
954                name: config.workflow_name.clone(),
955                kind: StepKind::Workflow,
956                position,
957                input: Some(serde_json::to_value(&config)?),
958            })
959            .await?;
960
961        self.start_step(step.id, Utc::now()).await?;
962
963        match self.execute_child_workflow(&config).await {
964            Ok(output) => {
965                self.total_cost_usd += output.cost_usd;
966                self.total_duration_ms += output.duration_ms;
967
968                let completed_at = Utc::now();
969                self.store
970                    .update_step(
971                        step.id,
972                        StepUpdate {
973                            status: Some(StepStatus::Completed),
974                            output: Some(output.output.clone()),
975                            duration_ms: Some(output.duration_ms),
976                            cost_usd: Some(output.cost_usd),
977                            completed_at: Some(completed_at),
978                            ..StepUpdate::default()
979                        },
980                    )
981                    .await?;
982
983                info!(
984                    run_id = %self.run_id,
985                    child_workflow = %config.workflow_name,
986                    duration_ms = output.duration_ms,
987                    "workflow step completed"
988                );
989
990                self.last_step_ids = vec![step.id];
991
992                Ok(output)
993            }
994            Err(err) => {
995                let completed_at = Utc::now();
996                if let Err(store_err) = self
997                    .store
998                    .update_step(
999                        step.id,
1000                        StepUpdate {
1001                            status: Some(StepStatus::Failed),
1002                            error: Some(err.to_string()),
1003                            completed_at: Some(completed_at),
1004                            ..StepUpdate::default()
1005                        },
1006                    )
1007                    .await
1008                {
1009                    error!(step_id = %step.id, error = %store_err, "failed to persist step failure");
1010                }
1011
1012                Err(err)
1013            }
1014        }
1015    }
1016
1017    /// Execute a child workflow and return aggregated output.
1018    async fn execute_child_workflow(
1019        &self,
1020        config: &WorkflowStepConfig,
1021    ) -> Result<StepOutput, EngineError> {
1022        let resolver = self.handler_resolver.as_ref().ok_or_else(|| {
1023            EngineError::InvalidWorkflow(
1024                "sub-workflow requires a handler resolver (use Engine to execute)".to_string(),
1025            )
1026        })?;
1027
1028        let handler = resolver(&config.workflow_name).ok_or_else(|| {
1029            EngineError::InvalidWorkflow(format!("no handler registered: {}", config.workflow_name))
1030        })?;
1031
1032        // A child run inherits both the parent labels and the parent author:
1033        // whoever triggered the parent workflow is accountable for its children.
1034        let parent = self.store.get_run(self.run_id).await?;
1035        let (parent_labels, parent_author) =
1036            parent.map(|r| (r.labels, r.created_by)).unwrap_or_default();
1037
1038        let child_run = self
1039            .store
1040            .create_run(NewRun {
1041                workflow_name: config.workflow_name.clone(),
1042                trigger: TriggerKind::Workflow,
1043                payload: config.payload.clone(),
1044                max_retries: 0,
1045                handler_version: None,
1046                labels: parent_labels,
1047                scheduled_at: None,
1048                created_by: parent_author,
1049                idempotency_key: None,
1050                // The child shares the parent's cap; it does not get its own budget.
1051                max_cost_usd: self.max_cost_usd,
1052            })
1053            .await?
1054            .into_run();
1055
1056        let child_run_id = child_run.id;
1057        info!(
1058            parent_run_id = %self.run_id,
1059            child_run_id = %child_run_id,
1060            workflow = %config.workflow_name,
1061            "child run created"
1062        );
1063
1064        self.store
1065            .update_run_status(child_run_id, RunStatus::Running)
1066            .await?;
1067
1068        let run_start = Instant::now();
1069        let mut child_ctx = WorkflowContext {
1070            run_id: child_run_id,
1071            store: self.store.clone(),
1072            provider: self.provider.clone(),
1073            handler_resolver: self.handler_resolver.clone(),
1074            position: 0,
1075            last_step_ids: Vec::new(),
1076            total_cost_usd: Decimal::ZERO,
1077            total_duration_ms: 0,
1078            max_cost_usd: self.max_cost_usd,
1079            // Everything the parent chain already spent counts against the
1080            // shared cap, so the child cannot restart the budget from zero.
1081            inherited_cost_usd: self.charged_cost_usd(),
1082            replay_steps: HashMap::new(),
1083            log_sender: self.log_sender.clone(),
1084        };
1085
1086        let result = handler.execute(&mut child_ctx).await;
1087        let total_duration = run_start.elapsed().as_millis() as u64;
1088        let completed_at = Utc::now();
1089
1090        match result {
1091            Ok(()) => {
1092                self.store
1093                    .update_run(
1094                        child_run_id,
1095                        RunUpdate {
1096                            status: Some(RunStatus::Completed),
1097                            cost_usd: Some(child_ctx.total_cost_usd),
1098                            duration_ms: Some(total_duration),
1099                            completed_at: Some(completed_at),
1100                            ..RunUpdate::default()
1101                        },
1102                    )
1103                    .await?;
1104
1105                Ok(StepOutput {
1106                    output: serde_json::json!({
1107                        "run_id": child_run_id,
1108                        "workflow_name": config.workflow_name,
1109                        "status": RunStatus::Completed,
1110                        "cost_usd": child_ctx.total_cost_usd,
1111                        "duration_ms": total_duration,
1112                    }),
1113                    duration_ms: total_duration,
1114                    cost_usd: child_ctx.total_cost_usd,
1115                    input_tokens: None,
1116                    output_tokens: None,
1117                    model: None,
1118                    debug_messages: None,
1119                })
1120            }
1121            Err(err) => {
1122                if let Err(store_err) = self
1123                    .store
1124                    .update_run(
1125                        child_run_id,
1126                        RunUpdate {
1127                            status: Some(RunStatus::Failed),
1128                            error: Some(err.to_string()),
1129                            cost_usd: Some(child_ctx.total_cost_usd),
1130                            duration_ms: Some(total_duration),
1131                            completed_at: Some(completed_at),
1132                            ..RunUpdate::default()
1133                        },
1134                    )
1135                    .await
1136                {
1137                    error!(
1138                        child_run_id = %child_run_id,
1139                        store_error = %store_err,
1140                        "failed to persist child run failure"
1141                    );
1142                }
1143
1144                Err(err)
1145            }
1146        }
1147    }
1148
1149    /// Try to replay a completed step from a previous execution.
1150    ///
1151    /// Returns `Some(StepOutput)` if a completed step exists at the given
1152    /// position, `None` otherwise.
1153    fn try_replay_step(&mut self, position: u32) -> Option<StepOutput> {
1154        let step = self.replay_steps.get(&position)?;
1155        if step.status.state != StepStatus::Completed {
1156            return None;
1157        }
1158        let output = StepOutput {
1159            output: step.output.clone().unwrap_or(Value::Null),
1160            duration_ms: step.duration_ms,
1161            cost_usd: step.cost_usd,
1162            input_tokens: step.input_tokens,
1163            output_tokens: step.output_tokens,
1164            model: None,
1165            debug_messages: None,
1166        };
1167        self.total_cost_usd += output.cost_usd;
1168        self.total_duration_ms += output.duration_ms;
1169        self.last_step_ids = vec![step.id];
1170        info!(
1171            run_id = %self.run_id,
1172            step = %step.name,
1173            position,
1174            "step replayed from previous execution"
1175        );
1176        Some(output)
1177    }
1178
1179    /// Internal: execute a step with full persistence lifecycle.
1180    async fn execute_step(
1181        &mut self,
1182        name: &str,
1183        kind: StepKind,
1184        config: StepConfig,
1185    ) -> Result<StepOutput, EngineError> {
1186        let position = self.position;
1187        self.position += 1;
1188
1189        // Replay: if this step already completed in a prior execution, return cached output.
1190        if let Some(output) = self.try_replay_step(position) {
1191            return Ok(output);
1192        }
1193
1194        // Cost cap: refuse before creating the step record, so a run that hits
1195        // its cap never launches the work it cannot afford.
1196        if let StepConfig::Agent(ref agent_config) = config {
1197            self.check_run_budget(step_budget_usd(agent_config.max_budget_usd))?;
1198        }
1199
1200        // Create step record in Pending.
1201        let step = self
1202            .store
1203            .create_step(NewStep {
1204                run_id: self.run_id,
1205                name: name.to_string(),
1206                kind,
1207                position,
1208                input: Some(serde_json::to_value(&config)?),
1209            })
1210            .await?;
1211
1212        self.start_step(step.id, Utc::now()).await?;
1213
1214        let step_log_sender = self
1215            .log_sender
1216            .as_ref()
1217            .map(|s| StepLogSender::new(s.clone(), self.run_id, step.id, name.to_string()));
1218
1219        match execute_step_config(&config, &self.provider, step_log_sender).await {
1220            Ok(output) => {
1221                self.total_cost_usd += output.cost_usd;
1222                self.total_duration_ms += output.duration_ms;
1223
1224                let debug_messages_json = output.debug_messages_json();
1225
1226                let completed_at = Utc::now();
1227                self.store
1228                    .update_step(
1229                        step.id,
1230                        StepUpdate {
1231                            status: Some(StepStatus::Completed),
1232                            output: Some(output.output.clone()),
1233                            duration_ms: Some(output.duration_ms),
1234                            cost_usd: Some(output.cost_usd),
1235                            input_tokens: output.input_tokens,
1236                            output_tokens: output.output_tokens,
1237                            completed_at: Some(completed_at),
1238                            debug_messages: debug_messages_json,
1239                            ..StepUpdate::default()
1240                        },
1241                    )
1242                    .await?;
1243
1244                info!(
1245                    run_id = %self.run_id,
1246                    step = %name,
1247                    duration_ms = output.duration_ms,
1248                    "step completed"
1249                );
1250
1251                self.last_step_ids = vec![step.id];
1252
1253                Ok(output)
1254            }
1255            Err(err) => {
1256                let completed_at = Utc::now();
1257                let debug_messages_json = extract_debug_messages_from_error(&err);
1258                let partial = extract_partial_usage_from_error(&err);
1259                let raw_response_output = extract_raw_response_from_error(&err);
1260
1261                if let Some(ref usage) = partial {
1262                    if let Some(cost) = usage.cost_usd {
1263                        self.total_cost_usd += cost;
1264                    }
1265                    if let Some(dur) = usage.duration_ms {
1266                        self.total_duration_ms += dur;
1267                    }
1268                }
1269
1270                if let Err(store_err) = self
1271                    .store
1272                    .update_step(
1273                        step.id,
1274                        StepUpdate {
1275                            status: Some(StepStatus::Failed),
1276                            error: Some(err.to_string()),
1277                            output: raw_response_output,
1278                            completed_at: Some(completed_at),
1279                            debug_messages: debug_messages_json,
1280                            duration_ms: partial.as_ref().and_then(|p| p.duration_ms),
1281                            cost_usd: partial.as_ref().and_then(|p| p.cost_usd),
1282                            input_tokens: partial.as_ref().and_then(|p| p.input_tokens),
1283                            output_tokens: partial.as_ref().and_then(|p| p.output_tokens),
1284                            ..StepUpdate::default()
1285                        },
1286                    )
1287                    .await
1288                {
1289                    tracing::error!(step_id = %step.id, error = %store_err, "failed to persist step failure");
1290                }
1291
1292                Err(err)
1293            }
1294        }
1295    }
1296
1297    /// Record dependency edges and transition a step to Running.
1298    ///
1299    /// Records edges from `step_id` to all `last_step_ids`, then
1300    /// transitions the step to `Running` with the given timestamp.
1301    async fn start_step(&self, step_id: Uuid, now: DateTime<Utc>) -> Result<(), EngineError> {
1302        if !self.last_step_ids.is_empty() {
1303            let deps: Vec<NewStepDependency> = self
1304                .last_step_ids
1305                .iter()
1306                .map(|&depends_on| NewStepDependency {
1307                    step_id,
1308                    depends_on,
1309                })
1310                .collect();
1311            self.store.create_step_dependencies(deps).await?;
1312        }
1313
1314        self.store
1315            .update_step(
1316                step_id,
1317                StepUpdate {
1318                    status: Some(StepStatus::Running),
1319                    started_at: Some(now),
1320                    ..StepUpdate::default()
1321                },
1322            )
1323            .await?;
1324
1325        Ok(())
1326    }
1327
1328    /// Access the store directly (advanced usage).
1329    pub fn store(&self) -> &Arc<dyn Store> {
1330        &self.store
1331    }
1332
1333    /// Access the payload that triggered this run.
1334    ///
1335    /// Fetches the run from the store and returns its payload.
1336    ///
1337    /// # Errors
1338    ///
1339    /// Returns [`EngineError::Store`] if the run is not found.
1340    pub async fn payload(&self) -> Result<Value, EngineError> {
1341        let run = self
1342            .store
1343            .get_run(self.run_id)
1344            .await?
1345            .ok_or(EngineError::Store(
1346                ironflow_store::error::StoreError::RunNotFound(self.run_id),
1347            ))?;
1348        Ok(run.payload)
1349    }
1350}
1351
1352impl fmt::Debug for WorkflowContext {
1353    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1354        f.debug_struct("WorkflowContext")
1355            .field("run_id", &self.run_id)
1356            .field("position", &self.position)
1357            .field("total_cost_usd", &self.total_cost_usd)
1358            .field("inherited_cost_usd", &self.inherited_cost_usd)
1359            .field("max_cost_usd", &self.max_cost_usd)
1360            .finish_non_exhaustive()
1361    }
1362}
1363
1364/// Extract debug messages from an engine error, if it wraps a schema validation
1365/// failure that carries a verbose conversation trace.
1366fn extract_debug_messages_from_error(err: &EngineError) -> Option<Value> {
1367    if let EngineError::Operation(OperationError::Agent(AgentError::SchemaValidation {
1368        debug_messages,
1369        ..
1370    })) = err
1371        && !debug_messages.is_empty()
1372    {
1373        return serde_json::to_value(debug_messages).ok();
1374    }
1375    None
1376}
1377
1378/// Partial usage with `Decimal` cost, converted from the `f64` in [`PartialUsage`].
1379///
1380/// Exists only because `ironflow-store` uses [`Decimal`] for monetary values
1381/// while `ironflow-core` uses `f64` (the CLI's native type). The conversion
1382/// happens here, at the engine/store boundary.
1383struct StepPartialUsage {
1384    cost_usd: Option<Decimal>,
1385    duration_ms: Option<u64>,
1386    input_tokens: Option<u64>,
1387    output_tokens: Option<u64>,
1388}
1389
1390/// Extract the raw response text from a schema validation error.
1391///
1392/// When the agent produced text but structured output extraction failed,
1393/// this returns the truncated raw text so it can be persisted as the
1394/// step output for dashboard visibility.
1395fn extract_raw_response_from_error(err: &EngineError) -> Option<Value> {
1396    if let EngineError::Operation(OperationError::Agent(AgentError::SchemaValidation {
1397        raw_response: Some(text),
1398        ..
1399    })) = err
1400    {
1401        return Some(Value::String(text.clone()));
1402    }
1403    None
1404}
1405
1406fn extract_partial_usage_from_error(err: &EngineError) -> Option<StepPartialUsage> {
1407    if let EngineError::Operation(OperationError::Agent(AgentError::SchemaValidation {
1408        partial_usage,
1409        ..
1410    })) = err
1411        && (partial_usage.cost_usd.is_some() || partial_usage.duration_ms.is_some())
1412    {
1413        return Some(StepPartialUsage {
1414            cost_usd: partial_usage
1415                .cost_usd
1416                .and_then(|c| Decimal::try_from(c).ok()),
1417            duration_ms: partial_usage.duration_ms,
1418            input_tokens: partial_usage.input_tokens,
1419            output_tokens: partial_usage.output_tokens,
1420        });
1421    }
1422    None
1423}
1424
1425#[cfg(test)]
1426mod tests {
1427    use super::*;
1428    use ironflow_core::providers::claude::ClaudeCodeProvider;
1429    use ironflow_core::providers::record_replay::RecordReplayProvider;
1430    use ironflow_store::memory::InMemoryStore;
1431    use ironflow_store::models::{Run, RunActor, RunFilter};
1432    use ironflow_store::store::RunStore;
1433    use serde_json::json;
1434    use std::sync::Arc;
1435    use std::sync::atomic::{AtomicBool, Ordering};
1436    use uuid::Uuid;
1437
1438    /// Helper to create a test provider with fixtures
1439    fn create_test_provider() -> Arc<dyn ironflow_core::provider::AgentProvider> {
1440        let inner = ClaudeCodeProvider::new();
1441        Arc::new(RecordReplayProvider::replay(
1442            inner,
1443            "/tmp/ironflow-fixtures",
1444        ))
1445    }
1446
1447    /// Helper to create a test context
1448    fn create_test_context() -> WorkflowContext {
1449        let store = Arc::new(InMemoryStore::new());
1450        let provider = create_test_provider();
1451        let run_id = Uuid::now_v7();
1452        WorkflowContext::new(run_id, store, provider)
1453    }
1454
1455    #[test]
1456    fn context_new_initializes_correctly() {
1457        let ctx = create_test_context();
1458        assert_eq!(ctx.position, 0);
1459        assert_eq!(ctx.total_cost_usd, Decimal::ZERO);
1460        assert_eq!(ctx.total_duration_ms, 0);
1461        assert!(ctx.last_step_ids.is_empty());
1462        assert!(ctx.replay_steps.is_empty());
1463        assert!(ctx.log_sender.is_none());
1464    }
1465
1466    #[test]
1467    fn context_run_id_returns_correct_id() {
1468        let run_id = Uuid::now_v7();
1469        let store = Arc::new(InMemoryStore::new());
1470        let provider = create_test_provider();
1471        let ctx = WorkflowContext::new(run_id, store, provider);
1472        assert_eq!(ctx.run_id(), run_id);
1473    }
1474
1475    #[test]
1476    fn context_total_cost_usd_initially_zero() {
1477        let ctx = create_test_context();
1478        assert_eq!(ctx.total_cost_usd(), Decimal::ZERO);
1479    }
1480
1481    #[test]
1482    fn context_total_duration_ms_initially_zero() {
1483        let ctx = create_test_context();
1484        assert_eq!(ctx.total_duration_ms(), 0);
1485    }
1486
1487    #[test]
1488    fn context_with_handler_resolver_creates_context_with_resolver() {
1489        let store = Arc::new(InMemoryStore::new());
1490        let provider = create_test_provider();
1491        let run_id = Uuid::now_v7();
1492
1493        let called = Arc::new(AtomicBool::new(false));
1494        let called_clone = called.clone();
1495
1496        let resolver: HandlerResolver = Arc::new(move |_name: &str| {
1497            called_clone.store(true, Ordering::SeqCst);
1498            None
1499        });
1500
1501        let ctx = WorkflowContext::with_handler_resolver(run_id, store, provider, resolver);
1502
1503        assert_eq!(ctx.run_id(), run_id);
1504        assert!(ctx.handler_resolver.is_some());
1505    }
1506
1507    #[tokio::test]
1508    async fn context_set_log_sender_attaches_sender() {
1509        let mut ctx = create_test_context();
1510        let (sender, _receiver) = crate::log_sender::channel();
1511        ctx.set_log_sender(sender);
1512        assert!(ctx.log_sender.is_some());
1513    }
1514
1515    #[tokio::test]
1516    async fn context_skip_creates_skipped_step() {
1517        let store = Arc::new(InMemoryStore::new());
1518        let provider = create_test_provider();
1519
1520        // Create the run first using RunStore trait
1521        store
1522            .create_run(NewRun {
1523                created_by: None,
1524                workflow_name: "test".to_string(),
1525                trigger: TriggerKind::Manual,
1526                payload: json!({}),
1527                max_retries: 0,
1528                handler_version: None,
1529                labels: Default::default(),
1530                scheduled_at: None,
1531                idempotency_key: None,
1532                max_cost_usd: None,
1533            })
1534            .await
1535            .expect("failed to create run")
1536            .into_run();
1537
1538        // Get the created run to extract its ID
1539        let runs = store
1540            .list_runs(RunFilter::default(), 1, 10)
1541            .await
1542            .expect("failed to list runs");
1543        let created_run_id = runs.items[0].id;
1544
1545        let mut ctx = WorkflowContext::new(created_run_id, store.clone(), provider);
1546        let initial_position = ctx.position;
1547
1548        ctx.skip("skip-step", "condition not met")
1549            .await
1550            .expect("skip failed");
1551
1552        assert_eq!(ctx.position, initial_position + 1);
1553        assert!(!ctx.last_step_ids.is_empty());
1554
1555        // Verify the step was recorded with Skipped status
1556        let steps = store
1557            .list_steps(created_run_id)
1558            .await
1559            .expect("failed to list steps");
1560        assert_eq!(steps.len(), 1);
1561        assert_eq!(steps[0].status.state, StepStatus::Skipped);
1562    }
1563
1564    /// Sub-workflow handler that records no steps, so the child run reaches a
1565    /// terminal state without touching the filesystem or the network.
1566    struct NoopSubWorkflow;
1567
1568    impl WorkflowHandler for NoopSubWorkflow {
1569        fn name(&self) -> &str {
1570            "noop-sub"
1571        }
1572
1573        fn execute<'a>(
1574            &'a self,
1575            _ctx: &'a mut WorkflowContext,
1576        ) -> crate::handler::HandlerFuture<'a> {
1577            Box::pin(async move { Ok(()) })
1578        }
1579    }
1580
1581    /// Run a parent workflow authored by `created_by` and return the child run
1582    /// created by its sub-workflow step.
1583    async fn child_run_of_parent_authored_by(created_by: Option<RunActor>) -> Run {
1584        let store = Arc::new(InMemoryStore::new());
1585        let provider = create_test_provider();
1586
1587        let parent = store
1588            .create_run(NewRun {
1589                workflow_name: "parent".to_string(),
1590                trigger: TriggerKind::Api,
1591                payload: json!({}),
1592                max_retries: 0,
1593                handler_version: None,
1594                labels: Default::default(),
1595                scheduled_at: None,
1596                created_by,
1597                idempotency_key: None,
1598                max_cost_usd: None,
1599            })
1600            .await
1601            .expect("failed to create parent run")
1602            .into_run();
1603
1604        let resolver: HandlerResolver = Arc::new(|name: &str| match name {
1605            "noop-sub" => Some(Arc::new(NoopSubWorkflow) as Arc<dyn WorkflowHandler>),
1606            _ => None,
1607        });
1608
1609        let mut ctx =
1610            WorkflowContext::with_handler_resolver(parent.id, store.clone(), provider, resolver);
1611        ctx.workflow(&NoopSubWorkflow, json!({}))
1612            .await
1613            .expect("sub-workflow failed");
1614
1615        let runs = store
1616            .list_runs(RunFilter::default(), 1, 10)
1617            .await
1618            .expect("failed to list runs");
1619        runs.items
1620            .into_iter()
1621            .find(|r| r.workflow_name == "noop-sub")
1622            .expect("child run was created")
1623    }
1624
1625    #[tokio::test]
1626    async fn child_run_inherits_the_parent_author() {
1627        let user_id = Uuid::now_v7();
1628        let child = child_run_of_parent_authored_by(Some(RunActor::User { user_id })).await;
1629
1630        assert_eq!(child.created_by, Some(RunActor::User { user_id }));
1631    }
1632
1633    #[tokio::test]
1634    async fn child_run_of_an_unattributed_parent_has_no_author() {
1635        let child = child_run_of_parent_authored_by(None).await;
1636
1637        assert!(child.created_by.is_none());
1638    }
1639
1640    #[tokio::test]
1641    async fn context_parallel_empty_steps_returns_empty_vec() {
1642        let mut ctx = create_test_context();
1643        let results = ctx
1644            .parallel(vec![], true)
1645            .await
1646            .expect("parallel should not fail on empty input");
1647        assert!(results.is_empty());
1648    }
1649
1650    #[tokio::test]
1651    async fn context_approval_first_execution_returns_error() {
1652        let store = Arc::new(InMemoryStore::new());
1653        let provider = create_test_provider();
1654
1655        // Create the run first
1656        store
1657            .create_run(NewRun {
1658                created_by: None,
1659                workflow_name: "test".to_string(),
1660                trigger: TriggerKind::Manual,
1661                payload: json!({}),
1662                max_retries: 0,
1663                handler_version: None,
1664                labels: Default::default(),
1665                scheduled_at: None,
1666                idempotency_key: None,
1667                max_cost_usd: None,
1668            })
1669            .await
1670            .expect("failed to create run")
1671            .into_run();
1672
1673        // Get the created run to extract its ID
1674        let runs = store
1675            .list_runs(RunFilter::default(), 1, 10)
1676            .await
1677            .expect("failed to list runs");
1678        let created_run_id = runs.items[0].id;
1679
1680        let mut ctx = WorkflowContext::new(created_run_id, store.clone(), provider);
1681
1682        let result = ctx
1683            .approval(
1684                "approve-step",
1685                crate::config::ApprovalConfig::new("Continue?"),
1686            )
1687            .await;
1688
1689        // First execution should return ApprovalRequired error
1690        assert!(matches!(result, Err(EngineError::ApprovalRequired { .. })));
1691
1692        // Verify position incremented
1693        assert_eq!(ctx.position, 1);
1694
1695        // Verify step was created with AwaitingApproval status
1696        let steps = store
1697            .list_steps(created_run_id)
1698            .await
1699            .expect("failed to list steps");
1700        assert_eq!(steps.len(), 1);
1701        assert_eq!(steps[0].status.state, StepStatus::AwaitingApproval);
1702    }
1703
1704    #[tokio::test]
1705    async fn context_approval_replay_returns_ok() {
1706        let store = Arc::new(InMemoryStore::new());
1707        let provider = create_test_provider();
1708
1709        // Create the run first
1710        store
1711            .create_run(NewRun {
1712                created_by: None,
1713                workflow_name: "test".to_string(),
1714                trigger: TriggerKind::Manual,
1715                payload: json!({}),
1716                max_retries: 0,
1717                handler_version: None,
1718                labels: Default::default(),
1719                scheduled_at: None,
1720                idempotency_key: None,
1721                max_cost_usd: None,
1722            })
1723            .await
1724            .expect("failed to create run")
1725            .into_run();
1726
1727        // Get the created run to extract its ID
1728        let runs = store
1729            .list_runs(RunFilter::default(), 1, 10)
1730            .await
1731            .expect("failed to list runs");
1732        let created_run_id = runs.items[0].id;
1733
1734        // Create an approval step that's already in AwaitingApproval state
1735        let step = store
1736            .create_step(NewStep {
1737                run_id: created_run_id,
1738                name: "approval".to_string(),
1739                kind: StepKind::Approval,
1740                position: 0,
1741                input: None,
1742            })
1743            .await
1744            .expect("failed to create step");
1745
1746        // Transition through proper states: Pending -> Running -> AwaitingApproval
1747        store
1748            .update_step(
1749                step.id,
1750                StepUpdate {
1751                    status: Some(StepStatus::Running),
1752                    started_at: Some(Utc::now()),
1753                    ..StepUpdate::default()
1754                },
1755            )
1756            .await
1757            .expect("failed to update step to Running");
1758
1759        store
1760            .update_step(
1761                step.id,
1762                StepUpdate {
1763                    status: Some(StepStatus::AwaitingApproval),
1764                    ..StepUpdate::default()
1765                },
1766            )
1767            .await
1768            .expect("failed to update step to AwaitingApproval");
1769
1770        // Create context and load replay steps
1771        let mut ctx = WorkflowContext::new(created_run_id, store.clone(), provider);
1772        ctx.load_replay_steps()
1773            .await
1774            .expect("failed to load replay steps");
1775
1776        // Now approval should succeed (replay)
1777        let result = ctx
1778            .approval("approval", crate::config::ApprovalConfig::new("Continue?"))
1779            .await;
1780
1781        assert!(result.is_ok());
1782
1783        // Verify the step was marked Completed
1784        let steps = store
1785            .list_steps(created_run_id)
1786            .await
1787            .expect("failed to list steps");
1788        assert_eq!(steps.len(), 1);
1789        assert_eq!(steps[0].status.state, StepStatus::Completed);
1790    }
1791
1792    #[tokio::test]
1793    async fn context_load_replay_steps_loads_completed_steps() {
1794        let store = Arc::new(InMemoryStore::new());
1795        let provider = create_test_provider();
1796
1797        // Create the run first
1798        store
1799            .create_run(NewRun {
1800                created_by: None,
1801                workflow_name: "test".to_string(),
1802                trigger: TriggerKind::Manual,
1803                payload: json!({}),
1804                max_retries: 0,
1805                handler_version: None,
1806                labels: Default::default(),
1807                scheduled_at: None,
1808                idempotency_key: None,
1809                max_cost_usd: None,
1810            })
1811            .await
1812            .expect("failed to create run")
1813            .into_run();
1814
1815        // Get the created run to extract its ID
1816        let runs = store
1817            .list_runs(RunFilter::default(), 1, 10)
1818            .await
1819            .expect("failed to list runs");
1820        let created_run_id = runs.items[0].id;
1821
1822        // Create multiple steps with different statuses
1823        let completed_step = store
1824            .create_step(NewStep {
1825                run_id: created_run_id,
1826                name: "completed".to_string(),
1827                kind: StepKind::Shell,
1828                position: 0,
1829                input: None,
1830            })
1831            .await
1832            .expect("failed to create step");
1833
1834        // Transition to Running then Completed
1835        store
1836            .update_step(
1837                completed_step.id,
1838                StepUpdate {
1839                    status: Some(StepStatus::Running),
1840                    started_at: Some(Utc::now()),
1841                    ..StepUpdate::default()
1842                },
1843            )
1844            .await
1845            .expect("failed to update step to Running");
1846
1847        store
1848            .update_step(
1849                completed_step.id,
1850                StepUpdate {
1851                    status: Some(StepStatus::Completed),
1852                    completed_at: Some(Utc::now()),
1853                    ..StepUpdate::default()
1854                },
1855            )
1856            .await
1857            .expect("failed to update step to Completed");
1858
1859        let _pending_step = store
1860            .create_step(NewStep {
1861                run_id: created_run_id,
1862                name: "pending".to_string(),
1863                kind: StepKind::Shell,
1864                position: 1,
1865                input: None,
1866            })
1867            .await
1868            .expect("failed to create step");
1869
1870        // Load replay steps
1871        let mut ctx = WorkflowContext::new(created_run_id, store, provider);
1872        ctx.load_replay_steps()
1873            .await
1874            .expect("failed to load replay steps");
1875
1876        // Only completed step should be in replay_steps
1877        assert_eq!(ctx.replay_steps.len(), 1);
1878        assert!(ctx.replay_steps.contains_key(&0));
1879        assert!(!ctx.replay_steps.contains_key(&1));
1880    }
1881
1882    #[tokio::test]
1883    async fn context_payload_returns_run_payload() {
1884        let store = Arc::new(InMemoryStore::new());
1885        let provider = create_test_provider();
1886        let test_payload = json!({"key": "value", "number": 42});
1887
1888        // Create the run first
1889        store
1890            .create_run(NewRun {
1891                created_by: None,
1892                workflow_name: "test".to_string(),
1893                trigger: TriggerKind::Manual,
1894                payload: test_payload.clone(),
1895                max_retries: 0,
1896                handler_version: None,
1897                labels: Default::default(),
1898                scheduled_at: None,
1899                idempotency_key: None,
1900                max_cost_usd: None,
1901            })
1902            .await
1903            .expect("failed to create run")
1904            .into_run();
1905
1906        // Get the created run to extract its ID
1907        let runs = store
1908            .list_runs(RunFilter::default(), 1, 10)
1909            .await
1910            .expect("failed to list runs");
1911        let created_run_id = runs.items[0].id;
1912
1913        let ctx = WorkflowContext::new(created_run_id, store, provider);
1914        let payload = ctx.payload().await.expect("failed to get payload");
1915
1916        assert_eq!(payload, test_payload);
1917    }
1918
1919    #[tokio::test]
1920    async fn context_payload_returns_error_for_nonexistent_run() {
1921        let store = Arc::new(InMemoryStore::new());
1922        let provider = create_test_provider();
1923        let run_id = Uuid::now_v7();
1924
1925        let ctx = WorkflowContext::new(run_id, store, provider);
1926        let result = ctx.payload().await;
1927
1928        assert!(result.is_err());
1929    }
1930
1931    #[tokio::test]
1932    async fn context_store_returns_reference() {
1933        let ctx = create_test_context();
1934        let _store = ctx.store();
1935        // store() returns a reference to the Arc<dyn Store>, which is always available
1936    }
1937
1938    #[test]
1939    fn context_debug_formatting() {
1940        let ctx = create_test_context();
1941        let debug_str = format!("{:?}", ctx);
1942        assert!(debug_str.contains("WorkflowContext"));
1943        assert!(debug_str.contains("run_id"));
1944    }
1945
1946    #[tokio::test]
1947    async fn context_last_step_ids_tracks_executed_steps() {
1948        let store = Arc::new(InMemoryStore::new());
1949        let provider = create_test_provider();
1950
1951        // Create the run first
1952        store
1953            .create_run(NewRun {
1954                created_by: None,
1955                workflow_name: "test".to_string(),
1956                trigger: TriggerKind::Manual,
1957                payload: json!({}),
1958                max_retries: 0,
1959                handler_version: None,
1960                labels: Default::default(),
1961                scheduled_at: None,
1962                idempotency_key: None,
1963                max_cost_usd: None,
1964            })
1965            .await
1966            .expect("failed to create run")
1967            .into_run();
1968
1969        // Get the created run to extract its ID
1970        let runs = store
1971            .list_runs(RunFilter::default(), 1, 10)
1972            .await
1973            .expect("failed to list runs");
1974        let created_run_id = runs.items[0].id;
1975
1976        let mut ctx = WorkflowContext::new(created_run_id, store, provider);
1977        assert!(ctx.last_step_ids.is_empty());
1978
1979        ctx.skip("step1", "reason").await.expect("skip failed");
1980
1981        assert_eq!(ctx.last_step_ids.len(), 1);
1982
1983        ctx.skip("step2", "reason").await.expect("skip failed");
1984
1985        // last_step_ids should now contain only step2's ID
1986        assert_eq!(ctx.last_step_ids.len(), 1);
1987    }
1988}
ironflow_engine/context.rs

ironflow_engine/
context.rs
