diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 4429858c91..aace2e7e1d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1663,6 +1663,8 @@ version = "0.0.0" dependencies = [ "anyhow", "portable-pty", + "pretty_assertions", + "tempfile", "tokio", ] diff --git a/codex-rs/utils/pty/Cargo.toml b/codex-rs/utils/pty/Cargo.toml index d640c71aa7..c47f34996f 100644 --- a/codex-rs/utils/pty/Cargo.toml +++ b/codex-rs/utils/pty/Cargo.toml @@ -10,4 +10,8 @@ workspace = true [dependencies] anyhow = { workspace = true } portable-pty = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync"] } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync", "time"] } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/utils/pty/src/lib.rs b/codex-rs/utils/pty/src/lib.rs index 14cc430760..6a2ff28ac1 100644 --- a/codex-rs/utils/pty/src/lib.rs +++ b/codex-rs/utils/pty/src/lib.rs @@ -6,16 +6,28 @@ use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::time::Duration; +#[cfg(unix)] +use std::ffi::OsStr; +#[cfg(unix)] +use std::os::unix::fs::symlink; + use anyhow::Result; use portable_pty::native_pty_system; use portable_pty::CommandBuilder; use portable_pty::PtySize; +#[cfg(unix)] +use tempfile::TempDir; use tokio::sync::broadcast; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::Mutex as TokioMutex; use tokio::task::JoinHandle; +#[cfg(unix)] +type Arg0TempDir = Option; +#[cfg(not(unix))] +type Arg0TempDir = (); + #[derive(Debug)] pub struct ExecCommandSession { writer_tx: mpsc::Sender>, @@ -26,6 +38,7 @@ pub struct ExecCommandSession { wait_handle: StdMutex>>, exit_status: Arc, exit_code: Arc>>, + _arg0_temp_dir: Arg0TempDir, } impl ExecCommandSession { @@ -39,6 +52,7 @@ impl ExecCommandSession { wait_handle: JoinHandle<()>, exit_status: Arc, exit_code: Arc>>, + _arg0_temp_dir: Arg0TempDir, ) -> (Self, broadcast::Receiver>) { let initial_output_rx = output_tx.subscribe(); ( @@ -51,6 +65,7 @@ impl ExecCommandSession { wait_handle: StdMutex::new(Some(wait_handle)), exit_status, exit_code, + _arg0_temp_dir, }, initial_output_rx, ) @@ -106,6 +121,43 @@ pub struct SpawnedPty { pub exit_rx: oneshot::Receiver, } +#[cfg(unix)] +fn program_for_command_builder( + program: &str, + arg0: &Option, +) -> Result<(String, Arg0TempDir)> { + let Some(arg0) = arg0.as_ref() else { + return Ok((program.to_string(), None)); + }; + + let program_path = Path::new(program); + if !program_path.is_absolute() { + return Ok((program.to_string(), None)); + } + + let Some(filename) = Path::new(arg0).file_name() else { + return Ok((program.to_string(), None)); + }; + + if filename == OsStr::new(".") || filename == OsStr::new("..") { + return Ok((program.to_string(), None)); + } + + let temp_dir = TempDir::new()?; + let link_path = temp_dir.path().join(filename); + symlink(program_path, &link_path)?; + + Ok((link_path.to_string_lossy().to_string(), Some(temp_dir))) +} + +#[cfg(not(unix))] +fn program_for_command_builder( + program: &str, + _arg0: &Option, +) -> Result<(String, Arg0TempDir)> { + Ok((program.to_string(), ())) +} + pub async fn spawn_pty_process( program: &str, args: &[String], @@ -117,6 +169,8 @@ pub async fn spawn_pty_process( anyhow::bail!("missing program for PTY spawn"); } + let (program_for_builder, arg0_temp_dir) = program_for_command_builder(program, arg0)?; + let pty_system = native_pty_system(); let pair = pty_system.openpty(PtySize { rows: 24, @@ -125,7 +179,7 @@ pub async fn spawn_pty_process( pixel_height: 0, })?; - let mut command_builder = CommandBuilder::new(arg0.as_ref().unwrap_or(&program.to_string())); + let mut command_builder = CommandBuilder::new(&program_for_builder); command_builder.cwd(cwd); command_builder.env_clear(); for arg in args { @@ -201,6 +255,7 @@ pub async fn spawn_pty_process( wait_handle, exit_status, exit_code, + arg0_temp_dir, ); Ok(SpawnedPty { diff --git a/codex-rs/utils/pty/tests/arg0.rs b/codex-rs/utils/pty/tests/arg0.rs new file mode 100644 index 0000000000..f1ed92d860 --- /dev/null +++ b/codex-rs/utils/pty/tests/arg0.rs @@ -0,0 +1,58 @@ +#[cfg(unix)] +use std::collections::HashMap; +#[cfg(unix)] +use std::time::Duration; + +#[cfg(unix)] +use pretty_assertions::assert_eq; +#[cfg(unix)] +use tokio::sync::broadcast::error::RecvError; + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn spawn_pty_preserves_arg0_without_path_lookup() -> anyhow::Result<()> { + let cwd = std::env::current_dir()?; + let mut env = HashMap::new(); + env.insert("PATH".to_string(), "/usr/bin:/bin".to_string()); + + let arg0 = Some("codex-linux-sandbox".to_string()); + let args = vec!["-c".to_string(), "echo $0".to_string()]; + + let spawned = codex_utils_pty::spawn_pty_process("/bin/sh", &args, &cwd, &env, &arg0).await?; + + let mut output_rx = spawned.output_rx; + let mut exit_rx = spawned.exit_rx; + let mut collected = Vec::new(); + + let exit_code = loop { + tokio::select! { + exit_code = &mut exit_rx => break exit_code?, + chunk = output_rx.recv() => match chunk { + Ok(chunk) => collected.extend_from_slice(&chunk), + Err(RecvError::Lagged(_)) => continue, + Err(RecvError::Closed) => break -1, + } + } + }; + assert_eq!(exit_code, 0); + + loop { + match tokio::time::timeout(Duration::from_millis(25), output_rx.recv()).await { + Ok(Ok(chunk)) => collected.extend_from_slice(&chunk), + Ok(Err(RecvError::Lagged(_))) => continue, + Ok(Err(RecvError::Closed)) | Err(_) => break, + } + } + + let output = String::from_utf8_lossy(&collected); + assert!( + output.contains("codex-linux-sandbox"), + "expected argv0 to include codex-linux-sandbox, got {output:?}" + ); + + Ok(()) +} + +#[cfg(not(unix))] +#[test] +fn spawn_pty_preserves_arg0_without_path_lookup() {}