Skip to main content

ironflow_engine/config/
agent.rs

1//! Agent step configuration -- re-exports [`AgentConfig`] from `ironflow-core`.
2//!
3//! [`AgentStepConfig`] is a type alias for [`AgentConfig`], keeping backward
4//! compatibility while eliminating the duplicated config struct.
5//! [`AgentStep`] ties each typestate of the builder to what
6//! [`WorkflowContext::agent`](crate::context::WorkflowContext::agent) returns.
7
8use serde::de::DeserializeOwned;
9
10pub use ironflow_core::provider::{AgentConfig, Tool, ToolProfile};
11use ironflow_core::provider::{NoSchema, NoTools, RawSchema, WithSchema, WithTools};
12
13use crate::error::EngineError;
14use crate::executor::StepOutput;
15
16/// Backward-compatible alias for [`AgentConfig`].
17///
18/// # Examples
19///
20/// ```
21/// use ironflow_engine::config::AgentStepConfig;
22///
23/// let config = AgentStepConfig::new("Review this code for security issues")
24///     .model("haiku")
25///     .max_budget_usd(0.10);
26/// ```
27pub type AgentStepConfig = AgentConfig;
28
29/// An agent configuration [`WorkflowContext::agent`](crate::context::WorkflowContext::agent)
30/// accepts, and the answer it returns for it.
31///
32/// A config built with [`output::<T>()`](AgentConfig::output) answers with the
33/// `T` itself; every other config answers with the raw [`StepOutput`].
34///
35/// # Examples
36///
37/// ```
38/// use ironflow_engine::config::{AgentStep, AgentStepConfig};
39/// use ironflow_engine::executor::{StepArtifacts, StepOutput};
40/// use rust_decimal::Decimal;
41/// use schemars::JsonSchema;
42/// use serde::Deserialize;
43/// use serde_json::json;
44///
45/// #[derive(Deserialize, JsonSchema)]
46/// struct Verdict {
47///     approved: bool,
48/// }
49///
50/// let config = AgentStepConfig::new("Review").max_turns(2).output::<Verdict>();
51/// let output = StepOutput {
52///     output: json!({"approved": true}),
53///     duration_ms: 0,
54///     cost_usd: Decimal::ZERO,
55///     input_tokens: None,
56///     cache_read_input_tokens: None,
57///     cache_creation_input_tokens: None,
58///     output_tokens: None,
59///     model: None,
60///     debug_messages: None,
61///     artifacts: StepArtifacts::default(),
62///     account_id: None,
63///     environment_id: None,
64/// };
65/// # fn answer<C: AgentStep>(_config: &C, output: StepOutput) -> Result<C::Answer, ironflow_engine::error::EngineError> {
66/// #     C::answer(output)
67/// # }
68/// let verdict: Verdict = answer(&config, output)?;
69/// assert!(verdict.approved);
70/// # Ok::<(), ironflow_engine::error::EngineError>(())
71/// ```
72pub trait AgentStep {
73    /// What the step returns to the handler.
74    type Answer;
75
76    /// The configuration, without its typestate.
77    fn into_config(self) -> AgentStepConfig;
78
79    /// Read the step output as the answer.
80    ///
81    /// # Errors
82    ///
83    /// Returns [`EngineError::Serialization`] when a typed answer does not
84    /// match its type.
85    fn answer(output: StepOutput) -> Result<Self::Answer, EngineError>;
86}
87
88impl AgentStep for AgentConfig<NoTools, NoSchema> {
89    type Answer = StepOutput;
90
91    fn into_config(self) -> AgentStepConfig {
92        self
93    }
94
95    fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
96        Ok(output)
97    }
98}
99
100impl AgentStep for AgentConfig<WithTools, NoSchema> {
101    type Answer = StepOutput;
102
103    fn into_config(self) -> AgentStepConfig {
104        self.into()
105    }
106
107    fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
108        Ok(output)
109    }
110}
111
112impl AgentStep for AgentConfig<NoTools, RawSchema> {
113    type Answer = StepOutput;
114
115    fn into_config(self) -> AgentStepConfig {
116        self.into()
117    }
118
119    fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
120        Ok(output)
121    }
122}
123
124impl<T: DeserializeOwned> AgentStep for AgentConfig<NoTools, WithSchema<T>> {
125    type Answer = T;
126
127    fn into_config(self) -> AgentStepConfig {
128        self.into()
129    }
130
131    fn answer(output: StepOutput) -> Result<T, EngineError> {
132        output.json()
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::executor::StepArtifacts;
140    use rust_decimal::Decimal;
141    use schemars::JsonSchema;
142    use serde::Deserialize;
143    use serde_json::{Value, json};
144
145    fn step_output(output: Value) -> StepOutput {
146        StepOutput {
147            output,
148            duration_ms: 0,
149            cost_usd: Decimal::ZERO,
150            input_tokens: None,
151            cache_read_input_tokens: None,
152            cache_creation_input_tokens: None,
153            output_tokens: None,
154            model: None,
155            debug_messages: None,
156            artifacts: StepArtifacts::default(),
157            account_id: None,
158            environment_id: None,
159        }
160    }
161
162    #[derive(Debug, PartialEq, Deserialize, JsonSchema)]
163    struct Verdict {
164        approved: bool,
165    }
166
167    #[test]
168    fn a_typed_config_answers_with_its_type() {
169        let answer = <AgentConfig<NoTools, WithSchema<Verdict>> as AgentStep>::answer(step_output(
170            json!({"approved": true}),
171        ))
172        .expect("matches Verdict");
173        assert_eq!(answer, Verdict { approved: true });
174    }
175
176    #[test]
177    fn a_typed_config_rejects_an_answer_of_another_shape() {
178        let err =
179            <AgentConfig<NoTools, WithSchema<Verdict>> as AgentStep>::answer(step_output(json!([
180                "approved"
181            ])))
182            .expect_err("not a Verdict");
183        assert!(matches!(err, EngineError::Serialization(_)));
184    }
185
186    #[test]
187    fn untyped_configs_answer_with_the_raw_output() {
188        let raw = <AgentConfig as AgentStep>::answer(step_output(json!("free text")))
189            .expect("raw output");
190        assert_eq!(raw.output, json!("free text"));
191    }
192
193    #[test]
194    fn into_config_keeps_the_settings() {
195        let config = AgentStepConfig::new("Review")
196            .max_turns(2)
197            .output::<Verdict>()
198            .into_config();
199        assert_eq!(config.max_turns, Some(2));
200        assert!(config.json_schema.is_some());
201
202        let config = AgentStepConfig::new("Explore")
203            .allow_tool(Tool::Grep)
204            .into_config();
205        assert_eq!(config.allowed_tools, vec!["Grep"]);
206    }
207
208    #[test]
209    fn builder() {
210        let config = AgentStepConfig::new("Review code")
211            .system_prompt("You are a code reviewer")
212            .model("haiku")
213            .max_budget_usd(0.50)
214            .max_turns(5)
215            .allow_tool(Tool::Read)
216            .working_dir("/repo")
217            .permission_mode(ironflow_core::operations::agent::PermissionMode::Auto);
218
219        assert_eq!(config.prompt, "Review code");
220        assert_eq!(config.system_prompt.unwrap(), "You are a code reviewer");
221        assert_eq!(config.model, "haiku");
222        assert_eq!(config.allowed_tools, vec!["Read"]);
223        assert!(config.json_schema.is_none());
224    }
225
226    #[test]
227    fn output_sets_schema_from_type() {
228        #[derive(serde::Deserialize, JsonSchema)]
229        #[allow(dead_code)]
230        struct Labels {
231            labels: Vec<String>,
232        }
233
234        let config = AgentStepConfig::new("Classify").output::<Labels>();
235
236        let schema = config.json_schema.expect("schema should be set");
237        assert!(schema.contains("labels"));
238    }
239
240    #[test]
241    fn output_schema_raw_sets_string() {
242        let raw = r#"{"type":"object"}"#;
243        let config = AgentStepConfig::new("Rate").output_schema_raw(raw);
244
245        assert_eq!(config.json_schema.as_deref(), Some(raw));
246    }
247
248    #[test]
249    fn output_overrides_previous_schema() {
250        #[derive(serde::Deserialize, JsonSchema)]
251        #[allow(dead_code)]
252        struct First {
253            a: String,
254        }
255
256        #[derive(serde::Deserialize, JsonSchema)]
257        #[allow(dead_code)]
258        struct Second {
259            b: i32,
260        }
261
262        let config = AgentStepConfig::new("Test")
263            .output::<First>()
264            .output::<Second>();
265
266        let schema = config.json_schema.expect("schema should be set");
267        assert!(!schema.contains("\"a\""));
268        assert!(schema.contains("\"b\""));
269    }
270
271    #[test]
272    fn output_schema_raw_overrides_typed_schema() {
273        #[derive(serde::Deserialize, JsonSchema)]
274        #[allow(dead_code)]
275        struct Typed {
276            field: String,
277        }
278
279        let raw = r#"{"type":"string"}"#;
280        let config = AgentStepConfig::new("Test")
281            .output::<Typed>()
282            .output_schema_raw(raw);
283
284        assert_eq!(config.json_schema.as_deref(), Some(raw));
285    }
286
287    #[test]
288    fn default_output_schema_is_none() {
289        let config = AgentStepConfig::new("Hello");
290        assert!(config.json_schema.is_none());
291    }
292
293    #[test]
294    fn serde_roundtrip_with_defaults() {
295        let json = r#"{"prompt":"hello"}"#;
296        let config: AgentConfig = serde_json::from_str(json).unwrap();
297        assert_eq!(config.prompt, "hello");
298        assert_eq!(config.model, "sonnet");
299        assert!(!config.verbose);
300    }
301
302    #[test]
303    fn serde_permission_mode_case_insensitive() {
304        let json = r#"{"prompt":"test","permission_mode":"auto"}"#;
305        let config: AgentConfig = serde_json::from_str(json).unwrap();
306        assert!(matches!(
307            config.permission_mode,
308            ironflow_core::operations::agent::PermissionMode::Auto
309        ));
310    }
311
312    #[test]
313    fn serde_output_schema_alias() {
314        let json = r#"{"prompt":"test","output_schema":"{\"type\":\"object\"}"}"#;
315        let config: AgentConfig = serde_json::from_str(json).unwrap();
316        assert_eq!(config.json_schema.as_deref(), Some(r#"{"type":"object"}"#));
317    }
318
319    #[test]
320    fn strict_mcp_config_defaults_to_false() {
321        let config = AgentStepConfig::new("test");
322        assert!(!config.strict_mcp_config);
323    }
324
325    #[test]
326    fn strict_mcp_config_builder_sets_flag() {
327        let config = AgentStepConfig::new("test").strict_mcp_config(true);
328        assert!(config.strict_mcp_config);
329    }
330
331    #[test]
332    fn strict_mcp_config_serde_default_when_missing() {
333        let json = r#"{"prompt":"test"}"#;
334        let config: AgentConfig = serde_json::from_str(json).unwrap();
335        assert!(!config.strict_mcp_config);
336    }
337
338    #[test]
339    fn strict_mcp_config_serde_roundtrip() {
340        let json = r#"{"prompt":"test","strict_mcp_config":true}"#;
341        let config: AgentConfig = serde_json::from_str(json).unwrap();
342        assert!(config.strict_mcp_config);
343    }
344}