Agentsd/core/agentsd/src/syscall/process.rs

use std::collections::HashMap;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::process::Command;
use tokio::sync::{mpsc, watch, Mutex, RwLock};
use tokio::time::Duration;

/// Per-process managed state, decoupled from the Child handle.
struct ManagedProcess {
    cmd: String,
    status: Arc<Mutex<String>>,
    stdin_tx: Option<mpsc::Sender<Vec<u8>>>,
    stdout_rx: Arc<Mutex<mpsc::Receiver<Vec<u8>>>>,
    /// Watch receiver: becomes `true` when process exits. Clone to wait.
    exit_rx: watch::Receiver<bool>,
    /// Handle to abort background tasks on kill.
    _tasks: Vec<tokio::task::JoinHandle<()>>,
}

/// Process syscall: spawn and manage subprocesses.
/// Uses the real OS PID as the identifier.
pub struct ProcessSyscall {
    procs: Arc<RwLock<HashMap<u64, Arc<ManagedProcess>>>>,
}

impl ProcessSyscall {
    pub fn new() -> Self {
        Self {
            procs: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn spawn(
        &self,
        cmd: &str,
        args: &[String],
        env: &HashMap<String, String>,
        cwd: &str,
    ) -> Result<u64, String> {
        let mut command = Command::new(cmd);
        command.args(args);
        if !cwd.is_empty() {
            command.current_dir(cwd);
        }
        for (k, v) in env {
            command.env(k, v);
        }
        command
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .kill_on_drop(true);

        let mut child = command.spawn().map_err(|e| e.to_string())?;
        let pid = child
            .id()
            .ok_or_else(|| "spawned process has no OS pid".to_string())? as u64;

        // Take ownership of stdio handles
        let child_stdin = child.stdin.take();
        let child_stdout = child.stdout.take();
        let child_stderr = child.stderr.take();

        let status = Arc::new(Mutex::new("running".to_string()));
        let (exit_tx, exit_rx) = watch::channel(false);

        // Stdin writer task
        let (stdin_tx, mut stdin_rx) = mpsc::channel::<Vec<u8>>(64);
        let stdin_task = tokio::spawn(async move {
            if let Some(mut stdin) = child_stdin {
                while let Some(data) = stdin_rx.recv().await {
                    if stdin.write_all(&data).await.is_err() {
                        break;
                    }
                    if stdin.flush().await.is_err() {
                        break;
                    }
                }
            }
        });

        // Stdout reader task
        let (stdout_tx, stdout_rx) = mpsc::channel::<Vec<u8>>(64);
        let stdout_task = tokio::spawn(async move {
            if let Some(mut stdout) = child_stdout {
                loop {
                    let mut buf = vec![0u8; 4096];
                    match stdout.read(&mut buf).await {
                        Ok(0) => break,
                        Ok(n) => {
                            buf.truncate(n);
                            if stdout_tx.send(buf).await.is_err() {
                                break;
                            }
                        }
                        Err(_) => break,
                    }
                }
            }
        });

        // Stderr drain task (prevent pipe fill-up)
        let stderr_task = tokio::spawn(async move {
            if let Some(mut stderr) = child_stderr {
                let mut buf = vec![0u8; 4096];
                loop {
                    match stderr.read(&mut buf).await {
                        Ok(0) | Err(_) => break,
                        Ok(_) => {} // discard
                    }
                }
            }
        });

        // Reaper task: waits for exit, updates status, signals watchers, delayed cleanup
        let status_clone = status.clone();
        let procs_clone = Arc::clone(&self.procs);
        let reaper_task = tokio::spawn(async move {
            let exit_status = child.wait().await;
            let code = exit_status.map(|s| s.code().unwrap_or(-1)).unwrap_or(-1);
            *status_clone.lock().await = format!("exited({})", code);
            let _ = exit_tx.send(true);
            // Delayed cleanup: give stdout readers time to drain buffered data
            tokio::time::sleep(Duration::from_secs(5)).await;
            procs_clone.write().await.remove(&pid);
        });

        let managed = Arc::new(ManagedProcess {
            cmd: cmd.to_string(),
            status,
            stdin_tx: Some(stdin_tx),
            stdout_rx: Arc::new(Mutex::new(stdout_rx)),
            exit_rx,
            _tasks: vec![stdin_task, stdout_task, stderr_task, reaper_task],
        });

        self.procs.write().await.insert(pid, managed);
        tracing::info!("process spawned: pid={} cmd={}", pid, cmd);
        Ok(pid)
    }

    pub async fn kill(&self, pid: u64) -> Result<(), String> {
        let proc = {
            let procs = self.procs.read().await;
            procs
                .get(&pid)
                .ok_or_else(|| format!("pid {} not found", pid))?
                .clone()
        };

        // Set status to killed first
        *proc.status.lock().await = "killed".to_string();

        // Abort all background tasks (stdin writer, stdout reader, stderr drain, reaper)
        // Aborting reaper drops exit_tx, so any wait() blocked on changed() will wake with Err.
        for task in proc._tasks.iter() {
            task.abort();
        }

        // Remove from map
        self.procs.write().await.remove(&pid);

        // Also send OS kill
        #[cfg(unix)]
        {
            let os_pid = pid as i32;
            unsafe { libc::kill(os_pid as libc::pid_t, 9) };
        }
        #[cfg(not(unix))]
        {
            // On Windows, aborting the reaper task triggers kill_on_drop on the Child
        }

        tracing::info!("process killed: pid={}", pid);
        Ok(())
    }

    pub async fn wait(&self, pid: u64) -> Result<i32, String> {
        let proc = {
            let procs = self.procs.read().await;
            procs
                .get(&pid)
                .ok_or_else(|| format!("pid {} not found", pid))?
                .clone()
        };

        // Use watch channel: no race between checking status and waiting.
        // watch retains the latest value, so even if send(true) happened before
        // we subscribe, borrow_and_update() will see `true` immediately.
        let mut rx = proc.exit_rx.clone();
        while !*rx.borrow_and_update() {
            if rx.changed().await.is_err() {
                break; // sender dropped (e.g. kill aborted reaper)
            }
        }

        let status = proc.status.lock().await.clone();
        parse_exit_code(&status)
    }

    pub async fn write_stdin(&self, pid: u64, data: &[u8]) -> Result<(), String> {
        let proc = {
            let procs = self.procs.read().await;
            procs
                .get(&pid)
                .ok_or_else(|| format!("pid {} not found", pid))?
                .clone()
        };

        if let Some(ref tx) = proc.stdin_tx {
            tx.send(data.to_vec())
                .await
                .map_err(|_| "stdin channel closed".to_string())
        } else {
            Err("stdin not available".to_string())
        }
    }

    pub async fn read_stdout(&self, pid: u64, buf: &mut [u8]) -> Result<usize, String> {
        let proc = {
            let procs = self.procs.read().await;
            procs
                .get(&pid)
                .ok_or_else(|| format!("pid {} not found", pid))?
                .clone()
        };

        let mut rx = proc.stdout_rx.lock().await;
        match rx.recv().await {
            Some(data) => {
                let n = std::cmp::min(data.len(), buf.len());
                buf[..n].copy_from_slice(&data[..n]);
                Ok(n)
            }
            None => Ok(0), // EOF
        }
    }

    pub async fn signal(&self, pid: u64, signal: i32) -> Result<(), String> {
        if signal == 9 {
            return self.kill(pid).await;
        }

        let procs = self.procs.read().await;
        if !procs.contains_key(&pid) {
            return Err(format!("pid {} not found", pid));
        }
        drop(procs);

        #[cfg(unix)]
        {
            let os_pid = i32::try_from(pid)
                .map_err(|_| format!("pid {} out of range for platform signal API", pid))?;
            let result = unsafe { libc::kill(os_pid as libc::pid_t, signal as libc::c_int) };
            if result == 0 {
                Ok(())
            } else {
                Err(format!(
                    "signal {} to pid {} failed: {}",
                    signal,
                    pid,
                    std::io::Error::last_os_error()
                ))
            }
        }

        #[cfg(not(unix))]
        {
            Err(format!(
                "process.signal only supports signal 9 on this platform, got {}",
                signal
            ))
        }
    }

    pub async fn list(&self) -> Vec<(u64, String, String)> {
        let snapshot: Vec<(u64, Arc<ManagedProcess>)> = {
            let procs = self.procs.read().await;
            procs.iter().map(|(pid, mp)| (*pid, mp.clone())).collect()
        };
        // procs read lock is dropped here
        let mut result = Vec::with_capacity(snapshot.len());
        for (pid, mp) in &snapshot {
            let status = mp.status.lock().await.clone();
            result.push((*pid, mp.cmd.clone(), status));
        }
        result
    }
}

fn parse_exit_code(status: &str) -> Result<i32, String> {
    if let Some(code_str) = status
        .strip_prefix("exited(")
        .and_then(|s| s.strip_suffix(')'))
    {
        code_str.parse::<i32>().map_err(|e| e.to_string())
    } else {
        Ok(-1)
    }
}