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 {