pitchfork-cli 2.30.1

//! Hook execution for daemon lifecycle events
//!
//! Hooks are fire-and-forget shell commands that run in response to daemon
//! lifecycle events (ready, fail, retry). They are configured in pitchfork.toml
//! under `[daemons.<name>.hooks]`.

use crate::Result;
use crate::daemon_id::DaemonId;
use crate::pitchfork_toml::PitchforkToml;
use crate::settings::resolve_shell;
use crate::shell::{HideConsoleWindow, ShellScript};
use crate::supervisor::SUPERVISOR;
use crate::{env, pitchfork_toml, template};
use indexmap::IndexMap;
use std::collections::HashMap;
use std::path::PathBuf;

/// The type of lifecycle hook to fire
#[allow(clippy::enum_variant_names)]
pub(crate) enum HookType {
    OnReady,
    OnFail,
    OnRetry,
    OnStop,
    OnExit,
}

impl std::fmt::Display for HookType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            HookType::OnReady => write!(f, "on_ready"),
            HookType::OnFail => write!(f, "on_fail"),
            HookType::OnRetry => write!(f, "on_retry"),
            HookType::OnStop => write!(f, "on_stop"),
            HookType::OnExit => write!(f, "on_exit"),
        }
    }
}

fn get_hook_cmd(
    hooks: &Option<pitchfork_toml::PitchforkTomlHooks>,
    hook_type: &HookType,
) -> Option<String> {
    hooks.as_ref().and_then(|h| match hook_type {
        HookType::OnReady => h.on_ready.clone(),
        HookType::OnFail => h.on_fail.clone(),
        HookType::OnRetry => h.on_retry.clone(),
        HookType::OnStop => h.on_stop.clone(),
        HookType::OnExit => h.on_exit.clone(),
    })
}

/// Where a run's hooks are read from and run, fixed when that run starts.
///
/// They are read from the directory of the config that defines the daemon,
/// rather than its working directory, which `dir` can put outside the project
/// where that config is not found. They run in the working directory. Fixed
/// per run rather than looked up from the daemon's record when a hook fires,
/// so a run's `on_stop` or `on_exit` is not read from a replacement's project.
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct HookDirs {
    run_dir: PathBuf,
    config_root: PathBuf,
}

impl HookDirs {
    /// `project_dir` is the directory of the config that defines the daemon,
    /// recorded as `watch_base_dir`. An ad-hoc daemon has none, so its hooks
    /// are read from its working directory.
    pub(crate) fn new(run_dir: PathBuf, project_dir: Option<PathBuf>) -> Self {
        let config_root = project_dir.unwrap_or_else(|| run_dir.clone());
        Self {
            run_dir,
            config_root,
        }
    }

    /// For a daemon acted on from its record rather than from a run's
    /// options: an adopted daemon, or one the retry checker starts again.
    ///
    /// A record made from config before `watch_base_dir` was stored has no
    /// project directory, and its `dir` may be outside the project, so the
    /// project is found from the config that defines the daemon now.
    pub(crate) async fn for_record(daemon: &crate::daemon::Daemon) -> Self {
        let run_dir = daemon.dir.clone().unwrap_or_else(|| env::CWD.clone());
        let project_dir = match &daemon.watch_base_dir {
            Some(dir) => Some(dir.clone()),
            None => {
                let id = daemon.id.clone();
                tokio::task::spawn_blocking(move || {
                    PitchforkToml::all_merged_all_namespaces_blocking()
                        .ok()
                        .and_then(|pt| defining_project_dir(&pt, &id))
                })
                .await
                .ok()
                .flatten()
            }
        };
        Self::new(run_dir, project_dir)
    }
}

/// The directory of the project whose config defines `id` in `pt`, if any.
fn defining_project_dir(pt: &PitchforkToml, id: &DaemonId) -> Option<PathBuf> {
    let path = pt.daemons.get(id)?.path.as_deref()?;
    pitchfork_toml::project_dir_for_config(path)
}

async fn load_hook_config(daemon_dir: PathBuf) -> Result<PitchforkToml> {
    PitchforkToml::all_merged_all_namespaces_from(&daemon_dir).await
}

