agent_trace/state/
manifest.rs

1use crate::config::StoreInfo;
2use crate::types::{DocId, DocType};
3use anyhow::{bail, Context, Result};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::path::{Path, PathBuf};
7
8// ── Document Entry ───────────────────────────────────────────────────────────
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
11pub struct DocumentEntry {
12    pub id: DocId,
13    pub path: PathBuf,
14    pub doc_type: DocType,
15    #[serde(default)]
16    pub tags: Vec<String>,
17    #[serde(default)]
18    pub description: String,
19    #[serde(default)]
20    pub agent_name: String,
21}
22
23impl DocumentEntry {
24    pub fn new(path: PathBuf, doc_type: DocType, agent_name: &str) -> Self {
25        Self {
26            id: DocId::new(),
27            path,
28            doc_type,
29            tags: Vec::new(),
30            description: String::new(),
31            agent_name: agent_name.to_string(),
32        }
33    }
34}
35
36// ── On-disk TOML structure ───────────────────────────────────────────────────
37
38/// The raw TOML representation persisted to disk.
39#[derive(Debug, Serialize, Deserialize)]
40struct ManifestFile {
41    store: StoreInfo,
42    #[serde(default, rename = "document")]
43    documents: Vec<DocumentEntry>,
44}
45
46// ── Manifest ─────────────────────────────────────────────────────────────────
47
48pub struct Manifest {
49    pub store: StoreInfo,
50    documents: Vec<DocumentEntry>,
51    /// path → index into `documents`
52    by_path: HashMap<PathBuf, usize>,
53    /// id → index into `documents`
54    by_id: HashMap<DocId, usize>,
55}
56
57impl Manifest {
58    // ── Construction ─────────────────────────────────────────────────────
59
60    fn build_indices(
61        documents: &[DocumentEntry],
62    ) -> (HashMap<PathBuf, usize>, HashMap<DocId, usize>) {
63        let mut by_path = HashMap::new();
64        let mut by_id = HashMap::new();
65        for (i, doc) in documents.iter().enumerate() {
66            by_path.insert(doc.path.clone(), i);
67            by_id.insert(doc.id.clone(), i);
68        }
69        (by_path, by_id)
70    }
71
72    pub fn from_parts(store: StoreInfo, documents: Vec<DocumentEntry>) -> Self {
73        let (by_path, by_id) = Self::build_indices(&documents);
74        Self {
75            store,
76            documents,
77            by_path,
78            by_id,
79        }
80    }
81
82    // ── I/O ──────────────────────────────────────────────────────────────
83
84    pub fn load(store_root: &Path) -> Result<Self> {
85        let path = manifest_path(store_root);
86
87        // Clean up any stale tmp file.
88        let tmp_path = tmp_manifest_path(store_root);
89        if tmp_path.exists() {
90            let _ = std::fs::remove_file(&tmp_path);
91        }
92
93        let contents = std::fs::read_to_string(&path)
94            .with_context(|| format!("Reading manifest: {}", path.display()))?;
95        let file: ManifestFile = toml::from_str(&contents)
96            .with_context(|| format!("Parsing manifest: {}", path.display()))?;
97
98        Ok(Self::from_parts(file.store, file.documents))
99    }
100
101    pub fn save(&self, store_root: &Path) -> Result<()> {
102        let path = manifest_path(store_root);
103        let tmp = tmp_manifest_path(store_root);
104
105        let file = ManifestFile {
106            store: self.store.clone(),
107            documents: self.documents.clone(),
108        };
109        let contents = toml::to_string_pretty(&file)?;
110
111        // Atomic write: tmp → rename
112        std::fs::write(&tmp, &contents)
113            .with_context(|| format!("Writing tmp manifest: {}", tmp.display()))?;
114        std::fs::rename(&tmp, &path)
115            .with_context(|| format!("Renaming manifest tmp to final: {}", path.display()))?;
116        Ok(())
117    }
118
119    pub fn create_empty(store_info: StoreInfo, store_root: &Path) -> Result<Self> {
120        let m = Self::from_parts(store_info, Vec::new());
121        m.save(store_root)?;
122        Ok(m)
123    }
124
125    // ── CRUD ─────────────────────────────────────────────────────────────
126
127    /// Register a new document. Returns an error if the path is already tracked.
128    pub fn register(
129        &mut self,
130        path: &Path,
131        doc_type: DocType,
132        agent_name: &str,
133    ) -> Result<&DocumentEntry> {
134        if self.by_path.contains_key(path) {
135            bail!("Path already tracked: {}", path.display());
136        }
137        let entry = DocumentEntry::new(path.to_path_buf(), doc_type, agent_name);
138        let idx = self.documents.len();
139        self.by_path.insert(entry.path.clone(), idx);
140        self.by_id.insert(entry.id.clone(), idx);
141        self.documents.push(entry);
142        Ok(&self.documents[idx])
143    }
144
145    pub fn find_by_path(&self, path: &Path) -> Option<&DocumentEntry> {
146        self.by_path.get(path).map(|&i| &self.documents[i])
147    }
148
149    pub fn find_by_id(&self, id: &DocId) -> Option<&DocumentEntry> {
150        self.by_id.get(id).map(|&i| &self.documents[i])
151    }
152
153    pub fn reclassify(&mut self, path: &Path, new_type: DocType) -> Result<()> {
154        let idx = *self
155            .by_path
156            .get(path)
157            .with_context(|| format!("Path not tracked: {}", path.display()))?;
158        self.documents[idx].doc_type = new_type;
159        Ok(())
160    }
161
162    pub fn update_path(&mut self, old_path: &Path, new_path: &Path) -> Result<()> {
163        let idx = *self
164            .by_path
165            .get(old_path)
166            .with_context(|| format!("Old path not tracked: {}", old_path.display()))?;
167        self.by_path.remove(old_path);
168        self.documents[idx].path = new_path.to_path_buf();
169        self.by_path.insert(new_path.to_path_buf(), idx);
170        Ok(())
171    }
172
173    pub fn untrack(&mut self, path: &Path) -> Result<()> {
174        let idx = *self
175            .by_path
176            .get(path)
177            .with_context(|| format!("Path not tracked: {}", path.display()))?;
178        let id = self.documents[idx].id.clone();
179        self.by_path.remove(path);
180        self.by_id.remove(&id);
181        self.documents.remove(idx);
182        // Rebuild indices since indices shifted after removal.
183        let (by_path, by_id) = Self::build_indices(&self.documents);
184        self.by_path = by_path;
185        self.by_id = by_id;
186        Ok(())
187    }
188
189    pub fn list(&self, type_filter: Option<&DocType>) -> Vec<&DocumentEntry> {
190        self.documents
191            .iter()
192            .filter(|d| type_filter.is_none_or(|t| &d.doc_type == t))
193            .collect()
194    }
195
196    pub fn is_tracked(&self, path: &Path) -> bool {
197        self.by_path.contains_key(path)
198    }
199
200    /// Update the description of a tracked document.
201    pub fn update_description(&mut self, path: &Path, description: &str) -> Result<()> {
202        let idx = *self
203            .by_path
204            .get(path)
205            .with_context(|| format!("Path not tracked: {}", path.display()))?;
206        self.documents[idx].description = description.to_string();
207        Ok(())
208    }
209
210    pub fn documents(&self) -> &[DocumentEntry] {
211        &self.documents
212    }
213
214    pub fn len(&self) -> usize {
215        self.documents.len()
216    }
217
218    pub fn is_empty(&self) -> bool {
219        self.documents.is_empty()
220    }
221}
222
223fn manifest_path(store_root: &Path) -> PathBuf {
224    store_root.join(".agent-trace").join("manifest.toml")
225}
226
227fn tmp_manifest_path(store_root: &Path) -> PathBuf {
228    store_root.join(".agent-trace").join(".manifest.toml.tmp")
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234    use crate::config::StoreInfo;
235    use tempfile::TempDir;
236
237    fn make_store(tmp: &TempDir) -> (PathBuf, StoreInfo) {
238        let root = tmp.path().to_path_buf();
239        std::fs::create_dir_all(root.join(".agent-trace")).unwrap();
240        let info = StoreInfo::new("test".into());
241        (root, info)
242    }
243
244    #[test]
245    fn test_empty_manifest_roundtrip() {
246        let tmp = TempDir::new().unwrap();
247        let (root, info) = make_store(&tmp);
248        let m = Manifest::create_empty(info, &root).unwrap();
249        assert!(m.is_empty());
250
251        let loaded = Manifest::load(&root).unwrap();
252        assert!(loaded.is_empty());
253        assert_eq!(loaded.store.name, "test");
254    }
255
256    #[test]
257    fn test_register_and_lookup() {
258        let tmp = TempDir::new().unwrap();
259        let (root, info) = make_store(&tmp);
260        let mut m = Manifest::create_empty(info, &root).unwrap();
261
262        let path = PathBuf::from("prd.md");
263        m.register(&path, DocType::Plan, "").unwrap();
264
265        let entry = m.find_by_path(&path).unwrap();
266        assert_eq!(entry.doc_type, DocType::Plan);
267        assert!(entry.id.0.parse::<uuid::Uuid>().is_ok());
268
269        let by_id = m.find_by_id(&entry.id.clone()).unwrap();
270        assert_eq!(by_id.path, path);
271    }
272
273    #[test]
274    fn test_register_duplicate_error() {
275        let tmp = TempDir::new().unwrap();
276        let (root, info) = make_store(&tmp);
277        let mut m = Manifest::create_empty(info, &root).unwrap();
278        let path = PathBuf::from("notes.md");
279        m.register(&path, DocType::Scratch, "").unwrap();
280        assert!(m.register(&path, DocType::Scratch, "").is_err());
281    }
282
283    #[test]
284    fn test_reclassify() {
285        let tmp = TempDir::new().unwrap();
286        let (root, info) = make_store(&tmp);
287        let mut m = Manifest::create_empty(info, &root).unwrap();
288        let path = PathBuf::from("notes.md");
289        m.register(&path, DocType::Scratch, "").unwrap();
290        m.reclassify(&path, DocType::Plan).unwrap();
291        assert_eq!(m.find_by_path(&path).unwrap().doc_type, DocType::Plan);
292    }
293
294    #[test]
295    fn test_update_path() {
296        let tmp = TempDir::new().unwrap();
297        let (root, info) = make_store(&tmp);
298        let mut m = Manifest::create_empty(info, &root).unwrap();
299        let old = PathBuf::from("old.md");
300        let new = PathBuf::from("new.md");
301        m.register(&old, DocType::Plan, "").unwrap();
302        m.update_path(&old, &new).unwrap();
303        assert!(m.find_by_path(&old).is_none());
304        assert!(m.find_by_path(&new).is_some());
305    }
306
307    #[test]
308    fn test_untrack() {
309        let tmp = TempDir::new().unwrap();
310        let (root, info) = make_store(&tmp);
311        let mut m = Manifest::create_empty(info, &root).unwrap();
312        let path = PathBuf::from("notes.md");
313        m.register(&path, DocType::Scratch, "").unwrap();
314        m.untrack(&path).unwrap();
315        assert!(!m.is_tracked(&path));
316        assert!(m.is_empty());
317    }
318
319    #[test]
320    fn test_list_filter() {
321        let tmp = TempDir::new().unwrap();
322        let (root, info) = make_store(&tmp);
323        let mut m = Manifest::create_empty(info, &root).unwrap();
324        m.register(&PathBuf::from("prd.md"), DocType::Plan, "")
325            .unwrap();
326        m.register(&PathBuf::from("notes.md"), DocType::Scratch, "")
327            .unwrap();
328        m.register(&PathBuf::from("plan2.md"), DocType::Plan, "")
329            .unwrap();
330
331        assert_eq!(m.list(None).len(), 3);
332        assert_eq!(m.list(Some(&DocType::Plan)).len(), 2);
333        assert_eq!(m.list(Some(&DocType::Scratch)).len(), 1);
334    }
335
336    #[test]
337    fn test_save_load_roundtrip() {
338        let tmp = TempDir::new().unwrap();
339        let (root, info) = make_store(&tmp);
340        let mut m = Manifest::create_empty(info, &root).unwrap();
341        m.register(&PathBuf::from("prd.md"), DocType::Plan, "agent-x")
342            .unwrap();
343        m.save(&root).unwrap();
344
345        let loaded = Manifest::load(&root).unwrap();
346        assert_eq!(loaded.len(), 1);
347        assert_eq!(loaded.documents()[0].doc_type, DocType::Plan);
348        assert_eq!(loaded.documents()[0].agent_name, "agent-x");
349    }
350
351    #[test]
352    fn test_stale_tmp_cleaned_on_load() {
353        let tmp = TempDir::new().unwrap();
354        let (root, info) = make_store(&tmp);
355        let m = Manifest::create_empty(info, &root).unwrap();
356        m.save(&root).unwrap();
357
358        // Simulate a stale tmp file.
359        let tmp_path = root.join(".agent-trace").join(".manifest.toml.tmp");
360        std::fs::write(&tmp_path, "garbage").unwrap();
361        assert!(tmp_path.exists());
362
363        // Load should clean it up.
364        Manifest::load(&root).unwrap();
365        assert!(!tmp_path.exists());
366    }
367}
agent_trace/state/manifest.rs

agent_trace/state/
manifest.rs
