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
This commit is contained in:
Eric Traut authored and copyberry committed 2026-09-09 01:03:42 +00:00
1 parent 4e09b0c1f1
commit c3eeaae9a3
8 files changed
+448 -33

No files matched your search

+6
View File
@@ -19,6 +19,7 @@ use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionRegistry; use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::ExtensionWarning; use codex_extension_api::ExtensionWarning;
use codex_extension_api::TurnStartAdmission;
use codex_goal_extension::GoalExtensionConfig; use codex_goal_extension::GoalExtensionConfig;
use codex_goal_extension::GoalService; use codex_goal_extension::GoalService;
use codex_http_client::HttpClientFactory; use codex_http_client::HttpClientFactory;
@@ -48,6 +49,7 @@ pub(crate) struct ThreadExtensionDependencies {
pub(crate) http_client_factory: HttpClientFactory, pub(crate) http_client_factory: HttpClientFactory,
/// Process-scoped queue shared by idle dispatch and app-server requests. /// Process-scoped queue shared by idle dispatch and app-server requests.
pub(crate) queue_service: Option<Arc<QueuedItemService>>, pub(crate) queue_service: Option<Arc<QueuedItemService>>,
pub(crate) turn_start_admission: Option<Arc<dyn TurnStartAdmission>>,
} }
pub(crate) fn thread_extensions<S>( pub(crate) fn thread_extensions<S>(
@@ -69,8 +71,12 @@ where
git_attribution_base_url, git_attribution_base_url,
http_client_factory, http_client_factory,
queue_service, queue_service,
turn_start_admission,
} = dependencies; } = dependencies;
let mut builder = ExtensionRegistryBuilder::<Config>::with_event_sink(Arc::clone(&event_sink)); let mut builder = ExtensionRegistryBuilder::<Config>::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 { if let Some(queue_service) = queue_service {
codex_queue_extension::install(&mut builder, queue_service); codex_queue_extension::install(&mut builder, queue_service);
} }
+61 -11
View File
@@ -132,6 +132,7 @@ mod skills_watcher;
mod thread_state; mod thread_state;
mod thread_status; mod thread_status;
mod transport; mod transport;
mod turn_admission;
mod turn_cost_worker; mod turn_cost_worker;
mod user_verification; mod user_verification;
mod user_verification_response; mod user_verification_response;
@@ -247,6 +248,7 @@ impl ShutdownState {
signal: ShutdownSignal, signal: ShutdownSignal,
connection_count: usize, connection_count: usize,
running_turn_count: usize, running_turn_count: usize,
turn_admission: &turn_admission::TurnAdmission,
) { ) {
if self.requested { if self.requested {
if matches!(signal, ShutdownSignal::Forceable) { if matches!(signal, ShutdownSignal::Forceable) {
@@ -255,20 +257,26 @@ impl ShutdownState {
return; return;
} }
turn_admission.begin_drain();
self.requested = true; self.requested = true;
self.last_logged_running_turn_count = None; self.last_logged_running_turn_count = None;
info!( 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, 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 { if !self.requested {
return ShutdownAction::Noop; return ShutdownAction::Noop;
} }
if self.forced || running_turn_count == 0 { if self.forced || (running_turn_count == 0 && active_admissions == 0) {
if self.forced { if self.forced {
info!( info!(
"received second shutdown signal; forcing restart with {running_turn_count} running assistant turn(s) and {connection_count} connection(s)" "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) { if self.last_logged_running_turn_count != Some(running_turn_count) {
info!( 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); 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 thread_created_rx = processor.thread_created_receiver();
let mut running_turn_count_rx = processor.subscribe_running_assistant_turn_count(); 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::<ConnectionId, ConnectionState>::new(); let mut connections = HashMap::<ConnectionId, ConnectionState>::new();
let mut connection_cleanup_tasks = ConnectionCleanupTasks::new(); let mut connection_cleanup_tasks = ConnectionCleanupTasks::new();
let mut remote_control_status_rx = remote_control_handle.status_receiver(); 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 listen_for_threads = true;
let mut shutdown_state = ShutdownState::default(); let mut shutdown_state = ShutdownState::default();
let exit_reason = loop { let exit_reason = loop {
let running_turn_count = { // Sample submissions first: one can publish a running turn before
let running_turn_count = running_turn_count_rx.borrow(); // releasing its permit, and shutdown must observe that new turn.
*running_turn_count let active_admissions = *active_admissions_rx.borrow_and_update();
}; let running_turn_count = *running_turn_count_rx.borrow_and_update();
if matches!( if matches!(
shutdown_state.update(running_turn_count, connections.len()), shutdown_state.update(running_turn_count, active_admissions, connections.len()),
ShutdownAction::Finish ShutdownAction::Finish
) { ) {
transport_shutdown_token.cancel(); 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(); 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() => { changed = running_turn_count_rx.changed(), if shutdown_state.requested() => {
if changed.is_err() { if changed.is_err() {
warn!("running-turn watcher closed during graceful restart drain"); 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() => { event = transport_event_rx.recv() => {
let Some(event) = event else { let Some(event) = event else {
break "transport_channel_closed"; break "transport_channel_closed";
}; };
match event { match event {
TransportEvent::DaemonShutdown => { 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 { TransportEvent::ConnectionOpened {
connection_id, connection_id,
@@ -1414,14 +1428,50 @@ fn analytics_rpc_transport(transport: &AppServerTransport) -> AppServerRpcTransp
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::LogFormat; use super::LogFormat;
use super::ShutdownAction;
use super::ShutdownSignal;
use super::ShutdownState;
#[cfg(debug_assertions)] #[cfg(debug_assertions)]
use super::loader_overrides_with_test_user_config_file; use super::loader_overrides_with_test_user_config_file;
use super::turn_admission::TurnAdmission;
#[cfg(debug_assertions)] #[cfg(debug_assertions)]
use codex_config::LoaderOverrides; use codex_config::LoaderOverrides;
#[cfg(debug_assertions)] #[cfg(debug_assertions)]
use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_absolute_path::AbsolutePathBuf;
use pretty_assertions::assert_eq; 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] #[test]
fn log_format_from_env_value_matches_json_values_case_insensitively() { fn log_format_from_env_value_matches_json_values_case_insensitively() {
assert_eq!(LogFormat::from_env_value(Some("json")), LogFormat::Json); assert_eq!(LogFormat::from_env_value(Some("json")), LogFormat::Json);
+1
View File
@@ -340,6 +340,7 @@ enabled = false
git_attribution_base_url: good_config.chatgpt_base_url.clone(), git_attribution_base_url: good_config.chatgpt_base_url.clone(),
http_client_factory: good_config.http_client_factory(), http_client_factory: good_config.http_client_factory(),
queue_service: None, queue_service: None,
turn_start_admission: None,
}, },
), ),
Arc::new(CodexHomeUserInstructionsProvider::new( Arc::new(CodexHomeUserInstructionsProvider::new(
@@ -78,6 +78,7 @@ use codex_core::ThreadManager;
use codex_core::config::Config; use codex_core::config::Config;
use codex_core::config::ThreadStoreConfig; use codex_core::config::ThreadStoreConfig;
use codex_exec_server::EnvironmentManager; use codex_exec_server::EnvironmentManager;
use codex_extension_api::TurnStartAdmission;
use codex_feedback::CodexFeedback; use codex_feedback::CodexFeedback;
use codex_goal_extension::GoalService; use codex_goal_extension::GoalService;
use codex_home::CodexHomeUserInstructionsProvider; use codex_home::CodexHomeUserInstructionsProvider;
@@ -101,6 +102,7 @@ use tokio_util::sync::CancellationToken;
use tracing::Instrument; use tracing::Instrument;
use crate::models_refresh_worker::ModelsRefreshWorker; use crate::models_refresh_worker::ModelsRefreshWorker;
use crate::turn_admission::TurnAdmission;
const CONNECTION_RPC_DRAIN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 30); 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) struct MessageProcessor {
pub(crate) turn_admission: TurnAdmission,
user_verification: Arc<crate::user_verification::Service>, user_verification: Arc<crate::user_verification::Service>,
outgoing: Arc<OutgoingMessageSender>, outgoing: Arc<OutgoingMessageSender>,
models_refresh_worker: ModelsRefreshWorker, models_refresh_worker: ModelsRefreshWorker,
@@ -311,6 +314,8 @@ impl MessageProcessor {
), ),
); );
let goal_service = Arc::new(GoalService::new()); let goal_service = Arc::new(GoalService::new());
let turn_admission = TurnAdmission::default();
let turn_start_admission: Arc<dyn TurnStartAdmission> = Arc::new(turn_admission.clone());
let extension_event_sink = let extension_event_sink =
app_server_extension_event_sink(outgoing.clone(), thread_state_manager.clone()); app_server_extension_event_sink(outgoing.clone(), thread_state_manager.clone());
let mut queue_service = None; let mut queue_service = None;
@@ -343,6 +348,7 @@ impl MessageProcessor {
git_attribution_base_url: config.chatgpt_base_url.clone(), git_attribution_base_url: config.chatgpt_base_url.clone(),
http_client_factory: config.http_client_factory(), http_client_factory: config.http_client_factory(),
queue_service: queue_service.clone(), queue_service: queue_service.clone(),
turn_start_admission: Some(Arc::clone(&turn_start_admission)),
}, },
), ),
Arc::new(CodexHomeUserInstructionsProvider::new( Arc::new(CodexHomeUserInstructionsProvider::new(
@@ -570,6 +576,7 @@ impl MessageProcessor {
); );
Self { Self {
turn_admission,
user_verification, user_verification,
outgoing, outgoing,
models_refresh_worker, models_refresh_worker,
@@ -939,6 +946,24 @@ impl MessageProcessor {
&codex_request, &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 { let event_stream_ready = match &codex_request {
ClientRequest::McpServerEventStreamStart { params, .. } => Some( ClientRequest::McpServerEventStreamStart { params, .. } => Some(
session session
@@ -956,6 +981,13 @@ impl MessageProcessor {
let request = QueuedInitializedRequest::new( let request = QueuedInitializedRequest::new(
rpc_gate, rpc_gate,
async move { 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 processor_for_request = Arc::clone(&processor);
let result = processor_for_request let result = processor_for_request
.handle_initialized_client_request( .handle_initialized_client_request(
@@ -491,7 +491,9 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn shutdown_of_live_gate_skips_already_queued_requests() { async fn shutdown_of_live_gate_skips_already_queued_requests() {
let queues = RequestSerializationQueues::default(); 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 key = RequestSerializationQueueKey::Global("test");
let live_gate = gate(); let live_gate = gate();
let (tx, mut rx) = mpsc::unbounded_channel(); let (tx, mut rx) = mpsc::unbounded_channel();
@@ -518,7 +520,7 @@ mod tests {
key, key,
RequestSerializationAccess::Exclusive, RequestSerializationAccess::Exclusive,
QueuedInitializedRequest::new(live_gate.clone(), async move { QueuedInitializedRequest::new(live_gate.clone(), async move {
let _queued_tx = queued_tx; let _permit = permit;
tx.send(SECOND_REQUEST_VALUE) tx.send(SECOND_REQUEST_VALUE)
.expect("receiver should be open"); .expect("receiver should be open");
}), }),
@@ -534,16 +536,11 @@ mod tests {
Some(FIRST_REQUEST_VALUE) Some(FIRST_REQUEST_VALUE)
); );
assert_eq!( admission.begin_drain();
queued_rx.try_recv(), assert_eq!(*active.borrow(), 1);
Err(oneshot::error::TryRecvError::Empty)
);
live_gate.close().await; live_gate.close().await;
queues.discard_closed().await; queues.discard_closed().await;
assert_eq!( assert_eq!(*active.borrow(), 0);
queued_rx.try_recv(),
Err(oneshot::error::TryRecvError::Closed)
);
let gate_for_shutdown = Arc::clone(&live_gate); let gate_for_shutdown = Arc::clone(&live_gate);
let shutdown_task = tokio::spawn(async move { let shutdown_task = tokio::spawn(async move {
+87
View File
@@ -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<Mutex<AdmissionState>>,
active_tx: watch::Sender<usize>,
}
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<usize> {
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<TurnPermit, JSONRPCErrorError> {
self.try_admit().ok_or_else(server_draining_error)
}
fn try_admit(&self) -> Option<TurnPermit> {
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<Box<dyn Send>> {
self.try_admit()
.map(|permit| Box::new(permit) as Box<dyn Send>)
}
}
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;
@@ -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);
}
@@ -4,8 +4,10 @@ use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT;
use super::connection_handling_websocket::WsClient; use super::connection_handling_websocket::WsClient;
use super::connection_handling_websocket::connect_websocket; use super::connection_handling_websocket::connect_websocket;
use super::connection_handling_websocket::create_config_toml; 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::read_response_for_id;
use super::connection_handling_websocket::send_initialize_request;
use super::connection_handling_websocket::send_request; use super::connection_handling_websocket::send_request;
use super::connection_handling_websocket::spawn_websocket_server; use super::connection_handling_websocket::spawn_websocket_server;
use anyhow::Context; use anyhow::Context;
@@ -13,19 +15,28 @@ use anyhow::Result;
use anyhow::bail; use anyhow::bail;
use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_final_assistant_message_sse_response;
use app_test_support::to_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::RequestId;
use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::ThreadStartResponse;
use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartParams;
use codex_app_server_protocol::TurnStartResponse;
use codex_app_server_protocol::UserInput as V2UserInput; use codex_app_server_protocol::UserInput as V2UserInput;
use core_test_support::responses; 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::SinkExt;
use futures::StreamExt; use futures::StreamExt;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
use serde_json::json;
#[cfg(unix)] #[cfg(unix)]
use std::process::Command as StdCommand; use std::process::Command as StdCommand;
use tempfile::TempDir; use tempfile::TempDir;
use tokio::process::Child; use tokio::process::Child;
use tokio::sync::oneshot;
use tokio::time::Duration; use tokio::time::Duration;
use tokio::time::Instant; use tokio::time::Instant;
use tokio::time::sleep; use tokio::time::sleep;
@@ -43,11 +54,60 @@ async fn websocket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Resu
_server, _server,
mut process, mut process,
mut ws, 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)?; send_sigint(&process)?;
assert_process_does_not_exit_within(&mut process, Duration::from_millis(300)).await?; 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( let status = wait_for_process_exit_within(
&mut process, &mut process,
Duration::from_secs(10), Duration::from_secs(10),
@@ -69,6 +129,7 @@ async fn websocket_transport_second_ctrl_c_forces_exit_while_turn_running() -> R
_server, _server,
mut process, mut process,
mut ws, mut ws,
..
} = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
send_sigint(&process)?; send_sigint(&process)?;
@@ -96,6 +157,7 @@ async fn websocket_transport_sigterm_waits_for_running_turn_before_exit() -> Res
_server, _server,
mut process, mut process,
mut ws, mut ws,
..
} = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
send_sigterm(&process, _codex_home.path())?; send_sigterm(&process, _codex_home.path())?;
@@ -122,6 +184,7 @@ async fn websocket_transport_second_sigterm_forces_exit_while_turn_running() ->
_server, _server,
mut process, mut process,
mut ws, mut ws,
..
} = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
send_sigterm(&process, _codex_home.path())?; send_sigterm(&process, _codex_home.path())?;
@@ -149,6 +212,7 @@ async fn websocket_transport_repeated_sighup_keeps_waiting_for_running_turn() ->
_server, _server,
mut process, mut process,
mut ws, mut ws,
..
} = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?;
send_sighup(&process)?; send_sighup(&process)?;
@@ -170,11 +234,140 @@ async fn websocket_transport_repeated_sighup_keeps_waiting_for_running_turn() ->
Ok(()) 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 { struct GracefulCtrlCFixture {
_codex_home: TempDir, _codex_home: TempDir,
_server: wiremock::MockServer, _server: wiremock::MockServer,
process: Child, process: Child,
ws: WsClient, ws: WsClient,
thread_id: String,
turn_id: String,
} }
async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result<GracefulCtrlCFixture> { 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) .mount(&server)
.await; .await;
let codex_home = TempDir::new()?; let (codex_home, process, mut ws) = start_ctrl_c_restart_client(&server.uri()).await?;
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));
send_thread_start_request(&mut ws, /*id*/ 2).await?; send_thread_start_request(&mut ws, /*id*/ 2).await?;
let thread_start_response = read_response_for_id(&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?; send_turn_start_request(&mut ws, /*id*/ 3, &thread.id).await?;
let turn_start_response = read_response_for_id(&mut ws, /*id*/ 3).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?; 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, _server: server,
process, process,
ws, 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<()> { async fn send_thread_start_request(stream: &mut WsClient, id: i64) -> Result<()> {
send_request( send_request(
stream, stream,