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