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