diff --git a/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs b/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs index e88f0caa74..88bf5d09ba 100644 --- a/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs +++ b/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs @@ -48,12 +48,23 @@ mod imp { use crate::unified_exec::SpawnLifecycle; use codex_shell_escalation::EscalationSession; + const ESCALATE_SOCKET_ENV_VAR: &str = "CODEX_ESCALATE_SOCKET"; + #[derive(Debug)] struct ZshForkSpawnLifecycle { escalation_session: EscalationSession, } impl SpawnLifecycle for ZshForkSpawnLifecycle { + fn inherited_fds(&self) -> Vec { + self.escalation_session + .env() + .get(ESCALATE_SOCKET_ENV_VAR) + .and_then(|fd| fd.parse().ok()) + .into_iter() + .collect() + } + fn after_spawn(&mut self) { self.escalation_session.close_client_socket(); } diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index ba5352ff69..d0d5e300b6 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -1,5 +1,6 @@ #![allow(clippy::module_inception)] +use std::os::fd::RawFd; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; @@ -25,6 +26,10 @@ use super::UnifiedExecError; use super::head_tail_buffer::HeadTailBuffer; pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync { + fn inherited_fds(&self) -> Vec { + Vec::new() + } + fn after_spawn(&mut self) {} } diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 4ba125009e..6f57ecf432 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -535,14 +535,16 @@ impl UnifiedExecProcessManager { .command .split_first() .ok_or(UnifiedExecError::MissingCommandLine)?; + let inherited_fds = spawn_lifecycle.inherited_fds(); let spawn_result = if tty { - codex_utils_pty::pty::spawn_process( + codex_utils_pty::pty::spawn_process_with_inherited_fds( program, args, env.cwd.as_path(), &env.env, &env.arg0, + &inherited_fds, ) .await } else { diff --git a/codex-rs/core/tests/common/process.rs b/codex-rs/core/tests/common/process.rs index 82e42aac95..787b67be7e 100644 --- a/codex-rs/core/tests/common/process.rs +++ b/codex-rs/core/tests/common/process.rs @@ -1,6 +1,7 @@ use anyhow::Context; use std::fs; use std::path::Path; +use std::process::Stdio; use std::time::Duration; pub async fn wait_for_pid_file(path: &Path) -> anyhow::Result { @@ -24,6 +25,7 @@ pub async fn wait_for_pid_file(path: &Path) -> anyhow::Result { pub fn process_is_alive(pid: &str) -> anyhow::Result { let status = std::process::Command::new("kill") .args(["-0", pid]) + .stderr(Stdio::null()) .status() .context("failed to probe process liveness with kill -0")?; Ok(status.success()) diff --git a/codex-rs/core/tests/suite/unified_exec.rs b/codex-rs/core/tests/suite/unified_exec.rs index d345153562..d00608d5a4 100644 --- a/codex-rs/core/tests/suite/unified_exec.rs +++ b/codex-rs/core/tests/suite/unified_exec.rs @@ -32,6 +32,8 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use core_test_support::wait_for_event_with_timeout; +use core_test_support::zsh_fork::build_unified_exec_zsh_fork_test; +use core_test_support::zsh_fork::zsh_fork_runtime; use pretty_assertions::assert_eq; use regex_lite::Regex; use serde_json::Value; @@ -1323,7 +1325,6 @@ async fn exec_command_reports_chunk_and_exit_metadata() -> Result<()> { .into_iter() .map(|request| request.body_json()) .collect::>(); - let outputs = collect_tool_outputs(&bodies)?; let metadata = outputs .get(call_id) @@ -1445,7 +1446,6 @@ async fn unified_exec_defaults_to_pipe() -> Result<()> { .into_iter() .map(|request| request.body_json()) .collect::>(); - let outputs = collect_tool_outputs(&bodies)?; let output = outputs .get(call_id) @@ -1539,7 +1539,6 @@ async fn unified_exec_can_enable_tty() -> Result<()> { .into_iter() .map(|request| request.body_json()) .collect::>(); - let outputs = collect_tool_outputs(&bodies)?; let output = outputs .get(call_id) @@ -1624,7 +1623,6 @@ async fn unified_exec_respects_early_exit_notifications() -> Result<()> { .into_iter() .map(|request| request.body_json()) .collect::>(); - let outputs = collect_tool_outputs(&bodies)?; let output = outputs .get(call_id) @@ -1823,6 +1821,227 @@ async fn write_stdin_returns_exit_metadata_and_clears_session() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[cfg(unix)] +async fn unified_exec_zsh_fork_keeps_python_repl_attached_to_zsh_session() -> Result<()> { + skip_if_no_network!(Ok(())); + + let Some(runtime) = zsh_fork_runtime("unified exec zsh-fork tty session test")? else { + return Ok(()); + }; + + let python = match which("python3") { + Ok(path) => path, + Err(_) => { + eprintln!("python3 not found in PATH, skipping zsh-fork python repl test."); + return Ok(()); + } + }; + + let server = start_mock_server().await; + let test = build_unified_exec_zsh_fork_test( + &server, + runtime, + AskForApproval::Never, + SandboxPolicy::new_workspace_write_policy(), + |_| {}, + ) + .await?; + + let start_call_id = "uexec-zsh-fork-python-start"; + let send_call_id = "uexec-zsh-fork-python-pid"; + let exit_call_id = "uexec-zsh-fork-python-exit"; + + let start_args = serde_json::json!({ + "cmd": python.display().to_string(), + "yield_time_ms": 500, + "tty": true, + }); + let send_args = serde_json::json!({ + "chars": "import os; print(os.getpid())\r\n", + "session_id": 1000, + "yield_time_ms": 500, + }); + let exit_args = serde_json::json!({ + "chars": "import sys; sys.exit(0)\r\n", + "session_id": 1000, + "yield_time_ms": 500, + }); + + let responses = vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + start_call_id, + "exec_command", + &serde_json::to_string(&start_args)?, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_function_call( + send_call_id, + "write_stdin", + &serde_json::to_string(&send_args)?, + ), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_assistant_message("msg-1", "python is running"), + ev_completed("resp-3"), + ]), + sse(vec![ + ev_response_created("resp-4"), + ev_function_call( + exit_call_id, + "write_stdin", + &serde_json::to_string(&exit_args)?, + ), + ev_completed("resp-4"), + ]), + sse(vec![ + ev_assistant_message("msg-2", "all done"), + ev_completed("resp-5"), + ]), + ]; + let request_log = mount_sse_sequence(&server, responses).await; + + test.codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "test unified exec zsh-fork tty behavior".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: test.cwd_path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_workspace_write_policy(), + model: test.session_configured.model.clone(), + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = request_log.requests(); + assert!(!requests.is_empty(), "expected at least one POST request"); + let bodies = requests + .into_iter() + .map(|request| request.body_json()) + .collect::>(); + let outputs = collect_tool_outputs(&bodies)?; + + let start_output = outputs + .get(start_call_id) + .expect("missing start output for exec_command"); + let process_id = start_output + .process_id + .clone() + .expect("expected process id from exec_command"); + assert!( + start_output.exit_code.is_none(), + "initial exec_command should leave the PTY session running" + ); + + let send_output = outputs + .get(send_call_id) + .expect("missing write_stdin output"); + let normalized = send_output.output.replace("\r\n", "\n"); + let python_pid = Regex::new(r"(?m)^(\d+)$") + .expect("valid python pid regex") + .captures(&normalized) + .and_then(|captures| captures.get(1)) + .map(|value| value.as_str().to_string()) + .with_context(|| format!("missing python pid in output {normalized:?}"))?; + assert!( + process_is_alive(&python_pid)?, + "python process should still be alive after printing its pid, got output {normalized:?}" + ); + assert_eq!(send_output.process_id.as_deref(), Some(process_id.as_str())); + assert!( + send_output.exit_code.is_none(), + "write_stdin should not report an exit code while the process is still running" + ); + + let zsh_pid = std::process::Command::new("ps") + .args(["-o", "ppid=", "-p", &python_pid]) + .output() + .context("failed to look up python parent pid")?; + let zsh_pid = String::from_utf8(zsh_pid.stdout) + .context("python parent pid output is not UTF-8")? + .trim() + .to_string(); + assert!( + !zsh_pid.is_empty(), + "expected python parent pid to identify the zsh session" + ); + assert!( + process_is_alive(&zsh_pid)?, + "expected zsh parent process {zsh_pid} to still be alive" + ); + + let zsh_command = std::process::Command::new("ps") + .args(["-o", "command=", "-p", &zsh_pid]) + .output() + .context("failed to look up zsh parent command")?; + let zsh_command = + String::from_utf8(zsh_command.stdout).context("zsh parent command output is not UTF-8")?; + assert!( + zsh_command.contains("zsh"), + "expected python parent command to be zsh, got {zsh_command:?}" + ); + + test.codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "shut down the python repl".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: test.cwd_path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_workspace_write_policy(), + model: test.session_configured.model.clone(), + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = request_log.requests(); + assert!(!requests.is_empty(), "expected at least one POST request"); + let bodies = requests + .into_iter() + .map(|request| request.body_json()) + .collect::>(); + let outputs = collect_tool_outputs(&bodies)?; + let exit_output = outputs + .get(exit_call_id) + .expect("missing exit output after requesting python shutdown"); + assert!( + exit_output.exit_code.is_none() || exit_output.exit_code == Some(0), + "exit request should either leave cleanup to the background watcher or report success directly, got {exit_output:?}" + ); + wait_for_process_exit(&python_pid).await?; + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn unified_exec_emits_end_event_when_session_dies_via_stdin() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/shell-escalation/src/unix/escalate_client.rs b/codex-rs/shell-escalation/src/unix/escalate_client.rs index f9c866df4e..43ae05624a 100644 --- a/codex-rs/shell-escalation/src/unix/escalate_client.rs +++ b/codex-rs/shell-escalation/src/unix/escalate_client.rs @@ -1,6 +1,6 @@ use std::io; +use std::os::fd::AsFd; use std::os::fd::AsRawFd; -use std::os::fd::FromRawFd as _; use std::os::fd::OwnedFd; use anyhow::Context as _; @@ -28,6 +28,12 @@ fn get_escalate_client() -> anyhow::Result { Ok(unsafe { AsyncDatagramSocket::from_raw_fd(client_fd) }?) } +fn duplicate_fd_for_transfer(fd: impl AsFd, name: &str) -> anyhow::Result { + fd.as_fd() + .try_clone_to_owned() + .with_context(|| format!("failed to duplicate {name} for escalation transfer")) +} + pub async fn run_shell_escalation_execve_wrapper( file: String, argv: Vec, @@ -62,11 +68,18 @@ pub async fn run_shell_escalation_execve_wrapper( .context("failed to receive EscalateResponse")?; match message.action { EscalateAction::Escalate => { - // TODO: maybe we should send ALL open FDs (except the escalate client)? + // Duplicate stdio before transferring ownership to the server. The + // wrapper must keep using its own stdin/stdout/stderr until the + // escalated child takes over. + let destination_fds = [ + io::stdin().as_raw_fd(), + io::stdout().as_raw_fd(), + io::stderr().as_raw_fd(), + ]; let fds_to_send = [ - unsafe { OwnedFd::from_raw_fd(io::stdin().as_raw_fd()) }, - unsafe { OwnedFd::from_raw_fd(io::stdout().as_raw_fd()) }, - unsafe { OwnedFd::from_raw_fd(io::stderr().as_raw_fd()) }, + duplicate_fd_for_transfer(io::stdin(), "stdin")?, + duplicate_fd_for_transfer(io::stdout(), "stdout")?, + duplicate_fd_for_transfer(io::stderr(), "stderr")?, ]; // TODO: also forward signals over the super-exec socket @@ -74,7 +87,7 @@ pub async fn run_shell_escalation_execve_wrapper( client .send_with_fds( SuperExecMessage { - fds: fds_to_send.iter().map(AsRawFd::as_raw_fd).collect(), + fds: destination_fds.into_iter().collect(), }, &fds_to_send, ) @@ -115,3 +128,23 @@ pub async fn run_shell_escalation_execve_wrapper( } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::os::fd::AsRawFd; + use std::os::unix::net::UnixStream; + + #[test] + fn duplicate_fd_for_transfer_does_not_close_original() { + let (left, _right) = UnixStream::pair().expect("socket pair"); + let original_fd = left.as_raw_fd(); + + let duplicate = duplicate_fd_for_transfer(&left, "test fd").expect("duplicate fd"); + assert_ne!(duplicate.as_raw_fd(), original_fd); + + drop(duplicate); + + assert_ne!(unsafe { libc::fcntl(original_fd, libc::F_GETFD) }, -1); + } +} diff --git a/codex-rs/utils/pty/src/process.rs b/codex-rs/utils/pty/src/process.rs index 5c487fd386..bd5ba84f9a 100644 --- a/codex-rs/utils/pty/src/process.rs +++ b/codex-rs/utils/pty/src/process.rs @@ -1,11 +1,10 @@ use core::fmt; +use std::any::Any; use std::io; use std::sync::atomic::AtomicBool; use std::sync::Arc; use std::sync::Mutex as StdMutex; -use portable_pty::MasterPty; -use portable_pty::SlavePty; use tokio::sync::broadcast; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -17,8 +16,8 @@ pub(crate) trait ChildTerminator: Send + Sync { } pub struct PtyHandles { - pub _slave: Option>, - pub _master: Box, + pub _slave: Option>, + pub _master: Box, } impl fmt::Debug for PtyHandles { diff --git a/codex-rs/utils/pty/src/pty.rs b/codex-rs/utils/pty/src/pty.rs index 367cb5b1dc..a48a20d902 100644 --- a/codex-rs/utils/pty/src/pty.rs +++ b/codex-rs/utils/pty/src/pty.rs @@ -1,6 +1,18 @@ use std::collections::HashMap; +#[cfg(unix)] +use std::fs::File; use std::io::ErrorKind; +#[cfg(unix)] +use std::os::fd::FromRawFd; +#[cfg(unix)] +use std::os::fd::RawFd; +#[cfg(unix)] +use std::os::unix::process::CommandExt; use std::path::Path; +#[cfg(unix)] +use std::process::Command as StdCommand; +#[cfg(unix)] +use std::process::Stdio; use std::sync::atomic::AtomicBool; use std::sync::Arc; use std::sync::Mutex as StdMutex; @@ -60,6 +72,18 @@ impl ChildTerminator for PtyChildTerminator { } } +#[cfg(unix)] +struct RawPidTerminator { + process_group_id: u32, +} + +#[cfg(unix)] +impl ChildTerminator for RawPidTerminator { + fn kill(&mut self) -> std::io::Result<()> { + crate::process_group::kill_process_group(self.process_group_id) + } +} + fn platform_native_pty_system() -> Box { #[cfg(windows)] { @@ -79,11 +103,39 @@ pub async fn spawn_process( cwd: &Path, env: &HashMap, arg0: &Option, +) -> Result { + spawn_process_with_inherited_fds(program, args, cwd, env, arg0, &[]).await +} + +/// Spawn a process attached to a PTY, preserving any inherited file +/// descriptors listed in `inherited_fds` across exec on Unix. +pub async fn spawn_process_with_inherited_fds( + program: &str, + args: &[String], + cwd: &Path, + env: &HashMap, + arg0: &Option, + inherited_fds: &[i32], ) -> Result { if program.is_empty() { anyhow::bail!("missing program for PTY spawn"); } + #[cfg(unix)] + if !inherited_fds.is_empty() { + return spawn_process_preserving_fds(program, args, cwd, env, arg0, inherited_fds).await; + } + + spawn_process_portable(program, args, cwd, env, arg0).await +} + +async fn spawn_process_portable( + program: &str, + args: &[String], + cwd: &Path, + env: &HashMap, + arg0: &Option, +) -> Result { let pty_system = platform_native_pty_system(); let pair = pty_system.openpty(PtySize { rows: 24, @@ -167,11 +219,11 @@ pub async fn spawn_process( let handles = PtyHandles { _slave: if cfg!(windows) { - Some(pair.slave) + Some(Box::new(pair.slave)) } else { None }, - _master: pair.master, + _master: Box::new(pair.master), }; let (handle, output_rx) = ProcessHandle::new( @@ -198,3 +250,214 @@ pub async fn spawn_process( exit_rx, }) } + +#[cfg(unix)] +async fn spawn_process_preserving_fds( + program: &str, + args: &[String], + cwd: &Path, + env: &HashMap, + arg0: &Option, + inherited_fds: &[RawFd], +) -> Result { + let (master, slave) = open_unix_pty()?; + let mut command = StdCommand::new(arg0.as_ref().unwrap_or(&program.to_string())); + command.current_dir(cwd); + command.env_clear(); + for arg in args { + command.arg(arg); + } + for (key, value) in env { + command.env(key, value); + } + + let stdin = slave.try_clone()?; + let stdout = slave.try_clone()?; + let stderr = slave.try_clone()?; + let inherited_fds = inherited_fds.to_vec(); + + unsafe { + command + .stdin(Stdio::from(stdin)) + .stdout(Stdio::from(stdout)) + .stderr(Stdio::from(stderr)) + .pre_exec(move || { + for signo in &[ + libc::SIGCHLD, + libc::SIGHUP, + libc::SIGINT, + libc::SIGQUIT, + libc::SIGTERM, + libc::SIGALRM, + ] { + libc::signal(*signo, libc::SIG_DFL); + } + + let empty_set: libc::sigset_t = std::mem::zeroed(); + libc::sigprocmask(libc::SIG_SETMASK, &empty_set, std::ptr::null_mut()); + + if libc::setsid() == -1 { + return Err(std::io::Error::last_os_error()); + } + + #[allow(clippy::cast_lossless)] + if libc::ioctl(0, libc::TIOCSCTTY as _, 0) == -1 { + return Err(std::io::Error::last_os_error()); + } + + close_random_fds_except(&inherited_fds); + Ok(()) + }); + } + + let mut child = command.spawn()?; + drop(slave); + let process_group_id = child.id(); + + let (writer_tx, mut writer_rx) = mpsc::channel::>(128); + let (output_tx, _) = broadcast::channel::>(256); + let initial_output_rx = output_tx.subscribe(); + + let mut reader = master.try_clone()?; + let output_tx_clone = output_tx.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 _ = output_tx_clone.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 (exit_tx, exit_rx) = oneshot::channel::(); + let exit_status = Arc::new(AtomicBool::new(false)); + let wait_exit_status = Arc::clone(&exit_status); + let exit_code = Arc::new(StdMutex::new(None)); + let wait_exit_code = Arc::clone(&exit_code); + let wait_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || { + let code = match child.wait() { + Ok(status) => status.code().unwrap_or(-1), + Err(_) => -1, + }; + wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst); + if let Ok(mut guard) = wait_exit_code.lock() { + *guard = Some(code); + } + let _ = exit_tx.send(code); + }); + + let handles = PtyHandles { + _slave: None, + _master: Box::new(master), + }; + + let (handle, output_rx) = ProcessHandle::new( + writer_tx, + output_tx, + initial_output_rx, + Box::new(RawPidTerminator { process_group_id }), + reader_handle, + Vec::new(), + writer_handle, + wait_handle, + exit_status, + exit_code, + Some(handles), + ); + + Ok(SpawnedProcess { + session: handle, + output_rx, + exit_rx, + }) +} + +#[cfg(unix)] +fn open_unix_pty() -> Result<(File, File)> { + let mut master: RawFd = -1; + let mut slave: RawFd = -1; + let mut size = libc::winsize { + ws_row: 24, + ws_col: 80, + ws_xpixel: 0, + ws_ypixel: 0, + }; + + let result = unsafe { + libc::openpty( + &mut master, + &mut slave, + std::ptr::null_mut(), + std::ptr::null_mut(), + &mut size, + ) + }; + if result != 0 { + anyhow::bail!("failed to openpty: {:?}", std::io::Error::last_os_error()); + } + + set_cloexec(master)?; + set_cloexec(slave)?; + + Ok(unsafe { (File::from_raw_fd(master), File::from_raw_fd(slave)) }) +} + +#[cfg(unix)] +fn set_cloexec(fd: RawFd) -> std::io::Result<()> { + let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) }; + if flags == -1 { + return Err(std::io::Error::last_os_error()); + } + let result = unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) }; + if result == -1 { + return Err(std::io::Error::last_os_error()); + } + Ok(()) +} + +#[cfg(unix)] +fn close_random_fds_except(preserved_fds: &[RawFd]) { + if let Ok(dir) = std::fs::read_dir("/dev/fd") { + let mut fds = Vec::new(); + for entry in dir { + let num = entry + .ok() + .map(|entry| entry.file_name()) + .and_then(|name| name.into_string().ok()) + .and_then(|name| name.parse::().ok()); + if let Some(num) = num { + if num <= 2 || preserved_fds.contains(&num) { + continue; + } + fds.push(num); + } + } + for fd in fds { + unsafe { + libc::close(fd); + } + } + } +} diff --git a/codex-rs/utils/pty/src/tests.rs b/codex-rs/utils/pty/src/tests.rs index ce50adddc9..b4264f3f25 100644 --- a/codex-rs/utils/pty/src/tests.rs +++ b/codex-rs/utils/pty/src/tests.rs @@ -3,6 +3,8 @@ use std::path::Path; use pretty_assertions::assert_eq; +#[cfg(unix)] +use crate::pty::spawn_process_with_inherited_fds; use crate::spawn_pipe_process; use crate::spawn_pty_process; @@ -183,15 +185,13 @@ async fn wait_for_marker_pid( collected.extend_from_slice(&chunk); let text = String::from_utf8_lossy(&collected); - if let Some(marker_idx) = text.find(marker) { - let suffix = &text[marker_idx + marker.len()..]; - let digits: String = suffix - .chars() - .skip_while(|ch| !ch.is_ascii_digit()) - .take_while(char::is_ascii_digit) - .collect(); - if !digits.is_empty() { - return Ok(digits.parse()?); + for line in text.lines() { + let line = line.trim_end_matches('\r').trim(); + let Some(suffix) = line.strip_prefix(marker) else { + continue; + }; + if !suffix.is_empty() && suffix.chars().all(|ch| ch.is_ascii_digit()) { + return Ok(suffix.parse()?); } } } @@ -444,3 +444,119 @@ async fn pty_terminate_kills_background_children_in_same_process_group() -> anyh Ok(()) } + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pty_spawn_can_preserve_inherited_fds() -> anyhow::Result<()> { + use std::io::Read; + use std::os::fd::AsRawFd; + use std::os::fd::FromRawFd; + + let mut fds = [0; 2]; + let result = unsafe { libc::pipe(fds.as_mut_ptr()) }; + if result != 0 { + return Err(std::io::Error::last_os_error().into()); + } + + let mut read_end = unsafe { std::fs::File::from_raw_fd(fds[0]) }; + let write_end = unsafe { std::fs::File::from_raw_fd(fds[1]) }; + + let mut env_map: HashMap = std::env::vars().collect(); + env_map.insert( + "PRESERVED_FD".to_string(), + write_end.as_raw_fd().to_string(), + ); + + let script = "printf __preserved__ >&$PRESERVED_FD"; + let spawned = spawn_process_with_inherited_fds( + "/bin/sh", + &["-c".to_string(), script.to_string()], + Path::new("."), + &env_map, + &None, + &[write_end.as_raw_fd()], + ) + .await?; + + drop(write_end); + + let (_, code) = collect_output_until_exit(spawned.output_rx, spawned.exit_rx, 2_000).await; + assert_eq!(code, 0, "expected preserved-fd PTY child to exit cleanly"); + + let mut pipe_output = String::new(); + read_end.read_to_string(&mut pipe_output)?; + assert_eq!(pipe_output, "__preserved__"); + + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pty_preserving_inherited_fds_keeps_python_repl_running() -> anyhow::Result<()> { + use std::os::fd::AsRawFd; + use std::os::fd::FromRawFd; + + let Some(python) = find_python() else { + eprintln!( + "python not found; skipping pty_preserving_inherited_fds_keeps_python_repl_running" + ); + return Ok(()); + }; + + let mut fds = [0; 2]; + let result = unsafe { libc::pipe(fds.as_mut_ptr()) }; + if result != 0 { + return Err(std::io::Error::last_os_error().into()); + } + + let read_end = unsafe { std::fs::File::from_raw_fd(fds[0]) }; + let preserved_fd = unsafe { std::fs::File::from_raw_fd(fds[1]) }; + + let mut env_map: HashMap = std::env::vars().collect(); + env_map.insert( + "PRESERVED_FD".to_string(), + preserved_fd.as_raw_fd().to_string(), + ); + + let spawned = spawn_process_with_inherited_fds( + &python, + &[], + Path::new("."), + &env_map, + &None, + &[preserved_fd.as_raw_fd()], + ) + .await?; + drop(read_end); + drop(preserved_fd); + + let writer = spawned.session.writer_sender(); + let mut output_rx = spawned.output_rx; + let newline = "\n"; + let mut output = wait_for_python_repl_ready(&writer, &mut output_rx, 5_000, newline).await?; + let marker = "__codex_preserved_py_pid:"; + writer + .send(format!("import os; print('{marker}' + str(os.getpid())){newline}").into_bytes()) + .await?; + + let python_pid = match wait_for_marker_pid(&mut output_rx, marker, 2_000).await { + Ok(pid) => pid, + Err(err) => { + spawned.session.terminate(); + return Err(err); + } + }; + assert!( + process_exists(python_pid)?, + "expected python pid {python_pid} to stay alive after prompt output" + ); + + writer.send(format!("exit(){newline}").into_bytes()).await?; + let (remaining_output, code) = + collect_output_until_exit(output_rx, spawned.exit_rx, 5_000).await; + output.extend_from_slice(&remaining_output); + + assert_eq!(code, 0, "expected python to exit cleanly"); + + Ok(()) +}