Skip to main content

aft/
checkpoint.rs

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