windows: add kill-on-close job process primitives

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,8 +38,14 @@ 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::JobProcess;
+#[cfg(windows)]
+pub use win::KillOnCloseJob;
+#[cfg(windows)]
 pub use win::PsuedoCon;
 #[cfg(windows)]
+pub use win::SuspendedProcess;
+#[cfg(windows)]
 pub use win::conpty::RawConPty;
 #[cfg(windows)]
 pub use windows_input::WindowsTtyInputNormalizer;

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

@@ -0,0 +1,285 @@
+use std::io;
+use std::mem;
+use std::os::windows::io::AsRawHandle;
+use std::os::windows::io::FromRawHandle;
+use std::os::windows::io::IntoRawHandle;
+use std::os::windows::io::OwnedHandle;
+use std::os::windows::io::RawHandle;
+use std::ptr;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::PoisonError;
+
+use anyhow::Context;
+use winapi::shared::minwindef::FALSE;
+use winapi::um::handleapi::CloseHandle;
+use winapi::um::jobapi2::AssignProcessToJobObject;
+use winapi::um::jobapi2::CreateJobObjectW;
+use winapi::um::jobapi2::SetInformationJobObject;
+use winapi::um::jobapi2::TerminateJobObject;
+use winapi::um::processthreadsapi::ResumeThread;
+use winapi::um::processthreadsapi::TerminateProcess;
+use winapi::um::synchapi::WaitForSingleObject;
+use winapi::um::winbase::INFINITE;
+use winapi::um::winnt::HANDLE;
+use winapi::um::winnt::JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+use winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION;
+use winapi::um::winnt::JobObjectExtendedLimitInformation;
+
+/// Injectable Win32 operations for the create/configure/assign/resume state
+/// machine. Production uses [`Win32JobObjectApi`]; tests replace individual
+/// stages so every failure path can be exercised deterministically.
+trait JobObjectApi {
+    fn create_job(&self) -> io::Result<OwnedHandle>;
+    fn configure_kill_on_close(&self, job: RawHandle) -> io::Result<()>;
+    fn assign_process(&self, job: RawHandle, process: RawHandle) -> io::Result<()>;
+    fn resume_thread(&self, primary_thread: RawHandle) -> io::Result<u32>;
+}
+
+struct Win32JobObjectApi;
+
+impl JobObjectApi for Win32JobObjectApi {
+    fn create_job(&self) -> io::Result<OwnedHandle> {
+        let raw_handle = unsafe { CreateJobObjectW(ptr::null_mut(), ptr::null()) };
+        if raw_handle.is_null() {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(unsafe { OwnedHandle::from_raw_handle(raw_handle.cast()) })
+        }
+    }
+
+    fn configure_kill_on_close(&self, job: RawHandle) -> io::Result<()> {
+        let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { mem::zeroed() };
+        limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+        let result = unsafe {
+            SetInformationJobObject(
+                job as HANDLE,
+                JobObjectExtendedLimitInformation,
+                ptr::addr_of_mut!(limits).cast(),
+                mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
+            )
+        };
+        if result == FALSE {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(())
+        }
+    }
+
+    fn assign_process(&self, job: RawHandle, process: RawHandle) -> io::Result<()> {
+        let result = unsafe { AssignProcessToJobObject(job as HANDLE, process as HANDLE) };
+        if result == FALSE {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(())
+        }
+    }
+
+    fn resume_thread(&self, primary_thread: RawHandle) -> io::Result<u32> {
+        let previous_suspend_count = unsafe { ResumeThread(primary_thread as HANDLE) };
+        if previous_suspend_count == u32::MAX {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(previous_suspend_count)
+        }
+    }
+}
+
+fn io_error_with_context(error: io::Error, context: &'static str) -> io::Error {
+    let kind = error.kind();
+    io::Error::new(kind, anyhow::Error::new(error).context(context))
+}
+
+/// Shared controller for a Windows job that terminates all members when closed.
+///
+/// Clones share one underlying job handle. Calling [`Self::close`] from any
+/// clone closes that handle exactly once, so process wait and cancellation paths
+/// can race safely without keeping the job alive through duplicated OS handles.
+#[derive(Clone, Debug)]
+pub struct KillOnCloseJob {
+    handle: Arc<Mutex<Option<OwnedHandle>>>,
+}
+
+impl KillOnCloseJob {
+    /// Create an unnamed, non-inheritable job configured to kill all members
+    /// when its sole operating-system handle is closed.
+    pub fn new() -> io::Result<Self> {
+        Self::new_with_api(&Win32JobObjectApi)
+    }
+
+    fn new_with_api(api: &impl JobObjectApi) -> io::Result<Self> {
+        let handle = api
+            .create_job()
+            .map_err(|err| io_error_with_context(err, "failed to create job object"))?;
+        api.configure_kill_on_close(handle.as_raw_handle())
+            .map_err(|err| {
+                io_error_with_context(err, "failed to configure kill-on-close job object")
+            })?;
+        Ok(Self {
+            handle: Arc::new(Mutex::new(Some(handle))),
+        })
+    }
+
+    fn assign_process_with_api(
+        &self,
+        process: RawHandle,
+        api: &impl JobObjectApi,
+    ) -> io::Result<()> {
+        let guard = self.handle.lock().unwrap_or_else(PoisonError::into_inner);
+        let Some(job) = guard.as_ref() else {
+            return Err(io::Error::new(
+                io::ErrorKind::BrokenPipe,
+                "job handle is already closed",
+            ));
+        };
+        api.assign_process(job.as_raw_handle(), process)
+    }
+
+    /// Close the shared job handle, terminating all processes in the job.
+    pub fn close(&self) -> io::Result<()> {
+        self.take_and_close(/*exit_code*/ None)
+    }
+
+    /// Terminate all job members with `exit_code`, then close the shared job
+    /// handle even if explicit termination reports an error.
+    pub fn terminate_and_close(&self, exit_code: u32) -> io::Result<()> {
+        self.take_and_close(Some(exit_code))
+    }
+
+    fn take_and_close(&self, exit_code: Option<u32>) -> io::Result<()> {
+        let mut guard = self.handle.lock().unwrap_or_else(PoisonError::into_inner);
+        let Some(handle) = guard.take() else {
+            return Ok(());
+        };
+        let raw_handle = handle.into_raw_handle();
+
+        let terminate_result = exit_code.map(|exit_code| {
+            let result = unsafe { TerminateJobObject(raw_handle.cast(), exit_code) };
+            if result == FALSE {
+                Err(io::Error::last_os_error())
+            } else {
+                Ok(())
+            }
+        });
+        let close_result = unsafe { CloseHandle(raw_handle.cast()) };
+        let close_result = if close_result == FALSE {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(())
+        };
+
+        terminate_result.unwrap_or(Ok(())).and(close_result)
+    }
+}
+
+/// Owns a newly created process and its suspended primary thread until the
+/// process has been assigned to a kill-on-close job and resumed.
+pub struct SuspendedProcess {
+    process: Option<OwnedHandle>,
+    primary_thread: Option<OwnedHandle>,
+    process_id: u32,
+}
+
+impl SuspendedProcess {
+    /// Take ownership of raw handles returned by a successful `CreateProcess*`
+    /// call made with `CREATE_SUSPENDED`.
+    ///
+    /// # Safety
+    ///
+    /// `process` and `primary_thread` must be valid, owned handles for the same
+    /// newly created process. The primary thread must still have its initial
+    /// suspension count, and ownership of both handles transfers to this value.
+    pub unsafe fn from_raw_handles(
+        process: RawHandle,
+        primary_thread: RawHandle,
+        process_id: u32,
+    ) -> Self {
+        Self {
+            process: Some(unsafe { OwnedHandle::from_raw_handle(process) }),
+            primary_thread: Some(unsafe { OwnedHandle::from_raw_handle(primary_thread) }),
+            process_id,
+        }
+    }
+
+    /// Assign the suspended process to `job`, then resume its primary thread.
+    /// Any failure leaves this guard armed, so the process is terminated and
+    /// waited before the error is returned.
+    pub fn assign_and_resume(self, job: KillOnCloseJob) -> anyhow::Result<JobProcess> {
+        self.assign_and_resume_with_api(job, &Win32JobObjectApi)
+    }
+
+    fn assign_and_resume_with_api(
+        mut self,
+        job: KillOnCloseJob,
+        api: &impl JobObjectApi,
+    ) -> anyhow::Result<JobProcess> {
+        let process = self
+            .process
+            .as_ref()
+            .ok_or_else(|| io::Error::other("suspended process is missing its process handle"))?;
+        job.assign_process_with_api(process.as_raw_handle(), api)
+            .context("failed to assign suspended process to job")?;
+
+        let primary_thread = self.primary_thread.as_ref().ok_or_else(|| {
+            io::Error::other("suspended process is missing its primary thread handle")
+        })?;
+        let previous_suspend_count = api
+            .resume_thread(primary_thread.as_raw_handle())
+            .context("failed to resume suspended process")?;
+        if previous_suspend_count != 1 {
+            return Err(io::Error::other(format!(
+                "expected suspended process thread count 1, got {previous_suspend_count}"
+            )))
+            .context("failed to resume suspended process");
+        }
+
+        drop(self.primary_thread.take());
+        let process = self
+            .process
+            .take()
+            .ok_or_else(|| io::Error::other("suspended process is missing its process handle"))?;
+        Ok(JobProcess {
+            process,
+            process_id: self.process_id,
+            controller: job,
+        })
+    }
+}
+
+impl Drop for SuspendedProcess {
+    fn drop(&mut self) {
+        let Some(process) = self.process.as_ref() else {
+            return;
+        };
+        unsafe {
+            let _ = TerminateProcess(process.as_raw_handle().cast(), /*uExitCode*/ 1);
+            let _ = WaitForSingleObject(process.as_raw_handle().cast(), INFINITE);
+        }
+    }
+}
+
+/// A running process contained by a kill-on-close Windows job.
+#[derive(Debug)]
+pub struct JobProcess {
+    process: OwnedHandle,
+    process_id: u32,
+    controller: KillOnCloseJob,
+}
+
+impl JobProcess {
+    pub fn process_id(&self) -> u32 {
+        self.process_id
+    }
+
+    pub fn as_raw_handle(&self) -> RawHandle {
+        self.process.as_raw_handle()
+    }
+
+    pub fn controller(&self) -> KillOnCloseJob {
+        self.controller.clone()
+    }
+}
+
+#[cfg(test)]
+#[path = "job_tests.rs"]
+mod tests;

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

