lean_ctx/tools/registered/
ctx_checkpoint.rs

1use std::path::Path;
2
3use rmcp::ErrorData;
4use rmcp::model::Tool;
5use serde_json::{Map, Value, json};
6
7use crate::core::git::shadow::{self, Checkpoint};
8use crate::core::tokens::count_tokens;
9use crate::server::tool_trait::{McpTool, ToolContext, ToolOutput, get_int, get_str};
10use crate::tool_defs::tool_def;
11
12const DEFAULT_LOG_LIMIT: usize = 20;
13const DIFF_MAX_TOKENS: usize = 8000;
14
15/// `ctx_checkpoint` — snapshot, review, diff, and revert the agent's code
16/// changes in a shadow git history kept outside the user's own `.git`.
17pub struct CtxCheckpointTool;
18
19impl McpTool for CtxCheckpointTool {
20    fn name(&self) -> &'static str {
21        "ctx_checkpoint"
22    }
23
24    fn tool_def(&self) -> Tool {
25        tool_def(
26            "ctx_checkpoint",
27            "Local shadow git history of the agent's changes — separate from the user's .git.\n\
28WORKFLOW: snapshot before+after changes to capture exactly what was modified.\n\
29Actions: snapshot (record current state), log (list checkpoints with SHAs),\n\
30diff from=... to=... (compare checkpoints), restore ref=... (revert files).\n\
31ANTIPATTERN: Never touches the user's repository — completely isolated shadow history.",
32            json!({
33                "type": "object",
34                "properties": {
35                    "action": {
36                        "type": "string",
37                        "enum": ["snapshot", "log", "diff", "restore"],
38                        "description": "snapshot|log|diff|restore (default: log)"
39                    },
40                    "message": { "type": "string", "description": "Label for the snapshot (action=snapshot only)" },
41                    "from": { "type": "string", "description": "Base checkpoint SHA for diff (default: HEAD~1, action=diff only)" },
42                    "to": { "type": "string", "description": "Target checkpoint SHA for diff (default: HEAD, action=diff only)" },
43                    "ref": { "type": "string", "description": "Checkpoint SHA to restore (required for action=restore)" },
44                    "path": { "type": "string", "description": "File/dir to restore (omit for full restore, action=restore only)" },
45                    "limit": { "type": "integer", "description": "Max checkpoints in log (default: 20, max: 200)" }
46                }
47            }),
48        )
49    }
50
51    fn handle(
52        &self,
53        args: &Map<String, Value>,
54        ctx: &ToolContext,
55    ) -> Result<ToolOutput, ErrorData> {
56        let action = get_str(args, "action").unwrap_or_else(|| "log".to_string());
57        let project = Path::new(&ctx.project_root).to_path_buf();
58        if ctx.project_root.is_empty() {
59            return Err(ErrorData::invalid_params(
60                "no project root resolved for ctx_checkpoint",
61                None,
62            ));
63        }
64
65        let result = tokio::task::block_in_place(|| match action.as_str() {
66            "snapshot" => {
67                let msg = get_str(args, "message").unwrap_or_default();
68                shadow::snapshot(&project, &msg).map(|c| render_checkpoint_line("Checkpoint", &c))
69            }
70            "log" => {
71                let limit =
72                    get_int(args, "limit").map_or(DEFAULT_LOG_LIMIT, |n| n.clamp(1, 200) as usize);
73                shadow::log(&project, limit).map(|cs| render_log(&cs))
74            }
75            "diff" => {
76                let from = get_str(args, "from");
77                let to = get_str(args, "to");
78                shadow::diff(&project, from.as_deref(), to.as_deref())
79                    .map(|d| budget(&d, DIFF_MAX_TOKENS))
80            }
81            "restore" => {
82                let git_ref = get_str(args, "ref").ok_or_else(|| {
83                    "restore requires 'ref' (a checkpoint sha from log)".to_string()
84                })?;
85                let path = get_str(args, "path");
86                shadow::restore(&project, &git_ref, path.as_deref())
87            }
88            other => Err(format!(
89                "invalid action '{other}' (use: snapshot, log, diff, restore)"
90            )),
91        });
92
93        match result {
94            Ok(text) => Ok(ToolOutput {
95                text,
96                original_tokens: 0,
97                saved_tokens: 0,
98                mode: Some(action),
99                path: None,
100                changed: matches!(
101                    args.get("action").and_then(Value::as_str),
102                    Some("snapshot" | "restore")
103                ),
104                shell_outcome: None,
105                content_blocks: None,
106            }),
107            Err(e) => Err(ErrorData::invalid_params(
108                format!("ctx_checkpoint failed: {e}"),
109                None,
110            )),
111        }
112    }
113}
114
115fn render_checkpoint_line(prefix: &str, c: &Checkpoint) -> String {
116    let files = c
117        .files_changed
118        .map(|n| format!(" · {n} file(s)"))
119        .unwrap_or_default();
120    format!("{prefix} {}{files} — {}", c.sha, c.message)
121}
122
123fn render_log(checkpoints: &[Checkpoint]) -> String {
124    if checkpoints.is_empty() {
125        return "No checkpoints yet. Run `ctx_checkpoint` with action=snapshot to record one."
126            .to_string();
127    }
128    let mut out = format!("{} checkpoint(s):\n", checkpoints.len());
129    for c in checkpoints {
130        out.push_str(&format!("- {} · {} — {}\n", c.sha, c.time, c.message));
131    }
132    out.trim_end().to_string()
133}
134
135fn budget(content: &str, max_tokens: usize) -> String {
136    if content.trim().is_empty() {
137        return "No differences.".to_string();
138    }
139    let tokens = count_tokens(content);
140    if tokens <= max_tokens {
141        return content.to_string();
142    }
143    let ratio = max_tokens as f64 / tokens as f64;
144    let keep = ((content.chars().count() as f64 * ratio) as usize).max(1);
145    let truncated: String = content.chars().take(keep).collect();
146    format!("{truncated}\n\n…[diff truncated to ~{max_tokens} tokens]")
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    fn cp(sha: &str, msg: &str, files: Option<usize>) -> Checkpoint {
154        Checkpoint {
155            sha: sha.to_string(),
156            time: "2026-06-07T10:00:00Z".to_string(),
157            message: msg.to_string(),
158            files_changed: files,
159        }
160    }
161
162    #[test]
163    fn renders_empty_log_with_hint() {
164        assert!(render_log(&[]).contains("snapshot"));
165    }
166
167    #[test]
168    fn renders_checkpoint_line_with_file_count() {
169        let line = render_checkpoint_line("Checkpoint", &cp("abc1234", "fix bug", Some(3)));
170        assert_eq!(line, "Checkpoint abc1234 · 3 file(s) — fix bug");
171    }
172
173    #[test]
174    fn renders_log_entries() {
175        let out = render_log(&[cp("a1", "one", None), cp("b2", "two", None)]);
176        assert!(out.contains("2 checkpoint(s)"));
177        assert!(out.contains("- a1 ·"));
178        assert!(out.contains("- b2 ·"));
179    }
180
181    #[test]
182    fn budget_handles_empty_diff() {
183        assert_eq!(budget("", 100), "No differences.");
184    }
185}
lean_ctx/tools/registered/ctx_checkpoint.rs

lean_ctx/tools/registered/
ctx_checkpoint.rs
