Merge branch 'main' into etraut/nested-exec-thread-id

codex-rs/core/src/tasks/user_shell.rs

@@ -123,11 +123,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();
@@ -155,10 +160,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.

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<String, String>,
+    env: &HashMap<String, String>,
 ) -> Vec<String> {
     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::<String>();
-    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}"

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])

codex-rs/core/src/tools/runtimes/shell.rs

@@ -226,6 +226,7 @@ impl ToolRuntime<ShellRequest, ExecToolCallOutput> 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)

codex-rs/core/src/tools/runtimes/unified_exec.rs

@@ -207,6 +207,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> 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)

codex-rs/exec/src/cli.rs

@@ -47,7 +47,7 @@ pub struct Cli {
     #[arg(long = "profile", short = 'p')]
     pub config_profile: Option<String>,
 
-    /// Convenience alias for low-friction sandboxed automatic execution (--sandbox workspace-write).
+    /// Convenience alias for low-friction sandboxed automatic execution (-a on-request, --sandbox workspace-write).
     #[arg(long = "full-auto", default_value_t = false, global = true)]
     pub full_auto: bool,