use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use tandem_types::EngineEvent;
const MAX_INLINE_ROLLBACK_SNAPSHOT_BYTES: usize = 16 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MutationCheckpointOutcome {
Succeeded,
Failed,
Blocked,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MutationCheckpointSnapshotStatus {
InlineText,
TooLarge,
Binary,
NotNeeded,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MutationCheckpointRollbackSnapshot {
pub status: MutationCheckpointSnapshotStatus,
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub byte_count: Option<usize>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MutationCheckpointFileRecord {
pub path: String,
pub resolved_path: String,
pub existed_before: bool,
pub existed_after: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub before_hash: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub after_hash: Option<String>,
pub changed: bool,
pub rollback_snapshot: MutationCheckpointRollbackSnapshot,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MutationCheckpointRecord {
pub session_id: String,
pub message_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
pub tool: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub workspace_root: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub effective_cwd: Option<String>,
pub outcome: MutationCheckpointOutcome,
pub files: Vec<MutationCheckpointFileRecord>,
pub file_count: usize,
pub changed_file_count: usize,
}
#[derive(Debug, Clone)]
pub struct MutationCheckpointBaseline {
tool: String,
workspace_root: Option<String>,
effective_cwd: Option<String>,
files: Vec<MutationCheckpointFileBaseline>,
}
#[derive(Debug, Clone)]
struct MutationCheckpointFileBaseline {
path: String,
resolved_path: PathBuf,
existed_before: bool,
before_hash: Option<String>,
rollback_snapshot: MutationCheckpointRollbackSnapshot,
}
pub fn mutation_checkpoint_event(record: MutationCheckpointRecord) -> EngineEvent {
EngineEvent::new(
"mutation.checkpoint.recorded",
json!({
"sessionID": record.session_id.clone(),
"messageID": record.message_id.clone(),
"tool": record.tool.clone(),
"record": record,
}),
)
}
pub fn prepare_mutation_checkpoint(tool: &str, args: &Value) -> Option<MutationCheckpointBaseline> {
if !is_mutating_tool(tool) {
return None;
}
let files = extract_mutation_target_paths(tool, args);
if files.is_empty() {
return None;
}
let workspace_root = string_field(args, "__workspace_root");
let effective_cwd = string_field(args, "__effective_cwd");
let files = files
.into_iter()
.map(|path| {
let resolved_path =
resolve_workspace_path(&path, effective_cwd.as_deref(), workspace_root.as_deref());
let (existed_before, before_hash, rollback_snapshot) =
snapshot_path_baseline(&resolved_path);
MutationCheckpointFileBaseline {
path,
resolved_path,
existed_before,
before_hash,
rollback_snapshot,
}
})
.collect::<Vec<_>>();
Some(MutationCheckpointBaseline {
tool: tool.to_string(),
workspace_root,
effective_cwd,
files,
})
}
pub fn finalize_mutation_checkpoint_record(
session_id: &str,
message_id: &str,
tool_call_id: Option<&str>,
baseline: &MutationCheckpointBaseline,
outcome: MutationCheckpointOutcome,
) -> MutationCheckpointRecord {
let files = baseline
.files
.iter()
.map(|file| {
let (existed_after, after_hash) = snapshot_path_state(&file.resolved_path);
let changed = file.existed_before != existed_after || file.before_hash != after_hash;
MutationCheckpointFileRecord {
path: file.path.clone(),
resolved_path: file.resolved_path.to_string_lossy().to_string(),
existed_before: file.existed_before,
existed_after,
before_hash: file.before_hash.clone(),
after_hash,
changed,
rollback_snapshot: file.rollback_snapshot.clone(),
}
})
.collect::<Vec<_>>();
let changed_file_count = files.iter().filter(|file| file.changed).count();
MutationCheckpointRecord {
session_id: session_id.to_string(),
message_id: message_id.to_string(),
tool_call_id: tool_call_id.map(str::to_string),
tool: baseline.tool.clone(),
workspace_root: baseline.workspace_root.clone(),
effective_cwd: baseline.effective_cwd.clone(),
file_count: files.len(),
changed_file_count,
files,
outcome,
}
}
fn is_mutating_tool(tool: &str) -> bool {
matches!(
tool.trim(),
"write" | "edit" | "apply_patch" | "delete" | "delete_file"
)
}
fn extract_mutation_target_paths(tool: &str, args: &Value) -> Vec<String> {
let mut paths = match tool.trim() {
"write" | "edit" | "delete" | "delete_file" => string_fields(args, &["path", "file_path"]),
"apply_patch" => args
.get("patchText")
.and_then(Value::as_str)
.map(extract_apply_patch_paths)
.unwrap_or_default(),
_ => Vec::new(),
};
paths.sort();
paths.dedup();
paths
}
fn extract_apply_patch_paths(patch: &str) -> Vec<String> {
let mut paths = HashSet::new();
for line in patch.lines() {
let trimmed = line.trim();
let marker = trimmed
.strip_prefix("*** Add File: ")
.or_else(|| trimmed.strip_prefix("*** Update File: "))
.or_else(|| trimmed.strip_prefix("*** Delete File: "));
if let Some(path) = marker.map(str::trim).filter(|value| !value.is_empty()) {
paths.insert(path.to_string());
}
}
let mut paths = paths.into_iter().collect::<Vec<_>>();
paths.sort();
paths
}
fn resolve_workspace_path(
path: &str,
effective_cwd: Option<&str>,
workspace_root: Option<&str>,
) -> PathBuf {
let candidate = PathBuf::from(path);
if candidate.is_absolute() {
return candidate;
}
if let Some(cwd) = effective_cwd {
return PathBuf::from(cwd).join(path);
}
if let Some(root) = workspace_root {
return PathBuf::from(root).join(path);
}
candidate
}
fn snapshot_path_state(path: &Path) -> (bool, Option<String>) {
match read_path_bytes(path) {
Ok(bytes) => (true, Some(hash_bytes(&bytes))),
Err(_) => (false, None),
}
}
fn snapshot_path_baseline(
path: &Path,
) -> (bool, Option<String>, MutationCheckpointRollbackSnapshot) {
match read_path_bytes(path) {
Ok(bytes) => (
true,
Some(hash_bytes(&bytes)),
rollback_snapshot_from_bytes(&bytes),
),
Err(_) => (
false,
None,
MutationCheckpointRollbackSnapshot {
status: MutationCheckpointSnapshotStatus::NotNeeded,
content: None,
byte_count: None,
},
),
}
}
fn rollback_snapshot_from_bytes(bytes: &[u8]) -> MutationCheckpointRollbackSnapshot {
if bytes.len() > MAX_INLINE_ROLLBACK_SNAPSHOT_BYTES {
return MutationCheckpointRollbackSnapshot {
status: MutationCheckpointSnapshotStatus::TooLarge,
content: None,
byte_count: Some(bytes.len()),
};
}
match String::from_utf8(bytes.to_vec()) {
Ok(content) => MutationCheckpointRollbackSnapshot {
status: MutationCheckpointSnapshotStatus::InlineText,
content: Some(content),
byte_count: Some(bytes.len()),
},
Err(_) => MutationCheckpointRollbackSnapshot {
status: MutationCheckpointSnapshotStatus::Binary,
content: None,
byte_count: Some(bytes.len()),
},
}
}
fn read_path_bytes(path: &Path) -> std::io::Result<Vec<u8>> {
std::fs::read(path)
}
fn hash_bytes(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
format!("{:x}", hasher.finalize())
}
fn string_field(args: &Value, key: &str) -> Option<String> {
args.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
}
fn string_fields(args: &Value, keys: &[&str]) -> Vec<String> {
keys.iter()
.filter_map(|key| string_field(args, key))
.collect::<Vec<_>>()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn prepare_mutation_checkpoint_extracts_apply_patch_paths() {
let baseline = prepare_mutation_checkpoint(
"apply_patch",
&json!({
"patchText": "*** Begin Patch\n*** Update File: src/lib.rs\n@@\n-a\n+b\n*** Add File: src/new.rs\n+hello\n*** End Patch\n",
"__workspace_root": "/workspace"
}),
)
.expect("baseline");
let paths = baseline
.files
.iter()
.map(|file| file.path.as_str())
.collect::<Vec<_>>();
assert_eq!(paths, vec!["src/lib.rs", "src/new.rs"]);
}
#[test]
fn finalize_mutation_checkpoint_detects_file_creation() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("created.txt");
let baseline = prepare_mutation_checkpoint(
"write",
&json!({
"path": "created.txt",
"__effective_cwd": dir.path().to_string_lossy().to_string()
}),
)
.expect("baseline");
std::fs::write(&path, "hello world").expect("write created file");
let record = finalize_mutation_checkpoint_record(
"session-1",
"message-1",
Some("call-1"),
&baseline,
MutationCheckpointOutcome::Succeeded,
);
assert_eq!(record.file_count, 1);
assert_eq!(record.changed_file_count, 1);
assert!(record.files[0].changed);
assert!(!record.files[0].existed_before);
assert!(record.files[0].existed_after);
assert_eq!(
record.files[0].rollback_snapshot.status,
MutationCheckpointSnapshotStatus::NotNeeded
);
}
#[test]
fn finalize_mutation_checkpoint_captures_inline_before_snapshot_for_existing_file() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("tracked.txt");
std::fs::write(&path, "before").expect("write initial file");
let baseline = prepare_mutation_checkpoint(
"edit",
&json!({
"path": "tracked.txt",
"__effective_cwd": dir.path().to_string_lossy().to_string()
}),
)
.expect("baseline");
std::fs::write(&path, "after").expect("write updated file");
let record = finalize_mutation_checkpoint_record(
"session-1",
"message-1",
Some("call-1"),
&baseline,
MutationCheckpointOutcome::Succeeded,
);
assert_eq!(
record.files[0].rollback_snapshot.status,
MutationCheckpointSnapshotStatus::InlineText
);
assert_eq!(
record.files[0].rollback_snapshot.content.as_deref(),
Some("before")
);
}
}