Cache remote Bash environment exports

codex-rs/core/src/unified_exec/process_manager.rs

@@ -21,6 +21,7 @@ use crate::exec_policy::ExecApprovalRequest;
 use crate::sandboxing::ExecOptions;
 use crate::sandboxing::ExecRequest;
 use crate::sandboxing::ExecServerEnvConfig;
+use crate::shell::ShellType;
 use crate::tools::context::ExecCommandToolOutput;
 use crate::tools::events::ToolEmitter;
 use crate::tools::events::ToolEventCtx;
@@ -56,6 +57,7 @@ use crate::unified_exec::process::OutputBuffer;
 use crate::unified_exec::process::OutputHandles;
 use crate::unified_exec::process::SpawnLifecycleHandle;
 use crate::unified_exec::process::UnifiedExecProcess;
+use codex_features::Feature;
 use codex_network_proxy::NetworkProxy;
 use codex_protocol::config_types::ShellEnvironmentPolicy;
 use codex_protocol::error::CodexErr;
@@ -129,6 +131,7 @@ fn exec_env_policy_from_shell_policy(
             .iter()
             .map(std::string::ToString::to_string)
             .collect(),
+        bash_env_cache_scope: None,
     }
 }
 
@@ -1122,10 +1125,19 @@ impl UnifiedExecProcessManager {
         let active_permission_profile = context.turn.config.permissions.active_permission_profile();
         inject_permission_profile_env(&mut env, active_permission_profile.as_ref());
         let env = apply_unified_exec_env(env);
+        let mut exec_server_policy = exec_env_policy_from_shell_policy(
+            &context.turn.config.permissions.shell_environment_policy,
+        );
+        if context.turn.config.features.enabled(Feature::ShellSnapshot)
+            && request.turn_environment.environment.is_remote()
+            && request.shell_type == ShellType::Bash
+            && !request.tty
+            && matches!(request.command.as_slice(), [_, option, _] if option == "-c")
+        {
+            exec_server_policy.bash_env_cache_scope = Some(request.turn_environment.cwd().clone());
+        }
         let exec_server_env_config = ExecServerEnvConfig {
-            policy: exec_env_policy_from_shell_policy(
-                &context.turn.config.permissions.shell_environment_policy,
-            ),
+            policy: exec_server_policy,
             local_policy_env,
         };
         let mut orchestrator = ToolOrchestrator::new();

codex-rs/core/src/unified_exec/process_manager_tests.rs

@@ -100,6 +100,7 @@ fn exec_env_policy_excludes_runtime_permission_profile() {
             exclude: vec![CODEX_PERMISSION_PROFILE_ENV_VAR.to_string()],
             r#set: HashMap::from([("KEEP".to_string(), "value".to_string())]),
             include_only: Vec::new(),
+            bash_env_cache_scope: None,
         }
     );
 }
