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}