1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
//! Custom operation step for [`WorkflowContext`].
use std::time::Instant;
use chrono::Utc;
use rust_decimal::Decimal;
use tracing::{error, info};
use ironflow_store::models::{NewStep, StepKind, StepStatus, StepUpdate, step_trace_id};
use crate::context::WorkflowContext;
use crate::error::EngineError;
use crate::executor::{StepArtifacts, StepOutput};
use crate::operation::Operation;
use crate::plan::{lock_plan, planned_custom_output};
impl WorkflowContext {
/// Execute a custom operation step.
///
/// Runs a user-defined [`Operation`] with full step lifecycle management:
/// creates the step record, transitions to Running, executes the operation,
/// persists the output and duration, and marks the step Completed or Failed.
///
/// The operation's [`kind()`](Operation::kind) is stored as
/// [`StepKind::Custom`].
///
/// On resume, a step that already completed in a prior execution of the
/// current attempt is replayed from the store instead of calling
/// [`Operation::execute`] again.
///
/// # Errors
///
/// Returns [`EngineError`] if the operation fails or the store errors.
///
/// # Examples
///
/// ```no_run
/// use async_trait::async_trait;
/// use ironflow_engine::context::WorkflowContext;
/// use ironflow_engine::operation::{Operation, OperationContext};
/// use ironflow_core::error::OperationError;
/// use ironflow_engine::error::EngineError;
/// use serde_json::{Value, json};
///
/// struct MyOp;
/// #[async_trait]
/// impl Operation for MyOp {
/// fn kind(&self) -> &str { "my-service" }
/// async fn execute(&self, _ctx: &OperationContext) -> Result<Value, OperationError> {
/// Ok(json!({"ok": true}))
/// }
/// }
///
/// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
/// let result = ctx.operation("call-service", &MyOp).await?;
/// println!("output: {}", result.output);
/// # Ok(())
/// # }
/// ```
pub async fn operation(
&mut self,
name: &str,
op: &dyn Operation,
) -> Result<StepOutput, EngineError> {
let kind = StepKind::Custom(op.kind().to_string());
// Plan mode: `op.execute` is never called, so no third-party API is
// touched while planning.
if let Some(plan) = self.plan().cloned() {
self.position += 1;
let mut recorder = lock_plan(&plan);
let estimate = recorder.estimate_for(name);
if recorder.record(name, kind, &self.workflow_name, None) {
recorder.set_last(vec![name.to_string()]);
}
let mut output = planned_custom_output(estimate);
output.artifacts = StepArtifacts::new(name, None, &[]);
return Ok(output);
}
let position = self.position;
self.position += 1;
// Replay: if this step already completed in a prior execution of the
// current attempt, return its cached output without calling
// `op.execute` or creating a new step.
if let Some(mut output) = self.try_replay_step(position) {
let step_id = self.last_step_ids.last().copied();
output.artifacts = StepArtifacts::new(name, step_id, &[]);
return Ok(output);
}
let trace_id = step_trace_id(self.run_id, name, position);
let step = self
.store
.create_step(NewStep {
run_id: self.run_id,
trace_id,
name: name.to_string(),
kind,
position,
input: op.input(),
is_error_handler: false,
})
.await?;
self.start_step(step.id, Utc::now()).await?;
let start = Instant::now();
let op_ctx = self.ensure_operation_ctx();
match op.execute(op_ctx).await {
Ok(output_value) => {
let duration_ms = start.elapsed().as_millis() as u64;
self.total_duration_ms += duration_ms;
let completed_at = Utc::now();
self.store
.update_step(
step.id,
StepUpdate {
status: Some(StepStatus::Completed),
output: Some(output_value.clone()),
duration_ms: Some(duration_ms),
cost_usd: Some(Decimal::ZERO),
completed_at: Some(completed_at),
..StepUpdate::default()
},
)
.await?;
info!(
run_id = %self.run_id,
step = %name,
kind = op.kind(),
duration_ms,
"operation step completed"
);
self.last_step_ids = vec![step.id];
Ok(StepOutput {
output: output_value,
duration_ms,
cost_usd: Decimal::ZERO,
input_tokens: None,
cache_read_input_tokens: None,
cache_creation_input_tokens: None,
output_tokens: None,
model: None,
debug_messages: None,
// An operation declares no output; its record is what
// `put_artifact` attaches bytes to.
artifacts: StepArtifacts::new(name, Some(step.id), &[]),
})
}
Err(err) => {
let completed_at = Utc::now();
let engine_err = EngineError::Operation(err);
if let Err(store_err) = self
.store
.update_step(
step.id,
StepUpdate {
status: Some(StepStatus::Failed),
error: Some(engine_err.to_string()),
completed_at: Some(completed_at),
..StepUpdate::default()
},
)
.await
{
error!(step_id = %step.id, error = %store_err, "failed to persist step failure");
}
Err(engine_err)
}
}
}
}