Skip to content

Commit b01e501

Browse files
committed
feat(checkpoint): opt-in pre-edit file backups for one-step undo
The agent edited files unattended with no rollback: a bad multi-edit had to be reconstructed by hand. Adds an opt-in checkpoint layer so an edit can be undone. Design — back up only the files the agent *touches*, not the whole tree. Before `edit_file`/`write_file` mutates a file, its prior content (or its absence, for a newly-created file) is saved under `.system_generated/checkpoints/<id>/`. Datasets and model weights in an ML workspace are never snapshotted, so this is safe at any workspace size — the reason a whole-tree shadow-git snapshot is impractical here. - `src/checkpoint.rs`: `CheckpointStore` (snapshot / list / restore) behind a process-wide `OnceLock`, like the other startup-config singletons. `snapshot_before` is a no-op when disabled and best-effort (a backup failure is logged, never breaks the edit). Restoring a creation removes the file; restoring an edit rewrites the prior bytes. - **Restore is symlink/TOCTOU-safe**: in restricted mode it refuses a target whose final component became a symlink after the snapshot and re-resolves the write through the file tools' own `resolve_path` (canonicalizing, so a swapped-in symlinked parent that resolves outside the sandbox is rejected) — closing an arbitrary-write-outside-sandbox vector that a lexical-only check would miss. - Retention: keep the most recent `MAX_CHECKPOINTS` (200), pruning the oldest on each snapshot, so disk can't grow without bound. - `checkpoint` tool (list/restore), registered only when enabled. - `edit_file` / `write_file` call `checkpoint::snapshot_before` before writing. - Config: `checkpoint_enabled` (default false); `main` initialises the store (restore confined to the sandbox in restricted mode) and registers the tool. Tests: snapshot→edit→restore recovers content; restoring a created file removes it; restore refuses a meta.path outside base and a symlinked target (no out-of-sandbox write); bad ids rejected; list newest-first.
1 parent e99e8b3 commit b01e501

5 files changed

Lines changed: 394 additions & 0 deletions

File tree

src/checkpoint.rs

Lines changed: 374 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,374 @@
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+
}

src/config.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,9 @@ pub struct AppConfig {
324324
pub max_iterations: Option<usize>,
325325
/// When true (default), detect repeated identical tool calls and inject a corrective user message.
326326
pub doom_loop_enabled: Option<bool>,
327+
/// When true, back up each file before `edit_file`/`write_file` mutates it and register the
328+
/// `checkpoint` tool for one-step undo. Default false. Only touched files are backed up.
329+
pub checkpoint_enabled: Option<bool>,
327330
pub max_tool_output_chars: Option<usize>,
328331
/// Max characters returned by `web_search` / `web_fetch` (default 50_000). Separate from
329332
/// `max_tool_output_chars`, which caps tool output when passed to the model.
@@ -553,6 +556,11 @@ impl AppConfig {
553556
self.doom_loop_enabled.unwrap_or(true)
554557
}
555558

559+
/// Pre-edit file checkpointing for one-step undo (default: disabled).
560+
pub fn checkpoint_enabled(&self) -> bool {
561+
self.checkpoint_enabled.unwrap_or(false)
562+
}
563+
556564
/// When true, `git_worktree` is registered (see `[harness.git_worktree]` in config).
557565
pub fn git_worktree_tool_enabled(&self) -> bool {
558566
self.harness

0 commit comments

Comments
 (0)