/// Create a tokio Command for a hook using the resolved shell for this platform
/// (the same shell used for daemon `run` commands). Returns an error if the
/// configured shell cannot be parsed, matching daemon startup validation —
/// callers should log and skip the hook.
fn hook_command(cmd: &str) -> Result<tokio::process::Command> {
    let parts = resolve_shell().map_err(|e| miette::miette!(e))?;
    let (program, args) = parts.split_first().unwrap();
    let mut command = tokio::process::Command::new(program);
    command.shell_script(program, args, cmd);
    command.hide_console_window();
    Ok(command)
}

/// Inject port env vars (`PITCHFORK_PORT` / `PITCHFORK_PORT0..N`) for the
/// daemon, mirroring the daemon's own spawn environment (`PORT` / `PORT0..N`,
/// see `apply_runtime_env` in lifecycle.rs) but namespaced to match the other
/// hook env vars.
fn inject_port_env(command: &mut tokio::process::Command, resolved_ports: &[u16]) {
    if let Some(port) = resolved_ports.first() {
        command.env("PITCHFORK_PORT", port.to_string());
        for (index, port) in resolved_ports.iter().enumerate() {
            command.env(format!("PITCHFORK_PORT{index}"), port.to_string());
        }
    }
}

/// Fire a hook command as a fire-and-forget tokio task.
///
/// Reads the hook command from fresh config rooted at the daemon's project,
/// then spawns it in the background. Errors are logged but never block the caller.
///
/// `resolved_ports` must be snapshotted by the caller before spawning this
/// task, so the hook observes the triggering attempt's ports even if a retry
/// or restart replaces the daemon's state afterwards.
///
/// The spawned task is also registered in `SUPERVISOR.hook_tasks` so that
/// supervisor shutdown (`close()`) can await all in-flight hooks before calling
/// `exit(0)`, ensuring hooks are not silently dropped during shutdown.
pub(crate) async fn fire_hook(
    hook_type: HookType,
    daemon_id: DaemonId,
    dirs: HookDirs,
    retry_count: u32,
    daemon_env: Option<IndexMap<String, String>>,
    resolved_ports: Vec<u16>,
    extra_env: Vec<(String, String)>,
) {
    let handle = tokio::spawn(async move {
        let pt = load_hook_config(dirs.config_root.clone())
            .await
            .unwrap_or_else(|e| {
                warn!("Failed to load config for hook '{hook_type}': {e}");
                PitchforkToml::default()
            });
        let hook_cmd = pt
            .daemons
            .get(&daemon_id)
            .and_then(|d| get_hook_cmd(&d.hooks, &hook_type));

        let Some(cmd) = hook_cmd else { return };

        // Render Tera templates in hook command with context from state file
        let cmd = match render_hook_template(&cmd, &daemon_id, &pt).await {
            Ok(cmd) => cmd,
            Err(e) => {
                warn!("{hook_type} hook template error for daemon {daemon_id}: {e}");
                return;
            }
        };

        info!("firing {hook_type} hook for daemon {daemon_id}: {cmd}");

        let mut command = match hook_command(&cmd) {
            Ok(c) => c,
            Err(e) => {
                warn!("{hook_type} hook for daemon {daemon_id}: {e}");
                return;
            }
        };
        command
            .current_dir(&dirs.run_dir)
            .stdout(std::process::Stdio::null())
            .stderr(std::process::Stdio::null());

        if let Some(ref path) = *env::ORIGINAL_PATH {
            command.env("PATH", path);
        }

        // Apply user env vars first
        if let Some(ref env_vars) = daemon_env {
            command.envs(env_vars);
        }

        // Inject pitchfork metadata env vars AFTER user env so they can't be overwritten
        command
            .env("PITCHFORK_DAEMON_ID", daemon_id.qualified())
            .env("PITCHFORK_DAEMON_NAMESPACE", daemon_id.namespace())
            .env("PITCHFORK_RETRY_COUNT", retry_count.to_string());
        inject_port_env(&mut command, &resolved_ports);

        for (key, value) in &extra_env {
            command.env(key, value);
        }

        match command.status().await {
            Ok(status) => {
                if !status.success() {
                    warn!("{hook_type} hook for daemon {daemon_id} exited with {status}");
                }
            }
            Err(e) => {
                error!("failed to execute {hook_type} hook for daemon {daemon_id}: {e}");
            }
        }
    });

    // Register the handle so supervisor shutdown can await it.
    // Use lock().await instead of try_lock() to guarantee registration — fire_hook
    // is called from async monitoring tasks (not latency-sensitive hot paths), so
    // awaiting the lock is safe and eliminates the silent-drop window where a
    // JoinHandle could be lost under lock contention.
    let mut tasks = SUPERVISOR.hook_tasks.lock().await;
    // Prune already-finished handles to avoid unbounded growth
    tasks.retain(|h| !h.is_finished());
    tasks.push(handle);
}

