lean_ctx_server/tools/registered/
ctx_checkpoint.rs

1// SPDX-License-Identifier: Apache-2.0
2use std::path::Path;
3
4use rmcp::ErrorData;
5use rmcp::model::Tool;
6use serde_json::{Map, Value, json};
7
8use crate::server::tool_trait::{McpTool, ToolContext, ToolOutput, get_int, get_str};
9use lean_ctx::core::git::shadow::{self, Checkpoint};
10use lean_ctx::core::tokens::count_tokens;
11use lean_ctx::tool_defs::tool_def;
12
13const DEFAULT_LOG_LIMIT: usize = 20;
14const DIFF_MAX_TOKENS: usize = 8000;
15
16/// `ctx_checkpoint` — snapshot, review, diff, and revert the agent's code
17/// changes in a shadow git history kept outside the user's own `.git`.
18pub struct CtxCheckpointTool;
19
20impl McpTool for CtxCheckpointTool {
21    fn name(&self) -> &'static str {
22        "ctx_checkpoint"
23    }
24
25    fn tool_def(&self) -> Tool {
26        tool_def(
27            "ctx_checkpoint",
28            "Local shadow git history of the agent's changes — separate from the user's .git.\n\
29WORKFLOW: snapshot before+after changes to capture exactly what was modified.\n\
30Actions: snapshot (record current state), log (list checkpoints with SHAs),\n\
31diff from=... to=... (compare checkpoints), restore ref=... (revert files).\n\
32ANTIPATTERN: Never touches the user's repository — completely isolated shadow history.",
33            json!({
34                "type": "object",
35                "properties": {
36                    "action": {
37                        "type": "string",
38                        "enum": ["snapshot", "log", "diff", "restore"],
39                        "description": "snapshot|log|diff|restore (default: log)"
40                    },
41                    "message": { "type": "string", "description": "Label for the snapshot (action=snapshot only)" },
42                    "from": { "type": "string", "description": "Base checkpoint SHA for diff (default: HEAD~1, action=diff only)" },
43                    "to": { "type": "string", "description": "Target checkpoint SHA for diff (default: HEAD, action=diff only)" },
44                    "ref": { "type": "string", "description": "Checkpoint SHA to restore (required for action=restore)" },
45                    "path": { "type": "string", "description": "File/dir to restore (omit for full restore, action=restore only)" },
46                    "limit": { "type": "integer", "description": "Max checkpoints in log (default: 20, max: 200)" }
47                },
48                "allOf": [
49                    {
50                        "if": { "properties": { "action": { "const": "restore" } }, "required": ["action"] },
51                        "then": { "required": ["action", "ref"] }
52                    }
53                ]
54            }),
55        )
56    }
57
58    fn handle(
59        &self,
60        args: &Map<String, Value>,
61        ctx: &ToolContext,
62    ) -> Result<ToolOutput, ErrorData> {
63        let action = get_str(args, "action").unwrap_or_else(|| "log".to_string());
64        let project = Path::new(&ctx.project_root).to_path_buf();
65        if ctx.project_root.is_empty() {
66            return Err(ErrorData::invalid_params(
67                "no project root resolved for ctx_checkpoint",
68                None,
69            ));
70        }
71
72        let result = tokio::task::block_in_place(|| match action.as_str() {
73            "snapshot" => {
74                let msg = get_str(args, "message").unwrap_or_default();
75                shadow::snapshot(&project, &msg).map(|c| render_checkpoint_line("Checkpoint", &c))
76            }
77            "log" => {
78                let limit =
79                    get_int(args, "limit").map_or(DEFAULT_LOG_LIMIT, |n| n.clamp(1, 200) as usize);
80                shadow::log(&project, limit).map(|cs| render_log(&cs))
81            }
82            "diff" => {
83                let from = get_str(args, "from");
84                let to = get_str(args, "to");
85                shadow::diff(&project, from.as_deref(), to.as_deref())
86                    .and_then(|d| admitted_diff(&d))
87                    .inspect(|_| {
88                        lean_ctx::core::context_admission::capture::note_source_id(&format!(
89                            "checkpoint-diff:{}..{}",
90                            from.as_deref().unwrap_or(""),
91                            to.as_deref().unwrap_or("")
92                        ));
93                    })
94                    .map(|d| budget(&d, DIFF_MAX_TOKENS))
95            }
96            "restore" => {
97                let git_ref = get_str(args, "ref").ok_or_else(|| {
98                    "restore requires 'ref' (a checkpoint sha from log)".to_string()
99                })?;
100                let path = get_str(args, "path")
101                    .map(|raw| {
102                        lean_ctx::core::workspace_authority::confine(&raw, &ctx.project_root, true)
103                    })
104                    .transpose()?;
105                shadow::restore(&project, &git_ref, path.as_deref())
106            }
107            other => Err(format!(
108                "invalid action '{other}' (use: snapshot, log, diff, restore)"
109            )),
110        });
111
112        match result {
113            Ok(text) => Ok(ToolOutput {
114                text,
115                original_tokens: 0,
116                saved_tokens: 0,
117                mode: Some(action),
118                path: None,
119                changed: matches!(
120                    args.get("action").and_then(Value::as_str),
121                    Some("snapshot" | "restore")
122                ),
123                shell_outcome: None,
124                content_blocks: None,
125            }),
126            Err(e) => Err(ErrorData::invalid_params(
127                format!("ctx_checkpoint failed: {e}"),
128                None,
129            )),
130        }
131    }
132}
133
134/// A checkpoint diff is content of project files taken from git objects, not
135/// from a rooted file read: it is admitted as a context source — project
136/// policy on the text, then the gateway — BEFORE it is cut to the token
137/// budget (G01). A diff that is withheld yields the content-free reason as
138/// the tool's error. Limit: a diff is an excerpt; a marking that sits in an
139/// unchanged part of a file does not travel with its hunks.
140fn admitted_diff(diff: &str) -> Result<String, String> {
141    if diff.trim().is_empty() {
142        return Ok(diff.to_string());
143    }
144    crate::tools::admitted_sources::admit_remote_text(diff, "ctx_checkpoint")
145        .map(|admission| admission.text)
146}
147
148fn render_checkpoint_line(prefix: &str, c: &Checkpoint) -> String {
149    let files = c
150        .files_changed
151        .map(|n| format!(" · {n} file(s)"))
152        .unwrap_or_default();
153    format!("{prefix} {}{files} — {}", c.sha, c.message)
154}
155
156fn render_log(checkpoints: &[Checkpoint]) -> String {
157    if checkpoints.is_empty() {
158        return "No checkpoints yet. Run `ctx_checkpoint` with action=snapshot to record one."
159            .to_string();
160    }
161    let mut out = format!("{} checkpoint(s):\n", checkpoints.len());
162    for c in checkpoints {
163        out.push_str(&format!("- {} · {} — {}\n", c.sha, c.time, c.message));
164    }
165    out.trim_end().to_string()
166}
167
168fn budget(content: &str, max_tokens: usize) -> String {
169    if content.trim().is_empty() {
170        return "No differences.".to_string();
171    }
172    let tokens = count_tokens(content);
173    if tokens <= max_tokens {
174        return content.to_string();
175    }
176    let ratio = max_tokens as f64 / tokens as f64;
177    let keep = ((content.chars().count() as f64 * ratio) as usize).max(1);
178    let truncated: String = content.chars().take(keep).collect();
179    lean_ctx::core::recovery::with_original_pointer(
180        "ctx_checkpoint",
181        "diff",
182        content,
183        format!("{truncated}\n\n…[diff truncated to ~{max_tokens} tokens]"),
184    )
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    fn cp(sha: &str, msg: &str, files: Option<usize>) -> Checkpoint {
192        Checkpoint {
193            sha: sha.to_string(),
194            time: "2026-06-07T10:00:00Z".to_string(),
195            message: msg.to_string(),
196            files_changed: files,
197        }
198    }
199
200    #[test]
201    fn renders_empty_log_with_hint() {
202        assert!(render_log(&[]).contains("snapshot"));
203    }
204
205    #[test]
206    fn renders_checkpoint_line_with_file_count() {
207        let line = render_checkpoint_line("Checkpoint", &cp("abc1234", "fix bug", Some(3)));
208        assert_eq!(line, "Checkpoint abc1234 · 3 file(s) — fix bug");
209    }
210
211    #[test]
212    fn renders_log_entries() {
213        let out = render_log(&[cp("a1", "one", None), cp("b2", "two", None)]);
214        assert!(out.contains("2 checkpoint(s)"));
215        assert!(out.contains("- a1 ·"));
216        assert!(out.contains("- b2 ·"));
217    }
218
219    #[test]
220    fn budget_handles_empty_diff() {
221        assert_eq!(budget("", 100), "No differences.");
222    }
223
224    /// G01: a checkpoint diff was returned as git printed it — a key in a
225    /// changed line reached the model, and a diff the project policy blocks
226    /// was delivered. A clean diff keeps its bytes.
227    #[test]
228    fn a_diff_is_admitted_before_it_is_budgeted() {
229        const AWS_KEY: &str = concat!("AK", "IAIOSFODNN7EXAMPLE");
230        let _data = lean_ctx::core::data_dir::isolated_data_dir();
231        let clean = "--- a/src/a.rs\n+++ b/src/a.rs\n@@ -1 +1 @@\n-fn a() {}\n+fn b() {}\n";
232        {
233            let _community = lean_ctx::core::policy::runtime::TestPolicyOverride::set(None);
234            assert_eq!(admitted_diff(clean).as_deref(), Ok(clean));
235            assert_eq!(admitted_diff("").as_deref(), Ok(""));
236            let keyed = format!("@@ -1 +1 @@\n+let aws_key = \"{AWS_KEY}\";\n");
237            assert!(admitted_diff(&keyed).is_ok_and(|diff| !diff.contains(AWS_KEY)));
238        }
239        let policy = lean_ctx::core::policy::load(
240            "name = \"ctx-checkpoint-test\"\nversion = \"1.0.0\"\ndescription = \"test\"\n\
241             [redaction]\ncustomer = \"CUS-[0-9]{4}\"\n",
242        )
243        .unwrap();
244        let _policy = lean_ctx::core::policy::runtime::TestPolicyOverride::set(Some(policy));
245        let masked = admitted_diff("@@ -1 +1 @@\n+let owner = \"CUS-1234\";\n").unwrap();
246        assert!(!masked.contains("CUS-1234"), "{masked}");
247    }
248}
lean_ctx_server/tools/registered/ctx_checkpoint.rs

lean_ctx_server/tools/registered/
ctx_checkpoint.rs
