Simplify session env file exports

codex-rs/core/src/shell_env_file.rs

@@ -1,19 +1,28 @@
 use std::collections::HashMap;
 use std::path::Path;
+use std::sync::Mutex;
 
-#[cfg(not(windows))]
 use anyhow::Context;
 use anyhow::Result;
+use anyhow::anyhow;
+use anyhow::bail;
 use codex_protocol::ThreadId;
 use codex_protocol::shell_environment::CODEX_ENV_FILE_ENV_VAR;
+use codex_protocol::shell_environment::CODEX_THREAD_ID_ENV_VAR;
 use tempfile::TempPath;
+use tokio::process::Command;
+
+use crate::shell::Shell;
+use crate::shell::ShellType;
 
 /// Session-owned script that hooks can populate with exported shell state.
 ///
-/// Local shell tool commands source this file before running so lifecycle hook
-/// setup remains scoped to the active session rather than persistent config.
+/// Only lifecycle hooks receive the writable file path. After SessionStart
+/// hooks finish, Codex captures supported exported variables and passes those
+/// values to later commands without exposing the writable path.
 pub(crate) struct ShellEnvFile {
     path: TempPath,
+    exports: Mutex<HashMap<String, Option<String>>>,
 }
 
 impl ShellEnvFile {
@@ -40,6 +49,7 @@ impl ShellEnvFile {
             .context("failed to create temporary shell env file")?;
         Ok(Self {
             path: file.into_temp_path(),
+            exports: Mutex::new(HashMap::new()),
         })
     }
 
@@ -47,12 +57,184 @@ impl ShellEnvFile {
         self.path.as_ref()
     }
 
-    pub(crate) fn insert_into_env(&self, env: &mut HashMap<String, String>) {
+    pub(crate) fn insert_path_into_env(&self, env: &mut HashMap<String, String>) {
         env.insert(
             CODEX_ENV_FILE_ENV_VAR.to_string(),
             self.path().to_string_lossy().to_string(),
         );
     }
+
+    /// Sources the hook-writable env file once and stores the resulting
+    /// environment diff.
+    ///
+    /// The temp file remains an input channel for SessionStart hooks, but later
+    /// commands receive only captured variable changes. Running both a baseline
+    /// environment dump and a sourced dump lets shell syntax such as command
+    /// substitution behave naturally without keeping `CODEX_ENV_FILE` available
+    /// after hook execution.
+    pub(crate) async fn capture_exports(
+        &self,
+        shell: &Shell,
+        cwd: &Path,
+        base_env: &HashMap<String, String>,
+    ) -> Result<()> {
+        if !matches!(
+            shell.shell_type,
+            ShellType::Zsh | ShellType::Bash | ShellType::Sh
+        ) {
+            return Ok(());
+        }
+
+        let mut capture_env = base_env.clone();
+        capture_env.remove(CODEX_ENV_FILE_ENV_VAR);
+        self.insert_path_into_env(&mut capture_env);
+
+        let baseline = capture_env_from_shell(shell, cwd, DUMP_ENV_SCRIPT, &capture_env).await?;
+        let captured = capture_env_from_shell(
+            shell,
+            cwd,
+            SOURCE_ENV_FILE_AND_DUMP_ENV_SCRIPT,
+            &capture_env,
+        )
+        .await?;
+        let exports = diff_env(&baseline, &captured);
+        *self
+            .exports
+            .lock()
+            .map_err(|_| anyhow!("shell env exports lock poisoned"))? = exports;
+        Ok(())
+    }
+
+    /// Applies captured SessionStart environment changes to a command
+    /// environment without exposing the writable env-file path.
+    ///
+    /// Explicit shell-environment policy values are layered back on top so
+    /// configured overrides keep precedence, and runtime-owned values such as
+    /// the Codex thread id are preserved rather than accepting hook-written
+    /// replacements.
+    pub(crate) fn apply_exports(
+        &self,
+        env: &mut HashMap<String, String>,
+        explicit_env_overrides: &HashMap<String, String>,
+    ) {
+        let thread_id = env.get(CODEX_THREAD_ID_ENV_VAR).cloned();
+        let exports = self
+            .exports
+            .lock()
+            .map(|exports| exports.clone())
+            .unwrap_or_default();
+        for (key, value) in exports {
+            if ignored_capture_key(&key) {
+                continue;
+            }
+            match value {
+                Some(value) => {
+                    env.insert(key, value);
+                }
+                None => {
+                    env.remove(&key);
+                }
+            }
+        }
+        for (key, value) in explicit_env_overrides {
+            env.insert(key.clone(), value.clone());
+        }
+        if let Some(thread_id) = thread_id {
+            env.insert(CODEX_THREAD_ID_ENV_VAR.to_string(), thread_id);
+        }
+        env.remove(CODEX_ENV_FILE_ENV_VAR);
+    }
+}
+
+const DUMP_ENV_SCRIPT: &str = r#"if [ -x /usr/bin/env ]; then
+  /usr/bin/env -0
+else
+  env -0
+fi"#;
+
+const SOURCE_ENV_FILE_AND_DUMP_ENV_SCRIPT: &str = r#"if [ -n "${CODEX_ENV_FILE:-}" ] && [ -f "$CODEX_ENV_FILE" ]; then
+  if . "$CODEX_ENV_FILE" >/dev/null 2>&1; then
+    :
+  fi
+fi
+if [ -x /usr/bin/env ]; then
+  /usr/bin/env -0
+else
+  env -0
+fi"#;
+
+async fn capture_env_from_shell(
+    shell: &Shell,
+    cwd: &Path,
+    script: &str,
+    env: &HashMap<String, String>,
+) -> Result<HashMap<String, String>> {
+    let output = Command::new(&shell.shell_path)
+        .current_dir(cwd)
+        .arg("-c")
+        .arg(script)
+        .env_clear()
+        .envs(env)
+        .output()
+        .await
+        .with_context(|| format!("failed to run {}", shell.shell_path.display()))?;
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        bail!(
+            "failed to capture shell environment with {}: {stderr}",
+            shell.shell_path.display()
+        );
+    }
+    parse_env_output(&output.stdout)
+}
+
+fn parse_env_output(output: &[u8]) -> Result<HashMap<String, String>> {
+    let mut env = HashMap::new();
+    for entry in output.split(|byte| *byte == 0) {
+        if entry.is_empty() {
+            continue;
+        }
+        let Some(separator) = entry.iter().position(|byte| *byte == b'=') else {
+            continue;
+        };
+        let key = String::from_utf8(entry[..separator].to_vec())
+            .context("captured shell environment key was not UTF-8")?;
+        let value = String::from_utf8(entry[separator + 1..].to_vec())
+            .context("captured shell environment value was not UTF-8")?;
+        env.insert(key, value);
+    }
+    Ok(env)
+}
+
+fn diff_env(
+    baseline: &HashMap<String, String>,
+    captured: &HashMap<String, String>,
+) -> HashMap<String, Option<String>> {
+    let mut exports = HashMap::new();
+    for (key, value) in captured {
+        if ignored_capture_key(key) {
+            continue;
+        }
+        if baseline.get(key) != Some(value) {
+            exports.insert(key.clone(), Some(value.clone()));
+        }
+    }
+    for key in baseline.keys() {
+        if ignored_capture_key(key) {
+            continue;
+        }
+        if !captured.contains_key(key) {
+            exports.insert(key.clone(), None);
+        }
+    }
+    exports
+}
+
+fn ignored_capture_key(key: &str) -> bool {
+    matches!(
+        key,
+        CODEX_ENV_FILE_ENV_VAR | CODEX_THREAD_ID_ENV_VAR | "PWD" | "OLDPWD" | "SHLVL" | "_"
+    )
 }
 
 #[cfg(all(test, not(windows)))]