github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 16:17:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve remote exec output until it is collected (#53622)

Browse Source 
## Why

Remote process exit can arrive before the final output, and inherited pipes can keep the output stream open after exit. Removing the session at exit can lose output that should remain available to `write_stdin` polls.

## What changed

- Retain remote sessions with open or unread output after exit, including early exits and cancelled initial collection.
- Collect pending output when `write_stdin` races with exit, skipping stdin writes once remote exit is known.
- After initial collection finishes and the process exits, allow a five-second grace period before a final background read. Respect active polling deadlines and keep later polls available while that read recovers retained output.
- Let explicit termination stop output collection even when the process has already exited, and install the exit watcher before cancellable cleanup.

## Testing

Extend remote process event tests to cover output arriving after the initial yield, polling during a delayed final read, active polls exceeding the background grace period, resumable cancellation, and termination before the initial yield.

GitOrigin-RevId: b1671fad8b61a2cdecf7cf9fdf92e57c028c7af8
This commit is contained in:
jif
2026-10-13 09:34:17 +00:00
committed by copyberry
parent 9167a0b2d4
commit b043762f9e
4 changed files with 363 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/unified_exec/process.rs
+55 -7
View File
@@ -37,6 +37,8 @@ use super::process_state::ProcessState;
use crate::shell_snapshot::ShellSnapshotFile;
const EARLY_EXIT_GRACE_PERIOD: Duration = Duration::from_millis(150);
const REMOTE_OUTPUT_DRAIN_GRACE: Duration = Duration::from_secs(5);
/// Output pending model polling and the retained completion transcript.
/// Append both under the same lock so cancellation cannot split an update.
#[derive(Default)]
@@ -60,6 +62,9 @@ pub(crate) struct OutputHandles<const MAX_BYTES: usize = UNIFIED_EXEC_OUTPUT_MAX
pub(crate) output_closed: Arc<AtomicBool>,
pub(crate) output_closed_notify: Arc<Notify>,
pub(crate) cancellation_token: CancellationToken,
pub(super) drain_policy: OutputDrainPolicy,
/// The initial exec call returned or was cancelled, so background draining may begin.
pub(crate) collection_finished: CancellationToken,
}
struct OutputTaskGuard {
@@ -102,6 +107,8 @@ pub(crate) struct UnifiedExecProcess {
output_task: Option<JoinHandle<()>>,
sandbox_type: Option<SandboxType>,
timed_out: AtomicBool,
/// Explicit termination overrides retention of remote output after exit.
pub(super) termination_requested: AtomicBool,
// The shell may still need to replay this file after process startup returns.
pub(crate) _shell_snapshot: Option<Arc<ShellSnapshotFile>>,
}
@@ -117,13 +124,19 @@ impl std::fmt::Debug for UnifiedExecProcess {
}
impl UnifiedExecProcess {
fn new(process_handle: ProcessHandle, sandbox_type: Option<SandboxType>) -> Self {
fn new(
process_handle: ProcessHandle,
sandbox_type: Option<SandboxType>,
drain_policy: OutputDrainPolicy,
) -> Self {
let output = OutputHandles {
output_buffer: Arc::new(Mutex::new(OutputBuffers::default())),
output_notify: Arc::new(Notify::new()),
output_closed: Arc::new(AtomicBool::new(false)),
output_closed_notify: Arc::new(Notify::new()),
cancellation_token: CancellationToken::new(),
drain_policy,
collection_finished: CancellationToken::new(),
};
let output_drained = Arc::new(Notify::new());
let (output_tx, _) = broadcast::channel(64);
@@ -140,6 +153,7 @@ impl UnifiedExecProcess {
output_task: None,
sandbox_type,
timed_out: AtomicBool::new(false),
termination_requested: AtomicBool::new(false),
_shell_snapshot: None,
}
}
@@ -219,6 +233,7 @@ impl UnifiedExecProcess {
}
fn finish_termination(&self) {
self.termination_requested.store(true, Ordering::Release);
self.output.cancellation_token.cancel();
if let Some(output_task) = &self.output_task {
output_task.abort();
@@ -351,6 +366,7 @@ impl UnifiedExecProcess {
let mut managed = Self::new(
ProcessHandle::Local(Box::new(process_handle)),
Some(sandbox_type),
OutputDrainPolicy::BoundedAfterExit,
);
managed.output_task = Some(Self::spawn_local_output_task(
output_rx,
@@ -400,7 +416,7 @@ impl UnifiedExecProcess {
// Older peers do not report this field. In that case, skip local
// classification rather than attributing a violation to a guessed backend.
let sandbox_type = started.sandbox_type;
let mut managed = Self::new(process_handle, sandbox_type);
let mut managed = Self::new(process_handle, sandbox_type, output_drain_policy);
let output_handles = managed.output_handles().clone();
managed.output_task = Some(Self::spawn_exec_server_output_task(
started,
@@ -408,6 +424,7 @@ impl UnifiedExecProcess {
output_handles,
managed.output_tx.clone(),
managed.state_tx.clone(),
managed.interaction_lock(),
));
let mut state_rx = managed.state_rx.clone();
@@ -437,6 +454,7 @@ impl UnifiedExecProcess {
output_handles: OutputHandles,
output_tx: broadcast::Sender<Vec<u8>>,
state_tx: watch::Sender<ProcessState>,
interaction_lock: Arc<Mutex<()>>,
) -> JoinHandle<()> {
let OutputHandles {
output_buffer,
@@ -444,20 +462,42 @@ impl UnifiedExecProcess {
output_closed,
output_closed_notify,
cancellation_token,
collection_finished,
..
} = output_handles;
let process = started.process;
let mut events = process.subscribe_events();
let mut state_rx = state_tx.subscribe();
tokio::spawn(async move {
let _output_task_guard = OutputTaskGuard {
output_closed: Arc::clone(&output_closed),
output_closed_notify: Arc::clone(&output_closed_notify),
};
let drain = async {
collection_finished.cancelled().await;
let _ = state_rx.wait_for(|state| state.has_exited).await;
tokio::time::sleep(REMOTE_OUTPUT_DRAIN_GRACE).await;
// An active poll keeps its own deadline, including transport recovery.
interaction_lock.lock_owned().await
};
tokio::pin!(drain);
let mut final_drain = false;
let mut last_seq: u64 = 0;
loop {
let event = match events.recv().await {
Ok(event) => Some(event),
Err(broadcast::error::RecvError::Lagged(_)) => None,
Err(broadcast::error::RecvError::Closed) => {
let received = tokio::select! {
received = events.recv() => Some(received),
guard = &mut drain, if output_drain_policy == OutputDrainPolicy::WaitForOutputClosure
&& !final_drain => {
// Later polls must remain able to collect output while this read recovers.
drop(guard);
final_drain = true;
None
}
};
let event = match received {
Some(Ok(event)) => Some(event),
Some(Err(broadcast::error::RecvError::Lagged(_))) | None => None,
Some(Err(broadcast::error::RecvError::Closed)) => {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(
state.failed("exec-server process event stream closed".to_string()),
@@ -486,6 +526,8 @@ impl UnifiedExecProcess {
|| event_seq.is_some_and(|seq| seq > last_seq.saturating_add(1))
|| missing_sandbox_denial
{
// The final drain read waits for transport recovery and collects retained
// output before closing an idle session with inherited pipes still open.
let response = match process
.read(
Some(last_seq),
@@ -542,7 +584,7 @@ impl UnifiedExecProcess {
if exited && output_drain_policy == OutputDrainPolicy::BoundedAfterExit {
cancellation_token.cancel();
}
if closed {
if closed || final_drain {
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
@@ -602,6 +644,12 @@ impl UnifiedExecProcess {
}
}
}
if final_drain {
tokio::spawn(async move {
let _ = process.terminate().await;
});
}
cancellation_token.cancel();
})
}
codex-rs/core/src/unified_exec/process_manager.rs
+89 -54
View File
@@ -271,6 +271,7 @@ struct PreparedProcessHandles {
struct InitialExecCommandGuard {
active: Option<Arc<AtomicBool>>,
metrics_sidecar: Option<PluginMetricsSidecar>,
collection_finished: CancellationToken,
}
impl InitialExecCommandGuard {
@@ -293,6 +294,7 @@ impl Drop for InitialExecCommandGuard {
if let Some(active) = self.active.as_ref() {
active.store(false, Ordering::Release);
}
self.collection_finished.cancel();
}
}
@@ -390,7 +392,7 @@ fn fail_process_with_message(process: &UnifiedExecProcess, message: String) -> U
#[allow(clippy::too_many_arguments)]
async fn emit_failed_initial_exec_end_if_unstored(
process_started_alive: bool,
process_stored: bool,
sandbox_type: Option<codex_protocol::sandbox::SandboxType>,
context: &UnifiedExecContext,
request: &ExecCommandRequest,
@@ -401,7 +403,7 @@ async fn emit_failed_initial_exec_end_if_unstored(
message: String,
wall_time: Duration,
) {
if process_started_alive {
if process_stored {
return;
}
@@ -592,13 +594,30 @@ impl UnifiedExecProcessManager {
);
emitter.emit(event_ctx, ToolEventStage::Begin).await;
// Serialize remote collection with later polls and final transcript publication.
let _interaction_guard =
if process.output_handles().drain_policy == OutputDrainPolicy::WaitForOutputClosure {
Some(process.interaction_lock().lock_owned().await)
} else {
None
};
start_streaming_output(&process, context);
let start = Instant::now();
// Persist live sessions before the initial yield wait so interrupting the
// turn cannot drop the last Arc and terminate the background process.
let process_started_alive = !process.has_exited() && process.exit_code().is_none();
let mut initial_exec_command_guard = if process_started_alive {
// Persist live sessions and remote output that remains pollable after an early exit.
let process_stored = (!process.has_exited() && process.exit_code().is_none())
|| (process.output_handles().drain_policy == OutputDrainPolicy::WaitForOutputClosure
&& !process
.output_handles()
.output_closed
.load(Ordering::Acquire));
let mut initial_exec_command_guard = InitialExecCommandGuard {
active: None,
metrics_sidecar,
collection_finished: process.output_handles().collection_finished.clone(),
};
if process_stored {
let initial_exec_command_active = Arc::new(AtomicBool::new(true));
initial_exec_command_guard.active = Some(Arc::clone(&initial_exec_command_active));
self.store_process(
Arc::clone(&process),
context,
@@ -613,21 +632,12 @@ impl UnifiedExecProcessManager {
request.tty,
deferred_network_approval.clone(),
network_denial_monitor,
metrics_sidecar,
initial_exec_command_guard.metrics_sidecar.take(),
Arc::clone(&output_buffer),
Arc::clone(&initial_exec_command_active),
)
.await;
InitialExecCommandGuard {
active: Some(initial_exec_command_active),
metrics_sidecar: None,
}
} else {
InitialExecCommandGuard {
active: None,
metrics_sidecar,
}
};
}
let yield_time_ms = clamp_yield_time(request.yield_time_ms);
// For the initial exec_command call, we both stream output to events
@@ -675,7 +685,7 @@ impl UnifiedExecProcessManager {
)
.await;
emit_failed_initial_exec_end_if_unstored(
process_started_alive,
process_stored,
process.sandbox_type(),
context,
&request,
@@ -697,7 +707,7 @@ impl UnifiedExecProcessManager {
)
.await;
emit_failed_initial_exec_end_if_unstored(
process_started_alive,
process_stored,
process.sandbox_type(),
context,
&request,
@@ -716,15 +726,19 @@ impl UnifiedExecProcessManager {
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 {
let (response_process_id, exit_code) = if process_stored {
let retention = if completion.is_some() {
OutputRetention::UntilExit
} else {
OutputRetention::UntilCollected
};
match self.refresh_process_state(process_id, retention).await {
ProcessStatus::Alive {
exit_code,
process_id,
..
} => (Some(process_id), exit_code),
ProcessStatus::Exited { exit_code, entry } => {
// Collection is finished even if the remote output stream remains open.
entry.process.cancellation_token().cancel();
if let Err(message) =
finish_deferred_network_approval_after_process_exit_for_session(
@@ -778,7 +792,7 @@ impl UnifiedExecProcessManager {
.await;
if let Err(message) = finish_result {
emit_failed_initial_exec_end_if_unstored(
process_started_alive,
process_stored,
process.sandbox_type(),
context,
&request,
@@ -981,9 +995,10 @@ impl UnifiedExecProcessManager {
} = self
.prepare_process_handles(process_id, &locked_process)
.await?;
let mut status_after_write = None;
if !request.input.is_empty() {
let remote_exited =
process.has_exited() && output.drain_policy == OutputDrainPolicy::WaitForOutputClosure;
// A write racing with remote exit still collects the final output.
if !request.input.is_empty() && !remote_exited {
if !tty {
if request.input == INTERRUPT {
process.interrupt().await?;
@@ -998,14 +1013,11 @@ impl UnifiedExecProcessManager {
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 {
if !process.has_exited() {
if matches!(err, UnifiedExecError::ProcessFailed { .. }) {
process.terminate();
self.release_process_id(process_id).await;
}
return Err(err);
}
}
@@ -1059,15 +1071,10 @@ impl UnifiedExecProcessManager {
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 or suppress TerminalInteraction
// with an appropriate process_id and exit_code.
let status = if let Some(status) = status_after_write {
status
} else {
self.refresh_process_state(process_id).await
};
// Only finalize remote sessions after both process exit and output collection.
let status = self
.refresh_process_state(process_id, OutputRetention::UntilCollected)
.await;
let (process_id, exit_code, event_call_id) = match status {
ProcessStatus::Alive {
exit_code,
@@ -1075,7 +1082,6 @@ impl UnifiedExecProcessManager {
process_id,
} => (Some(process_id), exit_code, call_id),
ProcessStatus::Exited { exit_code, entry } => {
// Release the exit watcher after the final poll, not on remote process exit.
entry.process.cancellation_token().cancel();
let call_id = entry.call_id.clone();
if let Err(message) =
@@ -1130,7 +1136,11 @@ impl UnifiedExecProcessManager {
Ok(response)
}
async fn refresh_process_state(&self, process_id: i32) -> ProcessStatus {
async fn refresh_process_state(
&self,
process_id: i32,
retention: OutputRetention,
) -> ProcessStatus {
let mut store = self.process_store.lock().await;
let Some(entry) = store.processes.get_mut(&process_id) else {
return ProcessStatus::Unknown;
@@ -1140,6 +1150,24 @@ impl UnifiedExecProcessManager {
let process_id = entry.process_id;
if entry.process.has_exited() {
let output = entry.process.output_handles();
// Never wait for output under the shared process-store lock. Retain busy output
// for another poll; closed, empty output cannot grow again.
if retention == OutputRetention::UntilCollected
&& output.drain_policy == OutputDrainPolicy::WaitForOutputClosure
&& !entry.process.termination_requested.load(Ordering::Acquire)
&& (!output.output_closed.load(Ordering::Acquire)
|| output
.output_buffer
.try_lock()
.map_or(true, |buffer| buffer.pending.total_bytes() > 0))
{
return ProcessStatus::Alive {
exit_code: None,
call_id: entry.call_id.clone(),
process_id,
};
}
let Some(entry) = store.remove(process_id) else {
return ProcessStatus::Unknown;
};
@@ -1233,13 +1261,6 @@ impl UnifiedExecProcessManager {
store.processes.insert(process_id, entry);
pruned_entry
};
// prune_processes_if_needed runs while holding process_store; do async
// network-approval cleanup only after dropping that lock.
if let Some(pruned_entry) = pruned_entry {
unregister_network_approval_for_entry(&pruned_entry).await;
pruned_entry.process.terminate();
}
spawn_exit_watcher(
Arc::clone(&process),
context,
@@ -1252,6 +1273,11 @@ impl UnifiedExecProcessManager {
network_denial_monitor,
plugin_metrics_sidecar,
);
// Install the watcher before cancellable cleanup, and outside the process-store lock.
if let Some(pruned_entry) = pruned_entry {
unregister_network_approval_for_entry(&pruned_entry).await;
pruned_entry.process.terminate();
}
}
#[allow(clippy::too_many_arguments)]
@@ -1589,6 +1615,7 @@ impl UnifiedExecProcessManager {
output_closed,
output_closed_notify,
cancellation_token,
..
} = output;
let mut collected = HeadTailBuffer::default();
let mut exit_signal_received = cancellation_token.is_cancelled();
@@ -1848,10 +1875,12 @@ impl UnifiedExecProcessManager {
(Arc::clone(&entry.process), entry.process.has_exited())
};
if !already_exited && process.terminate_confirmed().await.is_err() {
if already_exited {
// Stop the output producer even when only its inherited pipes remain open.
process.terminate();
} else if process.terminate_confirmed().await.is_err() {
return false;
}
process.cancellation_token().cancel();
let entry = {
let mut store = self.process_store.lock().await;
@@ -1875,6 +1904,12 @@ impl UnifiedExecProcessManager {
}
}
#[derive(PartialEq, Eq)]
enum OutputRetention {
UntilExit,
UntilCollected,
}
enum ProcessStatus {
Alive {
exit_code: Option<i32>,
codex-rs/core/src/unified_exec/process_manager_tests.rs
+5 -1
View File
@@ -355,6 +355,8 @@ async fn output_collection_stays_bounded_across_repeated_drains() {
output_closed: Arc::clone(&output_closed),
output_closed_notify: Arc::clone(&output_closed_notify),
cancellation_token: cancellation_token.clone(),
collection_finished: CancellationToken::new(),
drain_policy: OutputDrainPolicy::BoundedAfterExit,
};
let collect = UnifiedExecProcessManager::collect_output_until_deadline(
@@ -413,6 +415,8 @@ async fn output_collection_preserves_omissions_from_drained_buffer() {
output_closed,
output_closed_notify,
cancellation_token,
collection_finished: CancellationToken::new(),
drain_policy: OutputDrainPolicy::BoundedAfterExit,
};
let collected = UnifiedExecProcessManager::collect_output_until_deadline(
@@ -495,7 +499,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() {
output_buffer.lock().await.push_chunk(b"PARTIAL_TRANSCRIPT");
emit_failed_initial_exec_end_if_unstored(
/*process_started_alive*/ false,
/*process_stored*/ false,
Some(codex_protocol::sandbox::SandboxType::WindowsMxc),
&context,
&request,
codex-rs/core/tests/suite/unified_exec_process_events.rs
+214 -50
View File
@@ -64,6 +64,7 @@ use tokio_tungstenite::tungstenite::Message;
const CALL_ID: &str = "pushed-remote-process-events";
const POLL_CALL_ID: &str = "poll-remote-process-events";
const DRAIN_POLL_CALL_ID: &str = "poll-during-remote-drain";
const COMPLETE_OUTPUT: &str = "pushed remote output\n";
const RECOVERED_OUTPUT: &str = "recovered missing output\n";
const RETAINED_OUTPUT: &str = "retained output\n";
@@ -79,7 +80,9 @@ enum PushedExecScenario {
OutputOpenAtDeadline,
OutputOpenAtFinalPoll,
OutputOpenWhenCancelled,
OutputOpenWhenResumableCancelled,
OutputOpenWhenTerminated,
OutputOpenWhenTerminatedBeforeYield,
CodexPath,
CodexPathDefaultFeature,
CodexPathDisabledFeature,
@@ -240,6 +243,7 @@ async fn serve_exec_with_pushed_events(
scenario: PushedExecScenario,
process_started: Arc<AtomicBool>,
output_ready: Arc<tokio::sync::Notify>,
drain_read_ready: Arc<tokio::sync::Notify>,
) -> PushedExecServerResult {
let unrestricted_patch = matches!(
scenario,
@@ -434,9 +438,11 @@ async fn serve_exec_with_pushed_events(
PushedExecScenario::OutputOpenAtDeadline
| PushedExecScenario::OutputOpenAtFinalPoll
| PushedExecScenario::OutputOpenWhenCancelled
| PushedExecScenario::OutputOpenWhenResumableCancelled
| PushedExecScenario::OutputOpenWhenTerminated
| PushedExecScenario::OutputOpenWhenTerminatedBeforeYield
);
if output_stays_open {
if output_stays_open || matches!(scenario, PushedExecScenario::OutputAfterExit) {
output_ready.notified().await;
}
let mut process_read_requests = 0;
@@ -446,7 +452,9 @@ async fn serve_exec_with_pushed_events(
| PushedExecScenario::OutputOpenAtDeadline
| PushedExecScenario::OutputOpenAtFinalPoll
| PushedExecScenario::OutputOpenWhenCancelled
| PushedExecScenario::OutputOpenWhenTerminated => {
| PushedExecScenario::OutputOpenWhenResumableCancelled
| PushedExecScenario::OutputOpenWhenTerminated
| PushedExecScenario::OutputOpenWhenTerminatedBeforeYield => {
send_exec_server_json(
&mut websocket,
json!({
@@ -493,22 +501,36 @@ async fn serve_exec_with_pushed_events(
)
.await;
}
if matches!(scenario, PushedExecScenario::OutputAfterExit) {
// Keep the tail in transit until the initial collection yields.
output_ready.notified().await;
}
// Model foreground output delayed in transit beyond the local post-exit grace.
tokio::time::sleep(Duration::from_millis(250)).await;
send_exec_server_json(
&mut websocket,
json!({
"method": "process/output",
"params": {
"processId": &process_id,
"seq": 3,
"stream": "stdout",
"chunk": BASE64_STANDARD.encode("remote output\n"),
}
}),
// An active final poll must also outlive the five-second background drain.
tokio::time::sleep(
if matches!(scenario, PushedExecScenario::OutputOpenAtFinalPoll) {
Duration::from_secs(6)
} else {
Duration::from_millis(250)
},
)
.await;
if !output_stays_open {
if !matches!(scenario, PushedExecScenario::OutputOpenAtDeadline) {
send_exec_server_json(
&mut websocket,
json!({
"method": "process/output",
"params": {
"processId": &process_id,
"seq": 3,
"stream": "stdout",
"chunk": BASE64_STANDARD.encode("remote output\n"),
}
}),
)
.await;
}
if !output_stays_open || matches!(scenario, PushedExecScenario::OutputOpenAtFinalPoll) {
send_exec_server_json(
&mut websocket,
json!({
@@ -604,23 +626,72 @@ async fn serve_exec_with_pushed_events(
PushedExecScenario::ReplayGap => {}
}
let mut pending_drain_response = None;
loop {
let wait = if output_stays_open {
STARTUP_TIMEOUT
} else {
Duration::from_secs(/*secs*/ 5)
};
let request = read_exec_server_json(&mut websocket, wait).await;
let request = tokio::select! {
request = read_exec_server_json(&mut websocket, wait) => request,
_ = output_ready.notified(), if pending_drain_response.is_some() => {
send_exec_server_json(
&mut websocket,
pending_drain_response.take().expect("pending final read"),
).await;
continue;
}
};
match request["method"].as_str() {
Some("fs/getMetadata") => {
// Starting the polling turn still needs filesystem metadata while the read waits.
send_exec_server_json(
&mut websocket,
json!({
"id": request["id"],
"error": { "code": -32004, "message": "not found" }
}),
)
.await;
}
Some("process/read") => {
process_read_requests += 1;
let result = match scenario {
PushedExecScenario::OutputOpenAtDeadline => {
assert_eq!(request["params"]["afterSeq"], 2);
// Hold the final read until a later poll yields its buffered output.
// The tail remains recoverable even after that poll's deadline.
drain_read_ready.notify_one();
json!({
"chunks": [{
"seq": 3,
"stream": "stdout",
"chunk": BASE64_STANDARD.encode("remote output\n"),
}],
"nextSeq": 4,
"exited": true,
"exitCode": 0,
"closed": false,
"failure": null,
"sandboxDenied": false,
})
}
PushedExecScenario::OutputOpenWhenResumableCancelled => json!({
"chunks": [],
"nextSeq": 4,
"exited": true,
"exitCode": 0,
"closed": false,
"failure": null,
"sandboxDenied": false,
}),
PushedExecScenario::OutputAfterExit
| PushedExecScenario::OutputAfterRecoveredExit
| PushedExecScenario::OutputOpenAtDeadline
| PushedExecScenario::OutputOpenAtFinalPoll
| PushedExecScenario::OutputOpenWhenCancelled
| PushedExecScenario::OutputOpenWhenTerminated => {
| PushedExecScenario::OutputOpenWhenTerminated
| PushedExecScenario::OutputOpenWhenTerminatedBeforeYield => {
panic!("unexpected read after the delayed-output sequence")
}
PushedExecScenario::Complete
@@ -701,14 +772,12 @@ async fn serve_exec_with_pushed_events(
})
}
};
send_exec_server_json(
&mut websocket,
json!({
"id": request["id"],
"result": result,
}),
)
.await;
let response = json!({ "id": request["id"], "result": result });
if matches!(scenario, PushedExecScenario::OutputOpenAtDeadline) {
pending_drain_response = Some(response);
} else {
send_exec_server_json(&mut websocket, response).await;
}
if matches!(scenario, PushedExecScenario::ReplayGap) && process_read_requests == 1 {
send_exec_server_json(
&mut websocket,
@@ -748,7 +817,9 @@ async fn serve_exec_with_pushed_events(
#[test_case(PushedExecScenario::OutputOpenAtDeadline, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_at_deadline")]
#[test_case(PushedExecScenario::OutputOpenAtFinalPoll, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_at_final_poll")]
#[test_case(PushedExecScenario::OutputOpenWhenCancelled, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_when_cancelled")]
#[test_case(PushedExecScenario::OutputOpenWhenResumableCancelled, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_when_resumable_cancelled")]
#[test_case(PushedExecScenario::OutputOpenWhenTerminated, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_when_terminated")]
#[test_case(PushedExecScenario::OutputOpenWhenTerminatedBeforeYield, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_when_terminated_before_yield")]
#[test_case(PushedExecScenario::CodexPath, ManagedNetworkScenario::None, false ; "codex_path_is_exported_before_nested_command")]
#[test_case(PushedExecScenario::CodexPathDefaultFeature, ManagedNetworkScenario::None, false ; "codex_path_is_not_exported_by_default")]
#[test_case(PushedExecScenario::CodexPathDisabledFeature, ManagedNetworkScenario::None, false ; "disabled_codex_path_preserves_original_login_command")]
@@ -784,8 +855,19 @@ async fn exec_command_consumes_pushed_remote_process_events(
) -> Result<()> {
let finish_at_deadline = matches!(scenario, PushedExecScenario::OutputOpenAtDeadline);
let finish_at_final_poll = matches!(scenario, PushedExecScenario::OutputOpenAtFinalPoll);
let cancel_after_exit = matches!(scenario, PushedExecScenario::OutputOpenWhenCancelled);
let terminate_after_exit = matches!(scenario, PushedExecScenario::OutputOpenWhenTerminated);
let resumable_cancel = matches!(
scenario,
PushedExecScenario::OutputOpenWhenResumableCancelled
);
let cancel_after_exit =
resumable_cancel || matches!(scenario, PushedExecScenario::OutputOpenWhenCancelled);
let output_after_yield = matches!(scenario, PushedExecScenario::OutputAfterExit);
let terminate_initial_collection = matches!(
scenario,
PushedExecScenario::OutputOpenWhenTerminatedBeforeYield
);
let terminate_after_exit = terminate_initial_collection
|| matches!(scenario, PushedExecScenario::OutputOpenWhenTerminated);
let explicit_path = "/user/configured/bin:/user/other/bin";
let managed_network_configured = !matches!(managed_network, ManagedNetworkScenario::None);
let managed_network_enabled = matches!(managed_network, ManagedNetworkScenario::Enabled { .. });
@@ -856,7 +938,8 @@ async fn exec_command_consumes_pushed_remote_process_events(
| PushedExecScenario::CodexPathDisabledFeature => "sh -c 'command -v rg'".to_string(),
_ => "pwd".to_string(),
},
"yield_time_ms": 1_000,
"tty": output_after_yield,
"yield_time_ms": if resumable_cancel || terminate_initial_collection { 30_000 } else { 1_000 },
})
.to_string(),
),
@@ -876,7 +959,7 @@ async fn exec_command_consumes_pushed_remote_process_events(
if cancel_after_exit {
responses.truncate(/*len*/ 1);
}
if finish_at_final_poll {
if output_after_yield || finish_at_final_poll {
responses.insert(
/*index*/ 1,
sse(vec![
@@ -884,21 +967,48 @@ async fn exec_command_consumes_pushed_remote_process_events(
ev_function_call(
POLL_CALL_ID,
"write_stdin",
&json!({ "session_id": 1000, "chars": "", "yield_time_ms": 1_000 }).to_string(),
&json!({
"session_id": 1000,
// Exit is already observed; skip stdin but still collect the tail.
"chars": if output_after_yield { "ignored\n" } else { "" },
// The final poll receives output after the five-second background drain.
"yield_time_ms": if finish_at_final_poll { 10_000 } else { 1_000 },
})
.to_string(),
),
ev_completed("resp-poll"),
]),
);
}
if finish_at_deadline {
for call_id in [DRAIN_POLL_CALL_ID, POLL_CALL_ID] {
responses.push(sse(vec![
ev_response_created(call_id),
ev_function_call(
call_id,
"write_stdin",
&json!({ "session_id": 1000, "chars": "", "yield_time_ms": 1_000 }).to_string(),
),
ev_completed(call_id),
]));
}
responses.push(sse(vec![
ev_response_created("resp-drained"),
ev_assistant_message("msg-drained", "done"),
ev_completed("resp-drained"),
]));
}
let response_mock = mount_sse_sequence(&server, responses).await;
let exec_server_url = format!("ws://{}", listener.local_addr()?);
let process_started = Arc::new(AtomicBool::new(false));
let output_ready = Arc::new(tokio::sync::Notify::new());
let drain_read_ready = Arc::new(tokio::sync::Notify::new());
let exec_server = tokio::spawn(serve_exec_with_pushed_events(
listener,
scenario,
Arc::clone(&process_started),
Arc::clone(&output_ready),
Arc::clone(&drain_read_ready),
));
let mut builder = test_codex().with_exec_server_url(exec_server_url);
if managed_network_configured {
@@ -960,7 +1070,7 @@ timeout = 900
.expect("write managed-network test config");
});
}
if cancel_after_exit {
if matches!(scenario, PushedExecScenario::OutputOpenWhenCancelled) {
builder = builder.with_cloud_config_bundle(
CloudConfigBundleFixture::loader_with_enterprise_requirement(
"[features]\nunified_exec = false\nshell_tool = true\n",
@@ -1145,9 +1255,15 @@ timeout = 900
let mut command_end = None;
let mut stopped = false;
let mut turn_finished = false;
let mut drain_poll_started = false;
let wait_for_end =
finish_at_deadline || finish_at_final_poll || cancel_after_exit || terminate_after_exit;
let event_timeout = if finish_at_deadline || finish_at_final_poll || terminate_after_exit {
let event_timeout = if finish_at_deadline
|| finish_at_final_poll
|| terminate_after_exit
|| resumable_cancel
|| output_after_yield
{
STARTUP_TIMEOUT
} else {
Duration::from_secs(/*secs*/ 5)
@@ -1162,9 +1278,9 @@ timeout = 900
EventMsg::ExecCommandBegin(event) if event.call_id == CALL_ID => {
saw_exec_command_begin = true;
exec_command_begin_command = Some(event.command);
if finish_at_deadline || cancel_after_exit {
// Storage follows the streaming subscription. Wait for it before
// exiting so both deadline cleanup and cancellation see a live session.
if cancel_after_exit || output_after_yield || terminate_initial_collection {
// Storage follows the streaming subscription, so the fixture's
// exit and tail cannot race ahead of the cancellation trigger.
timeout(STARTUP_TIMEOUT, async {
while test.codex.list_background_terminals().await.is_empty() {
tokio::task::yield_now().await;
@@ -1227,10 +1343,19 @@ timeout = 900
stopped = true;
}
EventMsg::RawResponseItem(raw)
if (finish_at_final_poll || terminate_after_exit)
if finish_at_deadline
&& matches!(&raw.item, ResponseItem::FunctionCallOutput { call_id, .. } if call_id.as_deref() == Some(DRAIN_POLL_CALL_ID)) =>
{
output_ready.notify_one();
}
EventMsg::RawResponseItem(raw)
if ((terminate_after_exit && !terminate_initial_collection)
|| finish_at_deadline
|| finish_at_final_poll
|| output_after_yield)
&& matches!(&raw.item, ResponseItem::FunctionCallOutput { call_id, .. } if call_id.as_deref() == Some(CALL_ID)) =>
{
// Polling and explicit termination must observe exit after the initial call yielded.
// Release exit/output only after the initial exec_command yielded.
output_ready.notify_one();
}
EventMsg::ExecCommandEnd(event) if event.call_id == CALL_ID => {
@@ -1239,6 +1364,16 @@ timeout = 900
EventMsg::TurnAborted(_) if cancel_after_exit => {
turn_finished = true;
}
EventMsg::TurnComplete(_) if finish_at_deadline && !drain_poll_started => {
timeout(STARTUP_TIMEOUT, drain_read_ready.notified()).await?;
test.codex
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "Poll the command.".into(),
text_elements: Vec::new(),
}]))
.await?;
drain_poll_started = true;
}
EventMsg::TurnComplete(_) if !cancel_after_exit => turn_finished = true,
_ => {}
}
@@ -1421,10 +1556,24 @@ timeout = 900
"toolCallId": CALL_ID,
}),
);
if resumable_cancel {
assert_eq!(exec_server_result.process_read_requests, 1);
// The producer is closed without another poll; its unread output stays collectable.
assert!(
test.codex
.terminate_background_terminal(/*process_id*/ 1000)
.await
);
}
if cancel_after_exit || terminate_after_exit {
assert!(stopped);
assert_eq!(exec_server_result.process_read_requests, 0);
return Ok(());
assert_eq!(
exec_server_result.process_read_requests,
usize::from(resumable_cancel)
);
if !terminate_initial_collection {
return Ok(());
}
}
if foreign_cwd && !managed_network_configured && !deny_read {
let permissions = response_mock.requests()[0]
@@ -1575,15 +1724,15 @@ timeout = 900
let request = response_mock
.last_request()
.context("model should receive the exec_command output")?;
let output_call_id = if finish_at_final_poll {
let (initial_output, _) = request
let output_call_id = if output_after_yield || finish_at_final_poll || finish_at_deadline {
let (initial, _) = request
.function_call_output_content_and_success(CALL_ID)
.context("initial exec_command should yield a session")?;
assert!(
initial_output
.context("initial exec_command output should contain text")?
.contains("Process running with session ID 1000")
);
.context("initial collection")?;
let initial = initial.context("initial output")?;
assert!(initial.contains("Process running with session ID 1000"));
if output_after_yield {
assert!(initial.contains("pushed "));
}
assert!(
!test
.codex
@@ -1594,6 +1743,14 @@ timeout = 900
} else {
CALL_ID
};
if finish_at_deadline {
let (pending, _) = request
.function_call_output_content_and_success(DRAIN_POLL_CALL_ID)
.context("poll during the final read should yield")?;
let pending = pending.context("poll should contain buffered output")?;
assert!(pending.contains("Process running with session ID 1000"));
assert!(pending.contains("pushed "));
}
let (output, success) = request
.function_call_output_content_and_success(output_call_id)
.context("exec_command output should be model visible")?;
@@ -1605,12 +1762,18 @@ timeout = 900
| PushedExecScenario::OutputOpenAtDeadline
| PushedExecScenario::OutputOpenAtFinalPoll
| PushedExecScenario::OutputOpenWhenCancelled
| PushedExecScenario::OutputOpenWhenTerminated => {
| PushedExecScenario::OutputOpenWhenResumableCancelled
| PushedExecScenario::OutputOpenWhenTerminated
| PushedExecScenario::OutputOpenWhenTerminatedBeforeYield => {
assert_ne!(success, Some(false));
assert!(saw_exec_command_begin);
assert!(output.contains("Process exited with code 0"));
assert!(
output.contains(COMPLETE_OUTPUT),
output.contains(if output_after_yield || finish_at_deadline {
"remote output\n"
} else {
COMPLETE_OUTPUT
}),
"missing remote output tail: {output}"
);
assert_eq!(
@@ -1618,6 +1781,7 @@ timeout = 900
usize::from(matches!(
scenario,
PushedExecScenario::OutputAfterRecoveredExit
| PushedExecScenario::OutputOpenAtDeadline
))
);
}