app-server: reserve running resumes before idle unload

codex-rs/app-server/src/request_processors.rs

@@ -545,6 +545,8 @@ use crate::error_code::invalid_request;
 use crate::filters::compute_source_filters;
 use crate::filters::source_kind_matches;
 use crate::thread_state::ConnectionCapabilities;
+use crate::thread_state::ConnectionReservationAction;
+use crate::thread_state::ConnectionReservationOutcome;
 use crate::thread_state::ThreadListenerCommand;
 use crate::thread_state::ThreadListenerRegistration;
 use crate::thread_state::ThreadState;

codex-rs/app-server/src/request_processors/thread_lifecycle.rs

@@ -120,6 +120,51 @@ impl UnloadingState {
     }
 }
 
+#[expect(
+    clippy::await_holding_invalid_type,
+    reason = "idle unload admission holds both lifecycle guards through its final rechecks"
+)]
+async fn claim_idle_unload_if_permitted(
+    unloading_state: &mut UnloadingState,
+    thread_list_state_permit: Arc<Semaphore>,
+    pending_thread_unloads: &Arc<Mutex<HashSet<ThreadId>>>,
+    thread_state_manager: &ThreadStateManager,
+    conversation_id: ThreadId,
+    is_running: impl std::future::Future<Output = bool>,
+) -> bool {
+    let Ok(permit) = thread_list_state_permit.acquire_owned().await else {
+        return false;
+    };
+    let mut pending = pending_thread_unloads.lock().await;
+    if pending.contains(&conversation_id) {
+        return false;
+    }
+    if !unloading_state.should_unload_now() {
+        return false;
+    }
+    if is_running.await {
+        unloading_state.note_thread_activity_observed();
+        return false;
+    }
+    if thread_state_manager
+        .has_connections_or_reservations(conversation_id)
+        .await
+    {
+        return false;
+    }
+    if !unloading_state.should_unload_now() {
+        return false;
+    }
+    pending.insert(conversation_id);
+    drop(pending);
+    drop(permit);
+    true
+}
+
+#[cfg(test)]
+#[path = "thread_lifecycle_tests.rs"]
+mod tests;
+
 pub(super) enum ThreadShutdownResult {
     Complete,
     SubmitFailed,
@@ -384,15 +429,19 @@ pub(super) async fn ensure_listener_task_running(
                         unloading_state.note_thread_activity_observed();
                         continue;
                     }
+                    if !claim_idle_unload_if_permitted(
+                        &mut unloading_state,
+                        thread_list_state_permit.clone(),
+                        &pending_thread_unloads,
+                        &thread_state_manager,
+                        conversation_id,
+                        async {
+                            matches!(conversation.agent_status().await, AgentStatus::Running)
+                        },
+                    )
+                    .await
                     {
-                        let mut pending_thread_unloads = pending_thread_unloads.lock().await;
-                        if pending_thread_unloads.contains(&conversation_id) {
-                            continue;
-                        }
-                        if !unloading_state.should_unload_now() {
-                            continue;
-                        }
-                        pending_thread_unloads.insert(conversation_id);
+                        continue;
                     }
                     unload_thread_without_subscribers(
                         thread_manager.clone(),
@@ -519,7 +568,6 @@ pub(super) async fn handle_thread_listener_command(
                 conversation_id,
                 conversation,
                 codex_home,
-                thread_state_manager,
                 thread_state,
                 thread_watch_manager,
                 outgoing,
@@ -568,15 +616,10 @@ pub(super) async fn handle_thread_listener_command(
 }
 
 #[allow(clippy::too_many_arguments)]
-#[expect(
-    clippy::await_holding_invalid_type,
-    reason = "running-thread resume subscription must be serialized against pending unloads"
-)]
 pub(super) async fn handle_pending_thread_resume_request(
     conversation_id: ThreadId,
     conversation: &Arc<CodexThread>,
     _codex_home: &Path,
-    thread_state_manager: &ThreadStateManager,
     thread_state: &Arc<Mutex<ThreadState>>,
     thread_watch_manager: &ThreadWatchManager,
     outgoing: &Arc<OutgoingMessageSender>,
@@ -602,6 +645,7 @@ pub(super) async fn handle_pending_thread_resume_request(
                 .is_some_and(|turn| matches!(turn.status, TurnStatus::InProgress));
 
     let request_id = pending.request_id;
+    let resume_handled_tx = pending.resume_handled_tx;
     let connection_id = request_id.connection_id;
     let mut thread = pending.thread_summary;
     if pending.include_turns {
@@ -633,6 +677,7 @@ pub(super) async fn handle_pending_thread_resume_request(
             Ok(page) => Some(page),
             Err(error) => {
                 outgoing.send_error(request_id, error).await;
+                let _ = resume_handled_tx.send(ConnectionReservationAction::Cancel);
                 return;
             }
         }
@@ -646,31 +691,21 @@ pub(super) async fn handle_pending_thread_resume_request(
         }
     }
 
+    if pending_thread_unloads
+        .lock()
+        .await
+        .contains(&conversation_id)
     {
-        let pending_thread_unloads = pending_thread_unloads.lock().await;
-        if pending_thread_unloads.contains(&conversation_id) {
-            drop(pending_thread_unloads);
-            outgoing
-                .send_error(
-                    request_id,
-                    invalid_request(format!(
-                        "thread {conversation_id} is closing; retry thread/resume after the thread is closed"
-                    )),
-                )
-                .await;
-            return;
-        }
-        if !thread_state_manager
-            .try_add_connection_to_thread(conversation_id, connection_id)
-            .await
-        {
-            tracing::debug!(
-                thread_id = %conversation_id,
-                connection_id = ?connection_id,
-                "skipping running thread resume for closed connection"
-            );
-            return;
-        }
+        outgoing
+            .send_error(
+                request_id,
+                invalid_request(format!(
+                    "thread {conversation_id} is closing; retry thread/resume after the thread is closed"
+                )),
+            )
+            .await;
+        let _ = resume_handled_tx.send(ConnectionReservationAction::Cancel);
+        return;
     }
 
     let config_snapshot = pending.config_snapshot;
@@ -714,6 +749,7 @@ pub(super) async fn handle_pending_thread_resume_request(
     outgoing
         .send_response_with_thread_originator(request_id, response, originator)
         .await;
+    let _ = resume_handled_tx.send(ConnectionReservationAction::Commit);
     // Match cold resume: metadata-only resume should attach the listener without
     // paying the cost of turn reconstruction for historical usage replay.
     if let Some(token_usage_thread) = token_usage_thread {

codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs

@@ -0,0 +1,135 @@
+use super::*;
+use pretty_assertions::assert_eq;
+
+#[tokio::test]
+async fn resume_reservation_prevents_idle_unload_before_listener_commit() {
+    let now = Instant::now();
+    let (_has_subscribers_tx, has_subscribers_rx) = watch::channel(false);
+    let (_thread_status_tx, thread_status_rx) = watch::channel(ThreadStatus::Idle);
+    let mut unloading_state = UnloadingState {
+        delay: Duration::ZERO,
+        has_subscribers_rx,
+        has_subscribers: (false, now),
+        thread_status_rx,
+        is_active: (false, now),
+    };
+    let pending = Arc::new(Mutex::new(HashSet::new()));
+    let thread_state_manager = ThreadStateManager::new();
+    let thread_id = ThreadId::new();
+    let connection_id = ConnectionId(1);
+    thread_state_manager
+        .connection_initialized(connection_id, ConnectionCapabilities::default())
+        .await;
+    let permit = Arc::new(Semaphore::new(1));
+    let resume_permit = permit
+        .clone()
+        .acquire_owned()
+        .await
+        .expect("permit should be open");
+    let (resume_handled_tx, resume_outcome_rx) = thread_state_manager
+        .start_connection_reservation(thread_id, connection_id)
+        .await
+        .expect("live connection should reserve the thread");
+    drop(resume_permit);
+    assert!(
+        !claim_idle_unload_if_permitted(
+            &mut unloading_state,
+            permit,
+            &pending,
+            &thread_state_manager,
+            thread_id,
+            std::future::ready(false),
+        )
+        .await
+    );
+    assert!(pending.lock().await.is_empty());
+
+    assert!(
+        thread_state_manager
+            .unsubscribe_connection_from_thread(thread_id, connection_id)
+            .await
+    );
+    let (replacement_tx, replacement_outcome_rx) = thread_state_manager
+        .start_connection_reservation(thread_id, connection_id)
+        .await
+        .expect("replacement reservation should succeed");
+    resume_handled_tx
+        .send(ConnectionReservationAction::Commit)
+        .expect("reservation owner should still be running");
+    assert_eq!(
+        resume_outcome_rx
+            .await
+            .expect("reservation owner should report an outcome"),
+        ConnectionReservationOutcome::Handled,
+    );
+    assert!(
+        thread_state_manager
+            .subscribed_connection_ids(thread_id)
+            .await
+            .is_empty()
+    );
+    assert!(
+        thread_state_manager
+            .has_connections_or_reservations(thread_id)
+            .await
+    );
+
+    drop(replacement_tx);
+    assert_eq!(
+        replacement_outcome_rx
+            .await
+            .expect("reservation owner should report abandonment"),
+        ConnectionReservationOutcome::Abandoned,
+    );
+    assert!(
+        !thread_state_manager
+            .has_connections_or_reservations(thread_id)
+            .await
+    );
+    assert!(
+        thread_state_manager
+            .subscribed_connection_ids(thread_id)
+            .await
+            .is_empty()
+    );
+    thread_state_manager.clear_all_listeners().await;
+}
+
+#[tokio::test]
+async fn pending_resume_prevents_disconnect_from_reporting_thread_stale() {
+    let thread_state_manager = ThreadStateManager::new();
+    let thread_id = ThreadId::new();
+    let committed_connection_id = ConnectionId(1);
+    let pending_connection_id = ConnectionId(2);
+    for connection_id in [committed_connection_id, pending_connection_id] {
+        thread_state_manager
+            .connection_initialized(connection_id, ConnectionCapabilities::default())
+            .await;
+    }
+    assert!(
+        thread_state_manager
+            .try_add_connection_to_thread(thread_id, committed_connection_id)
+            .await
+    );
+    let (reservation_tx, reservation_outcome_rx) = thread_state_manager
+        .start_connection_reservation(thread_id, pending_connection_id)
+        .await
+        .expect("live connection should reserve the thread");
+
+    assert!(
+        thread_state_manager
+            .remove_connection(committed_connection_id)
+            .await
+            .is_empty(),
+        "a pending reservation must keep the thread from being reported stale"
+    );
+
+    drop(reservation_tx);
+    assert_eq!(
+        reservation_outcome_rx
+            .await
+            .expect("reservation owner should report abandonment"),
+        ConnectionReservationOutcome::Abandoned,
+    );
+    thread_state_manager.clear_all_listeners().await;
+}

codex-rs/app-server/src/request_processors/thread_processor.rs

@@ -371,6 +371,12 @@ pub(crate) struct ThreadRequestProcessor {
 enum RunningThreadResumeResult {
     /// The request was delegated to the loaded thread.
     Handled,
+    /// The listener owns the response after the caller releases the global
+    /// thread-list permit.
+    ListenerPending {
+        thread_id: ThreadId,
+        outcome_rx: oneshot::Receiver<ConnectionReservationOutcome>,
+    },
     /// No loaded thread handled the request.
     ///
     /// The optional stored thread contains the history-bearing probe that cold
@@ -2672,7 +2678,7 @@ impl ThreadRequestProcessor {
         let redact_resume_payloads =
             should_redact_thread_resume_payloads(app_server_client_name.as_deref());
 
-        let _thread_list_state_permit = match self.acquire_thread_list_state_permit().await {
+        let thread_list_state_permit = match self.acquire_thread_list_state_permit().await {
             Ok(permit) => permit,
             Err(error) => {
                 self.outgoing.send_error(request_id, error).await;
@@ -2689,6 +2695,23 @@ impl ThreadRequestProcessor {
             .await
         {
             Ok(RunningThreadResumeResult::Handled) => return Ok(()),
+            Ok(RunningThreadResumeResult::ListenerPending {
+                thread_id,
+                outcome_rx,
+            }) => {
+                drop(thread_list_state_permit);
+                if !matches!(outcome_rx.await, Ok(ConnectionReservationOutcome::Handled)) {
+                    self.outgoing
+                        .send_error(
+                            request_id,
+                            internal_error(format!(
+                                "failed to subscribe running thread resume for thread {thread_id}: thread listener closed before handling the request"
+                            )),
+                        )
+                        .await;
+                }
+                return Ok(());
+            }
             Ok(RunningThreadResumeResult::NotRunning(stored_thread)) => stored_thread,
             Err(error) => {
                 self.outgoing.send_error(request_id, error).await;
@@ -3131,10 +3154,18 @@ impl ThreadRequestProcessor {
                 .thread_goal_processor
                 .pending_resume_goal_state(existing_thread.as_ref())
                 .await;
+            let Some((resume_handled_tx, outcome_rx)) = self
+                .thread_state_manager
+                .start_connection_reservation(existing_thread_id, request_id.connection_id)
+                .await
+            else {
+                return Ok(RunningThreadResumeResult::Handled);
+            };
 
             let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse(
                 Box::new(crate::thread_state::PendingThreadResumeRequest {
                     request_id: request_id.clone(),
+                    resume_handled_tx,
                     history_items,
                     config_snapshot,
                     instruction_sources,
@@ -3151,7 +3182,10 @@ impl ThreadRequestProcessor {
                     "failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed"
                 )));
             }
-            return Ok(RunningThreadResumeResult::Handled);
+            return Ok(RunningThreadResumeResult::ListenerPending {
+                thread_id: existing_thread_id,
+                outcome_rx,
+            });
         }
         Ok(RunningThreadResumeResult::NotRunning(None))
     }