Skip to main content

echo_core/agent/
plan.rs

1//! Plan-and-Execute core types, Planner trait, and PlanStore trait
2
3use crate::error::Result;
4use futures::future::BoxFuture;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7
8// ── Compute simple similarity between two texts (0.0-1.0) ──────────────────
9
10fn text_similarity(a: &str, b: &str) -> f32 {
11    if a.is_empty() && b.is_empty() {
12        return 1.0;
13    }
14    if a.is_empty() || b.is_empty() {
15        return 0.0;
16    }
17    let set_a: std::collections::HashSet<char> = a.to_lowercase().chars().collect();
18    let set_b: std::collections::HashSet<char> = b.to_lowercase().chars().collect();
19    let intersection = set_a.intersection(&set_b).count() as f32;
20    let union = set_a.union(&set_b).count() as f32;
21    intersection / union
22}
23
24fn generate_plan_id() -> Option<String> {
25    Some(format!("plan_{}", uuid::Uuid::new_v4().as_simple()))
26}
27
28fn now_secs() -> u64 {
29    crate::utils::time::now_secs()
30}
31
32// ── Plan ────────────────────────────────────────────────────────────────────
33
34/// Execution plan
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct Plan {
37    /// Plan unique ID (auto-generated)
38    #[serde(default = "generate_plan_id")]
39    pub id: Option<String>,
40    /// Human-readable slug (e.g., "swift-fox")
41    #[serde(default)]
42    pub slug: Option<String>,
43    /// Version number (incremented on each replan)
44    #[serde(default)]
45    pub version: u32,
46    /// List of steps in the plan
47    pub steps: Vec<PlanStep>,
48    /// Overall goal description of the plan
49    pub goal: Option<String>,
50    /// Parent plan ID (for incremental replanning tracking)
51    #[serde(default)]
52    pub parent_plan_id: Option<String>,
53    /// Additional metadata
54    #[serde(default)]
55    pub metadata: HashMap<String, serde_json::Value>,
56    /// Creation timestamp (seconds)
57    #[serde(default)]
58    pub created_at: u64,
59    /// Last update timestamp (seconds)
60    #[serde(default)]
61    pub updated_at: u64,
62}
63
64impl Plan {
65    /// Create a new plan
66    pub fn new(steps: Vec<PlanStep>) -> Self {
67        let now = now_secs();
68        Self {
69            id: generate_plan_id(),
70            slug: None,
71            version: 1,
72            steps,
73            goal: None,
74            parent_plan_id: None,
75            metadata: HashMap::new(),
76            created_at: now,
77            updated_at: now,
78        }
79    }
80
81    /// Set the plan goal description
82    pub fn with_goal(mut self, goal: impl Into<String>) -> Self {
83        self.goal = Some(goal.into());
84        self
85    }
86
87    /// Set a human-readable slug
88    pub fn with_slug(mut self, slug: impl Into<String>) -> Self {
89        self.slug = Some(slug.into());
90        self
91    }
92
93    /// Set the parent plan ID (for incremental replanning)
94    pub fn with_parent(mut self, parent_id: impl Into<String>) -> Self {
95        self.parent_plan_id = Some(parent_id.into());
96        self
97    }
98
99    /// Add metadata
100    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
101        self.metadata.insert(key.into(), value);
102        self
103    }
104
105    /// Return the number of completed steps
106    pub fn completed_count(&self) -> usize {
107        self.steps
108            .iter()
109            .filter(|s| s.status == StepStatus::Completed)
110            .count()
111    }
112
113    /// Check if the plan is fully completed
114    pub fn is_completed(&self) -> bool {
115        self.steps.iter().all(|s| s.status == StepStatus::Completed)
116    }
117
118    /// Touch the update timestamp
119    pub fn touch(&mut self) {
120        self.updated_at = now_secs();
121    }
122
123    // ── Validation & Auto-Fix ──────────────────────────────────────────────
124
125    /// Validate the plan, returning all discovered issues
126    pub fn validate(&self) -> Vec<PlanValidationIssue> {
127        let mut issues = Vec::new();
128
129        if self.steps.is_empty() {
130            issues.push(PlanValidationIssue {
131                severity: IssueSeverity::Error,
132                message: "Plan has no steps".to_string(),
133                fix: Some("Add at least one step".to_string()),
134            });
135        }
136
137        let desc_set: std::collections::HashSet<&str> =
138            self.steps.iter().map(|s| s.description.as_str()).collect();
139
140        for (i, step) in self.steps.iter().enumerate() {
141            for dep in &step.dependencies {
142                if dep == &format!("step_{}", i) {
143                    issues.push(PlanValidationIssue {
144                        severity: IssueSeverity::Error,
145                        message: format!("Step {} depends on itself", i),
146                        fix: Some("Remove self-dependency".to_string()),
147                    });
148                }
149            }
150        }
151
152        for (i, step) in self.steps.iter().enumerate() {
153            for dep in &step.dependencies {
154                if !dep.starts_with("step_") {
155                    let matched = self.resolve_dependency(dep).is_some();
156                    if !matched && !desc_set.contains(dep.as_str()) {
157                        issues.push(PlanValidationIssue {
158                            severity: IssueSeverity::Warning,
159                            message: format!("Step {} has unresolvable dependency: {}", i, dep),
160                            fix: Some("Remove or fix the dependency reference".to_string()),
161                        });
162                    }
163                } else if let Ok(idx) = dep.trim_start_matches("step_").parse::<usize>()
164                    && idx >= self.steps.len()
165                {
166                    issues.push(PlanValidationIssue {
167                        severity: IssueSeverity::Error,
168                        message: format!("Step {} depends on non-existent step index {}", i, idx),
169                        fix: Some("Fix the dependency index".to_string()),
170                    });
171                }
172            }
173        }
174
175        for (i, step) in self.steps.iter().enumerate() {
176            if step.description.trim().is_empty() {
177                issues.push(PlanValidationIssue {
178                    severity: IssueSeverity::Warning,
179                    message: format!("Step {} has empty description", i),
180                    fix: Some("Provide a meaningful description".to_string()),
181                });
182            }
183        }
184
185        issues
186    }
187
188    /// Auto-fix fixable issues
189    pub fn auto_fix(&mut self) -> Vec<String> {
190        let mut fixes = Vec::new();
191        let steps_len = self.steps.len();
192
193        use std::collections::{HashMap, HashSet};
194        let mut all_deps = HashSet::new();
195        for step in self.steps.iter() {
196            for dep in &step.dependencies {
197                all_deps.insert(dep.clone());
198            }
199        }
200
201        let mut dependency_resolvable: HashMap<String, bool> = HashMap::new();
202        for dep in &all_deps {
203            let resolvable = self.resolve_dependency(dep).is_some();
204            dependency_resolvable.insert(dep.clone(), resolvable);
205        }
206
207        for (i, step) in self.steps.iter_mut().enumerate() {
208            let self_dep = format!("step_{}", i);
209            let before = step.dependencies.len();
210            step.dependencies.retain(|d| d != &self_dep);
211            if step.dependencies.len() < before {
212                fixes.push(format!("Removed self-dependency from step {}", i));
213            }
214
215            step.dependencies.retain(|d| {
216                if let Some(idx_str) = d.strip_prefix("step_")
217                    && let Ok(idx) = idx_str.parse::<usize>()
218                {
219                    return idx < steps_len;
220                }
221                *dependency_resolvable.get(d).unwrap_or(&false)
222            });
223
224            if step.description.trim().is_empty() {
225                step.description = format!("Unnamed step {}", i);
226                fixes.push(format!("Filled empty description for step {}", i));
227            }
228        }
229
230        if !fixes.is_empty() {
231            self.touch();
232        }
233
234        fixes
235    }
236
237    /// Resolve a dependency string to a step index, returning whether fuzzy matching was used
238    pub fn resolve_dependency_with_fuzzy(&self, dep: &str) -> (Option<usize>, bool) {
239        if let Some(idx_str) = dep.strip_prefix("step_")
240            && let Ok(idx) = idx_str.parse::<usize>()
241            && idx < self.steps.len()
242        {
243            return (Some(idx), false);
244        }
245        let result = self.resolve_dependency(dep);
246        (result, result.is_some())
247    }
248
249    /// Resolve a dependency string to a step index
250    ///
251    /// Resolution strategy (in order of priority):
252    /// 1. **Exact index reference** — `step_N` where N is a valid index
253    /// 2. **Exact description match** — dep string exactly equals a step description
254    /// 3. **Word-boundary substring** — dep is a word-bounded substring of a step description
255    /// 4. **Reverse containment** — step description is a substring of dep
256    /// 5. **Fuzzy text similarity** — Jaccard character-set similarity (threshold 0.6)
257    pub fn resolve_dependency(&self, dep: &str) -> Option<usize> {
258        // Strategy 1: Exact index reference (step_N)
259        if let Some(idx_str) = dep.strip_prefix("step_")
260            && let Ok(idx) = idx_str.parse::<usize>()
261            && idx < self.steps.len()
262        {
263            return Some(idx);
264        }
265
266        // Strategy 2: Exact description match (case-insensitive)
267        let dep_lower = dep.to_lowercase();
268        for (idx, step) in self.steps.iter().enumerate() {
269            if step.description.to_lowercase() == dep_lower {
270                return Some(idx);
271            }
272        }
273
274        // Strategies 3-5: fuzzy matching (with reduced confidence)
275        let mut candidates = Vec::new();
276
277        for (idx, step) in self.steps.iter().enumerate() {
278            if dep.len() >= 3 && step.description.contains(dep) {
279                let desc_lower = step.description.to_lowercase();
280
281                let mut positions = Vec::new();
282                let mut start = 0;
283                while let Some(pos) = desc_lower[start..].find(&dep_lower) {
284                    let actual_pos = start + pos;
285                    positions.push(actual_pos);
286                    start = actual_pos + 1;
287                }
288
289                let mut has_word_boundary = false;
290                for &pos in &positions {
291                    let prev_is_boundary = pos == 0
292                        || !desc_lower
293                            .chars()
294                            .nth(pos - 1)
295                            .is_some_and(|c| c.is_alphanumeric());
296                    let next_pos = pos + dep.len();
297                    let next_is_boundary = next_pos >= desc_lower.len()
298                        || !desc_lower
299                            .chars()
300                            .nth(next_pos)
301                            .is_some_and(|c| c.is_alphanumeric());
302
303                    if prev_is_boundary && next_is_boundary {
304                        has_word_boundary = true;
305                        break;
306                    }
307                }
308
309                if has_word_boundary {
310                    candidates.push((idx, 1.0));
311                } else if dep.len() >= 3 {
312                    candidates.push((idx, 0.8));
313                }
314                continue;
315            }
316
317            if step.description.len() >= 3 && dep.contains(&step.description) {
318                candidates.push((idx, 0.7));
319                continue;
320            }
321
322            if dep.len() >= 5 {
323                let similarity = text_similarity(dep, &step.description);
324                if similarity >= 0.6 {
325                    candidates.push((idx, similarity));
326                }
327            }
328        }
329
330        candidates.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
331        candidates.first().map(|&(idx, _)| idx)
332    }
333
334    /// Get all downstream steps of a given step (steps that depend on it)
335    pub fn downstream_steps(&self, step_idx: usize) -> Vec<usize> {
336        let step_id = format!("step_{}", step_idx);
337        self.steps
338            .iter()
339            .enumerate()
340            .filter(|(_, s)| s.dependencies.contains(&step_id))
341            .map(|(i, _)| i)
342            .collect()
343    }
344
345    /// Recursively get all downstream steps (including transitive dependencies)
346    pub fn downstream_steps_recursive(&self, step_idx: usize) -> Vec<usize> {
347        let mut result = Vec::new();
348        let mut visited = std::collections::HashSet::new();
349        self.collect_downstream(step_idx, &mut result, &mut visited);
350        result
351    }
352
353    fn collect_downstream(
354        &self,
355        step_idx: usize,
356        result: &mut Vec<usize>,
357        visited: &mut std::collections::HashSet<usize>,
358    ) {
359        if visited.contains(&step_idx) {
360            return;
361        }
362        visited.insert(step_idx);
363        for downstream in self.downstream_steps(step_idx) {
364            if !result.contains(&downstream) {
365                result.push(downstream);
366            }
367            self.collect_downstream(downstream, result, visited);
368        }
369    }
370}
371
372// ── Plan validation ─────────────────────────────────────────────────────────
373
374/// Plan validation issue
375#[derive(Debug, Clone, Serialize, Deserialize)]
376pub struct PlanValidationIssue {
377    /// Severity level
378    pub severity: IssueSeverity,
379    /// Issue description
380    pub message: String,
381    /// Suggested fix
382    #[serde(default)]
383    pub fix: Option<String>,
384}
385
386/// Issue severity level
387#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
388pub enum IssueSeverity {
389    /// Error: plan cannot be executed
390    Error,
391    /// Warning: plan is executable but may have issues
392    Warning,
393}
394
395// ── Plan step ───────────────────────────────────────────────────────────────
396
397/// A single step in a plan
398#[derive(Debug, Clone, Serialize, Deserialize)]
399pub struct PlanStep {
400    /// Step description
401    pub description: String,
402    /// Step status
403    pub status: StepStatus,
404    /// Expected input (dependency description on previous step results)
405    pub expected_input: Option<String>,
406    /// Expected output description
407    pub expected_output: Option<String>,
408    /// List of dependent step indices (generated from LLM planning output)
409    #[serde(default)]
410    pub dependencies: Vec<String>,
411}
412
413impl PlanStep {
414    /// Create a new step
415    pub fn new(description: impl Into<String>) -> Self {
416        Self {
417            description: description.into(),
418            status: StepStatus::Pending,
419            expected_input: None,
420            expected_output: None,
421            dependencies: Vec::new(),
422        }
423    }
424
425    /// Set expected input description
426    pub fn with_expected_input(mut self, input: impl Into<String>) -> Self {
427        self.expected_input = Some(input.into());
428        self
429    }
430
431    /// Set expected output description
432    pub fn with_expected_output(mut self, output: impl Into<String>) -> Self {
433        self.expected_output = Some(output.into());
434        self
435    }
436
437    /// Set step dependencies
438    pub fn with_dependencies(mut self, deps: Vec<String>) -> Self {
439        self.dependencies = deps;
440        self
441    }
442}
443
444/// Step execution status
445#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
446pub enum StepStatus {
447    /// Waiting to execute
448    Pending,
449    /// Currently executing
450    Running,
451    /// Completed
452    Completed,
453    /// Execution failed
454    Failed,
455}
456
457// ── LLM structured output types ─────────────────────────────────────────────
458
459/// Structured plan output returned by LLM
460#[derive(Debug, Clone, Serialize, Deserialize)]
461pub struct PlanOutput {
462    /// List of steps
463    pub steps: Vec<PlanStepOutput>,
464}
465
466/// Single step returned by LLM
467#[derive(Debug, Clone, Serialize, Deserialize)]
468pub struct PlanStepOutput {
469    /// Step description
470    pub description: String,
471    /// Dependent step references (prefer 'step_N' IDs; description keywords accepted as fallback)
472    #[serde(default)]
473    pub dependencies: Vec<String>,
474    /// Expected output
475    #[serde(default)]
476    pub expected_output: Option<String>,
477}
478
479/// Return JSON Schema for LLM structured output
480pub fn plan_output_schema() -> serde_json::Value {
481    serde_json::json!({
482        "type": "object",
483        "properties": {
484            "steps": {
485                "type": "array",
486                "items": {
487                    "type": "object",
488                    "properties": {
489                        "description": {
490                            "type": "string",
491                            "description": "Detailed description of the step"
492                        },
493                        "dependencies": {
494                            "type": "array",
495                            "items": { "type": "string" },
496                            "description": "Dependency references — prefer exact step IDs like 'step_0', 'step_1'. Only use description keywords when a step ID is unavailable."
497                        },
498                        "expected_output": {
499                            "type": "string",
500                            "description": "Expected output of the step"
501                        }
502                    },
503                    "required": ["description"]
504                },
505                "minItems": 1
506            }
507        },
508        "required": ["steps"]
509    })
510}
511
512/// Execution result of a single step
513#[derive(Debug, Clone, Serialize, Deserialize)]
514pub struct StepResult {
515    /// Step index in the plan
516    pub step_index: usize,
517    /// Step description
518    pub description: String,
519    /// Execution output
520    pub output: String,
521    /// Whether successful
522    pub success: bool,
523}
524
525// ── Planner trait ───────────────────────────────────────────────────────────
526
527/// Planner trait — accepts a task description and returns an execution plan
528pub trait Planner: Send + Sync {
529    /// Generate an execution plan based on the task description
530    fn plan<'a>(&'a self, task: &'a str) -> BoxFuture<'a, Result<Plan>>;
531}
532
533// ── StaticPlanner ───────────────────────────────────────────────────────────
534
535/// Static Planner: uses a pre-defined list of steps (for testing)
536pub struct StaticPlanner {
537    steps: Vec<String>,
538}
539
540impl StaticPlanner {
541    /// Create a StaticPlanner with predefined step descriptions
542    pub fn new(steps: Vec<impl Into<String>>) -> Self {
543        Self {
544            steps: steps.into_iter().map(|s| s.into()).collect(),
545        }
546    }
547}
548
549impl Planner for StaticPlanner {
550    fn plan<'a>(&'a self, _task: &'a str) -> BoxFuture<'a, Result<Plan>> {
551        Box::pin(async move {
552            let steps: Vec<PlanStep> = self
553                .steps
554                .iter()
555                .map(|desc| PlanStep::new(desc.as_str()))
556                .collect();
557            Ok(Plan::new(steps))
558        })
559    }
560}
561
562// ── PlanStore trait ─────────────────────────────────────────────────────────
563
564/// Lightweight plan summary for listing/search
565#[derive(Debug, Clone)]
566pub struct PlanSummary {
567    /// Plan unique identifier
568    pub id: String,
569    /// Human-readable short identifier (for URLs, etc.)
570    pub slug: Option<String>,
571    /// Plan goal description
572    pub goal: Option<String>,
573    /// Plan version number (for optimistic locking)
574    pub version: u32,
575    /// Total number of steps in the plan
576    pub total_steps: usize,
577    /// Number of completed steps
578    pub completed_steps: usize,
579}
580
581/// Trait for plan persistence operations
582pub trait PlanStore: Send + Sync {
583    /// Save plan to storage
584    fn save_plan<'a>(&'a self, plan: &'a Plan) -> BoxFuture<'a, Result<()>>;
585    /// Load plan by plan ID
586    fn load_plan<'a>(&'a self, plan_id: &'a str) -> BoxFuture<'a, Result<Option<Plan>>>;
587    /// Load plan by slug
588    fn load_plan_by_slug<'a>(&'a self, slug: &'a str) -> BoxFuture<'a, Result<Option<Plan>>>;
589    /// List plan summaries
590    fn list_plans<'a>(&'a self, limit: usize) -> BoxFuture<'a, Result<Vec<PlanSummary>>>;
591    /// Delete a plan
592    fn delete_plan<'a>(&'a self, plan_id: &'a str) -> BoxFuture<'a, Result<bool>>;
593    /// Search plans
594    fn search_plans<'a>(
595        &'a self,
596        query: &'a str,
597        limit: usize,
598    ) -> BoxFuture<'a, Result<Vec<PlanSummary>>>;
599}
600
601#[cfg(test)]
602mod tests {
603    use super::*;
604
605    #[test]
606    fn test_plan_step_status() {
607        let step = PlanStep::new("test step");
608        assert_eq!(step.status, StepStatus::Pending);
609        assert_eq!(step.description, "test step");
610    }
611
612    #[test]
613    fn test_plan() {
614        let plan = Plan::new(vec![
615            PlanStep::new("step 1"),
616            PlanStep::new("step 2"),
617            PlanStep::new("step 3"),
618        ]);
619        assert_eq!(plan.steps.len(), 3);
620        assert!(plan.id.is_some());
621        assert_eq!(plan.version, 1);
622    }
623
624    #[test]
625    fn test_plan_auto_id() {
626        let plan = Plan::new(vec![PlanStep::new("test")]);
627        assert!(plan.id.as_ref().unwrap().starts_with("plan_"));
628    }
629
630    #[test]
631    fn test_plan_with_metadata() {
632        let plan =
633            Plan::new(vec![PlanStep::new("test")]).with_metadata("key", serde_json::json!("value"));
634        assert_eq!(plan.metadata.get("key").unwrap(), "value");
635    }
636
637    #[test]
638    fn test_validate_empty_plan() {
639        let plan = Plan::new(vec![]);
640        let issues = plan.validate();
641        assert!(issues.iter().any(|i| i.message.contains("no steps")));
642    }
643
644    #[test]
645    fn test_validate_self_dependency() {
646        let plan = Plan::new(vec![
647            PlanStep::new("step 0"),
648            PlanStep::new("step 1").with_dependencies(vec!["step_1".to_string()]),
649        ]);
650        let issues = plan.validate();
651        assert!(
652            issues
653                .iter()
654                .any(|i| i.message.contains("depends on itself"))
655        );
656    }
657
658    #[test]
659    fn test_validate_invalid_index() {
660        let plan = Plan::new(vec![
661            PlanStep::new("step 0").with_dependencies(vec!["step_99".to_string()]),
662        ]);
663        let issues = plan.validate();
664        assert!(
665            issues
666                .iter()
667                .any(|i| i.message.contains("non-existent step index"))
668        );
669    }
670
671    #[test]
672    fn test_auto_fix_removes_self_dependency() {
673        let mut plan = Plan::new(vec![
674            PlanStep::new("step 0"),
675            PlanStep::new("step 1")
676                .with_dependencies(vec!["step_1".to_string(), "step_0".to_string()]),
677        ]);
678        let fixes = plan.auto_fix();
679        assert!(fixes.iter().any(|f| f.contains("self-dependency")));
680        assert_eq!(plan.steps[1].dependencies, vec!["step_0"]);
681    }
682
683    #[test]
684    fn test_auto_fix_empty_description() {
685        let mut plan = Plan::new(vec![PlanStep::new("")]);
686        let fixes = plan.auto_fix();
687        assert!(fixes.iter().any(|f| f.contains("empty description")));
688        assert!(!plan.steps[0].description.is_empty());
689    }
690
691    #[test]
692    fn test_downstream_steps() {
693        let plan = Plan::new(vec![
694            PlanStep::new("A"),
695            PlanStep::new("B").with_dependencies(vec!["step_0".to_string()]),
696            PlanStep::new("C").with_dependencies(vec!["step_0".to_string()]),
697            PlanStep::new("D").with_dependencies(vec!["step_1".to_string()]),
698        ]);
699        let downstream = plan.downstream_steps(0);
700        assert_eq!(downstream, vec![1, 2]);
701
702        let recursive = plan.downstream_steps_recursive(0);
703        assert!(recursive.contains(&1));
704        assert!(recursive.contains(&2));
705        assert!(recursive.contains(&3));
706    }
707
708    #[test]
709    fn test_plan_touch() {
710        let mut plan = Plan::new(vec![PlanStep::new("test")]);
711        let before = plan.updated_at;
712        std::thread::sleep(std::time::Duration::from_millis(10));
713        plan.touch();
714        assert!(plan.updated_at >= before);
715    }
716
717    #[test]
718    fn test_dependency_resolution_improved_matching() {
719        let plan = Plan::new(vec![
720            PlanStep::new("database migration"),
721            PlanStep::new("setup environment"),
722            PlanStep::new("group setup"),
723        ]);
724
725        let _result = plan.resolve_dependency("data");
726        let _result = plan.resolve_dependency("setup");
727        let result = plan.resolve_dependency("database migration");
728        assert_eq!(result, Some(0));
729        let result = plan.resolve_dependency("step_1");
730        assert_eq!(result, Some(1));
731        let result = plan.resolve_dependency("step_10");
732        assert_eq!(result, None);
733    }
734
735    #[tokio::test]
736    async fn test_static_planner() {
737        let planner = StaticPlanner::new(vec!["A", "B", "C"]);
738        let plan = planner.plan("test").await.unwrap();
739        assert_eq!(plan.steps.len(), 3);
740    }
741}