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, Mutex, RwLock};

/// 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>>>>,
    /// Notified when the process exits. Clone the receiver to wait.
    exit_notify: Arc<tokio::sync::Notify>,
    /// 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: RwLock<HashMap<u64, Arc<ManagedProcess>>>,
}

impl ProcessSyscall {
    pub fn new() -> Self {
        Self {
            procs: 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_notify = Arc::new(tokio::sync::Notify::new());

        // 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 and updates status
        let status_clone = status.clone();
        let exit_notify_clone = exit_notify.clone();
        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);
            exit_notify_clone.notify_waiters();
        });

        let managed = Arc::new(ManagedProcess {
            cmd: cmd.to_string(),
            status,
            stdin_tx: Some(stdin_tx),
            stdout_rx: Arc::new(Mutex::new(stdout_rx)),
            exit_notify,
            _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 procs = self.procs.read().await;
        let proc = procs
            .get(&pid)
            .ok_or_else(|| format!("pid {} not found", pid))?
            .clone();
        drop(procs);

        // Abort all background tasks (stdin writer, stdout reader, stderr drain, reaper)
        for task in proc._tasks.iter() {
            task.abort();
        }
        *proc.status.lock().await = "killed".to_string();
        proc.exit_notify.notify_waiters();

        // 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
        }

        self.procs.write().await.remove(&pid);
        tracing::info!("process killed: pid={}", pid);
        Ok(())
    }

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

        // Check if already exited before waiting
        {
            let status = proc.status.lock().await;
            if status.starts_with("exited") || *status == "killed" {
                let s = status.clone();
                drop(status);
                self.procs.write().await.remove(&pid);
                return parse_exit_code(&s);
            }
        }

        // Wait for exit notification
        proc.exit_notify.notified().await;
        let status = proc.status.lock().await.clone();

        // Clean up from map
        self.procs.write().await.remove(&pid);
        parse_exit_code(&status)
    }

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

        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 procs = self.procs.read().await;
        let proc = procs
            .get(&pid)
            .ok_or_else(|| format!("pid {} not found", pid))?
            .clone();
        drop(procs);

        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 procs = self.procs.read().await;
        let mut result = Vec::with_capacity(procs.len());
        for (pid, mp) in procs.iter() {
            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)
    }
}