fix(core): flaky test completed_commands_do_not_persist_sessions (#3596)

Fix flaky test:
```
        FAIL [   2.641s] codex-core unified_exec::tests::completed_commands_do_not_persist_sessions
  stdout ───

    running 1 test
    test unified_exec::tests::completed_commands_do_not_persist_sessions ... FAILED

    failures:

    failures:
        unified_exec::tests::completed_commands_do_not_persist_sessions

    test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 235 filtered out; finished in 2.63s
    
  stderr ───

    thread 'unified_exec::tests::completed_commands_do_not_persist_sessions' panicked at core/src/unified_exec/mod.rs:582:9:
    assertion failed: result.output.contains("codex")
```
This commit is contained in:
Fouad Matin
2025-09-14 18:04:05 -07:00
committed by GitHub
parent 4dffa496ac
commit 5185d69f13
3 changed files with 27 additions and 3 deletions

View File

@@ -11,6 +11,9 @@ pub(crate) struct ExecCommandSession {
/// Broadcast stream of output chunks read from the PTY. New subscribers /// Broadcast stream of output chunks read from the PTY. New subscribers
/// receive only chunks emitted after they subscribe. /// receive only chunks emitted after they subscribe.
output_tx: broadcast::Sender<Vec<u8>>, output_tx: broadcast::Sender<Vec<u8>>,
/// Receiver subscribed before the child process starts emitting output so
/// the first caller can consume any early data without races.
initial_output_rx: StdMutex<Option<broadcast::Receiver<Vec<u8>>>>,
/// Child killer handle for termination on drop (can signal independently /// Child killer handle for termination on drop (can signal independently
/// of a thread blocked in `.wait()`). /// of a thread blocked in `.wait()`).
@@ -42,6 +45,7 @@ impl ExecCommandSession {
Self { Self {
writer_tx, writer_tx,
output_tx, output_tx,
initial_output_rx: StdMutex::new(None),
killer: StdMutex::new(Some(killer)), killer: StdMutex::new(Some(killer)),
reader_handle: StdMutex::new(Some(reader_handle)), reader_handle: StdMutex::new(Some(reader_handle)),
writer_handle: StdMutex::new(Some(writer_handle)), writer_handle: StdMutex::new(Some(writer_handle)),
@@ -50,12 +54,26 @@ impl ExecCommandSession {
} }
} }
pub(crate) fn set_initial_output_receiver(&self, receiver: broadcast::Receiver<Vec<u8>>) {
if let Ok(mut guard) = self.initial_output_rx.lock()
&& guard.is_none()
{
*guard = Some(receiver);
}
}
pub(crate) fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> { pub(crate) fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> {
self.writer_tx.clone() self.writer_tx.clone()
} }
pub(crate) fn output_receiver(&self) -> broadcast::Receiver<Vec<u8>> { pub(crate) fn output_receiver(&self) -> broadcast::Receiver<Vec<u8>> {
self.output_tx.subscribe() if let Ok(mut guard) = self.initial_output_rx.lock()
&& let Some(receiver) = guard.take()
{
receiver
} else {
self.output_tx.subscribe()
}
} }
pub(crate) fn has_exited(&self) -> bool { pub(crate) fn has_exited(&self) -> bool {

View File

@@ -279,6 +279,7 @@ async fn create_exec_command_session(
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128); let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
// Broadcast for streaming PTY output to readers: subscribers receive from subscription time. // Broadcast for streaming PTY output to readers: subscribers receive from subscription time.
let (output_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(256); let (output_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(256);
let initial_output_rx = output_tx.subscribe();
// Reader task: drain PTY and forward chunks to output channel. // Reader task: drain PTY and forward chunks to output channel.
let mut reader = pair.master.try_clone_reader()?; let mut reader = pair.master.try_clone_reader()?;
@@ -350,6 +351,7 @@ async fn create_exec_command_session(
wait_handle, wait_handle,
exit_status, exit_status,
); );
session.set_initial_output_receiver(initial_output_rx);
Ok((session, exit_rx)) Ok((session, exit_rx))
} }

View File

@@ -327,6 +327,7 @@ async fn create_unified_exec_session(
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128); let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
let (output_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(256); let (output_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(256);
let initial_output_rx = output_tx.subscribe();
let mut reader = pair let mut reader = pair
.master .master
@@ -380,7 +381,7 @@ async fn create_unified_exec_session(
wait_exit_status.store(true, Ordering::SeqCst); wait_exit_status.store(true, Ordering::SeqCst);
}); });
Ok(ExecCommandSession::new( let session = ExecCommandSession::new(
writer_tx, writer_tx,
output_tx, output_tx,
killer, killer,
@@ -388,7 +389,10 @@ async fn create_unified_exec_session(
writer_handle, writer_handle,
wait_handle, wait_handle,
exit_status, exit_status,
)) );
session.set_initial_output_receiver(initial_output_rx);
Ok(session)
} }
#[cfg(test)] #[cfg(test)]