Cache stable active-cell layout measurements (#38242)

## What changed

- Reuse desired and rendered heights for active transcript cells that report a stable height.
- Invalidate the cached layout when the cell identity, revision, width, render mode, or syntax theme changes, and when the active cell is removed.
- Continue rebuilding display lines on every frame so animated content stays live, while dynamic-height cells bypass the persistent cache.

## Testing

- Add rendering tests for height reuse, bottom-anchored overflow, cache invalidation, animated display lines, and dynamic-height cells.

GitOrigin-RevId: 414cc4269bfb6b68105837c82e31787272d335a7

codex-rs/tui/src/chatwidget.rs

@@ -1219,7 +1219,7 @@ impl ChatWidget {
     }
 
     fn flush_active_cell(&mut self) {
-        if let Some(active) = self.transcript.active_cell.take() {
+        if let Some(active) = self.transcript.take_active_cell() {
             self.transcript.needs_final_message_separator = true;
             self.app_event_tx.send(AppEvent::InsertHistoryCell(active));
             self.request_pending_usage_output_insertion();
@@ -1378,7 +1378,7 @@ impl ChatWidget {
 
     /// Mark the active cell as failed (✗) and flush it into history.
     fn finalize_active_cell_as_failed(&mut self) {
-        if let Some(mut cell) = self.transcript.active_cell.take() {
+        if let Some(mut cell) = self.transcript.take_active_cell() {
             // Insert finalized cell into history and keep grouping consistent.
             if let Some(exec) = cell.as_any_mut().downcast_mut::<ExecCell>() {
                 exec.mark_failed();
@@ -1470,7 +1470,7 @@ impl ChatWidget {
     /// Merge the real session info cell with any placeholder header to avoid double boxes.
     fn apply_session_info_cell(&mut self, cell: history_cell::SessionInfoCell) {
         let mut session_info_cell = Some(Box::new(cell) as Box<dyn HistoryCell>);
-        let merged_header = if let Some(active) = self.transcript.active_cell.take() {
+        let merged_header = if let Some(active) = self.transcript.take_active_cell() {
             if active
                 .as_any()
                 .is::<history_cell::SessionHeaderHistoryCell>()

codex-rs/tui/src/chatwidget/rendering.rs

@@ -1,6 +1,9 @@
 //! Render composition for the main chat widget surface.
 
+use super::transcript::ActiveCellLayoutCache;
+use super::transcript::ActiveCellLayoutCacheKey;
 use super::*;
+use std::cell::Cell;
 
 impl ChatWidget {
     pub(crate) fn as_renderable(&self) -> RenderableItem<'_> {
@@ -10,6 +13,18 @@ impl ChatWidget {
                 child: cell.as_ref(),
                 top: 1,
                 right: active_cell_right_reserve,
+                // Externally backed transcript cells can also change viewport height without an
+                // active-cell revision. Spinner cells remain safe because their indicator width
+                // is stable and their display lines are still rebuilt on every frame.
+                persistent_layout: cell.has_stable_transcript_height().then_some(
+                    PersistentActiveCellLayout {
+                        cache: &self.transcript.active_cell_layout,
+                        cell_identity: cell.as_ref() as *const dyn HistoryCell as *const ()
+                            as usize,
+                        revision: self.transcript.active_cell_revision,
+                        render_mode: self.history_render_mode(),
+                    },
+                ),
             })),
             None => RenderableItem::Owned(Box::new(())),
         };
@@ -19,6 +34,7 @@ impl ChatWidget {
                     child: cell,
                     top: 1,
                     right: active_cell_right_reserve,
+                    persistent_layout: None,
                 }))
             }
             _ => RenderableItem::Owned(Box::new(())),
@@ -33,6 +49,7 @@ impl ChatWidget {
                     child: cell,
                     top: 1,
                     right: active_cell_right_reserve,
+                    persistent_layout: None,
                 })),
             );
         }
@@ -43,6 +60,7 @@ impl ChatWidget {
                     child: cell,
                     top: 1,
                     right: active_cell_right_reserve,
+                    persistent_layout: None,
                 })),
             );
         }
@@ -66,6 +84,14 @@ struct TranscriptAreaRenderable<'a> {
     child: &'a dyn HistoryCell,
     top: u16,
     right: u16,
+    persistent_layout: Option<PersistentActiveCellLayout<'a>>,
+}
+
+struct PersistentActiveCellLayout<'a> {
+    cache: &'a Cell<Option<ActiveCellLayoutCache>>,
+    cell_identity: usize,
+    revision: u64,
+    render_mode: HistoryRenderMode,
 }
 
 impl Renderable for TranscriptAreaRenderable<'_> {
