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;
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::timeout;
250    use uuid::Uuid;
251
252    use super::{AgentExecutor, StepExecutor};
253
254    /// Real provider returning a fixed output, used to exercise usage propagation.
255    struct FixedUsageProvider {
256        output: AgentOutput,
257    }
258
259    impl AgentProvider for FixedUsageProvider {
260        fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
261            Box::pin(async move { Ok(self.output.clone()) })
262        }
263    }
264
265    /// How [`SessionProvider`] fails a resume.
266    enum ResumeFailure {
267        /// The CLI exits with this stderr.
268        Stderr(String),
269        /// A json_schema run ends with a schema validation error carrying this `got`.
270        Schema(String),
271    }
272
273    /// Provider recording every config it receives, failing a resume with
274    /// `resume_error` when set.
275    struct SessionProvider {
276        seen: Mutex<Vec<AgentConfig>>,
277        resume_error: Option<ResumeFailure>,
278    }
279
280    impl SessionProvider {
281        fn new(resume_error: Option<&str>) -> Arc<Self> {
282            Arc::new(Self {
283                seen: Mutex::new(Vec::new()),
284                resume_error: resume_error.map(|s| ResumeFailure::Stderr(s.to_string())),
285            })
286        }
287
288        fn failing_schema(got: &str) -> Arc<Self> {
289            Arc::new(Self {
290                seen: Mutex::new(Vec::new()),
291                resume_error: Some(ResumeFailure::Schema(got.to_string())),
292            })
293        }
294
295        fn seen(&self) -> Vec<AgentConfig> {
296            self.seen.lock().unwrap().clone()
297        }
298    }
299
300    impl AgentProvider for SessionProvider {
301        fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a> {
302            Box::pin(async move {
303                self.seen.lock().unwrap().push(config.clone());
304                match (&config.resume_session_id, &self.resume_error) {
305                    (Some(_), Some(ResumeFailure::Stderr(stderr))) => {
306                        Err(AgentError::ProcessFailed {
307                            exit_code: 1,
308                            stderr: stderr.clone(),
309                        })
310                    }
311                    (Some(_), Some(ResumeFailure::Schema(got))) => {
312                        Err(AgentError::SchemaValidation {
313                            expected: "structured_output field".to_string(),
314                            got: got.clone(),
315                            debug_messages: Vec::new(),
316                            partial_usage: Box::default(),
317                            raw_response: None,
318                        })
319                    }
320                    _ => {
321                        let mut output = AgentOutput::new(json!("ok"));
322                        output.cost_usd = Some(0.02);
323                        Ok(output)
324                    }
325                }
326            })
327        }
328    }
329
330    const SID: &str = "0192f0c1-7d2e-7a4b-9c3d-1e2f3a4b5c6d";
331
332    #[tokio::test]
333    async fn agent_resume_executor_sends_resume_prompt() {
334        timeout(Duration::from_secs(10), async {
335            let recorder = SessionProvider::new(None);
336            let provider: Arc<dyn AgentProvider> = recorder.clone();
337            let config = budget_config().resume(SID).resume_prompt("go on");
338
339            AgentExecutor::new(&config)
340                .execute(&provider)
341                .await
342                .expect("resumed step succeeds");
343
344            let seen = recorder.seen();
345            assert_eq!(seen.len(), 1);
346            assert_eq!(seen[0].prompt, "go on");
347            assert_eq!(seen[0].resume_session_id.as_deref(), Some(SID));
348        })
349        .await
350        .expect("test timed out");
351    }
352
353    #[tokio::test]
354    async fn agent_resume_executor_falls_back_when_session_is_missing() {
355        timeout(Duration::from_secs(10), async {
356            let recorder =
357                SessionProvider::new(Some("No conversation found with session ID: 0192f0c1"));
358            let provider: Arc<dyn AgentProvider> = recorder.clone();
359            let config = budget_config().resume(SID).resume_prompt("go on");
360
361            AgentExecutor::new(&config)
362                .execute(&provider)
363                .await
364                .expect("a missing session never fails the step");
365
366            let seen = recorder.seen();
367            assert_eq!(seen.len(), 2);
368            assert_eq!(seen[1].prompt, "hi");
369            assert_eq!(seen[1].resume_session_id, None);
370            assert_eq!(seen[1].session_id.as_deref(), Some(SID));
371        })
372        .await
373        .expect("test timed out");
374    }
375
376    #[tokio::test]
377    async fn agent_resume_executor_restarts_json_schema_agent_when_session_is_missing() {
378        timeout(Duration::from_secs(10), async {
379            let recorder = SessionProvider::failing_schema(
380                "null (error_during_execution after 0 turns: the session to resume was not found)",
381            );
382            let provider: Arc<dyn AgentProvider> = recorder.clone();
383            let config: AgentConfig = budget_config()
384                .resume(SID)
385                .resume_prompt("go on")
386                .output_schema_raw(r#"{"type":"object"}"#)
387                .into();
388
389            AgentExecutor::new(&config)
390                .execute(&provider)
391                .await
392                .expect("a missing session never fails the step");
393
394            let seen = recorder.seen();
395            assert_eq!(seen.len(), 2, "the resume is not retried");
396            assert_eq!(seen[0].prompt, "go on");
397            assert_eq!(seen[1].prompt, "hi");
398            assert_eq!(seen[1].resume_session_id, None);
399            assert_eq!(seen[1].session_id.as_deref(), Some(SID));
400        })
401        .await
402        .expect("test timed out");
403    }
404
405    #[tokio::test]
406    async fn agent_resume_executor_keeps_schema_failure_with_turns() {
407        timeout(Duration::from_secs(10), async {
408            let recorder = SessionProvider::failing_schema("null");
409            let provider: Arc<dyn AgentProvider> = recorder.clone();
410            let config: AgentConfig = budget_config()
411                .resume(SID)
412                .resume_prompt("go on")
413                .output_schema_raw(r#"{"type":"object"}"#)
414                .into();
415
416            AgentExecutor::new(&config)
417                .execute(&provider)
418                .await
419                .expect_err("a schema failure of a live session is not masked");
420
421            assert!(
422                recorder
423                    .seen()
424                    .iter()
425                    .all(|c| c.resume_session_id.is_some()),
426                "no fresh run"
427            );
428        })
429        .await
430        .expect("test timed out");
431    }
432
433    #[tokio::test]
434    async fn agent_resume_executor_propagates_other_errors() {
435        timeout(Duration::from_secs(10), async {
436            let recorder = SessionProvider::new(Some("permission denied"));
437            let provider: Arc<dyn AgentProvider> = recorder.clone();
438            let config = budget_config().resume(SID).resume_prompt("go on");
439
440            let err = AgentExecutor::new(&config)
441                .execute(&provider)
442                .await
443                .expect_err("other errors are not masked");
444
445            assert!(err.to_string().contains("permission denied"), "{err}");
446            assert_eq!(recorder.seen().len(), 1);
447        })
448        .await
449        .expect("test timed out");
450    }
451
452    #[tokio::test]
453    async fn agent_resume_executor_leaves_user_resume_alone() {
454        timeout(Duration::from_secs(10), async {
455            let recorder = SessionProvider::new(None);
456            let provider: Arc<dyn AgentProvider> = recorder.clone();
457            // A user-set resume without resume prompt keeps the original prompt.
458            let config = budget_config().resume(SID);
459
460            AgentExecutor::new(&config)
461                .execute(&provider)
462                .await
463                .expect("step succeeds");
464
465            let seen = recorder.seen();
466            assert_eq!(seen.len(), 1);
467            assert_eq!(seen[0].prompt, "hi");
468            assert_eq!(seen[0].resume_session_id.as_deref(), Some(SID));
469        })
470        .await
471        .expect("test timed out");
472    }
473
474    fn budget_config() -> AgentConfig {
475        let mut config = AgentConfig::new("hi");
476        config.max_budget_usd = Some(0.10);
477        config
478    }
479
480    #[tokio::test]
481    async fn agent_executor_propagates_cache_tokens() {
482        timeout(Duration::from_secs(10), async {
483            let mut output = AgentOutput::new(json!("ok"));
484            output.input_tokens = Some(100);
485            output.cache_read_input_tokens = Some(5000);
486            output.cache_creation_input_tokens = Some(200);
487            output.output_tokens = Some(50);
488            output.cost_usd = Some(0.02);
489            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
490
491            let config = budget_config();
492            let step = AgentExecutor::new(&config)
493                .execute(&provider)
494                .await
495                .expect("agent step succeeds");
496
497            assert_eq!(step.input_tokens, Some(100));
498            assert_eq!(step.cache_read_input_tokens, Some(5000));
499            assert_eq!(step.cache_creation_input_tokens, Some(200));
500            assert_eq!(step.output_tokens, Some(50));
501            assert_eq!(step.total_tokens(), 5350);
502        })
503        .await
504        .expect("test timed out");
505    }
506
507    #[tokio::test]
508    async fn agent_executor_propagates_account_id() {
509        timeout(Duration::from_secs(10), async {
510            let account_id = Uuid::now_v7();
511            let mut output = AgentOutput::new(json!("ok"));
512            output.cost_usd = Some(0.02);
513            output.account_id = Some(account_id.to_string());
514            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
515
516            let step = AgentExecutor::new(&budget_config())
517                .execute(&provider)
518                .await
519                .expect("agent step succeeds");
520            assert_eq!(step.account_id, Some(account_id));
521
522            let mut invalid = AgentOutput::new(json!("ok"));
523            invalid.cost_usd = Some(0.02);
524            invalid.account_id = Some("not-a-uuid".to_string());
525            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output: invalid });
526            let step = AgentExecutor::new(&budget_config())
527                .execute(&provider)
528                .await
529                .expect("agent step succeeds");
530            assert_eq!(step.account_id, None);
531        })
532        .await
533        .expect("test timed out");
534    }
535
536    #[tokio::test]
537    async fn agent_executor_propagates_environment_id() {
538        timeout(Duration::from_secs(10), async {
539            let mut output = AgentOutput::new(json!("ok"));
540            output.cost_usd = Some(0.02);
541            output.environment_id = Some("ironflow-env-0a1b2c".to_string());
542            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
543            let step = AgentExecutor::new(&budget_config())
544                .execute(&provider)
545                .await
546                .expect("agent step succeeds");
547            assert_eq!(step.environment_id.as_deref(), Some("ironflow-env-0a1b2c"));
548
549            let mut output = AgentOutput::new(json!("ok"));
550            output.cost_usd = Some(0.02);
551            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
552            let step = AgentExecutor::new(&budget_config())
553                .execute(&provider)
554                .await
555                .expect("agent step succeeds");
556            assert_eq!(step.environment_id, None);
557        })
558        .await
559        .expect("test timed out");
560    }
561
562    #[tokio::test]
563    async fn agent_executor_without_cache_tokens_yields_none() {
564        timeout(Duration::from_secs(10), async {
565            let mut output = AgentOutput::new(json!("ok"));
566            output.input_tokens = Some(100);
567            output.output_tokens = Some(50);
568            output.cost_usd = Some(0.02);
569            let provider: Arc<dyn AgentProvider> = Arc::new(FixedUsageProvider { output });
570
571            let config = budget_config();
572            let step = AgentExecutor::new(&config)
573                .execute(&provider)
574                .await
575                .expect("agent step succeeds");
576
577            assert_eq!(step.input_tokens, Some(100));
578            assert!(step.cache_read_input_tokens.is_none());
579            assert!(step.cache_creation_input_tokens.is_none());
580            assert_eq!(step.total_tokens(), 150);
581        })
582        .await
583        .expect("test timed out");
584    }
585
586    #[test]
587    fn parse_permission_mode_via_serde() {
588        let json = r#""auto""#;
589        let mode: PermissionMode = serde_json::from_str(json).unwrap();
590        assert!(matches!(mode, PermissionMode::Auto));
591    }
592
593    #[test]
594    fn parse_permission_mode_dont_ask() {
595        let json = r#""dont_ask""#;
596        let mode: PermissionMode = serde_json::from_str(json).unwrap();
597        assert!(matches!(mode, PermissionMode::DontAsk));
598    }
599
600    #[test]
601    fn parse_permission_mode_bypass() {
602        let json = r#""bypass""#;
603        let mode: PermissionMode = serde_json::from_str(json).unwrap();
604        assert!(matches!(mode, PermissionMode::BypassPermissions));
605    }
606
607    #[test]
608    fn parse_permission_mode_case_insensitive() {
609        let json = r#""AUTO""#;
610        let mode: PermissionMode = serde_json::from_str(json).unwrap();
611        assert!(matches!(mode, PermissionMode::Auto));
612
613        let json = r#""DONT_ASK""#;
614        let mode: PermissionMode = serde_json::from_str(json).unwrap();
615        assert!(matches!(mode, PermissionMode::DontAsk));
616    }
617
618    #[test]
619    fn parse_permission_mode_unknown_defaults() {
620        let json = r#""unknown""#;
621        let mode: PermissionMode = serde_json::from_str(json).unwrap();
622        assert!(matches!(mode, PermissionMode::Default));
623    }
624}
ironflow_engine/executor/agent.rs

ironflow_engine/executor/
agent.rs
