Skip to main content

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}