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 type AgentStepConfig = AgentConfig;
28
29pub trait AgentStep {
72 type Answer;
74
75 fn into_config(self) -> AgentStepConfig;
77
78 fn answer(output: StepOutput) -> Result<Self::Answer, EngineError>;
85}
86
87impl AgentStep for AgentConfig<NoTools, NoSchema> {
88 type Answer = StepOutput;
89
90 fn into_config(self) -> AgentStepConfig {
91 self
92 }
93
94 fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
95 Ok(output)
96 }
97}
98
99impl AgentStep for AgentConfig<WithTools, NoSchema> {
100 type Answer = StepOutput;
101
102 fn into_config(self) -> AgentStepConfig {
103 self.into()
104 }
105
106 fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
107 Ok(output)
108 }
109}
110
111impl AgentStep for AgentConfig<NoTools, RawSchema> {
112 type Answer = StepOutput;
113
114 fn into_config(self) -> AgentStepConfig {
115 self.into()
116 }
117
118 fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
119 Ok(output)
120 }
121}
122
123impl<T: DeserializeOwned> AgentStep for AgentConfig<NoTools, WithSchema<T>> {
124 type Answer = T;
125
126 fn into_config(self) -> AgentStepConfig {
127 self.into()
128 }
129
130 fn answer(output: StepOutput) -> Result<T, EngineError> {
131 output.json()
132 }
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138 use crate::executor::StepArtifacts;
139 use rust_decimal::Decimal;
140 use schemars::JsonSchema;
141 use serde::Deserialize;
142 use serde_json::{Value, json};
143
144 fn step_output(output: Value) -> StepOutput {
145 StepOutput {
146 output,
147 duration_ms: 0,
148 cost_usd: Decimal::ZERO,
149 input_tokens: None,
150 cache_read_input_tokens: None,
151 cache_creation_input_tokens: None,
152 output_tokens: None,
153 model: None,
154 debug_messages: None,
155 artifacts: StepArtifacts::default(),
156 account_id: None,
157 }
158 }
159
160 #[derive(Debug, PartialEq, Deserialize, JsonSchema)]
161 struct Verdict {
162 approved: bool,
163 }
164
165 #[test]
166 fn a_typed_config_answers_with_its_type() {
167 let answer = <AgentConfig<NoTools, WithSchema<Verdict>> as AgentStep>::answer(step_output(
168 json!({"approved": true}),
169 ))
170 .expect("matches Verdict");
171 assert_eq!(answer, Verdict { approved: true });
172 }
173
174 #[test]
175 fn a_typed_config_rejects_an_answer_of_another_shape() {
176 let err =
177 <AgentConfig<NoTools, WithSchema<Verdict>> as AgentStep>::answer(step_output(json!([
178 "approved"
179 ])))
180 .expect_err("not a Verdict");
181 assert!(matches!(err, EngineError::Serialization(_)));
182 }
183
184 #[test]
185 fn untyped_configs_answer_with_the_raw_output() {
186 let raw = <AgentConfig as AgentStep>::answer(step_output(json!("free text")))
187 .expect("raw output");
188 assert_eq!(raw.output, json!("free text"));
189 }
190
191 #[test]
192 fn into_config_keeps_the_settings() {
193 let config = AgentStepConfig::new("Review")
194 .max_turns(2)
195 .output::<Verdict>()
196 .into_config();
197 assert_eq!(config.max_turns, Some(2));
198 assert!(config.json_schema.is_some());
199
200 let config = AgentStepConfig::new("Explore")
201 .allow_tool(Tool::Grep)
202 .into_config();
203 assert_eq!(config.allowed_tools, vec!["Grep"]);
204 }
205
206 #[test]
207 fn builder() {
208 let config = AgentStepConfig::new("Review code")
209 .system_prompt("You are a code reviewer")
210 .model("haiku")
211 .max_budget_usd(0.50)
212 .max_turns(5)
213 .allow_tool(Tool::Read)
214 .working_dir("/repo")
215 .permission_mode(ironflow_core::operations::agent::PermissionMode::Auto);
216
217 assert_eq!(config.prompt, "Review code");
218 assert_eq!(config.system_prompt.unwrap(), "You are a code reviewer");
219 assert_eq!(config.model, "haiku");
220 assert_eq!(config.allowed_tools, vec!["Read"]);
221 assert!(config.json_schema.is_none());
222 }
223
224 #[test]
225 fn output_sets_schema_from_type() {
226 #[derive(serde::Deserialize, JsonSchema)]
227 #[allow(dead_code)]
228 struct Labels {
229 labels: Vec<String>,
230 }
231
232 let config = AgentStepConfig::new("Classify").output::<Labels>();
233
234 let schema = config.json_schema.expect("schema should be set");
235 assert!(schema.contains("labels"));
236 }
237
238 #[test]
239 fn output_schema_raw_sets_string() {
240 let raw = r#"{"type":"object"}"#;
241 let config = AgentStepConfig::new("Rate").output_schema_raw(raw);
242
243 assert_eq!(config.json_schema.as_deref(), Some(raw));
244 }
245
246 #[test]
247 fn output_overrides_previous_schema() {
248 #[derive(serde::Deserialize, JsonSchema)]
249 #[allow(dead_code)]
250 struct First {
251 a: String,
252 }
253
254 #[derive(serde::Deserialize, JsonSchema)]
255 #[allow(dead_code)]
256 struct Second {
257 b: i32,
258 }
259
260 let config = AgentStepConfig::new("Test")
261 .output::<First>()
262 .output::<Second>();
263
264 let schema = config.json_schema.expect("schema should be set");
265 assert!(!schema.contains("\"a\""));
266 assert!(schema.contains("\"b\""));
267 }
268
269 #[test]
270 fn output_schema_raw_overrides_typed_schema() {
271 #[derive(serde::Deserialize, JsonSchema)]
272 #[allow(dead_code)]
273 struct Typed {
274 field: String,
275 }
276
277 let raw = r#"{"type":"string"}"#;
278 let config = AgentStepConfig::new("Test")
279 .output::<Typed>()
280 .output_schema_raw(raw);
281
282 assert_eq!(config.json_schema.as_deref(), Some(raw));
283 }
284
285 #[test]
286 fn default_output_schema_is_none() {
287 let config = AgentStepConfig::new("Hello");
288 assert!(config.json_schema.is_none());
289 }
290
291 #[test]
292 fn serde_roundtrip_with_defaults() {
293 let json = r#"{"prompt":"hello"}"#;
294 let config: AgentConfig = serde_json::from_str(json).unwrap();
295 assert_eq!(config.prompt, "hello");
296 assert_eq!(config.model, "sonnet");
297 assert!(!config.verbose);
298 }
299
300 #[test]
301 fn serde_permission_mode_case_insensitive() {
302 let json = r#"{"prompt":"test","permission_mode":"auto"}"#;
303 let config: AgentConfig = serde_json::from_str(json).unwrap();
304 assert!(matches!(
305 config.permission_mode,
306 ironflow_core::operations::agent::PermissionMode::Auto
307 ));
308 }
309
310 #[test]
311 fn serde_output_schema_alias() {
312 let json = r#"{"prompt":"test","output_schema":"{\"type\":\"object\"}"}"#;
313 let config: AgentConfig = serde_json::from_str(json).unwrap();
314 assert_eq!(config.json_schema.as_deref(), Some(r#"{"type":"object"}"#));
315 }
316
317 #[test]
318 fn strict_mcp_config_defaults_to_false() {
319 let config = AgentStepConfig::new("test");
320 assert!(!config.strict_mcp_config);
321 }
322
323 #[test]
324 fn strict_mcp_config_builder_sets_flag() {
325 let config = AgentStepConfig::new("test").strict_mcp_config(true);
326 assert!(config.strict_mcp_config);
327 }
328
329 #[test]
330 fn strict_mcp_config_serde_default_when_missing() {
331 let json = r#"{"prompt":"test"}"#;
332 let config: AgentConfig = serde_json::from_str(json).unwrap();
333 assert!(!config.strict_mcp_config);
334 }
335
336 #[test]
337 fn strict_mcp_config_serde_roundtrip() {
338 let json = r#"{"prompt":"test","strict_mcp_config":true}"#;
339 let config: AgentConfig = serde_json::from_str(json).unwrap();
340 assert!(config.strict_mcp_config);
341 }
342}