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(
60 project_path: &Path,
61 paths: &[PathBuf],
62 pending_action: Option<serde_json::Value>,
63) -> Result<CheckpointManifest> {
64 create_checkpoint_for_task(
65 project_path,
66 paths,
67 pending_action,
68 CheckpointOrigin::default(),
69 )
70}
71
72pub fn create_checkpoint_for_task(
84 project_path: &Path,
85 paths: &[PathBuf],
86 pending_action: Option<serde_json::Value>,
87 origin: CheckpointOrigin,
88) -> Result<CheckpointManifest> {
89 let id = crate::storage::fresh_id("checkpoint");
92 let root = data_dir()?.join("checkpoints").join(&id);
93 let files_dir = root.join("files");
94 std::fs::create_dir_all(&files_dir)
95 .with_context(|| format!("failed to create checkpoint dir {}", files_dir.display()))?;
96
97 let project_root = std::fs::canonicalize(project_path).unwrap_or_else(|_| project_path.into());
98 let mut files = Vec::new();
99 for path in paths {
100 let candidate = if path.is_absolute() {
101 path.clone()
102 } else {
103 project_path.join(path)
104 };
105 let normalized = std::fs::canonicalize(&candidate).unwrap_or(candidate.clone());
106 let display = normalized
107 .strip_prefix(&project_root)
108 .unwrap_or(&normalized)
109 .display()
110 .to_string();
111 if normalized.exists() && normalized.is_file() {
112 let safe_rel = sanitize_relpath(&display);
113 let dest = files_dir.join(&safe_rel);
114 if let Some(parent) = dest.parent() {
115 std::fs::create_dir_all(parent)?;
116 }
117 std::fs::copy(&normalized, &dest).with_context(|| {
118 format!(
119 "failed to copy checkpoint file {} -> {}",
120 normalized.display(),
121 dest.display()
122 )
123 })?;
124 files.push(CheckpointFile {
125 path: display,
126 existed: true,
127 snapshot_relpath: Some(format!("files/{safe_rel}")),
128 });
129 } else {
130 files.push(CheckpointFile {
131 path: display,
132 existed: false,
133 snapshot_relpath: None,
134 });
135 }
136 }
137
138 let shadow_git = snapshot_shadow_git(&project_root, &files, &id).ok();
139 let manifest = CheckpointManifest {
140 id: id.clone(),
141 task_id: origin.task_id.clone(),
142 session_id: origin.session_id.clone(),
143 message_index: origin.message_index,
144 project_path: project_path.display().to_string(),
145 files,
146 pending_action,
147 shadow_git_repo: shadow_git.as_ref().map(|snapshot| snapshot.repo.clone()),
148 shadow_git_commit: shadow_git.as_ref().map(|snapshot| snapshot.commit.clone()),
149 created_at: chrono::Utc::now().to_rfc3339(),
150 };
151 let manifest_path = root.join("manifest.json");
152 crate::write_atomic(&manifest_path, &serde_json::to_vec_pretty(&manifest)?)?;
155
156 if let Ok(store) = RuntimeStore::open_default() {
157 if let Err(error) = store.checkpoints().create(NewCheckpoint {
161 id: Some(id.clone()),
162 task_id: origin.task_id,
163 project_path: manifest.project_path.clone(),
164 snapshot_path: root.display().to_string(),
165 changed_files_json: serde_json::to_string(&manifest.files)?,
166 pending_action_json: manifest
167 .pending_action
168 .as_ref()
169 .map(serde_json::to_string)
170 .transpose()?,
171 approval_id: None,
172 session_id: manifest.session_id.clone(),
173 message_index: manifest.message_index,
174 }) {
175 let _ = std::fs::remove_dir_all(&root);
176 return Err(error)
177 .with_context(|| format!("failed to record checkpoint {id} in the runtime DB"));
178 }
179 }
180
181 let _ = crate::run_plugin_hooks(
182 "checkpoint",
183 &serde_json::json!({
184 "id": manifest.id.clone(),
185 "task_id": manifest.task_id.clone(),
186 "project_path": manifest.project_path.clone(),
187 "files": manifest.files.clone(),
188 "created_at": manifest.created_at.clone(),
189 }),
190 );
191
192 Ok(manifest)
193}
194
195pub fn restore_checkpoint(id: &str) -> Result<CheckpointManifest> {
207 let checkpoints_dir = data_dir()?.join("checkpoints");
210 let ckpt_dir = contain_within(&checkpoints_dir, id)
211 .with_context(|| format!("invalid checkpoint id: {id:?}"))?;
212 let manifest_path = ckpt_dir.join("manifest.json");
213 let raw = std::fs::read_to_string(&manifest_path)
214 .with_context(|| format!("failed to read {}", manifest_path.display()))?;
215 let manifest: CheckpointManifest = serde_json::from_str(&raw)?;
216 let project_root = resolve_restore_root(id, &manifest)?;
219
220 let mut writes: Vec<RestoreOp> = Vec::new();
228 let mut deletes: Vec<RestoreOp> = Vec::new();
229 for file in &manifest.files {
230 let target = match contain_within_canonical(&project_root, &file.path) {
236 Ok(target) => target,
237 Err(err) => {
238 tracing::warn!(
239 path = %file.path,
240 error = %err,
241 "skipping checkpoint entry that escapes the project root"
242 );
243 continue;
244 },
245 };
246 if file.existed {
247 let rel = file
248 .snapshot_relpath
249 .as_ref()
250 .context("checkpoint file missing snapshot_relpath")?;
251 let source = match contain_within(&ckpt_dir, rel) {
255 Ok(source) => source,
256 Err(err) => {
257 tracing::warn!(
258 relpath = %rel,
259 error = %err,
260 "skipping checkpoint entry with an escaping snapshot_relpath"
261 );
262 continue;
263 },
264 };
265 writes.push(RestoreOp::Write { target, source });
269 } else {
270 deletes.push(RestoreOp::Delete { target });
271 }
272 }
273
274 let staging = project_root.join(format!(
279 ".mermaid-restore.{}",
280 crate::storage::fresh_id("restore")
281 ));
282 std::fs::create_dir_all(&staging)
283 .with_context(|| format!("failed to create restore staging dir {}", staging.display()))?;
284
285 let mut applied: Vec<PriorState> = Vec::new();
286 if let Err(err) = apply_restore(&writes, &deletes, &staging, &mut applied) {
287 rollback_restore(&applied);
288 let _ = std::fs::remove_dir(&staging);
292 return Err(err.context(
293 "checkpoint restore failed; changes already applied were rolled back (best-effort)",
294 ));
295 }
296 let _ = std::fs::remove_dir_all(&staging);
298 if let Some(action) = manifest.pending_action.as_ref()
299 && action.get("tool").is_some()
300 && let Ok(store) = RuntimeStore::open_default()
301 {
302 let proposed_action = action
303 .get("tool")
304 .and_then(|value| value.as_str())
305 .unwrap_or("restored action")
306 .to_string();
307 let pending_action_json = serde_json::to_string(action).ok();
308 if let Ok(approval) = store.approvals().create(NewApproval {
309 task_id: manifest.task_id.clone(),
310 proposed_action: format!("restore replay: {proposed_action}"),
311 risk_classification: "restored_action".to_string(),
312 policy_decision: "ask".to_string(),
313 args_summary: pending_action_json.clone(),
314 checkpoint_id: Some(manifest.id.clone()),
315 pending_action_json,
316 }) {
317 let _ = store.checkpoints().set_approval(&manifest.id, &approval.id);
318 }
319 }
320 Ok(manifest)
321}
322
323enum RestoreOp {
329 Write { target: PathBuf, source: PathBuf },
330 Delete { target: PathBuf },
331}
332
333struct PriorState {
338 target: PathBuf,
339 staged: Option<PathBuf>,
342}
343
344fn stage_prior(target: &Path, staging: &Path, counter: &mut usize) -> Result<Option<PathBuf>> {
349 if !target.exists() {
350 return Ok(None);
351 }
352 let dest = staging.join(counter.to_string());
353 *counter += 1;
354 std::fs::rename(target, &dest)
355 .with_context(|| format!("failed to stage prior state of {}", target.display()))?;
356 Ok(Some(dest))
357}
358
359fn remove_path(path: &Path) {
363 match std::fs::symlink_metadata(path) {
364 Ok(meta) if meta.is_dir() => {
365 let _ = std::fs::remove_dir_all(path);
366 },
367 Ok(_) => {
368 let _ = std::fs::remove_file(path);
369 },
370 Err(_) => {},
371 }
372}
373
374fn apply_restore(
379 writes: &[RestoreOp],
380 deletes: &[RestoreOp],
381 staging: &Path,
382 applied: &mut Vec<PriorState>,
383) -> Result<()> {
384 let mut counter = 0usize;
385 for op in writes {
386 if let RestoreOp::Write { target, source } = op {
387 let bytes = std::fs::read(source).with_context(|| {
391 format!("failed to read checkpoint snapshot {}", source.display())
392 })?;
393 let staged = stage_prior(target, staging, &mut counter)?;
394 if let Some(parent) = target.parent() {
395 std::fs::create_dir_all(parent)?;
396 }
397 crate::write_atomic(target, &bytes).with_context(|| {
398 format!("failed to restore checkpoint file {}", target.display())
399 })?;
400 applied.push(PriorState {
401 target: target.clone(),
402 staged,
403 });
404 }
405 }
406 for op in deletes {
407 if let RestoreOp::Delete { target } = op
408 && target.exists()
409 {
410 let staged = stage_prior(target, staging, &mut counter)?;
413 applied.push(PriorState {
414 target: target.clone(),
415 staged,
416 });
417 }
418 }
419 Ok(())
420}
421
422fn rollback_restore(applied: &[PriorState]) {
427 for prior in applied.iter().rev() {
428 remove_path(&prior.target);
429 if let Some(staged) = &prior.staged {
430 if let Some(parent) = prior.target.parent() {
431 let _ = std::fs::create_dir_all(parent);
432 }
433 let _ = std::fs::rename(staged, &prior.target);
434 }
435 }
436}
437
438fn resolve_restore_root(id: &str, manifest: &CheckpointManifest) -> Result<PathBuf> {
445 let recorded = RuntimeStore::open_default()
446 .ok()
447 .and_then(|store| store.checkpoints().get(id).ok().flatten())
448 .map(|rec| rec.project_path);
449 let root_str = match recorded {
450 Some(db_path) => {
451 anyhow::ensure!(
452 db_path == manifest.project_path,
453 "checkpoint project_path does not match the recorded root (tampered manifest?)"
454 );
455 db_path
456 },
457 None => manifest.project_path.clone(),
458 };
459 let root = PathBuf::from(&root_str);
460 anyhow::ensure!(
461 root.is_absolute() && root.components().any(|c| matches!(c, Component::Normal(_))),
462 "unsafe checkpoint project root: {}",
463 root.display()
464 );
465 Ok(root)
466}
467
468fn sanitize_relpath(path: &str) -> String {
469 path.split(std::path::MAIN_SEPARATOR)
470 .flat_map(|part| part.split('/'))
471 .filter(|part| !part.is_empty() && *part != "." && *part != "..")
472 .collect::<Vec<_>>()
473 .join("__")
474}
475
476struct ShadowGitSnapshot {
477 repo: String,
478 commit: String,
479}
480
481fn snapshot_shadow_git(
482 project_root: &Path,
483 files: &[CheckpointFile],
484 checkpoint_id: &str,
485) -> Result<ShadowGitSnapshot> {
486 let repo_root = data_dir()?
487 .join("shadow-git")
488 .join(project_hash(project_root));
489 let worktree = repo_root.join("worktree");
490 std::fs::create_dir_all(&worktree)?;
491 if !worktree.join(".git").exists() {
492 git(&worktree).arg("init").run()?;
493 }
494
495 for file in files {
496 let rel = Path::new(&file.path);
509 if rel.is_absolute() || rel.components().any(|c| c == Component::ParentDir) {
510 continue;
511 }
512 let shadow_path = worktree.join(rel);
513 let project_path = project_root.join(rel);
514 if file.existed && project_path.is_file() {
515 if let Some(parent) = shadow_path.parent() {
516 std::fs::create_dir_all(parent)?;
517 }
518 std::fs::copy(&project_path, &shadow_path).with_context(|| {
519 format!(
520 "failed to update shadow checkpoint {} -> {}",
521 project_path.display(),
522 shadow_path.display()
523 )
524 })?;
525 } else if shadow_path.exists() {
526 if shadow_path.is_dir() {
527 std::fs::remove_dir_all(&shadow_path)?;
528 } else {
529 std::fs::remove_file(&shadow_path)?;
530 }
531 }
532 }
533
534 git(&worktree).args(["add", "-A"]).run()?;
535 if !git(&worktree)
538 .args(["diff", "--cached", "--quiet"])
539 .success()?
540 {
541 git(&worktree)
542 .args(["commit", "-m", &format!("checkpoint {checkpoint_id}")])
543 .run()?;
544 }
545 let commit = git(&worktree)
546 .args(["rev-parse", "HEAD"])
547 .output()
548 .unwrap_or_else(|_| "uncommitted".to_string());
549 Ok(ShadowGitSnapshot {
550 repo: worktree.display().to_string(),
551 commit,
552 })
553}
554
555pub(crate) fn project_hash(path: &Path) -> String {
556 let mut hasher = Sha256::new();
557 hasher.update(path.display().to_string().as_bytes());
558 crate::hex_lower(&hasher.finalize())
559}
560
561pub fn gc_old_checkpoint_dirs(retention_days: i64) -> Result<usize> {
582 let dir = data_dir()?.join("checkpoints");
583 let Ok(entries) = std::fs::read_dir(&dir) else {
584 return Ok(0);
585 };
586 let cutoff = std::time::SystemTime::now()
587 .checked_sub(std::time::Duration::from_secs(
588 retention_days.max(0) as u64 * 86_400,
589 ))
590 .unwrap_or(std::time::UNIX_EPOCH);
591 let store = RuntimeStore::open_default().ok();
592 let mut removed = 0;
593 for entry in entries.flatten() {
594 let path = entry.path();
595 if !path.is_dir() {
596 continue;
597 }
598 let too_old = entry
599 .metadata()
600 .and_then(|m| m.modified())
601 .map(|mtime| mtime < cutoff)
602 .unwrap_or(false);
603 if too_old && std::fs::remove_dir_all(&path).is_ok() {
604 removed += 1;
605 if let Some(store) = store.as_ref()
608 && let Some(id) = path.file_name().and_then(|name| name.to_str())
609 && let Err(error) = store.checkpoints().delete(id)
610 {
611 tracing::warn!(
612 id,
613 error = %error,
614 "failed to delete DB row for a GC'd checkpoint dir"
615 );
616 }
617 }
618 }
619 Ok(removed)
620}
621
622#[cfg(test)]
623mod tests {
624 use crate::*;
625
626 #[test]
627 fn checkpoint_restore_round_trips_file_and_created_file() {
628 let root = std::env::temp_dir().join("mermaid_checkpoint_test");
629 let _ = std::fs::remove_dir_all(&root);
630 std::fs::create_dir_all(&root).unwrap();
631 std::fs::write(root.join("a.txt"), "before").unwrap();
632 let manifest = create_checkpoint(
633 &root,
634 &[root.join("a.txt"), root.join("new.txt")],
635 Some(serde_json::json!({"tool": "write_file"})),
636 )
637 .unwrap();
638 std::fs::write(root.join("a.txt"), "after").unwrap();
639 std::fs::write(root.join("new.txt"), "created").unwrap();
640 let restored = restore_checkpoint(&manifest.id).unwrap();
641 assert_eq!(restored.id, manifest.id);
642 assert_eq!(
643 std::fs::read_to_string(root.join("a.txt")).unwrap(),
644 "before"
645 );
646 assert!(!root.join("new.txt").exists());
647 let _ = std::fs::remove_dir_all(&root);
648 }
649
650 #[test]
651 fn restore_rejects_paths_escaping_project_root() {
652 let pid = std::process::id();
656 let root = std::env::temp_dir().join(format!("mermaid_ckpt_escape_{pid}"));
657 let _ = std::fs::remove_dir_all(&root);
658 std::fs::create_dir_all(&root).unwrap();
659 std::fs::write(root.join("a.txt"), "before").unwrap();
660
661 let manifest = create_checkpoint(&root, &[root.join("a.txt")], None).unwrap();
662
663 let outside = std::env::temp_dir().join(format!("mermaid_ckpt_outside_{pid}.txt"));
665 std::fs::write(&outside, "do not delete").unwrap();
666 let outside_name = outside.file_name().unwrap().to_string_lossy().to_string();
667
668 let manifest_path = data_dir()
669 .unwrap()
670 .join("checkpoints")
671 .join(&manifest.id)
672 .join("manifest.json");
673 let mut tampered: CheckpointManifest =
674 serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
675 tampered.files.push(CheckpointFile {
677 path: format!("../{outside_name}"),
678 existed: false,
679 snapshot_relpath: None,
680 });
681 tampered.files.push(CheckpointFile {
682 path: outside.display().to_string(),
683 existed: false,
684 snapshot_relpath: None,
685 });
686 std::fs::write(
687 &manifest_path,
688 serde_json::to_vec_pretty(&tampered).unwrap(),
689 )
690 .unwrap();
691
692 let _ = restore_checkpoint(&manifest.id).unwrap();
693
694 assert!(
695 outside.exists(),
696 "restore must not delete a file outside the project root"
697 );
698 assert_eq!(std::fs::read_to_string(&outside).unwrap(), "do not delete");
699
700 let _ = std::fs::remove_file(&outside);
701 let _ = std::fs::remove_dir_all(&root);
702 }
703
704 #[test]
705 fn restore_rejects_tampered_project_root() {
706 let pid = std::process::id();
711 let root = std::env::temp_dir().join(format!("mermaid_ckpt_root_{pid}"));
712 let _ = std::fs::remove_dir_all(&root);
713 std::fs::create_dir_all(&root).unwrap();
714 std::fs::write(root.join("a.txt"), "before").unwrap();
715 let manifest = create_checkpoint(&root, &[root.join("a.txt")], None).unwrap();
716
717 let outside = std::env::temp_dir().join(format!("mermaid_ckpt_root_outside_{pid}.txt"));
718 std::fs::write(&outside, "do not delete").unwrap();
719
720 let manifest_path = data_dir()
721 .unwrap()
722 .join("checkpoints")
723 .join(&manifest.id)
724 .join("manifest.json");
725 let mut tampered: CheckpointManifest =
726 serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap();
727 tampered.project_path = "/".to_string();
728 tampered.files.push(CheckpointFile {
729 path: outside.display().to_string(),
730 existed: false,
731 snapshot_relpath: None,
732 });
733 std::fs::write(
734 &manifest_path,
735 serde_json::to_vec_pretty(&tampered).unwrap(),
736 )
737 .unwrap();
738
739 assert!(
740 restore_checkpoint(&manifest.id).is_err(),
741 "restore must reject a tampered project_path"
742 );
743 assert!(outside.exists(), "restore must not delete an outside file");
744 assert_eq!(std::fs::read_to_string(&outside).unwrap(), "do not delete");
745
746 let _ = std::fs::remove_file(&outside);
747 let _ = std::fs::remove_dir_all(&root);
748 }
749
750 #[test]
751 fn mid_restore_failure_restores_nonempty_prior_directory() {
752 use super::{PriorState, RestoreOp, apply_restore, rollback_restore};
758
759 let pid = std::process::id();
760 let root = std::env::temp_dir().join(format!("mermaid_ckpt_dirroll_{pid}"));
761 let _ = std::fs::remove_dir_all(&root);
762 std::fs::create_dir_all(&root).unwrap();
763
764 let victim = root.join("victim");
767 std::fs::create_dir_all(victim.join("sub")).unwrap();
768 std::fs::write(victim.join("inner.txt"), "precious").unwrap();
769 std::fs::write(victim.join("sub").join("deep.txt"), "deep").unwrap();
770
771 let src = root.join("snapshot.bin");
773 std::fs::write(&src, "new-content").unwrap();
774
775 let staging = root.join(".staging");
776 std::fs::create_dir_all(&staging).unwrap();
777
778 let writes = vec![
779 RestoreOp::Write {
780 target: victim.clone(),
781 source: src.clone(),
782 },
783 RestoreOp::Write {
786 target: root.join("other.txt"),
787 source: root.join("does-not-exist.bin"),
788 },
789 ];
790 let deletes: Vec<RestoreOp> = Vec::new();
791
792 let mut applied: Vec<PriorState> = Vec::new();
793 let result = apply_restore(&writes, &deletes, &staging, &mut applied);
794 assert!(
795 result.is_err(),
796 "a missing snapshot source must fail the restore"
797 );
798
799 rollback_restore(&applied);
800
801 assert!(victim.is_dir(), "prior directory subtree must be restored");
803 assert_eq!(
804 std::fs::read_to_string(victim.join("inner.txt")).unwrap(),
805 "precious"
806 );
807 assert_eq!(
808 std::fs::read_to_string(victim.join("sub").join("deep.txt")).unwrap(),
809 "deep"
810 );
811 assert!(!root.join("other.txt").exists());
813
814 let _ = std::fs::remove_dir_all(&root);
815 }
816
817 #[test]
818 fn shadow_git_ignores_absolute_paths_and_cannot_truncate_real_files() {
819 let tmp = std::env::temp_dir().join(format!(
824 "mermaid_shadow_abs_{}",
825 crate::storage::fresh_id("t")
826 ));
827 let project_root = tmp.join("project");
828 std::fs::create_dir_all(&project_root).unwrap();
829 let sentinel = tmp.join("outside.txt");
830 std::fs::write(&sentinel, "PRECIOUS").unwrap();
831
832 let files = vec![CheckpointFile {
833 path: sentinel.display().to_string(), existed: true,
835 snapshot_relpath: None,
836 }];
837 let _ = super::snapshot_shadow_git(&project_root, &files, "test-cp");
840 assert_eq!(
841 std::fs::read_to_string(&sentinel).unwrap(),
842 "PRECIOUS",
843 "shadow-git sync must not truncate a real out-of-tree file",
844 );
845 let _ = std::fs::remove_dir_all(&tmp);
846 }
847}