diff --git a/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs b/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs index bdbb283432..203194d616 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs @@ -11,7 +11,6 @@ use app_test_support::McpProcess; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_mock_responses_server_sequence_unchecked; -use app_test_support::create_shell_command_sse_response; use app_test_support::to_response; use codex_app_server_protocol::CommandAction; use codex_app_server_protocol::CommandExecutionApprovalDecision; @@ -35,9 +34,11 @@ use codex_core::features::Feature; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; +use serde_json::json; use std::collections::BTreeMap; use std::path::Path; use tempfile::TempDir; +use tokio::time::sleep; use tokio::time::timeout; #[cfg(windows)] @@ -61,10 +62,8 @@ async fn turn_start_shell_zsh_fork_executes_command_v2() -> Result<()> { }; eprintln!("using zsh path for zsh-fork test: {}", zsh_path.display()); - let responses = vec![create_shell_command_sse_response( - vec!["echo".to_string(), "hi".to_string()], - None, - Some(5000), + let responses = vec![create_zsh_fork_exec_command_sse_response( + "echo hi", "call-zsh-fork", )?]; let server = create_mock_responses_server_sequence(responses).await; @@ -74,7 +73,7 @@ async fn turn_start_shell_zsh_fork_executes_command_v2() -> Result<()> { "never", &BTreeMap::from([ (Feature::ShellZshFork, true), - (Feature::UnifiedExec, false), + (Feature::UnifiedExec, true), (Feature::ShellSnapshot, false), ]), &zsh_path, @@ -172,14 +171,8 @@ async fn turn_start_shell_zsh_fork_exec_approval_decline_v2() -> Result<()> { eprintln!("using zsh path for zsh-fork test: {}", zsh_path.display()); let responses = vec![ - create_shell_command_sse_response( - vec![ - "python3".to_string(), - "-c".to_string(), - "print(42)".to_string(), - ], - None, - Some(5000), + create_zsh_fork_exec_command_sse_response( + "python3 -c 'print(42)'", "call-zsh-fork-decline", )?, create_final_assistant_message_sse_response("done")?, @@ -191,7 +184,7 @@ async fn turn_start_shell_zsh_fork_exec_approval_decline_v2() -> Result<()> { "untrusted", &BTreeMap::from([ (Feature::ShellZshFork, true), - (Feature::UnifiedExec, false), + (Feature::UnifiedExec, true), (Feature::ShellSnapshot, false), ]), &zsh_path, @@ -307,14 +300,8 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { }; eprintln!("using zsh path for zsh-fork test: {}", zsh_path.display()); - let responses = vec![create_shell_command_sse_response( - vec![ - "python3".to_string(), - "-c".to_string(), - "print(42)".to_string(), - ], - None, - Some(5000), + let responses = vec![create_zsh_fork_exec_command_sse_response( + "python3 -c 'print(42)'", "call-zsh-fork-cancel", )?]; let server = create_mock_responses_server_sequence(responses).await; @@ -324,7 +311,7 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { "untrusted", &BTreeMap::from([ (Feature::ShellZshFork, true), - (Feature::UnifiedExec, false), + (Feature::UnifiedExec, true), (Feature::ShellSnapshot, false), ]), &zsh_path, @@ -422,6 +409,181 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { Ok(()) } +#[tokio::test] +async fn turn_start_shell_zsh_fork_interrupt_kills_approved_subcommand_v2() -> Result<()> { + skip_if_no_network!(Ok(())); + + let tmp = TempDir::new()?; + let codex_home = tmp.path().join("codex_home"); + std::fs::create_dir(&codex_home)?; + let workspace = tmp.path().join("workspace"); + std::fs::create_dir(&workspace)?; + let pid_file = workspace.join("approved-subcommand.pid"); + let pid_file_display = pid_file.display().to_string(); + assert!( + !pid_file_display.contains('\''), + "test workspace path should not contain single quotes: {pid_file_display}" + ); + + let Some(zsh_path) = find_test_zsh_path()? else { + eprintln!("skipping zsh fork interrupt cleanup test: no zsh executable found"); + return Ok(()); + }; + if !supports_exec_wrapper_intercept(&zsh_path) { + eprintln!( + "skipping zsh fork interrupt cleanup test: zsh does not support EXEC_WRAPPER intercepts ({})", + zsh_path.display() + ); + return Ok(()); + } + let zsh_path_display = zsh_path.display().to_string(); + eprintln!("using zsh path for zsh-fork test: {zsh_path_display}"); + + let shell_command = + format!("/bin/sh -c 'echo $$ > \"{pid_file_display}\" && exec /bin/sleep 100'"); + let tool_call_arguments = serde_json::to_string(&json!({ + "cmd": shell_command, + "yield_time_ms": 30_000, + }))?; + let response = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call( + "call-zsh-fork-interrupt-cleanup", + "exec_command", + &tool_call_arguments, + ), + responses::ev_completed("resp-1"), + ]); + let no_op_response = responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_completed("resp-2"), + ]); + let server = + create_mock_responses_server_sequence_unchecked(vec![response, no_op_response]).await; + create_config_toml( + &codex_home, + &server.uri(), + "untrusted", + &BTreeMap::from([ + (Feature::ShellZshFork, true), + (Feature::UnifiedExec, true), + (Feature::ShellSnapshot, false), + ]), + &zsh_path, + )?; + + let mut mcp = create_zsh_test_mcp_process(&codex_home, &workspace).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let start_id = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + cwd: Some(workspace.to_string_lossy().into_owned()), + ..Default::default() + }) + .await?; + let start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(start_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + + let turn_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "run the long-lived command".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(workspace.clone()), + approval_policy: Some(codex_app_server_protocol::AskForApproval::UnlessTrusted), + sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::WorkspaceWrite { + writable_roots: vec![workspace.clone().try_into()?], + read_only_access: codex_app_server_protocol::ReadOnlyAccess::FullAccess, + network_access: false, + exclude_tmpdir_env_var: false, + exclude_slash_tmp: false, + }), + model: Some("mock-model".to_string()), + effort: Some(codex_protocol::openai_models::ReasoningEffort::Medium), + summary: Some(codex_protocol::config_types::ReasoningSummary::Auto), + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), + ) + .await??; + let TurnStartResponse { turn } = to_response::(turn_resp)?; + + let mut saw_target_approval = false; + while !saw_target_approval { + let server_req = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::CommandExecutionRequestApproval { request_id, params } = server_req + else { + panic!("expected CommandExecutionRequestApproval request"); + }; + let approval_command = params.command.clone().unwrap_or_default(); + saw_target_approval = approval_command.contains("/bin/sh") + && approval_command.contains(&pid_file_display) + && !approval_command.contains(&zsh_path_display); + mcp.send_response( + request_id, + serde_json::to_value(CommandExecutionRequestApprovalResponse { + decision: CommandExecutionApprovalDecision::Accept, + })?, + ) + .await?; + } + + let pid = timeout(DEFAULT_READ_TIMEOUT, async { + loop { + if let Ok(contents) = std::fs::read_to_string(&pid_file) { + return Ok::(contents.trim().parse()?); + } + sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await??; + let still_running = std::process::Command::new("/bin/kill") + .args(["-0", &pid.to_string()]) + .status()? + .success(); + assert!( + still_running, + "expected approved intercepted subprocess pid {pid} to be running before interrupt" + ); + + mcp.interrupt_turn_and_wait_for_aborted( + thread.id.clone(), + turn.id.clone(), + DEFAULT_READ_TIMEOUT, + ) + .await?; + + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let still_running = std::process::Command::new("/bin/kill") + .args(["-0", &pid.to_string()]) + .status()? + .success(); + if !still_running { + return Ok::<(), anyhow::Error>(()); + } + sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await??; + + Ok(()) +} + #[tokio::test] async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() -> Result<()> { skip_if_no_network!(Ok(())); @@ -453,16 +615,15 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() first_file.display(), second_file.display() ); - let tool_call_arguments = serde_json::to_string(&serde_json::json!({ - "command": shell_command, - "workdir": serde_json::Value::Null, - "timeout_ms": 5000 + let tool_call_arguments = serde_json::to_string(&json!({ + "cmd": shell_command, + "yield_time_ms": 5000, }))?; let response = responses::sse(vec![ responses::ev_response_created("resp-1"), responses::ev_function_call( "call-zsh-fork-subcommand-decline", - "shell_command", + "exec_command", &tool_call_arguments, ), responses::ev_completed("resp-1"), @@ -483,7 +644,7 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() "untrusted", &BTreeMap::from([ (Feature::ShellZshFork, true), - (Feature::UnifiedExec, false), + (Feature::UnifiedExec, true), (Feature::ShellSnapshot, false), ]), &zsh_path, @@ -694,6 +855,21 @@ async fn create_zsh_test_mcp_process(codex_home: &Path, zdotdir: &Path) -> Resul McpProcess::new_with_env(codex_home, &[("ZDOTDIR", Some(zdotdir.as_str()))]).await } +fn create_zsh_fork_exec_command_sse_response( + command: &str, + call_id: &str, +) -> anyhow::Result { + let tool_call_arguments = serde_json::to_string(&json!({ + "cmd": command, + "yield_time_ms": 5000, + }))?; + Ok(responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, "exec_command", &tool_call_arguments), + responses::ev_completed("resp-1"), + ])) +} + fn create_config_toml( codex_home: &Path, server_uri: &str, diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index 101d9309fd..90a4651428 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/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 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!( diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 36f63be4fa..cc0c28de10 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/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; @@ -47,10 +49,14 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; -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,14 +101,11 @@ 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), ); - let exec_policy = Arc::new(RwLock::new( - ctx.session.services.exec_policy.current().as_ref().clone(), - )); let command_executor = CoreShellCommandExecutor { command, cwd: sandbox_cwd, @@ -147,7 +150,6 @@ pub(super) async fn try_run_zsh_fork( let stopwatch = Stopwatch::new(effective_timeout); let cancel_token = stopwatch.cancellation_token(); 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(), @@ -172,8 +174,100 @@ 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, 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 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 { + 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>, session: Arc, turn: Arc, call_id: String, @@ -493,18 +587,16 @@ impl EscalationPolicy for CoreShellActionProvider { .await; } - let evaluation = { - let policy = self.policy.read().await; - evaluate_intercepted_exec_policy( - &policy, - program, - argv, - self.approval_policy, - &self.sandbox_policy, - self.sandbox_permissions, - ENABLE_INTERCEPTED_EXEC_POLICY_SHELL_WRAPPER_PARSING, - ) - }; + let policy = self.session.services.exec_policy.current(); + let evaluation = evaluate_intercepted_exec_policy( + policy.as_ref(), + program, + argv, + self.approval_policy, + &self.sandbox_policy, + self.sandbox_permissions, + ENABLE_INTERCEPTED_EXEC_POLICY_SHELL_WRAPPER_PARSING, + ); // When true, means the Evaluation was due to *.rules, not the // fallback function. let decision_driven_by_policy = @@ -809,6 +901,7 @@ impl CoreShellCommandExecutor { #[derive(Debug, Eq, PartialEq)] struct ParsedShellCommand { + program: String, script: String, login: bool, } @@ -817,12 +910,20 @@ fn extract_shell_script(command: &[String]) -> Result 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( diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs index b88f5b0542..ad663a3fed 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs @@ -64,6 +64,7 @@ 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, } @@ -71,6 +72,7 @@ fn extract_shell_script_preserves_login_flag() { 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, } diff --git a/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs b/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs new file mode 100644 index 0000000000..e88f0caa74 --- /dev/null +++ b/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, 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, 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, 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, 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, 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, ToolError> { + let _ = (req, attempt, ctx, exec_request); + Ok(None) + } +} diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index f4d823bf0e..d557a4340b 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -16,6 +16,7 @@ use crate::tools::network_approval::NetworkApprovalMode; use crate::tools::network_approval::NetworkApprovalSpec; use crate::tools::runtimes::build_command_spec; use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot; +use crate::tools::runtimes::shell::zsh_fork_backend; use crate::tools::sandboxing::Approvable; use crate::tools::sandboxing::ApprovalCtx; use crate::tools::sandboxing::ExecApprovalRequirement; @@ -28,6 +29,8 @@ use crate::tools::sandboxing::ToolError; use crate::tools::sandboxing::ToolRuntime; use crate::tools::sandboxing::sandbox_override_for_first_attempt; use crate::tools::sandboxing::with_cached_approval; +use crate::tools::spec::UnifiedExecBackendConfig; +use crate::unified_exec::NoopSpawnLifecycle; use crate::unified_exec::UnifiedExecError; use crate::unified_exec::UnifiedExecProcess; use crate::unified_exec::UnifiedExecProcessManager; @@ -63,11 +66,12 @@ pub struct UnifiedExecApprovalKey { pub struct UnifiedExecRuntime<'a> { manager: &'a UnifiedExecProcessManager, + backend: UnifiedExecBackendConfig, } impl<'a> UnifiedExecRuntime<'a> { - pub fn new(manager: &'a UnifiedExecProcessManager) -> Self { - Self { manager } + pub fn new(manager: &'a UnifiedExecProcessManager, backend: UnifiedExecBackendConfig) -> Self { + Self { manager, backend } } } @@ -184,6 +188,47 @@ impl<'a> ToolRuntime for UnifiedExecRunt if let Some(network) = req.network.as_ref() { network.apply_to_env(&mut env); } + if self.backend == UnifiedExecBackendConfig::ZshFork { + let spec = build_command_spec( + &command, + &req.cwd, + &env, + ExecExpiration::DefaultTimeout, + req.sandbox_permissions, + req.additional_permissions.clone(), + req.justification.clone(), + ) + .map_err(|_| ToolError::Rejected("missing command line for PTY".to_string()))?; + let exec_env = attempt + .env_for(spec, req.network.as_ref()) + .map_err(|err| ToolError::Codex(err.into()))?; + match zsh_fork_backend::maybe_prepare_unified_exec(req, attempt, ctx, exec_env).await? { + Some(prepared) => { + return self + .manager + .open_session_with_exec_env( + &prepared.exec_request, + req.tty, + prepared.spawn_lifecycle, + ) + .await + .map_err(|err| match err { + UnifiedExecError::SandboxDenied { output, .. } => { + ToolError::Codex(CodexErr::Sandbox(SandboxErr::Denied { + output: Box::new(output), + network_policy_decision: None, + })) + } + other => ToolError::Rejected(other.to_string()), + }); + } + None => { + tracing::warn!( + "UnifiedExec ZshFork backend specified, but conditions for using it were not met, falling back to direct execution", + ); + } + } + } let spec = build_command_spec( &command, &req.cwd, @@ -198,7 +243,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt .env_for(spec, req.network.as_ref()) .map_err(|err| ToolError::Codex(err.into()))?; self.manager - .open_session_with_exec_env(&exec_env, req.tty) + .open_session_with_exec_env(&exec_env, req.tty, Box::new(NoopSpawnLifecycle)) .await .map_err(|err| match err { UnifiedExecError::SandboxDenied { output, .. } => { diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index dcdb48500e..3b60923736 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -42,10 +42,17 @@ pub enum ShellCommandBackendConfig { ZshFork, } +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum UnifiedExecBackendConfig { + Direct, + ZshFork, +} + #[derive(Debug, Clone)] pub(crate) struct ToolsConfig { pub shell_type: ConfigShellToolType, shell_command_backend: ShellCommandBackendConfig, + pub unified_exec_backend: UnifiedExecBackendConfig, pub allow_login_shell: bool, pub apply_patch_tool_type: Option, pub web_search_mode: Option, @@ -99,11 +106,15 @@ impl ToolsConfig { } else { ShellCommandBackendConfig::Classic }; + let unified_exec_backend = + if features.enabled(Feature::ShellTool) && features.enabled(Feature::ShellZshFork) { + UnifiedExecBackendConfig::ZshFork + } else { + UnifiedExecBackendConfig::Direct + }; let shell_type = if !features.enabled(Feature::ShellTool) { ConfigShellToolType::Disabled - } else if features.enabled(Feature::ShellZshFork) { - ConfigShellToolType::ShellCommand } else if features.enabled(Feature::UnifiedExec) { // If ConPTY not supported (for old Windows versions), fallback on ShellCommand. if codex_utils_pty::conpty_supported() { @@ -111,6 +122,8 @@ impl ToolsConfig { } else { ConfigShellToolType::ShellCommand } + } else if features.enabled(Feature::ShellZshFork) { + ConfigShellToolType::ShellCommand } else { model_info.shell_type }; @@ -137,6 +150,7 @@ impl ToolsConfig { Self { shell_type, shell_command_backend, + unified_exec_backend, allow_login_shell: true, apply_patch_tool_type, web_search_mode: *web_search_mode, @@ -2751,7 +2765,7 @@ mod tests { } #[test] - fn shell_zsh_fork_prefers_shell_command_over_unified_exec() { + fn shell_zsh_fork_uses_unified_exec_when_enabled() { let config = test_config(); let model_info = ModelsManager::construct_model_info_offline_for_tests("o3", &config); let mut features = Features::with_defaults(); @@ -2765,11 +2779,15 @@ mod tests { session_source: SessionSource::Cli, }); - assert_eq!(tools_config.shell_type, ConfigShellToolType::ShellCommand); + assert_eq!(tools_config.shell_type, ConfigShellToolType::UnifiedExec); assert_eq!( tools_config.shell_command_backend, ShellCommandBackendConfig::ZshFork ); + assert_eq!( + tools_config.unified_exec_backend, + UnifiedExecBackendConfig::ZshFork + ); } #[test] diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 1d949dc2b9..6323857150 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -49,6 +49,10 @@ pub(crate) fn set_deterministic_process_ids_for_tests(enabled: bool) { } pub(crate) use errors::UnifiedExecError; +pub(crate) use process::NoopSpawnLifecycle; +#[cfg(unix)] +pub(crate) use process::SpawnLifecycle; +pub(crate) use process::SpawnLifecycleHandle; pub(crate) use process::UnifiedExecProcess; pub(crate) const MIN_YIELD_TIME_MS: u64 = 250; diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 69790dbad2..ba5352ff69 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -24,6 +24,17 @@ use super::UNIFIED_EXEC_OUTPUT_MAX_TOKENS; use super::UnifiedExecError; use super::head_tail_buffer::HeadTailBuffer; +pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync { + fn after_spawn(&mut self) {} +} + +pub(crate) type SpawnLifecycleHandle = Box; + +#[derive(Debug, Default)] +pub(crate) struct NoopSpawnLifecycle; + +impl SpawnLifecycle for NoopSpawnLifecycle {} + pub(crate) type OutputBuffer = Arc>; pub(crate) struct OutputHandles { pub(crate) output_buffer: OutputBuffer, @@ -44,6 +55,7 @@ pub(crate) struct UnifiedExecProcess { output_drained: Arc, output_task: JoinHandle<()>, sandbox_type: SandboxType, + _spawn_lifecycle: SpawnLifecycleHandle, } impl UnifiedExecProcess { @@ -51,6 +63,7 @@ impl UnifiedExecProcess { process_handle: ExecCommandSession, initial_output_rx: tokio::sync::broadcast::Receiver>, sandbox_type: SandboxType, + spawn_lifecycle: SpawnLifecycleHandle, ) -> Self { let output_buffer = Arc::new(Mutex::new(HeadTailBuffer::default())); let output_notify = Arc::new(Notify::new()); @@ -92,6 +105,7 @@ impl UnifiedExecProcess { output_drained, output_task, sandbox_type, + _spawn_lifecycle: spawn_lifecycle, } } @@ -196,13 +210,14 @@ impl UnifiedExecProcess { pub(super) async fn from_spawned( spawned: SpawnedPty, sandbox_type: SandboxType, + spawn_lifecycle: SpawnLifecycleHandle, ) -> Result { let SpawnedPty { session: process_handle, output_rx, mut exit_rx, } = spawned; - let managed = Self::new(process_handle, output_rx, sandbox_type); + let managed = Self::new(process_handle, output_rx, sandbox_type, spawn_lifecycle); let exit_ready = matches!(exit_rx.try_recv(), Ok(_) | Err(TryRecvError::Closed)); diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 821abc683e..4ba125009e 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -49,6 +49,7 @@ use crate::unified_exec::generate_chunk_id; use crate::unified_exec::head_tail_buffer::HeadTailBuffer; use crate::unified_exec::process::OutputBuffer; use crate::unified_exec::process::OutputHandles; +use crate::unified_exec::process::SpawnLifecycleHandle; use crate::unified_exec::process::UnifiedExecProcess; use crate::unified_exec::resolve_max_tokens; @@ -528,6 +529,7 @@ impl UnifiedExecProcessManager { &self, env: &ExecRequest, tty: bool, + mut spawn_lifecycle: SpawnLifecycleHandle, ) -> Result { let (program, args) = env .command @@ -555,7 +557,8 @@ impl UnifiedExecProcessManager { }; let spawned = spawn_result.map_err(|err| UnifiedExecError::create_process(err.to_string()))?; - UnifiedExecProcess::from_spawned(spawned, env.sandbox).await + spawn_lifecycle.after_spawn(); + UnifiedExecProcess::from_spawned(spawned, env.sandbox, spawn_lifecycle).await } pub(super) async fn open_session_with_sandbox( @@ -569,7 +572,8 @@ impl UnifiedExecProcessManager { Some(context.session.conversation_id), )); let mut orchestrator = ToolOrchestrator::new(); - let mut runtime = UnifiedExecRuntime::new(self); + let mut runtime = + UnifiedExecRuntime::new(self, context.turn.tools_config.unified_exec_backend); let exec_approval_requirement = context .session .services diff --git a/codex-rs/core/tests/common/zsh_fork.rs b/codex-rs/core/tests/common/zsh_fork.rs index fbc46421d9..29db068ef5 100644 --- a/codex-rs/core/tests/common/zsh_fork.rs +++ b/codex-rs/core/tests/common/zsh_fork.rs @@ -91,6 +91,29 @@ where builder.build(server).await } +pub async fn build_unified_exec_zsh_fork_test( + server: &wiremock::MockServer, + runtime: ZshForkRuntime, + approval_policy: AskForApproval, + sandbox_policy: SandboxPolicy, + pre_build_hook: F, +) -> Result +where + F: FnOnce(&Path) + Send + 'static, +{ + let mut builder = test_codex() + .with_pre_build_hook(pre_build_hook) + .with_config(move |config| { + runtime.apply_to_config(config, approval_policy, sandbox_policy); + config.use_experimental_unified_exec_tool = true; + config + .features + .enable(Feature::UnifiedExec) + .expect("test config should allow feature update"); + }); + builder.build(server).await +} + fn find_test_zsh_path() -> Result> { let repo_root = codex_utils_cargo_bin::repo_root()?; let dotslash_zsh = repo_root.join("codex-rs/app-server/tests/suite/zsh"); diff --git a/codex-rs/core/tests/suite/approvals.rs b/codex-rs/core/tests/suite/approvals.rs index 9ca159965f..9bc8a9185b 100644 --- a/codex-rs/core/tests/suite/approvals.rs +++ b/codex-rs/core/tests/suite/approvals.rs @@ -35,6 +35,7 @@ use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_event_with_timeout; +use core_test_support::zsh_fork::build_unified_exec_zsh_fork_test; use core_test_support::zsh_fork::build_zsh_fork_test; use core_test_support::zsh_fork::restrictive_workspace_write_policy; use core_test_support::zsh_fork::zsh_fork_runtime; @@ -1985,6 +1986,158 @@ async fn approving_execpolicy_amendment_persists_policy_and_skips_future_prompts Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[cfg(unix)] +async fn unified_exec_zsh_fork_execpolicy_amendment_skips_later_subcommands() -> Result<()> { + skip_if_no_network!(Ok(())); + + let Some(runtime) = zsh_fork_runtime("unified exec zsh-fork execpolicy amendment test")? else { + return Ok(()); + }; + + let approval_policy = AskForApproval::UnlessTrusted; + let sandbox_policy = SandboxPolicy::new_read_only_policy(); + let server = start_mock_server().await; + let test = build_unified_exec_zsh_fork_test( + &server, + runtime, + approval_policy, + sandbox_policy.clone(), + |_| {}, + ) + .await?; + let allow_prefix_path = test.cwd.path().join("allow-prefix-zsh-fork.txt"); + let _ = fs::remove_file(&allow_prefix_path); + + let call_id = "allow-prefix-zsh-fork"; + let command = "touch allow-prefix-zsh-fork.txt && touch allow-prefix-zsh-fork.txt"; + let event = exec_command_event( + call_id, + command, + Some(1_000), + SandboxPermissions::UseDefault, + None, + )?; + let _ = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-zsh-fork-allow-prefix-1"), + event, + ev_completed("resp-zsh-fork-allow-prefix-1"), + ]), + ) + .await; + let results = mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-zsh-fork-allow-prefix-1", "done"), + ev_completed("resp-zsh-fork-allow-prefix-2"), + ]), + ) + .await; + + submit_turn( + &test, + "allow-prefix-zsh-fork", + approval_policy, + sandbox_policy, + ) + .await?; + + let expected_execpolicy_amendment = ExecPolicyAmendment::new(vec![ + "touch".to_string(), + "allow-prefix-zsh-fork.txt".to_string(), + ]); + let mut saw_parent_approval = false; + let mut saw_subcommand_approval = false; + loop { + let event = wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::ExecApprovalRequest(_) | EventMsg::TurnComplete(_) + ) + }) + .await; + + match event { + EventMsg::TurnComplete(_) => break, + EventMsg::ExecApprovalRequest(approval) => { + let command_parts = approval.command.clone(); + let last_arg = command_parts.last().map(String::as_str).unwrap_or_default(); + if last_arg == command { + assert!( + !saw_parent_approval, + "unexpected duplicate parent approval: {command_parts:?}" + ); + saw_parent_approval = true; + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: None, + decision: ReviewDecision::Approved, + }) + .await?; + continue; + } + + let is_touch_subcommand = command_parts + .iter() + .any(|part| part == "allow-prefix-zsh-fork.txt") + && command_parts + .first() + .is_some_and(|part| part.ends_with("/touch") || part == "touch"); + if is_touch_subcommand { + assert!( + !saw_subcommand_approval, + "execpolicy amendment should suppress later matching subcommand approvals: {command_parts:?}" + ); + saw_subcommand_approval = true; + assert_eq!( + approval.proposed_execpolicy_amendment, + Some(expected_execpolicy_amendment.clone()) + ); + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: None, + decision: ReviewDecision::ApprovedExecpolicyAmendment { + proposed_execpolicy_amendment: expected_execpolicy_amendment + .clone(), + }, + }) + .await?; + continue; + } + + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: None, + decision: ReviewDecision::Approved, + }) + .await?; + } + other => panic!("unexpected event: {other:?}"), + } + } + + assert!(saw_parent_approval, "expected parent unified-exec approval"); + assert!( + saw_subcommand_approval, + "expected at least one intercepted touch approval" + ); + + let result = parse_result(&results.single_request().function_call_output(call_id)); + assert_eq!(result.exit_code.unwrap_or(0), 0); + assert!( + allow_prefix_path.exists(), + "expected touch command to complete after approving the first intercepted subcommand; output: {}", + result.stdout + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[cfg(unix)] async fn matched_prefix_rule_runs_unsandboxed_under_zsh_fork() -> Result<()> { diff --git a/codex-rs/core/tests/suite/skill_approval.rs b/codex-rs/core/tests/suite/skill_approval.rs index d81cbb5c66..6822791093 100644 --- a/codex-rs/core/tests/suite/skill_approval.rs +++ b/codex-rs/core/tests/suite/skill_approval.rs @@ -18,6 +18,7 @@ use core_test_support::skip_if_no_network; use core_test_support::test_codex::TestCodex; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; +use core_test_support::zsh_fork::build_unified_exec_zsh_fork_test; use core_test_support::zsh_fork::build_zsh_fork_test; use core_test_support::zsh_fork::restrictive_workspace_write_policy; use core_test_support::zsh_fork::zsh_fork_runtime; @@ -48,6 +49,13 @@ fn shell_command_arguments(command: &str) -> Result { }))?) } +fn exec_command_arguments(command: &str) -> Result { + Ok(serde_json::to_string(&json!({ + "cmd": command, + "yield_time_ms": 500, + }))?) +} + async fn submit_turn_with_policies( test: &TestCodex, prompt: &str, @@ -87,6 +95,38 @@ echo 'zsh-fork-stderr' >&2 ) } +#[cfg(unix)] +fn write_repo_skill_with_shell_script_contents( + repo_root: &Path, + name: &str, + script_name: &str, + script_contents: &str, +) -> Result { + use std::os::unix::fs::PermissionsExt; + + let skill_dir = repo_root.join(".agents").join("skills").join(name); + let scripts_dir = skill_dir.join("scripts"); + fs::create_dir_all(&scripts_dir)?; + fs::write(repo_root.join(".git"), "gitdir: here")?; + fs::write( + skill_dir.join("SKILL.md"), + format!( + r#"--- +name: {name} +description: {name} skill +--- +"# + ), + )?; + + let script_path = scripts_dir.join(script_name); + fs::write(&script_path, script_contents)?; + let mut permissions = fs::metadata(&script_path)?.permissions(); + permissions.set_mode(0o755); + fs::set_permissions(&script_path, permissions)?; + Ok(script_path) +} + #[cfg(unix)] fn write_skill_with_shell_script_contents( home: &Path, @@ -265,6 +305,168 @@ permissions: Ok(()) } +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unified_exec_zsh_fork_prompts_for_skill_script_execution() -> Result<()> { + skip_if_no_network!(Ok(())); + + let Some(runtime) = zsh_fork_runtime("unified exec zsh-fork skill prompt test")? else { + return Ok(()); + }; + + let server = start_mock_server().await; + let tool_call_id = "uexec-zsh-fork-skill-call"; + let test = build_unified_exec_zsh_fork_test( + &server, + runtime, + AskForApproval::OnRequest, + SandboxPolicy::new_workspace_write_policy(), + |home| { + write_skill_with_shell_script(home, "mbolin-test-skill", "hello-mbolin.sh").unwrap(); + write_skill_metadata( + home, + "mbolin-test-skill", + r#" +permissions: + file_system: + read: + - "./data" + write: + - "./output" +"#, + ) + .unwrap(); + }, + ) + .await?; + + let (script_path_str, command) = skill_script_command(&test, "hello-mbolin.sh")?; + let arguments = exec_command_arguments(&command)?; + let mocks = + mount_function_call_agent_response(&server, tool_call_id, &arguments, "exec_command").await; + + submit_turn_with_policies( + &test, + "use $mbolin-test-skill", + AskForApproval::OnRequest, + SandboxPolicy::new_workspace_write_policy(), + ) + .await?; + + let approval = wait_for_exec_approval_request(&test) + .await + .expect("expected exec approval request before completion"); + assert_eq!(approval.call_id, tool_call_id); + assert_eq!(approval.command, vec![script_path_str.clone()]); + assert_eq!( + approval.available_decisions, + Some(vec![ + ReviewDecision::Approved, + ReviewDecision::ApprovedForSession, + ReviewDecision::Abort, + ]) + ); + assert_eq!( + approval.additional_permissions, + Some(PermissionProfile { + file_system: Some(FileSystemPermissions { + read: Some(vec![absolute_path( + &test.codex_home_path().join("skills/mbolin-test-skill/data"), + )]), + write: Some(vec![absolute_path( + &test + .codex_home_path() + .join("skills/mbolin-test-skill/output"), + )]), + }), + ..Default::default() + }) + ); + + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: None, + decision: ReviewDecision::Denied, + }) + .await?; + + wait_for_turn_complete(&test).await; + + let call_output = mocks + .completion + .single_request() + .function_call_output(tool_call_id); + let output = call_output["output"].as_str().unwrap_or_default(); + assert!( + output.contains("Execution denied: User denied execution"), + "expected rejection marker in function_call_output: {output:?}" + ); + + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unified_exec_zsh_fork_keeps_skill_loading_pinned_to_turn_cwd() -> Result<()> { + skip_if_no_network!(Ok(())); + + let Some(runtime) = zsh_fork_runtime("unified exec zsh-fork turn cwd skill test")? else { + return Ok(()); + }; + + let server = start_mock_server().await; + let tool_call_id = "uexec-zsh-fork-repo-skill-call"; + let test = build_unified_exec_zsh_fork_test( + &server, + runtime, + AskForApproval::OnRequest, + SandboxPolicy::new_workspace_write_policy(), + |_| {}, + ) + .await?; + + let repo_root = test.cwd_path().join("repo"); + let script_path = write_repo_skill_with_shell_script_contents( + &repo_root, + "repo-skill", + "repo-skill.sh", + "#!/bin/sh\necho 'repo-skill-output'\n", + )?; + let script_path_quoted = shlex::try_join([script_path.to_string_lossy().as_ref()])?; + let repo_root_quoted = shlex::try_join([repo_root.to_string_lossy().as_ref()])?; + let command = format!("cd {repo_root_quoted} && {script_path_quoted}"); + let arguments = exec_command_arguments(&command)?; + let mocks = + mount_function_call_agent_response(&server, tool_call_id, &arguments, "exec_command").await; + + submit_turn_with_policies( + &test, + "run the repo skill after changing directories", + AskForApproval::OnRequest, + SandboxPolicy::new_workspace_write_policy(), + ) + .await?; + + let approval = wait_for_exec_approval_request(&test).await; + assert!( + approval.is_none(), + "changing directories inside unified exec should not load repo-local skills from the shell cwd", + ); + + let call_output = mocks + .completion + .single_request() + .function_call_output(tool_call_id); + let output = call_output["output"].as_str().unwrap_or_default(); + assert!( + output.contains("repo-skill-output"), + "expected repo skill script to run without skill-specific approval when only the shell cwd changes: {output:?}" + ); + + Ok(()) +} + /// Look for `additional_permissions == None`, then verify that both the first /// run and the cached session-approval rerun stay inside the turn sandbox. #[cfg(unix)] diff --git a/codex-rs/shell-escalation/src/lib.rs b/codex-rs/shell-escalation/src/lib.rs index 1cc42a46db..6d08a78716 100644 --- a/codex-rs/shell-escalation/src/lib.rs +++ b/codex-rs/shell-escalation/src/lib.rs @@ -14,6 +14,8 @@ pub use unix::EscalationPermissions; #[cfg(unix)] pub use unix::EscalationPolicy; #[cfg(unix)] +pub use unix::EscalationSession; +#[cfg(unix)] pub use unix::ExecParams; #[cfg(unix)] pub use unix::ExecResult; diff --git a/codex-rs/shell-escalation/src/unix/escalate_server.rs b/codex-rs/shell-escalation/src/unix/escalate_server.rs index 92d9adb9d3..3c5b4f5ed2 100644 --- a/codex-rs/shell-escalation/src/unix/escalate_server.rs +++ b/codex-rs/shell-escalation/src/unix/escalate_server.rs @@ -7,7 +7,9 @@ use std::time::Duration; use anyhow::Context as _; use codex_utils_absolute_path::AbsolutePathBuf; +use socket2::Socket; use tokio::process::Command; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use crate::unix::escalate_protocol::ESCALATE_SOCKET_ENV_VAR; @@ -82,6 +84,31 @@ pub struct PreparedExec { pub arg0: Option, } +#[derive(Debug)] +pub struct EscalationSession { + env: HashMap, + task: JoinHandle>, + client_socket: Option, + cancellation_token: CancellationToken, +} + +impl EscalationSession { + pub fn env(&self) -> &HashMap { + &self.env + } + + pub fn close_client_socket(&mut self) { + self.client_socket.take(); + } +} + +impl Drop for EscalationSession { + fn drop(&mut self) { + self.cancellation_token.cancel(); + self.task.abort(); + } +} + pub struct EscalateServer { bash_path: PathBuf, execve_wrapper: PathBuf, @@ -106,16 +133,51 @@ impl EscalateServer { cancel_rx: CancellationToken, command_executor: Arc, ) -> anyhow::Result { + let session = self.start_session(cancel_rx.clone(), Arc::clone(&command_executor))?; + let command = vec![ + self.bash_path.to_string_lossy().to_string(), + if params.login == Some(false) { + "-c".to_string() + } else { + "-lc".to_string() + }, + params.command, + ]; + let workdir = AbsolutePathBuf::try_from(params.workdir)?; + let result = command_executor + .run( + command, + workdir.to_path_buf(), + session.env().clone(), + cancel_rx, + ) + .await?; + Ok(result) + } + + /// Starts an escalation session and returns the environment overlay a shell + /// needs in order to route intercepted execs through this server. + /// + /// This does not spawn the shell itself. Callers own process creation and + /// only use the returned environment plus the session lifetime handle. + pub fn start_session( + &self, + parent_cancellation_token: CancellationToken, + command_executor: Arc, + ) -> anyhow::Result { + let cancellation_token = CancellationToken::new(); let (escalate_server, escalate_client) = AsyncDatagramSocket::pair()?; let client_socket = escalate_client.into_inner(); // Only the client endpoint should cross exec into the wrapper process. client_socket.set_cloexec(false)?; - let escalate_task = tokio::spawn(escalate_task( + let task = tokio::spawn(escalate_task( escalate_server, Arc::clone(&self.policy), Arc::clone(&command_executor), + parent_cancellation_token, + cancellation_token.clone(), )); - let mut env = std::env::vars().collect::>(); + let mut env = HashMap::new(); env.insert( ESCALATE_SOCKET_ENV_VAR.to_string(), client_socket.as_raw_fd().to_string(), @@ -128,22 +190,12 @@ impl EscalateServer { LEGACY_BASH_EXEC_WRAPPER_ENV_VAR.to_string(), self.execve_wrapper.to_string_lossy().to_string(), ); - - let command = vec![ - self.bash_path.to_string_lossy().to_string(), - if params.login == Some(false) { - "-c".to_string() - } else { - "-lc".to_string() - }, - params.command, - ]; - let workdir = AbsolutePathBuf::try_from(params.workdir)?; - let result = command_executor - .run(command, workdir.to_path_buf(), env, cancel_rx) - .await?; - escalate_task.abort(); - Ok(result) + Ok(EscalationSession { + env, + task, + client_socket: Some(client_socket), + cancellation_token, + }) } } @@ -151,9 +203,15 @@ async fn escalate_task( socket: AsyncDatagramSocket, policy: Arc, command_executor: Arc, + parent_cancellation_token: CancellationToken, + session_cancellation_token: CancellationToken, ) -> anyhow::Result<()> { loop { - let (_, mut fds) = socket.receive_with_fds().await?; + let (_, mut fds) = tokio::select! { + received = socket.receive_with_fds() => received?, + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; if fds.len() != 1 { tracing::error!("expected 1 fd in datagram handshake, got {}", fds.len()); continue; @@ -161,9 +219,17 @@ async fn escalate_task( let stream_socket = AsyncSocket::from_fd(fds.remove(0))?; let policy = Arc::clone(&policy); let command_executor = Arc::clone(&command_executor); + let parent_cancellation_token = parent_cancellation_token.clone(); + let session_cancellation_token = session_cancellation_token.clone(); tokio::spawn(async move { - if let Err(err) = - handle_escalate_session_with_policy(stream_socket, policy, command_executor).await + if let Err(err) = handle_escalate_session_with_policy( + stream_socket, + policy, + command_executor, + parent_cancellation_token, + session_cancellation_token, + ) + .await { tracing::error!("escalate session failed: {err:?}"); } @@ -175,18 +241,27 @@ async fn handle_escalate_session_with_policy( socket: AsyncSocket, policy: Arc, command_executor: Arc, + parent_cancellation_token: CancellationToken, + session_cancellation_token: CancellationToken, ) -> anyhow::Result<()> { let EscalateRequest { file, argv, workdir, env, - } = socket.receive::().await?; + } = tokio::select! { + request = socket.receive::() => request?, + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; let program = AbsolutePathBuf::resolve_path_against_base(file, workdir.as_path())?; - let decision = policy - .determine_action(&program, &argv, &workdir) - .await - .context("failed to determine escalation action")?; + let decision = tokio::select! { + decision = policy.determine_action(&program, &argv, &workdir) => { + decision.context("failed to determine escalation action")? + } + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; tracing::debug!("decided {decision:?} for {program:?} {argv:?} {workdir:?}"); @@ -204,10 +279,13 @@ async fn handle_escalate_session_with_policy( action: EscalateAction::Escalate, }) .await?; - let (msg, fds) = socket - .receive_with_fds::() - .await - .context("failed to receive SuperExecMessage")?; + let (msg, fds) = tokio::select! { + message = socket.receive_with_fds::() => { + message.context("failed to receive SuperExecMessage")? + } + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; if fds.len() != msg.fds.len() { return Err(anyhow::anyhow!( "mismatched number of fds in SuperExecMessage: {} in the message, {} from the control message", @@ -231,9 +309,11 @@ async fn handle_escalate_session_with_policy( cwd, env, arg0, - } = command_executor - .prepare_escalated_exec(&program, &argv, &workdir, env, execution) - .await?; + } = tokio::select! { + prepared = command_executor.prepare_escalated_exec(&program, &argv, &workdir, env, execution) => prepared?, + _ = parent_cancellation_token.cancelled() => return Ok(()), + _ = session_cancellation_token.cancelled() => return Ok(()), + }; let (program, args) = command .split_first() .ok_or_else(|| anyhow::anyhow!("prepared escalated command must not be empty"))?; @@ -245,7 +325,8 @@ async fn handle_escalate_session_with_policy( .current_dir(&cwd) .stdin(Stdio::null()) .stdout(Stdio::null()) - .stderr(Stdio::null()); + .stderr(Stdio::null()) + .kill_on_drop(true); unsafe { command.pre_exec(move || { for (dst_fd, src_fd) in msg.fds.iter().zip(&fds) { @@ -255,7 +336,17 @@ async fn handle_escalate_session_with_policy( }); } let mut child = command.spawn()?; - let exit_status = child.wait().await?; + let exit_status = tokio::select! { + status = child.wait() => status?, + _ = parent_cancellation_token.cancelled() => { + let _ = child.start_kill(); + child.wait().await? + } + _ = session_cancellation_token.cancelled() => { + let _ = child.start_kill(); + child.wait().await? + } + }; socket .send(SuperExecResult { exit_code: exit_status.code().unwrap_or(127), @@ -283,6 +374,13 @@ mod tests { use pretty_assertions::assert_eq; use std::collections::HashMap; use std::path::PathBuf; + use std::sync::LazyLock; + use tempfile::TempDir; + use tokio::time::Instant; + use tokio::time::sleep; + + static ESCALATE_SERVER_TEST_LOCK: LazyLock> = + LazyLock::new(|| tokio::sync::Mutex::new(())); struct DeterministicEscalationPolicy { decision: EscalationDecision, @@ -391,8 +489,89 @@ mod tests { } } + async fn wait_for_pid_file(pid_file: &std::path::Path) -> anyhow::Result { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if let Ok(contents) = std::fs::read_to_string(pid_file) { + return Ok(contents.trim().parse()?); + } + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for pid file {}", + pid_file.display() + )); + } + sleep(Duration::from_millis(20)).await; + } + } + + fn process_exists(pid: i32) -> bool { + let rc = unsafe { libc::kill(pid, 0) }; + if rc == 0 { + return true; + } + std::io::Error::last_os_error().raw_os_error() != Some(libc::ESRCH) + } + + async fn wait_for_process_exit(pid: i32) -> anyhow::Result<()> { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if !process_exists(pid) { + return Ok(()); + } + if Instant::now() >= deadline { + return Err(anyhow::anyhow!("timed out waiting for pid {pid} to exit")); + } + sleep(Duration::from_millis(20)).await; + } + } + + /// Verifies that `start_session()` returns only the wrapper/socket env + /// overlay and does not need to touch the configured shell or wrapper + /// executable paths. + /// + /// The `/bin/bash` and `/tmp/codex-execve-wrapper` values here are + /// intentionally fake sentinels: this test asserts that the paths are + /// copied into the exported environment and that the socket fd stays valid + /// until `close_client_socket()` is called. + #[tokio::test] + async fn start_session_exposes_wrapper_env_overlay() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; + let execve_wrapper = PathBuf::from("/tmp/codex-execve-wrapper"); + let execve_wrapper_str = execve_wrapper.to_string_lossy().to_string(); + let server = EscalateServer::new( + PathBuf::from("/bin/bash"), + execve_wrapper.clone(), + DeterministicEscalationPolicy { + decision: EscalationDecision::run(), + }, + ); + + let mut session = server.start_session( + CancellationToken::new(), + Arc::new(ForwardingShellCommandExecutor), + )?; + let env = session.env(); + assert_eq!(env.get(EXEC_WRAPPER_ENV_VAR), Some(&execve_wrapper_str)); + assert_eq!( + env.get(LEGACY_BASH_EXEC_WRAPPER_ENV_VAR), + Some(&execve_wrapper_str) + ); + let socket_fd = env + .get(ESCALATE_SOCKET_ENV_VAR) + .expect("session should export shell escalation socket"); + let socket_fd = socket_fd.parse::()?; + assert!(socket_fd >= 0); + assert_ne!(unsafe { libc::fcntl(socket_fd, libc::F_GETFD) }, -1); + session.close_client_socket(); + assert_eq!(unsafe { libc::fcntl(socket_fd, libc::F_GETFD) }, -1); + + Ok(()) + } + #[tokio::test] async fn handle_escalate_session_respects_run_in_sandbox_decision() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; let (server, client) = AsyncSocket::pair()?; let server_task = tokio::spawn(handle_escalate_session_with_policy( server, @@ -400,6 +579,8 @@ mod tests { decision: EscalationDecision::run(), }), Arc::new(ForwardingShellCommandExecutor), + CancellationToken::new(), + CancellationToken::new(), )); let mut env = HashMap::new(); @@ -430,6 +611,7 @@ mod tests { #[tokio::test] async fn handle_escalate_session_resolves_relative_file_against_request_workdir() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; let (server, client) = AsyncSocket::pair()?; let tmp = tempfile::TempDir::new()?; let workdir = tmp.path().join("workspace"); @@ -443,6 +625,8 @@ mod tests { expected_workdir: workdir.clone(), }), Arc::new(ForwardingShellCommandExecutor), + CancellationToken::new(), + CancellationToken::new(), )); client @@ -466,6 +650,7 @@ mod tests { #[tokio::test] async fn handle_escalate_session_executes_escalated_command() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; let (server, client) = AsyncSocket::pair()?; let server_task = tokio::spawn(handle_escalate_session_with_policy( server, @@ -473,6 +658,8 @@ mod tests { decision: EscalationDecision::escalate(EscalationExecution::Unsandboxed), }), Arc::new(ForwardingShellCommandExecutor), + CancellationToken::new(), + CancellationToken::new(), )); client @@ -508,6 +695,7 @@ mod tests { #[tokio::test] async fn handle_escalate_session_passes_permissions_to_executor() -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; let (server, client) = AsyncSocket::pair()?; let server_task = tokio::spawn(handle_escalate_session_with_policy( server, @@ -529,6 +717,8 @@ mod tests { ..Default::default() }), }), + CancellationToken::new(), + CancellationToken::new(), )); client @@ -557,4 +747,85 @@ mod tests { server_task.await? } + + #[tokio::test] + async fn dropping_session_aborts_intercept_workers_and_kills_spawned_child() + -> anyhow::Result<()> { + let _guard = ESCALATE_SERVER_TEST_LOCK.lock().await; + let tmp = TempDir::new()?; + let pid_file = tmp.path().join("escalated-child.pid"); + let pid_file_display = pid_file.display().to_string(); + assert!( + !pid_file_display.contains('\''), + "test temp path should not contain single quotes: {pid_file_display}" + ); + let server = EscalateServer::new( + PathBuf::from("/bin/bash"), + PathBuf::from("/tmp/codex-execve-wrapper"), + DeterministicEscalationPolicy { + decision: EscalationDecision::escalate(EscalationExecution::Unsandboxed), + }, + ); + + let session = server.start_session( + CancellationToken::new(), + Arc::new(ForwardingShellCommandExecutor), + )?; + let socket_fd = session + .env() + .get(ESCALATE_SOCKET_ENV_VAR) + .expect("session should export shell escalation socket") + .parse::()?; + let dup_socket_fd = unsafe { libc::dup(socket_fd) }; + assert!(dup_socket_fd >= 0, "expected dup() to succeed"); + let handshake_client = unsafe { AsyncDatagramSocket::from_raw_fd(dup_socket_fd) }?; + let (server_stream, client_stream) = AsyncSocket::pair()?; + let server_stream_fd = server_stream.into_inner().into(); + handshake_client + .send_with_fds(&[0], &[server_stream_fd]) + .await + .context("failed to send handshake datagram")?; + + client_stream + .send(EscalateRequest { + file: PathBuf::from("/bin/sh"), + argv: vec![ + "sh".to_string(), + "-c".to_string(), + format!("echo $$ > '{pid_file_display}' && exec /bin/sleep 100"), + ], + workdir: AbsolutePathBuf::current_dir()?, + env: HashMap::new(), + }) + .await + .context("failed to send EscalateRequest")?; + + let response = client_stream + .receive::() + .await + .context("failed to receive EscalateResponse")?; + assert_eq!( + EscalateResponse { + action: EscalateAction::Escalate, + }, + response + ); + + client_stream + .send_with_fds(SuperExecMessage { fds: Vec::new() }, &[]) + .await + .context("failed to send SuperExecMessage")?; + + let pid = wait_for_pid_file(&pid_file).await?; + assert!( + process_exists(pid), + "expected spawned child pid {pid} to exist" + ); + + drop(session); + + wait_for_process_exit(pid).await?; + + Ok(()) + } } diff --git a/codex-rs/shell-escalation/src/unix/mod.rs b/codex-rs/shell-escalation/src/unix/mod.rs index 6de12297a4..13b7146ab4 100644 --- a/codex-rs/shell-escalation/src/unix/mod.rs +++ b/codex-rs/shell-escalation/src/unix/mod.rs @@ -66,6 +66,7 @@ pub use self::escalate_protocol::EscalateAction; pub use self::escalate_protocol::EscalationDecision; pub use self::escalate_protocol::EscalationExecution; pub use self::escalate_server::EscalateServer; +pub use self::escalate_server::EscalationSession; pub use self::escalate_server::ExecParams; pub use self::escalate_server::ExecResult; pub use self::escalate_server::PreparedExec; diff --git a/codex-rs/shell-escalation/src/unix/stopwatch.rs b/codex-rs/shell-escalation/src/unix/stopwatch.rs index f564434813..19821d8be2 100644 --- a/codex-rs/shell-escalation/src/unix/stopwatch.rs +++ b/codex-rs/shell-escalation/src/unix/stopwatch.rs @@ -9,7 +9,7 @@ use tokio_util::sync::CancellationToken; #[derive(Clone, Debug)] pub struct Stopwatch { - limit: Duration, + limit: Option, inner: Arc>, notify: Arc, } @@ -30,13 +30,27 @@ impl Stopwatch { active_pauses: 0, })), notify: Arc::new(Notify::new()), - limit, + limit: Some(limit), + } + } + + pub fn unlimited() -> Self { + Self { + inner: Arc::new(Mutex::new(StopwatchState { + elapsed: Duration::ZERO, + running_since: Some(Instant::now()), + active_pauses: 0, + })), + notify: Arc::new(Notify::new()), + limit: None, } } pub fn cancellation_token(&self) -> CancellationToken { - let limit = self.limit; let token = CancellationToken::new(); + let Some(limit) = self.limit else { + return token; + }; let cancel = token.clone(); let inner = Arc::clone(&self.inner); let notify = Arc::clone(&self.notify); @@ -208,4 +222,16 @@ mod tests { // Now the stopwatch should resume and hit the limit shortly after. token.cancelled().await; } + + #[tokio::test] + async fn unlimited_stopwatch_never_cancels() { + let stopwatch = Stopwatch::unlimited(); + let token = stopwatch.cancellation_token(); + + assert!( + timeout(Duration::from_millis(30), token.cancelled()) + .await + .is_err() + ); + } }