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