codex-rs/app-server: graceful websocket restart on Ctrl-C (#12517)

## Summary
- add graceful websocket app-server restart on Ctrl-C by draining until
no assistant turns are running
- stop the websocket acceptor and disconnect existing connections once
the drain condition is met
- add a websocket integration test that verifies Ctrl-C waits for an
in-flight turn before exit

## Verification
- `cargo check -p codex-app-server --quiet`
- `cargo test -p codex-app-server --test all
suite::v2::connection_handling_websocket`
- I (maxj) tested remote and local Codex.app

---------

Co-authored-by: Codex <noreply@openai.com>

codex-rs/app-server/src/codex_message_processor.rs

@@ -269,6 +269,7 @@ use std::time::SystemTime;
 use tokio::sync::Mutex;
 use tokio::sync::broadcast;
 use tokio::sync::oneshot;
+use tokio::sync::watch;
 use toml::Value as TomlValue;
 use tracing::error;
 use tracing::info;
@@ -2796,6 +2797,10 @@ impl CodexMessageProcessor {
             .await;
     }
 
+    pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver<usize> {
+        self.thread_watch_manager.subscribe_running_turn_count()
+    }
+
     /// Best-effort: ensure initialized connections are subscribed to this thread.
     pub(crate) async fn try_attach_thread_listener(
         &mut self,

codex-rs/app-server/src/lib.rs

@@ -98,6 +98,72 @@ enum OutboundControlEvent {
     },
     /// Remove state for a closed/disconnected connection.
     Closed { connection_id: ConnectionId },
+    /// Disconnect all connection-oriented clients during graceful restart.
+    DisconnectAll,
+}
+
+#[derive(Default)]
+struct ShutdownState {
+    requested: bool,
+    forced: bool,
+    last_logged_running_turn_count: Option<usize>,
+}
+
+enum ShutdownAction {
+    Noop,
+    Finish,
+}
+
+impl ShutdownState {
+    fn requested(&self) -> bool {
+        self.requested
+    }
+
+    fn forced(&self) -> bool {
+        self.forced
+    }
+
+    fn on_ctrl_c(&mut self, connection_count: usize, running_turn_count: usize) {
+        if self.requested {
+            self.forced = true;
+            return;
+        }
+
+        self.requested = true;
+        self.last_logged_running_turn_count = None;
+        info!(
+            "received Ctrl-C; entering graceful restart drain (connections={}, runningAssistantTurns={}, requests still accepted until no assistant turns are running)",
+            connection_count, running_turn_count,
+        );
+    }
+
+    fn update(&mut self, running_turn_count: usize, connection_count: usize) -> ShutdownAction {
+        if !self.requested {
+            return ShutdownAction::Noop;
+        }
+
+        if self.forced || running_turn_count == 0 {
+            if self.forced {
+                info!(
+                    "received second Ctrl-C; forcing restart with {running_turn_count} running assistant turn(s) and {connection_count} connection(s)"
+                );
+            } else {
+                info!(
+                    "Ctrl-C restart: no assistant turns running; stopping acceptor and disconnecting {connection_count} connection(s)"
+                );
+            }
+            return ShutdownAction::Finish;
+        }
+
+        if self.last_logged_running_turn_count != Some(running_turn_count) {
+            info!(
+                "Ctrl-C restart: waiting for {running_turn_count} running assistant turn(s) to finish"
+            );
+            self.last_logged_running_turn_count = Some(running_turn_count);
+        }
+
+        ShutdownAction::Noop
+    }
 }
 
 fn config_warning_from_error(
@@ -253,19 +319,37 @@ pub async fn run_main_with_transport(
     let (outbound_control_tx, mut outbound_control_rx) =
         mpsc::channel::<OutboundControlEvent>(CHANNEL_CAPACITY);
 
+    enum TransportRuntime {
+        Stdio,
+        WebSocket {
+            accept_handle: JoinHandle<()>,
+            shutdown_token: CancellationToken,
+        },
+    }
+
     let mut stdio_handles = Vec::<JoinHandle<()>>::new();
-    let mut websocket_accept_handle = None;
-    match transport {
+    let transport_runtime = match transport {
         AppServerTransport::Stdio => {
             start_stdio_connection(transport_event_tx.clone(), &mut stdio_handles).await?;
+            TransportRuntime::Stdio
         }
         AppServerTransport::WebSocket { bind_address } => {
-            websocket_accept_handle =
-                Some(start_websocket_acceptor(bind_address, transport_event_tx.clone()).await?);
+            let shutdown_token = CancellationToken::new();
+            let accept_handle = start_websocket_acceptor(
+                bind_address,
+                transport_event_tx.clone(),
+                shutdown_token.clone(),
+            )
+            .await?;
+            TransportRuntime::WebSocket {
+                accept_handle,
+                shutdown_token,
+            }
         }
-    }
-    let single_client_mode = matches!(transport, AppServerTransport::Stdio);
+    };
+    let single_client_mode = matches!(&transport_runtime, TransportRuntime::Stdio);
     let shutdown_when_no_connections = single_client_mode;
+    let graceful_ctrl_c_restart_enabled = !single_client_mode;
 
     // Parse CLI overrides once and derive the base Config eagerly so later
     // components do not need to work with raw TOML values.
@@ -434,6 +518,16 @@ pub async fn run_main_with_transport(
                             OutboundControlEvent::Closed { connection_id } => {
                                 outbound_connections.remove(&connection_id);
                             }
+                            OutboundControlEvent::DisconnectAll => {
+                                info!(
+                                    "disconnecting {} outbound websocket connection(s) for graceful restart",
+                                    outbound_connections.len()
+                                );
+                                for connection_state in outbound_connections.values() {
+                                    connection_state.request_disconnect();
+                                }
+                                outbound_connections.clear();
+                            }
                         }
                     }
                     envelope = outgoing_rx.recv() => {
@@ -464,11 +558,46 @@ pub async fn run_main_with_transport(
             config_warnings,
         });
         let mut thread_created_rx = processor.thread_created_receiver();
+        let mut running_turn_count_rx = processor.subscribe_running_assistant_turn_count();
         let mut connections = HashMap::<ConnectionId, ConnectionState>::new();
+        let websocket_accept_shutdown = match &transport_runtime {
+            TransportRuntime::WebSocket { shutdown_token, .. } => Some(shutdown_token.clone()),
+            TransportRuntime::Stdio => None,
+        };
         async move {
             let mut listen_for_threads = true;
+            let mut shutdown_state = ShutdownState::default();
             loop {
+                let running_turn_count = {
+                    let running_turn_count = running_turn_count_rx.borrow();
+                    *running_turn_count
+                };
+                if matches!(
+                    shutdown_state.update(running_turn_count, connections.len()),
+                    ShutdownAction::Finish
+                ) {
+                    if let Some(shutdown_token) = &websocket_accept_shutdown {
+                        shutdown_token.cancel();
+                    }
+                    let _ = outbound_control_tx
+                        .send(OutboundControlEvent::DisconnectAll)
+                        .await;
+                    break;
+                }
+
                 tokio::select! {
+                    ctrl_c_result = tokio::signal::ctrl_c(), if graceful_ctrl_c_restart_enabled && !shutdown_state.forced() => {
+                        if let Err(err) = ctrl_c_result {
+                            warn!("failed to listen for Ctrl-C during graceful restart drain: {err}");
+                        }
+                        let running_turn_count = *running_turn_count_rx.borrow();
+                        shutdown_state.on_ctrl_c(connections.len(), running_turn_count);
+                    }
+                    changed = running_turn_count_rx.changed(), if graceful_ctrl_c_restart_enabled && shutdown_state.requested() => {
+                        if changed.is_err() {
+                            warn!("running-turn watcher closed during graceful restart drain");
+                        }
+                    }
                     event = transport_event_rx.recv() => {
                         let Some(event) = event else {
                             break;
@@ -619,8 +748,13 @@ pub async fn run_main_with_transport(
     let _ = processor_handle.await;
     let _ = outbound_handle.await;
 
-    if let Some(handle) = websocket_accept_handle {
-        handle.abort();
+    if let TransportRuntime::WebSocket {
+        accept_handle,
+        shutdown_token,
+    } = transport_runtime
+    {
+        shutdown_token.cancel();
+        let _ = accept_handle.await;
     }
 
     for handle in stdio_handles {

codex-rs/app-server/src/message_processor.rs

@@ -51,6 +51,7 @@ use codex_protocol::ThreadId;
 use codex_protocol::protocol::SessionSource;
 use futures::FutureExt;
 use tokio::sync::broadcast;
+use tokio::sync::watch;
 use tokio::time::Duration;
 use tokio::time::timeout;
 use toml::Value as TomlValue;
@@ -438,6 +439,11 @@ impl MessageProcessor {
             .await;
     }
 
+    pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver<usize> {
+        self.codex_message_processor
+            .subscribe_running_assistant_turn_count()
+    }
+
     /// Handle a standalone JSON-RPC response originating from the peer.
     pub(crate) async fn process_response(&mut self, response: JSONRPCResponse) {
         tracing::info!("<- response: {:?}", response);

codex-rs/app-server/src/thread_status.rs

@@ -15,11 +15,13 @@ use std::sync::Arc;
 use tokio::sync::Mutex;
 #[cfg(test)]
 use tokio::sync::mpsc;
+use tokio::sync::watch;
 
 #[derive(Clone)]
 pub(crate) struct ThreadWatchManager {
     state: Arc<Mutex<ThreadWatchState>>,
     outgoing: Option<Arc<OutgoingMessageSender>>,
+    running_turn_count_tx: watch::Sender<usize>,
 }
 
 pub(crate) struct ThreadWatchActiveGuard {
@@ -71,16 +73,20 @@ impl Default for ThreadWatchManager {
 
 impl ThreadWatchManager {
     pub(crate) fn new() -> Self {
+        let (running_turn_count_tx, _running_turn_count_rx) = watch::channel(0);
         Self {
             state: Arc::new(Mutex::new(ThreadWatchState::default())),
             outgoing: None,
+            running_turn_count_tx,
         }
     }
 
     pub(crate) fn new_with_outgoing(outgoing: Arc<OutgoingMessageSender>) -> Self {
+        let (running_turn_count_tx, _running_turn_count_rx) = watch::channel(0);
         Self {
             state: Arc::new(Mutex::new(ThreadWatchState::default())),
             outgoing: Some(outgoing),
+            running_turn_count_tx,
         }
     }
 
@@ -113,6 +119,21 @@ impl ThreadWatchManager {
             .collect()
     }
 
+    #[cfg(test)]
+    pub(crate) async fn running_turn_count(&self) -> usize {
+        self.state
+            .lock()
+            .await
+            .runtime_by_thread_id
+            .values()
+            .filter(|runtime| runtime.running)
+            .count()
+    }
+
+    pub(crate) fn subscribe_running_turn_count(&self) -> watch::Receiver<usize> {
+        self.running_turn_count_tx.subscribe()
+    }
+
     pub(crate) async fn note_turn_started(&self, thread_id: &str) {
         self.update_runtime_for_thread(thread_id, |runtime| {
             runtime.is_loaded = true;
@@ -193,10 +214,17 @@ impl ThreadWatchManager {
     where
         F: FnOnce(&mut ThreadWatchState) -> Option<ThreadStatusChangedNotification>,
     {
-        let notification = {
+        let (notification, running_turn_count) = {
             let mut state = self.state.lock().await;
-            mutate(&mut state)
+            let notification = mutate(&mut state);
+            let running_turn_count = state
+                .runtime_by_thread_id
+                .values()
+                .filter(|runtime| runtime.running)
+                .count();
+            (notification, running_turn_count)
         };
+        let _ = self.running_turn_count_tx.send(running_turn_count);
 
         if let Some(notification) = notification
             && let Some(outgoing) = &self.outgoing
@@ -588,6 +616,32 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn has_running_turns_tracks_runtime_running_flag_only() {
+        let manager = ThreadWatchManager::new();
+        manager
+            .upsert_thread(test_thread(
+                INTERACTIVE_THREAD_ID,
+                codex_app_server_protocol::SessionSource::Cli,
+            ))
+            .await;
+
+        assert_eq!(manager.running_turn_count().await, 0);
+
+        let _permission_guard = manager
+            .note_permission_requested(INTERACTIVE_THREAD_ID)
+            .await;
+        assert_eq!(manager.running_turn_count().await, 0);
+
+        manager.note_turn_started(INTERACTIVE_THREAD_ID).await;
+        assert_eq!(manager.running_turn_count().await, 1);
+
+        manager
+            .note_turn_completed(INTERACTIVE_THREAD_ID, false)
+            .await;
+        assert_eq!(manager.running_turn_count().await, 0);
+    }
+
     #[tokio::test]
     async fn status_change_emits_notification() {
         let (outgoing_tx, mut outgoing_rx) = mpsc::channel(8);

codex-rs/app-server/src/transport.rs

@@ -68,12 +68,6 @@ fn print_websocket_startup_banner(addr: SocketAddr) {
     }
 }
 
-#[allow(clippy::print_stderr)]
-fn print_websocket_connection(peer_addr: SocketAddr) {
-    let connected_label = colorize("websocket client connected from", Style::new().dimmed());
-    eprintln!("{connected_label} {peer_addr}");
-}
-
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum AppServerTransport {
     Stdio,
@@ -193,7 +187,7 @@ impl OutboundConnectionState {
         self.disconnect_sender.is_some()
     }
 
-    fn request_disconnect(&self) {
+    pub(crate) fn request_disconnect(&self) {
         if let Some(disconnect_sender) = &self.disconnect_sender {
             disconnect_sender.cancel();
         }
@@ -271,6 +265,7 @@ pub(crate) async fn start_stdio_connection(
 pub(crate) async fn start_websocket_acceptor(
     bind_address: SocketAddr,
     transport_event_tx: mpsc::Sender<TransportEvent>,
+    shutdown_token: CancellationToken,
 ) -> IoResult<JoinHandle<()>> {
     let listener = TcpListener::bind(bind_address).await?;
     let local_addr = listener.local_addr()?;
@@ -280,23 +275,31 @@ pub(crate) async fn start_websocket_acceptor(
     let connection_counter = Arc::new(AtomicU64::new(1));
     Ok(tokio::spawn(async move {
         loop {
-            match listener.accept().await {
-                Ok((stream, peer_addr)) => {
-                    print_websocket_connection(peer_addr);
-                    let connection_id =
-                        ConnectionId(connection_counter.fetch_add(1, Ordering::Relaxed));
-                    let transport_event_tx_for_connection = transport_event_tx.clone();
-                    tokio::spawn(async move {
-                        run_websocket_connection(
-                            connection_id,
-                            stream,
-                            transport_event_tx_for_connection,
-                        )
-                        .await;
-                    });
+            tokio::select! {
+                _ = shutdown_token.cancelled() => {
+                    info!("websocket acceptor shutting down");
+                    break;
                 }
-                Err(err) => {
-                    error!("failed to accept websocket connection: {err}");
+                accept_result = listener.accept() => {
+                    match accept_result {
+                        Ok((stream, peer_addr)) => {
+                            info!(%peer_addr, "websocket client connected");
+                            let connection_id =
+                                ConnectionId(connection_counter.fetch_add(1, Ordering::Relaxed));
+                            let transport_event_tx_for_connection = transport_event_tx.clone();
+                            tokio::spawn(async move {
+                                run_websocket_connection(
+                                    connection_id,
+                                    stream,
+                                    transport_event_tx_for_connection,
+                                )
+                                .await;
+                            });
+                        }
+                        Err(err) => {
+                            error!("failed to accept websocket connection: {err}");
+                        }
+                    }
                 }
             }
         }

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

@@ -28,9 +28,9 @@ use tokio_tungstenite::WebSocketStream;
 use tokio_tungstenite::connect_async;
 use tokio_tungstenite::tungstenite::Message as WebSocketMessage;
 
-const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(5);
+pub(super) const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(5);
 
-type WsClient = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
+pub(super) type WsClient = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
 
 #[tokio::test]
 async fn websocket_transport_routes_per_connection_handshake_and_responses() -> Result<()> {
@@ -78,7 +78,10 @@ async fn websocket_transport_routes_per_connection_handshake_and_responses() ->
     Ok(())
 }
 
-async fn spawn_websocket_server(codex_home: &Path, bind_addr: SocketAddr) -> Result<Child> {
+pub(super) async fn spawn_websocket_server(
+    codex_home: &Path,
+    bind_addr: SocketAddr,
+) -> Result<Child> {
     let program = codex_utils_cargo_bin::cargo_bin("codex-app-server")
         .context("should find app-server binary")?;
     let mut cmd = Command::new(program);
@@ -106,14 +109,14 @@ async fn spawn_websocket_server(codex_home: &Path, bind_addr: SocketAddr) -> Res
     Ok(process)
 }
 
-fn reserve_local_addr() -> Result<SocketAddr> {
+pub(super) fn reserve_local_addr() -> Result<SocketAddr> {
     let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
     let addr = listener.local_addr()?;
     drop(listener);
     Ok(addr)
 }
 
-async fn connect_websocket(bind_addr: SocketAddr) -> Result<WsClient> {
+pub(super) async fn connect_websocket(bind_addr: SocketAddr) -> Result<WsClient> {
     let url = format!("ws://{bind_addr}");
     let deadline = Instant::now() + Duration::from_secs(10);
     loop {
@@ -129,7 +132,11 @@ async fn connect_websocket(bind_addr: SocketAddr) -> Result<WsClient> {
     }
 }
 
-async fn send_initialize_request(stream: &mut WsClient, id: i64, client_name: &str) -> Result<()> {
+pub(super) async fn send_initialize_request(
+    stream: &mut WsClient,
+    id: i64,
+    client_name: &str,
+) -> Result<()> {
     let params = InitializeParams {
         client_info: ClientInfo {
             name: client_name.to_string(),
@@ -157,7 +164,7 @@ async fn send_config_read_request(stream: &mut WsClient, id: i64) -> Result<()>
     .await
 }
 
-async fn send_request(
+pub(super) async fn send_request(
     stream: &mut WsClient,
     method: &str,
     id: i64,
@@ -179,7 +186,10 @@ async fn send_jsonrpc(stream: &mut WsClient, message: JSONRPCMessage) -> Result<
         .context("failed to send websocket frame")
 }
 
-async fn read_response_for_id(stream: &mut WsClient, id: i64) -> Result<JSONRPCResponse> {
+pub(super) async fn read_response_for_id(
+    stream: &mut WsClient,
+    id: i64,
+) -> Result<JSONRPCResponse> {
     let target_id = RequestId::Integer(id);
     loop {
         let message = read_jsonrpc_message(stream).await?;
@@ -235,7 +245,7 @@ async fn assert_no_message(stream: &mut WsClient, wait_for: Duration) -> Result<
     }
 }
 
-fn create_config_toml(
+pub(super) fn create_config_toml(
     codex_home: &Path,
     server_uri: &str,
     approval_policy: &str,

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

@@ -0,0 +1,237 @@
+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_response_for_id;
+use super::connection_handling_websocket::reserve_local_addr;
+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;
+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::RequestId;
+use codex_app_server_protocol::ThreadStartParams;
+use codex_app_server_protocol::ThreadStartResponse;
+use codex_app_server_protocol::TurnStartParams;
+use codex_app_server_protocol::UserInput as V2UserInput;
+use core_test_support::responses;
+use futures::SinkExt;
+use futures::StreamExt;
+use std::process::Command as StdCommand;
+use tempfile::TempDir;
+use tokio::process::Child;
+use tokio::time::Duration;
+use tokio::time::Instant;
+use tokio::time::sleep;
+use tokio::time::timeout;
+use tokio_tungstenite::tungstenite::Message as WebSocketMessage;
+use wiremock::Mock;
+use wiremock::matchers::method;
+use wiremock::matchers::path_regex;
+
+#[tokio::test]
+async fn websocket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Result<()> {
+    let GracefulCtrlCFixture {
+        _codex_home,
+        _server,
+        mut process,
+        mut ws,
+    } = 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?;
+
+    let status = wait_for_process_exit_within(
+        &mut process,
+        Duration::from_secs(10),
+        "timed out waiting for graceful Ctrl-C restart shutdown",
+    )
+    .await?;
+    assert!(status.success(), "expected graceful exit, got {status}");
+
+    expect_websocket_disconnect(&mut ws).await?;
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn websocket_transport_second_ctrl_c_forces_exit_while_turn_running() -> Result<()> {
+    let GracefulCtrlCFixture {
+        _codex_home,
+        _server,
+        mut process,
+        mut ws,
+    } = 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_sigint(&process)?;
+    let status = wait_for_process_exit_within(
+        &mut process,
+        Duration::from_secs(2),
+        "timed out waiting for forced Ctrl-C restart shutdown",
+    )
+    .await?;
+    assert!(status.success(), "expected graceful exit, got {status}");
+
+    expect_websocket_disconnect(&mut ws).await?;
+
+    Ok(())
+}
+
+struct GracefulCtrlCFixture {
+    _codex_home: TempDir,
+    _server: wiremock::MockServer,
+    process: Child,
+    ws: WsClient,
+}
+
+async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result<GracefulCtrlCFixture> {
+    let server = responses::start_mock_server().await;
+    let delayed_turn_response = create_final_assistant_message_sse_response("Done")?;
+    Mock::given(method("POST"))
+        .and(path_regex(".*/responses$"))
+        .respond_with(responses::sse_response(delayed_turn_response).set_delay(turn_delay))
+        .up_to_n_times(1)
+        .mount(&server)
+        .await;
+
+    let codex_home = TempDir::new()?;
+    create_config_toml(codex_home.path(), &server.uri(), "never")?;
+
+    let bind_addr = reserve_local_addr()?;
+    let process = spawn_websocket_server(codex_home.path(), bind_addr).await?;
+    let mut ws = connect_websocket(bind_addr).await?;
+
+    send_initialize_request(&mut ws, 1, "ws_graceful_shutdown").await?;
+    let init_response = read_response_for_id(&mut ws, 1).await?;
+    assert_eq!(init_response.id, RequestId::Integer(1));
+
+    send_thread_start_request(&mut ws, 2).await?;
+    let thread_start_response = read_response_for_id(&mut ws, 2).await?;
+    let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?;
+
+    send_turn_start_request(&mut ws, 3, &thread.id).await?;
+    let turn_start_response = read_response_for_id(&mut ws, 3).await?;
+    assert_eq!(turn_start_response.id, RequestId::Integer(3));
+
+    wait_for_responses_post(&server, Duration::from_secs(5)).await?;
+
+    Ok(GracefulCtrlCFixture {
+        _codex_home: codex_home,
+        _server: server,
+        process,
+        ws,
+    })
+}
+
+async fn send_thread_start_request(stream: &mut WsClient, id: i64) -> Result<()> {
+    send_request(
+        stream,
+        "thread/start",
+        id,
+        Some(serde_json::to_value(ThreadStartParams {
+            model: Some("mock-model".to_string()),
+            ..Default::default()
+        })?),
+    )
+    .await
+}
+
+async fn send_turn_start_request(stream: &mut WsClient, id: i64, thread_id: &str) -> Result<()> {
+    send_request(
+        stream,
+        "turn/start",
+        id,
+        Some(serde_json::to_value(TurnStartParams {
+            thread_id: thread_id.to_string(),
+            input: vec![V2UserInput::Text {
+                text: "Hello".to_string(),
+                text_elements: Vec::new(),
+            }],
+            ..Default::default()
+        })?),
+    )
+    .await
+}
+
+async fn wait_for_responses_post(server: &wiremock::MockServer, wait_for: Duration) -> Result<()> {
+    let deadline = Instant::now() + wait_for;
+    loop {
+        let requests = server
+            .received_requests()
+            .await
+            .context("failed to read mock server requests")?;
+        if requests
+            .iter()
+            .any(|request| request.method == "POST" && request.url.path().ends_with("/responses"))
+        {
+            return Ok(());
+        }
+        if Instant::now() >= deadline {
+            bail!("timed out waiting for /responses request");
+        }
+        sleep(Duration::from_millis(10)).await;
+    }
+}
+
+fn send_sigint(process: &Child) -> Result<()> {
+    let pid = process
+        .id()
+        .context("websocket app-server process has no pid")?;
+    let status = StdCommand::new("kill")
+        .arg("-INT")
+        .arg(pid.to_string())
+        .status()
+        .context("failed to invoke kill -INT")?;
+    if !status.success() {
+        bail!("kill -INT exited with {status}");
+    }
+    Ok(())
+}
+
+async fn assert_process_does_not_exit_within(process: &mut Child, window: Duration) -> Result<()> {
+    match timeout(window, process.wait()).await {
+        Err(_) => Ok(()),
+        Ok(Ok(status)) => bail!("process exited too early during graceful drain: {status}"),
+        Ok(Err(err)) => Err(err).context("failed waiting for process"),
+    }
+}
+
+async fn wait_for_process_exit_within(
+    process: &mut Child,
+    window: Duration,
+    timeout_context: &'static str,
+) -> Result<std::process::ExitStatus> {
+    timeout(window, process.wait())
+        .await
+        .context(timeout_context)?
+        .context("failed waiting for websocket app-server process exit")
+}
+
+async fn expect_websocket_disconnect(stream: &mut WsClient) -> Result<()> {
+    loop {
+        let frame = timeout(DEFAULT_READ_TIMEOUT, stream.next())
+            .await
+            .context("timed out waiting for websocket disconnect")?;
+        match frame {
+            None => return Ok(()),
+            Some(Ok(WebSocketMessage::Close(_))) => return Ok(()),
+            Some(Ok(WebSocketMessage::Ping(payload))) => {
+                stream
+                    .send(WebSocketMessage::Pong(payload))
+                    .await
+                    .context("failed to reply to ping while waiting for disconnect")?;
+            }
+            Some(Ok(WebSocketMessage::Pong(_))) => {}
+            Some(Ok(WebSocketMessage::Frame(_))) => {}
+            Some(Ok(WebSocketMessage::Text(_))) => {}
+            Some(Ok(WebSocketMessage::Binary(_))) => {}
+            Some(Err(_)) => return Ok(()),
+        }
+    }
+}

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

@@ -5,6 +5,8 @@ mod collaboration_mode_list;
 mod compaction;
 mod config_rpc;
 mod connection_handling_websocket;
+#[cfg(unix)]
+mod connection_handling_websocket_unix;
 mod dynamic_tools;
 mod experimental_api;
 mod experimental_feature_list;