scv-tools 0.4.7

//! Long-running delegated agents: one child process per conversation, spoken
//! to in newline-delimited JSON over its stdin and stdout.
//!
//! [`LiveChild`] is protocol-neutral. It starts the child in its adapter
//! environment and process group, records it as a delegation (so `scv agents
//! ps`, `kill`, reconcile, and orphan reaping apply), frames its stdout into
//! bounded lines, and shuts it down: closing stdin, a grace period, then a
//! group kill. A protocol client such as the SCV one in `scv_agent` (or an
//! ACP JSON-RPC client) runs on top of it.
//!
//! A conversation keeps its child between turns, and its delegation record
//! says so ([`LiveChild::begin_turn`]); the protocol client decides whether
//! anything reads from the child then.
//! A reaper task owns the process and waits for it from the start: whenever
//! the child exits, whether it ends by itself or `scv agents kill` stops it,
//! the reaper collects it at once, stops what is left of its group, and
//! removes its delegation record.

use std::{
    ffi::OsString,
    os::unix::process::CommandExt as _,
    path::PathBuf,
    sync::{
        Arc, Mutex as StdMutex,
        atomic::{AtomicBool, Ordering},
    },
    time::Duration,
};

use scv_core::ToolError;
use scv_protocol::{Frame, FrameDecoder, Overflow};
use serde::Serialize;
use tokio::{
    io::{AsyncWriteExt as _, BufReader},
    process::{Child, ChildStdin, ChildStdout, Command},
    sync::{Mutex, mpsc, watch},
    time::Instant,
};
use tokio_util::sync::CancellationToken;

use crate::{
    delegate::{
        activity::RunActivity,
        output::TailBuffer,
        records::{
            self, DelegationGuard, DelegationRegistry, PendingDelegation, STOP_GRACE, group_exists,
        },
    },
    process::{ProcessGroup, apply_agent_environment, drain_output, until},
    sync::lock,
};

/// Last bytes of a live child's stderr kept for failure reports.
const STDERR_TAIL_BYTES: usize = 4096;
/// Lines read ahead of the protocol client, a burst of events within one
/// turn. A child that speaks between turns needs a client reading then.
const LINE_QUEUE: usize = 256;
/// How long a message the protocol owes the child may take to write before
/// the child counts as not reading its input.
pub(crate) const SEND_GRACE: Duration = Duration::from_secs(2);

/// How to start a live child.
pub(crate) struct LiveSpec {
    pub(crate) executable: OsString,
    pub(crate) args: Vec<OsString>,
    pub(crate) cwd: PathBuf,
    /// The adapter environment plus the delegation tags.
    pub(crate) environment: Vec<(OsString, OsString)>,
    /// Longest stdout line accepted; a longer one ends the child.
    pub(crate) max_line_bytes: usize,
    /// The run's activity log, ended when the child exits.
    pub(crate) activity: Option<RunActivity>,
}

/// One line from a live child's stdout.
#[derive(Debug)]
pub(crate) enum LiveLine {
    Line(Vec<u8>),
    /// A line longer than [`LiveSpec::max_line_bytes`]; nothing more is read.
    TooLong,
}

/// Where a live child's process is in its life.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Life {
    Running,
    /// The leader exited and was collected; cleanup is under way.
    Exited,
    /// Its group is stopped and its delegation record removed.
    Finished,
}

/// A running live child. Dropping the last reference shuts it down in the
/// background; [`LiveChild::close`] does the same and waits for it.
pub(crate) struct LiveChild {
    pid: u32,
    stdin: Mutex<Option<ChildStdin>>,
    lines: Mutex<mpsc::Receiver<LiveLine>>,
    stderr: Arc<Mutex<TailBuffer>>,
    /// Shared with the reaper, which finishes it when the child exits.
    guard: Arc<StdMutex<Option<DelegationGuard>>>,
    life: watch::Receiver<Life>,
    closed: AtomicBool,
    activity: Option<RunActivity>,
}

impl std::fmt::Debug for LiveChild {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("LiveChild")
            .field("pid", &self.pid)
            .field("closed", &self.closed.load(Ordering::Relaxed))
            .finish()
    }
}

