1use std::path::{Path, PathBuf};
2use std::process::{Command, Stdio};
3
4use anyhow::{Context, Result};
5use serde::{Deserialize, Serialize};
6use sha2::{Digest, Sha256};
7
8use crate::{NewApproval, NewCheckpoint, RuntimeStore, data_dir};
9
10#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
11pub struct CheckpointFile {
12 pub path: String,
13 pub existed: bool,
14 pub snapshot_relpath: Option<String>,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
18pub struct CheckpointManifest {
19 pub id: String,
20 #[serde(default)]
21 pub task_id: Option<String>,
22 pub project_path: String,
23 pub files: Vec<CheckpointFile>,
24 pub pending_action: Option<serde_json::Value>,
25 #[serde(default)]
26 pub shadow_git_repo: Option<String>,
27 #[serde(default)]
28 pub shadow_git_commit: Option<String>,
29 pub created_at: String,
30}
31
32pub fn create_checkpoint(
33 project_path: &Path,
34 paths: &[PathBuf],
35 pending_action: Option<serde_json::Value>,
36) -> Result<CheckpointManifest> {
37 create_checkpoint_for_task(project_path, paths, pending_action, None)
38}
39
40pub fn create_checkpoint_for_task(
41 project_path: &Path,
42 paths: &[PathBuf],
43 pending_action: Option<serde_json::Value>,
44 task_id: Option<String>,
45) -> Result<CheckpointManifest> {
46 let id = fresh_checkpoint_id();
47 let root = data_dir()?.join("checkpoints").join(&id);
48 let files_dir = root.join("files");
49 std::fs::create_dir_all(&files_dir)
50 .with_context(|| format!("failed to create checkpoint dir {}", files_dir.display()))?;
51
52 let project_root = std::fs::canonicalize(project_path).unwrap_or_else(|_| project_path.into());
53 let mut files = Vec::new();
54 for path in paths {
55 let candidate = if path.is_absolute() {
56 path.clone()
57 } else {
58 project_path.join(path)
59 };
60 let normalized = std::fs::canonicalize(&candidate).unwrap_or(candidate.clone());
61 let display = normalized
62 .strip_prefix(&project_root)
63 .unwrap_or(&normalized)
64 .display()
65 .to_string();
66 if normalized.exists() && normalized.is_file() {
67 let safe_rel = sanitize_relpath(&display);
68 let dest = files_dir.join(&safe_rel);
69 if let Some(parent) = dest.parent() {
70 std::fs::create_dir_all(parent)?;
71 }
72 std::fs::copy(&normalized, &dest).with_context(|| {
73 format!(
74 "failed to copy checkpoint file {} -> {}",
75 normalized.display(),
76 dest.display()
77 )
78 })?;
79 files.push(CheckpointFile {
80 path: display,
81 existed: true,
82 snapshot_relpath: Some(format!("files/{}", safe_rel)),
83 });
84 } else {
85 files.push(CheckpointFile {
86 path: display,
87 existed: false,
88 snapshot_relpath: None,
89 });
90 }
91 }
92
93 let shadow_git = snapshot_shadow_git(&project_root, &files, &id).ok();
94 let manifest = CheckpointManifest {
95 id: id.clone(),
96 task_id: task_id.clone(),
97 project_path: project_path.display().to_string(),
98 files,
99 pending_action,
100 shadow_git_repo: shadow_git.as_ref().map(|snapshot| snapshot.repo.clone()),
101 shadow_git_commit: shadow_git.as_ref().map(|snapshot| snapshot.commit.clone()),
102 created_at: chrono::Utc::now().to_rfc3339(),
103 };
104 let manifest_path = root.join("manifest.json");
105 std::fs::write(&manifest_path, serde_json::to_vec_pretty(&manifest)?)?;
106
107 if let Ok(store) = RuntimeStore::open_default() {
108 let _ = store.checkpoints().create(NewCheckpoint {
109 id: Some(id),
110 task_id,
111 project_path: manifest.project_path.clone(),
112 snapshot_path: root.display().to_string(),
113 changed_files_json: serde_json::to_string(&manifest.files)?,
114 pending_action_json: manifest
115 .pending_action
116 .as_ref()
117 .map(serde_json::to_string)
118 .transpose()?,
119 approval_id: None,
120 });
121 }
122
123 let _ = crate::run_plugin_hooks(
124 "checkpoint",
125 &serde_json::json!({
126 "id": manifest.id.clone(),
127 "task_id": manifest.task_id.clone(),
128 "project_path": manifest.project_path.clone(),
129 "files": manifest.files.clone(),
130 "created_at": manifest.created_at.clone(),
131 }),
132 );
133
134 Ok(manifest)
135}
136
137pub fn restore_checkpoint(id: &str) -> Result<CheckpointManifest> {
138 let root = data_dir()?.join("checkpoints").join(id);
139 let manifest_path = root.join("manifest.json");
140 let raw = std::fs::read_to_string(&manifest_path)
141 .with_context(|| format!("failed to read {}", manifest_path.display()))?;
142 let manifest: CheckpointManifest = serde_json::from_str(&raw)?;
143 let project_root = PathBuf::from(&manifest.project_path);
144 for file in &manifest.files {
145 let target = project_root.join(&file.path);
146 if file.existed {
147 let rel = file
148 .snapshot_relpath
149 .as_ref()
150 .context("checkpoint file missing snapshot_relpath")?;
151 let source = root.join(rel);
152 if let Some(parent) = target.parent() {
153 std::fs::create_dir_all(parent)?;
154 }
155 std::fs::copy(&source, &target).with_context(|| {
156 format!(
157 "failed to restore checkpoint file {} -> {}",
158 source.display(),
159 target.display()
160 )
161 })?;
162 } else if target.exists() {
163 if target.is_file() {
164 std::fs::remove_file(&target)?;
165 } else if target.is_dir() {
166 std::fs::remove_dir_all(&target)?;
167 }
168 }
169 }
170 if let Some(action) = manifest.pending_action.as_ref()
171 && action.get("tool").is_some()
172 && let Ok(store) = RuntimeStore::open_default()
173 {
174 let proposed_action = action
175 .get("tool")
176 .and_then(|value| value.as_str())
177 .unwrap_or("restored action")
178 .to_string();
179 let pending_action_json = serde_json::to_string(action).ok();
180 if let Ok(approval) = store.approvals().create(NewApproval {
181 task_id: manifest.task_id.clone(),
182 proposed_action: format!("restore replay: {}", proposed_action),
183 risk_classification: "restored_action".to_string(),
184 policy_decision: "ask".to_string(),
185 args_summary: pending_action_json.clone(),
186 checkpoint_id: Some(manifest.id.clone()),
187 pending_action_json,
188 }) {
189 let _ = store.checkpoints().set_approval(&manifest.id, &approval.id);
190 }
191 }
192 Ok(manifest)
193}
194
195fn sanitize_relpath(path: &str) -> String {
196 path.split(std::path::MAIN_SEPARATOR)
197 .flat_map(|part| part.split('/'))
198 .filter(|part| !part.is_empty() && *part != "." && *part != "..")
199 .collect::<Vec<_>>()
200 .join("__")
201}
202
203struct ShadowGitSnapshot {
204 repo: String,
205 commit: String,
206}
207
208fn snapshot_shadow_git(
209 project_root: &Path,
210 files: &[CheckpointFile],
211 checkpoint_id: &str,
212) -> Result<ShadowGitSnapshot> {
213 let repo_root = data_dir()?
214 .join("shadow-git")
215 .join(project_hash(project_root));
216 let worktree = repo_root.join("worktree");
217 std::fs::create_dir_all(&worktree)?;
218 if !worktree.join(".git").exists() {
219 run_git(&worktree, ["init"])?;
220 }
221
222 run_git(&worktree, ["config", "user.name", "Mermaid Checkpoints"])?;
223 run_git(
224 &worktree,
225 ["config", "user.email", "mermaid-checkpoints@localhost"],
226 )?;
227
228 for file in files {
229 let shadow_path = worktree.join(&file.path);
230 let project_path = project_root.join(&file.path);
231 if file.existed && project_path.is_file() {
232 if let Some(parent) = shadow_path.parent() {
233 std::fs::create_dir_all(parent)?;
234 }
235 std::fs::copy(&project_path, &shadow_path).with_context(|| {
236 format!(
237 "failed to update shadow checkpoint {} -> {}",
238 project_path.display(),
239 shadow_path.display()
240 )
241 })?;
242 } else if shadow_path.exists() {
243 if shadow_path.is_dir() {
244 std::fs::remove_dir_all(&shadow_path)?;
245 } else {
246 std::fs::remove_file(&shadow_path)?;
247 }
248 }
249 }
250
251 run_git(&worktree, ["add", "-A"])?;
252 let status = Command::new("git")
253 .arg("diff")
254 .arg("--cached")
255 .arg("--quiet")
256 .current_dir(&worktree)
257 .status()?;
258 if !status.success() {
259 run_git_with_env(
260 &worktree,
261 ["commit", "-m", &format!("checkpoint {checkpoint_id}")],
262 )?;
263 }
264 let commit =
265 git_output(&worktree, ["rev-parse", "HEAD"]).unwrap_or_else(|_| "uncommitted".to_string());
266 Ok(ShadowGitSnapshot {
267 repo: worktree.display().to_string(),
268 commit,
269 })
270}
271
272fn run_git<const N: usize>(cwd: &Path, args: [&str; N]) -> Result<()> {
273 run_git_with_env(cwd, args)
274}
275
276fn run_git_with_env<const N: usize>(cwd: &Path, args: [&str; N]) -> Result<()> {
277 let status = Command::new("git")
278 .args(args)
279 .current_dir(cwd)
280 .env("GIT_AUTHOR_NAME", "Mermaid Checkpoints")
281 .env("GIT_AUTHOR_EMAIL", "mermaid-checkpoints@localhost")
282 .env("GIT_COMMITTER_NAME", "Mermaid Checkpoints")
283 .env("GIT_COMMITTER_EMAIL", "mermaid-checkpoints@localhost")
284 .stdout(Stdio::null())
285 .stderr(Stdio::null())
286 .status()?;
287 anyhow::ensure!(
288 status.success(),
289 "shadow git command failed in {}",
290 cwd.display()
291 );
292 Ok(())
293}
294
295fn git_output<const N: usize>(cwd: &Path, args: [&str; N]) -> Result<String> {
296 let output = Command::new("git").args(args).current_dir(cwd).output()?;
297 anyhow::ensure!(output.status.success(), "shadow git command failed");
298 Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
299}
300
301fn project_hash(path: &Path) -> String {
302 let mut hasher = Sha256::new();
303 hasher.update(path.display().to_string().as_bytes());
304 hex_lower(&hasher.finalize())
305}
306
307fn hex_lower(bytes: &[u8]) -> String {
308 const HEX: &[u8; 16] = b"0123456789abcdef";
309 let mut out = String::with_capacity(bytes.len() * 2);
310 for byte in bytes {
311 out.push(HEX[(byte >> 4) as usize] as char);
312 out.push(HEX[(byte & 0x0f) as usize] as char);
313 }
314 out
315}
316
317fn fresh_checkpoint_id() -> String {
318 let nanos = std::time::SystemTime::now()
319 .duration_since(std::time::UNIX_EPOCH)
320 .map(|d| d.as_nanos())
321 .unwrap_or_default();
322 format!("checkpoint-{nanos:x}")
323}
324
325#[cfg(test)]
326mod tests {
327 use crate::*;
328
329 #[test]
330 fn checkpoint_restore_round_trips_file_and_created_file() {
331 let root = std::env::temp_dir().join("mermaid_checkpoint_test");
332 let _ = std::fs::remove_dir_all(&root);
333 std::fs::create_dir_all(&root).unwrap();
334 std::fs::write(root.join("a.txt"), "before").unwrap();
335 let manifest = create_checkpoint(
336 &root,
337 &[root.join("a.txt"), root.join("new.txt")],
338 Some(serde_json::json!({"tool": "write_file"})),
339 )
340 .unwrap();
341 std::fs::write(root.join("a.txt"), "after").unwrap();
342 std::fs::write(root.join("new.txt"), "created").unwrap();
343 let restored = restore_checkpoint(&manifest.id).unwrap();
344 assert_eq!(restored.id, manifest.id);
345 assert_eq!(
346 std::fs::read_to_string(root.join("a.txt")).unwrap(),
347 "before"
348 );
349 assert!(!root.join("new.txt").exists());
350 let _ = std::fs::remove_dir_all(&root);
351 }
352}