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