Add patch checkpointing

2 files changed, 301 insertions, 0 deletions

diff --git a/makima/src/daemon/storage/mod.rs b/makima/src/daemon/storage/mod.rs
new file mode 100644
index 0000000..cc5441a
--- /dev/null
+++ b/makima/src/daemon/storage/mod.rs
@@ -0,0 +1,8 @@
+//! Checkpoint storage for task recovery.
+//!
+//! This module provides functionality to store and restore git patches
+//! in PostgreSQL for recovering task worktrees when they are lost.
+
+mod patch;
+
+pub use patch::{create_patch, apply_patch, PatchError};
diff --git a/makima/src/daemon/storage/patch.rs b/makima/src/daemon/storage/patch.rs
new file mode 100644
index 0000000..45624b5
--- /dev/null
+++ b/makima/src/daemon/storage/patch.rs
@@ -0,0 +1,293 @@
+//! Git patch creation and application for checkpoint recovery.
+
+use flate2::read::GzDecoder;
+use flate2::write::GzEncoder;
+use flate2::Compression;
+use std::io::{Read, Write};
+use std::path::Path;
+use thiserror::Error;
+use tokio::process::Command;
+
+/// Errors that can occur during patch operations.
+#[derive(Error, Debug)]
+pub enum PatchError {
+    #[error("Git command failed: {0}")]
+    GitCommand(String),
+
+    #[error("Compression error: {0}")]
+    Compression(#[from] std::io::Error),
+
+    #[error("Patch too large: {size} bytes (max: {max} bytes)")]
+    TooLarge { size: usize, max: usize },
+
+    #[error("Empty patch (no changes)")]
+    EmptyPatch,
+
+    #[error("Failed to apply patch: {0}")]
+    ApplyFailed(String),
+}
+
+/// Create a compressed git diff from worktree changes.
+///
+/// Generates a diff between `base_sha` and HEAD, then compresses it with gzip.
+/// Returns the compressed patch bytes and the number of files changed.
+pub async fn create_patch(
+    worktree_path: &Path,
+    base_sha: &str,
+) -> Result<(Vec<u8>, usize), PatchError> {
+    // Get the diff between base commit and HEAD
+    let output = Command::new("git")
+        .current_dir(worktree_path)
+        .args(["diff", base_sha, "HEAD", "--binary"])
+        .output()
+        .await
+        .map_err(|e| PatchError::GitCommand(format!("Failed to run git diff: {}", e)))?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(PatchError::GitCommand(format!("git diff failed: {}", stderr)));
+    }
+
+    let diff_data = output.stdout;
+    if diff_data.is_empty() {
+        return Err(PatchError::EmptyPatch);
+    }
+
+    // Count files changed
+    let files_output = Command::new("git")
+        .current_dir(worktree_path)
+        .args(["diff", base_sha, "HEAD", "--name-only"])
+        .output()
+        .await
+        .map_err(|e| PatchError::GitCommand(format!("Failed to count files: {}", e)))?;
+
+    let files_count = if files_output.status.success() {
+        String::from_utf8_lossy(&files_output.stdout)
+            .lines()
+            .filter(|l| !l.is_empty())
+            .count()
+    } else {
+        0
+    };
+
+    // Compress with gzip
+    let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
+    encoder.write_all(&diff_data)?;
+    let compressed = encoder.finish()?;
+
+    Ok((compressed, files_count))
+}
+
+/// Apply a compressed patch to restore worktree state.
+///
+/// The worktree should already be checked out at `base_sha` before calling this.
+pub async fn apply_patch(worktree_path: &Path, patch_data: &[u8]) -> Result<(), PatchError> {
+    // Decompress gzip
+    let mut decoder = GzDecoder::new(patch_data);
+    let mut decompressed = Vec::new();
+    decoder.read_to_end(&mut decompressed)?;
+
+    if decompressed.is_empty() {
+        return Err(PatchError::EmptyPatch);
+    }
+
+    // Apply the patch using git apply
+    let mut child = Command::new("git")
+        .current_dir(worktree_path)
+        .args(["apply", "--binary", "-"])
+        .stdin(std::process::Stdio::piped())
+        .stdout(std::process::Stdio::piped())
+        .stderr(std::process::Stdio::piped())
+        .spawn()
+        .map_err(|e| PatchError::GitCommand(format!("Failed to spawn git apply: {}", e)))?;
+
+    // Write patch to stdin
+    if let Some(mut stdin) = child.stdin.take() {
+        use tokio::io::AsyncWriteExt;
+        stdin.write_all(&decompressed).await?;
+        drop(stdin); // Close stdin to signal EOF
+    }
+
+    let output = child
+        .wait_with_output()
+        .await
+        .map_err(|e| PatchError::GitCommand(format!("Failed to wait for git apply: {}", e)))?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(PatchError::ApplyFailed(stderr.to_string()));
+    }
+
+    Ok(())
+}
+
+/// Get the parent commit SHA (HEAD~1) from a worktree.
+pub async fn get_parent_sha(worktree_path: &Path) -> Result<String, PatchError> {
+    let output = Command::new("git")
+        .current_dir(worktree_path)
+        .args(["rev-parse", "HEAD~1"])
+        .output()
+        .await
+        .map_err(|e| PatchError::GitCommand(format!("Failed to get parent SHA: {}", e)))?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(PatchError::GitCommand(format!(
+            "git rev-parse HEAD~1 failed: {}",
+            stderr
+        )));
+    }
+
+    Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
+}
+
+/// Checkout a specific commit in the worktree.
+pub async fn checkout_commit(worktree_path: &Path, sha: &str) -> Result<(), PatchError> {
+    let output = Command::new("git")
+        .current_dir(worktree_path)
+        .args(["checkout", sha])
+        .output()
+        .await
+        .map_err(|e| PatchError::GitCommand(format!("Failed to checkout: {}", e)))?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(PatchError::GitCommand(format!(
+            "git checkout {} failed: {}",
+            sha, stderr
+        )));
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs;
+    use tempfile::TempDir;
+
+    async fn setup_test_repo() -> TempDir {
+        let dir = TempDir::new().unwrap();
+        let path = dir.path();
+
+        // Initialize git repo
+        Command::new("git")
+            .current_dir(path)
+            .args(["init"])
+            .output()
+            .await
+            .unwrap();
+
+        // Configure git user
+        Command::new("git")
+            .current_dir(path)
+            .args(["config", "user.email", "test@test.com"])
+            .output()
+            .await
+            .unwrap();
+        Command::new("git")
+            .current_dir(path)
+            .args(["config", "user.name", "Test"])
+            .output()
+            .await
+            .unwrap();
+
+        // Create initial commit
+        fs::write(path.join("file.txt"), "initial").unwrap();
+        Command::new("git")
+            .current_dir(path)
+            .args(["add", "."])
+            .output()
+            .await
+            .unwrap();
+        Command::new("git")
+            .current_dir(path)
+            .args(["commit", "-m", "initial"])
+            .output()
+            .await
+            .unwrap();
+
+        dir
+    }
+
+    #[tokio::test]
+    async fn test_create_and_apply_patch() {
+        let dir = setup_test_repo().await;
+        let path = dir.path();
+
+        // Get base SHA
+        let base_sha = get_parent_sha(path).await;
+        // This will fail since there's only one commit
+        assert!(base_sha.is_err());
+
+        // Make another commit first
+        fs::write(path.join("file.txt"), "modified").unwrap();
+        Command::new("git")
+            .current_dir(path)
+            .args(["add", "."])
+            .output()
+            .await
+            .unwrap();
+        Command::new("git")
+            .current_dir(path)
+            .args(["commit", "-m", "modified"])
+            .output()
+            .await
+            .unwrap();
+
+        // Now get the base SHA
+        let base_sha = get_parent_sha(path).await.unwrap();
+
+        // Create patch
+        let (patch_data, files_count) = create_patch(path, &base_sha).await.unwrap();
+        assert!(!patch_data.is_empty());
+        assert_eq!(files_count, 1);
+
+        // Reset to base and apply patch
+        checkout_commit(path, &base_sha).await.unwrap();
+        assert_eq!(fs::read_to_string(path.join("file.txt")).unwrap(), "initial");
+
+        apply_patch(path, &patch_data).await.unwrap();
+        assert_eq!(
+            fs::read_to_string(path.join("file.txt")).unwrap(),
+            "modified"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_empty_patch() {
+        let dir = setup_test_repo().await;
+        let path = dir.path();
+
+        // Make another commit
+        fs::write(path.join("file.txt"), "modified").unwrap();
+        Command::new("git")
+            .current_dir(path)
+            .args(["add", "."])
+            .output()
+            .await
+            .unwrap();
+        Command::new("git")
+            .current_dir(path)
+            .args(["commit", "-m", "modified"])
+            .output()
+            .await
+            .unwrap();
+
+        // Get current HEAD
+        let head_output = Command::new("git")
+            .current_dir(path)
+            .args(["rev-parse", "HEAD"])
+            .output()
+            .await
+            .unwrap();
+        let head_sha = String::from_utf8_lossy(&head_output.stdout)
+            .trim()
+            .to_string();
+
+        // Try to create patch from HEAD to HEAD (no changes)
+        let result = create_patch(path, &head_sha).await;
+        assert!(matches!(result, Err(PatchError::EmptyPatch)));
+    }
+}