agent_base/engine/
plan.rs

1use async_trait::async_trait;
2use serde::de::DeserializeOwned;
3use serde_json::Value;
4
5use crate::types::{AgentResult, ExecutionPlan, PlanStep, PlanStoreData, RecoveryAction, StepResult};
6
7// ---------------------------------------------------------------------------
8// Traits
9// ---------------------------------------------------------------------------
10
11/// Generates an `ExecutionPlan` from a high-level objective.
12///
13/// The generator may use LLM prompting, rule engines, or templates.
14#[async_trait]
15pub trait PlanGenerator: Send + Sync {
16    async fn generate_plan(
17        &self,
18        objective: &str,
19        context: &str,
20        tools: &[Value],
21    ) -> AgentResult<ExecutionPlan>;
22
23    async fn generate_plan_streaming(
24        &self,
25        objective: &str,
26        context: &str,
27        tools: &[Value],
28        on_generating: Box<dyn Fn() + Send>,
29        on_step_parsed: Box<dyn Fn(usize, String, String) + Send>,
30        on_raw_chunk: Box<dyn Fn(String) + Send>,
31    ) -> AgentResult<ExecutionPlan> {
32        let plan = self.generate_plan(objective, context, tools).await?;
33        on_generating();
34        for (i, step) in plan.steps.iter().enumerate() {
35            on_step_parsed(i, step.id.clone(), step.description.clone());
36        }
37        let plan_json = serde_json::to_string(&plan).unwrap_or_default();
38        on_raw_chunk(plan_json);
39        Ok(plan)
40    }
41}
42
43/// Executes a single `PlanStep` and returns its result.
44///
45/// Implementors know how to interpret `step.payload` for their domain.
46#[async_trait]
47pub trait StepExecutor: Send + Sync {
48    async fn execute_step(
49        &self,
50        step: &PlanStep,
51        plan_context: &Value,
52    ) -> AgentResult<StepResult>;
53}
54
55/// Decides whether the plan should continue executing a given step.
56#[async_trait]
57pub trait StepContinuePolicy: Send + Sync {
58    async fn should_continue(
59        &self,
60        plan: &ExecutionPlan,
61        current_step: &PlanStep,
62    ) -> AgentResult<bool>;
63}
64
65/// Decides what to do when a step fails.
66#[async_trait]
67pub trait RecoveryStrategy: Send + Sync {
68    async fn handle_step_failure(
69        &self,
70        step: &PlanStep,
71        error: &str,
72        retry_count: usize,
73    ) -> AgentResult<RecoveryAction>;
74}
75
76// ---------------------------------------------------------------------------
77// Default / convenience implementations
78// ---------------------------------------------------------------------------
79
80/// Always continues.
81pub struct AlwaysContinue;
82
83#[async_trait]
84impl StepContinuePolicy for AlwaysContinue {
85    async fn should_continue(
86        &self,
87        _plan: &ExecutionPlan,
88        _current_step: &PlanStep,
89    ) -> AgentResult<bool> {
90        Ok(true)
91    }
92}
93
94/// Always aborts on failure.
95pub struct AbortOnFailure;
96
97#[async_trait]
98impl RecoveryStrategy for AbortOnFailure {
99    async fn handle_step_failure(
100        &self,
101        _step: &PlanStep,
102        _error: &str,
103        _retry_count: usize,
104    ) -> AgentResult<RecoveryAction> {
105        Ok(RecoveryAction::Abort)
106    }
107}
108
109// ---------------------------------------------------------------------------
110// Streaming JSON parser (generic)
111// ---------------------------------------------------------------------------
112
113/// Parses JSON objects of type `T` from a stream of text chunks.
114///
115/// It scans for objects inside a JSON array (by default) and yields each
116/// fully-formed object as soon as braces are balanced. Useful when an LLM
117/// streams a JSON plan and you want to display / process steps incrementally.
118#[derive(Debug)]
119pub struct StreamingJsonParser<T> {
120    buffer: String,
121    scan_offset: usize,
122    items: Vec<T>,
123    items_start_byte: usize,
124    in_items: bool,
125    in_string: bool,
126    escape_next: bool,
127    array_key: Option<String>,
128}
129
130impl<T: DeserializeOwned + Clone> StreamingJsonParser<T> {
131    pub fn new() -> Self {
132        Self {
133            buffer: String::new(),
134            scan_offset: 0,
135            items: Vec::new(),
136            items_start_byte: 0,
137            in_items: false,
138            in_string: false,
139            escape_next: false,
140            array_key: None,
141        }
142    }
143
144    /// Set the array key to look for. e.g. `with_key("steps")` will look for
145    /// `"steps":[...]` in the JSON.
146    pub fn with_key(mut self, key: impl Into<String>) -> Self {
147        self.array_key = Some(key.into());
148        self
149    }
150
151    /// Append a new chunk and return any newly parsed items.
152    pub fn process_chunk(&mut self, chunk: &str) -> Vec<T> {
153        let mut new_items = Vec::new();
154        self.buffer.push_str(chunk);
155
156        if !self.in_items {
157            if let Some(pos) = self.find_items_array_start() {
158                self.items_start_byte = pos + 1;
159                self.scan_offset = 0;
160                self.in_items = true;
161            }
162        }
163
164        if self.in_items {
165            new_items = self.extract_items();
166            self.items.extend(new_items.clone());
167        }
168
169        new_items
170    }
171
172    /// Return all accumulated items so far.
173    pub fn accumulated(&self) -> &[T] {
174        &self.items
175    }
176
177    /// Consume parser and return the full raw text.
178    pub fn into_buffer(self) -> String {
179        self.buffer
180    }
181
182    fn find_items_array_start(&self) -> Option<usize> {
183        if let Some(ref key) = self.array_key {
184            if let Some(pos) = self.buffer.find(&format!("\"{}\"", key)) {
185                let after = &self.buffer[pos..];
186                if let Some(bracket_pos) = after.find('[') {
187                    return Some(pos + bracket_pos);
188                }
189            }
190        } else {
191            // Fallback: look for any quoted key followed by '['
192            if let Some(pos) = self.buffer.find('"') {
193                let after = &self.buffer[pos..];
194                if let Some(bracket_pos) = after.find('[') {
195                    return Some(pos + bracket_pos);
196                }
197            }
198        }
199        // Last fallback: raw array
200        self.buffer.find('[')
201    }
202
203    fn extract_items(&mut self) -> Vec<T> {
204        let mut results = Vec::new();
205        let slice = &self.buffer[self.items_start_byte..];
206        let mut brace_depth: i32 = 0;
207        let mut item_start_byte: Option<usize> = None;
208
209        for (byte_offset, ch) in slice.char_indices().skip(self.scan_offset) {
210            if self.escape_next {
211                self.escape_next = false;
212                self.scan_offset = byte_offset + ch.len_utf8();
213                continue;
214            }
215
216            if self.in_string {
217                if ch == '\\' {
218                    self.escape_next = true;
219                } else if ch == '"' {
220                    self.in_string = false;
221                }
222                self.scan_offset = byte_offset + ch.len_utf8();
223                continue;
224            }
225
226            match ch {
227                '"' => self.in_string = true,
228                '{' => {
229                    if brace_depth == 0 {
230                        let abs_byte = self.items_start_byte + byte_offset;
231                        item_start_byte = Some(abs_byte);
232                    }
233                    brace_depth += 1;
234                }
235                '}' => {
236                    brace_depth -= 1;
237                    if brace_depth == 0 {
238                        if let Some(start) = item_start_byte.take() {
239                            let end = self.items_start_byte + byte_offset + ch.len_utf8();
240                            let item_json = &self.buffer[start..end];
241                            if let Ok(item) = serde_json::from_str::<T>(item_json) {
242                                results.push(item);
243                            }
244                        }
245                    }
246                }
247                _ => {}
248            }
249
250            self.scan_offset = byte_offset + ch.len_utf8();
251        }
252
253        results
254    }
255}
256
257impl<T: DeserializeOwned + Clone> Default for StreamingJsonParser<T> {
258    fn default() -> Self {
259        Self::new()
260    }
261}
262
263// ---------------------------------------------------------------------------
264// PlanStore
265// ---------------------------------------------------------------------------
266
267#[async_trait]
268pub trait PlanStore: Send + Sync {
269    async fn save_plan(&self, plan: &ExecutionPlan, metadata: Value) -> AgentResult<()>;
270
271    async fn load_plan(&self, plan_id: &str) -> AgentResult<Option<PlanStoreData>>;
272
273    async fn delete_plan(&self, plan_id: &str) -> AgentResult<()>;
274
275    async fn list_plans(&self) -> AgentResult<Vec<String>>;
276}
277
278pub struct InMemoryPlanStore {
279    plans: tokio::sync::RwLock<std::collections::HashMap<String, PlanStoreData>>,
280}
281
282impl InMemoryPlanStore {
283    pub fn new() -> Self {
284        Self {
285            plans: tokio::sync::RwLock::new(std::collections::HashMap::new()),
286        }
287    }
288}
289
290impl Default for InMemoryPlanStore {
291    fn default() -> Self {
292        Self::new()
293    }
294}
295
296#[async_trait]
297impl PlanStore for InMemoryPlanStore {
298    async fn save_plan(&self, plan: &ExecutionPlan, metadata: Value) -> AgentResult<()> {
299        let mut plans = self.plans.write().await;
300        plans.insert(
301            plan.id.clone(),
302            PlanStoreData {
303                plan: plan.clone(),
304                metadata,
305            },
306        );
307        Ok(())
308    }
309
310    async fn load_plan(&self, plan_id: &str) -> AgentResult<Option<PlanStoreData>> {
311        let plans = self.plans.read().await;
312        Ok(plans.get(plan_id).cloned())
313    }
314
315    async fn delete_plan(&self, plan_id: &str) -> AgentResult<()> {
316        let mut plans = self.plans.write().await;
317        plans.remove(plan_id);
318        Ok(())
319    }
320
321    async fn list_plans(&self) -> AgentResult<Vec<String>> {
322        let plans = self.plans.read().await;
323        Ok(plans.keys().cloned().collect())
324    }
325}
agent_base/engine/plan.rs

agent_base/engine/
plan.rs
