huskies/server/src/agents/pool/worktree_lifecycle.rs

//! TransitionFired subscribers for worktree and feature-branch lifecycle.
//!
//! `spawn_worktree_create_subscriber` creates worktrees when stories enter
//! `Stage::Coding`. `spawn_worktree_cleanup_subscriber` removes worktrees
//! when stories reach terminal stages (Done, Archived, Abandoned, Superseded).

use std::path::{Path, PathBuf};

use crate::pipeline_state::Stage;
use crate::slog;
use crate::slog_warn;

/// Spawn a background task that creates a git worktree when a story enters `Stage::Coding`.
///
/// Subscribes to the pipeline transition broadcast channel. On each
/// `Stage::Coding` transition, creates the worktree and feature branch for the
/// story at `project_root` using `port` for the `.mcp.json` server URL.
/// The create is idempotent — if the worktree already exists it is reused.
pub(crate) fn spawn_worktree_create_subscriber(project_root: PathBuf, port: u16) {
    let mut rx = crate::pipeline_state::subscribe_transitions();
    tokio::spawn(async move {
        loop {
            match rx.recv().await {
                Ok(fired) => {
                    if matches!(fired.after, Stage::Coding { .. }) {
                        on_coding_transition(&project_root, port, &fired.story_id.0).await;
                    }
                }
                Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                    slog_warn!(
                        "[worktree-create-sub] Subscriber lagged, skipped {n} event(s). \
                         Some worktrees may need manual creation."
                    );
                }
                Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
            }
        }
    });
}

/// Spawn a background task that removes git worktrees when stories reach terminal stages.
///
/// Subscribes to the pipeline transition broadcast channel. On transitions into
/// `Stage::Done`, `Stage::Archived`, `Stage::Abandoned`, or `Stage::Superseded`,
/// removes the worktree and feature branch for the story.
/// Non-fatal if the worktree does not exist.
pub(crate) fn spawn_worktree_cleanup_subscriber(project_root: PathBuf) {
    let mut rx = crate::pipeline_state::subscribe_transitions();
    tokio::spawn(async move {
        loop {
            match rx.recv().await {
                Ok(fired) => {
                    if matches!(
                        fired.after,
                        Stage::Done { .. }
                            | Stage::Archived { .. }
                            | Stage::Abandoned { .. }
                            | Stage::Superseded { .. }
                    ) {
                        on_terminal_transition(&project_root, &fired.story_id.0).await;
                    }
                }
                Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                    slog_warn!(
                        "[worktree-cleanup-sub] Subscriber lagged, skipped {n} event(s). \
                         Some worktrees may need manual removal."
                    );
                }
                Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
            }
        }
    });
}

/// Reconcile worktree creation: for each story currently in `Stage::Coding`, ensure its worktree exists.
///
/// Idempotent — creates worktrees for Coding stories that have no worktree yet, and is
/// a no-op for stories whose worktree already exists.  Called by the periodic reconciler
/// so that Lagged events on the broadcast channel never leave Coding stories without worktrees.
pub(crate) async fn reconcile_worktree_create(project_root: &Path, port: u16) {
    for item in crate::pipeline_state::read_all_typed() {
        if matches!(item.stage, crate::pipeline_state::Stage::Coding { .. }) {
            on_coding_transition(project_root, port, &item.story_id.0).await;
        }
    }
}

/// Reconcile worktree cleanup: for each story in a terminal stage, ensure its worktree is removed.
///
/// Idempotent — removes worktrees for terminal stories that still have one, and is a no-op
/// for stories with no worktree.  Called by the periodic reconciler so that Lagged events on
/// the broadcast channel never leave terminal stories with dangling worktrees.
pub(crate) async fn reconcile_worktree_cleanup(project_root: &Path) {
    for item in crate::pipeline_state::read_all_typed() {
        if matches!(
            item.stage,
            crate::pipeline_state::Stage::Done { .. }
                | crate::pipeline_state::Stage::Archived { .. }
                | crate::pipeline_state::Stage::Abandoned { .. }
                | crate::pipeline_state::Stage::Superseded { .. }
        ) {
            on_terminal_transition(project_root, &item.story_id.0).await;
        }
    }
}