@@ -76,9 +102,19 @@ impl Renderable for TranscriptAreaRenderable<'_> {
         let y = if area.height == 0 {
             0
         } else {
-            let overflow = paragraph
-                .line_count(area.width)
-                .saturating_sub(usize::from(area.height));
+            let rendered_height = if let Some((cache, mut layout)) = self.layout(area.width) {
+                if let Some(height) = layout.rendered_height {
+                    height
+                } else {
+                    let height = paragraph.line_count(area.width);
+                    layout.rendered_height = Some(height);
+                    cache.set(Some(layout));
+                    height
+                }
+            } else {
+                paragraph.line_count(area.width)
+            };
+            let overflow = rendered_height.saturating_sub(usize::from(area.height));
             u16::try_from(overflow).unwrap_or(u16::MAX)
         };
         Clear.render(area, buf);
@@ -87,11 +123,47 @@ impl Renderable for TranscriptAreaRenderable<'_> {
 
     fn desired_height(&self, width: u16) -> u16 {
         let child_width = width.saturating_sub(self.right).max(1);
-        HistoryCell::desired_height(self.child, child_width) + self.top
+        let desired_height = if let Some((cache, mut layout)) = self.layout(child_width) {
+            if let Some(height) = layout.desired_height {
+                height
+            } else {
+                let height = HistoryCell::desired_height(self.child, child_width);
+                layout.desired_height = Some(height);
+                cache.set(Some(layout));
+                height
+            }
+        } else {
+            HistoryCell::desired_height(self.child, child_width)
+        };
+        desired_height + self.top
     }
 }
 
 impl TranscriptAreaRenderable<'_> {
+    fn layout(
+        &self,
+        width: u16,
+    ) -> Option<(&Cell<Option<ActiveCellLayoutCache>>, ActiveCellLayoutCache)> {
+        let persistent = self.persistent_layout.as_ref()?;
+        let key = ActiveCellLayoutCacheKey {
+            cell_identity: persistent.cell_identity,
+            revision: persistent.revision,
+            width,
+            render_mode: persistent.render_mode,
+            syntax_theme_revision: crate::render::highlight::syntax_theme_revision(),
+        };
+        let layout = persistent
+            .cache
+            .get()
+            .filter(|layout| layout.key == key)
+            .unwrap_or(ActiveCellLayoutCache {
+                key,
+                desired_height: None,
+                rendered_height: None,
+            });
+        Some((persistent.cache, layout))
+    }
+
     fn child_area(&self, area: Rect) -> Rect {
         let y = area.y.saturating_add(self.top);
         let height = area.height.saturating_sub(self.top);
@@ -104,6 +176,10 @@ impl TranscriptAreaRenderable<'_> {
     }
 }
 
