R2

codex-rs/core/src/pty.rs

@@ -0,0 +1,130 @@
+use std::collections::HashMap;
+use std::io::ErrorKind;
+use std::sync::Arc;
+use std::sync::Mutex as StdMutex;
+use std::sync::atomic::AtomicBool;
+
+use portable_pty::CommandBuilder;
+use portable_pty::PtySize;
+use portable_pty::native_pty_system;
+use tokio::sync::broadcast;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio::task::JoinHandle;
+use tokio::time::Duration;
+
+use crate::exec_command::ExecCommandSession;
+
+#[derive(Debug)]
+pub(crate) struct SpawnedPty {
+    pub session: ExecCommandSession,
+    pub output_rx: broadcast::Receiver<Vec<u8>>,
+    pub exit_rx: oneshot::Receiver<i32>,
+}
+
+/// Spawn a PTY-based process and return the interactive session along with
+/// receivers for streaming output and exit status.
+pub(crate) async fn spawn_pty_process(
+    program: &str,
+    args: &[String],
+    env: &HashMap<String, String>,
+) -> anyhow::Result<SpawnedPty> {
+    if program.is_empty() {
+        anyhow::bail!("missing program for PTY spawn");
+    }
+
+    let pty_system = 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(program);
+    for arg in args {
+        command_builder.arg(arg.clone());
+    }
+    for (key, value) in env {
+        command_builder.env(key.clone(), value.clone());
+    }
+
+    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 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; 8192];
+        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(StdMutex::new(writer));
+    let writer_handle: JoinHandle<()> = tokio::spawn({
+        let writer = writer.clone();
+        async move {
+            while let Some(bytes) = writer_rx.recv().await {
+                let writer = writer.clone();
+                let _ = tokio::task::spawn_blocking(move || {
+                    if let Ok(mut guard) = writer.lock() {
+                        use std::io::Write;
+                        let _ = guard.write_all(&bytes);
+                        let _ = guard.flush();
+                    }
+                })
+                .await;
+            }
+        }
+    });
+
+    let (exit_tx, exit_rx) = oneshot::channel::<i32>();
+    let exit_status = Arc::new(AtomicBool::new(false));
+    let wait_exit_status = exit_status.clone();
+    let exit_code = Arc::new(StdMutex::new(None));
+    let wait_exit_code = exit_code.clone();
+    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 (session, output_rx) = ExecCommandSession::new(
+        writer_tx,
+        output_tx,
+        killer,
+        reader_handle,
+        writer_handle,
+        wait_handle,
+        exit_status,
+        exit_code,
+    );
+
+    Ok(SpawnedPty {
+        session,
+        output_rx,
+        exit_rx,
+    })
+}