dot_agent_core/history/
graph.rs1use 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#[derive(Debug, Clone, Default, Serialize, Deserialize)]
12pub struct OperationGraph {
13 pub version: u32,
15 #[serde(default)]
17 pub operations: HashMap<String, Operation>,
18 #[serde(default)]
20 pub roots: Vec<String>,
21 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 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 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 pub fn add_operation(&mut self, operation: Operation) {
63 let id = operation.id.as_str().to_string();
64
65 if operation.parent.is_none() {
67 self.roots.push(id.clone());
68 }
69
70 self.head = Some(id.clone());
72
73 self.operations.insert(id, operation);
75 }
76
77 pub fn get_operation(&self, id: &str) -> Option<&Operation> {
79 self.operations.get(id)
80 }
81
82 pub fn get_head(&self) -> Option<&Operation> {
84 self.head.as_ref().and_then(|id| self.operations.get(id))
85 }
86
87 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 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 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 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 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 pub fn len(&self) -> usize {
135 self.operations.len()
136 }
137
138 pub fn is_empty(&self) -> bool {
140 self.operations.is_empty()
141 }
142
143 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 let children: Vec<_> = self
151 .operations
152 .values()
153 .filter(|op| op.parent.as_ref().map(|p| p.as_str()) == Some(¤t_id))
154 .map(|op| op.id.as_str().to_string())
155 .collect();
156
157 to_remove.extend(children);
158
159 if let Some(op) = self.operations.remove(¤t_id) {
161 removed.push(op);
162 }
163
164 self.roots.retain(|r| r != ¤t_id);
166 }
167
168 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
185pub 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 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 #[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}