Preserve TUI event ordering during active-thread draining (#39991)

## Why

Draining every queued active-thread event in one foreground pass can overrun the
frame budget. Buffered thread events can also race ahead of queued replay and
startup app events, causing the TUI to switch views before the corresponding
history and operations have been applied.

## What changed

- Bound active-thread draining to one target frame interval and leave remaining
  notifications queued for a later frame.
- Prioritize queued app events over active-thread events so replay state is
  applied before buffered closures or other thread notifications.
- Block startup input only for pending requests, rather than ordinary queued
  notifications whose draining yielded at the frame deadline.
- Leave selected side-thread closure handling to the foreground event loop.

## Testing

Added tests covering deadline-limited notification draining and foreground
handling of a selected side thread closing.

GitOrigin-RevId: 41080d3df13ca308eec1ea0fbb0ba067a4640c4c

codex-rs/tui/src/app/session_lifecycle.rs

@@ -561,7 +561,6 @@ impl App {
             };
             self.chat_widget.add_info_message(message, /*hint*/ None);
         }
-        self.drain_active_thread_events(tui).await?;
         self.refresh_pending_thread_approvals().await;
 
         Ok(())

codex-rs/tui/src/app/startup.rs

@@ -49,6 +49,16 @@ impl App {
                 .active_thread_rx
                 .as_ref()
                 .is_some_and(|receiver| !receiver.is_empty())
+                && self
+                    .active_thread_id
+                    .and_then(|thread_id| self.thread_event_channels.get(&thread_id))
+                    .is_none_or(|channel| {
+                        // A bounded drain can leave ordinary notifications queued. Only protect
+                        // input for pending requests, or when their state cannot be inspected.
+                        channel.store.try_lock().map_or(/*default*/ true, |store| {
+                            store.side_parent_pending_status().is_some()
+                        })
+                    })
                 || self
                     .pending_primary_events
                     .iter()
@@ -651,12 +661,15 @@ See the Codex keymap documentation for supported actions and examples."
             Ok(exit_reason)
         } else {
             loop {
+                // Replay queues history and operations. A buffered closure must not switch
+                // widgets before those app events have been applied.
+                let has_pending_app_events = !app_event_rx.is_empty();
                 let initial_session_header_pending = waiting_for_initial_session_header
                     && app.primary_session_configured.is_some()
-                    && !app_event_rx.is_empty();
+                    && has_pending_app_events;
                 let block_terminal_input_for_pending_startup_events = initial_session_header_pending
                     || (pending_startup_draft.is_some() || app.startup_protected_input_boundary)
-                        && !app_event_rx.is_empty()
+                        && has_pending_app_events
                     || (!waiting_for_initial_session_configured
                         && app.has_queued_startup_protected_request());
                 let control = select! {
@@ -693,7 +706,7 @@ See the Codex keymap documentation for supported actions and examples."
                     }, if App::should_handle_active_thread_events(
                         waiting_for_initial_session_configured,
                         app.active_thread_rx.is_some()
-                    ) => {
+                    ) && !has_pending_app_events => {
                         if let Some(event) = active {
                             if let Err(err) = app.handle_active_thread_event(tui, &mut app_server, event).await {
                                 break Err(err);

codex-rs/tui/src/app/tests.rs

@@ -755,6 +755,117 @@ async fn enqueue_thread_event_does_not_block_when_channel_full() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn active_thread_drain_yields_after_frame_deadline_without_dropping_events() -> Result<()> {
+    let mut app = make_test_app().await;
+    app.startup_protected_input_boundary = true;
+    let thread_id = ThreadId::new();
+    app.thread_event_channels
+        .insert(thread_id, ThreadEventChannel::new(/*capacity*/ 2));
+    app.activate_thread_channel(thread_id).await;
+
+    let first_event = token_usage_notification(thread_id, "turn-1", Some(100));
+    let mut second_event = token_usage_notification(thread_id, "turn-1", Some(100));
+    if let ServerNotification::ThreadTokenUsageUpdated(notification) = &mut second_event {
+        notification.token_usage.total.output_tokens = 10;
+        notification.token_usage.total.total_tokens = 15;
+    }
+    for event in [first_event, second_event] {
+        app.enqueue_thread_notification(thread_id, event).await?;
+    }
+
+    let mut tui = crate::tui::test_support::make_test_tui()?;
+    app.drain_active_thread_events_until(&mut tui, Instant::now())
+        .await?;
+    assert_eq!(
+        app.chat_widget.token_usage(),
+        crate::token_usage::TokenUsage {
+            input_tokens: 4,
+            cached_input_tokens: 1,
+            output_tokens: 5,
+            reasoning_output_tokens: 0,
+            total_tokens: 10,
+        }
+    );
+    assert!(
+        app.active_thread_rx
+            .as_ref()
+            .is_some_and(|receiver| !receiver.is_empty()),
+        "an expired frame deadline should preserve queued notifications"
+    );
+    assert!(
+        !app.has_queued_startup_protected_request(),
+        "ordinary notifications left by the frame deadline must not block terminal input"
+    );
+
+    app.drain_active_thread_events(&mut tui).await?;
+    assert!(
+        app.active_thread_rx
+            .as_ref()
+            .is_some_and(tokio::sync::mpsc::Receiver::is_empty),
+        "the next foreground frame should deliver the preserved notification"
+    );
+    assert_eq!(
+        app.chat_widget.token_usage(),
+        crate::token_usage::TokenUsage {
+            input_tokens: 4,
+            cached_input_tokens: 1,
+            output_tokens: 10,
+            reasoning_output_tokens: 0,
+            total_tokens: 15,
+        }
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn selected_side_thread_close_is_handled_by_foreground_event_owner() -> Result<()> {
+    let mut app = make_test_app().await;
+    let mut app_server = crate::start_embedded_app_server_for_picker(&app.config).await?;
+    let primary = app_server.start_thread(&app.config).await?;
+    let primary_thread_id = primary.session.thread_id;
+    app.enqueue_primary_thread_session(primary.session, primary.turns)
+        .await?;
+
+    let side_thread_id = ThreadId::new();
+    let side_channel = ThreadEventChannel::new(/*capacity*/ 1);
+    side_channel.store.lock().await.set_session(
+        test_thread_session(side_thread_id, app.config.cwd.to_path_buf()),
+        Vec::new(),
+    );
+    side_channel
+        .sender
+        .try_send(ThreadBufferedEvent::Notification(Box::new(
+            thread_closed_notification(side_thread_id),
+        )))
+        .expect("closed side-thread notification should fit in its saved receiver");
+    app.thread_event_channels
+        .insert(side_thread_id, side_channel);
+    app.side_threads
+        .insert(side_thread_id, SideThreadState::new(primary_thread_id));
+
+    let mut tui = crate::tui::test_support::make_test_tui()?;
+    app.select_agent_thread(&mut tui, &mut app_server, side_thread_id)
+        .await?;
+    assert_eq!(app.active_thread_id, Some(side_thread_id));
+    let event = app
+        .active_thread_rx
+        .as_mut()
+        .expect("selected side thread should retain its receiver")
+        .try_recv()
+        .expect("thread closure should remain queued for the foreground loop");
+    app.handle_active_thread_event(&mut tui, &mut app_server, event)
+        .await?;
+
+    assert_eq!(app.active_thread_id, Some(primary_thread_id));
+    assert!(!app.side_threads.contains_key(&side_thread_id));
+    assert!(!app.thread_event_channels.contains_key(&side_thread_id));
+    assert!(app.active_thread_rx.is_some());
+    app_server.shutdown().await?;
+    Ok(())
+}
+
 #[tokio::test]
 async fn active_history_batch_is_delivered_without_replay_buffering() -> Result<()> {
     let mut app = make_test_app().await;

codex-rs/tui/src/app/thread_routing.rs

@@ -1448,6 +1448,16 @@ impl App {
     /// historical id now" and converted into closed picker entries instead of deleting them, so
     /// the stable traversal order remains intact for review and keyboard navigation.
     pub(super) async fn drain_active_thread_events(&mut self, tui: &mut tui::Tui) -> Result<()> {
+        let frame_deadline = Instant::now() + tui::TARGET_FRAME_INTERVAL;
+        self.drain_active_thread_events_until(tui, frame_deadline)
+            .await
+    }
+
+    pub(super) async fn drain_active_thread_events_until(
+        &mut self,
+        tui: &mut tui::Tui,
+        frame_deadline: Instant,
+    ) -> Result<()> {
         let Some(mut rx) = self.active_thread_rx.take() else {
             return Ok(());
         };
@@ -1465,6 +1475,9 @@ impl App {
                     break;
                 }
             }
+            if Instant::now() >= frame_deadline {
+                break;
+            }
         }
 
         if !disconnected {