feat: adding piped process to replace PTY when needed (#8797)

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

@@ -0,0 +1,174 @@
+use std::collections::HashMap;
+use std::io::ErrorKind;
+use std::path::Path;
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+use std::sync::Mutex as StdMutex;
+use std::time::Duration;
+
+use anyhow::Result;
+#[cfg(not(windows))]
+use portable_pty::native_pty_system;
+use portable_pty::CommandBuilder;
+use portable_pty::PtySize;
+use tokio::sync::broadcast;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio::task::JoinHandle;
+
+use crate::process::ChildTerminator;
+use crate::process::ProcessHandle;
+use crate::process::PtyHandles;
+use crate::process::SpawnedProcess;
+
+/// Returns true when ConPTY support is available (Windows only).
+#[cfg(windows)]
+pub fn conpty_supported() -> bool {
+    crate::win::conpty_supported()
+}
+
+/// Returns true when ConPTY support is available (non-Windows always true).
+#[cfg(not(windows))]
+pub fn conpty_supported() -> bool {
+    true
+}
+
+struct PtyChildTerminator {
+    killer: Box<dyn portable_pty::ChildKiller + Send + Sync>,
+}
+
+impl ChildTerminator for PtyChildTerminator {
+    fn kill(&mut self) -> std::io::Result<()> {
+        self.killer.kill()
+    }
+}
+
+fn platform_native_pty_system() -> Box<dyn portable_pty::PtySystem + Send> {
+    #[cfg(windows)]
+    {
+        Box::new(crate::win::ConPtySystem::default())
+    }
+
+    #[cfg(not(windows))]
+    {
+        native_pty_system()
+    }
+}
+
+/// Spawn a process attached to a PTY, returning handles for stdin, output, and exit.
+pub async fn spawn_process(
+    program: &str,
+    args: &[String],
+    cwd: &Path,
+    env: &HashMap<String, String>,
+    arg0: &Option<String>,
+) -> Result<SpawnedProcess> {
+    if program.is_empty() {
+        anyhow::bail!("missing program for PTY spawn");
+    }
+
+    let pty_system = platform_native_pty_system();
+    let pair = pty_system.openpty(PtySize {
+        rows: 24,
+        cols: 80,
+        pixel_width: 0,
+        pixel_height: 0,
+    })?;
+
+    let mut command_builder = CommandBuilder::new(arg0.as_ref().unwrap_or(&program.to_string()));
+    command_builder.cwd(cwd);
+    command_builder.env_clear();
+    for arg in args {
+        command_builder.arg(arg);
+    }
+    for (key, value) in env {
+        command_builder.env(key, value);
+    }
+
+    let mut child = pair.slave.spawn_command(command_builder)?;
+    let killer = child.clone_killer();
+
+    let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
+    let (output_tx, _) = broadcast::channel::<Vec<u8>>(256);
+    let initial_output_rx = output_tx.subscribe();
+
+    let mut reader = pair.master.try_clone_reader()?;
+    let output_tx_clone = output_tx.clone();
+    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 _ = 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 = 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 {
+            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::<i32>();
+    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.exit_code() as i32,
+            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: if cfg!(windows) {
+            Some(pair.slave)
+        } else {
+            None
+        },
+        _master: pair.master,
+    };
+
+    let (handle, output_rx) = ProcessHandle::new(
+        writer_tx,
+        output_tx,
+        initial_output_rx,
+        Box::new(PtyChildTerminator { killer }),
+        reader_handle,
+        Vec::new(),
+        writer_handle,
+        wait_handle,
+        exit_status,
+        exit_code,
+        Some(handles),
+    );
+
+    Ok(SpawnedProcess {
+        session: handle,
+        output_rx,
+        exit_rx,
+    })
+}