1use chrono::{DateTime, Utc};
27use serde::{Deserialize, Serialize};
28use similar::TextDiff;
29use std::collections::HashMap;
30use std::sync::RwLock;
31
32use crate::error::{read_or_recover, write_or_recover};
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct FileSnapshot {
37 pub version: usize,
39 pub path: String,
41 pub content: String,
43 pub timestamp: DateTime<Utc>,
45 pub tool_name: String,
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct VersionSummary {
52 pub version: usize,
54 pub path: String,
56 pub timestamp: DateTime<Utc>,
58 pub tool_name: String,
60 pub size: usize,
62}
63
64impl From<&FileSnapshot> for VersionSummary {
65 fn from(snapshot: &FileSnapshot) -> Self {
66 Self {
67 version: snapshot.version,
68 path: snapshot.path.clone(),
69 timestamp: snapshot.timestamp,
70 tool_name: snapshot.tool_name.clone(),
71 size: snapshot.content.len(),
72 }
73 }
74}
75
76pub struct FileHistory {
81 snapshots: RwLock<HashMap<String, Vec<FileSnapshot>>>,
83 max_snapshots: usize,
85}
86
87impl FileHistory {
88 pub fn new(max_snapshots: usize) -> Self {
93 Self {
94 snapshots: RwLock::new(HashMap::new()),
95 max_snapshots,
96 }
97 }
98
99 pub fn save_snapshot(&self, path: &str, content: &str, tool_name: &str) -> usize {
103 let mut snapshots = write_or_recover(&self.snapshots);
104
105 let file_versions = snapshots.entry(path.to_string()).or_default();
106 let version = file_versions.len();
107
108 file_versions.push(FileSnapshot {
109 version,
110 path: path.to_string(),
111 content: content.to_string(),
112 timestamp: Utc::now(),
113 tool_name: tool_name.to_string(),
114 });
115
116 self.evict_if_needed(&mut snapshots);
118
119 version
120 }
121
122 pub fn list_versions(&self, path: &str) -> Vec<VersionSummary> {
124 let snapshots = read_or_recover(&self.snapshots);
125 snapshots
126 .get(path)
127 .map(|versions| versions.iter().map(VersionSummary::from).collect())
128 .unwrap_or_default()
129 }
130
131 pub fn list_files(&self) -> Vec<(String, usize)> {
133 let snapshots = read_or_recover(&self.snapshots);
134 snapshots
135 .iter()
136 .map(|(path, versions)| (path.clone(), versions.len()))
137 .collect()
138 }
139
140 pub fn get_version(&self, path: &str, version: usize) -> Option<FileSnapshot> {
142 let snapshots = read_or_recover(&self.snapshots);
143 snapshots
144 .get(path)
145 .and_then(|versions| versions.get(version).cloned())
146 }
147
148 pub fn get_latest(&self, path: &str) -> Option<FileSnapshot> {
150 let snapshots = read_or_recover(&self.snapshots);
151 snapshots
152 .get(path)
153 .and_then(|versions| versions.last().cloned())
154 }
155
156 pub fn diff(&self, path: &str, from_version: usize, to_version: usize) -> Option<String> {
160 let snapshots = read_or_recover(&self.snapshots);
161 let versions = snapshots.get(path)?;
162
163 let from = versions.get(from_version)?;
164 let to = versions.get(to_version)?;
165
166 Some(generate_unified_diff(
167 &from.content,
168 &to.content,
169 path,
170 from_version,
171 to_version,
172 ))
173 }
174
175 pub fn diff_with_current(
177 &self,
178 path: &str,
179 version: usize,
180 current_content: &str,
181 ) -> Option<String> {
182 let snapshots = read_or_recover(&self.snapshots);
183 let versions = snapshots.get(path)?;
184 let from = versions.get(version)?;
185
186 Some(generate_unified_diff(
187 &from.content,
188 current_content,
189 path,
190 version,
191 versions.len(), ))
193 }
194
195 pub fn total_snapshots(&self) -> usize {
197 let snapshots = read_or_recover(&self.snapshots);
198 snapshots.values().map(|v| v.len()).sum()
199 }
200
201 pub fn clear_file(&self, path: &str) {
203 let mut snapshots = write_or_recover(&self.snapshots);
204 snapshots.remove(path);
205 }
206
207 pub fn clear_all(&self) {
209 let mut snapshots = write_or_recover(&self.snapshots);
210 snapshots.clear();
211 }
212
213 fn evict_if_needed(&self, snapshots: &mut HashMap<String, Vec<FileSnapshot>>) {
215 let total: usize = snapshots.values().map(|v| v.len()).sum();
216 if total <= self.max_snapshots {
217 return;
218 }
219
220 let to_remove = total - self.max_snapshots;
221
222 let mut all_entries: Vec<(String, usize, DateTime<Utc>)> = Vec::new();
224 for (path, versions) in snapshots.iter() {
225 for snapshot in versions {
226 all_entries.push((path.clone(), snapshot.version, snapshot.timestamp));
227 }
228 }
229 all_entries.sort_by_key(|e| e.2);
230
231 for (path, version, _) in all_entries.into_iter().take(to_remove) {
233 if let Some(versions) = snapshots.get_mut(&path) {
234 versions.retain(|s| s.version != version);
235 if versions.is_empty() {
236 snapshots.remove(&path);
237 }
238 }
239 }
240 }
241}
242
243fn generate_unified_diff(
245 old: &str,
246 new: &str,
247 path: &str,
248 from_version: usize,
249 to_version: usize,
250) -> String {
251 let diff = TextDiff::from_lines(old, new);
252 let mut output = String::new();
253
254 output.push_str(&format!("--- a/{} (version {})\n", path, from_version));
255 output.push_str(&format!("+++ b/{} (version {})\n", path, to_version));
256
257 for hunk in diff.unified_diff().context_radius(3).iter_hunks() {
258 output.push_str(&format!("{}", hunk));
259 }
260
261 output
262}
263
264pub fn is_file_modifying_tool(tool_name: &str) -> bool {
266 matches!(tool_name, "write" | "edit" | "patch")
267}
268
269pub fn extract_file_path(tool_name: &str, args: &serde_json::Value) -> Option<String> {
271 if is_file_modifying_tool(tool_name) {
272 args.get("file_path")
273 .and_then(|v| v.as_str())
274 .map(|s| s.to_string())
275 } else {
276 None
277 }
278}
279
280#[cfg(test)]
281mod tests {
282 use super::*;
283
284 #[test]
289 fn test_new_history() {
290 let history = FileHistory::new(100);
291 assert_eq!(history.total_snapshots(), 0);
292 assert!(history.list_files().is_empty());
293 }
294
295 #[test]
296 fn test_save_snapshot() {
297 let history = FileHistory::new(100);
298 let v = history.save_snapshot("test.rs", "fn main() {}", "write");
299 assert_eq!(v, 0);
300 assert_eq!(history.total_snapshots(), 1);
301 }
302
303 #[test]
304 fn test_save_multiple_snapshots() {
305 let history = FileHistory::new(100);
306 let v0 = history.save_snapshot("test.rs", "version 0", "write");
307 let v1 = history.save_snapshot("test.rs", "version 1", "edit");
308 let v2 = history.save_snapshot("test.rs", "version 2", "patch");
309 assert_eq!(v0, 0);
310 assert_eq!(v1, 1);
311 assert_eq!(v2, 2);
312 assert_eq!(history.total_snapshots(), 3);
313 }
314
315 #[test]
316 fn test_save_multiple_files() {
317 let history = FileHistory::new(100);
318 history.save_snapshot("a.rs", "content a", "write");
319 history.save_snapshot("b.rs", "content b", "write");
320 assert_eq!(history.total_snapshots(), 2);
321 assert_eq!(history.list_files().len(), 2);
322 }
323
324 #[test]
329 fn test_list_versions_empty() {
330 let history = FileHistory::new(100);
331 assert!(history.list_versions("nonexistent.rs").is_empty());
332 }
333
334 #[test]
335 fn test_list_versions() {
336 let history = FileHistory::new(100);
337 history.save_snapshot("test.rs", "v0", "write");
338 history.save_snapshot("test.rs", "v1", "edit");
339
340 let versions = history.list_versions("test.rs");
341 assert_eq!(versions.len(), 2);
342 assert_eq!(versions[0].version, 0);
343 assert_eq!(versions[0].tool_name, "write");
344 assert_eq!(versions[0].size, 2);
345 assert_eq!(versions[1].version, 1);
346 assert_eq!(versions[1].tool_name, "edit");
347 }
348
349 #[test]
354 fn test_get_version() {
355 let history = FileHistory::new(100);
356 history.save_snapshot("test.rs", "original", "write");
357 history.save_snapshot("test.rs", "modified", "edit");
358
359 let v0 = history.get_version("test.rs", 0).unwrap();
360 assert_eq!(v0.content, "original");
361 assert_eq!(v0.tool_name, "write");
362
363 let v1 = history.get_version("test.rs", 1).unwrap();
364 assert_eq!(v1.content, "modified");
365 }
366
367 #[test]
368 fn test_get_version_nonexistent() {
369 let history = FileHistory::new(100);
370 assert!(history.get_version("test.rs", 0).is_none());
371
372 history.save_snapshot("test.rs", "content", "write");
373 assert!(history.get_version("test.rs", 99).is_none());
374 }
375
376 #[test]
377 fn test_get_latest() {
378 let history = FileHistory::new(100);
379 assert!(history.get_latest("test.rs").is_none());
380
381 history.save_snapshot("test.rs", "v0", "write");
382 history.save_snapshot("test.rs", "v1", "edit");
383
384 let latest = history.get_latest("test.rs").unwrap();
385 assert_eq!(latest.content, "v1");
386 assert_eq!(latest.version, 1);
387 }
388
389 #[test]
394 fn test_diff_between_versions() {
395 let history = FileHistory::new(100);
396 history.save_snapshot("test.rs", "line1\nline2\nline3\n", "write");
397 history.save_snapshot("test.rs", "line1\nmodified\nline3\n", "edit");
398
399 let diff = history.diff("test.rs", 0, 1).unwrap();
400 assert!(diff.contains("--- a/test.rs (version 0)"));
401 assert!(diff.contains("+++ b/test.rs (version 1)"));
402 assert!(diff.contains("-line2"));
403 assert!(diff.contains("+modified"));
404 }
405
406 #[test]
407 fn test_diff_nonexistent_version() {
408 let history = FileHistory::new(100);
409 history.save_snapshot("test.rs", "content", "write");
410 assert!(history.diff("test.rs", 0, 5).is_none());
411 assert!(history.diff("nonexistent.rs", 0, 1).is_none());
412 }
413
414 #[test]
415 fn test_diff_same_version() {
416 let history = FileHistory::new(100);
417 history.save_snapshot("test.rs", "same content\n", "write");
418
419 let diff = history.diff("test.rs", 0, 0).unwrap();
420 assert!(diff.contains("--- a/test.rs"));
422 assert!(!diff.contains("-same content"));
423 }
424
425 #[test]
426 fn test_diff_with_current() {
427 let history = FileHistory::new(100);
428 history.save_snapshot("test.rs", "old\n", "write");
429
430 let diff = history.diff_with_current("test.rs", 0, "new\n").unwrap();
431 assert!(diff.contains("-old"));
432 assert!(diff.contains("+new"));
433 }
434
435 #[test]
440 fn test_list_files() {
441 let history = FileHistory::new(100);
442 history.save_snapshot("a.rs", "a", "write");
443 history.save_snapshot("b.rs", "b1", "write");
444 history.save_snapshot("b.rs", "b2", "edit");
445
446 let files = history.list_files();
447 assert_eq!(files.len(), 2);
448
449 let a_count = files.iter().find(|(p, _)| p == "a.rs").unwrap().1;
450 let b_count = files.iter().find(|(p, _)| p == "b.rs").unwrap().1;
451 assert_eq!(a_count, 1);
452 assert_eq!(b_count, 2);
453 }
454
455 #[test]
460 fn test_clear_file() {
461 let history = FileHistory::new(100);
462 history.save_snapshot("a.rs", "a", "write");
463 history.save_snapshot("b.rs", "b", "write");
464
465 history.clear_file("a.rs");
466 assert_eq!(history.total_snapshots(), 1);
467 assert!(history.list_versions("a.rs").is_empty());
468 assert_eq!(history.list_versions("b.rs").len(), 1);
469 }
470
471 #[test]
472 fn test_clear_all() {
473 let history = FileHistory::new(100);
474 history.save_snapshot("a.rs", "a", "write");
475 history.save_snapshot("b.rs", "b", "write");
476
477 history.clear_all();
478 assert_eq!(history.total_snapshots(), 0);
479 assert!(history.list_files().is_empty());
480 }
481
482 #[test]
487 fn test_eviction_when_over_limit() {
488 let history = FileHistory::new(3);
489 history.save_snapshot("test.rs", "v0", "write");
490 history.save_snapshot("test.rs", "v1", "edit");
491 history.save_snapshot("test.rs", "v2", "edit");
492 assert_eq!(history.total_snapshots(), 3);
494
495 history.save_snapshot("test.rs", "v3", "edit");
497 assert!(history.total_snapshots() <= 3);
498 }
499
500 #[test]
501 fn test_eviction_across_files() {
502 let history = FileHistory::new(3);
503 history.save_snapshot("a.rs", "a0", "write");
504 history.save_snapshot("b.rs", "b0", "write");
505 history.save_snapshot("c.rs", "c0", "write");
506
507 history.save_snapshot("d.rs", "d0", "write");
509 assert!(history.total_snapshots() <= 3);
510 }
511
512 #[test]
517 fn test_version_summary_from_snapshot() {
518 let snapshot = FileSnapshot {
519 version: 5,
520 path: "test.rs".to_string(),
521 content: "hello world".to_string(),
522 timestamp: Utc::now(),
523 tool_name: "edit".to_string(),
524 };
525 let summary = VersionSummary::from(&snapshot);
526 assert_eq!(summary.version, 5);
527 assert_eq!(summary.path, "test.rs");
528 assert_eq!(summary.tool_name, "edit");
529 assert_eq!(summary.size, 11); }
531
532 #[test]
537 fn test_is_file_modifying_tool() {
538 assert!(is_file_modifying_tool("write"));
539 assert!(is_file_modifying_tool("edit"));
540 assert!(is_file_modifying_tool("patch"));
541 assert!(!is_file_modifying_tool("read"));
542 assert!(!is_file_modifying_tool("bash"));
543 assert!(!is_file_modifying_tool("grep"));
544 assert!(!is_file_modifying_tool("glob"));
545 assert!(!is_file_modifying_tool("ls"));
546 }
547
548 #[test]
549 fn test_extract_file_path() {
550 let args = serde_json::json!({"file_path": "src/main.rs", "content": "hello"});
551 assert_eq!(
552 extract_file_path("write", &args),
553 Some("src/main.rs".to_string())
554 );
555 assert_eq!(
556 extract_file_path("edit", &args),
557 Some("src/main.rs".to_string())
558 );
559 assert_eq!(
560 extract_file_path("patch", &args),
561 Some("src/main.rs".to_string())
562 );
563 assert_eq!(extract_file_path("read", &args), None);
564 assert_eq!(extract_file_path("bash", &args), None);
565 }
566
567 #[test]
568 fn test_extract_file_path_missing() {
569 let args = serde_json::json!({"content": "hello"});
570 assert_eq!(extract_file_path("write", &args), None);
571 }
572
573 #[test]
578 fn test_generate_unified_diff() {
579 let old = "line1\nline2\nline3\n";
580 let new = "line1\nchanged\nline3\n";
581 let diff = generate_unified_diff(old, new, "test.rs", 0, 1);
582 assert!(diff.contains("--- a/test.rs (version 0)"));
583 assert!(diff.contains("+++ b/test.rs (version 1)"));
584 assert!(diff.contains("-line2"));
585 assert!(diff.contains("+changed"));
586 }
587
588 #[test]
589 fn test_generate_unified_diff_no_changes() {
590 let content = "same\n";
591 let diff = generate_unified_diff(content, content, "test.rs", 0, 0);
592 assert!(diff.contains("--- a/test.rs"));
593 assert!(!diff.contains("@@"));
595 }
596
597 #[test]
598 fn test_generate_unified_diff_addition() {
599 let old = "line1\nline3\n";
600 let new = "line1\nline2\nline3\n";
601 let diff = generate_unified_diff(old, new, "test.rs", 0, 1);
602 assert!(diff.contains("+line2"));
603 }
604
605 #[test]
606 fn test_generate_unified_diff_deletion() {
607 let old = "line1\nline2\nline3\n";
608 let new = "line1\nline3\n";
609 let diff = generate_unified_diff(old, new, "test.rs", 0, 1);
610 assert!(diff.contains("-line2"));
611 }
612}