From f6519a355a309155c6c9d8f33fc22ab4d4a9ffdb Mon Sep 17 00:00:00 2001 From: Tamir Duberstein Date: Fri, 21 Aug 2026 18:02:16 +0000 Subject: [PATCH] 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 | 1 - codex-rs/tui/src/app/startup.rs | 19 +++- codex-rs/tui/src/app/tests.rs | 111 ++++++++++++++++++++++ codex-rs/tui/src/app/thread_routing.rs | 13 +++ 4 files changed, 140 insertions(+), 4 deletions(-) diff --git a/codex-rs/tui/src/app/session_lifecycle.rs b/codex-rs/tui/src/app/session_lifecycle.rs index 60e6c54004..8178adaf83 100644 --- a/codex-rs/tui/src/app/session_lifecycle.rs +++ b/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(()) diff --git a/codex-rs/tui/src/app/startup.rs b/codex-rs/tui/src/app/startup.rs index ee9b713a76..615398e352 100644 --- a/codex-rs/tui/src/app/startup.rs +++ b/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); diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 71b2c47337..ca3cf6cb35 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/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; diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 28a003a66f..861fbc1396 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/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 {