impl LiveChild {
    /// Start the child in its own process group with the agent environment,
    /// recording it under `registration` when given. Recording never fails
    /// the start: an unrecorded child is still tagged for later sweeps.
    pub(crate) async fn spawn(
        mut spec: LiveSpec,
        registration: Option<(Arc<DelegationRegistry>, PendingDelegation)>,
        cancellation: tokio_util::sync::CancellationToken,
    ) -> Result<Arc<Self>, ToolError> {
        let project = crate::project_environment::resolve(
            &spec.cwd,
            &spec.environment,
            cancellation.clone().cancelled_owned(),
        )
        .await
        .map_err(|error| {
            if cancellation.is_cancelled() {
                ToolError::cancelled("Rust preflight cancelled")
            } else {
                ToolError::unavailable(format!("{error:#}"))
            }
        })?;
        if cancellation.is_cancelled() {
            return Err(ToolError::cancelled("agent launch cancelled"));
        }
        spec.environment.extend(project.environment);
        let mut command = Command::new(&spec.executable);
        apply_agent_environment(command.as_std_mut(), &spec.environment);
        command
            .args(&spec.args)
            .current_dir(&spec.cwd)
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .kill_on_drop(true);
        command.as_std_mut().process_group(0);
        let mut child = command.spawn().map_err(|error| {
            ToolError::unavailable(format!("launch {:?}: {error}", spec.executable))
        })?;
        let pid = child
            .id()
            .ok_or_else(|| ToolError::failed("child process has no pid"))?;
        records::track_spawned(pid);
        let guard = Arc::new(StdMutex::new(
            registration.and_then(|(registry, pending)| registry.register(pending, pid).ok()),
        ));
        let stdin = child.stdin.take();
        let stdout = child.stdout.take();
        let stderr = Arc::new(Mutex::new(TailBuffer::new(STDERR_TAIL_BYTES)));
        if let Some(reader) = child.stderr.take() {
            tokio::spawn(drain_output(reader, Arc::clone(&stderr)));
        }
        let (sender, receiver) = mpsc::channel(LINE_QUEUE);
        if let Some(stdout) = stdout {
            tokio::spawn(read_lines(stdout, sender, spec.max_line_bytes));
        }
        let (life_tx, life) = watch::channel(Life::Running);
        tokio::spawn(reap(
            pid,
            child,
            Arc::clone(&guard),
            life_tx,
            spec.activity.clone(),
        ));
        Ok(Arc::new(Self {
            pid,
            stdin: Mutex::new(stdin),
            lines: Mutex::new(receiver),
            stderr,
            guard,
            life,
            closed: AtomicBool::new(false),
            activity: spec.activity,
        }))
    }

    /// Write `message` as one JSON line that the protocol owes the child,
    /// such as an answer to its request or a cancellation: within
    /// [`SEND_GRACE`] of now, whatever the turn's deadline, and never given
    /// up for the turn's cancellation, which is often what made it due.
    pub(crate) async fn send(&self, message: &impl Serialize) -> Result<(), ToolError> {
        self.send_until(
            message,
            Some(Instant::now() + SEND_GRACE),
            &CancellationToken::new(),
        )
        .await
    }

    /// Write `message` as one JSON line, giving up at `deadline` or on
    /// `cancellation`, for a request that starts work. The write is tried
    /// before either is checked, so one the child can take at once is not
    /// lost to a deadline or token that is also due. A write given up
    /// part-way would leave the child half a line, so its input is closed
    /// instead: the child is lost rather than sent a garbled message.
    pub(crate) async fn send_until(
        &self,
        message: &impl Serialize,
        deadline: Option<Instant>,
        cancellation: &CancellationToken,
    ) -> Result<(), ToolError> {
        let mut bytes = serde_json::to_vec(message)
            .map_err(|error| ToolError::failed(format!("encode message: {error}")))?;
        bytes.push(b'\n');
        let mut stdin = tokio::select! {
            biased;
            stdin = self.stdin.lock() => stdin,
            () = until(deadline) => return Err(ToolError::limit("write to child timed out")),
            () = cancellation.cancelled() => return Err(ToolError::cancelled("write to child cancelled")),
        };
        let writer = stdin
            .as_mut()
            .ok_or_else(|| ToolError::failed("the child's input is closed"))?;
        let mut written = 0;
        let write = async {
            while written < bytes.len() {
                match writer.write(&bytes[written..]).await {
                    Ok(0) => return Err(std::io::ErrorKind::WriteZero.into()),
                    Ok(count) => written += count,
                    Err(error) => return Err(error),
                }
            }
            writer.flush().await
        };
        let interrupted = tokio::select! {
            biased;
            result = write => {
                return result.map_err(|error| ToolError::failed(format!("write to child: {error}")));
            }
            () = until(deadline) => ToolError::limit("write to child timed out"),
            () = cancellation.cancelled() => ToolError::cancelled("write to child cancelled"),
        };
        if written > 0 {
            // Anything sent after half a line would be misread.
            *stdin = None;
        }
        Err(interrupted)
    }

    /// The next stdout line, or `None` once the child closed its output.
    pub(crate) async fn recv(&self) -> Option<LiveLine> {
        self.lines.lock().await.recv().await
    }

    /// The end of the child's stderr.
    pub(crate) async fn stderr_tail(&self) -> String {
        self.stderr.lock().await.text()
    }

    /// Whether the child process is still running.
    pub(crate) fn is_running(&self) -> bool {
        !self.closed.load(Ordering::Acquire) && *self.life.borrow() == Life::Running
    }

    /// The child's activity log, when it keeps one.
    pub(crate) fn activity(&self) -> Option<&RunActivity> {
        self.activity.as_ref()
    }

