Skip to main content

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};
8
9use ironflow_core::operations::agent::Agent;
10use ironflow_core::pricing::{CostBreakdown, StaticPricing, spawn_log};
11use ironflow_core::provider::{AgentConfig, AgentProvider, LogSink};
12
13use crate::error::EngineError;
14use crate::log_sender::StepLogSender;
15use crate::notify::LogStream;
16
17use super::{StepExecutor, StepOutput};
18
19/// Executor for agent (AI) steps.
20///
21/// Runs an AI agent with the given prompt and configuration, capturing
22/// the response value, cost, and token counts. When a [`StepLogSender`]
23/// is attached, emits system log lines for step start/end.
24pub struct AgentExecutor<'a> {
25    config: &'a AgentConfig,
26    log_sender: Option<StepLogSender>,
27}
28
29impl<'a> AgentExecutor<'a> {
30    /// Create a new agent executor from a config reference.
31    pub fn new(config: &'a AgentConfig) -> Self {
32        Self {
33            config,
34            log_sender: None,
35        }
36    }
37
38    /// Attach a log sender for system-level log lines.
39    pub fn with_log_sender(mut self, sender: StepLogSender) -> Self {
40        self.log_sender = Some(sender);
41        self
42    }
43}
44
45impl StepExecutor for AgentExecutor<'_> {
46    async fn execute(&self, provider: &Arc<dyn AgentProvider>) -> Result<StepOutput, EngineError> {
47        let start = Instant::now();
48
49        if let Some(ref sender) = self.log_sender {
50            sender.emit(
51                LogStream::System,
52                &format!("agent step started (model={})", self.config.model),
53            );
54        }
55
56        if self.config.json_schema.is_some() && self.config.max_turns == Some(1) {
57            warn!(
58                "structured output (json_schema) requires max_turns >= 2; \
59                 max_turns is set to 1, the agent will likely fail with error_max_turns"
60            );
61        }
62
63        let mut agent = Agent::from_config(self.config.clone());
64        if let Some(ref sender) = self.log_sender {
65            agent = agent.log_sink(Arc::new(sender.clone()) as Arc<dyn LogSink>);
66        }
67        let result = agent.run(provider.as_ref()).await?;
68
69        let duration_ms = start.elapsed().as_millis() as u64;
70        let cost = Decimal::try_from(result.cost_usd().unwrap_or(0.0)).unwrap_or(Decimal::ZERO);
71        let input_tokens = result.input_tokens();
72        let output_tokens = result.output_tokens();
73
74        info!(
75            step_kind = "agent",
76            model = %self.config.model,
77            cost_usd = %cost,
78            input_tokens = ?input_tokens,
79            output_tokens = ?output_tokens,
80            duration_ms,
81            "agent step completed"
82        );
83
84        let pricing = StaticPricing::new();
85        let breakdown = CostBreakdown::compute(
86            &pricing,
87            &self.config.model,
88            input_tokens.unwrap_or(0),
89            output_tokens.unwrap_or(0),
90        );
91        spawn_log("agent", &self.config.model, breakdown);
92
93        #[cfg(feature = "prometheus")]
94        {
95            use ironflow_core::metric_names::{
96                AGENT_COST_USD_TOTAL, AGENT_DURATION_SECONDS, AGENT_TOKENS_INPUT_TOTAL,
97                AGENT_TOKENS_OUTPUT_TOTAL, AGENT_TOTAL, STATUS_SUCCESS,
98            };
99            use metrics::{counter, gauge, histogram};
100            let model_label = self.config.model.clone();
101            counter!(AGENT_TOTAL, "model" => model_label.clone(), "status" => STATUS_SUCCESS)
102                .increment(1);
103            histogram!(AGENT_DURATION_SECONDS, "model" => model_label.clone())
104                .record(duration_ms as f64 / 1000.0);
105            gauge!(AGENT_COST_USD_TOTAL, "model" => model_label.clone())
106                .increment(cost.to_string().parse::<f64>().unwrap_or(0.0));
107            if let Some(inp) = input_tokens {
108                counter!(AGENT_TOKENS_INPUT_TOTAL, "model" => model_label.clone()).increment(inp);
109            }
110            if let Some(out) = output_tokens {
111                counter!(AGENT_TOKENS_OUTPUT_TOTAL, "model" => model_label).increment(out);
112            }
113        }
114
115        if let Some(ref sender) = self.log_sender {
116            sender.emit(
117                LogStream::System,
118                &format!(
119                    "agent step completed (cost=${cost}, tokens_in={}, tokens_out={})",
120                    input_tokens.unwrap_or(0),
121                    output_tokens.unwrap_or(0),
122                ),
123            );
124        }
125
126        let debug_messages = result.debug_messages().map(|msgs| msgs.to_vec());
127
128        Ok(StepOutput {
129            output: result.value().clone(),
130            duration_ms,
131            cost_usd: cost,
132            input_tokens,
133            output_tokens,
134            model: result.model().map(String::from),
135            debug_messages,
136        })
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use ironflow_core::operations::agent::PermissionMode;
143
144    #[test]
145    fn parse_permission_mode_via_serde() {
146        let json = r#""auto""#;
147        let mode: PermissionMode = serde_json::from_str(json).unwrap();
148        assert!(matches!(mode, PermissionMode::Auto));
149    }
150
151    #[test]
152    fn parse_permission_mode_dont_ask() {
153        let json = r#""dont_ask""#;
154        let mode: PermissionMode = serde_json::from_str(json).unwrap();
155        assert!(matches!(mode, PermissionMode::DontAsk));
156    }
157
158    #[test]
159    fn parse_permission_mode_bypass() {
160        let json = r#""bypass""#;
161        let mode: PermissionMode = serde_json::from_str(json).unwrap();
162        assert!(matches!(mode, PermissionMode::BypassPermissions));
163    }
164
165    #[test]
166    fn parse_permission_mode_case_insensitive() {
167        let json = r#""AUTO""#;
168        let mode: PermissionMode = serde_json::from_str(json).unwrap();
169        assert!(matches!(mode, PermissionMode::Auto));
170
171        let json = r#""DONT_ASK""#;
172        let mode: PermissionMode = serde_json::from_str(json).unwrap();
173        assert!(matches!(mode, PermissionMode::DontAsk));
174    }
175
176    #[test]
177    fn parse_permission_mode_unknown_defaults() {
178        let json = r#""unknown""#;
179        let mode: PermissionMode = serde_json::from_str(json).unwrap();
180        assert!(matches!(mode, PermissionMode::Default));
181    }
182}