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