use std::path::Path;
use rmcp::ErrorData;
use rmcp::model::Tool;
use serde_json::{Map, Value, json};
use crate::core::git::shadow::{self, Checkpoint};
use crate::core::tokens::count_tokens;
use crate::server::tool_trait::{McpTool, ToolContext, ToolOutput, get_int, get_str};
use crate::tool_defs::tool_def;
const DEFAULT_LOG_LIMIT: usize = 20;
const DIFF_MAX_TOKENS: usize = 8000;
pub struct CtxCheckpointTool;
impl McpTool for CtxCheckpointTool {
fn name(&self) -> &'static str {
"ctx_checkpoint"
}
fn tool_def(&self) -> Tool {
tool_def(
"ctx_checkpoint",
"Local shadow git history of the agent's changes — separate from the user's .git.\n\
WORKFLOW: snapshot before+after changes to capture exactly what was modified.\n\
Actions: snapshot (record current state), log (list checkpoints with SHAs),\n\
diff from=... to=... (compare checkpoints), restore ref=... (revert files).\n\
ANTIPATTERN: Never touches the user's repository — completely isolated shadow history.",
json!({
"type": "object",
"properties": {
"action": {
"type": "string",
"enum": ["snapshot", "log", "diff", "restore"],
"description": "snapshot|log|diff|restore (default: log)"
},
"message": { "type": "string", "description": "Label for the snapshot (action=snapshot only)" },
"from": { "type": "string", "description": "Base checkpoint SHA for diff (default: HEAD~1, action=diff only)" },
"to": { "type": "string", "description": "Target checkpoint SHA for diff (default: HEAD, action=diff only)" },
"ref": { "type": "string", "description": "Checkpoint SHA to restore (required for action=restore)" },
"path": { "type": "string", "description": "File/dir to restore (omit for full restore, action=restore only)" },
"limit": { "type": "integer", "description": "Max checkpoints in log (default: 20, max: 200)" }
},
"allOf": [
{
"if": { "properties": { "action": { "const": "restore" } }, "required": ["action"] },
"then": { "required": ["action", "ref"] }
}
]
}),
)
}
fn handle(
&self,
args: &Map<String, Value>,
ctx: &ToolContext,
) -> Result<ToolOutput, ErrorData> {
let action = get_str(args, "action").unwrap_or_else(|| "log".to_string());
let project = Path::new(&ctx.project_root).to_path_buf();
if ctx.project_root.is_empty() {
return Err(ErrorData::invalid_params(
"no project root resolved for ctx_checkpoint",
None,
));
}
let result = tokio::task::block_in_place(|| match action.as_str() {
"snapshot" => {
let msg = get_str(args, "message").unwrap_or_default();
shadow::snapshot(&project, &msg).map(|c| render_checkpoint_line("Checkpoint", &c))
}
"log" => {
let limit =
get_int(args, "limit").map_or(DEFAULT_LOG_LIMIT, |n| n.clamp(1, 200) as usize);
shadow::log(&project, limit).map(|cs| render_log(&cs))
}
"diff" => {
let from = get_str(args, "from");
let to = get_str(args, "to");
shadow::diff(&project, from.as_deref(), to.as_deref())
.map(|d| budget(&d, DIFF_MAX_TOKENS))
}
"restore" => {
let git_ref = get_str(args, "ref").ok_or_else(|| {
"restore requires 'ref' (a checkpoint sha from log)".to_string()
})?;
let path = get_str(args, "path");
shadow::restore(&project, &git_ref, path.as_deref())
}
other => Err(format!(
"invalid action '{other}' (use: snapshot, log, diff, restore)"
)),
});
match result {
Ok(text) => Ok(ToolOutput {
text,
original_tokens: 0,
saved_tokens: 0,
mode: Some(action),
path: None,
changed: matches!(
args.get("action").and_then(Value::as_str),
Some("snapshot" | "restore")
),
shell_outcome: None,
content_blocks: None,
}),
Err(e) => Err(ErrorData::invalid_params(
format!("ctx_checkpoint failed: {e}"),
None,
)),
}
}
}
fn render_checkpoint_line(prefix: &str, c: &Checkpoint) -> String {
let files = c
.files_changed
.map(|n| format!(" · {n} file(s)"))
.unwrap_or_default();
format!("{prefix} {}{files} — {}", c.sha, c.message)
}
fn render_log(checkpoints: &[Checkpoint]) -> String {
if checkpoints.is_empty() {
return "No checkpoints yet. Run `ctx_checkpoint` with action=snapshot to record one."
.to_string();
}
let mut out = format!("{} checkpoint(s):\n", checkpoints.len());
for c in checkpoints {
out.push_str(&format!("- {} · {} — {}\n", c.sha, c.time, c.message));
}
out.trim_end().to_string()
}
fn budget(content: &str, max_tokens: usize) -> String {
if content.trim().is_empty() {
return "No differences.".to_string();
}
let tokens = count_tokens(content);
if tokens <= max_tokens {
return content.to_string();
}
let ratio = max_tokens as f64 / tokens as f64;
let keep = ((content.chars().count() as f64 * ratio) as usize).max(1);
let truncated: String = content.chars().take(keep).collect();
format!("{truncated}\n\n…[diff truncated to ~{max_tokens} tokens]")
}
#[cfg(test)]
mod tests {
use super::*;
fn cp(sha: &str, msg: &str, files: Option<usize>) -> Checkpoint {
Checkpoint {
sha: sha.to_string(),
time: "2026-06-07T10:00:00Z".to_string(),
message: msg.to_string(),
files_changed: files,
}
}
#[test]
fn renders_empty_log_with_hint() {
assert!(render_log(&[]).contains("snapshot"));
}
#[test]
fn renders_checkpoint_line_with_file_count() {
let line = render_checkpoint_line("Checkpoint", &cp("abc1234", "fix bug", Some(3)));
assert_eq!(line, "Checkpoint abc1234 · 3 file(s) — fix bug");
}
#[test]
fn renders_log_entries() {
let out = render_log(&[cp("a1", "one", None), cp("b2", "two", None)]);
assert!(out.contains("2 checkpoint(s)"));
assert!(out.contains("- a1 ·"));
assert!(out.contains("- b2 ·"));
}
#[test]
fn budget_handles_empty_diff() {
assert_eq!(budget("", 100), "No differences.");
}
}