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