Preserve the app-server shutdown signal future across loop iterations (#45849)

## What changed

Keep `shutdown_signal()` pinned across `tokio::select!` iterations and reset it only when it completes, so other events do not cancel and recreate the pending signal listener.

## Testing

Update the WebSocket drain interruption test to use gated responses and keep a second turn active until the `turn/interrupt` reply arrives, preventing shutdown from racing the reply.

GitOrigin-RevId: c939d84ed0e09c2ee49229e81a3becdfa6a4c7f1
This commit is contained in:
jgershen-oai
2026-09-16 03:30:44 +00:00
committed by copyberry
parent ffae979216
commit 90f7b37d23
2 changed files with 42 additions and 9 deletions
+3 -1
View File
@@ -1022,6 +1022,7 @@ pub async fn run_main_with_transport_options(
let mut snapshot_finished = !managed_daemon;
let mut clients_disconnected = false;
let mut shutdown_state = ShutdownState::default();
let mut shutdown_signal_future = Box::pin(shutdown_signal());
let exit_reason = loop {
// Sample submissions first: one can publish a running turn before
// releasing its permit, and shutdown must observe that new turn.
@@ -1057,7 +1058,8 @@ pub async fn run_main_with_transport_options(
_ = &mut snapshot, if shutdown_state.requested() && active_admissions == 0 && !snapshot_finished => {
snapshot_finished = true;
}
shutdown_signal_result = shutdown_signal(), if graceful_signal_restart_enabled && !shutdown_state.forced() => {
shutdown_signal_result = &mut shutdown_signal_future, if graceful_signal_restart_enabled && !shutdown_state.forced() => {
shutdown_signal_future.set(shutdown_signal());
let signal = match shutdown_signal_result {
Ok(signal) => signal,
Err(err) => {
@@ -242,14 +242,44 @@ async fn websocket_transport_repeated_sighup_keeps_waiting_for_running_turn() ->
#[cfg(unix)]
#[tokio::test]
async fn websocket_transport_allows_turn_interrupt_during_drain() -> Result<()> {
let GracefulCtrlCFixture {
_codex_home,
_server,
mut process,
mut ws,
thread_id,
turn_id,
} = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
let (_release_target, target_gate) = oneshot::channel();
let (release_guard, guard_gate) = oneshot::channel();
let (server, _completions) = start_streaming_sse_server(vec![
vec![StreamingSseChunk {
gate: Some(target_gate),
body: create_final_assistant_message_sse_response("Done")?,
}],
vec![StreamingSseChunk {
gate: Some(guard_gate),
body: create_final_assistant_message_sse_response("Guard done")?,
}],
])
.await;
let (_codex_home, mut process, mut ws) = start_ctrl_c_restart_client(server.uri()).await?;
send_thread_start_request(&mut ws, /*id*/ 2).await?;
let ThreadStartResponse { thread, .. } =
to_response(read_response_for_id(&mut ws, /*id*/ 2).await?)?;
let thread_id = thread.id;
send_turn_start_request(&mut ws, /*id*/ 3, &thread_id).await?;
let TurnStartResponse { turn } = to_response(read_response_for_id(&mut ws, /*id*/ 3).await?)?;
let turn_id = turn.id;
timeout(
DEFAULT_READ_TIMEOUT,
server.wait_for_request_count(/*count*/ 1),
)
.await?;
// Keep another turn active so shutdown cannot race the interrupt reply.
send_thread_start_request(&mut ws, /*id*/ 10).await?;
let ThreadStartResponse { thread, .. } =
to_response(read_response_for_id(&mut ws, /*id*/ 10).await?)?;
send_turn_start_request(&mut ws, /*id*/ 11, &thread.id).await?;
read_response_for_id(&mut ws, /*id*/ 11).await?;
timeout(
DEFAULT_READ_TIMEOUT,
server.wait_for_request_count(/*count*/ 2),
)
.await?;
send_sigint(&process)?;
assert_process_does_not_exit_within(&mut process, Duration::from_millis(300)).await?;
@@ -266,6 +296,7 @@ async fn websocket_transport_allows_turn_interrupt_during_drain() -> Result<()>
)
.await?;
read_response_for_id(&mut ws, /*id*/ 4).await?;
release_guard.send(()).expect("guard response is waiting");
let status = wait_for_process_exit_within(
&mut process,
Duration::from_secs(10),