From 28dff23b2f794bc80bc9ed8243366aa40f741c03 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 26 May 2026 11:26:33 -0700 Subject: [PATCH] Refine env file snapshot restore --- codex-rs/core/src/tasks/user_shell.rs | 1 + codex-rs/core/src/tools/handlers/shell.rs | 14 ++--- .../src/tools/handlers/shell/shell_command.rs | 48 ++++++--------- .../core/src/tools/handlers/shell_tests.rs | 11 +--- codex-rs/core/src/tools/runtimes/mod.rs | 26 ++++++--- codex-rs/core/src/tools/runtimes/mod_tests.rs | 58 +++++++++++++++++++ codex-rs/core/src/tools/runtimes/shell.rs | 2 + .../core/src/tools/runtimes/unified_exec.rs | 3 + .../core/src/unified_exec/process_manager.rs | 8 ++- 9 files changed, 112 insertions(+), 59 deletions(-) diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 396aecbeea..0e0f831d00 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -152,6 +152,7 @@ pub(crate) async fn execute_user_shell_command( #[allow(deprecated)] &turn_context.cwd, &turn_context.shell_environment_policy.r#set, + &[], &exec_env_map, ); diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 1a49817c82..ee02b36ce3 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -53,13 +53,12 @@ struct RunExecLikeArgs { tracker: crate::tools::context::SharedTurnDiffTracker, call_id: String, shell_runtime_backend: ShellRuntimeBackend, - snapshot_restore_env_keys: Vec, } async fn run_exec_like(args: RunExecLikeArgs) -> Result { let RunExecLikeArgs { tool_name, - exec_params, + mut exec_params, hook_command, shell_type, additional_permissions, @@ -69,7 +68,6 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result Result Result Result, -} - #[derive(Clone, Copy)] pub(crate) struct ShellCommandHandlerOptions { pub(crate) backend_config: ShellCommandBackendConfig, @@ -93,34 +88,28 @@ impl ShellCommandHandler { turn_context: &TurnContext, thread_id: ThreadId, allow_login_shell: bool, - ) -> Result { + ) -> Result { let shell = session.user_shell(); let use_login_shell = Self::resolve_use_login_shell(params.login, allow_login_shell)?; let command = Self::base_command(shell.as_ref(), ¶ms.command, use_login_shell); #[allow(deprecated)] let cwd = turn_context.resolve_path(params.workdir.clone()); - let mut env = create_env(&turn_context.shell_environment_policy, Some(thread_id)); - let snapshot_restore_env_keys = session.apply_hook_env_file(&mut env); - - Ok(ShellCommandExecParams { - exec_params: ExecParams { - command, - cwd, - expiration: params.timeout_ms.into(), - capture_policy: ExecCapturePolicy::ShellTool, - env, - network: turn_context.network.clone(), - sandbox_permissions: params.sandbox_permissions.unwrap_or_default(), - windows_sandbox_level: turn_context.windows_sandbox_level, - windows_sandbox_private_desktop: turn_context - .config - .permissions - .windows_sandbox_private_desktop, - justification: params.justification.clone(), - arg0: None, - }, - snapshot_restore_env_keys, + Ok(ExecParams { + command, + cwd, + expiration: params.timeout_ms.into(), + capture_policy: ExecCapturePolicy::ShellTool, + env: create_env(&turn_context.shell_environment_policy, Some(thread_id)), + network: turn_context.network.clone(), + sandbox_permissions: params.sandbox_permissions.unwrap_or_default(), + windows_sandbox_level: turn_context.windows_sandbox_level, + windows_sandbox_private_desktop: turn_context + .config + .permissions + .windows_sandbox_private_desktop, + justification: params.justification.clone(), + arg0: None, }) } } @@ -185,7 +174,7 @@ impl ToolExecutor for ShellCommandHandler { ) .await; let prefix_rule = params.prefix_rule.clone(); - let shell_exec_params = Self::to_exec_params( + let exec_params = Self::to_exec_params( ¶ms, session.as_ref(), turn.as_ref(), @@ -195,7 +184,7 @@ impl ToolExecutor for ShellCommandHandler { let shell_type = Some(session.user_shell().shell_type.clone()); run_exec_like(RunExecLikeArgs { tool_name, - exec_params: shell_exec_params.exec_params, + exec_params, hook_command: params.command, shell_type, additional_permissions: params.additional_permissions.clone(), @@ -205,7 +194,6 @@ impl ToolExecutor for ShellCommandHandler { tracker, call_id, shell_runtime_backend: self.shell_runtime_backend(), - snapshot_restore_env_keys: shell_exec_params.snapshot_restore_env_keys, }) .await .map(boxed_tool_output) diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index 3c6f25f81b..eda23533a3 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -106,7 +106,7 @@ async fn shell_command_handler_to_exec_params_uses_session_shell_and_turn_contex justification: justification.clone(), }; - let shell_exec_params = ShellCommandHandler::to_exec_params( + let exec_params = ShellCommandHandler::to_exec_params( ¶ms, &session, &turn_context, @@ -114,7 +114,6 @@ async fn shell_command_handler_to_exec_params_uses_session_shell_and_turn_contex /*allow_login_shell*/ true, ) .expect("login shells should be allowed"); - let exec_params = shell_exec_params.exec_params; // ExecParams cannot derive Eq due to the CancellationToken field, so we manually compare the fields. assert_eq!(exec_params.command, expected_command); @@ -125,10 +124,6 @@ async fn shell_command_handler_to_exec_params_uses_session_shell_and_turn_contex assert_eq!(exec_params.sandbox_permissions, sandbox_permissions); assert_eq!(exec_params.justification, justification); assert_eq!(exec_params.arg0, None); - assert_eq!( - shell_exec_params.snapshot_restore_env_keys, - Vec::::new() - ); } #[test] @@ -178,7 +173,7 @@ async fn shell_command_handler_defaults_to_non_login_when_disallowed() { justification: None, }; - let shell_exec_params = ShellCommandHandler::to_exec_params( + let exec_params = ShellCommandHandler::to_exec_params( ¶ms, &session, &turn_context, @@ -188,7 +183,7 @@ async fn shell_command_handler_defaults_to_non_login_when_disallowed() { .expect("non-login shells should still be allowed"); assert_eq!( - shell_exec_params.exec_params.command, + exec_params.command, session .user_shell() .derive_exec_args("echo hello", /*use_login_shell*/ false) diff --git a/codex-rs/core/src/tools/runtimes/mod.rs b/codex-rs/core/src/tools/runtimes/mod.rs index bba1c572ef..41a77d64a2 100644 --- a/codex-rs/core/src/tools/runtimes/mod.rs +++ b/codex-rs/core/src/tools/runtimes/mod.rs @@ -116,13 +116,17 @@ pub(crate) fn disable_powershell_profile_for_elevated_windows_sandbox( /// `explicit_env_overrides` contains policy-driven shell env overrides that /// should win after the snapshot is sourced, while `env` is the full live exec /// environment. We need access to both so snapshot restore logic can preserve -/// runtime-only vars like `CODEX_THREAD_ID` without pretending they came from -/// the explicit override policy. +/// runtime-only vars like `CODEX_THREAD_ID`. +/// +/// `snapshot_restore_env_keys` covers additional live env values, such as +/// `CODEX_ENV_FILE` updates, that should survive snapshot sourcing without +/// pretending they came from the explicit override policy. pub(crate) fn maybe_wrap_shell_lc_with_snapshot( command: &[String], session_shell: &Shell, cwd: &AbsolutePathBuf, explicit_env_overrides: &HashMap, + snapshot_restore_env_keys: &[String], env: &HashMap, ) -> Vec { if cfg!(windows) { @@ -159,11 +163,8 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot( .iter() .map(|arg| format!(" '{}'", shell_single_quote(arg))) .collect::(); - let mut override_env = explicit_env_overrides.clone(); - if let Some(thread_id) = env.get(CODEX_THREAD_ID_ENV_VAR) { - override_env.insert(CODEX_THREAD_ID_ENV_VAR.to_string(), thread_id.clone()); - } - let (override_captures, override_exports) = build_override_exports(&override_env); + let (override_captures, override_exports) = + build_override_exports(explicit_env_overrides, snapshot_restore_env_keys, env); let (proxy_captures, proxy_exports) = build_proxy_env_exports(); let override_captures = join_shell_blocks([override_captures, proxy_captures]); let override_exports = join_shell_blocks([override_exports, proxy_exports]); @@ -180,13 +181,22 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot( vec![shell_path.to_string(), "-c".to_string(), rewritten_script] } -fn build_override_exports(explicit_env_overrides: &HashMap) -> (String, String) { +fn build_override_exports( + explicit_env_overrides: &HashMap, + snapshot_restore_env_keys: &[String], + env: &HashMap, +) -> (String, String) { let mut keys = explicit_env_overrides .keys() .map(String::as_str) + .chain(snapshot_restore_env_keys.iter().map(String::as_str)) .filter(|key| is_valid_shell_variable_name(key)) .collect::>(); + if env.contains_key(CODEX_THREAD_ID_ENV_VAR) { + keys.push(CODEX_THREAD_ID_ENV_VAR); + } keys.sort_unstable(); + keys.dedup(); build_override_exports_for_keys("__CODEX_SNAPSHOT_OVERRIDE", &keys) } diff --git a/codex-rs/core/src/tools/runtimes/mod_tests.rs b/codex-rs/core/src/tools/runtimes/mod_tests.rs index fd10f22242..4f94b971ad 100644 --- a/codex-rs/core/src/tools/runtimes/mod_tests.rs +++ b/codex-rs/core/src/tools/runtimes/mod_tests.rs @@ -184,6 +184,7 @@ fn maybe_wrap_shell_lc_with_snapshot_bootstraps_in_user_shell() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); @@ -215,6 +216,7 @@ fn maybe_wrap_shell_lc_with_snapshot_escapes_single_quotes() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); @@ -243,6 +245,7 @@ fn maybe_wrap_shell_lc_with_snapshot_uses_bash_bootstrap_shell() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); @@ -274,6 +277,7 @@ fn maybe_wrap_shell_lc_with_snapshot_uses_sh_bootstrap_shell() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); @@ -307,6 +311,7 @@ fn maybe_wrap_shell_lc_with_snapshot_preserves_trailing_args() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); @@ -342,6 +347,7 @@ fn maybe_wrap_shell_lc_with_snapshot_skips_when_cwd_mismatch() { &session_shell, &command_cwd.abs(), &HashMap::new(), + &[], &HashMap::new(), ); @@ -371,6 +377,7 @@ fn maybe_wrap_shell_lc_with_snapshot_accepts_dot_alias_cwd() { &session_shell, &command_cwd.abs(), &HashMap::new(), + &[], &HashMap::new(), ); @@ -407,6 +414,7 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_explicit_override_precedence() { &session_shell, &dir.path().abs(), &explicit_env_overrides, + &[], &HashMap::from([("TEST_ENV_SNAPSHOT".to_string(), "worktree".to_string())]), ); let output = Command::new(&rewritten[0]) @@ -447,6 +455,7 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_codex_thread_id_from_env() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::from([("CODEX_THREAD_ID".to_string(), "nested-thread".to_string())]), ); let output = Command::new(&rewritten[0]) @@ -489,6 +498,7 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_proxy_env_from_process_env() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); let output = Command::new(&rewritten[0]) @@ -546,6 +556,7 @@ fn maybe_wrap_shell_lc_with_snapshot_refreshes_codex_proxy_git_ssh_command() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); let output = Command::new(&rewritten[0]) @@ -591,6 +602,7 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_custom_git_ssh_command() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); let output = Command::new(&rewritten[0]) @@ -637,6 +649,7 @@ fn maybe_wrap_shell_lc_with_snapshot_clears_stale_codex_git_ssh_command_without_ &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); let output = Command::new(&rewritten[0]) @@ -674,6 +687,7 @@ fn maybe_wrap_shell_lc_with_snapshot_keeps_user_proxy_env_when_proxy_inactive() &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); let mut command = Command::new(&rewritten[0]); @@ -724,6 +738,7 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_live_env_when_snapshot_proxy_activ &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::from([( "HTTP_PROXY".to_string(), "http://user.proxy:8080".to_string(), @@ -769,6 +784,7 @@ fn maybe_wrap_shell_lc_with_snapshot_keeps_snapshot_path_without_override() { &session_shell, &dir.path().abs(), &HashMap::new(), + &[], &HashMap::new(), ); let output = Command::new(&rewritten[0]) @@ -806,6 +822,46 @@ fn maybe_wrap_shell_lc_with_snapshot_applies_explicit_path_override() { &session_shell, &dir.path().abs(), &explicit_env_overrides, + &[], + &HashMap::from([("PATH".to_string(), "/worktree/bin".to_string())]), + ); + let output = Command::new(&rewritten[0]) + .args(&rewritten[1..]) + .env("PATH", "/worktree/bin") + .output() + .expect("run rewritten command"); + + assert!(output.status.success(), "command failed: {output:?}"); + assert_eq!(String::from_utf8_lossy(&output.stdout), "/worktree/bin"); +} + +#[test] +fn maybe_wrap_shell_lc_with_snapshot_restores_extra_live_path_key() { + let dir = tempdir().expect("create temp dir"); + let snapshot_path = dir.path().join("snapshot.sh"); + std::fs::write( + &snapshot_path, + "# Snapshot file\nexport PATH='/snapshot/bin'\n", + ) + .expect("write snapshot"); + let session_shell = shell_with_snapshot( + ShellType::Bash, + "/bin/bash", + snapshot_path.abs(), + dir.path().abs(), + ); + let command = vec![ + "/bin/bash".to_string(), + "-lc".to_string(), + "printf '%s' \"$PATH\"".to_string(), + ]; + let snapshot_restore_env_keys = vec!["PATH".to_string()]; + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + &dir.path().abs(), + &HashMap::new(), + &snapshot_restore_env_keys, &HashMap::from([("PATH".to_string(), "/worktree/bin".to_string())]), ); let output = Command::new(&rewritten[0]) @@ -847,6 +903,7 @@ fn maybe_wrap_shell_lc_with_snapshot_does_not_embed_override_values_in_argv() { &session_shell, &dir.path().abs(), &explicit_env_overrides, + &[], &HashMap::from([( "OPENAI_API_KEY".to_string(), "super-secret-value".to_string(), @@ -895,6 +952,7 @@ fn maybe_wrap_shell_lc_with_snapshot_preserves_unset_override_variables() { &session_shell, &dir.path().abs(), &explicit_env_overrides, + &[], &HashMap::new(), ); diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index 6251267773..77a3e5909f 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -54,6 +54,7 @@ pub struct ShellRequest { pub timeout_ms: Option, pub env: HashMap, pub explicit_env_overrides: HashMap, + pub snapshot_restore_env_keys: Vec, pub network: Option, pub sandbox_permissions: SandboxPermissions, pub additional_permissions: Option, @@ -237,6 +238,7 @@ impl ToolRuntime for ShellRuntime { session_shell.as_ref(), &req.cwd, &req.explicit_env_overrides, + &req.snapshot_restore_env_keys, &env, ); let command = disable_powershell_profile_for_elevated_windows_sandbox( diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index c613f198e0..224148dd98 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -66,6 +66,7 @@ pub struct UnifiedExecRequest { pub env: HashMap, pub exec_server_env_config: Option, pub explicit_env_overrides: HashMap, + pub snapshot_restore_env_keys: Vec, pub network: Option, pub tty: bool, pub sandbox_permissions: SandboxPermissions, @@ -270,6 +271,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt session_shell.as_ref(), &req.cwd, &req.explicit_env_overrides, + &req.snapshot_restore_env_keys, &env, ) }; @@ -418,6 +420,7 @@ mod tests { env: HashMap::new(), exec_server_env_config: None, explicit_env_overrides: HashMap::new(), + snapshot_restore_env_keys: Vec::new(), network: None, tty: false, sandbox_permissions: SandboxPermissions::UseDefault, diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index ce0f273f61..34e7828744 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -1006,13 +1006,14 @@ impl UnifiedExecProcessManager { /*thread_id*/ None, ); let mut env = local_policy_env.clone(); - if request.environment.is_remote() { + let snapshot_restore_env_keys = if request.environment.is_remote() { tracing::debug!( "CODEX_ENV_FILE overlays are local-only; skipping persisted hook env for remote exec", ); + Vec::new() } else { - let _ = context.session.apply_hook_env_file(&mut env); - } + context.session.apply_hook_env_file(&mut env) + }; env.insert( CODEX_THREAD_ID_ENV_VAR.to_string(), context.session.conversation_id.to_string(), @@ -1057,6 +1058,7 @@ impl UnifiedExecProcessManager { env, exec_server_env_config: Some(exec_server_env_config), explicit_env_overrides: context.turn.shell_environment_policy.r#set.clone(), + snapshot_restore_env_keys, network: request.network.clone(), tty: request.tty, sandbox_permissions: request.sandbox_permissions,