@@ -0,0 +1,351 @@
+use std::cell::Cell;
+use std::ffi::OsStr;
+use std::fs;
+use std::io;
+use std::mem;
+use std::os::windows::ffi::OsStrExt;
+use std::os::windows::io::AsRawHandle;
+use std::os::windows::io::OwnedHandle;
+use std::os::windows::io::RawHandle;
+use std::path::Path;
+use std::path::PathBuf;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
+use std::thread;
+use std::time::Duration;
+use std::time::Instant;
+
+use pretty_assertions::assert_eq;
+use winapi::shared::minwindef::FALSE;
+use winapi::um::minwinbase::STILL_ACTIVE;
+use winapi::um::processthreadsapi::CreateProcessW;
+use winapi::um::processthreadsapi::GetExitCodeProcess;
+use winapi::um::processthreadsapi::PROCESS_INFORMATION;
+use winapi::um::processthreadsapi::STARTUPINFOW;
+use winapi::um::synchapi::WaitForSingleObject;
+use winapi::um::winbase::CREATE_NO_WINDOW;
+use winapi::um::winbase::CREATE_SUSPENDED;
+use winapi::um::winbase::WAIT_OBJECT_0;
+
+use super::JobObjectApi;
+use super::KillOnCloseJob;
+use super::SuspendedProcess;
+use super::Win32JobObjectApi;
+
+const PROCESS_WAIT_MS: u32 = 5_000;
+
+struct TestDirectory {
+    path: PathBuf,
+}
+
+impl TestDirectory {
+    fn new() -> io::Result<Self> {
+        static NEXT_DIRECTORY: AtomicU64 = AtomicU64::new(0);
+
+        let path = std::env::temp_dir().join(format!(
+            "codex-utils-pty-job-{}-{}",
+            std::process::id(),
+            NEXT_DIRECTORY.fetch_add(1, Ordering::Relaxed)
+        ));
+        fs::create_dir(&path)?;
+        Ok(Self { path })
+    }
+
+    fn join(&self, name: &str) -> PathBuf {
+        self.path.join(name)
+    }
+}
+
+impl Drop for TestDirectory {
+    fn drop(&mut self) {
+        let _ = fs::remove_dir_all(&self.path);
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum FailureStage {
+    Create,
+    Configure,
+    Assign,
+    Resume,
+    ResumeCount(u32),
+}
+
+struct FailingJobObjectApi {
+    stage: FailureStage,
+}
+
+impl FailingJobObjectApi {
+    fn failure(&self, stage: FailureStage) -> io::Result<()> {
+        if self.stage == stage {
+            Err(io::Error::other(format!("forced {stage:?} failure")))
+        } else {
+            Ok(())
+        }
+    }
+}
+
+impl JobObjectApi for FailingJobObjectApi {
+    fn create_job(&self) -> io::Result<OwnedHandle> {
+        self.failure(FailureStage::Create)?;
+        Win32JobObjectApi.create_job()
+    }
+
+    fn configure_kill_on_close(&self, job: RawHandle) -> io::Result<()> {
+        self.failure(FailureStage::Configure)?;
+        Win32JobObjectApi.configure_kill_on_close(job)
+    }
+
+    fn assign_process(&self, job: RawHandle, process: RawHandle) -> io::Result<()> {
+        self.failure(FailureStage::Assign)?;
+        Win32JobObjectApi.assign_process(job, process)
+    }
+
+    fn resume_thread(&self, primary_thread: RawHandle) -> io::Result<u32> {
+        if self.stage == FailureStage::Resume {
+            return Err(io::Error::other("forced Resume failure"));
+        }
+        if let FailureStage::ResumeCount(count) = self.stage {
+            return Ok(count);
+        }
+        Win32JobObjectApi.resume_thread(primary_thread)
+    }
+}
+
+fn spawn_suspended(command_line: &str) -> io::Result<SuspendedProcess> {
+    let mut command_line = OsStr::new(command_line)
+        .encode_wide()
+        .chain(Some(0))
+        .collect::<Vec<_>>();
+    let mut startup_info: STARTUPINFOW = unsafe { mem::zeroed() };
+    startup_info.cb = mem::size_of::<STARTUPINFOW>() as u32;
+    let mut process_info: PROCESS_INFORMATION = unsafe { mem::zeroed() };
+    let result = unsafe {
+        CreateProcessW(
+            std::ptr::null(),
+            command_line.as_mut_ptr(),
+            std::ptr::null_mut(),
+            std::ptr::null_mut(),
+            FALSE,
+            CREATE_NO_WINDOW | CREATE_SUSPENDED,
+            std::ptr::null_mut(),
+            std::ptr::null(),
+            &mut startup_info,
+            &mut process_info,
+        )
+    };
+    if result == FALSE {
+        return Err(io::Error::last_os_error());
+    }
+    Ok(unsafe {
+        SuspendedProcess::from_raw_handles(
+            process_info.hProcess.cast(),
+            process_info.hThread.cast(),
+            process_info.dwProcessId,
+        )
+    })
+}
+
+fn create_job_then_spawn<T>(
+    api: &impl JobObjectApi,
+    spawn: impl FnOnce(KillOnCloseJob) -> io::Result<T>,
+) -> io::Result<T> {
+    KillOnCloseJob::new_with_api(api).and_then(spawn)
+}
+
+fn batch_command(script: &Path) -> String {
+    format!("cmd.exe /d /c call \"{}\"", script.display())
+}
+
+fn spawn_marker_process(directory: &TestDirectory) -> io::Result<(SuspendedProcess, PathBuf)> {
+    let script = directory.join("marker.cmd");
+    let marker = directory.join("marker");
+    fs::write(&script, "@echo off\r\necho ran>\"%~dp0marker\"\r\n")?;
+    Ok((spawn_suspended(&batch_command(&script))?, marker))
+}
+
+fn wait_for_file(path: &Path) -> io::Result<()> {
+    let start = Instant::now();
+    while !path.exists() {
+        if start.elapsed() >= Duration::from_secs(5) {
+            return Err(io::Error::new(
+                io::ErrorKind::TimedOut,
+                format!("timed out waiting for {}", path.display()),
+            ));
+        }
+        thread::sleep(Duration::from_millis(10));
+    }
+    Ok(())
+}
+
+fn process_observer(process: &SuspendedProcess) -> io::Result<OwnedHandle> {
+    let Some(handle) = process.process.as_ref() else {
+        return Err(io::Error::other("suspended process has no process handle"));
+    };
+    handle.try_clone()
+}
+
+fn assert_process_exited(process: &OwnedHandle) {
+    let result = unsafe {
+        WaitForSingleObject(process.as_raw_handle().cast(), /*dwMilliseconds*/ 0)
+    };
+    assert_eq!(result, WAIT_OBJECT_0);
+}
+
+#[test]
+fn direct_job_close_terminates_root_and_grandchild() -> anyhow::Result<()> {
+    let directory = TestDirectory::new()?;
+    let root_script = directory.join("root.cmd");
+    let grandchild_script = directory.join("grandchild.cmd");
+    let grandchild_ready = directory.join("grandchild-ready");
+    let grandchild_escaped = directory.join("grandchild-escaped");
+    fs::write(
+        &root_script,
+        "@echo off\r\nstart \"\" /b cmd.exe /d /c call \"%~dp0grandchild.cmd\"\r\nping.exe -n 30 127.0.0.1 >NUL\r\n",
+    )?;
+    fs::write(
+        &grandchild_script,
+        "@echo off\r\necho ready>\"%~dp0grandchild-ready\"\r\nping.exe -n 3 127.0.0.1 >NUL\r\necho escaped>\"%~dp0grandchild-escaped\"\r\n",
+    )?;
+
+    let job = KillOnCloseJob::new()?;
+    let process = spawn_suspended(&batch_command(&root_script))?.assign_and_resume(job)?;
+
+    wait_for_file(&grandchild_ready)?;
+
+    let mut exit_code = 0;
+    let result = unsafe { GetExitCodeProcess(process.as_raw_handle().cast(), &mut exit_code) };
+    assert_ne!(result, FALSE);
+    assert_eq!(exit_code, STILL_ACTIVE);
+
+    process.controller().close()?;
+    let result = unsafe { WaitForSingleObject(process.as_raw_handle().cast(), PROCESS_WAIT_MS) };
+    assert_eq!(result, WAIT_OBJECT_0);
+
+    thread::sleep(Duration::from_secs(4));
+    assert!(grandchild_ready.exists());
+    assert!(!grandchild_escaped.exists());
+
+    Ok(())
+}
+
+#[test]
+fn closing_job_is_idempotent() -> anyhow::Result<()> {
+    let job = KillOnCloseJob::new()?;
+    job.close()?;
+    job.close()?;
+    Ok(())
+}
+
+#[test]
+fn job_creation_failure_retains_stage_context() {
+    let spawn_reached = Cell::new(false);
+    let result = create_job_then_spawn(
+        &FailingJobObjectApi {
+            stage: FailureStage::Create,
+        },
+        |_| {
+            spawn_reached.set(true);
+            Ok(())
+        },
+    );
+    let Err(err) = result else {
+        panic!("forced job creation should fail");
+    };
+    assert!(err.to_string().contains("failed to create job object"));
+    assert!(!spawn_reached.get());
+}
+
+#[test]
+fn job_configuration_failure_retains_stage_context() {
+    let spawn_reached = Cell::new(false);
+    let result = create_job_then_spawn(
+        &FailingJobObjectApi {
+            stage: FailureStage::Configure,
+        },
+        |_| {
+            spawn_reached.set(true);
+            Ok(())
+        },
+    );
+    let Err(err) = result else {
+        panic!("forced job configuration should fail");
+    };
+    assert!(
+        err.to_string()
+            .contains("failed to configure kill-on-close job object")
+    );
+    assert!(!spawn_reached.get());
+}
+
+#[test]
+fn assignment_failure_terminates_suspended_process() -> anyhow::Result<()> {
+    let directory = TestDirectory::new()?;
+    let (process, marker) = spawn_marker_process(&directory)?;
+    let observer = process_observer(&process)?;
+    let job = KillOnCloseJob::new()?;
+    let result = process.assign_and_resume_with_api(
+        job,
+        &FailingJobObjectApi {
+            stage: FailureStage::Assign,
+        },
+    );
+    let Err(err) = result else {
+        panic!("forced job assignment should fail");
+    };
+    assert!(
+        err.to_string()
+            .contains("failed to assign suspended process to job")
+    );
+    assert_process_exited(&observer);
+    assert!(!marker.exists());
+    Ok(())
+}
+
+#[test]
+fn resume_failure_terminates_suspended_process() -> anyhow::Result<()> {
+    let directory = TestDirectory::new()?;
+    let (process, marker) = spawn_marker_process(&directory)?;
+    let observer = process_observer(&process)?;
+    let job = KillOnCloseJob::new()?;
+    let result = process.assign_and_resume_with_api(
+        job,
+        &FailingJobObjectApi {
+            stage: FailureStage::Resume,
+        },
+    );
+    let Err(err) = result else {
+        panic!("forced process resume should fail");
+    };
+    assert!(
+        err.to_string()
+            .contains("failed to resume suspended process")
+    );
+    assert_process_exited(&observer);
+    assert!(!marker.exists());
+    Ok(())
+}
+
+#[test]
+fn unexpected_resume_count_terminates_suspended_process() -> anyhow::Result<()> {
+    let directory = TestDirectory::new()?;
+    let (process, marker) = spawn_marker_process(&directory)?;
+    let observer = process_observer(&process)?;
+    let job = KillOnCloseJob::new()?;
+    let result = process.assign_and_resume_with_api(
+        job,
+        &FailingJobObjectApi {
+            stage: FailureStage::ResumeCount(0),
+        },
+    );
+    let Err(err) = result else {
+        panic!("unexpected resume count should fail");
+    };
+    assert!(
+        err.to_string()
+            .contains("expected suspended process thread count 1")
+    );
+    assert_process_exited(&observer);
+    assert!(!marker.exists());
+    Ok(())
+}

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

@@ -45,10 +45,14 @@ use winapi::um::synchapi::WaitForSingleObject;
 use winapi::um::winbase::INFINITE;
 
 pub(crate) mod conpty;
+mod job;
 mod procthreadattr;
 mod psuedocon;
 
 pub use conpty::ConPtySystem;
+pub use job::JobProcess;
+pub use job::KillOnCloseJob;
+pub use job::SuspendedProcess;
 pub use psuedocon::PsuedoCon;
 pub use psuedocon::conpty_supported;