openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-16 16:58:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve early unified exec output in completion events (#47665)

Browse Source 
## Why

Output produced before the streaming subscriber attaches can be missing from command completion events, even when it was returned to the model.

## What changed

Record the bounded completion transcript directly in `UnifiedExecProcess` as local and exec-server output arrives. Keep it independent of output polling and streaming delta delivery, and use it for completion events.

## Testing

Add regression coverage for output collected before subscription, both with and without later output. Strengthen the lifecycle test to assert that `aggregated_output` includes both early and late output.

GitOrigin-RevId: 9825acbd706dc4845837c8b4ab1347d8267ce0af
This commit is contained in:
Steve Coffey
2026-09-23 20:56:32 +00:00
committed by copyberry
parent 8e17909b27
commit d47b9a8c00
5 changed files with 104 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/unified_exec/async_watcher.rs
+3 -18
View File
@@ -50,18 +50,11 @@ struct Emitter {
struct Buffer<const MAX_BYTES: usize = UNIFIED_EXEC_OUTPUT_DELTA_MAX_BYTES> {
pending: Vec<u8>,
transcript: Arc<Mutex<HeadTailBuffer>>,
emitter: Emitter,
}
/// Spawn a background task that continuously reads from the PTY, appends to the
/// shared transcript, and emits ExecCommandOutputDelta events on UTF‑8
/// boundaries.
pub(crate) fn start_streaming_output(
process: &UnifiedExecProcess,
context: &UnifiedExecContext,
transcript: Arc<Mutex<HeadTailBuffer>>,
) {
/// Spawn a background task that emits ExecCommandOutputDelta events on UTF‑8 boundaries.
pub(crate) fn start_streaming_output(process: &UnifiedExecProcess, context: &UnifiedExecContext) {
let mut receiver = process.output_receiver();
let output_drained = process.output_drained_notify();
let exit_token = process.cancellation_token();
@@ -83,7 +76,6 @@ pub(crate) fn start_streaming_output(
let mut output: Buffer = Buffer {
pending: Vec::new(),
transcript,
emitter,
};
@@ -258,13 +250,7 @@ impl<const MAX_BYTES: usize> Buffer<MAX_BYTES> {
"a frame must fit one UTF-8 scalar"
)
};
let Self {
pending,
transcript,
emitter,
} = self;
transcript.lock().await.push_chunk(&bytes);
let Self { pending, emitter } = self;
// Reuse a producer chunk when it fits, retaining only an incomplete
// UTF-8 suffix for the next push.
@@ -299,7 +285,6 @@ impl<const MAX_BYTES: usize> Buffer<MAX_BYTES> {
async fn finish(self) {
let Self {
pending,
transcript: _,
mut emitter,
} = self;
debug_assert!(
codex-rs/core/src/unified_exec/async_watcher_tests.rs
+80 -14
View File
@@ -8,6 +8,7 @@ use super::start_streaming_output;
use super::utf8_boundary;
use crate::session::tests::make_session_and_context_with_rx;
use crate::unified_exec::UnifiedExecContext;
use crate::unified_exec::UnifiedExecProcessManager;
use crate::unified_exec::head_tail_buffer::HeadTailBuffer;
use crate::unified_exec::process::NoopSpawnLifecycle;
use crate::unified_exec::process::UnifiedExecProcess;
@@ -58,8 +59,7 @@ async fn streaming_output_harness() -> anyhow::Result<StreamingOutputHarness> {
tokio_util::sync::CancellationToken::new(),
"streaming-output-test".to_string(),
);
let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
start_streaming_output(&process, &context, Arc::clone(&transcript));
let transcript = process.transcript();
Ok(StreamingOutputHarness {
process,
@@ -71,6 +71,76 @@ async fn streaming_output_harness() -> anyhow::Result<StreamingOutputHarness> {
})
}
#[test_case::test_case(b""; "no_late_output")]
#[test_case::test_case(b"late\n"; "with_late_output")]
#[tokio::test]
async fn completed_output_preserves_bytes_before_subscription(
late_output: &[u8],
) -> anyhow::Result<()> {
let StreamingOutputHarness {
process,
stdout_tx,
exit_tx,
transcript,
context,
rx_event,
} = streaming_output_harness().await?;
let collected = process.output_handles().output_notify.notified();
tokio::pin!(collected);
collected.as_mut().enable();
stdout_tx.send(b"early\n".to_vec())?;
collected.await;
let model_output = UnifiedExecProcessManager::collect_output_until_deadline(
process.output_handles(),
/*pause_state*/ None,
Instant::now(),
)
.await;
assert_eq!(model_output.to_bytes(), b"early\n");
start_streaming_output(&process, &context);
#[allow(deprecated)]
let cwd = context.step_context.turn.cwd.clone().into();
spawn_exit_watcher(
Arc::clone(&process),
&context,
vec!["proof".to_string()],
cwd,
/*process_id*/ 123,
/*plugin_attribution*/ None,
transcript,
Instant::now(),
/*network_denial_monitor*/ None,
/*plugin_metrics_sidecar*/ None,
);
stdout_tx.send(late_output.to_vec())?;
drop(stdout_tx);
exit_tx.send(0).expect("send exit");
let item = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if let EventMsg::ItemCompleted(completed) = rx_event.recv().await?.msg
&& let TurnItem::CommandExecution(item) = completed.item
{
return Ok::<_, async_channel::RecvError>(item);
}
}
})
.await??;
let expected_output = String::from_utf8([b"early\n".as_slice(), late_output].concat())?;
assert_eq!(
(item.status, item.exit_code, item.aggregated_output),
(
CommandExecutionStatus::Completed,
Some(0),
Some(expected_output)
)
);
Ok(())
}
#[tokio::test]
async fn streaming_output_preserves_multibyte_characters_across_chunks() -> anyhow::Result<()> {
let StreamingOutputHarness {
@@ -79,8 +149,9 @@ async fn streaming_output_preserves_multibyte_characters_across_chunks() -> anyh
exit_tx,
transcript,
rx_event,
..
context,
} = streaming_output_harness().await?;
start_streaming_output(&process, &context);
let output_drained = process.output_drained_notify();
let drained = output_drained.notified();
tokio::pin!(drained);
@@ -121,8 +192,10 @@ async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyho
stdout_tx,
exit_tx,
transcript,
context,
..
} = streaming_output_harness().await?;
start_streaming_output(&process, &context);
let output_drained = process.output_drained_notify();
let drained = output_drained.notified();
tokio::pin!(drained);
@@ -161,8 +234,9 @@ async fn streaming_output_keeps_grace_as_fallback_without_close() -> anyhow::Res
exit_tx,
transcript,
rx_event,
..
context,
} = streaming_output_harness().await?;
start_streaming_output(&process, &context);
let output_drained = process.output_drained_notify();
let drained = output_drained.notified();
tokio::pin!(drained);
@@ -211,6 +285,7 @@ async fn exit_watcher_waits_for_late_network_denial_before_classifying_end() ->
mut context,
rx_event,
} = streaming_output_harness().await?;
start_streaming_output(&process, &context);
tokio::time::pause();
let process_for_late_denial = Arc::clone(&process);
@@ -311,12 +386,10 @@ fn utf8_boundary_batches_malformed_output() {
}
#[tokio::test]
async fn streaming_output_bounds_invalid_bytes_and_keeps_the_full_transcript() {
async fn streaming_output_bounds_invalid_bytes() {
let (session, turn, rx_event) = make_session_and_context_with_rx().await;
let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
let mut output = Buffer::<8> {
pending: Vec::new(),
transcript: Arc::clone(&transcript),
emitter: Emitter {
remaining_deltas: 2,
session,
@@ -345,11 +418,4 @@ async fn streaming_output_bounds_invalid_bytes_and_keeps_the_full_transcript() {
b"\xf0\x9f\x98\x80\xff\xff\xff".to_vec(),
]
);
let mut expected_transcript = bytes.to_vec();
expected_transcript.extend([0xfe, 0xfe]);
assert_eq!(
transcript.lock().await.to_bytes_with_omission_marker(),
expected_transcript
);
}
codex-rs/core/src/unified_exec/process.rs
+14
View File
@@ -92,6 +92,8 @@ pub(crate) struct UnifiedExecProcess {
process_handle: ProcessHandle,
output_tx: broadcast::Sender<Vec<u8>>,
output: OutputHandles,
// Keep completion output independent of polling and lossy delta delivery.
transcript: Arc<Mutex<HeadTailBuffer>>,
output_drained: Arc<Notify>,
interaction_lock: Arc<Mutex<()>>,
state_tx: watch::Sender<ProcessState>,
@@ -135,6 +137,7 @@ impl UnifiedExecProcess {
process_handle,
output_tx,
output,
transcript: Arc::new(Mutex::new(HeadTailBuffer::default())),
output_drained,
interaction_lock: Arc::new(Mutex::new(())),
state_tx,
@@ -176,6 +179,10 @@ impl UnifiedExecProcess {
&self.output
}
pub(super) fn transcript(&self) -> Arc<Mutex<HeadTailBuffer>> {
Arc::clone(&self.transcript)
}
pub(super) fn output_receiver(&self) -> tokio::sync::broadcast::Receiver<Vec<u8>> {
self.output_tx.subscribe()
}
@@ -360,6 +367,7 @@ impl UnifiedExecProcess {
managed.output_task = Some(Self::spawn_local_output_task(
output_rx,
managed.output_handles().clone(),
managed.transcript(),
managed.output_tx.clone(),
));
@@ -409,6 +417,7 @@ impl UnifiedExecProcess {
managed.output_task = Some(Self::spawn_exec_server_output_task(
started,
output_handles,
managed.transcript(),
managed.output_tx.clone(),
managed.state_tx.clone(),
));
@@ -437,6 +446,7 @@ impl UnifiedExecProcess {
fn spawn_exec_server_output_task(
started: StartedExecProcess,
output_handles: OutputHandles,
transcript: Arc<Mutex<HeadTailBuffer>>,
output_tx: broadcast::Sender<Vec<u8>>,
state_tx: watch::Sender<ProcessState>,
) -> JoinHandle<()> {
@@ -517,6 +527,7 @@ impl UnifiedExecProcess {
} = response;
for chunk in chunks.into_iter().filter(|chunk| chunk.seq > last_seq) {
let bytes = chunk.chunk.into_inner();
transcript.lock().await.push_chunk(&bytes);
let mut guard = output_buffer.lock().await;
guard.push_chunk(&bytes);
drop(guard);
@@ -560,6 +571,7 @@ impl UnifiedExecProcess {
}
last_seq = chunk.seq;
let bytes = chunk.chunk.into_inner();
transcript.lock().await.push_chunk(&bytes);
let mut guard = output_buffer.lock().await;
guard.push_chunk(&bytes);
drop(guard);
@@ -604,6 +616,7 @@ impl UnifiedExecProcess {
fn spawn_local_output_task(
mut receiver: tokio::sync::broadcast::Receiver<Vec<u8>>,
output_handles: OutputHandles,
transcript: Arc<Mutex<HeadTailBuffer>>,
output_tx: broadcast::Sender<Vec<u8>>,
) -> JoinHandle<()> {
let OutputHandles {
@@ -621,6 +634,7 @@ impl UnifiedExecProcess {
loop {
match receiver.recv().await {
Ok(chunk) => {
transcript.lock().await.push_chunk(&chunk);
let mut guard = output_buffer.lock().await;
guard.push_chunk(&chunk);
drop(guard);
codex-rs/core/src/unified_exec/process_manager.rs
+2 -2
View File
@@ -557,7 +557,7 @@ impl UnifiedExecProcessManager {
)
});
let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
let transcript = process.transcript();
let model_context = context.step_context.model_context();
let mut event_ctx = ToolEventCtx::new(
context.session.as_ref(),
@@ -595,7 +595,7 @@ impl UnifiedExecProcessManager {
);
emitter.emit(event_ctx, ToolEventStage::Begin).await;
start_streaming_output(&process, context, Arc::clone(&transcript));
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.
codex-rs/core/tests/suite/unified_exec.rs
+5 -6
View File
@@ -1136,9 +1136,9 @@ async fn unified_exec_full_lifecycle_with_background_end_event() -> Result<()> {
let test = builder.build_with_auto_env(&server).await?;
let call_id = "uexec-full-lifecycle";
// This timing force the long-standing PTY
// Print before the subscriber attaches, then keep the process alive.
let args = json!({
"cmd": "sleep 0.5; printf 'HELLO-FULL-LIFECYCLE'",
"cmd": "printf 'EARLY-OUTPUT'; sleep 0.5; printf 'HELLO-FULL-LIFECYCLE'",
"yield_time_ms": 1000,
});
@@ -1205,10 +1205,9 @@ async fn unified_exec_full_lifecycle_with_background_end_event() -> Result<()> {
end_event.process_id.is_some(),
"end event should include process_id emitted by background watcher"
);
assert!(
end_event.aggregated_output.contains("HELLO-FULL-LIFECYCLE"),
"aggregated_output should contain the full PTY transcript; got {:?}",
end_event.aggregated_output
assert_eq!(
end_event.aggregated_output,
"EARLY-OUTPUTHELLO-FULL-LIFECYCLE"
);
Ok(())
}