use std::path::{Path, PathBuf};
use anyhow::{Context as _, Result};
use crate::utils::files::{copy_dir_recursive, copy_profile_files};
pub struct ProfileTransaction {
target_dir: PathBuf,
staging_dir: PathBuf,
original_dir: Option<PathBuf>,
staged: bool,
committed: bool,
}
impl ProfileTransaction {
pub fn new(target_dir: impl Into<PathBuf>) -> Result<Self> {
let target_dir = target_dir.into();
let parent = target_dir
.parent()
.unwrap_or_else(|| Path::new("."))
.to_path_buf();
std::fs::create_dir_all(&parent)
.with_context(|| format!("Failed to create parent directory: {}", parent.display()))?;
let staging_dir = unique_sibling_path(&parent, ".codex_txn_staging")?;
std::fs::create_dir_all(&staging_dir).with_context(|| {
format!(
"Failed to create staging directory: {}",
staging_dir.display()
)
})?;
Ok(Self {
target_dir,
staging_dir,
original_dir: None,
staged: false,
committed: false,
})
}
pub fn stage_profile(&mut self, src: &Path, files: &[&str]) -> Result<Vec<String>> {
let copied = copy_profile_files(src, &self.staging_dir, files)
.context("Failed to stage profile files")?;
self.staged = true;
Ok(copied)
}
#[allow(dead_code)]
pub fn stage_dir(&mut self, src: &Path) -> Result<()> {
copy_dir_recursive(src, &self.staging_dir).context("Failed to stage directory")?;
self.staged = true;
Ok(())
}
pub fn commit(&mut self) -> Result<()> {
if !self.staged {
anyhow::bail!("Cannot commit: no profile has been staged");
}
let parent = self.target_dir.parent().unwrap_or_else(|| Path::new("."));
if self.target_dir.exists() {
let orig_path = unique_sibling_path(parent, ".codex_txn_orig")?;
std::fs::rename(&self.target_dir, &orig_path).with_context(|| {
format!(
"Failed to move current profile aside: {}",
self.target_dir.display()
)
})?;
self.original_dir = Some(orig_path);
}
std::fs::rename(&self.staging_dir, &self.target_dir).with_context(|| {
format!(
"Failed to rename staging dir to target: {}",
self.target_dir.display()
)
})?;
self.committed = true;
Ok(())
}
pub fn rollback(&mut self) -> Result<()> {
if self.committed {
if self.target_dir.exists() {
std::fs::remove_dir_all(&self.target_dir)
.context("Failed to remove committed profile during rollback")?;
}
if let Some(ref orig) = self.original_dir.take()
&& orig.exists()
{
std::fs::rename(orig, &self.target_dir)
.context("Failed to restore original profile during rollback")?;
}
self.committed = false;
} else {
if self.staging_dir.exists() {
std::fs::remove_dir_all(&self.staging_dir)
.context("Failed to clean up staging directory during rollback")?;
}
self.staged = false;
}
Ok(())
}
pub fn cleanup_original(&self) -> Result<()> {
if let Some(ref orig) = self.original_dir
&& orig.exists()
{
std::fs::remove_dir_all(orig)
.context("Failed to remove original backup after commit")?;
}
Ok(())
}
#[must_use]
#[cfg(test)]
pub fn original_backup_path(&self) -> Option<&Path> {
self.original_dir.as_deref()
}
#[must_use]
pub fn staging_dir(&self) -> &Path {
&self.staging_dir
}
}
impl Drop for ProfileTransaction {
fn drop(&mut self) {
if !self.committed && self.staging_dir.exists() {
let _ = std::fs::remove_dir_all(&self.staging_dir);
}
}
}
#[allow(clippy::unnecessary_wraps)]
fn unique_sibling_path(parent: &Path, prefix: &str) -> Result<PathBuf> {
use std::time::{SystemTime, UNIX_EPOCH};
let ts = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| d.as_nanos());
let rand_suffix: u64 = rand::random::<u64>();
let path = parent.join(format!("{prefix}_{ts}_{rand_suffix:016x}"));
Ok(path)
}
#[cfg(test)]
mod tests {
use tempfile::TempDir;
use super::*;
fn make_profile(dir: &TempDir) -> PathBuf {
let profile = dir.path().join("profile");
std::fs::create_dir_all(&profile).unwrap();
std::fs::write(profile.join("auth.json"), r#"{"token":"test"}"#).unwrap();
std::fs::write(profile.join("config.toml"), "model = \"o4-mini\"").unwrap();
profile
}
#[test]
fn test_stage_and_commit_creates_target() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
let profile_src = make_profile(&tmp);
let mut txn = ProfileTransaction::new(&target).unwrap();
txn.stage_profile(&profile_src, &["auth.json", "config.toml"])
.unwrap();
txn.commit().unwrap();
assert!(target.exists(), "target should exist after commit");
assert!(target.join("auth.json").exists());
assert!(target.join("config.toml").exists());
}
#[test]
fn test_commit_clears_stale_files() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
std::fs::create_dir_all(&target).unwrap();
std::fs::write(target.join("stale.json"), "old").unwrap();
let profile_src = make_profile(&tmp);
let mut txn = ProfileTransaction::new(&target).unwrap();
txn.stage_profile(&profile_src, &["auth.json"]).unwrap();
txn.commit().unwrap();
assert!(target.join("auth.json").exists());
assert!(
!target.join("stale.json").exists(),
"stale file must be gone after atomic switch"
);
}
#[test]
fn test_rollback_before_commit_removes_staging() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
let profile_src = make_profile(&tmp);
let mut txn = ProfileTransaction::new(&target).unwrap();
let staging = txn.staging_dir().to_path_buf();
txn.stage_profile(&profile_src, &["auth.json"]).unwrap();
assert!(staging.exists(), "staging should exist before rollback");
txn.rollback().unwrap();
assert!(!staging.exists(), "staging should be gone after rollback");
assert!(!target.exists(), "target should not have been created");
}
#[test]
fn test_rollback_after_commit_restores_original() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
std::fs::create_dir_all(&target).unwrap();
std::fs::write(target.join("auth.json"), r#"{"token":"original"}"#).unwrap();
let profile_src = make_profile(&tmp);
let mut txn = ProfileTransaction::new(&target).unwrap();
txn.stage_profile(&profile_src, &["auth.json"]).unwrap();
txn.commit().unwrap();
let committed = std::fs::read_to_string(target.join("auth.json")).unwrap();
assert!(committed.contains("test"));
txn.rollback().unwrap();
let restored = std::fs::read_to_string(target.join("auth.json")).unwrap();
assert!(
restored.contains("original"),
"original content should be restored after rollback"
);
}
#[test]
fn test_rollback_after_commit_no_original() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
let profile_src = make_profile(&tmp);
let mut txn = ProfileTransaction::new(&target).unwrap();
txn.stage_profile(&profile_src, &["auth.json"]).unwrap();
txn.commit().unwrap();
assert!(target.exists());
txn.rollback().unwrap();
assert!(
!target.exists(),
"target should be removed when there was no original"
);
}
#[test]
fn test_cleanup_original_removes_backup() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
std::fs::create_dir_all(&target).unwrap();
std::fs::write(target.join("auth.json"), "old").unwrap();
let profile_src = make_profile(&tmp);
let mut txn = ProfileTransaction::new(&target).unwrap();
txn.stage_profile(&profile_src, &["auth.json"]).unwrap();
txn.commit().unwrap();
let orig_path = txn.original_backup_path().map(PathBuf::from);
assert!(orig_path.is_some());
let orig_path = orig_path.unwrap();
assert!(
orig_path.exists(),
"original backup should exist after commit"
);
txn.cleanup_original().unwrap();
assert!(
!orig_path.exists(),
"original backup should be gone after cleanup"
);
}
#[test]
fn test_drop_cleans_up_uncommitted_staging() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
let profile_src = make_profile(&tmp);
let staging_path = {
let mut txn = ProfileTransaction::new(&target).unwrap();
txn.stage_profile(&profile_src, &["auth.json"]).unwrap();
txn.staging_dir().to_path_buf()
};
assert!(
!staging_path.exists(),
"staging dir should be removed on drop"
);
}
#[test]
fn test_commit_without_stage_returns_error() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
let mut txn = ProfileTransaction::new(&target).unwrap();
let result = txn.commit();
assert!(result.is_err());
}
#[test]
fn test_stage_dir_stages_entire_tree() {
let tmp = TempDir::new().unwrap();
let target = tmp.path().join("codex");
let src = tmp.path().join("src");
std::fs::create_dir_all(&src).unwrap();
std::fs::write(src.join("auth.json"), "a").unwrap();
std::fs::create_dir_all(src.join("sessions")).unwrap();
std::fs::write(src.join("sessions").join("s1.json"), "s").unwrap();
let mut txn = ProfileTransaction::new(&target).unwrap();
txn.stage_dir(&src).unwrap();
txn.commit().unwrap();
assert!(target.join("auth.json").exists());
assert!(target.join("sessions").join("s1.json").exists());
}
}