Skip to main content

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::AgentStepConfig;
64    /// use ironflow_engine::error::EngineError;
65    /// use schemars::JsonSchema;
66    /// use serde::Deserialize;
67    ///
68    /// #[derive(Deserialize, JsonSchema)]
69    /// struct Review {
70    ///     approved: bool,
71    ///     comments: Vec<String>,
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::{AgentStepConfig, Tool};
115    /// use ironflow_engine::error::EngineError;
116    /// use schemars::JsonSchema;
117    /// use serde::Deserialize;
118    ///
119    /// #[derive(Deserialize, JsonSchema)]
120    /// struct Triage {
121    ///     test: String,
122    /// }
123    ///
124    /// # async fn example(ctx: &mut WorkflowContext) -> Result<(), EngineError> {
125    /// // A structured step cannot use tools: it only reads its prompt.
126    /// let reply = ctx
127    ///     .agent_with_meta(
128    ///         "triage",
129    ///         AgentStepConfig::new("Name the failing test in this CI log: ...")
130    ///             .max_budget_usd(0.50)
131    ///             .output::<Triage>(),
132    ///     )
133    ///     .await?;
134    /// if let Some(environment) = reply.environment_id.as_deref() {
135    ///     let prompt = format!("Fix the test {} in /workspace", reply.answer.test);
136    ///     ctx.agent(
137    ///         "fix",
138    ///         AgentStepConfig::new(&prompt)
139    ///             .allow_tool(Tool::Bash)
140    ///             .max_budget_usd(0.50)
141    ///             .resume_environment(environment),
142    ///     )
143    ///     .await?;
144    /// }
145    /// # Ok(())
146    /// # }
147    /// ```
148    pub async fn agent_with_meta<C: AgentStep>(
149        &mut self,
150        name: &str,
151        config: C,
152    ) -> Result<AgentReply<C::Answer>, EngineError> {
153        let output = self
154            .execute_step(
155                name,
156                StepKind::Agent,
157                StepConfig::Agent(config.into_config()),
158            )
159            .await?;
160        let environment_id = output.environment_id.clone();
161        let account_id = output.account_id;
162        Ok(AgentReply {
163            answer: C::answer(output)?,
164            environment_id,
165            account_id,
166        })
167    }
168}