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