1use std::path::{Component, Path, PathBuf};
2use std::process::{Command, Stdio};
3
4use anyhow::{Context, Result};
5use serde::{Deserialize, Serialize};
6use sha2::{Digest, Sha256};
7
8use crate::pathguard::{contain_within, contain_within_canonical};
9use crate::{NewApproval, NewCheckpoint, RuntimeStore, data_dir};
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
12pub struct CheckpointFile {
13 pub path: String,
14 pub existed: bool,
15 pub snapshot_relpath: Option<String>,
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
19pub struct CheckpointManifest {
20 pub id: String,
21 #[serde(default)]
22 pub task_id: Option<String>,
23 pub project_path: String,
24 pub files: Vec<CheckpointFile>,
25 pub pending_action: Option<serde_json::Value>,
26 #[serde(default)]
27 pub shadow_git_repo: Option<String>,
28 #[serde(default)]
29 pub shadow_git_commit: Option<String>,
30 pub created_at: String,
31}
32
33pub fn create_checkpoint(
34 project_path: &Path,
35 paths: &[PathBuf],
36 pending_action: Option<serde_json::Value>,
37) -> Result<CheckpointManifest> {
38 create_checkpoint_for_task(project_path, paths, pending_action, None)
39}
40
41pub fn create_checkpoint_for_task(
42 project_path: &Path,
43 paths: &[PathBuf],
44 pending_action: Option<serde_json::Value>,
45 task_id: Option<String>,
46) -> Result<CheckpointManifest> {
47 let id = crate::storage::fresh_id("checkpoint");
50 let root = data_dir()?.join("checkpoints").join(&id);
51 let files_dir = root.join("files");
52 std::fs::create_dir_all(&files_dir)
53 .with_context(|| format!("failed to create checkpoint dir {}", files_dir.display()))?;
54
55 let project_root = std::fs::canonicalize(project_path).unwrap_or_else(|_| project_path.into());
56 let mut files = Vec::new();
57 for path in paths {
58 let candidate = if path.is_absolute() {
59 path.clone()
60 } else {
61 project_path.join(path)
62 };
63 let normalized = std::fs::canonicalize(&candidate).unwrap_or(candidate.clone());
64 let display = normalized
65 .strip_prefix(&project_root)
66 .unwrap_or(&normalized)
67 .display()
68 .to_string();
69 if normalized.exists() && normalized.is_file() {
70 let safe_rel = sanitize_relpath(&display);
71 let dest = files_dir.join(&safe_rel);
72 if let Some(parent) = dest.parent() {
73 std::fs::create_dir_all(parent)?;
74 }
75 std::fs::copy(&normalized, &dest).with_context(|| {
76 format!(
77 "failed to copy checkpoint file {} -> {}",
78 normalized.display(),
79 dest.display()
80 )
81 })?;
82 files.push(CheckpointFile {
83 path: display,
84 existed: true,
85 snapshot_relpath: Some(format!("files/{}", safe_rel)),
86 });
87 } else {
88 files.push(CheckpointFile {
89 path: display,
90 existed: false,
91 snapshot_relpath: None,
92 });
93 }
94 }
95
96 let shadow_git = snapshot_shadow_git(&project_root, &files, &id).ok();
97 let manifest = CheckpointManifest {
98 id: id.clone(),
99 task_id: task_id.clone(),
100 project_path: project_path.display().to_string(),
101 files,
102 pending_action,
103 shadow_git_repo: shadow_git.as_ref().map(|snapshot| snapshot.repo.clone()),
104 shadow_git_commit: shadow_git.as_ref().map(|snapshot| snapshot.commit.clone()),
105 created_at: chrono::Utc::now().to_rfc3339(),
106 };
107 let manifest_path = root.join("manifest.json");
108 crate::write_atomic(&manifest_path, &serde_json::to_vec_pretty(&manifest)?)?;
111
112 if let Ok(store) = RuntimeStore::open_default() {
113 if let Err(error) = store.checkpoints().create(NewCheckpoint {
117 id: Some(id.clone()),
118 task_id,
119 project_path: manifest.project_path.clone(),
120 snapshot_path: root.display().to_string(),
121 changed_files_json: serde_json::to_string(&manifest.files)?,
122 pending_action_json: manifest
123 .pending_action
124 .as_ref()
125 .map(serde_json::to_string)
126 .transpose()?,
127 approval_id: None,
128 }) {
129 let _ = std::fs::remove_dir_all(&root);
130 return Err(error)
131 .with_context(|| format!("failed to record checkpoint {id} in the runtime DB"));
132 }
133 }
134
135 let _ = crate::run_plugin_hooks(
136 "checkpoint",
137 &serde_json::json!({
138 "id": manifest.id.clone(),
139 "task_id": manifest.task_id.clone(),
140 "project_path": manifest.project_path.clone(),
141 "files": manifest.files.clone(),
142 "created_at": manifest.created_at.clone(),
143 }),
144 );
145
146 Ok(manifest)
147}
148
149pub fn restore_checkpoint(id: &str) -> Result<CheckpointManifest> {
150 let checkpoints_dir = data_dir()?.join("checkpoints");
153 let ckpt_dir = contain_within(&checkpoints_dir, id)
154 .with_context(|| format!("invalid checkpoint id: {id:?}"))?;
155 let manifest_path = ckpt_dir.join("manifest.json");
156 let raw = std::fs::read_to_string(&manifest_path)
157 .with_context(|| format!("failed to read {}", manifest_path.display()))?;
158 let manifest: CheckpointManifest = serde_json::from_str(&raw)?;
159 let project_root = resolve_restore_root(id, &manifest)?;
162
163 let mut writes: Vec<RestoreOp> = Vec::new();
171 let mut deletes: Vec<RestoreOp> = Vec::new();
172 for file in &manifest.files {
173 let target = match contain_within_canonical(&project_root, &file.path) {
179 Ok(target) => target,
180 Err(err) => {
181 tracing::warn!(
182 path = %file.path,
183 error = %err,
184 "skipping checkpoint entry that escapes the project root"
185 );
186 continue;
187 },
188 };
189 if file.existed {
190 let rel = file
191 .snapshot_relpath
192 .as_ref()
193 .context("checkpoint file missing snapshot_relpath")?;
194 let source = match contain_within(&ckpt_dir, rel) {
198 Ok(source) => source,
199 Err(err) => {
200 tracing::warn!(
201 relpath = %rel,
202 error = %err,
203 "skipping checkpoint entry with an escaping snapshot_relpath"
204 );
205 continue;
206 },
207 };
208 writes.push(RestoreOp::Write { target, source });
212 } else {
213 deletes.push(RestoreOp::Delete { target });
214 }
215 }
216
217 let staging = project_root.join(format!(
222 ".mermaid-restore.{}",
223 crate::storage::fresh_id("restore")
224 ));
225 std::fs::create_dir_all(&staging)
226 .with_context(|| format!("failed to create restore staging dir {}", staging.display()))?;
227
228 let mut applied: Vec<PriorState> = Vec::new();
229 if let Err(err) = apply_restore(&writes, &deletes, &staging, &mut applied) {
230 rollback_restore(&applied);
231 let _ = std::fs::remove_dir(&staging);
235 return Err(err.context(
236 "checkpoint restore failed; changes already applied were rolled back (best-effort)",
237 ));
238 }
239 let _ = std::fs::remove_dir_all(&staging);
241 if let Some(action) = manifest.pending_action.as_ref()
242 && action.get("tool").is_some()
243 && let Ok(store) = RuntimeStore::open_default()
244 {
245 let proposed_action = action
246 .get("tool")
247 .and_then(|value| value.as_str())
248 .unwrap_or("restored action")
249 .to_string();
250 let pending_action_json = serde_json::to_string(action).ok();
251 if let Ok(approval) = store.approvals().create(NewApproval {
252 task_id: manifest.task_id.clone(),
253 proposed_action: format!("restore replay: {}", proposed_action),
254 risk_classification: "restored_action".to_string(),
255 policy_decision: "ask".to_string(),
256 args_summary: pending_action_json.clone(),
257 checkpoint_id: Some(manifest.id.clone()),
258 pending_action_json,
259 }) {
260 let _ = store.checkpoints().set_approval(&manifest.id, &approval.id);
261 }
262 }
263 Ok(manifest)
264}
265
266enum RestoreOp {
272 Write { target: PathBuf, source: PathBuf },
273 Delete { target: PathBuf },
274}
275
276struct PriorState {
281 target: PathBuf,
282 staged: Option<PathBuf>,
285}
286
287fn stage_prior(target: &Path, staging: &Path, counter: &mut usize) -> Result<Option<PathBuf>> {
292 if !target.exists() {
293 return Ok(None);
294 }
295 let dest = staging.join(counter.to_string());
296 *counter += 1;
297 std::fs::rename(target, &dest)
298 .with_context(|| format!("failed to stage prior state of {}", target.display()))?;
299 Ok(Some(dest))
300}
301
302fn remove_path(path: &Path) {
306 match std::fs::symlink_metadata(path) {
307 Ok(meta) if meta.is_dir() => {
308 let _ = std::fs::remove_dir_all(path);
309 },
310 Ok(_) => {
311 let _ = std::fs::remove_file(path);
312 },
313 Err(_) => {},
314 }
315}
316
317fn apply_restore(
322 writes: &[RestoreOp],
323 deletes: &[RestoreOp],
324 staging: &Path,
325 applied: &mut Vec<PriorState>,
326) -> Result<()> {
327 let mut counter = 0usize;
328 for op in writes {
329 if let RestoreOp::Write { target, source } = op {
330 let bytes = std::fs::read(source).with_context(|| {
334 format!("failed to read checkpoint snapshot {}", source.display())
335 })?;
336 let staged = stage_prior(target, staging, &mut counter)?;
337 if let Some(parent) = target.parent() {
338 std::fs::create_dir_all(parent)?;
339 }
340 crate::write_atomic(target, &bytes).with_context(|| {
341 format!("failed to restore checkpoint file {}", target.display())
342 })?;
343 applied.push(PriorState {
344 target: target.clone(),
345 staged,
346 });
347 }
348 }
349 for op in deletes {
350 if let RestoreOp::Delete { target } = op
351 && target.exists()
352 {
353 let staged = stage_prior(target, staging, &mut counter)?;
356 applied.push(PriorState {
357 target: target.clone(),
358 staged,
359 });
360 }
361 }
362 Ok(())
363}
364
365fn rollback_restore(applied: &[PriorState]) {
370 for prior in applied.iter().rev() {
371 remove_path(&prior.target);
372 if let Some(staged) = &prior.staged {
373 if let Some(parent) = prior.target.parent() {
374 let _ = std::fs::create_dir_all(parent);
375 }
376 let _ = std::fs::rename(staged, &prior.target);
377 }
378 }
379}
380
381fn resolve_restore_root(id: &str, manifest: &CheckpointManifest) -> Result<PathBuf> {
388 let recorded = RuntimeStore::open_default()
389 .ok()
390 .and_then(|store| store.checkpoints().get(id).ok().flatten())
391 .map(|rec| rec.project_path);
392 let root_str = match recorded {
393 Some(db_path) => {
394 anyhow::ensure!(
395 db_path == manifest.project_path,
396 "checkpoint project_path does not match the recorded root (tampered manifest?)"
397 );
398 db_path
399 },
400 None => manifest.project_path.clone(),
401 };
402 let root = PathBuf::from(&root_str);
403 anyhow::ensure!(
404 root.is_absolute() && root.components().any(|c| matches!(c, Component::Normal(_))),
405 "unsafe checkpoint project root: {}",
406 root.display()
407 );
408 Ok(root)
409}
410
411fn sanitize_relpath(path: &str) -> String {
412 path.split(std::path::MAIN_SEPARATOR)
413 .flat_map(|part| part.split('/'))
414 .filter(|part| !part.is_empty() && *part != "." && *part != "..")
415 .collect::<Vec<_>>()
416 .join("__")
417}
418
419struct ShadowGitSnapshot {
420 repo: String,
421 commit: String,
422}
423
424fn snapshot_shadow_git(
425 project_root: &Path,
426 files: &[CheckpointFile],
427 checkpoint_id: &str,
428) -> Result<ShadowGitSnapshot> {
429 let repo_root = data_dir()?
430 .join("shadow-git")
431 .join(project_hash(project_root));
432 let worktree = repo_root.join("worktree");
433 std::fs::create_dir_all(&worktree)?;
434 if !worktree.join(".git").exists() {
435 run_git(&worktree, ["init"])?;
436 }
437
438 run_git(&worktree, ["config", "user.name", "Mermaid Checkpoints"])?;
439 run_git(
440 &worktree,
441 ["config", "user.email", "mermaid-checkpoints@localhost"],
442 )?;
443
444 for file in files {
445 let shadow_path = worktree.join(&file.path);
446 let project_path = project_root.join(&file.path);
447 if file.existed && project_path.is_file() {
448 if let Some(parent) = shadow_path.parent() {
449 std::fs::create_dir_all(parent)?;
450 }
451 std::fs::copy(&project_path, &shadow_path).with_context(|| {
452 format!(
453 "failed to update shadow checkpoint {} -> {}",
454 project_path.display(),
455 shadow_path.display()
456 )
457 })?;
458 } else if shadow_path.exists() {
459 if shadow_path.is_dir() {
460 std::fs::remove_dir_all(&shadow_path)?;
461 } else {
462 std::fs::remove_file(&shadow_path)?;
463 }
464 }
465 }
466
467 run_git(&worktree, ["add", "-A"])?;
468 let status = Command::new("git")
469 .args(HOOKS_OFF)
470 .arg("diff")
471 .arg("--cached")
472 .arg("--quiet")
473 .current_dir(&worktree)
474 .status()?;
475 if !status.success() {
476 run_git_with_env(
477 &worktree,
478 ["commit", "-m", &format!("checkpoint {checkpoint_id}")],
479 )?;
480 }
481 let commit =
482 git_output(&worktree, ["rev-parse", "HEAD"]).unwrap_or_else(|_| "uncommitted".to_string());
483 Ok(ShadowGitSnapshot {
484 repo: worktree.display().to_string(),
485 commit,
486 })
487}
488
489const HOOKS_OFF: [&str; 2] = ["-c", "core.hooksPath=/dev/null"];
495
496fn run_git<const N: usize>(cwd: &Path, args: [&str; N]) -> Result<()> {
497 run_git_with_env(cwd, args)
498}
499
500fn run_git_with_env<const N: usize>(cwd: &Path, args: [&str; N]) -> Result<()> {
501 let status = Command::new("git")
502 .args(HOOKS_OFF)
503 .args(args)
504 .current_dir(cwd)
505 .env("GIT_AUTHOR_NAME", "Mermaid Checkpoints")
506 .env("GIT_AUTHOR_EMAIL", "mermaid-checkpoints@localhost")
507 .env("GIT_COMMITTER_NAME", "Mermaid Checkpoints")
508 .env("GIT_COMMITTER_EMAIL", "mermaid-checkpoints@localhost")
509 .stdout(Stdio::null())
510 .stderr(Stdio::null())
511 .status()?;
512 anyhow::ensure!(
513 status.success(),
514 "shadow git command failed in {}",
515 cwd.display()
516 );
517 Ok(())
518}
519
520fn git_output<const N: usize>(cwd: &Path, args: [&str; N]) -> Result<String> {
521 let output = Command::new("git")
522 .args(HOOKS_OFF)
523 .args(args)
524 .current_dir(cwd)
525 .output()?;
526 anyhow::ensure!(output.status.success(), "shadow git command failed");
527 Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
528}
529
530fn project_hash(path: &Path) -> String {
531 let mut hasher = Sha256::new();
532 hasher.update(path.display().to_string().as_bytes());
533 hex_lower(&hasher.finalize())
534}
535
536fn hex_lower(bytes: &[u8]) -> String {
537 const HEX: &[u8; 16] = b"0123456789abcdef";
538 let mut out = String::with_capacity(bytes.len() * 2);
539 for byte in bytes {
540 out.push(HEX[(byte >> 4) as usize] as char);
541 out.push(HEX[(byte & 0x0f) as usize] as char);
542 }
543 out
544}
545
546pub fn gc_old_checkpoint_dirs(retention_days: i64) -> Result<usize> {
559 let dir = data_dir()?.join("checkpoints");
560 let Ok(entries) = std::fs::read_dir(&dir) else {
561 return Ok(0);
562 };
563 let cutoff = std::time::SystemTime::now()
564 .checked_sub(std::time::Duration::from_secs(
565 retention_days.max(0) as u64 * 86_400,
566 ))
567 .unwrap_or(std::time::UNIX_EPOCH);
568 let store = RuntimeStore::open_default().ok();
569 let mut removed = 0;
570 for entry in entries.flatten() {
571 let path = entry.path();
572 if !path.is_dir() {
573 continue;
574 }
575 let too_old = entry
576 .metadata()
577 .and_then(|m| m.modified())
578 .map(|mtime| mtime < cutoff)
579 .unwrap_or(false);
580 if too_old && std::fs::remove_dir_all(&path).is_ok() {
581 removed += 1;
582 if let Some(store) = store.as_ref()
585 && let Some(id) = path.file_name().and_then(|name| name.to_str())
586 && let Err(error) = store.checkpoints().delete(id)
587 {
588 tracing::warn!(
589 id,
590 error = %error,
591 "failed to delete DB row for a GC'd checkpoint dir"
592 );
593 }
594 }
595 }
596 Ok(removed)
597}
598
599#[cfg(test)]
600mod tests {
601 use crate::*;
602
603 #[test]
604 fn checkpoint_restore_round_trips_file_and_created_file() {
605 let root = std::env::temp_dir().join("mermaid_checkpoint_test");
606 let _ = std::fs::remove_dir_all(&root);
607 std::fs::create_dir_all(&root).unwrap();
608 std::fs::write(root.join("a.txt"), "before").unwrap();
609 let manifest = create_checkpoint(
610 &root,
611 &[root.join("a.txt"), root.join("new.txt")],
612 Some(serde_json::json!({"tool": "write_file"})),
613 )
614 .unwrap();
615 std::fs::write(root.join("a.txt"), "after").unwrap();
616 std::fs::write(root.join("new.txt"), "created").unwrap();
617 let restored = restore_checkpoint(&manifest.id).unwrap();
618 assert_eq!(restored.id, manifest.id);
619 assert_eq!(
620 std::fs::read_to_string(root.join("a.txt")).unwrap(),
621 "before"
622 );
623 assert!(!root.join("new.txt").exists());
624 let _ = std::fs::remove_dir_all(&root);
625 }
626
627 #[test]
628 fn restore_rejects_paths_escaping_project_root() {
629 let pid = std::process::id();
633 let root = std::env::temp_dir().join(format!("mermaid_ckpt_escape_{pid}"));
634 let _ = std::fs::remove_dir_all(&root);
635 std::fs::create_dir_all(&root).unwrap();
636 std::fs::write(root.join("a.txt"), "before").unwrap();
637
638 let manifest = create_checkpoint(&root, &[root.join("a.txt")], None).unwrap();
639
640 let outside = std::env::temp_dir().join(format!("mermaid_ckpt_outside_{pid}.txt"));
642 std::fs::write(&outside, "do not delete").unwrap();
643 let outside_name = outside.file_name().unwrap().to_string_lossy().to_string();
644
645 let manifest_path = data_dir()
646 .unwrap()
647 .join("checkpoints")
648 .join(&manifest.id)
649 .join("manifest.json");
650 let mut tampered: CheckpointManifest =
651 serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
652 tampered.files.push(CheckpointFile {
654 path: format!("../{outside_name}"),
655 existed: false,
656 snapshot_relpath: None,
657 });
658 tampered.files.push(CheckpointFile {
659 path: outside.display().to_string(),
660 existed: false,
661 snapshot_relpath: None,
662 });
663 std::fs::write(
664 &manifest_path,
665 serde_json::to_vec_pretty(&tampered).unwrap(),
666 )
667 .unwrap();
668
669 let _ = restore_checkpoint(&manifest.id).unwrap();
670
671 assert!(
672 outside.exists(),
673 "restore must not delete a file outside the project root"
674 );
675 assert_eq!(std::fs::read_to_string(&outside).unwrap(), "do not delete");
676
677 let _ = std::fs::remove_file(&outside);
678 let _ = std::fs::remove_dir_all(&root);
679 }
680
681 #[test]
682 fn restore_rejects_tampered_project_root() {
683 let pid = std::process::id();
688 let root = std::env::temp_dir().join(format!("mermaid_ckpt_root_{pid}"));
689 let _ = std::fs::remove_dir_all(&root);
690 std::fs::create_dir_all(&root).unwrap();
691 std::fs::write(root.join("a.txt"), "before").unwrap();
692 let manifest = create_checkpoint(&root, &[root.join("a.txt")], None).unwrap();
693
694 let outside = std::env::temp_dir().join(format!("mermaid_ckpt_root_outside_{pid}.txt"));
695 std::fs::write(&outside, "do not delete").unwrap();
696
697 let manifest_path = data_dir()
698 .unwrap()
699 .join("checkpoints")
700 .join(&manifest.id)
701 .join("manifest.json");
702 let mut tampered: CheckpointManifest =
703 serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
704 tampered.project_path = "/".to_string();
705 tampered.files.push(CheckpointFile {
706 path: outside.display().to_string(),
707 existed: false,
708 snapshot_relpath: None,
709 });
710 std::fs::write(
711 &manifest_path,
712 serde_json::to_vec_pretty(&tampered).unwrap(),
713 )
714 .unwrap();
715
716 assert!(
717 restore_checkpoint(&manifest.id).is_err(),
718 "restore must reject a tampered project_path"
719 );
720 assert!(outside.exists(), "restore must not delete an outside file");
721 assert_eq!(std::fs::read_to_string(&outside).unwrap(), "do not delete");
722
723 let _ = std::fs::remove_file(&outside);
724 let _ = std::fs::remove_dir_all(&root);
725 }
726
727 #[test]
728 fn mid_restore_failure_restores_nonempty_prior_directory() {
729 use super::{PriorState, RestoreOp, apply_restore, rollback_restore};
735
736 let pid = std::process::id();
737 let root = std::env::temp_dir().join(format!("mermaid_ckpt_dirroll_{pid}"));
738 let _ = std::fs::remove_dir_all(&root);
739 std::fs::create_dir_all(&root).unwrap();
740
741 let victim = root.join("victim");
744 std::fs::create_dir_all(victim.join("sub")).unwrap();
745 std::fs::write(victim.join("inner.txt"), "precious").unwrap();
746 std::fs::write(victim.join("sub").join("deep.txt"), "deep").unwrap();
747
748 let src = root.join("snapshot.bin");
750 std::fs::write(&src, "new-content").unwrap();
751
752 let staging = root.join(".staging");
753 std::fs::create_dir_all(&staging).unwrap();
754
755 let writes = vec![
756 RestoreOp::Write {
757 target: victim.clone(),
758 source: src.clone(),
759 },
760 RestoreOp::Write {
763 target: root.join("other.txt"),
764 source: root.join("does-not-exist.bin"),
765 },
766 ];
767 let deletes: Vec<RestoreOp> = Vec::new();
768
769 let mut applied: Vec<PriorState> = Vec::new();
770 let result = apply_restore(&writes, &deletes, &staging, &mut applied);
771 assert!(
772 result.is_err(),
773 "a missing snapshot source must fail the restore"
774 );
775
776 rollback_restore(&applied);
777
778 assert!(victim.is_dir(), "prior directory subtree must be restored");
780 assert_eq!(
781 std::fs::read_to_string(victim.join("inner.txt")).unwrap(),
782 "precious"
783 );
784 assert_eq!(
785 std::fs::read_to_string(victim.join("sub").join("deep.txt")).unwrap(),
786 "deep"
787 );
788 assert!(!root.join("other.txt").exists());
790
791 let _ = std::fs::remove_dir_all(&root);
792 }
793}