ironflow_engine/config/
agent.rs1use 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
16pub const DEFAULT_RESUME_PROMPT: &str = "The previous run of this task was interrupted. Continue from where you stopped and finish the task.";
28
29pub type AgentStepConfig = AgentConfig;
41
42pub trait AgentStep {
86 type Answer;
88
89 fn into_config(self) -> AgentStepConfig;
91
92 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}