/// Fire the `on_output` hook for a daemon as a fire-and-forget task.
///
/// `cmd` is the hook command string resolved at call time (from `on_output.run`).
/// `matched_line` is exposed to the command via `PITCHFORK_MATCHED_LINE`.
/// `resolved_ports` must be snapshotted by the caller before spawning this task
/// (same contract as `fire_hook`).
pub(crate) async fn fire_output_hook(
    daemon_id: DaemonId,
    dirs: HookDirs,
    retry_count: u32,
    daemon_env: Option<IndexMap<String, String>>,
    resolved_ports: Vec<u16>,
    cmd: String,
    matched_line: String,
) {
    let handle = tokio::spawn(async move {
        // Render Tera templates in output hook command
        let pt = load_hook_config(dirs.config_root.clone())
            .await
            .unwrap_or_default();
        let cmd = match render_hook_template(&cmd, &daemon_id, &pt).await {
            Ok(cmd) => cmd,
            Err(e) => {
                warn!("on_output hook template error for daemon {daemon_id}: {e}");
                return;
            }
        };

        info!("firing on_output hook for daemon {daemon_id}: {cmd}");

        let mut command = match hook_command(&cmd) {
            Ok(c) => c,
            Err(e) => {
                warn!("on_output hook for daemon {daemon_id}: {e}");
                return;
            }
        };
        command
            .current_dir(&dirs.run_dir)
            .stdout(std::process::Stdio::null())
            .stderr(std::process::Stdio::null());

        if let Some(ref path) = *env::ORIGINAL_PATH {
            command.env("PATH", path);
        }

        if let Some(ref env_vars) = daemon_env {
            command.envs(env_vars);
        }

        command
            .env("PITCHFORK_DAEMON_ID", daemon_id.qualified())
            .env("PITCHFORK_DAEMON_NAMESPACE", daemon_id.namespace())
            .env("PITCHFORK_RETRY_COUNT", retry_count.to_string())
            .env("PITCHFORK_MATCHED_LINE", &matched_line);
        inject_port_env(&mut command, &resolved_ports);

        match command.status().await {
            Ok(status) => {
                if !status.success() {
                    warn!("on_output hook for daemon {daemon_id} exited with {status}");
                }
            }
            Err(e) => {
                error!("failed to execute on_output hook for daemon {daemon_id}: {e}");
            }
        }
    });

    let mut tasks = SUPERVISOR.hook_tasks.lock().await;
    tasks.retain(|h| !h.is_finished());
    tasks.push(handle);
}

