openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-04 15:08:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add generic process lifecycle hooks

Browse Source
This commit is contained in:
Kyle Brown
2026-06-12 23:02:56 +00:00
parent ce4a157adf
commit 6aedd64e5b
10 changed files with 248 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
codex-rs/core/src/exec.rs
View File

@@ -503,6 +503,7 @@ async fn get_raw_output_result(
windows_sandbox_policy_cwd,
windows_sandbox_workspace_roots,
windows_sandbox_filesystem_overrides,
after_spawn,
)
.await;
}
@@ -583,6 +584,7 @@ async fn exec_windows_sandbox(
windows_sandbox_policy_cwd: &AbsolutePathBuf,
windows_sandbox_workspace_roots: &[AbsolutePathBuf],
windows_sandbox_filesystem_overrides: Option<&WindowsSandboxFilesystemOverrides>,
after_spawn: Option<Box<dyn FnOnce() + Send>>,
) -> Result<RawExecToolCallOutput> {
use crate::config::find_codex_home;
use codex_windows_sandbox::run_windows_sandbox_capture_for_permission_profile_elevated;
@@ -662,6 +664,7 @@ async fn exec_windows_sandbox(
write_roots_override: elevated_write_roots_override.as_deref(),
deny_read_paths_override: &additional_deny_read_paths,
deny_write_paths_override: &additional_deny_write_paths,
after_spawn,
},
)
} else {
@@ -677,6 +680,7 @@ async fn exec_windows_sandbox(
&additional_deny_read_paths,
&additional_deny_write_paths,
windows_sandbox_private_desktop,
after_spawn,
)
}
})

2
codex-rs/core/src/unified_exec/async_watcher.rs
View File

