diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 18dfaff6d0..af8f3639a0 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3176,6 +3176,7 @@ dependencies = [ "codex-shell-command", "codex-skills", "codex-test-binary-support", + "codex-uds", "codex-utils-absolute-path", "codex-utils-cargo-bin", "codex-utils-home-dir", diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index c390da762f..fd9928db36 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -73,8 +73,10 @@ url = { workspace = true } uuid = { workspace = true, features = ["v4"] } [target.'cfg(unix)'.dependencies] +codex-uds = { workspace = true } libc = { workspace = true } rustix = { workspace = true, features = ["fs"] } +tempfile = { workspace = true } [target.'cfg(windows)'.dependencies] windows-sys = { version = "0.52", features = [ diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index be0b6f31fe..0851fd918e 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -46,6 +46,8 @@ mod sandboxed_file_system; mod server; #[cfg(unix)] mod shell_snapshot; +#[cfg(unix)] +mod shell_snapshot_file; mod telemetry; mod trace_context; mod websocket_pong_watchdog; diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index cddd17052a..a84ad8e329 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -379,7 +379,8 @@ impl LocalProcess { #[cfg(unix)] let mut prepared = prepared; #[cfg(unix)] - self.inner + let snapshot_file = self + .inner .shell_snapshots .prepare( ¶ms, @@ -413,6 +414,13 @@ impl LocalProcess { ); } + #[cfg(unix)] + let inherited_fds = snapshot_file + .iter() + .map(std::os::fd::AsRawFd::as_raw_fd) + .collect::>(); + #[cfg(not(unix))] + let inherited_fds = Vec::new(); let spawned_result = codex_sandboxing::spawn_process(codex_sandboxing::SpawnRequest { command: &prepared.command, cwd: prepared.cwd.as_path(), @@ -422,9 +430,11 @@ impl LocalProcess { windows_sandbox: prepared.windows_sandbox_spawn_request(), tty: params.tty, stdin_open: params.tty || params.pipe_stdin, - inherited_fds: codex_utils_pty::ChildFds::Inherited(&[]), + inherited_fds: codex_utils_pty::ChildFds::Attached(&inherited_fds), }) .await; + #[cfg(unix)] + drop(snapshot_file); let spawned = match spawned_result { Ok(spawned) => spawned, Err(err) => { @@ -827,6 +837,7 @@ impl ExecBackend for LocalProcess { CapturePurpose::Prewarm, ) .await + .map(|_| ()) .map_err(map_handler_error) }) } diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs index 0ba65b9dc7..b0003f6ea8 100644 --- a/codex-rs/exec-server/src/shell_snapshot.rs +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -1,5 +1,10 @@ +//! Cache captured shell state in executor memory. Replay through a per-launch +//! unnamed reader when the capture sandbox permits it; otherwise retain env replay. + use std::collections::HashMap; use std::collections::VecDeque; +use std::fs::File; +use std::os::fd::AsRawFd; use std::process::Stdio; use std::sync::Arc; use std::time::Duration; @@ -13,7 +18,7 @@ use codex_shell_command::shell_detect::ShellType; use codex_shell_command::shell_snapshot::CapturedSnapshot; use codex_shell_command::shell_snapshot::SnapshotCaptureOptions; use codex_shell_command::shell_snapshot::SnapshotStartup; -use codex_shell_command::shell_snapshot::snapshot_capture_script; +use codex_shell_command::shell_snapshot::snapshot_source_capture_script; use codex_utils_path_uri::PathUri; use tokio::io::AsyncReadExt; use tokio::process::Command; @@ -58,6 +63,7 @@ struct CachedShellSnapshot { struct ShellSnapshot { state: String, + file_source: bool, environment: HashMap, } @@ -85,9 +91,9 @@ impl ShellSnapshotCache { prepared: &mut PreparedExecRequest, telemetry: &ExecServerTelemetry, purpose: CapturePurpose, - ) -> Result<(), JSONRPCErrorError> { + ) -> Result, JSONRPCErrorError> { let Some(request) = params.shell_snapshot.as_ref() else { - return Ok(()); + return Ok(None); }; if request.scope_id.is_empty() || request.scope_id.len() > MAX_SNAPSHOT_SCOPE_BYTES { return Err(invalid_params(format!( @@ -100,7 +106,7 @@ impl ShellSnapshotCache { || params.argv[1] != "-lc" || !prepared.command.ends_with(¶ms.argv) { - return Ok(()); + return Ok(None); } let shell_type = match request.shell.name.as_str() { @@ -199,12 +205,34 @@ impl ShellSnapshotCache { } }; let Ok(snapshot) = snapshot else { - return Ok(()); + return Ok(None); }; if purpose == CapturePurpose::Prewarm { - return Ok(()); + return Ok(None); } + // POSIX sh cannot portably close arbitrary descriptors above 9. Keep its + // existing replay until the launcher supports remapping child descriptors. + let reader = if snapshot.file_source { + let state = snapshot.state.clone(); + match tokio::task::spawn_blocking(move || { + crate::shell_snapshot_file::materialize(shell_type, &state) + }) + .await + { + Ok(Ok(reader)) => Some(reader), + error => { + tracing::warn!( + ?error, + "cannot prepare shell snapshot transport; using normal startup" + ); + return Ok(None); + } + } + } else { + None + }; + let request_overrides = params .env .iter() @@ -226,24 +254,29 @@ impl ShellSnapshotCache { .env .retain(|name, _| !shell_environment::is_non_inheritable_env_var(name)); - let mut state = snapshot.state.as_str(); - let mut state_variables = Vec::new(); - while !state.is_empty() { - let mut end = state.len().min(MAX_SNAPSHOT_ENV_VALUE_BYTES); - while !state.is_char_boundary(end) { - end -= 1; + let restore = if let Some(reader) = &reader { + format!(". /dev/fd/{}", reader.as_raw_fd()) + } else { + let mut state = snapshot.state.as_str(); + let mut state_variables = Vec::new(); + while !state.is_empty() { + let mut end = state.len().min(MAX_SNAPSHOT_ENV_VALUE_BYTES); + while !state.is_char_boundary(end) { + end -= 1; + } + let (chunk, remaining) = state.split_at(end); + let name = format!("__CODEX_SHELL_SNAPSHOT_STATE_{}", state_variables.len()); + prepared.env.insert(name.clone(), chunk.to_string()); + state_variables.push(name); + state = remaining; } - let (chunk, remaining) = state.split_at(end); - let name = format!("__CODEX_SHELL_SNAPSHOT_STATE_{}", state_variables.len()); - prepared.env.insert(name.clone(), chunk.to_string()); - state_variables.push(name); - state = remaining; - } - let state_expansion = state_variables - .iter() - .map(|name| format!("${{{name}}}")) - .collect::(); - let state_variables = state_variables.join(" "); + let state_expansion = state_variables + .iter() + .map(|name| format!("${{{name}}}")) + .collect::(); + let state_variables = state_variables.join(" "); + format!("eval \"unset {state_variables}\n{state_expansion}\"") + }; let shell_start = prepared.command.len() - params.argv.len(); // Automatic startup files run before the restoration script and could // reintroduce environment variables that the snapshot already filtered. @@ -255,11 +288,11 @@ impl ShellSnapshotCache { }; prepared.command[shell_start + 1] = shell_flag.to_string(); prepared.command[shell_start + 2] = format!( - "{startup}if ! eval \"unset {state_variables}\n{state_expansion}\" >/dev/null; then printf 'failed to restore shell snapshot\\n' >&2; fi\n{}", + "{startup}if ! {restore} >/dev/null; then printf 'failed to restore shell snapshot\\n' >&2; fi\n{}", params.argv[2] ); - Ok(()) + Ok(reader) } } @@ -273,7 +306,7 @@ async fn capture_snapshot( prepared: &PreparedExecRequest, shell_type: ShellType, ) -> CaptureResult { - let script = snapshot_capture_script( + let mut script = snapshot_source_capture_script( shell_type, SnapshotCaptureOptions { startup: SnapshotStartup::Interactive, @@ -287,6 +320,27 @@ async fn capture_snapshot( invalid_params("unsupported shell snapshot script".to_string()), ) })?; + // Probe under the exact capture sandbox, including its /proc fallback and + // read policy. Keep cached env replay when descriptor paths cannot be opened. + let probe = if shell_type == ShellType::Sh { + None + } else { + tokio::task::spawn_blocking(move || crate::shell_snapshot_file::materialize(shell_type, "")) + .await + .ok() + .and_then(Result::ok) + }; + const SOURCE_MARKER: &str = "CODEX_SNAPSHOT_SOURCE"; + script = if let Some(probe) = &probe { + let fd = probe.as_raw_fd(); + // Automatic startup has already run: do not call exec/unset here, + // which may be user functions. Capture exit closes this harmless probe. + format!( + "if case '' in '') ;; esac 2>/dev/null io::Result { + // Mode 0700 alone cannot stop a same-user sandboxed command from replacing + // the pathname or retaining a writable handle before unlink. Reuse the root + // that every filesystem-restricted sandbox hides, independent of CODEX_HOME. + let root = codex_uds::prepare_shared_daemon_socket_directory()?; + let directory = tempfile::tempdir_in(root)?; + let writer = tempfile::NamedTempFile::new_in(directory.path())?; + // Open independently, rather than dup: writing must not advance the reader. + let reader = File::open(writer.path())?; + let reader = File::from(rustix::io::fcntl_dupfd_cloexec(&reader, /*min*/ 10)?); + let (mut writer, path) = writer.into_parts(); + path.close()?; + directory.close()?; + if reader.metadata()?.nlink() != 0 { + return Err(io::Error::other("shell snapshot transport is still linked")); + } + + let fd = reader.as_raw_fd(); + // Source has opened its own reader before this runs. Close the carrier before + // restoring functions that may shadow exec/unset, without borrowing stdin. + let close = match shell { + ShellType::Bash => format!("\\exec {fd}<&-\n"), + ShellType::Zsh => format!( + "__codex_snapshot_fd={fd}\nexec {{__codex_snapshot_fd}}<&-\nunset __codex_snapshot_fd\n" + ), + ShellType::Sh | ShellType::PowerShell | ShellType::Cmd => { + return Err(io::Error::other("unsupported snapshot shell")); + } + }; + writer.write_all(close.as_bytes())?; + writer.write_all(state.as_bytes())?; + Ok(reader) +} + +#[cfg(test)] +#[path = "shell_snapshot_file_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/shell_snapshot_file_tests.rs b/codex-rs/exec-server/src/shell_snapshot_file_tests.rs new file mode 100644 index 0000000000..b03a77eb63 --- /dev/null +++ b/codex-rs/exec-server/src/shell_snapshot_file_tests.rs @@ -0,0 +1,74 @@ +use super::materialize; +use codex_shell_command::shell_detect::ShellType; +use pretty_assertions::assert_eq; +use std::collections::HashMap; +use std::io::Write; +use std::os::fd::AsRawFd; +use std::os::unix::fs::MetadataExt; +use std::path::Path; +use test_case::test_case; + +#[test_case("/bin/bash", ShellType::Bash; "bash")] +#[cfg_attr(target_os = "macos", test_case("/bin/zsh", ShellType::Zsh; "zsh"))] +#[tokio::test] +async fn unnamed_snapshot_preserves_stdin_and_closes_read_only_carrier( + shell: &str, + shell_type: ShellType, +) -> anyhow::Result<()> { + let state = "helper() { printf 'restored:%s' \"$1\"; }\nexec() { exit 42; }\n"; + let mut reader = materialize(shell_type, state)?; + let fd = reader.as_raw_fd(); + #[cfg(target_os = "linux")] + assert!( + std::fs::read_link(format!("/proc/self/fd/{fd}"))? + .starts_with(codex_uds::shared_daemon_socket_directory()?), + "the transport must be protected even before it is unlinked" + ); + let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) }; + assert_eq!( + ( + reader.metadata()?.nlink(), + flags & libc::FD_CLOEXEC, + reader.write(b"x").unwrap_err().raw_os_error() + ), + (0, libc::FD_CLOEXEC, Some(libc::EBADF)), + ); + let closed_check = format!("if [ -e /dev/fd/{fd} ]; then exit 43; fi\n"); + let args = vec![ + "-c".to_string(), + format!(". /dev/fd/{fd}\n{closed_check}IFS= read -r line; helper \"$line\"; exit 7"), + ]; + let env = HashMap::from([("PATH".to_string(), "/usr/bin:/bin".to_string())]); + let spawned = + codex_utils_pty::spawn_pipe_process(shell, &args, Path::new("/"), &env, &None, &[fd]) + .await?; + assert_eq!(unsafe { libc::fcntl(fd, libc::F_GETFD) }, flags); + drop(reader); + let codex_utils_pty::SpawnedProcess { + session, + mut stdout_rx, + mut stderr_rx, + exit_rx, + } = spawned; + session.writer_sender().send(b"input\n".to_vec()).await?; + let result = tokio::time::timeout(std::time::Duration::from_secs(10), async { + let mut output = Vec::new(); + while let Some(chunk) = stdout_rx.recv().await { + output.extend(chunk); + } + let mut errors = Vec::new(); + while let Some(chunk) = stderr_rx.recv().await { + errors.extend(chunk); + } + Ok::<_, anyhow::Error>((String::from_utf8(output)?, errors, exit_rx.await?)) + }) + .await; + session.terminate(); + let (output, errors, status) = result??; + assert!( + output.ends_with("restored:input"), + "stdout={output:?}, stderr={errors:?}, status={status}" + ); + assert_eq!((errors, status), (Vec::new(), 7)); + Ok(()) +} diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs index cdab0f7508..c79eb72589 100644 --- a/codex-rs/exec-server/src/shell_snapshot_tests.rs +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -123,12 +123,20 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( cache.prepare(¶ms, &mut prepared, &telemetry, purpose), cache.prepare(¶ms, &mut concurrent, &telemetry, CapturePurpose::Execution), ); - if prewarming { + let first = if prewarming { first.expect_err("prewarm must report failure without caching it"); + None } else { - first.expect("capture failure must preserve command fallback"); + first.expect("capture failure must preserve command fallback") + }; + let second = second.expect("waiting command must complete even when prewarm fails"); + for (prepared, reader) in [(&mut prepared, first), (&mut concurrent, second)] { + if let Some(reader) = reader { + let fd = std::os::fd::AsRawFd::as_raw_fd(&reader); + prepared.command[2] = + prepared.command[2].replace(&format!("/dev/fd/{fd}"), "/dev/fd/SNAPSHOT"); + } } - second.expect("waiting command must complete even when prewarm fails"); assert_eq!( (&prepared.command, &prepared.env), (&concurrent.command, &concurrent.env) diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index 095cbe25e1..a07512282e 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -1,4 +1,7 @@ mod common; +#[cfg(unix)] +#[path = "exec_process/shell_snapshot.rs"] +mod shell_snapshot; #[path = "exec_process/windows_sandbox.rs"] mod windows_sandbox; @@ -250,7 +253,7 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( let shadowed_builtins = if posix_shell { "" } else { - "unset() { exit 41; }\nbuiltin() { :; }\n" + "unset() { exit 41; }\nbuiltin() { exit 41; }\nexec() { exit 41; }\n" }; std::fs::write( &profile_path, diff --git a/codex-rs/exec-server/tests/exec_process/shell_snapshot.rs b/codex-rs/exec-server/tests/exec_process/shell_snapshot.rs new file mode 100644 index 0000000000..9f1396bf79 --- /dev/null +++ b/codex-rs/exec-server/tests/exec_process/shell_snapshot.rs @@ -0,0 +1,174 @@ +use super::*; +use pretty_assertions::assert_eq; +use test_case::test_case; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum SnapshotSandbox { + None, + WorkspaceWrite, + DenyFdPath, +} + +#[test_case("bash", false, SnapshotSandbox::None; "bash")] +#[test_case("bash", true, SnapshotSandbox::None; "bash_tty")] +#[cfg_attr(target_os = "macos", test_case("zsh", false, SnapshotSandbox::None; "zsh"))] +#[cfg_attr(target_os = "macos", test_case("zsh", true, SnapshotSandbox::None; "zsh_tty"))] +#[test_case("bash", false, SnapshotSandbox::DenyFdPath; "blocked_descriptor_path")] +#[test_case("bash", false, SnapshotSandbox::WorkspaceWrite; "bash_protected_transport")] +#[test_case("bash", true, SnapshotSandbox::WorkspaceWrite; "bash_protected_transport_tty")] +#[cfg_attr(target_os = "macos", test_case("zsh", false, SnapshotSandbox::WorkspaceWrite; "zsh_protected_transport"))] +#[cfg_attr(target_os = "macos", test_case("zsh", true, SnapshotSandbox::WorkspaceWrite; "zsh_protected_transport_tty"))] +#[serial_test::serial(remote_exec_server)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_concurrent_replays_keep_independent_readers( + shell: &str, + tty: bool, + sandbox_mode: SnapshotSandbox, +) -> Result<()> { + let deny_fd_path = sandbox_mode == SnapshotSandbox::DenyFdPath; + let use_sandbox = sandbox_mode != SnapshotSandbox::None; + if use_sandbox + && let Some(warning) = + codex_sandboxing::system_bwrap_warning(&PermissionProfile::read_only()) + { + eprintln!("skipping sandbox test: {warning}"); + return Ok(()); + } + let context = if sandbox_mode == SnapshotSandbox::WorkspaceWrite { + let (exe, linux_sandbox) = current_test_binary_helper_paths()?; + let environment = Environment::create( + /*exec_server_url*/ None, + codex_exec_server::ExecServerRuntimePaths::new(exe, linux_sandbox)?, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + )?; + ProcessContext { + backend: environment.get_exec_backend(), + _server: None, + } + } else { + create_process_context(deny_fd_path).await? + }; + let home = TempDir::new()?; + let cwd = PathUri::from_host_native_path(home.path())?; + // Keep full-size state in a function, below the 512 KiB state + environment cap. + // The blocked-path case still exercises the smaller environment fallback. + let payload_len = if deny_fd_path { 1 } else { 480 * 1024 }; + let payload = "x".repeat(payload_len); + let source = if shell == "bash" { + "${BASH_SOURCE[0]}" + } else { + "${(%):-%x}" + }; + let source_check = if deny_fd_path { + String::new() + } else { + format!("case \"{source}\" in /dev/fd/*) ;; *) return 42 ;; esac; ") + }; + std::fs::write( + home.path().join(format!(".{shell}rc")), + format!( + "printf x >> \"$HOME/captures\"\nprofile_helper() {{ {source_check}local payload='{payload}'; [ \"${{#payload}}\" = {payload_len} ] || return 43; printf 'restored:%s' \"$1\"; }}\nexec() {{ exit 41; }}\nset -u\n" + ), + )?; + let protected_file = + tempfile::NamedTempFile::new_in(codex_uds::prepare_shared_daemon_socket_directory()?)?; + std::fs::write(protected_file.path(), "protected")?; + let command = if sandbox_mode == SnapshotSandbox::WorkspaceWrite { + // Broad /tmp grants must not expose the file while it still has a name. + "if /bin/cat \"$SNAPSHOT_TRANSPORT_PROBE\" >/dev/null 2>&1 || (printf poisoned >> \"$SNAPSHOT_TRANSPORT_PROBE\") 2>/dev/null || /bin/mv \"$SNAPSHOT_TRANSPORT_PROBE\" \"$HOME/stolen\" 2>/dev/null; then exit 44; fi\n" + } else { + "" + }; + let command = format!("{command}IFS= read -r line; profile_helper \"$line\"; exit 7"); + let sandbox = if use_sandbox { + let mut policy = FileSystemSandboxPolicy::read_only(); + policy.entries.push(FileSystemSandboxEntry::new( + cwd.clone().into(), + FileSystemAccessMode::Write, + )); + if deny_fd_path { + policy.entries.push(FileSystemSandboxEntry::new( + PathUri::from_host_native_path("/dev/fd")?.into(), + FileSystemAccessMode::Deny, + )); + } else { + policy.entries.push(FileSystemSandboxEntry::new( + PathUri::from_host_native_path(std::fs::canonicalize("/tmp")?)?.into(), + FileSystemAccessMode::Write, + )); + } + Some(FileSystemSandboxContext::from_permission_profile( + PermissionProfile::from_runtime_permissions(&policy, NetworkSandboxPolicy::Restricted), + cwd.clone(), + )) + } else { + None + }; + let commands = (0..8).map(|index| { + let backend = &context.backend; + let cwd = cwd.clone(); + let home = home.path(); + let sandbox = sandbox.clone(); + let command = command.clone(); + let protected_path = protected_file.path(); + async move { + let started = backend + .start(ExecParams { + metadata: Default::default(), + process_id: format!("parallel-{index}").into(), + argv: vec![format!("/bin/{shell}"), "-lc".to_string(), command], + cwd, + env: HashMap::new(), + env_policy: Some(ExecEnvPolicy { + inherit: ShellEnvironmentPolicyInherit::None, + ignore_default_excludes: false, + exclude: Vec::new(), + r#set: HashMap::from([ + ("HOME".to_string(), home.to_string_lossy().into_owned()), + ("PATH".to_string(), "/usr/bin:/bin".to_string()), + ( + "SNAPSHOT_TRANSPORT_PROBE".to_string(), + protected_path.to_string_lossy().into_owned(), + ), + ]), + include_only: Vec::new(), + }), + shell_snapshot: Some(ShellSnapshotRequest { + scope_id: "parallel".to_string(), + shell: ShellInfo { + name: shell.to_string(), + path: format!("/bin/{shell}"), + }, + }), + tty, + pipe_stdin: true, + arg0: None, + sandbox, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }) + .await?; + started + .process + .write(format!("input-{index}\n").into_bytes()) + .await?; + let (output, errors, status, closed) = + collect_process_output_from_events(started.process).await?; + assert!( + output.ends_with(&format!("restored:input-{index}")), + "{output:?}" + ); + assert_eq!((errors, status, closed), (String::new(), Some(7), true)); + Ok::<_, anyhow::Error>(()) + } + }); + for result in futures::future::join_all(commands).await { + result?; + } + assert_eq!(std::fs::read_to_string(home.path().join("captures"))?, "x"); + assert_eq!(std::fs::read_to_string(protected_file.path())?, "protected"); + Ok(()) +} diff --git a/codex-rs/shell-command/src/shell_snapshot.rs b/codex-rs/shell-command/src/shell_snapshot.rs index 5ebdd569af..a1fa1740a1 100644 --- a/codex-rs/shell-command/src/shell_snapshot.rs +++ b/codex-rs/shell-command/src/shell_snapshot.rs @@ -15,6 +15,7 @@ pub use capture::CapturedSnapshot; pub use capture::CapturedStartupEnvironment; pub use capture::SnapshotCaptureOptions; pub use capture::snapshot_capture_script; +pub use capture::snapshot_source_capture_script; pub use credentials::PreparedSnapshot; pub use credentials::SnapshotCredentialEnvironment; pub use credentials::prepare_snapshot_credentials; diff --git a/codex-rs/shell-command/src/shell_snapshot_capture.rs b/codex-rs/shell-command/src/shell_snapshot_capture.rs index 80042fd626..547f146edc 100644 --- a/codex-rs/shell-command/src/shell_snapshot_capture.rs +++ b/codex-rs/shell-command/src/shell_snapshot_capture.rs @@ -42,6 +42,36 @@ pub struct SnapshotCaptureOptions { pub fn snapshot_capture_script( shell_type: ShellType, options: SnapshotCaptureOptions, +) -> Option { + capture_script(shell_type, options, Replay::Eval) +} + +/// Capture state for incremental sourcing. Parse options and aliases together so +/// restored aliases and quoting options cannot reinterpret later declarations. +pub fn snapshot_source_capture_script( + shell_type: ShellType, + options: SnapshotCaptureOptions, +) -> Option { + capture_script( + shell_type, + options, + if shell_type == ShellType::Sh { + Replay::Eval + } else { + Replay::Source + }, + ) +} + +enum Replay { + Eval, + Source, +} + +fn capture_script( + shell_type: ShellType, + options: SnapshotCaptureOptions, + replay: Replay, ) -> Option { let script = match shell_type { ShellType::Zsh => zsh_snapshot_script(options.startup), @@ -56,6 +86,8 @@ pub fn snapshot_capture_script( }; Some( script + .replace("SNAPSHOT_OPTIONS_BEGIN", match replay { Replay::Eval => "", Replay::Source => "printf '{\\n'" }) + .replace("SNAPSHOT_ALIASES_END", match replay { Replay::Eval => "", Replay::Source => "printf \"case '' in '') ;; esac\\n}\\n\"" }) .replace("SNAPSHOT_EXPORTS", &declarations) .replace( "SNAPSHOT_DECLARATION_ENVIRONMENT", @@ -92,16 +124,18 @@ print 'unalias -a 2>/dev/null || true' print '# Functions' functions print '' +SNAPSHOT_OPTIONS_BEGIN SNAPSHOT_COMMAND_HELPER setopt_count=$(setopt | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') print "# setopts $setopt_count" setopt | __codex_snapshot_command sed 's/^/setopt /' print '' printf '\0' -alias_count=$(alias -L | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') +alias_count=$(\alias -L | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') print "# aliases $alias_count" -alias -L +\alias -L print '' +SNAPSHOT_ALIASES_END printf '\0' SNAPSHOT_EXPORTS printf '\0' @@ -122,6 +156,7 @@ shopt -p || true echo '# Functions' declare -f echo '' +SNAPSHOT_OPTIONS_BEGIN SNAPSHOT_COMMAND_HELPER bash_opts=$(set -o | __codex_snapshot_command awk '$2=="on"{print $1}') bash_opt_count=$(printf '%s\n' "$bash_opts" | __codex_snapshot_command sed '/^$/d' | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') @@ -131,10 +166,11 @@ if [ -n "$bash_opts" ]; then fi echo '' printf '\0' -alias_count=$(alias -p | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') +alias_count=$(\alias -p | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') echo "# aliases $alias_count" -alias -p +\alias -p echo '' +SNAPSHOT_ALIASES_END printf '\0' SNAPSHOT_EXPORTS printf '\0' @@ -180,6 +216,7 @@ elif command -v declare >/dev/null 2>&1; then declare -f fi echo '' +SNAPSHOT_OPTIONS_BEGIN SNAPSHOT_COMMAND_HELPER if set -o >/dev/null 2>&1; then sh_opts=$(set -o | __codex_snapshot_command awk '$2=="on"{print $1}') @@ -201,6 +238,7 @@ if alias >/dev/null 2>&1; then else echo '# aliases 0' fi +SNAPSHOT_ALIASES_END printf '\0' SNAPSHOT_EXPORTS printf '\0' diff --git a/codex-rs/shell-command/src/shell_snapshot_tests.rs b/codex-rs/shell-command/src/shell_snapshot_tests.rs index c8b877d40d..3ab21128e3 100644 --- a/codex-rs/shell-command/src/shell_snapshot_tests.rs +++ b/codex-rs/shell-command/src/shell_snapshot_tests.rs @@ -2233,3 +2233,53 @@ fn captured_script(shell_type: ShellType, source: &str) -> Result { CapturedSnapshot::parse(shell_type, source.as_bytes()).context("invalid native capture")?; Ok(captured.render_script()) } + +#[test] +fn streamed_snapshot_preserves_alias_and_option_parse_order() -> Result<()> { + for (shell, shell_type, options) in [ + ("/bin/bash", ShellType::Bash, "shopt -s expand_aliases"), + ("/bin/zsh", ShellType::Zsh, "setopt RC_QUOTES"), + ] { + if !std::path::Path::new(shell).exists() { + continue; + } + let dir = tempdir()?; + let mut setup = format!( + "{options}\nhelper() {{ printf helper; }}\nfunction : () {{ exit 41; }}\nalias snapshot_alias=\"printf 'one''two'\"\n" + ); + if shell_type == ShellType::Zsh { + setup.push_str("alias alias=false\n"); + } + let mut results = Vec::new(); + for source in [false, true] { + let capture = if source { + super::snapshot_source_capture_script(shell_type, CAPTURE_NON_INTERACTIVE) + } else { + snapshot_capture_script(shell_type, CAPTURE_NON_INTERACTIVE) + } + .unwrap(); + let output = Command::new(shell) + .args(["-c", &format!("{setup}{capture}")]) + .output()?; + assert!(output.status.success(), "{shell}: {output:?}"); + let captured = CapturedSnapshot::parse(shell_type, &output.stdout).unwrap(); + let path = dir.path().join("state"); + std::fs::write(&path, captured.render_state())?; + let restore = if source { + ". \"$1\"" + } else { + "eval \"$(cat \"$1\")\"" + }; + let command = format!("{restore}\n\\alias snapshot_alias\nhelper\n"); + let output = Command::new(shell) + .args(["-c", &command, "snapshot-test"]) + .arg(&path) + .output()?; + assert!(output.status.success(), "{shell}: {output:?}"); + assert!(String::from_utf8_lossy(&output.stdout).contains("snapshot_alias=")); + results.push((output.stdout, output.stderr)); + } + assert_eq!(results[0], results[1], "{shell}"); + } + Ok(()) +}