agent_air_runtime/controller/tools/
plan_store.rs

1//! Shared state and utilities for plan tools.
2//!
3//! Provides the `PlanStore` struct that manages the plans directory,
4//! per-file locking, sequential plan ID generation, and status marker
5//! conversion helpers used by both `MarkdownPlanTool` and `UpdatePlanStepTool`.
6
7use std::collections::HashMap;
8use std::path::{Path, PathBuf};
9use std::sync::Arc;
10
11use tokio::sync::{Mutex, RwLock};
12
13/// Directory name under workspace root where plans are stored.
14const PLANS_DIR_NAME: &str = ".agent-air/plans";
15
16/// Status marker for pending steps.
17pub const PENDING_MARKER: &str = " ";
18/// Status marker for in-progress steps.
19pub const IN_PROGRESS_MARKER: &str = "~";
20/// Status marker for completed steps.
21pub const COMPLETED_MARKER: &str = "x";
22/// Status marker for skipped steps.
23pub const SKIPPED_MARKER: &str = "-";
24
25/// Shared state for plan tools, managing the plans directory and per-file locks.
26pub struct PlanStore {
27    /// Resolved path to the plans directory.
28    plans_dir: PathBuf,
29    /// Per-file mutexes to prevent concurrent writes to the same plan file.
30    file_locks: RwLock<HashMap<PathBuf, Arc<Mutex<()>>>>,
31}
32
33impl PlanStore {
34    /// Create a new `PlanStore` rooted at the given workspace directory.
35    pub fn new(workspace_root: PathBuf) -> Self {
36        Self {
37            plans_dir: workspace_root.join(PLANS_DIR_NAME),
38            file_locks: RwLock::new(HashMap::new()),
39        }
40    }
41
42    /// Returns the resolved plans directory path.
43    pub fn plans_dir(&self) -> &Path {
44        &self.plans_dir
45    }
46
47    /// Returns a per-file mutex for the given path, creating one if it does not exist.
48    pub async fn acquire_lock(&self, path: &Path) -> Arc<Mutex<()>> {
49        // Fast path: check if lock already exists.
50        {
51            let locks = self.file_locks.read().await;
52            if let Some(lock) = locks.get(path) {
53                return lock.clone();
54            }
55        }
56
57        // Slow path: create a new lock.
58        let mut locks = self.file_locks.write().await;
59        locks
60            .entry(path.to_path_buf())
61            .or_insert_with(|| Arc::new(Mutex::new(())))
62            .clone()
63    }
64
65    /// Scans the plans directory for existing `plan-NNN.md` files and returns
66    /// the next sequential plan ID (e.g., `plan-001`, `plan-002`).
67    pub async fn get_next_plan_id(&self) -> Result<String, String> {
68        let plans_dir = &self.plans_dir;
69
70        // If the directory doesn't exist yet, start at 001.
71        if !plans_dir.exists() {
72            return Ok("plan-001".to_string());
73        }
74
75        let mut max_num: u32 = 0;
76
77        let mut entries = tokio::fs::read_dir(plans_dir)
78            .await
79            .map_err(|e| format!("Failed to read plans directory: {}", e))?;
80
81        while let Some(entry) = entries
82            .next_entry()
83            .await
84            .map_err(|e| format!("Failed to read directory entry: {}", e))?
85        {
86            let file_name = entry.file_name();
87            let name = file_name.to_string_lossy();
88            if let Some(num_str) = name
89                .strip_prefix("plan-")
90                .and_then(|s| s.strip_suffix(".md"))
91                && let Ok(num) = num_str.parse::<u32>()
92                && num > max_num
93            {
94                max_num = num;
95            }
96        }
97
98        Ok(format!("plan-{:03}", max_num + 1))
99    }
100
101    /// Converts a step status string to its markdown checkbox marker.
102    pub fn status_to_marker(status: &str) -> Result<&'static str, String> {
103        match status {
104            "pending" => Ok(PENDING_MARKER),
105            "in_progress" => Ok(IN_PROGRESS_MARKER),
106            "completed" => Ok(COMPLETED_MARKER),
107            "skipped" => Ok(SKIPPED_MARKER),
108            _ => Err(format!(
109                "Invalid step status '{}'. Must be one of: pending, in_progress, completed, skipped",
110                status
111            )),
112        }
113    }
114
115    /// Converts a markdown checkbox marker to its status string.
116    pub fn marker_to_status(marker: &str) -> &'static str {
117        match marker {
118            " " => "pending",
119            "~" => "in_progress",
120            "x" => "completed",
121            "-" => "skipped",
122            _ => "pending",
123        }
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use tempfile::TempDir;
131
132    #[tokio::test]
133    async fn test_get_next_plan_id_empty_dir() {
134        let temp_dir = TempDir::new().unwrap();
135        let store = PlanStore::new(temp_dir.path().to_path_buf());
136
137        // Plans directory doesn't exist yet — should return plan-001.
138        let id = store.get_next_plan_id().await.unwrap();
139        assert_eq!(id, "plan-001");
140    }
141
142    #[tokio::test]
143    async fn test_get_next_plan_id_existing_plans() {
144        let temp_dir = TempDir::new().unwrap();
145        let store = PlanStore::new(temp_dir.path().to_path_buf());
146
147        // Create the plans directory and some plan files.
148        let plans_dir = temp_dir.path().join(PLANS_DIR_NAME);
149        tokio::fs::create_dir_all(&plans_dir).await.unwrap();
150        tokio::fs::write(plans_dir.join("plan-001.md"), "# Plan 1")
151            .await
152            .unwrap();
153        tokio::fs::write(plans_dir.join("plan-003.md"), "# Plan 3")
154            .await
155            .unwrap();
156        // Non-matching files should be ignored.
157        tokio::fs::write(plans_dir.join("notes.md"), "# Notes")
158            .await
159            .unwrap();
160
161        let id = store.get_next_plan_id().await.unwrap();
162        assert_eq!(id, "plan-004");
163    }
164
165    #[test]
166    fn test_plans_dir_derived_from_workspace_root() {
167        let store = PlanStore::new(PathBuf::from("/workspace/root"));
168        assert_eq!(
169            store.plans_dir(),
170            Path::new("/workspace/root/.agent-air/plans")
171        );
172    }
173
174    #[test]
175    fn test_status_to_marker() {
176        assert_eq!(PlanStore::status_to_marker("pending").unwrap(), " ");
177        assert_eq!(PlanStore::status_to_marker("in_progress").unwrap(), "~");
178        assert_eq!(PlanStore::status_to_marker("completed").unwrap(), "x");
179        assert_eq!(PlanStore::status_to_marker("skipped").unwrap(), "-");
180        assert!(PlanStore::status_to_marker("invalid").is_err());
181    }
182
183    #[test]
184    fn test_marker_to_status() {
185        assert_eq!(PlanStore::marker_to_status(" "), "pending");
186        assert_eq!(PlanStore::marker_to_status("~"), "in_progress");
187        assert_eq!(PlanStore::marker_to_status("x"), "completed");
188        assert_eq!(PlanStore::marker_to_status("-"), "skipped");
189        // Unknown markers default to pending.
190        assert_eq!(PlanStore::marker_to_status("?"), "pending");
191    }
192
193    #[tokio::test]
194    async fn test_acquire_lock_returns_same_lock_for_same_path() {
195        let temp_dir = TempDir::new().unwrap();
196        let store = PlanStore::new(temp_dir.path().to_path_buf());
197
198        let path = PathBuf::from("/some/plan.md");
199        let lock1 = store.acquire_lock(&path).await;
200        let lock2 = store.acquire_lock(&path).await;
201
202        // Both should point to the same underlying mutex.
203        assert!(Arc::ptr_eq(&lock1, &lock2));
204    }
205
206    #[tokio::test]
207    async fn test_acquire_lock_returns_different_locks_for_different_paths() {
208        let temp_dir = TempDir::new().unwrap();
209        let store = PlanStore::new(temp_dir.path().to_path_buf());
210
211        let path_a = PathBuf::from("/some/plan-a.md");
212        let path_b = PathBuf::from("/some/plan-b.md");
213        let lock_a = store.acquire_lock(&path_a).await;
214        let lock_b = store.acquire_lock(&path_b).await;
215
216        assert!(!Arc::ptr_eq(&lock_a, &lock_b));
217    }
218}
agent_air_runtime/controller/tools/plan_store.rs

agent_air_runtime/controller/tools/
plan_store.rs
