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