Add in-memory shell snapshots to unified exec (#39957)

## What changed

- Add the under-development `shell_snapshot_v2` feature and advertise executor support through environment capabilities.
- Use executor-managed, in-memory snapshots for eligible direct `bash`, `zsh`, and `sh` login commands, while applying the configured shell environment policy and avoiding duplicate automatic startup-file effects.
- Keep file-backed snapshots available for user-shell commands and fall back to the existing execution path when in-memory snapshots are unsupported or inapplicable.

## Testing

- Cover policy filtering, snapshot reuse without snapshot files, automatic startup files, local and remote execution, and legacy user-shell snapshots.

GitOrigin-RevId: 4ad6cdf13824913ac8c393ba38b9844230677579

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

@@ -51,6 +51,7 @@ mod head_tail_buffer;
 mod process;
 mod process_manager;
 mod process_state;
+mod shell_snapshot;
 
 pub(crate) fn set_deterministic_process_ids_for_tests(enabled: bool) {
     process_manager::set_deterministic_process_ids_for_tests(enabled);

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

@@ -60,6 +60,7 @@ use crate::unified_exec::head_tail_buffer::HeadTailBuffer;
 use crate::unified_exec::process::OutputHandles;
 use crate::unified_exec::process::SpawnLifecycleHandle;
 use crate::unified_exec::process::UnifiedExecProcess;
+use crate::unified_exec::shell_snapshot::shell_snapshot_request;
 use crate::unified_exec::take_plugin_metrics_sidecar;
 use codex_core_plugins::PLUGIN_METRICS_OUTPUT_ENV_VAR;
 use codex_core_plugins::PluginCommandAttribution;
@@ -186,6 +187,11 @@ fn exec_server_env_for_request(
     if let Some(exec_server_env_config) = &request.exec_server_env_config {
         let mut env =
             env_overlay_for_exec_server(&request.env, &exec_server_env_config.local_policy_env);
+        if request.exec_server_shell_snapshot.is_some()
+            && !exec_server_env_config.policy.r#set.contains_key("PATH")
+        {
+            env.remove("PATH");
+        }
         if request.exec_server_managed_network.is_some() {
             for (key, value) in &request.env {
                 if is_managed_proxy_env_var(key, value) {
@@ -210,18 +216,20 @@ fn exec_server_params_for_request(
         sandbox.windows_sandbox_proxy_settings_mode = Some(windows_sandbox_proxy_settings_mode);
         sandbox
     });
-    // Sandbox retries reuse the unified-exec ID but start a distinct executor process.
-    let exec_server_process_id = if request.exec_server_sandbox.is_some() {
-        format!("{process_id}-{}", Uuid::new_v4())
-    } else {
-        process_id.to_string()
-    };
+    // Sandbox retries and memory-backed local launches can reuse a unified-exec
+    // ID while the executor still retains the previous process.
+    let exec_server_process_id =
+        if request.exec_server_sandbox.is_some() || request.exec_server_shell_snapshot.is_some() {
+            format!("{process_id}-{}", Uuid::new_v4())
+        } else {
+            process_id.to_string()
+        };
     codex_exec_server::ExecParams {
         process_id: exec_server_process_id.into(),
         argv: request.command.clone(),
         cwd: request.cwd.clone(),
         env_policy,
-        shell_snapshot: None,
+        shell_snapshot: request.exec_server_shell_snapshot.clone(),
         env,
         tty,
         pipe_stdin: false,
@@ -1060,12 +1068,13 @@ impl UnifiedExecProcessManager {
         network_proxy_launch: Option<codex_network_proxy::RemoteNetworkProxyLaunchConfig>,
         environment_id: Option<&str>,
         exec_server_env_config: Option<ExecServerEnvConfig>,
+        shell_snapshot: Option<codex_exec_server::ShellSnapshotRequest>,
         windows_sandbox_proxy_settings_mode: codex_sandboxing::WindowsSandboxProxySettingsMode,
         tty: bool,
         spawn_lifecycle: SpawnLifecycleHandle,
         environment: &codex_exec_server::Environment,
     ) -> Result<UnifiedExecProcess, ToolError> {
-        let mut request = if environment.is_remote() {
+        let mut request = if environment.is_remote() || shell_snapshot.is_some() {
             attempt.env_for_exec_server(command, options)
         } else {
             attempt.env_for(command, options, network, environment_id)
@@ -1077,6 +1086,7 @@ impl UnifiedExecProcessManager {
             .and_then(|_| network.and_then(NetworkProxy::remote_policy_decider));
         request.exec_server_network_proxy = network_proxy_launch;
         request.exec_server_env_config = exec_server_env_config;
+        request.exec_server_shell_snapshot = shell_snapshot;
         self.open_session_with_prepared_exec_env(
             process_id,
             &request,
@@ -1111,7 +1121,7 @@ impl UnifiedExecProcessManager {
     ) -> Result<UnifiedExecProcess, UnifiedExecError> {
         let inherited_fds = spawn_lifecycle.inherited_fds();
 
-        if environment.is_remote() {
+        if environment.is_remote() || request.exec_server_shell_snapshot.is_some() {
             if !inherited_fds.is_empty() {
                 return Err(UnifiedExecError::create_process(
                     "remote exec-server does not support inherited file descriptors".to_string(),
@@ -1237,6 +1247,7 @@ impl UnifiedExecProcessManager {
             policy: exec_env_policy_from_shell_policy(shell_environment_policy),
             local_policy_env,
         };
+        let shell_snapshot = shell_snapshot_request(request, &cwd, context);
         let mut orchestrator = ToolOrchestrator::new();
         let mut runtime = UnifiedExecRuntime::new(self, request.shell_mode.clone());
         let session_shell = context.session.user_shell();
@@ -1278,6 +1289,7 @@ impl UnifiedExecProcessManager {
             turn_environment: request.turn_environment.clone(),
             env,
             exec_server_env_config: Some(exec_server_env_config),
+            shell_snapshot,
             explicit_env_overrides,
             network: request.network.clone(),
             tty: request.tty,

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

@@ -184,6 +184,7 @@ fn exec_server_params_use_path_uri_and_env_policy_overlay_contract() {
                 ),
             ]),
         }),
+        exec_server_shell_snapshot: None,
         network: None,
         network_environment_id: None,
         expiration: crate::exec::ExecExpiration::DefaultTimeout,
@@ -230,6 +231,18 @@ fn exec_server_params_use_path_uri_and_env_policy_overlay_contract() {
             ("CODEX_NETWORK_PROXY_ACTIVE".to_string(), "1".to_string(),),
         ])
     );
+    request.exec_server_shell_snapshot = Some(codex_exec_server::ShellSnapshotRequest {
+        scope_id: "attachment-1".to_string(),
+        shell: codex_exec_server::ShellInfo {
+            name: "bash".to_string(),
+            path: "/bin/bash".to_string(),
+        },
+    });
+    let mut snapshot_env = params.env;
+    snapshot_env.remove("PATH");
+    assert_eq!(params_for_request(&request).env, snapshot_env);
+    request.exec_server_shell_snapshot = None;
+
     request.exec_server_sandbox = Some(
         codex_exec_server::FileSystemSandboxContext::from_permission_profile(permission_profile),
     );

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

@@ -0,0 +1,40 @@
+use codex_exec_server::ShellInfo;
+use codex_exec_server::ShellSnapshotRequest;
+use codex_features::Feature;
+use codex_tools::UnifiedExecShellMode;
+use codex_utils_path_uri::PathUri;
+
+use super::ExecCommandRequest;
+use super::UnifiedExecContext;
+use crate::shell::ShellType;
+
+pub(super) fn shell_snapshot_request(
+    request: &ExecCommandRequest,
+    cwd: &PathUri,
+    context: &UnifiedExecContext,
+) -> Option<ShellSnapshotRequest> {
+    if !context.session.features().enabled(Feature::ShellSnapshotV2)
+        || !request.turn_environment.shell_snapshot_v2_supported
+        || request.turn_environment.selection.cwd != *cwd
+        || !matches!(request.shell_mode, UnifiedExecShellMode::Direct)
+        || !matches!(
+            request.shell_type,
+            ShellType::Bash | ShellType::Zsh | ShellType::Sh
+        )
+        || request.command.get(1).is_none_or(|flag| flag != "-lc")
+    {
+        return None;
+    }
+
+    Some(ShellSnapshotRequest {
+        scope_id: format!(
+            "{}:{}",
+            context.session.thread_id(),
+            request.turn_environment.selection.environment_id
+        ),
+        shell: ShellInfo {
+            name: request.shell_type.name().to_string(),
+            path: request.command.first()?.clone(),
+        },
+    })
+}