From eceb3eeaf3a68d732596fd8c0e8a6807f9166770 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 20 Jul 2026 13:50:38 +0000 Subject: [PATCH] Cache TUI flex heights across frame passes (#34348) ## Why Sizing, rendering, and cursor placement can query the same chat widget layout multiple times in one frame, repeatedly measuring active transcript cells. ## What changed - Build one chat widget renderable tree per frame and reuse it for sizing, rendering, and cursor placement. - Cache each flex child's desired height by width for the lifetime of that tree. - Reuse the bottom pane's renderable directly instead of forwarding each renderable operation through a wrapper. ## Testing - Verify flex layouts measure a child once across frame passes and remeasure it when the width changes. - Verify a chat widget frame measures its active transcript cell once. GitOrigin-RevId: 5ad1a6711f4011c699b5d002b13dc3319cb4db8e --- codex-rs/tui/src/app.rs | 36 ++++++++++----- codex-rs/tui/src/app/tests.rs | 47 ++++++++++++++++++++ codex-rs/tui/src/bottom_pane/mod.rs | 39 +--------------- codex-rs/tui/src/chatwidget/rendering.rs | 43 ++++-------------- codex-rs/tui/src/chatwidget/tests/helpers.rs | 4 ++ codex-rs/tui/src/render/renderable.rs | 22 ++++++--- codex-rs/tui/src/render/renderable_tests.rs | 40 +++++++++++++++++ 7 files changed, 141 insertions(+), 90 deletions(-) diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 3066688dd5f7..c70d1ea5d6fb 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -1364,18 +1364,30 @@ See the Codex keymap documentation for supported actions and examples." } fn render_chat_widget_frame(&mut self, tui: &mut tui::Tui) -> Result { - let desired_height = self.chat_widget.desired_height(tui.terminal.size()?.width); - let mut rendered_area = Rect::default(); - tui.draw_with_resize_reflow(desired_height, |frame| { - let area = frame.area(); - rendered_area = area; - self.chat_widget.render(area, frame.buffer); - if let Some((x, y)) = self.chat_widget.cursor_pos(area) { - frame.set_cursor_style(self.chat_widget.cursor_style(area)); - frame.set_cursor_position((x, y)); - } - })?; - Ok(rendered_area) + let width = tui.terminal.size()?.width; + self.with_chat_widget_frame(width, |desired_height, chat_widget| { + let mut rendered_area = Rect::default(); + tui.draw_with_resize_reflow(desired_height, |frame| { + let area = frame.area(); + rendered_area = area; + chat_widget.render(area, frame.buffer); + self.chat_widget.note_rendered_width(area.width); + if let Some((x, y)) = chat_widget.cursor_pos(area) { + frame.set_cursor_style(chat_widget.cursor_style(area)); + frame.set_cursor_position((x, y)); + } + })?; + Ok(rendered_area) + }) + } + + fn with_chat_widget_frame( + &self, + width: u16, + render: impl FnOnce(u16, &dyn Renderable) -> T, + ) -> T { + let chat_widget = self.chat_widget.as_renderable(); + render(chat_widget.desired_height(width), &chat_widget) } } diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 381df134cda9..f044c45afb17 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -20,6 +20,7 @@ use crate::app_event::HistoryBatchEntryResponse; use crate::chatwidget::ChatWidgetInit; use crate::chatwidget::create_initial_user_message; use crate::chatwidget::tests::helpers::render_bottom_popup; +use crate::chatwidget::tests::helpers::set_active_cell; use crate::chatwidget::tests::make_chatwidget_manual_with_sender; use crate::chatwidget::tests::set_chatgpt_auth; use crate::chatwidget::tests::set_fast_mode_test_catalog; @@ -116,11 +117,14 @@ use codex_utils_absolute_path::AbsolutePathBuf; use crossterm::event::KeyModifiers; use insta::assert_snapshot; use pretty_assertions::assert_eq; +use ratatui::buffer::Buffer; use ratatui::prelude::Line; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use tempfile::tempdir; use tokio::time; @@ -136,6 +140,49 @@ fn test_absolute_path(path: &str) -> AbsolutePathBuf { AbsolutePathBuf::try_from(PathBuf::from(path)).expect("absolute test path") } +#[tokio::test] +async fn chat_widget_frame_reuses_active_cell_height_across_frame_passes() { + #[derive(Debug)] + struct CountingHistoryCell { + desired_height_calls: Arc, + } + + impl HistoryCell for CountingHistoryCell { + fn display_lines(&self, _width: u16) -> Vec> { + vec![Line::from("active cell")] + } + + fn raw_lines(&self) -> Vec> { + vec![Line::from("active cell")] + } + + fn desired_height(&self, _width: u16) -> u16 { + self.desired_height_calls.fetch_add(1, Ordering::Relaxed); + 1 + } + } + + let mut app = make_test_app().await; + let desired_height_calls = Arc::new(AtomicUsize::new(0)); + set_active_cell( + &mut app.chat_widget, + Box::new(CountingHistoryCell { + desired_height_calls: Arc::clone(&desired_height_calls), + }), + ); + let width = 80; + app.with_chat_widget_frame(width, |desired_height, chat_widget| { + let area = Rect::new(/*x*/ 0, /*y*/ 0, width, desired_height); + let mut buffer = Buffer::empty(area); + + chat_widget.render(area, &mut buffer); + assert!(chat_widget.cursor_pos(area).is_some()); + let _ = chat_widget.cursor_style(area); + }); + + assert_eq!(desired_height_calls.load(Ordering::Relaxed), 1); +} + async fn next_thread_settings_updated( app_server: &mut AppServerSession, thread_id: ThreadId, diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index 944a2b893113..ee9b8051c72f 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -1692,7 +1692,7 @@ impl BottomPane { self.as_renderable_with_composer_right_reserve(/*composer_right_reserve*/ 0) } - fn as_renderable_with_composer_right_reserve( + pub(crate) fn as_renderable_with_composer_right_reserve( &'_ self, composer_right_reserve: u16, ) -> RenderableItem<'_> { @@ -1750,43 +1750,6 @@ impl BottomPane { } } - pub(crate) fn render_with_composer_right_reserve( - &self, - area: Rect, - buf: &mut Buffer, - composer_right_reserve: u16, - ) { - self.as_renderable_with_composer_right_reserve(composer_right_reserve) - .render(area, buf); - } - - pub(crate) fn desired_height_with_composer_right_reserve( - &self, - width: u16, - composer_right_reserve: u16, - ) -> u16 { - self.as_renderable_with_composer_right_reserve(composer_right_reserve) - .desired_height(width) - } - - pub(crate) fn cursor_pos_with_composer_right_reserve( - &self, - area: Rect, - composer_right_reserve: u16, - ) -> Option<(u16, u16)> { - self.as_renderable_with_composer_right_reserve(composer_right_reserve) - .cursor_pos(area) - } - - pub(crate) fn cursor_style_with_composer_right_reserve( - &self, - area: Rect, - composer_right_reserve: u16, - ) -> crossterm::cursor::SetCursorStyle { - self.as_renderable_with_composer_right_reserve(composer_right_reserve) - .cursor_style(area) - } - pub(crate) fn set_status_line(&mut self, status_line: Option>) { if self.composer.set_status_line(status_line) { self.request_redraw(); diff --git a/codex-rs/tui/src/chatwidget/rendering.rs b/codex-rs/tui/src/chatwidget/rendering.rs index 710af6b7ea42..44aef7050d61 100644 --- a/codex-rs/tui/src/chatwidget/rendering.rs +++ b/codex-rs/tui/src/chatwidget/rendering.rs @@ -3,7 +3,7 @@ use super::*; impl ChatWidget { - pub(super) fn as_renderable(&self) -> RenderableItem<'_> { + pub(crate) fn as_renderable(&self) -> RenderableItem<'_> { let active_cell_right_reserve = self.ambient_pet_wrap_reserved_cols(); let active_cell_renderable = match &self.transcript.active_cell { Some(cell) => RenderableItem::Owned(Box::new(TranscriptAreaRenderable { @@ -48,42 +48,17 @@ impl ChatWidget { } flex.push( /*flex*/ 0, - RenderableItem::Owned(Box::new(BottomPaneComposerReserveRenderable { - bottom_pane: &self.bottom_pane, - right_reserve: active_cell_right_reserve, - })) - .inset(Insets::tlbr( - /*top*/ 1, /*left*/ 0, /*bottom*/ 0, /*right*/ 0, - )), + self.bottom_pane + .as_renderable_with_composer_right_reserve(active_cell_right_reserve) + .inset(Insets::tlbr( + /*top*/ 1, /*left*/ 0, /*bottom*/ 0, /*right*/ 0, + )), ); RenderableItem::Owned(Box::new(flex)) } -} - -struct BottomPaneComposerReserveRenderable<'a> { - bottom_pane: &'a BottomPane, - right_reserve: u16, -} -impl Renderable for BottomPaneComposerReserveRenderable<'_> { - fn render(&self, area: Rect, buf: &mut Buffer) { - self.bottom_pane - .render_with_composer_right_reserve(area, buf, self.right_reserve); - } - - fn desired_height(&self, width: u16) -> u16 { - self.bottom_pane - .desired_height_with_composer_right_reserve(width, self.right_reserve) - } - - fn cursor_pos(&self, area: Rect) -> Option<(u16, u16)> { - self.bottom_pane - .cursor_pos_with_composer_right_reserve(area, self.right_reserve) - } - - fn cursor_style(&self, area: Rect) -> crossterm::cursor::SetCursorStyle { - self.bottom_pane - .cursor_style_with_composer_right_reserve(area, self.right_reserve) + pub(crate) fn note_rendered_width(&self, width: u16) { + self.last_rendered_width.set(Some(width as usize)); } } @@ -132,7 +107,7 @@ impl TranscriptAreaRenderable<'_> { impl Renderable for ChatWidget { fn render(&self, area: Rect, buf: &mut Buffer) { self.as_renderable().render(area, buf); - self.last_rendered_width.set(Some(area.width as usize)); + self.note_rendered_width(area.width); } fn desired_height(&self, width: u16) -> u16 { diff --git a/codex-rs/tui/src/chatwidget/tests/helpers.rs b/codex-rs/tui/src/chatwidget/tests/helpers.rs index b12d3ddce93b..fe392481f645 100644 --- a/codex-rs/tui/src/chatwidget/tests/helpers.rs +++ b/codex-rs/tui/src/chatwidget/tests/helpers.rs @@ -215,6 +215,10 @@ pub(super) async fn make_chatwidget_manual_with_auth( (widget, rx, op_rx) } +pub(crate) fn set_active_cell(chat: &mut ChatWidget, cell: Box) { + chat.transcript.active_cell = Some(cell); +} + // ChatWidget may emit other `Op`s (e.g. history/logging updates) on the same channel; this helper // filters until we see a submission op. pub(super) fn next_submit_op(op_rx: &mut tokio::sync::mpsc::UnboundedReceiver) -> Op { diff --git a/codex-rs/tui/src/render/renderable.rs b/codex-rs/tui/src/render/renderable.rs index bd22353d05f6..a1792fe3fd84 100644 --- a/codex-rs/tui/src/render/renderable.rs +++ b/codex-rs/tui/src/render/renderable.rs @@ -1,3 +1,4 @@ +use std::cell::Cell; use std::sync::Arc; use crossterm::cursor::SetCursorStyle; @@ -247,6 +248,7 @@ impl<'a> ColumnRenderable<'a> { pub struct FlexChild<'a> { flex: i32, child: RenderableItem<'a>, + cached_height: Cell>, } pub struct FlexRenderable<'a> { @@ -266,6 +268,7 @@ impl<'a> FlexRenderable<'a> { self.children.push(FlexChild { flex, child: child.into(), + cached_height: Cell::new(None), }); } @@ -280,13 +283,20 @@ impl<'a> FlexRenderable<'a> { // 1. Allocate space to non-flex children. let max_size = area.height; - for (i, FlexChild { flex, child }) in self.children.iter().enumerate() { - if *flex > 0 { - flex_children.push((i, *flex as u16, child.desired_height(area.width))); + for (i, child) in self.children.iter().enumerate() { + let desired_height = if let Some((width, height)) = child.cached_height.get() + && width == area.width + { + height + } else { + let height = child.child.desired_height(area.width); + child.cached_height.set(Some((area.width, height))); + height + }; + if child.flex > 0 { + flex_children.push((i, child.flex as u16, desired_height)); } else { - child_sizes[i] = child - .desired_height(area.width) - .min(max_size.saturating_sub(allocated_size)); + child_sizes[i] = desired_height.min(max_size.saturating_sub(allocated_size)); allocated_size += child_sizes[i]; } } diff --git a/codex-rs/tui/src/render/renderable_tests.rs b/codex-rs/tui/src/render/renderable_tests.rs index 9c14130a9f88..7d95d4785ab4 100644 --- a/codex-rs/tui/src/render/renderable_tests.rs +++ b/codex-rs/tui/src/render/renderable_tests.rs @@ -1,5 +1,6 @@ use super::*; use pretty_assertions::assert_eq; +use std::cell::Cell; struct HeightRenderable(u16); @@ -70,3 +71,42 @@ fn flex_reserves_non_flex_space_before_flexible_children() { vec![4, 2, 4], ); } + +#[test] +fn flex_caches_child_height_across_frame_passes() { + struct CountingRenderable<'a>(&'a Cell); + + impl Renderable for CountingRenderable<'_> { + fn render(&self, _area: Rect, _buf: &mut Buffer) {} + + fn desired_height(&self, _width: u16) -> u16 { + self.0.set(self.0.get() + 1); + 1 + } + + fn cursor_pos(&self, area: Rect) -> Option<(u16, u16)> { + Some((area.x, area.y)) + } + } + + let calls = Cell::new(0); + let renderable = CountingRenderable(&calls); + let mut flex = FlexRenderable::new(); + flex.push(/*flex*/ 1, RenderableItem::Borrowed(&renderable)); + let area = Rect::new( + /*x*/ 0, /*y*/ 0, /*width*/ 80, /*height*/ 10, + ); + let mut buf = Buffer::empty(area); + + assert_eq!(flex.desired_height(area.width), 1); + flex.render(area, &mut buf); + assert_eq!(flex.cursor_pos(area), Some((0, 0))); + assert!(matches!( + flex.cursor_style(area), + crossterm::cursor::SetCursorStyle::DefaultUserShape + )); + assert_eq!(calls.get(), 1); + + assert_eq!(flex.desired_height(/*width*/ 100), 1); + assert_eq!(calls.get(), 2); +}