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