app-server: add ctrl-c restart websocket test

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

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

@@ -1,7 +1,9 @@
 use anyhow::Context;
 use anyhow::Result;
 use anyhow::bail;
+use app_test_support::create_final_assistant_message_sse_response;
 use app_test_support::create_mock_responses_server_sequence_unchecked;
+use app_test_support::to_response;
 use codex_app_server_protocol::ClientInfo;
 use codex_app_server_protocol::InitializeParams;
 use codex_app_server_protocol::JSONRPCError;
@@ -9,11 +11,18 @@ use codex_app_server_protocol::JSONRPCMessage;
 use codex_app_server_protocol::JSONRPCRequest;
 use codex_app_server_protocol::JSONRPCResponse;
 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 serde_json::json;
 use std::net::SocketAddr;
 use std::path::Path;
+#[cfg(unix)]
+use std::process::Command as StdCommand;
 use std::process::Stdio;
 use tempfile::TempDir;
 use tokio::io::AsyncBufReadExt;
@@ -27,6 +36,12 @@ use tokio_tungstenite::MaybeTlsStream;
 use tokio_tungstenite::WebSocketStream;
 use tokio_tungstenite::connect_async;
 use tokio_tungstenite::tungstenite::Message as WebSocketMessage;
+#[cfg(unix)]
+use wiremock::Mock;
+#[cfg(unix)]
+use wiremock::matchers::method;
+#[cfg(unix)]
+use wiremock::matchers::path_regex;
 
 const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(5);
 
@@ -78,6 +93,55 @@ async fn websocket_transport_routes_per_connection_handshake_and_responses() ->
     Ok(())
 }
 
+#[cfg(unix)]
+#[tokio::test]
+async fn websocket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Result<()> {
+    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(Duration::from_secs(3)),
+        )
+        .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 mut 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?;
+
+    send_sigint(&process)?;
+    assert_process_does_not_exit_within(&mut process, Duration::from_millis(300)).await?;
+
+    let status = timeout(Duration::from_secs(10), process.wait())
+        .await
+        .context("timed out waiting for graceful Ctrl-C restart shutdown")?
+        .context("failed waiting for websocket app-server process exit")?;
+    assert!(status.success(), "expected graceful exit, got {status}");
+
+    expect_websocket_disconnect(&mut ws).await?;
+
+    Ok(())
+}
+
 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")?;
@@ -157,6 +221,36 @@ async fn send_config_read_request(stream: &mut WsClient, id: i64) -> Result<()>
     .await
 }
 
+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 send_request(
     stream: &mut WsClient,
     method: &str,
@@ -235,6 +329,76 @@ async fn assert_no_message(stream: &mut WsClient, wait_for: Duration) -> Result<
     }
 }
 
+#[cfg(unix)]
+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;
+    }
+}
+
+#[cfg(unix)]
+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(())
+}
+
+#[cfg(unix)]
+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"),
+    }
+}
+
+#[cfg(unix)]
+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(()),
+        }
+    }
+}
+
 fn create_config_toml(
     codex_home: &Path,
     server_uri: &str,