ironflow_engine/config/agent.rs
1//! [`AgentStepConfig`] — serializable configuration for an agent step.
2
3use schemars::JsonSchema;
4use serde::{Deserialize, Serialize};
5
6/// Serializable configuration for an agent step.
7///
8/// # Examples
9///
10/// ```
11/// use ironflow_engine::config::AgentStepConfig;
12///
13/// let config = AgentStepConfig::new("Review this code for security issues")
14/// .model("haiku")
15/// .max_budget_usd(0.10);
16/// ```
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct AgentStepConfig {
19 /// The user prompt.
20 pub prompt: String,
21 /// Optional system prompt.
22 pub system_prompt: Option<String>,
23 /// Model name (e.g. "sonnet", "opus", "haiku").
24 pub model: Option<String>,
25 /// Maximum budget in USD.
26 pub max_budget_usd: Option<f64>,
27 /// Maximum number of agentic turns.
28 pub max_turns: Option<u32>,
29 /// Tool allowlist.
30 pub allowed_tools: Vec<String>,
31 /// Working directory for the agent.
32 pub working_dir: Option<String>,
33 /// Permission mode (e.g. "auto", "dont_ask").
34 pub permission_mode: Option<String>,
35 /// Optional JSON Schema string for structured output.
36 ///
37 /// When set, the agent provider will request typed output conforming to this schema.
38 /// The result value is guaranteed to be valid JSON matching the schema.
39 pub output_schema: Option<String>,
40}
41
42impl AgentStepConfig {
43 /// Create a new agent config with the given prompt.
44 ///
45 /// # Examples
46 ///
47 /// ```
48 /// use ironflow_engine::config::AgentStepConfig;
49 ///
50 /// let config = AgentStepConfig::new("Summarize this file");
51 /// assert_eq!(config.prompt, "Summarize this file");
52 /// ```
53 pub fn new(prompt: &str) -> Self {
54 Self {
55 prompt: prompt.to_string(),
56 system_prompt: None,
57 model: None,
58 max_budget_usd: None,
59 max_turns: None,
60 allowed_tools: Vec::new(),
61 working_dir: None,
62 permission_mode: None,
63 output_schema: None,
64 }
65 }
66
67 /// Set the system prompt.
68 pub fn system_prompt(mut self, prompt: &str) -> Self {
69 self.system_prompt = Some(prompt.to_string());
70 self
71 }
72
73 /// Set the model name.
74 pub fn model(mut self, model: &str) -> Self {
75 self.model = Some(model.to_string());
76 self
77 }
78
79 /// Set the maximum budget in USD.
80 pub fn max_budget_usd(mut self, budget: f64) -> Self {
81 self.max_budget_usd = Some(budget);
82 self
83 }
84
85 /// Set the maximum number of turns.
86 pub fn max_turns(mut self, turns: u32) -> Self {
87 self.max_turns = Some(turns);
88 self
89 }
90
91 /// Add an allowed tool.
92 pub fn allow_tool(mut self, tool: &str) -> Self {
93 self.allowed_tools.push(tool.to_string());
94 self
95 }
96
97 /// Set the working directory.
98 pub fn working_dir(mut self, dir: &str) -> Self {
99 self.working_dir = Some(dir.to_string());
100 self
101 }
102
103 /// Set the permission mode.
104 pub fn permission_mode(mut self, mode: &str) -> Self {
105 self.permission_mode = Some(mode.to_string());
106 self
107 }
108
109 /// Set structured output from a Rust type implementing [`JsonSchema`].
110 ///
111 /// The schema is serialized once at build time. When set, the agent provider
112 /// will request typed output conforming to this schema.
113 ///
114 /// **Important:** structured output requires `max_turns >= 2`. The Claude CLI
115 /// uses the first turn for reasoning and a second turn to produce the
116 /// schema-conforming JSON. If `max_turns` is set to `1`, the agent will
117 /// fail at runtime with an `error_max_turns` error.
118 ///
119 /// # Examples
120 ///
121 /// ```
122 /// use ironflow_engine::config::AgentStepConfig;
123 /// use schemars::JsonSchema;
124 /// use serde::Deserialize;
125 ///
126 /// #[derive(Deserialize, JsonSchema)]
127 /// struct Labels {
128 /// labels: Vec<String>,
129 /// }
130 ///
131 /// let config = AgentStepConfig::new("Classify this email")
132 /// .output::<Labels>()
133 /// .max_turns(2);
134 ///
135 /// assert!(config.output_schema.is_some());
136 /// ```
137 pub fn output<T: JsonSchema>(mut self) -> Self {
138 let schema = schemars::schema_for!(T);
139 self.output_schema = match serde_json::to_string(&schema) {
140 Ok(s) => Some(s),
141 Err(e) => {
142 tracing::warn!(
143 error = %e,
144 type_name = std::any::type_name::<T>(),
145 "failed to serialize JSON schema, structured output disabled"
146 );
147 None
148 }
149 };
150 self
151 }
152
153 /// Set structured output from a pre-serialized JSON Schema string.
154 ///
155 /// Use this when the schema comes from configuration (e.g. YAML/JSON files)
156 /// rather than a Rust type. For type-safe schema generation, prefer
157 /// [`output`](AgentStepConfig::output).
158 ///
159 /// **Important:** structured output requires `max_turns >= 2`. See
160 /// [`output`](AgentStepConfig::output) for details.
161 ///
162 /// # Examples
163 ///
164 /// ```
165 /// use ironflow_engine::config::AgentStepConfig;
166 ///
167 /// let schema = r#"{"type":"object","properties":{"score":{"type":"integer"}}}"#;
168 /// let config = AgentStepConfig::new("Rate this PR")
169 /// .output_schema_raw(schema.to_string());
170 ///
171 /// assert_eq!(config.output_schema.as_deref(), Some(schema));
172 /// ```
173 pub fn output_schema_raw(mut self, schema: String) -> Self {
174 self.output_schema = Some(schema);
175 self
176 }
177}
178
179#[cfg(test)]
180mod tests {
181 use super::*;
182
183 #[test]
184 fn builder() {
185 let config = AgentStepConfig::new("Review code")
186 .system_prompt("You are a code reviewer")
187 .model("haiku")
188 .max_budget_usd(0.50)
189 .max_turns(5)
190 .allow_tool("read")
191 .working_dir("/repo")
192 .permission_mode("auto");
193
194 assert_eq!(config.prompt, "Review code");
195 assert_eq!(config.system_prompt.unwrap(), "You are a code reviewer");
196 assert_eq!(config.model.unwrap(), "haiku");
197 assert_eq!(config.allowed_tools, vec!["read"]);
198 assert!(config.output_schema.is_none());
199 }
200
201 #[test]
202 fn output_sets_schema_from_type() {
203 #[derive(serde::Deserialize, JsonSchema)]
204 #[allow(dead_code)]
205 struct Labels {
206 labels: Vec<String>,
207 }
208
209 let config = AgentStepConfig::new("Classify").output::<Labels>();
210
211 let schema = config.output_schema.expect("schema should be set");
212 assert!(schema.contains("labels"));
213 }
214
215 #[test]
216 fn output_schema_raw_sets_string() {
217 let raw = r#"{"type":"object"}"#;
218 let config = AgentStepConfig::new("Rate").output_schema_raw(raw.to_string());
219
220 assert_eq!(config.output_schema.as_deref(), Some(raw));
221 }
222
223 #[test]
224 fn output_overrides_previous_schema() {
225 #[derive(serde::Deserialize, JsonSchema)]
226 #[allow(dead_code)]
227 struct First {
228 a: String,
229 }
230
231 #[derive(serde::Deserialize, JsonSchema)]
232 #[allow(dead_code)]
233 struct Second {
234 b: i32,
235 }
236
237 let config = AgentStepConfig::new("Test")
238 .output::<First>()
239 .output::<Second>();
240
241 let schema = config.output_schema.expect("schema should be set");
242 assert!(!schema.contains("\"a\""));
243 assert!(schema.contains("\"b\""));
244 }
245
246 #[test]
247 fn output_schema_raw_overrides_typed_schema() {
248 #[derive(serde::Deserialize, JsonSchema)]
249 #[allow(dead_code)]
250 struct Typed {
251 field: String,
252 }
253
254 let raw = r#"{"type":"string"}"#;
255 let config = AgentStepConfig::new("Test")
256 .output::<Typed>()
257 .output_schema_raw(raw.to_string());
258
259 assert_eq!(config.output_schema.as_deref(), Some(raw));
260 }
261
262 #[test]
263 fn default_output_schema_is_none() {
264 let config = AgentStepConfig::new("Hello");
265 assert!(config.output_schema.is_none());
266 }
267}