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}