ironflow_engine/config/
agent.rs1use serde::de::DeserializeOwned;
9
10pub use ironflow_core::prompt::{DEFAULT_UNTRUSTED_INSTRUCTION, Prompt};
11pub use ironflow_core::provider::{AgentConfig, Tool, ToolProfile};
12use ironflow_core::provider::{NoSchema, NoTools, RawSchema, WithSchema, WithTools};
13
14use crate::error::EngineError;
15use crate::executor::StepOutput;
16
17pub const DEFAULT_RESUME_PROMPT: &str = "The previous run of this task was interrupted. Continue from where you stopped and finish the task.";
29
30pub type AgentStepConfig = AgentConfig;
42
43pub trait AgentStep {
88 type Answer;
90
91 fn into_config(self) -> AgentStepConfig;
93
94 fn answer(output: StepOutput) -> Result<Self::Answer, EngineError>;
101}
102
103impl AgentStep for AgentConfig<NoTools, NoSchema> {
104 type Answer = StepOutput;
105
106 fn into_config(self) -> AgentStepConfig {
107 self
108 }
109
110 fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
111 Ok(output)
112 }
113}
114
115impl AgentStep for AgentConfig<WithTools, NoSchema> {
116 type Answer = StepOutput;
117
118 fn into_config(self) -> AgentStepConfig {
119 self.into()
120 }
121
122 fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
123 Ok(output)
124 }
125}
126
127impl AgentStep for AgentConfig<NoTools, RawSchema> {
128 type Answer = StepOutput;
129
130 fn into_config(self) -> AgentStepConfig {
131 self.into()
132 }
133
134 fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
135 Ok(output)
136 }
137}
138
139impl<T: DeserializeOwned> AgentStep for AgentConfig<NoTools, WithSchema<T>> {
140 type Answer = T;
141
142 fn into_config(self) -> AgentStepConfig {
143 self.into()
144 }
145
146 fn answer(output: StepOutput) -> Result<T, EngineError> {
147 output.json()
148 }
149}
150
151impl AgentStep for AgentConfig<WithTools, RawSchema> {
152 type Answer = StepOutput;
153
154 fn into_config(self) -> AgentStepConfig {
155 self.into()
156 }
157
158 fn answer(output: StepOutput) -> Result<StepOutput, EngineError> {
159 Ok(output)
160 }
161}
162
163impl<T: DeserializeOwned> AgentStep for AgentConfig<WithTools, WithSchema<T>> {
164 type Answer = T;
165
166 fn into_config(self) -> AgentStepConfig {
167 self.into()
168 }
169
170 fn answer(output: StepOutput) -> Result<T, EngineError> {
171 output.json()
172 }
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178 use crate::executor::StepArtifacts;
179 use ironflow_core::operations::agent::Effort;
180 use rust_decimal::Decimal;
181 use schemars::JsonSchema;
182 use serde::Deserialize;
183 use serde_json::{Value, from_value, json};
184
185 #[test]
186 fn step_config_effort_builder_keeps_value() {
187 let config = AgentStepConfig::new("x").effort(Effort::Low);
188 assert_eq!(config.effort, Some(Effort::Low));
189 }
190
191 #[test]
192 fn step_config_effort_deserializes_and_rejects_invalid() {
193 let config: AgentStepConfig = from_value(json!({"prompt": "x", "effort": "low"})).unwrap();
194 assert_eq!(config.effort, Some(Effort::Low));
195 assert!(from_value::<AgentStepConfig>(json!({"prompt": "x", "effort": "bogus"})).is_err());
196 }
197
198 fn step_output(output: Value) -> StepOutput {
199 StepOutput {
200 output,
201 duration_ms: 0,
202 cost_usd: Decimal::ZERO,
203 input_tokens: None,
204 cache_read_input_tokens: None,
205 cache_creation_input_tokens: None,
206 output_tokens: None,
207 model: None,
208 debug_messages: None,
209 artifacts: StepArtifacts::default(),
210 account_id: None,
211 environment_id: None,
212 }
213 }
214
215 #[derive(Debug, PartialEq, Deserialize, JsonSchema)]
216 struct Verdict {
217 approved: bool,
218 }
219
220 #[test]
221 fn a_typed_config_answers_with_its_type() {
222 let answer = <AgentConfig<NoTools, WithSchema<Verdict>> as AgentStep>::answer(step_output(
223 json!({"approved": true}),
224 ))
225 .expect("matches Verdict");
226 assert_eq!(answer, Verdict { approved: true });
227 }
228
229 #[test]
230 fn a_typed_config_rejects_an_answer_of_another_shape() {
231 let err =
232 <AgentConfig<NoTools, WithSchema<Verdict>> as AgentStep>::answer(step_output(json!([
233 "approved"
234 ])))
235 .expect_err("not a Verdict");
236 assert!(matches!(err, EngineError::Serialization(_)));
237 }
238
239 #[test]
240 fn untyped_configs_answer_with_the_raw_output() {
241 let raw = <AgentConfig as AgentStep>::answer(step_output(json!("free text")))
242 .expect("raw output");
243 assert_eq!(raw.output, json!("free text"));
244 }
245
246 #[test]
247 fn into_config_keeps_the_settings() {
248 let config = AgentStepConfig::new("Review")
249 .max_turns(2)
250 .output::<Verdict>()
251 .into_config();
252 assert_eq!(config.max_turns, Some(2));
253 assert!(config.json_schema.is_some());
254
255 let config = AgentStepConfig::new("Explore")
256 .allow_tool(Tool::Grep)
257 .into_config();
258 assert_eq!(config.allowed_tools, vec!["Grep"]);
259 }
260
261 #[test]
262 fn builder() {
263 let config = AgentStepConfig::new("Review code")
264 .system_prompt("You are a code reviewer")
265 .model("haiku")
266 .max_budget_usd(0.50)
267 .max_turns(5)
268 .allow_tool(Tool::Read)
269 .working_dir("/repo")
270 .permission_mode(ironflow_core::operations::agent::PermissionMode::Auto);
271
272 assert_eq!(config.prompt, "Review code");
273 assert_eq!(config.system_prompt.unwrap(), "You are a code reviewer");
274 assert_eq!(config.model, "haiku");
275 assert_eq!(config.allowed_tools, vec!["Read"]);
276 assert!(config.json_schema.is_none());
277 }
278
279 #[test]
280 fn output_sets_schema_from_type() {
281 #[derive(serde::Deserialize, JsonSchema)]
282 #[allow(dead_code)]
283 struct Labels {
284 labels: Vec<String>,
285 }
286
287 let config = AgentStepConfig::new("Classify").output::<Labels>();
288
289 let schema = config.json_schema.expect("schema should be set");
290 assert!(schema.contains("labels"));
291 }
292
293 #[test]
294 fn output_schema_raw_sets_string() {
295 let raw = r#"{"type":"object"}"#;
296 let config = AgentStepConfig::new("Rate").output_schema_raw(raw);
297
298 assert_eq!(config.json_schema.as_deref(), Some(raw));
299 }
300
301 #[test]
302 fn output_overrides_previous_schema() {
303 #[derive(serde::Deserialize, JsonSchema)]
304 #[allow(dead_code)]
305 struct First {
306 a: String,
307 }
308
309 #[derive(serde::Deserialize, JsonSchema)]
310 #[allow(dead_code)]
311 struct Second {
312 b: i32,
313 }
314
315 let config = AgentStepConfig::new("Test")
316 .output::<First>()
317 .output::<Second>();
318
319 let schema = config.json_schema.expect("schema should be set");
320 assert!(!schema.contains("\"a\""));
321 assert!(schema.contains("\"b\""));
322 }
323
324 #[test]
325 fn output_schema_raw_overrides_typed_schema() {
326 #[derive(serde::Deserialize, JsonSchema)]
327 #[allow(dead_code)]
328 struct Typed {
329 field: String,
330 }
331
332 let raw = r#"{"type":"string"}"#;
333 let config = AgentStepConfig::new("Test")
334 .output::<Typed>()
335 .output_schema_raw(raw);
336
337 assert_eq!(config.json_schema.as_deref(), Some(raw));
338 }
339
340 #[test]
341 fn default_output_schema_is_none() {
342 let config = AgentStepConfig::new("Hello");
343 assert!(config.json_schema.is_none());
344 }
345
346 #[test]
347 fn serde_roundtrip_with_defaults() {
348 let json = r#"{"prompt":"hello"}"#;
349 let config: AgentConfig = serde_json::from_str(json).unwrap();
350 assert_eq!(config.prompt, "hello");
351 assert_eq!(config.model, "sonnet");
352 assert!(!config.verbose);
353 }
354
355 #[test]
356 fn serde_permission_mode_case_insensitive() {
357 let json = r#"{"prompt":"test","permission_mode":"auto"}"#;
358 let config: AgentConfig = serde_json::from_str(json).unwrap();
359 assert!(matches!(
360 config.permission_mode,
361 ironflow_core::operations::agent::PermissionMode::Auto
362 ));
363 }
364
365 #[test]
366 fn serde_output_schema_alias() {
367 let json = r#"{"prompt":"test","output_schema":"{\"type\":\"object\"}"}"#;
368 let config: AgentConfig = serde_json::from_str(json).unwrap();
369 assert_eq!(config.json_schema.as_deref(), Some(r#"{"type":"object"}"#));
370 }
371
372 #[test]
373 fn strict_mcp_config_defaults_to_false() {
374 let config = AgentStepConfig::new("test");
375 assert!(!config.strict_mcp_config);
376 }
377
378 #[test]
379 fn strict_mcp_config_builder_sets_flag() {
380 let config = AgentStepConfig::new("test").strict_mcp_config(true);
381 assert!(config.strict_mcp_config);
382 }
383
384 #[test]
385 fn strict_mcp_config_serde_default_when_missing() {
386 let json = r#"{"prompt":"test"}"#;
387 let config: AgentConfig = serde_json::from_str(json).unwrap();
388 assert!(!config.strict_mcp_config);
389 }
390
391 #[test]
392 fn strict_mcp_config_serde_roundtrip() {
393 let json = r#"{"prompt":"test","strict_mcp_config":true}"#;
394 let config: AgentConfig = serde_json::from_str(json).unwrap();
395 assert!(config.strict_mcp_config);
396 }
397}