@@ -124,6 +124,7 @@ pub(crate) fn spawn_exit_watcher(
let duration = Instant::now().saturating_duration_since(started_at);
if let Some(message) = process.failure_message() {
process.notify_lifecycle_finished(/*exit_code*/ None, /*failed*/ true);
emit_failed_exec_end_for_unified_exec(
session_ref,
turn_ref,
@@ -139,6 +140,7 @@ pub(crate) fn spawn_exit_watcher(
.await;
} else {
let exit_code = process.exit_code().unwrap_or(-1);
process.notify_lifecycle_finished(Some(exit_code), /*failed*/ false);
emit_exec_end_for_unified_exec(
session_ref,
turn_ref,

1
codex-rs/core/src/unified_exec/mod.rs
View File

@@ -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;

59
codex-rs/core/src/unified_exec/mod_tests.rs
View File

@@ -30,6 +30,8 @@ use pretty_assertions::assert_eq;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tokio::sync::Notify;
use tokio::sync::watch;
use tokio::time::Duration;
@@ -203,12 +205,17 @@ async fn exec_command_with_tty(
#[derive(Debug)]
struct TestSpawnLifecycle {
inherited_fds: Vec<i32>,
after_spawn: Arc<AtomicBool>,
}
impl SpawnLifecycle for TestSpawnLifecycle {
fn inherited_fds(&self) -> Vec<i32> {
self.inherited_fds.clone()
}
fn after_spawn(&mut self) {
self.after_spawn.store(true, Ordering::SeqCst);
}
}
struct BlockingTerminateExecProcess {
@@ -295,6 +302,7 @@ async fn blocking_terminate_unified_process(
}),
},
SandboxType::None,
Box::new(NoopSpawnLifecycle),
)
.await?,
))
@@ -796,6 +804,54 @@ async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Resu
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn local_spawn_lifecycle_runs_only_after_successful_creation() -> anyhow::Result<()> {
let (_, turn) = make_session_and_context().await;
#[allow(deprecated)]
let cwd = turn.cwd.clone();
let environment = codex_exec_server::Environment::default_for_tests();
let manager = UnifiedExecProcessManager::default();
let after_spawn = Arc::new(AtomicBool::new(false));
let request = test_exec_request(
&turn,
vec!["bash".to_string(), "-lc".to_string(), "true".to_string()],
cwd.clone(),
shell_env(),
);
manager
.open_session_with_exec_env(
/*process_id*/ 1234,
&request,
/*tty*/ false,
Box::new(TestSpawnLifecycle {
inherited_fds: Vec::new(),
after_spawn: Arc::clone(&after_spawn),
}),
&environment,
)
.await?;
assert!(after_spawn.load(Ordering::SeqCst));
let after_failed_spawn = Arc::new(AtomicBool::new(false));
let request = test_exec_request(&turn, Vec::new(), cwd, shell_env());
manager
.open_session_with_exec_env(
/*process_id*/ 1235,
&request,
/*tty*/ false,
Box::new(TestSpawnLifecycle {
inherited_fds: Vec::new(),
after_spawn: Arc::clone(&after_failed_spawn),
}),
&environment,
)
.await
.expect_err("missing command should fail before spawn");
assert!(!after_failed_spawn.load(Ordering::SeqCst));
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn completed_pipe_commands_preserve_exit_code() -> anyhow::Result<()> {
let (_, turn) = make_session_and_context().await;
@@ -908,6 +964,7 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<()
);
let manager = UnifiedExecProcessManager::default();
let after_spawn = Arc::new(AtomicBool::new(false));
let err = manager
.open_session_with_exec_env(
/*process_id*/ 1234,
@@ -915,6 +972,7 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<()
/*tty*/ true,
Box::new(TestSpawnLifecycle {
inherited_fds: vec![42],
after_spawn: Arc::clone(&after_spawn),
}),
turn.environments
.primary()
@@ -929,5 +987,6 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<()
err.to_string(),
"Failed to create unified exec process: remote exec-server does not support inherited file descriptors"
);
assert!(!after_spawn.load(Ordering::SeqCst));
Ok(())
}

38
codex-rs/core/src/unified_exec/process.rs
View File

@@ -33,6 +33,11 @@ use super::head_tail_buffer::HeadTailBuffer;
use super::process_state::ProcessState;
const EARLY_EXIT_GRACE_PERIOD: Duration = Duration::from_millis(150);
/// Process lifecycle hooks for successful spawn and terminal observation.
///
/// Default methods are no-op. `after_spawn` runs only after child creation
/// succeeds. The manager, watcher, and `Drop` can all observe terminal state,
/// so implementations must make terminal emission idempotent.
pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync {
/// Returns file descriptors that must stay open across the child `exec()`.
///
@@ -43,7 +48,14 @@ pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync {
Vec::new()
}
/// Releases parent-only spawn resources after successful child creation.
fn after_spawn(&mut self) {}
/// Records cancellation when orchestration stops a still-live process.
fn mark_cancelled(&self) {}
/// Records the process terminal state; implementations must be idempotent.
fn finish(&self, _exit_code: Option<i32>, _failed: bool) {}
}
pub(crate) type SpawnLifecycleHandle = Box<dyn SpawnLifecycle>;
@@ -85,7 +97,7 @@ pub(crate) struct UnifiedExecProcess {
state_rx: watch::Receiver<ProcessState>,
output_task: Option<JoinHandle<()>>,
sandbox_type: SandboxType,
_spawn_lifecycle: Option<SpawnLifecycleHandle>,
spawn_lifecycle: Option<SpawnLifecycleHandle>,
}
impl std::fmt::Debug for UnifiedExecProcess {
@@ -126,7 +138,19 @@ impl UnifiedExecProcess {
state_rx,
output_task: None,
sandbox_type,
_spawn_lifecycle: spawn_lifecycle,
spawn_lifecycle,
}
}
pub(super) fn notify_lifecycle_cancelled(&self) {
if let Some(spawn_lifecycle) = self.spawn_lifecycle.as_ref() {
spawn_lifecycle.mark_cancelled();
}
}
pub(super) fn notify_lifecycle_finished(&self, exit_code: Option<i32>, failed: bool) {
if let Some(spawn_lifecycle) = self.spawn_lifecycle.as_ref() {
spawn_lifecycle.finish(exit_code, failed);
}
}
@@ -249,6 +273,7 @@ impl UnifiedExecProcess {
if state.failure_message.is_none() {
let _ = self.state_tx.send_replace(state.failed(message));
}
self.notify_lifecycle_finished(/*exit_code*/ None, /*failed*/ true);
self.terminate();
}
@@ -375,9 +400,10 @@ impl UnifiedExecProcess {
pub(super) async fn from_exec_server_started(
started: StartedExecProcess,
sandbox_type: SandboxType,
spawn_lifecycle: SpawnLifecycleHandle,
) -> Result<Self, UnifiedExecError> {
let process_handle = ProcessHandle::ExecServer(Arc::clone(&started.process));
let mut managed = Self::new(process_handle, sandbox_type, /*spawn_lifecycle*/ None);
let mut managed = Self::new(process_handle, sandbox_type, Some(spawn_lifecycle));
let output_handles = managed.output_handles();
managed.output_task = Some(Self::spawn_exec_server_output_task(
started,
@@ -534,6 +560,12 @@ impl UnifiedExecProcess {
impl Drop for UnifiedExecProcess {
fn drop(&mut self) {
let has_exited = self.has_exited();
if !has_exited {
self.notify_lifecycle_cancelled();
}
let failed = self.state_rx.borrow().failure_message.is_some();
self.notify_lifecycle_finished(self.exit_code(), failed || !has_exited);
self.terminate();
}
}

31
codex-rs/core/src/unified_exec/process_manager.rs
View File

@@ -280,6 +280,7 @@ async fn finish_deferred_network_approval_after_process_exit_for_session(
fn fail_process_with_message(process: &UnifiedExecProcess, message: String) -> UnifiedExecError {
if let Some(message) = process.failure_message() {
process.notify_lifecycle_finished(/*exit_code*/ None, /*failed*/ true);
process.terminate();
return UnifiedExecError::process_failed(message);
}
@@ -503,6 +504,7 @@ impl UnifiedExecProcessManager {
return Err(fail_process_with_message(process.as_ref(), message));
}
if let Some(message) = process.failure_message() {
process.notify_lifecycle_finished(/*exit_code*/ None, /*failed*/ true);
let finish_result = finish_deferred_network_approval_for_session(
Some(&context.session),
deferred_network_approval.take(),
@@ -576,6 +578,7 @@ impl UnifiedExecProcessManager {
}
let exit_code = process.exit_code();
let exit = exit_code.unwrap_or(-1);
process.notify_lifecycle_finished(exit_code, /*failed*/ false);
emit_exec_end_for_unified_exec(
Arc::clone(&context.session),
Arc::clone(&context.turn),
@@ -639,6 +642,7 @@ impl UnifiedExecProcessManager {
if !request.input.is_empty() {
if !tty {
if request.input == INTERRUPT {
process.notify_lifecycle_cancelled();
process.interrupt().await?;
} else {
return Err(UnifiedExecError::StdinClosed);
@@ -655,6 +659,9 @@ impl UnifiedExecProcessManager {
if matches!(status, ProcessStatus::Exited { .. }) {
status_after_write = Some(status);
} else if matches!(err, UnifiedExecError::ProcessFailed { .. }) {
process.notify_lifecycle_finished(
/*exit_code*/ None, /*failed*/ true,
);
process.terminate();
self.release_process_id(process_id).await;
return Err(err);
@@ -704,6 +711,7 @@ impl UnifiedExecProcessManager {
return Err(fail_process_with_message(process.as_ref(), message));
}
if let Some(message) = process.failure_message() {
process.notify_lifecycle_finished(/*exit_code*/ None, /*failed*/ true);
let finish_result = finish_deferred_network_approval_for_session(
session.as_ref(),
network_approval.clone(),
@@ -870,6 +878,9 @@ impl UnifiedExecProcessManager {
// network-approval cleanup only after dropping that lock.
if let Some(pruned_entry) = pruned_entry {
unregister_network_approval_for_entry(&pruned_entry).await;
if !pruned_entry.process.has_exited() {
pruned_entry.process.notify_lifecycle_cancelled();
}
pruned_entry.process.terminate();
}
@@ -974,13 +985,11 @@ impl UnifiedExecProcessManager {
.await
}
};
let spawned =
spawned.map_err(|err| UnifiedExecError::create_process(err.to_string()))?;
spawn_lifecycle.after_spawn();
return UnifiedExecProcess::from_spawned(
spawned.map_err(|err| UnifiedExecError::create_process(err.to_string()))?,
request.sandbox,
spawn_lifecycle,
)
.await;
return UnifiedExecProcess::from_spawned(spawned, request.sandbox, spawn_lifecycle)
.await;
}
if environment.is_remote() {
if !inherited_fds.is_empty() {
@@ -995,7 +1004,12 @@ impl UnifiedExecProcessManager {
.await
.map_err(|err| UnifiedExecError::create_process(err.to_string()))?;
spawn_lifecycle.after_spawn();
return UnifiedExecProcess::from_exec_server_started(started, request.sandbox).await;
return UnifiedExecProcess::from_exec_server_started(
started,
request.sandbox,
spawn_lifecycle,
)
.await;
}
// TODO(anp): Keep PathUri through the local PTY/process launch boundary.
@@ -1322,6 +1336,9 @@ impl UnifiedExecProcessManager {
for entry in entries {
unregister_network_approval_for_entry(&entry).await;
if !entry.process.has_exited() {
entry.process.notify_lifecycle_cancelled();
}
entry.process.terminate();
}
}

109
codex-rs/core/src/unified_exec/process_tests.rs
View File

@@ -1,3 +1,5 @@
use super::process::NoopSpawnLifecycle;
use super::process::SpawnLifecycle;
use super::process::UnifiedExecProcess;
use crate::unified_exec::UnifiedExecError;
use codex_exec_server::ExecProcess;
@@ -14,10 +16,42 @@ use codex_sandboxing::SandboxType;
use pretty_assertions::assert_eq;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tokio::sync::Mutex;
use tokio::sync::watch;
use tokio::time::Duration;
#[derive(Debug, Default)]
struct RecordingLifecycleState {
cancelled: AtomicBool,
finishes: std::sync::Mutex<Vec<(Option<i32>, bool)>>,
}
#[derive(Debug)]
struct RecordingLifecycle {
state: Arc<RecordingLifecycleState>,
}
impl SpawnLifecycle for RecordingLifecycle {
fn mark_cancelled(&self) {
self.state.cancelled.store(true, Ordering::SeqCst);
}
fn finish(&self, exit_code: Option<i32>, failed: bool) {
self.state
.finishes
.lock()
.expect("finish state")
.push((exit_code, failed));
}
}
fn recording_lifecycle() -> (Arc<RecordingLifecycleState>, Box<RecordingLifecycle>) {
let state = Arc::new(RecordingLifecycleState::default());
(Arc::clone(&state), Box::new(RecordingLifecycle { state }))
}
struct MockExecProcess {
process_id: ProcessId,
write_response: WriteResponse,
@@ -86,9 +120,10 @@ impl ExecProcess for MockExecProcess {
}
}
async fn remote_process(
async fn remote_process_with_lifecycle(
write_status: WriteStatus,
terminate_error: Option<String>,
spawn_lifecycle: Box<dyn SpawnLifecycle>,
) -> UnifiedExecProcess {
let (wake_tx, _wake_rx) = watch::channel(0);
let started = StartedExecProcess {
@@ -103,11 +138,18 @@ async fn remote_process(
}),
};
UnifiedExecProcess::from_exec_server_started(started, SandboxType::None)
UnifiedExecProcess::from_exec_server_started(started, SandboxType::None, spawn_lifecycle)
.await
.expect("remote process should start")
}
async fn remote_process(
write_status: WriteStatus,
terminate_error: Option<String>,
) -> UnifiedExecProcess {
remote_process_with_lifecycle(write_status, terminate_error, Box::new(NoopSpawnLifecycle)).await
}
#[tokio::test]
async fn remote_write_unknown_process_marks_process_exited() {
let process = remote_process(WriteStatus::UnknownProcess, /*terminate_error*/ None).await;
@@ -148,6 +190,59 @@ async fn fail_and_terminate_preserves_failure_message() {
);
}
#[tokio::test]
async fn fail_and_terminate_forwards_terminal_failure() {
let (state, lifecycle) = recording_lifecycle();
let process = remote_process_with_lifecycle(WriteStatus::Accepted, None, lifecycle).await;
process.fail_and_terminate("network denied".to_string());
assert_eq!(
*state.finishes.lock().expect("finish state"),
vec![(None, true)]
);
}
#[tokio::test]
async fn dropping_live_process_marks_cancelled_and_failed() {
let (state, lifecycle) = recording_lifecycle();
let process = remote_process_with_lifecycle(WriteStatus::Accepted, None, lifecycle).await;
drop(process);
assert!(state.cancelled.load(Ordering::SeqCst));
assert_eq!(
*state.finishes.lock().expect("finish state"),
vec![(None, true)]
);
}
#[tokio::test]
async fn dropping_exited_process_does_not_mark_cancelled() {
let (state, lifecycle) = recording_lifecycle();
let process = remote_process_with_lifecycle(WriteStatus::UnknownProcess, None, lifecycle).await;
process
.write(b"hello")
.await
.expect_err("expected write failure");
drop(process);
assert!(!state.cancelled.load(Ordering::SeqCst));
}
#[tokio::test]
async fn noop_spawn_lifecycle_preserves_process_behavior() {
let process = remote_process(WriteStatus::Accepted, None).await;
process.fail_and_terminate("network denied".to_string());
assert_eq!(
process.failure_message(),
Some("network denied".to_string())
);
}
#[tokio::test]
async fn remote_terminate_confirmed_updates_state_on_success_only() {
let process = remote_process(
@@ -201,9 +296,13 @@ async fn remote_process_waits_for_early_exit_event() {
let _ = wake_tx.send(1);
});
let process = UnifiedExecProcess::from_exec_server_started(started, SandboxType::None)
.await
.expect("remote process should observe early exit");
let process = UnifiedExecProcess::from_exec_server_started(
started,
SandboxType::None,
Box::new(NoopSpawnLifecycle),
)
.await
.expect("remote process should observe early exit");
assert!(process.has_exited());
assert_eq!(process.exit_code(), Some(17));

5
codex-rs/windows-sandbox-rs/src/elevated_impl.rs
View File

@@ -20,6 +20,7 @@ pub struct ElevatedSandboxProfileCaptureRequest<'a> {
pub write_roots_override: Option<&'a [PathBuf]>,
pub deny_read_paths_override: &'a [AbsolutePathBuf],
pub deny_write_paths_override: &'a [AbsolutePathBuf],
pub after_spawn: Option<Box<dyn FnOnce() + Send>>,
}
mod windows_impl {
@@ -115,6 +116,7 @@ mod windows_impl {
write_roots_override,
deny_read_paths_override,
deny_write_paths_override,
after_spawn,
} = request;
let permissions =
ResolvedWindowsSandboxPermissions::try_from_permission_profile_for_workspace_roots(
@@ -225,6 +227,9 @@ mod windows_impl {
}
Err(err) => return Err(err),
};
if let Some(after_spawn) = after_spawn {
after_spawn();
}
let (pipe_write, mut pipe_read) = transport.into_files();
let cancel_writer = spawn_cancel_writer(&pipe_write, cancellation)?;

5
codex-rs/windows-sandbox-rs/src/lib.rs
View File

@@ -475,6 +475,7 @@ mod windows_impl {
&[],
&[],
use_private_desktop,
/*after_spawn*/ None,
)
}
@@ -491,6 +492,7 @@ mod windows_impl {
additional_deny_read_paths: &[AbsolutePathBuf],
additional_deny_write_paths: &[AbsolutePathBuf],
use_private_desktop: bool,
after_spawn: Option<Box<dyn FnOnce() + Send>>,
) -> Result<CaptureResult> {
let additional_deny_read_paths = additional_deny_read_paths
.iter()
@@ -576,6 +578,9 @@ mod windows_impl {
return Err(err);
}
};
if let Some(after_spawn) = after_spawn {
after_spawn();
}
let pi = created.process_info;
let _desktop = created;

11
codex-rs/windows-sandbox-rs/src/unified_exec/tests.rs
View File

@@ -6,6 +6,7 @@ use crate::ipc_framed::Message;
use crate::ipc_framed::decode_bytes;
use crate::ipc_framed::read_frame;
use crate::run_windows_sandbox_capture;
use crate::run_windows_sandbox_capture_with_filesystem_overrides;
use codex_protocol::models::PermissionProfile;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_pty::ProcessDriver;
@@ -425,7 +426,9 @@ fn legacy_capture_powershell_emits_output() {
let codex_home = sandbox_home("legacy-capture-pwsh");
println!("capture pwsh codex_home={}", codex_home.path().display());
let permission_profile = PermissionProfile::workspace_write();
let result = run_windows_sandbox_capture(
let spawned = Arc::new(AtomicBool::new(false));
let spawned_for_callback = Arc::clone(&spawned);
let result = run_windows_sandbox_capture_with_filesystem_overrides(
&permission_profile,
workspace_roots_for(cwd.as_path()).as_slice(),
codex_home.path(),
@@ -439,9 +442,15 @@ fn legacy_capture_powershell_emits_output() {
HashMap::new(),
Some(10_000),
/*cancellation*/ None,
/*additional_deny_read_paths*/ &[],
/*additional_deny_write_paths*/ &[],
/*use_private_desktop*/ true,
Some(Box::new(move || {
spawned_for_callback.store(true, Ordering::SeqCst);
})),
)
.expect("run legacy capture powershell");
assert!(spawned.load(Ordering::SeqCst));
println!("capture pwsh exit_code={}", result.exit_code);
println!("capture pwsh timed_out={}", result.timed_out);
let stdout = String::from_utf8_lossy(&result.stdout);