sloop/runner/

local.rs

use std::fs;
use std::io::{self, Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::process::CommandExt;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::sync::Arc;

use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD as BASE64_TOKEN;
use tokio::net::UnixListener;

use super::{
    AgentProcessCheckpoint, ExecProcessCheckpoint, ExecutionEvidence, ExecutionFailure,
    ProcessIdentity, RunnerError, StageExecution, StageHooks, StageOrder, WorkerCredentials,
    WorkerScope,
};
use crate::clock::Clock;
use crate::run_log::{OutputSource, OutputStream, RunLogWriter};

/// An agent process under supervision: spawned, checkpointed, and being
/// drained of output until it exits.
pub struct LaunchedAgent {
    child: Child,
    readers: Vec<std::thread::JoinHandle<bool>>,
    process: ProcessIdentity,
    started_at_ms: i64,
}

pub struct AgentCompletion {
    pub evidence: ExecutionEvidence,
    pub wait_error: Option<String>,
}

impl LaunchedAgent {
    pub fn process(&self) -> ProcessIdentity {
        self.process
    }

    pub fn wait(mut self, clock: &dyn Clock) -> AgentCompletion {
        let (exit_code, wait_error) = match self.child.wait() {
            Ok(status) => (status.code(), None),
            Err(error) => (None, Some(error.to_string())),
        };
        let stragglers_killed = kill_straggler_process_group(self.process.process_group_id);
        let mut output_capture_complete = true;
        for reader in self.readers {
            output_capture_complete &= reader.join().unwrap_or(false);
        }
        AgentCompletion {
            evidence: ExecutionEvidence {
                exit_code,
                started_at_ms: self.started_at_ms,
                finished_at_ms: clock.now_ms(),
                process: Some(self.process),
                output_capture_complete,
                stragglers_killed,
            },
            wait_error,
        }
    }
}

/// Creates the isolated workspace a run executes every one of its stages in.
/// The driver does this once, before the first stage, because a flow need not
/// open with an agent action.
pub fn create_run_worktree(repository: &Path, worktree: &Path, branch: &str) -> Result<(), String> {
    if worktree.exists() {
        return Ok(());
    }
    fs::create_dir_all(
        worktree
            .parent()
            .ok_or_else(|| "worktree has no parent".to_owned())?,
    )
    .map_err(|error| error.to_string())?;
    let git = Command::new("git")
        .args(["worktree", "add", "--quiet", "-b", branch])
        .arg(worktree)
        .current_dir(repository)
        .output()
        .map_err(|error| error.to_string())?;
    if git.status.success() {
        Ok(())
    } else {
        Err(format!(
            "git worktree add failed: {}",
            String::from_utf8_lossy(&git.stderr).trim()
        ))
    }
}

/// Binds a worker socket and mints the token that authenticates against it,
/// scoped to what the token may do. Returned unregistered: the caller decides
/// who serves the socket, and no request can be answered before it does.
pub fn mint_worker_credentials(
    socket_path: &Path,
    scope: WorkerScope,
) -> Result<(WorkerCredentials, UnixListener), String> {
    let token = generate_worker_token()?;
    fs::create_dir_all(socket_path.parent().expect("worker sockets have a parent"))
        .map_err(|error| error.to_string())?;
    let _ = fs::remove_file(socket_path);
    let listener = UnixListener::bind(socket_path).map_err(|error| error.to_string())?;
    fs::set_permissions(socket_path, fs::Permissions::from_mode(0o600))
        .map_err(|error| error.to_string())?;
    Ok((
        WorkerCredentials {
            socket: socket_path.to_path_buf(),
            token,
            scope,
        },
        listener,
    ))
}

pub fn launch_agent<H: StageHooks>(
    order: StageOrder,
    hooks: &H,
    clock: Arc<dyn Clock>,
    output_observer: Arc<dyn Fn() + Send + Sync>,
) -> Result<LaunchedAgent, RunnerError<H::Error>> {
    let StageExecution::Agent(launch) = &order.execution else {
        return Err(RunnerError::Execution(
            "agent launch requires an agent stage order".into(),
        ));
    };
    let Some(program) = launch.argv.first() else {
        return Err(RunnerError::Execution("agent argv is empty".into()));
    };

    let output_log = RunLogWriter::open(&order.output_path)
        .map_err(|error| RunnerError::Execution(error.to_string()))?;
    let worker = launch.worker.clone();

    let mut command = Command::new(program);
    command
        .args(&launch.argv[1..])
        .current_dir(&order.worktree)
        .envs(launch.environment.iter().cloned())
        .env("SLOOP_SOCKET", &worker.socket)
        .env("SLOOP_TOKEN", &worker.token)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0);
    let mut child = command
        .spawn()
        .map_err(|error| RunnerError::Execution(error.to_string()))?;
    let readers = vec![
        spawn_output_reader(
            child.stdout.take().expect("stdout was piped"),
            output_log.clone(),
            OutputSource::Agent,
            Some(order.stage.clone()),
            order.attempt,
            OutputStream::Stdout,
            Some(clock.clone()),
            Some(output_observer.clone()),
        ),
        spawn_output_reader(
            child.stderr.take().expect("stderr was piped"),
            output_log,
            OutputSource::Agent,
            Some(order.stage.clone()),
            order.attempt,
            OutputStream::Stderr,
            Some(clock.clone()),
            Some(output_observer),
        ),
    ];

    let pid = child.id();
    let process = ProcessIdentity {
        pid,
        start_time: process_start_time(pid),
        process_group_id: pid,
    };
    let started_at_ms = clock.now_ms();
    let checkpoint = AgentProcessCheckpoint {
        run_id: order.run_id,
        attempt: order.attempt,
        stage: order.stage,
        branch: order.branch,
        worktree: order.worktree,
        process,
        worker,
        started_at_ms,
    };
    wait_for_test_hook(&format!(
        "before-stage-process-checkpoint-{}",
        checkpoint.stage
    ));
    if let Err(error) = hooks.record_agent_process(&checkpoint) {
        kill_process_group(process);
        let _ = child.wait();
        for reader in readers {
            let _ = reader.join();
        }
        return Err(RunnerError::Hook(error));
    }

    Ok(LaunchedAgent {
        child,
        readers,
        process,
        started_at_ms,
    })
}

