megacode-agent-sdk 0.31.0

//! The session's process (spec slice-ar R1, R2.1, R8.6): the options file,
//! the start, the stderr tail, and the stop.

use std::collections::VecDeque;
use std::ffi::OsString;
use std::io::Write as _;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::{mpsc, watch};

use crate::error::{Diagnostics, Error};

/// The environment variable naming a harness binary (R12.1).
pub(crate) const BINARY_ENV: &str = "MEGACODE_AGENT_SDK_BINARY";

/// The lines of stderr kept for diagnostics (R1.4).
const STDERR_TAIL: usize = 200;

/// The options file in its own private folder (R2.1): removed, file then
/// folder, at `init` or at the process's exit, whichever is first.
pub(crate) struct OptionsFile {
    dir: PathBuf,
    file: PathBuf,
    removed: AtomicBool,
}

impl OptionsFile {
    /// Writes `bytes` to `options.json` in a new folder (mode `0700`) under
    /// the temporary folder; the file is created exclusively, mode `0600`.
    pub(crate) fn write(bytes: &[u8]) -> std::io::Result<Self> {
        let base = std::env::temp_dir();
        let mut tries = 0;
        let dir = loop {
            let dir = base.join(format!(
                "megacode-agent-sdk-{}-{:016x}",
                std::process::id(),
                unique()
            ));
            let mut builder = std::fs::DirBuilder::new();
            #[cfg(unix)]
            std::os::unix::fs::DirBuilderExt::mode(&mut builder, 0o700);
            match builder.create(&dir) {
                Ok(()) => break dir,
                Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists && tries < 100 => {
                    tries += 1;
                }
                Err(e) => return Err(e),
            }
        };
        let file = dir.join("options.json");
        let made = Self {
            dir,
            file,
            removed: AtomicBool::new(false),
        };
        let mut open = std::fs::OpenOptions::new();
        open.write(true).create_new(true);
        #[cfg(unix)]
        std::os::unix::fs::OpenOptionsExt::mode(&mut open, 0o600);
        // On failure `made` is dropped, which removes what was made.
        open.open(&made.file)?.write_all(bytes)?;
        Ok(made)
    }

    pub(crate) fn path(&self) -> &Path {
        &self.file
    }

    /// Removes the file and its folder; later calls do nothing.
    pub(crate) fn remove(&self) {
        if !self.removed.swap(true, Ordering::SeqCst) {
            let _ = std::fs::remove_file(&self.file);
            let _ = std::fs::remove_dir(&self.dir);
        }
    }
}

impl Drop for OptionsFile {
    fn drop(&mut self) {
        self.remove();
    }
}

/// A value no other call in this process returns, and unlikely to be
/// another process's: a counter mixed with std's per-process random keys.
pub(crate) fn unique() -> u64 {
    use std::hash::{BuildHasher, Hasher};
    static COUNT: AtomicU64 = AtomicU64::new(0);
    let mut h = std::collections::hash_map::RandomState::new().build_hasher();
    h.write_u64(COUNT.fetch_add(1, Ordering::Relaxed));
    h.finish()
}

/// How the process ended.
#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct Exit {
    pub code: Option<i32>,
    pub signal: Option<i32>,
}

/// A request to the task that owns the child.
#[derive(Debug)]
pub(crate) enum Signal {
    Term,
    Kill,
}

/// The parts of a started process the session keeps.
pub(crate) struct Process {
    pub stdin: ChildStdin,
    pub stdout: BufReader<ChildStdout>,
    /// Signals for the task that owns the child; closing it kills the child.
    pub ctl: mpsc::UnboundedSender<Signal>,
    /// `Some` once the child was reaped.
    pub exit: watch::Receiver<Option<Exit>>,
    pub stderr: Arc<Mutex<VecDeque<String>>>,
    /// `true` once stderr reached its end.
    pub stderr_done: watch::Receiver<bool>,
}

/// What to start.
pub(crate) struct Launch<'a> {
    pub binary: &'a Path,
    pub resume: Option<&'a str>,
    pub mock: Option<&'a Path>,
    pub cwd: Option<&'a Path>,
    pub env: &'a [(OsString, OsString)],
    pub options: Arc<OptionsFile>,
}

