vtcode_core/tools/
plan.rs

1use std::sync::Arc;
2
3use anyhow::{Result, anyhow, ensure};
4use chrono::{DateTime, Utc};
5use parking_lot::RwLock;
6use serde::{Deserialize, Serialize};
7
8const PLAN_UPDATE_PROGRESS: &str = "Plan updated. Continue working through TODOs.";
9const PLAN_UPDATE_COMPLETE: &str = "Plan completed. All TODOs are done.";
10const PLAN_UPDATE_CLEARED: &str = "Plan cleared. Start a new TODO list.";
11const MAX_PLAN_STEPS: usize = 12;
12const MIN_PLAN_STEPS: usize = 1;
13const CHECKBOX_PENDING: &str = "[ ]";
14const CHECKBOX_IN_PROGRESS: &str = "[ ]";
15const CHECKBOX_COMPLETED: &str = "[x]";
16const IN_PROGRESS_NOTE: &str = " _(in progress)_";
17
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
19#[serde(rename_all = "snake_case")]
20pub enum StepStatus {
21    Pending,
22    InProgress,
23    Completed,
24}
25
26impl StepStatus {
27    pub fn label(&self) -> &'static str {
28        match self {
29            StepStatus::Pending => "pending",
30            StepStatus::InProgress => "in_progress",
31            StepStatus::Completed => "completed",
32        }
33    }
34
35    pub fn checkbox(&self) -> &'static str {
36        match self {
37            StepStatus::Pending => CHECKBOX_PENDING,
38            StepStatus::InProgress => CHECKBOX_IN_PROGRESS,
39            StepStatus::Completed => CHECKBOX_COMPLETED,
40        }
41    }
42
43    pub fn status_note(&self) -> Option<&'static str> {
44        match self {
45            StepStatus::InProgress => Some(IN_PROGRESS_NOTE),
46            StepStatus::Pending | StepStatus::Completed => None,
47        }
48    }
49
50    pub fn is_complete(&self) -> bool {
51        matches!(self, StepStatus::Completed)
52    }
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
56pub struct PlanStep {
57    pub step: String,
58    pub status: StepStatus,
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
62#[serde(rename_all = "snake_case")]
63pub enum PlanCompletionState {
64    Empty,
65    InProgress,
66    Done,
67}
68
69impl PlanCompletionState {
70    pub fn label(&self) -> &'static str {
71        match self {
72            PlanCompletionState::Empty => "no_todos",
73            PlanCompletionState::InProgress => "todos_remaining",
74            PlanCompletionState::Done => "done",
75        }
76    }
77
78    pub fn description(&self) -> &'static str {
79        match self {
80            PlanCompletionState::Empty => "No TODOs recorded in the current plan.",
81            PlanCompletionState::InProgress => "TODOs remain in the current plan.",
82            PlanCompletionState::Done => "All TODOs have been completed.",
83        }
84    }
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
88pub struct PlanSummary {
89    pub total_steps: usize,
90    pub completed_steps: usize,
91    pub status: PlanCompletionState,
92}
93
94impl Default for PlanSummary {
95    fn default() -> Self {
96        Self {
97            total_steps: 0,
98            completed_steps: 0,
99            status: PlanCompletionState::Empty,
100        }
101    }
102}
103
104impl PlanSummary {
105    pub fn from_steps(steps: &[PlanStep]) -> Self {
106        if steps.is_empty() {
107            return Self::default();
108        }
109
110        let total_steps = steps.len();
111        let completed_steps = steps
112            .iter()
113            .filter(|step| step.status.is_complete())
114            .count();
115        let status = if completed_steps == total_steps {
116            PlanCompletionState::Done
117        } else {
118            PlanCompletionState::InProgress
119        };
120
121        Self {
122            total_steps,
123            completed_steps,
124            status,
125        }
126    }
127}
128
129#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
130pub struct TaskPlan {
131    #[serde(skip_serializing_if = "Option::is_none")]
132    pub explanation: Option<String>,
133    pub steps: Vec<PlanStep>,
134    pub summary: PlanSummary,
135    pub version: u64,
136    pub updated_at: DateTime<Utc>,
137}
138
139impl Default for TaskPlan {
140    fn default() -> Self {
141        Self {
142            explanation: None,
143            steps: Vec::new(),
144            summary: PlanSummary::default(),
145            version: 0,
146            updated_at: Utc::now(),
147        }
148    }
149}
150
151#[derive(Debug, Clone, Deserialize)]
152pub struct UpdatePlanArgs {
153    #[serde(default, skip_serializing_if = "Option::is_none")]
154    pub explanation: Option<String>,
155    pub plan: Vec<PlanStep>,
156}
157
158#[derive(Debug, Clone, Serialize)]
159pub struct PlanUpdateResult {
160    pub success: bool,
161    pub message: String,
162    pub plan: TaskPlan,
163}
164
165impl PlanUpdateResult {
166    pub fn success(plan: TaskPlan) -> Self {
167        let message = match plan.summary.status {
168            PlanCompletionState::Done => PLAN_UPDATE_COMPLETE.to_string(),
169            PlanCompletionState::InProgress => PLAN_UPDATE_PROGRESS.to_string(),
170            PlanCompletionState::Empty => PLAN_UPDATE_CLEARED.to_string(),
171        };
172        Self {
173            success: true,
174            message,
175            plan,
176        }
177    }
178}
179
180#[derive(Debug, Clone)]
181pub struct PlanManager {
182    inner: Arc<RwLock<TaskPlan>>,
183}
184
185impl Default for PlanManager {
186    fn default() -> Self {
187        Self {
188            inner: Arc::new(RwLock::new(TaskPlan::default())),
189        }
190    }
191}
192
193impl PlanManager {
194    pub fn new() -> Self {
195        Self::default()
196    }
197
198    pub fn snapshot(&self) -> TaskPlan {
199        self.inner.read().clone()
200    }
201
202    pub fn update_plan(&self, update: UpdatePlanArgs) -> Result<TaskPlan> {
203        validate_plan(&update)?;
204
205        let sanitized_explanation = update
206            .explanation
207            .as_ref()
208            .map(|text| text.trim().to_string())
209            .filter(|text| !text.is_empty());
210
211        let mut in_progress_count = 0usize;
212        let mut sanitized_steps: Vec<PlanStep> = Vec::with_capacity(update.plan.len());
213        for (index, mut step) in update.plan.into_iter().enumerate() {
214            let trimmed = step.step.trim();
215            if trimmed.is_empty() {
216                return Err(anyhow!("Plan step {} cannot be empty", index + 1));
217            }
218            if matches!(step.status, StepStatus::InProgress) {
219                in_progress_count += 1;
220            }
221            step.step = trimmed.to_string();
222            sanitized_steps.push(step);
223        }
224
225        ensure!(
226            in_progress_count <= 1,
227            "At most one plan step can be in_progress"
228        );
229
230        let mut guard = self.inner.write();
231        let version = guard.version.saturating_add(1);
232        let summary = PlanSummary::from_steps(&sanitized_steps);
233        let updated_plan = TaskPlan {
234            explanation: sanitized_explanation,
235            steps: sanitized_steps,
236            summary,
237            version,
238            updated_at: Utc::now(),
239        };
240        *guard = updated_plan.clone();
241        Ok(updated_plan)
242    }
243}
244
245fn validate_plan(update: &UpdatePlanArgs) -> Result<()> {
246    let step_count = update.plan.len();
247    ensure!(
248        step_count >= MIN_PLAN_STEPS,
249        "Plan must contain at least {} step(s)",
250        MIN_PLAN_STEPS
251    );
252    ensure!(
253        step_count <= MAX_PLAN_STEPS,
254        "Plan must not exceed {} steps",
255        MAX_PLAN_STEPS
256    );
257
258    for (index, step) in update.plan.iter().enumerate() {
259        ensure!(
260            !step.step.trim().is_empty(),
261            "Plan step {} cannot be empty",
262            index + 1
263        );
264    }
265
266    Ok(())
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn initializes_with_default_state() {
275        let manager = PlanManager::new();
276        let snapshot = manager.snapshot();
277        assert_eq!(snapshot.steps.len(), 0);
278        assert_eq!(snapshot.version, 0);
279        assert_eq!(snapshot.summary.status, PlanCompletionState::Empty);
280        assert_eq!(snapshot.summary.total_steps, 0);
281    }
282
283    #[test]
284    fn rejects_empty_plan() {
285        let manager = PlanManager::new();
286        let args = UpdatePlanArgs {
287            explanation: None,
288            plan: Vec::new(),
289        };
290        assert!(manager.update_plan(args).is_err());
291    }
292
293    #[test]
294    fn rejects_multiple_in_progress_steps() {
295        let manager = PlanManager::new();
296        let args = UpdatePlanArgs {
297            explanation: None,
298            plan: vec![
299                PlanStep {
300                    step: "Step one".to_string(),
301                    status: StepStatus::InProgress,
302                },
303                PlanStep {
304                    step: "Step two".to_string(),
305                    status: StepStatus::InProgress,
306                },
307            ],
308        };
309        assert!(manager.update_plan(args).is_err());
310    }
311
312    #[test]
313    fn updates_plan_successfully() {
314        let manager = PlanManager::new();
315        let args = UpdatePlanArgs {
316            explanation: Some("Focus on API layer".to_string()),
317            plan: vec![
318                PlanStep {
319                    step: "Audit handlers".to_string(),
320                    status: StepStatus::Pending,
321                },
322                PlanStep {
323                    step: "Add tests".to_string(),
324                    status: StepStatus::Pending,
325                },
326            ],
327        };
328        let result = manager.update_plan(args).expect("plan should update");
329        assert_eq!(result.steps.len(), 2);
330        assert_eq!(result.version, 1);
331        assert_eq!(result.steps[0].status, StepStatus::Pending);
332        assert_eq!(result.summary.total_steps, 2);
333        assert_eq!(result.summary.completed_steps, 0);
334        assert_eq!(result.summary.status, PlanCompletionState::InProgress);
335    }
336
337    #[test]
338    fn marks_plan_done_when_all_completed() {
339        let manager = PlanManager::new();
340        let args = UpdatePlanArgs {
341            explanation: None,
342            plan: vec![PlanStep {
343                step: "Finalize deployment".to_string(),
344                status: StepStatus::Completed,
345            }],
346        };
347        let result = manager.update_plan(args).expect("plan should update");
348        assert_eq!(result.summary.total_steps, 1);
349        assert_eq!(result.summary.completed_steps, 1);
350        assert_eq!(result.summary.status, PlanCompletionState::Done);
351    }
352}