make it more robust

Author: jif-oai

codex-rs/utils/pty/src/pipe.rs

@@ -26,18 +26,29 @@ use crate::process::SpawnedProcess;
 use libc;
 
 struct PipeChildTerminator {
+    #[cfg(windows)]
     pid: u32,
+    #[cfg(unix)]
+    process_group_id: u32,
 }
 
 impl ChildTerminator for PipeChildTerminator {
     fn kill(&mut self) -> io::Result<()> {
-        kill_process(self.pid)
-    }
-}
+        #[cfg(unix)]
+        {
+            return crate::process_group::kill_process_group(self.process_group_id);
+        }
 
-#[cfg(unix)]
-fn kill_process(pid: u32) -> io::Result<()> {
-    crate::process_group::kill_process_group_by_pid(pid)
+        #[cfg(windows)]
+        {
+            return kill_process(self.pid);
+        }
+
+        #[cfg(not(any(unix, windows)))]
+        {
+            Ok(())
+        }
+    }
 }
 
 #[cfg(windows)]
@@ -125,6 +136,8 @@ pub async fn spawn_process(
     let pid = child
         .id()
         .ok_or_else(|| io::Error::other("missing child pid"))?;
+    #[cfg(unix)]
+    let process_group_id = pid;
 
     let stdin = child.stdin.take();
     let stdout = child.stdout.take();
@@ -159,6 +172,13 @@ pub async fn spawn_process(
             read_output_stream(BufReader::new(stderr), output_tx).await;
         })
     });
