agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: exec-server prep for unified exec (#15691)

Browse Source 
This PR partially rebase `unified_exec` on the `exec-server` and adapt
the `exec-server` accordingly.

## What changed in `exec-server`

1. Replaced the old "broadcast-driven; process-global" event model with
process-scoped session events. The goal is to be able to have dedicated
handler for each process.
2. Add to protocol contract to support explicit lifecycle status and
stream ordering:
- `WriteResponse` now returns `WriteStatus` (Accepted, UnknownProcess,
StdinClosed, Starting) instead of a bool.
  - Added seq fields to output/exited notifications.
  - Added terminal process/closed notification.
3. Demultiplexed remote notifications into per-process channels. Same as
for the event sys
4. Local and remote backends now both implement ExecBackend.
5. Local backend wraps internal process ID/operations into per-process
ExecProcess objects.
6. Remote backend registers a session channel before launch and
unregisters on failed launch.

## What changed in `unified_exec`

1. Added unified process-state model and backend-neutral process
wrapper. This will probably disappear in the future, but it makes it
easier to keep the work flowing on both side.
- `UnifiedExecProcess` now handles both local PTY sessions and remote
exec-server processes through a shared `ProcessHandle`.
- Added `ProcessState` to track has_exited, exit_code, and terminal
failure message consistently across backends.
2. Routed write and lifecycle handling through process-level methods.

## Some rationals

1. The change centralizes execution transport in exec-server while
preserving policy and orchestration ownership in core, avoiding
duplicated launch approval logic. This comes from internal discussion.
2. Session-scoped events remove coupling/cross-talk between processes
and make stream ordering and terminal state explicit (seq, closed,
failed).
3. The failure-path surfacing (remote launch failures, write failures,
transport disconnects) makes command tool output and cleanup behavior
deterministic

## Follow-ups:
* Unify the concept of thread ID behind an obfuscated struct
* FD handling
* Full zsh-fork compatibility
* Full network sandboxing compatibility
* Handle ws disconnection
This commit is contained in:
jif-oai
2026-03-26 14:22:34 +00:00
committed by GitHub
Unverified
parent 4a5635b5a0
commit 7dac332c93
24 changed files with 1933 additions and 325 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/unified_exec/async_watcher.rs
+76 -14
View File
@@ -20,6 +20,7 @@ use crate::protocol::ExecCommandSource;
use crate::protocol::ExecOutputStream;
use crate::tools::events::ToolEmitter;
use crate::tools::events::ToolEventCtx;
use crate::tools::events::ToolEventFailure;
use crate::tools::events::ToolEventStage;
use crate::unified_exec::head_tail_buffer::HeadTailBuffer;
@@ -121,21 +122,36 @@ pub(crate) fn spawn_exit_watcher(
exit_token.cancelled().await;
output_drained.notified().await;
let exit_code = process.exit_code().unwrap_or(-1);
let duration = Instant::now().saturating_duration_since(started_at);
emit_exec_end_for_unified_exec(
session_ref,
turn_ref,
call_id,
command,
cwd,
Some(process_id.to_string()),
transcript,
String::new(),
exit_code,
duration,
)
.await;
if let Some(message) = process.failure_message() {
emit_failed_exec_end_for_unified_exec(
session_ref,
turn_ref,
call_id,
command,
cwd,
Some(process_id.to_string()),
transcript,
message,
duration,
)
.await;
} else {
let exit_code = process.exit_code().unwrap_or(-1);
emit_exec_end_for_unified_exec(
session_ref,
turn_ref,
call_id,
command,
cwd,
Some(process_id.to_string()),
transcript,
String::new(),
exit_code,
duration,
)
.await;
}
});
}
@@ -213,6 +229,52 @@ pub(crate) async fn emit_exec_end_for_unified_exec(
.await;
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn emit_failed_exec_end_for_unified_exec(
session_ref: Arc<Session>,
turn_ref: Arc<TurnContext>,
call_id: String,
command: Vec<String>,
cwd: PathBuf,
process_id: Option<String>,
transcript: Arc<Mutex<HeadTailBuffer>>,
message: String,
duration: Duration,
) {
let stdout = resolve_aggregated_output(&transcript, String::new()).await;
let aggregated_output = if stdout.is_empty() {
message.clone()
} else {
format!("{stdout}\n{message}")
};
let output = ExecToolCallOutput {
exit_code: -1,
stdout: StreamOutput::new(stdout),
stderr: StreamOutput::new(message),
aggregated_output: StreamOutput::new(aggregated_output),
duration,
timed_out: false,
};
let event_ctx = ToolEventCtx::new(
session_ref.as_ref(),
turn_ref.as_ref(),
&call_id,
/*turn_diff_tracker*/ None,
);
let emitter = ToolEmitter::unified_exec(
&command,
cwd,
ExecCommandSource::UnifiedExecStartup,
process_id,
);
emitter
.emit(
event_ctx,
ToolEventStage::Failure(ToolEventFailure::Output(output)),
)
.await;
}
fn split_valid_utf8_prefix(buffer: &mut Vec<u8>) -> Option<Vec<u8>> {
split_valid_utf8_prefix_with_max(buffer, UNIFIED_EXEC_OUTPUT_DELTA_MAX_BYTES)
}
codex-rs/core/src/unified_exec/errors.rs
+6
View File
@@ -5,6 +5,8 @@ use thiserror::Error;
pub(crate) enum UnifiedExecError {
#[error("Failed to create unified exec process: {message}")]
CreateProcess { message: String },
#[error("Unified exec process failed: {message}")]
ProcessFailed { message: String },
// The model is trained on `session_id`, but internally we track a `process_id`.
#[error("Unknown process id {process_id}")]
UnknownProcessId { process_id: i32 },
@@ -28,6 +30,10 @@ impl UnifiedExecError {
Self::CreateProcess { message }
}
pub(crate) fn process_failed(message: String) -> Self {
Self::ProcessFailed { message }
}
pub(crate) fn sandbox_denied(message: String, output: ExecToolCallOutput) -> Self {
Self::SandboxDenied { message, output }
}
codex-rs/core/src/unified_exec/mod.rs
+6
View File
@@ -19,6 +19,7 @@
//! This keeps policy logic and user interaction centralized while the PTY/process
//! concerns remain isolated here. The implementation is split between:
//! - `process.rs`: PTY process lifecycle + output buffering.
//! - `process_state.rs`: shared exit/failure state for local and remote processes.
//! - `process_manager.rs`: orchestration (approvals, sandboxing, reuse) and request handling.
use std::collections::HashMap;
@@ -42,6 +43,7 @@ mod errors;
mod head_tail_buffer;
mod process;
mod process_manager;
mod process_state;
pub(crate) fn set_deterministic_process_ids_for_tests(enabled: bool) {
process_manager::set_deterministic_process_ids_for_tests(enabled);
@@ -167,6 +169,10 @@ pub(crate) fn generate_chunk_id() -> String {
.collect()
}
#[cfg(test)]
#[cfg(unix)]
#[path = "process_tests.rs"]
mod process_tests;
#[cfg(test)]
#[cfg(unix)]
#[path = "mod_tests.rs"]
codex-rs/core/src/unified_exec/mod_tests.rs
+295 -48
View File
@@ -3,27 +3,26 @@ use super::*;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::codex::make_session_and_context;
use crate::protocol::AskForApproval;
use crate::protocol::SandboxPolicy;
use crate::exec::ExecCapturePolicy;
use crate::exec::ExecExpiration;
use crate::sandboxing::ExecRequest;
use crate::tools::context::ExecCommandToolOutput;
use crate::unified_exec::ExecCommandRequest;
use crate::unified_exec::WriteStdinRequest;
use crate::unified_exec::process::OutputHandles;
use codex_sandboxing::SandboxType;
use codex_utils_output_truncation::approx_token_count;
use core_test_support::get_remote_test_env;
use core_test_support::skip_if_sandbox;
use core_test_support::test_codex::test_env as remote_test_env;
use pretty_assertions::assert_eq;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::time::Duration;
use tokio::time::Instant;
async fn test_session_and_turn() -> (Arc<Session>, Arc<TurnContext>) {
let (session, mut turn) = make_session_and_context().await;
turn.approval_policy
.set(AskForApproval::Never)
.expect("test setup should allow updating approval policy");
turn.sandbox_policy
.set(SandboxPolicy::DangerFullAccess)
.expect("test setup should allow updating sandbox policy");
turn.file_system_sandbox_policy =
codex_protocol::permissions::FileSystemSandboxPolicy::from(turn.sandbox_policy.get());
turn.network_sandbox_policy =
codex_protocol::permissions::NetworkSandboxPolicy::from(turn.sandbox_policy.get());
let (session, turn) = make_session_and_context().await;
(Arc::new(session), Arc::new(turn))
}
@@ -32,36 +31,143 @@ async fn exec_command(
turn: &Arc<TurnContext>,
cmd: &str,
yield_time_ms: u64,
workdir: Option<PathBuf>,
) -> Result<ExecCommandToolOutput, UnifiedExecError> {
exec_command_with_tty(session, turn, cmd, yield_time_ms, workdir, true).await
}
fn shell_env() -> HashMap<String, String> {
std::env::vars().collect()
}
fn test_exec_request(
turn: &TurnContext,
command: Vec<String>,
cwd: PathBuf,
env: HashMap<String, String>,
) -> ExecRequest {
let windows_sandbox_private_desktop = false;
let sandbox_policy = turn.sandbox_policy.get().clone();
let file_system_sandbox_policy = turn.file_system_sandbox_policy.clone();
let network_sandbox_policy = turn.network_sandbox_policy;
let network = None;
let arg0 = None;
ExecRequest::new(
command,
cwd,
env,
network,
ExecExpiration::DefaultTimeout,
ExecCapturePolicy::ShellTool,
SandboxType::None,
turn.windows_sandbox_level,
windows_sandbox_private_desktop,
sandbox_policy,
file_system_sandbox_policy,
network_sandbox_policy,
arg0,
)
}
async fn exec_command_with_tty(
session: &Arc<Session>,
turn: &Arc<TurnContext>,
cmd: &str,
yield_time_ms: u64,
workdir: Option<PathBuf>,
tty: bool,
) -> Result<ExecCommandToolOutput, UnifiedExecError> {
let manager = &session.services.unified_exec_manager;
let process_id = manager.allocate_process_id().await;
let cwd = workdir.unwrap_or_else(|| turn.cwd.clone().to_path_buf());
let command = vec!["bash".to_string(), "-lc".to_string(), cmd.to_string()];
let request = test_exec_request(turn, command.clone(), cwd.clone(), shell_env());
let process = Arc::new(
manager
.open_session_with_exec_env(
process_id,
&request,
tty,
Box::new(NoopSpawnLifecycle),
turn.environment.as_ref(),
)
.await?,
);
let context =
UnifiedExecContext::new(Arc::clone(session), Arc::clone(turn), "call".to_string());
let process_id = session
.services
.unified_exec_manager
.allocate_process_id()
.await;
let started_at = Instant::now();
let process_started_alive = !process.has_exited() && process.exit_code().is_none();
if process_started_alive {
let entry = ProcessEntry {
process: Arc::clone(&process),
call_id: context.call_id.clone(),
process_id,
command: command.clone(),
tty,
network_approval_id: None,
session: Arc::downgrade(session),
last_used: started_at,
};
manager
.process_store
.lock()
.await
.processes
.insert(process_id, entry);
}
session
.services
.unified_exec_manager
.exec_command(
ExecCommandRequest {
command: vec!["bash".to_string(), "-lc".to_string(), cmd.to_string()],
process_id,
yield_time_ms,
max_output_tokens: None,
workdir: None,
network: None,
tty: true,
sandbox_permissions: SandboxPermissions::UseDefault,
additional_permissions: None,
additional_permissions_preapproved: false,
justification: None,
prefix_rule: None,
},
&context,
)
.await
let OutputHandles {
output_buffer,
output_notify,
output_closed,
output_closed_notify,
cancellation_token,
} = process.output_handles();
let deadline = started_at + Duration::from_millis(yield_time_ms);
let collected = UnifiedExecProcessManager::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
&output_closed_notify,
&cancellation_token,
Some(session.subscribe_out_of_band_elicitation_pause_state()),
deadline,
)
.await;
let wall_time = Instant::now().saturating_duration_since(started_at);
let text = String::from_utf8_lossy(&collected).to_string();
let has_exited = process.has_exited();
let exit_code = process.exit_code();
let response_process_id = if process_started_alive && !has_exited {
Some(process_id)
} else {
manager.release_process_id(process_id).await;
None
};
Ok(ExecCommandToolOutput {
event_call_id: context.call_id,
chunk_id: generate_chunk_id(),
wall_time,
raw_output: collected,
max_output_tokens: None,
process_id: response_process_id,
exit_code,
original_token_count: Some(approx_token_count(&text)),
session_command: Some(command),
})
}
#[derive(Debug)]
struct TestSpawnLifecycle {
inherited_fds: Vec<i32>,
}
impl SpawnLifecycle for TestSpawnLifecycle {
fn inherited_fds(&self) -> Vec<i32> {
self.inherited_fds.clone()
}
}
async fn write_stdin(
@@ -121,7 +227,7 @@ async fn unified_exec_persists_across_requests() -> anyhow::Result<()> {
let (session, turn) = test_session_and_turn().await;
let open_shell = exec_command(&session, &turn, "bash -i", 2_500).await?;
let open_shell = exec_command(&session, &turn, "bash -i", 2_500, None).await?;
let process_id = open_shell.process_id.expect("expected process_id");
write_stdin(
@@ -153,7 +259,7 @@ async fn multi_unified_exec_sessions() -> anyhow::Result<()> {
let (session, turn) = test_session_and_turn().await;
let shell_a = exec_command(&session, &turn, "bash -i", 2_500).await?;
let shell_a = exec_command(&session, &turn, "bash -i", 2_500, None).await?;
let session_a = shell_a.process_id.expect("expected process id");
write_stdin(
@@ -164,7 +270,14 @@ async fn multi_unified_exec_sessions() -> anyhow::Result<()> {
)
.await?;
let out_2 = exec_command(&session, &turn, "echo $CODEX_INTERACTIVE_SHELL_VAR", 2_500).await?;
let out_2 = exec_command(
&session,
&turn,
"echo $CODEX_INTERACTIVE_SHELL_VAR",
2_500,
None,
)
.await?;
tokio::time::sleep(Duration::from_secs(2)).await;
assert!(
out_2.process_id.is_none(),
@@ -198,7 +311,7 @@ async fn unified_exec_timeouts() -> anyhow::Result<()> {
let (session, turn) = test_session_and_turn().await;
let open_shell = exec_command(&session, &turn, "bash -i", 2_500).await?;
let open_shell = exec_command(&session, &turn, "bash -i", 2_500, None).await?;
let process_id = open_shell.process_id.expect("expected process id");
write_stdin(
@@ -247,7 +360,14 @@ async fn unified_exec_pause_blocks_yield_timeout() -> anyhow::Result<()> {
});
let started = tokio::time::Instant::now();
let response = exec_command(&session, &turn, "sleep 1 && echo unified-exec-done", 250).await?;
let response = exec_command(
&session,
&turn,
"sleep 1 && echo unified-exec-done",
250,
None,
)
.await?;
assert!(
started.elapsed() >= Duration::from_secs(2),
@@ -270,7 +390,7 @@ async fn unified_exec_pause_blocks_yield_timeout() -> anyhow::Result<()> {
async fn requests_with_large_timeout_are_capped() -> anyhow::Result<()> {
let (session, turn) = test_session_and_turn().await;
let result = exec_command(&session, &turn, "echo codex", 120_000).await?;
let result = exec_command(&session, &turn, "echo codex", 120_000, None).await?;
assert!(result.process_id.is_some());
assert!(result.truncated_output().contains("codex"));
@@ -282,7 +402,7 @@ async fn requests_with_large_timeout_are_capped() -> anyhow::Result<()> {
#[ignore] // Ignored while we have a better way to test this.
async fn completed_commands_do_not_persist_sessions() -> anyhow::Result<()> {
let (session, turn) = test_session_and_turn().await;
let result = exec_command(&session, &turn, "echo codex", 2_500).await?;
let result = exec_command(&session, &turn, "echo codex", 2_500, None).await?;
assert!(
result.process_id.is_some(),
@@ -310,7 +430,7 @@ async fn reusing_completed_process_returns_unknown_process() -> anyhow::Result<(
let (session, turn) = test_session_and_turn().await;
let open_shell = exec_command(&session, &turn, "bash -i", 2_500).await?;
let open_shell = exec_command(&session, &turn, "bash -i", 2_500, None).await?;
let process_id = open_shell.process_id.expect("expected process id");
write_stdin(&session, process_id, "exit\n", 2_500).await?;
@@ -341,3 +461,130 @@ async fn reusing_completed_process_returns_unknown_process() -> anyhow::Result<(
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;
let request = test_exec_request(
&turn,
vec!["bash".to_string(), "-lc".to_string(), "exit 17".to_string()],
PathBuf::from("/tmp"),
shell_env(),
);
let environment = codex_exec_server::Environment::default();
let process = UnifiedExecProcessManager::default()
.open_session_with_exec_env(
1234,
&request,
false,
Box::new(NoopSpawnLifecycle),
&environment,
)
.await?;
if !process.has_exited() {
let exit_signal = process.cancellation_token();
assert!(
tokio::time::timeout(Duration::from_secs(2), exit_signal.cancelled())
.await
.is_ok(),
"process did not report exit within timeout"
);
}
assert!(process.has_exited());
assert_eq!(process.exit_code(), Some(17));
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unified_exec_uses_remote_exec_server_when_configured() -> anyhow::Result<()> {
skip_if_sandbox!(Ok(()));
let Some(_remote_env) = get_remote_test_env() else {
return Ok(());
};
let remote_test_env = remote_test_env().await?;
let (_, turn) = make_session_and_context().await;
let request = test_exec_request(
&turn,
vec!["bash".to_string(), "-i".to_string()],
PathBuf::from("/tmp"),
shell_env(),
);
let manager = UnifiedExecProcessManager::default();
let process = manager
.open_session_with_exec_env(
1234,
&request,
true,
Box::new(NoopSpawnLifecycle),
remote_test_env.environment(),
)
.await?;
process.write(b"printf 'remote-unified-exec\\n'\n").await?;
tokio::time::sleep(Duration::from_millis(100)).await;
let crate::unified_exec::process::OutputHandles {
output_buffer,
output_notify,
output_closed,
output_closed_notify,
cancellation_token,
} = process.output_handles();
let collected = UnifiedExecProcessManager::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
&output_closed_notify,
&cancellation_token,
None,
Instant::now() + Duration::from_millis(2_500),
)
.await;
assert!(String::from_utf8_lossy(&collected).contains("remote-unified-exec"));
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<()> {
skip_if_sandbox!(Ok(()));
let Some(_remote_env) = get_remote_test_env() else {
return Ok(());
};
let remote_test_env = remote_test_env().await?;
let (_, mut turn) = make_session_and_context().await;
turn.environment = Arc::new(remote_test_env.environment().clone());
let request = test_exec_request(
&turn,
vec!["bash".to_string(), "-lc".to_string(), "echo ok".to_string()],
PathBuf::from("/tmp"),
shell_env(),
);
let manager = UnifiedExecProcessManager::default();
let err = manager
.open_session_with_exec_env(
1234,
&request,
true,
Box::new(TestSpawnLifecycle {
inherited_fds: vec![42],
}),
turn.environment.as_ref(),
)
.await
.expect_err("expected inherited fd rejection");
assert_eq!(
err.to_string(),
"Failed to create unified exec process: remote exec-server does not support inherited file descriptors"
);
Ok(())
}
codex-rs/core/src/unified_exec/process.rs
+282 -58
View File
@@ -6,8 +6,8 @@ use std::sync::atomic::Ordering;
use tokio::sync::Mutex;
use tokio::sync::Notify;
use tokio::sync::broadcast;
use tokio::sync::mpsc;
use tokio::sync::oneshot::error::TryRecvError;
use tokio::sync::watch;
use tokio::task::JoinHandle;
use tokio::time::Duration;
use tokio_util::sync::CancellationToken;
@@ -15,8 +15,12 @@ use tokio_util::sync::CancellationToken;
use crate::exec::ExecToolCallOutput;
use crate::exec::StreamOutput;
use crate::exec::is_likely_sandbox_denied;
use codex_exec_server::ExecProcess;
use codex_exec_server::ReadResponse as ExecReadResponse;
use codex_exec_server::StartedExecProcess;
use codex_exec_server::WriteStatus;
use codex_protocol::protocol::TruncationPolicy;
use codex_sandboxing::SandboxType;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::formatted_truncate_text;
use codex_utils_pty::ExecCommandSession;
use codex_utils_pty::SpawnedPty;
@@ -24,6 +28,9 @@ use codex_utils_pty::SpawnedPty;
use super::UNIFIED_EXEC_OUTPUT_MAX_TOKENS;
use super::UnifiedExecError;
use super::head_tail_buffer::HeadTailBuffer;
use super::process_state::ProcessState;
const EARLY_EXIT_GRACE_PERIOD: Duration = Duration::from_millis(150);
pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync {
/// Returns file descriptors that must stay open across the child `exec()`.
@@ -41,11 +48,13 @@ pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync {
pub(crate) type SpawnLifecycleHandle = Box<dyn SpawnLifecycle>;
#[derive(Debug, Default)]
/// Spawn lifecycle that performs no extra setup around process launch.
pub(crate) struct NoopSpawnLifecycle;
impl SpawnLifecycle for NoopSpawnLifecycle {}
pub(crate) type OutputBuffer = Arc<Mutex<HeadTailBuffer>>;
/// Shared output state exposed to polling and streaming consumers.
pub(crate) struct OutputHandles {
pub(crate) output_buffer: OutputBuffer,
pub(crate) output_notify: Arc<Notify>,
@@ -54,27 +63,44 @@ pub(crate) struct OutputHandles {
pub(crate) cancellation_token: CancellationToken,
}
#[derive(Debug)]
/// Transport-specific process handle used by unified exec.
enum ProcessHandle {
Local(Box<ExecCommandSession>),
Remote(Arc<dyn ExecProcess>),
}
/// Unified wrapper over local PTY sessions and exec-server-backed processes.
pub(crate) struct UnifiedExecProcess {
process_handle: ExecCommandSession,
output_rx: broadcast::Receiver<Vec<u8>>,
process_handle: ProcessHandle,
output_tx: broadcast::Sender<Vec<u8>>,
output_buffer: OutputBuffer,
output_notify: Arc<Notify>,
output_closed: Arc<AtomicBool>,
output_closed_notify: Arc<Notify>,
cancellation_token: CancellationToken,
output_drained: Arc<Notify>,
output_task: JoinHandle<()>,
state_tx: watch::Sender<ProcessState>,
state_rx: watch::Receiver<ProcessState>,
output_task: Option<JoinHandle<()>>,
sandbox_type: SandboxType,
_spawn_lifecycle: SpawnLifecycleHandle,
_spawn_lifecycle: Option<SpawnLifecycleHandle>,
}
impl std::fmt::Debug for UnifiedExecProcess {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("UnifiedExecProcess")
.field("has_exited", &self.has_exited())
.field("exit_code", &self.exit_code())
.field("sandbox_type", &self.sandbox_type)
.finish_non_exhaustive()
}
}
impl UnifiedExecProcess {
pub(super) fn new(
process_handle: ExecCommandSession,
initial_output_rx: tokio::sync::broadcast::Receiver<Vec<u8>>,
fn new(
process_handle: ProcessHandle,
sandbox_type: SandboxType,
spawn_lifecycle: SpawnLifecycleHandle,
spawn_lifecycle: Option<SpawnLifecycleHandle>,
) -> Self {
let output_buffer = Arc::new(Mutex::new(HeadTailBuffer::default()));
let output_notify = Arc::new(Notify::new());
@@ -82,48 +108,49 @@ impl UnifiedExecProcess {
let output_closed_notify = Arc::new(Notify::new());
let cancellation_token = CancellationToken::new();
let output_drained = Arc::new(Notify::new());
let mut receiver = initial_output_rx;
let output_rx = receiver.resubscribe();
let buffer_clone = Arc::clone(&output_buffer);
let notify_clone = Arc::clone(&output_notify);
let output_closed_clone = Arc::clone(&output_closed);
let output_closed_notify_clone = Arc::clone(&output_closed_notify);
let output_task = tokio::spawn(async move {
loop {
match receiver.recv().await {
Ok(chunk) => {
let mut guard = buffer_clone.lock().await;
guard.push_chunk(chunk);
drop(guard);
notify_clone.notify_waiters();
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
output_closed_clone.store(true, Ordering::Release);
output_closed_notify_clone.notify_waiters();
break;
}
};
}
});
let (output_tx, _) = broadcast::channel(64);
let (state_tx, state_rx) = watch::channel(ProcessState::default());
Self {
process_handle,
output_rx,
output_tx,
output_buffer,
output_notify,
output_closed,
output_closed_notify,
cancellation_token,
output_drained,
output_task,
state_tx,
state_rx,
output_task: None,
sandbox_type,
_spawn_lifecycle: spawn_lifecycle,
}
}
pub(super) fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> {
self.process_handle.writer_sender()
pub(super) async fn write(&self, data: &[u8]) -> Result<(), UnifiedExecError> {
match &self.process_handle {
ProcessHandle::Local(process_handle) => process_handle
.writer_sender()
.send(data.to_vec())
.await
.map_err(|_| UnifiedExecError::WriteToStdin),
ProcessHandle::Remote(process_handle) => {
match process_handle.write(data.to_vec()).await {
Ok(response) => match response.status {
WriteStatus::Accepted => Ok(()),
WriteStatus::UnknownProcess | WriteStatus::StdinClosed => {
let state = self.state_rx.borrow().clone();
let _ = self.state_tx.send_replace(state.exited(state.exit_code));
self.cancellation_token.cancel();
Err(UnifiedExecError::WriteToStdin)
}
WriteStatus::Starting => Err(UnifiedExecError::WriteToStdin),
},
Err(err) => Err(UnifiedExecError::process_failed(err.to_string())),
}
}
}
}
pub(super) fn output_handles(&self) -> OutputHandles {
@@ -137,7 +164,7 @@ impl UnifiedExecProcess {
}
pub(super) fn output_receiver(&self) -> tokio::sync::broadcast::Receiver<Vec<u8>> {
self.output_rx.resubscribe()
self.output_tx.subscribe()
}
pub(super) fn cancellation_token(&self) -> CancellationToken {
@@ -149,19 +176,39 @@ impl UnifiedExecProcess {
}
pub(super) fn has_exited(&self) -> bool {
self.process_handle.has_exited()
let state = self.state_rx.borrow().clone();
match &self.process_handle {
ProcessHandle::Local(process_handle) => state.has_exited || process_handle.has_exited(),
ProcessHandle::Remote(_) => state.has_exited,
}
}
pub(super) fn exit_code(&self) -> Option<i32> {
self.process_handle.exit_code()
let state = self.state_rx.borrow().clone();
match &self.process_handle {
ProcessHandle::Local(process_handle) => {
state.exit_code.or_else(|| process_handle.exit_code())
}
ProcessHandle::Remote(_) => state.exit_code,
}
}
pub(super) fn terminate(&self) {
self.output_closed.store(true, Ordering::Release);
self.output_closed_notify.notify_waiters();
self.process_handle.terminate();
match &self.process_handle {
ProcessHandle::Local(process_handle) => process_handle.terminate(),
ProcessHandle::Remote(process_handle) => {
let process_handle = Arc::clone(process_handle);
tokio::spawn(async move {
let _ = process_handle.terminate().await;
});
}
}
self.cancellation_token.cancel();
self.output_task.abort();
if let Some(output_task) = &self.output_task {
output_task.abort();
}
}
async fn snapshot_output(&self) -> Vec<Vec<u8>> {
@@ -173,6 +220,10 @@ impl UnifiedExecProcess {
self.sandbox_type
}
pub(super) fn failure_message(&self) -> Option<String> {
self.state_rx.borrow().failure_message.clone()
}
pub(super) async fn check_for_sandbox_denial(&self) -> Result<(), UnifiedExecError> {
let _ =
tokio::time::timeout(Duration::from_millis(20), self.output_notify.notified()).await;
@@ -232,29 +283,47 @@ impl UnifiedExecProcess {
mut exit_rx,
} = spawned;
let output_rx = codex_utils_pty::combine_output_receivers(stdout_rx, stderr_rx);
let managed = Self::new(process_handle, output_rx, sandbox_type, spawn_lifecycle);
let mut managed = Self::new(
ProcessHandle::Local(Box::new(process_handle)),
sandbox_type,
Some(spawn_lifecycle),
);
managed.output_task = Some(Self::spawn_local_output_task(
output_rx,
Arc::clone(&managed.output_buffer),
Arc::clone(&managed.output_notify),
Arc::clone(&managed.output_closed),
Arc::clone(&managed.output_closed_notify),
managed.output_tx.clone(),
));
let exit_ready = matches!(exit_rx.try_recv(), Ok(_) | Err(TryRecvError::Closed));
if exit_ready {
managed.signal_exit();
managed.check_for_sandbox_denial().await?;
return Ok(managed);
match exit_rx.try_recv() {
Ok(exit_code) => {
managed.signal_exit(Some(exit_code));
managed.check_for_sandbox_denial().await?;
return Ok(managed);
}
Err(TryRecvError::Closed) => {
managed.signal_exit(/*exit_code*/ None);
managed.check_for_sandbox_denial().await?;
return Ok(managed);
}
Err(TryRecvError::Empty) => {}
}
if tokio::time::timeout(Duration::from_millis(150), &mut exit_rx)
.await
.is_ok()
{
managed.signal_exit();
if let Ok(exit_result) = tokio::time::timeout(EARLY_EXIT_GRACE_PERIOD, &mut exit_rx).await {
managed.signal_exit(exit_result.ok());
managed.check_for_sandbox_denial().await?;
return Ok(managed);
}
tokio::spawn({
let state_tx = managed.state_tx.clone();
let cancellation_token = managed.cancellation_token.clone();
async move {
let _ = exit_rx.await;
let exit_code = exit_rx.await.ok();
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.exited(exit_code));
cancellation_token.cancel();
}
});
@@ -262,7 +331,162 @@ impl UnifiedExecProcess {
Ok(managed)
}
fn signal_exit(&self) {
pub(super) async fn from_remote_started(
started: StartedExecProcess,
sandbox_type: SandboxType,
) -> Result<Self, UnifiedExecError> {
let process_handle = ProcessHandle::Remote(Arc::clone(&started.process));
let mut managed = Self::new(process_handle, sandbox_type, /*spawn_lifecycle*/ None);
let output_handles = managed.output_handles();
managed.output_task = Some(Self::spawn_remote_output_task(
started,
output_handles,
managed.output_tx.clone(),
managed.state_tx.clone(),
));
let mut state_rx = managed.state_rx.clone();
if tokio::time::timeout(EARLY_EXIT_GRACE_PERIOD, async {
loop {
let state = state_rx.borrow().clone();
if state.has_exited || state.failure_message.is_some() {
break;
}
if state_rx.changed().await.is_err() {
break;
}
}
})
.await
.is_ok()
{
managed.check_for_sandbox_denial().await?;
}
Ok(managed)
}
fn spawn_remote_output_task(
started: StartedExecProcess,
output_handles: OutputHandles,
output_tx: broadcast::Sender<Vec<u8>>,
state_tx: watch::Sender<ProcessState>,
) -> JoinHandle<()> {
let OutputHandles {
output_buffer,
output_notify,
output_closed,
output_closed_notify,
cancellation_token,
} = output_handles;
let process = started.process;
let mut wake_rx = process.subscribe_wake();
tokio::spawn(async move {
let mut after_seq = None;
loop {
match process
.read(after_seq, /*max_bytes*/ None, /*wait_ms*/ Some(0))
.await
{
Ok(response) => {
let ExecReadResponse {
chunks,
next_seq,
exited,
exit_code,
closed,
failure,
} = response;
for chunk in chunks {
let bytes = chunk.chunk.into_inner();
let mut guard = output_buffer.lock().await;
guard.push_chunk(bytes.clone());
drop(guard);
let _ = output_tx.send(bytes);
output_notify.notify_waiters();
}
if let Some(message) = failure {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.failed(message));
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
break;
}
if exited {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.exited(exit_code));
}
if closed {
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
}
after_seq = next_seq.checked_sub(1);
if output_closed.load(Ordering::Acquire) {
break;
}
}
Err(err) => {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.failed(err.to_string()));
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
break;
}
}
if wake_rx.changed().await.is_err() {
let state = state_tx.borrow().clone();
let _ = state_tx
.send_replace(state.failed("exec-server wake channel closed".to_string()));
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
break;
}
}
})
}
fn spawn_local_output_task(
mut receiver: tokio::sync::broadcast::Receiver<Vec<u8>>,
buffer: OutputBuffer,
output_notify: Arc<Notify>,
output_closed: Arc<AtomicBool>,
output_closed_notify: Arc<Notify>,
output_tx: broadcast::Sender<Vec<u8>>,
) -> JoinHandle<()> {
tokio::spawn(async move {
loop {
match receiver.recv().await {
Ok(chunk) => {
let mut guard = buffer.lock().await;
guard.push_chunk(chunk.clone());
drop(guard);
let _ = output_tx.send(chunk);
output_notify.notify_waiters();
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
break;
}
};
}
})
}
fn signal_exit(&self, exit_code: Option<i32>) {
let state = self.state_rx.borrow().clone();
let _ = self.state_tx.send_replace(state.exited(exit_code));
self.cancellation_token.cancel();
}
}
codex-rs/core/src/unified_exec/process_manager.rs
+86 -19
View File
@@ -7,7 +7,6 @@ use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tokio::sync::Notify;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::time::Duration;
use tokio::time::Instant;
@@ -40,6 +39,7 @@ use crate::unified_exec::UnifiedExecProcessManager;
use crate::unified_exec::WARNING_UNIFIED_EXEC_PROCESSES;
use crate::unified_exec::WriteStdinRequest;
use crate::unified_exec::async_watcher::emit_exec_end_for_unified_exec;
use crate::unified_exec::async_watcher::emit_failed_exec_end_for_unified_exec;
use crate::unified_exec::async_watcher::spawn_exit_watcher;
use crate::unified_exec::async_watcher::start_streaming_output;
use crate::unified_exec::clamp_yield_time;
@@ -89,8 +89,9 @@ fn apply_unified_exec_env(mut env: HashMap<String, String>) -> HashMap<String, S
env
}
/// Borrowed process state prepared for a `write_stdin` or poll operation.
struct PreparedProcessHandles {
writer_tx: mpsc::Sender<Vec<u8>>,
process: Arc<UnifiedExecProcess>,
output_buffer: OutputBuffer,
output_notify: Arc<Notify>,
output_closed: Arc<AtomicBool>,
@@ -102,6 +103,10 @@ struct PreparedProcessHandles {
tty: bool,
}
fn exec_server_process_id(process_id: i32) -> String {
process_id.to_string()
}
impl UnifiedExecProcessManager {
pub(crate) async fn allocate_process_id(&self) -> i32 {
loop {
@@ -243,6 +248,29 @@ impl UnifiedExecProcessManager {
let text = String::from_utf8_lossy(&collected).to_string();
let chunk_id = generate_chunk_id();
if let Some(message) = process.failure_message() {
if !process_started_alive {
emit_failed_exec_end_for_unified_exec(
Arc::clone(&context.session),
Arc::clone(&context.turn),
context.call_id.clone(),
request.command.clone(),
cwd.clone(),
Some(request.process_id.to_string()),
Arc::clone(&transcript),
message.clone(),
wall_time,
)
.await;
}
self.release_process_id(request.process_id).await;
finish_deferred_network_approval(
context.session.as_ref(),
deferred_network_approval.take(),
)
.await;
return Err(UnifiedExecError::process_failed(message));
}
let process_id = request.process_id;
let (response_process_id, exit_code) = if process_started_alive {
match self.refresh_process_state(process_id).await {
@@ -312,7 +340,7 @@ impl UnifiedExecProcessManager {
let process_id = request.process_id;
let PreparedProcessHandles {
writer_tx,
process,
output_buffer,
output_notify,
output_closed,
@@ -324,15 +352,31 @@ impl UnifiedExecProcessManager {
tty,
..
} = self.prepare_process_handles(process_id).await?;
let mut status_after_write = None;
if !request.input.is_empty() {
if !tty {
return Err(UnifiedExecError::StdinClosed);
}
Self::send_input(&writer_tx, request.input.as_bytes()).await?;
// Give the remote process a brief window to react so that we are
// more likely to capture its output in the poll below.
tokio::time::sleep(Duration::from_millis(100)).await;
match process.write(request.input.as_bytes()).await {
Ok(()) => {
// Give the remote process a brief window to react so that we are
// more likely to capture its output in the poll below.
tokio::time::sleep(Duration::from_millis(100)).await;
}
Err(err) => {
let status = self.refresh_process_state(process_id).await;
if matches!(status, ProcessStatus::Exited { .. }) {
status_after_write = Some(status);
} else if matches!(err, UnifiedExecError::ProcessFailed { .. }) {
process.terminate();
self.release_process_id(process_id).await;
return Err(err);
} else {
return Err(err);
}
}
}
}
let yield_time_ms = {
@@ -362,12 +406,20 @@ impl UnifiedExecProcessManager {
let text = String::from_utf8_lossy(&collected).to_string();
let original_token_count = approx_token_count(&text);
let chunk_id = generate_chunk_id();
if let Some(message) = process.failure_message() {
self.release_process_id(process_id).await;
return Err(UnifiedExecError::process_failed(message));
}
// After polling, refresh_process_state tells us whether the PTY is
// still alive or has exited and been removed from the store; we thread
// that through so the handler can tag TerminalInteraction with an
// appropriate process_id and exit_code.
let status = self.refresh_process_state(process_id).await;
let status = if let Some(status) = status_after_write {
status
} else {
self.refresh_process_state(process_id).await
};
let (process_id, exit_code, event_call_id) = match status {
ProcessStatus::Alive {
exit_code,
@@ -455,7 +507,7 @@ impl UnifiedExecProcessManager {
.map(|session| session.subscribe_out_of_band_elicitation_pause_state());
Ok(PreparedProcessHandles {
writer_tx: entry.process.writer_sender(),
process: Arc::clone(&entry.process),
output_buffer,
output_notify,
output_closed,
@@ -468,16 +520,6 @@ impl UnifiedExecProcessManager {
})
}
async fn send_input(
writer_tx: &mpsc::Sender<Vec<u8>>,
data: &[u8],
) -> Result<(), UnifiedExecError> {
writer_tx
.send(data.to_vec())
.await
.map_err(|_| UnifiedExecError::WriteToStdin)
}
#[allow(clippy::too_many_arguments)]
async fn store_process(
&self,
@@ -539,9 +581,11 @@ impl UnifiedExecProcessManager {
pub(crate) async fn open_session_with_exec_env(
&self,
process_id: i32,
env: &ExecRequest,
tty: bool,
mut spawn_lifecycle: SpawnLifecycleHandle,
environment: &codex_exec_server::Environment,
) -> Result<UnifiedExecProcess, UnifiedExecError> {
let (program, args) = env
.command
@@ -549,6 +593,28 @@ impl UnifiedExecProcessManager {
.ok_or(UnifiedExecError::MissingCommandLine)?;
let inherited_fds = spawn_lifecycle.inherited_fds();
if environment.exec_server_url().is_some() {
if !inherited_fds.is_empty() {
return Err(UnifiedExecError::create_process(
"remote exec-server does not support inherited file descriptors".to_string(),
));
}
let started = environment
.get_exec_backend()
.start(codex_exec_server::ExecParams {
process_id: exec_server_process_id(process_id),
argv: env.command.clone(),
cwd: env.cwd.clone(),
env: env.env.clone(),
tty,
arg0: env.arg0.clone(),
})
.await
.map_err(|err| UnifiedExecError::create_process(err.to_string()))?;
return UnifiedExecProcess::from_remote_started(started, env.sandbox).await;
}
let spawn_result = if tty {
codex_utils_pty::pty::spawn_process_with_inherited_fds(
program,
@@ -611,6 +677,7 @@ impl UnifiedExecProcessManager {
.await;
let req = UnifiedExecToolRequest {
command: request.command.clone(),
process_id: request.process_id,
cwd,
env,
explicit_env_overrides: context.turn.shell_environment_policy.r#set.clone(),
codex-rs/core/src/unified_exec/process_manager_tests.rs
+5
View File
@@ -34,6 +34,11 @@ fn unified_exec_env_overrides_existing_values() {
assert_eq!(env.get("PATH"), Some(&"/usr/bin".to_string()));
}
#[test]
fn exec_server_process_id_matches_unified_exec_process_id() {
assert_eq!(exec_server_process_id(4321), "4321");
}
#[test]
fn pruning_prefers_exited_processes_outside_recently_used() {
let now = Instant::now();
codex-rs/core/src/unified_exec/process_state.rs
+24
View File
@@ -0,0 +1,24 @@
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub(crate) struct ProcessState {
pub(crate) has_exited: bool,
pub(crate) exit_code: Option<i32>,
pub(crate) failure_message: Option<String>,
}
impl ProcessState {
pub(crate) fn exited(&self, exit_code: Option<i32>) -> Self {
Self {
has_exited: true,
exit_code,
failure_message: self.failure_message.clone(),
}
}
pub(crate) fn failed(&self, message: String) -> Self {
Self {
has_exited: true,
exit_code: self.exit_code,
failure_message: Some(message),
}
}
}
codex-rs/core/src/unified_exec/process_tests.rs
+142
View File
@@ -0,0 +1,142 @@
use super::process::UnifiedExecProcess;
use crate::unified_exec::UnifiedExecError;
use async_trait::async_trait;
use codex_exec_server::ExecProcess;
use codex_exec_server::ExecServerError;
use codex_exec_server::ProcessId;
use codex_exec_server::ReadResponse;
use codex_exec_server::StartedExecProcess;
use codex_exec_server::WriteResponse;
use codex_exec_server::WriteStatus;
use codex_sandboxing::SandboxType;
use pretty_assertions::assert_eq;
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::watch;
use tokio::time::Duration;
struct MockExecProcess {
process_id: ProcessId,
write_response: WriteResponse,
read_responses: Mutex<VecDeque<ReadResponse>>,
wake_tx: watch::Sender<u64>,
}
#[async_trait]
impl ExecProcess for MockExecProcess {
fn process_id(&self) -> &ProcessId {
&self.process_id
}
fn subscribe_wake(&self) -> watch::Receiver<u64> {
self.wake_tx.subscribe()
}
async fn read(
&self,
_after_seq: Option<u64>,
_max_bytes: Option<usize>,
_wait_ms: Option<u64>,
) -> Result<ReadResponse, ExecServerError> {
Ok(self
.read_responses
.lock()
.await
.pop_front()
.unwrap_or(ReadResponse {
chunks: Vec::new(),
next_seq: 1,
exited: false,
exit_code: None,
closed: false,
failure: None,
}))
}
async fn write(&self, _chunk: Vec<u8>) -> Result<WriteResponse, ExecServerError> {
Ok(self.write_response.clone())
}
async fn terminate(&self) -> Result<(), ExecServerError> {
Ok(())
}
}
async fn remote_process(write_status: WriteStatus) -> UnifiedExecProcess {
let (wake_tx, _wake_rx) = watch::channel(0);
let started = StartedExecProcess {
process: Arc::new(MockExecProcess {
process_id: "test-process".to_string().into(),
write_response: WriteResponse {
status: write_status,
},
read_responses: Mutex::new(VecDeque::new()),
wake_tx,
}),
};
UnifiedExecProcess::from_remote_started(started, SandboxType::None)
.await
.expect("remote process should start")
}
#[tokio::test]
async fn remote_write_unknown_process_marks_process_exited() {
let process = remote_process(WriteStatus::UnknownProcess).await;
let err = process
.write(b"hello")
.await
.expect_err("expected write failure");
assert!(matches!(err, UnifiedExecError::WriteToStdin));
assert!(process.has_exited());
}
#[tokio::test]
async fn remote_write_closed_stdin_marks_process_exited() {
let process = remote_process(WriteStatus::StdinClosed).await;
let err = process
.write(b"hello")
.await
.expect_err("expected write failure");
assert!(matches!(err, UnifiedExecError::WriteToStdin));
assert!(process.has_exited());
}
#[tokio::test]
async fn remote_process_waits_for_early_exit_event() {
let (wake_tx, _wake_rx) = watch::channel(0);
let started = StartedExecProcess {
process: Arc::new(MockExecProcess {
process_id: "test-process".to_string().into(),
write_response: WriteResponse {
status: WriteStatus::Accepted,
},
read_responses: Mutex::new(VecDeque::from([ReadResponse {
chunks: Vec::new(),
next_seq: 2,
exited: true,
exit_code: Some(17),
closed: true,
failure: None,
}])),
wake_tx: wake_tx.clone(),
}),
};
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await;
let _ = wake_tx.send(1);
});
let process = UnifiedExecProcess::from_remote_started(started, SandboxType::None)
.await
.expect("remote process should observe early exit");
assert!(process.has_exited());
assert_eq!(process.exit_code(), Some(17));
}