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

//! Process management — kills orphaned PTY child processes on server shutdown.
use crate::slog;

use super::AgentPool;

impl AgentPool {
    /// Kill all active PTY child processes.
    ///
    /// Called on server shutdown to prevent orphaned Claude Code processes from
    /// continuing to run after the server exits. Each registered killer is called
    /// once, then the registry is cleared.
    pub fn kill_all_children(&self) {
        if let Ok(mut killers) = self.child_killers.lock() {
            for (key, killer) in killers.iter_mut() {
                slog!("[agents] Killing child process for {key} on shutdown");
                let _ = killer.kill();
            }
            killers.clear();
        }
    }

    /// Kill and deregister the child process for a specific agent key.
    ///
    /// Used by `stop_agent` to ensure the PTY child is terminated even though
    /// aborting a `spawn_blocking` task handle does not interrupt the blocking thread.
    pub(super) fn kill_child_for_key(&self, key: &str) {
        if let Ok(mut killers) = self.child_killers.lock()
            && let Some(mut killer) = killers.remove(key)
        {
            slog!("[agents] Killing child process for {key} on stop");
            let _ = killer.kill();
        }
    }

    /// Test helper: inject a child killer into the registry.
    #[cfg(test)]
    pub fn inject_child_killer(
        &self,
        key: &str,
        killer: Box<dyn portable_pty::ChildKiller + Send + Sync>,
    ) {
        let mut killers = self.child_killers.lock().unwrap();
        killers.insert(key.to_string(), killer);
    }

    /// Test helper: return the number of registered child killers.
    #[cfg(test)]
    pub fn child_killer_count(&self) -> usize {
        self.child_killers.lock().unwrap().len()
    }
}

#[cfg(test)]
mod tests {
    use super::super::AgentPool;
    use portable_pty::{CommandBuilder, PtySize, native_pty_system};
    use std::process::Command;

    /// Returns true if a process with the given PID is currently running.
    fn process_is_running(pid: u32) -> bool {
        Command::new("ps")
            .args(["-p", &pid.to_string()])
            .output()
            .map(|o| o.status.success())
            .unwrap_or(false)
    }

    #[test]
    fn kill_all_children_is_safe_on_empty_pool() {
        let pool = AgentPool::new_test(3001);
        pool.kill_all_children();
        assert_eq!(pool.child_killer_count(), 0);
    }

    #[test]
    fn kill_all_children_kills_real_process() {
        let pool = AgentPool::new_test(3001);

        let pty_system = native_pty_system();
        let pair = pty_system
            .openpty(PtySize {
                rows: 24,
                cols: 80,
                pixel_width: 0,
                pixel_height: 0,
            })
            .expect("failed to open pty");

        let mut cmd = CommandBuilder::new("sleep");
        cmd.arg("100");
        let mut child = pair
            .slave
            .spawn_command(cmd)
            .expect("failed to spawn sleep");
        let pid = child.process_id().expect("no pid");

        pool.inject_child_killer("story:agent", child.clone_killer());

        assert!(
            process_is_running(pid),
            "process {pid} should be running before kill_all_children"
        );

        pool.kill_all_children();
        let _ = child.wait();

        assert!(
            !process_is_running(pid),
            "process {pid} should have been killed by kill_all_children"
        );
    }

    #[test]
    fn kill_all_children_clears_registry() {
        let pool = AgentPool::new_test(3001);

        let pty_system = native_pty_system();
        let pair = pty_system
            .openpty(PtySize {
                rows: 24,
                cols: 80,
                pixel_width: 0,
                pixel_height: 0,
            })
            .expect("failed to open pty");

        let mut cmd = CommandBuilder::new("sleep");
        cmd.arg("1");
        let mut child = pair
            .slave
            .spawn_command(cmd)
            .expect("failed to spawn sleep");

        pool.inject_child_killer("story:agent", child.clone_killer());
        assert_eq!(pool.child_killer_count(), 1);

        pool.kill_all_children();
        let _ = child.wait();

        assert_eq!(
            pool.child_killer_count(),
            0,
            "child_killers should be cleared after kill_all_children"
        );
    }
}