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