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

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;

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;

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

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<Option<ClientResponsePayload>, JSONRPCErrorError> {
         let mut prepared_config = None;
-        while self
+        // Keep the resume future off the request handler's stack.
-            .thread_resume_inner(
+        while Box::pin(self.thread_resume_inner(
-                request_id.clone(),
+            request_id.clone(),
-                &params,
+            &params,
-                app_server_client_name.clone(),
+            app_server_client_name.clone(),
-                app_server_client_version.clone(),
+            app_server_client_version.clone(),
-                client_mcp_extensions.clone(),
+            client_mcp_extensions.clone(),
-                &mut prepared_config,
+            &mut prepared_config,
-            )
+        ))
-            .await?
+        .await?
-            .is_continue()
+        .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(
+            let (completion_tx, completion_rx) = tokio::sync::oneshot::channel();
-                Box::new(crate::thread_state::PendingThreadResumeRequest {
+            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))
     }

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;

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<PendingThreadResumeRequest>),
+    SendThreadResumeResponse {
+        request: Box<PendingThreadResumeRequest>,
+        completion_tx: oneshot::Sender<()>,
+    },
     // EmitThreadGoalUpdated is used to order goal updates with running-thread resume responses and goal clears.
     EmitThreadGoalUpdated {
         turn_id: Option<String>,
@@ -95,7 +98,7 @@ pub(crate) struct TurnSummary {
 #[derive(Default)]
 pub(crate) struct ThreadState {
     pub(crate) pending_interrupts: PendingInterruptQueue,
-    pub(crate) pending_rollbacks: Option<ConnectionRequestId>,
+    pub(crate) pending_rollbacks: Option<(ConnectionRequestId, oneshot::Sender<()>)>,
     pub(crate) turn_summary: TurnSummary,
     pub(crate) last_terminal_turn_id: Option<String>,
     /// 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();

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 [
+    // Startup may still be finishing after its response was queued.
-        ServerDiagnosticsGauge {
+    assert!(matches!(
-            name: "app.requests.in_flight".to_string(),
+        diagnostics
-            value: 1,
+            .gauges
-        },
+            .iter()
-        ServerDiagnosticsGauge {
+            .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(())
 }

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,