Codex Handbook
utils/pty/src/process.rs 456 lines
use core::fmt;use std::io;#[cfg(unix)]use std::os::fd::RawFd;use std::process::ExitStatus;use std::sync::Arc;use std::sync::Mutex as StdMutex;use std::sync::atomic::AtomicBool;use anyhow::anyhow;use portable_pty::MasterPty;use portable_pty::PtySize;use portable_pty::SlavePty;use tokio::sync::broadcast;use tokio::sync::mpsc;use tokio::sync::oneshot;use tokio::sync::watch;use tokio::task::AbortHandle;use tokio::task::JoinHandle;#[derive(Clone, Copy, Debug, PartialEq, Eq)]pub enum ProcessSignal {    Interrupt,}pub(crate) fn unsupported_signal(signal: ProcessSignal) -> io::Error {    match signal {        ProcessSignal::Interrupt => io::Error::new(            io::ErrorKind::Unsupported,            "process interrupt is not supported by this process backend",        ),    }}pub(crate) fn exit_code_from_status(status: ExitStatus) -> i32 {    if let Some(code) = status.code() {        return code;    }    #[cfg(unix)]    {        use std::os::unix::process::ExitStatusExt;        if let Some(signal) = status.signal() {            return 128 + signal;        }    }    -1}pub(crate) trait ChildTerminator: Send + Sync {    fn signal(&mut self, signal: ProcessSignal) -> io::Result<()>;    fn kill(&mut self) -> io::Result<()>;}#[derive(Clone, Copy, Debug, PartialEq, Eq)]pub struct TerminalSize {    pub rows: u16,    pub cols: u16,}impl Default for TerminalSize {    fn default() -> Self {        Self { rows: 24, cols: 80 }    }}impl From<TerminalSize> for PtySize {    fn from(value: TerminalSize) -> Self {        Self {            rows: value.rows,            cols: value.cols,            pixel_width: 0,            pixel_height: 0,        }    }}#[cfg(unix)]pub(crate) trait PtyHandleKeepAlive: Send {}#[cfg(unix)]impl<T: Send + ?Sized> PtyHandleKeepAlive for T {}pub(crate) enum PtyMasterHandle {    Resizable(Box<dyn MasterPty + Send>),    #[cfg(unix)]    Opaque {        raw_fd: RawFd,        _handle: Box<dyn PtyHandleKeepAlive>,    },}pub struct PtyHandles {    pub _slave: Option<Box<dyn SlavePty + Send>>,    pub(crate) _master: PtyMasterHandle,}impl fmt::Debug for PtyHandles {    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {        f.debug_struct("PtyHandles").finish()    }}/// Callback used by driver-backed sessions to resize a PTY-like backend when/// there is no local `PtyHandles` instance to resize directly.type ResizeFn = Box<dyn FnMut(TerminalSize) -> anyhow::Result<()> + Send>;/// Handle for driving an interactive process (PTY or pipe).pub struct ProcessHandle {    writer_tx: StdMutex<Option<mpsc::Sender<Vec<u8>>>>,    killer: StdMutex<Option<Box<dyn ChildTerminator>>>,    reader_handle: StdMutex<Option<JoinHandle<()>>>,    reader_abort_handles: StdMutex<Vec<AbortHandle>>,    writer_handle: StdMutex<Option<JoinHandle<()>>>,    wait_handle: StdMutex<Option<JoinHandle<()>>>,    exit_status: Arc<AtomicBool>,    exit_code: Arc<StdMutex<Option<i32>>>,    // PtyHandles must be preserved because the process will receive Control+C if the    // slave is closed    _pty_handles: StdMutex<Option<PtyHandles>>,    // Optional resize hook for driver-backed sessions that proxy PTY control to    // another backend instead of owning local PTY handles.    resizer: StdMutex<Option<ResizeFn>>,}impl fmt::Debug for ProcessHandle {    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {        f.debug_struct("ProcessHandle").finish()    }}impl ProcessHandle {    #[allow(clippy::too_many_arguments)]    pub(crate) fn new(        writer_tx: mpsc::Sender<Vec<u8>>,        killer: Box<dyn ChildTerminator>,        reader_handle: JoinHandle<()>,        reader_abort_handles: Vec<AbortHandle>,        writer_handle: JoinHandle<()>,        wait_handle: JoinHandle<()>,        exit_status: Arc<AtomicBool>,        exit_code: Arc<StdMutex<Option<i32>>>,        pty_handles: Option<PtyHandles>,        resizer: Option<ResizeFn>,    ) -> Self {        Self {            writer_tx: StdMutex::new(Some(writer_tx)),            killer: StdMutex::new(Some(killer)),            reader_handle: StdMutex::new(Some(reader_handle)),            reader_abort_handles: StdMutex::new(reader_abort_handles),            writer_handle: StdMutex::new(Some(writer_handle)),            wait_handle: StdMutex::new(Some(wait_handle)),            exit_status,            exit_code,            _pty_handles: StdMutex::new(pty_handles),            resizer: StdMutex::new(resizer),        }    }    /// Returns a channel sender for writing raw bytes to the child stdin.    pub fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> {        if let Ok(writer_tx) = self.writer_tx.lock()            && let Some(writer_tx) = writer_tx.as_ref()        {            return writer_tx.clone();        }        let (writer_tx, writer_rx) = mpsc::channel(1);        drop(writer_rx);        writer_tx    }    /// True if the child process has exited.    pub fn has_exited(&self) -> bool {        self.exit_status.load(std::sync::atomic::Ordering::SeqCst)    }    /// Returns the exit code if known.    pub fn exit_code(&self) -> Option<i32> {        self.exit_code.lock().ok().and_then(|guard| *guard)    }    /// Resize the PTY in character cells.    pub fn resize(&self, size: TerminalSize) -> anyhow::Result<()> {        {            let handles = self                ._pty_handles                .lock()                .map_err(|_| anyhow!("failed to lock PTY handles"))?;            if let Some(handles) = handles.as_ref() {                return match &handles._master {                    PtyMasterHandle::Resizable(master) => master.resize(size.into()),                    #[cfg(unix)]                    PtyMasterHandle::Opaque { raw_fd, .. } => resize_raw_pty(*raw_fd, size),                };            }        }        let mut resizer = self            .resizer            .lock()            .map_err(|_| anyhow!("failed to lock PTY resizer"))?;        if let Some(resizer) = resizer.as_mut() {            resizer(size)        } else {            Err(anyhow!("process is not attached to a PTY"))        }    }    /// Close the child's stdin channel.    pub fn close_stdin(&self) {        if let Ok(mut writer_tx) = self.writer_tx.lock() {            writer_tx.take();        }    }    /// Attempts to kill the child while leaving the reader/writer tasks alive    /// so callers can still drain output until EOF.    pub fn request_terminate(&self) {        if let Ok(mut killer_opt) = self.killer.lock()            && let Some(mut killer) = killer_opt.take()        {            let _ = killer.kill();        }    }    pub fn signal(&self, signal: ProcessSignal) -> io::Result<()> {        let Ok(mut killer_opt) = self.killer.lock() else {            return Ok(());        };        let Some(killer) = killer_opt.as_mut() else {            return Ok(());        };        killer.signal(signal)    }    /// Attempts to kill the child and abort helper tasks.    pub fn terminate(&self) {        self.request_terminate();        if let Ok(mut h) = self.reader_handle.lock()            && let Some(handle) = h.take()        {            handle.abort();        }        if let Ok(mut handles) = self.reader_abort_handles.lock() {            for handle in handles.drain(..) {                handle.abort();            }        }        if let Ok(mut h) = self.writer_handle.lock()            && let Some(handle) = h.take()        {            handle.abort();        }        if let Ok(mut h) = self.wait_handle.lock()            && let Some(handle) = h.take()        {            handle.abort();        }    }}impl Drop for ProcessHandle {    fn drop(&mut self) {        self.terminate();    }}/// Adapts a closure into a `ChildTerminator` implementation.struct ClosureTerminator {    inner: Option<Box<dyn FnMut() + Send + Sync>>,}impl ChildTerminator for ClosureTerminator {    fn signal(&mut self, signal: ProcessSignal) -> io::Result<()> {        Err(unsupported_signal(signal))    }    fn kill(&mut self) -> io::Result<()> {        if let Some(inner) = self.inner.as_mut() {            (inner)();        }        Ok(())    }}#[cfg(unix)]fn resize_raw_pty(raw_fd: RawFd, size: TerminalSize) -> anyhow::Result<()> {    let mut winsize = libc::winsize {        ws_row: size.rows,        ws_col: size.cols,        ws_xpixel: 0,        ws_ypixel: 0,    };    let result = unsafe { libc::ioctl(raw_fd, libc::TIOCSWINSZ, &mut winsize) };    if result == -1 {        return Err(std::io::Error::last_os_error().into());    }    Ok(())}/// Combine split stdout/stderr receivers into a single broadcast receiver.pub fn combine_output_receivers(    mut stdout_rx: mpsc::Receiver<Vec<u8>>,    mut stderr_rx: mpsc::Receiver<Vec<u8>>,) -> broadcast::Receiver<Vec<u8>> {    let (combined_tx, combined_rx) = broadcast::channel(256);    tokio::spawn(async move {        let mut stdout_open = true;        let mut stderr_open = true;        loop {            tokio::select! {                stdout = stdout_rx.recv(), if stdout_open => match stdout {                    Some(chunk) => {                        let _ = combined_tx.send(chunk);                    }                    None => {                        stdout_open = false;                    }                },                stderr = stderr_rx.recv(), if stderr_open => match stderr {                    Some(chunk) => {                        let _ = combined_tx.send(chunk);                    }                    None => {                        stderr_open = false;                    }                },                else => break,            }        }    });    combined_rx}/// Return value from PTY or pipe spawn helpers.#[derive(Debug)]pub struct SpawnedProcess {    pub session: ProcessHandle,    pub stdout_rx: mpsc::Receiver<Vec<u8>>,    pub stderr_rx: mpsc::Receiver<Vec<u8>>,    pub exit_rx: oneshot::Receiver<i32>,}/// Driver-backed process handles for non-standard spawn backends.pub struct ProcessDriver {    pub writer_tx: mpsc::Sender<Vec<u8>>,    pub stdout_rx: broadcast::Receiver<Vec<u8>>,    pub stderr_rx: Option<broadcast::Receiver<Vec<u8>>>,    pub exit_rx: oneshot::Receiver<i32>,    pub terminator: Option<Box<dyn FnMut() + Send + Sync>>,    pub writer_handle: Option<JoinHandle<()>>,    pub resizer: Option<ResizeFn>,}/// Build a `SpawnedProcess` from a driver that supplies stdin/output/exit channels.pub fn spawn_from_driver(driver: ProcessDriver) -> SpawnedProcess {    let ProcessDriver {        writer_tx,        stdout_rx: stdout_driver_rx,        stderr_rx: mut stderr_driver_rx,        exit_rx,        terminator,        writer_handle,        resizer,    } = driver;    let (stdout_tx, stdout_rx) = mpsc::channel::<Vec<u8>>(256);    let (stderr_tx, stderr_rx) = mpsc::channel::<Vec<u8>>(256);    let (exit_seen_tx, exit_seen_rx) = watch::channel(false);    let spawn_stream_reader =        |mut output_rx: broadcast::Receiver<Vec<u8>>,         output_tx: mpsc::Sender<Vec<u8>>,         mut exit_seen_rx: watch::Receiver<bool>| {            tokio::spawn(async move {                loop {                    let recv_result = if *exit_seen_rx.borrow() {                        // Once exit has been observed, we no longer want a timer here. Some                        // backends publish the exit code before their final stdout/stderr bytes                        // have been forwarded through the broadcast channel, so a fixed grace                        // period can still drop the tail of the stream under load.                        //                        // Instead, keep waiting until the driver closes the broadcast sender.                        // That makes the shutdown contract explicit: the backend is responsible                        // for dropping its sender when it has truly finished forwarding output.                        output_rx.recv().await                    } else {                        tokio::select! {                            _ = exit_seen_rx.changed() => {                                continue;                            }                            result = output_rx.recv() => result,                        }                    };                    match recv_result {                        Ok(chunk) => {                            if output_tx.send(chunk).await.is_err() {                                break;                            }                        }                        Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,                        Err(tokio::sync::broadcast::error::RecvError::Closed) => break,                    }                }            })        };    let reader_handle = spawn_stream_reader(stdout_driver_rx, stdout_tx, exit_seen_rx.clone());    let stderr_reader_handle = stderr_driver_rx        .take()        .map(|rx| spawn_stream_reader(rx, stderr_tx, exit_seen_rx));    let writer_handle = writer_handle.unwrap_or_else(|| tokio::spawn(async {}));    let (exit_tx, exit_rx_out) = oneshot::channel::<i32>();    let exit_status = Arc::new(AtomicBool::new(false));    let wait_exit_status = Arc::clone(&exit_status);    let exit_code = Arc::new(StdMutex::new(None));    let wait_exit_code = Arc::clone(&exit_code);    let wait_handle = tokio::spawn(async move {        let code = exit_rx.await.unwrap_or(-1);        wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst);        if let Ok(mut guard) = wait_exit_code.lock() {            *guard = Some(code);        }        let _ = exit_seen_tx.send(true);        let _ = exit_tx.send(code);    });    let handle = ProcessHandle::new(        writer_tx,        Box::new(ClosureTerminator { inner: terminator }),        reader_handle,        stderr_reader_handle            .map(|handle| handle.abort_handle())            .into_iter()            .collect(),        writer_handle,        wait_handle,        exit_status,        exit_code,        /*pty_handles*/ None,        resizer,    );    SpawnedProcess {        session: handle,        stdout_rx,        stderr_rx,        exit_rx: exit_rx_out,    }}