use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use supercode_harness::{
Agent, ChatMessage, ChatRequest, Config, FunctionCall, Provider, Role, ToolCall, Usage,
};
fn tool_call(id: &str, name: &str, args: serde_json::Value) -> ToolCall {
ToolCall {
id: id.to_string(),
kind: "function".to_string(),
function: FunctionCall {
name: name.to_string(),
arguments: args.to_string(),
},
}
}
fn assistant_with_call(call: ToolCall) -> ChatMessage {
ChatMessage {
role: Role::Assistant,
content: None,
content_parts: None,
tool_calls: Some(vec![call]),
tool_call_id: None,
name: None,
metadata: Default::default(),
}
}
fn tmp(tag: &str) -> std::path::PathBuf {
let dir = std::env::temp_dir().join(format!(
"supercode-checkpoint-engine-{tag}-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&dir).unwrap();
dir
}
struct ScriptedWrites {
calls: AtomicUsize,
}
#[async_trait]
impl Provider for ScriptedWrites {
async fn complete(
&self,
req: &ChatRequest,
_on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
) -> supercode_harness::Result<(ChatMessage, Usage)> {
match self.calls.fetch_add(1, Ordering::SeqCst) {
0 => Ok((
assistant_with_call(tool_call(
"call_1",
"write_file",
serde_json::json!({"path": "a.txt", "content": "hello"}),
)),
Usage::default(),
)),
1 => Ok((ChatMessage::assistant("created a.txt"), Usage::default())),
2 => Ok((
assistant_with_call(tool_call(
"call_2",
"edit_file",
serde_json::json!({"path": "a.txt", "old_string": "hello", "new_string": "world"}),
)),
Usage::default(),
)),
3 => Ok((
assistant_with_call(tool_call(
"call_3",
"write_file",
serde_json::json!({"path": "b.txt", "content": "brand new"}),
)),
Usage::default(),
)),
4 => Ok((
ChatMessage::assistant("edited and added b.txt"),
Usage::default(),
)),
n => panic!("unexpected provider call {n}: {:?}", req.messages.last()),
}
}
}
async fn run_script(config: Config) -> Agent {
let mut agent = Agent::with_provider(
config,
Box::new(ScriptedWrites {
calls: AtomicUsize::new(0),
}),
);
let r1 = agent.send("create a.txt").await.unwrap();
assert_eq!(r1, "created a.txt");
let r2 = agent.send("edit it and add b.txt").await.unwrap();
assert_eq!(r2, "edited and added b.txt");
agent
}
fn history_signature(agent: &Agent) -> Vec<String> {
agent.history().iter().map(|m| format!("{m:?}")).collect()
}
#[tokio::test]
async fn checkpoint_off_by_default_touches_no_disk_and_is_a_true_no_op() {
let project = tmp("off-project");
let shadow = tmp("off-shadow");
std::fs::remove_dir_all(&shadow).unwrap();
let config = Config::builder()
.cwd(project.clone())
.checkpoint_dir(shadow.clone()) .build();
assert!(!config.checkpoint_enabled, "checkpoint must default off");
let agent = run_script(config).await;
assert!(
agent.checkpoint_observer().is_none(),
"no observer should exist when the module is off"
);
assert!(
!shadow.exists(),
"checkpoint_enabled=false must never create the shadow directory"
);
std::fs::remove_dir_all(&project).ok();
}
#[tokio::test]
async fn checkpoint_on_or_off_produces_byte_identical_conversation_history() {
let project_off = tmp("parity-off-project");
let project_on = tmp("parity-on-project");
let shadow_on = tmp("parity-on-shadow");
let config_off = Config::builder().cwd(project_off.clone()).build();
let config_on = Config::builder()
.cwd(project_on.clone())
.checkpoint_enabled(true)
.checkpoint_dir(shadow_on.clone())
.build();
let agent_off = run_script(config_off).await;
let agent_on = run_script(config_on).await;
assert_eq!(
history_signature(&agent_off),
history_signature(&agent_on),
"checkpoint on/off must never change the observable conversation"
);
assert!(agent_on.checkpoint_observer().is_some());
assert!(
shadow_on.exists(),
"checkpoint on DOES touch its own shadow dir"
);
std::fs::remove_dir_all(&project_off).ok();
std::fs::remove_dir_all(&project_on).ok();
std::fs::remove_dir_all(&shadow_on).ok();
}
#[tokio::test]
async fn enabled_checkpoint_captures_and_restore_undoes_a_whole_turn() {
let project = tmp("restore-project");
let shadow = tmp("restore-shadow");
let config = Config::builder()
.cwd(project.clone())
.checkpoint_enabled(true)
.checkpoint_dir(shadow.clone())
.build();
let agent = run_script(config).await;
assert_eq!(
std::fs::read_to_string(project.join("a.txt")).unwrap(),
"world"
);
assert_eq!(
std::fs::read_to_string(project.join("b.txt")).unwrap(),
"brand new"
);
let cp = agent.checkpoint_observer().expect("checkpoint is enabled");
let checkpoints = cp.list().unwrap();
assert_eq!(checkpoints.len(), 2, "one checkpoint per turn");
let turn2 = &checkpoints[0]; let turn1 = &checkpoints[1];
let mut diff1 = cp.turn_diff(&turn1.id).unwrap();
diff1.sort();
assert_eq!(diff1, vec!["a.txt".to_string()]);
let mut diff2 = cp.turn_diff(&turn2.id).unwrap();
diff2.sort();
assert_eq!(diff2, vec!["a.txt".to_string(), "b.txt".to_string()]);
let report = cp.restore(&turn2.id).unwrap();
assert!(report.refused.is_empty(), "{:?}", report.refused);
assert_eq!(
std::fs::read_to_string(project.join("a.txt")).unwrap(),
"hello"
);
assert!(!project.join("b.txt").exists());
let report1 = cp.restore(&turn1.id).unwrap();
assert!(report1.refused.is_empty(), "{:?}", report1.refused);
assert!(!project.join("a.txt").exists());
std::fs::remove_dir_all(&project).ok();
std::fs::remove_dir_all(&shadow).ok();
}
#[tokio::test]
async fn enabled_checkpoint_never_touches_the_projects_real_git() {
let project = tmp("realgit-project");
let shadow = tmp("realgit-shadow");
let git_dir = project.join(".git");
std::fs::create_dir_all(&git_dir).unwrap();
std::fs::write(git_dir.join("HEAD"), "ref: refs/heads/main\n").unwrap();
let head_before = std::fs::read(git_dir.join("HEAD")).unwrap();
let config = Config::builder()
.cwd(project.clone())
.checkpoint_enabled(true)
.checkpoint_dir(shadow.clone())
.build();
let agent = run_script(config).await;
assert_eq!(
std::fs::read(git_dir.join("HEAD")).unwrap(),
head_before,
"the project's real .git/HEAD must be byte-identical"
);
let mut git_entries: Vec<_> = std::fs::read_dir(&git_dir)
.unwrap()
.map(|e| e.unwrap().file_name())
.collect();
git_entries.sort();
assert_eq!(
git_entries,
vec![std::ffi::OsString::from("HEAD")],
"checkpoint must never write anything into the project's .git"
);
let cp = agent.checkpoint_observer().unwrap();
assert_eq!(cp.list().unwrap().len(), 2);
std::fs::remove_dir_all(&project).ok();
std::fs::remove_dir_all(&shadow).ok();
}
#[tokio::test]
async fn enabled_checkpoint_prunes_old_checkpoints_past_the_retain_bound() {
let project = tmp("prune-project");
let shadow = tmp("prune-shadow");
let config = Config::builder()
.cwd(project.clone())
.checkpoint_enabled(true)
.checkpoint_dir(shadow.clone())
.checkpoint_retain(1)
.build();
let agent = run_script(config).await; let cp = agent.checkpoint_observer().unwrap();
let checkpoints = cp.list().unwrap();
assert_eq!(
checkpoints.len(),
1,
"retain=1 must prune down to the single newest checkpoint"
);
std::fs::remove_dir_all(&project).ok();
std::fs::remove_dir_all(&shadow).ok();
}