/// Starts `harness run|resume <id> [--mock <file>] --input-format
/// stream-json --options <file>` (R1.1).
pub(crate) async fn spawn(launch: Launch<'_>) -> Result<Process, Error> {
    let mut cmd = Command::new(launch.binary);
    match launch.resume {
        Some(id) => cmd.arg("resume").arg(id),
        None => cmd.arg("run"),
    };
    if let Some(mock) = launch.mock {
        cmd.arg("--mock").arg(mock);
    }
    cmd.args(["--input-format", "stream-json", "--options"])
        .arg(launch.options.path())
        .envs(launch.env.iter().map(|(k, v)| (k, v)))
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .kill_on_drop(true);
    if let Some(cwd) = launch.cwd {
        cmd.current_dir(cwd);
    }
    // A binary written moments ago can still be open for writing in a child
    // another thread forked in between (`ETXTBSY`); that passes.
    let mut tries = 0;
    let mut child = loop {
        match cmd.spawn() {
            Ok(child) => break child,
            #[cfg(unix)]
            Err(e) if e.raw_os_error() == Some(libc::ETXTBSY) && tries < 50 => {
                tries += 1;
                tokio::time::sleep(Duration::from_millis(20)).await;
            }
            Err(e) => {
                return Err(Error::BinaryUnavailable(format!(
                    "cannot run {}: {e}",
                    launch.binary.display()
                )));
            }
        }
    };
    let stdin = child.stdin.take().expect("stdin is piped");
    let stdout = BufReader::new(child.stdout.take().expect("stdout is piped"));
    let stderr_pipe = child.stderr.take().expect("stderr is piped");

    let stderr = Arc::new(Mutex::new(VecDeque::new()));
    let (stderr_tx, stderr_done) = watch::channel(false);
    let tail = stderr.clone();
    tokio::spawn(async move {
        let mut lines = BufReader::new(stderr_pipe);
        let mut buf = Vec::new();
        loop {
            buf.clear();
            match lines.read_until(b'\n', &mut buf).await {
                Ok(0) | Err(_) => break,
                Ok(_) => {
                    let line = String::from_utf8_lossy(&buf).trim_end().to_string();
                    let mut tail = tail.lock().expect("stderr tail lock");
                    if tail.len() == STDERR_TAIL {
                        tail.pop_front();
                    }
                    tail.push_back(line);
                }
            }
        }
        let _ = stderr_tx.send(true);
    });

    let (ctl, ctl_rx) = mpsc::unbounded_channel();
    let (exit_tx, exit) = watch::channel(None);
    tokio::spawn(own(child, ctl_rx, exit_tx, launch.options));
    Ok(Process {
        stdin,
        stdout,
        ctl,
        exit,
        stderr,
        stderr_done,
    })
}

/// Owns the child until it is reaped: signals it on request, kills it when
/// every sender is gone, removes the options file at its exit.
async fn own(
    mut child: Child,
    mut ctl: mpsc::UnboundedReceiver<Signal>,
    exit: watch::Sender<Option<Exit>>,
    options: Arc<OptionsFile>,
) {
    let status = loop {
        tokio::select! {
            status = child.wait() => break status.ok(),
            signal = ctl.recv() => match signal {
                Some(Signal::Term) => terminate(&mut child),
                Some(Signal::Kill) => { let _ = child.start_kill(); }
                None => {
                    let _ = child.start_kill();
                    break child.wait().await.ok();
                }
            },
        }
    };
    options.remove();
    let mut out = Exit::default();
    if let Some(status) = status {
        out.code = status.code();
        #[cfg(unix)]
        {
            out.signal = std::os::unix::process::ExitStatusExt::signal(&status);
        }
    }
    let _ = exit.send(Some(out));
}

/// SIGTERM to a child not yet reaped (its pid is still its own).
fn terminate(child: &mut Child) {
    #[cfg(unix)]
    if let Some(pid) = child.id() {
        // SAFETY: `kill` takes plain integers; the pid is this unreaped child's.
        unsafe {
            libc::kill(pid as libc::pid_t, libc::SIGTERM);
        }
        return;
    }
    let _ = child.start_kill();
}

/// Waits until the process was reaped, at most `limit`; `true` if it was.
pub(crate) async fn wait_exit(exit: &mut watch::Receiver<Option<Exit>>, limit: Duration) -> bool {
    tokio::time::timeout(limit, exit.wait_for(Option::is_some))
        .await
        .is_ok_and(|r| r.is_ok())
}

/// The stop (R8.6.2), after the caller sent `quit` or closed stdin: up to
/// 5 s for the exit, then SIGTERM and up to 5 s, then SIGKILL; reaped.
pub(crate) async fn stop(
    ctl: &mpsc::UnboundedSender<Signal>,
    exit: &mut watch::Receiver<Option<Exit>>,
    first_wait: Duration,
) {
    if wait_exit(exit, first_wait).await {
        return;
    }
    let _ = ctl.send(Signal::Term);
    if wait_exit(exit, Duration::from_secs(5)).await {
        return;
    }
    let _ = ctl.send(Signal::Kill);
    wait_exit(exit, Duration::from_secs(4)).await;
}

/// The diagnostics of an ended process: its exit and stderr tail, each
/// waited for briefly.
pub(crate) async fn diagnostics(
    exit: &mut watch::Receiver<Option<Exit>>,
    stderr: &Mutex<VecDeque<String>>,
    stderr_done: &mut watch::Receiver<bool>,
) -> Diagnostics {
    wait_exit(exit, Duration::from_secs(2)).await;
    let _ = tokio::time::timeout(Duration::from_secs(1), stderr_done.wait_for(|d| *d)).await;
    let ended = (*exit.borrow()).unwrap_or_default();
    Diagnostics {
        exit_code: ended.code,
        signal: ended.signal,
        stderr_tail: stderr
            .lock()
            .expect("stderr tail lock")
            .iter()
            .cloned()
            .collect(),
    }
}