pub fn run_exec_stage<H: StageHooks>(
    order: &StageOrder,
    hooks: &H,
    clock: &dyn Clock,
) -> Result<ExecutionEvidence, ExecutionFailure<H::Error>> {
    let started_at_ms = clock.now_ms();
    let failed = |process, error| ExecutionFailure {
        evidence: ExecutionEvidence {
            exit_code: None,
            started_at_ms,
            finished_at_ms: clock.now_ms(),
            process,
            output_capture_complete: false,
            stragglers_killed: false,
        },
        error,
    };
    let StageExecution::Exec(launch) = &order.execution else {
        return Err(failed(
            None,
            RunnerError::Execution("exec launch requires an exec stage order".into()),
        ));
    };
    let Some(program) = launch.argv.first() else {
        return Err(failed(
            None,
            RunnerError::Execution("exec argv is empty".into()),
        ));
    };
    if hooks.cancellation_requested(&order.run_id) {
        return Err(failed(
            None,
            RunnerError::Execution("stage cancelled before launch".into()),
        ));
    }
    let output_log = RunLogWriter::open(&order.output_path).map_err(|error| {
        failed(
            None,
            RunnerError::Execution(format!("cannot open run output: {error}")),
        )
    })?;

    let mut command = if launch.worker.is_some() {
        let mut command = Command::new("sh");
        command
            .args([
                "-c",
                "IFS= read -r _ || exit 125; exec \"$@\"",
                "sloop-stage",
            ])
            .args(&launch.argv);
        command
    } else {
        let mut command = Command::new(program);
        command.args(&launch.argv[1..]);
        command
    };
    command
        .current_dir(&order.worktree)
        .envs(launch.environment.iter().cloned())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0);
    if let Some(worker) = &launch.worker {
        command
            .env("SLOOP_SOCKET", &worker.socket)
            .env("SLOOP_TOKEN", &worker.token)
            .stdin(Stdio::piped());
    } else {
        command
            .env_remove("SLOOP_SOCKET")
            .env_remove("SLOOP_TOKEN")
            .stdin(Stdio::null());
    }
    let mut child = command
        .spawn()
        .map_err(|error| failed(None, RunnerError::Execution(error.to_string())))?;
    let mut gate = launch
        .worker
        .as_ref()
        .map(|_| child.stdin.take().expect("reported stage stdin was piped"));
    let pid = child.id();
    let process = ProcessIdentity {
        pid,
        start_time: process_start_time(pid),
        process_group_id: pid,
    };
    let Some(start_time) = process.start_time else {
        kill_process_group(process);
        let _ = child.wait();
        return Err(failed(
            Some(process),
            RunnerError::Execution("cannot identify exec process".into()),
        ));
    };
    let readers = vec![
        spawn_output_reader(
            child.stdout.take().expect("stdout was piped"),
            output_log.clone(),
            OutputSource::Stage,
            Some(order.stage.clone()),
            order.attempt,
            OutputStream::Stdout,
            None,
            None,
        ),
        spawn_output_reader(
            child.stderr.take().expect("stderr was piped"),
            output_log,
            OutputSource::Stage,
            Some(order.stage.clone()),
            order.attempt,
            OutputStream::Stderr,
            None,
            None,
        ),
    ];
    if order.stage == "test" {
        wait_for_test_hook("before-test-process-checkpoint");
    }
    wait_for_test_hook(&format!("before-stage-process-checkpoint-{}", order.stage));
    let checkpoint = ExecProcessCheckpoint {
        run_id: order.run_id.clone(),
        stage: order.stage.clone(),
        attempt: order.attempt,
        process: ProcessIdentity {
            start_time: Some(start_time),
            ..process
        },
        started_at_ms: clock.now_ms(),
    };
    if let Err(error) = hooks.record_exec_process(&checkpoint) {
        kill_process_group_if_matches(checkpoint.process);
        let _ = child.wait();
        join_readers(readers);
        return Err(failed(Some(process), RunnerError::Hook(error)));
    }
    wait_for_test_hook(&format!("after-stage-process-checkpoint-{}", order.stage));
    if hooks.cancellation_requested(&order.run_id) {
        kill_process_group_if_matches(checkpoint.process);
        let _ = child.wait();
        join_readers(readers);
        return Err(failed(
            Some(process),
            RunnerError::Execution("stage cancelled after launch".into()),
        ));
    }
    if let Some(gate) = gate.as_mut()
        && gate.write_all(b"run\n").is_err()
    {
        kill_process_group_if_matches(checkpoint.process);
        let _ = child.wait();
        join_readers(readers);
        return Err(failed(
            Some(process),
            RunnerError::Execution("cannot release exec process".into()),
        ));
    }
    drop(gate);

    let status = child.wait();
    let stragglers_killed = kill_straggler_process_group(pid);
    let output_capture_complete = join_readers(readers);
    if !output_capture_complete {
        return Err(ExecutionFailure {
            evidence: ExecutionEvidence {
                exit_code: None,
                started_at_ms,
                finished_at_ms: clock.now_ms(),
                process: Some(process),
                output_capture_complete: false,
                stragglers_killed,
            },
            error: RunnerError::Execution("exec output capture was incomplete".into()),
        });
    }
    let status = match status {
        Ok(status) => status,
        Err(error) => {
            return Err(ExecutionFailure {
                evidence: ExecutionEvidence {
                    exit_code: None,
                    started_at_ms,
                    finished_at_ms: clock.now_ms(),
                    process: Some(process),
                    output_capture_complete: true,
                    stragglers_killed,
                },
                error: RunnerError::Execution(error.to_string()),
            });
        }
    };
    Ok(ExecutionEvidence {
        exit_code: status.code(),
        started_at_ms,
        finished_at_ms: clock.now_ms(),
        process: Some(process),
        output_capture_complete: true,
        stragglers_killed,
    })
}

