agent_base/engine/runtime/
plan.rs

1use serde_json::json;
2use tokio::sync::broadcast;
3
4use crate::types::{
5    AgentError, AgentEvent, AgentResult, ExecutionPlan, PlanStatus, RecoveryAction,
6    RunOutcome, SessionId, StepStatus,
7};
8use crate::engine::plan::{
9    AlwaysContinue, AbortOnFailure, PlanGenerator, PlanStore, RecoveryStrategy,
10    StepContinuePolicy, StepExecutor,
11};
12use crate::engine::runtime::event_bus::EventBus;
13use super::AgentRuntime;
14use std::sync::Arc;
15
16impl AgentRuntime {
17    /// Run a plan in **agentic** mode: each step becomes an agent turn.
18    ///
19    /// The agent receives step instructions as user input and decides autonomously
20    /// which tools to call. No `StepExecutor` is needed.
21    pub async fn run_plan_agentic<F>(
22        &self,
23        session_id: SessionId,
24        objective: &str,
25        generator: Arc<dyn PlanGenerator>,
26        plan_store: Option<Arc<dyn PlanStore>>,
27        mut on_event: F,
28    ) -> AgentResult<RunOutcome>
29    where
30        F: FnMut(AgentEvent) -> AgentResult<()> + Send,
31    {
32        let tool_definitions = self.tool_engine.definitions();
33        let mut event_rx = self.subscribe_events();
34
35        let mut plan = generator
36            .generate_plan(objective, "", &tool_definitions)
37            .await
38            .map_err(|e| AgentError::plan_generation(e.to_string()))?;
39
40        self.emit_and_drain(
41            AgentEvent::PlanGenerated {
42                session_id: session_id.clone(),
43                plan: plan.clone(),
44            },
45            &mut event_rx,
46            &mut on_event,
47        );
48
49        if let Some(store) = &plan_store {
50            store
51                .save_plan(&plan, json!({}))
52                .await
53                .map_err(|e| AgentError::plan_storage(e.to_string()))?;
54        }
55
56        plan.status = PlanStatus::Executing;
57
58        let result = self
59            .run_plan_steps(
60                &session_id,
61                &mut plan,
62                None::<Arc<dyn StepExecutor>>,
63                None::<Arc<dyn StepContinuePolicy>>,
64                None::<Arc<dyn RecoveryStrategy>>,
65                &mut event_rx,
66                &mut on_event,
67            )
68            .await;
69
70        if let Some(store) = &plan_store {
71            let _ = store.save_plan(&plan, json!({})).await;
72        }
73
74        result
75    }
76
77    /// Run a plan in **deterministic** mode: each step is executed directly
78    /// through the provided `StepExecutor`.
79    ///
80    /// Use this when you want the plan to be executed without LLM turn
81    /// overhead (e.g. predetermined SSH commands).
82    #[allow(clippy::too_many_arguments)]
83    pub async fn run_plan_deterministic<F>(
84        &self,
85        session_id: SessionId,
86        objective: &str,
87        generator: Arc<dyn PlanGenerator>,
88        executor: Arc<dyn StepExecutor>,
89        policy: Option<Arc<dyn StepContinuePolicy>>,
90        recovery: Option<Arc<dyn RecoveryStrategy>>,
91        plan_store: Option<Arc<dyn PlanStore>>,
92        mut on_event: F,
93    ) -> AgentResult<RunOutcome>
94    where
95        F: FnMut(AgentEvent) -> AgentResult<()> + Send,
96    {
97        let tool_definitions = self.tool_engine.definitions();
98        let mut event_rx = self.subscribe_events();
99
100        let mut plan = generator
101            .generate_plan(objective, "", &tool_definitions)
102            .await
103            .map_err(|e| AgentError::plan_generation(e.to_string()))?;
104
105        self.emit_and_drain(
106            AgentEvent::PlanGenerated {
107                session_id: session_id.clone(),
108                plan: plan.clone(),
109            },
110            &mut event_rx,
111            &mut on_event,
112        );
113
114        if let Some(store) = &plan_store {
115            store
116                .save_plan(&plan, json!({}))
117                .await
118                .map_err(|e| AgentError::plan_storage(e.to_string()))?;
119        }
120
121        plan.status = PlanStatus::Executing;
122
123        let result = self
124            .run_plan_steps(
125                &session_id,
126                &mut plan,
127                Some(executor),
128                policy.or_else(|| Some(Arc::new(AlwaysContinue))),
129                recovery.or_else(|| Some(Arc::new(AbortOnFailure))),
130                &mut event_rx,
131                &mut on_event,
132            )
133            .await;
134
135        if let Some(store) = &plan_store {
136            let _ = store.save_plan(&plan, json!({})).await;
137        }
138
139        result
140    }
141
142    /// Internal: shared plan-step execution loop.
143    ///
144    /// - If `executor` is `None` → agentic mode (step becomes agent turn).
145    /// - If `executor` is `Some` → deterministic mode (step goes to executor).
146    async fn run_plan_steps<F>(
147        &self,
148        session_id: &SessionId,
149        plan: &mut ExecutionPlan,
150        executor: Option<Arc<dyn StepExecutor>>,
151        policy: Option<Arc<dyn StepContinuePolicy>>,
152        recovery: Option<Arc<dyn RecoveryStrategy>>,
153        event_rx: &mut broadcast::Receiver<AgentEvent>,
154        on_event: &mut F,
155    ) -> AgentResult<RunOutcome>
156    where
157        F: FnMut(AgentEvent) -> AgentResult<()> + Send,
158    {
159        let mut i = 0usize;
160        while i < plan.steps.len() {
161            // Check dependencies before running
162            if plan.steps[i].status == StepStatus::Pending
163                && !self.check_dependencies_met(plan, i)
164            {
165                plan.steps[i].status = StepStatus::Skipped;
166                i += 1;
167                continue;
168            }
169
170            plan.steps[i].status = StepStatus::Running;
171
172            self.emit_and_drain(
173                AgentEvent::PlanStepStarted {
174                    session_id: session_id.clone(),
175                    step_id: plan.steps[i].id.clone(),
176                    step_description: plan.steps[i].description.clone(),
177                },
178                event_rx,
179                on_event,
180            );
181
182            // Step execution
183            let step = &plan.steps[i];
184            let step_result = if let Some(exec) = &executor {
185                // Deterministic mode
186                let should_continue = if let Some(p) = &policy {
187                    p.should_continue(plan, step)
188                        .await
189                        .unwrap_or(true)
190                } else {
191                    true
192                };
193
194                if !should_continue {
195                    Ok(crate::types::StepResult::success("Skipped", 0))
196                } else {
197                    exec.execute_step(step, &plan.context).await
198                }
199            } else {
200                // Agentic mode: run as a full agent turn
201                let mut step_events = Vec::new();
202                let outcome = self
203                    .run(session_id.clone(), |event| {
204                        step_events.push(event.clone());
205                        on_event(event)
206                    })
207                    .await;
208
209                let _ = EventBus::drain_async_events(event_rx, on_event);
210
211                match outcome {
212                    Ok(RunOutcome::Completed) => Ok(crate::types::StepResult::success("Step completed", 0)),
213                    Ok(RunOutcome::Failed { error }) => Ok(crate::types::StepResult::failure(error, 0)),
214                    Ok(RunOutcome::MaxTurnsExceeded { .. }) => Ok(crate::types::StepResult::failure("Max turns exceeded".to_string(), 0)),
215                    Ok(RunOutcome::Cancelled) => Ok(crate::types::StepResult::failure("Cancelled".to_string(), 0)),
216                    Err(e) => Err(e),
217                }
218            };
219
220            match step_result {
221                Ok(result) => {
222                    let error = result.error.clone().unwrap_or_default();
223                    let success = result.success;
224                    plan.steps[i].result = Some(result);
225
226                    if success {
227                        plan.steps[i].status = StepStatus::Completed;
228
229                        self.emit_and_drain(
230                            AgentEvent::PlanStepCompleted {
231                                session_id: session_id.clone(),
232                                step_id: plan.steps[i].id.clone(),
233                                success: true,
234                                result: plan.steps[i].result.as_ref().unwrap().output.clone(),
235                            },
236                            event_rx,
237                            on_event,
238                        );
239
240                        i += 1; // Move to next step
241                    } else {
242                        let action: RecoveryAction = if let Some(r) = &recovery {
243                            r.handle_step_failure(&plan.steps[i], &error, 0)
244                                .await
245                                .unwrap_or(RecoveryAction::Abort)
246                        } else {
247                            RecoveryAction::Abort
248                        };
249
250                        match action {
251                            RecoveryAction::Retry => {
252                                plan.steps[i].status = StepStatus::Pending;
253                                plan.steps[i].result = None;
254                                // i is NOT incremented, so this step will be retried
255                            }
256                            RecoveryAction::Skip => {
257                                plan.steps[i].status = StepStatus::Skipped;
258
259                                self.emit_and_drain(
260                                    AgentEvent::PlanStepCompleted {
261                                        session_id: session_id.clone(),
262                                        step_id: plan.steps[i].id.clone(),
263                                        success: false,
264                                        result: Some(format!("Skipped: {}", error)),
265                                    },
266                                    event_rx,
267                                    on_event,
268                                );
269
270                                i += 1; // Move to next step
271                            }
272                            RecoveryAction::Abort => {
273                                plan.steps[i].status = StepStatus::Failed;
274                                plan.status = PlanStatus::Failed;
275
276                                self.emit_and_drain(
277                                    AgentEvent::PlanStepCompleted {
278                                        session_id: session_id.clone(),
279                                        step_id: plan.steps[i].id.clone(),
280                                        success: false,
281                                        result: Some(error.clone()),
282                                    },
283                                    event_rx,
284                                    on_event,
285                                );
286
287                                self.emit_and_drain(
288                                    AgentEvent::PlanCompleted {
289                                        session_id: session_id.clone(),
290                                        plan_id: plan.id.clone(),
291                                        success: false,
292                                    },
293                                    event_rx,
294                                    on_event,
295                                );
296
297                                return Ok(RunOutcome::Failed {
298                                    error: format!("Step '{}' failed: {}", plan.steps[i].id, error),
299                                });
300                            }
301                        }
302                    }
303                }
304                Err(e) => {
305                    plan.steps[i].status = StepStatus::Failed;
306                    plan.status = PlanStatus::Failed;
307
308                    self.emit_and_drain(
309                        AgentEvent::PlanStepCompleted {
310                            session_id: session_id.clone(),
311                            step_id: plan.steps[i].id.clone(),
312                            success: false,
313                            result: Some(e.to_string()),
314                        },
315                        event_rx,
316                        on_event,
317                    );
318
319                    self.emit_and_drain(
320                        AgentEvent::PlanCompleted {
321                            session_id: session_id.clone(),
322                            plan_id: plan.id.clone(),
323                            success: false,
324                        },
325                        event_rx,
326                        on_event,
327                    );
328
329                    return Err(e);
330                }
331            }
332        }
333
334        plan.status = PlanStatus::Completed;
335
336        self.emit_and_drain(
337            AgentEvent::PlanCompleted {
338                session_id: session_id.clone(),
339                plan_id: plan.id.clone(),
340                success: true,
341            },
342            event_rx,
343            on_event,
344        );
345
346        Ok(RunOutcome::Completed)
347    }
348
349    fn check_dependencies_met(&self, plan: &ExecutionPlan, step_index: usize) -> bool {
350        let step = &plan.steps[step_index];
351        if step.dependencies.is_empty() {
352            return true;
353        }
354
355        step.dependencies.iter().all(|dep_id: &String| {
356            plan.steps
357                .iter()
358                .find(|s| s.id == *dep_id)
359                .map(|s| matches!(s.status, StepStatus::Completed | StepStatus::Skipped))
360                .unwrap_or(false)
361        })
362    }
363
364    fn emit_and_drain<F>(
365        &self,
366        event: AgentEvent,
367        event_rx: &mut broadcast::Receiver<AgentEvent>,
368        on_event: &mut F,
369    ) where
370        F: FnMut(AgentEvent) -> AgentResult<()>,
371    {
372        self.emit_event(event);
373        let _ = EventBus::drain_async_events(event_rx, on_event);
374    }
375}
agent_base/engine/runtime/plan.rs

agent_base/engine/runtime/
plan.rs
