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 rust_decimal::Decimal;
154    use schemars::JsonSchema;
155    use serde::Deserialize;
156    use serde_json::{Value, json};
157
158    fn step_output(output: Value) -> StepOutput {
159        StepOutput {
160            output,
161            duration_ms: 0,
162            cost_usd: Decimal::ZERO,
163            input_tokens: None,
164            cache_read_input_tokens: None,
165            cache_creation_input_tokens: None,
166            output_tokens: None,
167            model: None,
168            debug_messages: None,
169            artifacts: StepArtifacts::default(),
170            account_id: None,
171            environment_id: None,
172        }
173    }
174
175    #[derive(Debug, PartialEq, Deserialize, JsonSchema)]
176    struct Verdict {
177        approved: bool,
178    }
179
180    #[test]
181    fn a_typed_config_answers_with_its_type() {
182        let answer = <AgentConfig<NoTools, WithSchema<Verdict>> as AgentStep>::answer(step_output(
183            json!({"approved": true}),
184        ))
185        .expect("matches Verdict");
186        assert_eq!(answer, Verdict { approved: true });
187    }
188
189    #[test]
190    fn a_typed_config_rejects_an_answer_of_another_shape() {
191        let err =
192            <AgentConfig<NoTools, WithSchema<Verdict>> as AgentStep>::answer(step_output(json!([
193                "approved"
194            ])))
195            .expect_err("not a Verdict");
196        assert!(matches!(err, EngineError::Serialization(_)));
197    }
198
199    #[test]
200    fn untyped_configs_answer_with_the_raw_output() {
201        let raw = <AgentConfig as AgentStep>::answer(step_output(json!("free text")))
202            .expect("raw output");
203        assert_eq!(raw.output, json!("free text"));
204    }
205
206    #[test]
207    fn into_config_keeps_the_settings() {
208        let config = AgentStepConfig::new("Review")
209            .max_turns(2)
210            .output::<Verdict>()
211            .into_config();
212        assert_eq!(config.max_turns, Some(2));
213        assert!(config.json_schema.is_some());
214
215        let config = AgentStepConfig::new("Explore")
216            .allow_tool(Tool::Grep)
217            .into_config();
218        assert_eq!(config.allowed_tools, vec!["Grep"]);
219    }
220
221    #[test]
222    fn builder() {
223        let config = AgentStepConfig::new("Review code")
224            .system_prompt("You are a code reviewer")
225            .model("haiku")
226            .max_budget_usd(0.50)
227            .max_turns(5)
228            .allow_tool(Tool::Read)
229            .working_dir("/repo")
230            .permission_mode(ironflow_core::operations::agent::PermissionMode::Auto);
231
232        assert_eq!(config.prompt, "Review code");
233        assert_eq!(config.system_prompt.unwrap(), "You are a code reviewer");
234        assert_eq!(config.model, "haiku");
235        assert_eq!(config.allowed_tools, vec!["Read"]);
236        assert!(config.json_schema.is_none());
237    }
238
239    #[test]
240    fn output_sets_schema_from_type() {
241        #[derive(serde::Deserialize, JsonSchema)]
242        #[allow(dead_code)]
243        struct Labels {
244            labels: Vec<String>,
245        }
246
247        let config = AgentStepConfig::new("Classify").output::<Labels>();
248
249        let schema = config.json_schema.expect("schema should be set");
250        assert!(schema.contains("labels"));
251    }
252
253    #[test]
254    fn output_schema_raw_sets_string() {
255        let raw = r#"{"type":"object"}"#;
256        let config = AgentStepConfig::new("Rate").output_schema_raw(raw);
257
258        assert_eq!(config.json_schema.as_deref(), Some(raw));
259    }
260
261    #[test]
262    fn output_overrides_previous_schema() {
263        #[derive(serde::Deserialize, JsonSchema)]
264        #[allow(dead_code)]
265        struct First {
266            a: String,
267        }
268
269        #[derive(serde::Deserialize, JsonSchema)]
270        #[allow(dead_code)]
271        struct Second {
272            b: i32,
273        }
274
275        let config = AgentStepConfig::new("Test")
276            .output::<First>()
277            .output::<Second>();
278
279        let schema = config.json_schema.expect("schema should be set");
280        assert!(!schema.contains("\"a\""));
281        assert!(schema.contains("\"b\""));
282    }
283
284    #[test]
285    fn output_schema_raw_overrides_typed_schema() {
286        #[derive(serde::Deserialize, JsonSchema)]
287        #[allow(dead_code)]
288        struct Typed {
289            field: String,
290        }
291
292        let raw = r#"{"type":"string"}"#;
293        let config = AgentStepConfig::new("Test")
294            .output::<Typed>()
295            .output_schema_raw(raw);
296
297        assert_eq!(config.json_schema.as_deref(), Some(raw));
298    }
299
300    #[test]
301    fn default_output_schema_is_none() {
302        let config = AgentStepConfig::new("Hello");
303        assert!(config.json_schema.is_none());
304    }
305
306    #[test]
307    fn serde_roundtrip_with_defaults() {
308        let json = r#"{"prompt":"hello"}"#;
309        let config: AgentConfig = serde_json::from_str(json).unwrap();
310        assert_eq!(config.prompt, "hello");
311        assert_eq!(config.model, "sonnet");
312        assert!(!config.verbose);
313    }
314
315    #[test]
316    fn serde_permission_mode_case_insensitive() {
317        let json = r#"{"prompt":"test","permission_mode":"auto"}"#;
318        let config: AgentConfig = serde_json::from_str(json).unwrap();
319        assert!(matches!(
320            config.permission_mode,
321            ironflow_core::operations::agent::PermissionMode::Auto
322        ));
323    }
324
325    #[test]
326    fn serde_output_schema_alias() {
327        let json = r#"{"prompt":"test","output_schema":"{\"type\":\"object\"}"}"#;
328        let config: AgentConfig = serde_json::from_str(json).unwrap();
329        assert_eq!(config.json_schema.as_deref(), Some(r#"{"type":"object"}"#));
330    }
331
332    #[test]
333    fn strict_mcp_config_defaults_to_false() {
334        let config = AgentStepConfig::new("test");
335        assert!(!config.strict_mcp_config);
336    }
337
338    #[test]
339    fn strict_mcp_config_builder_sets_flag() {
340        let config = AgentStepConfig::new("test").strict_mcp_config(true);
341        assert!(config.strict_mcp_config);
342    }
343
344    #[test]
345    fn strict_mcp_config_serde_default_when_missing() {
346        let json = r#"{"prompt":"test"}"#;
347        let config: AgentConfig = serde_json::from_str(json).unwrap();
348        assert!(!config.strict_mcp_config);
349    }
350
351    #[test]
352    fn strict_mcp_config_serde_roundtrip() {
353        let json = r#"{"prompt":"test","strict_mcp_config":true}"#;
354        let config: AgentConfig = serde_json::from_str(json).unwrap();
355        assert!(config.strict_mcp_config);
356    }
357}