Gate new app-server work during graceful shutdown (#43959)

## Why

New client requests and automatic continuations can start more work while the app-server is draining. Shutdown also needs to account for requests still preparing or submitting work before they appear as running turns.

## What changed

- Close a shared admission gate when shutdown begins, rejecting new turn work and thread lifecycle changes with the server-draining error.
- Apply the gate to automatic turn starts, including queued work and goal continuations, and recheck serialized turn requests before execution.
- Wait for admitted requests and running turns to finish, while preserving forced shutdown. Keep reads and `turn/interrupt` available during drain.

## Testing

Add coverage for admission permit tracking, discarded queued requests, forced shutdown, and WebSocket drain behavior, including request rejection, interruption, and suppression of automatic continuations.

GitOrigin-RevId: 1bc108aa00a427cfa41974fa9ec60754c2b968c2

codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs

@@ -4,8 +4,10 @@ use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT;
 use super::connection_handling_websocket::WsClient;
 use super::connection_handling_websocket::connect_websocket;
 use super::connection_handling_websocket::create_config_toml;
+use super::connection_handling_websocket::read_error_for_id;
+use super::connection_handling_websocket::read_jsonrpc_message;
+use super::connection_handling_websocket::read_notification_for_method;
 use super::connection_handling_websocket::read_response_for_id;
-use super::connection_handling_websocket::send_initialize_request;
 use super::connection_handling_websocket::send_request;
 use super::connection_handling_websocket::spawn_websocket_server;
 use anyhow::Context;
@@ -13,19 +15,28 @@ use anyhow::Result;
 use anyhow::bail;
 use app_test_support::create_final_assistant_message_sse_response;
 use app_test_support::to_response;
+use codex_app_server_protocol::ClientInfo;
+use codex_app_server_protocol::InitializeCapabilities;
+use codex_app_server_protocol::InitializeParams;
+use codex_app_server_protocol::JSONRPCMessage;
 use codex_app_server_protocol::RequestId;
 use codex_app_server_protocol::ThreadStartParams;
 use codex_app_server_protocol::ThreadStartResponse;
 use codex_app_server_protocol::TurnStartParams;
+use codex_app_server_protocol::TurnStartResponse;
 use codex_app_server_protocol::UserInput as V2UserInput;
 use core_test_support::responses;
+use core_test_support::streaming_sse::StreamingSseChunk;
+use core_test_support::streaming_sse::start_streaming_sse_server;
 use futures::SinkExt;
 use futures::StreamExt;
 use pretty_assertions::assert_eq;
+use serde_json::json;
 #[cfg(unix)]
 use std::process::Command as StdCommand;
 use tempfile::TempDir;
 use tokio::process::Child;
+use tokio::sync::oneshot;
 use tokio::time::Duration;
 use tokio::time::Instant;
 use tokio::time::sleep;
@@ -43,11 +54,60 @@ async fn websocket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Resu
         _server,
         mut process,
         mut ws,
-    } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
+        thread_id,
+        turn_id,
+    } = start_ctrl_c_restart_fixture(Duration::from_secs(6)).await?;
 
     send_sigint(&process)?;
     assert_process_does_not_exit_within(&mut process, Duration::from_millis(300)).await?;
 