+#[cfg(test)]
+#[path = "rendering_tests.rs"]
+mod tests;
+
 impl Renderable for ChatWidget {
     fn render(&self, area: Rect, buf: &mut Buffer) {
         self.as_renderable().render(area, buf);

codex-rs/tui/src/chatwidget/rendering_tests.rs

@@ -0,0 +1,213 @@
+use super::*;
+use crate::chatwidget::tests::make_chatwidget_manual_with_sender;
+use pretty_assertions::assert_eq;
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+#[derive(Debug)]
+struct CountingHistoryCell {
+    desired_height_calls: Arc<AtomicUsize>,
+    display_lines_calls: Arc<AtomicUsize>,
+    desired_height: u16,
+    line_count: usize,
+    stable_height: bool,
+}
+
+impl HistoryCell for CountingHistoryCell {
+    fn display_lines(&self, _width: u16) -> Vec<Line<'static>> {
+        let frame = self.display_lines_calls.fetch_add(1, Ordering::Relaxed) + 1;
+        (0..self.line_count)
+            .map(|row| Line::from(format!("frame {frame} row {row}")))
+            .collect()
+    }
+
+    fn raw_lines(&self) -> Vec<Line<'static>> {
+        Vec::new()
+    }
+
+    fn desired_height(&self, _width: u16) -> u16 {
+        self.desired_height_calls.fetch_add(1, Ordering::Relaxed);
+        self.desired_height
+    }
+
+    fn has_stable_transcript_height(&self) -> bool {
+        self.stable_height
+    }
+}
+
+async fn widget_with_counting_cell(
+    desired_height: u16,
+    line_count: usize,
+    stable_height: bool,
+) -> (ChatWidget, Arc<AtomicUsize>, Arc<AtomicUsize>) {
+    let (mut widget, _sender, _events, _operations) = make_chatwidget_manual_with_sender().await;
+    let desired_height_calls = Arc::new(AtomicUsize::new(0));
+    let display_lines_calls = Arc::new(AtomicUsize::new(0));
+    widget.transcript.active_cell = Some(Box::new(CountingHistoryCell {
+        desired_height_calls: Arc::clone(&desired_height_calls),
+        display_lines_calls: Arc::clone(&display_lines_calls),
+        desired_height,
+        line_count,
+        stable_height,
+    }));
+    (widget, desired_height_calls, display_lines_calls)
+}
+
+fn render_frame(widget: &ChatWidget, width: u16) -> Buffer {
+    let renderable = widget.as_renderable();
+    let height = renderable.desired_height(width);
+    let area = Rect::new(/*x*/ 0, /*y*/ 0, width, height);
+    let mut buffer = Buffer::empty(area);
+    renderable.render(area, &mut buffer);
+    buffer
+}
+
+fn contains_text(buffer: &Buffer, text: &str) -> bool {
+    buffer
+        .content
+        .chunks(usize::from(buffer.area.width))
+        .any(|row| {
+            row.iter()
+                .map(ratatui::buffer::Cell::symbol)
+                .collect::<String>()
+                .contains(text)
+        })
+}
+
+#[tokio::test]
+async fn active_cell_layout_reuses_heights_without_freezing_animation() {
+    let (widget, desired_height_calls, display_lines_calls) = widget_with_counting_cell(
+        /*desired_height*/ 2, /*line_count*/ 2, /*stable_height*/ true,
+    )
+    .await;
+
+    let first = render_frame(&widget, /*width*/ 80);
+    let second = render_frame(&widget, /*width*/ 80);
+
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 1);
+    assert_eq!(display_lines_calls.load(Ordering::Relaxed), 2);
+    assert!(contains_text(&first, "frame 1 row 0"));
+    assert!(contains_text(&second, "frame 2 row 0"));
+    let cached = widget
+        .transcript
+        .active_cell_layout
+        .get()
+        .expect("active-cell layout should be cached");
+    assert_eq!(cached.desired_height, Some(2));
+    assert_eq!(cached.rendered_height, Some(2));
+}
+
+#[tokio::test]
+async fn active_cell_layout_preserves_custom_height_and_bottom_scroll() {
+    let (widget, desired_height_calls, _display_lines_calls) = widget_with_counting_cell(
+        /*desired_height*/ 1, /*line_count*/ 3, /*stable_height*/ true,
+    )
+    .await;
+
+    let first = render_frame(&widget, /*width*/ 80);
+    let second = render_frame(&widget, /*width*/ 80);
+
+    let mut transcript = second.clone();
+    transcript.resize(Rect {
+        height: 2,
+        ..second.area
+    });
+    insta::assert_snapshot!(
+        "cached_active_cell_bottom_scroll",
+        format!("{transcript:?}")
+    );
+    assert!(contains_text(&first, "frame 1 row 2"));
+    assert!(!contains_text(&first, "frame 1 row 0"));
+    assert!(contains_text(&second, "frame 2 row 2"));
+    assert!(!contains_text(&second, "frame 2 row 0"));
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 1);
+    let cached = widget
+        .transcript
+        .active_cell_layout
+        .get()
+        .expect("active-cell layout should be cached");
+    assert_eq!(cached.desired_height, Some(1));
+    assert_eq!(cached.rendered_height, Some(3));
+}
+
+#[tokio::test]
+async fn active_cell_layout_invalidates_width_revision_mode_theme_and_identity() {
+    let (mut widget, desired_height_calls, _display_lines_calls) = widget_with_counting_cell(
+        /*desired_height*/ 1, /*line_count*/ 1, /*stable_height*/ true,
+    )
+    .await;
+
+    render_frame(&widget, /*width*/ 80);
+    render_frame(&widget, /*width*/ 80);
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 1);
+
+    render_frame(&widget, /*width*/ 81);
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 2);
+
+    widget.transcript.bump_active_cell_revision();
+    render_frame(&widget, /*width*/ 81);
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 3);
+
+    widget.raw_output_mode = true;
+    render_frame(&widget, /*width*/ 81);
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 4);
+
+    let mut cached = widget
+        .transcript
+        .active_cell_layout
+        .get()
+        .expect("active-cell layout should be cached");
+    cached.key.syntax_theme_revision = cached.key.syntax_theme_revision.wrapping_sub(1);
+    widget.transcript.active_cell_layout.set(Some(cached));
+    render_frame(&widget, /*width*/ 81);
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 5);
+
+    let replacement_height_calls = Arc::new(AtomicUsize::new(0));
+    widget.transcript.active_cell = Some(Box::new(CountingHistoryCell {
+        desired_height_calls: Arc::clone(&replacement_height_calls),
+        display_lines_calls: Arc::new(AtomicUsize::new(0)),
+        desired_height: 1,
+        line_count: 1,
+        stable_height: true,
+    }));
+    render_frame(&widget, /*width*/ 81);
+    assert_eq!(replacement_height_calls.load(Ordering::Relaxed), 1);
+}
+
+#[tokio::test]
+async fn externally_mutable_active_cells_bypass_persistent_layout_cache() {
+    let (widget, desired_height_calls, display_lines_calls) = widget_with_counting_cell(
+        /*desired_height*/ 2, /*line_count*/ 2, /*stable_height*/ false,
+    )
+    .await;
+
+    render_frame(&widget, /*width*/ 80);
+    render_frame(&widget, /*width*/ 80);
+
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 2);
+    assert_eq!(display_lines_calls.load(Ordering::Relaxed), 2);
+    assert_eq!(widget.transcript.active_cell_layout.get(), None);
+}
+
+#[tokio::test]
+async fn removing_active_cell_invalidates_layout_before_reusing_its_identity() {
+    let (mut widget, desired_height_calls, _display_lines_calls) = widget_with_counting_cell(
+        /*desired_height*/ 2, /*line_count*/ 2, /*stable_height*/ true,
+    )
+    .await;
+
+    render_frame(&widget, /*width*/ 80);
+    let revision = widget.transcript.active_cell_revision;
+    let active_cell = widget
+        .transcript
+        .take_active_cell()
+        .expect("active cell should exist");
+
+    assert_eq!(widget.transcript.active_cell_layout.get(), None);
+    widget.transcript.active_cell = Some(active_cell);
+    assert_eq!(widget.transcript.active_cell_revision, revision);
+
+    render_frame(&widget, /*width*/ 80);
+    assert_eq!(desired_height_calls.load(Ordering::Relaxed), 2);
+}

codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__rendering__tests__cached_active_cell_bottom_scroll.snap

@@ -0,0 +1,14 @@
+---
+source: tui/src/chatwidget/rendering_tests.rs
+expression: "format!(\"{transcript:?}\")"
+---
+Buffer {
+    area: Rect { x: 0, y: 0, width: 80, height: 2 },
+    content: [
+        "                                                                                ",
+        "frame 2 row 2                                                                   ",
+    ],
+    styles: [
+        x: 0, y: 0, fg: Reset, bg: Reset, underline: Reset, modifier: NONE,
+    ]
+}

codex-rs/tui/src/chatwidget/transcript.rs

@@ -1,12 +1,37 @@
 //! Transcript and active-cell bookkeeping for `ChatWidget`.
 
 use super::HistoryCell;
+use super::HistoryRenderMode;
+use std::cell::Cell;
+
+/// Identifies the render state that determines an active cell's viewport height.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(super) struct ActiveCellLayoutCacheKey {
+    pub(super) cell_identity: usize,
+    pub(super) revision: u64,
+    pub(super) width: u16,
+    pub(super) render_mode: HistoryRenderMode,
+    pub(super) syntax_theme_revision: u64,
+}
+
+/// Retains the active cell's semantic and actual wrapped heights independently.
+///
+/// History cells may override their desired height, so it cannot be substituted for the rendered
+/// row count used to keep overflowing content anchored to the bottom of the viewport.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(super) struct ActiveCellLayoutCache {
+    pub(super) key: ActiveCellLayoutCacheKey,
+    pub(super) desired_height: Option<u16>,
+    pub(super) rendered_height: Option<usize>,
+}
 
 #[derive(Default)]
 pub(super) struct TranscriptState {
     pub(super) active_cell: Option<Box<dyn HistoryCell>>,
     /// Monotonic-ish counter used to invalidate transcript overlay caching.
     pub(super) active_cell_revision: u64,
+    /// One bounded entry shared by layout and paint across unchanged active-cell frames.
+    pub(super) active_cell_layout: Cell<Option<ActiveCellLayoutCache>>,
     /// Raw markdown of the most recently completed agent response.
     pub(super) last_agent_markdown: Option<String>,
     pub(super) last_completed_agent_message: Option<(String, String)>,
@@ -43,6 +68,16 @@ impl TranscriptState {
         // Wrapping avoids overflow; wraparound would require 2^64 bumps and at
         // worst causes a one-time cache-key collision.
         self.active_cell_revision = self.active_cell_revision.wrapping_add(1);
+        self.active_cell_layout.set(None);
+    }
+
+    /// Remove the active cell and invalidate its layout before its address can be reused.
+    pub(super) fn take_active_cell(&mut self) -> Option<Box<dyn HistoryCell>> {
+        let active_cell = self.active_cell.take();
+        if active_cell.is_some() {
+            self.active_cell_layout.set(None);
+        }
+        active_cell
     }
 
     pub(super) fn record_agent_markdown(&mut self, markdown: String) {