diff --git a/codex-rs/app-server/tests/suite/v2/analytics.rs b/codex-rs/app-server/tests/suite/v2/analytics.rs index ea876940ab..645e963a80 100644 --- a/codex-rs/app-server/tests/suite/v2/analytics.rs +++ b/codex-rs/app-server/tests/suite/v2/analytics.rs @@ -125,10 +125,50 @@ pub(crate) async fn wait_for_analytics_event( read_timeout: Duration, event_type: &str, ) -> Result { - wait_for_matching_analytics_event(server, read_timeout, |event| { - event["event_type"] == event_type + wait_for_analytics_events(server, read_timeout, event_type, /*expected_count*/ 1) + .await? + .pop() + .ok_or_else(|| anyhow::anyhow!("matching analytics event should be present")) +} + +pub(crate) async fn wait_for_analytics_events( + server: &MockServer, + read_timeout: Duration, + event_type: &str, + expected_count: usize, +) -> Result> { + timeout(read_timeout, async { + loop { + let events = captured_analytics_events(server) + .await? + .into_iter() + .filter(|event| event["event_type"] == event_type) + .collect::>(); + if events.len() >= expected_count { + return Ok::, anyhow::Error>(events); + } + tokio::time::sleep(Duration::from_millis(25)).await; + } }) - .await + .await? +} + +pub(crate) async fn captured_analytics_events(server: &MockServer) -> Result> { + let Some(requests) = server.received_requests().await else { + return Ok(Vec::new()); + }; + let mut events = Vec::new(); + for request in &requests { + if request.method != "POST" || request.url.path() != "/codex/analytics-events/events" { + continue; + } + let payload: Value = serde_json::from_slice(&request.body) + .map_err(|err| anyhow::anyhow!("invalid analytics payload: {err}"))?; + if let Some(batch) = payload["events"].as_array() { + events.extend(batch.iter().cloned()); + } + } + Ok(events) } pub(crate) async fn wait_for_goal_event( diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 9d07da0660..55e6a69ea9 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -38,6 +38,7 @@ mod plan_item; mod plugin_install; mod plugin_list; mod plugin_read; +mod plugin_script_lifecycle; mod plugin_share; mod plugin_uninstall; mod process_exec; diff --git a/codex-rs/app-server/tests/suite/v2/plugin_script_lifecycle.rs b/codex-rs/app-server/tests/suite/v2/plugin_script_lifecycle.rs new file mode 100644 index 0000000000..5464847fcc --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/plugin_script_lifecycle.rs @@ -0,0 +1,456 @@ +#![cfg(unix)] + +use anyhow::Result; +use app_test_support::TestAppServer; +use app_test_support::create_final_assistant_message_sse_response; +use app_test_support::create_mock_responses_server_sequence; +use app_test_support::create_shell_command_sse_response; +use app_test_support::to_response; +use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; +use codex_app_server_protocol::AskForApproval; +use codex_app_server_protocol::CommandExecutionApprovalDecision; +use codex_app_server_protocol::CommandExecutionRequestApprovalResponse; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::SandboxPolicy; +use codex_app_server_protocol::ServerRequest; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnInterruptParams; +use codex_app_server_protocol::TurnInterruptResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::UserInput as V2UserInput; +use core_test_support::skip_if_no_network; +use pretty_assertions::assert_eq; +use serde_json::Value; +use std::os::unix::fs::PermissionsExt; +use tempfile::TempDir; +use tokio::time::timeout; + +use super::analytics::captured_analytics_events; +use super::analytics::mount_analytics_capture; +use super::analytics::wait_for_analytics_event; +use super::analytics::wait_for_analytics_events; + +const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const CURATED_PLUGIN_SHA: &str = "0123456789abcdef0123456789abcdef01234567"; +const CURATED_PLUGIN_VERSION: &str = "01234567"; +const PLUGIN_NAME: &str = "lifecycle"; +const PLUGIN_CONFIG_NAME: &str = "lifecycle@openai-curated"; + +#[derive(Clone, Copy)] +enum LifecycleBackend { + Classic, + ZshFork, +} + +struct LifecycleRun { + backend: LifecycleBackend, + sandbox_policy: SandboxPolicy, + approval_policy: AskForApproval, + accept_approval: bool, + expected_events: usize, + executable: Option<&'static str>, +} + +impl LifecycleRun { + fn classic() -> Self { + Self { + backend: LifecycleBackend::Classic, + sandbox_policy: SandboxPolicy::DangerFullAccess, + approval_policy: AskForApproval::Never, + accept_approval: false, + expected_events: 2, + executable: Some("sh"), + } + } +} + +struct LifecycleFixture { + events: Vec, + thread_id: String, + session_id: String, + turn_id: String, + plugin_root: std::path::PathBuf, +} + +#[tokio::test] +async fn plugin_script_emits_started_and_completed_lifecycle_analytics() -> Result<()> { + let fixture = run_lifecycle_fixture( + "printf 'sensitive-script-output\\n'\n", + &["secret-argument"], + /*interrupt*/ false, + LifecycleRun::classic(), + ) + .await?; + assert_eq!(fixture.events.len(), 2); + + let started = event_with_status(&fixture.events, "started")?; + let completed = event_with_status(&fixture.events, "completed")?; + let started_params = &started["event_params"]; + let completed_params = &completed["event_params"]; + + assert_eq!(started_params["version"], 1); + assert_eq!(started_params["thread_id"], fixture.thread_id); + assert_eq!(started_params["session_id"], fixture.session_id); + assert_eq!(started_params["turn_id"], fixture.turn_id); + assert_eq!(started_params["plugin_id"], PLUGIN_CONFIG_NAME); + assert_eq!(started_params["script_path"], "scripts/run.sh"); + assert!(started_params["timestamp"].as_str().is_some()); + assert!(started_params.get("duration_ms").is_none()); + assert!(started_params.get("exit_code").is_none()); + assert_eq!(started_params["skill_id"], Value::Null); + assert_eq!( + completed_params["execution_id"], + started_params["execution_id"] + ); + assert!(completed_params["duration_ms"].as_u64().is_some()); + assert_eq!(completed_params["exit_code"], 0); + + let serialized_events = serde_json::to_string(&fixture.events)?; + assert!(!serialized_events.contains("secret-argument")); + assert!(!serialized_events.contains("sensitive-script-output")); + assert!(!serialized_events.contains(fixture.plugin_root.to_string_lossy().as_ref())); + + Ok(()) +} + +#[tokio::test] +async fn zsh_fork_plugin_script_emits_terminal_lifecycle_analytics() -> Result<()> { + skip_if_no_network!(Ok(())); + let Some(zsh_path) = super::turn_start_zsh_fork::find_test_zsh_path()? else { + return Ok(()); + }; + if !super::turn_start_zsh_fork::supports_exec_wrapper_intercept(&zsh_path) { + return Ok(()); + } + for (script, interrupt, status, exit_code) in [ + ("exit 0\n", false, "completed", Some(0)), + ("exit 7\n", false, "failed", Some(7)), + ("sleep 30\n", true, "cancelled", None), + ] { + let fixture = run_lifecycle_fixture( + script, + &[], + interrupt, + LifecycleRun { + backend: LifecycleBackend::ZshFork, + ..LifecycleRun::classic() + }, + ) + .await?; + assert_eq!(fixture.events.len(), 2); + let started = event_with_status(&fixture.events, "started")?; + let terminal = event_with_status(&fixture.events, status)?; + assert_eq!( + terminal["event_params"]["execution_id"], + started["event_params"]["execution_id"] + ); + assert!(terminal["event_params"]["duration_ms"].as_u64().is_some()); + if let Some(exit_code) = exit_code { + assert_eq!(terminal["event_params"]["exit_code"], exit_code); + } + } + Ok(()) +} + +#[tokio::test] +async fn sandbox_denial_terminates_each_plugin_script_attempt() -> Result<()> { + let denial = run_lifecycle_fixture( + "printf denied > denied.txt\n", + &[], + /*interrupt*/ false, + LifecycleRun { + sandbox_policy: SandboxPolicy::ReadOnly { + network_access: false, + }, + ..LifecycleRun::classic() + }, + ) + .await?; + assert_eq!(statuses(&denial.events), vec!["started", "failed"]); + + let retry = run_lifecycle_fixture( + "printf denied > denied.txt\n", + &[], + /*interrupt*/ false, + LifecycleRun { + sandbox_policy: SandboxPolicy::ReadOnly { + network_access: false, + }, + approval_policy: AskForApproval::OnFailure, + accept_approval: true, + expected_events: 4, + ..LifecycleRun::classic() + }, + ) + .await?; + assert_eq!( + statuses(&retry.events), + vec!["started", "failed", "started", "completed"] + ); + assert_ne!( + retry.events[0]["event_params"]["execution_id"], + retry.events[2]["event_params"]["execution_id"] + ); + Ok(()) +} + +#[tokio::test] +async fn pre_spawn_plugin_script_failure_emits_no_lifecycle_analytics() -> Result<()> { + let fixture = run_lifecycle_fixture( + "exit 0\n", + &[], + /*interrupt*/ false, + LifecycleRun { + expected_events: 0, + executable: None, + ..LifecycleRun::classic() + }, + ) + .await?; + assert!(fixture.events.is_empty()); + Ok(()) +} + +async fn run_lifecycle_fixture( + script: &str, + script_args: &[&str], + interrupt: bool, + run: LifecycleRun, +) -> Result { + let temp = TempDir::new()?; + let codex_home = temp.path().join("codex-home"); + let working_directory = temp.path().join("workdir"); + std::fs::create_dir_all(&working_directory)?; + write_curated_provenance(&codex_home)?; + let plugin_root = codex_home + .join("plugins/cache/openai-curated") + .join(PLUGIN_NAME) + .join(CURATED_PLUGIN_VERSION); + std::fs::create_dir_all(plugin_root.join(".codex-plugin"))?; + std::fs::create_dir_all(plugin_root.join("scripts"))?; + std::fs::write( + plugin_root.join(".codex-plugin/plugin.json"), + serde_json::to_vec_pretty(&serde_json::json!({ + "name": PLUGIN_NAME, + "interface": { "developerName": "OpenAI" }, + }))?, + )?; + let script_path = plugin_root.join("scripts/run.sh"); + std::fs::write(&script_path, script)?; + #[cfg(unix)] + assert_eq!( + std::fs::metadata(&script_path)?.permissions().mode() & 0o111, + 0 + ); + let script_path = script_path.to_string_lossy().into_owned(); + let mut command = match run.executable { + Some(executable) => vec![executable.to_string(), script_path], + None => vec![script_path], + }; + command.extend(script_args.iter().map(|arg| (*arg).to_string())); + let mut responses = vec![create_shell_command_sse_response( + command, + Some(&working_directory), + Some(60_000), + "plugin-script-call", + )?]; + if !interrupt { + responses.push(create_final_assistant_message_sse_response("done")?); + } + let server = create_mock_responses_server_sequence(responses).await; + write_mock_responses_config_toml_with_chatgpt_base_url( + &codex_home, + &server.uri(), + &server.uri(), + )?; + let config_path = codex_home.join("config.toml"); + let config = std::fs::read_to_string(&config_path)?; + std::fs::write( + &config_path, + format!( + r#"{config} +[features] +plugins = true +plugin_script_lifecycle = true +{} + +[plugins."{PLUGIN_CONFIG_NAME}"] +enabled = true +"#, + match run.backend { + LifecycleBackend::Classic => "", + LifecycleBackend::ZshFork => + "shell_zsh_fork = true\nshell_snapshot = false\nunified_exec = false", + }, + ), + )?; + mount_analytics_capture(&server, &codex_home).await?; + let isolated_home = codex_home.to_string_lossy(); + let mut mcp = match run.backend { + LifecycleBackend::Classic => { + TestAppServer::new_with_env( + &codex_home, + &[ + ("HOME", Some(isolated_home.as_ref())), + ("USERPROFILE", Some(isolated_home.as_ref())), + ], + ) + .await? + } + LifecycleBackend::ZshFork => { + let zsh_path = super::turn_start_zsh_fork::find_test_zsh_path()? + .ok_or_else(|| anyhow::anyhow!("zsh-fork lifecycle fixture requires zsh"))?; + super::turn_start_zsh_fork::create_zsh_test_mcp_process( + &codex_home, + &working_directory, + &zsh_path, + ) + .await? + } + }; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let thread_request = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_request)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_response)?; + if run.executable.is_none() { + // Keep the attributed direct plugin-script command absolute so resolver attribution + // succeeds even though this missing cwd makes the actual shell process fail to spawn. + std::fs::remove_dir(&working_directory)?; + } + let turn_request = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "run the plugin script".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(working_directory), + sandbox_policy: Some(run.sandbox_policy), + approval_policy: Some(run.approval_policy), + ..Default::default() + }) + .await?; + let turn_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_request)), + ) + .await??; + let TurnStartResponse { turn } = to_response::(turn_response)?; + if run.accept_approval { + let ServerRequest::CommandExecutionRequestApproval { request_id, .. } = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await?? + else { + panic!("expected command approval request"); + }; + mcp.send_response( + request_id, + serde_json::to_value(CommandExecutionRequestApprovalResponse { + decision: CommandExecutionApprovalDecision::Accept, + })?, + ) + .await?; + } + if interrupt { + wait_for_analytics_events( + &server, + DEFAULT_READ_TIMEOUT, + "codex_plugin_lifecycle_event", + /*expected_count*/ 1, + ) + .await?; + let interrupt_request = mcp + .send_turn_interrupt_request(TurnInterruptParams { + thread_id: thread.id.clone(), + turn_id: turn.id.clone(), + }) + .await?; + let interrupt_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(interrupt_request)), + ) + .await??; + let _: TurnInterruptResponse = to_response(interrupt_response)?; + } + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + // The terminal turn event is reduced after the tool result that emits plugin lifecycle + // facts, so observing this matching marker proves analytics drained those facts without + // relying on a sleep that could miss late duplicate lifecycle events. + let terminal_turn_event = + wait_for_analytics_event(&server, DEFAULT_READ_TIMEOUT, "codex_turn_event").await?; + assert_eq!(terminal_turn_event["event_params"]["thread_id"], thread.id); + assert_eq!(terminal_turn_event["event_params"]["turn_id"], turn.id); + let events = lifecycle_events(&server).await?; + assert_eq!(events.len(), run.expected_events); + Ok(LifecycleFixture { + events, + thread_id: thread.id, + session_id: thread.session_id, + turn_id: turn.id, + plugin_root, + }) +} + +fn write_curated_provenance(codex_home: &std::path::Path) -> Result<()> { + let marketplace_path = codex_home.join(".tmp/plugins/.agents/plugins/marketplace.json"); + std::fs::create_dir_all( + marketplace_path + .parent() + .ok_or_else(|| anyhow::anyhow!("marketplace path should have a parent"))?, + )?; + std::fs::write( + marketplace_path, + serde_json::to_vec_pretty(&serde_json::json!({ + "name": "openai-curated", + "plugins": [{ + "name": PLUGIN_NAME, + "source": { + "source": "local", + "path": "./plugins/lifecycle", + }, + }], + }))?, + )?; + std::fs::write(codex_home.join(".tmp/plugins.sha"), CURATED_PLUGIN_SHA)?; + Ok(()) +} + +async fn lifecycle_events(server: &wiremock::MockServer) -> Result> { + Ok(captured_analytics_events(server) + .await? + .into_iter() + .filter(|event| event["event_type"] == "codex_plugin_lifecycle_event") + .collect()) +} + +fn statuses(events: &[Value]) -> Vec<&str> { + events + .iter() + .filter_map(|event| event["event_params"]["status"].as_str()) + .collect() +} + +fn event_with_status<'a>(events: &'a [Value], status: &str) -> Result<&'a Value> { + events + .iter() + .find(|event| event["event_params"]["status"] == status) + .ok_or_else(|| anyhow::anyhow!("missing plugin lifecycle status {status}")) +} 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 15072c1ccf..d78f67a985 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 @@ -739,7 +739,7 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() Ok(()) } -async fn create_zsh_test_mcp_process( +pub(super) async fn create_zsh_test_mcp_process( codex_home: &Path, zdotdir: &Path, zsh_path: &Path, @@ -840,7 +840,7 @@ stream_max_retries = 0 ) } -fn find_test_zsh_path() -> Result> { +pub(super) 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"); if !dotslash_zsh.is_file() { @@ -860,7 +860,7 @@ fn find_test_zsh_path() -> Result> { Ok(None) } -fn supports_exec_wrapper_intercept(zsh_path: &Path) -> bool { +pub(super) fn supports_exec_wrapper_intercept(zsh_path: &Path) -> bool { let status = std::process::Command::new(zsh_path) .arg("-fc") .arg("/usr/bin/true") diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index f88c2be3d4..b1410569db 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -13,9 +13,10 @@ use crate::exec::ExecCapturePolicy; use crate::guardian::GuardianApprovalRequest; use crate::guardian::GuardianNetworkAccessTrigger; use crate::guardian::review_approval_request; +use crate::plugin_script_lifecycle::PluginScriptExecution; use crate::sandboxing::ExecOptions; use crate::sandboxing::SandboxPermissions; -use crate::sandboxing::execute_env; +use crate::sandboxing::execute_exec_request_with_after_spawn; use crate::session::turn_context::TurnEnvironment; use crate::shell::ShellType; use crate::tools::flat_tool_name; @@ -41,6 +42,8 @@ use crate::tools::sandboxing::managed_network_for_sandbox_permissions; use crate::tools::sandboxing::sandbox_permissions_preserving_denied_reads; use crate::tools::sandboxing::with_cached_approval; use codex_network_proxy::NetworkProxy; +use codex_protocol::error::CodexErr; +use codex_protocol::error::SandboxErr; use codex_protocol::exec_output::ExecToolCallOutput; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::protocol::ReviewDecision; @@ -49,6 +52,7 @@ use codex_shell_command::powershell::prefix_powershell_script_with_utf8; use codex_utils_absolute_path::AbsolutePathBuf; use futures::future::BoxFuture; use std::collections::HashMap; +use std::sync::Arc; use tokio_util::sync::CancellationToken; #[derive(Clone, Debug)] @@ -89,6 +93,30 @@ pub(crate) enum ShellRuntimeBackend { pub struct ShellRuntime { backend: ShellRuntimeBackend, + plugin_script: Option>>, +} + +fn finish_plugin_script( + result: Result, + plugin_script: Option<&Arc>, + cancellation_token: &CancellationToken, +) -> Result { + if let Some(execution) = plugin_script { + if cancellation_token.is_cancelled() { + execution.mark_cancelled(); + } + match &result { + Ok(out) => execution.finish(Some(out.exit_code), out.timed_out), + Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Timeout { output }))) => { + execution.finish(Some(output.exit_code), /*failed*/ true) + } + Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Denied { output, .. }))) => { + execution.finish(Some(output.exit_code), /*failed*/ true) + } + Err(_) => execution.finish(/*exit_code*/ None, /*failed*/ true), + } + } + result } #[derive(serde::Serialize, Clone, Debug, Eq, PartialEq, Hash)] @@ -102,7 +130,10 @@ pub(crate) struct ApprovalKey { impl ShellRuntime { pub(crate) fn for_shell_command(backend: ShellRuntimeBackend) -> Self { - Self { backend } + Self { + backend, + plugin_script: None, + } } fn stdout_stream(ctx: &ToolCtx) -> Option { @@ -112,6 +143,24 @@ impl ShellRuntime { tx_event: ctx.session.get_tx_event(), }) } + fn finish_plugin_script_attempt( + &mut self, + result: Result, + plugin_script: Option<&Arc>, + cancellation_token: &CancellationToken, + ) -> Result { + let denied = matches!( + &result, + Err(ToolError::Codex(CodexErr::Sandbox( + SandboxErr::Denied { .. } + ))) + ); + let result = finish_plugin_script(result, plugin_script, cancellation_token); + if denied { + self.plugin_script = None; + } + result + } } impl Sandboxable for ShellRuntime { @@ -252,6 +301,18 @@ impl ToolRuntime for ShellRuntime { .as_ref() .unwrap_or(session_shell.as_ref()); let shell_snapshot_location = req.turn_environment.shell_snapshot(&req.cwd); + let plugin_script = self + .plugin_script + .get_or_insert_with(|| { + PluginScriptExecution::resolve( + ctx.session.as_ref(), + ctx.turn.as_ref(), + &req.hook_command, + &req.cwd, + req.shell_type.unwrap_or(session_shell.shell_type), + ) + }) + .clone(); let (file_system_sandbox_policy, _) = attempt.permissions.to_runtime_permissions(); let sandbox_permissions = sandbox_permissions_preserving_denied_reads( req.sandbox_permissions, @@ -299,9 +360,30 @@ impl ToolRuntime for ShellRuntime { }; if self.backend == ShellRuntimeBackend::ShellCommandZshFork { - match zsh_fork_backend::maybe_run_shell_command(req, attempt, ctx, &command).await? { - Some(out) => return Ok(out), - None => { + match zsh_fork_backend::maybe_run_shell_command( + req, + attempt, + ctx, + &command, + plugin_script.clone(), + ) + .await + { + Ok(Some(out)) => { + return self.finish_plugin_script_attempt( + Ok(out), + plugin_script.as_ref(), + &req.cancellation_token, + ); + } + Err(err) => { + return self.finish_plugin_script_attempt( + Err(err), + plugin_script.as_ref(), + &req.cancellation_token, + ); + } + Ok(None) => { tracing::warn!( "ZshFork backend specified, but conditions for using it were not met, falling back to normal execution", ); @@ -323,10 +405,15 @@ impl ToolRuntime for ShellRuntime { let env = attempt .env_for(command, options, managed_network) .map_err(ToolError::Codex)?; - let out = execute_env(env, Self::stdout_stream(ctx)) - .await - .map_err(ToolError::Codex)?; - Ok(out) + let after_spawn = plugin_script.as_ref().map(|plugin_script| { + let plugin_script = Arc::clone(plugin_script); + Box::new(move || plugin_script.mark_started()) as Box + }); + let result = + execute_exec_request_with_after_spawn(env, Self::stdout_stream(ctx), after_spawn) + .await + .map_err(ToolError::Codex); + self.finish_plugin_script_attempt(result, plugin_script.as_ref(), &req.cancellation_token) } } 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 b8e417f856..2414ebe67f 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -107,6 +107,7 @@ pub(super) async fn try_run_zsh_fork( attempt: &SandboxAttempt<'_>, ctx: &ToolCtx, command: &[String], + plugin_script: Option>, ) -> 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."); @@ -195,6 +196,7 @@ pub(super) async fn try_run_zsh_fork( windows_sandbox_workspace_roots, codex_linux_sandbox_exe: ctx.turn.config.codex_linux_sandbox_exe.clone(), use_legacy_landlock: ctx.turn.config.features.use_legacy_landlock(), + plugin_script, }; let main_execve_wrapper_exe = ctx .session @@ -218,6 +220,7 @@ pub(super) async fn try_run_zsh_fork( // escalation server. let stopwatch = Stopwatch::new(effective_timeout); let mut cancel_token = stopwatch.cancellation_token(); + cancel_token = cancel_when_either(cancel_token, req.cancellation_token.clone()); if let Some(cancellation) = attempt.network_denial_cancellation_token.clone() { cancel_token = cancel_when_either(cancel_token, cancellation); } @@ -252,7 +255,7 @@ pub(super) async fn try_run_zsh_fork( .await .map_err(|err| ToolError::Rejected(err.to_string()))?; - map_exec_result(attempt.sandbox, exec_result).map(Some) + map_exec_result(attempt.sandbox, exec_result, &req.cancellation_token).map(Some) } pub(crate) async fn prepare_unified_exec_zsh_fork( @@ -307,6 +310,7 @@ pub(crate) async fn prepare_unified_exec_zsh_fork( windows_sandbox_workspace_roots: exec_request.windows_sandbox_workspace_roots.clone(), codex_linux_sandbox_exe: ctx.turn.config.codex_linux_sandbox_exe.clone(), use_legacy_landlock: ctx.turn.config.features.use_legacy_landlock(), + plugin_script: None, }; let escalation_policy = CoreShellActionProvider { policy: Arc::clone(&exec_policy), @@ -815,6 +819,7 @@ struct CoreShellCommandExecutor { windows_sandbox_workspace_roots: Vec, codex_linux_sandbox_exe: Option, use_legacy_landlock: bool, + plugin_script: Option>, } struct PrepareSandboxedExecParams<'a> { @@ -872,6 +877,18 @@ impl CoreShellCommandExecutor { } } + let after_spawn = match (after_spawn, self.plugin_script.clone()) { + (Some(after_spawn), Some(plugin_script)) => Some(Box::new(move || { + after_spawn(); + plugin_script.mark_started(); + }) + as Box), + (Some(after_spawn), None) => Some(after_spawn), + (None, Some(plugin_script)) => { + Some(Box::new(move || plugin_script.mark_started()) as Box) + } + (None, None) => None, + }; let result = crate::sandboxing::execute_exec_request_with_after_spawn( crate::sandboxing::ExecRequest { command: self.command.clone(), @@ -1071,6 +1088,7 @@ fn extract_shell_script(command: &[String]) -> Result Result { let output = ExecToolCallOutput { exit_code: result.exit_code, @@ -1081,6 +1099,10 @@ fn map_exec_result( timed_out: result.timed_out, }; + if user_cancellation_token.is_cancelled() { + return Err(ToolError::Rejected("rejected by user".to_string())); + } + if result.timed_out { return Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Timeout { output: Box::new(output), 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 2b00147612..9c5fc1cd8a 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 @@ -46,6 +46,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; fn host_absolute_path(segments: &[&str]) -> String { let mut path = if cfg!(windows) { @@ -274,6 +275,7 @@ fn map_exec_result_preserves_stdout_and_stderr() { duration: Duration::from_millis(1), timed_out: false, }, + &CancellationToken::new(), ) .unwrap(); @@ -282,6 +284,31 @@ fn map_exec_result_preserves_stdout_and_stderr() { assert_eq!(out.aggregated_output.text, "outerr"); } +#[test] +fn map_exec_result_maps_user_cancelled_exit_to_rejected_by_user() { + let user_cancellation_token = CancellationToken::new(); + user_cancellation_token.cancel(); + + let result = map_exec_result( + SandboxType::None, + ExecResult { + exit_code: 1, + stdout: String::new(), + stderr: String::new(), + output: String::new(), + duration: Duration::from_millis(1), + timed_out: false, + }, + &user_cancellation_token, + ); + + assert!(matches!( + result, + Err(crate::tools::sandboxing::ToolError::Rejected(message)) + if message == "rejected by user" + )); +} + #[test] fn shell_request_escalation_execution_is_explicit() { let requested_permissions = AdditionalPermissionProfile { @@ -372,6 +399,7 @@ async fn unsandboxed_intercepted_exec_strips_managed_network_env() -> anyhow::Re windows_sandbox_workspace_roots: vec![workdir.clone()], codex_linux_sandbox_exe: None, use_legacy_landlock: false, + plugin_script: None, }; let mut env = HashMap::new(); env.insert(PROXY_ACTIVE_ENV_KEY.to_string(), "1".to_string()); 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 index 819658ecca..805a81b6f4 100644 --- a/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs +++ b/codex-rs/core/src/tools/runtimes/shell/zsh_fork_backend.rs @@ -1,4 +1,5 @@ use super::ShellRequest; +use crate::plugin_script_lifecycle::PluginScriptExecution; use crate::sandboxing::ExecRequest; use crate::tools::runtimes::unified_exec::UnifiedExecRequest; use crate::tools::sandboxing::SandboxAttempt; @@ -7,6 +8,7 @@ use crate::tools::sandboxing::ToolError; use crate::unified_exec::SpawnLifecycleHandle; use codex_protocol::exec_output::ExecToolCallOutput; use codex_tools::ZshForkConfig; +use std::sync::Arc; pub(crate) struct PreparedUnifiedExecSpawn { pub(crate) exec_request: ExecRequest, @@ -23,8 +25,9 @@ pub(crate) async fn maybe_run_shell_command( attempt: &SandboxAttempt<'_>, ctx: &ToolCtx, command: &[String], + plugin_script: Option>, ) -> Result, ToolError> { - imp::maybe_run_shell_command(req, attempt, ctx, command).await + imp::maybe_run_shell_command(req, attempt, ctx, command, plugin_script).await } /// Prepares unified exec to launch through the zsh-fork backend when the @@ -76,8 +79,9 @@ mod imp { attempt: &SandboxAttempt<'_>, ctx: &ToolCtx, command: &[String], + plugin_script: Option>, ) -> Result, ToolError> { - unix_escalation::try_run_zsh_fork(req, attempt, ctx, command).await + unix_escalation::try_run_zsh_fork(req, attempt, ctx, command, plugin_script).await } pub(super) async fn maybe_prepare_unified_exec( @@ -118,8 +122,9 @@ mod imp { attempt: &SandboxAttempt<'_>, ctx: &ToolCtx, command: &[String], + plugin_script: Option>, ) -> Result, ToolError> { - let _ = (req, attempt, ctx, command); + let _ = (req, attempt, ctx, command, plugin_script); Ok(None) }