From 102e1763b90e70d79a0bba5d262cf71f904583c1 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Wed, 9 Sep 2026 00:11:14 +0000 Subject: [PATCH] Keep app-server thread RPCs active until delegated work completes (#43950) ## Why Thread startup, running-thread resume, and rollback could outlive their request handlers, allowing connection draining and request serialization to finish too early. ## What changed - Await thread startup and listener completion for running-thread resume and rollback. - Release the thread-list permit before waiting for resume completion so the listener can finish rollback responses. - Release pending rollback waiters when clearing a listener, and discard queued requests when their connection closes. ## Testing Add coverage for releasing queued requests and rollback waiters, and pipeline resume during rollback to check that both responses complete. Update diagnostics expectations to allow startup to remain in flight after its response is queued. GitOrigin-RevId: 138156179afbc40de29965820db7fa9cd3669cc7 --- .../app-server/src/bespoke_event_handling.rs | 4 +- codex-rs/app-server/src/message_processor.rs | 1 + .../request_processors/thread_lifecycle.rs | 6 +- .../request_processors/thread_processor.rs | 57 ++++++++++++------- .../app-server/src/request_serialization.rs | 31 ++++++++++ codex-rs/app-server/src/thread_state.rs | 30 +++++++++- .../tests/suite/v2/server_diagnostics.rs | 33 +++++------ .../tests/suite/v2/thread_rollback.rs | 14 ++++- 8 files changed, 134 insertions(+), 42 deletions(-) diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 2d3ff975f3..be146d82c7 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -1245,7 +1245,7 @@ pub(crate) async fn apply_bespoke_event_handling( state.pending_rollbacks.take() }; - if let Some(request_id) = pending { + if let Some((request_id, _completion_tx)) = pending { let _thread_list_state_permit = match thread_list_state_permit.acquire().await { Ok(permit) => permit, Err(err) => { @@ -1643,7 +1643,7 @@ async fn handle_thread_rollback_failed( ) { let pending_rollback = thread_state.lock().await.pending_rollbacks.take(); - if let Some(request_id) = pending_rollback { + if let Some((request_id, _completion_tx)) = pending_rollback { outgoing .send_error(request_id, invalid_request(message)) .await; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index aff6d9ad6c..511ea54587 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -818,6 +818,7 @@ impl MessageProcessor { session_state: &ConnectionSessionState, ) { session_state.rpc_gate.close().await; + self.request_serialization_queues.discard_closed().await; self.outgoing .disconnect_user_verification_connection(connection_id) .await; diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 4758d01947..07e5145980 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -487,7 +487,10 @@ pub(super) async fn handle_thread_listener_command( listener_command: ThreadListenerCommand, ) { match listener_command { - ThreadListenerCommand::SendThreadResumeResponse(resume_request) => { + ThreadListenerCommand::SendThreadResumeResponse { + request: resume_request, + completion_tx, + } => { handle_pending_thread_resume_request( conversation_id, conversation, @@ -500,6 +503,7 @@ pub(super) async fn handle_thread_listener_command( *resume_request, ) .await; + let _ = completion_tx.send(()); } ThreadListenerCommand::EmitThreadGoalUpdated { turn_id, goal } => { outgoing diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index df4ffe73ab..c4cce3011d 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -460,7 +460,7 @@ pub(crate) struct ThreadRequestProcessor { /// Outcome of trying to satisfy a resume request from an already loaded thread. enum RunningThreadResumeResult { /// The request was delegated to the loaded thread. - Handled, + Handled(tokio::sync::oneshot::Receiver<()>), /// No loaded thread handled the request. /// /// The optional stored thread contains the history-bearing probe that cold @@ -551,17 +551,17 @@ impl ThreadRequestProcessor { client_mcp_extensions: ClientMcpExtensions, ) -> Result, JSONRPCErrorError> { let mut prepared_config = None; - while self - .thread_resume_inner( - request_id.clone(), - ¶ms, - app_server_client_name.clone(), - app_server_client_version.clone(), - client_mcp_extensions.clone(), - &mut prepared_config, - ) - .await? - .is_continue() + // Keep the resume future off the request handler's stack. + while Box::pin(self.thread_resume_inner( + request_id.clone(), + ¶ms, + app_server_client_name.clone(), + app_server_client_version.clone(), + client_mcp_extensions.clone(), + &mut prepared_config, + )) + .await? + .is_continue() {} Ok(None) } @@ -1267,7 +1267,9 @@ impl ThreadRequestProcessor { } }; self.background_tasks - .spawn(thread_start_task.instrument(request_context.span())); + .spawn(thread_start_task.instrument(request_context.span())) + .await + .map_err(|_| internal_error("thread startup task stopped before completing"))?; Ok(()) } @@ -2353,6 +2355,7 @@ impl ThreadRequestProcessor { } let request = request_id.clone(); + let (completion_tx, completion_rx) = tokio::sync::oneshot::channel(); let rollback_already_in_progress = { let thread_state = self.thread_state_manager.thread_state(thread_id).await; @@ -2360,7 +2363,7 @@ impl ThreadRequestProcessor { if thread_state.pending_rollbacks.is_some() { true } else { - thread_state.pending_rollbacks = Some(request.clone()); + thread_state.pending_rollbacks = Some((request, completion_tx)); false } }; @@ -2385,6 +2388,9 @@ impl ThreadRequestProcessor { return Err(internal_error(format!("failed to start rollback: {err}"))); } + // The listener drops the sender after queuing the response, including errors. + // Keep the RPC's drain admission alive until then, without holding thread state. + let _ = completion_rx.await; Ok(()) } @@ -3678,7 +3684,12 @@ impl ThreadRequestProcessor { ) .await { - Ok(RunningThreadResumeResult::Handled) => return Ok(ControlFlow::Break(())), + Ok(RunningThreadResumeResult::Handled(completion)) => { + // The listener may need this permit to finish the response. + drop(_thread_list_state_permit); + let _ = completion.await; + return Ok(ControlFlow::Break(())); + } Ok(RunningThreadResumeResult::NotRunning(stored_thread)) => stored_thread, Err(error) => { self.outgoing.send_error(request_id, error).await; @@ -3801,7 +3812,11 @@ impl ThreadRequestProcessor { ) .await? { - RunningThreadResumeResult::Handled => Ok(ControlFlow::Break(())), + RunningThreadResumeResult::Handled(completion) => { + drop(_thread_list_state_permit); + let _ = completion.await; + Ok(ControlFlow::Break(())) + } RunningThreadResumeResult::NotRunning(_) => Err(invalid_request( "cannot resume an unloaded multi-agent v2 sub-agent through its parent; resume the parent first, or use thread/read to inspect it", )), @@ -4496,8 +4511,9 @@ impl ThreadRequestProcessor { }; let resume_cursor_store = paginated_resume.then(|| Arc::clone(&self.thread_store)); - let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse( - Box::new(crate::thread_state::PendingThreadResumeRequest { + let (completion_tx, completion_rx) = tokio::sync::oneshot::channel(); + let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse { + request: Box::new(crate::thread_state::PendingThreadResumeRequest { request_id: request_id.clone(), history_items, cold_resume_token_usage_turn_id, @@ -4514,13 +4530,14 @@ impl ThreadRequestProcessor { resume_cursor_store, redact_resume_payloads, }), - ); + completion_tx, + }; if listener_command_tx.send(command).is_err() { return Err(internal_error(format!( "failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed" ))); } - return Ok(RunningThreadResumeResult::Handled); + return Ok(RunningThreadResumeResult::Handled(completion_rx)); } Ok(RunningThreadResumeResult::NotRunning(None)) } diff --git a/codex-rs/app-server/src/request_serialization.rs b/codex-rs/app-server/src/request_serialization.rs index ae897271ca..a55b737460 100644 --- a/codex-rs/app-server/src/request_serialization.rs +++ b/codex-rs/app-server/src/request_serialization.rs @@ -158,6 +158,24 @@ pub(crate) struct RequestSerializationQueues { } impl RequestSerializationQueues { + /// Release requests whose connection closed while they waited in a queue. + pub(crate) async fn discard_closed(&self) { + let mut queues = self.inner.lock().await; + for queue in queues.values_mut() { + let previous_len = queue.requests.len(); + queue.requests.retain(|request| { + !request + .request + .gate + .as_ref() + .is_some_and(|gate| gate.is_closed()) + }); + if queue.requests.len() != previous_len { + queue.changed.notify_one(); + } + } + } + /// Enqueue app-owned work alongside RPCs that mutate the same serialized resource. pub(crate) async fn enqueue_background( &self, @@ -473,6 +491,7 @@ 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 key = RequestSerializationQueueKey::Global("test"); let live_gate = gate(); let (tx, mut rx) = mpsc::unbounded_channel(); @@ -499,6 +518,7 @@ mod tests { key, RequestSerializationAccess::Exclusive, QueuedInitializedRequest::new(live_gate.clone(), async move { + let _queued_tx = queued_tx; tx.send(SECOND_REQUEST_VALUE) .expect("receiver should be open"); }), @@ -514,6 +534,17 @@ mod tests { Some(FIRST_REQUEST_VALUE) ); + assert_eq!( + queued_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + live_gate.close().await; + queues.discard_closed().await; + assert_eq!( + queued_rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + ); + let gate_for_shutdown = Arc::clone(&live_gate); let shutdown_task = tokio::spawn(async move { gate_for_shutdown.shutdown().await; diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 3795638cdd..ba011dd73a 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -57,7 +57,10 @@ pub(crate) struct PendingThreadResumeRequest { // ThreadListenerCommand is used to perform operations in the context of the thread listener, for serialization purposes. pub(crate) enum ThreadListenerCommand { // SendThreadResumeResponse is used to resume an already running thread by sending the thread's history to the client and atomically subscribing for new updates. - SendThreadResumeResponse(Box), + SendThreadResumeResponse { + request: Box, + completion_tx: oneshot::Sender<()>, + }, // EmitThreadGoalUpdated is used to order goal updates with running-thread resume responses and goal clears. EmitThreadGoalUpdated { turn_id: Option, @@ -95,7 +98,7 @@ pub(crate) struct TurnSummary { #[derive(Default)] pub(crate) struct ThreadState { pub(crate) pending_interrupts: PendingInterruptQueue, - pub(crate) pending_rollbacks: Option, + pub(crate) pending_rollbacks: Option<(ConnectionRequestId, oneshot::Sender<()>)>, pub(crate) turn_summary: TurnSummary, pub(crate) last_terminal_turn_id: Option, /// Lets an internal runtime replacement wait until the old listener has processed Core's @@ -143,6 +146,7 @@ impl ThreadState { let _ = cancel_tx.send(()); } self.shutdown_drain_waiter = None; + self.pending_rollbacks = None; self.listener_command_tx = None; self.current_turn_history.reset(); self.listener_thread = None; @@ -247,6 +251,28 @@ mod tests { use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; + #[test] + fn clear_listener_releases_pending_rollback() { + let (completion_tx, mut completion_rx) = oneshot::channel(); + let mut state = ThreadState { + pending_rollbacks: Some(( + ConnectionRequestId { + connection_id: ConnectionId(1), + request_id: RequestId::Integer(1), + }, + completion_tx, + )), + ..Default::default() + }; + + state.clear_listener(); + + assert_eq!( + completion_rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + ); + } + #[test] fn note_thread_settings_reports_only_effective_changes() { let mut state = ThreadState::default(); diff --git a/codex-rs/app-server/tests/suite/v2/server_diagnostics.rs b/codex-rs/app-server/tests/suite/v2/server_diagnostics.rs index 3fba0104f7..83432e883d 100644 --- a/codex-rs/app-server/tests/suite/v2/server_diagnostics.rs +++ b/codex-rs/app-server/tests/suite/v2/server_diagnostics.rs @@ -47,24 +47,25 @@ async fn server_diagnostics_exposes_process_and_registered_thread_gauge() -> Res assert!(diagnostics.process.physical_footprint_bytes.is_some()); #[cfg(not(target_os = "macos"))] assert_eq!(diagnostics.process.physical_footprint_bytes, None); - for expected_gauge in [ - ServerDiagnosticsGauge { - name: "app.requests.in_flight".to_string(), - value: 1, - }, - ServerDiagnosticsGauge { + // Startup may still be finishing after its response was queued. + assert!(matches!( + diagnostics + .gauges + .iter() + .find(|gauge| gauge.name == "app.requests.in_flight") + .map(|gauge| gauge.value), + Some(1 | 2) + )); + assert_eq!( + diagnostics + .gauges + .iter() + .find(|gauge| gauge.name == "core.threads.live"), + Some(&ServerDiagnosticsGauge { name: "core.threads.live".to_string(), value: 1, - }, - ] { - assert_eq!( - diagnostics - .gauges - .iter() - .find(|gauge| gauge.name == expected_gauge.name), - Some(&expected_gauge) - ); - } + }) + ); Ok(()) } diff --git a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs index 8bad7219dd..48c48d8ed1 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs @@ -230,6 +230,15 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<() num_turns: 1, }) .await?; + // Pipeline resume while the listener may still be handling rollback. Waiting for + // its response must not retain the thread-list permit needed by that listener. + let pipelined_resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: thread.id.clone(), + exclude_turns: true, + ..Default::default() + }) + .await?; let deprecation_notice = timeout(DEFAULT_READ_TIMEOUT, mcp.read_next_message()).await??; let JSONRPCMessage::Notification(deprecation_notice) = deprecation_notice else { panic!("thread/rollback should emit deprecationNotice before its response"); @@ -297,7 +306,10 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<() other => panic!("expected user message item, got {other:?}"), } - // Resume and confirm the history is pruned. + let _: ThreadResumeResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(pipelined_resume_id)).await??; + + // Resume after rollback completes to verify the pruned history. let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id: thread.id,