ironflow_engine/context/steps/operation.rs
1//! Custom operation step for [`WorkflowContext`].
2
3use std::time::Instant;
4
5use chrono::Utc;
6use rust_decimal::Decimal;
7use tracing::{error, info};
8
9use ironflow_store::models::{NewStep, StepKind, StepStatus, StepUpdate, step_trace_id};
10
11use crate::context::WorkflowContext;
12use crate::context::failure::allowed_failure_output;
13use crate::error::EngineError;
14use crate::executor::{StepArtifacts, StepOutput};
15use crate::operation::Operation;
16use crate::plan::{lock_plan, planned_custom_output};
17
18impl WorkflowContext {
19 /// Execute a custom operation step.
20 ///
21 /// Runs a user-defined [`Operation`] with full step lifecycle management:
22 /// creates the step record, transitions to Running, executes the operation,
23 /// persists the output and duration, and marks the step Completed or Failed.
24 ///
25 /// The operation's [`kind()`](Operation::kind) is stored as
26 /// [`StepKind::Custom`].
27 ///
28 /// On resume, a step that already completed in a prior execution of the
29 /// current attempt is replayed from the store instead of calling
30 /// [`Operation::execute`] again.
31 ///
32 /// # Errors
33 ///
34 /// Returns [`EngineError`] if the operation fails or the store errors.
35 /// Returns [`EngineError::ReplayDivergence`] when the step recorded at
36 /// this position has a different name or kind.
37 ///
38 /// # Examples
39 ///
40 /// ```no_run
41 /// use async_trait::async_trait;
42 /// use ironflow_engine::context::WorkflowContext;
43 /// use ironflow_engine::operation::{Operation, OperationContext};
44 /// use ironflow_core::error::OperationError;
45 /// use ironflow_engine::error::EngineError;
46 /// use serde_json::{Value, json};
47 ///
48 /// struct MyOp;
49 /// #[async_trait]
50 /// impl Operation for MyOp {
51 /// fn kind(&self) -> &str { "my-service" }
52 /// async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
53 /// Ok(json!({"ok": true}))
54 /// }
55 /// }
56 ///
57 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
58 /// let result = ctx.operation("call-service", &MyOp).await?;
59 /// println!("output: {}", result.output);
60 /// # Ok(())
61 /// # }
62 /// ```
63 pub async fn operation(
64 &mut self,
65 name: &str,
66 op: &dyn Operation,
67 ) -> Result<StepOutput, EngineError> {
68 self.run_operation(name, op, false).await
69 }
70
71 /// Execute a custom operation step whose failure does not stop the run.
72 ///
73 /// Same as [`operation`](Self::operation), except when
74 /// [`Operation::execute`] fails: the step is recorded as Failed, the run
75 /// ends as `Warning` instead of `Completed`, and the handler receives an
76 /// output whose [`error()`](StepOutput::error) is the failure. On resume,
77 /// the tolerated failure is replayed with that output, not executed again.
78 ///
79 /// # Errors
80 ///
81 /// Returns [`EngineError`] if the store errors, or
82 /// [`EngineError::ReplayDivergence`] when the step recorded at this
83 /// position has a different name or kind. Never for a failure of the
84 /// operation itself.
85 ///
86 /// # Examples
87 ///
88 /// ```no_run
89 /// use ironflow_engine::context::WorkflowContext;
90 /// use ironflow_engine::error::EngineError;
91 /// use ironflow_engine::operation::Operation;
92 ///
93 /// # async fn example(ctx: &mut WorkflowContext, op: &dyn Operation) -> Result<(), EngineError> {
94 /// let result = ctx.operation_allow_failure("regenerate", op).await?;
95 /// if let Some(error) = result.error() {
96 /// println!("regeneration failed, carrying on: {error}");
97 /// }
98 /// # Ok(())
99 /// # }
100 /// ```
101 pub async fn operation_allow_failure(
102 &mut self,
103 name: &str,
104 op: &dyn Operation,
105 ) -> Result<StepOutput, EngineError> {
106 self.run_operation(name, op, true).await
107 }
108
109 async fn run_operation(
110 &mut self,
111 name: &str,
112 op: &dyn Operation,
113 allow_failure: bool,
114 ) -> Result<StepOutput, EngineError> {
115 let kind = StepKind::Custom(op.kind().to_string());
116
117 // Plan mode: `op.execute` is never called, so no third-party API is
118 // touched while planning.
119 if let Some(plan) = self.plan().cloned() {
120 self.position += 1;
121 let mut recorder = lock_plan(&plan);
122 let estimate = recorder.estimate_for(name);
123 recorder.record_or_stop(name, kind, &self.workflow_name, None)?;
124 recorder.set_last(vec![name.to_string()]);
125 let mut output = planned_custom_output(estimate);
126 output.artifacts = StepArtifacts::new(name, None, &[]);
127 return Ok(output);
128 }
129
130 let position = self.position;
131 self.position += 1;
132
133 // Replay: if this step already completed, or failed while tolerated,
134 // in a prior execution of the current attempt, return its recorded
135 // output without calling `op.execute` or creating a new step.
136 if let Some(mut output) = self.try_replay_step(position, name, &kind, allow_failure)? {
137 let step_id = self.last_step_ids.last().copied();
138 output.artifacts = StepArtifacts::new(name, step_id, &[]);
139 return Ok(output);
140 }
141
142 let trace_id = step_trace_id(self.run_id, name, position);
143 let step = self
144 .store
145 .create_step(NewStep {
146 run_id: self.run_id,
147 trace_id,
148 name: name.to_string(),
149 kind,
150 position,
151 input: op.input(),
152 is_error_handler: false,
153 })
154 .await?;
155
156 self.start_step(step.id, Utc::now()).await?;
157
158 let start = Instant::now();
159
160 let op_ctx = self.ensure_operation_ctx();
161
162 match op.execute(op_ctx).await {
163 Ok(output_value) => {
164 let duration_ms = start.elapsed().as_millis() as u64;
165 self.total_duration_ms += duration_ms;
166
167 let completed_at = Utc::now();
168 let step_id = self
169 .complete_step(
170 &step,
171 StepUpdate {
172 status: Some(StepStatus::Completed),
173 output: Some(output_value.clone()),
174 duration_ms: Some(duration_ms),
175 cost_usd: Some(Decimal::ZERO),
176 completed_at: Some(completed_at),
177 ..StepUpdate::default()
178 },
179 )
180 .await?;
181
182 info!(
183 run_id = %self.run_id,
184 step = %name,
185 kind = op.kind(),
186 duration_ms,
187 "operation step completed"
188 );
189
190 self.last_step_ids = vec![step_id];
191
192 Ok(StepOutput {
193 output: output_value,
194 duration_ms,
195 cost_usd: Decimal::ZERO,
196 input_tokens: None,
197 cache_read_input_tokens: None,
198 cache_creation_input_tokens: None,
199 output_tokens: None,
200 model: None,
201 debug_messages: None,
202 // An operation declares no output; its record is what
203 // `put_artifact` attaches bytes to.
204 artifacts: StepArtifacts::new(name, Some(step_id), &[]),
205 account_id: None,
206 environment_id: None,
207 })
208 }
209 Err(err) => {
210 let completed_at = Utc::now();
211 let engine_err = EngineError::Operation(err);
212 if let Err(store_err) = self
213 .store
214 .update_step(
215 step.id,
216 StepUpdate {
217 status: Some(StepStatus::Failed),
218 error: Some(engine_err.to_string()),
219 completed_at: Some(completed_at),
220 ..StepUpdate::default()
221 },
222 )
223 .await
224 {
225 error!(step_id = %step.id, error = %store_err, "failed to persist step failure");
226 }
227
228 if !allow_failure {
229 return Err(engine_err);
230 }
231 self.has_allowed_failure = true;
232 self.last_step_ids = vec![step.id];
233 info!(
234 run_id = %self.run_id,
235 step = %name,
236 kind = op.kind(),
237 error = %engine_err,
238 "operation step failed but allow_failure is set, continuing"
239 );
240 let mut output = allowed_failure_output(&engine_err.to_string(), None, None);
241 output.duration_ms = start.elapsed().as_millis() as u64;
242 self.total_duration_ms += output.duration_ms;
243 output.artifacts = StepArtifacts::new(name, Some(step.id), &[]);
244 Ok(output)
245 }
246 }
247 }
248}