Skip to main content

aft/
checkpoint.rs

1use std::collections::HashMap;
2use std::fs;
3use std::io;
4use std::path::{Path, PathBuf};
5use std::time::Duration;
6
7use crate::backup::BackupStore;
8use crate::error::AftError;
9use crate::fs_lock;
10
11const CHECKPOINT_LOCK_TIMEOUT: Duration = Duration::from_secs(30);
12
13/// Metadata about a checkpoint, returned by list/create/restore.
14#[derive(Debug, Clone)]
15pub struct CheckpointInfo {
16    pub name: String,
17    pub file_count: usize,
18    pub created_at: u64,
19    /// Paths that could not be snapshotted (e.g. deleted since last edit),
20    /// paired with the OS-level error that stopped us from reading them.
21    /// Empty on successful round-trips. Populated only on `create()` — the
22    /// `list()` / `restore()` paths leave it empty.
23    pub skipped: Vec<(PathBuf, String)>,
24}
25
26/// A stored checkpoint: a snapshot of multiple file contents and metadata.
27#[derive(Debug, Clone)]
28struct Checkpoint {
29    name: String,
30    file_contents: HashMap<PathBuf, CheckpointFile>,
31    created_at: u64,
32}
33
34#[derive(Debug, Clone)]
35struct CheckpointFile {
36    metadata: fs::Metadata,
37    kind: CheckpointFileKind,
38}
39
40#[derive(Debug, Clone)]
41enum CheckpointFileKind {
42    Regular {
43        bytes: Vec<u8>,
44    },
45    Symlink {
46        target: PathBuf,
47        target_is_dir: bool,
48    },
49}
50
51impl CheckpointFile {
52    fn read(path: &Path) -> io::Result<Self> {
53        let metadata = fs::symlink_metadata(path)?;
54        let file_type = metadata.file_type();
55        if file_type.is_symlink() {
56            let target = fs::read_link(path)?;
57            let target_is_dir = fs::metadata(path)
58                .map(|target_metadata| target_metadata.is_dir())
59                .unwrap_or(false);
60            return Ok(Self {
61                metadata,
62                kind: CheckpointFileKind::Symlink {
63                    target,
64                    target_is_dir,
65                },
66            });
67        }
68
69        if metadata.is_file() {
70            let bytes = fs::read(path)?;
71            return Ok(Self {
72                metadata,
73                kind: CheckpointFileKind::Regular { bytes },
74            });
75        }
76
77        Err(io::Error::new(
78            io::ErrorKind::InvalidInput,
79            "not a regular file or symlink",
80        ))
81    }
82
83    fn read_optional(path: &Path) -> io::Result<Option<Self>> {
84        match Self::read(path) {
85            Ok(snapshot) => Ok(Some(snapshot)),
86            Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None),
87            Err(error) => Err(error),
88        }
89    }
90}
91
92/// Workspace-wide, per-session checkpoint store.
93///
94/// Partitioned by session (issue #14): two OpenCode sessions sharing one bridge
95/// can both create checkpoints named `snap1` without collision, and restoring
96/// from one session does not leak the other's file set. Checkpoints are kept
97/// in memory only — a bridge crash drops all of them, which is a deliberate
98/// trade-off to keep this refactor bounded. Durable checkpoints are a possible
99/// follow-up.
100#[derive(Debug)]
101pub struct CheckpointStore {
102    /// session -> name -> checkpoint
103    checkpoints: HashMap<String, HashMap<String, Checkpoint>>,
104    lock_path: PathBuf,
105    lock_timeout: Duration,
106}
107
108impl CheckpointStore {
109    pub fn new() -> Self {
110        let project_root = std::env::current_dir().unwrap_or_else(|_| std::env::temp_dir());
111        let project_key = crate::search_index::project_cache_key(&project_root);
112        let lock_path = crate::bash_background::storage_dir(None)
113            .join("checkpoints")
114            .join(project_key)
115            .join("checkpoint.lock");
116        Self::with_lock_path(lock_path, CHECKPOINT_LOCK_TIMEOUT)
117    }
118
119    fn with_lock_path(lock_path: PathBuf, lock_timeout: Duration) -> Self {
120        CheckpointStore {
121            checkpoints: HashMap::new(),
122            lock_path,
123            lock_timeout,
124        }
125    }
126
127    fn acquire_mutation_lock(&self) -> Result<fs_lock::LockGuard, AftError> {
128        if let Some(parent) = self.lock_path.parent() {
129            fs::create_dir_all(parent).map_err(|error| AftError::IoError {
130                path: parent.display().to_string(),
131                message: format!("failed to create checkpoint lock directory: {error}"),
132            })?;
133        }
134
135        fs_lock::try_acquire(&self.lock_path, self.lock_timeout).map_err(|error| match error {
136            fs_lock::AcquireError::Timeout => AftError::IoError {
137                path: self.lock_path.display().to_string(),
138                message: "timed out acquiring checkpoint mutation lock".to_string(),
139            },
140            fs_lock::AcquireError::Io(error) => AftError::IoError {
141                path: self.lock_path.display().to_string(),
142                message: format!("failed to acquire checkpoint mutation lock: {error}"),
143            },
144        })
145    }
146
147    /// Create a checkpoint by reading the given files, scoped to `session`.
148    ///
149    /// If `files` is empty, snapshots all tracked files for **that session**
150    /// from the BackupStore (other sessions' tracked files are not visible).
151    /// Overwrites any existing checkpoint with the same name in this session.
152    ///
153    /// Unreadable paths (e.g. deleted since their last edit) are skipped with
154    /// a warning instead of failing the whole checkpoint. The paths and their
155    /// errors are returned via `CheckpointInfo::skipped` so callers can
156    /// surface them. A checkpoint is only rejected outright when *every*
157    /// requested path fails — that case still returns a `FileNotFound`
158    /// error so callers can distinguish "partial success" from "nothing
159    /// snapshotted at all".
160    pub fn create(
161        &mut self,
162        session: &str,
163        name: &str,
164        files: Vec<PathBuf>,
165        backup_store: &BackupStore,
166    ) -> Result<CheckpointInfo, AftError> {
167        let _mutation_lock = self.acquire_mutation_lock()?;
168        let explicit_request = !files.is_empty();
169        let file_list = if files.is_empty() {
170            backup_store.tracked_files(session)
171        } else {
172            files
173        };
174
175        let mut file_contents = HashMap::new();
176        let mut skipped: Vec<(PathBuf, String)> = Vec::new();
177        for path in &file_list {
178            match CheckpointFile::read(path) {
179                Ok(snapshot) => {
180                    file_contents.insert(path.clone(), snapshot);
181                }
182                Err(e) => {
183                    crate::slog_warn!(
184                        "checkpoint {}: skipping unreadable file {}: {}",
185                        name,
186                        path.display(),
187                        e
188                    );
189                    skipped.push((path.clone(), e.to_string()));
190                }
191            }
192        }
193
194        // If the caller explicitly named a single file and it was unreadable,
195        // that's a real error — surface it rather than silently returning an
196        // empty checkpoint. For empty `files` (tracked-file fallback) with no
197        // readable files at all, the empty-file checkpoint is a legitimate
198        // "nothing to snapshot" outcome and we keep it.
199        if explicit_request && file_contents.is_empty() && !skipped.is_empty() {
200            let (path, err) = &skipped[0];
201            return Err(AftError::FileNotFound {
202                path: format!("{}: {}", path.display(), err),
203            });
204        }
205
206        let created_at = current_timestamp();
207        let file_count = file_contents.len();
208
209        let checkpoint = Checkpoint {
210            name: name.to_string(),
211            file_contents,
212            created_at,
213        };
214
215        self.checkpoints
216            .entry(session.to_string())
217            .or_default()
218            .insert(name.to_string(), checkpoint);
219
220        if skipped.is_empty() {
221            crate::slog_info!("checkpoint created: {} ({} files)", name, file_count);
222        } else {
223            crate::slog_info!(
224                "checkpoint created: {} ({} files, {} skipped)",
225                name,
226                file_count,
227                skipped.len()
228            );
229        }
230
231        Ok(CheckpointInfo {
232            name: name.to_string(),
233            file_count,
234            created_at,
235            skipped,
236        })
237    }
238
239    /// Restore a checkpoint by overwriting files with stored content.
240    pub fn restore(&self, session: &str, name: &str) -> Result<CheckpointInfo, AftError> {
241        let _mutation_lock = self.acquire_mutation_lock()?;
242        let checkpoint = self.get(session, name)?;
243        let mut paths = checkpoint.file_contents.keys().cloned().collect::<Vec<_>>();
244        paths.sort();
245
246        restore_paths_atomically(checkpoint, &paths)?;
247
248        crate::slog_info!("checkpoint restored: {}", name);
249
250        Ok(CheckpointInfo {
251            name: checkpoint.name.clone(),
252            file_count: checkpoint.file_contents.len(),
253            created_at: checkpoint.created_at,
254            skipped: Vec::new(),
255        })
256    }
257
258    /// Restore a checkpoint using a caller-validated path list.
259    pub fn restore_validated(
260        &self,
261        session: &str,
262        name: &str,
263        validated_paths: &[PathBuf],
264    ) -> Result<CheckpointInfo, AftError> {
265        let _mutation_lock = self.acquire_mutation_lock()?;
266        let checkpoint = self.get(session, name)?;
267
268        for path in validated_paths {
269            checkpoint
270                .file_contents
271                .get(path)
272                .ok_or_else(|| AftError::FileNotFound {
273                    path: path.display().to_string(),
274                })?;
275        }
276        restore_paths_atomically(checkpoint, validated_paths)?;
277
278        crate::slog_info!("checkpoint restored: {}", name);
279
280        Ok(CheckpointInfo {
281            name: checkpoint.name.clone(),
282            file_count: checkpoint.file_contents.len(),
283            created_at: checkpoint.created_at,
284            skipped: Vec::new(),
285        })
286    }
287
288    /// Return the file paths stored for a checkpoint.
289    pub fn file_paths(&self, session: &str, name: &str) -> Result<Vec<PathBuf>, AftError> {
290        let checkpoint = self.get(session, name)?;
291        Ok(checkpoint.file_contents.keys().cloned().collect())
292    }
293
294    /// Delete a checkpoint from a session. Returns true when a checkpoint was removed.
295    pub fn delete(&mut self, session: &str, name: &str) -> bool {
296        let Some(session_checkpoints) = self.checkpoints.get_mut(session) else {
297            return false;
298        };
299        let removed = session_checkpoints.remove(name).is_some();
300        if session_checkpoints.is_empty() {
301            self.checkpoints.remove(session);
302        }
303        removed
304    }
305
306    /// List all checkpoints for this session with metadata.
307    pub fn list(&self, session: &str) -> Vec<CheckpointInfo> {
308        self.checkpoints
309            .get(session)
310            .map(|s| {
311                s.values()
312                    .map(|cp| CheckpointInfo {
313                        name: cp.name.clone(),
314                        file_count: cp.file_contents.len(),
315                        created_at: cp.created_at,
316                        skipped: Vec::new(),
317                    })
318                    .collect()
319            })
320            .unwrap_or_default()
321    }
322
323    /// Total checkpoint count across all sessions (for `/aft-status`).
324    pub fn total_count(&self) -> usize {
325        self.checkpoints.values().map(|s| s.len()).sum()
326    }
327
328    /// Remove checkpoints older than `ttl_hours` across all sessions.
329    /// Empty session entries are pruned after cleanup.
330    pub fn cleanup(&mut self, ttl_hours: u32) {
331        let now = current_timestamp();
332        let ttl_secs = ttl_hours as u64 * 3600;
333        self.checkpoints.retain(|_, session_cps| {
334            session_cps.retain(|_, cp| now.saturating_sub(cp.created_at) < ttl_secs);
335            !session_cps.is_empty()
336        });
337    }
338
339    fn get(&self, session: &str, name: &str) -> Result<&Checkpoint, AftError> {
340        self.checkpoints
341            .get(session)
342            .and_then(|s| s.get(name))
343            .ok_or_else(|| AftError::CheckpointNotFound {
344                name: name.to_string(),
345            })
346    }
347}
348
349fn restore_paths_atomically(checkpoint: &Checkpoint, paths: &[PathBuf]) -> Result<(), AftError> {
350    let mut pre_restore_snapshot: HashMap<PathBuf, Option<CheckpointFile>> = HashMap::new();
351    for path in paths {
352        let current = CheckpointFile::read_optional(path).map_err(|error| AftError::IoError {
353            path: path.display().to_string(),
354            message: format!("failed to snapshot pre-restore file metadata: {error}"),
355        })?;
356        pre_restore_snapshot.insert(path.clone(), current);
357    }
358
359    let mut restored_paths: Vec<PathBuf> = Vec::new();
360    let mut created_dirs: Vec<PathBuf> = Vec::new();
361    for path in paths {
362        let snapshot =
363            checkpoint
364                .file_contents
365                .get(path)
366                .ok_or_else(|| AftError::FileNotFound {
367                    path: path.display().to_string(),
368                })?;
369        if let Err(e) = write_restored_file(path, snapshot, &mut created_dirs) {
370            let mut rollback_errors = Vec::new();
371            if let Some(snapshot) = pre_restore_snapshot.get(path) {
372                if let Err(rollback_error) = restore_snapshot_file(path, snapshot.as_ref()) {
373                    rollback_errors.push(format!("{}: {}", path.display(), rollback_error));
374                }
375            }
376            for restored_path in restored_paths.iter().rev() {
377                if let Some(snapshot) = pre_restore_snapshot.get(restored_path) {
378                    if let Err(rollback_error) =
379                        restore_snapshot_file(restored_path, snapshot.as_ref())
380                    {
381                        rollback_errors.push(format!(
382                            "{}: {}",
383                            restored_path.display(),
384                            rollback_error
385                        ));
386                    }
387                }
388            }
389            let dirs_rollback_ok = rollback_created_dirs(&created_dirs);
390            if rollback_errors.is_empty() && dirs_rollback_ok {
391                return Err(e);
392            }
393            return Err(AftError::IoError {
394                path: path.display().to_string(),
395                message: format!(
396                    "{}; restore_checkpoint rollback_succeeded: {}; rollback_errors: {}",
397                    e,
398                    rollback_errors.is_empty() && dirs_rollback_ok,
399                    if rollback_errors.is_empty() {
400                        "none".to_string()
401                    } else {
402                        rollback_errors.join("; ")
403                    }
404                ),
405            });
406        }
407        restored_paths.push(path.clone());
408    }
409
410    Ok(())
411}
412
413fn restore_snapshot_file(path: &Path, snapshot: Option<&CheckpointFile>) -> Result<(), AftError> {
414    match snapshot {
415        Some(snapshot) => write_restored_file(path, snapshot, &mut Vec::new()),
416        None => remove_file_if_exists(path).map_err(|error| AftError::IoError {
417            path: path.display().to_string(),
418            message: format!("failed to remove file during checkpoint restore rollback: {error}"),
419        }),
420    }
421}
422
423fn write_restored_file(
424    path: &Path,
425    snapshot: &CheckpointFile,
426    created_dirs: &mut Vec<PathBuf>,
427) -> Result<(), AftError> {
428    create_parent_dirs(path, created_dirs)?;
429
430    match &snapshot.kind {
431        CheckpointFileKind::Regular { bytes } => {
432            if path_is_symlink(path) {
433                remove_file_if_exists(path).map_err(|error| AftError::IoError {
434                    path: path.display().to_string(),
435                    message: format!("failed to replace symlink with regular file: {error}"),
436                })?;
437            }
438            fs::write(path, bytes).map_err(|error| AftError::IoError {
439                path: path.display().to_string(),
440                message: format!("failed to restore checkpoint file contents: {error}"),
441            })?;
442            fs::set_permissions(path, snapshot.metadata.permissions()).map_err(|error| {
443                AftError::IoError {
444                    path: path.display().to_string(),
445                    message: format!("failed to restore checkpoint file permissions: {error}"),
446                }
447            })
448        }
449        CheckpointFileKind::Symlink {
450            target,
451            target_is_dir,
452        } => {
453            remove_file_if_exists(path).map_err(|error| AftError::IoError {
454                path: path.display().to_string(),
455                message: format!("failed to replace file with checkpoint symlink: {error}"),
456            })?;
457            create_symlink(target, path, *target_is_dir).map_err(|error| AftError::IoError {
458                path: path.display().to_string(),
459                message: format!("failed to restore checkpoint symlink: {error}"),
460            })
461        }
462    }
463}
464
465fn create_parent_dirs(path: &Path, created_dirs: &mut Vec<PathBuf>) -> Result<(), AftError> {
466    if let Some(parent) = path.parent() {
467        let missing_dirs = missing_parent_dirs(parent);
468        fs::create_dir_all(parent).map_err(|error| AftError::IoError {
469            path: parent.display().to_string(),
470            message: format!("failed to create checkpoint restore parent directories: {error}"),
471        })?;
472        created_dirs.extend(missing_dirs);
473    }
474    Ok(())
475}
476
477fn path_is_symlink(path: &Path) -> bool {
478    fs::symlink_metadata(path)
479        .map(|metadata| metadata.file_type().is_symlink())
480        .unwrap_or(false)
481}
482
483fn remove_file_if_exists(path: &Path) -> io::Result<()> {
484    match fs::remove_file(path) {
485        Ok(()) => Ok(()),
486        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
487        Err(error) => Err(error),
488    }
489}
490
491#[cfg(unix)]
492fn create_symlink(target: &Path, link: &Path, target_is_dir: bool) -> io::Result<()> {
493    let _ = target_is_dir;
494    std::os::unix::fs::symlink(target, link)
495}
496
497#[cfg(windows)]
498fn create_symlink(target: &Path, link: &Path, target_is_dir: bool) -> io::Result<()> {
499    if target_is_dir {
500        std::os::windows::fs::symlink_dir(target, link)
501    } else {
502        std::os::windows::fs::symlink_file(target, link)
503    }
504}
505
506#[cfg(not(any(unix, windows)))]
507fn create_symlink(_target: &Path, _link: &Path, _target_is_dir: bool) -> io::Result<()> {
508    Err(io::Error::new(
509        io::ErrorKind::Unsupported,
510        "checkpoint symlink restore is unsupported on this platform",
511    ))
512}
513
514fn missing_parent_dirs(parent: &Path) -> Vec<PathBuf> {
515    let mut dirs = Vec::new();
516    let mut current = Some(parent);
517
518    while let Some(dir) = current {
519        if dir.as_os_str().is_empty() || dir.exists() {
520            break;
521        }
522        dirs.push(dir.to_path_buf());
523        current = dir.parent();
524    }
525
526    dirs
527}
528
529fn rollback_created_dirs(dirs: &[PathBuf]) -> bool {
530    let mut dirs = dirs.to_vec();
531    dirs.sort_by_key(|dir| std::cmp::Reverse(dir.components().count()));
532    dirs.dedup();
533
534    let mut ok = true;
535    for dir in dirs {
536        match std::fs::remove_dir(&dir) {
537            Ok(()) => {}
538            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
539            Err(_) => ok = false,
540        }
541    }
542    ok
543}
544
545fn current_timestamp() -> u64 {
546    std::time::SystemTime::now()
547        .duration_since(std::time::UNIX_EPOCH)
548        .unwrap_or_default()
549        .as_secs()
550}
551
552#[cfg(test)]
553mod tests {
554    use super::*;
555    use crate::protocol::DEFAULT_SESSION_ID;
556    use std::fs;
557
558    fn temp_file(name: &str, content: &str) -> PathBuf {
559        let dir = std::env::temp_dir().join("aft_checkpoint_tests");
560        fs::create_dir_all(&dir).unwrap();
561        let path = dir.join(name);
562        fs::write(&path, content).unwrap();
563        path
564    }
565
566    fn checkpoint_file(content: &str) -> CheckpointFile {
567        let file = tempfile::NamedTempFile::new().unwrap();
568        fs::write(file.path(), content).unwrap();
569        CheckpointFile::read(file.path()).unwrap()
570    }
571
572    #[test]
573    fn create_and_restore_round_trip() {
574        let path1 = temp_file("cp_rt1.txt", "hello");
575        let path2 = temp_file("cp_rt2.txt", "world");
576
577        let backup_store = BackupStore::new();
578        let mut store = CheckpointStore::new();
579
580        let info = store
581            .create(
582                DEFAULT_SESSION_ID,
583                "snap1",
584                vec![path1.clone(), path2.clone()],
585                &backup_store,
586            )
587            .unwrap();
588        assert_eq!(info.name, "snap1");
589        assert_eq!(info.file_count, 2);
590
591        // Modify files
592        fs::write(&path1, "changed1").unwrap();
593        fs::write(&path2, "changed2").unwrap();
594
595        // Restore
596        let info = store.restore(DEFAULT_SESSION_ID, "snap1").unwrap();
597        assert_eq!(info.file_count, 2);
598        assert_eq!(fs::read_to_string(&path1).unwrap(), "hello");
599        assert_eq!(fs::read_to_string(&path2).unwrap(), "world");
600    }
601
602    #[test]
603    fn overwrite_existing_name() {
604        let path = temp_file("cp_overwrite.txt", "v1");
605        let backup_store = BackupStore::new();
606        let mut store = CheckpointStore::new();
607
608        store
609            .create(DEFAULT_SESSION_ID, "dup", vec![path.clone()], &backup_store)
610            .unwrap();
611        fs::write(&path, "v2").unwrap();
612        store
613            .create(DEFAULT_SESSION_ID, "dup", vec![path.clone()], &backup_store)
614            .unwrap();
615
616        // Restore should give v2 (the overwritten checkpoint)
617        fs::write(&path, "v3").unwrap();
618        store.restore(DEFAULT_SESSION_ID, "dup").unwrap();
619        assert_eq!(fs::read_to_string(&path).unwrap(), "v2");
620    }
621
622    #[test]
623    fn list_returns_metadata_scoped_to_session() {
624        let path = temp_file("cp_list.txt", "data");
625        let backup_store = BackupStore::new();
626        let mut store = CheckpointStore::new();
627
628        store
629            .create(DEFAULT_SESSION_ID, "a", vec![path.clone()], &backup_store)
630            .unwrap();
631        store
632            .create(DEFAULT_SESSION_ID, "b", vec![path.clone()], &backup_store)
633            .unwrap();
634        store
635            .create("other_session", "c", vec![path.clone()], &backup_store)
636            .unwrap();
637
638        let default_list = store.list(DEFAULT_SESSION_ID);
639        assert_eq!(default_list.len(), 2);
640        let names: Vec<&str> = default_list.iter().map(|i| i.name.as_str()).collect();
641        assert!(names.contains(&"a"));
642        assert!(names.contains(&"b"));
643
644        let other_list = store.list("other_session");
645        assert_eq!(other_list.len(), 1);
646        assert_eq!(other_list[0].name, "c");
647    }
648
649    #[test]
650    fn sessions_isolate_checkpoint_names() {
651        // Same checkpoint name in two sessions does not collide on restore.
652        let path_a = temp_file("cp_isolated_a.txt", "a-original");
653        let path_b = temp_file("cp_isolated_b.txt", "b-original");
654        let backup_store = BackupStore::new();
655        let mut store = CheckpointStore::new();
656
657        // Both sessions create a checkpoint with the same name but different files.
658        store
659            .create("session_a", "snap", vec![path_a.clone()], &backup_store)
660            .unwrap();
661        store
662            .create("session_b", "snap", vec![path_b.clone()], &backup_store)
663            .unwrap();
664
665        fs::write(&path_a, "a-modified").unwrap();
666        fs::write(&path_b, "b-modified").unwrap();
667
668        // Restoring session A's "snap" only touches path_a.
669        store.restore("session_a", "snap").unwrap();
670        assert_eq!(fs::read_to_string(&path_a).unwrap(), "a-original");
671        assert_eq!(fs::read_to_string(&path_b).unwrap(), "b-modified");
672
673        // Restoring session B's "snap" only touches path_b.
674        fs::write(&path_a, "a-modified").unwrap();
675        store.restore("session_b", "snap").unwrap();
676        assert_eq!(fs::read_to_string(&path_a).unwrap(), "a-modified");
677        assert_eq!(fs::read_to_string(&path_b).unwrap(), "b-original");
678    }
679
680    #[test]
681    fn cleanup_removes_expired_across_sessions() {
682        let path = temp_file("cp_cleanup.txt", "data");
683        let backup_store = BackupStore::new();
684        let mut store = CheckpointStore::new();
685
686        store
687            .create(
688                DEFAULT_SESSION_ID,
689                "recent",
690                vec![path.clone()],
691                &backup_store,
692            )
693            .unwrap();
694
695        // Manually insert an expired checkpoint in another session.
696        store
697            .checkpoints
698            .entry("other".to_string())
699            .or_default()
700            .insert(
701                "old".to_string(),
702                Checkpoint {
703                    name: "old".to_string(),
704                    file_contents: HashMap::new(),
705                    created_at: 1000, // far in the past
706                },
707            );
708
709        assert_eq!(store.total_count(), 2);
710        store.cleanup(24); // 24 hours
711        assert_eq!(store.total_count(), 1);
712        assert_eq!(store.list(DEFAULT_SESSION_ID)[0].name, "recent");
713        assert!(store.list("other").is_empty());
714    }
715
716    #[test]
717    fn restore_nonexistent_returns_error() {
718        let store = CheckpointStore::new();
719        let result = store.restore(DEFAULT_SESSION_ID, "nope");
720        assert!(result.is_err());
721        match result.unwrap_err() {
722            AftError::CheckpointNotFound { name } => {
723                assert_eq!(name, "nope");
724            }
725            other => panic!("expected CheckpointNotFound, got: {:?}", other),
726        }
727    }
728
729    #[test]
730    fn restore_nonexistent_in_other_session_returns_error() {
731        // A "snap" that exists in session A must NOT be visible from session B.
732        let path = temp_file("cp_cross_session.txt", "data");
733        let backup_store = BackupStore::new();
734        let mut store = CheckpointStore::new();
735        store
736            .create("session_a", "only_a", vec![path], &backup_store)
737            .unwrap();
738        assert!(store.restore("session_b", "only_a").is_err());
739    }
740
741    #[test]
742    fn create_skips_missing_files_from_backup_tracked_set() {
743        // Simulate the reported issue #15-follow-up: an agent deletes a
744        // previously-edited file, then calls checkpoint with no explicit
745        // file list. Before the fix, the stale backup-tracked entry caused
746        // the whole checkpoint to fail on the missing path. Now the checkpoint
747        // succeeds with the readable file and reports the skipped one.
748        let readable = temp_file("cp_skip_readable.txt", "still_here");
749        let deleted = temp_file("cp_skip_deleted.txt", "about_to_vanish");
750
751        // Backup store canonicalizes keys, so the skipped path in the
752        // checkpoint result is the canonical form, not the raw temp path.
753        let deleted_canonical = fs::canonicalize(&deleted).unwrap();
754
755        let mut backup_store = BackupStore::new();
756        backup_store
757            .snapshot(DEFAULT_SESSION_ID, &readable, "auto")
758            .unwrap();
759        backup_store
760            .snapshot(DEFAULT_SESSION_ID, &deleted, "auto")
761            .unwrap();
762
763        fs::remove_file(&deleted).unwrap();
764
765        let mut store = CheckpointStore::new();
766        let info = store
767            .create(DEFAULT_SESSION_ID, "partial", vec![], &backup_store)
768            .expect("checkpoint should succeed despite one missing file");
769        assert_eq!(info.file_count, 1);
770        assert_eq!(info.skipped.len(), 1);
771        assert_eq!(info.skipped[0].0, deleted_canonical);
772        assert!(!info.skipped[0].1.is_empty());
773    }
774
775    #[test]
776    fn create_with_explicit_single_missing_file_errors() {
777        // When the caller names a single file explicitly and it can't be read,
778        // fail loudly — an empty checkpoint isn't what the caller asked for.
779        let missing = std::env::temp_dir()
780            .join("aft_checkpoint_tests/cp_explicit_missing_does_not_exist.txt");
781        let _ = fs::remove_file(&missing);
782
783        let backup_store = BackupStore::new();
784        let mut store = CheckpointStore::new();
785        let result = store.create(
786            DEFAULT_SESSION_ID,
787            "explicit",
788            vec![missing.clone()],
789            &backup_store,
790        );
791
792        assert!(result.is_err());
793        match result.unwrap_err() {
794            AftError::FileNotFound { path } => {
795                assert!(path.contains(&missing.display().to_string()));
796            }
797            other => panic!("expected FileNotFound, got: {:?}", other),
798        }
799    }
800
801    #[test]
802    fn create_with_explicit_mixed_files_keeps_readable_and_reports_skipped() {
803        // Explicit file list with one readable + one missing: keep the
804        // readable one in the checkpoint, report the missing one under
805        // `skipped` instead of failing outright.
806        let good = temp_file("cp_mixed_good.txt", "ok");
807        let missing = std::env::temp_dir().join("aft_checkpoint_tests/cp_mixed_missing.txt");
808        let _ = fs::remove_file(&missing);
809
810        let backup_store = BackupStore::new();
811        let mut store = CheckpointStore::new();
812        let info = store
813            .create(
814                DEFAULT_SESSION_ID,
815                "mixed",
816                vec![good.clone(), missing.clone()],
817                &backup_store,
818            )
819            .expect("mixed checkpoint should succeed when any file is readable");
820        assert_eq!(info.file_count, 1);
821        assert_eq!(info.skipped.len(), 1);
822        assert_eq!(info.skipped[0].0, missing);
823    }
824
825    #[test]
826    fn create_with_empty_files_uses_backup_tracked() {
827        let path = temp_file("cp_tracked.txt", "tracked_content");
828        let mut backup_store = BackupStore::new();
829        backup_store
830            .snapshot(DEFAULT_SESSION_ID, &path, "auto")
831            .unwrap();
832
833        let mut store = CheckpointStore::new();
834        let info = store
835            .create(DEFAULT_SESSION_ID, "from_tracked", vec![], &backup_store)
836            .unwrap();
837        assert!(info.file_count >= 1);
838
839        // Modify and restore
840        fs::write(&path, "modified").unwrap();
841        store.restore(DEFAULT_SESSION_ID, "from_tracked").unwrap();
842        assert_eq!(fs::read_to_string(&path).unwrap(), "tracked_content");
843    }
844
845    #[test]
846    fn restore_recreates_missing_parent_directories() {
847        let dir = tempfile::tempdir().unwrap();
848        let path = dir.path().join("nested").join("deeper").join("file.txt");
849        fs::create_dir_all(path.parent().unwrap()).unwrap();
850        fs::write(&path, "original nested content").unwrap();
851
852        let backup_store = BackupStore::new();
853        let mut store = CheckpointStore::new();
854        store
855            .create(
856                DEFAULT_SESSION_ID,
857                "nested",
858                vec![path.clone()],
859                &backup_store,
860            )
861            .unwrap();
862
863        fs::remove_dir_all(dir.path().join("nested")).unwrap();
864
865        store.restore(DEFAULT_SESSION_ID, "nested").unwrap();
866        assert_eq!(
867            fs::read_to_string(&path).unwrap(),
868            "original nested content"
869        );
870    }
871
872    #[cfg(unix)]
873    #[test]
874    fn checkpoint_restore_rolls_back_on_partial_failure() {
875        use std::os::unix::fs::PermissionsExt;
876
877        let dir = tempfile::tempdir().unwrap();
878        let path_a = dir.path().join("a.txt");
879        let path_b = dir.path().join("b.txt");
880        fs::write(&path_a, "checkpoint-a").unwrap();
881        fs::write(&path_b, "checkpoint-b").unwrap();
882
883        let backup_store = BackupStore::new();
884        let mut store = CheckpointStore::new();
885        store
886            .create(
887                DEFAULT_SESSION_ID,
888                "partial_failure",
889                vec![path_a.clone(), path_b.clone()],
890                &backup_store,
891            )
892            .unwrap();
893
894        fs::write(&path_a, "pre-restore-a").unwrap();
895        fs::write(&path_b, "pre-restore-b").unwrap();
896        let mut readonly = fs::metadata(&path_b).unwrap().permissions();
897        readonly.set_mode(0o444);
898        fs::set_permissions(&path_b, readonly).unwrap();
899
900        let result = store.restore(DEFAULT_SESSION_ID, "partial_failure");
901        let mut writable = fs::metadata(&path_b).unwrap().permissions();
902        writable.set_mode(0o644);
903        fs::set_permissions(&path_b, writable).unwrap();
904
905        assert!(result.is_err(), "restore should surface write failure");
906        assert_eq!(fs::read_to_string(&path_a).unwrap(), "pre-restore-a");
907        assert_eq!(fs::read_to_string(&path_b).unwrap(), "pre-restore-b");
908    }
909
910    #[test]
911    fn checkpoint_create_and_restore_use_mutation_lock() {
912        let dir = tempfile::tempdir().unwrap();
913        let lock_path = dir.path().join("locks").join("checkpoint.lock");
914        fs::create_dir_all(lock_path.parent().unwrap()).unwrap();
915        let mut store =
916            CheckpointStore::with_lock_path(lock_path.clone(), Duration::from_millis(50));
917        let backup_store = BackupStore::new();
918        let path = dir.path().join("locked.txt");
919        fs::write(&path, "original").unwrap();
920
921        let held_lock =
922            fs_lock::try_acquire(&lock_path, Duration::from_secs(1)).expect("hold checkpoint lock");
923        let create_result = store.create(
924            DEFAULT_SESSION_ID,
925            "locked",
926            vec![path.clone()],
927            &backup_store,
928        );
929        assert!(matches!(create_result, Err(AftError::IoError { .. })));
930        drop(held_lock);
931
932        store
933            .create(
934                DEFAULT_SESSION_ID,
935                "locked",
936                vec![path.clone()],
937                &backup_store,
938            )
939            .unwrap();
940        fs::write(&path, "changed").unwrap();
941
942        let held_lock =
943            fs_lock::try_acquire(&lock_path, Duration::from_secs(1)).expect("hold checkpoint lock");
944        let restore_result = store.restore(DEFAULT_SESSION_ID, "locked");
945        assert!(matches!(restore_result, Err(AftError::IoError { .. })));
946        drop(held_lock);
947
948        store.restore(DEFAULT_SESSION_ID, "locked").unwrap();
949        assert_eq!(fs::read_to_string(&path).unwrap(), "original");
950    }
951
952    #[cfg(unix)]
953    #[test]
954    fn checkpoint_restore_preserves_regular_file_permissions() {
955        use std::os::unix::fs::PermissionsExt;
956
957        let dir = tempfile::tempdir().unwrap();
958        let path = dir.path().join("mode.txt");
959        fs::write(&path, "original").unwrap();
960        let mut original_permissions = fs::metadata(&path).unwrap().permissions();
961        original_permissions.set_mode(0o600);
962        fs::set_permissions(&path, original_permissions).unwrap();
963
964        let backup_store = BackupStore::new();
965        let mut store = CheckpointStore::new();
966        store
967            .create(
968                DEFAULT_SESSION_ID,
969                "mode",
970                vec![path.clone()],
971                &backup_store,
972            )
973            .unwrap();
974
975        fs::write(&path, "changed").unwrap();
976        let mut changed_permissions = fs::metadata(&path).unwrap().permissions();
977        changed_permissions.set_mode(0o644);
978        fs::set_permissions(&path, changed_permissions).unwrap();
979
980        store.restore(DEFAULT_SESSION_ID, "mode").unwrap();
981
982        assert_eq!(fs::read_to_string(&path).unwrap(), "original");
983        let restored_mode = fs::metadata(&path).unwrap().permissions().mode() & 0o777;
984        assert_eq!(restored_mode, 0o600);
985    }
986
987    #[cfg(unix)]
988    #[test]
989    fn checkpoint_restore_recreates_symlink() {
990        let dir = tempfile::tempdir().unwrap();
991        let target = dir.path().join("target.txt");
992        let link = dir.path().join("link.txt");
993        fs::write(&target, "target content").unwrap();
994        std::os::unix::fs::symlink(&target, &link).unwrap();
995
996        let backup_store = BackupStore::new();
997        let mut store = CheckpointStore::new();
998        store
999            .create(
1000                DEFAULT_SESSION_ID,
1001                "symlink",
1002                vec![link.clone()],
1003                &backup_store,
1004            )
1005            .unwrap();
1006
1007        fs::remove_file(&link).unwrap();
1008        fs::write(&link, "plain file").unwrap();
1009
1010        store.restore(DEFAULT_SESSION_ID, "symlink").unwrap();
1011
1012        assert!(fs::symlink_metadata(&link)
1013            .unwrap()
1014            .file_type()
1015            .is_symlink());
1016        assert_eq!(fs::read_link(&link).unwrap(), target);
1017        assert_eq!(fs::read_to_string(&link).unwrap(), "target content");
1018    }
1019
1020    #[test]
1021    fn checkpoint_restore_failure_removes_created_parent_dirs() {
1022        let dir = tempfile::tempdir().unwrap();
1023        let missing_root = dir.path().join("created");
1024        let path_a = missing_root.join("nested").join("a.txt");
1025        let path_b = dir.path().join("blocking-dir");
1026        fs::create_dir(&path_b).unwrap();
1027
1028        let checkpoint = Checkpoint {
1029            name: "dir-cleanup".to_string(),
1030            file_contents: HashMap::from([
1031                (path_a.clone(), checkpoint_file("checkpoint-a")),
1032                (path_b.clone(), checkpoint_file("checkpoint-b")),
1033            ]),
1034            created_at: current_timestamp(),
1035        };
1036
1037        let result = restore_paths_atomically(&checkpoint, &[path_a.clone(), path_b.clone()]);
1038
1039        assert!(
1040            result.is_err(),
1041            "second restore write should fail on directory"
1042        );
1043        assert!(!path_a.exists(), "restored file should be rolled back");
1044        assert!(
1045            !missing_root.exists(),
1046            "new parent directories should be removed on rollback"
1047        );
1048        assert!(path_b.is_dir(), "pre-existing blocking directory remains");
1049    }
1050}