+    send_turn_start_request(&mut ws, /*id*/ 4, &thread_id).await?;
+    let rejected = read_error_for_id(&mut ws, /*id*/ 4).await?;
+    assert_eq!(rejected.error.code, -32600);
+    assert_eq!(rejected.error.data, None);
+
+    send_request(
+        &mut ws,
+        "thread/read",
+        /*id*/ 5,
+        Some(json!({"threadId": thread_id, "includeTurns": false})),
+    )
+    .await?;
+    read_response_for_id(&mut ws, /*id*/ 5).await?;
+
+    for (method, params) in [
+        (
+            "thread/shellCommand",
+            json!({"threadId": thread_id, "command": "echo blocked"}),
+        ),
+        (
+            "turn/steer",
+            json!({"threadId": thread_id, "expectedTurnId": turn_id, "input": []}),
+        ),
+        ("thread/queue/start", json!({"threadId": thread_id})),
+        ("thread/start", json!({})),
+        ("thread/fork", json!({"threadId": thread_id})),
+        ("thread/resume", json!({"threadId": thread_id})),
+        (
+            "thread/rollback",
+            json!({"threadId": thread_id, "numTurns": 1}),
+        ),
+        (
+            "thread/revert",
+            json!({"threadId": thread_id, "beforeTurnId": "unused"}),
+        ),
+        (
+            "review/start",
+            json!({"threadId": thread_id, "target": {"type": "uncommittedChanges"}}),
+        ),
+    ] {
+        send_request(&mut ws, method, /*id*/ 9, Some(params)).await?;
+        assert_eq!(
+            read_error_for_id(&mut ws, /*id*/ 9).await?.error,
+            rejected.error
+        );
+    }
+
     let status = wait_for_process_exit_within(
         &mut process,
         Duration::from_secs(10),
@@ -69,6 +129,7 @@ async fn websocket_transport_second_ctrl_c_forces_exit_while_turn_running() -> R
         _server,
         mut process,
         mut ws,
+        ..
     } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
 
     send_sigint(&process)?;
@@ -96,6 +157,7 @@ async fn websocket_transport_sigterm_waits_for_running_turn_before_exit() -> Res
         _server,
         mut process,
         mut ws,
+        ..
     } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
 
     send_sigterm(&process, _codex_home.path())?;
@@ -122,6 +184,7 @@ async fn websocket_transport_second_sigterm_forces_exit_while_turn_running() ->
         _server,
         mut process,
         mut ws,
+        ..
     } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
 
     send_sigterm(&process, _codex_home.path())?;
@@ -149,6 +212,7 @@ async fn websocket_transport_repeated_sighup_keeps_waiting_for_running_turn() ->
         _server,
         mut process,
         mut ws,
+        ..
     } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
 
     send_sighup(&process)?;
@@ -170,11 +234,140 @@ async fn websocket_transport_repeated_sighup_keeps_waiting_for_running_turn() ->
     Ok(())
 }
 