    /// Record that this child serves `turn` of its conversation until the
    /// returned guard drops, which records it idle between turns. A planned
    /// restart waits for a live child only while a turn runs.
    pub(crate) fn begin_turn(&self, turn: u32) -> LiveTurn<'_> {
        // Held while the record is rewritten, so the reaper, which takes the
        // guard before removing the record, never sees it written back.
        if let Some(guard) = lock(&self.guard).as_ref() {
            guard.set_turn(turn);
        }
        LiveTurn(self)
    }

    /// Record how many of the child's own background jobs still run or wait
    /// to be reported to it. Between turns, a planned restart waits for them.
    pub(crate) fn set_background_jobs(&self, jobs: usize) {
        if let Some(guard) = lock(&self.guard).as_ref() {
            guard.set_background_jobs(jobs);
        }
    }

    /// Shut the child down and wait for it: close its input, give it
    /// [`STOP_GRACE`] to exit, then kill its process group and anything
    /// still tagged with its delegation.
    pub(crate) async fn close(&self) {
        if self.closed.swap(true, Ordering::AcqRel) {
            return;
        }
        let stdin = self.stdin.lock().await.take();
        shut_down(self.pid, stdin, self.life.clone()).await;
    }
}

/// A turn a live child is serving, from [`LiveChild::begin_turn`]. However
/// the turn ends, dropping this records the child idle; a child already
/// stopped has no record left to change.
#[must_use = "the turn ends when this guard drops"]
pub(crate) struct LiveTurn<'a>(&'a LiveChild);

impl Drop for LiveTurn<'_> {
    fn drop(&mut self) {
        if let Some(guard) = lock(&self.0.guard).as_ref() {
            guard.set_idle();
        }
        if let Some(activity) = &self.0.activity {
            if self.0.is_running() {
                activity.idle();
            } else {
                activity.turn_released();
            }
        }
    }
}

impl Drop for LiveChild {
    fn drop(&mut self) {
        if self.closed.swap(true, Ordering::AcqRel) {
            return;
        }
        let stdin = self.stdin.get_mut().take();
        let pid = self.pid;
        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
            runtime.spawn(shut_down(pid, stdin, self.life.clone()));
        } else {
            // No runtime to wait in, so the reaper is gone too: kill at
            // once. The guard's drop sweeps tagged leftovers.
            kill_group(pid);
            records::untrack_spawned(pid);
            drop(lock(&self.guard).take());
        }
    }
}

/// Close the child's input, give it [`STOP_GRACE`] to exit, then kill its
/// group, and wait for the reaper to finish cleaning up.
async fn shut_down(pid: u32, stdin: Option<ChildStdin>, mut life: watch::Receiver<Life>) {
    // Closing input asks the child to finish; an SCV server exits on EOF.
    drop(stdin);
    if tokio::time::timeout(STOP_GRACE, life.wait_for(|state| *state != Life::Running))
        .await
        .is_err()
    {
        kill_group(pid);
    }
    // The reaper then stops leftovers, giving tagged ones their own grace.
    let _ = tokio::time::timeout(
        STOP_GRACE * 2 + Duration::from_millis(500),
        life.wait_for(|state| *state == Life::Finished),
    )
    .await;
}

/// Own the child for its whole life: collect it the moment it exits, then
/// stop the rest of its group and forget its delegation, so an exit between
/// turns leaves neither a zombie nor a stale `scv agents ps` entry.
async fn reap(
    pid: u32,
    mut child: Child,
    guard: Arc<StdMutex<Option<DelegationGuard>>>,
    life: watch::Sender<Life>,
    activity: Option<RunActivity>,
) {
    let _ = child.wait().await;
    let _ = life.send(Life::Exited);
    let killed = lock(&guard)
        .as_ref()
        .is_some_and(DelegationGuard::was_killed);
    if let Some(activity) = activity {
        activity.run_ended(if killed { "killed" } else { "exited" });
    }
    // Descendants may outlive the leader, so the group always goes.
    if group_exists(pid) {
        kill_group(pid);
    }
    records::untrack_spawned(pid);
    let guard = lock(&guard).take();
    if let Some(guard) = guard {
        guard.finish().await;
    }
    let _ = life.send(Life::Finished);
}

/// Split `stdout` into lines of at most `max_bytes`, counting the line
/// ending, for the protocol client.
async fn read_lines(stdout: ChildStdout, sender: mpsc::Sender<LiveLine>, max_bytes: usize) {
    let mut reader = BufReader::new(stdout);
    let mut decoder = FrameDecoder::new(max_bytes, Overflow::Stop);
    loop {
        let line = match scv_client::read_frame(&mut reader, &mut decoder).await {
            Ok(Frame::Line(line)) => LiveLine::Line(line),
            Ok(Frame::TooLarge) => {
                let _ = sender.send(LiveLine::TooLong).await;
                return;
            }
            // End of output; a partial last line is not a message.
            Ok(Frame::End | Frame::Truncated(_)) | Err(_) => return,
        };
        if sender.send(line).await.is_err() {
            return;
        }
    }
}

/// Kill every process in the child's group.
fn kill_group(pid: u32) {
    if let Some(group) = ProcessGroup::new(pid) {
        group.signal(libc::SIGKILL);
    }
}

#[cfg(test)]
mod tests;