refactor: prepare unified exec for zsh-fork backend

Author: bolinfest

codex-rs/core/src/exec.rs

@@ -235,6 +235,15 @@ pub(crate) async fn execute_exec_env(
     env: ExecRequest,
     sandbox_policy: &SandboxPolicy,
     stdout_stream: Option<StdoutStream>,
+) -> Result<ExecToolCallOutput> {
+    execute_exec_env_after_spawn(env, sandbox_policy, stdout_stream, Box::new(|| {})).await
+}
+
+pub(crate) async fn execute_exec_env_after_spawn(
+    env: ExecRequest,
+    sandbox_policy: &SandboxPolicy,
+    stdout_stream: Option<StdoutStream>,
+    after_spawn: Box<dyn FnOnce() + Send>,
 ) -> Result<ExecToolCallOutput> {
     let ExecRequest {
         command,
@@ -263,7 +272,7 @@ pub(crate) async fn execute_exec_env(
     };
 
     let start = Instant::now();
-    let raw_output_result = exec(params, sandbox, sandbox_policy, stdout_stream).await;
+    let raw_output_result = exec(params, sandbox, sandbox_policy, stdout_stream, after_spawn).await;
     let duration = start.elapsed();
     finalize_exec_result(raw_output_result, sandbox, duration)
 }
@@ -693,6 +702,7 @@ async fn exec(
     sandbox: SandboxType,
     sandbox_policy: &SandboxPolicy,
     stdout_stream: Option<StdoutStream>,
+    after_spawn: Box<dyn FnOnce() + Send>,
 ) -> Result<RawExecToolCallOutput> {
     #[cfg(target_os = "windows")]
     if sandbox == SandboxType::WindowsRestrictedToken
@@ -738,6 +748,7 @@ async fn exec(
         env,
     })
     .await?;
+    after_spawn();
     consume_truncated_output(child, expiration, stdout_stream).await
 }
 
@@ -1136,6 +1147,7 @@ mod tests {
             SandboxType::None,
             &SandboxPolicy::new_read_only_policy(),
             None,
+            Box::new(|| {}),
         )
         .await?;
         assert!(output.timed_out);

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

@@ -11,6 +11,7 @@ use crate::exec::ExecToolCallOutput;
 use crate::exec::SandboxType;
 use crate::exec::StdoutStream;
 use crate::exec::execute_exec_env;
+use crate::exec::execute_exec_env_after_spawn;
 use crate::landlock::allow_network_for_proxy;
 use crate::landlock::create_linux_sandbox_command_args;
 use crate::protocol::SandboxPolicy;
@@ -425,6 +426,15 @@ pub async fn execute_env(
     execute_exec_env(env, &effective_policy, stdout_stream).await
 }
 
