From 5cdff40076ddbd5bb6df8fcefbfaca320851b21f Mon Sep 17 00:00:00 2001 From: Kyle Brown Date: Sat, 13 Jun 2026 05:30:22 +0000 Subject: [PATCH] Add unified exec plugin lifecycle adapter --- .../core/src/tools/runtimes/unified_exec.rs | 282 +++++++++++++++++- codex-rs/core/src/unified_exec/mod.rs | 1 - 2 files changed, 280 insertions(+), 3 deletions(-) diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 80f0990b5c..ed813f3f27 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -10,6 +10,7 @@ use crate::exec::ExecExpiration; 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::ExecServerEnvConfig; use crate::sandboxing::SandboxPermissions; @@ -38,6 +39,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 crate::unified_exec::NoopSpawnLifecycle; +use crate::unified_exec::SpawnLifecycle; +use crate::unified_exec::SpawnLifecycleHandle; use crate::unified_exec::UnifiedExecError; use crate::unified_exec::UnifiedExecProcess; use crate::unified_exec::UnifiedExecProcessManager; @@ -53,6 +56,7 @@ use codex_tools::UnifiedExecShellMode; use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; use std::collections::HashMap; +use std::sync::Arc; use tokio_util::sync::CancellationToken; use tracing::error; @@ -99,6 +103,116 @@ pub struct UnifiedExecRuntime<'a> { shell_mode: UnifiedExecShellMode, } +trait PluginScriptLifecycle: Send + Sync { + fn mark_started(&self); + fn mark_cancelled(&self); + fn finish(&self, exit_code: Option, failed: bool); +} + +impl PluginScriptLifecycle for PluginScriptExecution { + fn mark_started(&self) { + self.mark_started(); + } + + fn mark_cancelled(&self) { + self.mark_cancelled(); + } + + fn finish(&self, exit_code: Option, failed: bool) { + self.finish(exit_code, failed); + } +} + +/// Adds plugin attribution without replacing unified exec's generic spawn work. +struct PluginScriptSpawnLifecycle { + inner: SpawnLifecycleHandle, + plugin_script: Arc, +} + +impl std::fmt::Debug for PluginScriptSpawnLifecycle { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PluginScriptSpawnLifecycle") + .field("inner", &self.inner) + .field("plugin_script", &"tracked") + .finish() + } +} + +impl SpawnLifecycle for PluginScriptSpawnLifecycle { + #[cfg(unix)] + fn inherited_fds(&self) -> Vec { + self.inner.inherited_fds() + } + + fn after_spawn(&mut self) { + self.inner.after_spawn(); + self.plugin_script.mark_started(); + } + + fn mark_cancelled(&self) { + self.inner.mark_cancelled(); + self.plugin_script.mark_cancelled(); + } + + fn finish(&self, exit_code: Option, failed: bool) { + self.inner.finish(exit_code, failed); + self.plugin_script.finish(exit_code, failed); + } +} + +impl Drop for PluginScriptSpawnLifecycle { + fn drop(&mut self) { + self.plugin_script + .finish(/*exit_code*/ None, /*failed*/ true); + } +} + +fn wrap_plugin_script_lifecycle( + inner: SpawnLifecycleHandle, + plugin_script: Option>, +) -> SpawnLifecycleHandle { + match plugin_script { + Some(plugin_script) => Box::new(PluginScriptSpawnLifecycle { + inner, + plugin_script, + }), + None => inner, + } +} + +fn wrap_spawn_lifecycle( + inner: SpawnLifecycleHandle, + plugin_script: Option<&Arc>, +) -> SpawnLifecycleHandle { + wrap_plugin_script_lifecycle( + inner, + plugin_script + .map(|plugin_script| Arc::clone(plugin_script) as Arc), + ) +} + +fn should_track_plugin_script(req: &UnifiedExecRequest) -> bool { + !req.turn_environment.environment.is_remote() +} + +fn resolve_plugin_script_execution( + req: &UnifiedExecRequest, + ctx: &ToolCtx, +) -> Option> { + should_track_plugin_script(req) + .then(|| { + PluginScriptExecution::resolve( + ctx.session.as_ref(), + ctx.turn.as_ref(), + &req.hook_command, + &req.cwd, + req.shell_type, + ) + }) + .flatten() +} + fn unified_exec_options( network_denial_cancellation_token: Option, ) -> ExecOptions { @@ -323,6 +437,10 @@ impl<'a> ToolRuntime for UnifiedExecRunt if let Some(network) = managed_network { network.apply_to_env(&mut env); } + let environment_is_remote = req.turn_environment.environment.is_remote(); + // Resolve before snapshot, sandbox, PowerShell, or exec-server rewriting: + // only the original local request has safe plugin roots and cwd. + let plugin_script = resolve_plugin_script_execution(req, ctx); let explicit_env_overrides = req.explicit_env_overrides.clone(); #[cfg(unix)] let runtime_path_prepends = { @@ -408,7 +526,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt req.process_id, &prepared.exec_request, req.tty, - prepared.spawn_lifecycle, + wrap_spawn_lifecycle(prepared.spawn_lifecycle, plugin_script.as_ref()), req.turn_environment.environment.as_ref(), ) .await @@ -446,12 +564,14 @@ impl<'a> ToolRuntime for UnifiedExecRunt .env_for(command, options, managed_network) .map_err(ToolError::Codex)?; exec_env.exec_server_env_config = req.exec_server_env_config.clone(); + let spawn_lifecycle = + wrap_spawn_lifecycle(Box::new(NoopSpawnLifecycle), plugin_script.as_ref()); self.manager .open_session_with_exec_env( req.process_id, &exec_env, req.tty, - Box::new(NoopSpawnLifecycle), + spawn_lifecycle, req.turn_environment.environment.as_ref(), ) .await @@ -665,6 +785,164 @@ mod tests { } } + #[derive(Debug)] + struct RecordingLifecycle { + #[cfg(unix)] + inherited_fds: Vec, + events: Arc>>, + } + + impl SpawnLifecycle for RecordingLifecycle { + #[cfg(unix)] + fn inherited_fds(&self) -> Vec { + self.inherited_fds.clone() + } + + fn after_spawn(&mut self) { + self.events.lock().unwrap().push("inner_started"); + } + + fn mark_cancelled(&self) { + self.events.lock().unwrap().push("inner_cancelled"); + } + + fn finish(&self, _exit_code: Option, _failed: bool) { + self.events.lock().unwrap().push("inner_finished"); + } + } + + struct RecordingPluginLifecycle { + events: Arc>>, + terminal_emitted: std::sync::atomic::AtomicBool, + } + + impl PluginScriptLifecycle for RecordingPluginLifecycle { + fn mark_started(&self) { + self.events.lock().unwrap().push("plugin_started"); + } + + fn mark_cancelled(&self) { + self.events.lock().unwrap().push("plugin_cancelled"); + } + + fn finish(&self, _exit_code: Option, _failed: bool) { + if !self + .terminal_emitted + .swap(true, std::sync::atomic::Ordering::AcqRel) + { + self.events.lock().unwrap().push("plugin_finished"); + } + } + } + + fn recording_lifecycle() -> ( + PluginScriptSpawnLifecycle, + Arc>>, + ) { + let events = Arc::new(std::sync::Mutex::new(Vec::new())); + let lifecycle = PluginScriptSpawnLifecycle { + inner: Box::new(RecordingLifecycle { + #[cfg(unix)] + inherited_fds: vec![7, 11], + events: Arc::clone(&events), + }), + plugin_script: Arc::new(RecordingPluginLifecycle { + events: Arc::clone(&events), + terminal_emitted: std::sync::atomic::AtomicBool::new(false), + }), + }; + (lifecycle, events) + } + + #[test] + fn plugin_lifecycle_starts_after_inner() { + let (mut lifecycle, events) = recording_lifecycle(); + + lifecycle.after_spawn(); + + assert_eq!( + *events.lock().unwrap(), + vec!["inner_started", "plugin_started"] + ); + } + + #[cfg(unix)] + #[test] + fn plugin_lifecycle_preserves_inherited_fds() { + let (lifecycle, _events) = recording_lifecycle(); + + assert_eq!(lifecycle.inherited_fds(), vec![7, 11]); + } + + #[test] + fn plugin_lifecycle_wraps_local_shell_trackers() { + for shell_type in [ShellType::PowerShell, ShellType::Cmd, ShellType::Bash] { + let events = Arc::new(std::sync::Mutex::new(Vec::new())); + let mut lifecycle = wrap_plugin_script_lifecycle( + Box::new(RecordingLifecycle { + #[cfg(unix)] + inherited_fds: Vec::new(), + events: Arc::clone(&events), + }), + Some(Arc::new(RecordingPluginLifecycle { + events: Arc::clone(&events), + terminal_emitted: std::sync::atomic::AtomicBool::new(false), + })), + ); + + lifecycle.after_spawn(); + lifecycle.finish(Some(0), /*failed*/ false); + assert_eq!( + *events.lock().unwrap(), + vec![ + "inner_started", + "plugin_started", + "inner_finished", + "plugin_finished" + ], + "{shell_type:?} should keep plugin attribution", + ); + } + } + + #[test] + fn plugin_lifecycle_forwards_cancellation_and_idempotent_terminal_state() { + let (lifecycle, events) = recording_lifecycle(); + + lifecycle.mark_cancelled(); + lifecycle.finish(Some(0), /*failed*/ false); + lifecycle.finish(Some(9), /*failed*/ true); + drop(lifecycle); + + assert_eq!( + *events.lock().unwrap(), + vec![ + "inner_cancelled", + "plugin_cancelled", + "inner_finished", + "plugin_finished", + "inner_finished", + ] + ); + } + + #[tokio::test] + async fn remote_unified_exec_requests_do_not_track_plugin_scripts() { + let mut request = test_request( + SandboxPermissions::UseDefault, + ExecApprovalRequirement::Skip { + bypass_sandbox: false, + proposed_execpolicy_amendment: None, + }, + ); + request.environment = Arc::new( + Environment::create_for_tests(Some("ws://127.0.0.1:1/remote-exec-server".to_string())) + .expect("remote environment"), + ); + + assert!(!should_track_plugin_script(&request)); + } + fn zsh_fork_mode() -> UnifiedExecShellMode { let cwd = std::env::current_dir().expect("read current dir"); UnifiedExecShellMode::ZshFork(ZshForkConfig { diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 34d88c8149..7fb8c518c4 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -56,7 +56,6 @@ 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;