Skip to main content

recursive/tools/
checkpoint.rs

1//! Read-only and save checkpoint tools that surface the session's
2//! checkpoint chain.
3//!
4//! `checkpoint_list` and `checkpoint_diff` are read-only observer tools.
5//! `checkpoint_save` (Goal 284) lets the agent create explicit restore
6//! points on demand — the only path that creates new shadow-git objects
7//! during a run.
8//!
9//! Restoration remains runtime/CLI-driven (see `recursive sessions rewind`).
10
11use async_trait::async_trait;
12use serde_json::{json, Value};
13use std::path::PathBuf;
14use std::sync::{Arc, Mutex};
15
16use super::Tool;
17use crate::checkpoint::{CheckpointId, ShadowRepo};
18use crate::checkpoint_log::{read_log, CheckpointLogWriter, CheckpointRecord, TouchedVia};
19use crate::error::{Error, Result};
20use crate::llm::ToolSpec;
21use crate::tools::TouchedFiles;
22
23/// Shared state for the read-only checkpoint tools.
24///
25/// `session_id` is fixed at registration time so the agent can only
26/// inspect its own checkpoint chain — never another session's.
27#[derive(Clone)]
28pub struct CheckpointToolCtx {
29    pub repo: Arc<Mutex<ShadowRepo>>,
30    pub session_id: String,
31}
32
33// ── checkpoint_list ───────────────────────────────────────────────────────────
34
35pub struct CheckpointList {
36    ctx: CheckpointToolCtx,
37}
38
39impl CheckpointList {
40    pub fn new(ctx: CheckpointToolCtx) -> Self {
41        Self { ctx }
42    }
43}
44
45#[async_trait]
46impl Tool for CheckpointList {
47    fn is_deferred(&self) -> bool {
48        true
49    }
50
51    fn spec(&self) -> ToolSpec {
52        ToolSpec {
53            name: "checkpoint_list".into(),
54            description: "List checkpoints for the current session, newest first. \
55                 Checkpoints are saved by the agent via `checkpoint_save` or \
56                 by the runtime at session boundaries. Use `checkpoint_diff` \
57                 to inspect the changes captured in a checkpoint."
58                .into(),
59            parameters: json!({"type": "object", "properties": {}}),
60        }
61    }
62
63    fn side_effect_class(&self) -> crate::tools::ToolSideEffect {
64        crate::tools::ToolSideEffect::ReadOnly
65    }
66
67    async fn execute(&self, _args: Value) -> Result<String> {
68        let infos = {
69            let repo = self.ctx.repo.lock().map_err(|_| Error::Tool {
70                name: "checkpoint_list".into(),
71                call_id: None,
72                message: "lock poisoned".into(),
73            })?;
74            repo.list_for_session(&self.ctx.session_id)?
75        };
76        if infos.is_empty() {
77            return Ok("no checkpoints in this session yet".to_string());
78        }
79        let mut out = String::new();
80        for i in &infos {
81            out.push_str(&format!(
82                "[{}]  {}  ({} file(s) changed)  ts={}\n",
83                i.id, i.message, i.files_changed, i.timestamp
84            ));
85        }
86        Ok(out.trim_end().to_string())
87    }
88}
89
90// ── checkpoint_diff ───────────────────────────────────────────────────────────
91
92pub struct CheckpointDiff {
93    ctx: CheckpointToolCtx,
94}
95
96impl CheckpointDiff {
97    pub fn new(ctx: CheckpointToolCtx) -> Self {
98        Self { ctx }
99    }
100}
101
102#[async_trait]
103impl Tool for CheckpointDiff {
104    fn is_deferred(&self) -> bool {
105        true
106    }
107
108    fn spec(&self) -> ToolSpec {
109        ToolSpec {
110            name: "checkpoint_diff".into(),
111            description: "Show a unified diff between two checkpoints in this session. \
112                 If `b` is omitted, diffs `a` against the current workspace state."
113                .into(),
114            parameters: json!({
115                "type": "object",
116                "properties": {
117                    "a": {"type": "string", "description": "Older checkpoint id."},
118                    "b": {"type": "string", "description": "Optional newer checkpoint id; omit for current workspace."}
119                },
120                "required": ["a"]
121            }),
122        }
123    }
124
125    fn side_effect_class(&self) -> crate::tools::ToolSideEffect {
126        crate::tools::ToolSideEffect::ReadOnly
127    }
128
129    async fn execute(&self, args: Value) -> Result<String> {
130        let a = args["a"].as_str().ok_or_else(|| Error::BadToolArgs {
131            name: "checkpoint_diff".into(),
132            message: "missing `a`".into(),
133        })?;
134        let b = args["b"].as_str();
135        let aid = CheckpointId(a.to_string());
136        let bid = b.map(|s| CheckpointId(s.to_string()));
137        let diff = {
138            let repo = self.ctx.repo.lock().map_err(|_| Error::Tool {
139                name: "checkpoint_diff".into(),
140                call_id: None,
141                message: "lock poisoned".into(),
142            })?;
143            repo.diff(&aid, bid.as_ref(), &[])?
144        };
145        if diff.trim().is_empty() {
146            Ok("no differences".to_string())
147        } else {
148            Ok(diff)
149        }
150    }
151}
152
153// ── checkpoint_save (Goal 284) ────────────────────────────────────────────────
154
155/// Context for the on-demand `checkpoint_save` tool.
156///
157/// Unlike the read-only tools, this holds the touched-files collector,
158/// log writer, turn index, and log path so it can snapshot the workspace,
159/// attribute file changes, and append a [`CheckpointRecord`].
160#[derive(Clone)]
161pub struct CheckpointSaveCtx {
162    pub repo: Arc<Mutex<ShadowRepo>>,
163    pub session_id: String,
164    pub touched_files: Option<Arc<Mutex<TouchedFiles>>>,
165    pub writer: Arc<Mutex<CheckpointLogWriter>>,
166    pub turn_index: Arc<std::sync::atomic::AtomicUsize>,
167    pub log_path: PathBuf,
168}
169
170pub struct CheckpointSave {
171    ctx: CheckpointSaveCtx,
172}
173
174impl CheckpointSave {
175    pub fn new(ctx: CheckpointSaveCtx) -> Self {
176        Self { ctx }
177    }
178}
179
180#[async_trait]
181impl Tool for CheckpointSave {
182    fn is_deferred(&self) -> bool {
183        true
184    }
185
186    fn spec(&self) -> ToolSpec {
187        ToolSpec {
188            name: "checkpoint_save".into(),
189            description: "Save an explicit restore point for the current session. \
190                 Call this before making a risky batch of changes, or after \
191                 completing a logical unit of work you might want to revert \
192                 to. Unlike automatic checkpoints (which no longer exist), \
193                 this runs only when you call it."
194                .into(),
195            parameters: json!({
196                "type": "object",
197                "properties": {
198                    "message": {
199                        "type": "string",
200                        "description": "A short label for this checkpoint, e.g. 'before refactor' or 'after adding tests'. Defaults to the current turn number if omitted."
201                    }
202                }
203            }),
204        }
205    }
206
207    fn side_effect_class(&self) -> crate::tools::ToolSideEffect {
208        crate::tools::ToolSideEffect::ReadOnly
209    }
210
211    async fn execute(&self, args: Value) -> Result<String> {
212        let turn = self
213            .ctx
214            .turn_index
215            .load(std::sync::atomic::Ordering::Relaxed);
216        let message = args
217            .get("message")
218            .and_then(|v| v.as_str())
219            .map(|s| s.to_string())
220            .unwrap_or_else(|| format!("turn {turn}"));
221
222        // Collect touched files before snapshotting.
223        let (mut paths, saw_shell) = match &self.ctx.touched_files {
224            Some(slot) => match slot.lock() {
225                Ok(t) => (t.paths_sorted(), t.saw_shell),
226                Err(_) => (vec![], false),
227            },
228            None => (vec![], true),
229        };
230
231        // Find the last checkpoint id for shell-diff attribution.
232        let last_id: Option<CheckpointId> = {
233            match read_log(&self.ctx.log_path) {
234                Ok(recs) => recs.last().map(|r| r.id.clone()),
235                Err(_) => None,
236            }
237        };
238
239        // Snapshot the workspace.
240        let id = {
241            let repo = self.ctx.repo.lock().map_err(|_| Error::Tool {
242                name: "checkpoint_save".into(),
243                call_id: None,
244                message: "lock poisoned".into(),
245            })?;
246            repo.snapshot_for_session(&self.ctx.session_id, &message)?
247        };
248
249        // If the agent ran a shell command this turn, diff against the
250        // last checkpoint to capture files created/modified.
251        let mut via = TouchedVia::Structured;
252        if saw_shell {
253            via = TouchedVia::ShellDiff;
254            if let Some(ref last) = last_id {
255                let repo = self.ctx.repo.lock().map_err(|_| Error::Tool {
256                    name: "checkpoint_save".into(),
257                    call_id: None,
258                    message: "lock poisoned".into(),
259                })?;
260                if let Ok(diff_paths) = repo.changed_paths(last, &id) {
261                    let mut set: std::collections::HashSet<String> = paths.drain(..).collect();
262                    for p in diff_paths {
263                        set.insert(p);
264                    }
265                    paths = {
266                        let mut v: Vec<String> = set.into_iter().collect();
267                        v.sort();
268                        v
269                    };
270                }
271            }
272        }
273
274        // Write the log record.
275        {
276            let writer = self.ctx.writer.lock().map_err(|_| Error::Tool {
277                name: "checkpoint_save".into(),
278                call_id: None,
279                message: "writer lock poisoned".into(),
280            })?;
281            let rec = CheckpointRecord {
282                turn,
283                pre: last_id,
284                id: id.clone(),
285                message: Some(message),
286                touched_files: paths,
287                touched_via: via,
288                started_at: 0,
289                finished_at: 0,
290                saved_at: unix_now(),
291            };
292            writer.append(&rec).map_err(|e| Error::Tool {
293                name: "checkpoint_save".into(),
294                call_id: None,
295                message: format!("failed to append checkpoint record: {e}"),
296            })?;
297        }
298
299        Ok(id.0)
300    }
301}
302
303/// Build a [`CheckpointSave`] tool wired to the given shared state.
304pub fn build_checkpoint_save_tool(
305    repo: Arc<Mutex<ShadowRepo>>,
306    session_id: String,
307    touched_files: Option<Arc<Mutex<TouchedFiles>>>,
308    writer: Arc<Mutex<CheckpointLogWriter>>,
309    turn_index: Arc<std::sync::atomic::AtomicUsize>,
310    log_path: PathBuf,
311) -> CheckpointSave {
312    CheckpointSave::new(CheckpointSaveCtx {
313        repo,
314        session_id,
315        touched_files,
316        writer,
317        turn_index,
318        log_path,
319    })
320}
321
322/// Current Unix timestamp in seconds.
323fn unix_now() -> i64 {
324    std::time::SystemTime::now()
325        .duration_since(std::time::UNIX_EPOCH)
326        .map(|d| d.as_secs() as i64)
327        .unwrap_or(0)
328}
329
330// ── helper ────────────────────────────────────────────────────────────────────
331
332/// Build the read-only checkpoint tools for a given session.
333/// Returns `None` if the shadow repo cannot be opened (e.g. git missing).
334pub fn build_checkpoint_tools(
335    workspace: impl Into<std::path::PathBuf>,
336    session_id: impl Into<String>,
337) -> Option<(CheckpointList, CheckpointDiff, Arc<Mutex<ShadowRepo>>)> {
338    let repo = match ShadowRepo::open(workspace) {
339        Ok(r) => Arc::new(Mutex::new(r)),
340        Err(e) => {
341            tracing::warn!("checkpoint tools unavailable: {e}");
342            return None;
343        }
344    };
345    build_checkpoint_tools_inner(repo, session_id.into())
346}
347
348/// Like [`build_checkpoint_tools`], but with an explicit `shadow_dir`,
349/// for tests that want to bypass `paths::user_data_dir()` resolution.
350pub fn build_checkpoint_tools_at(
351    workspace: impl Into<std::path::PathBuf>,
352    shadow_dir: impl Into<std::path::PathBuf>,
353    session_id: impl Into<String>,
354) -> Option<(CheckpointList, CheckpointDiff, Arc<Mutex<ShadowRepo>>)> {
355    let repo = match ShadowRepo::open_at(workspace, shadow_dir) {
356        Ok(r) => Arc::new(Mutex::new(r)),
357        Err(e) => {
358            tracing::warn!("checkpoint tools unavailable: {e}");
359            return None;
360        }
361    };
362    build_checkpoint_tools_inner(repo, session_id.into())
363}
364
365fn build_checkpoint_tools_inner(
366    repo: Arc<Mutex<ShadowRepo>>,
367    session_id: String,
368) -> Option<(CheckpointList, CheckpointDiff, Arc<Mutex<ShadowRepo>>)> {
369    let ctx = CheckpointToolCtx {
370        repo: Arc::clone(&repo),
371        session_id,
372    };
373    Some((
374        CheckpointList::new(ctx.clone()),
375        CheckpointDiff::new(ctx),
376        repo,
377    ))
378}
379
380// ── tests ─────────────────────────────────────────────────────────────────────
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385    use std::fs;
386    use std::process::Command;
387    use tempfile::TempDir;
388
389    fn has_git() -> bool {
390        Command::new("git").arg("--version").output().is_ok()
391    }
392
393    /// Test workspace bundle: workspace tempdir + sibling shadow tempdir.
394    /// Tests use `build_checkpoint_tools_at` to bypass `paths::user_data_dir()`,
395    /// avoiding any need to mutate `RECURSIVE_HOME` (and thus the global
396    /// env lock). This lets the suite run fully in parallel.
397    struct TestWs {
398        workspace: TempDir,
399        shadow: TempDir,
400    }
401
402    impl TestWs {
403        fn path(&self) -> &std::path::Path {
404            self.workspace.path()
405        }
406        fn shadow_dir(&self) -> std::path::PathBuf {
407            self.shadow.path().join("shadow-git")
408        }
409    }
410
411    fn ws() -> TestWs {
412        TestWs {
413            workspace: tempfile::tempdir().expect("workspace tempdir"),
414            shadow: tempfile::tempdir().expect("shadow tempdir"),
415        }
416    }
417
418    #[test]
419    fn unix_now_is_positive() {
420        // kills `map(|d| d.as_secs() as i64)` → `map(|_| 0)` or `unwrap_or(0)` mutations
421        let ts = unix_now();
422        // Unix epoch was in 1970; any value < 1_000_000_000 is clearly wrong
423        assert!(
424            ts > 1_000_000_000,
425            "unix_now() returned suspicious value: {ts}"
426        );
427    }
428
429    #[tokio::test]
430    async fn list_tool_is_deferred() {
431        // kills `replace CheckpointList::is_deferred -> bool with false`
432        if !has_git() {
433            return;
434        }
435        let w = ws();
436        let (list, _, _) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
437        assert!(
438            list.is_deferred(),
439            "checkpoint_list must be a deferred tool"
440        );
441    }
442
443    #[tokio::test]
444    async fn diff_tool_is_deferred() {
445        // kills `replace CheckpointDiff::is_deferred -> bool with false`
446        if !has_git() {
447            return;
448        }
449        let w = ws();
450        let (_, diff, _) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
451        assert!(
452            diff.is_deferred(),
453            "checkpoint_diff must be a deferred tool"
454        );
455    }
456
457    #[tokio::test]
458    async fn list_tool_empty_session_returns_sentinel() {
459        // kills `if infos.is_empty()` branch-swap mutations
460        if !has_git() {
461            return;
462        }
463        let w = ws();
464        let (list, _, _) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
465        let out = list.execute(json!({})).await.unwrap();
466        assert_eq!(out, "no checkpoints in this session yet");
467    }
468
469    #[tokio::test]
470    async fn list_tool_shows_session_checkpoints() {
471        if !has_git() {
472            return;
473        }
474        let w = ws();
475        fs::write(w.path().join("a.txt"), "hi").unwrap();
476        let (list, _, repo) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "alpha").unwrap();
477        repo.lock()
478            .unwrap()
479            .snapshot_for_session("alpha", "t0")
480            .unwrap();
481        let out = list.execute(json!({})).await.unwrap();
482        assert!(out.contains("t0"));
483    }
484
485    #[tokio::test]
486    async fn list_tool_only_shows_own_session() {
487        if !has_git() {
488            return;
489        }
490        let w = ws();
491        fs::write(w.path().join("a.txt"), "hi").unwrap();
492        let (list_a, _, repo) =
493            build_checkpoint_tools_at(w.path(), w.shadow_dir(), "alpha").unwrap();
494        repo.lock()
495            .unwrap()
496            .snapshot_for_session("alpha", "tA")
497            .unwrap();
498        repo.lock()
499            .unwrap()
500            .snapshot_for_session("beta", "tB")
501            .unwrap();
502        let out_a = list_a.execute(json!({})).await.unwrap();
503        assert!(out_a.contains("tA"));
504        assert!(
505            !out_a.contains("tB"),
506            "alpha must not see beta's checkpoints"
507        );
508    }
509
510    #[tokio::test]
511    async fn diff_tool_returns_empty_for_no_change() {
512        if !has_git() {
513            return;
514        }
515        let w = ws();
516        fs::write(w.path().join("a.txt"), "x").unwrap();
517        let (_, diff, repo) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
518        let id = repo
519            .lock()
520            .unwrap()
521            .snapshot_for_session("s", "snap")
522            .unwrap();
523        let out = diff.execute(json!({"a": id.0})).await.unwrap();
524        assert_eq!(out, "no differences");
525    }
526
527    /// Goal 284: verify that `checkpoint_save` creates a shadow-git
528    /// commit and a `checkpoints.jsonl` entry.
529    #[tokio::test]
530    async fn checkpoint_save_tool_creates_entry() {
531        if !has_git() {
532            return;
533        }
534        let w = ws();
535        let log_path = w.path().join("checkpoints.jsonl");
536        fs::write(w.path().join("a.txt"), "hello").unwrap();
537
538        let repo = Arc::new(Mutex::new(
539            ShadowRepo::open_at(w.path(), w.shadow_dir()).unwrap(),
540        ));
541        let touched = Arc::new(Mutex::new(TouchedFiles::new()));
542        {
543            let mut t = touched.lock().unwrap();
544            t.paths.insert("a.txt".to_string());
545        }
546        let writer = Arc::new(Mutex::new(CheckpointLogWriter::open(&log_path).unwrap()));
547        let turn_index = Arc::new(std::sync::atomic::AtomicUsize::new(3));
548
549        let tool = build_checkpoint_save_tool(
550            repo.clone(),
551            "sess".into(),
552            Some(touched.clone()),
553            writer.clone(),
554            turn_index.clone(),
555            log_path.clone(),
556        );
557
558        let id_str = tool
559            .execute(json!({"message": "my save point"}))
560            .await
561            .unwrap();
562        assert!(!id_str.is_empty(), "should return a checkpoint id");
563
564        // Verify the checkpoint is in the shadow repo.
565        let cid = CheckpointId(id_str.clone());
566        let data = repo.lock().unwrap().read_file_at(&cid, "a.txt").unwrap();
567        assert_eq!(data, Some(b"hello".to_vec()));
568
569        // Verify the log entry.
570        let recs = read_log(&log_path).unwrap();
571        assert_eq!(recs.len(), 1);
572        assert_eq!(recs[0].turn, 3);
573        assert_eq!(recs[0].id.0, id_str);
574        assert_eq!(recs[0].message.as_deref(), Some("my save point"));
575        assert_eq!(recs[0].touched_files, vec!["a.txt".to_string()]);
576        assert_eq!(recs[0].touched_via, TouchedVia::Structured);
577    }
578
579    #[tokio::test]
580    async fn diff_tool_missing_a_arg_errors() {
581        // kills "missing `a`" guard removal mutation
582        if !has_git() {
583            return;
584        }
585        let w = ws();
586        let (_, diff, _) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
587        let res = diff.execute(json!({})).await;
588        assert!(res.is_err(), "diff without 'a' argument must fail");
589    }
590
591    #[tokio::test]
592    async fn checkpoint_save_defaults_message_to_turn_index() {
593        // kills `unwrap_or_else(|| format!("turn {turn}"))` mutation
594        if !has_git() {
595            return;
596        }
597        let w = ws();
598        let log_path = w.path().join("checkpoints.jsonl");
599        fs::write(w.path().join("b.txt"), "world").unwrap();
600
601        let repo = Arc::new(Mutex::new(
602            ShadowRepo::open_at(w.path(), w.shadow_dir()).unwrap(),
603        ));
604        let writer = Arc::new(Mutex::new(CheckpointLogWriter::open(&log_path).unwrap()));
605        let turn_index = Arc::new(std::sync::atomic::AtomicUsize::new(7));
606
607        let tool = build_checkpoint_save_tool(
608            repo.clone(),
609            "sess2".into(),
610            None,
611            writer.clone(),
612            turn_index.clone(),
613            log_path.clone(),
614        );
615
616        // No "message" argument — should default to "turn 7"
617        tool.execute(json!({})).await.unwrap();
618
619        let recs = read_log(&log_path).unwrap();
620        assert_eq!(recs.len(), 1);
621        assert_eq!(
622            recs[0].message.as_deref(),
623            Some("turn 7"),
624            "default message must be 'turn N' where N is the turn index"
625        );
626    }
627}