path: root/crates/asist-core/src/proc.rs

//! Running a program with a deadline.
//!
//! Three places needed it with the same care (never leave a process hanging,
//! never go through a shell, tell «failed» from «took too long») and each had
//! its own version. Here it is, once.
//!
//! None of this goes through `sh`: the arguments go to `execve` as they are.
//! It matters more than it seems, because what ends up in them comes,
//! ultimately, from what the microphone heard.
//!
//! The error messages are in Spanish on purpose: they end up in tool results
//! that the model reads and speaks.

use std::ffi::OsStr;
use std::io::Read;
use std::path::Path;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

/// How it ended.
#[derive(Debug)]
pub struct Output {
    pub stdout: Vec<u8>,
    pub stderr: String,
    pub status: Option<i32>,
    pub took: Duration,
}

impl Output {
    pub fn success(&self) -> bool {
        self.status == Some(0)
    }

    /// Last line of stderr, which is where programs usually put the real
    /// reason.
    pub fn last_error_line(&self) -> &str {
        self.stderr
            .trim()
            .lines()
            .next_back()
            .unwrap_or("sin detalles")
    }
}

#[derive(Debug, thiserror::Error)]
pub enum ProcError {
    #[error("no se pudo ejecutar {program}: {source}")]
    Spawn {
        program: String,
        #[source]
        source: std::io::Error,
    },
    #[error("{program} tardó más de {} s", timeout.as_secs())]
    Timeout { program: String, timeout: Duration },
    #[error("fallo esperando a {program}: {source}")]
    Wait {
        program: String,
        #[source]
        source: std::io::Error,
    },
}

/// Launches `program`, waits up to `timeout` and returns whatever it wrote.
///
/// If the deadline passes, it kills the process before giving up: an ffmpeg
/// against a busy camera, or a search engine that does not answer, would stay
/// there forever.
pub fn run(
    program: impl AsRef<Path>,
    args: impl IntoIterator<Item = impl AsRef<OsStr>>,
    timeout: Duration,
    working_dir: Option<&Path>,
) -> Result<Output, ProcError> {
    let program = program.as_ref();
    let name = program.display().to_string();
    let started = Instant::now();

    let mut command = Command::new(program);
    command
        .args(args)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    if let Some(dir) = working_dir {
        command.current_dir(dir);
    }

    let mut child = command.spawn().map_err(|source| ProcError::Spawn {
        program: name.clone(),
        source,
    })?;

    // The pipes are drained on separate threads, not after waiting.
    //
    // This is not an optimization: a child that writes more than fits in the
    // pipe buffer (64 KB on Linux) blocks writing, never finishes, and the wait
    // loop ends up killing it for missing the deadline even though it was doing
    // its job. It was hard to find because the first three uses (a date, a
    // search JSON, a 9 KB frame) fit easily; the first 180 KB screenshot did not.
    let mut stdout = child.stdout.take();
    let mut stderr = child.stderr.take();
    let stdout_reader = std::thread::spawn(move || {
        let mut buf = Vec::new();
        if let Some(pipe) = stdout.as_mut() {
            let _ = pipe.read_to_end(&mut buf);
        }
        buf
    });
    let stderr_reader = std::thread::spawn(move || {
        let mut buf = Vec::new();
        if let Some(pipe) = stderr.as_mut() {
            let _ = pipe.read_to_end(&mut buf);
        }
        buf
    });

    let deadline = started + timeout;
    let status = loop {
        match child.try_wait() {
            Ok(Some(status)) => break status,
            Ok(None) if Instant::now() >= deadline => {
                // Killing closes the pipes, so the readers finish and can be
                // joined without hanging.
                let _ = child.kill();
                let _ = child.wait();
                let _ = stdout_reader.join();
                let _ = stderr_reader.join();
                return Err(ProcError::Timeout {
                    program: name,
                    timeout,
                });
            }
            Ok(None) => std::thread::sleep(Duration::from_millis(20)),
            Err(source) => {
                return Err(ProcError::Wait {
                    program: name,
                    source,
                })
            }
        }
    };

    let stdout = stdout_reader.join().unwrap_or_default();
    let stderr = stderr_reader.join().unwrap_or_default();

    Ok(Output {
        stdout,
        stderr: String::from_utf8_lossy(&stderr).into_owned(),
        status: status.code(),
        took: started.elapsed(),
    })
}

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

    #[test]
    fn stdout_is_collected() {
        let out = run("echo", ["hola"], Duration::from_secs(5), None).unwrap();
        assert!(out.success());
        assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "hola");
    }

    #[test]
    fn hung_process_is_killed_at_the_deadline() {
        let started = Instant::now();
        let err = run("sleep", ["30"], Duration::from_millis(300), None).unwrap_err();
        assert!(matches!(err, ProcError::Timeout { .. }), "{err}");
        assert!(
            started.elapsed() < Duration::from_secs(3),
            "should have stopped right away, took {:?}",
            started.elapsed()
        );
    }

    #[test]
    fn large_output_does_not_block_the_child() {
        // The regression that hung the screenshot: without draining the pipe
        // while waiting, a child that writes more than 64 KB blocks and ends up
        // dying for missing the deadline.
        let out = run(
            "dd",
            ["if=/dev/zero", "bs=1024", "count=512", "status=none"],
            Duration::from_secs(10),
            None,
        )
        .expect("the deadline should not have passed");
        assert!(out.success());
        assert_eq!(out.stdout.len(), 512 * 1024);
    }

    #[test]
    fn large_stderr_does_not_block_either() {
        // Same problem on the other pipe: lots of noise on stderr and little
        // output is exactly what ffmpeg does when something goes wrong.
        let out = run(
            "sh",
            ["-c", "yes error | head -c 200000 >&2; echo ok"],
            Duration::from_secs(10),
            None,
        )
        .expect("the deadline should not have passed");
        assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "ok");
        assert!(out.stderr.len() > 100_000);
    }

    #[test]
    fn missing_program_error_includes_its_name() {
        let err = run("/no/existe/nada", ["x"], Duration::from_secs(1), None).unwrap_err();
        assert!(err.to_string().contains("/no/existe/nada"), "{err}");
    }

    #[test]
    fn failure_keeps_the_code_and_message() {
        let out = run("ls", ["/no/existe"], Duration::from_secs(5), None).unwrap();
        assert!(!out.success());
        assert!(!out.last_error_line().is_empty());
    }

    #[test]
    fn no_shell_interprets_metacharacters() {
        let out = run("echo", ["a; echo b"], Duration::from_secs(5), None).unwrap();
        assert_eq!(
            String::from_utf8_lossy(&out.stdout).trim(),
            "a; echo b",
            "the semicolon must arrive as text, not as a command separator"
        );
    }
}