Codex Handbook
windows-sandbox-rs/src/stdio_bridge.rs 126 lines
use std::io::Read;use std::io::Write;use std::sync::Arc;use std::time::Duration;use codex_utils_pty::SpawnedProcess;use tokio::sync::mpsc;use tokio::sync::oneshot;/// Forwards this process' stdio to a Windows sandbox session and returns the/// session exit code.pub async fn forward_sandbox_session_stdio(spawned: SpawnedProcess) -> i32 {    let session = Arc::new(spawned.session);    let tokio_runtime = tokio::runtime::Handle::current();    // Give large or slow tail output a better chance to finish draining without    // letting rare EOF issues hang the wrapper indefinitely.    let output_drain_timeout = Duration::from_secs(5);    // A helper thread watches our stdin. When the input source closes it, the    // thread tells the main async code so we can also close stdin for the    // sandboxed child process.    let (stdin_eof_tx, stdin_eof_rx) = oneshot::channel();    // Start background threads that copy stdin/stdout/stderr. We intentionally    // do not keep their JoinHandles; dropping the handle does not stop the    // thread, it just means we are not going to wait on it later.    drop(spawn_input_forwarder(        std::io::stdin(),        session.writer_sender(),        stdin_eof_tx,    ));    let (stdout_forwarder, stdout_forwarder_done_rx) =        spawn_output_forwarder(tokio_runtime.clone(), spawned.stdout_rx, std::io::stdout());    drop(stdout_forwarder);    let (stderr_forwarder, stderr_forwarder_done_rx) =        spawn_output_forwarder(tokio_runtime.clone(), spawned.stderr_rx, std::io::stderr());    drop(stderr_forwarder);    let stdin_close_task = tokio::spawn({        let session = Arc::clone(&session);        async move {            let _ = stdin_eof_rx.await;            session.close_stdin();        }    });    let mut exit_rx = spawned.exit_rx;    let exit_code = tokio::select! {        res = &mut exit_rx => res.unwrap_or(-1),        res = tokio::signal::ctrl_c() => {            if let Ok(()) = res {                session.request_terminate();            }            exit_rx.await.unwrap_or(-1)        }    };    stdin_close_task.abort();    let _ = tokio::time::timeout(output_drain_timeout, async {        let _ = stdout_forwarder_done_rx.await;        let _ = stderr_forwarder_done_rx.await;    })    .await;    exit_code}fn spawn_input_forwarder<R>(    mut input: R,    writer_tx: mpsc::Sender<Vec<u8>>,    stdin_eof_tx: oneshot::Sender<()>,) -> std::thread::JoinHandle<()>where    R: Read + Send + 'static,{    const STDIN_FORWARD_CHUNK_SIZE: usize = 8 * 1024;    std::thread::spawn(move || {        let mut buffer = [0_u8; STDIN_FORWARD_CHUNK_SIZE];        loop {            match input.read(&mut buffer) {                Ok(0) => break,                Ok(n) => {                    if writer_tx.blocking_send(buffer[..n].to_vec()).is_err() {                        break;                    }                }                Err(err) if err.kind() == std::io::ErrorKind::Interrupted => continue,                Err(err) => {                    eprintln!("windows sandbox stdin forwarder failed: {err}");                    break;                }            }        }        let _ = stdin_eof_tx.send(());    })}fn spawn_output_forwarder<W>(    tokio_runtime: tokio::runtime::Handle,    output_rx: mpsc::Receiver<Vec<u8>>,    mut writer: W,) -> (std::thread::JoinHandle<()>, oneshot::Receiver<()>)where    W: Write + Send + 'static,{    let (done_tx, done_rx) = oneshot::channel();    // The sandbox session emits output on Tokio channels, but writing to the    // caller's stdio is simplest from a dedicated blocking thread.    let handle = std::thread::spawn(move || {        let mut output_rx = output_rx;        while let Some(chunk) = tokio_runtime.block_on(output_rx.recv()) {            if let Err(err) = writer.write_all(&chunk) {                eprintln!("windows sandbox output forwarder failed to write: {err}");                break;            }            if let Err(err) = writer.flush() {                eprintln!("windows sandbox output forwarder failed to flush: {err}");                break;            }        }        let _ = done_tx.send(());    });    (handle, done_rx)}#[cfg(test)]#[path = "stdio_bridge_tests.rs"]mod tests;