Protect shell snapshots when credential brokerage is enabled (#43909)

## Why

Shell startup files can introduce credentials or overwrite brokered dummy values. Snapshot capture and replay need to respect the command's sandbox and environment policy without persisting real credentials.

## What changed

- Capture and validate brokered snapshots per command using its requested shell, working directory, login mode, sandbox, and environment policy.
- Discover shell-initialized credentials and prepare snapshot exports for brokered replay. Preserve shell functions and aliases while preventing startup files from overwriting dummy credentials during replay.
- Rebuild snapshots as credential broker state changes and retain each command's snapshot for the lifetime of its process.
- Reject brokered POSIX shell execution when a protected snapshot cannot be created, and reject brokerage with `shell_zsh_fork`. Restore real credentials and remove brokerage markers for unsupported PowerShell and Cmd execution; escalated commands bypass protected snapshots.

## Testing

Add regression coverage for credential discovery and redaction, environment filtering, readonly exports, shell function and heredoc replay, command working directories and login modes, sandbox enforcement, and missing or disabled snapshots.

GitOrigin-RevId: 685356cb9e80f11cf69d053e4d239ee6ccf43a9e

codex-rs/core/src/environment_selection.rs

@@ -15,6 +15,7 @@ use codex_exec_server::ExecutorFileSystem;
 use codex_exec_server::SelectedCapabilityRootsStatus;
 use codex_protocol::capabilities::CapabilityRootLocation;
 use codex_protocol::capabilities::SelectedCapabilityRoot;
+use codex_protocol::config_types::ShellEnvironmentPolicy;
 use codex_protocol::models::PermissionProfile;
 use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::EnvironmentConfig;
@@ -37,6 +38,7 @@ use crate::session::turn_context::ShellSnapshotTask;
 use crate::session::turn_context::TurnEnvironment;
 use crate::shell::Shell;
 use crate::shell_snapshot::ShellSnapshot;
+use crate::shell_snapshot::SnapshotCredentialBrokerState;
 
 /// Records whether a normalized config should follow later thread setting updates.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
@@ -142,6 +144,7 @@ struct ResolvedEnvironment {
     executor_platform_os: Option<String>,
     temporary_directories: Option<Vec<PathUri>>,
     shell_snapshot: ShellSnapshotTask,
+    shell_snapshot_builder: ShellSnapshot,
     shell_snapshot_v2_supported: bool,
     installed_config: Option<EnvironmentConfig>,
 }
@@ -256,6 +259,13 @@ impl ThreadEnvironments {
                 let selection = environment.selection;
                 let config_origin = environment.config_origin;
                 let selected_environment = Arc::clone(&environment.environment);
+                let inherited_snapshot = if !selected_environment.is_remote()
+                    && shell_snapshot.should_rebuild_inherited()
+                {
+                    futures::future::ready(None).boxed().shared()
+                } else {
+                    environment.shell_snapshot
+                };
                 let resolution: TurnEnvironmentResolution =
                     futures::future::ready(Ok(ResolvedEnvironment {
                         environment: environment.environment,
@@ -263,8 +273,11 @@ impl ThreadEnvironments {
                         user_home_dir: environment.user_home_dir,
                         executor_platform_os: environment.executor_platform_os,
                         temporary_directories: environment.temporary_directories,
-                        shell_snapshot: environment.shell_snapshot,
-                        shell_snapshot_v2_supported: environment.shell_snapshot_v2_supported,
+                        shell_snapshot: inherited_snapshot,
+                        shell_snapshot_builder: shell_snapshot.clone(),
+                        shell_snapshot_v2_supported: environment.shell_snapshot_v2_supported
+                            && (selected_environment.is_remote()
+                                || !shell_snapshot.should_rebuild_inherited()),
                         installed_config: None,
                     }))
                     .boxed()
@@ -292,6 +305,36 @@ impl ThreadEnvironments {
         }
     }
 
+    fn start_shell_snapshot_task(
+        shell_snapshot: ShellSnapshot,
+        environment: Arc<Environment>,
+        cwd: PathUri,
+        shell: Option<Shell>,
+    ) -> ShellSnapshotTask {
+        // Protected snapshots are captured by the command, using its sandbox.
+        if shell_snapshot.should_rebuild_inherited() {
+            return futures::future::ready(None).boxed().shared();
+        }
+        let shell_snapshot = shell_snapshot
+            .build(
+                environment,
+                cwd,
+                shell,
+                /*allow_login_shell*/ true,
+                ShellEnvironmentPolicy::default(),
+                /*sandbox*/ None,
+            )
+            .boxed()
+            .shared();
+        drop(tokio::spawn(
+            shell_snapshot
+                .clone()
+                .in_current_span()
+                .with_current_subscriber(),
+        ));
+        shell_snapshot
+    }
+
     pub(crate) fn update_selections(
         &self,
         environments: &[TurnEnvironmentSelection],
@@ -482,6 +525,47 @@ impl ThreadEnvironments {
         self.environments.store(Arc::new(environments));
     }
 
+    pub(crate) fn set_snapshot_credential_broker(&self, state: SnapshotCredentialBrokerState) {
+        if !self.shell_snapshot.set_credential_broker(state) {
+            return;
+        }
+
+        let mut environments = Vec::clone(&self.environments.load());
+        let mut changed = false;
+        for selected in &mut environments {
+            if !selected.environment.is_remote()
+                && let Some(Ok(resolved)) = selected.resolution.clone().now_or_never()
+            {
+                self.restart_shell_snapshot(selected, resolved);
+                changed = true;
+            }
+        }
+        if changed {
+            self.environments.store(Arc::new(environments));
+        }
+    }
+
+    fn restart_shell_snapshot(
+        &self,
+        selected: &mut SelectedTurnEnvironment,
+        resolved: ResolvedEnvironment,
+    ) {
+        let shell_snapshot = Self::start_shell_snapshot_task(
+            self.shell_snapshot.clone(),
+            Arc::clone(&resolved.environment),
+            selected.selection.cwd.clone(),
+            resolved.shell.clone(),
+        );
+        selected.resolution = futures::future::ready(Ok(ResolvedEnvironment {
+            shell_snapshot,
+            shell_snapshot_v2_supported: cfg!(unix)
+                && !self.shell_snapshot.should_rebuild_inherited(),
+            ..resolved
+        }))
+        .boxed()
+        .shared();
+    }
+
     /// Combines persisted thread roots with installed attachment roots, keeping
     /// thread roots first and hiding attachments that are not ready yet.
     pub(crate) fn inspect_selected_capability_roots(
@@ -667,13 +751,15 @@ impl ThreadEnvironments {
                 cfg!(unix),
             )
         };
-        let task = shell_snapshot
-            .build(Arc::clone(&environment), selection.cwd, shell.clone())
-            .boxed()
-            .shared();
-        drop(tokio::spawn(
-            task.clone().in_current_span().with_current_subscriber(),
-        ));
+        let shell_snapshot_builder = shell_snapshot.clone();
+        let task = Self::start_shell_snapshot_task(
+            shell_snapshot,
+            Arc::clone(&environment),
+            selection.cwd,
+            shell.clone(),
+        );
+        let shell_snapshot_v2_supported = snapshot_v2
+            && (environment.is_remote() || !shell_snapshot_builder.should_rebuild_inherited());
         Ok(ResolvedEnvironment {
             environment,
             shell,
@@ -681,7 +767,8 @@ impl ThreadEnvironments {
             executor_platform_os,
             temporary_directories: temporary_dirs,
             shell_snapshot: task,
-            shell_snapshot_v2_supported: snapshot_v2,
+            shell_snapshot_builder,
+            shell_snapshot_v2_supported,
             installed_config,
         })
     }
@@ -762,6 +849,8 @@ impl TurnEnvironmentState {
                 );
                 turn_environment.executor_platform_os = environment.executor_platform_os;
                 turn_environment.shell_snapshot = environment.shell_snapshot;
+                turn_environment.shell_snapshot_builder =
+                    Some(Box::new(environment.shell_snapshot_builder));
                 turn_environment.shell_snapshot_v2_supported =
                     environment.shell_snapshot_v2_supported;
                 turn_environment.user_home_dir = environment.user_home_dir;