| 1 | //! Plan tool implementation with step tracking and validation |
| 2 | |
| 3 | use std::sync::Arc; |
| 4 | use std::time::{Duration, Instant}; |
| 5 | use tokio::sync::Mutex; |
| 6 | |
| 7 | use async_trait::async_trait; |
| 8 | use serde::{Deserialize, Serialize}; |
| 9 | use serde_json::json; |
| 10 | |
| 11 | use crate::tools::spec::{ |
| 12 | ApprovalRequirement, ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, |
| 13 | }; |
| 14 | |
| 15 | // === Types === |
| 16 | |
| 17 | /// Status of a plan step. |
| 18 | #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] |
| 19 | #[serde(rename_all = "snake_case")] |
| 20 | pub enum StepStatus { |
| 21 | Pending, |
| 22 | InProgress, |
| 23 | Completed, |
| 24 | } |
| 25 | |
| 26 | impl StepStatus { |
| 27 | #[allow(dead_code)] |
| 28 | #[must_use] |
| 29 | pub fn from_str(value: &str) -> Option<Self> { |
| 30 | match value.trim().to_lowercase().as_str() { |
| 31 | "pending" => Some(StepStatus::Pending), |
| 32 | "in_progress" | "inprogress" => Some(StepStatus::InProgress), |
| 33 | "completed" | "done" => Some(StepStatus::Completed), |
| 34 | _ => None, |
| 35 | } |
| 36 | } |
| 37 | |
| 38 | #[allow(dead_code)] |
| 39 | #[must_use] |
| 40 | pub fn symbol(&self) -> &'static str { |
| 41 | match self { |
| 42 | StepStatus::Pending => "○", |
| 43 | StepStatus::InProgress => "◎", |
| 44 | StepStatus::Completed => "●", |
| 45 | } |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | /// Input representation for a plan item. |
| 50 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 51 | pub struct PlanItemArg { |
| 52 | pub step: String, |
| 53 | pub status: StepStatus, |
| 54 | } |
| 55 | |
| 56 | /// Update payload used by the plan tool. |
| 57 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 58 | pub struct UpdatePlanArgs { |
| 59 | #[serde(default)] |
| 60 | pub explanation: Option<String>, |
| 61 | pub plan: Vec<PlanItemArg>, |
| 62 | } |
| 63 | |
| 64 | // === Plan State === |
| 65 | |
| 66 | /// A plan step with timing information |
| 67 | #[derive(Debug, Clone)] |
| 68 | pub struct PlanStep { |
| 69 | pub text: String, |
| 70 | pub status: StepStatus, |
| 71 | /// When the step was started (transitioned to `InProgress`) |
| 72 | pub started_at: Option<Instant>, |
| 73 | /// When the step was completed |
| 74 | pub completed_at: Option<Instant>, |
| 75 | } |
| 76 | |
| 77 | impl PlanStep { |
| 78 | /// Create a new plan step. |
| 79 | pub fn new(text: String, status: StepStatus) -> Self { |
| 80 | Self { |
| 81 | text, |
| 82 | status, |
| 83 | started_at: None, |
| 84 | completed_at: None, |
| 85 | } |
| 86 | } |
| 87 | |
| 88 | /// Get the elapsed time if the step has timing info |
| 89 | #[must_use] |
| 90 | pub fn elapsed(&self) -> Option<Duration> { |
| 91 | match (self.started_at, self.completed_at) { |
| 92 | (Some(start), Some(end)) => Some(end.duration_since(start)), |
| 93 | (Some(start), None) if self.status == StepStatus::InProgress => Some(start.elapsed()), |
| 94 | _ => None, |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | /// Format elapsed time for display |
| 99 | #[must_use] |
| 100 | pub fn elapsed_str(&self) -> String { |
| 101 | match self.elapsed() { |
| 102 | Some(d) => { |
| 103 | let secs = d.as_secs(); |
| 104 | if secs < 60 { |
| 105 | format!("{secs}s") |
| 106 | } else if secs < 3600 { |
| 107 | format!("{}m {}s", secs / 60, secs % 60) |
| 108 | } else { |
| 109 | format!("{}h {}m", secs / 3600, (secs % 3600) / 60) |
| 110 | } |
| 111 | } |
| 112 | None => String::new(), |
| 113 | } |
| 114 | } |
| 115 | } |
| 116 | |
| 117 | /// Serializable snapshot for display |
| 118 | #[derive(Debug, Clone, Serialize)] |
| 119 | pub struct PlanSnapshot { |
| 120 | pub explanation: Option<String>, |
| 121 | pub items: Vec<PlanItemArg>, |
| 122 | } |
| 123 | |
| 124 | /// State tracking for the current plan |
| 125 | #[derive(Debug, Clone, Default)] |
| 126 | pub struct PlanState { |
| 127 | explanation: Option<String>, |
| 128 | steps: Vec<PlanStep>, |
| 129 | } |
| 130 | |
| 131 | impl PlanState { |
| 132 | /// Check whether the plan is empty. |
| 133 | #[must_use] |
| 134 | pub fn is_empty(&self) -> bool { |
| 135 | self.steps.is_empty() && self.explanation.as_deref().unwrap_or("").is_empty() |
| 136 | } |
| 137 | |
| 138 | pub fn update(&mut self, args: UpdatePlanArgs) { |
| 139 | self.explanation = args.explanation.filter(|s| !s.trim().is_empty()); |
| 140 | |
| 141 | let now = Instant::now(); |
| 142 | let mut new_steps = Vec::new(); |
| 143 | let mut in_progress_seen = false; |
| 144 | |
| 145 | for item in args.plan { |
| 146 | // Try to find existing step to preserve timing |
| 147 | let existing = self.steps.iter().find(|s| s.text == item.step); |
| 148 | |
| 149 | let mut status = item.status; |
| 150 | // Enforce single in_progress |
| 151 | if status == StepStatus::InProgress { |
| 152 | if in_progress_seen { |
| 153 | status = StepStatus::Pending; |
| 154 | } else { |
| 155 | in_progress_seen = true; |
| 156 | } |
| 157 | } |
| 158 | |
| 159 | let step = if let Some(old) = existing { |
| 160 | let mut s = old.clone(); |
| 161 | let old_status = s.status.clone(); |
| 162 | s.status = status.clone(); |
| 163 | |
| 164 | // Track timing transitions |
| 165 | if old_status == StepStatus::Pending && status == StepStatus::InProgress { |
| 166 | s.started_at = Some(now); |
| 167 | } |
| 168 | if old_status == StepStatus::InProgress && status == StepStatus::Completed { |
| 169 | s.completed_at = Some(now); |
| 170 | } |
| 171 | |
| 172 | s |
| 173 | } else { |
| 174 | let mut s = PlanStep::new(item.step, status.clone()); |
| 175 | if status == StepStatus::InProgress { |
| 176 | s.started_at = Some(now); |
| 177 | } |
| 178 | s |
| 179 | }; |
| 180 | |
| 181 | new_steps.push(step); |
| 182 | } |
| 183 | |
| 184 | self.steps = new_steps; |
| 185 | } |
| 186 | |
| 187 | pub fn snapshot(&self) -> PlanSnapshot { |
| 188 | PlanSnapshot { |
| 189 | explanation: self.explanation.clone(), |
| 190 | items: self |
| 191 | .steps |
| 192 | .iter() |
| 193 | .map(|s| PlanItemArg { |
| 194 | step: s.text.clone(), |
| 195 | status: s.status.clone(), |
| 196 | }) |
| 197 | .collect(), |
| 198 | } |
| 199 | } |
| 200 | |
| 201 | pub fn explanation(&self) -> Option<&str> { |
| 202 | self.explanation.as_deref() |
| 203 | } |
| 204 | |
| 205 | pub fn steps(&self) -> &[PlanStep] { |
| 206 | &self.steps |
| 207 | } |
| 208 | |
| 209 | /// Get counts of steps by status |
| 210 | pub fn counts(&self) -> (usize, usize, usize) { |
| 211 | let mut pending = 0; |
| 212 | let mut in_progress = 0; |
| 213 | let mut completed = 0; |
| 214 | for s in &self.steps { |
| 215 | match s.status { |
| 216 | StepStatus::Pending => pending += 1, |
| 217 | StepStatus::InProgress => in_progress += 1, |
| 218 | StepStatus::Completed => completed += 1, |
| 219 | } |
| 220 | } |
| 221 | (pending, in_progress, completed) |
| 222 | } |
| 223 | |
| 224 | /// Get progress as a percentage |
| 225 | pub fn progress_percent(&self) -> u8 { |
| 226 | if self.steps.is_empty() { |
| 227 | return 0; |
| 228 | } |
| 229 | let completed = self |
| 230 | .steps |
| 231 | .iter() |
| 232 | .filter(|s| s.status == StepStatus::Completed) |
| 233 | .count(); |
| 234 | let percent = completed.saturating_mul(100) / self.steps.len(); |
| 235 | u8::try_from(percent).unwrap_or(u8::MAX) |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | /// Validation result for plan transitions |
| 240 | #[derive(Debug)] |
| 241 | #[allow(dead_code)] |
| 242 | pub enum PlanValidation { |
| 243 | Ok, |
| 244 | Warning(String), |
| 245 | Error(String), |
| 246 | } |
| 247 | |
| 248 | /// Validate a plan update |
| 249 | #[allow(dead_code)] |
| 250 | pub fn validate_plan_update(current: &PlanState, update: &UpdatePlanArgs) -> PlanValidation { |
| 251 | let current_steps: std::collections::HashMap<_, _> = current |
| 252 | .steps() |
| 253 | .iter() |
| 254 | .map(|s| (s.text.clone(), &s.status)) |
| 255 | .collect(); |
| 256 | |
| 257 | for item in &update.plan { |
| 258 | if let Some(old_status) = current_steps.get(&item.step) { |
| 259 | // Check for invalid transitions |
| 260 | match (old_status, &item.status) { |
| 261 | (StepStatus::Completed, StepStatus::Pending) => { |
| 262 | return PlanValidation::Warning(format!( |
| 263 | "Step '{}' was completed but is now pending", |
| 264 | item.step |
| 265 | )); |
| 266 | } |
| 267 | (StepStatus::Completed, StepStatus::InProgress) => { |
| 268 | return PlanValidation::Warning(format!( |
| 269 | "Step '{}' was completed but is now in progress", |
| 270 | item.step |
| 271 | )); |
| 272 | } |
| 273 | _ => {} |
| 274 | } |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | PlanValidation::Ok |
| 279 | } |
| 280 | |
| 281 | // === UpdatePlanTool - ToolSpec implementation === |
| 282 | |
| 283 | /// Shared reference to `PlanState` for use across tools |
| 284 | pub type SharedPlanState = Arc<Mutex<PlanState>>; |
| 285 | |
| 286 | /// Create a new shared `PlanState` |
| 287 | pub fn new_shared_plan_state() -> SharedPlanState { |
| 288 | Arc::new(Mutex::new(PlanState::default())) |
| 289 | } |
| 290 | |
| 291 | /// Tool for updating the implementation plan |
| 292 | pub struct UpdatePlanTool { |
| 293 | plan_state: SharedPlanState, |
| 294 | } |
| 295 | |
| 296 | impl UpdatePlanTool { |
| 297 | pub fn new(plan_state: SharedPlanState) -> Self { |
| 298 | Self { plan_state } |
| 299 | } |
| 300 | } |
| 301 | |
| 302 | #[async_trait] |
| 303 | impl ToolSpec for UpdatePlanTool { |
| 304 | fn name(&self) -> &'static str { |
| 305 | "update_plan" |
| 306 | } |
| 307 | |
| 308 | fn description(&self) -> &'static str { |
| 309 | "Update the implementation plan with steps and their status. Use this to track progress on implementation tasks. Each step has a description and status (pending, in_progress, completed). Optionally include an explanation of the overall approach." |
| 310 | } |
| 311 | |
| 312 | fn input_schema(&self) -> serde_json::Value { |
| 313 | json!({ |
| 314 | "type": "object", |
| 315 | "properties": { |
| 316 | "explanation": { |
| 317 | "type": "string", |
| 318 | "description": "Optional high-level explanation of the plan or approach" |
| 319 | }, |
| 320 | "plan": { |
| 321 | "type": "array", |
| 322 | "description": "List of plan steps", |
| 323 | "items": { |
| 324 | "type": "object", |
| 325 | "properties": { |
| 326 | "step": { |
| 327 | "type": "string", |
| 328 | "description": "Description of the step" |
| 329 | }, |
| 330 | "status": { |
| 331 | "type": "string", |
| 332 | "enum": ["pending", "in_progress", "completed"], |
| 333 | "description": "Step status" |
| 334 | } |
| 335 | }, |
| 336 | "required": ["step", "status"] |
| 337 | } |
| 338 | } |
| 339 | }, |
| 340 | "required": ["plan"] |
| 341 | }) |
| 342 | } |
| 343 | |
| 344 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 345 | vec![ToolCapability::WritesFiles] |
| 346 | } |
| 347 | |
| 348 | fn approval_requirement(&self) -> ApprovalRequirement { |
| 349 | ApprovalRequirement::Auto |
| 350 | } |
| 351 | |
| 352 | async fn execute( |
| 353 | &self, |
| 354 | input: serde_json::Value, |
| 355 | _context: &ToolContext, |
| 356 | ) -> Result<ToolResult, ToolError> { |
| 357 | let explanation = input |
| 358 | .get("explanation") |
| 359 | .and_then(|v| v.as_str()) |
| 360 | .map(std::string::ToString::to_string); |
| 361 | |
| 362 | let plan_items = input |
| 363 | .get("plan") |
| 364 | .and_then(|v| v.as_array()) |
| 365 | .ok_or_else(|| ToolError::invalid_input("Missing or invalid 'plan' array"))?; |
| 366 | |
| 367 | let mut plan_args = Vec::new(); |
| 368 | for item in plan_items { |
| 369 | let step = item |
| 370 | .get("step") |
| 371 | .and_then(|v| v.as_str()) |
| 372 | .ok_or_else(|| ToolError::invalid_input("Plan item missing 'step'"))?; |
| 373 | |
| 374 | let status_str = item |
| 375 | .get("status") |
| 376 | .and_then(|v| v.as_str()) |
| 377 | .unwrap_or("pending"); |
| 378 | |
| 379 | let status = StepStatus::from_str(status_str).unwrap_or(StepStatus::Pending); |
| 380 | |
| 381 | plan_args.push(PlanItemArg { |
| 382 | step: step.to_string(), |
| 383 | status, |
| 384 | }); |
| 385 | } |
| 386 | |
| 387 | let args = UpdatePlanArgs { |
| 388 | explanation, |
| 389 | plan: plan_args, |
| 390 | }; |
| 391 | |
| 392 | let mut state = self.plan_state.lock().await; |
| 393 | |
| 394 | state.update(args); |
| 395 | |
| 396 | let snapshot = state.snapshot(); |
| 397 | let (pending, in_progress, completed) = state.counts(); |
| 398 | let progress = state.progress_percent(); |
| 399 | |
| 400 | let result = serde_json::to_string_pretty(&snapshot).unwrap_or_else(|_| "{}".to_string()); |
| 401 | |
| 402 | Ok(ToolResult::success(format!( |
| 403 | "Plan updated: {pending} pending, {in_progress} in progress, {completed} completed ({progress}% done)\n{result}" |
| 404 | ))) |
| 405 | } |
| 406 | } |
| 407 |