From d0cc662b8c9239e5886acef9d6762bb33df9a058 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 20 Aug 2026 17:13:03 +0000 Subject: [PATCH] Cache shell snapshots in the exec server (#39756) ## What changed - Add the `shellSnapshotV2` executor capability and an optional shell snapshot request to `ExecParams`. - Capture and restore Unix shell state and profile exports from an in-memory, attachment-scoped cache for `bash`, `zsh`, and `sh`. - Apply environment policies, runtime `PATH` entries, sandbox context, and live managed-proxy settings when preparing restored commands. - Bound snapshot size, capture time, scope length, and cache capacity, and fall back to the original command when capture fails. ## Testing - Cover local, remote, TTY, sandboxed, and supported-shell execution, plus environment filtering, proxy handling, in-memory reuse, and capture failure fallback. GitOrigin-RevId: 624f747972c249c88c6f10f42cf0af97b75b5541 --- codex-rs/Cargo.lock | 1 + codex-rs/cli/tests/exec_server.rs | 1 + .../core/src/unified_exec/process_manager.rs | 1 + codex-rs/exec-server-protocol/src/protocol.rs | 23 ++ codex-rs/exec-server/Cargo.toml | 1 + codex-rs/exec-server/src/client.rs | 1 + .../src/client/tests/network_policy_tests.rs | 1 + codex-rs/exec-server/src/environment.rs | 2 + codex-rs/exec-server/src/lib.rs | 3 + codex-rs/exec-server/src/local_process.rs | 20 +- .../exec-server/src/process_sandbox_tests.rs | 7 + .../exec-server/src/server/handler/tests.rs | 1 + codex-rs/exec-server/src/server/processor.rs | 1 + codex-rs/exec-server/src/shell_snapshot.rs | 317 ++++++++++++++++++ .../exec-server/src/shell_snapshot_tests.rs | 87 +++++ codex-rs/exec-server/tests/exec_process.rs | 316 +++++++++++++++++ codex-rs/exec-server/tests/forward.rs | 1 + codex-rs/exec-server/tests/relay.rs | 1 + .../rmcp-client/src/stdio_server_launcher.rs | 1 + codex-rs/shell-command/src/shell_snapshot.rs | 14 + 20 files changed, 799 insertions(+), 1 deletion(-) create mode 100644 codex-rs/exec-server/src/shell_snapshot.rs create mode 100644 codex-rs/exec-server/src/shell_snapshot_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 471c342ff9..d6ff1d2061 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3014,6 +3014,7 @@ dependencies = [ "codex-otel", "codex-protocol", "codex-sandboxing", + "codex-shell-command", "codex-test-binary-support", "codex-utils-absolute-path", "codex-utils-home-dir", diff --git a/codex-rs/cli/tests/exec_server.rs b/codex-rs/cli/tests/exec_server.rs index 42bf2bc5cf..b54b3159a3 100644 --- a/codex-rs/cli/tests/exec_server.rs +++ b/codex-rs/cli/tests/exec_server.rs @@ -251,6 +251,7 @@ metrics_exporter = {{ otlp-http = {{ endpoint = "{collector_url}/v1/metrics", pr process_id: ProcessId::from("parent-lifetime-process"), argv: argv.into_iter().map(str::to_string).collect(), cwd: cwd.as_str().parse()?, + shell_snapshot: None, env_policy: Some(codex_exec_server::ExecEnvPolicy { inherit: codex_protocol::config_types::ShellEnvironmentPolicyInherit::All, ignore_default_excludes: false, diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 8edc5a5675..db28caad1f 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -221,6 +221,7 @@ fn exec_server_params_for_request( argv: request.command.clone(), cwd: request.cwd.clone(), env_policy, + shell_snapshot: None, env, tty, pipe_stdin: false, diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index 016283e290..ad557d8de6 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -121,6 +121,9 @@ pub struct EnvironmentCapabilities { /// Whether filesystem streams can use the requested platform sandbox. #[serde(default)] pub sandboxed_file_streaming: bool, + /// Whether shell state can be cached and restored entirely inside the executor. + #[serde(default)] + pub shell_snapshot_v2: bool, } /// Status returned by an initialized exec-server connection. @@ -183,6 +186,7 @@ impl EnvironmentInfo { capability_discovery_sandbox: true, environment_config_read: true, sandboxed_file_streaming: true, + shell_snapshot_v2: cfg!(unix), }, } } @@ -219,6 +223,9 @@ pub struct ExecParams { pub cwd: PathUri, #[serde(default)] pub env_policy: Option, + /// Optional request to restore executor-owned, attachment-scoped shell state. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub shell_snapshot: Option, pub env: HashMap, pub tty: bool, /// Keep non-tty stdin writable through `process/write`. @@ -254,6 +261,19 @@ pub struct ExecEnvPolicy { pub include_only: Vec, } +/// Identifies shell state owned by one attachment within an executor session. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ShellSnapshotRequest { + /// Attachment identity; executor sessions independently scope every cache. + pub scope_id: String, + /// Executor-native shell used to capture and restore the snapshot. + pub shell: ShellInfo, + /// Runtime-owned PATH entries to replay after restoring profile state. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub runtime_path_prepends: Vec, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExecResponse { @@ -849,6 +869,7 @@ mod tests { argv: vec!["true".to_string()], cwd, env_policy: None, + shell_snapshot: None, env: HashMap::new(), tty: false, pipe_stdin: false, @@ -943,6 +964,7 @@ mod tests { capability_discovery_sandbox: true, environment_config_read: false, sandboxed_file_streaming: false, + shell_snapshot_v2: false, } ); } @@ -958,6 +980,7 @@ mod tests { "capabilityDiscoverySandbox": false, "environmentConfigRead": false, "sandboxedFileStreaming": false, + "shellSnapshotV2": false, }, }); let info: EnvironmentInfo = serde_json::from_value(expected.clone()) diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 20ecc463d6..746bfb231d 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -25,6 +25,7 @@ codex-network-proxy = { workspace = true } codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-sandboxing = { workspace = true } +codex-shell-command = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-home-dir = { workspace = true } codex-utils-path-uri = { workspace = true } diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 313e082529..af0d47514c 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -1850,6 +1850,7 @@ mod tests { argv: vec!["true".to_string()], cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) .expect("cwd URI"), + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, diff --git a/codex-rs/exec-server/src/client/tests/network_policy_tests.rs b/codex-rs/exec-server/src/client/tests/network_policy_tests.rs index 5173e95bb9..468c3dfc7a 100644 --- a/codex-rs/exec-server/src/client/tests/network_policy_tests.rs +++ b/codex-rs/exec-server/src/client/tests/network_policy_tests.rs @@ -270,6 +270,7 @@ async fn abandoned_process_start_unregisters_and_cleans_up() { argv: vec!["true".to_string()], cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) .expect("cwd URI"), + shell_snapshot: None, env_policy: None, env: Default::default(), tty: false, diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 20ada8d8fb..b1e73dca3b 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -1737,6 +1737,7 @@ mod tests { std::env::current_dir().expect("read current dir"), ) .expect("cwd URI"), + shell_snapshot: None, env_policy: None, env: Default::default(), tty: false, @@ -1778,6 +1779,7 @@ mod tests { std::env::current_dir().expect("read current dir"), ) .expect("cwd URI"), + shell_snapshot: None, env_policy: None, env: Default::default(), tty: false, diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index a840bf8fef..f0e23038bb 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -37,6 +37,8 @@ mod runtime_paths; mod sandboxed_file_open; mod sandboxed_file_system; mod server; +#[cfg(unix)] +mod shell_snapshot; mod telemetry; mod trace_context; mod websocket_pong_watchdog; @@ -186,6 +188,7 @@ pub use protocol::ProcessSignal; pub use protocol::ReadParams; pub use protocol::ReadResponse; pub use protocol::ShellInfo; +pub use protocol::ShellSnapshotRequest; pub use protocol::SignalParams; pub use protocol::SignalResponse; pub use protocol::TerminateParams; diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 016ad2ae65..927276f36e 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -158,6 +158,8 @@ struct Inner { notifications: std::sync::RwLock>, requests: Arc>>, processes: Mutex>, + #[cfg(unix)] + shell_snapshots: crate::shell_snapshot::ShellSnapshotCache, telemetry: ExecServerTelemetry, } @@ -215,6 +217,8 @@ impl LocalProcess { notifications: std::sync::RwLock::new(Some(notifications)), requests: Arc::new(std::sync::RwLock::new(Some(requests))), processes: Mutex::new(HashMap::new()), + #[cfg(unix)] + shell_snapshots: crate::shell_snapshot::ShellSnapshotCache::default(), telemetry, }), runtime_paths, @@ -332,6 +336,12 @@ impl LocalProcess { let _ = notifications.try_notify(NETWORK_POLICY_DECISION_METHOD, ¬ification); }) as NetworkPolicyAuditObserver }); + #[cfg(not(unix))] + if params.shell_snapshot.is_some() { + return Err(invalid_params( + "shell snapshots are unsupported on this platform".to_string(), + )); + } let prepared = prepare_exec_request( ¶ms, child_env(¶ms), @@ -340,6 +350,13 @@ impl LocalProcess { network_policy_audit_observer, ) .await?; + #[cfg(unix)] + let mut prepared = prepared; + #[cfg(unix)] + self.inner + .shell_snapshots + .prepare(¶ms, &mut prepared) + .await?; if prepared.command.is_empty() { return Err(invalid_params("argv must not be empty".to_string())); } @@ -692,7 +709,7 @@ fn child_env(params: &ExecParams) -> HashMap { env } -fn shell_environment_policy(env_policy: &ExecEnvPolicy) -> ShellEnvironmentPolicy { +pub(crate) fn shell_environment_policy(env_policy: &ExecEnvPolicy) -> ShellEnvironmentPolicy { ShellEnvironmentPolicy { inherit: env_policy.inherit.clone(), ignore_default_excludes: env_policy.ignore_default_excludes, @@ -1139,6 +1156,7 @@ mod tests { argv: vec!["true".to_string()], cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) .expect("cwd URI"), + shell_snapshot: None, env_policy: None, env, tty: false, diff --git a/codex-rs/exec-server/src/process_sandbox_tests.rs b/codex-rs/exec-server/src/process_sandbox_tests.rs index 218b760fa8..a66387ee0d 100644 --- a/codex-rs/exec-server/src/process_sandbox_tests.rs +++ b/codex-rs/exec-server/src/process_sandbox_tests.rs @@ -56,6 +56,7 @@ async fn sandbox_request_wraps_native_argv_on_executor() { "pwd".to_string(), ], cwd: cwd_uri, + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, @@ -125,6 +126,7 @@ async fn sandbox_request_routes_custom_arg0_to_inner_helper() { process_id: ProcessId::from("process-custom-arg0"), argv: vec!["/bin/sh".to_string(), "-c".to_string(), "true".to_string()], cwd: cwd_uri, + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, @@ -186,6 +188,7 @@ async fn sandbox_request_allows_prepared_managed_proxy_port() { process_id: ProcessId::from("process-managed-network"), argv: vec!["/usr/bin/true".to_string()], cwd: cwd_uri, + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, @@ -230,6 +233,7 @@ async fn native_request_preserves_native_launch_fields() { process_id: ProcessId::from("process-1"), argv: vec!["echo".to_string(), "hello".to_string()], cwd: cwd_uri, + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, @@ -274,6 +278,7 @@ async fn native_request_handles_remote_proxy_config_for_platform() { process_id: ProcessId::from("process-remote-proxy"), argv: vec!["echo".to_string(), "hello".to_string()], cwd: PathUri::from_abs_path(&cwd), + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, @@ -356,6 +361,7 @@ async fn disabled_remote_proxy_config_is_rejected_before_exporting_ports() { process_id: ProcessId::from("process-disabled-remote-proxy"), argv: vec!["echo".to_string(), "hello".to_string()], cwd: PathUri::from_abs_path(&cwd), + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, @@ -416,6 +422,7 @@ async fn managed_network_honors_windows_sandbox_level(windows_sandbox_level: Win process_id: ProcessId::from("process-managed-network"), argv: vec!["cmd.exe".to_string(), "/c".to_string(), "exit".to_string()], cwd: cwd_uri, + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs index 3ecfd7dddc..1fb96850ed 100644 --- a/codex-rs/exec-server/src/server/handler/tests.rs +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -31,6 +31,7 @@ fn exec_params_with_argv(process_id: &str, argv: Vec) -> ExecParams { argv, cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) .expect("cwd URI"), + shell_snapshot: None, env_policy: None, env: inherited_path_env(), tty: false, diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index ae98877d76..edc8d41c46 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -575,6 +575,7 @@ mod tests { argv: sleep_then_print_argv(), cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) .expect("cwd URI"), + shell_snapshot: None, env_policy: None, env, tty: false, diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs new file mode 100644 index 0000000000..f85ca4a2c5 --- /dev/null +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -0,0 +1,317 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::process::Stdio; +use std::sync::Arc; +use std::time::Duration; + +use codex_exec_server_protocol::JSONRPCErrorError; +use codex_network_proxy::PROXY_ACTIVE_ENV_KEY; +use codex_network_proxy::strip_managed_proxy_env; +use codex_protocol::config_types::ShellEnvironmentPolicyInherit; +use codex_protocol::shell_environment; +use codex_shell_command::shell_detect::ShellType; +use codex_shell_command::shell_snapshot::snapshot_state_and_environment_script; +use codex_utils_path_uri::PathUri; +use tokio::io::AsyncReadExt; +use tokio::process::Command; +use tokio::sync::Mutex; +use tokio::sync::OnceCell; + +use crate::FileSystemSandboxContext; +use crate::local_process::shell_environment_policy; +use crate::process_sandbox::PreparedExecRequest; +use crate::protocol::ExecEnvPolicy; +use crate::protocol::ExecParams; +use crate::protocol::ShellSnapshotRequest; +use crate::rpc::internal_error; +use crate::rpc::invalid_params; + +const MAX_CACHED_SNAPSHOTS: usize = 16; +const MAX_SNAPSHOT_BYTES: usize = 512 * 1024; +const MAX_SNAPSHOT_ENV_VALUE_BYTES: usize = 60 * 1024; +const MAX_SNAPSHOT_SCOPE_BYTES: usize = 256; +const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Default)] +pub(crate) struct ShellSnapshotCache { + entries: Mutex>, +} + +struct CachedShellSnapshot { + request: ShellSnapshotRequest, + cwd: PathUri, + env_policy: Option, + sandbox: Option, + snapshot: Arc>>, +} + +struct ShellSnapshot { + state: String, + environment: HashMap, +} + +impl ShellSnapshotCache { + pub(crate) async fn prepare( + &self, + params: &ExecParams, + prepared: &mut PreparedExecRequest, + ) -> Result<(), JSONRPCErrorError> { + let Some(request) = params.shell_snapshot.as_ref() else { + return Ok(()); + }; + if request.scope_id.is_empty() || request.scope_id.len() > MAX_SNAPSHOT_SCOPE_BYTES { + return Err(invalid_params(format!( + "shell snapshot scope must be non-empty and at most {MAX_SNAPSHOT_SCOPE_BYTES} bytes" + ))); + } + + if params.argv.len() < 3 + || params.argv[0] != request.shell.path + || params.argv[1] != "-lc" + || !prepared.command.ends_with(¶ms.argv) + { + return Ok(()); + } + + let shell_type = match request.shell.name.as_str() { + "bash" => ShellType::Bash, + "zsh" => ShellType::Zsh, + "sh" => ShellType::Sh, + name => { + return Err(invalid_params(format!( + "shell snapshots are unsupported for shell `{name}`" + ))); + } + }; + + let snapshot = { + let mut entries = self.entries.lock().await; + let position = entries.iter().position(|entry| { + &entry.request == request + && entry.cwd == params.cwd + && entry.env_policy == params.env_policy + && entry.sandbox == params.sandbox + }); + let cached = position.and_then(|position| { + let entry = entries.remove(position)?; + let snapshot = Arc::clone(&entry.snapshot); + entries.push_back(entry); + Some(snapshot) + }); + if let Some(snapshot) = cached { + snapshot + } else { + let snapshot = Arc::new(OnceCell::new()); + let entry = CachedShellSnapshot { + request: request.clone(), + cwd: params.cwd.clone(), + env_policy: params.env_policy.clone(), + sandbox: params.sandbox.clone(), + snapshot: Arc::clone(&snapshot), + }; + entries.push_back(entry); + if entries.len() > MAX_CACHED_SNAPSHOTS { + entries.pop_front(); + } + + snapshot + } + }; + let Some(snapshot) = snapshot + .get_or_init(|| async { + match capture_snapshot(params, prepared, shell_type).await { + Ok(snapshot) => Some(snapshot), + Err(err) => { + tracing::warn!("failed to capture shell snapshot: {err:?}"); + None + } + } + }) + .await + else { + return Ok(()); + }; + + let restore_snapshot_path = snapshot.environment.contains_key("PATH") + && params + .env_policy + .as_ref() + .is_none_or(|policy| !policy.r#set.contains_key("PATH")); + let request_overrides = params + .env + .iter() + .filter(|(name, _)| name.as_str() != "PATH" || !restore_snapshot_path) + .map(|(name, value)| { + ( + name.clone(), + prepared.env.get(name).unwrap_or(value).clone(), + ) + }) + .collect::>(); + prepared.env.extend( + snapshot + .environment + .iter() + .map(|(name, value)| (name.clone(), value.clone())), + ); + prepared.env.extend(request_overrides); + if restore_snapshot_path && let Some(path) = prepared.env.get_mut("PATH") { + for entry in &request.runtime_path_prepends { + if entry.is_empty() { + continue; + } + *path = std::iter::once(entry.as_str()) + .chain( + path.split(':') + .filter(|existing| !existing.is_empty() && *existing != entry), + ) + .collect::>() + .join(":"); + } + } + prepared + .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 (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 shell_start = prepared.command.len() - params.argv.len(); + prepared.command[shell_start + 1] = "-c".to_string(); + prepared.command[shell_start + 2] = format!( + "if ! eval \"unset {state_variables}\n{state_expansion}\" >/dev/null; then printf 'failed to restore shell snapshot\\n' >&2; fi\n{}", + params.argv[2] + ); + + Ok(()) + } +} + +async fn capture_snapshot( + params: &ExecParams, + prepared: &PreparedExecRequest, + shell_type: ShellType, +) -> Result { + let script = snapshot_state_and_environment_script(shell_type) + .ok_or_else(|| invalid_params("unsupported shell snapshot script".to_string()))?; + let shell_start = prepared.command.len() - params.argv.len(); + let mut argv = prepared.command.clone(); + argv[shell_start + 2] = script; + let (program, args) = argv + .split_first() + .ok_or_else(|| internal_error("missing shell snapshot command".to_string()))?; + + let mut command = Command::new(program); + command + .args(args) + .current_dir(prepared.cwd.as_path()) + .env_clear() + .envs(&prepared.env) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .kill_on_drop(true); + if let Some(arg0) = &prepared.arg0 { + command.arg0(arg0); + } + let mut child = command + .spawn() + .map_err(|err| internal_error(format!("cannot capture shell snapshot: {err}")))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| internal_error("missing shell snapshot output".to_string()))?; + let capture = async { + let mut output = Vec::new(); + stdout + .take((MAX_SNAPSHOT_BYTES + 1) as u64) + .read_to_end(&mut output) + .await + .map_err(|err| internal_error(format!("cannot read shell snapshot: {err}")))?; + if output.len() > MAX_SNAPSHOT_BYTES { + return Err(internal_error(format!( + "shell snapshot exceeds {MAX_SNAPSHOT_BYTES} bytes" + ))); + } + let status = child + .wait() + .await + .map_err(|err| internal_error(format!("cannot finish shell snapshot: {err}")))?; + if !status.success() { + return Err(internal_error(format!( + "shell snapshot capture exited with {status}" + ))); + } + Ok(output) + }; + let output = tokio::time::timeout(SNAPSHOT_TIMEOUT, capture) + .await + .map_err(|_| internal_error("shell snapshot capture timed out".to_string()))??; + + parse_snapshot(&output, params.env_policy.as_ref()) +} + +fn parse_snapshot( + output: &[u8], + env_policy: Option<&ExecEnvPolicy>, +) -> Result { + let separator = output + .iter() + .position(|byte| *byte == 0) + .ok_or_else(|| internal_error("shell snapshot is missing its environment".to_string()))?; + let state = &output[..separator]; + let marker = b"# Snapshot file"; + let start = state + .windows(marker.len()) + .position(|window| window == marker) + .ok_or_else(|| internal_error("shell snapshot is missing its state marker".to_string()))?; + let state = std::str::from_utf8(&state[start..]) + .map_err(|err| internal_error(format!("shell snapshot state is not UTF-8: {err}")))?; + + let mut environment = output[separator + 1..] + .split(|byte| *byte == 0) + .filter(|entry| !entry.is_empty()) + .filter_map(|entry| { + let (name, value) = std::str::from_utf8(entry).ok()?.split_once('=')?; + Some((name.to_string(), value.to_string())) + }) + .collect::>(); + if environment.contains_key(PROXY_ACTIVE_ENV_KEY) { + strip_managed_proxy_env(&mut environment); + } + let mut environment = match env_policy { + Some(policy) => { + let mut policy = shell_environment_policy(policy); + policy.inherit = ShellEnvironmentPolicyInherit::All; + shell_environment::create_env_from_vars(environment, &policy, /*thread_id*/ None) + } + None => environment, + }; + environment.remove("PWD"); + environment.remove("OLDPWD"); + environment.retain(|name, _| !shell_environment::is_non_inheritable_env_var(name)); + + Ok(ShellSnapshot { + state: state.to_string(), + environment, + }) +} + +#[cfg(test)] +#[path = "shell_snapshot_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs new file mode 100644 index 0000000000..8e6aeda771 --- /dev/null +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -0,0 +1,87 @@ +use std::collections::HashMap; + +use codex_protocol::config_types::ShellEnvironmentPolicyInherit; +use pretty_assertions::assert_eq; + +use super::parse_snapshot; +use crate::protocol::ExecEnvPolicy; + +#[test] +fn snapshot_filters_profile_exports_after_capture() { + let policy = ExecEnvPolicy { + inherit: ShellEnvironmentPolicyInherit::All, + ignore_default_excludes: false, + exclude: vec!["PROFILE_DENIED".to_string()], + r#set: HashMap::from([("PROFILE_ALLOWED".to_string(), "override".to_string())]), + include_only: vec!["PROFILE_*".to_string()], + }; + let snapshot = parse_snapshot( + b"profile noise\n# Snapshot file\nfunction profile_helper() { :; }\n\0PROFILE_ALLOWED=profile\0PROFILE_DENIED=denied\0PROFILE_SECRET=secret\0PWD=/tmp\0", + Some(&policy), + ) + .expect("snapshot should parse"); + + assert_eq!( + snapshot.environment, + HashMap::from([("PROFILE_ALLOWED".to_string(), "override".to_string())]) + ); + assert_eq!( + snapshot.state, + "# Snapshot file\nfunction profile_helper() { :; }\n" + ); +} + +#[test] +fn snapshot_preserves_profile_exports_with_restrictive_inheritance() { + for inherit in [ + ShellEnvironmentPolicyInherit::None, + ShellEnvironmentPolicyInherit::Core, + ] { + let policy = ExecEnvPolicy { + inherit, + ignore_default_excludes: false, + exclude: vec!["PROFILE_DENIED".to_string()], + r#set: HashMap::new(), + include_only: Vec::new(), + }; + let snapshot = parse_snapshot( + b"# Snapshot file\n\0PROFILE_ALLOWED=profile\0SDKROOT=/sdk\0PROFILE_SECRET=secret\0PROFILE_DENIED=denied\0", + Some(&policy), + ) + .expect("snapshot should parse"); + + assert_eq!( + snapshot.environment, + HashMap::from([ + ("PROFILE_ALLOWED".to_string(), "profile".to_string()), + ("SDKROOT".to_string(), "/sdk".to_string()), + ]) + ); + } +} + +#[test] +fn snapshot_caches_only_unmanaged_proxy_state() { + for (exports, expected) in [ + ( + "PROFILE_ALLOWED=profile\0HTTP_PROXY=http://127.0.0.1:4321\0CODEX_NETWORK_PROXY_ACTIVE=1\0CODEX_NETWORK_PROXY_CREDENTIAL_BROKER_ACTIVE=1\0", + HashMap::from([("PROFILE_ALLOWED".to_string(), "profile".to_string())]), + ), + ( + "PROFILE_ALLOWED=profile\0HTTP_PROXY=http://user-proxy.example\0", + HashMap::from([ + ("PROFILE_ALLOWED".to_string(), "profile".to_string()), + ( + "HTTP_PROXY".to_string(), + "http://user-proxy.example".to_string(), + ), + ]), + ), + ] { + let output = format!("# Snapshot file\n\0{exports}"); + let snapshot = + parse_snapshot(output.as_bytes(), /*env_policy*/ None).expect("snapshot should parse"); + + assert_eq!(snapshot.environment, expected); + } +} diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index 743ffdd354..4f1126a097 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -9,6 +9,8 @@ use anyhow::Context; use anyhow::Result; use codex_exec_server::Environment; use codex_exec_server::ExecBackend; +#[cfg(unix)] +use codex_exec_server::ExecEnvPolicy; use codex_exec_server::ExecOutputStream; use codex_exec_server::ExecParams; use codex_exec_server::ExecProcess; @@ -18,8 +20,20 @@ use codex_exec_server::FileSystemSandboxContext; use codex_exec_server::ProcessId; use codex_exec_server::ProcessSignal; use codex_exec_server::ReadResponse; +#[cfg(unix)] +use codex_exec_server::ShellInfo; +#[cfg(unix)] +use codex_exec_server::ShellSnapshotRequest; use codex_exec_server::StartedExecProcess; use codex_exec_server::WriteStatus; +#[cfg(unix)] +use codex_network_proxy::NetworkProxyConfig; +#[cfg(unix)] +use codex_network_proxy::RemoteNetworkProxyConfig; +#[cfg(unix)] +use codex_network_proxy::RemoteNetworkProxyLaunchConfig; +#[cfg(unix)] +use codex_protocol::config_types::ShellEnvironmentPolicyInherit; use codex_protocol::config_types::WindowsSandboxLevel; #[cfg(unix)] use codex_protocol::models::PermissionProfile; @@ -88,6 +102,291 @@ async fn create_process_context(use_remote: bool) -> Result { } } +#[cfg(unix)] +#[test_case(false, false, false, "bash"; "local_pipe")] +#[test_case(false, true, false, "bash"; "local_tty")] +#[test_case(true, false, false, "bash"; "remote_pipe")] +#[test_case(true, true, false, "bash"; "remote_tty")] +#[test_case(true, false, true, "bash"; "remote_sandbox")] +#[test_case(false, false, false, "sh"; "local_sh_pipe")] +#[cfg_attr( + target_os = "macos", + test_case(false, false, false, "zsh"; "local_zsh_pipe") +)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +// Serialize tests that launch a real exec-server process through the full CLI. +#[serial_test::serial(remote_exec_server)] +async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( + use_remote: bool, + tty: bool, + use_sandbox: bool, + shell_name: &str, +) -> Result<()> { + if use_sandbox + && let Some(warning) = + codex_sandboxing::system_bwrap_warning(&PermissionProfile::read_only()) + { + eprintln!("skipping sandbox test: {warning}"); + return Ok(()); + } + let context = create_process_context(use_remote).await?; + let home = TempDir::new()?; + let cwd = PathUri::from_host_native_path(home.path())?; + let (shell_path, profile_name) = match shell_name { + "bash" => ("/bin/bash", ".bashrc"), + "sh" => ("/bin/sh", ".snapshot-env"), + "zsh" => ("/bin/zsh", ".zshrc"), + name => anyhow::bail!("unsupported test shell {name}"), + }; + let profile_path = home.path().join(profile_name); + let profile_path_entry = home.path().join("profile-bin"); + let runtime_path_entry = home.path().join("runtime-bin"); + let inherited_path = std::env::var("PATH")?; + let runtime_path = format!("{}:{inherited_path}", runtime_path_entry.display()); + let padding = if !use_remote && !tty && shell_name == "bash" { + format!( + "snapshot_padding() {{ printf '%s' '{}'; }}\n", + "🦀".repeat(20_000) + ) + } else { + String::new() + }; + let shadowed_builtins = if shell_name == "sh" { + "" + } else { + "unset() { exit 41; }\nbuiltin() { :; }\n" + }; + std::fs::write( + &profile_path, + format!( + "printf x >> \"$HOME/captures\"\nexport PATH=\"$HOME/profile-bin:/usr/bin:/bin\"\nexport PROFILE_ALLOWED=profile\nexport PROFILE_SECRET=secret\nexport PROFILE_DENIED=denied\nprofile_helper() {{ printf helper; }}\n{shadowed_builtins}{padding}" + ), + )?; + let mut configured_environment = HashMap::from([( + "HOME".to_string(), + home.path().to_string_lossy().into_owned(), + )]); + if shell_name == "sh" { + configured_environment.insert( + "ENV".to_string(), + profile_path.to_string_lossy().into_owned(), + ); + } + let policy = ExecEnvPolicy { + inherit: ShellEnvironmentPolicyInherit::All, + ignore_default_excludes: false, + exclude: vec!["PROFILE_DENIED".to_string()], + r#set: configured_environment, + include_only: vec![ + "ENV".to_string(), + "HOME".to_string(), + "PATH".to_string(), + "PROFILE_*".to_string(), + ], + }; + let (command_prefix, expected_prefix) = if shell_name == "sh" { + ("", "") + } else { + ("profile_helper; ", "helper") + }; + let command = format!( + "{command_prefix}printf '|%s|%s|%s|%s|%s|%s' \"$PROFILE_ALLOWED\" \"${{PROFILE_SECRET-missing}}\" \"${{PROFILE_DENIED-missing}}\" \"$PATH\" \"${{__CODEX_SHELL_SNAPSHOT_STATE_0-missing}}\" \"${{__CODEX_SHELL_SNAPSHOT_STATE_1-missing}}\"" + ); + let expected_stdout = format!( + "{expected_prefix}|profile|missing|missing|{}:{}:/usr/bin:/bin|missing|missing", + runtime_path_entry.display(), + profile_path_entry.display(), + ); + + for attempt in 0..2 { + let started = context + .backend + .start(ExecParams { + process_id: ProcessId::from(format!("snapshot-{attempt}")), + argv: vec![shell_path.to_string(), "-lc".to_string(), command.clone()], + cwd: cwd.clone(), + env_policy: Some(policy.clone()), + shell_snapshot: Some(ShellSnapshotRequest { + scope_id: "attachment-1".to_string(), + shell: ShellInfo { + name: shell_name.to_string(), + path: shell_path.to_string(), + }, + runtime_path_prepends: vec![runtime_path_entry.to_string_lossy().into_owned()], + }), + env: HashMap::from([("PATH".to_string(), runtime_path.clone())]), + tty, + pipe_stdin: false, + arg0: None, + sandbox: (use_sandbox && attempt == 0).then(|| { + FileSystemSandboxContext::from_permission_profile_with_cwd( + PermissionProfile::read_only(), + cwd.clone(), + ) + }), + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }) + .await?; + let (stdout, stderr, status, closed) = + collect_process_output_from_events(started.process).await?; + assert_eq!( + (stdout, stderr, status, closed), + (expected_stdout.clone(), String::new(), Some(0), true,) + ); + } + + assert_eq!(std::fs::read_to_string(home.path().join("captures"))?, "x"); + if let Some(server) = context._server { + assert!(!server.codex_home().join("shell_snapshots").exists()); + } + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(remote_exec_server)] +async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context() -> Result<()> { + let context = create_process_context(/*use_remote*/ true).await?; + let home = TempDir::new()?; + let cwd = PathUri::from_host_native_path(home.path())?; + std::fs::write( + home.path().join(".bashrc"), + "printf '%s\\n' \"$HTTP_PROXY\" >> \"$HOME/captures\"\ntest \"$CODEX_NETWORK_PROXY_ACTIVE\" = 1 || exit 41\nexport PROFILE_ALLOWED=profile\nprofile_helper() { printf helper; }\n", + )?; + let policy = ExecEnvPolicy { + inherit: ShellEnvironmentPolicyInherit::All, + ignore_default_excludes: false, + exclude: Vec::new(), + r#set: HashMap::from([( + "HOME".to_string(), + home.path().to_string_lossy().into_owned(), + )]), + include_only: vec![ + "HOME".to_string(), + "PATH".to_string(), + "PROFILE_*".to_string(), + ], + }; + let proxy_config = RemoteNetworkProxyConfig::from_effective_config(&NetworkProxyConfig { + enabled: true, + ..NetworkProxyConfig::default() + })?; + let mut proxy_addresses = Vec::new(); + + for attempt in 0..2 { + let started = context + .backend + .start(ExecParams { + process_id: ProcessId::from(format!("managed-snapshot-{attempt}")), + argv: vec![ + "/bin/bash".to_string(), + "-lc".to_string(), + "profile_helper; printf '|%s|%s|%s' \"$PROFILE_ALLOWED\" \"$CODEX_NETWORK_PROXY_ACTIVE\" \"$HTTP_PROXY\"".to_string(), + ], + cwd: cwd.clone(), + env_policy: Some(policy.clone()), + shell_snapshot: Some(ShellSnapshotRequest { + scope_id: "managed-attachment".to_string(), + shell: ShellInfo { + name: "bash".to_string(), + path: "/bin/bash".to_string(), + }, + runtime_path_prepends: Vec::new(), + }), + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: true, + managed_network: None, + network_proxy: Some( + RemoteNetworkProxyLaunchConfig::new(proxy_config.clone()).for_execution( + "remote-environment".to_string(), + format!("managed-snapshot-{attempt}"), + ), + ), + }) + .await?; + let (stdout, stderr, status, closed) = + collect_process_output_from_events(started.process).await?; + let proxy_address = stdout + .strip_prefix("helper|profile|1|") + .context("snapshot should restore profile functions and live proxy state")?; + assert!(proxy_address.starts_with("http://127.0.0.1:")); + assert_eq!((stderr, status, closed), (String::new(), Some(0), true)); + proxy_addresses.push(proxy_address.to_string()); + } + + assert_eq!( + std::fs::read_to_string(home.path().join("captures"))?, + format!("{}\n", proxy_addresses[0]) + ); + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_v2_capture_failure_falls_back_to_original_command() -> Result<()> { + let context = create_process_context(/*use_remote*/ false).await?; + let home = TempDir::new()?; + let cwd = PathUri::from_host_native_path(home.path())?; + std::fs::write( + home.path().join(".bashrc"), + "printf x >> \"$HOME/captures\"\nexit 7\n", + )?; + let policy = ExecEnvPolicy { + inherit: ShellEnvironmentPolicyInherit::All, + ignore_default_excludes: false, + exclude: Vec::new(), + r#set: HashMap::from([( + "HOME".to_string(), + home.path().to_string_lossy().into_owned(), + )]), + include_only: vec!["HOME".to_string(), "PATH".to_string()], + }; + let mut params = ExecParams { + process_id: ProcessId::from("snapshot-first"), + argv: vec![ + "/bin/bash".to_string(), + "-lc".to_string(), + "printf original".to_string(), + ], + cwd, + env_policy: Some(policy), + shell_snapshot: Some(ShellSnapshotRequest { + scope_id: "attachment-1".to_string(), + shell: ShellInfo { + name: "bash".to_string(), + path: "/bin/bash".to_string(), + }, + runtime_path_prepends: Vec::new(), + }), + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }; + + for attempt in 0..2 { + params.process_id = ProcessId::from(format!("snapshot-fallback-{attempt}")); + let fallback = context.backend.start(params.clone()).await?; + let fallback_output = collect_process_output_from_events(fallback.process).await?; + assert_eq!( + fallback_output, + ("original".to_string(), String::new(), Some(0), true) + ); + } + assert_eq!(std::fs::read_to_string(home.path().join("captures"))?, "x"); + Ok(()) +} + #[cfg(unix)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_sandboxed_process_preserves_custom_arg0() -> Result<()> { @@ -133,6 +432,7 @@ async fn remote_sandboxed_process_preserves_custom_arg0() -> Result<()> { .to_string(), ], cwd, + shell_snapshot: None, env_policy: None, env: HashMap::from([ ("PATH".to_string(), std::env::var("PATH")?), @@ -167,6 +467,7 @@ async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> { process_id: ProcessId::from("proc-1"), argv: vec!["true".to_string()], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -228,6 +529,7 @@ async fn remote_process_keeps_sandbox_helper_visible_with_restricted_reads() -> process_id: ProcessId::from("proc-restricted-helper"), argv: vec!["/bin/cat".to_string(), file.to_string_lossy().into_owned()], cwd, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: HashMap::from([("PATH".to_string(), std::env::var("PATH")?)]), tty: false, @@ -297,6 +599,7 @@ async fn remote_tty_process_uses_configured_sandbox_helper_with_hostile_path() - process_id: ProcessId::from("proc-hostile-helper-path"), argv: vec!["/bin/cat".to_string(), file.to_string_lossy().into_owned()], cwd, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: HashMap::from([( "PATH".to_string(), @@ -352,6 +655,7 @@ async fn remote_process_preserves_empty_workspace_roots() -> Result<()> { process_id: ProcessId::from("proc-empty-workspace-roots"), argv: vec!["/bin/cat".to_string(), file.to_string_lossy().into_owned()], cwd, + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, @@ -497,6 +801,7 @@ async fn assert_exec_process_streams_output(use_remote: bool) -> Result<()> { "sleep 0.05; printf 'session output\\n'".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -532,6 +837,7 @@ async fn assert_exec_process_pushes_events(use_remote: bool) -> Result<()> { "printf 'event output\\n'; sleep 0.1; printf 'event err\\n' >&2; sleep 0.1; exit 7".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -583,6 +889,7 @@ async fn assert_exec_process_replays_events_after_close(use_remote: bool) -> Res "printf 'late one\\n'; printf 'late two\\n'".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -635,6 +942,7 @@ async fn assert_exec_process_retains_output_after_exit_until_streams_close( release_path.to_string_lossy().into_owned(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -712,6 +1020,7 @@ async fn assert_exec_process_write_then_read(use_remote: bool) -> Result<()> { "IFS= read line; printf 'from-stdin:%s\\n' \"$line\"".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: true, @@ -753,6 +1062,7 @@ async fn assert_exec_process_write_then_read_without_tty(use_remote: bool) -> Re "IFS= read line; printf 'from-stdin:%s\\n' \"$line\"".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -804,6 +1114,7 @@ async fn assert_remote_windows_sandbox_process_write() -> Result<()> { ), ], cwd, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -849,6 +1160,7 @@ async fn assert_exec_process_rejects_write_without_pipe_stdin(use_remote: bool) "sleep 0.3; if IFS= read -r line; then printf 'read:%s\\n' \"$line\"; else printf 'eof\\n'; fi".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -887,6 +1199,7 @@ async fn assert_exec_process_signal_interrupts_process(use_remote: bool) -> Resu "trap 'printf \"signal:2\\n\"; exit 7' INT; printf 'ready\\n'; while :; do :; done".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -944,6 +1257,7 @@ async fn assert_exec_process_signal_terminates_on_windows(use_remote: bool) -> R "echo ready && ping -n 30 127.0.0.1 >NUL".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -980,6 +1294,7 @@ async fn assert_exec_process_preserves_queued_events_before_subscribe( "printf 'queued output\\n'".to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: Default::default(), tty: false, @@ -1033,6 +1348,7 @@ async fn remote_exec_process_recovers_after_transport_disconnect() -> Result<()> .to_string(), ], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: /*env_policy*/ None, env: HashMap::from([ ( diff --git a/codex-rs/exec-server/tests/forward.rs b/codex-rs/exec-server/tests/forward.rs index 33f7bec91e..02c272bdae 100644 --- a/codex-rs/exec-server/tests/forward.rs +++ b/codex-rs/exec-server/tests/forward.rs @@ -65,6 +65,7 @@ async fn forwarder_runs_commands_and_transfers_files() -> Result<()> { process_id: process_id.clone(), argv, cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 119fe9f2cd..e54dd2e9f0 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -329,6 +329,7 @@ async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> { process_id: ProcessId::from("proc-1"), argv: vec!["true".to_string()], cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, env_policy: None, env: HashMap::new(), tty: false, diff --git a/codex-rs/rmcp-client/src/stdio_server_launcher.rs b/codex-rs/rmcp-client/src/stdio_server_launcher.rs index 9eecfadce7..1f6ab22be0 100644 --- a/codex-rs/rmcp-client/src/stdio_server_launcher.rs +++ b/codex-rs/rmcp-client/src/stdio_server_launcher.rs @@ -605,6 +605,7 @@ impl ExecutorStdioServerLauncher { process_id, argv, cwd, + shell_snapshot: None, env_policy: Some(Self::remote_env_policy(&remote_env_vars)), env, tty: false, diff --git a/codex-rs/shell-command/src/shell_snapshot.rs b/codex-rs/shell-command/src/shell_snapshot.rs index 62e4af7837..d0cf01c3b3 100644 --- a/codex-rs/shell-command/src/shell_snapshot.rs +++ b/codex-rs/shell-command/src/shell_snapshot.rs @@ -5,6 +5,7 @@ use crate::shell_detect::ShellType; mod tests; const EXCLUDED_EXPORT_VARS: &[&str] = &["PWD", "OLDPWD"]; +const EXPORT_CAPTURE_MARKER: &str = "# Capture exported variables"; /// Returns the shell-native script used to capture restorable shell state. /// @@ -20,6 +21,16 @@ pub fn snapshot_script(shell_type: ShellType) -> Option { } } +/// Captures shell state and a separate NUL-delimited exported environment. +/// +/// Keeping exports outside the restorable script lets executors apply their +/// environment policy after shell profiles have run. +pub fn snapshot_state_and_environment_script(shell_type: ShellType) -> Option { + let script = snapshot_script(shell_type)?; + let (state, _) = script.split_once(EXPORT_CAPTURE_MARKER)?; + Some(format!("{state}printf '\\0'\n/usr/bin/env -0\n")) +} + fn excluded_exports_regex() -> String { EXCLUDED_EXPORT_VARS.join("|") } @@ -46,6 +57,7 @@ alias_count=$(alias -L | wc -l | tr -d ' ') print "# aliases $alias_count" alias -L print '' +# Capture exported variables export_lines=$(export -p | awk ' /^(export|declare -x|typeset -x) / { line=$0 @@ -97,6 +109,7 @@ alias_count=$(alias -p | wc -l | tr -d ' ') echo "# aliases $alias_count" alias -p echo '' +# Capture exported variables export_lines=$( while IFS= read -r name; do if [[ "$name" =~ ^(EXCLUDED_EXPORTS)$ ]]; then @@ -151,6 +164,7 @@ if alias >/dev/null 2>&1; then else echo '# aliases 0' fi +# Capture exported variables if export -p >/dev/null 2>&1; then export_lines=$(export -p | awk ' /^(export|declare -x|typeset -x) / {