refactor: prepare unified exec for zsh-fork backend

Author: bolinfest

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(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>,
@@ -809,6 +913,7 @@ impl CoreShellCommandExecutor {
 
 #[derive(Debug, Eq, PartialEq)]
 struct ParsedShellCommand {
+    program: String,
     script: String,
     login: bool,
 }
@@ -817,12 +922,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,
         }

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

@@ -0,0 +1,115 @@
+use super::ShellRequest;
+use crate::exec::ExecToolCallOutput;
+use crate::sandboxing::ExecRequest;
+use crate::tools::runtimes::unified_exec::UnifiedExecRequest;
+use crate::tools::sandboxing::SandboxAttempt;
+use crate::tools::sandboxing::ToolCtx;
+use crate::tools::sandboxing::ToolError;
+use crate::unified_exec::SpawnLifecycleHandle;
+
+pub(crate) struct PreparedUnifiedExecSpawn {
+    pub(crate) exec_request: ExecRequest,
+    pub(crate) spawn_lifecycle: SpawnLifecycleHandle,
+}
+
+/// Runs the zsh-fork shell-command backend when this request should be handled
+/// by executable-level escalation instead of the default shell runtime.
+///
+/// Returns `Ok(None)` when the current platform or request shape should fall
+/// back to the normal shell-command path.
+pub(crate) async fn maybe_run_shell_command(
+    req: &ShellRequest,
+    attempt: &SandboxAttempt<'_>,
+    ctx: &ToolCtx,
+    command: &[String],
+) -> Result<Option<ExecToolCallOutput>, ToolError> {
+    imp::maybe_run_shell_command(req, attempt, ctx, command).await
+}
+
+/// Prepares unified exec to launch through the zsh-fork backend when the
+/// request matches a wrapped `zsh -c/-lc` command on a supported platform.
+///
+/// Returns the transformed `ExecRequest` plus a spawn lifecycle that keeps the
+/// escalation server alive for the session and performs post-spawn cleanup.
+/// Returns `Ok(None)` when unified exec should use its normal spawn path.
+pub(crate) async fn maybe_prepare_unified_exec(
+    req: &UnifiedExecRequest,
+    attempt: &SandboxAttempt<'_>,
+    ctx: &ToolCtx,
+    exec_request: ExecRequest,
+) -> Result<Option<PreparedUnifiedExecSpawn>, ToolError> {
+    imp::maybe_prepare_unified_exec(req, attempt, ctx, exec_request).await
+}
+
+#[cfg(unix)]
+mod imp {
+    use super::*;
+    use crate::tools::runtimes::shell::unix_escalation;
+    use crate::unified_exec::SpawnLifecycle;
+    use codex_shell_escalation::EscalationSession;
+
+    #[derive(Debug)]
+    struct ZshForkSpawnLifecycle {
+        escalation_session: EscalationSession,
+    }
+
+    impl SpawnLifecycle for ZshForkSpawnLifecycle {
+        fn after_spawn(&mut self) {
+            self.escalation_session.close_client_socket();
+        }
+    }
+
+    pub(super) async fn maybe_run_shell_command(
+        req: &ShellRequest,
+        attempt: &SandboxAttempt<'_>,
+        ctx: &ToolCtx,
+        command: &[String],
+    ) -> Result<Option<ExecToolCallOutput>, ToolError> {
+        unix_escalation::try_run_zsh_fork(req, attempt, ctx, command).await
+    }
+
+    pub(super) async fn maybe_prepare_unified_exec(
+        req: &UnifiedExecRequest,
+        attempt: &SandboxAttempt<'_>,
+        ctx: &ToolCtx,
+        exec_request: ExecRequest,
+    ) -> Result<Option<PreparedUnifiedExecSpawn>, ToolError> {
+        let Some(prepared) =
+            unix_escalation::prepare_unified_exec_zsh_fork(req, attempt, ctx, exec_request).await?
+        else {
+            return Ok(None);
+        };
+
+        Ok(Some(PreparedUnifiedExecSpawn {
+            exec_request: prepared.exec_request,
+            spawn_lifecycle: Box::new(ZshForkSpawnLifecycle {
+                escalation_session: prepared.escalation_session,
+            }),
+        }))
+    }
+}
+
+#[cfg(not(unix))]
+mod imp {
+    use super::*;
+
+    pub(super) async fn maybe_run_shell_command(
+        req: &ShellRequest,
+        attempt: &SandboxAttempt<'_>,
+        ctx: &ToolCtx,
+        command: &[String],
+    ) -> Result<Option<ExecToolCallOutput>, ToolError> {
+        let _ = (req, attempt, ctx, command);
+        Ok(None)
+    }
+
+    pub(super) async fn maybe_prepare_unified_exec(
+        req: &UnifiedExecRequest,
+        attempt: &SandboxAttempt<'_>,
+        ctx: &ToolCtx,
+        exec_request: ExecRequest,
+    ) -> Result<Option<PreparedUnifiedExecSpawn>, ToolError> {
+        let _ = (req, attempt, ctx, exec_request);
+        Ok(None)
+    }
+}