diff --git a/codex-rs/core/src/pty.rs b/codex-rs/core/src/pty.rs new file mode 100644 index 0000000000..7e4382f90f --- /dev/null +++ b/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>, + pub exit_rx: oneshot::Receiver, +} + +/// 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, +) -> anyhow::Result { + 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::>(128); + let (output_tx, _) = broadcast::channel::>(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::(); + 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, + }) +}