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}