+#[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?;
+
+    send_sigint(&process)?;
+    assert_process_does_not_exit_within(&mut process, Duration::from_millis(300)).await?;
+    send_turn_start_request(&mut ws, /*id*/ 5, &thread_id).await?;
+    assert_eq!(
+        read_error_for_id(&mut ws, /*id*/ 5).await?.error.code,
+        -32600
+    );
+    send_request(
+        &mut ws,
+        "turn/interrupt",
+        /*id*/ 4,
+        Some(json!({"threadId": thread_id, "turnId": turn_id})),
+    )
+    .await?;
+    read_response_for_id(&mut ws, /*id*/ 4).await?;
+    let status = wait_for_process_exit_within(
+        &mut process,
+        Duration::from_secs(10),
+        "timed out waiting for shutdown after turn interruption",
+    )
+    .await?;
+    assert!(status.success());
+    Ok(())
+}
+
+#[cfg(unix)]
+#[test_case::test_case("thread/queue/add", json!({
+    "clientUserMessageId": "queued-before-drain",
+    "input": [{"type": "text", "text": "continue from queue"}],
+}); "persisted queue")]
+#[test_case::test_case("thread/goal/set", json!({
+    "objective": "continue the goal",
+}); "goal continuation")]
+#[tokio::test]
+async fn websocket_transport_drain_stops_automatic_turns(
+    rpc_method: &str,
+    mut params: serde_json::Value,
+) -> Result<()> {
+    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?)?;
+    send_turn_start_request(&mut ws, /*id*/ 3, &thread.id).await?;
+    read_response_for_id(&mut ws, /*id*/ 3).await?;
+    timeout(
+        DEFAULT_READ_TIMEOUT,
+        server.wait_for_request_count(/*count*/ 1),
+    )
+    .await?;
+
+    params["threadId"] = thread.id.into();
+    send_request(&mut ws, rpc_method, /*id*/ 4, Some(params)).await?;
+    read_response_for_id(&mut ws, /*id*/ 4).await?;
+    // Keep another turn active while the target turn runs its idle hooks.
+    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)?;
+    // Observe admission closing before allowing either response to finish.
+    timeout(DEFAULT_READ_TIMEOUT, async {
+        loop {
+            send_thread_start_request(&mut ws, /*id*/ 12).await?;
+            loop {
+                match read_jsonrpc_message(&mut ws).await? {
+                    JSONRPCMessage::Error(error) if error.id == RequestId::Integer(12) => {
+                        assert_eq!(error.error.code, -32600);
+                        return Ok::<_, anyhow::Error>(());
+                    }
+                    JSONRPCMessage::Response(response) if response.id == RequestId::Integer(12) => {
+                        break;
+                    }
+                    _ => {}
+                }
+            }
+        }
+    })
+    .await??;
+    release_target.send(()).expect("target response is waiting");
+    read_notification_for_method(&mut ws, "turn/completed").await?;
+    assert_process_does_not_exit_within(&mut process, Duration::from_millis(300)).await?;
+    release_guard.send(()).expect("guard response is waiting");
+    assert!(
+        wait_for_process_exit_within(
+            &mut process,
+            Duration::from_secs(15),
+            "automatic continuation kept the daemon alive during drain",
+        )
+        .await?
+        .success()
+    );
+    assert_eq!(2, server.requests().await.len());
+    Ok(())
+}
+
 struct GracefulCtrlCFixture {
     _codex_home: TempDir,
     _server: wiremock::MockServer,
     process: Child,
     ws: WsClient,
+    thread_id: String,
+    turn_id: String,
 }
 
 async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result<GracefulCtrlCFixture> {
@@ -187,15 +380,7 @@ async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result<GracefulCt
         .mount(&server)
         .await;
 
-    let codex_home = TempDir::new()?;
-    create_config_toml(codex_home.path(), &server.uri(), "never")?;
-
-    let (process, bind_addr) = spawn_websocket_server(codex_home.path()).await?;
-    let mut ws = connect_websocket(bind_addr).await?;
-
-    send_initialize_request(&mut ws, /*id*/ 1, "ws_graceful_shutdown").await?;
-    let init_response = read_response_for_id(&mut ws, /*id*/ 1).await?;
-    assert_eq!(init_response.id, RequestId::Integer(1));
+    let (codex_home, process, mut ws) = start_ctrl_c_restart_client(&server.uri()).await?;
 
     send_thread_start_request(&mut ws, /*id*/ 2).await?;
     let thread_start_response = read_response_for_id(&mut ws, /*id*/ 2).await?;
@@ -203,7 +388,7 @@ async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result<GracefulCt
 
     send_turn_start_request(&mut ws, /*id*/ 3, &thread.id).await?;
     let turn_start_response = read_response_for_id(&mut ws, /*id*/ 3).await?;
-    assert_eq!(turn_start_response.id, RequestId::Integer(3));
+    let TurnStartResponse { turn } = to_response(turn_start_response)?;
 
     wait_for_responses_post(&server, Duration::from_secs(5)).await?;
 
@@ -212,9 +397,41 @@ async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result<GracefulCt
         _server: server,
         process,
         ws,
+        thread_id: thread.id,
+        turn_id: turn.id,
     })
 }
 
+async fn start_ctrl_c_restart_client(server_uri: &str) -> Result<(TempDir, Child, WsClient)> {
+    let codex_home = TempDir::new()?;
+    create_config_toml(codex_home.path(), server_uri, "never")?;
+
+    let (process, bind_addr) = spawn_websocket_server(codex_home.path()).await?;
+    let mut ws = connect_websocket(bind_addr).await?;
+
+    send_request(
+        &mut ws,
+        "initialize",
+        /*id*/ 1,
+        Some(serde_json::to_value(InitializeParams {
+            client_info: ClientInfo {
+                name: "ws_graceful_shutdown".to_string(),
+                title: None,
+                version: "0.1.0".to_string(),
+            },
+            capabilities: Some(InitializeCapabilities {
+                experimental_api: true,
+                ..Default::default()
+            }),
+        })?),
+    )
+    .await?;
+    let init_response = read_response_for_id(&mut ws, /*id*/ 1).await?;
+    assert_eq!(init_response.id, RequestId::Integer(1));
+
+    Ok((codex_home, process, ws))
+}
+
 async fn send_thread_start_request(stream: &mut WsClient, id: i64) -> Result<()> {
     send_request(
         stream,