1use std::path::{Component, Path, PathBuf};
2
3use anyhow::{Context, Result};
4use serde::{Deserialize, Serialize};
5use sha2::{Digest, Sha256};
6
7use crate::git::git;
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, Default, PartialEq, Eq)]
22pub struct CheckpointOrigin {
23 pub task_id: Option<String>,
25 pub session_id: Option<String>,
27 pub message_index: Option<i64>,
31}
32
33#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
34pub struct CheckpointManifest {
35 pub id: String,
36 #[serde(default)]
37 pub task_id: Option<String>,
38 #[serde(default)]
41 pub session_id: Option<String>,
42 #[serde(default)]
43 pub message_index: Option<i64>,
44 pub project_path: String,
45 pub files: Vec<CheckpointFile>,
46 pub pending_action: Option<serde_json::Value>,
47 #[serde(default)]
48 pub shadow_git_repo: Option<String>,
49 #[serde(default)]
50 pub shadow_git_commit: Option<String>,
51 pub created_at: String,
52}
53
54pub fn create_checkpoint(
55 project_path: &Path,
56 paths: &[PathBuf],
57 pending_action: Option<serde_json::Value>,
58) -> Result<CheckpointManifest> {
59 create_checkpoint_for_task(
60 project_path,
61 paths,
62 pending_action,
63 CheckpointOrigin::default(),
64 )
65}
66
67pub fn create_checkpoint_for_task(
68 project_path: &Path,
69 paths: &[PathBuf],
70 pending_action: Option<serde_json::Value>,
71 origin: CheckpointOrigin,
72) -> Result<CheckpointManifest> {
73 let id = crate::storage::fresh_id("checkpoint");
76 let root = data_dir()?.join("checkpoints").join(&id);
77 let files_dir = root.join("files");
78 std::fs::create_dir_all(&files_dir)
79 .with_context(|| format!("failed to create checkpoint dir {}", files_dir.display()))?;
80
81 let project_root = std::fs::canonicalize(project_path).unwrap_or_else(|_| project_path.into());
82 let mut files = Vec::new();
83 for path in paths {
84 let candidate = if path.is_absolute() {
85 path.clone()
86 } else {
87 project_path.join(path)
88 };
89 let normalized = std::fs::canonicalize(&candidate).unwrap_or(candidate.clone());
90 let display = normalized
91 .strip_prefix(&project_root)
92 .unwrap_or(&normalized)
93 .display()
94 .to_string();
95 if normalized.exists() && normalized.is_file() {
96 let safe_rel = sanitize_relpath(&display);
97 let dest = files_dir.join(&safe_rel);
98 if let Some(parent) = dest.parent() {
99 std::fs::create_dir_all(parent)?;
100 }
101 std::fs::copy(&normalized, &dest).with_context(|| {
102 format!(
103 "failed to copy checkpoint file {} -> {}",
104 normalized.display(),
105 dest.display()
106 )
107 })?;
108 files.push(CheckpointFile {
109 path: display,
110 existed: true,
111 snapshot_relpath: Some(format!("files/{}", safe_rel)),
112 });
113 } else {
114 files.push(CheckpointFile {
115 path: display,
116 existed: false,
117 snapshot_relpath: None,
118 });
119 }
120 }
121
122 let shadow_git = snapshot_shadow_git(&project_root, &files, &id).ok();
123 let manifest = CheckpointManifest {
124 id: id.clone(),
125 task_id: origin.task_id.clone(),
126 session_id: origin.session_id.clone(),
127 message_index: origin.message_index,
128 project_path: project_path.display().to_string(),
129 files,
130 pending_action,
131 shadow_git_repo: shadow_git.as_ref().map(|snapshot| snapshot.repo.clone()),
132 shadow_git_commit: shadow_git.as_ref().map(|snapshot| snapshot.commit.clone()),
133 created_at: chrono::Utc::now().to_rfc3339(),
134 };
135 let manifest_path = root.join("manifest.json");
136 crate::write_atomic(&manifest_path, &serde_json::to_vec_pretty(&manifest)?)?;
139
140 if let Ok(store) = RuntimeStore::open_default() {
141 if let Err(error) = store.checkpoints().create(NewCheckpoint {
145 id: Some(id.clone()),
146 task_id: origin.task_id,
147 project_path: manifest.project_path.clone(),
148 snapshot_path: root.display().to_string(),
149 changed_files_json: serde_json::to_string(&manifest.files)?,
150 pending_action_json: manifest
151 .pending_action
152 .as_ref()
153 .map(serde_json::to_string)
154 .transpose()?,
155 approval_id: None,
156 session_id: manifest.session_id.clone(),
157 message_index: manifest.message_index,
158 }) {
159 let _ = std::fs::remove_dir_all(&root);
160 return Err(error)
161 .with_context(|| format!("failed to record checkpoint {id} in the runtime DB"));
162 }
163 }
164
165 let _ = crate::run_plugin_hooks(
166 "checkpoint",
167 &serde_json::json!({
168 "id": manifest.id.clone(),
169 "task_id": manifest.task_id.clone(),
170 "project_path": manifest.project_path.clone(),
171 "files": manifest.files.clone(),
172 "created_at": manifest.created_at.clone(),
173 }),
174 );
175
176 Ok(manifest)
177}
178
179pub fn restore_checkpoint(id: &str) -> Result<CheckpointManifest> {
180 let checkpoints_dir = data_dir()?.join("checkpoints");
183 let ckpt_dir = contain_within(&checkpoints_dir, id)
184 .with_context(|| format!("invalid checkpoint id: {id:?}"))?;
185 let manifest_path = ckpt_dir.join("manifest.json");
186 let raw = std::fs::read_to_string(&manifest_path)
187 .with_context(|| format!("failed to read {}", manifest_path.display()))?;
188 let manifest: CheckpointManifest = serde_json::from_str(&raw)?;
189 let project_root = resolve_restore_root(id, &manifest)?;
192
193 let mut writes: Vec<RestoreOp> = Vec::new();
201 let mut deletes: Vec<RestoreOp> = Vec::new();
202 for file in &manifest.files {
203 let target = match contain_within_canonical(&project_root, &file.path) {
209 Ok(target) => target,
210 Err(err) => {
211 tracing::warn!(
212 path = %file.path,
213 error = %err,
214 "skipping checkpoint entry that escapes the project root"
215 );
216 continue;
217 },
218 };
219 if file.existed {
220 let rel = file
221 .snapshot_relpath
222 .as_ref()
223 .context("checkpoint file missing snapshot_relpath")?;
224 let source = match contain_within(&ckpt_dir, rel) {
228 Ok(source) => source,
229 Err(err) => {
230 tracing::warn!(
231 relpath = %rel,
232 error = %err,
233 "skipping checkpoint entry with an escaping snapshot_relpath"
234 );
235 continue;
236 },
237 };
238 writes.push(RestoreOp::Write { target, source });
242 } else {
243 deletes.push(RestoreOp::Delete { target });
244 }
245 }
246
247 let staging = project_root.join(format!(
252 ".mermaid-restore.{}",
253 crate::storage::fresh_id("restore")
254 ));
255 std::fs::create_dir_all(&staging)
256 .with_context(|| format!("failed to create restore staging dir {}", staging.display()))?;
257
258 let mut applied: Vec<PriorState> = Vec::new();
259 if let Err(err) = apply_restore(&writes, &deletes, &staging, &mut applied) {
260 rollback_restore(&applied);
261 let _ = std::fs::remove_dir(&staging);
265 return Err(err.context(
266 "checkpoint restore failed; changes already applied were rolled back (best-effort)",
267 ));
268 }
269 let _ = std::fs::remove_dir_all(&staging);
271 if let Some(action) = manifest.pending_action.as_ref()
272 && action.get("tool").is_some()
273 && let Ok(store) = RuntimeStore::open_default()
274 {
275 let proposed_action = action
276 .get("tool")
277 .and_then(|value| value.as_str())
278 .unwrap_or("restored action")
279 .to_string();
280 let pending_action_json = serde_json::to_string(action).ok();
281 if let Ok(approval) = store.approvals().create(NewApproval {
282 task_id: manifest.task_id.clone(),
283 proposed_action: format!("restore replay: {}", proposed_action),
284 risk_classification: "restored_action".to_string(),
285 policy_decision: "ask".to_string(),
286 args_summary: pending_action_json.clone(),
287 checkpoint_id: Some(manifest.id.clone()),
288 pending_action_json,
289 }) {
290 let _ = store.checkpoints().set_approval(&manifest.id, &approval.id);
291 }
292 }
293 Ok(manifest)
294}
295
296enum RestoreOp {
302 Write { target: PathBuf, source: PathBuf },
303 Delete { target: PathBuf },
304}
305
306struct PriorState {
311 target: PathBuf,
312 staged: Option<PathBuf>,
315}
316
317fn stage_prior(target: &Path, staging: &Path, counter: &mut usize) -> Result<Option<PathBuf>> {
322 if !target.exists() {
323 return Ok(None);
324 }
325 let dest = staging.join(counter.to_string());
326 *counter += 1;
327 std::fs::rename(target, &dest)
328 .with_context(|| format!("failed to stage prior state of {}", target.display()))?;
329 Ok(Some(dest))
330}
331
332fn remove_path(path: &Path) {
336 match std::fs::symlink_metadata(path) {
337 Ok(meta) if meta.is_dir() => {
338 let _ = std::fs::remove_dir_all(path);
339 },
340 Ok(_) => {
341 let _ = std::fs::remove_file(path);
342 },
343 Err(_) => {},
344 }
345}
346
347fn apply_restore(
352 writes: &[RestoreOp],
353 deletes: &[RestoreOp],
354 staging: &Path,
355 applied: &mut Vec<PriorState>,
356) -> Result<()> {
357 let mut counter = 0usize;
358 for op in writes {
359 if let RestoreOp::Write { target, source } = op {
360 let bytes = std::fs::read(source).with_context(|| {
364 format!("failed to read checkpoint snapshot {}", source.display())
365 })?;
366 let staged = stage_prior(target, staging, &mut counter)?;
367 if let Some(parent) = target.parent() {
368 std::fs::create_dir_all(parent)?;
369 }
370 crate::write_atomic(target, &bytes).with_context(|| {
371 format!("failed to restore checkpoint file {}", target.display())
372 })?;
373 applied.push(PriorState {
374 target: target.clone(),
375 staged,
376 });
377 }
378 }
379 for op in deletes {
380 if let RestoreOp::Delete { target } = op
381 && target.exists()
382 {
383 let staged = stage_prior(target, staging, &mut counter)?;
386 applied.push(PriorState {
387 target: target.clone(),
388 staged,
389 });
390 }
391 }
392 Ok(())
393}
394
395fn rollback_restore(applied: &[PriorState]) {
400 for prior in applied.iter().rev() {
401 remove_path(&prior.target);
402 if let Some(staged) = &prior.staged {
403 if let Some(parent) = prior.target.parent() {
404 let _ = std::fs::create_dir_all(parent);
405 }
406 let _ = std::fs::rename(staged, &prior.target);
407 }
408 }
409}
410
411fn resolve_restore_root(id: &str, manifest: &CheckpointManifest) -> Result<PathBuf> {
418 let recorded = RuntimeStore::open_default()
419 .ok()
420 .and_then(|store| store.checkpoints().get(id).ok().flatten())
421 .map(|rec| rec.project_path);
422 let root_str = match recorded {
423 Some(db_path) => {
424 anyhow::ensure!(
425 db_path == manifest.project_path,
426 "checkpoint project_path does not match the recorded root (tampered manifest?)"
427 );
428 db_path
429 },
430 None => manifest.project_path.clone(),
431 };
432 let root = PathBuf::from(&root_str);
433 anyhow::ensure!(
434 root.is_absolute() && root.components().any(|c| matches!(c, Component::Normal(_))),
435 "unsafe checkpoint project root: {}",
436 root.display()
437 );
438 Ok(root)
439}
440
441fn sanitize_relpath(path: &str) -> String {
442 path.split(std::path::MAIN_SEPARATOR)
443 .flat_map(|part| part.split('/'))
444 .filter(|part| !part.is_empty() && *part != "." && *part != "..")
445 .collect::<Vec<_>>()
446 .join("__")
447}
448
449struct ShadowGitSnapshot {
450 repo: String,
451 commit: String,
452}
453
454fn snapshot_shadow_git(
455 project_root: &Path,
456 files: &[CheckpointFile],
457 checkpoint_id: &str,
458) -> Result<ShadowGitSnapshot> {
459 let repo_root = data_dir()?
460 .join("shadow-git")
461 .join(project_hash(project_root));
462 let worktree = repo_root.join("worktree");
463 std::fs::create_dir_all(&worktree)?;
464 if !worktree.join(".git").exists() {
465 git(&worktree).arg("init").run()?;
466 }
467
468 for file in files {
469 let rel = Path::new(&file.path);
482 if rel.is_absolute() || rel.components().any(|c| c == Component::ParentDir) {
483 continue;
484 }
485 let shadow_path = worktree.join(rel);
486 let project_path = project_root.join(rel);
487 if file.existed && project_path.is_file() {
488 if let Some(parent) = shadow_path.parent() {
489 std::fs::create_dir_all(parent)?;
490 }
491 std::fs::copy(&project_path, &shadow_path).with_context(|| {
492 format!(
493 "failed to update shadow checkpoint {} -> {}",
494 project_path.display(),
495 shadow_path.display()
496 )
497 })?;
498 } else if shadow_path.exists() {
499 if shadow_path.is_dir() {
500 std::fs::remove_dir_all(&shadow_path)?;
501 } else {
502 std::fs::remove_file(&shadow_path)?;
503 }
504 }
505 }
506
507 git(&worktree).args(["add", "-A"]).run()?;
508 if !git(&worktree)
511 .args(["diff", "--cached", "--quiet"])
512 .success()?
513 {
514 git(&worktree)
515 .args(["commit", "-m", &format!("checkpoint {checkpoint_id}")])
516 .run()?;
517 }
518 let commit = git(&worktree)
519 .args(["rev-parse", "HEAD"])
520 .output()
521 .unwrap_or_else(|_| "uncommitted".to_string());
522 Ok(ShadowGitSnapshot {
523 repo: worktree.display().to_string(),
524 commit,
525 })
526}
527
528pub(crate) fn project_hash(path: &Path) -> String {
529 let mut hasher = Sha256::new();
530 hasher.update(path.display().to_string().as_bytes());
531 crate::hex_lower(&hasher.finalize())
532}
533
534pub fn gc_old_checkpoint_dirs(retention_days: i64) -> Result<usize> {
547 let dir = data_dir()?.join("checkpoints");
548 let Ok(entries) = std::fs::read_dir(&dir) else {
549 return Ok(0);
550 };
551 let cutoff = std::time::SystemTime::now()
552 .checked_sub(std::time::Duration::from_secs(
553 retention_days.max(0) as u64 * 86_400,
554 ))
555 .unwrap_or(std::time::UNIX_EPOCH);
556 let store = RuntimeStore::open_default().ok();
557 let mut removed = 0;
558 for entry in entries.flatten() {
559 let path = entry.path();
560 if !path.is_dir() {
561 continue;
562 }
563 let too_old = entry
564 .metadata()
565 .and_then(|m| m.modified())
566 .map(|mtime| mtime < cutoff)
567 .unwrap_or(false);
568 if too_old && std::fs::remove_dir_all(&path).is_ok() {
569 removed += 1;
570 if let Some(store) = store.as_ref()
573 && let Some(id) = path.file_name().and_then(|name| name.to_str())
574 && let Err(error) = store.checkpoints().delete(id)
575 {
576 tracing::warn!(
577 id,
578 error = %error,
579 "failed to delete DB row for a GC'd checkpoint dir"
580 );
581 }
582 }
583 }
584 Ok(removed)
585}
586
587#[cfg(test)]
588mod tests {
589 use crate::*;
590
591 #[test]
592 fn checkpoint_restore_round_trips_file_and_created_file() {
593 let root = std::env::temp_dir().join("mermaid_checkpoint_test");
594 let _ = std::fs::remove_dir_all(&root);
595 std::fs::create_dir_all(&root).unwrap();
596 std::fs::write(root.join("a.txt"), "before").unwrap();
597 let manifest = create_checkpoint(
598 &root,
599 &[root.join("a.txt"), root.join("new.txt")],
600 Some(serde_json::json!({"tool": "write_file"})),
601 )
602 .unwrap();
603 std::fs::write(root.join("a.txt"), "after").unwrap();
604 std::fs::write(root.join("new.txt"), "created").unwrap();
605 let restored = restore_checkpoint(&manifest.id).unwrap();
606 assert_eq!(restored.id, manifest.id);
607 assert_eq!(
608 std::fs::read_to_string(root.join("a.txt")).unwrap(),
609 "before"
610 );
611 assert!(!root.join("new.txt").exists());
612 let _ = std::fs::remove_dir_all(&root);
613 }
614
615 #[test]
616 fn restore_rejects_paths_escaping_project_root() {
617 let pid = std::process::id();
621 let root = std::env::temp_dir().join(format!("mermaid_ckpt_escape_{pid}"));
622 let _ = std::fs::remove_dir_all(&root);
623 std::fs::create_dir_all(&root).unwrap();
624 std::fs::write(root.join("a.txt"), "before").unwrap();
625
626 let manifest = create_checkpoint(&root, &[root.join("a.txt")], None).unwrap();
627
628 let outside = std::env::temp_dir().join(format!("mermaid_ckpt_outside_{pid}.txt"));
630 std::fs::write(&outside, "do not delete").unwrap();
631 let outside_name = outside.file_name().unwrap().to_string_lossy().to_string();
632
633 let manifest_path = data_dir()
634 .unwrap()
635 .join("checkpoints")
636 .join(&manifest.id)
637 .join("manifest.json");
638 let mut tampered: CheckpointManifest =
639 serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
640 tampered.files.push(CheckpointFile {
642 path: format!("../{outside_name}"),
643 existed: false,
644 snapshot_relpath: None,
645 });
646 tampered.files.push(CheckpointFile {
647 path: outside.display().to_string(),
648 existed: false,
649 snapshot_relpath: None,
650 });
651 std::fs::write(
652 &manifest_path,
653 serde_json::to_vec_pretty(&tampered).unwrap(),
654 )
655 .unwrap();
656
657 let _ = restore_checkpoint(&manifest.id).unwrap();
658
659 assert!(
660 outside.exists(),
661 "restore must not delete a file outside the project root"
662 );
663 assert_eq!(std::fs::read_to_string(&outside).unwrap(), "do not delete");
664
665 let _ = std::fs::remove_file(&outside);
666 let _ = std::fs::remove_dir_all(&root);
667 }
668
669 #[test]
670 fn restore_rejects_tampered_project_root() {
671 let pid = std::process::id();
676 let root = std::env::temp_dir().join(format!("mermaid_ckpt_root_{pid}"));
677 let _ = std::fs::remove_dir_all(&root);
678 std::fs::create_dir_all(&root).unwrap();
679 std::fs::write(root.join("a.txt"), "before").unwrap();
680 let manifest = create_checkpoint(&root, &[root.join("a.txt")], None).unwrap();
681
682 let outside = std::env::temp_dir().join(format!("mermaid_ckpt_root_outside_{pid}.txt"));
683 std::fs::write(&outside, "do not delete").unwrap();
684
685 let manifest_path = data_dir()
686 .unwrap()
687 .join("checkpoints")
688 .join(&manifest.id)
689 .join("manifest.json");
690 let mut tampered: CheckpointManifest =
691 serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
692 tampered.project_path = "/".to_string();
693 tampered.files.push(CheckpointFile {
694 path: outside.display().to_string(),
695 existed: false,
696 snapshot_relpath: None,
697 });
698 std::fs::write(
699 &manifest_path,
700 serde_json::to_vec_pretty(&tampered).unwrap(),
701 )
702 .unwrap();
703
704 assert!(
705 restore_checkpoint(&manifest.id).is_err(),
706 "restore must reject a tampered project_path"
707 );
708 assert!(outside.exists(), "restore must not delete an outside file");
709 assert_eq!(std::fs::read_to_string(&outside).unwrap(), "do not delete");
710
711 let _ = std::fs::remove_file(&outside);
712 let _ = std::fs::remove_dir_all(&root);
713 }
714
715 #[test]
716 fn mid_restore_failure_restores_nonempty_prior_directory() {
717 use super::{PriorState, RestoreOp, apply_restore, rollback_restore};
723
724 let pid = std::process::id();
725 let root = std::env::temp_dir().join(format!("mermaid_ckpt_dirroll_{pid}"));
726 let _ = std::fs::remove_dir_all(&root);
727 std::fs::create_dir_all(&root).unwrap();
728
729 let victim = root.join("victim");
732 std::fs::create_dir_all(victim.join("sub")).unwrap();
733 std::fs::write(victim.join("inner.txt"), "precious").unwrap();
734 std::fs::write(victim.join("sub").join("deep.txt"), "deep").unwrap();
735
736 let src = root.join("snapshot.bin");
738 std::fs::write(&src, "new-content").unwrap();
739
740 let staging = root.join(".staging");
741 std::fs::create_dir_all(&staging).unwrap();
742
743 let writes = vec![
744 RestoreOp::Write {
745 target: victim.clone(),
746 source: src.clone(),
747 },
748 RestoreOp::Write {
751 target: root.join("other.txt"),
752 source: root.join("does-not-exist.bin"),
753 },
754 ];
755 let deletes: Vec<RestoreOp> = Vec::new();
756
757 let mut applied: Vec<PriorState> = Vec::new();
758 let result = apply_restore(&writes, &deletes, &staging, &mut applied);
759 assert!(
760 result.is_err(),
761 "a missing snapshot source must fail the restore"
762 );
763
764 rollback_restore(&applied);
765
766 assert!(victim.is_dir(), "prior directory subtree must be restored");
768 assert_eq!(
769 std::fs::read_to_string(victim.join("inner.txt")).unwrap(),
770 "precious"
771 );
772 assert_eq!(
773 std::fs::read_to_string(victim.join("sub").join("deep.txt")).unwrap(),
774 "deep"
775 );
776 assert!(!root.join("other.txt").exists());
778
779 let _ = std::fs::remove_dir_all(&root);
780 }
781
782 #[test]
783 fn shadow_git_ignores_absolute_paths_and_cannot_truncate_real_files() {
784 let tmp = std::env::temp_dir().join(format!(
789 "mermaid_shadow_abs_{}",
790 crate::storage::fresh_id("t")
791 ));
792 let project_root = tmp.join("project");
793 std::fs::create_dir_all(&project_root).unwrap();
794 let sentinel = tmp.join("outside.txt");
795 std::fs::write(&sentinel, "PRECIOUS").unwrap();
796
797 let files = vec![CheckpointFile {
798 path: sentinel.display().to_string(), existed: true,
800 snapshot_relpath: None,
801 }];
802 let _ = super::snapshot_shadow_git(&project_root, &files, "test-cp");
805 assert_eq!(
806 std::fs::read_to_string(&sentinel).unwrap(),
807 "PRECIOUS",
808 "shadow-git sync must not truncate a real out-of-tree file",
809 );
810 let _ = std::fs::remove_dir_all(&tmp);
811 }
812}