From c3eeaae9a3d401f7e7ade0817b7b27299af7d2a1 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Wed, 9 Sep 2026 00:52:51 +0000 Subject: [PATCH] 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/src/extensions.rs | 6 + codex-rs/app-server/src/lib.rs | 72 +++++- codex-rs/app-server/src/mcp_refresh.rs | 1 + codex-rs/app-server/src/message_processor.rs | 32 +++ .../app-server/src/request_serialization.rs | 17 +- codex-rs/app-server/src/turn_admission.rs | 87 +++++++ .../app-server/src/turn_admission_tests.rs | 25 ++ .../v2/connection_handling_websocket_unix.rs | 241 +++++++++++++++++- 8 files changed, 448 insertions(+), 33 deletions(-) create mode 100644 codex-rs/app-server/src/turn_admission.rs create mode 100644 codex-rs/app-server/src/turn_admission_tests.rs diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index fba9272779..46af9c9a9c 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -19,6 +19,7 @@ use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistry; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionWarning; +use codex_extension_api::TurnStartAdmission; use codex_goal_extension::GoalExtensionConfig; use codex_goal_extension::GoalService; use codex_http_client::HttpClientFactory; @@ -48,6 +49,7 @@ pub(crate) struct ThreadExtensionDependencies { pub(crate) http_client_factory: HttpClientFactory, /// Process-scoped queue shared by idle dispatch and app-server requests. pub(crate) queue_service: Option>, + pub(crate) turn_start_admission: Option>, } pub(crate) fn thread_extensions( @@ -69,8 +71,12 @@ where git_attribution_base_url, http_client_factory, queue_service, + turn_start_admission, } = dependencies; let mut builder = ExtensionRegistryBuilder::::with_event_sink(Arc::clone(&event_sink)); + if let Some(admission) = turn_start_admission { + builder.turn_start_admission(admission); + } if let Some(queue_service) = queue_service { codex_queue_extension::install(&mut builder, queue_service); } diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 7bb96f6ad0..06269bfd5c 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -132,6 +132,7 @@ mod skills_watcher; mod thread_state; mod thread_status; mod transport; +mod turn_admission; mod turn_cost_worker; mod user_verification; mod user_verification_response; @@ -247,6 +248,7 @@ impl ShutdownState { signal: ShutdownSignal, connection_count: usize, running_turn_count: usize, + turn_admission: &turn_admission::TurnAdmission, ) { if self.requested { if matches!(signal, ShutdownSignal::Forceable) { @@ -255,20 +257,26 @@ impl ShutdownState { return; } + turn_admission.begin_drain(); self.requested = true; self.last_logged_running_turn_count = None; info!( - "received shutdown signal; entering graceful restart drain (connections={}, runningAssistantTurns={}, requests still accepted until no assistant turns are running)", + "received shutdown signal; entering graceful restart drain (connections={}, runningAssistantTurns={}, new client turns rejected)", connection_count, running_turn_count, ); } - fn update(&mut self, running_turn_count: usize, connection_count: usize) -> ShutdownAction { + fn update( + &mut self, + running_turn_count: usize, + active_admissions: usize, + connection_count: usize, + ) -> ShutdownAction { if !self.requested { return ShutdownAction::Noop; } - if self.forced || running_turn_count == 0 { + if self.forced || (running_turn_count == 0 && active_admissions == 0) { if self.forced { info!( "received second shutdown signal; forcing restart with {running_turn_count} running assistant turn(s) and {connection_count} connection(s)" @@ -283,7 +291,7 @@ impl ShutdownState { if self.last_logged_running_turn_count != Some(running_turn_count) { info!( - "shutdown signal restart: waiting for {running_turn_count} running assistant turn(s) to finish" + "shutdown signal restart: waiting for {running_turn_count} running assistant turn(s) and {active_admissions} admitted request(s) to finish" ); self.last_logged_running_turn_count = Some(running_turn_count); } @@ -954,6 +962,7 @@ pub async fn run_main_with_transport_options( })); let mut thread_created_rx = processor.thread_created_receiver(); let mut running_turn_count_rx = processor.subscribe_running_assistant_turn_count(); + let mut active_admissions_rx = processor.turn_admission.subscribe_active(); let mut connections = HashMap::::new(); let mut connection_cleanup_tasks = ConnectionCleanupTasks::new(); let mut remote_control_status_rx = remote_control_handle.status_receiver(); @@ -963,12 +972,12 @@ pub async fn run_main_with_transport_options( let mut listen_for_threads = true; let mut shutdown_state = ShutdownState::default(); let exit_reason = loop { - let running_turn_count = { - let running_turn_count = running_turn_count_rx.borrow(); - *running_turn_count - }; + // Sample submissions first: one can publish a running turn before + // releasing its permit, and shutdown must observe that new turn. + let active_admissions = *active_admissions_rx.borrow_and_update(); + let running_turn_count = *running_turn_count_rx.borrow_and_update(); if matches!( - shutdown_state.update(running_turn_count, connections.len()), + shutdown_state.update(running_turn_count, active_admissions, connections.len()), ShutdownAction::Finish ) { transport_shutdown_token.cancel(); @@ -988,20 +997,25 @@ pub async fn run_main_with_transport_options( } }; let running_turn_count = *running_turn_count_rx.borrow(); - shutdown_state.on_signal(signal, connections.len(), running_turn_count); + shutdown_state.on_signal(signal, connections.len(), running_turn_count, &processor.turn_admission); } changed = running_turn_count_rx.changed(), if shutdown_state.requested() => { if changed.is_err() { warn!("running-turn watcher closed during graceful restart drain"); } } + changed = active_admissions_rx.changed(), if shutdown_state.requested() => { + if changed.is_err() { + warn!("turn admission watcher closed during graceful restart drain"); + } + } event = transport_event_rx.recv() => { let Some(event) = event else { break "transport_channel_closed"; }; match event { TransportEvent::DaemonShutdown => { - shutdown_state.on_signal(ShutdownSignal::Forceable, connections.len(), *running_turn_count_rx.borrow()); + shutdown_state.on_signal(ShutdownSignal::Forceable, connections.len(), *running_turn_count_rx.borrow(), &processor.turn_admission); } TransportEvent::ConnectionOpened { connection_id, @@ -1414,14 +1428,50 @@ fn analytics_rpc_transport(transport: &AppServerTransport) -> AppServerRpcTransp #[cfg(test)] mod tests { use super::LogFormat; + use super::ShutdownAction; + use super::ShutdownSignal; + use super::ShutdownState; #[cfg(debug_assertions)] use super::loader_overrides_with_test_user_config_file; + use super::turn_admission::TurnAdmission; #[cfg(debug_assertions)] use codex_config::LoaderOverrides; #[cfg(debug_assertions)] use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; + #[test] + fn shutdown_waits_for_admitted_requests_but_force_remains_available() { + let admission = TurnAdmission::default(); + let active = admission.subscribe_active(); + let in_flight = admission.admit().expect("request admitted"); + let mut shutdown = ShutdownState::default(); + let connection_count = 1; + let running_turn_count = 0; + shutdown.on_signal( + ShutdownSignal::Forceable, + connection_count, + running_turn_count, + &admission, + ); + assert!(matches!( + shutdown.update(running_turn_count, *active.borrow(), connection_count), + ShutdownAction::Noop + )); + shutdown.on_signal( + ShutdownSignal::Forceable, + connection_count, + running_turn_count, + &admission, + ); + assert!(shutdown.forced()); + assert!(matches!( + shutdown.update(running_turn_count, *active.borrow(), connection_count), + ShutdownAction::Finish + )); + drop(in_flight); + } + #[test] fn log_format_from_env_value_matches_json_values_case_insensitively() { assert_eq!(LogFormat::from_env_value(Some("json")), LogFormat::Json); diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 8a2e571a52..f2c446b3a4 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -340,6 +340,7 @@ enabled = false git_attribution_base_url: good_config.chatgpt_base_url.clone(), http_client_factory: good_config.http_client_factory(), queue_service: None, + turn_start_admission: None, }, ), Arc::new(CodexHomeUserInstructionsProvider::new( diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 511ea54587..0721f00be9 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -78,6 +78,7 @@ use codex_core::ThreadManager; use codex_core::config::Config; use codex_core::config::ThreadStoreConfig; use codex_exec_server::EnvironmentManager; +use codex_extension_api::TurnStartAdmission; use codex_feedback::CodexFeedback; use codex_goal_extension::GoalService; use codex_home::CodexHomeUserInstructionsProvider; @@ -101,6 +102,7 @@ use tokio_util::sync::CancellationToken; use tracing::Instrument; use crate::models_refresh_worker::ModelsRefreshWorker; +use crate::turn_admission::TurnAdmission; const CONNECTION_RPC_DRAIN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 30); @@ -136,6 +138,7 @@ fn reject_removed_permission_profile(request: &JSONRPCRequest) -> Result<(), JSO } pub(crate) struct MessageProcessor { + pub(crate) turn_admission: TurnAdmission, user_verification: Arc, outgoing: Arc, models_refresh_worker: ModelsRefreshWorker, @@ -311,6 +314,8 @@ impl MessageProcessor { ), ); let goal_service = Arc::new(GoalService::new()); + let turn_admission = TurnAdmission::default(); + let turn_start_admission: Arc = Arc::new(turn_admission.clone()); let extension_event_sink = app_server_extension_event_sink(outgoing.clone(), thread_state_manager.clone()); let mut queue_service = None; @@ -343,6 +348,7 @@ impl MessageProcessor { git_attribution_base_url: config.chatgpt_base_url.clone(), http_client_factory: config.http_client_factory(), queue_service: queue_service.clone(), + turn_start_admission: Some(Arc::clone(&turn_start_admission)), }, ), Arc::new(CodexHomeUserInstructionsProvider::new( @@ -570,6 +576,7 @@ impl MessageProcessor { ); Self { + turn_admission, user_verification, outgoing, models_refresh_worker, @@ -939,6 +946,24 @@ impl MessageProcessor { &codex_request, ); + let (turn_admission, recheck_turn_admission) = match &codex_request { + ClientRequest::ThreadStart { .. } + | ClientRequest::ThreadFork { .. } + | ClientRequest::ThreadResume { .. } + | ClientRequest::ThreadRollback { .. } + | ClientRequest::ThreadRevert { .. } => (Some(self.turn_admission.admit()?), false), + ClientRequest::TurnStart { .. } + | ClientRequest::TurnSteer { .. } + | ClientRequest::ReviewStart { .. } + | ClientRequest::ThreadCompactStart { .. } + | ClientRequest::ThreadShellCommand { .. } + | ClientRequest::ThreadQueueStart { .. } + | ClientRequest::ThreadRealtimeStart { .. } => { + (Some(self.turn_admission.admit()?), true) + } + _ => (None, false), + }; + let event_stream_ready = match &codex_request { ClientRequest::McpServerEventStreamStart { params, .. } => Some( session @@ -956,6 +981,13 @@ impl MessageProcessor { let request = QueuedInitializedRequest::new( rpc_gate, async move { + let _turn_admission = turn_admission; + // Runtime changes already admitted before drain finish normally. Turn work + // still waiting in serialization must observe the newly closed gate. + if recheck_turn_admission && let Err(error) = processor.turn_admission.admit() { + processor.outgoing.send_error(error_request_id, error).await; + return; + } let processor_for_request = Arc::clone(&processor); let result = processor_for_request .handle_initialized_client_request( diff --git a/codex-rs/app-server/src/request_serialization.rs b/codex-rs/app-server/src/request_serialization.rs index a55b737460..390896a6cc 100644 --- a/codex-rs/app-server/src/request_serialization.rs +++ b/codex-rs/app-server/src/request_serialization.rs @@ -491,7 +491,9 @@ mod tests { #[tokio::test] async fn shutdown_of_live_gate_skips_already_queued_requests() { let queues = RequestSerializationQueues::default(); - let (queued_tx, mut queued_rx) = oneshot::channel::<()>(); + let admission = crate::turn_admission::TurnAdmission::default(); + let active = admission.subscribe_active(); + let permit = admission.admit().expect("admit queued request"); let key = RequestSerializationQueueKey::Global("test"); let live_gate = gate(); let (tx, mut rx) = mpsc::unbounded_channel(); @@ -518,7 +520,7 @@ mod tests { key, RequestSerializationAccess::Exclusive, QueuedInitializedRequest::new(live_gate.clone(), async move { - let _queued_tx = queued_tx; + let _permit = permit; tx.send(SECOND_REQUEST_VALUE) .expect("receiver should be open"); }), @@ -534,16 +536,11 @@ mod tests { Some(FIRST_REQUEST_VALUE) ); - assert_eq!( - queued_rx.try_recv(), - Err(oneshot::error::TryRecvError::Empty) - ); + admission.begin_drain(); + assert_eq!(*active.borrow(), 1); live_gate.close().await; queues.discard_closed().await; - assert_eq!( - queued_rx.try_recv(), - Err(oneshot::error::TryRecvError::Closed) - ); + assert_eq!(*active.borrow(), 0); let gate_for_shutdown = Arc::clone(&live_gate); let shutdown_task = tokio::spawn(async move { diff --git a/codex-rs/app-server/src/turn_admission.rs b/codex-rs/app-server/src/turn_admission.rs new file mode 100644 index 0000000000..6c40044596 --- /dev/null +++ b/codex-rs/app-server/src/turn_admission.rs @@ -0,0 +1,87 @@ +//! Tracks client runtime changes and turn starts admitted before a shutdown drain. + +use codex_app_server_protocol::JSONRPCErrorError; +use codex_extension_api::TurnStartAdmission; +use std::sync::Arc; +use std::sync::Mutex; +use tokio::sync::watch; + +use crate::error_code::server_draining_error; + +#[derive(Debug, Default)] +struct AdmissionState { + closed: bool, + active: usize, +} + +#[derive(Clone, Debug)] +pub(crate) struct TurnAdmission { + state: Arc>, + active_tx: watch::Sender, +} + +impl Default for TurnAdmission { + fn default() -> Self { + Self { + state: Arc::new(Mutex::new(AdmissionState::default())), + active_tx: watch::channel(0).0, + } + } +} + +pub(crate) struct TurnPermit(TurnAdmission); + +impl TurnAdmission { + pub(crate) fn begin_drain(&self) { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .closed = true; + } + + pub(crate) fn subscribe_active(&self) -> watch::Receiver { + self.active_tx.subscribe() + } + + // Admit and close take the same short lock. The permit keeps shutdown from + // finishing while an earlier request is still preparing or submitting work. + pub(crate) fn admit(&self) -> Result { + self.try_admit().ok_or_else(server_draining_error) + } + + fn try_admit(&self) -> Option { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.closed { + return None; + } + state.active += 1; + self.active_tx.send_replace(state.active); + Some(TurnPermit(self.clone())) + } +} + +impl TurnStartAdmission for TurnAdmission { + fn admit_turn_start(&self) -> Option> { + self.try_admit() + .map(|permit| Box::new(permit) as Box) + } +} + +impl Drop for TurnPermit { + fn drop(&mut self) { + let mut state = self + .0 + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.active -= 1; + self.0.active_tx.send_replace(state.active); + } +} + +#[cfg(test)] +#[path = "turn_admission_tests.rs"] +mod tests; diff --git a/codex-rs/app-server/src/turn_admission_tests.rs b/codex-rs/app-server/src/turn_admission_tests.rs new file mode 100644 index 0000000000..b5cf7bf41f --- /dev/null +++ b/codex-rs/app-server/src/turn_admission_tests.rs @@ -0,0 +1,25 @@ +use super::TurnAdmission; +use crate::error_code::INVALID_REQUEST_ERROR_CODE; +use codex_extension_api::TurnStartAdmission; +use pretty_assertions::assert_eq; + +#[test] +fn drain_rejects_new_work_without_waiting_for_admitted_work() { + let admission = TurnAdmission::default(); + let active = admission.subscribe_active(); + let in_flight = admission.admit().expect("initial request admitted"); + let automatic = admission + .admit_turn_start() + .expect("automatic start admitted"); + admission.begin_drain(); + assert_eq!(*active.borrow(), 2); + assert!(admission.admit_turn_start().is_none()); + + let error = admission.admit().err().expect("new request rejected"); + assert_eq!(error.code, INVALID_REQUEST_ERROR_CODE); + assert_eq!(error.data, None); + drop(in_flight); + assert_eq!(*active.borrow(), 1); + drop(automatic); + assert_eq!(*active.borrow(), 0); +} diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs index c94656a068..095d3e3514 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs +++ b/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 { @@ -187,15 +380,7 @@ async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result Result Result 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,