#[allow(clippy::too_many_arguments)]
fn spawn_output_reader(
    pipe: impl Read + Send + 'static,
    log: RunLogWriter,
    source: OutputSource,
    stage: Option<String>,
    attempt: u32,
    stream: OutputStream,
    clock: Option<Arc<dyn Clock>>,
    output_observer: Option<Arc<dyn Fn() + Send + Sync>>,
) -> std::thread::JoinHandle<bool> {
    std::thread::spawn(move || {
        let mut pipe = pipe;
        let mut buffer = [0u8; 8192];
        loop {
            match pipe.read(&mut buffer) {
                Ok(0) => return true,
                Ok(read) => {
                    let appended = match clock.as_ref() {
                        Some(clock) => log.append_at(
                            clock.now_ms(),
                            source,
                            stage.as_deref(),
                            Some(attempt),
                            stream,
                            &buffer[..read],
                        ),
                        None => log.append(
                            source,
                            stage.as_deref(),
                            Some(attempt),
                            stream,
                            &buffer[..read],
                        ),
                    };
                    if appended.is_err() {
                        return false;
                    }
                    if let Some(observer) = &output_observer {
                        observer();
                    }
                }
                Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
                Err(_) => return false,
            }
        }
    })
}

fn join_readers(readers: Vec<std::thread::JoinHandle<bool>>) -> bool {
    readers.into_iter().fold(true, |complete, reader| {
        complete & reader.join().unwrap_or(false)
    })
}