+pub async fn execute_env_after_spawn(
+    env: ExecRequest,
+    stdout_stream: Option<StdoutStream>,
+    after_spawn: Box<dyn FnOnce() + Send>,
+) -> crate::error::Result<ExecToolCallOutput> {
+    let effective_policy = env.sandbox_policy.clone();
+    execute_exec_env_after_spawn(env, &effective_policy, stdout_stream, after_spawn).await
+}
+
 #[cfg(test)]
 mod tests {
     use super::SandboxManager;

codex-rs/core/src/tools/runtimes/shell.rs

@@ -5,7 +5,8 @@ Executes shell requests under the orchestrator: asks for approval when needed,
 builds a CommandSpec, and runs it under the current SandboxAttempt.
 */
 #[cfg(unix)]
-mod unix_escalation;
+pub(crate) mod unix_escalation;
+pub(crate) mod zsh_fork_backend;
 
 use crate::command_canonicalization::canonicalize_command_for_approval;
 use crate::exec::ExecToolCallOutput;
@@ -80,7 +81,6 @@ pub(crate) enum ShellRuntimeBackend {
 
 #[derive(Default)]
 pub struct ShellRuntime {
-    #[cfg_attr(not(unix), allow(dead_code))]
     backend: ShellRuntimeBackend,
 }
 
@@ -215,9 +215,8 @@ impl ToolRuntime<ShellRequest, ExecToolCallOutput> for ShellRuntime {
             command
         };
 
-        #[cfg(unix)]
         if self.backend == ShellRuntimeBackend::ShellCommandZshFork {
-            match unix_escalation::try_run_zsh_fork(req, attempt, ctx, &command).await? {
+            match zsh_fork_backend::maybe_run_shell_command(req, attempt, ctx, &command).await? {
                 Some(out) => return Ok(out),
                 None => {
                     tracing::warn!(

codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs

@@ -6,6 +6,7 @@ use crate::exec::ExecToolCallOutput;
 use crate::exec::SandboxType;
 use crate::exec::is_likely_sandbox_denied;
 use crate::features::Feature;
+use crate::sandboxing::ExecRequest;
 use crate::sandboxing::SandboxPermissions;
 use crate::shell::ShellType;
 use crate::skills::SkillMetadata;
@@ -36,6 +37,7 @@ use codex_shell_escalation::EscalationDecision;
 use codex_shell_escalation::EscalationExecution;
 use codex_shell_escalation::EscalationPermissions;
 use codex_shell_escalation::EscalationPolicy;
+use codex_shell_escalation::EscalationSession;
 use codex_shell_escalation::ExecParams;
 use codex_shell_escalation::ExecResult;
 use codex_shell_escalation::Permissions as EscalatedPermissions;
@@ -51,6 +53,11 @@ use tokio::sync::RwLock;
 use tokio_util::sync::CancellationToken;
 use uuid::Uuid;
 
+pub(crate) struct PreparedUnifiedExecZshFork {
+    pub(crate) exec_request: ExecRequest,
+    pub(crate) escalation_session: EscalationSession,
+}
+
 pub(super) async fn try_run_zsh_fork(
     req: &ShellRequest,
     attempt: &SandboxAttempt<'_>,
@@ -95,7 +102,7 @@ pub(super) async fn try_run_zsh_fork(
         justification,
         arg0,
     } = sandbox_exec_request;
-    let ParsedShellCommand { script, login } = extract_shell_script(&command)?;
+    let ParsedShellCommand { script, login, .. } = extract_shell_script(&command)?;
     let effective_timeout = Duration::from_millis(
         req.timeout_ms
             .unwrap_or(crate::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS),
@@ -172,6 +179,103 @@ pub(super) async fn try_run_zsh_fork(
     map_exec_result(attempt.sandbox, exec_result).map(Some)
 }
 
+pub(crate) async fn prepare_unified_exec_zsh_fork(
+    req: &crate::tools::runtimes::unified_exec::UnifiedExecRequest,
+    attempt: &SandboxAttempt<'_>,
+    ctx: &ToolCtx,
+    exec_request: ExecRequest,
+) -> Result<Option<PreparedUnifiedExecZshFork>, ToolError> {
+    let Some(shell_zsh_path) = ctx.session.services.shell_zsh_path.as_ref() else {
+        tracing::warn!("ZshFork backend specified, but shell_zsh_path is not configured.");
+        return Ok(None);
+    };
+    if !ctx.session.features().enabled(Feature::ShellZshFork) {
+        tracing::warn!("ZshFork backend specified, but ShellZshFork feature is not enabled.");
+        return Ok(None);
+    }
+    if !matches!(ctx.session.user_shell().shell_type, ShellType::Zsh) {
+        tracing::warn!("ZshFork backend specified, but user shell is not Zsh.");
+        return Ok(None);
+    }
+
+    let parsed = match extract_shell_script(&exec_request.command) {
+        Ok(parsed) => parsed,
+        Err(err) => {
+            tracing::warn!("ZshFork unified exec fallback: {err:?}");
+            return Ok(None);
+        }
+    };
+    if parsed.program != shell_zsh_path.to_string_lossy() {
+        tracing::warn!(
+            "ZshFork backend specified, but unified exec command targets `{}` instead of `{}`.",
+            parsed.program,
+            shell_zsh_path.display(),
+        );
+        return Ok(None);
+    }
+
+    let exec_policy = Arc::new(RwLock::new(
+        ctx.session.services.exec_policy.current().as_ref().clone(),
+    ));
+    let command_executor = CoreShellCommandExecutor {
+        command: exec_request.command.clone(),
+        cwd: exec_request.cwd.clone(),
+        sandbox_policy: exec_request.sandbox_policy.clone(),
+        sandbox: exec_request.sandbox,
+        env: exec_request.env.clone(),
+        network: exec_request.network.clone(),
+        windows_sandbox_level: exec_request.windows_sandbox_level,
+        sandbox_permissions: exec_request.sandbox_permissions,
+        justification: exec_request.justification.clone(),
+        arg0: exec_request.arg0.clone(),
+        sandbox_policy_cwd: ctx.turn.cwd.clone(),
+        macos_seatbelt_profile_extensions: ctx
+            .turn
+            .config
+            .permissions
+            .macos_seatbelt_profile_extensions
+            .clone(),
+        codex_linux_sandbox_exe: ctx.turn.codex_linux_sandbox_exe.clone(),
+        use_linux_sandbox_bwrap: ctx.turn.features.enabled(Feature::UseLinuxSandboxBwrap),
+    };
+    let main_execve_wrapper_exe = ctx
+        .session
+        .services
+        .main_execve_wrapper_exe
+        .clone()
+        .ok_or_else(|| {
+            ToolError::Rejected(
+                "zsh fork feature enabled, but execve wrapper is not configured".to_string(),
+            )
+        })?;
+    let escalation_policy = CoreShellActionProvider {
+        policy: Arc::clone(&exec_policy),
+        session: Arc::clone(&ctx.session),
+        turn: Arc::clone(&ctx.turn),
+        call_id: ctx.call_id.clone(),
+        approval_policy: ctx.turn.approval_policy.value(),
+        sandbox_policy: attempt.policy.clone(),
+        sandbox_permissions: req.sandbox_permissions,
+        prompt_permissions: req.additional_permissions.clone(),
+        stopwatch: Stopwatch::unlimited(),
+    };
+
+    let escalate_server = EscalateServer::new(
+        shell_zsh_path.clone(),
+        main_execve_wrapper_exe,
+        escalation_policy,
+    );
+    let escalation_session = escalate_server
+        .start_session(CancellationToken::new(), Arc::new(command_executor))
+        .map_err(|err| ToolError::Rejected(err.to_string()))?;
+    let mut exec_request = exec_request;
+    exec_request.env.extend(escalation_session.env().clone());
+    Ok(Some(PreparedUnifiedExecZshFork {
+        exec_request,
+        escalation_session,
+    }))
+}
+
 struct CoreShellActionProvider {
     policy: Arc<RwLock<Policy>>,
     session: Arc<crate::codex::Session>,
@@ -648,17 +752,20 @@ impl ShellCommandExecutor for CoreShellCommandExecutor {
         &self,
         _command: Vec<String>,
         _cwd: PathBuf,
-        env: HashMap<String, String>,
+        env_overlay: HashMap<String, String>,
         cancel_rx: CancellationToken,
+        after_spawn: Option<Box<dyn FnOnce() + Send>>,
     ) -> anyhow::Result<ExecResult> {
         let mut exec_env = self.env.clone();
+        // `env_overlay` comes from `EscalationSession::env()`, so merge only the
+        // wrapper/socket variables into the base shell environment.
         for var in ["CODEX_ESCALATE_SOCKET", "EXEC_WRAPPER", "BASH_EXEC_WRAPPER"] {
-            if let Some(value) = env.get(var) {
+            if let Some(value) = env_overlay.get(var) {
                 exec_env.insert(var.to_string(), value.clone());
             }
         }
 
-        let result = crate::sandboxing::execute_env(
+        let result = crate::sandboxing::execute_env_after_spawn(
             crate::sandboxing::ExecRequest {
                 command: self.command.clone(),
                 cwd: self.cwd.clone(),
@@ -673,6 +780,7 @@ impl ShellCommandExecutor for CoreShellCommandExecutor {
                 arg0: self.arg0.clone(),
             },
             None,
+            after_spawn.unwrap_or_else(|| Box::new(|| {})),
         )
         .await?;
 
@@ -809,6 +917,7 @@ impl CoreShellCommandExecutor {
 
 #[derive(Debug, Eq, PartialEq)]
 struct ParsedShellCommand {
+    program: String,
     script: String,
     login: bool,
 }
@@ -817,12 +926,20 @@ fn extract_shell_script(command: &[String]) -> Result<ParsedShellCommand, ToolEr
     // Commands reaching zsh-fork can be wrapped by environment/sandbox helpers, so
     // we search for the first `-c`/`-lc` triple anywhere in the argv rather
     // than assuming it is the first positional form.
-    if let Some((script, login)) = command.windows(3).find_map(|parts| match parts {
-        [_, flag, script] if flag == "-c" => Some((script.to_owned(), false)),
-        [_, flag, script] if flag == "-lc" => Some((script.to_owned(), true)),
+    if let Some((program, script, login)) = command.windows(3).find_map(|parts| match parts {
+        [program, flag, script] if flag == "-c" => {
+            Some((program.to_owned(), script.to_owned(), false))
+        }
+        [program, flag, script] if flag == "-lc" => {
+            Some((program.to_owned(), script.to_owned(), true))
+        }
         _ => None,
     }) {
-        return Ok(ParsedShellCommand { script, login });
+        return Ok(ParsedShellCommand {
+            program,
+            script,
+            login,
+        });
     }
 
     Err(ToolError::Rejected(

codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs

@@ -64,13 +64,15 @@ fn extract_shell_script_preserves_login_flag() {
     assert_eq!(
         extract_shell_script(&["/bin/zsh".into(), "-lc".into(), "echo hi".into()]).unwrap(),
         ParsedShellCommand {
+            program: "/bin/zsh".to_string(),
             script: "echo hi".to_string(),
             login: true,
         }
     );
     assert_eq!(
         extract_shell_script(&["/bin/zsh".into(), "-c".into(), "echo hi".into()]).unwrap(),
         ParsedShellCommand {
+            program: "/bin/zsh".to_string(),
             script: "echo hi".to_string(),
             login: false,
         }
@@ -89,6 +91,7 @@ fn extract_shell_script_supports_wrapped_command_prefixes() {
         ])
         .unwrap(),
         ParsedShellCommand {
+            program: "/bin/zsh".to_string(),
             script: "echo hello".to_string(),
             login: true,
         }
@@ -105,6 +108,7 @@ fn extract_shell_script_supports_wrapped_command_prefixes() {
         ])
         .unwrap(),
         ParsedShellCommand {
+            program: "/bin/zsh".to_string(),
             script: "pwd".to_string(),
             login: false,
         }