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