@@ -142,6 +143,7 @@ fn exec_server_params_use_path_uri_and_env_policy_overlay_contract() {
                 exclude: Vec::new(),
                 r#set: HashMap::new(),
                 include_only: Vec::new(),
+                bash_env_cache_scope: None,
             },
             local_policy_env: HashMap::from([
                 ("HOME".to_string(), "/client-home".to_string()),

codex-rs/exec-server-protocol/src/protocol.rs

@@ -154,6 +154,10 @@ pub struct ExecEnvPolicy {
     pub exclude: Vec<String>,
     pub r#set: HashMap<String, String>,
     pub include_only: Vec<String>,
+    /// Optional workspace scope for caching exports produced by Bash's `BASH_ENV`.
+    /// Older exec-server peers ignore this field, so reuse is opportunistic.
+    #[serde(default)]
+    pub bash_env_cache_scope: Option<PathUri>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -566,6 +570,7 @@ mod base64_bytes {
 #[cfg(test)]
 mod tests {
     use super::EnvironmentInfo;
+    use super::ExecEnvPolicy;
     use super::ExecExitedNotification;
     use super::ExecParams;
     use super::FsReadFileParams;
@@ -574,21 +579,29 @@ mod tests {
     use super::ShellInfo;
     use codex_file_system::FileSystemSandboxContext;
     use codex_network_proxy::ManagedNetworkSandboxContext;
+    use codex_protocol::config_types::ShellEnvironmentPolicyInherit;
     use codex_protocol::models::PermissionProfile;
     use codex_utils_path_uri::PathUri;
     use pretty_assertions::assert_eq;
     use std::collections::HashMap;
 
     #[test]
-    fn exec_params_managed_network_context_round_trips_and_defaults_for_legacy_peers() {
+    fn exec_params_optional_context_round_trips_and_defaults_for_legacy_peers() {
         let cwd =
             PathUri::from_host_native_path(std::env::current_dir().expect("current directory"))
                 .expect("cwd URI");
         let params = ExecParams {
             process_id: ProcessId::from("managed-network"),
             argv: vec!["true".to_string()],
-            cwd,
-            env_policy: None,
+            cwd: cwd.clone(),
+            env_policy: Some(ExecEnvPolicy {
+                inherit: ShellEnvironmentPolicyInherit::Core,
+                ignore_default_excludes: false,
+                exclude: Vec::new(),
+                r#set: HashMap::new(),
+                include_only: Vec::new(),
+                bash_env_cache_scope: Some(cwd.clone()),
+            }),
             env: HashMap::new(),
             tty: false,
             pipe_stdin: false,
@@ -609,6 +622,10 @@ mod tests {
                 "allowLocalBinding": false,
             })
         );
+        assert_eq!(
+            serialized["envPolicy"]["bashEnvCacheScope"],
+            serde_json::to_value(cwd).expect("serialize cwd")
+        );
         let round_trip: ExecParams =
             serde_json::from_value(serialized.clone()).expect("deserialize exec params");
         assert_eq!(round_trip, params);
@@ -617,10 +634,18 @@ mod tests {
             .as_object_mut()
             .expect("exec params object")
             .remove("managedNetwork");
+        serialized["envPolicy"]
+            .as_object_mut()
+            .expect("env policy object")
+            .remove("bashEnvCacheScope");
         let legacy: ExecParams =
             serde_json::from_value(serialized).expect("deserialize legacy exec params");
         assert!(legacy.enforce_managed_network);
         assert_eq!(legacy.managed_network, None);
+        assert_eq!(
+            legacy.env_policy.expect("env policy").bash_env_cache_scope,
+            None
+        );
     }
 
     #[test]

codex-rs/exec-server/src/bash_env_cache.rs

@@ -0,0 +1,242 @@
+use std::collections::HashMap;
+
+use crate::ExecServerRuntimePaths;
+use crate::protocol::ExecParams;
+
+#[cfg(unix)]
+mod imp {
+    use std::path::Path;
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    use codex_file_system::FileSystemSandboxContext;
+    use codex_network_proxy::ManagedNetworkSandboxContext;
+    use codex_utils_path_uri::PathUri;
+    use tokio::sync::Mutex;
+    use tokio::sync::OnceCell;
+    use tokio::time::timeout;
+    use uuid::Uuid;
+
+    use super::*;
+    use crate::process_sandbox::prepare_exec_request;
+
+    const CAPTURE_TIMEOUT: Duration = Duration::from_secs(10);
+    const MAX_CAPTURE_BYTES: usize = 1024 * 1024;
+
+    #[derive(Default)]
+    pub(crate) struct BashEnvCache {
+        entry: Mutex<Option<CacheEntry>>,
+    }
+
+    struct CacheEntry {
+        key: CacheKey,
+        environment: Arc<OnceCell<Option<HashMap<String, String>>>>,
+    }
+
+    #[derive(PartialEq, Eq)]
+    struct CacheKey {
+        scope: PathUri,
+        shell: String,
+        cwd: PathUri,
+        environment: HashMap<String, String>,
+        sandbox: Option<FileSystemSandboxContext>,
+        enforce_managed_network: bool,
+        managed_network: Option<ManagedNetworkSandboxContext>,
+    }
+
+    impl BashEnvCache {
+        pub(crate) async fn environment_for_launch(
+            &self,
+            params: &ExecParams,
+            environment: HashMap<String, String>,
+            runtime_paths: Option<&ExecServerRuntimePaths>,
+        ) -> HashMap<String, String> {
+            let Some(key) = CacheKey::new(params, &environment) else {
+                return environment;
+            };
+
+            let cached_environment = {
+                let mut entry = self.entry.lock().await;
+                if let Some(cached) = entry.as_ref()
+                    && cached.key == key
+                {
+                    Arc::clone(&cached.environment)
+                } else {
+                    let environment = Arc::new(OnceCell::new());
+                    *entry = Some(CacheEntry {
+                        key,
+                        environment: Arc::clone(&environment),
+                    });
+                    environment
+                }
+            };
+            cached_environment
+                .get_or_init(|| capture_environment(params, &environment, runtime_paths))
+                .await
+                .clone()
+                .unwrap_or(environment)
+        }
+    }
+
+    impl CacheKey {
+        fn new(params: &ExecParams, environment: &HashMap<String, String>) -> Option<Self> {
+            let scope = params.env_policy.as_ref()?.bash_env_cache_scope.as_ref()?;
+            let [shell, option, _script] = params.argv.as_slice() else {
+                return None;
+            };
+            if params.tty
+                || params.pipe_stdin
+                || params.arg0.is_some()
+                || option != "-c"
+                || Path::new(shell).file_name()?.to_str()? != "bash"
+                || !params.cwd.starts_with(scope)
+                || environment.get("BASH_ENV").is_none_or(String::is_empty)
+            {
+                return None;
+            }
+
+            Some(Self {
+                scope: scope.clone(),
+                shell: shell.clone(),
+                cwd: params.cwd.clone(),
+                environment: environment.clone(),
+                sandbox: params.sandbox.clone(),
+                enforce_managed_network: params.enforce_managed_network,
+                managed_network: params.managed_network.clone(),
+            })
+        }
+    }
+
+    async fn capture_environment(
+        params: &ExecParams,
+        environment: &HashMap<String, String>,
+        runtime_paths: Option<&ExecServerRuntimePaths>,
+    ) -> Option<HashMap<String, String>> {
+        let nonce = Uuid::new_v4().simple().to_string();
+        let start_marker = format!("__CODEX_BASH_ENV_START_{nonce}__");
+        let end_marker = format!("__CODEX_BASH_ENV_END_{nonce}__");
+        let mut capture = params.clone();
+        capture.argv = vec![
+            params.argv.first()?.clone(),
+            "-c".to_string(),
+            format!(
+                "builtin printf '%s' '{start_marker}'; if builtin command env -0; then builtin printf '%s' '{end_marker}'; else builtin exit 125; fi"
+            ),
+        ];
+        let prepared = prepare_exec_request(&capture, environment.clone(), runtime_paths).ok()?;
+        let (program, args) = prepared.command.split_first()?;
+        let spawned = codex_utils_pty::spawn_pipe_process_no_stdin(
+            program,
+            args,
+            prepared.cwd.as_path(),
+            &prepared.env,
+            &prepared.arg0,
+        )
+        .await
+        .ok()?;
+
+        let captured = timeout(CAPTURE_TIMEOUT, async move {
+            let _session = spawned.session;
+            let (stdout, stderr, exit_code) = tokio::join!(
+                collect_output(spawned.stdout_rx),
+                collect_output(spawned.stderr_rx),
+                spawned.exit_rx,
+            );
+            (exit_code.ok()? == 0 && stderr?.is_empty()).then_some(parse_capture(
+                stdout?,
+                &start_marker,
+                &end_marker,
+            )?)
+        })
+        .await
+        .ok()
+        .flatten()?;
+
+        Some(sanitize_environment(captured, environment))
+    }
+
+    async fn collect_output(mut receiver: tokio::sync::mpsc::Receiver<Vec<u8>>) -> Option<Vec<u8>> {
+        let mut output = Vec::new();
+        while let Some(chunk) = receiver.recv().await {
+            if output.len().checked_add(chunk.len())? > MAX_CAPTURE_BYTES {
+                return None;
+            }
+            output.extend_from_slice(&chunk);
+        }
+        Some(output)
+    }
+
+    fn parse_capture(
+        output: Vec<u8>,
+        start_marker: &str,
+        end_marker: &str,
+    ) -> Option<HashMap<String, String>> {
+        let payload = output.strip_prefix(start_marker.as_bytes())?;
+        let payload = payload.strip_suffix(end_marker.as_bytes())?;
+        if !payload.is_empty() && !payload.ends_with(&[0]) {
+            return None;
+        }
+
+        payload
+            .split(|byte| *byte == 0)
+            .filter(|entry| !entry.is_empty())
+            .map(|entry| {
+                let separator = entry.iter().position(|byte| *byte == b'=')?;
+                let key = std::str::from_utf8(&entry[..separator]).ok()?.to_string();
+                let value = std::str::from_utf8(&entry[separator + 1..])
+                    .ok()?
+                    .to_string();
+                Some((key, value))
+            })
+            .collect()
+    }
+
+    fn sanitize_environment(
+        mut captured: HashMap<String, String>,
+        original: &HashMap<String, String>,
+    ) -> HashMap<String, String> {
+        captured.remove("BASH_ENV");
+        captured.retain(|key, _| !shell_managed(key));
+        captured.extend(
+            original
+                .iter()
+                .filter(|(key, _)| shell_managed(key))
+                .map(|(key, value)| (key.clone(), value.clone())),
+        );
+        captured
+    }
+
+    fn shell_managed(key: &str) -> bool {
+        key.starts_with("BASH_FUNC_")
+            || matches!(
+                key,
+                "BASHOPTS"
+                    | "BASH_ARGV0"
+                    | "BASH_EXECUTION_STRING"
+                    | "OLDPWD"
+                    | "PWD"
+                    | "SHELLOPTS"
+                    | "SHLVL"
+                    | "_"
+            )
+    }
+}
+
+#[cfg(unix)]
+pub(crate) use imp::BashEnvCache;
+
+#[cfg(not(unix))]
+#[derive(Default)]
+pub(crate) struct BashEnvCache;
+
+#[cfg(not(unix))]
+impl BashEnvCache {
+    pub(crate) async fn environment_for_launch(
+        &self,
+        _params: &ExecParams,
+        environment: HashMap<String, String>,
+        _runtime_paths: Option<&ExecServerRuntimePaths>,
+    ) -> HashMap<String, String> {
+        environment
+    }
+}

codex-rs/exec-server/src/lib.rs

@@ -1,3 +1,4 @@
+mod bash_env_cache;
 mod client;
 mod client_api;
 mod client_transport;