agent_air_runtime/controller/tools/
update_plan_step.rs

1//! UpdatePlanStep tool implementation.
2//!
3//! Updates the status of a single step in an existing markdown plan file.
4//! Finds the step by its 1-indexed number and replaces the checkbox marker
5//! (`[ ]`, `[~]`, `[x]`, or `[-]`) with the requested status.
6//!
7//! Plan files are internal agent artifacts so this tool handles its own
8//! permissions and never prompts the user.
9
10use std::collections::HashMap;
11use std::future::Future;
12use std::pin::Pin;
13use std::sync::Arc;
14
15use regex::Regex;
16use tokio::fs;
17
18use super::plan_store::PlanStore;
19use super::types::{
20    DisplayConfig, DisplayResult, Executable, ResultContentType, ToolContext, ToolType,
21};
22
23/// UpdatePlanStep tool name constant.
24pub const UPDATE_PLAN_STEP_TOOL_NAME: &str = "update_plan_step";
25
26/// UpdatePlanStep tool description constant.
27pub const UPDATE_PLAN_STEP_TOOL_DESCRIPTION: &str = r#"Updates the status of a single step in an existing plan.
28
29Usage:
30- Specify the plan_id and 1-indexed step number
31- Status can be: pending, in_progress, completed, or skipped
32- The plan file must already exist in .agent-air/plans/
33
34Returns:
35- The updated step line and its new status"#;
36
37/// UpdatePlanStep tool JSON schema constant.
38pub const UPDATE_PLAN_STEP_TOOL_SCHEMA: &str = r#"{
39    "type": "object",
40    "properties": {
41        "plan_id": {
42            "type": "string",
43            "description": "Plan ID"
44        },
45        "step": {
46            "type": "integer",
47            "description": "Step number (1-indexed)",
48            "minimum": 1
49        },
50        "status": {
51            "type": "string",
52            "enum": ["pending", "in_progress", "completed", "skipped"]
53        }
54    },
55    "required": ["plan_id", "step", "status"]
56}"#;
57
58/// Tool that updates the status of a single step in a plan file.
59///
60/// Plan files live in `.agent-air/plans/` and are internal agent artifacts,
61/// so no user permission is required.
62pub struct UpdatePlanStepTool {
63    /// Shared plan store for directory paths and file locking.
64    plan_store: Arc<PlanStore>,
65}
66
67impl UpdatePlanStepTool {
68    /// Create a new UpdatePlanStepTool.
69    pub fn new(plan_store: Arc<PlanStore>) -> Self {
70        Self { plan_store }
71    }
72}
73
74impl Executable for UpdatePlanStepTool {
75    fn name(&self) -> &str {
76        UPDATE_PLAN_STEP_TOOL_NAME
77    }
78
79    fn description(&self) -> &str {
80        UPDATE_PLAN_STEP_TOOL_DESCRIPTION
81    }
82
83    fn input_schema(&self) -> &str {
84        UPDATE_PLAN_STEP_TOOL_SCHEMA
85    }
86
87    fn tool_type(&self) -> ToolType {
88        ToolType::Custom
89    }
90
91    fn execute(
92        &self,
93        _context: ToolContext,
94        input: HashMap<String, serde_json::Value>,
95    ) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> {
96        let plan_store = self.plan_store.clone();
97
98        Box::pin(async move {
99            // Step 1: Extract and validate parameters.
100            let plan_id = input
101                .get("plan_id")
102                .and_then(|v| v.as_str())
103                .ok_or_else(|| "Missing required 'plan_id' parameter".to_string())?;
104
105            let step_num = input.get("step").and_then(|v| v.as_u64()).ok_or_else(|| {
106                "Missing required 'step' parameter (must be a positive integer)".to_string()
107            })? as usize;
108
109            if step_num == 0 {
110                return Err("'step' must be >= 1 (1-indexed)".to_string());
111            }
112
113            let status = input
114                .get("status")
115                .and_then(|v| v.as_str())
116                .ok_or_else(|| "Missing required 'status' parameter".to_string())?;
117
118            let new_marker = PlanStore::status_to_marker(status)?;
119
120            // Step 2: Build file path and validate existence.
121            let plans_dir = plan_store.plans_dir();
122            let file_path = plans_dir.join(format!("{}.md", plan_id));
123            let file_path_str = file_path.to_string_lossy().to_string();
124
125            if !file_path.exists() {
126                return Err(format!(
127                    "Plan file not found: {}. Create the plan first using markdown_plan.",
128                    file_path_str
129                ));
130            }
131
132            // Step 3: Acquire per-file lock.
133            let lock = plan_store.acquire_lock(&file_path).await;
134            let _guard = lock.lock().await;
135
136            // Step 4: Read and parse the plan file.
137            let content = fs::read_to_string(&file_path)
138                .await
139                .map_err(|e| format!("Failed to read plan file: {}", e))?;
140
141            // Step 5: Find and update the step line.
142            // Step lines match the pattern: `N. [marker] description`
143            let step_pattern = Regex::new(r"^(\d+)\. \[[ x~-]\] ")
144                .map_err(|e| format!("Internal regex error: {}", e))?;
145            let marker_re =
146                Regex::new(r"\[[ x~-]\]").map_err(|e| format!("Internal regex error: {}", e))?;
147
148            let mut lines: Vec<String> = content.lines().map(String::from).collect();
149            let mut step_count: usize = 0;
150            let mut updated_line: Option<String> = None;
151
152            for line in &mut lines {
153                if step_pattern.is_match(line) {
154                    step_count += 1;
155                    if step_count == step_num {
156                        // Replace the marker inside the brackets.
157                        let replacement = format!("[{}]", new_marker);
158                        *line = marker_re
159                            .replace(line.as_str(), replacement.as_str())
160                            .to_string();
161                        updated_line = Some(line.clone());
162                        break;
163                    }
164                }
165            }
166
167            match updated_line {
168                None => {
169                    if step_count == 0 {
170                        Err(format!(
171                            "No steps found in plan '{}'. The plan file may be malformed.",
172                            plan_id
173                        ))
174                    } else {
175                        Err(format!(
176                            "Step {} is out of range. Plan '{}' has {} step(s).",
177                            step_num, plan_id, step_count
178                        ))
179                    }
180                }
181                Some(updated) => {
182                    // Step 6: Write the updated content back.
183                    let new_content = lines.join("\n");
184                    fs::write(&file_path, &new_content)
185                        .await
186                        .map_err(|e| format!("Failed to write updated plan file: {}", e))?;
187
188                    Ok(format!(
189                        "Step {} updated to {}: {}",
190                        step_num, status, updated
191                    ))
192                }
193            }
194        })
195    }
196
197    fn handles_own_permissions(&self) -> bool {
198        true
199    }
200
201    fn display_config(&self) -> DisplayConfig {
202        DisplayConfig {
203            display_name: "Update Plan Step".to_string(),
204            display_title: Box::new(|input| {
205                let plan_id = input
206                    .get("plan_id")
207                    .and_then(|v| v.as_str())
208                    .unwrap_or("unknown");
209                let step = input.get("step").and_then(|v| v.as_u64()).unwrap_or(0);
210                format!("{} step {}", plan_id, step)
211            }),
212            display_content: Box::new(|_input, result| DisplayResult {
213                content: result.to_string(),
214                content_type: ResultContentType::PlainText,
215                is_truncated: false,
216                full_length: result.lines().count(),
217            }),
218        }
219    }
220
221    fn compact_summary(&self, input: &HashMap<String, serde_json::Value>, _result: &str) -> String {
222        let plan_id = input
223            .get("plan_id")
224            .and_then(|v| v.as_str())
225            .unwrap_or("unknown");
226        let step = input.get("step").and_then(|v| v.as_u64()).unwrap_or(0);
227        let status = input
228            .get("status")
229            .and_then(|v| v.as_str())
230            .unwrap_or("unknown");
231
232        format!("[UpdatePlanStep: {} step {} -> {}]", plan_id, step, status)
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239    use tempfile::TempDir;
240
241    /// Helper to create a plan file with the given steps for testing.
242    async fn create_test_plan(plans_dir: &std::path::Path, plan_id: &str, steps: &[&str]) {
243        fs::create_dir_all(plans_dir).await.unwrap();
244        let mut content = format!(
245            "# Plan: Test\n\n**ID**: {}\n**Status**: active\n**Created**: 2025-01-01\n\n## Steps\n\n",
246            plan_id
247        );
248        for (i, step) in steps.iter().enumerate() {
249            content.push_str(&format!("{}. [ ] {}\n", i + 1, step));
250        }
251        fs::write(plans_dir.join(format!("{}.md", plan_id)), &content)
252            .await
253            .unwrap();
254    }
255
256    fn make_context(tool_use_id: &str) -> ToolContext {
257        ToolContext {
258            session_id: 1,
259            tool_use_id: tool_use_id.to_string(),
260            turn_id: None,
261            permissions_pre_approved: false,
262        }
263    }
264
265    #[tokio::test]
266    async fn test_update_step_to_completed() {
267        let temp_dir = TempDir::new().unwrap();
268        let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
269        let plans_dir = temp_dir.path().join(".agent-air/plans");
270
271        create_test_plan(&plans_dir, "plan-001", &["First step", "Second step"]).await;
272
273        let tool = UpdatePlanStepTool::new(plan_store);
274
275        let mut input = HashMap::new();
276        input.insert(
277            "plan_id".to_string(),
278            serde_json::Value::String("plan-001".to_string()),
279        );
280        input.insert("step".to_string(), serde_json::json!(1));
281        input.insert(
282            "status".to_string(),
283            serde_json::Value::String("completed".to_string()),
284        );
285
286        let result = tool.execute(make_context("test-complete"), input).await;
287        assert!(result.is_ok());
288        let output = result.unwrap();
289        assert!(output.contains("Step 1 updated to completed"));
290        assert!(output.contains("[x]"));
291
292        // Verify file on disk.
293        let content = fs::read_to_string(plans_dir.join("plan-001.md"))
294            .await
295            .unwrap();
296        assert!(content.contains("1. [x] First step"));
297        assert!(content.contains("2. [ ] Second step"));
298    }
299
300    #[tokio::test]
301    async fn test_update_step_to_in_progress() {
302        let temp_dir = TempDir::new().unwrap();
303        let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
304        let plans_dir = temp_dir.path().join(".agent-air/plans");
305
306        create_test_plan(&plans_dir, "plan-001", &["Step A"]).await;
307
308        let tool = UpdatePlanStepTool::new(plan_store);
309
310        let mut input = HashMap::new();
311        input.insert(
312            "plan_id".to_string(),
313            serde_json::Value::String("plan-001".to_string()),
314        );
315        input.insert("step".to_string(), serde_json::json!(1));
316        input.insert(
317            "status".to_string(),
318            serde_json::Value::String("in_progress".to_string()),
319        );
320
321        let result = tool.execute(make_context("test-progress"), input).await;
322        assert!(result.is_ok());
323        let output = result.unwrap();
324        assert!(output.contains("[~]"));
325    }
326
327    #[tokio::test]
328    async fn test_update_step_to_skipped() {
329        let temp_dir = TempDir::new().unwrap();
330        let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
331        let plans_dir = temp_dir.path().join(".agent-air/plans");
332
333        create_test_plan(&plans_dir, "plan-001", &["Skip me"]).await;
334
335        let tool = UpdatePlanStepTool::new(plan_store);
336
337        let mut input = HashMap::new();
338        input.insert(
339            "plan_id".to_string(),
340            serde_json::Value::String("plan-001".to_string()),
341        );
342        input.insert("step".to_string(), serde_json::json!(1));
343        input.insert(
344            "status".to_string(),
345            serde_json::Value::String("skipped".to_string()),
346        );
347
348        let result = tool.execute(make_context("test-skip"), input).await;
349        assert!(result.is_ok());
350        let output = result.unwrap();
351        assert!(output.contains("[-]"));
352    }
353
354    #[tokio::test]
355    async fn test_step_out_of_range() {
356        let temp_dir = TempDir::new().unwrap();
357        let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
358        let plans_dir = temp_dir.path().join(".agent-air/plans");
359
360        create_test_plan(&plans_dir, "plan-001", &["Only step"]).await;
361
362        let tool = UpdatePlanStepTool::new(plan_store);
363
364        let mut input = HashMap::new();
365        input.insert(
366            "plan_id".to_string(),
367            serde_json::Value::String("plan-001".to_string()),
368        );
369        input.insert("step".to_string(), serde_json::json!(5));
370        input.insert(
371            "status".to_string(),
372            serde_json::Value::String("completed".to_string()),
373        );
374
375        let result = tool.execute(make_context("test-range"), input).await;
376        assert!(result.is_err());
377        let err = result.unwrap_err();
378        assert!(err.contains("out of range"));
379        assert!(err.contains("1 step(s)"));
380    }
381
382    #[tokio::test]
383    async fn test_plan_file_not_found() {
384        let temp_dir = TempDir::new().unwrap();
385        let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
386
387        let tool = UpdatePlanStepTool::new(plan_store);
388
389        let mut input = HashMap::new();
390        input.insert(
391            "plan_id".to_string(),
392            serde_json::Value::String("plan-999".to_string()),
393        );
394        input.insert("step".to_string(), serde_json::json!(1));
395        input.insert(
396            "status".to_string(),
397            serde_json::Value::String("completed".to_string()),
398        );
399
400        let result = tool.execute(make_context("test-notfound"), input).await;
401        assert!(result.is_err());
402        assert!(result.unwrap_err().contains("Plan file not found"));
403    }
404
405    #[tokio::test]
406    async fn test_missing_required_fields() {
407        let temp_dir = TempDir::new().unwrap();
408        let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
409        let tool = UpdatePlanStepTool::new(plan_store);
410
411        // Missing plan_id.
412        let mut input = HashMap::new();
413        input.insert("step".to_string(), serde_json::json!(1));
414        input.insert(
415            "status".to_string(),
416            serde_json::Value::String("completed".to_string()),
417        );
418
419        let result = tool.execute(make_context("test-missing"), input).await;
420        assert!(result.is_err());
421        assert!(result.unwrap_err().contains("Missing required 'plan_id'"));
422
423        // Missing step.
424        let mut input = HashMap::new();
425        input.insert(
426            "plan_id".to_string(),
427            serde_json::Value::String("plan-001".to_string()),
428        );
429        input.insert(
430            "status".to_string(),
431            serde_json::Value::String("completed".to_string()),
432        );
433
434        let result = tool.execute(make_context("test-missing-step"), input).await;
435        assert!(result.is_err());
436        assert!(result.unwrap_err().contains("Missing required 'step'"));
437
438        // Missing status.
439        let mut input = HashMap::new();
440        input.insert(
441            "plan_id".to_string(),
442            serde_json::Value::String("plan-001".to_string()),
443        );
444        input.insert("step".to_string(), serde_json::json!(1));
445
446        let result = tool
447            .execute(make_context("test-missing-status"), input)
448            .await;
449        assert!(result.is_err());
450        assert!(result.unwrap_err().contains("Missing required 'status'"));
451    }
452
453    #[test]
454    fn test_compact_summary() {
455        let plan_store = Arc::new(PlanStore::new(std::path::PathBuf::from("/tmp")));
456        let tool = UpdatePlanStepTool::new(plan_store);
457
458        let mut input = HashMap::new();
459        input.insert(
460            "plan_id".to_string(),
461            serde_json::Value::String("plan-001".to_string()),
462        );
463        input.insert("step".to_string(), serde_json::json!(3));
464        input.insert(
465            "status".to_string(),
466            serde_json::Value::String("completed".to_string()),
467        );
468
469        let summary = tool.compact_summary(&input, "");
470        assert_eq!(summary, "[UpdatePlanStep: plan-001 step 3 -> completed]");
471    }
472
473    #[tokio::test]
474    async fn test_update_already_completed_step() {
475        let temp_dir = TempDir::new().unwrap();
476        let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
477        let plans_dir = temp_dir.path().join(".agent-air/plans");
478
479        // Create a plan with a completed step.
480        fs::create_dir_all(&plans_dir).await.unwrap();
481        let content = "# Plan: Test\n\n**ID**: plan-001\n**Status**: active\n**Created**: 2025-01-01\n\n## Steps\n\n1. [x] Already done\n";
482        fs::write(plans_dir.join("plan-001.md"), content)
483            .await
484            .unwrap();
485
486        let tool = UpdatePlanStepTool::new(plan_store);
487
488        let mut input = HashMap::new();
489        input.insert(
490            "plan_id".to_string(),
491            serde_json::Value::String("plan-001".to_string()),
492        );
493        input.insert("step".to_string(), serde_json::json!(1));
494        input.insert(
495            "status".to_string(),
496            serde_json::Value::String("pending".to_string()),
497        );
498
499        let result = tool.execute(make_context("test-revert"), input).await;
500        assert!(result.is_ok());
501        let output = result.unwrap();
502        assert!(output.contains("[ ]"));
503
504        // Verify on disk.
505        let updated = fs::read_to_string(plans_dir.join("plan-001.md"))
506            .await
507            .unwrap();
508        assert!(updated.contains("1. [ ] Already done"));
509    }
510
511    #[test]
512    fn test_handles_own_permissions() {
513        let plan_store = Arc::new(PlanStore::new(std::path::PathBuf::from("/tmp")));
514        let tool = UpdatePlanStepTool::new(plan_store);
515        assert!(tool.handles_own_permissions());
516    }
517}
agent_air_runtime/controller/tools/update_plan_step.rs

agent_air_runtime/controller/tools/
update_plan_step.rs
