openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-09 15:58:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refine env file snapshot restore

Browse Source
This commit is contained in:
Abhinav Vedmala
2026-05-26 11:26:33 -07:00
parent 16190659f5
commit 28dff23b2f
9 changed files with 112 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
codex-rs/core/src/tasks/user_shell.rs
View File

@@ -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,
);

14
codex-rs/core/src/tools/handlers/shell.rs
View File

@@ -53,13 +53,12 @@ struct RunExecLikeArgs {
tracker: crate::tools::context::SharedTurnDiffTracker,
call_id: String,
shell_runtime_backend: ShellRuntimeBackend,
snapshot_restore_env_keys: Vec<String>,
}
async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, FunctionCallError> {
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<FunctionToolOutput, Func
tracker,
call_id,
shell_runtime_backend,
snapshot_restore_env_keys,
} = args;
let Some(turn_environment) = turn.environments.primary() else {
@@ -79,12 +77,6 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, Func
};
let fs = turn_environment.environment.get_filesystem();
let mut explicit_env_overrides = turn.shell_environment_policy.r#set.clone();
for key in snapshot_restore_env_keys {
// The snapshot wrapper restores the live values for keys listed here after sourcing a
// shell snapshot. Include hook env-file keys so persisted updates survive the snapshot.
explicit_env_overrides.entry(key).or_default();
}
let exec_permission_approvals_enabled =
session.features().enabled(Feature::ExecPermissionApprovals);
let requested_additional_permissions = additional_permissions.clone();
@@ -164,6 +156,7 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, Func
);
emitter.begin(event_ctx).await;
let snapshot_restore_env_keys = session.apply_hook_env_file(&mut exec_params.env);
let file_system_sandbox_policy = turn.file_system_sandbox_policy();
let exec_approval_requirement = session
.services
@@ -191,7 +184,8 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, Func
cwd: exec_params.cwd.clone(),
timeout_ms: exec_params.expiration.timeout_ms(),
env: exec_params.env.clone(),
explicit_env_overrides,
explicit_env_overrides: turn.shell_environment_policy.r#set.clone(),
snapshot_restore_env_keys,
network: exec_params.network.clone(),
sandbox_permissions: effective_additional_permissions.sandbox_permissions,
additional_permissions: normalized_additional_permissions,

48
codex-rs/core/src/tools/handlers/shell/shell_command.rs
View File

@@ -42,11 +42,6 @@ pub struct ShellCommandHandler {
options: ShellCommandHandlerOptions,
}
pub(super) struct ShellCommandExecParams {
pub(super) exec_params: ExecParams,
pub(super) snapshot_restore_env_keys: Vec<String>,
}
#[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<ShellCommandExecParams, FunctionCallError> {
) -> Result<ExecParams, FunctionCallError> {
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(), &params.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<ToolInvocation> 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(
&params,
session.as_ref(),
turn.as_ref(),
@@ -195,7 +184,7 @@ impl ToolExecutor<ToolInvocation> 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<ToolInvocation> 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)

11
codex-rs/core/src/tools/handlers/shell_tests.rs
View File

@@ -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(
&params,
&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::<String>::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(
&params,
&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)

26
codex-rs/core/src/tools/runtimes/mod.rs
View File

@@ -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<String, String>,
snapshot_restore_env_keys: &[String],
env: &HashMap<String, String>,
) -> Vec<String> {
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::<String>();
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>) -> (String, String) {
fn build_override_exports(
explicit_env_overrides: &HashMap<String, String>,
snapshot_restore_env_keys: &[String],
env: &HashMap<String, String>,
) -> (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::<Vec<_>>();
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)
}

58
codex-rs/core/src/tools/runtimes/mod_tests.rs
View File

@@ -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(),
);

2
codex-rs/core/src/tools/runtimes/shell.rs
View File

@@ -54,6 +54,7 @@ pub struct ShellRequest {
pub timeout_ms: Option<u64>,
pub env: HashMap<String, String>,
pub explicit_env_overrides: HashMap<String, String>,
pub snapshot_restore_env_keys: Vec<String>,
pub network: Option<NetworkProxy>,
pub sandbox_permissions: SandboxPermissions,
pub additional_permissions: Option<AdditionalPermissionProfile>,
@@ -237,6 +238,7 @@ impl ToolRuntime<ShellRequest, ExecToolCallOutput> 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(

3
codex-rs/core/src/tools/runtimes/unified_exec.rs
View File

@@ -66,6 +66,7 @@ pub struct UnifiedExecRequest {
pub env: HashMap<String, String>,
pub exec_server_env_config: Option<ExecServerEnvConfig>,
pub explicit_env_overrides: HashMap<String, String>,
pub snapshot_restore_env_keys: Vec<String>,
pub network: Option<NetworkProxy>,
pub tty: bool,
pub sandbox_permissions: SandboxPermissions,
@@ -270,6 +271,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> 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,

8
codex-rs/core/src/unified_exec/process_manager.rs
View File

@@ -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,