use super::store::CheckpointStore;
use crate::config::McPaths;
use serde::{Deserialize, Serialize};
use std::{
fs,
io::ErrorKind,
path::{Path, PathBuf},
};
pub(crate) const CHECKPOINT_SCHEMA_VERSION: u64 = 1;
pub(crate) const MAX_SNAPSHOT_BYTES: u64 = 1024 * 1024;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct SnapshotBlobRef {
pub(crate) sha256: String,
pub(crate) bytes: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct FileSnapshotEvent {
pub(crate) session_id: String,
pub(crate) user_turn: u64,
pub(crate) tool: SnapshotTool,
pub(crate) cwd: PathBuf,
pub(crate) relative_path: PathBuf,
pub(crate) pre: Option<SnapshotBlobRef>,
pub(crate) post: Option<SnapshotBlobRef>,
pub(crate) status: SnapshotCaptureStatus,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct SnapshotLedgerRecord {
pub(crate) schema_version: u64,
pub(crate) event: FileSnapshotEvent,
}
impl SnapshotLedgerRecord {
pub(crate) fn new(event: FileSnapshotEvent) -> Self {
Self {
schema_version: CHECKPOINT_SCHEMA_VERSION,
event,
}
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
pub(crate) enum SnapshotTool {
WriteFile,
Edit,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum SnapshotCaptureStatus {
Captured,
Excluded { reason: SnapshotExclusionReason },
Oversized { bytes: u64, max_bytes: u64 },
Unavailable { reason: String },
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum SnapshotExclusionReason {
PathEscape,
SecretPath,
CheckpointStorage,
SessionStorage,
GitStorage,
NotFile,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum SnapshotEligibility {
Eligible,
Excluded(SnapshotExclusionReason),
Oversized { bytes: u64, max_bytes: u64 },
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct SnapshotContext {
pub(crate) store: CheckpointStore,
pub(crate) paths: McPaths,
pub(crate) session_id: String,
pub(crate) user_turn: u64,
}
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub(crate) struct LedgerRead {
pub(crate) records: Vec<SnapshotLedgerRecord>,
pub(crate) diagnostics: Vec<String>,
}
pub(crate) fn classify_snapshot_eligibility(
path: &Path,
cwd_canonical: &Path,
paths: &McPaths,
) -> SnapshotEligibility {
if !path.starts_with(cwd_canonical) {
return SnapshotEligibility::Excluded(SnapshotExclusionReason::PathEscape);
}
if path.starts_with(&paths.checkpoints) {
return SnapshotEligibility::Excluded(SnapshotExclusionReason::CheckpointStorage);
}
if path.starts_with(&paths.sessions) {
return SnapshotEligibility::Excluded(SnapshotExclusionReason::SessionStorage);
}
if path == paths.auth_file
|| path == paths.settings_file
|| path == paths.project_settings_file
|| paths.local_settings_file.as_deref() == Some(path)
|| secret_like_path(path)
{
return SnapshotEligibility::Excluded(SnapshotExclusionReason::SecretPath);
}
if path
.components()
.any(|component| component.as_os_str() == ".git")
{
return SnapshotEligibility::Excluded(SnapshotExclusionReason::GitStorage);
}
let metadata = match fs::metadata(path) {
Ok(metadata) => metadata,
Err(error) if error.kind() == ErrorKind::NotFound => return SnapshotEligibility::Eligible,
Err(_) => return SnapshotEligibility::Excluded(SnapshotExclusionReason::NotFile),
};
if !metadata.is_file() {
return SnapshotEligibility::Excluded(SnapshotExclusionReason::NotFile);
}
if metadata.len() > MAX_SNAPSHOT_BYTES {
return SnapshotEligibility::Oversized {
bytes: metadata.len(),
max_bytes: MAX_SNAPSHOT_BYTES,
};
}
SnapshotEligibility::Eligible
}
pub(crate) fn capture_file_snapshot(
context: &SnapshotContext,
cwd_canonical: &Path,
tool: SnapshotTool,
path: &Path,
pre_bytes: Option<&[u8]>,
post_bytes: Option<&[u8]>,
) -> anyhow::Result<()> {
capture_file_snapshot_with_hook(
context,
cwd_canonical,
tool,
path,
pre_bytes,
post_bytes,
|| {},
)
}
pub(super) fn capture_file_snapshot_with_hook(
context: &SnapshotContext,
cwd_canonical: &Path,
tool: SnapshotTool,
path: &Path,
pre_bytes: Option<&[u8]>,
post_bytes: Option<&[u8]>,
after_blob_writes: impl FnOnce(),
) -> anyhow::Result<()> {
let Ok(relative_path) = sanitize_relative_path(path, cwd_canonical) else {
return Ok(());
};
let _mutation_lock = context.store.lock_mutations()?;
let event = match classify_snapshot_eligibility(path, cwd_canonical, &context.paths) {
SnapshotEligibility::Eligible => FileSnapshotEvent {
session_id: context.session_id.clone(),
user_turn: context.user_turn,
tool,
cwd: cwd_canonical.to_path_buf(),
relative_path,
pre: pre_bytes
.map(|bytes| context.store.write_blob(bytes))
.transpose()?,
post: post_bytes
.map(|bytes| context.store.write_blob(bytes))
.transpose()?,
status: SnapshotCaptureStatus::Captured,
},
SnapshotEligibility::Excluded(reason) => FileSnapshotEvent {
session_id: context.session_id.clone(),
user_turn: context.user_turn,
tool,
cwd: cwd_canonical.to_path_buf(),
relative_path,
pre: None,
post: None,
status: SnapshotCaptureStatus::Excluded { reason },
},
SnapshotEligibility::Oversized { bytes, max_bytes } => FileSnapshotEvent {
session_id: context.session_id.clone(),
user_turn: context.user_turn,
tool,
cwd: cwd_canonical.to_path_buf(),
relative_path,
pre: None,
post: None,
status: SnapshotCaptureStatus::Oversized { bytes, max_bytes },
},
};
after_blob_writes();
context
.store
.append_record(&SnapshotLedgerRecord::new(event))
}
pub(super) fn secret_like_path(path: &Path) -> bool {
path.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| {
name == "auth.json"
|| name == "id_rsa"
|| name == "id_ed25519"
|| name.starts_with(".env")
|| name.ends_with(".pem")
|| name.ends_with(".key")
})
}
fn sanitize_relative_path(path: &Path, cwd: &Path) -> anyhow::Result<PathBuf> {
let relative = path.strip_prefix(cwd).map_err(|_| {
anyhow::anyhow!("path '{}' escapes cwd '{}'", path.display(), cwd.display())
})?;
if relative.as_os_str().is_empty()
|| relative
.components()
.any(|component| matches!(component, std::path::Component::ParentDir))
{
anyhow::bail!("unsafe checkpoint path");
}
Ok(relative.to_path_buf())
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn context(temp: &TempDir) -> SnapshotContext {
let paths = McPaths::from_root(temp.path().join("mc"));
SnapshotContext {
store: CheckpointStore::from_paths(&paths),
paths,
session_id: "session".to_string(),
user_turn: 1,
}
}
#[test]
fn snapshot_excludes_auth_and_secret_paths() {
let temp = TempDir::new().unwrap();
let paths = McPaths::from_root(temp.path().join("mc"));
fs::create_dir_all(&paths.root).unwrap();
fs::write(&paths.auth_file, "secret").unwrap();
assert_eq!(
classify_snapshot_eligibility(&paths.auth_file, temp.path(), &paths),
SnapshotEligibility::Excluded(SnapshotExclusionReason::SecretPath)
);
let env_file = temp.path().join(".env");
fs::write(&env_file, "secret").unwrap();
assert_eq!(
classify_snapshot_eligibility(&env_file, temp.path(), &paths),
SnapshotEligibility::Excluded(SnapshotExclusionReason::SecretPath)
);
}
#[test]
fn snapshot_silently_skips_path_escape() {
let temp = TempDir::new().unwrap();
let outside = TempDir::new().unwrap();
let context = context(&temp);
let file = outside.path().join("file.txt");
fs::write(&file, "x").unwrap();
assert_eq!(
classify_snapshot_eligibility(&file, temp.path(), &context.paths),
SnapshotEligibility::Excluded(SnapshotExclusionReason::PathEscape)
);
capture_file_snapshot(
&context,
temp.path(),
SnapshotTool::WriteFile,
&file,
Some(b"before"),
Some(b"after"),
)
.unwrap();
let read = context.store.read_records("session");
assert!(read.records.is_empty());
assert!(read.diagnostics.is_empty());
assert!(!context.store.blobs_dir().exists());
}
#[test]
fn snapshot_records_oversized_as_non_rewindable() {
let temp = TempDir::new().unwrap();
let context = context(&temp);
let file = temp.path().join("big.txt");
fs::write(&file, vec![b'x'; MAX_SNAPSHOT_BYTES as usize + 1]).unwrap();
capture_file_snapshot(
&context,
temp.path(),
SnapshotTool::WriteFile,
&file,
None,
None,
)
.unwrap();
let read = context.store.read_records("session");
assert!(matches!(
read.records[0].event.status,
SnapshotCaptureStatus::Oversized { .. }
));
}
}