jettypod 4.4.118 → 4.4.121

package/crates/jettypod-core/src/worktree.rs

@@ -0,0 +1,628 @@
+//! Worktree lifecycle management.
+//!
+//! Manages the full lifecycle of git worktrees used for isolated work:
+//! creating worktrees (with DB tracking), cleaning them up (multi-stage
+//! resilient removal), and querying their state.
+//!
+//! **Database is the single source of truth.** Every worktree operation
+//! persists state to the `worktrees` table before touching the filesystem.
+
+use std::fs;
+use std::path::{Path, PathBuf};
+
+use rusqlite::params;
+
+use crate::config;
+use crate::db::Database;
+use crate::error::{CoreError, GitError, Result, WorkError};
+use crate::git;
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+/// A worktree record as stored in the database.
+#[derive(Debug, Clone)]
+pub struct Worktree {
+    pub id: i64,
+    pub work_item_id: i64,
+    pub worktree_path: String,
+    pub branch_name: String,
+    pub status: WorktreeStatus,
+    pub created_at: String,
+    pub updated_at: String,
+}
+
+/// Valid worktree lifecycle states.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum WorktreeStatus {
+    Active,
+    Merging,
+    CleanupPending,
+    Cleaned,
+}
+
+impl WorktreeStatus {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Self::Active => "active",
+            Self::Merging => "merging",
+            Self::CleanupPending => "cleanup_pending",
+            Self::Cleaned => "cleaned",
+        }
+    }
+}
+
+impl std::fmt::Display for WorktreeStatus {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.as_str())
+    }
+}
+
+impl std::str::FromStr for WorktreeStatus {
+    type Err = CoreError;
+    fn from_str(s: &str) -> Result<Self> {
+        match s {
+            "active" => Ok(Self::Active),
+            "merging" => Ok(Self::Merging),
+            "cleanup_pending" => Ok(Self::CleanupPending),
+            "cleaned" => Ok(Self::Cleaned),
+            _ => Err(WorkError::InvalidState(format!("invalid worktree status: {s}")).into()),
+        }
+    }
+}
+
+fn row_to_worktree(row: &rusqlite::Row) -> rusqlite::Result<Worktree> {
+    let status_str: String = row.get("status")?;
+    Ok(Worktree {
+        id: row.get("id")?,
+        work_item_id: row.get("work_item_id")?,
+        worktree_path: row.get("worktree_path")?,
+        branch_name: row.get("branch_name")?,
+        status: status_str.parse().unwrap_or(WorktreeStatus::Active),
+        created_at: row.get("created_at")?,
+        updated_at: row.get("updated_at")?,
+    })
+}
+
+// ---------------------------------------------------------------------------
+// Queries
+// ---------------------------------------------------------------------------
+
+/// Get the most recent worktree for a work item (any status).
+pub fn get_for_work_item(db: &Database, work_item_id: i64) -> Result<Option<Worktree>> {
+    let mut stmt = db.conn().prepare(
+        "SELECT * FROM worktrees WHERE work_item_id = ? ORDER BY created_at DESC LIMIT 1",
+    )?;
+    let mut rows = stmt.query_map(params![work_item_id], row_to_worktree)?;
+    match rows.next() {
+        Some(Ok(wt)) => Ok(Some(wt)),
+        Some(Err(e)) => Err(e.into()),
+        None => Ok(None),
+    }
+}
+
+/// Get all worktrees with `status = 'active'`.
+pub fn get_all_active(db: &Database) -> Result<Vec<Worktree>> {
+    let mut stmt = db.conn().prepare(
+        "SELECT * FROM worktrees WHERE status = 'active' ORDER BY created_at DESC",
+    )?;
+    let rows = stmt.query_map([], row_to_worktree)?;
+    rows.map(|r| r.map_err(CoreError::from)).collect()
+}
+
+/// Update a worktree's status.
+pub fn mark_status(db: &Database, worktree_id: i64, status: WorktreeStatus) -> Result<()> {
+    let changed = db.conn().execute(
+        "UPDATE worktrees SET status = ?, updated_at = datetime('now') WHERE id = ?",
+        params![status.as_str(), worktree_id],
+    )?;
+    if changed == 0 {
+        return Err(WorkError::Other(format!("worktree not found: {worktree_id}")).into());
+    }
+    Ok(())
+}
+
+// ---------------------------------------------------------------------------
+// Create worktree
+// ---------------------------------------------------------------------------
+
+/// Options for creating a worktree.
+pub struct CreateWorktreeOpts<'a> {
+    /// The work item id and title (used to generate branch name + path).
+    pub work_item_id: i64,
+    pub title: &'a str,
+    /// Absolute path to the **main** repository root.
+    pub repo_root: &'a Path,
+}
+
+/// Create a new worktree for a work item.
+///
+/// 1. Inserts a DB record with status='active'.
+/// 2. Runs `git worktree add`.
+/// 3. Creates `.jettypod` symlink for shared DB access.
+/// 4. Symlinks `.env*` files.
+///
+/// On failure, rolls back the DB record and cleans up any partial git state.
+pub fn create(db: &Database, opts: &CreateWorktreeOpts) -> Result<Worktree> {
+    let title_slug = git::slugify(opts.title);
+    let branch_name = format!("feature/work-{}-{}", opts.work_item_id, title_slug);
+    let worktree_base = opts.repo_root.join(".jettypod-work");
+    let worktree_path = worktree_base.join(format!("{}-{}", opts.work_item_id, title_slug));
+    let worktree_path_str = worktree_path.to_string_lossy().to_string();
+
+    // Resolve default branch
+    let default_branch = config::resolve_default_branch(opts.repo_root);
+
+    // Auto-commit untracked BDD files so the worktree branch gets them
+    auto_commit_bdd_files(opts.repo_root, opts.work_item_id);
+
+    let mut worktree_id: Option<i64> = None;
+    let mut git_worktree_created = false;
+
+    // --- transactional creation ---
+    let result = (|| -> Result<Worktree> {
+        // Step 1: Insert DB record
+        db.conn().execute(
+            "INSERT INTO worktrees (work_item_id, branch_name, worktree_path, status) \
+             VALUES (?, ?, ?, 'active')",
+            params![opts.work_item_id, branch_name, worktree_path_str],
+        )?;
+        worktree_id = Some(db.conn().last_insert_rowid());
+
+        // Step 2: Create git worktree
+        git::worktree_add(opts.repo_root, &branch_name, &worktree_path, &default_branch)?;
+        git_worktree_created = true;
+
+        // Step 3: Verify directory exists
+        if !worktree_path.exists() {
+            return Err(GitError::Other(format!(
+                "worktree directory was not created: {}",
+                worktree_path.display()
+            )).into());
+        }
+
+        // Step 4: Create .jettypod symlink
+        setup_jettypod_symlink(opts.repo_root, &worktree_path)?;
+
+        // Step 5: Symlink .env files
+        symlink_env_files(opts.repo_root, &worktree_path);
+
+        // Step 6: Return the created record
+        let wt = db.conn().query_row(
+            "SELECT * FROM worktrees WHERE id = ?",
+            params![worktree_id.unwrap()],
+            row_to_worktree,
+        )?;
+        Ok(wt)
+    })();
+
+    // --- rollback on error ---
+    if let Err(ref _e) = result {
+        // Clean up DB
+        if let Some(id) = worktree_id {
+            if let Err(e) = db.conn().execute("DELETE FROM worktrees WHERE id = ?", params![id]) {
+                eprintln!("Warning: rollback cleanup failed (delete worktree record): {e}");
+            }
+        }
+        // Clean up git worktree
+        if git_worktree_created {
+            if let Err(e) = git::worktree_remove(opts.repo_root, &worktree_path) {
+                eprintln!("Warning: rollback cleanup failed (git worktree remove): {e}");
+            }
+        }
+        // Clean up leftover directory
+        if worktree_path.exists() {
+            if let Err(e) = fs::remove_dir_all(&worktree_path) {
+                eprintln!("Warning: rollback cleanup failed (remove directory): {e}");
+            }
+        }
+    }
+
+    result
+}
+
+// ---------------------------------------------------------------------------
+// Cleanup worktree
+// ---------------------------------------------------------------------------
+
+/// Options for cleaning up a worktree.
+pub struct CleanupWorktreeOpts<'a> {
+    /// Absolute path to the **main** repository root.
+    pub repo_root: &'a Path,
+    /// Whether to delete the git branch after cleanup.
+    pub delete_branch: bool,
+}
+
+/// Clean up a worktree: remove from filesystem, update DB status.
+///
+/// Uses a resilient multi-stage removal strategy:
+/// 1. `git worktree remove`
+/// 2. `git worktree remove --force`
+/// 3. `rm -rf` + `git worktree prune`
+pub fn cleanup(db: &Database, worktree_id: i64, opts: &CleanupWorktreeOpts) -> Result<()> {
+    // Fetch the worktree record
+    let wt = db.conn().query_row(
+        "SELECT * FROM worktrees WHERE id = ?",
+        params![worktree_id],
+        row_to_worktree,
+    )?;
+
+    // Mark as cleanup_pending
+    mark_status(db, worktree_id, WorktreeStatus::CleanupPending)?;
+
+    let cleanup_result = (|| -> Result<()> {
+        // Remove the worktree directory (resilient multi-stage)
+        let wt_path = PathBuf::from(&wt.worktree_path);
+
+        // Safety: only remove paths inside .jettypod-work
+        let worktree_base = opts.repo_root.join(".jettypod-work");
+        if !wt_path.starts_with(&worktree_base) {
+            return Err(GitError::Other(format!(
+                "SAFETY: can only remove directories within .jettypod-work/. Got: {}",
+                wt_path.display()
+            )).into());
+        }
+
+        if wt_path.exists() {
+            git::worktree_remove(opts.repo_root, &wt_path)?;
+        }
+
+        // Optionally delete the branch
+        if opts.delete_branch {
+            // Non-fatal if branch deletion fails
+            if let Err(e) = git::delete_branch(opts.repo_root, &wt.branch_name) {
+                eprintln!("Warning: Failed to delete branch {}: {e}", wt.branch_name);
+            }
+        }
+
+        Ok(())
+    })();
+
+    // Always mark as cleaned (terminal state) regardless of success
+    if let Err(e) = mark_status(db, worktree_id, WorktreeStatus::Cleaned) {
+        eprintln!("Warning: Failed to mark worktree as cleaned: {e}");
+    }
+
+    cleanup_result
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/// Auto-commit untracked BDD files (.feature, steps/) so the worktree gets them.
+fn auto_commit_bdd_files(repo_root: &Path, work_item_id: i64) {
+    let untracked = match git::untracked_files(repo_root) {
+        Ok(f) => f,
+        Err(_) => return,
+    };
+
+    let bdd_files: Vec<&str> = untracked
+        .iter()
+        .filter(|f| {
+            f.ends_with(".feature")
+                || f.contains("step_definitions/")
+                || f.contains("steps/")
+        })
+        .map(|s| s.as_str())
+        .collect();
+
+    if bdd_files.is_empty() {
+        return;
+    }
+
+    for f in &bdd_files {
+        if let Err(e) = git::run(repo_root, &["add", f]) {
+            eprintln!("Warning: Failed to stage backup file {}: {e}", f);
+        }
+    }
+    let msg = format!("Add BDD scenarios for work item #{work_item_id}");
+    if let Err(e) = git::run(repo_root, &["commit", "-m", &msg]) {
+        eprintln!("Warning: Failed to commit backup: {e}");
+    }
+}
+
+/// Create a `.jettypod` symlink inside the worktree pointing to the main repo's `.jettypod`.
+fn setup_jettypod_symlink(repo_root: &Path, worktree_path: &Path) -> Result<()> {
+    let link_path = worktree_path.join(".jettypod");
+    let target = repo_root.join(".jettypod");
+
+    if link_path.exists() || link_path.symlink_metadata().is_ok() {
+        // Check if it's already a correct symlink
+        if link_path.symlink_metadata().map(|m| m.file_type().is_symlink()).unwrap_or(false) {
+            let current_target = fs::read_link(&link_path)?;
+            let resolved = worktree_path.join(&current_target);
+            let resolved_canon = resolved.canonicalize().unwrap_or(resolved);
+            let target_canon = target.canonicalize().unwrap_or_else(|_| target.clone());
+            if resolved_canon == target_canon {
+                return Ok(()); // already correct
+            }
+            fs::remove_file(&link_path)?;
+        } else {
+            // It's a directory (git may create one) — remove it
+            fs::remove_dir_all(&link_path)?;
+        }
+    }
+
+    #[cfg(unix)]
+    std::os::unix::fs::symlink(&target, &link_path)?;
+
+    #[cfg(not(unix))]
+    return Err(GitError::Other("symlinks not supported on this platform".into()).into());
+
+    Ok(())
+}
+
+/// Symlink `.env*` files from the main repo into the worktree.
+fn symlink_env_files(repo_root: &Path, worktree_path: &Path) {
+    let entries = match fs::read_dir(repo_root) {
+        Ok(e) => e,
+        Err(_) => return,
+    };
+
+    for entry in entries.flatten() {
+        let name = entry.file_name();
+        let name_str = name.to_string_lossy();
+        if !name_str.starts_with(".env") {
+            continue;
+        }
+        let source = entry.path();
+        if !source.is_file() {
+            continue;
+        }
+        let dest = worktree_path.join(&name);
+        if dest.exists() {
+            continue;
+        }
+
+        #[cfg(unix)]
+        {
+            if let Err(e) = std::os::unix::fs::symlink(&source, &dest) {
+                eprintln!("Warning: Failed to create CLAUDE.md symlink: {e}");
+            }
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Merge lock
+// ---------------------------------------------------------------------------
+
+/// Acquire an exclusive merge lock. Returns the lock id.
+/// Blocks for up to `timeout_secs`, polling every 2 seconds.
+///
+/// Schema: merge_locks(id, locked_by, locked_at, operation, work_item_id, heartbeat_at)
+pub fn acquire_merge_lock(
+    db: &Database,
+    work_item_id: i64,
+    locked_by: &str,
+    timeout_secs: u64,
+) -> Result<i64> {
+    let start = std::time::Instant::now();
+    let poll_interval = std::time::Duration::from_secs(2);
+
+    loop {
+        // Clean up stale locks (no heartbeat in 2+ minutes)
+        db.conn().execute(
+            "DELETE FROM merge_locks WHERE heartbeat_at < datetime('now', '-2 minutes')",
+            [],
+        )?;
+
+        // Try to acquire
+        let existing: Option<i64> = db
+            .conn()
+            .query_row(
+                "SELECT id FROM merge_locks LIMIT 1",
+                [],
+                |row| row.get(0),
+            )
+            .ok();
+
+        if existing.is_none() {
+            db.conn().execute(
+                "INSERT INTO merge_locks (work_item_id, locked_by, operation) \
+                 VALUES (?, ?, 'merging')",
+                params![work_item_id, locked_by],
+            )?;
+            let lock_id = db.conn().last_insert_rowid();
+            return Ok(lock_id);
+        }
+
+        if start.elapsed().as_secs() >= timeout_secs {
+            return Err(WorkError::Other(
+                "timed out waiting for merge lock".into(),
+            ).into());
+        }
+
+        std::thread::sleep(poll_interval);
+    }
+}
+
+/// Release a merge lock.
+pub fn release_merge_lock(db: &Database, lock_id: i64) -> Result<()> {
+    db.conn()
+        .execute("DELETE FROM merge_locks WHERE id = ?", params![lock_id])?;
+    Ok(())
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::db::Database;
+    use tempfile::TempDir;
+
+    fn setup() -> (TempDir, Database) {
+        let dir = TempDir::new().unwrap();
+        // Init git repo
+        git::run(dir.path(), &["init"]).unwrap();
+        git::run(dir.path(), &["config", "user.email", "test@test.com"]).unwrap();
+        git::run(dir.path(), &["config", "user.name", "Test"]).unwrap();
+        std::fs::write(dir.path().join("README.md"), "init").unwrap();
+        git::run(dir.path(), &["add", "."]).unwrap();
+        git::run(dir.path(), &["commit", "-m", "init"]).unwrap();
+
+        // Create .jettypod dir + DB
+        let jettypod = dir.path().join(".jettypod");
+        std::fs::create_dir_all(&jettypod).unwrap();
+        let db_path = jettypod.join("work.db");
+        let db = Database::open_path_unchecked(&db_path).unwrap();
+
+        // Insert a work item to reference
+        db.conn()
+            .execute(
+                "INSERT INTO work_items (id, type, title, status) VALUES (1, 'chore', 'Test chore', 'todo')",
+                [],
+            )
+            .unwrap();
+
+        // Commit .jettypod so git worktree add doesn't complain about untracked
+        git::run(dir.path(), &["add", "."]).unwrap();
+        git::run(dir.path(), &["commit", "-m", "add jettypod"]).unwrap();
+
+        (dir, db)
+    }
+
+    #[test]
+    fn create_and_get_worktree() {
+        let (dir, db) = setup();
+
+        let wt = create(
+            &db,
+            &CreateWorktreeOpts {
+                work_item_id: 1,
+                title: "Test chore",
+                repo_root: dir.path(),
+            },
+        )
+        .unwrap();
+
+        assert_eq!(wt.work_item_id, 1);
+        assert_eq!(wt.status, WorktreeStatus::Active);
+        assert!(wt.branch_name.contains("test-chore"));
+        assert!(PathBuf::from(&wt.worktree_path).exists());
+
+        // .jettypod symlink should exist
+        let jettypod_link = PathBuf::from(&wt.worktree_path).join(".jettypod");
+        assert!(jettypod_link.exists());
+
+        // Query it back
+        let found = get_for_work_item(&db, 1).unwrap().unwrap();
+        assert_eq!(found.id, wt.id);
+    }
+
+    #[test]
+    fn cleanup_removes_worktree() {
+        let (dir, db) = setup();
+
+        let wt = create(
+            &db,
+            &CreateWorktreeOpts {
+                work_item_id: 1,
+                title: "Test chore",
+                repo_root: dir.path(),
+            },
+        )
+        .unwrap();
+
+        let wt_path = PathBuf::from(&wt.worktree_path);
+        assert!(wt_path.exists());
+
+        cleanup(
+            &db,
+            wt.id,
+            &CleanupWorktreeOpts {
+                repo_root: dir.path(),
+                delete_branch: true,
+            },
+        )
+        .unwrap();
+
+        assert!(!wt_path.exists());
+
+        let found = get_for_work_item(&db, 1).unwrap().unwrap();
+        assert_eq!(found.status, WorktreeStatus::Cleaned);
+    }
+
+    #[test]
+    fn get_all_active_works() {
+        let (dir, db) = setup();
+
+        // Insert a second work item
+        db.conn()
+            .execute(
+                "INSERT INTO work_items (id, type, title, status) VALUES (2, 'chore', 'Second', 'todo')",
+                [],
+            )
+            .unwrap();
+
+        let _wt1 = create(
+            &db,
+            &CreateWorktreeOpts {
+                work_item_id: 1,
+                title: "First",
+                repo_root: dir.path(),
+            },
+        )
+        .unwrap();
+
+        let _wt2 = create(
+            &db,
+            &CreateWorktreeOpts {
+                work_item_id: 2,
+                title: "Second",
+                repo_root: dir.path(),
+            },
+        )
+        .unwrap();
+
+        let active = get_all_active(&db).unwrap();
+        assert_eq!(active.len(), 2);
+    }
+
+    #[test]
+    fn mark_status_works() {
+        let (dir, db) = setup();
+
+        let wt = create(
+            &db,
+            &CreateWorktreeOpts {
+                work_item_id: 1,
+                title: "Test",
+                repo_root: dir.path(),
+            },
+        )
+        .unwrap();
+
+        mark_status(&db, wt.id, WorktreeStatus::Merging).unwrap();
+        let found = get_for_work_item(&db, 1).unwrap().unwrap();
+        assert_eq!(found.status, WorktreeStatus::Merging);
+    }
+
+    #[test]
+    fn merge_lock_acquire_release() {
+        let (_dir, db) = setup();
+
+        let lock_id = acquire_merge_lock(&db, 1, "session-1", 5).unwrap();
+        assert!(lock_id > 0);
+
+        release_merge_lock(&db, lock_id).unwrap();
+
+        // Should be able to acquire again
+        let lock_id2 = acquire_merge_lock(&db, 1, "session-2", 5).unwrap();
+        release_merge_lock(&db, lock_id2).unwrap();
+    }
+
+    #[test]
+    fn worktree_status_roundtrip() {
+        assert_eq!("active".parse::<WorktreeStatus>().unwrap(), WorktreeStatus::Active);
+        assert_eq!("merging".parse::<WorktreeStatus>().unwrap(), WorktreeStatus::Merging);
+        assert_eq!("cleanup_pending".parse::<WorktreeStatus>().unwrap(), WorktreeStatus::CleanupPending);
+        assert_eq!("cleaned".parse::<WorktreeStatus>().unwrap(), WorktreeStatus::Cleaned);
+        assert!("invalid".parse::<WorktreeStatus>().is_err());
+    }
+}