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::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::RunStore;
45
46use crate::config::{
47 AgentStepConfig, ApprovalConfig, HttpConfig, ShellConfig, StepConfig, WorkflowStepConfig,
48};
49use crate::error::EngineError;
50use crate::executor::{ParallelStepResult, StepOutput, execute_step_config};
51use crate::handler::WorkflowHandler;
52use crate::operation::Operation;
53
54/// Callback type for resolving workflow handlers by name.
55pub(crate) type HandlerResolver =
56 Arc<dyn Fn(&str) -> Option<Arc<dyn WorkflowHandler>> + Send + Sync>;
57
58/// Execution context for a single workflow run.
59///
60/// Tracks the current step position and provides convenience methods
61/// for executing operations with automatic persistence.
62///
63/// # Examples
64///
65/// ```no_run
66/// use ironflow_engine::context::WorkflowContext;
67/// use ironflow_engine::config::ShellConfig;
68/// use ironflow_engine::error::EngineError;
69///
70/// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
71/// let result = ctx.shell("greet", ShellConfig::new("echo hello")).await?;
72/// assert!(result.output["stdout"].as_str().unwrap().contains("hello"));
73/// # Ok(())
74/// # }
75/// ```
76pub struct WorkflowContext {
77 run_id: Uuid,
78 store: Arc<dyn RunStore>,
79 provider: Arc<dyn AgentProvider>,
80 handler_resolver: Option<HandlerResolver>,
81 position: u32,
82 /// IDs of the last executed step(s) -- used to record DAG dependencies.
83 last_step_ids: Vec<Uuid>,
84 /// Accumulated cost across all steps in this run.
85 total_cost_usd: Decimal,
86 /// Accumulated duration across all steps.
87 total_duration_ms: u64,
88 /// Steps from a previous execution, keyed by position.
89 /// Used when resuming after approval to replay completed steps.
90 replay_steps: HashMap<u32, Step>,
91}
92
93impl WorkflowContext {
94 /// Create a new context for a run.
95 ///
96 /// Not typically called directly — the [`Engine`](crate::engine::Engine)
97 /// creates this when executing a [`WorkflowHandler`].
98 pub fn new(run_id: Uuid, store: Arc<dyn RunStore>, provider: Arc<dyn AgentProvider>) -> Self {
99 Self {
100 run_id,
101 store,
102 provider,
103 handler_resolver: None,
104 position: 0,
105 last_step_ids: Vec::new(),
106 total_cost_usd: Decimal::ZERO,
107 total_duration_ms: 0,
108 replay_steps: HashMap::new(),
109 }
110 }
111
112 /// Create a new context with a handler resolver for sub-workflow support.
113 ///
114 /// The resolver is called when [`workflow`](Self::workflow) is invoked to
115 /// look up registered handlers by name.
116 pub(crate) fn with_handler_resolver(
117 run_id: Uuid,
118 store: Arc<dyn RunStore>,
119 provider: Arc<dyn AgentProvider>,
120 resolver: HandlerResolver,
121 ) -> Self {
122 Self {
123 run_id,
124 store,
125 provider,
126 handler_resolver: Some(resolver),
127 position: 0,
128 last_step_ids: Vec::new(),
129 total_cost_usd: Decimal::ZERO,
130 total_duration_ms: 0,
131 replay_steps: HashMap::new(),
132 }
133 }
134
135 /// Load existing steps from the store for replay after approval.
136 ///
137 /// Called by the engine when resuming a run. All completed steps
138 /// and the approved approval step are indexed by position so that
139 /// `execute_step` and `approval` can skip them.
140 pub(crate) async fn load_replay_steps(&mut self) -> Result<(), EngineError> {
141 let steps = self.store.list_steps(self.run_id).await?;
142 for step in steps {
143 let dominated = matches!(
144 step.status.state,
145 StepStatus::Completed | StepStatus::Running | StepStatus::AwaitingApproval
146 );
147 if dominated {
148 self.replay_steps.insert(step.position, step);
149 }
150 }
151 Ok(())
152 }
153
154 /// The run ID this context is executing for.
155 pub fn run_id(&self) -> Uuid {
156 self.run_id
157 }
158
159 /// Accumulated cost across all executed steps so far.
160 pub fn total_cost_usd(&self) -> Decimal {
161 self.total_cost_usd
162 }
163
164 /// Accumulated duration across all executed steps so far.
165 pub fn total_duration_ms(&self) -> u64 {
166 self.total_duration_ms
167 }
168
169 /// Execute multiple steps concurrently (wait-all model).
170 ///
171 /// All steps in the batch execute in parallel via `tokio::JoinSet`.
172 /// Each step is recorded with the same `position` (execution wave).
173 /// Dependencies on previous steps are recorded automatically.
174 ///
175 /// When `fail_fast` is true, remaining steps are aborted on the first
176 /// failure. When false, all steps run to completion and the first
177 /// error is returned.
178 ///
179 /// # Errors
180 ///
181 /// Returns [`EngineError`] if any step fails.
182 ///
183 /// # Examples
184 ///
185 /// ```no_run
186 /// use ironflow_engine::context::WorkflowContext;
187 /// use ironflow_engine::config::{StepConfig, ShellConfig};
188 /// use ironflow_engine::error::EngineError;
189 ///
190 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
191 /// let results = ctx.parallel(
192 /// vec![
193 /// ("test-unit", StepConfig::Shell(ShellConfig::new("cargo test --lib"))),
194 /// ("lint", StepConfig::Shell(ShellConfig::new("cargo clippy"))),
195 /// ],
196 /// true,
197 /// ).await?;
198 ///
199 /// for r in &results {
200 /// println!("{}: {:?}", r.name, r.output.output);
201 /// }
202 /// # Ok(())
203 /// # }
204 /// ```
205 pub async fn parallel(
206 &mut self,
207 steps: Vec<(&str, StepConfig)>,
208 fail_fast: bool,
209 ) -> Result<Vec<ParallelStepResult>, EngineError> {
210 if steps.is_empty() {
211 return Ok(Vec::new());
212 }
213
214 let wave_position = self.position;
215 self.position += 1;
216
217 let now = Utc::now();
218 let mut step_records: Vec<(Uuid, String, StepConfig)> = Vec::with_capacity(steps.len());
219
220 for (name, config) in &steps {
221 let kind = config.kind();
222 let step = self
223 .store
224 .create_step(NewStep {
225 run_id: self.run_id,
226 name: name.to_string(),
227 kind,
228 position: wave_position,
229 input: Some(serde_json::to_value(config)?),
230 })
231 .await?;
232
233 self.start_step(step.id, now).await?;
234
235 step_records.push((step.id, name.to_string(), config.clone()));
236 }
237
238 let mut join_set = JoinSet::new();
239 for (idx, (_id, _name, config)) in step_records.iter().enumerate() {
240 let provider = self.provider.clone();
241 let config = config.clone();
242 join_set.spawn(async move { (idx, execute_step_config(&config, &provider).await) });
243 }
244
245 // JoinSet returns in completion order; indexed_results restores input order.
246 let mut indexed_results: Vec<Option<Result<StepOutput, String>>> =
247 vec![None; step_records.len()];
248 let mut first_error: Option<EngineError> = None;
249
250 while let Some(join_result) = join_set.join_next().await {
251 let (idx, step_result) = match join_result {
252 Ok(r) => r,
253 Err(e) => {
254 if first_error.is_none() {
255 first_error = Some(EngineError::StepConfig(format!("join error: {e}")));
256 }
257 if fail_fast {
258 join_set.abort_all();
259 }
260 continue;
261 }
262 };
263
264 let (step_id, step_name, _) = &step_records[idx];
265 let completed_at = Utc::now();
266
267 match step_result {
268 Ok(output) => {
269 self.total_cost_usd += output.cost_usd;
270 self.total_duration_ms += output.duration_ms;
271
272 let debug_messages_json = output.debug_messages_json();
273
274 self.store
275 .update_step(
276 *step_id,
277 StepUpdate {
278 status: Some(StepStatus::Completed),
279 output: Some(output.output.clone()),
280 duration_ms: Some(output.duration_ms),
281 cost_usd: Some(output.cost_usd),
282 input_tokens: output.input_tokens,
283 output_tokens: output.output_tokens,
284 completed_at: Some(completed_at),
285 debug_messages: debug_messages_json,
286 ..StepUpdate::default()
287 },
288 )
289 .await?;
290
291 info!(
292 run_id = %self.run_id,
293 step = %step_name,
294 duration_ms = output.duration_ms,
295 "parallel step completed"
296 );
297
298 indexed_results[idx] = Some(Ok(output));
299 }
300 Err(err) => {
301 let err_msg = err.to_string();
302 let debug_messages_json = extract_debug_messages_from_error(&err);
303 let partial = extract_partial_usage_from_error(&err);
304
305 if let Some(ref usage) = partial {
306 if let Some(cost) = usage.cost_usd {
307 self.total_cost_usd += cost;
308 }
309 if let Some(dur) = usage.duration_ms {
310 self.total_duration_ms += dur;
311 }
312 }
313
314 if let Err(store_err) = self
315 .store
316 .update_step(
317 *step_id,
318 StepUpdate {
319 status: Some(StepStatus::Failed),
320 error: Some(err_msg.clone()),
321 completed_at: Some(completed_at),
322 debug_messages: debug_messages_json,
323 duration_ms: partial.as_ref().and_then(|p| p.duration_ms),
324 cost_usd: partial.as_ref().and_then(|p| p.cost_usd),
325 input_tokens: partial.as_ref().and_then(|p| p.input_tokens),
326 output_tokens: partial.as_ref().and_then(|p| p.output_tokens),
327 ..StepUpdate::default()
328 },
329 )
330 .await
331 {
332 tracing::error!(
333 step_id = %step_id,
334 error = %store_err,
335 "failed to persist parallel step failure"
336 );
337 }
338
339 indexed_results[idx] = Some(Err(err_msg.clone()));
340
341 if first_error.is_none() {
342 first_error = Some(err);
343 }
344
345 if fail_fast {
346 join_set.abort_all();
347 }
348 }
349 }
350 }
351
352 if let Some(err) = first_error {
353 return Err(err);
354 }
355
356 self.last_step_ids = step_records.iter().map(|(id, _, _)| *id).collect();
357
358 // Build results in original order.
359 let results: Vec<ParallelStepResult> = step_records
360 .iter()
361 .enumerate()
362 .map(|(idx, (step_id, name, _))| {
363 let output = match indexed_results[idx].take() {
364 Some(Ok(o)) => o,
365 _ => unreachable!("all steps succeeded if no error returned"),
366 };
367 ParallelStepResult {
368 name: name.clone(),
369 output,
370 step_id: *step_id,
371 }
372 })
373 .collect();
374
375 Ok(results)
376 }
377
378 /// Execute a shell step.
379 ///
380 /// Creates the step record, runs the command, persists the result,
381 /// and returns the output for use in subsequent steps.
382 ///
383 /// # Errors
384 ///
385 /// Returns [`EngineError`] if the command fails or the store errors.
386 ///
387 /// # Examples
388 ///
389 /// ```no_run
390 /// use ironflow_engine::context::WorkflowContext;
391 /// use ironflow_engine::config::ShellConfig;
392 /// use ironflow_engine::error::EngineError;
393 ///
394 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
395 /// let files = ctx.shell("list", ShellConfig::new("ls -la")).await?;
396 /// println!("stdout: {}", files.output["stdout"]);
397 /// # Ok(())
398 /// # }
399 /// ```
400 pub async fn shell(
401 &mut self,
402 name: &str,
403 config: ShellConfig,
404 ) -> Result<StepOutput, EngineError> {
405 self.execute_step(name, StepKind::Shell, StepConfig::Shell(config))
406 .await
407 }
408
409 /// Execute an HTTP step.
410 ///
411 /// # Errors
412 ///
413 /// Returns [`EngineError`] if the request fails or the store errors.
414 ///
415 /// # Examples
416 ///
417 /// ```no_run
418 /// use ironflow_engine::context::WorkflowContext;
419 /// use ironflow_engine::config::HttpConfig;
420 /// use ironflow_engine::error::EngineError;
421 ///
422 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
423 /// let resp = ctx.http("health", HttpConfig::get("https://api.example.com/health")).await?;
424 /// println!("status: {}", resp.output["status"]);
425 /// # Ok(())
426 /// # }
427 /// ```
428 pub async fn http(
429 &mut self,
430 name: &str,
431 config: HttpConfig,
432 ) -> Result<StepOutput, EngineError> {
433 self.execute_step(name, StepKind::Http, StepConfig::Http(config))
434 .await
435 }
436
437 /// Execute an agent step.
438 ///
439 /// # Errors
440 ///
441 /// Returns [`EngineError`] if the agent invocation fails or the store errors.
442 ///
443 /// # Examples
444 ///
445 /// ```no_run
446 /// use ironflow_engine::context::WorkflowContext;
447 /// use ironflow_engine::config::AgentStepConfig;
448 /// use ironflow_engine::error::EngineError;
449 ///
450 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
451 /// let review = ctx.agent("review", AgentStepConfig::new("Review the code")).await?;
452 /// println!("review: {}", review.output["value"]);
453 /// # Ok(())
454 /// # }
455 /// ```
456 pub async fn agent(
457 &mut self,
458 name: &str,
459 config: impl Into<AgentStepConfig>,
460 ) -> Result<StepOutput, EngineError> {
461 self.execute_step(name, StepKind::Agent, StepConfig::Agent(config.into()))
462 .await
463 }
464
465 /// Create a human approval gate.
466 ///
467 /// On first execution, records an approval step and returns
468 /// [`EngineError::ApprovalRequired`] to suspend the run. The engine
469 /// transitions the run to `AwaitingApproval`.
470 ///
471 /// On resume (after a human approved via the API), the approval step
472 /// is replayed: it is marked as `Completed` and execution continues
473 /// past it. Multiple approval gates in the same handler work -- each
474 /// one pauses and resumes independently.
475 ///
476 /// # Errors
477 ///
478 /// Returns [`EngineError::ApprovalRequired`] to pause the run on
479 /// first execution. Returns other [`EngineError`] variants on store
480 /// failures.
481 ///
482 /// # Examples
483 ///
484 /// ```no_run
485 /// use ironflow_engine::context::WorkflowContext;
486 /// use ironflow_engine::config::ApprovalConfig;
487 /// use ironflow_engine::error::EngineError;
488 ///
489 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
490 /// ctx.approval("deploy-gate", ApprovalConfig::new("Approve deployment?")).await?;
491 /// // Execution continues here after approval
492 /// # Ok(())
493 /// # }
494 /// ```
495 pub async fn approval(
496 &mut self,
497 name: &str,
498 config: ApprovalConfig,
499 ) -> Result<(), EngineError> {
500 let position = self.position;
501 self.position += 1;
502
503 // Replay: if this approval step exists from a prior execution,
504 // the run was approved -- mark it completed (if not already) and continue.
505 if let Some(existing) = self.replay_steps.get(&position)
506 && existing.kind == StepKind::Approval
507 {
508 if existing.status.state == StepStatus::AwaitingApproval {
509 self.store
510 .update_step(
511 existing.id,
512 StepUpdate {
513 status: Some(StepStatus::Completed),
514 completed_at: Some(Utc::now()),
515 ..StepUpdate::default()
516 },
517 )
518 .await?;
519 }
520
521 self.last_step_ids = vec![existing.id];
522 info!(
523 run_id = %self.run_id,
524 step = %name,
525 position,
526 "approval step replayed (approved)"
527 );
528 return Ok(());
529 }
530
531 // First execution: create the approval step and suspend.
532 let step = self
533 .store
534 .create_step(NewStep {
535 run_id: self.run_id,
536 name: name.to_string(),
537 kind: StepKind::Approval,
538 position,
539 input: Some(serde_json::to_value(&config)?),
540 })
541 .await?;
542
543 self.start_step(step.id, Utc::now()).await?;
544
545 // Transition the step to AwaitingApproval so it reflects
546 // the suspended state on the dashboard.
547 self.store
548 .update_step(
549 step.id,
550 StepUpdate {
551 status: Some(StepStatus::AwaitingApproval),
552 ..StepUpdate::default()
553 },
554 )
555 .await?;
556
557 self.last_step_ids = vec![step.id];
558
559 Err(EngineError::ApprovalRequired {
560 run_id: self.run_id,
561 step_id: step.id,
562 message: config.message().to_string(),
563 })
564 }
565
566 /// Execute a custom operation step.
567 ///
568 /// Runs a user-defined [`Operation`] with full step lifecycle management:
569 /// creates the step record, transitions to Running, executes the operation,
570 /// persists the output and duration, and marks the step Completed or Failed.
571 ///
572 /// The operation's [`kind()`](Operation::kind) is stored as
573 /// [`StepKind::Custom`].
574 ///
575 /// # Errors
576 ///
577 /// Returns [`EngineError`] if the operation fails or the store errors.
578 ///
579 /// # Examples
580 ///
581 /// ```no_run
582 /// use ironflow_engine::context::WorkflowContext;
583 /// use ironflow_engine::operation::Operation;
584 /// use ironflow_engine::error::EngineError;
585 /// use serde_json::{Value, json};
586 /// use std::pin::Pin;
587 /// use std::future::Future;
588 ///
589 /// struct MyOp;
590 /// impl Operation for MyOp {
591 /// fn kind(&self) -> &str { "my-service" }
592 /// fn execute(&self) -> Pin<Box<dyn Future<Output = Result<Value, EngineError>> + Send + '_>> {
593 /// Box::pin(async { Ok(json!({"ok": true})) })
594 /// }
595 /// }
596 ///
597 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
598 /// let result = ctx.operation("call-service", &MyOp).await?;
599 /// println!("output: {}", result.output);
600 /// # Ok(())
601 /// # }
602 /// ```
603 pub async fn operation(
604 &mut self,
605 name: &str,
606 op: &dyn Operation,
607 ) -> Result<StepOutput, EngineError> {
608 let kind = StepKind::Custom(op.kind().to_string());
609 let position = self.position;
610 self.position += 1;
611
612 let step = self
613 .store
614 .create_step(NewStep {
615 run_id: self.run_id,
616 name: name.to_string(),
617 kind,
618 position,
619 input: op.input(),
620 })
621 .await?;
622
623 self.start_step(step.id, Utc::now()).await?;
624
625 let start = Instant::now();
626
627 match op.execute().await {
628 Ok(output_value) => {
629 let duration_ms = start.elapsed().as_millis() as u64;
630 self.total_duration_ms += duration_ms;
631
632 let completed_at = Utc::now();
633 self.store
634 .update_step(
635 step.id,
636 StepUpdate {
637 status: Some(StepStatus::Completed),
638 output: Some(output_value.clone()),
639 duration_ms: Some(duration_ms),
640 cost_usd: Some(Decimal::ZERO),
641 completed_at: Some(completed_at),
642 ..StepUpdate::default()
643 },
644 )
645 .await?;
646
647 info!(
648 run_id = %self.run_id,
649 step = %name,
650 kind = op.kind(),
651 duration_ms,
652 "operation step completed"
653 );
654
655 self.last_step_ids = vec![step.id];
656
657 Ok(StepOutput {
658 output: output_value,
659 duration_ms,
660 cost_usd: Decimal::ZERO,
661 input_tokens: None,
662 output_tokens: None,
663 debug_messages: None,
664 })
665 }
666 Err(err) => {
667 let completed_at = Utc::now();
668 if let Err(store_err) = self
669 .store
670 .update_step(
671 step.id,
672 StepUpdate {
673 status: Some(StepStatus::Failed),
674 error: Some(err.to_string()),
675 completed_at: Some(completed_at),
676 ..StepUpdate::default()
677 },
678 )
679 .await
680 {
681 error!(step_id = %step.id, error = %store_err, "failed to persist step failure");
682 }
683
684 Err(err)
685 }
686 }
687 }
688
689 /// Execute a sub-workflow step.
690 ///
691 /// Creates a child run for the named workflow handler, executes it with
692 /// its own steps and lifecycle, and returns a [`StepOutput`] containing
693 /// the child run ID and aggregated metrics.
694 ///
695 /// Requires the context to be created with
696 /// `with_handler_resolver`.
697 ///
698 /// # Errors
699 ///
700 /// Returns [`EngineError::InvalidWorkflow`] if no handler is registered
701 /// with the given name, or if no handler resolver is available.
702 ///
703 /// # Examples
704 ///
705 /// ```no_run
706 /// use ironflow_engine::context::WorkflowContext;
707 /// use ironflow_engine::error::EngineError;
708 /// use serde_json::json;
709 ///
710 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
711 /// // let result = ctx.workflow(&MySubWorkflow, json!({})).await?;
712 /// # Ok(())
713 /// # }
714 /// ```
715 pub async fn workflow(
716 &mut self,
717 handler: &dyn WorkflowHandler,
718 payload: Value,
719 ) -> Result<StepOutput, EngineError> {
720 let config = WorkflowStepConfig::new(handler.name(), payload);
721 let position = self.position;
722 self.position += 1;
723
724 let step = self
725 .store
726 .create_step(NewStep {
727 run_id: self.run_id,
728 name: config.workflow_name.clone(),
729 kind: StepKind::Workflow,
730 position,
731 input: Some(serde_json::to_value(&config)?),
732 })
733 .await?;
734
735 self.start_step(step.id, Utc::now()).await?;
736
737 match self.execute_child_workflow(&config).await {
738 Ok(output) => {
739 self.total_cost_usd += output.cost_usd;
740 self.total_duration_ms += output.duration_ms;
741
742 let completed_at = Utc::now();
743 self.store
744 .update_step(
745 step.id,
746 StepUpdate {
747 status: Some(StepStatus::Completed),
748 output: Some(output.output.clone()),
749 duration_ms: Some(output.duration_ms),
750 cost_usd: Some(output.cost_usd),
751 completed_at: Some(completed_at),
752 ..StepUpdate::default()
753 },
754 )
755 .await?;
756
757 info!(
758 run_id = %self.run_id,
759 child_workflow = %config.workflow_name,
760 duration_ms = output.duration_ms,
761 "workflow step completed"
762 );
763
764 self.last_step_ids = vec![step.id];
765
766 Ok(output)
767 }
768 Err(err) => {
769 let completed_at = Utc::now();
770 if let Err(store_err) = self
771 .store
772 .update_step(
773 step.id,
774 StepUpdate {
775 status: Some(StepStatus::Failed),
776 error: Some(err.to_string()),
777 completed_at: Some(completed_at),
778 ..StepUpdate::default()
779 },
780 )
781 .await
782 {
783 error!(step_id = %step.id, error = %store_err, "failed to persist step failure");
784 }
785
786 Err(err)
787 }
788 }
789 }
790
791 /// Execute a child workflow and return aggregated output.
792 async fn execute_child_workflow(
793 &self,
794 config: &WorkflowStepConfig,
795 ) -> Result<StepOutput, EngineError> {
796 let resolver = self.handler_resolver.as_ref().ok_or_else(|| {
797 EngineError::InvalidWorkflow(
798 "sub-workflow requires a handler resolver (use Engine to execute)".to_string(),
799 )
800 })?;
801
802 let handler = resolver(&config.workflow_name).ok_or_else(|| {
803 EngineError::InvalidWorkflow(format!("no handler registered: {}", config.workflow_name))
804 })?;
805
806 let child_run = self
807 .store
808 .create_run(NewRun {
809 workflow_name: config.workflow_name.clone(),
810 trigger: TriggerKind::Workflow,
811 payload: config.payload.clone(),
812 max_retries: 0,
813 })
814 .await?;
815
816 let child_run_id = child_run.id;
817 info!(
818 parent_run_id = %self.run_id,
819 child_run_id = %child_run_id,
820 workflow = %config.workflow_name,
821 "child run created"
822 );
823
824 self.store
825 .update_run_status(child_run_id, RunStatus::Running)
826 .await?;
827
828 let run_start = Instant::now();
829 let mut child_ctx = WorkflowContext {
830 run_id: child_run_id,
831 store: self.store.clone(),
832 provider: self.provider.clone(),
833 handler_resolver: self.handler_resolver.clone(),
834 position: 0,
835 last_step_ids: Vec::new(),
836 total_cost_usd: Decimal::ZERO,
837 total_duration_ms: 0,
838 replay_steps: HashMap::new(),
839 };
840
841 let result = handler.execute(&mut child_ctx).await;
842 let total_duration = run_start.elapsed().as_millis() as u64;
843 let completed_at = Utc::now();
844
845 match result {
846 Ok(()) => {
847 self.store
848 .update_run(
849 child_run_id,
850 RunUpdate {
851 status: Some(RunStatus::Completed),
852 cost_usd: Some(child_ctx.total_cost_usd),
853 duration_ms: Some(total_duration),
854 completed_at: Some(completed_at),
855 ..RunUpdate::default()
856 },
857 )
858 .await?;
859
860 Ok(StepOutput {
861 output: serde_json::json!({
862 "run_id": child_run_id,
863 "workflow_name": config.workflow_name,
864 "status": RunStatus::Completed,
865 "cost_usd": child_ctx.total_cost_usd,
866 "duration_ms": total_duration,
867 }),
868 duration_ms: total_duration,
869 cost_usd: child_ctx.total_cost_usd,
870 input_tokens: None,
871 output_tokens: None,
872 debug_messages: None,
873 })
874 }
875 Err(err) => {
876 if let Err(store_err) = self
877 .store
878 .update_run(
879 child_run_id,
880 RunUpdate {
881 status: Some(RunStatus::Failed),
882 error: Some(err.to_string()),
883 cost_usd: Some(child_ctx.total_cost_usd),
884 duration_ms: Some(total_duration),
885 completed_at: Some(completed_at),
886 ..RunUpdate::default()
887 },
888 )
889 .await
890 {
891 error!(
892 child_run_id = %child_run_id,
893 store_error = %store_err,
894 "failed to persist child run failure"
895 );
896 }
897
898 Err(err)
899 }
900 }
901 }
902
903 /// Try to replay a completed step from a previous execution.
904 ///
905 /// Returns `Some(StepOutput)` if a completed step exists at the given
906 /// position, `None` otherwise.
907 fn try_replay_step(&mut self, position: u32) -> Option<StepOutput> {
908 let step = self.replay_steps.get(&position)?;
909 if step.status.state != StepStatus::Completed {
910 return None;
911 }
912 let output = StepOutput {
913 output: step.output.clone().unwrap_or(Value::Null),
914 duration_ms: step.duration_ms,
915 cost_usd: step.cost_usd,
916 input_tokens: step.input_tokens,
917 output_tokens: step.output_tokens,
918 debug_messages: None,
919 };
920 self.total_cost_usd += output.cost_usd;
921 self.total_duration_ms += output.duration_ms;
922 self.last_step_ids = vec![step.id];
923 info!(
924 run_id = %self.run_id,
925 step = %step.name,
926 position,
927 "step replayed from previous execution"
928 );
929 Some(output)
930 }
931
932 /// Internal: execute a step with full persistence lifecycle.
933 async fn execute_step(
934 &mut self,
935 name: &str,
936 kind: StepKind,
937 config: StepConfig,
938 ) -> Result<StepOutput, EngineError> {
939 let position = self.position;
940 self.position += 1;
941
942 // Replay: if this step already completed in a prior execution, return cached output.
943 if let Some(output) = self.try_replay_step(position) {
944 return Ok(output);
945 }
946
947 // Create step record in Pending.
948 let step = self
949 .store
950 .create_step(NewStep {
951 run_id: self.run_id,
952 name: name.to_string(),
953 kind,
954 position,
955 input: Some(serde_json::to_value(&config)?),
956 })
957 .await?;
958
959 self.start_step(step.id, Utc::now()).await?;
960
961 match execute_step_config(&config, &self.provider).await {
962 Ok(output) => {
963 self.total_cost_usd += output.cost_usd;
964 self.total_duration_ms += output.duration_ms;
965
966 let debug_messages_json = output.debug_messages_json();
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 input_tokens: output.input_tokens,
978 output_tokens: output.output_tokens,
979 completed_at: Some(completed_at),
980 debug_messages: debug_messages_json,
981 ..StepUpdate::default()
982 },
983 )
984 .await?;
985
986 info!(
987 run_id = %self.run_id,
988 step = %name,
989 duration_ms = output.duration_ms,
990 "step completed"
991 );
992
993 self.last_step_ids = vec![step.id];
994
995 Ok(output)
996 }
997 Err(err) => {
998 let completed_at = Utc::now();
999 let debug_messages_json = extract_debug_messages_from_error(&err);
1000 let partial = extract_partial_usage_from_error(&err);
1001
1002 if let Some(ref usage) = partial {
1003 if let Some(cost) = usage.cost_usd {
1004 self.total_cost_usd += cost;
1005 }
1006 if let Some(dur) = usage.duration_ms {
1007 self.total_duration_ms += dur;
1008 }
1009 }
1010
1011 if let Err(store_err) = self
1012 .store
1013 .update_step(
1014 step.id,
1015 StepUpdate {
1016 status: Some(StepStatus::Failed),
1017 error: Some(err.to_string()),
1018 completed_at: Some(completed_at),
1019 debug_messages: debug_messages_json,
1020 duration_ms: partial.as_ref().and_then(|p| p.duration_ms),
1021 cost_usd: partial.as_ref().and_then(|p| p.cost_usd),
1022 input_tokens: partial.as_ref().and_then(|p| p.input_tokens),
1023 output_tokens: partial.as_ref().and_then(|p| p.output_tokens),
1024 ..StepUpdate::default()
1025 },
1026 )
1027 .await
1028 {
1029 tracing::error!(step_id = %step.id, error = %store_err, "failed to persist step failure");
1030 }
1031
1032 Err(err)
1033 }
1034 }
1035 }
1036
1037 /// Record dependency edges and transition a step to Running.
1038 ///
1039 /// Records edges from `step_id` to all `last_step_ids`, then
1040 /// transitions the step to `Running` with the given timestamp.
1041 async fn start_step(&self, step_id: Uuid, now: DateTime<Utc>) -> Result<(), EngineError> {
1042 if !self.last_step_ids.is_empty() {
1043 let deps: Vec<NewStepDependency> = self
1044 .last_step_ids
1045 .iter()
1046 .map(|&depends_on| NewStepDependency {
1047 step_id,
1048 depends_on,
1049 })
1050 .collect();
1051 self.store.create_step_dependencies(deps).await?;
1052 }
1053
1054 self.store
1055 .update_step(
1056 step_id,
1057 StepUpdate {
1058 status: Some(StepStatus::Running),
1059 started_at: Some(now),
1060 ..StepUpdate::default()
1061 },
1062 )
1063 .await?;
1064
1065 Ok(())
1066 }
1067
1068 /// Access the store directly (advanced usage).
1069 pub fn store(&self) -> &Arc<dyn RunStore> {
1070 &self.store
1071 }
1072
1073 /// Access the payload that triggered this run.
1074 ///
1075 /// Fetches the run from the store and returns its payload.
1076 ///
1077 /// # Errors
1078 ///
1079 /// Returns [`EngineError::Store`] if the run is not found.
1080 pub async fn payload(&self) -> Result<Value, EngineError> {
1081 let run = self
1082 .store
1083 .get_run(self.run_id)
1084 .await?
1085 .ok_or(EngineError::Store(
1086 ironflow_store::error::StoreError::RunNotFound(self.run_id),
1087 ))?;
1088 Ok(run.payload)
1089 }
1090}
1091
1092impl fmt::Debug for WorkflowContext {
1093 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1094 f.debug_struct("WorkflowContext")
1095 .field("run_id", &self.run_id)
1096 .field("position", &self.position)
1097 .field("total_cost_usd", &self.total_cost_usd)
1098 .finish_non_exhaustive()
1099 }
1100}
1101
1102/// Extract debug messages from an engine error, if it wraps a schema validation
1103/// failure that carries a verbose conversation trace.
1104fn extract_debug_messages_from_error(err: &EngineError) -> Option<Value> {
1105 if let EngineError::Operation(OperationError::Agent(AgentError::SchemaValidation {
1106 debug_messages,
1107 ..
1108 })) = err
1109 && !debug_messages.is_empty()
1110 {
1111 return serde_json::to_value(debug_messages).ok();
1112 }
1113 None
1114}
1115
1116/// Partial usage with `Decimal` cost, converted from the `f64` in [`PartialUsage`].
1117///
1118/// Exists only because `ironflow-store` uses [`Decimal`] for monetary values
1119/// while `ironflow-core` uses `f64` (the CLI's native type). The conversion
1120/// happens here, at the engine/store boundary.
1121struct StepPartialUsage {
1122 cost_usd: Option<Decimal>,
1123 duration_ms: Option<u64>,
1124 input_tokens: Option<u64>,
1125 output_tokens: Option<u64>,
1126}
1127
1128fn extract_partial_usage_from_error(err: &EngineError) -> Option<StepPartialUsage> {
1129 if let EngineError::Operation(OperationError::Agent(AgentError::SchemaValidation {
1130 partial_usage,
1131 ..
1132 })) = err
1133 && (partial_usage.cost_usd.is_some() || partial_usage.duration_ms.is_some())
1134 {
1135 return Some(StepPartialUsage {
1136 cost_usd: partial_usage
1137 .cost_usd
1138 .and_then(|c| Decimal::try_from(c).ok()),
1139 duration_ms: partial_usage.duration_ms,
1140 input_tokens: partial_usage.input_tokens,
1141 output_tokens: partial_usage.output_tokens,
1142 });
1143 }
1144 None
1145}