use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
use serde::{Deserialize, Serialize};
use crate::error::{Error, Result};
use crate::tools::WriteObserver;
pub const DEFAULT_RETAIN: usize = 50;
pub type CheckpointId = String;
fn now_ms() -> u128 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis())
.unwrap_or(0)
}
fn mint_id() -> CheckpointId {
use std::hash::BuildHasher;
let millis = now_ms();
let salt = std::collections::hash_map::RandomState::new().hash_one((
std::process::id(),
millis,
std::time::Instant::now(),
));
format!("{millis:020}-{:08x}", salt as u32)
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CheckpointFileEntry {
pub path: String,
pub blob: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointManifest {
pub id: CheckpointId,
pub created_at_ms: u128,
pub label: String,
pub files: Vec<CheckpointFileEntry>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointMeta {
pub id: CheckpointId,
pub created_at_ms: u128,
pub label: String,
pub file_count: usize,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RestoreReport {
pub restored: Vec<String>,
pub refused: Vec<(String, String)>,
}
#[derive(Debug)]
pub struct CheckpointStore {
root: PathBuf,
}
impl CheckpointStore {
pub fn open(root: impl Into<PathBuf>) -> Result<Self> {
let root = root.into();
std::fs::create_dir_all(root.join("objects"))?;
std::fs::create_dir_all(root.join("checkpoints"))?;
Ok(CheckpointStore { root })
}
fn objects_dir(&self) -> PathBuf {
self.root.join("objects")
}
fn checkpoints_dir(&self) -> PathBuf {
self.root.join("checkpoints")
}
fn manifest_path(&self, id: &str) -> PathBuf {
self.checkpoints_dir().join(format!("{id}.json"))
}
fn blob_path(&self, hash: &str) -> PathBuf {
let prefix = &hash[..hash.len().min(2)];
self.objects_dir().join(prefix).join(hash)
}
fn validate_id(id: &str) -> Result<()> {
let bad = id.is_empty()
|| id.contains('/')
|| id.contains('\\')
|| id.contains('\0')
|| id.split(['/', '\\']).any(|c| c == ".." || c == ".")
|| Path::new(id).is_absolute()
|| id.trim() != id;
if bad {
return Err(Error::Other(format!("invalid checkpoint id: `{id}`")));
}
Ok(())
}
fn write_blob(&self, content: &[u8]) -> Result<String> {
let hash = crate::reduce::content_hash(content);
let dest = self.blob_path(&hash);
if dest.exists() {
return Ok(hash);
}
let Some(parent) = dest.parent() else {
return Err(Error::Other("blob path has no parent".to_string()));
};
std::fs::create_dir_all(parent)?;
let tmp = parent.join(format!(".tmp-{}-{}", std::process::id(), mint_id()));
std::fs::write(&tmp, content)?;
match std::fs::rename(&tmp, &dest) {
Ok(()) => {}
Err(e) if dest.exists() => {
let _ = std::fs::remove_file(&tmp);
let _ = e;
}
Err(e) => return Err(e.into()),
}
Ok(hash)
}
fn read_blob(&self, hash: &str) -> Result<Vec<u8>> {
std::fs::read(self.blob_path(hash)).map_err(Into::into)
}
fn write_manifest(&self, m: &CheckpointManifest) -> Result<()> {
Self::validate_id(&m.id)?;
let dest = self.manifest_path(&m.id);
let json = serde_json::to_vec_pretty(m).map_err(|e| Error::Other(e.to_string()))?;
let dir = self.checkpoints_dir();
std::fs::create_dir_all(&dir)?;
let tmp = dir.join(format!(".tmp-{}-{}", std::process::id(), mint_id()));
std::fs::write(&tmp, &json)?;
std::fs::rename(&tmp, &dest)?;
Ok(())
}
pub fn create_checkpoint(&self, label: &str) -> Result<CheckpointId> {
let id = mint_id();
let manifest = CheckpointManifest {
id: id.clone(),
created_at_ms: now_ms(),
label: label.to_string(),
files: Vec::new(),
};
self.write_manifest(&manifest)?;
Ok(id)
}
pub fn record_pre_image(&self, id: &str, rel: &str, content: Option<Vec<u8>>) -> Result<()> {
let mut manifest = self.manifest(id)?;
if manifest.files.iter().any(|f| f.path == rel) {
return Ok(());
}
let blob = match content {
Some(bytes) => Some(self.write_blob(&bytes)?),
None => None,
};
manifest.files.push(CheckpointFileEntry {
path: rel.to_string(),
blob,
});
self.write_manifest(&manifest)
}
pub fn manifest(&self, id: &str) -> Result<CheckpointManifest> {
Self::validate_id(id)?;
let text = std::fs::read_to_string(self.manifest_path(id))
.map_err(|e| Error::Other(format!("checkpoint `{id}` not found: {e}")))?;
serde_json::from_str(&text)
.map_err(|e| Error::Other(format!("checkpoint `{id}` manifest is corrupt: {e}")))
}
pub fn list(&self) -> Result<Vec<CheckpointMeta>> {
let dir = self.checkpoints_dir();
let mut metas = Vec::new();
if !dir.exists() {
return Ok(metas);
}
for entry in std::fs::read_dir(&dir)? {
let entry = entry?;
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("json") {
continue;
}
let Ok(text) = std::fs::read_to_string(&path) else {
continue;
};
if let Ok(m) = serde_json::from_str::<CheckpointManifest>(&text) {
metas.push(CheckpointMeta {
id: m.id,
created_at_ms: m.created_at_ms,
label: m.label,
file_count: m.files.len(),
});
}
}
metas.sort_by(|a, b| b.id.cmp(&a.id));
Ok(metas)
}
pub fn turn_diff(&self, id: &str) -> Result<Vec<String>> {
Ok(self
.manifest(id)?
.files
.into_iter()
.map(|f| f.path)
.collect())
}
pub fn turn_patch(&self, id: &str, project_root: &Path) -> Result<String> {
let mut entries = self.manifest(id)?.files;
entries.sort_by(|a, b| a.path.cmp(&b.path));
let mut out = String::new();
for entry in entries {
let before: Option<Vec<u8>> = match &entry.blob {
Some(blob) => Some(self.read_blob(blob)?),
None => None,
};
let after = std::fs::read(project_root.join(&entry.path)).ok();
if before == after {
continue;
}
let (Some(before_text), Some(after_text)) = (
decode_side(before.as_deref()),
decode_side(after.as_deref()),
) else {
out.push_str(&format!(
"Binary files a/{0} and b/{0} differ
",
entry.path
));
continue;
};
let patch = diffy::create_patch(&before_text, &after_text);
out.push_str(&format!(
"--- a/{0}
+++ b/{0}
",
entry.path
));
for line in patch.to_string().lines().skip(2) {
out.push_str(line);
out.push('\n');
}
}
Ok(out)
}
pub fn restore(
&self,
id: &str,
project_root: &Path,
protected_globs: &[String],
) -> Result<RestoreReport> {
let manifest = self.manifest(id)?;
let mut report = RestoreReport::default();
for entry in &manifest.files {
if let Some(reason) = reject_unsafe_manifest_path(&entry.path) {
report.refused.push((entry.path.clone(), reason));
continue;
}
let target = project_root.join(&entry.path);
let Some(normalized_rel) = normalized_project_rel(project_root, &target) else {
report.refused.push((
entry.path.clone(),
"refused: escapes the project root".to_string(),
));
continue;
};
if is_protected(&normalized_rel, protected_globs) {
report
.refused
.push((entry.path.clone(), "refused: protected path".to_string()));
continue;
}
if !contained(project_root, &target) {
report.refused.push((
entry.path.clone(),
"refused: escapes the project root".to_string(),
));
continue;
}
if let Some(resolved_rel) = resolved_project_rel(project_root, &target) {
if is_protected(&resolved_rel, protected_globs) {
report
.refused
.push((entry.path.clone(), "refused: protected path".to_string()));
continue;
}
}
match &entry.blob {
Some(hash) => {
let bytes = match self.read_blob(hash) {
Ok(b) => b,
Err(e) => {
report
.refused
.push((entry.path.clone(), format!("blob unreadable: {e}")));
continue;
}
};
if let Some(parent) = target.parent() {
let _ = std::fs::create_dir_all(parent);
}
if let Err(e) = std::fs::write(&target, &bytes) {
report
.refused
.push((entry.path.clone(), format!("write failed: {e}")));
continue;
}
}
None if target.exists() => {
if let Err(e) = std::fs::remove_file(&target) {
report
.refused
.push((entry.path.clone(), format!("delete failed: {e}")));
continue;
}
}
None => {}
}
report.restored.push(entry.path.clone());
}
Ok(report)
}
pub fn prune(&self, keep: usize) -> Result<usize> {
let mut metas = self.list()?; if metas.len() <= keep {
return Ok(0);
}
let stale = metas.split_off(keep);
let removed = stale.len();
for m in stale {
let _ = std::fs::remove_file(self.manifest_path(&m.id));
}
self.gc_unreferenced_blobs()?;
Ok(removed)
}
fn gc_unreferenced_blobs(&self) -> Result<()> {
let mut referenced: HashSet<String> = HashSet::new();
for meta in self.list()? {
if let Ok(m) = self.manifest(&meta.id) {
for f in m.files {
if let Some(b) = f.blob {
referenced.insert(b);
}
}
}
}
let objects = self.objects_dir();
if !objects.exists() {
return Ok(());
}
for entry in std::fs::read_dir(&objects)? {
let entry = entry?;
if !entry.file_type()?.is_dir() {
continue;
}
for inner in std::fs::read_dir(entry.path())? {
let inner = inner?;
let name = inner.file_name();
let Some(name) = name.to_str() else {
continue;
};
if name.starts_with(".tmp-") {
continue; }
if !referenced.contains(name) {
let _ = std::fs::remove_file(inner.path());
}
}
}
Ok(())
}
}
fn is_protected(rel: &str, extra_globs: &[String]) -> bool {
crate::safe_path::is_protected(rel, extra_globs)
}
fn reject_unsafe_manifest_path(rel: &str) -> Option<String> {
crate::safe_path::reject_unsafe_rel_path(rel)
}
fn normalized_project_rel(root: &Path, path: &Path) -> Option<String> {
crate::safe_path::normalized_project_rel(root, path)
}
fn contained(root: &Path, path: &Path) -> bool {
crate::safe_path::contained(root, path)
}
fn resolved_project_rel(root: &Path, path: &Path) -> Option<String> {
crate::safe_path::resolved_project_rel(root, path)
}
#[derive(Debug, Default)]
struct ObserverState {
current: Option<CheckpointId>,
captured: HashSet<String>,
}
#[derive(Debug)]
pub struct CheckpointObserver {
store: CheckpointStore,
project_root: PathBuf,
retain: usize,
protected: Vec<String>,
restore_enabled: bool,
state: Mutex<ObserverState>,
disabled: AtomicBool,
}
impl CheckpointObserver {
pub fn new(
store: CheckpointStore,
project_root: PathBuf,
retain: usize,
protected: Vec<String>,
) -> Self {
Self::with_restore(store, project_root, retain, protected, true)
}
pub fn with_restore(
store: CheckpointStore,
project_root: PathBuf,
retain: usize,
protected: Vec<String>,
restore_enabled: bool,
) -> Self {
CheckpointObserver {
store,
project_root,
retain: retain.max(1),
protected,
restore_enabled,
state: Mutex::new(ObserverState::default()),
disabled: AtomicBool::new(false),
}
}
pub fn store(&self) -> &CheckpointStore {
&self.store
}
pub fn list(&self) -> Result<Vec<CheckpointMeta>> {
self.store.list()
}
pub fn turn_diff(&self, id: &str) -> Result<Vec<String>> {
self.store.turn_diff(id)
}
pub fn turn_patch(&self, id: &str) -> Result<String> {
self.store.turn_patch(id, &self.project_root)
}
pub fn restore_enabled(&self) -> bool {
self.restore_enabled
}
pub fn restore(&self, id: &str) -> Result<RestoreReport> {
if !self.restore_enabled {
return Err(Error::Other(
"checkpoint restore is off for this harness (`[capabilities.checkpoint] restore = false`) — this preset tracks each turn's diff but has no code-restore surface"
.to_string(),
));
}
self.store.restore(id, &self.project_root, &self.protected)
}
pub fn is_disabled(&self) -> bool {
self.disabled.load(Ordering::Relaxed)
}
pub fn begin_turn(&self, label: &str) -> Option<CheckpointId> {
if self.disabled.load(Ordering::Relaxed) {
return None;
}
let short: String = label.chars().take(120).collect();
match self.store.create_checkpoint(&short) {
Ok(id) => {
if let Ok(mut st) = self.state.lock() {
st.current = Some(id.clone());
st.captured.clear();
}
if let Err(e) = self.store.prune(self.retain) {
eprintln!("warning: checkpoint: prune failed: {e}");
}
Some(id)
}
Err(e) => {
eprintln!(
"warning: checkpoint disabled for the rest of this session — \
failed to open a new checkpoint: {e}"
);
self.disabled.store(true, Ordering::Relaxed);
None
}
}
}
pub fn current(&self) -> Option<CheckpointId> {
self.state.lock().ok().and_then(|s| s.current.clone())
}
}
#[async_trait::async_trait]
impl WriteObserver for CheckpointObserver {
async fn before_write(&self, path: &Path) {
if self.disabled.load(Ordering::Relaxed) {
return;
}
if !contained(&self.project_root, path) {
return;
}
let Some(normalized) = crate::tools::normalize(path) else {
return;
};
let Some(root_normalized) = crate::tools::normalize(&self.project_root) else {
return;
};
let Ok(rel_path) = normalized.strip_prefix(&root_normalized) else {
return;
};
let rel = rel_path.to_string_lossy().replace('\\', "/");
if rel.is_empty() || rel == ".git" || rel.starts_with(".git/") {
return;
}
let mut st = match self.state.lock() {
Ok(g) => g,
Err(poisoned) => poisoned.into_inner(),
};
if st.current.is_none() {
drop(st);
self.begin_turn("untracked");
st = match self.state.lock() {
Ok(g) => g,
Err(poisoned) => poisoned.into_inner(),
};
}
if st.captured.contains(&rel) {
return;
}
let Some(id) = st.current.clone() else {
return; };
let content = std::fs::read(path).ok(); match self.store.record_pre_image(&id, &rel, content) {
Ok(()) => {
st.captured.insert(rel);
}
Err(e) => {
eprintln!(
"warning: checkpoint disabled for the rest of this session — \
failed to record a snapshot: {e}"
);
self.disabled.store(true, Ordering::Relaxed);
}
}
}
async fn after_write(&self, _path: &Path) -> Option<String> {
None
}
}
pub(crate) fn project_tag(cwd: &Path) -> String {
use std::hash::{Hash, Hasher};
let canon = std::fs::canonicalize(cwd).unwrap_or_else(|_| cwd.to_path_buf());
let mut h = std::collections::hash_map::DefaultHasher::new();
canon.hash(&mut h);
format!("{:016x}", h.finish())
}
fn default_shadow_root(cwd: &Path) -> PathBuf {
crate::agent::global_instructions_dir()
.join("checkpoints")
.join(project_tag(cwd))
}
fn decode_side(bytes: Option<&[u8]>) -> Option<String> {
match bytes {
None => Some(String::new()),
Some(b) => String::from_utf8(b.to_vec()).ok(),
}
}
pub fn observer_for_config(config: &crate::Config) -> Option<std::sync::Arc<CheckpointObserver>> {
if !config.checkpoint_enabled {
return None;
}
let root = config
.checkpoint_dir
.clone()
.unwrap_or_else(|| default_shadow_root(&config.cwd));
match CheckpointStore::open(&root) {
Ok(store) => Some(std::sync::Arc::new(CheckpointObserver::with_restore(
store,
config.cwd.clone(),
config.checkpoint_retain,
config.permissions_protected_paths.clone(),
config.checkpoint_restore,
))),
Err(e) => {
eprintln!(
"warning: [capabilities.checkpoint] is enabled but the shadow store at \
{} could not be opened — checkpoint is disabled for this session: {e}",
root.display()
);
None
}
}
}