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