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    // A diff of this checkout: the project's own text.
145    lean_ctx::core::context_admission::admit_instance_context_source(diff)
146        .map_err(|error| error.to_string())
147        .map(|admission| admission.text)
148}
149
150fn render_checkpoint_line(prefix: &str, c: &Checkpoint) -> String {
151    let files = c
152        .files_changed
153        .map(|n| format!(" · {n} file(s)"))
154        .unwrap_or_default();
155    format!("{prefix} {}{files} — {}", c.sha, c.message)
156}
157
158fn render_log(checkpoints: &[Checkpoint]) -> String {
159    if checkpoints.is_empty() {
160        return "No checkpoints yet. Run `ctx_checkpoint` with action=snapshot to record one."
161            .to_string();
162    }
163    let mut out = format!("{} checkpoint(s):\n", checkpoints.len());
164    for c in checkpoints {
165        out.push_str(&format!("- {} · {} — {}\n", c.sha, c.time, c.message));
166    }
167    out.trim_end().to_string()
168}
169
170fn budget(content: &str, max_tokens: usize) -> String {
171    if content.trim().is_empty() {
172        return "No differences.".to_string();
173    }
174    let tokens = count_tokens(content);
175    if tokens <= max_tokens {
176        return content.to_string();
177    }
178    let ratio = max_tokens as f64 / tokens as f64;
179    let keep = ((content.chars().count() as f64 * ratio) as usize).max(1);
180    let truncated: String = content.chars().take(keep).collect();
181    lean_ctx::core::recovery::with_original_pointer(
182        "ctx_checkpoint",
183        "diff",
184        content,
185        format!("{truncated}\n\n…[diff truncated to ~{max_tokens} tokens]"),
186    )
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    fn cp(sha: &str, msg: &str, files: Option<usize>) -> Checkpoint {
194        Checkpoint {
195            sha: sha.to_string(),
196            time: "2026-06-07T10:00:00Z".to_string(),
197            message: msg.to_string(),
198            files_changed: files,
199        }
200    }
201
202    #[test]
203    fn renders_empty_log_with_hint() {
204        assert!(render_log(&[]).contains("snapshot"));
205    }
206
207    #[test]
208    fn renders_checkpoint_line_with_file_count() {
209        let line = render_checkpoint_line("Checkpoint", &cp("abc1234", "fix bug", Some(3)));
210        assert_eq!(line, "Checkpoint abc1234 · 3 file(s) — fix bug");
211    }
212
213    #[test]
214    fn renders_log_entries() {
215        let out = render_log(&[cp("a1", "one", None), cp("b2", "two", None)]);
216        assert!(out.contains("2 checkpoint(s)"));
217        assert!(out.contains("- a1 ·"));
218        assert!(out.contains("- b2 ·"));
219    }
220
221    #[test]
222    fn budget_handles_empty_diff() {
223        assert_eq!(budget("", 100), "No differences.");
224    }
225
226    /// G01: a checkpoint diff was returned as git printed it — a key in a
227    /// changed line reached the model, and a diff the project policy blocks
228    /// was delivered. A clean diff keeps its bytes.
229    #[test]
230    fn a_diff_is_admitted_before_it_is_budgeted() {
231        const AWS_KEY: &str = concat!("AK", "IAIOSFODNN7EXAMPLE");
232        let _data = lean_ctx::core::data_dir::isolated_data_dir();
233        let clean = "--- a/src/a.rs\n+++ b/src/a.rs\n@@ -1 +1 @@\n-fn a() {}\n+fn b() {}\n";
234        {
235            let _community = lean_ctx::core::policy::runtime::TestPolicyOverride::set(None);
236            assert_eq!(admitted_diff(clean).as_deref(), Ok(clean));
237            assert_eq!(admitted_diff("").as_deref(), Ok(""));
238            let keyed = format!("@@ -1 +1 @@\n+let aws_key = \"{AWS_KEY}\";\n");
239            assert!(admitted_diff(&keyed).is_ok_and(|diff| !diff.contains(AWS_KEY)));
240        }
241        let policy = lean_ctx::core::policy::load(
242            "name = \"ctx-checkpoint-test\"\nversion = \"1.0.0\"\ndescription = \"test\"\n\
243             [redaction]\ncustomer = \"CUS-[0-9]{4}\"\n",
244        )
245        .unwrap();
246        let _policy = lean_ctx::core::policy::runtime::TestPolicyOverride::set(Some(policy));
247        let masked = admitted_diff("@@ -1 +1 @@\n+let owner = \"CUS-1234\";\n").unwrap();
248        assert!(!masked.contains("CUS-1234"), "{masked}");
249    }
250}
lean_ctx_server/tools/registered/ctx_checkpoint.rs

lean_ctx_server/tools/registered/
ctx_checkpoint.rs
