ironflow_engine/context/steps/agent.rs
1//! Agent step for [`WorkflowContext`].
2
3use uuid::Uuid;
4
5use ironflow_store::models::StepKind;
6
7use crate::config::{AgentStep, StepConfig};
8use crate::context::WorkflowContext;
9use crate::error::EngineError;
10
11/// What an agent step answered, with the metadata the provider attached to it.
12///
13/// Returned by [`WorkflowContext::agent_with_meta`]. It keeps the typed answer
14/// of a step built with
15/// [`output::<T>()`](crate::config::AgentStepConfig::output) next to the
16/// environment and account the step ran on.
17///
18/// # Examples
19///
20/// ```
21/// use ironflow_engine::context::AgentReply;
22///
23/// let reply = AgentReply {
24/// answer: 42_u32,
25/// environment_id: Some("ironflow-env-0a1b2c".to_string()),
26/// account_id: None,
27/// };
28/// assert_eq!(reply.environment_id.as_deref(), Some("ironflow-env-0a1b2c"));
29/// ```
30#[derive(Debug, Clone, PartialEq)]
31pub struct AgentReply<A> {
32 /// The answer of the step: the typed `T` with `.output::<T>()`, the raw
33 /// [`StepOutput`](crate::executor::StepOutput) otherwise.
34 pub answer: A,
35 /// Claim name to pass to
36 /// [`AgentStepConfig::resume_environment`](crate::config::AgentStepConfig::resume_environment).
37 /// `None` on providers without persistent environments.
38 pub environment_id: Option<String>,
39 /// Provider Account the step ran on. `None` when the worker environment
40 /// was used.
41 pub account_id: Option<Uuid>,
42}
43
44impl WorkflowContext {
45 /// Execute an agent step.
46 ///
47 /// With [`output::<T>()`](crate::config::AgentStepConfig::output) on the
48 /// config, the step returns the `T` the agent answered; otherwise it
49 /// returns the raw [`StepOutput`](crate::executor::StepOutput). See
50 /// [`AgentStep`]. To read the environment or account of the step as well,
51 /// use [`agent_with_meta`](WorkflowContext::agent_with_meta).
52 ///
53 /// # Errors
54 ///
55 /// Returns [`EngineError`] if the agent invocation fails or the store
56 /// errors, and [`EngineError::Serialization`] if a typed answer does not
57 /// match its type.
58 ///
59 /// # Examples
60 ///
61 /// ```no_run
62 /// use ironflow_engine::context::WorkflowContext;
63 /// use ironflow_engine::config::{AgentAnswer, AgentStepConfig, AgentText};
64 /// use ironflow_engine::error::EngineError;
65 /// use schemars::JsonSchema;
66 /// use serde::Deserialize;
67 ///
68 /// #[derive(Deserialize, JsonSchema, AgentAnswer)]
69 /// struct Review {
70 /// approved: bool,
71 /// comments: Vec<AgentText>,
72 /// }
73 ///
74 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
75 /// let review = ctx
76 /// .agent(
77 /// "review",
78 /// AgentStepConfig::new("Review the code").max_turns(2).output::<Review>(),
79 /// )
80 /// .await?;
81 /// if !review.approved {
82 /// println!("{} comments", review.comments.len());
83 /// }
84 /// # Ok(())
85 /// # }
86 /// ```
87 pub async fn agent<C: AgentStep>(
88 &mut self,
89 name: &str,
90 config: C,
91 ) -> Result<C::Answer, EngineError> {
92 self.agent_with_meta(name, config)
93 .await
94 .map(|reply| reply.answer)
95 }
96
97 /// Execute an agent step and return its answer with the step metadata.
98 ///
99 /// Behaves like [`agent`](WorkflowContext::agent), but the
100 /// [`AgentReply`] also carries the `environment_id` and `account_id` of
101 /// the step, which a typed answer would otherwise hide. On replay the
102 /// reply carries the ids persisted with the step.
103 ///
104 /// # Errors
105 ///
106 /// Returns [`EngineError`] if the agent invocation fails or the store
107 /// errors, and [`EngineError::Serialization`] if a typed answer does not
108 /// match its type.
109 ///
110 /// # Examples
111 ///
112 /// ```no_run
113 /// use ironflow_engine::context::WorkflowContext;
114 /// use ironflow_engine::config::{AgentAnswer, AgentStepConfig, AgentText, Prompt, Tool};
115 /// use ironflow_engine::error::EngineError;
116 /// use schemars::JsonSchema;
117 /// use serde::Deserialize;
118 ///
119 /// #[derive(Deserialize, JsonSchema, AgentAnswer)]
120 /// struct Triage {
121 /// test: AgentText,
122 /// }
123 ///
124 /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
125 /// let reply = ctx
126 /// .agent_with_meta(
127 /// "triage",
128 /// AgentStepConfig::new("Name the failing test in this CI log: ...")
129 /// .max_budget_usd(0.50)
130 /// .output::<Triage>(),
131 /// )
132 /// .await?;
133 /// if let Some(environment) = reply.environment_id.as_deref() {
134 /// // The test name was written by an agent: it enters the prompt as data.
135 /// let prompt = Prompt::new()
136 /// .text("Fix this test in /workspace:")
137 /// .untrusted(&reply.answer.test);
138 /// ctx.agent(
139 /// "fix",
140 /// AgentStepConfig::new(&prompt.to_string())
141 /// .allow_tool(Tool::Bash)
142 /// .max_budget_usd(0.50)
143 /// .resume_environment(environment),
144 /// )
145 /// .await?;
146 /// }
147 /// # Ok(())
148 /// # }
149 /// ```
150 pub async fn agent_with_meta<C: AgentStep>(
151 &mut self,
152 name: &str,
153 config: C,
154 ) -> Result<AgentReply<C::Answer>, EngineError> {
155 let output = self
156 .execute_step(
157 name,
158 StepKind::Agent,
159 StepConfig::Agent(config.into_config()),
160 )
161 .await?;
162 let environment_id = output.environment_id.clone();
163 let account_id = output.account_id;
164 Ok(AgentReply {
165 answer: C::answer(output)?,
166 environment_id,
167 account_id,
168 })
169 }
170}