Terminate Windows process trees with job objects

codex-rs/utils/pty/Cargo.toml

@@ -22,6 +22,7 @@ log = { workspace = true }
 shared_library = "0.1.9"
 winapi = { version = "0.3.9", features = [
     "handleapi",
+    "jobapi2",
     "minwinbase",
     "processthreadsapi",
     "synchapi",

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

@@ -38,6 +38,8 @@ pub use pty::conpty_supported;
 /// Spawn a process attached to a PTY for interactive use.
 pub use pty::spawn_process as spawn_pty_process;
 #[cfg(windows)]
+pub use win::JobObject;
+#[cfg(windows)]
 pub use win::PsuedoCon;
 #[cfg(windows)]
 pub use win::conpty::RawConPty;

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

@@ -26,9 +26,15 @@ use crate::process::exit_code_from_status;
 #[cfg(target_os = "linux")]
 use libc;
 
+#[cfg(windows)]
+enum WindowsChildTerminator {
+    Job(Arc<crate::win::JobObject>),
+    Process(u32),
+}
+
 struct PipeChildTerminator {
     #[cfg(windows)]
-    pid: u32,
+    windows: WindowsChildTerminator,
     #[cfg(unix)]
     process_group_id: u32,
 }
@@ -58,7 +64,10 @@ impl ChildTerminator for PipeChildTerminator {
 
         #[cfg(windows)]
         {
-            kill_process(self.pid)
+            match &self.windows {
+                WindowsChildTerminator::Job(job) => job.terminate(),
+                WindowsChildTerminator::Process(pid) => kill_process(*pid),
+            }
         }
 
         #[cfg(not(any(unix, windows)))]
@@ -109,6 +118,8 @@ enum PipeStdinMode {
     Null,
 }
 
+/// On Windows, process-tree containment is best-effort because Tokio returns
+/// only after the root process starts, so job assignment cannot be atomic.
 async fn spawn_process_with_stdin_mode(
     program: &str,
     args: &[String],
@@ -165,12 +176,37 @@ async fn spawn_process_with_stdin_mode(
     command.stdout(Stdio::piped());
     command.stderr(Stdio::piped());
 
+    #[cfg(windows)]
+    let job = crate::win::JobObject::create().map(Arc::new);
     let mut child = command.spawn()?;
-    let pid = child
+    #[cfg(windows)]
+    let windows_terminator = {
+        // Accept the small race: a descendant created between spawn and
+        // assignment is not guaranteed to join the job and can escape termination.
+        let pid = child
+            .id()
+            .ok_or_else(|| io::Error::other("missing child pid"))?;
+        let assigned_job = job.and_then(|job| {
+            let process_handle = child
+                .raw_handle()
+                .ok_or_else(|| io::Error::other("missing child process handle"))?;
+            job.assign_process(process_handle)?;
+            Ok(job)
+        });
+        match assigned_job {
+            Ok(job) => WindowsChildTerminator::Job(job),
+            Err(err) => {
+                log::warn!(
+                    "Windows pipe process tree containment unavailable for pid {pid}: {err}"
+                );
+                WindowsChildTerminator::Process(pid)
+            }
+        }
+    };
+    #[cfg(unix)]
+    let process_group_id = 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();
@@ -225,9 +261,24 @@ async fn spawn_process_with_stdin_mode(
     let wait_exit_status = Arc::clone(&exit_status);
     let exit_code = Arc::new(StdMutex::new(None));
     let wait_exit_code = Arc::clone(&exit_code);
+    #[cfg(windows)]
+    let wait_job = match &windows_terminator {
+        WindowsChildTerminator::Job(job) => Some(Arc::clone(job)),
+        WindowsChildTerminator::Process(_) => None,
+    };
     let wait_handle: JoinHandle<()> = tokio::spawn(async move {
         let code = match child.wait().await {
-            Ok(status) => exit_code_from_status(status),
+            Ok(status) => {
+                #[cfg(windows)]
+                if let Some(job) = wait_job
+                    && let Err(err) = job.preserve_descendants()
+                {
+                    log::warn!(
+                        "Windows pipe failed to preserve descendants after root exit: {err}"
+                    );
+                }
+                exit_code_from_status(status)
+            }
             Err(_) => -1,
         };
         wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst);
@@ -241,7 +292,7 @@ async fn spawn_process_with_stdin_mode(
         writer_tx,
         Box::new(PipeChildTerminator {
             #[cfg(windows)]
-            pid,
+            windows: windows_terminator,
             #[cfg(unix)]
             process_group_id,
         }),
@@ -306,3 +357,7 @@ pub async fn spawn_process_no_stdin(
     )
     .await
 }
+
+#[cfg(all(test, windows))]
+#[path = "pipe_tests.rs"]
+mod tests;

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

@@ -0,0 +1,19 @@
+use super::*;
+
+#[test]
+fn process_fallback_terminates_root() -> anyhow::Result<()> {
+    let mut child = std::process::Command::new("ping.exe")
+        .args(["-n", "60", "127.0.0.1"])
+        .stdin(Stdio::null())
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .spawn()?;
+    let mut terminator = PipeChildTerminator {
+        windows: WindowsChildTerminator::Process(child.id()),
+    };
+
+    terminator.kill()?;
+
+    assert!(!child.wait()?.success());
+    Ok(())
+}

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

@@ -0,0 +1,122 @@
+use filedescriptor::OwnedHandle;
+use std::io;
+use std::os::windows::io::AsRawHandle;
+use std::os::windows::io::FromRawHandle;
+use std::os::windows::io::RawHandle;
+use std::sync::Mutex;
+use winapi::um::jobapi2::AssignProcessToJobObject;
+use winapi::um::jobapi2::CreateJobObjectW;
+use winapi::um::jobapi2::SetInformationJobObject;
+use winapi::um::jobapi2::TerminateJobObject;
+use winapi::um::winnt::JOB_OBJECT_LIMIT_BREAKAWAY_OK;
+use winapi::um::winnt::JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+use winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION;
+use winapi::um::winnt::JobObjectExtendedLimitInformation;
+
+/// Owns a Windows Job Object used to terminate a spawned process tree.
+#[derive(Debug)]
+pub struct JobObject {
+    handle: OwnedHandle,
+    // A mutex makes the state check, Job Object API call, and state update
+    // atomic with respect to concurrent preserve and terminate requests.
+    preserve_descendants: Mutex<bool>,
+}
+
+impl JobObject {
+    /// Creates a Job Object configured to terminate all members when its last handle closes.
+    pub fn create() -> io::Result<Self> {
+        let handle = unsafe { CreateJobObjectW(std::ptr::null_mut(), std::ptr::null()) };
+        if handle.is_null() {
+            return Err(io::Error::last_os_error());
+        }
+        let handle = unsafe { OwnedHandle::from_raw_handle(handle.cast()) };
+
+        Self::set_limit_flags(
+            &handle,
+            JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE | JOB_OBJECT_LIMIT_BREAKAWAY_OK,
+        )?;
+
+        Ok(Self {
+            handle,
+            preserve_descendants: Mutex::new(false),
+        })
+    }
+
+    fn set_limit_flags(handle: &OwnedHandle, flags: u32) -> io::Result<()> {
+        let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { std::mem::zeroed() };
+        limits.BasicLimitInformation.LimitFlags = flags;
+        let configured = unsafe {
+            SetInformationJobObject(
+                handle.as_raw_handle().cast(),
+                JobObjectExtendedLimitInformation,
+                std::ptr::addr_of_mut!(limits).cast(),
+                std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
+            )
+        };
+        if configured == 0 {
+            return Err(io::Error::last_os_error());
+        }
+        Ok(())
+    }
+
+    /// Assigns a running process to this job.
+    ///
+    /// Assignment is not retroactive: descendants created before this call
+    /// completes are not guaranteed to become members of the job.
+    pub(crate) fn assign_process(&self, process_handle: RawHandle) -> io::Result<()> {
+        let assigned = unsafe {
+            AssignProcessToJobObject(self.handle.as_raw_handle().cast(), process_handle.cast())
+        };
+        if assigned == 0 {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Allows contained descendants to keep running after the root exits normally.
+    ///
+    /// This disables both explicit job termination and kill-on-close for this
+    /// object. Calls race safely with [`Self::terminate`]: whichever operation
+    /// acquires the state lock first determines whether the process tree is
+    /// preserved or terminated.
+    pub fn preserve_descendants(&self) -> io::Result<()> {
+        let mut preserve_descendants = self
+            .preserve_descendants
+            .lock()
+            .map_err(|_| io::Error::other("job state lock poisoned"))?;
+        if *preserve_descendants {
+            return Ok(());
+        }
+
+        Self::set_limit_flags(&self.handle, JOB_OBJECT_LIMIT_BREAKAWAY_OK)?;
+        *preserve_descendants = true;
+        Ok(())
+    }
+
+    /// Terminates every process currently assigned to the job.
+    pub fn terminate(&self) -> io::Result<()> {
+        let preserve_descendants = self
+            .preserve_descendants
+            .lock()
+            .map_err(|_| io::Error::other("job state lock poisoned"))?;
+        if *preserve_descendants {
+            return Ok(());
+        }
+
+        let terminated = unsafe {
+            TerminateJobObject(self.handle.as_raw_handle().cast(), /*uExitCode*/ 1)
+        };
+        if terminated == 0 {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(())
+        }
+    }
+}
+
+impl AsRawHandle for JobObject {
+    fn as_raw_handle(&self) -> RawHandle {
+        self.handle.as_raw_handle()
+    }
+}