fn generate_worker_token() -> Result<String, String> {
    let mut bytes = [0u8; 32];
    let mut urandom = fs::File::open("/dev/urandom").map_err(|error| error.to_string())?;
    urandom
        .read_exact(&mut bytes)
        .map_err(|error| error.to_string())?;
    Ok(BASE64_TOKEN.encode(bytes))
}

pub fn run_output_path(state_dir: &Path, run_id: &str) -> PathBuf {
    state_dir.join("runs").join(run_id).join("output.ndjson")
}

pub fn worker_socket_path(runtime_dir: &Path, run_id: &str) -> PathBuf {
    // macOS caps Unix socket paths at 104 bytes, and the runtime root under
    // `$TMPDIR` is long there, so worker sockets sit directly in the runtime
    // directory and use the short run id. The full 32-hex id pushed a macOS
    // path past the cap and every agent spawn failed at bind.
    runtime_dir.join(format!("w{}.sock", crate::run_ref::short(run_id)))
}

pub fn process_start_time(pid: u32) -> Option<i64> {
    #[cfg(target_os = "linux")]
    {
        let stat = fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;
        let after_command = &stat[stat.rfind(')')? + 1..];
        after_command
            .split_whitespace()
            .nth(19)
            .and_then(|field| field.parse().ok())
    }

    #[cfg(not(target_os = "linux"))]
    {
        let output = Command::new("ps")
            .args(["-o", "lstart=", "-p", &pid.to_string()])
            .output()
            .ok()?;
        if !output.status.success() || output.stdout.iter().all(u8::is_ascii_whitespace) {
            return None;
        }
        let mut hash = 0xcbf29ce484222325_u64;
        for byte in output.stdout {
            hash ^= u64::from(byte);
            hash = hash.wrapping_mul(0x100000001b3);
        }
        Some(hash as i64)
    }
}

pub fn process_identity_matches(pid: u32, expected_start_time: Option<i64>) -> bool {
    matches!(
        (expected_start_time, process_start_time(pid)),
        (Some(expected), Some(actual)) if expected == actual
    )
}

pub fn kill_straggler_process_group(group: u32) -> bool {
    let group = -(group as libc::pid_t);
    let stragglers_present = unsafe { libc::kill(group, 0) } == 0;
    if stragglers_present {
        unsafe {
            libc::kill(group, libc::SIGKILL);
        }
    }
    stragglers_present
}

fn kill_process_group(process: ProcessIdentity) {
    unsafe {
        libc::kill(-(process.process_group_id as libc::pid_t), libc::SIGKILL);
    }
}

fn kill_process_group_if_matches(process: ProcessIdentity) {
    if process_identity_matches(process.pid, process.start_time) {
        kill_process_group(process);
    }
}

#[cfg(debug_assertions)]
pub fn wait_for_test_hook(name: &str) {
    use std::time::Duration;

    let Some(directory) = std::env::var_os("SLOOP_TEST_HOOK_DIR").map(PathBuf::from) else {
        return;
    };
    let armed = directory.join(format!("{name}.armed"));
    if !armed.is_file() {
        return;
    }
    let reached = directory.join(format!("{name}.reached"));
    let release = directory.join(format!("{name}.release"));
    if fs::write(&reached, b"").is_err() {
        return;
    }
    while !release.is_file() {
        std::thread::sleep(Duration::from_millis(10));
    }
}

#[cfg(not(debug_assertions))]
pub fn wait_for_test_hook(_name: &str) {}

