ironflow_engine/config/agent_role.rs
1//! Reusable agent with a fixed role.
2//!
3//! An [`AgentRole`] is an agent defined once: its persona, tools, rights and
4//! output type. A workflow calls it with
5//! `ctx.agent("step", Role.task(input))`. The task is an [`AgentStep`], so it is
6//! persisted, replayed and run by the worker exactly like a hand-built
7//! [`AgentStepConfig`].
8
9use std::any::type_name;
10use std::collections::BTreeMap;
11use std::fmt;
12use std::marker::PhantomData;
13use std::time::Duration;
14
15use ironflow_core::operations::agent::{Effort, PermissionMode};
16use ironflow_core::retry::RetryPolicy;
17use schemars::JsonSchema;
18use serde::de::DeserializeOwned;
19
20use super::agent::{AgentStep, AgentStepConfig, Tool};
21use super::agent_answer::AgentAnswer;
22use crate::error::EngineError;
23use crate::executor::StepOutput;
24
25/// Longest [`AgentRole::name`] accepted, in bytes.
26///
27/// [`AgentRole::task`] panics on a longer name, and the API rejects a longer
28/// `agent_role` filter.
29///
30/// # Examples
31///
32/// ```
33/// use ironflow_engine::config::MAX_ROLE_NAME_LEN;
34///
35/// assert_eq!(MAX_ROLE_NAME_LEN, 128);
36/// ```
37pub const MAX_ROLE_NAME_LEN: usize = 128;
38
39/// A reusable agent with a fixed role: persona, tools, rights and output type.
40///
41/// Only [`prompt`](Self::prompt) is required. The persona, tools, permission
42/// mode, MCP config and output type are locked by the role: the caller of
43/// [`task`](Self::task) cannot change them (a role with more rights is another
44/// role). The model, budget, turn limit, effort, timeout and per-call settings
45/// are defaults the caller may override on the [`AgentTask`].
46///
47/// # Examples
48///
49/// ```
50/// use ironflow_engine::config::{AgentAnswer, AgentRole, AgentStep, Tool};
51/// use schemars::JsonSchema;
52/// use serde::Deserialize;
53///
54/// #[derive(Deserialize, JsonSchema, AgentAnswer)]
55/// struct Verdict {
56/// approved: bool,
57/// }
58///
59/// struct Reviewer;
60///
61/// impl AgentRole for Reviewer {
62/// type Input = String;
63/// type Output = Verdict;
64///
65/// fn prompt(&self, input: String) -> String {
66/// format!("Review: {input}")
67/// }
68///
69/// fn append_system_prompt(&self) -> Option<&str> {
70/// Some("You are a strict reviewer.")
71/// }
72///
73/// fn tools(&self) -> Vec<Tool> {
74/// vec![Tool::Read]
75/// }
76///
77/// fn max_turns(&self) -> Option<u32> {
78/// Some(4)
79/// }
80/// }
81///
82/// let config = Reviewer.task("diff".to_string()).max_budget_usd(0.10).into_config();
83/// assert_eq!(config.append_system_prompt.as_deref(), Some("You are a strict reviewer."));
84/// assert_eq!(config.allowed_tools, vec!["Read"]);
85/// assert!(config.json_schema.is_some());
86/// assert_eq!(config.max_budget_usd, Some(0.10));
87/// ```
88///
89/// Calling it from a workflow:
90///
91/// ```no_run
92/// # use ironflow_engine::config::AgentRole;
93/// # use ironflow_engine::context::WorkflowContext;
94/// # use ironflow_engine::error::EngineError;
95/// # use schemars::JsonSchema;
96/// # use ironflow_engine::config::AgentAnswer;
97/// # use serde::Deserialize;
98/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
99/// # struct Verdict { approved: bool }
100/// # struct Reviewer;
101/// # impl AgentRole for Reviewer {
102/// # type Input = String;
103/// # type Output = Verdict;
104/// # fn prompt(&self, input: String) -> String { input }
105/// # }
106/// # async fn example(ctx: &mut WorkflowContext, diff: String) -> Result<(), EngineError> {
107/// let verdict: Verdict = ctx.agent("review", Reviewer.task(diff)).await?;
108/// # Ok(())
109/// # }
110/// ```
111pub trait AgentRole {
112 /// What the caller passes to [`task`](Self::task) to build the prompt.
113 type Input;
114
115 /// What the role answers; `ctx.agent` returns it. Its free-text fields
116 /// are [`AgentText`](super::AgentText): see [`AgentAnswer`].
117 type Output: DeserializeOwned + JsonSchema + AgentAnswer;
118
119 /// Name of the role, recorded on every step it runs and shown by the API,
120 /// CLI and dashboard. Defaults to the type name without its module path
121 /// (`Reviewer`). Override it to keep the name stable across a type rename.
122 ///
123 /// Renaming the type changes the default name: the steps and costs of the
124 /// role are then split between the old and the new name. The caller of
125 /// [`task`](Self::task) cannot override it.
126 ///
127 /// # Examples
128 ///
129 /// ```
130 /// # use ironflow_engine::config::{AgentRole, AgentStep};
131 /// # use schemars::JsonSchema;
132 /// # use ironflow_engine::config::AgentAnswer;
133 /// # use serde::Deserialize;
134 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
135 /// # struct Verdict { approved: bool }
136 /// struct Reviewer;
137 ///
138 /// impl AgentRole for Reviewer {
139 /// type Input = String;
140 /// type Output = Verdict;
141 ///
142 /// fn prompt(&self, input: String) -> String {
143 /// input
144 /// }
145 /// }
146 ///
147 /// struct Named;
148 ///
149 /// impl AgentRole for Named {
150 /// type Input = String;
151 /// type Output = Verdict;
152 ///
153 /// fn name(&self) -> &str {
154 /// "code-reviewer"
155 /// }
156 ///
157 /// fn prompt(&self, input: String) -> String {
158 /// input
159 /// }
160 /// }
161 ///
162 /// assert_eq!(Reviewer.name(), "Reviewer");
163 /// assert_eq!(Named.name(), "code-reviewer");
164 /// let config = Named.task("diff".to_string()).into_config();
165 /// assert_eq!(config.role.as_deref(), Some("code-reviewer"));
166 /// ```
167 fn name(&self) -> &str {
168 short_type_name(type_name::<Self>())
169 }
170
171 /// The prompt of one task. Required.
172 fn prompt(&self, input: Self::Input) -> String;
173
174 /// Replaces the provider system prompt. Claude Code loses its own (tools,
175 /// skills) when it is replaced. `None` by default.
176 fn system_prompt(&self) -> Option<&str> {
177 None
178 }
179
180 /// Persona appended to the provider system prompt. Prefer it to
181 /// [`system_prompt`](Self::system_prompt) for a role with tools. `None` by
182 /// default.
183 fn append_system_prompt(&self) -> Option<&str> {
184 None
185 }
186
187 /// Tools the role may use. None by default.
188 fn tools(&self) -> Vec<Tool> {
189 Vec::new()
190 }
191
192 /// Permission mode of the role; `None` keeps the
193 /// `AgentConfig::new` default.
194 fn permission_mode(&self) -> Option<PermissionMode> {
195 None
196 }
197
198 /// MCP configuration, as JSON or a path. `None` by default.
199 fn mcp_config(&self) -> Option<&str> {
200 None
201 }
202
203 /// Restrict the MCP servers to [`mcp_config`](Self::mcp_config). `false` by
204 /// default.
205 fn strict_mcp_config(&self) -> bool {
206 false
207 }
208
209 /// Default model; `None` keeps the `AgentConfig::new` default. The
210 /// caller may override it.
211 fn model(&self) -> Option<&str> {
212 None
213 }
214
215 /// Default budget in USD. The caller may override it.
216 fn max_budget_usd(&self) -> Option<f64> {
217 None
218 }
219
220 /// Default turn limit. The caller may override it.
221 ///
222 /// A structured output needs at least 2 turns; like [`AgentStepConfig`],
223 /// the role does not validate it.
224 fn max_turns(&self) -> Option<u32> {
225 None
226 }
227
228 /// Default effort. The caller may override it.
229 fn effort(&self) -> Option<Effort> {
230 None
231 }
232
233 /// Default step timeout, see
234 /// [`AgentConfig::timeout`](ironflow_core::provider::AgentConfig::timeout).
235 /// The caller may override it.
236 fn timeout(&self) -> Option<Duration> {
237 None
238 }
239
240 /// Build one task of this role.
241 ///
242 /// # Panics
243 ///
244 /// Panics if [`name`](Self::name) is empty or blank, has leading or
245 /// trailing whitespace, or is longer than
246 /// [`MAX_ROLE_NAME_LEN`] bytes.
247 ///
248 /// # Examples
249 ///
250 /// ```
251 /// # use ironflow_engine::config::{AgentRole, AgentStep};
252 /// # use schemars::JsonSchema;
253 /// # use ironflow_engine::config::AgentAnswer;
254 /// # use serde::Deserialize;
255 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
256 /// # struct Verdict { approved: bool }
257 /// # struct Reviewer;
258 /// # impl AgentRole for Reviewer {
259 /// # type Input = String;
260 /// # type Output = Verdict;
261 /// # fn prompt(&self, input: String) -> String { format!("Review: {input}") }
262 /// # }
263 /// let config = Reviewer.task("diff".to_string()).into_config();
264 /// assert_eq!(config.prompt, "Review: diff");
265 /// ```
266 fn task(&self, input: Self::Input) -> AgentTask<Self::Output> {
267 AgentTask::from_role(self, input)
268 }
269}
270
271/// Type name without its module path and generic arguments:
272/// `a::b::Wrapper<c::Inner>` gives `Wrapper`.
273fn short_type_name(full: &'static str) -> &'static str {
274 let without_generics = full.split('<').next().unwrap_or(full);
275 without_generics
276 .rsplit("::")
277 .next()
278 .unwrap_or(without_generics)
279}
280
281/// One call of an [`AgentRole`], answering with the `T` of the role.
282///
283/// The role locks its persona, tools, permission mode, MCP config and output
284/// type: none of them has a setter here, and each of these calls does not
285/// compile.
286///
287/// The system prompt is locked:
288///
289/// ```compile_fail,E0599
290/// # use ironflow_engine::config::AgentRole;
291/// # use schemars::JsonSchema;
292/// # use ironflow_engine::config::AgentAnswer;
293/// # use serde::Deserialize;
294/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
295/// # struct Verdict { approved: bool }
296/// # struct Reviewer;
297/// # impl AgentRole for Reviewer {
298/// # type Input = String;
299/// # type Output = Verdict;
300/// # fn prompt(&self, input: String) -> String { input }
301/// # }
302/// let _ = Reviewer.task("diff".to_string()).system_prompt("x");
303/// ```
304///
305/// The persona is locked:
306///
307/// ```compile_fail,E0599
308/// # use ironflow_engine::config::AgentRole;
309/// # use schemars::JsonSchema;
310/// # use ironflow_engine::config::AgentAnswer;
311/// # use serde::Deserialize;
312/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
313/// # struct Verdict { approved: bool }
314/// # struct Reviewer;
315/// # impl AgentRole for Reviewer {
316/// # type Input = String;
317/// # type Output = Verdict;
318/// # fn prompt(&self, input: String) -> String { input }
319/// # }
320/// let _ = Reviewer.task("diff".to_string()).append_system_prompt("x");
321/// ```
322///
323/// The tools are locked:
324///
325/// ```compile_fail,E0599
326/// # use ironflow_engine::config::{AgentRole, Tool};
327/// # use schemars::JsonSchema;
328/// # use ironflow_engine::config::AgentAnswer;
329/// # use serde::Deserialize;
330/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
331/// # struct Verdict { approved: bool }
332/// # struct Reviewer;
333/// # impl AgentRole for Reviewer {
334/// # type Input = String;
335/// # type Output = Verdict;
336/// # fn prompt(&self, input: String) -> String { input }
337/// # }
338/// let _ = Reviewer.task("diff".to_string()).allow_tool(Tool::Bash);
339/// ```
340///
341/// The permission mode is locked:
342///
343/// ```compile_fail,E0599
344/// # use ironflow_core::operations::agent::PermissionMode;
345/// # use ironflow_engine::config::AgentRole;
346/// # use schemars::JsonSchema;
347/// # use ironflow_engine::config::AgentAnswer;
348/// # use serde::Deserialize;
349/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
350/// # struct Verdict { approved: bool }
351/// # struct Reviewer;
352/// # impl AgentRole for Reviewer {
353/// # type Input = String;
354/// # type Output = Verdict;
355/// # fn prompt(&self, input: String) -> String { input }
356/// # }
357/// let _ = Reviewer
358/// .task("diff".to_string())
359/// .permission_mode(PermissionMode::BypassPermissions);
360/// ```
361///
362/// The MCP config is locked:
363///
364/// ```compile_fail,E0599
365/// # use ironflow_engine::config::AgentRole;
366/// # use schemars::JsonSchema;
367/// # use ironflow_engine::config::AgentAnswer;
368/// # use serde::Deserialize;
369/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
370/// # struct Verdict { approved: bool }
371/// # struct Reviewer;
372/// # impl AgentRole for Reviewer {
373/// # type Input = String;
374/// # type Output = Verdict;
375/// # fn prompt(&self, input: String) -> String { input }
376/// # }
377/// let _ = Reviewer.task("diff".to_string()).mcp_config("{}");
378/// ```
379///
380/// The output type is locked:
381///
382/// ```compile_fail,E0599
383/// # use ironflow_engine::config::AgentRole;
384/// # use schemars::JsonSchema;
385/// # use ironflow_engine::config::AgentAnswer;
386/// # use serde::Deserialize;
387/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
388/// # struct Verdict { approved: bool }
389/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
390/// # struct Other { n: u8 }
391/// # struct Reviewer;
392/// # impl AgentRole for Reviewer {
393/// # type Input = String;
394/// # type Output = Verdict;
395/// # fn prompt(&self, input: String) -> String { input }
396/// # }
397/// let _ = Reviewer.task("diff".to_string()).output::<Other>();
398/// ```
399///
400/// The role name is locked:
401///
402/// ```compile_fail,E0599
403/// # use ironflow_engine::config::AgentRole;
404/// # use schemars::JsonSchema;
405/// # use ironflow_engine::config::AgentAnswer;
406/// # use serde::Deserialize;
407/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
408/// # struct Verdict { approved: bool }
409/// # struct Reviewer;
410/// # impl AgentRole for Reviewer {
411/// # type Input = String;
412/// # type Output = Verdict;
413/// # fn prompt(&self, input: String) -> String { input }
414/// # }
415/// let _ = Reviewer.task("diff".to_string()).role("x");
416/// ```
417///
418/// Taking the raw config with [`AgentStep::into_config`] and editing it loses
419/// the typed answer (`ctx.agent` then returns a [`StepOutput`]), so there is no
420/// silent override.
421///
422/// A call-site override changes the step config, so it changes the hash of the
423/// `RecordReplayProvider` fixture: a fixture recorded for the role defaults
424/// does not match an overridden task.
425///
426/// # Examples
427///
428/// ```
429/// # use ironflow_engine::config::{AgentRole, AgentStep};
430/// # use schemars::JsonSchema;
431/// # use ironflow_engine::config::AgentAnswer;
432/// # use serde::Deserialize;
433/// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
434/// # struct Verdict { approved: bool }
435/// # struct Reviewer;
436/// # impl AgentRole for Reviewer {
437/// # type Input = String;
438/// # type Output = Verdict;
439/// # fn prompt(&self, input: String) -> String { input }
440/// # }
441/// let task = Reviewer.task("diff".to_string()).model("haiku").allow_failure();
442/// let config = task.into_config();
443/// assert_eq!(config.model, "haiku");
444/// assert!(config.allow_failure);
445/// ```
446pub struct AgentTask<T> {
447 config: AgentStepConfig,
448 _output: PhantomData<fn() -> T>,
449}
450
451// Manual impls: a derive would require `T: Debug` / `T: Clone`, though `T` is
452// only a marker (same reasoning as `WithSchema` in `ironflow-core`).
453impl<T> fmt::Debug for AgentTask<T> {
454 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
455 f.debug_struct("AgentTask")
456 .field("config", &self.config)
457 .finish()
458 }
459}
460
461impl<T> Clone for AgentTask<T> {
462 fn clone(&self) -> Self {
463 Self {
464 config: self.config.clone(),
465 _output: PhantomData,
466 }
467 }
468}
469
470impl<T: JsonSchema> AgentTask<T> {
471 /// Builds the task from the role defaults. `pub(crate)`: no other crate can
472 /// wrap an arbitrary config in an [`AgentTask`].
473 ///
474 /// # Panics
475 ///
476 /// Panics if [`AgentRole::name`] is empty or blank, has leading or
477 /// trailing whitespace, or is longer than
478 /// [`MAX_ROLE_NAME_LEN`] bytes.
479 pub(crate) fn from_role<R: AgentRole<Output = T> + ?Sized>(role: &R, input: R::Input) -> Self {
480 let name = role.name();
481 assert!(!name.trim().is_empty(), "AgentRole::name must not be empty");
482 assert!(
483 name.trim() == name,
484 "AgentRole::name must not have leading or trailing whitespace"
485 );
486 assert!(
487 name.len() <= MAX_ROLE_NAME_LEN,
488 "AgentRole::name must be at most {MAX_ROLE_NAME_LEN} bytes"
489 );
490 let mut config = AgentStepConfig::new(&role.prompt(input));
491 if let Some(prompt) = role.system_prompt() {
492 config = config.system_prompt(prompt);
493 }
494 if let Some(prompt) = role.append_system_prompt() {
495 config = config.append_system_prompt(prompt);
496 }
497 if let Some(mcp) = role.mcp_config() {
498 config = config.mcp_config(mcp);
499 }
500 if role.strict_mcp_config() {
501 config = config.strict_mcp_config(true);
502 }
503 if let Some(model) = role.model() {
504 config = config.model(model);
505 }
506 if let Some(budget) = role.max_budget_usd() {
507 config = config.max_budget_usd(budget);
508 }
509 if let Some(turns) = role.max_turns() {
510 config = config.max_turns(turns);
511 }
512 if let Some(limit) = role.timeout() {
513 config = config.timeout(limit);
514 }
515 if let Some(effort) = role.effort() {
516 config = config.effort(effort);
517 }
518 if let Some(mode) = role.permission_mode() {
519 config = config.permission_mode(mode);
520 }
521 for tool in role.tools() {
522 config = config.allow_tool(tool).into();
523 }
524 config = config.output::<T>().into();
525 config.role = Some(name.to_string());
526 Self {
527 config,
528 _output: PhantomData,
529 }
530 }
531}
532
533impl<T> AgentTask<T> {
534 /// Override the model of the role for this call.
535 ///
536 /// # Examples
537 ///
538 /// ```
539 /// # use ironflow_engine::config::{AgentRole, AgentStep};
540 /// # use schemars::JsonSchema;
541 /// # use ironflow_engine::config::AgentAnswer;
542 /// # use serde::Deserialize;
543 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
544 /// # struct Verdict { approved: bool }
545 /// # struct Reviewer;
546 /// # impl AgentRole for Reviewer {
547 /// # type Input = String;
548 /// # type Output = Verdict;
549 /// # fn prompt(&self, input: String) -> String { input }
550 /// # }
551 /// let config = Reviewer.task("diff".to_string()).model("haiku").into_config();
552 /// assert_eq!(config.model, "haiku");
553 /// ```
554 pub fn model(mut self, model: &str) -> Self {
555 self.config = self.config.model(model);
556 self
557 }
558
559 /// Override the budget of the role for this call.
560 ///
561 /// # Examples
562 ///
563 /// ```
564 /// # use ironflow_engine::config::{AgentRole, AgentStep};
565 /// # use schemars::JsonSchema;
566 /// # use ironflow_engine::config::AgentAnswer;
567 /// # use serde::Deserialize;
568 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
569 /// # struct Verdict { approved: bool }
570 /// # struct Reviewer;
571 /// # impl AgentRole for Reviewer {
572 /// # type Input = String;
573 /// # type Output = Verdict;
574 /// # fn prompt(&self, input: String) -> String { input }
575 /// # }
576 /// let config = Reviewer.task("diff".to_string()).max_budget_usd(0.10).into_config();
577 /// assert_eq!(config.max_budget_usd, Some(0.10));
578 /// ```
579 pub fn max_budget_usd(mut self, budget: f64) -> Self {
580 self.config = self.config.max_budget_usd(budget);
581 self
582 }
583
584 /// Override the turn limit of the role for this call.
585 ///
586 /// # Examples
587 ///
588 /// ```
589 /// # use ironflow_engine::config::{AgentRole, AgentStep};
590 /// # use schemars::JsonSchema;
591 /// # use ironflow_engine::config::AgentAnswer;
592 /// # use serde::Deserialize;
593 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
594 /// # struct Verdict { approved: bool }
595 /// # struct Reviewer;
596 /// # impl AgentRole for Reviewer {
597 /// # type Input = String;
598 /// # type Output = Verdict;
599 /// # fn prompt(&self, input: String) -> String { input }
600 /// # }
601 /// let config = Reviewer.task("diff".to_string()).max_turns(6).into_config();
602 /// assert_eq!(config.max_turns, Some(6));
603 /// ```
604 pub fn max_turns(mut self, turns: u32) -> Self {
605 self.config = self.config.max_turns(turns);
606 self
607 }
608
609 /// Override the effort of the role for this call.
610 ///
611 /// # Examples
612 ///
613 /// ```
614 /// # use ironflow_engine::config::{AgentRole, AgentStep};
615 /// # use ironflow_core::operations::agent::Effort;
616 /// # use schemars::JsonSchema;
617 /// # use ironflow_engine::config::AgentAnswer;
618 /// # use serde::Deserialize;
619 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
620 /// # struct Verdict { approved: bool }
621 /// # struct Reviewer;
622 /// # impl AgentRole for Reviewer {
623 /// # type Input = String;
624 /// # type Output = Verdict;
625 /// # fn prompt(&self, input: String) -> String { input }
626 /// # }
627 /// let config = Reviewer.task("diff".to_string()).effort(Effort::High).into_config();
628 /// assert!(matches!(config.effort, Some(Effort::High)));
629 /// ```
630 pub fn effort(mut self, effort: Effort) -> Self {
631 self.config = self.config.effort(effort);
632 self
633 }
634
635 /// Override the step timeout of the role for this call, see
636 /// [`AgentConfig::timeout`](ironflow_core::provider::AgentConfig::timeout).
637 ///
638 /// # Panics
639 ///
640 /// Panics if `limit` is zero.
641 ///
642 /// # Examples
643 ///
644 /// ```
645 /// # use std::time::Duration;
646 /// # use ironflow_engine::config::{AgentRole, AgentStep};
647 /// # use schemars::JsonSchema;
648 /// # use ironflow_engine::config::AgentAnswer;
649 /// # use serde::Deserialize;
650 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
651 /// # struct Verdict { approved: bool }
652 /// # struct Reviewer;
653 /// # impl AgentRole for Reviewer {
654 /// # type Input = String;
655 /// # type Output = Verdict;
656 /// # fn prompt(&self, input: String) -> String { input }
657 /// # }
658 /// let task = Reviewer.task("diff".to_string());
659 /// let config = task.timeout(Duration::from_secs(600)).into_config();
660 /// assert_eq!(config.timeout, Some(Duration::from_secs(600)));
661 /// ```
662 pub fn timeout(mut self, limit: Duration) -> Self {
663 self.config = self.config.timeout(limit);
664 self
665 }
666
667 /// Set the working directory of the call.
668 ///
669 /// # Examples
670 ///
671 /// ```
672 /// # use ironflow_engine::config::{AgentRole, AgentStep};
673 /// # use schemars::JsonSchema;
674 /// # use ironflow_engine::config::AgentAnswer;
675 /// # use serde::Deserialize;
676 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
677 /// # struct Verdict { approved: bool }
678 /// # struct Reviewer;
679 /// # impl AgentRole for Reviewer {
680 /// # type Input = String;
681 /// # type Output = Verdict;
682 /// # fn prompt(&self, input: String) -> String { input }
683 /// # }
684 /// let config = Reviewer.task("diff".to_string()).working_dir("/repo").into_config();
685 /// assert_eq!(config.working_dir.as_deref(), Some("/repo"));
686 /// ```
687 pub fn working_dir(mut self, dir: &str) -> Self {
688 self.config = self.config.working_dir(dir);
689 self
690 }
691
692 /// Resume a Claude Code session.
693 ///
694 /// # Examples
695 ///
696 /// ```
697 /// # use ironflow_engine::config::{AgentRole, AgentStep};
698 /// # use schemars::JsonSchema;
699 /// # use ironflow_engine::config::AgentAnswer;
700 /// # use serde::Deserialize;
701 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
702 /// # struct Verdict { approved: bool }
703 /// # struct Reviewer;
704 /// # impl AgentRole for Reviewer {
705 /// # type Input = String;
706 /// # type Output = Verdict;
707 /// # fn prompt(&self, input: String) -> String { input }
708 /// # }
709 /// let config = Reviewer.task("diff".to_string()).resume("sess").into_config();
710 /// assert_eq!(config.resume_session_id.as_deref(), Some("sess"));
711 /// ```
712 pub fn resume(mut self, session_id: &str) -> Self {
713 self.config = self.config.resume(session_id);
714 self
715 }
716
717 /// Pin the session id of the call.
718 ///
719 /// # Examples
720 ///
721 /// ```
722 /// # use ironflow_engine::config::{AgentRole, AgentStep};
723 /// # use schemars::JsonSchema;
724 /// # use ironflow_engine::config::AgentAnswer;
725 /// # use serde::Deserialize;
726 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
727 /// # struct Verdict { approved: bool }
728 /// # struct Reviewer;
729 /// # impl AgentRole for Reviewer {
730 /// # type Input = String;
731 /// # type Output = Verdict;
732 /// # fn prompt(&self, input: String) -> String { input }
733 /// # }
734 /// let config = Reviewer.task("diff".to_string()).session_id("sid").into_config();
735 /// assert_eq!(config.session_id.as_deref(), Some("sid"));
736 /// ```
737 pub fn session_id(mut self, session_id: &str) -> Self {
738 self.config = self.config.session_id(session_id);
739 self
740 }
741
742 /// Set the prompt sent when the step is resumed after a lost lease.
743 ///
744 /// # Examples
745 ///
746 /// ```
747 /// # use ironflow_engine::config::{AgentRole, AgentStep};
748 /// # use schemars::JsonSchema;
749 /// # use ironflow_engine::config::AgentAnswer;
750 /// # use serde::Deserialize;
751 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
752 /// # struct Verdict { approved: bool }
753 /// # struct Reviewer;
754 /// # impl AgentRole for Reviewer {
755 /// # type Input = String;
756 /// # type Output = Verdict;
757 /// # fn prompt(&self, input: String) -> String { input }
758 /// # }
759 /// let config = Reviewer.task("diff".to_string()).resume_prompt("go on").into_config();
760 /// assert_eq!(config.resume_prompt.as_deref(), Some("go on"));
761 /// ```
762 pub fn resume_prompt(mut self, prompt: &str) -> Self {
763 self.config = self.config.resume_prompt(prompt);
764 self
765 }
766
767 /// Resume the environment handed out by an earlier step.
768 ///
769 /// # Examples
770 ///
771 /// ```
772 /// # use ironflow_engine::config::{AgentRole, AgentStep};
773 /// # use schemars::JsonSchema;
774 /// # use ironflow_engine::config::AgentAnswer;
775 /// # use serde::Deserialize;
776 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
777 /// # struct Verdict { approved: bool }
778 /// # struct Reviewer;
779 /// # impl AgentRole for Reviewer {
780 /// # type Input = String;
781 /// # type Output = Verdict;
782 /// # fn prompt(&self, input: String) -> String { input }
783 /// # }
784 /// let config = Reviewer.task("diff".to_string()).resume_environment("env-1").into_config();
785 /// assert_eq!(config.resume_environment_id.as_deref(), Some("env-1"));
786 /// ```
787 pub fn resume_environment(mut self, environment_id: &str) -> Self {
788 self.config = self.config.resume_environment(environment_id);
789 self
790 }
791
792 /// Set the retry policy of the step.
793 ///
794 /// # Examples
795 ///
796 /// ```
797 /// # use ironflow_engine::config::{AgentRole, AgentStep};
798 /// # use ironflow_core::retry::RetryPolicy;
799 /// # use schemars::JsonSchema;
800 /// # use ironflow_engine::config::AgentAnswer;
801 /// # use serde::Deserialize;
802 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
803 /// # struct Verdict { approved: bool }
804 /// # struct Reviewer;
805 /// # impl AgentRole for Reviewer {
806 /// # type Input = String;
807 /// # type Output = Verdict;
808 /// # fn prompt(&self, input: String) -> String { input }
809 /// # }
810 /// let config = Reviewer.task("diff".to_string()).retry_policy(RetryPolicy::new(2)).into_config();
811 /// assert_eq!(config.retry.map(|r| r.max_retries()), Some(2));
812 /// ```
813 pub fn retry_policy(mut self, policy: RetryPolicy) -> Self {
814 self.config = self.config.retry_policy(policy);
815 self
816 }
817
818 /// Let the run go on when this step fails.
819 ///
820 /// # Examples
821 ///
822 /// ```
823 /// # use ironflow_engine::config::{AgentRole, AgentStep};
824 /// # use schemars::JsonSchema;
825 /// # use ironflow_engine::config::AgentAnswer;
826 /// # use serde::Deserialize;
827 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
828 /// # struct Verdict { approved: bool }
829 /// # struct Reviewer;
830 /// # impl AgentRole for Reviewer {
831 /// # type Input = String;
832 /// # type Output = Verdict;
833 /// # fn prompt(&self, input: String) -> String { input }
834 /// # }
835 /// let config = Reviewer.task("diff".to_string()).allow_failure().into_config();
836 /// assert!(config.allow_failure);
837 /// ```
838 pub fn allow_failure(mut self) -> Self {
839 self.config = self.config.allow_failure();
840 self
841 }
842
843 /// Record the agent conversation of the call.
844 ///
845 /// # Examples
846 ///
847 /// ```
848 /// # use ironflow_engine::config::{AgentRole, AgentStep};
849 /// # use schemars::JsonSchema;
850 /// # use ironflow_engine::config::AgentAnswer;
851 /// # use serde::Deserialize;
852 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
853 /// # struct Verdict { approved: bool }
854 /// # struct Reviewer;
855 /// # impl AgentRole for Reviewer {
856 /// # type Input = String;
857 /// # type Output = Verdict;
858 /// # fn prompt(&self, input: String) -> String { input }
859 /// # }
860 /// let config = Reviewer.task("diff".to_string()).verbose(true).into_config();
861 /// assert!(config.verbose);
862 /// ```
863 pub fn verbose(mut self, enabled: bool) -> Self {
864 self.config = self.config.verbose(enabled);
865 self
866 }
867
868 /// Add a label to the pod running the call.
869 ///
870 /// # Examples
871 ///
872 /// ```
873 /// # use ironflow_engine::config::{AgentRole, AgentStep};
874 /// # use schemars::JsonSchema;
875 /// # use ironflow_engine::config::AgentAnswer;
876 /// # use serde::Deserialize;
877 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
878 /// # struct Verdict { approved: bool }
879 /// # struct Reviewer;
880 /// # impl AgentRole for Reviewer {
881 /// # type Input = String;
882 /// # type Output = Verdict;
883 /// # fn prompt(&self, input: String) -> String { input }
884 /// # }
885 /// let config = Reviewer.task("diff".to_string()).pod_label("team", "qa").into_config();
886 /// assert_eq!(config.pod_labels.get("team").map(String::as_str), Some("qa"));
887 /// ```
888 pub fn pod_label(mut self, key: &str, value: &str) -> Self {
889 self.config = self.config.pod_label(key, value);
890 self
891 }
892
893 /// Set the labels of the pod running the call.
894 ///
895 /// # Examples
896 ///
897 /// ```
898 /// # use ironflow_engine::config::{AgentRole, AgentStep};
899 /// # use std::collections::BTreeMap;
900 /// # use schemars::JsonSchema;
901 /// # use ironflow_engine::config::AgentAnswer;
902 /// # use serde::Deserialize;
903 /// # #[derive(Deserialize, JsonSchema, AgentAnswer)]
904 /// # struct Verdict { approved: bool }
905 /// # struct Reviewer;
906 /// # impl AgentRole for Reviewer {
907 /// # type Input = String;
908 /// # type Output = Verdict;
909 /// # fn prompt(&self, input: String) -> String { input }
910 /// # }
911 /// let labels = BTreeMap::from([("team".to_string(), "qa".to_string())]);
912 /// let config = Reviewer.task("diff".to_string()).pod_labels(labels.clone()).into_config();
913 /// assert_eq!(config.pod_labels, labels);
914 /// ```
915 pub fn pod_labels(mut self, labels: BTreeMap<String, String>) -> Self {
916 self.config = self.config.pod_labels(labels);
917 self
918 }
919}
920
921impl<T: DeserializeOwned + AgentAnswer> AgentStep for AgentTask<T> {
922 type Answer = T;
923
924 fn into_config(self) -> AgentStepConfig {
925 self.config
926 }
927
928 fn answer(output: StepOutput) -> Result<T, EngineError> {
929 output.json()
930 }
931}
932
933#[cfg(test)]
934mod tests {
935 use super::*;
936 use crate::executor::StepArtifacts;
937 use rust_decimal::Decimal;
938 use serde::Deserialize;
939 use serde_json::{Value, json, to_value};
940
941 #[derive(Debug, PartialEq, Deserialize, JsonSchema, AgentAnswer)]
942 struct Verdict {
943 approved: bool,
944 }
945
946 /// Implements only the required method.
947 struct Bare;
948
949 impl AgentRole for Bare {
950 type Input = String;
951 type Output = Verdict;
952
953 fn prompt(&self, input: String) -> String {
954 input
955 }
956 }
957
958 /// Locks MCP in strict mode.
959 struct Strict;
960
961 impl AgentRole for Strict {
962 type Input = String;
963 type Output = Verdict;
964
965 fn prompt(&self, input: String) -> String {
966 input
967 }
968
969 fn mcp_config(&self) -> Option<&str> {
970 Some("{}")
971 }
972
973 fn strict_mcp_config(&self) -> bool {
974 true
975 }
976 }
977
978 /// Sets every default the other test roles leave out.
979 struct Full;
980
981 impl AgentRole for Full {
982 type Input = String;
983 type Output = Verdict;
984
985 fn prompt(&self, input: String) -> String {
986 input
987 }
988
989 fn system_prompt(&self) -> Option<&str> {
990 Some("base prompt")
991 }
992
993 fn append_system_prompt(&self) -> Option<&str> {
994 Some("persona")
995 }
996
997 fn max_budget_usd(&self) -> Option<f64> {
998 Some(0.5)
999 }
1000
1001 fn effort(&self) -> Option<Effort> {
1002 Some(Effort::Low)
1003 }
1004
1005 fn timeout(&self) -> Option<Duration> {
1006 Some(Duration::from_secs(300))
1007 }
1008 }
1009
1010 fn step_output(output: Value) -> StepOutput {
1011 StepOutput {
1012 output,
1013 duration_ms: 0,
1014 cost_usd: Decimal::ZERO,
1015 input_tokens: None,
1016 cache_read_input_tokens: None,
1017 cache_creation_input_tokens: None,
1018 output_tokens: None,
1019 model: None,
1020 debug_messages: None,
1021 artifacts: StepArtifacts::default(),
1022 account_id: None,
1023 environment_id: None,
1024 }
1025 }
1026
1027 #[test]
1028 fn a_role_with_only_a_prompt_takes_the_agent_config_new_defaults() {
1029 let task = Bare.task("p".to_string());
1030 let mut expected = AgentStepConfig::new("p").output::<Verdict>().into_config();
1031 expected.role = Some("Bare".to_string());
1032 assert_eq!(
1033 to_value(task.into_config()).expect("serialize"),
1034 to_value(expected).expect("serialize")
1035 );
1036 }
1037
1038 #[test]
1039 fn a_role_without_tools_has_no_allowed_tools_and_keeps_its_schema() {
1040 let config = Bare.task("p".to_string()).into_config();
1041 assert!(config.allowed_tools.is_empty());
1042 assert!(
1043 config
1044 .json_schema
1045 .as_deref()
1046 .is_some_and(|schema| schema.contains("approved"))
1047 );
1048 }
1049
1050 #[test]
1051 fn strict_mcp_config_of_the_role_is_applied() {
1052 let config = Strict.task("p".to_string()).into_config();
1053 assert!(config.strict_mcp_config);
1054 assert_eq!(config.mcp_config.as_deref(), Some("{}"));
1055 assert!(!Bare.task("p".to_string()).into_config().strict_mcp_config);
1056 }
1057
1058 #[test]
1059 fn role_defaults_for_prompts_budget_and_effort_are_applied() {
1060 let config = Full.task("p".to_string()).into_config();
1061 assert_eq!(config.system_prompt.as_deref(), Some("base prompt"));
1062 assert_eq!(config.append_system_prompt.as_deref(), Some("persona"));
1063 assert_eq!(config.max_budget_usd, Some(0.5));
1064 assert_eq!(config.effort, Some(Effort::Low));
1065 let overridden = Full
1066 .task("p".to_string())
1067 .max_budget_usd(1.5)
1068 .effort(Effort::High)
1069 .into_config();
1070 assert_eq!(overridden.max_budget_usd, Some(1.5));
1071 assert_eq!(overridden.effort, Some(Effort::High));
1072 assert_eq!(overridden.system_prompt.as_deref(), Some("base prompt"));
1073 }
1074
1075 #[test]
1076 fn role_timeout_is_a_default_the_task_overrides() {
1077 let config = Full.task("p".to_string()).into_config();
1078 assert_eq!(config.timeout, Some(Duration::from_secs(300)));
1079 let overridden = Full
1080 .task("p".to_string())
1081 .timeout(Duration::from_secs(1800))
1082 .into_config();
1083 assert_eq!(overridden.timeout, Some(Duration::from_secs(1800)));
1084 assert_eq!(Bare.task("p".to_string()).into_config().timeout, None);
1085 }
1086
1087 #[test]
1088 fn task_answer_reads_the_typed_output() {
1089 let answer =
1090 <AgentTask<Verdict> as AgentStep>::answer(step_output(json!({"approved": true})))
1091 .expect("matches Verdict");
1092 assert_eq!(answer, Verdict { approved: true });
1093 }
1094
1095 #[test]
1096 fn task_answer_rejects_another_shape() {
1097 let err = <AgentTask<Verdict> as AgentStep>::answer(step_output(json!(["approved"])))
1098 .expect_err("not a Verdict");
1099 assert!(matches!(err, EngineError::Serialization(_)));
1100 }
1101
1102 #[test]
1103 fn retry_policy_and_resume_settings_are_applied() {
1104 let mut labels = BTreeMap::new();
1105 labels.insert("team".to_string(), "qa".to_string());
1106 let config = Bare
1107 .task("p".to_string())
1108 .retry_policy(RetryPolicy::new(2))
1109 .resume("sess")
1110 .session_id("sid")
1111 .resume_prompt("go on")
1112 .resume_environment("env-1")
1113 .verbose(true)
1114 .pod_labels(labels.clone())
1115 .into_config();
1116 assert_eq!(config.retry.map(|r| r.max_retries()), Some(2));
1117 assert_eq!(config.resume_session_id.as_deref(), Some("sess"));
1118 assert_eq!(config.session_id.as_deref(), Some("sid"));
1119 assert_eq!(config.resume_prompt.as_deref(), Some("go on"));
1120 assert_eq!(config.resume_environment_id.as_deref(), Some("env-1"));
1121 assert!(config.verbose);
1122 assert_eq!(config.pod_labels, labels);
1123 }
1124
1125 #[test]
1126 fn clone_keeps_the_config() {
1127 let task = Bare
1128 .task("p".to_string())
1129 .model("haiku")
1130 .pod_label("a", "b");
1131 let copy = task.clone();
1132 assert_eq!(
1133 to_value(task.into_config()).expect("serialize"),
1134 to_value(copy.into_config()).expect("serialize")
1135 );
1136 }
1137
1138 /// Overrides the default name.
1139 struct Named;
1140
1141 impl AgentRole for Named {
1142 type Input = String;
1143 type Output = Verdict;
1144
1145 fn name(&self) -> &str {
1146 "code-reviewer"
1147 }
1148
1149 fn prompt(&self, input: String) -> String {
1150 input
1151 }
1152 }
1153
1154 /// Returns a blank name.
1155 struct Blank;
1156
1157 impl AgentRole for Blank {
1158 type Input = String;
1159 type Output = Verdict;
1160
1161 fn name(&self) -> &str {
1162 " "
1163 }
1164
1165 fn prompt(&self, input: String) -> String {
1166 input
1167 }
1168 }
1169
1170 /// Returns a name one byte over the limit.
1171 struct TooLong(String);
1172
1173 impl AgentRole for TooLong {
1174 type Input = String;
1175 type Output = Verdict;
1176
1177 fn name(&self) -> &str {
1178 &self.0
1179 }
1180
1181 fn prompt(&self, input: String) -> String {
1182 input
1183 }
1184 }
1185
1186 #[test]
1187 fn the_default_name_is_the_short_type_name() {
1188 assert_eq!(Bare.name(), "Bare");
1189 let config = Bare.task("p".to_string()).into_config();
1190 assert_eq!(config.role.as_deref(), Some("Bare"));
1191 }
1192
1193 #[test]
1194 fn an_overridden_name_is_recorded() {
1195 assert_eq!(Named.name(), "code-reviewer");
1196 let config = Named.task("p".to_string()).into_config();
1197 assert_eq!(config.role.as_deref(), Some("code-reviewer"));
1198 }
1199
1200 #[test]
1201 fn short_type_name_strips_path_and_generics() {
1202 assert_eq!(short_type_name("a::b::Wrapper<c::Inner>"), "Wrapper");
1203 assert_eq!(short_type_name("Plain"), "Plain");
1204 }
1205
1206 #[test]
1207 fn call_site_overrides_keep_the_role_name() {
1208 let config = Bare
1209 .task("p".to_string())
1210 .model("haiku")
1211 .allow_failure()
1212 .into_config();
1213 assert_eq!(config.model, "haiku");
1214 assert!(config.allow_failure);
1215 assert_eq!(config.role.as_deref(), Some("Bare"));
1216 }
1217
1218 #[test]
1219 #[should_panic(expected = "must not be empty")]
1220 fn an_empty_role_name_panics() {
1221 let _ = Blank.task("p".to_string());
1222 }
1223
1224 #[test]
1225 #[should_panic(expected = "must be at most")]
1226 fn a_too_long_role_name_panics() {
1227 let _ = TooLong("r".repeat(MAX_ROLE_NAME_LEN + 1)).task("p".to_string());
1228 }
1229
1230 #[test]
1231 #[should_panic(expected = "must not have leading or trailing whitespace")]
1232 fn a_role_name_with_surrounding_whitespace_panics() {
1233 let _ = TooLong(" Reviewer".to_string()).task("p".to_string());
1234 }
1235
1236 #[test]
1237 fn a_role_name_at_the_limit_is_accepted() {
1238 let name = "r".repeat(MAX_ROLE_NAME_LEN);
1239 let config = TooLong(name.clone()).task("p".to_string()).into_config();
1240 assert_eq!(config.role, Some(name));
1241 }
1242}