diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 3181e73698..744d8d797c 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -124,11 +124,16 @@ pub(crate) async fn execute_user_shell_command( let use_login_shell = true; let session_shell = session.user_shell(); let display_command = session_shell.derive_exec_args(&command, use_login_shell); + let exec_env_map = create_env( + &turn_context.shell_environment_policy, + Some(session.conversation_id), + ); let exec_command = maybe_wrap_shell_lc_with_snapshot( &display_command, session_shell.as_ref(), turn_context.cwd.as_path(), &turn_context.shell_environment_policy.r#set, + &exec_env_map, ); let call_id = Uuid::new_v4().to_string(); @@ -156,10 +161,7 @@ pub(crate) async fn execute_user_shell_command( let exec_env = ExecRequest { command: exec_command.clone(), cwd: cwd.to_path_buf(), - env: create_env( - &turn_context.shell_environment_policy, - Some(session.conversation_id), - ), + env: exec_env_map, network: turn_context.network.clone(), // TODO(zhao-oai): Now that we have ExecExpiration::Cancellation, we // should use that instead of an "arbitrarily large" timeout here. diff --git a/codex-rs/core/src/tools/runtimes/mod.rs b/codex-rs/core/src/tools/runtimes/mod.rs index a4b44ff0ec..99b5371f17 100644 --- a/codex-rs/core/src/tools/runtimes/mod.rs +++ b/codex-rs/core/src/tools/runtimes/mod.rs @@ -4,6 +4,7 @@ Module: runtimes Concrete ToolRuntime implementations for specific tools. Each runtime stays small and focused and reuses the orchestrator for approvals + sandbox + retry. */ +use crate::exec_env::CODEX_THREAD_ID_ENV_VAR; use crate::path_utils; use crate::shell::Shell; use crate::tools::sandboxing::ToolError; @@ -48,11 +49,19 @@ pub(crate) fn build_sandbox_command( /// This wrapper script uses POSIX constructs (`if`, `.`, `exec`) so it can /// be run by Bash/Zsh/sh. On non-matching commands, or when command cwd does /// not match the snapshot cwd, this is a no-op. +/// +/// `explicit_env_overrides` and `env` are intentionally separate inputs. +/// `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. pub(crate) fn maybe_wrap_shell_lc_with_snapshot( command: &[String], session_shell: &Shell, cwd: &Path, explicit_env_overrides: &HashMap, + env: &HashMap, ) -> Vec { if cfg!(windows) { return command.to_vec(); @@ -95,7 +104,11 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot( .iter() .map(|arg| format!(" '{}'", shell_single_quote(arg))) .collect::(); - let (override_captures, override_exports) = build_override_exports(explicit_env_overrides); + 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 rewritten_script = if override_exports.is_empty() { format!( "if . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\nexec '{original_shell}' -c '{original_script}'{trailing_args}" diff --git a/codex-rs/core/src/tools/runtimes/mod_tests.rs b/codex-rs/core/src/tools/runtimes/mod_tests.rs index dbc341d1de..22c22770ab 100644 --- a/codex-rs/core/src/tools/runtimes/mod_tests.rs +++ b/codex-rs/core/src/tools/runtimes/mod_tests.rs @@ -42,8 +42,13 @@ fn maybe_wrap_shell_lc_with_snapshot_bootstraps_in_user_shell() { "echo hello".to_string(), ]; - let rewritten = - maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path(), &HashMap::new()); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + dir.path(), + &HashMap::new(), + &HashMap::new(), + ); assert_eq!(rewritten[0], "/bin/zsh"); assert_eq!(rewritten[1], "-c"); @@ -68,8 +73,13 @@ fn maybe_wrap_shell_lc_with_snapshot_escapes_single_quotes() { "echo 'hello'".to_string(), ]; - let rewritten = - maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path(), &HashMap::new()); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + dir.path(), + &HashMap::new(), + &HashMap::new(), + ); assert!(rewritten[2].contains(r#"exec '/bin/bash' -c 'echo '"'"'hello'"'"''"#)); } @@ -91,8 +101,13 @@ fn maybe_wrap_shell_lc_with_snapshot_uses_bash_bootstrap_shell() { "echo hello".to_string(), ]; - let rewritten = - maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path(), &HashMap::new()); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + dir.path(), + &HashMap::new(), + &HashMap::new(), + ); assert_eq!(rewritten[0], "/bin/bash"); assert_eq!(rewritten[1], "-c"); @@ -117,8 +132,13 @@ fn maybe_wrap_shell_lc_with_snapshot_uses_sh_bootstrap_shell() { "echo hello".to_string(), ]; - let rewritten = - maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path(), &HashMap::new()); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + dir.path(), + &HashMap::new(), + &HashMap::new(), + ); assert_eq!(rewritten[0], "/bin/sh"); assert_eq!(rewritten[1], "-c"); @@ -145,8 +165,13 @@ fn maybe_wrap_shell_lc_with_snapshot_preserves_trailing_args() { "arg1".to_string(), ]; - let rewritten = - maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path(), &HashMap::new()); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + dir.path(), + &HashMap::new(), + &HashMap::new(), + ); assert!( rewritten[2] @@ -171,8 +196,13 @@ fn maybe_wrap_shell_lc_with_snapshot_skips_when_cwd_mismatch() { "echo hello".to_string(), ]; - let rewritten = - maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, &command_cwd, &HashMap::new()); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + &command_cwd, + &HashMap::new(), + &HashMap::new(), + ); assert_eq!(rewritten, command); } @@ -195,8 +225,13 @@ fn maybe_wrap_shell_lc_with_snapshot_accepts_dot_alias_cwd() { ]; let command_cwd = dir.path().join("."); - let rewritten = - maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, &command_cwd, &HashMap::new()); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + &command_cwd, + &HashMap::new(), + &HashMap::new(), + ); assert_eq!(rewritten[0], "/bin/zsh"); assert_eq!(rewritten[1], "-c"); @@ -231,6 +266,7 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_explicit_override_precedence() { &session_shell, dir.path(), &explicit_env_overrides, + &HashMap::from([("TEST_ENV_SNAPSHOT".to_string(), "worktree".to_string())]), ); let output = Command::new(&rewritten[0]) .args(&rewritten[1..]) @@ -245,6 +281,43 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_explicit_override_precedence() { ); } +#[test] +fn maybe_wrap_shell_lc_with_snapshot_restores_codex_thread_id_from_env() { + let dir = tempdir().expect("create temp dir"); + let snapshot_path = dir.path().join("snapshot.sh"); + std::fs::write( + &snapshot_path, + "# Snapshot file\nexport CODEX_THREAD_ID='parent-thread'\n", + ) + .expect("write snapshot"); + let session_shell = shell_with_snapshot( + ShellType::Bash, + "/bin/bash", + snapshot_path, + dir.path().to_path_buf(), + ); + let command = vec![ + "/bin/bash".to_string(), + "-lc".to_string(), + "printf '%s' \"$CODEX_THREAD_ID\"".to_string(), + ]; + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + dir.path(), + &HashMap::new(), + &HashMap::from([("CODEX_THREAD_ID".to_string(), "nested-thread".to_string())]), + ); + let output = Command::new(&rewritten[0]) + .args(&rewritten[1..]) + .env("CODEX_THREAD_ID", "nested-thread") + .output() + .expect("run rewritten command"); + + assert!(output.status.success(), "command failed: {output:?}"); + assert_eq!(String::from_utf8_lossy(&output.stdout), "nested-thread"); +} + #[test] fn maybe_wrap_shell_lc_with_snapshot_keeps_snapshot_path_without_override() { let dir = tempdir().expect("create temp dir"); @@ -265,8 +338,13 @@ fn maybe_wrap_shell_lc_with_snapshot_keeps_snapshot_path_without_override() { "-lc".to_string(), "printf '%s' \"$PATH\"".to_string(), ]; - let rewritten = - maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path(), &HashMap::new()); + let rewritten = maybe_wrap_shell_lc_with_snapshot( + &command, + &session_shell, + dir.path(), + &HashMap::new(), + &HashMap::new(), + ); let output = Command::new(&rewritten[0]) .args(&rewritten[1..]) .output() @@ -302,6 +380,7 @@ fn maybe_wrap_shell_lc_with_snapshot_applies_explicit_path_override() { &session_shell, dir.path(), &explicit_env_overrides, + &HashMap::from([("PATH".to_string(), "/worktree/bin".to_string())]), ); let output = Command::new(&rewritten[0]) .args(&rewritten[1..]) @@ -342,6 +421,10 @@ fn maybe_wrap_shell_lc_with_snapshot_does_not_embed_override_values_in_argv() { &session_shell, dir.path(), &explicit_env_overrides, + &HashMap::from([( + "OPENAI_API_KEY".to_string(), + "super-secret-value".to_string(), + )]), ); assert!(!rewritten[2].contains("super-secret-value")); @@ -386,6 +469,7 @@ fn maybe_wrap_shell_lc_with_snapshot_preserves_unset_override_variables() { &session_shell, dir.path(), &explicit_env_overrides, + &HashMap::new(), ); let output = Command::new(&rewritten[0]) diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index ba8713254a..7b0c81ac60 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -226,6 +226,7 @@ impl ToolRuntime for ShellRuntime { session_shell.as_ref(), &req.cwd, &req.explicit_env_overrides, + &req.env, ); let command = if matches!(session_shell.shell_type, ShellType::PowerShell) { prefix_powershell_script_with_utf8(&command) diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 3858150da5..7b271363b2 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -207,6 +207,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt session_shell.as_ref(), &req.cwd, &req.explicit_env_overrides, + &req.env, ); let command = if matches!(session_shell.shell_type, ShellType::PowerShell) { prefix_powershell_script_with_utf8(&command)