/// Render Tera templates in a hook command using the current state file data.
///
/// In the supervisor, daemons are already running, so `resolved_port` and
/// `active_port` are available from the state file.
async fn render_hook_template(
    template_str: &str,
    daemon_id: &DaemonId,
    pt: &PitchforkToml,
) -> Result<String, template::RenderError> {
    // Collect resolved ports from all running daemons in the state file
    let resolved_daemons: HashMap<DaemonId, Vec<u16>> = {
        let state_file = SUPERVISOR.state_file.lock().await;
        state_file
            .daemons
            .iter()
            .filter_map(|(id, d)| {
                if d.resolved_port.is_empty() {
                    None
                } else {
                    Some((id.clone(), d.resolved_port.clone()))
                }
            })
            .collect()
    };

    // Building the context reads the slug registry and derives the daemon's
    // hostname, which walks the project's checkouts. Hooks fire on the
    // supervisor's runtime, so that work goes to a blocking worker.
    let pt = pt.clone();
    let daemon_id = daemon_id.clone();
    let template_str = template_str.to_string();
    let template_for_error = template_str.clone();
    tokio::task::spawn_blocking(move || {
        let daemon_config = pt.daemons.get(&daemon_id);
        let mut ctx = template::TemplateContext::new(
            &daemon_id,
            daemon_config.unwrap_or(&Default::default()),
            &resolved_daemons,
            &pt.daemons,
        );

        // Merge top-level env with per-daemon env (per-daemon wins), render the
        // values, and expose them as `{{ env.X }}` for hook templates.
        if let Some(rendered_env) = template::render_env(
            pt.env.as_ref(),
            daemon_config.and_then(|d| d.env.as_ref()),
            &ctx,
        )? {
            ctx.set_env(rendered_env);
        }

        template::render_template(&template_str, &ctx)
    })
    .await
    .unwrap_or_else(|e| {
        Err(template::RenderError::RenderFailed {
            template: template_for_error,
            source: tera::Error::message(format!("hook template rendering task failed: {e}")),
        })
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn hooks_are_read_from_the_project_of_a_daemon_running_elsewhere() {
        let dirs = HookDirs::new(PathBuf::from("/elsewhere"), Some(PathBuf::from("/project")));
        assert_eq!(dirs.config_root, PathBuf::from("/project"));
        assert_eq!(dirs.run_dir, PathBuf::from("/elsewhere"));
        // An ad-hoc daemon has no project; its working directory is used.
        let dirs = HookDirs::new(PathBuf::from("/elsewhere"), None);
        assert_eq!(dirs.config_root, PathBuf::from("/elsewhere"));
    }

    #[test]
    fn a_record_without_its_project_finds_it_from_the_config_defining_it() {
        let temp = tempfile::tempdir().unwrap();
        let project = temp.path().join("project");
        std::fs::create_dir(&project).unwrap();
        std::fs::write(
            project.join("pitchfork.toml"),
            "[daemons.api]\nrun = \"true\"\ndir = \"../elsewhere\"\n",
        )
        .unwrap();
        let pt = PitchforkToml::all_merged_from(&project).unwrap();
        let id = pt.daemons.keys().next().unwrap().clone();

        let project = project.canonicalize().unwrap();
        let found = defining_project_dir(&pt, &id).map(|d| d.canonicalize().unwrap());
        assert_eq!(found, Some(project));
        // Not defined by any config, like a daemon from `pitchfork run`.
        assert_eq!(
            defining_project_dir(&pt, &DaemonId::new("global", "adhoc")),
            None
        );
    }

    #[tokio::test]
    async fn hook_config_is_loaded_from_each_daemon_directory() {
        let temp = tempfile::tempdir().unwrap();
        let first = temp.path().join("worktree-one");
        let second = temp.path().join("worktree-two");
        tokio::fs::create_dir(&first).await.unwrap();
        tokio::fs::create_dir(&second).await.unwrap();
        tokio::fs::write(
            first.join("pitchfork.toml"),
            "[daemons.api]\nrun = \"true\"\n[daemons.api.hooks]\non_ready = \"echo one\"\n",
        )
        .await
        .unwrap();
        tokio::fs::write(
            second.join("pitchfork.toml"),
            "[daemons.api]\nrun = \"true\"\n[daemons.api.hooks]\non_ready = \"echo two\"\n",
        )
        .await
        .unwrap();

        let first_id = DaemonId::new("worktree-one", "api");
        let second_id = DaemonId::new("worktree-two", "api");
        let first_config = load_hook_config(first).await.unwrap();
        let second_config = load_hook_config(second).await.unwrap();

        assert_eq!(
            get_hook_cmd(&first_config.daemons[&first_id].hooks, &HookType::OnReady),
            Some("echo one".to_string())
        );
        assert_eq!(
            get_hook_cmd(&second_config.daemons[&second_id].hooks, &HookType::OnReady),
            Some("echo two".to_string())
        );
        assert!(!first_config.daemons.contains_key(&second_id));
        assert!(!second_config.daemons.contains_key(&first_id));
    }
}