/// Create the worktree and feature branch for `story_id` when it enters `Stage::Coding`.
pub(crate) async fn on_coding_transition(project_root: &Path, port: u16, story_id: &str) {
    let config = match crate::config::ProjectConfig::load(project_root) {
        Ok(c) => c,
        Err(e) => {
            slog_warn!("[worktree-create-sub] Failed to load config for '{story_id}': {e}");
            return;
        }
    };
    slog!("[worktree-create-sub] Story '{story_id}' entered Coding; ensuring worktree exists.");
    match crate::worktree::create_worktree(project_root, story_id, &config, port).await {
        Ok(info) => {
            slog!(
                "[worktree-create-sub] Worktree ready for '{story_id}' at {}",
                info.path.display()
            );
            if let Err(e) = crate::worktree::install_pre_commit_hook(&info.path) {
                slog_warn!(
                    "[worktree-create-sub] Pre-commit hook install failed for '{story_id}': {e}"
                );
            }
        }
        Err(e) => {
            slog_warn!("[worktree-create-sub] Failed to create worktree for '{story_id}': {e}");
        }
    }
}

/// Remove the worktree and feature branch for `story_id` after it reaches a terminal stage.
pub(crate) async fn on_terminal_transition(project_root: &Path, story_id: &str) {
    let config = match crate::config::ProjectConfig::load(project_root) {
        Ok(c) => c,
        Err(e) => {
            slog_warn!("[worktree-cleanup-sub] Failed to load config for '{story_id}': {e}");
            return;
        }
    };
    slog!("[worktree-cleanup-sub] Story '{story_id}' reached terminal stage; removing worktree.");
    match crate::worktree::remove_worktree_by_story_id(project_root, story_id, &config).await {
        Ok(()) => slog!("[worktree-cleanup-sub] Worktree removed for '{story_id}'."),
        Err(e) => {
            // Non-fatal — worktree may not exist (story never had one, or already removed).
            slog!("[worktree-cleanup-sub] Worktree removal for '{story_id}': {e}");
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::process::Command;
    use tempfile::TempDir;

    fn init_git_repo(dir: &Path) {
        Command::new("git")
            .args(["init"])
            .current_dir(dir)
            .output()
            .expect("git init");
        Command::new("git")
            .args(["commit", "--allow-empty", "-m", "init"])
            .current_dir(dir)
            .output()
            .expect("git commit");
    }

    fn setup_project(tmp: &TempDir) -> PathBuf {
        let root = tmp.path().join("project");
        fs::create_dir_all(root.join(".huskies")).unwrap();
        std::fs::write(root.join(".huskies").join("project.toml"), "").unwrap();
        init_git_repo(&root);
        root
    }

    /// AC1: on_coding_transition creates the worktree and feature branch.
    #[tokio::test]
    async fn coding_transition_creates_worktree() {
        let tmp = TempDir::new().unwrap();
        let root = setup_project(&tmp);
        let story_id = "1006_test_create";

        on_coding_transition(&root, 3001, story_id).await;

        let wt_path = crate::worktree::worktree_path(&root, story_id);
        assert!(
            wt_path.exists(),
            "worktree directory must exist after Coding transition"
        );
        assert!(
            wt_path.join(".mcp.json").exists(),
            ".mcp.json must be written in the worktree"
        );
        // Verify the feature branch was created.
        let branch_output = Command::new("git")
            .args(["branch", "--list", "feature/story-1006_test_create"])
            .current_dir(&root)
            .output()
            .expect("git branch --list");
        assert!(
            !String::from_utf8_lossy(&branch_output.stdout)
                .trim()
                .is_empty(),
            "feature branch must exist after Coding transition"
        );
    }

    /// AC1: calling on_coding_transition twice is idempotent.
    #[tokio::test]
    async fn coding_transition_is_idempotent() {
        let tmp = TempDir::new().unwrap();
        let root = setup_project(&tmp);
        let story_id = "1006_test_idempotent";

        on_coding_transition(&root, 3001, story_id).await;
        on_coding_transition(&root, 3001, story_id).await;

        let wt_path = crate::worktree::worktree_path(&root, story_id);
        assert!(
            wt_path.exists(),
            "worktree must still exist after second Coding transition"
        );
    }

    /// AC2: on_terminal_transition removes the worktree after it was created.
    #[tokio::test]
    async fn terminal_transition_removes_worktree() {
        let tmp = TempDir::new().unwrap();
        let root = setup_project(&tmp);
        let story_id = "1006_test_remove";

        on_coding_transition(&root, 3001, story_id).await;
        let wt_path = crate::worktree::worktree_path(&root, story_id);
        assert!(wt_path.exists(), "worktree must exist before cleanup");

        on_terminal_transition(&root, story_id).await;
        assert!(
            !wt_path.exists(),
            "worktree must be removed after terminal transition"
        );
    }

    /// AC2: on_terminal_transition is a no-op (non-fatal) when no worktree exists.
    #[tokio::test]
    async fn terminal_transition_noop_when_no_worktree() {
        let tmp = TempDir::new().unwrap();
        let root = setup_project(&tmp);

        // Should not panic or error — just log and return.
        on_terminal_transition(&root, "1006_test_no_wt").await;
    }

    /// AC1+AC4: spawn_worktree_create_subscriber reacts to a real Coding transition.
    #[tokio::test]
    async fn create_subscriber_reacts_to_coding_transition() {
        crate::crdt_state::init_for_test();
        crate::db::ensure_content_store();

        let tmp = TempDir::new().unwrap();
        let root = setup_project(&tmp);
        let story_id = "1006_test_sub_create";

        crate::db::write_item_with_content(
            story_id,
            "1_backlog",
            "---\nname: Test\n---\n",
            crate::db::ItemMeta::named("Test"),
        );

        spawn_worktree_create_subscriber(root.clone(), 3001);

        // Trigger the Coding transition.
        crate::agents::lifecycle::move_story_to_current(story_id)
            .expect("move to current must succeed");

        // Give the subscriber task time to run and create the worktree.
        tokio::time::sleep(std::time::Duration::from_millis(500)).await;

        let wt_path = crate::worktree::worktree_path(&root, story_id);
        assert!(
            wt_path.exists(),
            "worktree must exist after Coding transition via subscriber"
        );
    }

    /// AC2+AC4: spawn_worktree_cleanup_subscriber reacts to terminal transitions.
    #[tokio::test]
    async fn cleanup_subscriber_reacts_to_terminal_transition() {
        crate::crdt_state::init_for_test();
        crate::db::ensure_content_store();

        let tmp = TempDir::new().unwrap();
        let root = setup_project(&tmp);
        let story_id = "1006_test_sub_cleanup";

        crate::db::write_item_with_content(
            story_id,
            "2_current",
            "---\nname: Test\n---\n",
            crate::db::ItemMeta::named("Test"),
        );

        // Create the worktree manually so the cleanup subscriber has something to remove.
        let config = crate::config::ProjectConfig::load(&root).unwrap();
        crate::worktree::create_worktree(&root, story_id, &config, 3001)
            .await
            .expect("create worktree must succeed");

        let wt_path = crate::worktree::worktree_path(&root, story_id);
        assert!(
            wt_path.exists(),
            "worktree must exist before subscriber test"
        );

        spawn_worktree_cleanup_subscriber(root.clone());

        // Trigger a terminal transition (Abandoned).
        crate::agents::lifecycle::abandon_story(story_id).expect("abandon must succeed");

        // Give the subscriber time to process and remove the worktree.
        tokio::time::sleep(std::time::Duration::from_millis(500)).await;

        assert!(
            !wt_path.exists(),
            "worktree must be removed after terminal transition via subscriber"
        );
    }
}