Skip to main content

agent_sandbox/fs/
overlay.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use sha2::{Digest, Sha256};
5
6use crate::error::Result;
7
8/// Tracks filesystem changes by comparing against initial snapshots.
9#[derive(Debug, Clone, PartialEq)]
10pub enum FsChangeKind {
11    Created,
12    Modified,
13    Deleted,
14}
15
16#[derive(Debug, Clone)]
17pub struct FsChange {
18    pub path: String,
19    pub kind: FsChangeKind,
20}
21
22/// Filesystem overlay that tracks changes to the work directory.
23pub struct FsOverlay {
24    root: PathBuf,
25    /// SHA-256 hashes of files at snapshot time.
26    snapshot: HashMap<PathBuf, Vec<u8>>,
27}
28
29impl FsOverlay {
30    /// Create a new overlay and snapshot the current state of the root directory.
31    pub fn new(root: &Path) -> Result<Self> {
32        let root = root.canonicalize()?;
33        let mut snapshot = HashMap::new();
34        snapshot_dir(&root, &mut snapshot)?;
35
36        Ok(Self { root, snapshot })
37    }
38
39    /// Compare the current state against the snapshot and return changes.
40    pub fn diff(&self) -> Result<Vec<FsChange>> {
41        let mut changes = Vec::new();
42        let mut current_files = HashMap::new();
43
44        // Walk current state
45        snapshot_dir(&self.root, &mut current_files)?;
46
47        // Find created and modified files
48        for (path, hash) in &current_files {
49            let rel = path
50                .strip_prefix(&self.root)
51                .unwrap_or(path)
52                .to_string_lossy()
53                .to_string();
54
55            match self.snapshot.get(path) {
56                None => {
57                    changes.push(FsChange {
58                        path: rel,
59                        kind: FsChangeKind::Created,
60                    });
61                }
62                Some(old_hash) if old_hash != hash => {
63                    changes.push(FsChange {
64                        path: rel,
65                        kind: FsChangeKind::Modified,
66                    });
67                }
68                _ => {}
69            }
70        }
71
72        // Find deleted files
73        for path in self.snapshot.keys() {
74            if !current_files.contains_key(path) {
75                let rel = path
76                    .strip_prefix(&self.root)
77                    .unwrap_or(path)
78                    .to_string_lossy()
79                    .to_string();
80                changes.push(FsChange {
81                    path: rel,
82                    kind: FsChangeKind::Deleted,
83                });
84            }
85        }
86
87        changes.sort_by(|a, b| a.path.cmp(&b.path));
88        Ok(changes)
89    }
90}
91
92/// Maximum file size to snapshot (50 MB). Larger files are skipped to prevent OOM.
93const MAX_SNAPSHOT_FILE_SIZE: u64 = 50 * 1024 * 1024;
94
95fn snapshot_dir(dir: &Path, snapshot: &mut HashMap<PathBuf, Vec<u8>>) -> Result<()> {
96    if !dir.is_dir() {
97        return Ok(());
98    }
99
100    for entry in std::fs::read_dir(dir)? {
101        let entry = entry?;
102        let path = entry.path();
103
104        if path.is_dir() {
105            snapshot_dir(&path, snapshot)?;
106        } else if path.is_file() {
107            let metadata = std::fs::metadata(&path)?;
108            if metadata.len() > MAX_SNAPSHOT_FILE_SIZE {
109                continue;
110            }
111            let content = std::fs::read(&path)?;
112            let hash = Sha256::digest(&content).to_vec();
113            snapshot.insert(path, hash);
114        }
115    }
116
117    Ok(())
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn test_detect_created_file() {
126        let tmp = tempfile::tempdir().unwrap();
127        let root = tmp.path();
128
129        std::fs::write(root.join("existing.txt"), "hello").unwrap();
130
131        let overlay = FsOverlay::new(root).unwrap();
132
133        // Create a new file
134        std::fs::write(root.join("new.txt"), "world").unwrap();
135
136        let changes = overlay.diff().unwrap();
137        assert_eq!(changes.len(), 1);
138        assert_eq!(changes[0].path, "new.txt");
139        assert_eq!(changes[0].kind, FsChangeKind::Created);
140    }
141
142    #[test]
143    fn test_detect_modified_file() {
144        let tmp = tempfile::tempdir().unwrap();
145        let root = tmp.path();
146
147        std::fs::write(root.join("file.txt"), "original").unwrap();
148
149        let overlay = FsOverlay::new(root).unwrap();
150
151        // Modify the file
152        std::fs::write(root.join("file.txt"), "modified").unwrap();
153
154        let changes = overlay.diff().unwrap();
155        assert_eq!(changes.len(), 1);
156        assert_eq!(changes[0].path, "file.txt");
157        assert_eq!(changes[0].kind, FsChangeKind::Modified);
158    }
159
160    #[test]
161    fn test_detect_deleted_file() {
162        let tmp = tempfile::tempdir().unwrap();
163        let root = tmp.path();
164
165        std::fs::write(root.join("file.txt"), "content").unwrap();
166
167        let overlay = FsOverlay::new(root).unwrap();
168
169        // Delete the file
170        std::fs::remove_file(root.join("file.txt")).unwrap();
171
172        let changes = overlay.diff().unwrap();
173        assert_eq!(changes.len(), 1);
174        assert_eq!(changes[0].path, "file.txt");
175        assert_eq!(changes[0].kind, FsChangeKind::Deleted);
176    }
177
178    #[test]
179    fn test_no_changes() {
180        let tmp = tempfile::tempdir().unwrap();
181        let root = tmp.path();
182
183        std::fs::write(root.join("file.txt"), "content").unwrap();
184
185        let overlay = FsOverlay::new(root).unwrap();
186        let changes = overlay.diff().unwrap();
187        assert!(changes.is_empty());
188    }
189}