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. This also covers a json_schema
70    /// agent whose resume fails at 0 turns.
71    async fn resume(
72        &self,
73        provider: &Arc<dyn AgentProvider>,
74        session_id: &str,
75        resume_prompt: &str,
76    ) -> Result<AgentResult, OperationError> {
77        info!(session_id, "agent step resumed from session");
78        self.emit_system(&format!("agent step resumed from session {session_id}"));
79
80        let mut resumed = self.config.clone();
81        resumed.prompt = resume_prompt.to_string();
82        match self.run_agent(resumed, provider).await {
83            Err(OperationError::Agent(ref err)) if is_session_not_found(err) => {
84                warn!(
85                    session_id,
86                    error = %err,
87                    "session not found, restarting the agent from scratch"
88                );
89                self.emit_system(&format!(
90                    "session {session_id} not found, restarting the agent from scratch"
91                ));
92                let mut fresh = self.config.clone();
93                fresh.resume_session_id = None;
94                fresh.session_id = Some(session_id.to_string());
95                self.run_agent(fresh, provider).await
96            }
97            other => other,
98        }
99    }
100}
101
102impl StepExecutor for AgentExecutor<'_> {
103    fn kind(&self) -> StepKind {
104        StepKind::Agent
105    }
106
107    async fn execute(&self, provider: &Arc<dyn AgentProvider>) -> Result<StepOutput, EngineError> {
108        let start = Instant::now();
109
110        if let Some(ref sender) = self.log_sender {
111            let effort = self
112                .config
113                .effort
114                .map(|e| format!(", effort={e}"))
115                .unwrap_or_default();
116            sender.emit(
117                LogStream::System,
118                &format!("agent step started (model={}{})", self.config.model, effort),
119            );
120        }
121
122        if self.config.json_schema.is_some() && self.config.max_turns == Some(1) {
123            warn!(
124                "structured output (json_schema) requires max_turns >= 2; \
125                 max_turns is set to 1, the agent will likely fail with error_max_turns"
126            );
127        }
128
129        let result = match (&self.config.resume_session_id, &self.config.resume_prompt) {
130            (Some(session_id), Some(resume_prompt)) => {
131                self.resume(provider, session_id, resume_prompt).await?
132            }
133            _ => self.run_agent(self.config.clone(), provider).await?,
134        };
135
136        let duration_ms = start.elapsed().as_millis() as u64;
137        let cost = Decimal::try_from(result.cost_usd().unwrap_or(0.0)).unwrap_or(Decimal::ZERO);
138        let input_tokens = result.input_tokens();
139        let cache_read_tokens = result.cache_read_input_tokens();
140        let cache_creation_tokens = result.cache_creation_input_tokens();
141        let output_tokens = result.output_tokens();
142
143        info!(
144            step_kind = "agent",
145            model = %self.config.model,
146            effort = ?self.config.effort,
147            cost_usd = %cost,
148            input_tokens = ?input_tokens,
149            cache_read_input_tokens = ?cache_read_tokens,
150            cache_creation_input_tokens = ?cache_creation_tokens,
151            output_tokens = ?output_tokens,
152            duration_ms,
153            "agent step completed"
154        );
155
156        let pricing = StaticPricing::new();
157        let breakdown = CostBreakdown::compute_with_cache(
158            &pricing,
159            &self.config.model,
160            input_tokens.unwrap_or(0),
161            cache_read_tokens.unwrap_or(0),
162            cache_creation_tokens.unwrap_or(0),
163            output_tokens.unwrap_or(0),
164        );
165        spawn_log("agent", &self.config.model, breakdown);
166
167        #[cfg(feature = "prometheus")]
168        {
169            use ironflow_core::metric_names::{
170                AGENT_COST_USD_TOTAL, AGENT_DURATION_SECONDS, AGENT_TOKENS_CACHE_READ_TOTAL,
171                AGENT_TOKENS_CACHE_WRITE_TOTAL, AGENT_TOKENS_INPUT_TOTAL,
172                AGENT_TOKENS_OUTPUT_TOTAL, AGENT_TOTAL, STATUS_SUCCESS,
173            };
174            use metrics::{counter, gauge, histogram};
175            let model_label = self.config.model.clone();
176            counter!(AGENT_TOTAL, "model" => model_label.clone(), "status" => STATUS_SUCCESS)
177                .increment(1);
178            histogram!(AGENT_DURATION_SECONDS, "model" => model_label.clone())
179                .record(duration_ms as f64 / 1000.0);
180            gauge!(AGENT_COST_USD_TOTAL, "model" => model_label.clone())
181                .increment(cost.to_string().parse::<f64>().unwrap_or(0.0));
182            if let Some(inp) = input_tokens {
183                counter!(AGENT_TOKENS_INPUT_TOTAL, "model" => model_label.clone()).increment(inp);
184            }
185            if let Some(t) = cache_read_tokens {
186                counter!(AGENT_TOKENS_CACHE_READ_TOTAL, "model" => model_label.clone())
187                    .increment(t);
188            }
189            if let Some(t) = cache_creation_tokens {
190                counter!(AGENT_TOKENS_CACHE_WRITE_TOTAL, "model" => model_label.clone())
191                    .increment(t);
192            }
193            if let Some(out) = output_tokens {
194                counter!(AGENT_TOKENS_OUTPUT_TOTAL, "model" => model_label).increment(out);
195            }
196        }
197
198        if let Some(ref sender) = self.log_sender {
199            sender.emit(
200                LogStream::System,
201                &format!(
202                    "agent step completed (cost=${cost}, tokens_in={}, cache_read={}, cache_write={}, tokens_out={})",
203                    input_tokens.unwrap_or(0),
204                    cache_read_tokens.unwrap_or(0),
205                    cache_creation_tokens.unwrap_or(0),
206                    output_tokens.unwrap_or(0),
207                ),
208            );
209        }
210
211        let debug_messages = result.debug_messages().map(|msgs| msgs.to_vec());
212        let account_id = match result.account_id() {
213            Some(raw) => match Uuid::parse_str(raw) {
214                Ok(id) => Some(id),
215                Err(e) => {
216                    warn!(account_id = raw, error = %e, "agent output carries an invalid account id");
217                    None
218                }
219            },
220            None => None,
221        };
222
223        Ok(StepOutput {
224            output: result.value().clone(),
225            duration_ms,
226            cost_usd: cost,
227            input_tokens,
228            cache_read_input_tokens: cache_read_tokens,
229            cache_creation_input_tokens: cache_creation_tokens,
230            output_tokens,
231            model: result.model().map(String::from),
232            debug_messages,
233            artifacts: StepArtifacts::default(),
234            account_id,
235            environment_id: result.environment_id().map(String::from),
236        })
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use std::sync::{Arc, Mutex};
243    use std::time::Duration;
244
245    use ironflow_core::error::{AgentError, OperationError};
246    use ironflow_core::operations::agent::PermissionMode;
247    use ironflow_core::provider::{AgentConfig, AgentOutput, AgentProvider, InvokeFuture};
248    use serde_json::json;
249    use tokio::time::{sleep, timeout};
250    use uuid::Uuid;
251
252    use super::{AgentExecutor, StepExecutor};
253    use crate::error::EngineError;
254
255    /// Real provider returning a fixed output, used to exercise usage propagation.
256    struct FixedUsageProvider {
257        output: AgentOutput,
258    }
259
260    impl AgentProvider for FixedUsageProvider {
261        fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
262            Box::pin(async move { Ok(self.output.clone()) })
263        }
264    }
265
266    /// How [`SessionProvider`] fails a resume.
267    enum ResumeFailure {
268        /// The CLI exits with this stderr.
269        Stderr(String),
270        /// A json_schema run ends with a schema validation error carrying this `got`.
271        Schema(String),
272    }
273
274    /// Provider recording every config it receives, failing a resume with
275    /// `resume_error` when set.
276    struct SessionProvider {
277        seen: Mutex<Vec<AgentConfig>>,
278        resume_error: Option<ResumeFailure>,
279    }
280
281    impl SessionProvider {
282        fn new(resume_error: Option<&str>) -> Arc<Self> {
283            Arc::new(Self {
284                seen: Mutex::new(Vec::new()),
285                resume_error: resume_error.map(|s| ResumeFailure::Stderr(s.to_string())),
286            })
287        }
288
289        fn failing_schema(got: &str) -> Arc<Self> {
290            Arc::new(Self {
291                seen: Mutex::new(Vec::new()),
292                resume_error: Some(ResumeFailure::Schema(got.to_string())),
293            })
294        }
295
296        fn seen(&self) -> Vec<AgentConfig> {
297            self.seen.lock().unwrap().clone()
298        }
299    }
300
301    impl AgentProvider for SessionProvider {
302        fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a> {
303            Box::pin(async move {
304                self.seen.lock().unwrap().push(config.clone());
305                match (&config.resume_session_id, &self.resume_error) {
306                    (Some(_), Some(ResumeFailure::Stderr(stderr))) => {
307                        Err(AgentError::ProcessFailed {
308                            exit_code: 1,
309                            stderr: stderr.clone(),
310                        })
311                    }
312                    (Some(_), Some(ResumeFailure::Schema(got))) => {
313                        Err(AgentError::SchemaValidation {
314                            expected: "structured_output field".to_string(),
315                            got: got.clone(),
316                            debug_messages: Vec::new(),
317                            partial_usage: Box::default(),
318                            raw_response: None,
319                        })
320                    }
321                    _ => {
322                        let mut output = AgentOutput::new(json!("ok"));
323                        output.cost_usd = Some(0.02);
324                        Ok(output)
325                    }
326                }
327            })
328        }
329    }
330
331    const SID: &str = "0192f0c1-7d2e-7a4b-9c3d-1e2f3a4b5c6d";
332
333    #[tokio::test]
334    async fn agent_resume_executor_sends_resume_prompt() {
335        timeout(Duration::from_secs(10), async {
336            let recorder = SessionProvider::new(None);
337            let provider: Arc<dyn AgentProvider> = recorder.clone();
338            let config = budget_config().resume(SID).resume_prompt("go on");
339
340            AgentExecutor::new(&config)
341                .execute(&provider)
342                .await
343                .expect("resumed step succeeds");
344
345            let seen = recorder.seen();
346            assert_eq!(seen.len(), 1);
347            assert_eq!(seen[0].prompt, "go on");
348            assert_eq!(seen[0].resume_session_id.as_deref(), Some(SID));
349        })
350        .await
351        .expect("test timed out");
352    }
353
354    #[tokio::test]
355    async fn agent_resume_executor_falls_back_when_session_is_missing() {
356        timeout(Duration::from_secs(10), async {
357            let recorder =
358                SessionProvider::new(Some("No conversation found with session ID: 0192f0c1"));
359            let provider: Arc<dyn AgentProvider> = recorder.clone();
360            let config = budget_config().resume(SID).resume_prompt("go on");
361
362            AgentExecutor::new(&config)
363                .execute(&provider)
364                .await
365                .expect("a missing session never fails the step");
366
367            let seen = recorder.seen();
368            assert_eq!(seen.len(), 2);
369            assert_eq!(seen[1].prompt, "hi");
370            assert_eq!(seen[1].resume_session_id, None);
371            assert_eq!(seen[1].session_id.as_deref(), Some(SID));
372        })
373        .await
374        .expect("test timed out");
375    }
376
377    #[tokio::test]
378    async fn agent_resume_executor_restarts_json_schema_agent_when_session_is_missing() {
379        timeout(Duration::from_secs(10), async {
380            let recorder = SessionProvider::failing_schema(
381                "null (error_during_execution after 0 turns: the session to resume was not found)",
382            );
383            let provider: Arc<dyn AgentProvider> = recorder.clone();
384            let config: AgentConfig = budget_config()
385                .resume(SID)
386                .resume_prompt("go on")
387                .output_schema_raw(r#"{"type":"object"}"#)
388                .into();
389
390            AgentExecutor::new(&config)
391                .execute(&provider)
392                .await
393                .expect("a missing session never fails the step");
394
395            let seen = recorder.seen();
396            assert_eq!(seen.len(), 2, "the resume is not retried");
397            assert_eq!(seen[0].prompt, "go on");
398            assert_eq!(seen[1].prompt, "hi");
399            assert_eq!(seen[1].resume_session_id, None);
400            assert_eq!(seen[1].session_id.as_deref(), Some(SID));
401        })
402        .await
403        .expect("test timed out");
404    }
405
406    #[tokio::test]
407    async fn agent_resume_executor_keeps_schema_failure_with_turns() {
408        timeout(Duration::from_secs(10), async {
409            let recorder = SessionProvider::failing_schema("null");
410            let provider: Arc<dyn AgentProvider> = recorder.clone();
411            let config: AgentConfig = budget_config()
412                .resume(SID)
413                .resume_prompt("go on")
414                .output_schema_raw(r#"{"type":"object"}"#)
415                .into();
416
417            AgentExecutor::new(&config)
418                .execute(&provider)
419                .await
420                .expect_err("a schema failure of a live session is not masked");
421
422            assert!(
423                recorder
424                    .seen()
425                    .iter()
426                    .all(|c| c.resume_session_id.is_some()),
427                "no fresh run"
428            );
429        })
430        .await
431        .expect("test timed out");
432    }
433
434    #[tokio::test]
435    async fn agent_resume_executor_propagates_other_errors() {
436        timeout(Duration::from_secs(10), async {
437            let recorder = SessionProvider::new(Some("permission denied"));
438            let provider: Arc<dyn AgentProvider> = recorder.clone();
439            let config = budget_config().resume(SID).resume_prompt("go on");
440
441            let err = AgentExecutor::new(&config)
442                .execute(&provider)
443                .await
444                .expect_err("other errors are not masked");
445
446            assert!(err.to_string().contains("permission denied"), "{err}");
447            assert_eq!(recorder.seen().len(), 1);
448        })
449        .await
450        .expect("test timed out");
451    }
452
453    #[tokio::test]
454    async fn agent_resume_executor_leaves_user_resume_alone() {
455        timeout(Duration::from_secs(10), async {
456            let recorder = SessionProvider::new(None);
457            let provider: Arc<dyn AgentProvider> = recorder.clone();
458            // A user-set resume without resume prompt keeps the original prompt.
459            let config = budget_config().resume(SID);
460
461            AgentExecutor::new(&config)
462                .execute(&provider)
463                .await
464                .expect("step succeeds");
465
466            let seen = recorder.seen();
467            assert_eq!(seen.len(), 1);
468            assert_eq!(seen[0].prompt, "hi");
469            assert_eq!(seen[0].resume_session_id.as_deref(), Some(SID));
470        })
471        .await
472        .expect("test timed out");
473    }
474
475    fn budget_config() -> AgentConfig {
476        let mut config = AgentConfig::new("hi");
477        config.max_budget_usd = Some(0.10);
478        config
479    }
480
481    /// Real provider that answers after `delay`, as a slow agent would.
482    struct SlowProvider {
483        delay: Duration,
484    }
485
486    impl AgentProvider for SlowProvider {
487        fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
488            Box::pin(async move {
489                sleep(self.delay).await;
490                Ok(AgentOutput::new(json!("ok")))
491            })
492        }
493    }
494
495    #[tokio::test]
496    async fn agent_executor_fails_a_step_past_its_timeout() {
497        timeout(Duration::from_secs(10), async {
498            let provider: Arc<dyn AgentProvider> = Arc::new(SlowProvider {
499                delay: Duration::from_secs(30),
500            });
501            let config = budget_config().timeout(Duration::from_millis(100));
502
503            let err = AgentExecutor::new(&config)
504                .execute(&provider)
505                .await
506                .expect_err("the step outlives its timeout");
507
508            assert!(
509                matches!(
510                    err,
511                    EngineError::Operation(OperationError::Agent(AgentError::Timeout { limit }))
512                        if limit == Duration::from_millis(100)
513                ),
514                "expected an agent timeout, got {err:?}"
515            );
516            assert!(err.to_string().contains("timed out"), "{err}");
517        })
518        .await
519        .expect("test timed out");
520    }
521
522    #[tokio::test]
523    async fn agent_executor_without_timeout_runs_the_step() {
524        timeout(Duration::from_secs(10), async {
525            let provider: Arc<dyn AgentProvider> = Arc::new(SlowProvider {
526                delay: Duration::from_millis(50),
527            });
528
529            let step = AgentExecutor::new(&budget_config())
530                .execute(&provider)
531                .await
532                .expect("agent step succeeds");
533
534            assert_eq!(step.output, json!("ok"));
535        })
536        .await
537        .expect("test timed out");
538    }
539
540    #[tokio::test]
541    async fn agent_executor_propagates_cache_tokens() {
542        timeout(Duration::from_secs(10), async {
543            let mut output = AgentOutput::new(json!("ok"));
544            output.input_tokens = Some(100);
545            output.cache_read_input_tokens = Some(5000);
546            output.cache_creation_input_tokens = Some(200);
547            output.output_tokens = Some(50);
548            output.cost_usd = Some(0.02);
549            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
550
551            let config = budget_config();
552            let step = AgentExecutor::new(&config)
553                .execute(&provider)
554                .await
555                .expect("agent step succeeds");
556
557            assert_eq!(step.input_tokens, Some(100));
558            assert_eq!(step.cache_read_input_tokens, Some(5000));
559            assert_eq!(step.cache_creation_input_tokens, Some(200));
560            assert_eq!(step.output_tokens, Some(50));
561            assert_eq!(step.total_tokens(), 5350);
562        })
563        .await
564        .expect("test timed out");
565    }
566
567    #[tokio::test]
568    async fn agent_executor_propagates_account_id() {
569        timeout(Duration::from_secs(10), async {
570            let account_id = Uuid::now_v7();
571            let mut output = AgentOutput::new(json!("ok"));
572            output.cost_usd = Some(0.02);
573            output.account_id = Some(account_id.to_string());
574            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
575
576            let step = AgentExecutor::new(&budget_config())
577                .execute(&provider)
578                .await
579                .expect("agent step succeeds");
580            assert_eq!(step.account_id, Some(account_id));
581
582            let mut invalid = AgentOutput::new(json!("ok"));
583            invalid.cost_usd = Some(0.02);
584            invalid.account_id = Some("not-a-uuid".to_string());
585            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output: invalid });
586            let step = AgentExecutor::new(&budget_config())
587                .execute(&provider)
588                .await
589                .expect("agent step succeeds");
590            assert_eq!(step.account_id, None);
591        })
592        .await
593        .expect("test timed out");
594    }
595
596    #[tokio::test]
597    async fn agent_executor_propagates_environment_id() {
598        timeout(Duration::from_secs(10), async {
599            let mut output = AgentOutput::new(json!("ok"));
600            output.cost_usd = Some(0.02);
601            output.environment_id = Some("ironflow-env-0a1b2c".to_string());
602            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
603            let step = AgentExecutor::new(&budget_config())
604                .execute(&provider)
605                .await
606                .expect("agent step succeeds");
607            assert_eq!(step.environment_id.as_deref(), Some("ironflow-env-0a1b2c"));
608
609            let mut output = AgentOutput::new(json!("ok"));
610            output.cost_usd = Some(0.02);
611            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
612            let step = AgentExecutor::new(&budget_config())
613                .execute(&provider)
614                .await
615                .expect("agent step succeeds");
616            assert_eq!(step.environment_id, None);
617        })
618        .await
619        .expect("test timed out");
620    }
621
622    #[tokio::test]
623    async fn agent_executor_without_cache_tokens_yields_none() {
624        timeout(Duration::from_secs(10), async {
625            let mut output = AgentOutput::new(json!("ok"));
626            output.input_tokens = Some(100);
627            output.output_tokens = Some(50);
628            output.cost_usd = Some(0.02);
629            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
630
631            let config = budget_config();
632            let step = AgentExecutor::new(&config)
633                .execute(&provider)
634                .await
635                .expect("agent step succeeds");
636
637            assert_eq!(step.input_tokens, Some(100));
638            assert!(step.cache_read_input_tokens.is_none());
639            assert!(step.cache_creation_input_tokens.is_none());
640            assert_eq!(step.total_tokens(), 150);
641        })
642        .await
643        .expect("test timed out");
644    }
645
646    #[test]
647    fn parse_permission_mode_via_serde() {
648        let json = r#""auto""#;
649        let mode: PermissionMode = serde_json::from_str(json).unwrap();
650        assert!(matches!(mode, PermissionMode::Auto));
651    }
652
653    #[test]
654    fn parse_permission_mode_dont_ask() {
655        let json = r#""dont_ask""#;
656        let mode: PermissionMode = serde_json::from_str(json).unwrap();
657        assert!(matches!(mode, PermissionMode::DontAsk));
658    }
659
660    #[test]
661    fn parse_permission_mode_bypass() {
662        let json = r#""bypass""#;
663        let mode: PermissionMode = serde_json::from_str(json).unwrap();
664        assert!(matches!(mode, PermissionMode::BypassPermissions));
665    }
666
667    #[test]
668    fn parse_permission_mode_case_insensitive() {
669        let json = r#""AUTO""#;
670        let mode: PermissionMode = serde_json::from_str(json).unwrap();
671        assert!(matches!(mode, PermissionMode::Auto));
672
673        let json = r#""DONT_ASK""#;
674        let mode: PermissionMode = serde_json::from_str(json).unwrap();
675        assert!(matches!(mode, PermissionMode::DontAsk));
676    }
677
678    #[test]
679    fn parse_permission_mode_unknown_defaults() {
680        let json = r#""unknown""#;
681        let mode: PermissionMode = serde_json::from_str(json).unwrap();
682        assert!(matches!(mode, PermissionMode::Default));
683    }
684}
ironflow_engine/executor/agent.rs

ironflow_engine/executor/
agent.rs