+    let mut reader_abort_handles = Vec::new();
+    if let Some(handle) = stdout_handle.as_ref() {
+        reader_abort_handles.push(handle.abort_handle());
+    }
+    if let Some(handle) = stderr_handle.as_ref() {
+        reader_abort_handles.push(handle.abort_handle());
+    }
     let reader_handle = tokio::spawn(async move {
         if let Some(handle) = stdout_handle {
             let _ = handle.await;
@@ -189,8 +209,14 @@ pub async fn spawn_process(
         writer_tx,
         output_tx,
         initial_output_rx,
-        Box::new(PipeChildTerminator { pid }),
+        Box::new(PipeChildTerminator {
+            #[cfg(windows)]
+            pid,
+            #[cfg(unix)]
+            process_group_id,
+        }),
         reader_handle,
+        reader_abort_handles,
         writer_handle,
         wait_handle,
         exit_status,

codex-rs/utils/pty/src/process.rs

@@ -9,6 +9,7 @@ use portable_pty::SlavePty;
 use tokio::sync::broadcast;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
+use tokio::task::AbortHandle;
 use tokio::task::JoinHandle;
 
 pub(crate) trait ChildTerminator: Send + Sync {
@@ -32,6 +33,7 @@ pub struct ProcessHandle {
     output_tx: broadcast::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>,
@@ -55,6 +57,7 @@ impl ProcessHandle {
         initial_output_rx: broadcast::Receiver<Vec<u8>>,
         killer: Box<dyn ChildTerminator>,
         reader_handle: JoinHandle<()>,
+        reader_abort_handles: Vec<AbortHandle>,
         writer_handle: JoinHandle<()>,
         wait_handle: JoinHandle<()>,
         exit_status: Arc<AtomicBool>,
@@ -67,6 +70,7 @@ impl ProcessHandle {
                 output_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,
@@ -110,6 +114,11 @@ impl ProcessHandle {
                 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() {
             if let Some(handle) = h.take() {
                 handle.abort();

codex-rs/utils/pty/src/process_group.rs

@@ -5,6 +5,7 @@
 //! - `set_process_group` is called in `pre_exec` so the child starts its own
 //!   process group.
 //! - `kill_process_group_by_pid` targets the whole group (children/grandchildren)
+//! - `kill_process_group` targets a known process group ID directly
 //!   instead of a single PID.
 //! - `set_parent_death_signal` (Linux only) arranges for the child to receive a
 //!   `SIGTERM` when the parent exits, and re-checks the parent PID to avoid
@@ -94,6 +95,29 @@ pub fn kill_process_group_by_pid(_pid: u32) -> io::Result<()> {
     Ok(())
 }
 
+#[cfg(unix)]
+/// Kill a specific process group ID (best-effort).
+pub fn kill_process_group(process_group_id: u32) -> io::Result<()> {
+    use std::io::ErrorKind;
+
+    let pgid = process_group_id as libc::pid_t;
+    let result = unsafe { libc::killpg(pgid, libc::SIGKILL) };
+    if result == -1 {
+        let err = io::Error::last_os_error();
+        if err.kind() != ErrorKind::NotFound {
+            return Err(err);
+        }
+    }
+
+    Ok(())
+}
+
+#[cfg(not(unix))]
+/// No-op on non-Unix platforms.
+pub fn kill_process_group(_process_group_id: u32) -> io::Result<()> {
+    Ok(())
+}
+
 #[cfg(unix)]
 /// Kill the process group for a tokio child (best-effort).
 pub fn kill_child_process_group(child: &mut Child) -> io::Result<()> {

codex-rs/utils/pty/src/pty.rs

@@ -158,6 +158,7 @@ pub async fn spawn_process(
         initial_output_rx,
         Box::new(PtyChildTerminator { killer }),
         reader_handle,
+        Vec::new(),
         writer_handle,
         wait_handle,
         exit_status,

codex-rs/utils/pty/src/tests.rs

@@ -1,10 +1,15 @@
 use std::collections::HashMap;
 use std::path::Path;
 
+#[cfg(unix)]
+use std::io::ErrorKind;
+
 use pretty_assertions::assert_eq;
 
 use crate::spawn_pipe_process;
 use crate::spawn_pty_process;
+#[cfg(unix)]
+use crate::SpawnedProcess;
 
 fn find_python() -> Option<String> {
     for candidate in ["python3", "python"] {
@@ -20,6 +25,27 @@ fn find_python() -> Option<String> {
     None
 }
 
+fn setsid_available() -> bool {
+    if cfg!(windows) {
+        return false;
+    }
+    std::process::Command::new("setsid")
+        .arg("true")
+        .status()
+        .map(|status| status.success())
+        .unwrap_or(false)
+}
+
+#[cfg(unix)]
+fn process_exists(pid: i32) -> bool {
+    let result = unsafe { libc::kill(pid, 0) };
+    if result == 0 {
+        return true;
+    }
+    let err = std::io::Error::last_os_error();
+    err.kind() != ErrorKind::NotFound
+}
+
 fn shell_command(program: &str) -> (String, Vec<String>) {
     if cfg!(windows) {
         let cmd = std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string());
@@ -190,3 +216,83 @@ async fn pipe_drains_stderr_without_stdout_activity() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn pipe_terminate_aborts_detached_readers() -> anyhow::Result<()> {
+    if !setsid_available() {
+        eprintln!("setsid not available; skipping pipe_terminate_aborts_detached_readers");
+        return Ok(());
+    }
+
+    let env_map: HashMap<String, String> = std::env::vars().collect();
+    let script =
+        "setsid sh -c 'i=0; while [ $i -lt 200 ]; do echo tick; sleep 0.01; i=$((i+1)); done' &";
+    let (program, args) = shell_command(script);
+    let mut spawned = spawn_pipe_process(&program, &args, Path::new("."), &env_map, &None).await?;
+
+    let _ = tokio::time::timeout(
+        tokio::time::Duration::from_millis(500),
+        spawned.output_rx.recv(),
+    )
+    .await
+    .map_err(|_| anyhow::anyhow!("expected detached output before terminate"))??;
+
+    spawned.session.terminate();
+    let mut post_rx = spawned.session.output_receiver();
+
+    let post_terminate =
+        tokio::time::timeout(tokio::time::Duration::from_millis(200), post_rx.recv()).await;
+
+    match post_terminate {
+        Err(_) => Ok(()),
+        Ok(Err(tokio::sync::broadcast::error::RecvError::Closed)) => Ok(()),
+        Ok(Err(tokio::sync::broadcast::error::RecvError::Lagged(_))) => {
+            anyhow::bail!("unexpected output after terminate (lagged)")
+        }
+        Ok(Ok(chunk)) => anyhow::bail!(
+            "unexpected output after terminate: {:?}",
+            String::from_utf8_lossy(&chunk)
+        ),
+    }
+}
+
+#[cfg(unix)]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn pipe_terminate_kills_process_group_after_exit() -> anyhow::Result<()> {
+    let env_map: HashMap<String, String> = std::env::vars().collect();
+    let (program, args) = shell_command("sleep 9999 & echo $!");
+    let SpawnedProcess {
+        session,
+        output_rx,
+        exit_rx,
+    } = spawn_pipe_process(&program, &args, Path::new("."), &env_map, &None).await?;
+
+    let (output, code) = collect_output_until_exit(output_rx, exit_rx, 2_000).await;
+    assert_eq!(code, 0, "expected shell to exit cleanly");
+
+    let pid_line = String::from_utf8_lossy(&output)
+        .lines()
+        .find(|line| !line.trim().is_empty())
+        .unwrap_or("")
+        .trim()
+        .to_string();
+    let pid: i32 = pid_line
+        .parse()
+        .map_err(|_| anyhow::anyhow!("failed to parse background pid from {pid_line:?}"))?;
+
+    session.terminate();
+
+    let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_millis(500);
+    while tokio::time::Instant::now() < deadline && process_exists(pid) {
+        tokio::time::sleep(tokio::time::Duration::from_millis(20)).await;
+    }
+
+    if process_exists(pid) {
+        unsafe {
+            libc::kill(pid, libc::SIGKILL);
+        }
+        anyhow::bail!("background pid still alive after terminate: {pid}");
+    }
+
+    Ok(())
+}