#[cfg(test)]
mod tests {
    use std::convert::Infallible;

    use tempfile::tempdir;

    use super::run_exec_stage;
    use crate::clock::SystemClock;
    use crate::config::{AgentTarget, expand_agent_cmd};
    use crate::runner::{
        AgentProcessCheckpoint, ExecLaunch, ExecProcessCheckpoint, StageExecution, StageHooks,
        StageOrder,
    };

    struct NoopHooks;

    impl StageHooks for NoopHooks {
        type Error = Infallible;

        fn cancellation_requested(&self, _run_id: &str) -> bool {
            false
        }

        fn record_agent_process(
            &self,
            _checkpoint: &AgentProcessCheckpoint,
        ) -> Result<(), Self::Error> {
            Ok(())
        }

        fn record_exec_process(
            &self,
            _checkpoint: &ExecProcessCheckpoint,
        ) -> Result<(), Self::Error> {
            Ok(())
        }
    }

    fn target(cmd: &[&str], model: Option<&str>, effort: Option<&str>) -> AgentTarget {
        AgentTarget {
            cmd: cmd.iter().map(|argument| (*argument).to_owned()).collect(),
            model: model.map(str::to_owned),
            effort: effort.map(str::to_owned),
        }
    }

    #[test]
    fn agent_command_expands_ticket_model_and_effort() {
        let template = target(
            &[
                "agent",
                "--model={model}",
                "--effort",
                "{effort}",
                "prompt={prompt}",
            ],
            None,
            None,
        );

        assert_eq!(
            expand_agent_cmd(&template, Some("sonnet"), Some("medium"), "assignment").unwrap(),
            [
                "agent",
                "--model=sonnet",
                "--effort",
                "medium",
                "prompt=assignment"
            ]
        );
    }

    #[test]
    fn agent_command_rejects_a_missing_ticket_field() {
        let template = target(&["agent", "{model}"], None, Some("medium"));

        assert_eq!(
            expand_agent_cmd(&template, None, Some("medium"), "assignment"),
            Err("does not specify `model`".to_owned())
        );
    }

    #[test]
    fn agent_command_falls_back_to_target_defaults() {
        let template = target(
            &["agent", "{model}", "{effort}", "{prompt}"],
            Some("opus"),
            Some("high"),
        );

        assert_eq!(
            expand_agent_cmd(&template, None, None, "assignment").unwrap(),
            ["agent", "opus", "high", "assignment"]
        );
    }

    #[test]
    fn agent_command_prefers_ticket_values_over_target_defaults() {
        let template = target(
            &["agent", "{model}", "{effort}", "{prompt}"],
            Some("opus"),
            Some("high"),
        );

        assert_eq!(
            expand_agent_cmd(&template, Some("haiku"), Some("low"), "assignment").unwrap(),
            ["agent", "haiku", "low", "assignment"]
        );
    }

    #[test]
    fn scripted_exec_returns_factual_evidence_without_a_daemon_or_store() {
        let directory = tempdir().unwrap();
        let order = StageOrder {
            run_id: "R1".into(),
            stage: "check".into(),
            attempt: 1,
            execution: StageExecution::Exec(ExecLaunch {
                argv: vec!["sh".into(), "-c".into(), "printf runner-output".into()],
                worker: None,
                environment: Vec::new(),
            }),
            worktree: directory.path().into(),
            branch: "sloop/T1-a1-R1".into(),
            output_path: directory.path().join("output.ndjson"),
        };

        let evidence = run_exec_stage(&order, &NoopHooks, &SystemClock).unwrap();

        assert_eq!(evidence.exit_code, Some(0));
        assert!(evidence.output_capture_complete);
        assert!(evidence.process.is_some());
    }

    #[test]
    fn worker_socket_paths_fit_the_macos_socket_length_cap() {
        let runtime_dir = std::path::Path::new(
            "/var/folders/24/8k48jl6d249_n_qfxwsl6xvm0000gn/T/.tmpAbC123/sloop/0123456789abcdef",
        );
        let socket = super::worker_socket_path(runtime_dir, &"ab".repeat(16));

        assert!(
            socket.as_os_str().len() <= 104,
            "worker socket path is {} bytes, over the 104-byte macOS cap: {}",
            socket.as_os_str().len(),
            socket.display()
        );
    }
}