ironflow_engine/executor/
agent.rs

1//! Agent step executor.
2
3use std::sync::Arc;
4use std::time::Instant;
5
6use rust_decimal::Decimal;
7use tracing::{info, warn};
8use uuid::Uuid;
9
10use ironflow_core::error::OperationError;
11use ironflow_core::operations::agent::{Agent, AgentResult};
12use ironflow_core::pricing::{CostBreakdown, StaticPricing, spawn_log};
13use ironflow_core::provider::{AgentConfig, AgentProvider, LogSink};
14use ironflow_core::providers::claude::is_session_not_found;
15use ironflow_store::entities::StepKind;
16
17use crate::error::EngineError;
18use crate::log_sender::StepLogSender;
19use crate::notify::LogStream;
20
21use super::{StepArtifacts, StepExecutor, StepOutput};
22
23/// Executor for agent (AI) steps.
24///
25/// Runs an AI agent with the given prompt and configuration, capturing
26/// the response value, cost, and token counts. When a [`StepLogSender`]
27/// is attached, emits system log lines for step start/end.
28pub struct AgentExecutor<'a> {
29    config: &'a AgentConfig,
30    log_sender: Option<StepLogSender>,
31}
32
33impl<'a> AgentExecutor<'a> {
34    /// Create a new agent executor from a config reference.
35    pub fn new(config: &'a AgentConfig) -> Self {
36        Self {
37            config,
38            log_sender: None,
39        }
40    }
41
42    /// Attach a log sender for system-level log lines.
43    pub fn with_log_sender(mut self, sender: StepLogSender) -> Self {
44        self.log_sender = Some(sender);
45        self
46    }
47
48    fn emit_system(&self, line: &str) {
49        if let Some(ref sender) = self.log_sender {
50            sender.emit(LogStream::System, line);
51        }
52    }
53
54    async fn run_agent(
55        &self,
56        config: AgentConfig,
57        provider: &Arc<dyn AgentProvider>,
58    ) -> Result<AgentResult, OperationError> {
59        let mut agent = Agent::from_config(config);
60        if let Some(ref sender) = self.log_sender {
61            agent = agent.log_sink(Arc::new(sender.clone()) as Arc<dyn LogSink>);
62        }
63        agent.run(provider.as_ref()).await
64    }
65
66    /// Resume `session_id` with `resume_prompt`. When the session does not
67    /// exist anymore (ephemeral HOME, other machine, other `cwd`), run the
68    /// original prompt from scratch in a session of the same id: the step
69    /// never fails because its session is gone.
70    async fn resume(
71        &self,
72        provider: &Arc<dyn AgentProvider>,
73        session_id: &str,
74        resume_prompt: &str,
75    ) -> Result<AgentResult, OperationError> {
76        info!(session_id, "agent step resumed from session");
77        self.emit_system(&format!("agent step resumed from session {session_id}"));
78
79        let mut resumed = self.config.clone();
80        resumed.prompt = resume_prompt.to_string();
81        match self.run_agent(resumed, provider).await {
82            Err(OperationError::Agent(ref err)) if is_session_not_found(err) => {
83                warn!(
84                    session_id,
85                    error = %err,
86                    "session not found, restarting the agent from scratch"
87                );
88                self.emit_system(&format!(
89                    "session {session_id} not found, restarting the agent from scratch"
90                ));
91                let mut fresh = self.config.clone();
92                fresh.resume_session_id = None;
93                fresh.session_id = Some(session_id.to_string());
94                self.run_agent(fresh, provider).await
95            }
96            other => other,
97        }
98    }
99}
100
101impl StepExecutor for AgentExecutor<'_> {
102    fn kind(&self) -> StepKind {
103        StepKind::Agent
104    }
105
106    async fn execute(&self, provider: &Arc<dyn AgentProvider>) -> Result<StepOutput, EngineError> {
107        let start = Instant::now();
108
109        if let Some(ref sender) = self.log_sender {
110            let effort = self
111                .config
112                .effort
113                .map(|e| format!(", effort={e}"))
114                .unwrap_or_default();
115            sender.emit(
116                LogStream::System,
117                &format!("agent step started (model={}{})", self.config.model, effort),
118            );
119        }
120
121        if self.config.json_schema.is_some() && self.config.max_turns == Some(1) {
122            warn!(
123                "structured output (json_schema) requires max_turns >= 2; \
124                 max_turns is set to 1, the agent will likely fail with error_max_turns"
125            );
126        }
127
128        let result = match (&self.config.resume_session_id, &self.config.resume_prompt) {
129            (Some(session_id), Some(resume_prompt)) => {
130                self.resume(provider, session_id, resume_prompt).await?
131            }
132            _ => self.run_agent(self.config.clone(), provider).await?,
133        };
134
135        let duration_ms = start.elapsed().as_millis() as u64;
136        let cost = Decimal::try_from(result.cost_usd().unwrap_or(0.0)).unwrap_or(Decimal::ZERO);
137        let input_tokens = result.input_tokens();
138        let cache_read_tokens = result.cache_read_input_tokens();
139        let cache_creation_tokens = result.cache_creation_input_tokens();
140        let output_tokens = result.output_tokens();
141
142        info!(
143            step_kind = "agent",
144            model = %self.config.model,
145            effort = ?self.config.effort,
146            cost_usd = %cost,
147            input_tokens = ?input_tokens,
148            cache_read_input_tokens = ?cache_read_tokens,
149            cache_creation_input_tokens = ?cache_creation_tokens,
150            output_tokens = ?output_tokens,
151            duration_ms,
152            "agent step completed"
153        );
154
155        let pricing = StaticPricing::new();
156        let breakdown = CostBreakdown::compute_with_cache(
157            &pricing,
158            &self.config.model,
159            input_tokens.unwrap_or(0),
160            cache_read_tokens.unwrap_or(0),
161            cache_creation_tokens.unwrap_or(0),
162            output_tokens.unwrap_or(0),
163        );
164        spawn_log("agent", &self.config.model, breakdown);
165
166        #[cfg(feature = "prometheus")]
167        {
168            use ironflow_core::metric_names::{
169                AGENT_COST_USD_TOTAL, AGENT_DURATION_SECONDS, AGENT_TOKENS_CACHE_READ_TOTAL,
170                AGENT_TOKENS_CACHE_WRITE_TOTAL, AGENT_TOKENS_INPUT_TOTAL,
171                AGENT_TOKENS_OUTPUT_TOTAL, AGENT_TOTAL, STATUS_SUCCESS,
172            };
173            use metrics::{counter, gauge, histogram};
174            let model_label = self.config.model.clone();
175            counter!(AGENT_TOTAL, "model" => model_label.clone(), "status" => STATUS_SUCCESS)
176                .increment(1);
177            histogram!(AGENT_DURATION_SECONDS, "model" => model_label.clone())
178                .record(duration_ms as f64 / 1000.0);
179            gauge!(AGENT_COST_USD_TOTAL, "model" => model_label.clone())
180                .increment(cost.to_string().parse::<f64>().unwrap_or(0.0));
181            if let Some(inp) = input_tokens {
182                counter!(AGENT_TOKENS_INPUT_TOTAL, "model" => model_label.clone()).increment(inp);
183            }
184            if let Some(t) = cache_read_tokens {
185                counter!(AGENT_TOKENS_CACHE_READ_TOTAL, "model" => model_label.clone())
186                    .increment(t);
187            }
188            if let Some(t) = cache_creation_tokens {
189                counter!(AGENT_TOKENS_CACHE_WRITE_TOTAL, "model" => model_label.clone())
190                    .increment(t);
191            }
192            if let Some(out) = output_tokens {
193                counter!(AGENT_TOKENS_OUTPUT_TOTAL, "model" => model_label).increment(out);
194            }
195        }
196
197        if let Some(ref sender) = self.log_sender {
198            sender.emit(
199                LogStream::System,
200                &format!(
201                    "agent step completed (cost=${cost}, tokens_in={}, cache_read={}, cache_write={}, tokens_out={})",
202                    input_tokens.unwrap_or(0),
203                    cache_read_tokens.unwrap_or(0),
204                    cache_creation_tokens.unwrap_or(0),
205                    output_tokens.unwrap_or(0),
206                ),
207            );
208        }
209
210        let debug_messages = result.debug_messages().map(|msgs| msgs.to_vec());
211        let account_id = match result.account_id() {
212            Some(raw) => match Uuid::parse_str(raw) {
213                Ok(id) => Some(id),
214                Err(e) => {
215                    warn!(account_id = raw, error = %e, "agent output carries an invalid account id");
216                    None
217                }
218            },
219            None => None,
220        };
221
222        Ok(StepOutput {
223            output: result.value().clone(),
224            duration_ms,
225            cost_usd: cost,
226            input_tokens,
227            cache_read_input_tokens: cache_read_tokens,
228            cache_creation_input_tokens: cache_creation_tokens,
229            output_tokens,
230            model: result.model().map(String::from),
231            debug_messages,
232            artifacts: StepArtifacts::default(),
233            account_id,
234            environment_id: result.environment_id().map(String::from),
235        })
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use std::sync::{Arc, Mutex};
242    use std::time::Duration;
243
244    use ironflow_core::error::AgentError;
245    use ironflow_core::operations::agent::PermissionMode;
246    use ironflow_core::provider::{AgentConfig, AgentOutput, AgentProvider, InvokeFuture};
247    use serde_json::json;
248    use tokio::time::timeout;
249    use uuid::Uuid;
250
251    use super::{AgentExecutor, StepExecutor};
252
253    /// Real provider returning a fixed output, used to exercise usage propagation.
254    struct FixedUsageProvider {
255        output: AgentOutput,
256    }
257
258    impl AgentProvider for FixedUsageProvider {
259        fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
260            Box::pin(async move { Ok(self.output.clone()) })
261        }
262    }
263
264    /// Provider recording every config it receives, failing a resume with
265    /// `resume_error` when set.
266    struct SessionProvider {
267        seen: Mutex<Vec<AgentConfig>>,
268        resume_error: Option<String>,
269    }
270
271    impl SessionProvider {
272        fn new(resume_error: Option<&str>) -> Arc<Self> {
273            Arc::new(Self {
274                seen: Mutex::new(Vec::new()),
275                resume_error: resume_error.map(String::from),
276            })
277        }
278
279        fn seen(&self) -> Vec<AgentConfig> {
280            self.seen.lock().unwrap().clone()
281        }
282    }
283
284    impl AgentProvider for SessionProvider {
285        fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a> {
286            Box::pin(async move {
287                self.seen.lock().unwrap().push(config.clone());
288                match (&config.resume_session_id, &self.resume_error) {
289                    (Some(_), Some(stderr)) => Err(AgentError::ProcessFailed {
290                        exit_code: 1,
291                        stderr: stderr.clone(),
292                    }),
293                    _ => {
294                        let mut output = AgentOutput::new(json!("ok"));
295                        output.cost_usd = Some(0.02);
296                        Ok(output)
297                    }
298                }
299            })
300        }
301    }
302
303    const SID: &str = "0192f0c1-7d2e-7a4b-9c3d-1e2f3a4b5c6d";
304
305    #[tokio::test]
306    async fn agent_resume_executor_sends_resume_prompt() {
307        timeout(Duration::from_secs(10), async {
308            let recorder = SessionProvider::new(None);
309            let provider: Arc<dyn AgentProvider> = recorder.clone();
310            let config = budget_config().resume(SID).resume_prompt("go on");
311
312            AgentExecutor::new(&config)
313                .execute(&provider)
314                .await
315                .expect("resumed step succeeds");
316
317            let seen = recorder.seen();
318            assert_eq!(seen.len(), 1);
319            assert_eq!(seen[0].prompt, "go on");
320            assert_eq!(seen[0].resume_session_id.as_deref(), Some(SID));
321        })
322        .await
323        .expect("test timed out");
324    }
325
326    #[tokio::test]
327    async fn agent_resume_executor_falls_back_when_session_is_missing() {
328        timeout(Duration::from_secs(10), async {
329            let recorder =
330                SessionProvider::new(Some("No conversation found with session ID: 0192f0c1"));
331            let provider: Arc<dyn AgentProvider> = recorder.clone();
332            let config = budget_config().resume(SID).resume_prompt("go on");
333
334            AgentExecutor::new(&config)
335                .execute(&provider)
336                .await
337                .expect("a missing session never fails the step");
338
339            let seen = recorder.seen();
340            assert_eq!(seen.len(), 2);
341            assert_eq!(seen[1].prompt, "hi");
342            assert_eq!(seen[1].resume_session_id, None);
343            assert_eq!(seen[1].session_id.as_deref(), Some(SID));
344        })
345        .await
346        .expect("test timed out");
347    }
348
349    #[tokio::test]
350    async fn agent_resume_executor_propagates_other_errors() {
351        timeout(Duration::from_secs(10), async {
352            let recorder = SessionProvider::new(Some("permission denied"));
353            let provider: Arc<dyn AgentProvider> = recorder.clone();
354            let config = budget_config().resume(SID).resume_prompt("go on");
355
356            let err = AgentExecutor::new(&config)
357                .execute(&provider)
358                .await
359                .expect_err("other errors are not masked");
360
361            assert!(err.to_string().contains("permission denied"), "{err}");
362            assert_eq!(recorder.seen().len(), 1);
363        })
364        .await
365        .expect("test timed out");
366    }
367
368    #[tokio::test]
369    async fn agent_resume_executor_leaves_user_resume_alone() {
370        timeout(Duration::from_secs(10), async {
371            let recorder = SessionProvider::new(None);
372            let provider: Arc<dyn AgentProvider> = recorder.clone();
373            // A user-set resume without resume prompt keeps the original prompt.
374            let config = budget_config().resume(SID);
375
376            AgentExecutor::new(&config)
377                .execute(&provider)
378                .await
379                .expect("step succeeds");
380
381            let seen = recorder.seen();
382            assert_eq!(seen.len(), 1);
383            assert_eq!(seen[0].prompt, "hi");
384            assert_eq!(seen[0].resume_session_id.as_deref(), Some(SID));
385        })
386        .await
387        .expect("test timed out");
388    }
389
390    fn budget_config() -> AgentConfig {
391        let mut config = AgentConfig::new("hi");
392        config.max_budget_usd = Some(0.10);
393        config
394    }
395
396    #[tokio::test]
397    async fn agent_executor_propagates_cache_tokens() {
398        timeout(Duration::from_secs(10), async {
399            let mut output = AgentOutput::new(json!("ok"));
400            output.input_tokens = Some(100);
401            output.cache_read_input_tokens = Some(5000);
402            output.cache_creation_input_tokens = Some(200);
403            output.output_tokens = Some(50);
404            output.cost_usd = Some(0.02);
405            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
406
407            let config = budget_config();
408            let step = AgentExecutor::new(&config)
409                .execute(&provider)
410                .await
411                .expect("agent step succeeds");
412
413            assert_eq!(step.input_tokens, Some(100));
414            assert_eq!(step.cache_read_input_tokens, Some(5000));
415            assert_eq!(step.cache_creation_input_tokens, Some(200));
416            assert_eq!(step.output_tokens, Some(50));
417            assert_eq!(step.total_tokens(), 5350);
418        })
419        .await
420        .expect("test timed out");
421    }
422
423    #[tokio::test]
424    async fn agent_executor_propagates_account_id() {
425        timeout(Duration::from_secs(10), async {
426            let account_id = Uuid::now_v7();
427            let mut output = AgentOutput::new(json!("ok"));
428            output.cost_usd = Some(0.02);
429            output.account_id = Some(account_id.to_string());
430            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
431
432            let step = AgentExecutor::new(&budget_config())
433                .execute(&provider)
434                .await
435                .expect("agent step succeeds");
436            assert_eq!(step.account_id, Some(account_id));
437
438            let mut invalid = AgentOutput::new(json!("ok"));
439            invalid.cost_usd = Some(0.02);
440            invalid.account_id = Some("not-a-uuid".to_string());
441            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output: invalid });
442            let step = AgentExecutor::new(&budget_config())
443                .execute(&provider)
444                .await
445                .expect("agent step succeeds");
446            assert_eq!(step.account_id, None);
447        })
448        .await
449        .expect("test timed out");
450    }
451
452    #[tokio::test]
453    async fn agent_executor_propagates_environment_id() {
454        timeout(Duration::from_secs(10), async {
455            let mut output = AgentOutput::new(json!("ok"));
456            output.cost_usd = Some(0.02);
457            output.environment_id = Some("ironflow-env-0a1b2c".to_string());
458            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
459            let step = AgentExecutor::new(&budget_config())
460                .execute(&provider)
461                .await
462                .expect("agent step succeeds");
463            assert_eq!(step.environment_id.as_deref(), Some("ironflow-env-0a1b2c"));
464
465            let mut output = AgentOutput::new(json!("ok"));
466            output.cost_usd = Some(0.02);
467            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
468            let step = AgentExecutor::new(&budget_config())
469                .execute(&provider)
470                .await
471                .expect("agent step succeeds");
472            assert_eq!(step.environment_id, None);
473        })
474        .await
475        .expect("test timed out");
476    }
477
478    #[tokio::test]
479    async fn agent_executor_without_cache_tokens_yields_none() {
480        timeout(Duration::from_secs(10), async {
481            let mut output = AgentOutput::new(json!("ok"));
482            output.input_tokens = Some(100);
483            output.output_tokens = Some(50);
484            output.cost_usd = Some(0.02);
485            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
486
487            let config = budget_config();
488            let step = AgentExecutor::new(&config)
489                .execute(&provider)
490                .await
491                .expect("agent step succeeds");
492
493            assert_eq!(step.input_tokens, Some(100));
494            assert!(step.cache_read_input_tokens.is_none());
495            assert!(step.cache_creation_input_tokens.is_none());
496            assert_eq!(step.total_tokens(), 150);
497        })
498        .await
499        .expect("test timed out");
500    }
501
502    #[test]
503    fn parse_permission_mode_via_serde() {
504        let json = r#""auto""#;
505        let mode: PermissionMode = serde_json::from_str(json).unwrap();
506        assert!(matches!(mode, PermissionMode::Auto));
507    }
508
509    #[test]
510    fn parse_permission_mode_dont_ask() {
511        let json = r#""dont_ask""#;
512        let mode: PermissionMode = serde_json::from_str(json).unwrap();
513        assert!(matches!(mode, PermissionMode::DontAsk));
514    }
515
516    #[test]
517    fn parse_permission_mode_bypass() {
518        let json = r#""bypass""#;
519        let mode: PermissionMode = serde_json::from_str(json).unwrap();
520        assert!(matches!(mode, PermissionMode::BypassPermissions));
521    }
522
523    #[test]
524    fn parse_permission_mode_case_insensitive() {
525        let json = r#""AUTO""#;
526        let mode: PermissionMode = serde_json::from_str(json).unwrap();
527        assert!(matches!(mode, PermissionMode::Auto));
528
529        let json = r#""DONT_ASK""#;
530        let mode: PermissionMode = serde_json::from_str(json).unwrap();
531        assert!(matches!(mode, PermissionMode::DontAsk));
532    }
533
534    #[test]
535    fn parse_permission_mode_unknown_defaults() {
536        let json = r#""unknown""#;
537        let mode: PermissionMode = serde_json::from_str(json).unwrap();
538        assert!(matches!(mode, PermissionMode::Default));
539    }
540}
ironflow_engine/executor/agent.rs

ironflow_engine/executor/
agent.rs
