diff --git a/codex-rs/utils/pty/src/tests.rs b/codex-rs/utils/pty/src/tests.rs index f9255176f4..7da410e3e4 100644 --- a/codex-rs/utils/pty/src/tests.rs +++ b/codex-rs/utils/pty/src/tests.rs @@ -20,6 +20,14 @@ use crate::spawn_pty_process; #[path = "windows_tests.rs"] mod windows_tests; +#[cfg(windows)] +#[path = "windows_job_test_support.rs"] +mod windows_job_test_support; + +#[cfg(windows)] +#[path = "windows_conpty_job_tests.rs"] +mod windows_conpty_job_tests; + fn find_python() -> Option { for candidate in ["python3", "python"] { if let Ok(output) = std::process::Command::new(candidate) diff --git a/codex-rs/utils/pty/src/win/job.rs b/codex-rs/utils/pty/src/win/job.rs index 536c169807..bc783bf9a0 100644 --- a/codex-rs/utils/pty/src/win/job.rs +++ b/codex-rs/utils/pty/src/win/job.rs @@ -278,6 +278,10 @@ impl JobProcess { pub fn controller(&self) -> KillOnCloseJob { self.controller.clone() } + + pub(crate) fn try_clone_process_handle(&self) -> io::Result { + self.process.try_clone() + } } #[cfg(test)] diff --git a/codex-rs/utils/pty/src/win/mod.rs b/codex-rs/utils/pty/src/win/mod.rs index 70c4d07b7b..33af20d8e1 100644 --- a/codex-rs/utils/pty/src/win/mod.rs +++ b/codex-rs/utils/pty/src/win/mod.rs @@ -19,15 +19,10 @@ // SOFTWARE. // Local modifications: -// - Fix Codex bug #13945 in the Windows PTY kill path. The vendored code treated -// `TerminateProcess`'s nonzero success return as failure and `0` as success, -// which inverts kill outcomes for both `WinChild::do_kill` and -// `WinChildKiller::kill`. -// - This bug still exists in the original WezTerm source as of 2026-03-08, so -// this is an intentional divergence from upstream. +// - Keep each Windows PTY process in a kill-on-close job so terminating the +// session also terminates descendants. use anyhow::Context as _; -use filedescriptor::OwnedHandle; use portable_pty::Child; use portable_pty::ChildKiller; use portable_pty::ExitStatus; @@ -39,10 +34,11 @@ use std::sync::Mutex; use std::task::Context; use std::task::Poll; use winapi::shared::minwindef::DWORD; -use winapi::um::minwinbase::STILL_ACTIVE; +use winapi::shared::winerror::WAIT_TIMEOUT; use winapi::um::processthreadsapi::*; use winapi::um::synchapi::WaitForSingleObject; use winapi::um::winbase::INFINITE; +use winapi::um::winbase::WAIT_OBJECT_0; pub(crate) mod conpty; mod job; @@ -58,68 +54,73 @@ pub use psuedocon::conpty_supported; #[derive(Debug)] pub struct WinChild { - proc: Mutex, + process: Mutex, } impl WinChild { fn is_complete(&mut self) -> IoResult> { - let mut status: DWORD = 0; - let proc = self.proc.lock().unwrap().try_clone().unwrap(); - let res = unsafe { GetExitCodeProcess(proc.as_raw_handle() as _, &mut status) }; - if res != 0 { - if status == STILL_ACTIVE { - Ok(None) - } else { - Ok(Some(ExitStatus::with_exit_code(status))) + let process = self.process.lock().unwrap(); + let proc = process.try_clone_process_handle()?; + let controller = process.controller(); + drop(process); + + match unsafe { + WaitForSingleObject(proc.as_raw_handle() as _, /*dwMilliseconds*/ 0) + } { + WAIT_TIMEOUT => Ok(None), + WAIT_OBJECT_0 => { + let mut status: DWORD = 0; + let res = unsafe { GetExitCodeProcess(proc.as_raw_handle() as _, &mut status) }; + let status = if res != 0 { + Ok(Some(ExitStatus::with_exit_code(status))) + } else { + Err(IoError::last_os_error()) + }; + controller.close()?; + status + } + _ => { + let err = IoError::last_os_error(); + let _ = controller.terminate_and_close(/*exit_code*/ 1); + Err(err) } - } else { - Ok(None) } } fn do_kill(&mut self) -> IoResult<()> { - let proc = self.proc.lock().unwrap().try_clone().unwrap(); - let res = unsafe { TerminateProcess(proc.as_raw_handle() as _, 1) }; - // Codex bug #13945: Win32 returns nonzero on success, so only `0` is an error. - if res == 0 { - Err(IoError::last_os_error()) - } else { - Ok(()) - } + self.process + .lock() + .unwrap() + .controller() + .terminate_and_close(/*exit_code*/ 1) } } impl ChildKiller for WinChild { fn kill(&mut self) -> IoResult<()> { - self.do_kill().ok(); - Ok(()) + self.do_kill() } fn clone_killer(&self) -> Box { - let proc = self.proc.lock().unwrap().try_clone().unwrap(); - Box::new(WinChildKiller { proc }) + let controller = self.process.lock().unwrap().controller(); + Box::new(WinChildKiller { controller }) } } #[derive(Debug)] pub struct WinChildKiller { - proc: OwnedHandle, + controller: KillOnCloseJob, } impl ChildKiller for WinChildKiller { fn kill(&mut self) -> IoResult<()> { - let res = unsafe { TerminateProcess(self.proc.as_raw_handle() as _, 1) }; - // Codex bug #13945: Win32 returns nonzero on success, so only `0` is an error. - if res == 0 { - Err(IoError::last_os_error()) - } else { - Ok(()) - } + self.controller.terminate_and_close(/*exit_code*/ 1) } fn clone_killer(&self) -> Box { - let proc = self.proc.try_clone().unwrap(); - Box::new(WinChildKiller { proc }) + Box::new(WinChildKiller { + controller: self.controller.clone(), + }) } } @@ -129,30 +130,36 @@ impl Child for WinChild { } fn wait(&mut self) -> IoResult { - if let Ok(Some(status)) = self.try_wait() { + if let Some(status) = self.try_wait()? { return Ok(status); } - let proc = self.proc.lock().unwrap().try_clone().unwrap(); - unsafe { - WaitForSingleObject(proc.as_raw_handle() as _, INFINITE); + let process = self.process.lock().unwrap(); + let proc = process.try_clone_process_handle()?; + let controller = process.controller(); + drop(process); + let wait_result = unsafe { WaitForSingleObject(proc.as_raw_handle() as _, INFINITE) }; + if wait_result != WAIT_OBJECT_0 { + let err = IoError::last_os_error(); + let _ = controller.terminate_and_close(/*exit_code*/ 1); + return Err(err); } let mut status: DWORD = 0; let res = unsafe { GetExitCodeProcess(proc.as_raw_handle() as _, &mut status) }; - if res != 0 { + let status = if res != 0 { Ok(ExitStatus::with_exit_code(status)) } else { Err(IoError::last_os_error()) - } + }; + controller.close()?; + status } fn process_id(&self) -> Option { - let res = unsafe { GetProcessId(self.proc.lock().unwrap().as_raw_handle() as _) }; - if res == 0 { None } else { Some(res) } + Some(self.process.lock().unwrap().process_id()) } fn as_raw_handle(&self) -> Option { - let proc = self.proc.lock().unwrap(); - Some(proc.as_raw_handle()) + Some(self.process.lock().unwrap().as_raw_handle()) } } @@ -164,7 +171,7 @@ impl std::future::Future for WinChild { Ok(Some(status)) => Poll::Ready(Ok(status)), Err(err) => Poll::Ready(Err(err).context("Failed to retrieve process exit status")), Ok(None) => { - let proc = self.proc.lock().unwrap().try_clone()?; + let proc = self.process.lock().unwrap().try_clone_process_handle()?; let waker = cx.waker().clone(); std::thread::spawn(move || { unsafe { diff --git a/codex-rs/utils/pty/src/win/psuedocon.rs b/codex-rs/utils/pty/src/win/psuedocon.rs index f235c70278..d6f8fd2c96 100644 --- a/codex-rs/utils/pty/src/win/psuedocon.rs +++ b/codex-rs/utils/pty/src/win/psuedocon.rs @@ -20,12 +20,14 @@ // SOFTWARE. use super::WinChild; +use crate::win::KillOnCloseJob; +use crate::win::SuspendedProcess; use crate::win::procthreadattr::ProcThreadAttributeList; +use anyhow::Context as _; use anyhow::Error; use anyhow::bail; use anyhow::ensure; use filedescriptor::FileDescriptor; -use filedescriptor::OwnedHandle; use lazy_static::lazy_static; use portable_pty::cmdbuilder::CommandBuilder; use shared_library::shared_library; @@ -37,7 +39,6 @@ use std::mem; use std::os::windows::ffi::OsStrExt; use std::os::windows::ffi::OsStringExt; use std::os::windows::io::AsRawHandle; -use std::os::windows::io::FromRawHandle; use std::path::Path; use std::ptr; use std::sync::Mutex; @@ -48,6 +49,7 @@ use winapi::shared::winerror::HRESULT; use winapi::shared::winerror::S_OK; use winapi::um::handleapi::*; use winapi::um::processthreadsapi::*; +use winapi::um::winbase::CREATE_SUSPENDED; use winapi::um::winbase::CREATE_UNICODE_ENVIRONMENT; use winapi::um::winbase::EXTENDED_STARTUPINFO_PRESENT; use winapi::um::winbase::STARTF_USESTDHANDLES; @@ -173,6 +175,7 @@ impl PsuedoCon { } pub fn spawn_command(&self, cmd: CommandBuilder) -> anyhow::Result { + let job = KillOnCloseJob::new().context("failed to create process job")?; let mut si: STARTUPINFOEXW = unsafe { mem::zeroed() }; si.StartupInfo.cb = mem::size_of::() as u32; si.StartupInfo.dwFlags = STARTF_USESTDHANDLES; @@ -199,7 +202,7 @@ impl PsuedoCon { ptr::null_mut(), ptr::null_mut(), 0, - EXTENDED_STARTUPINFO_PRESENT | CREATE_UNICODE_ENVIRONMENT, + EXTENDED_STARTUPINFO_PRESENT | CREATE_UNICODE_ENVIRONMENT | CREATE_SUSPENDED, env_block.as_mut_ptr() as *mut _, cwd.as_ref().map_or(ptr::null(), std::vec::Vec::as_ptr), &mut si.StartupInfo, @@ -218,11 +221,17 @@ impl PsuedoCon { bail!("{msg}"); } - let _main_thread = unsafe { OwnedHandle::from_raw_handle(pi.hThread as _) }; - let proc = unsafe { OwnedHandle::from_raw_handle(pi.hProcess as _) }; + let suspended = unsafe { + SuspendedProcess::from_raw_handles( + pi.hProcess.cast(), + pi.hThread.cast(), + pi.dwProcessId, + ) + }; + let process = suspended.assign_and_resume(job)?; Ok(WinChild { - proc: Mutex::new(proc), + process: Mutex::new(process), }) } } diff --git a/codex-rs/utils/pty/src/windows_conpty_job_tests.rs b/codex-rs/utils/pty/src/windows_conpty_job_tests.rs new file mode 100644 index 0000000000..52dbda1104 --- /dev/null +++ b/codex-rs/utils/pty/src/windows_conpty_job_tests.rs @@ -0,0 +1,113 @@ +use super::collect_output_until_exit; +use super::combine_spawned_output; +use super::windows_job_test_support::TestDirectory; +use super::windows_job_test_support::wait_for_path; +use super::windows_job_test_support::write_descendant_scripts; +use crate::SpawnedProcess; +use crate::TerminalSize; +use crate::spawn_pty_process; +use std::collections::HashMap; +use std::path::Path; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn conpty_preserves_exit_code_that_matches_still_active() -> anyhow::Result<()> { + let program = std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string()); + let args = vec![ + "/D".to_string(), + "/Q".to_string(), + "/C".to_string(), + "exit /b 259".to_string(), + ]; + let env: HashMap = std::env::vars().collect(); + let spawned = spawn_pty_process( + &program, + &args, + Path::new("."), + &env, + &None, + TerminalSize::default(), + ) + .await?; + let (_session, output_rx, exit_rx) = combine_spawned_output(spawned); + let (_output, exit_code) = + collect_output_until_exit(output_rx, exit_rx, /*timeout_ms*/ 10_000).await; + + assert_eq!(exit_code, 259); + Ok(()) +} + +async fn spawn_conpty_script( + script: &Path, + cwd: &Path, + env: &HashMap, +) -> anyhow::Result { + let command_interpreter = std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string()); + let args = vec![ + "/D".to_string(), + "/Q".to_string(), + "/C".to_string(), + format!("call \"{}\"", script.display()), + ]; + spawn_pty_process( + &command_interpreter, + &args, + cwd, + env, + &None, + TerminalSize::default(), + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn conpty_explicit_termination_kills_descendants() -> anyhow::Result<()> { + let directory = TestDirectory::new("conpty-terminate")?; + let (root, ready, escaped) = write_descendant_scripts(&directory, false)?; + let env: HashMap = std::env::vars().collect(); + let spawned = spawn_conpty_script(&root, &directory.path, &env).await?; + let (session, output_rx, exit_rx) = combine_spawned_output(spawned); + wait_for_path(&ready).await?; + session.request_terminate(); + let (output, _exit_code) = + collect_output_until_exit(output_rx, exit_rx, /*timeout_ms*/ 10_000).await; + + assert!(!escaped.exists()); + assert!(String::from_utf8_lossy(&output).contains("inherited-grandchild-ready")); + assert!(!String::from_utf8_lossy(&output).contains("inherited-grandchild-escaped")); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn conpty_natural_root_exit_closes_job() -> anyhow::Result<()> { + let directory = TestDirectory::new("conpty-root-exit")?; + let (root, ready, escaped) = write_descendant_scripts(&directory, true)?; + let env: HashMap = std::env::vars().collect(); + let spawned = spawn_conpty_script(&root, &directory.path, &env).await?; + let (_session, output_rx, exit_rx) = combine_spawned_output(spawned); + let (output, exit_code) = + collect_output_until_exit(output_rx, exit_rx, /*timeout_ms*/ 10_000).await; + + assert_eq!(exit_code, 37); + assert!(ready.exists()); + assert!(!escaped.exists()); + assert!(String::from_utf8_lossy(&output).contains("inherited-grandchild-ready")); + assert!(!String::from_utf8_lossy(&output).contains("inherited-grandchild-escaped")); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn dropping_conpty_session_kills_descendants() -> anyhow::Result<()> { + let directory = TestDirectory::new("conpty-drop")?; + let (root, ready, escaped) = write_descendant_scripts(&directory, false)?; + let env: HashMap = std::env::vars().collect(); + let spawned = spawn_conpty_script(&root, &directory.path, &env).await?; + wait_for_path(&ready).await?; + drop(spawned.session); + drop(spawned.stdout_rx); + drop(spawned.stderr_rx); + drop(spawned.exit_rx); + tokio::time::sleep(tokio::time::Duration::from_secs(4)).await; + + assert!(!escaped.exists()); + Ok(()) +} diff --git a/codex-rs/utils/pty/src/windows_job_test_support.rs b/codex-rs/utils/pty/src/windows_job_test_support.rs new file mode 100644 index 0000000000..09c5f405a0 --- /dev/null +++ b/codex-rs/utils/pty/src/windows_job_test_support.rs @@ -0,0 +1,71 @@ +use std::fs; +use std::io; +use std::path::Path; +use std::path::PathBuf; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +pub(super) struct TestDirectory { + pub(super) path: PathBuf, +} + +impl TestDirectory { + pub(super) fn new(label: &str) -> io::Result { + static NEXT_DIRECTORY: AtomicU64 = AtomicU64::new(0); + + let path = std::env::temp_dir().join(format!( + "codex-utils-pty-{label}-{}-{}", + std::process::id(), + NEXT_DIRECTORY.fetch_add(1, Ordering::Relaxed) + )); + fs::create_dir(&path)?; + Ok(Self { path }) + } + + pub(super) 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); + } +} + +pub(super) fn write_descendant_scripts( + directory: &TestDirectory, + root_exits: bool, +) -> anyhow::Result<(PathBuf, PathBuf, PathBuf)> { + let root = directory.join("root.cmd"); + let grandchild = directory.join("grandchild.cmd"); + let ready = directory.join("grandchild-ready"); + let escaped = directory.join("grandchild-escaped"); + fs::write( + &grandchild, + "@echo off\r\necho inherited-grandchild-ready\r\necho ready>\"%~dp0grandchild-ready\"\r\nping.exe -n 4 127.0.0.1 >NUL\r\necho inherited-grandchild-escaped\r\necho escaped>\"%~dp0grandchild-escaped\"\r\n", + )?; + let final_command = if root_exits { + "exit /b 37" + } else { + "ping.exe -n 30 127.0.0.1 >NUL" + }; + fs::write( + &root, + format!( + "@echo off\r\nstart \"\" /b cmd.exe /d /q /c call \"%~dp0grandchild.cmd\"\r\n:wait\r\nif not exist \"%~dp0grandchild-ready\" (\r\n ping.exe -n 2 127.0.0.1 >NUL\r\n goto wait\r\n)\r\n{final_command}\r\n" + ), + )?; + Ok((root, ready, escaped)) +} + +pub(super) async fn wait_for_path(path: &Path) -> anyhow::Result<()> { + let timeout = tokio::time::Duration::from_secs(10); + tokio::time::timeout(timeout, async { + while !path.exists() { + tokio::time::sleep(tokio::time::Duration::from_millis(20)).await; + } + }) + .await + .map_err(|_| anyhow::anyhow!("timed out waiting for {}", path.display())) +}