feat: shell snapshotting (#7641)

codex-rs/core/src/shell.rs

@@ -1,6 +1,9 @@
 use serde::Deserialize;
 use serde::Serialize;
 use std::path::PathBuf;
+use std::sync::Arc;
+
+use crate::shell_snapshot::ShellSnapshot;
 
 #[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
 pub enum ShellType {
@@ -15,6 +18,8 @@ pub enum ShellType {
 pub struct Shell {
     pub(crate) shell_type: ShellType,
     pub(crate) shell_path: PathBuf,
+    #[serde(skip_serializing, skip_deserializing, default)]
+    pub(crate) shell_snapshot: Option<Arc<ShellSnapshot>>,
 }
 
 impl Shell {
@@ -58,6 +63,33 @@ impl Shell {
             }
         }
     }
+
+    pub(crate) fn wrap_command_with_snapshot(&self, command: &[String]) -> Vec<String> {
+        let Some(snapshot) = &self.shell_snapshot else {
+            return command.to_vec();
+        };
+
+        if command.is_empty() {
+            return command.to_vec();
+        }
+
+        match self.shell_type {
+            ShellType::Zsh | ShellType::Bash | ShellType::Sh => {
+                let mut args = self.derive_exec_args(". \"$0\" && exec \"$@\"", false);
+                args.push(snapshot.path.to_string_lossy().to_string());
+                args.extend_from_slice(command);
+                args
+            }
+            ShellType::PowerShell => {
+                let mut args =
+                    self.derive_exec_args("param($snapshot) . $snapshot; & @args", false);
+                args.push(snapshot.path.to_string_lossy().to_string());
+                args.extend_from_slice(command);
+                args
+            }
+            ShellType::Cmd => command.to_vec(),
+        }
+    }
 }
 
 #[cfg(unix)]
@@ -134,6 +166,7 @@ fn get_zsh_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::Zsh,
         shell_path,
+        shell_snapshot: None,
     })
 }
 
@@ -143,6 +176,7 @@ fn get_bash_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::Bash,
         shell_path,
+        shell_snapshot: None,
     })
 }
 
@@ -152,6 +186,7 @@ fn get_sh_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::Sh,
         shell_path,
+        shell_snapshot: None,
     })
 }
 
@@ -167,6 +202,7 @@ fn get_powershell_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::PowerShell,
         shell_path,
+        shell_snapshot: None,
     })
 }
 
@@ -176,6 +212,7 @@ fn get_cmd_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::Cmd,
         shell_path,
+        shell_snapshot: None,
     })
 }
 
@@ -184,11 +221,13 @@ fn ultimate_fallback_shell() -> Shell {
         Shell {
             shell_type: ShellType::Cmd,
             shell_path: PathBuf::from("cmd.exe"),
+            shell_snapshot: None,
         }
     } else {
         Shell {
             shell_type: ShellType::Sh,
             shell_path: PathBuf::from("/bin/sh"),
+            shell_snapshot: None,
         }
     }
 }
@@ -413,6 +452,7 @@ mod tests {
         let test_bash_shell = Shell {
             shell_type: ShellType::Bash,
             shell_path: PathBuf::from("/bin/bash"),
+            shell_snapshot: None,
         };
         assert_eq!(
             test_bash_shell.derive_exec_args("echo hello", false),
@@ -426,6 +466,7 @@ mod tests {
         let test_zsh_shell = Shell {
             shell_type: ShellType::Zsh,
             shell_path: PathBuf::from("/bin/zsh"),
+            shell_snapshot: None,
         };
         assert_eq!(
             test_zsh_shell.derive_exec_args("echo hello", false),
@@ -439,6 +480,7 @@ mod tests {
         let test_powershell_shell = Shell {
             shell_type: ShellType::PowerShell,
             shell_path: PathBuf::from("pwsh.exe"),
+            shell_snapshot: None,
         };
         assert_eq!(
             test_powershell_shell.derive_exec_args("echo hello", false),
@@ -465,6 +507,7 @@ mod tests {
                 Shell {
                     shell_type: ShellType::Zsh,
                     shell_path: PathBuf::from(shell_path),
+                    shell_snapshot: None,
                 }
             );
         }