|
| 1 | +//! Pre-edit file backups for one-step undo. |
| 2 | +//! |
| 3 | +//! Before `edit_file` / `write_file` mutate a file, the prior content (or its absence, for a |
| 4 | +//! newly-created file) is saved here, so an edit can be rolled back via the `checkpoint` tool. |
| 5 | +//! **Only files the agent actually touches are backed up** — datasets/models in an ML workspace are |
| 6 | +//! never snapshotted — so this is safe regardless of workspace size, unlike a whole-tree snapshot. |
| 7 | +//! |
| 8 | +//! Opt-in via `checkpoint_enabled = true`; entirely inert otherwise (the global store is unset, so |
| 9 | +//! [`snapshot_before`] is a no-op and the `checkpoint` tool reports that it's disabled). |
| 10 | +
|
| 11 | +use async_trait::async_trait; |
| 12 | +use serde_json::{json, Value}; |
| 13 | +use std::path::{Path, PathBuf}; |
| 14 | +use std::sync::OnceLock; |
| 15 | + |
| 16 | +use crate::traits::Tool; |
| 17 | + |
| 18 | +static STORE: OnceLock<CheckpointStore> = OnceLock::new(); |
| 19 | + |
| 20 | +/// Cap on retained checkpoints; the oldest beyond this are pruned on each new snapshot so an |
| 21 | +/// always-on agent doing thousands of edits can't grow `.system_generated/checkpoints` without bound. |
| 22 | +const MAX_CHECKPOINTS: usize = 200; |
| 23 | + |
| 24 | +/// Initialise the process-wide checkpoint store. `root` holds the backups; `base`, when `Some`, |
| 25 | +/// confines restores to within that directory (the sandbox, when the file tools are workspace- |
| 26 | +/// restricted). Call once at startup when checkpointing is enabled. |
| 27 | +pub fn init(root: PathBuf, base: Option<PathBuf>) { |
| 28 | + let _ = STORE.set(CheckpointStore::new(root, base)); |
| 29 | +} |
| 30 | + |
| 31 | +/// Back up `path` before it is mutated. No-op when checkpointing is disabled. Best-effort: a backup |
| 32 | +/// failure is logged, never propagated, so it can't break the edit. |
| 33 | +pub fn snapshot_before(path: &Path, label: &str) { |
| 34 | + if let Some(store) = STORE.get() { |
| 35 | + if let Err(e) = store.snapshot(path, label) { |
| 36 | + log::warn!("checkpoint snapshot failed for {}: {}", path.display(), e); |
| 37 | + } |
| 38 | + } |
| 39 | +} |
| 40 | + |
| 41 | +/// The store, for the `checkpoint` tool. `None` when disabled. |
| 42 | +pub fn store() -> Option<&'static CheckpointStore> { |
| 43 | + STORE.get() |
| 44 | +} |
| 45 | + |
| 46 | +#[derive(serde::Serialize, serde::Deserialize)] |
| 47 | +struct Meta { |
| 48 | + /// Absolute original path that was (or would be) mutated. |
| 49 | + path: String, |
| 50 | + /// The tool that triggered the snapshot (e.g. `edit_file`). |
| 51 | + label: String, |
| 52 | + /// `false` when the file did not exist pre-edit — restoring such an entry removes the file. |
| 53 | + existed: bool, |
| 54 | + created_ms: u128, |
| 55 | +} |
| 56 | + |
| 57 | +/// One backed-up pre-edit state. |
| 58 | +pub struct CheckpointEntry { |
| 59 | + pub id: String, |
| 60 | + pub path: String, |
| 61 | + pub label: String, |
| 62 | + pub created_ms: u128, |
| 63 | + pub existed: bool, |
| 64 | +} |
| 65 | + |
| 66 | +pub struct CheckpointStore { |
| 67 | + root: PathBuf, |
| 68 | + /// When `Some`, restores are confined to within this directory (workspace-restricted mode). |
| 69 | + base: Option<PathBuf>, |
| 70 | +} |
| 71 | + |
| 72 | +impl CheckpointStore { |
| 73 | + pub fn new(root: PathBuf, base: Option<PathBuf>) -> Self { |
| 74 | + Self { root, base } |
| 75 | + } |
| 76 | + |
| 77 | + fn snapshot(&self, path: &Path, label: &str) -> std::io::Result<()> { |
| 78 | + let id = uuid::Uuid::new_v4().to_string(); |
| 79 | + let dir = self.root.join(&id); |
| 80 | + std::fs::create_dir_all(&dir)?; |
| 81 | + let existed = path.exists(); |
| 82 | + if existed { |
| 83 | + std::fs::copy(path, dir.join("content"))?; |
| 84 | + } |
| 85 | + let meta = Meta { |
| 86 | + path: path.to_string_lossy().into_owned(), |
| 87 | + label: label.to_string(), |
| 88 | + existed, |
| 89 | + created_ms: now_ms(), |
| 90 | + }; |
| 91 | + let bytes = serde_json::to_vec(&meta) |
| 92 | + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; |
| 93 | + std::fs::write(dir.join("meta.json"), bytes)?; |
| 94 | + self.prune(); |
| 95 | + Ok(()) |
| 96 | + } |
| 97 | + |
| 98 | + /// Bound disk growth: keep only the most recent [`MAX_CHECKPOINTS`], pruning the oldest. |
| 99 | + /// Best-effort — a prune failure never fails the snapshot. |
| 100 | + fn prune(&self) { |
| 101 | + let entries = self.list(); // newest first |
| 102 | + for e in entries.iter().skip(MAX_CHECKPOINTS) { |
| 103 | + let _ = std::fs::remove_dir_all(self.root.join(&e.id)); |
| 104 | + } |
| 105 | + } |
| 106 | + |
| 107 | + /// All checkpoints, newest first. |
| 108 | + pub fn list(&self) -> Vec<CheckpointEntry> { |
| 109 | + let mut entries = Vec::new(); |
| 110 | + let Ok(rd) = std::fs::read_dir(&self.root) else { |
| 111 | + return entries; |
| 112 | + }; |
| 113 | + for e in rd.flatten() { |
| 114 | + let Ok(bytes) = std::fs::read(e.path().join("meta.json")) else { |
| 115 | + continue; |
| 116 | + }; |
| 117 | + let Ok(m) = serde_json::from_slice::<Meta>(&bytes) else { |
| 118 | + continue; |
| 119 | + }; |
| 120 | + if let Some(id) = e.file_name().to_str().map(String::from) { |
| 121 | + entries.push(CheckpointEntry { |
| 122 | + id, |
| 123 | + path: m.path, |
| 124 | + label: m.label, |
| 125 | + created_ms: m.created_ms, |
| 126 | + existed: m.existed, |
| 127 | + }); |
| 128 | + } |
| 129 | + } |
| 130 | + entries.sort_by_key(|e| std::cmp::Reverse(e.created_ms)); |
| 131 | + entries |
| 132 | + } |
| 133 | + |
| 134 | + /// Restore the file recorded by checkpoint `id` to its pre-edit state. |
| 135 | + pub fn restore(&self, id: &str) -> Result<String, String> { |
| 136 | + // Confine the lookup to a direct child of the store. |
| 137 | + if id.is_empty() || id.contains('/') || id.contains('\\') || id.contains("..") { |
| 138 | + return Err("invalid checkpoint id".to_string()); |
| 139 | + } |
| 140 | + let dir = self.root.join(id); |
| 141 | + let meta_bytes = |
| 142 | + std::fs::read(dir.join("meta.json")).map_err(|e| format!("checkpoint {id} not found: {e}"))?; |
| 143 | + let m: Meta = |
| 144 | + serde_json::from_slice(&meta_bytes).map_err(|e| format!("checkpoint meta parse: {e}"))?; |
| 145 | + let target = PathBuf::from(&m.path); |
| 146 | + if m.existed { |
| 147 | + // SECURITY: re-validate the write target against the LIVE filesystem, not just the |
| 148 | + // recorded string. The agent could have swapped the path (or a parent) for a symlink |
| 149 | + // after the snapshot, and `fs::copy` follows symlinks — a lexical check on the unchanged |
| 150 | + // string would let that redirect the write outside the sandbox (TOCTOU). So in restricted |
| 151 | + // mode we refuse a symlinked final component and re-resolve through the same boundary the |
| 152 | + // file tools use (which canonicalizes, catching a symlinked parent that resolves out). |
| 153 | + let dest = self.safe_write_target(&target, &m.path)?; |
| 154 | + std::fs::copy(dir.join("content"), &dest).map_err(|e| format!("restore copy: {e}"))?; |
| 155 | + Ok(format!("Restored {} from checkpoint {}.", m.path, id)) |
| 156 | + } else { |
| 157 | + // The snapshotted edit created the file; undo = remove it. `remove_file` unlinks a |
| 158 | + // symlink itself (not its target), so a lexical containment check is sufficient here. |
| 159 | + if let Some(base) = &self.base { |
| 160 | + if !target.starts_with(base) { |
| 161 | + return Err( |
| 162 | + "checkpoint target is outside the workspace; refusing to restore".to_string(), |
| 163 | + ); |
| 164 | + } |
| 165 | + } |
| 166 | + match std::fs::remove_file(&target) { |
| 167 | + Ok(()) => Ok(format!( |
| 168 | + "Removed {} (it was created after checkpoint {}).", |
| 169 | + m.path, id |
| 170 | + )), |
| 171 | + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { |
| 172 | + Ok(format!("{} is already absent.", m.path)) |
| 173 | + } |
| 174 | + Err(e) => Err(format!("restore remove: {e}")), |
| 175 | + } |
| 176 | + } |
| 177 | + } |
| 178 | + |
| 179 | + /// Resolve a safe destination for a restore *write*, closing the symlink/TOCTOU escape. In |
| 180 | + /// unrestricted mode (`base == None`) edits already go anywhere, so the recorded path is used |
| 181 | + /// as-is. In restricted mode the final component must not be a symlink, and the path is |
| 182 | + /// re-resolved through `resolve_path` (the same canonicalizing boundary edits use). |
| 183 | + fn safe_write_target(&self, target: &Path, raw: &str) -> Result<PathBuf, String> { |
| 184 | + let Some(base) = &self.base else { |
| 185 | + return Ok(target.to_path_buf()); |
| 186 | + }; |
| 187 | + if std::fs::symlink_metadata(target) |
| 188 | + .map(|m| m.file_type().is_symlink()) |
| 189 | + .unwrap_or(false) |
| 190 | + { |
| 191 | + return Err("checkpoint target is now a symlink; refusing to restore".to_string()); |
| 192 | + } |
| 193 | + crate::tools::builtin::resolve_path(raw, base, true) |
| 194 | + .map_err(|e| format!("checkpoint target rejected: {e}")) |
| 195 | + } |
| 196 | +} |
| 197 | + |
| 198 | +fn now_ms() -> u128 { |
| 199 | + std::time::SystemTime::now() |
| 200 | + .duration_since(std::time::UNIX_EPOCH) |
| 201 | + .map(|d| d.as_millis()) |
| 202 | + .unwrap_or(0) |
| 203 | +} |
| 204 | + |
| 205 | +/// `checkpoint` tool: list pre-edit backups and restore one (one-step undo for `edit_file` / |
| 206 | +/// `write_file`). Registered only when checkpointing is enabled. |
| 207 | +pub struct CheckpointTool; |
| 208 | + |
| 209 | +#[async_trait] |
| 210 | +impl Tool for CheckpointTool { |
| 211 | + fn name(&self) -> &str { |
| 212 | + "checkpoint" |
| 213 | + } |
| 214 | + |
| 215 | + fn description(&self) -> &str { |
| 216 | + "List or restore pre-edit file checkpoints — a one-step undo for edit_file/write_file. \ |
| 217 | + action 'list' shows recent checkpoints (newest first); action 'restore' with an 'id' rolls \ |
| 218 | + that file back to its state before the edit (removing it if the edit had created it)." |
| 219 | + } |
| 220 | + |
| 221 | + fn parameters(&self) -> Value { |
| 222 | + json!({ |
| 223 | + "type": "object", |
| 224 | + "properties": { |
| 225 | + "action": { "type": "string", "enum": ["list", "restore"], "description": "list (default) or restore" }, |
| 226 | + "id": { "type": "string", "description": "checkpoint id to restore (required when action=restore)" } |
| 227 | + } |
| 228 | + }) |
| 229 | + } |
| 230 | + |
| 231 | + async fn execute(&self, args: Value) -> Result<String, String> { |
| 232 | + let Some(store) = store() else { |
| 233 | + return Ok("Checkpointing is disabled (set checkpoint_enabled = true).".to_string()); |
| 234 | + }; |
| 235 | + let action = args.get("action").and_then(|v| v.as_str()).unwrap_or("list"); |
| 236 | + match action { |
| 237 | + "list" => { |
| 238 | + let entries = store.list(); |
| 239 | + if entries.is_empty() { |
| 240 | + return Ok("No checkpoints.".to_string()); |
| 241 | + } |
| 242 | + let mut out = String::from("Checkpoints (newest first):\n"); |
| 243 | + for e in entries.iter().take(50) { |
| 244 | + out.push_str(&format!( |
| 245 | + "- {} [{}] {}{}\n", |
| 246 | + e.id, |
| 247 | + e.label, |
| 248 | + e.path, |
| 249 | + if e.existed { "" } else { " (created)" } |
| 250 | + )); |
| 251 | + } |
| 252 | + Ok(out) |
| 253 | + } |
| 254 | + "restore" => { |
| 255 | + let id = args |
| 256 | + .get("id") |
| 257 | + .and_then(|v| v.as_str()) |
| 258 | + .ok_or("restore requires 'id'")?; |
| 259 | + store.restore(id) |
| 260 | + } |
| 261 | + other => Err(format!("unknown action '{other}' (use list or restore)")), |
| 262 | + } |
| 263 | + } |
| 264 | +} |
| 265 | + |
| 266 | +#[cfg(test)] |
| 267 | +mod tests { |
| 268 | + use super::*; |
| 269 | + |
| 270 | + fn temp() -> (PathBuf, CheckpointStore) { |
| 271 | + let base = std::env::temp_dir().join(format!("isan_ckpt_{}", uuid::Uuid::new_v4())); |
| 272 | + std::fs::create_dir_all(&base).unwrap(); |
| 273 | + let store = CheckpointStore::new(base.join(".checkpoints"), Some(base.clone())); |
| 274 | + (base, store) |
| 275 | + } |
| 276 | + |
| 277 | + #[test] |
| 278 | + fn snapshot_then_restore_recovers_prior_content() { |
| 279 | + let (base, store) = temp(); |
| 280 | + let file = base.join("a.txt"); |
| 281 | + std::fs::write(&file, "v1").unwrap(); |
| 282 | + |
| 283 | + store.snapshot(&file, "edit_file").unwrap(); |
| 284 | + std::fs::write(&file, "v2-broken").unwrap(); // simulate a bad edit |
| 285 | + |
| 286 | + let entries = store.list(); |
| 287 | + assert_eq!(entries.len(), 1); |
| 288 | + assert_eq!(entries[0].label, "edit_file"); |
| 289 | + assert!(entries[0].existed); |
| 290 | + |
| 291 | + store.restore(&entries[0].id).unwrap(); |
| 292 | + assert_eq!(std::fs::read_to_string(&file).unwrap(), "v1"); |
| 293 | + let _ = std::fs::remove_dir_all(&base); |
| 294 | + } |
| 295 | + |
| 296 | + #[test] |
| 297 | + fn restoring_a_created_file_removes_it() { |
| 298 | + let (base, store) = temp(); |
| 299 | + let file = base.join("new.txt"); |
| 300 | + // File does not exist yet -> snapshot records a creation. |
| 301 | + store.snapshot(&file, "write_file").unwrap(); |
| 302 | + std::fs::write(&file, "created").unwrap(); |
| 303 | + |
| 304 | + let id = store.list()[0].id.clone(); |
| 305 | + assert!(!store.list()[0].existed); |
| 306 | + store.restore(&id).unwrap(); |
| 307 | + assert!(!file.exists(), "restoring a creation should remove the file"); |
| 308 | + let _ = std::fs::remove_dir_all(&base); |
| 309 | + } |
| 310 | + |
| 311 | + #[test] |
| 312 | + fn restore_rejects_bad_ids() { |
| 313 | + let (base, store) = temp(); |
| 314 | + assert!(store.restore("../etc").is_err()); |
| 315 | + assert!(store.restore("a/b").is_err()); |
| 316 | + assert!(store.restore("missing-id").is_err()); |
| 317 | + let _ = std::fs::remove_dir_all(&base); |
| 318 | + } |
| 319 | + |
| 320 | + #[test] |
| 321 | + fn restore_refuses_meta_path_outside_base() { |
| 322 | + let (base, store) = temp(); |
| 323 | + // Craft a checkpoint whose meta.path points OUTSIDE the base (a tampered/forged meta). |
| 324 | + let outside = |
| 325 | + std::env::temp_dir().join(format!("isan_outside_{}.txt", uuid::Uuid::new_v4())); |
| 326 | + let dir = store.root.join("crafted"); |
| 327 | + std::fs::create_dir_all(&dir).unwrap(); |
| 328 | + std::fs::write(dir.join("content"), b"payload").unwrap(); |
| 329 | + let meta = format!( |
| 330 | + r#"{{"path":{:?},"label":"edit_file","existed":true,"created_ms":1}}"#, |
| 331 | + outside.to_string_lossy() |
| 332 | + ); |
| 333 | + std::fs::write(dir.join("meta.json"), meta).unwrap(); |
| 334 | + |
| 335 | + assert!(store.restore("crafted").is_err(), "must refuse out-of-base meta"); |
| 336 | + assert!(!outside.exists(), "must not write the payload outside base"); |
| 337 | + let _ = std::fs::remove_dir_all(&base); |
| 338 | + } |
| 339 | + |
| 340 | + #[cfg(unix)] |
| 341 | + #[test] |
| 342 | + fn restore_refuses_symlinked_target() { |
| 343 | + let (base, store) = temp(); |
| 344 | + let file = base.join("real.txt"); |
| 345 | + std::fs::write(&file, "v1").unwrap(); |
| 346 | + store.snapshot(&file, "edit_file").unwrap(); |
| 347 | + std::fs::write(&file, "v2").unwrap(); |
| 348 | + let id = store.list()[0].id.clone(); |
| 349 | + |
| 350 | + // Agent swaps the target for a symlink pointing outside the base (TOCTOU). |
| 351 | + let outside = std::env::temp_dir().join(format!("isan_symout_{}", uuid::Uuid::new_v4())); |
| 352 | + std::fs::remove_file(&file).unwrap(); |
| 353 | + std::os::unix::fs::symlink(&outside, &file).unwrap(); |
| 354 | + |
| 355 | + let res = store.restore(&id); |
| 356 | + assert!(res.is_err(), "must refuse a symlinked target: {res:?}"); |
| 357 | + assert!(!outside.exists(), "must not write through the symlink"); |
| 358 | + let _ = std::fs::remove_dir_all(&base); |
| 359 | + } |
| 360 | + |
| 361 | + #[test] |
| 362 | + fn list_is_newest_first() { |
| 363 | + let (base, store) = temp(); |
| 364 | + let f = base.join("x"); |
| 365 | + std::fs::write(&f, "1").unwrap(); |
| 366 | + store.snapshot(&f, "edit_file").unwrap(); |
| 367 | + std::thread::sleep(std::time::Duration::from_millis(2)); |
| 368 | + store.snapshot(&f, "write_file").unwrap(); |
| 369 | + let entries = store.list(); |
| 370 | + assert_eq!(entries.len(), 2); |
| 371 | + assert!(entries[0].created_ms >= entries[1].created_ms); |
| 372 | + let _ = std::fs::remove_dir_all(&base); |
| 373 | + } |
| 374 | +} |
0 commit comments