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            }),
106            Err(e) => Err(ErrorData::invalid_params(
107                format!("ctx_checkpoint failed: {e}"),
108                None,
109            )),
110        }
111    }
112}
113
114fn render_checkpoint_line(prefix: &str, c: &Checkpoint) -> String {
115    let files = c
116        .files_changed
117        .map(|n| format!(" · {n} file(s)"))
118        .unwrap_or_default();
119    format!("{prefix} {}{files} — {}", c.sha, c.message)
120}
121
122fn render_log(checkpoints: &[Checkpoint]) -> String {
123    if checkpoints.is_empty() {
124        return "No checkpoints yet. Run `ctx_checkpoint` with action=snapshot to record one."
125            .to_string();
126    }
127    let mut out = format!("{} checkpoint(s):\n", checkpoints.len());
128    for c in checkpoints {
129        out.push_str(&format!("- {} · {} — {}\n", c.sha, c.time, c.message));
130    }
131    out.trim_end().to_string()
132}
133
134fn budget(content: &str, max_tokens: usize) -> String {
135    if content.trim().is_empty() {
136        return "No differences.".to_string();
137    }
138    let tokens = count_tokens(content);
139    if tokens <= max_tokens {
140        return content.to_string();
141    }
142    let ratio = max_tokens as f64 / tokens as f64;
143    let keep = ((content.chars().count() as f64 * ratio) as usize).max(1);
144    let truncated: String = content.chars().take(keep).collect();
145    format!("{truncated}\n\n…[diff truncated to ~{max_tokens} tokens]")
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    fn cp(sha: &str, msg: &str, files: Option<usize>) -> Checkpoint {
153        Checkpoint {
154            sha: sha.to_string(),
155            time: "2026-06-07T10:00:00Z".to_string(),
156            message: msg.to_string(),
157            files_changed: files,
158        }
159    }
160
161    #[test]
162    fn renders_empty_log_with_hint() {
163        assert!(render_log(&[]).contains("snapshot"));
164    }
165
166    #[test]
167    fn renders_checkpoint_line_with_file_count() {
168        let line = render_checkpoint_line("Checkpoint", &cp("abc1234", "fix bug", Some(3)));
169        assert_eq!(line, "Checkpoint abc1234 · 3 file(s) — fix bug");
170    }
171
172    #[test]
173    fn renders_log_entries() {
174        let out = render_log(&[cp("a1", "one", None), cp("b2", "two", None)]);
175        assert!(out.contains("2 checkpoint(s)"));
176        assert!(out.contains("- a1 ·"));
177        assert!(out.contains("- b2 ·"));
178    }
179
180    #[test]
181    fn budget_handles_empty_diff() {
182        assert_eq!(budget("", 100), "No differences.");
183    }
184}
lean_ctx/tools/registered/ctx_checkpoint.rs

lean_ctx/tools/registered/
ctx_checkpoint.rs
