use anyhow::Result;
use std::path::{Path, PathBuf};
use tracing::warn;
use super::manifest::{compute_checksum, AssetManifest};
#[derive(Debug, Clone, Default)]
pub struct RollbackReport {
pub manifest_missing: bool,
pub restored: Vec<PathBuf>,
pub removed: Vec<PathBuf>,
pub skipped: Vec<PathBuf>,
pub errors: Vec<String>,
}
pub async fn rollback(project_dir: &Path, dry_run: bool) -> Result<RollbackReport> {
let manifest = match AssetManifest::load(project_dir).await? {
Some(m) => m,
None => {
return Ok(RollbackReport {
manifest_missing: true,
..RollbackReport::default()
});
}
};
let mut report = RollbackReport::default();
for entry in &manifest.files {
let abs = project_dir.join(&entry.path);
let backup = find_backup_for_entry(&manifest, project_dir, &entry.path, &abs).await;
if let Some(ref backup_path) = backup {
if dry_run {
report.restored.push(entry.path.clone());
} else {
match tokio::fs::copy(backup_path, &abs).await {
Ok(_) => {
if let Err(e) = tokio::fs::remove_file(backup_path).await {
warn!(
path = %backup_path.display(),
error = %e,
"Failed to remove backup file"
);
}
report.restored.push(entry.path.clone());
}
Err(e) => {
report.errors.push(format!(
"restore {} from backup {}: {}",
abs.display(),
backup_path.display(),
e
));
report.skipped.push(entry.path.clone());
}
}
}
} else if abs.exists() {
let should_remove = match entry.checksum {
Some(ref expected) => match tokio::fs::read_to_string(&abs).await {
Ok(content) => {
let actual = compute_checksum(&content);
actual == *expected
}
Err(_) => false,
},
None => false,
};
if should_remove {
if dry_run {
report.removed.push(entry.path.clone());
} else {
match tokio::fs::remove_file(&abs).await {
Ok(()) => report.removed.push(entry.path.clone()),
Err(e) => {
report
.errors
.push(format!("remove {}: {}", abs.display(), e));
}
}
}
} else {
report.skipped.push(entry.path.clone());
}
} else {
report.skipped.push(entry.path.clone());
}
}
let mut dirs: Vec<_> = manifest.directories.clone();
dirs.sort_by_key(|d| std::cmp::Reverse(d.components().count()));
for dir in &dirs {
let abs = project_dir.join(dir);
if abs.exists() {
if dry_run {
match tokio::fs::read_dir(&abs).await {
Ok(mut rd) => match rd.next_entry().await {
Ok(None) => report.removed.push(dir.clone()),
Ok(Some(_)) => report.skipped.push(dir.clone()),
Err(e) => report
.errors
.push(format!("read dir {}: {}", abs.display(), e)),
},
Err(e) => {
report
.errors
.push(format!("read dir {}: {}", abs.display(), e));
}
}
} else {
match tokio::fs::read_dir(&abs).await {
Ok(mut rd) => {
if rd.next_entry().await?.is_none() {
match tokio::fs::remove_dir(&abs).await {
Ok(()) => report.removed.push(dir.clone()),
Err(e) => {
report.errors.push(format!(
"remove dir {}: {}",
abs.display(),
e
));
}
}
} else {
report.skipped.push(dir.clone());
}
}
Err(e) => {
report
.errors
.push(format!("read dir {}: {}", abs.display(), e));
}
}
}
} else {
report.skipped.push(dir.clone());
}
}
if !dry_run {
let manifest_path = AssetManifest::manifest_path(project_dir);
if manifest_path.exists() {
let _ = tokio::fs::remove_file(&manifest_path).await;
}
}
Ok(report)
}
async fn find_backup_for(path: &Path) -> Option<PathBuf> {
let parent = path.parent()?;
let file_name = path.file_name()?.to_string_lossy();
let prefix = format!("{}.omk-backup-", file_name);
let mut entries = tokio::fs::read_dir(parent).await.ok()?;
let mut backups = Vec::new();
while let Ok(Some(entry)) = entries.next_entry().await {
let name = entry.file_name().to_string_lossy().to_string();
if name.starts_with(&prefix) {
backups.push(entry.path());
}
}
backups.sort();
backups.into_iter().next_back()
}
async fn find_backup_for_entry(
manifest: &AssetManifest,
project_dir: &Path,
managed_rel_path: &Path,
managed_abs_path: &Path,
) -> Option<PathBuf> {
if let Some(indexed_rel) = manifest.latest_backup_for(managed_rel_path) {
let indexed_abs = project_dir.join(&indexed_rel);
if indexed_abs.exists() {
return Some(indexed_abs);
}
}
find_backup_for(managed_abs_path).await
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_rollback_no_manifest() {
let dir = TempDir::new().unwrap();
let report = rollback(dir.path(), false).await.unwrap();
assert!(report.manifest_missing);
assert!(report.restored.is_empty());
assert!(report.removed.is_empty());
assert!(report.skipped.is_empty());
assert!(report.errors.is_empty());
}
#[tokio::test]
async fn test_rollback_removes_omk_created_file() {
let dir = TempDir::new().unwrap();
tokio::fs::create_dir_all(dir.path().join(".kimi"))
.await
.unwrap();
let mut manifest = AssetManifest::new(dir.path());
let file_path = dir.path().join("test.txt");
tokio::fs::write(&file_path, "omk-content").await.unwrap();
manifest
.add_file(
std::path::Path::new("test.txt"),
super::super::manifest::EntryKind::Other,
)
.await;
manifest.save(dir.path()).await.unwrap();
let report = rollback(dir.path(), false).await.unwrap();
assert!(!file_path.exists());
assert_eq!(report.removed.len(), 1);
assert_eq!(report.removed[0], PathBuf::from("test.txt"));
}
#[tokio::test]
async fn test_rollback_skips_user_modified_file() {
let dir = TempDir::new().unwrap();
tokio::fs::create_dir_all(dir.path().join(".kimi"))
.await
.unwrap();
let mut manifest = AssetManifest::new(dir.path());
let file_path = dir.path().join("test.txt");
tokio::fs::write(&file_path, "original").await.unwrap();
manifest
.add_file(
std::path::Path::new("test.txt"),
super::super::manifest::EntryKind::Other,
)
.await;
manifest.save(dir.path()).await.unwrap();
tokio::fs::write(&file_path, "modified").await.unwrap();
let report = rollback(dir.path(), false).await.unwrap();
assert!(file_path.exists());
assert_eq!(report.skipped.len(), 1);
assert_eq!(report.skipped[0], PathBuf::from("test.txt"));
}
#[tokio::test]
async fn test_rollback_restores_from_backup() {
let dir = TempDir::new().unwrap();
tokio::fs::create_dir_all(dir.path().join(".kimi"))
.await
.unwrap();
let mut manifest = AssetManifest::new(dir.path());
let file_path = dir.path().join("test.txt");
tokio::fs::write(&file_path, "user-content").await.unwrap();
let backup_path = format!("{}.omk-backup-1234567890", file_path.display());
tokio::fs::write(&backup_path, "user-content")
.await
.unwrap();
tokio::fs::write(&file_path, "omk-content").await.unwrap();
manifest
.add_file(
std::path::Path::new("test.txt"),
super::super::manifest::EntryKind::Other,
)
.await;
manifest.save(dir.path()).await.unwrap();
let report = rollback(dir.path(), false).await.unwrap();
assert_eq!(report.restored.len(), 1);
assert_eq!(report.restored[0], PathBuf::from("test.txt"));
let content = tokio::fs::read_to_string(&file_path).await.unwrap();
assert_eq!(content, "user-content");
assert!(!std::path::Path::new(&backup_path).exists());
}
#[tokio::test]
async fn test_rollback_dry_run_no_changes() {
let dir = TempDir::new().unwrap();
tokio::fs::create_dir_all(dir.path().join(".kimi"))
.await
.unwrap();
let mut manifest = AssetManifest::new(dir.path());
let file_path = dir.path().join("test.txt");
tokio::fs::write(&file_path, "omk-content").await.unwrap();
manifest
.add_file(
std::path::Path::new("test.txt"),
super::super::manifest::EntryKind::Other,
)
.await;
manifest.save(dir.path()).await.unwrap();
let report = rollback(dir.path(), true).await.unwrap();
assert_eq!(report.removed.len(), 1);
assert!(file_path.exists()); }
#[tokio::test]
async fn test_rollback_partial_failure_with_corrupt_backup_keeps_other_files_safe() {
let dir = TempDir::new().unwrap();
tokio::fs::create_dir_all(dir.path().join(".kimi"))
.await
.unwrap();
let mut manifest = AssetManifest::new(dir.path());
let removable_path = dir.path().join("remove-me.txt");
tokio::fs::write(&removable_path, "omk-owned")
.await
.unwrap();
manifest
.add_file(
std::path::Path::new("remove-me.txt"),
super::super::manifest::EntryKind::Other,
)
.await;
let restore_path = dir.path().join("restore-me.txt");
tokio::fs::write(&restore_path, "omk-content")
.await
.unwrap();
manifest
.add_file(
std::path::Path::new("restore-me.txt"),
super::super::manifest::EntryKind::Other,
)
.await;
manifest.save(dir.path()).await.unwrap();
let corrupt_backup = format!("{}.omk-backup-9999999999", restore_path.display());
tokio::fs::create_dir_all(&corrupt_backup).await.unwrap();
let report = rollback(dir.path(), false).await.unwrap();
assert!(!removable_path.exists());
assert!(report.removed.contains(&PathBuf::from("remove-me.txt")));
let restore_content = tokio::fs::read_to_string(&restore_path).await.unwrap();
assert_eq!(restore_content, "omk-content");
assert!(std::path::Path::new(&corrupt_backup).exists());
assert!(report.skipped.contains(&PathBuf::from("restore-me.txt")));
assert_eq!(report.errors.len(), 1);
assert!(report.errors[0].contains("restore"));
assert!(report.errors[0].contains("restore-me.txt"));
}
}