fix: preserve zsh-fork escalation fds in unified-exec PTYs

Author: bolinfest

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<i32> {
+            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();
         }

codex-rs/core/src/unified_exec/process.rs

@@ -25,6 +25,10 @@ use super::UnifiedExecError;
 use super::head_tail_buffer::HeadTailBuffer;
 
 pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync {
+    fn inherited_fds(&self) -> Vec<i32> {
+        Vec::new()
+    }
+
     fn after_spawn(&mut self) {}
 }
 

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 {

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<AsyncDatagramSocket> {
     Ok(unsafe { AsyncDatagramSocket::from_raw_fd(client_fd) }?)
 }
 
+fn duplicate_fd_for_transfer(fd: impl AsFd, name: &str) -> anyhow::Result<OwnedFd> {
+    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<String>,
@@ -62,19 +68,26 @@ 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
 
             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);
+    }
+}

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<Box<dyn SlavePty + Send>>,
-    pub _master: Box<dyn MasterPty + Send>,
+    pub _slave: Option<Box<dyn Any + Send>>,
+    pub _master: Box<dyn Any + Send>,
 }
 
 impl fmt::Debug for PtyHandles {