diff --git a/codex-rs/utils/pty/Cargo.toml b/codex-rs/utils/pty/Cargo.toml index 84801f4f1192..2a27ab72ea29 100644 --- a/codex-rs/utils/pty/Cargo.toml +++ b/codex-rs/utils/pty/Cargo.toml @@ -10,7 +10,7 @@ workspace = true [dependencies] anyhow = { workspace = true } portable-pty = { workspace = true } -tokio = { workspace = true, features = ["io-util", "macros", "process", "rt-multi-thread", "sync", "time"] } +tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt-multi-thread", "sync", "time"] } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/utils/pty/src/lib.rs b/codex-rs/utils/pty/src/lib.rs index 24c667a0ac0c..1200c1143803 100644 --- a/codex-rs/utils/pty/src/lib.rs +++ b/codex-rs/utils/pty/src/lib.rs @@ -4,6 +4,8 @@ pub mod process_group; pub mod pty; #[cfg(test)] mod tests; +#[cfg(unix)] +mod unix_io; #[cfg(windows)] mod win; #[cfg(windows)] diff --git a/codex-rs/utils/pty/src/pty.rs b/codex-rs/utils/pty/src/pty.rs index 182f96f4e33b..1b848b33b716 100644 --- a/codex-rs/utils/pty/src/pty.rs +++ b/codex-rs/utils/pty/src/pty.rs @@ -18,6 +18,7 @@ use std::process::Stdio; use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicBool; +#[cfg(not(unix))] use std::time::Duration; use anyhow::Result; @@ -159,6 +160,12 @@ async fn spawn_process_portable( ) -> Result { let pty_system = platform_native_pty_system(); let pair = pty_system.openpty(size.into())?; + #[cfg(unix)] + let io = crate::unix_io::PtyIo::new( + pair.master + .as_raw_fd() + .ok_or_else(|| anyhow::anyhow!("PTY master has no file descriptor"))?, + )?; let mut command_builder = CommandBuilder::new(arg0.as_ref().unwrap_or(&program.to_string())); command_builder.cwd(cwd); @@ -178,45 +185,56 @@ async fn spawn_process_portable( let process_group_id = child.process_id(); let killer = child.clone_killer(); - let (writer_tx, mut writer_rx) = mpsc::channel::>(128); + let (writer_tx, writer_rx) = mpsc::channel::>(128); let (stdout_tx, stdout_rx) = mpsc::channel::>(128); let (_stderr_tx, stderr_rx) = mpsc::channel::>(1); - let mut reader = pair.master.try_clone_reader()?; - let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || { - let mut buf = [0u8; 8_192]; - loop { - match reader.read(&mut buf) { - Ok(0) => break, - Ok(n) => { - let _ = stdout_tx.blocking_send(buf[..n].to_vec()); - } - Err(ref e) if e.kind() == ErrorKind::Interrupted => continue, - Err(ref e) if e.kind() == ErrorKind::WouldBlock => { - std::thread::sleep(Duration::from_millis(5)); - continue; + #[cfg(unix)] + let (reader_handle, writer_handle) = io.spawn( + stdout_tx, + writer_rx, + crate::unix_io::StdinCloseBehavior::SendEof, + ); + #[cfg(not(unix))] + let (reader_handle, writer_handle) = { + let mut reader = pair.master.try_clone_reader()?; + let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || { + let mut buf = [0u8; 8_192]; + loop { + match reader.read(&mut buf) { + Ok(0) => break, + Ok(n) => { + let _ = stdout_tx.blocking_send(buf[..n].to_vec()); + } + Err(ref e) if e.kind() == ErrorKind::Interrupted => continue, + Err(ref e) if e.kind() == ErrorKind::WouldBlock => { + std::thread::sleep(Duration::from_millis(5)); + continue; + } + Err(_) => break, } - Err(_) => break, } - } - }); - - let writer = pair.master.take_writer()?; - let writer = Arc::new(tokio::sync::Mutex::new(writer)); - let writer_handle: JoinHandle<()> = tokio::spawn({ - let writer = Arc::clone(&writer); - async move { - #[cfg(windows)] - let mut windows_input = crate::WindowsTtyInputNormalizer::default(); - while let Some(bytes) = writer_rx.recv().await { + }); + + let mut writer_rx = writer_rx; + let writer = pair.master.take_writer()?; + let writer = Arc::new(tokio::sync::Mutex::new(writer)); + let writer_handle: JoinHandle<()> = tokio::spawn({ + let writer = Arc::clone(&writer); + async move { #[cfg(windows)] - let bytes = windows_input.normalize(&bytes); - let mut guard = writer.lock().await; - use std::io::Write; - let _ = guard.write_all(&bytes); - let _ = guard.flush(); + let mut windows_input = crate::WindowsTtyInputNormalizer::default(); + while let Some(bytes) = writer_rx.recv().await { + #[cfg(windows)] + let bytes = windows_input.normalize(&bytes); + let mut guard = writer.lock().await; + use std::io::Write; + let _ = guard.write_all(&bytes); + let _ = guard.flush(); + } } - } - }); + }); + (reader_handle, writer_handle) + }; let (exit_tx, exit_rx) = oneshot::channel::(); let exit_status = Arc::new(AtomicBool::new(false)); @@ -280,6 +298,7 @@ async fn spawn_process_preserving_fds( inherited_fds: &[RawFd], ) -> Result { let (master, slave) = open_unix_pty(size)?; + let io = crate::unix_io::PtyIo::new(master.as_raw_fd())?; let mut command = StdCommand::new(program); if let Some(arg0) = arg0 { command.arg0(arg0); @@ -342,40 +361,14 @@ async fn spawn_process_preserving_fds( drop(slave); let process_group_id = child.id(); - let (writer_tx, mut writer_rx) = mpsc::channel::>(128); + let (writer_tx, writer_rx) = mpsc::channel::>(128); let (stdout_tx, stdout_rx) = mpsc::channel::>(128); let (_stderr_tx, stderr_rx) = mpsc::channel::>(1); - let mut reader = master.try_clone()?; - let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || { - let mut buf = [0u8; 8_192]; - loop { - match std::io::Read::read(&mut reader, &mut buf) { - Ok(0) => break, - Ok(n) => { - let _ = stdout_tx.blocking_send(buf[..n].to_vec()); - } - Err(ref e) if e.kind() == ErrorKind::Interrupted => continue, - Err(ref e) if e.kind() == ErrorKind::WouldBlock => { - std::thread::sleep(Duration::from_millis(5)); - continue; - } - Err(_) => break, - } - } - }); - - let writer = Arc::new(tokio::sync::Mutex::new(master.try_clone()?)); - let writer_handle: JoinHandle<()> = tokio::spawn({ - let writer = Arc::clone(&writer); - async move { - while let Some(bytes) = writer_rx.recv().await { - let mut guard = writer.lock().await; - use std::io::Write; - let _ = guard.write_all(&bytes); - let _ = guard.flush(); - } - } - }); + let (reader_handle, writer_handle) = io.spawn( + stdout_tx, + writer_rx, + crate::unix_io::StdinCloseBehavior::NoEof, + ); let (exit_tx, exit_rx) = oneshot::channel::(); let exit_status = Arc::new(AtomicBool::new(false)); diff --git a/codex-rs/utils/pty/src/tests.rs b/codex-rs/utils/pty/src/tests.rs index 8a7ddbebfa50..b3d732bf854f 100644 --- a/codex-rs/utils/pty/src/tests.rs +++ b/codex-rs/utils/pty/src/tests.rs @@ -873,6 +873,314 @@ async fn pipe_drop_reaps_child() -> anyhow::Result<()> { Ok(()) } +#[cfg(unix)] +enum PtyShutdown { + Terminate, + Drop, +} + +#[cfg(unix)] +fn assert_pty_shutdown_with_detached_child( + inherited_fds: &[i32], + shutdown: PtyShutdown, +) -> anyhow::Result<()> { + let Some(python) = find_python() else { + eprintln!("python not found; skipping detached PTY shutdown test"); + return Ok(()); + }; + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let (session, child_pid) = runtime.block_on(async { + let marker = "__codex_detached_pid:"; + // The detached child retains the slave descriptors but outlives the + // original process group. Its alarm bounds a regressed runtime drop. + let script = format!( + r"import os, select, signal, time, tty +signal.alarm(10) +if os.fork() == 0: + os.setsid() + signal.signal(signal.SIGHUP, signal.SIG_IGN) + signal.alarm(10) + tty.setraw(0) + print('{marker}' + str(os.getpid()), flush=True) + select.select([0], [], []) + print('__codex_input_ready__', flush=True) + time.sleep(10) + os._exit(0) +time.sleep(10) +" + ); + let env_map: HashMap = std::env::vars().collect(); + let spawned = spawn_pty_process( + &python, + &["-c".to_string(), script], + Path::new("."), + &env_map, + &None, + TerminalSize::default(), + inherited_fds, + ) + .await?; + let (session, mut output_rx, exit_rx) = combine_spawned_output(spawned); + let child_pid = wait_for_marker_pid(&mut output_rx, marker, /*timeout_ms*/ 2_000).await?; + let writer = session.writer_sender(); + writer.send(vec![b'x'; 1024 * 1024]).await?; + writer.send(b"pending".to_vec()).await?; + // The slave reports readable input without consuming it; the first write + // exceeds terminal capacity, so the second chunk must remain queued. + wait_for_output_contains( + &mut output_rx, + "__codex_input_ready__", + /*timeout_ms*/ 2_000, + ) + .await?; + assert_eq!(writer.capacity(), writer.max_capacity() - 1); + drop(writer); + let session = match shutdown { + PtyShutdown::Terminate => { + session.terminate(); + Some(session) + } + PtyShutdown::Drop => { + drop(session); + None + } + }; + tokio::time::timeout(std::time::Duration::from_secs(2), exit_rx).await??; + Ok::<_, anyhow::Error>((session, child_pid)) + })?; + + let started = std::time::Instant::now(); + drop(runtime); + let elapsed = started.elapsed(); + let detached_child_alive = process_exists(child_pid)?; + let _ = unsafe { libc::kill(child_pid, libc::SIGKILL) }; + drop(session); + assert!( + elapsed < std::time::Duration::from_secs(2), + "runtime shutdown waited for detached PTY child: {elapsed:?}" + ); + assert!( + detached_child_alive, + "detached child exited before runtime shutdown" + ); + Ok(()) +} + +#[cfg(unix)] +#[test] +fn pty_spawn_without_io_driver_returns_error() -> anyhow::Result<()> { + use std::os::fd::AsRawFd; + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build()?; + let preserved_fd = std::fs::File::open("/dev/null")?; + let env_map: HashMap = std::env::vars().collect(); + let (program, args) = shell_command("true"); + for inherited_fds in [&[][..], &[preserved_fd.as_raw_fd()][..]] { + let Err(error) = runtime.block_on(spawn_pty_process( + &program, + &args, + Path::new("."), + &env_map, + &None, + TerminalSize::default(), + inherited_fds, + )) else { + anyhow::bail!("PTY spawn succeeded without a Tokio I/O driver"); + }; + assert_eq!( + error.to_string(), + "PTY I/O requires a Tokio runtime with I/O enabled" + ); + } + Ok(()) +} + +#[cfg(unix)] +#[test] +fn pty_terminate_allows_runtime_shutdown_with_detached_child() -> anyhow::Result<()> { + assert_pty_shutdown_with_detached_child(&[], PtyShutdown::Terminate) +} + +#[cfg(unix)] +#[test] +fn pty_drop_allows_runtime_shutdown_with_detached_child() -> anyhow::Result<()> { + assert_pty_shutdown_with_detached_child(&[], PtyShutdown::Drop) +} + +#[cfg(unix)] +#[test] +fn pty_preserving_fds_terminate_allows_runtime_shutdown_with_detached_child() -> anyhow::Result<()> +{ + use std::os::fd::AsRawFd; + + let preserved_fd = std::fs::File::open("/dev/null")?; + assert_pty_shutdown_with_detached_child(&[preserved_fd.as_raw_fd()], PtyShutdown::Terminate) +} + +#[cfg(unix)] +#[test] +fn pty_preserving_fds_drop_allows_runtime_shutdown_with_detached_child() -> anyhow::Result<()> { + use std::os::fd::AsRawFd; + + let preserved_fd = std::fs::File::open("/dev/null")?; + assert_pty_shutdown_with_detached_child(&[preserved_fd.as_raw_fd()], PtyShutdown::Drop) +} + +#[cfg(unix)] +#[test] +fn pty_terminate_allows_runtime_shutdown_with_full_output_channel() -> anyhow::Result<()> { + use std::os::fd::AsRawFd; + + let Some(python) = find_python() else { + eprintln!("python not found; skipping PTY output backpressure test"); + return Ok(()); + }; + let preserved_fd = std::fs::File::open("/dev/null")?; + for inherited_fds in [&[][..], &[preserved_fd.as_raw_fd()][..]] { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let (session, stdout_rx) = runtime.block_on(async { + let env_map: HashMap = std::env::vars().collect(); + let script = + "import os, signal\nsignal.alarm(10)\nwhile True:\n os.write(1, b'x' * 8192)\n"; + let SpawnedProcess { + session, + stdout_rx, + exit_rx, + .. + } = spawn_pty_process( + &python, + &["-c".to_string(), script.to_string()], + Path::new("."), + &env_map, + &None, + TerminalSize::default(), + inherited_fds, + ) + .await?; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while stdout_rx.len() < stdout_rx.max_capacity() { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await?; + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + session.terminate(); + tokio::time::timeout(std::time::Duration::from_secs(2), exit_rx).await??; + Ok::<_, anyhow::Error>((session, stdout_rx)) + })?; + + // Keep the receiver alive through runtime shutdown. An OS-thread + // watchdog releases a regressed blocking_send without a Tokio timer. + let (shutdown_tx, shutdown_rx) = std::sync::mpsc::channel(); + let receiver_guard = std::thread::spawn(move || { + let _ = shutdown_rx.recv_timeout(std::time::Duration::from_secs(3)); + drop(stdout_rx); + }); + let started = std::time::Instant::now(); + drop(runtime); + let elapsed = started.elapsed(); + let _ = shutdown_tx.send(()); + receiver_guard + .join() + .map_err(|_| anyhow::anyhow!("output receiver watchdog panicked"))?; + drop(session); + assert!( + elapsed < std::time::Duration::from_secs(2), + "runtime shutdown waited on a full PTY output channel: {elapsed:?}" + ); + } + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pty_dropped_output_receiver_keeps_draining_child() -> anyhow::Result<()> { + use std::os::fd::AsRawFd; + + let Some(python) = find_python() else { + eprintln!("python not found; skipping PTY dropped output receiver test"); + return Ok(()); + }; + let preserved_fd = std::fs::File::open("/dev/null")?; + let env_map: HashMap = std::env::vars().collect(); + let script = + "import signal, sys\nsignal.alarm(10)\nsys.stdout.buffer.write(b'x' * (4 * 1024 * 1024))"; + for inherited_fds in [&[][..], &[preserved_fd.as_raw_fd()][..]] { + let SpawnedProcess { + session: _session, + stdout_rx, + exit_rx, + .. + } = spawn_pty_process( + &python, + &["-c".to_string(), script.to_string()], + Path::new("."), + &env_map, + &None, + TerminalSize::default(), + inherited_fds, + ) + .await?; + drop(stdout_rx); + let code = tokio::time::timeout(std::time::Duration::from_secs(2), exit_rx).await??; + assert_eq!(code, 0, "child should finish even when output is discarded"); + } + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pty_close_stdin_preserves_large_input_and_delivers_eof() -> anyhow::Result<()> { + let Some(python) = find_python() else { + eprintln!("python not found; skipping PTY input backpressure and EOF test"); + return Ok(()); + }; + let script = r"import signal, sys, termios, time +signal.alarm(10) +attrs = termios.tcgetattr(0) +attrs[3] &= ~termios.ECHO +termios.tcsetattr(0, termios.TCSANOW, attrs) +print('__ready__', flush=True) +time.sleep(0.2) +data = sys.stdin.buffer.read() +# Closing the portable writer appends a newline before VEOF. +assert data == b'0123456789abcdefghijklmnopqrstuvwxyz\n' * 32768 + b'\n', len(data) +print('__complete__', flush=True) +"; + let env_map: HashMap = std::env::vars().collect(); + let spawned = spawn_pty_process( + &python, + &["-c".to_string(), script.to_string()], + Path::new("."), + &env_map, + &None, + TerminalSize::default(), + &[], + ) + .await?; + let (session, mut output_rx, exit_rx) = combine_spawned_output(spawned); + wait_for_output_contains(&mut output_rx, "__ready__", /*timeout_ms*/ 2_000).await?; + let writer = session.writer_sender(); + writer + .send(b"0123456789abcdefghijklmnopqrstuvwxyz\n".repeat(32_768)) + .await?; + drop(writer); + session.close_stdin(); + + let (output, code) = collect_output_until_exit(output_rx, exit_rx, /*timeout_ms*/ 5_000).await; + assert_eq!( + (code, String::from_utf8_lossy(&output).trim()), + (0, "__complete__") + ); + Ok(()) +} + #[cfg(unix)] #[test] fn pty_terminate_reaps_child_when_waiter_is_queued() -> anyhow::Result<()> { diff --git a/codex-rs/utils/pty/src/unix_io.rs b/codex-rs/utils/pty/src/unix_io.rs new file mode 100644 index 000000000000..315c2b656afd --- /dev/null +++ b/codex-rs/utils/pty/src/unix_io.rs @@ -0,0 +1,109 @@ +use std::fs::File; +use std::io; +use std::io::Read; +use std::io::Write; +use std::os::fd::AsRawFd; +use std::os::fd::BorrowedFd; +use std::os::fd::RawFd; +use std::sync::Arc; + +use tokio::io::Interest; +use tokio::io::unix::AsyncFd; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +pub(crate) enum StdinCloseBehavior { + SendEof, + NoEof, +} + +pub(crate) struct PtyIo { + fd: Arc>, +} + +impl PtyIo { + pub(crate) fn new(master_fd: RawFd) -> io::Result { + // The PTY owner remains alive while its descriptor is duplicated. + let fd = unsafe { BorrowedFd::borrow_raw(master_fd) }.try_clone_to_owned()?; + let flags = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFL) }; + if flags == -1 + || unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFL, flags | libc::O_NONBLOCK) } == -1 + { + return Err(io::Error::last_os_error()); + } + // Duplicates share O_NONBLOCK, so both reads and writes must use readiness. + // Async tasks can be aborted even when a detached child keeps the PTY open + // or the output channel is full; blocking read/send tasks cannot. + // AsyncFd panics if the current runtime has no I/O driver. + let fd = std::panic::catch_unwind(move || AsyncFd::new(File::from(fd))) + .map_err(|_| io::Error::other("PTY I/O requires a Tokio runtime with I/O enabled"))??; + Ok(Self { fd: Arc::new(fd) }) + } + + pub(crate) fn spawn( + self, + stdout_tx: mpsc::Sender>, + mut writer_rx: mpsc::Receiver>, + stdin_close: StdinCloseBehavior, + ) -> (JoinHandle<()>, JoinHandle<()>) { + let fd = self.fd; + let reader = Arc::clone(&fd); + let reader_handle = tokio::spawn(async move { + let mut buf = [0u8; 8_192]; + loop { + match reader + .async_io(Interest::READABLE, |mut file| file.read(&mut buf)) + .await + { + Ok(0) => break, + Ok(count) => { + // Output caps may close the receiver before the child exits. + let _ = stdout_tx.send(buf[..count].to_vec()).await; + } + Err(err) if err.kind() == io::ErrorKind::Interrupted => continue, + // Unix PTYs may report EIO rather than EOF when the slave closes. + Err(_) => break, + } + } + }); + let writer_handle = tokio::spawn(async move { + while let Some(bytes) = writer_rx.recv().await { + if write_all(&fd, &bytes).await.is_err() { + return; + } + } + match stdin_close { + StdinCloseBehavior::SendEof => { + // Preserve portable-pty's newline + current VEOF on stdin close, + // including when the terminal input queue requires waiting. + let mut termios = std::mem::MaybeUninit::::uninit(); + if unsafe { libc::tcgetattr(fd.get_ref().as_raw_fd(), termios.as_mut_ptr()) } + == 0 + { + let eof = unsafe { termios.assume_init() }.c_cc[libc::VEOF]; + if eof != 0 { + let _ = write_all(&fd, &[b'\n', eof]).await; + } + } + } + StdinCloseBehavior::NoEof => {} + } + }); + (reader_handle, writer_handle) + } +} + +async fn write_all(fd: &AsyncFd, mut bytes: &[u8]) -> io::Result<()> { + while !bytes.is_empty() { + match fd + .async_io(Interest::WRITABLE, |mut file| file.write(bytes)) + .await + { + Ok(0) => return Err(io::ErrorKind::WriteZero.into()), + Ok(count) => bytes = &bytes[count..], + Err(err) if err.kind() == io::ErrorKind::Interrupted => continue, + Err(err) => return Err(err), + } + } + Ok(()) +}