Skip to main content

dot_agent_core/history/
graph.rs

1use std::collections::HashMap;
2use std::fs;
3use std::path::{Path, PathBuf};
4
5use serde::{Deserialize, Serialize};
6
7use super::operation::Operation;
8use crate::error::Result;
9
10/// Persistent operation graph (DAG)
11#[derive(Debug, Clone, Default, Serialize, Deserialize)]
12pub struct OperationGraph {
13    /// Graph version for future migrations
14    pub version: u32,
15    /// All operations indexed by ID
16    #[serde(default)]
17    pub operations: HashMap<String, Operation>,
18    /// Root operations (no parent)
19    #[serde(default)]
20    pub roots: Vec<String>,
21    /// Latest operation ID (head)
22    pub head: Option<String>,
23}
24
25impl OperationGraph {
26    pub fn new() -> Self {
27        Self {
28            version: 1,
29            operations: HashMap::new(),
30            roots: Vec::new(),
31            head: None,
32        }
33    }
34
35    /// Load graph from file
36    pub fn load(path: &Path) -> Result<Self> {
37        if !path.exists() {
38            return Ok(Self::new());
39        }
40
41        let content = fs::read_to_string(path)?;
42        let graph: Self = toml::from_str(&content)
43            .map_err(|e| crate::error::DotAgentError::Internal(e.to_string()))?;
44
45        Ok(graph)
46    }
47
48    /// Save graph to file
49    pub fn save(&self, path: &Path) -> Result<()> {
50        if let Some(parent) = path.parent() {
51            fs::create_dir_all(parent)?;
52        }
53
54        let content = toml::to_string_pretty(self)
55            .map_err(|e| crate::error::DotAgentError::Internal(e.to_string()))?;
56
57        fs::write(path, content)?;
58        Ok(())
59    }
60
61    /// Add an operation to the graph
62    pub fn add_operation(&mut self, operation: Operation) {
63        let id = operation.id.as_str().to_string();
64
65        // Update roots if this is a root operation
66        if operation.parent.is_none() {
67            self.roots.push(id.clone());
68        }
69
70        // Update head
71        self.head = Some(id.clone());
72
73        // Store operation
74        self.operations.insert(id, operation);
75    }
76
77    /// Get an operation by ID
78    pub fn get_operation(&self, id: &str) -> Option<&Operation> {
79        self.operations.get(id)
80    }
81
82    /// Get the latest operation
83    pub fn get_head(&self) -> Option<&Operation> {
84        self.head.as_ref().and_then(|id| self.operations.get(id))
85    }
86
87    /// Get all operations in chronological order
88    pub fn operations_chronological(&self) -> Vec<&Operation> {
89        let mut ops: Vec<_> = self.operations.values().collect();
90        ops.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
91        ops
92    }
93
94    /// Get operations in reverse chronological order
95    pub fn operations_reverse_chronological(&self) -> Vec<&Operation> {
96        let mut ops: Vec<_> = self.operations.values().collect();
97        ops.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));
98        ops
99    }
100
101    /// Get the ancestry chain from an operation back to root
102    pub fn get_ancestry(&self, id: &str) -> Vec<&Operation> {
103        let mut chain = Vec::new();
104        let mut current_id = Some(id.to_string());
105
106        while let Some(ref op_id) = current_id {
107            if let Some(op) = self.operations.get(op_id) {
108                chain.push(op);
109                current_id = op.parent.as_ref().map(|p| p.as_str().to_string());
110            } else {
111                break;
112            }
113        }
114
115        chain
116    }
117
118    /// Get children of an operation
119    pub fn get_children(&self, id: &str) -> Vec<&Operation> {
120        self.operations
121            .values()
122            .filter(|op| op.parent.as_ref().map(|p| p.as_str()) == Some(id))
123            .collect()
124    }
125
126    /// Find operation by checkpoint ID
127    pub fn find_by_checkpoint(&self, checkpoint_id: &str) -> Option<&Operation> {
128        self.operations
129            .values()
130            .find(|op| op.checkpoint_id == checkpoint_id)
131    }
132
133    /// Get operations count
134    pub fn len(&self) -> usize {
135        self.operations.len()
136    }
137
138    /// Check if graph is empty
139    pub fn is_empty(&self) -> bool {
140        self.operations.is_empty()
141    }
142
143    /// Remove an operation and its descendants
144    pub fn remove_operation_tree(&mut self, id: &str) -> Vec<Operation> {
145        let mut removed = Vec::new();
146        let mut to_remove = vec![id.to_string()];
147
148        while let Some(current_id) = to_remove.pop() {
149            // Find children before removing
150            let children: Vec<_> = self
151                .operations
152                .values()
153                .filter(|op| op.parent.as_ref().map(|p| p.as_str()) == Some(&current_id))
154                .map(|op| op.id.as_str().to_string())
155                .collect();
156
157            to_remove.extend(children);
158
159            // Remove the operation
160            if let Some(op) = self.operations.remove(&current_id) {
161                removed.push(op);
162            }
163
164            // Update roots
165            self.roots.retain(|r| r != &current_id);
166        }
167
168        // Update head if removed
169        if self
170            .head
171            .as_ref()
172            .map(|h| removed.iter().any(|op| op.id.as_str() == h))
173            == Some(true)
174        {
175            self.head = self
176                .operations_reverse_chronological()
177                .first()
178                .map(|op| op.id.as_str().to_string());
179        }
180
181        removed
182    }
183}
184
185/// Manages the operation graph file
186pub struct GraphManager {
187    graph: OperationGraph,
188    path: PathBuf,
189}
190
191impl GraphManager {
192    pub fn new(base_dir: &Path) -> Result<Self> {
193        let path = base_dir.join("graph.toml");
194        let graph = OperationGraph::load(&path)?;
195
196        Ok(Self { graph, path })
197    }
198
199    pub fn graph(&self) -> &OperationGraph {
200        &self.graph
201    }
202
203    pub fn graph_mut(&mut self) -> &mut OperationGraph {
204        &mut self.graph
205    }
206
207    pub fn add_operation(&mut self, operation: Operation) -> Result<()> {
208        self.graph.add_operation(operation);
209        self.save()
210    }
211
212    pub fn save(&self) -> Result<()> {
213        self.graph.save(&self.path)
214    }
215
216    /// Get operations since a specific operation (exclusive)
217    pub fn operations_since(&self, since_id: &str) -> Vec<&Operation> {
218        let since_op = match self.graph.get_operation(since_id) {
219            Some(op) => op,
220            None => return Vec::new(),
221        };
222
223        self.graph
224            .operations
225            .values()
226            .filter(|op| op.timestamp > since_op.timestamp)
227            .collect()
228    }
229
230    /// Find the common ancestor of two operations (reserved for merge/rebase features)
231    #[allow(dead_code)]
232    pub fn common_ancestor(&self, id1: &str, id2: &str) -> Option<&Operation> {
233        let ancestry1: std::collections::HashSet<_> = self
234            .graph
235            .get_ancestry(id1)
236            .iter()
237            .map(|op| op.id.as_str())
238            .collect();
239
240        self.graph
241            .get_ancestry(id2)
242            .into_iter()
243            .find(|op| ancestry1.contains(op.id.as_str()))
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use crate::history::operation::{InstallOperationOptions, OperationId, OperationType};
251    use std::path::PathBuf;
252    use tempfile::TempDir;
253
254    fn create_test_operation(id: &str, parent: Option<&str>, checkpoint: &str) -> Operation {
255        Operation {
256            id: OperationId::from_string(id),
257            operation_type: OperationType::Install {
258                profile: "test".into(),
259                source: None,
260                target: PathBuf::from("/test"),
261                options: InstallOperationOptions::default(),
262            },
263            parent: parent.map(|p| OperationId::from_string(p)),
264            checkpoint_id: checkpoint.into(),
265            timestamp: chrono::Utc::now(),
266            description: None,
267        }
268    }
269
270    #[test]
271    fn graph_add_operation() {
272        let mut graph = OperationGraph::new();
273
274        let op1 = create_test_operation("op-001", None, "cp-001");
275        graph.add_operation(op1);
276
277        assert_eq!(graph.len(), 1);
278        assert_eq!(graph.roots.len(), 1);
279        assert!(graph.get_operation("op-001").is_some());
280    }
281
282    #[test]
283    fn graph_ancestry() {
284        let mut graph = OperationGraph::new();
285
286        let op1 = create_test_operation("op-001", None, "cp-001");
287        let op2 = create_test_operation("op-002", Some("op-001"), "cp-002");
288        let op3 = create_test_operation("op-003", Some("op-002"), "cp-003");
289
290        graph.add_operation(op1);
291        graph.add_operation(op2);
292        graph.add_operation(op3);
293
294        let ancestry = graph.get_ancestry("op-003");
295        assert_eq!(ancestry.len(), 3);
296        assert_eq!(ancestry[0].id.as_str(), "op-003");
297        assert_eq!(ancestry[1].id.as_str(), "op-002");
298        assert_eq!(ancestry[2].id.as_str(), "op-001");
299    }
300
301    #[test]
302    fn graph_persistence() -> Result<()> {
303        let temp = TempDir::new()?;
304        let path = temp.path().join("graph.toml");
305
306        let mut graph = OperationGraph::new();
307        let op = create_test_operation("op-001", None, "cp-001");
308        graph.add_operation(op);
309        graph.save(&path)?;
310
311        let loaded = OperationGraph::load(&path)?;
312        assert_eq!(loaded.len(), 1);
313        assert!(loaded.get_operation("op-001").is_some());
314
315        Ok(())
316    }
317
318    #[test]
319    fn graph_children() {
320        let mut graph = OperationGraph::new();
321
322        let op1 = create_test_operation("op-001", None, "cp-001");
323        let op2 = create_test_operation("op-002", Some("op-001"), "cp-002");
324        let op3 = create_test_operation("op-003", Some("op-001"), "cp-003");
325
326        graph.add_operation(op1);
327        graph.add_operation(op2);
328        graph.add_operation(op3);
329
330        let children = graph.get_children("op-001");
331        assert_eq!(children.len(), 2);
332    }
333
334    #[test]
335    fn graph_remove_tree() {
336        let mut graph = OperationGraph::new();
337
338        let op1 = create_test_operation("op-001", None, "cp-001");
339        let op2 = create_test_operation("op-002", Some("op-001"), "cp-002");
340        let op3 = create_test_operation("op-003", Some("op-002"), "cp-003");
341
342        graph.add_operation(op1);
343        graph.add_operation(op2);
344        graph.add_operation(op3);
345
346        let removed = graph.remove_operation_tree("op-002");
347        assert_eq!(removed.len(), 2);
348        assert_eq!(graph.len(), 1);
349    }
350}