From edc9effc019ba4814fdd723d8b5535ccbc5dad65 Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Sun, 31 May 2026 11:35:54 -0300 Subject: [PATCH] perf(tui): load large transcripts asynchronously --- codex-rs/tui/src/app/event_dispatch.rs | 7 + codex-rs/tui/src/app_backtrack.rs | 4 +- codex-rs/tui/src/app_event.rs | 4 + codex-rs/tui/src/lib.rs | 1 + codex-rs/tui/src/pager_overlay.rs | 295 +++++++++++++++++- codex-rs/tui/src/session_log.rs | 13 + ...renders_loading_until_layout_is_ready.snap | 13 + codex-rs/tui/src/transcript_layout.rs | 158 ++++++++++ 8 files changed, 488 insertions(+), 7 deletions(-) create mode 100644 codex-rs/tui/src/snapshots/codex_tui__pager_overlay__tests__large_transcript_renders_loading_until_layout_is_ready.snap create mode 100644 codex-rs/tui/src/transcript_layout.rs diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index 2e1b97aace..c9fa579d2e 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -219,6 +219,13 @@ impl App { ); } } + AppEvent::TranscriptLayoutReady(result) => { + if let Some(Overlay::Transcript(transcript)) = &mut self.overlay + && transcript.apply_layout_result(result) + { + tui.frame_requester().schedule_frame(); + } + } AppEvent::EndInitialHistoryReplayBuffer => { self.finish_initial_history_replay_buffer(tui); } diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index f692c9ab05..52ef930601 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -443,8 +443,10 @@ impl App { { let active_key = self.chat_widget.active_cell_transcript_key(); let chat_widget = &self.chat_widget; + let app_event_tx = self.app_event_tx.clone(); tui.draw(u16::MAX, |frame| { - let width = frame.area().width.max(1); + let width = frame.area().width.max(/*other*/ 1); + t.ensure_async_layout(width, app_event_tx.clone()); t.sync_live_tail(width, active_key, |w| { chat_widget.active_cell_transcript_hyperlink_lines(w) }); diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index f1e612dd58..81c782bf81 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -38,6 +38,7 @@ use crate::bottom_pane::ApprovalRequest; use crate::bottom_pane::StatusLineItem; use crate::bottom_pane::TerminalTitleItem; use crate::chatwidget::UserMessage; +use crate::transcript_layout::TranscriptLayoutResult; use codex_app_server_protocol::AskForApproval; use codex_config::types::ApprovalsReviewer; use codex_features::Feature; @@ -576,6 +577,9 @@ pub(crate) enum AppEvent { InsertHistoryCell(Box), + /// Result of measuring the transcript overlay layout off the UI thread. + TranscriptLayoutReady(TranscriptLayoutResult), + /// Finish buffering initial resume replay after all replay events have been queued. EndInitialHistoryReplayBuffer, diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 061dedb672..4cb2ac8832 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -192,6 +192,7 @@ mod text_formatting; mod theme_picker; mod token_usage; mod tooltips; +mod transcript_layout; mod transcript_reflow; mod tui; mod ui_consts; diff --git a/codex-rs/tui/src/pager_overlay.rs b/codex-rs/tui/src/pager_overlay.rs index 6c2da9f403..edf2f22ad1 100644 --- a/codex-rs/tui/src/pager_overlay.rs +++ b/codex-rs/tui/src/pager_overlay.rs @@ -23,6 +23,7 @@ use std::sync::Arc; use std::time::Duration; use std::time::Instant; +use crate::app_event_sender::AppEventSender; use crate::chatwidget::ActiveCellTranscriptKey; use crate::chatwidget::CopyStatus; use crate::footer_hints::FooterHint; @@ -45,6 +46,9 @@ use crate::style::user_message_style; use crate::terminal_hyperlinks::HyperlinkLine; use crate::terminal_hyperlinks::mark_buffer_hyperlinks; use crate::terminal_hyperlinks::visible_lines; +use crate::transcript_layout::TranscriptLayoutKey; +use crate::transcript_layout::TranscriptLayoutResult; +use crate::transcript_layout::spawn_transcript_layout_worker; use crate::tui; use crate::tui::TuiEvent; use crossterm::event::KeyCode; @@ -200,6 +204,21 @@ impl PagerView { self.layout(width).total_height } + fn install_precomputed_layout(&mut self, width: u16, heights: Vec) { + let mut offsets = Vec::with_capacity(heights.len()); + let mut total_height = 0usize; + for height in &heights { + offsets.push(total_height); + total_height = total_height.saturating_add(*height); + } + self.layout = Some(PagerLayout { + width, + offsets: offsets.into(), + heights: heights.into(), + total_height, + }); + } + fn layout(&mut self, width: u16) -> &PagerLayout { let needs_rebuild = self.layout.as_ref().is_none_or(|layout| { layout.width != width || layout.heights.len() != self.renderables.len() @@ -550,10 +569,12 @@ pub(crate) struct TranscriptOverlay { footer_status: Option, /// Cache key for the render-only live tail appended after committed cells. live_tail_key: Option, + async_layout: AsyncTranscriptLayout, is_done: bool, } const FOOTER_STATUS_TTL: Duration = Duration::from_secs(2); +const ASYNC_TRANSCRIPT_LAYOUT_CELL_THRESHOLD: usize = 2_000; #[derive(Clone)] struct FooterStatus { @@ -576,6 +597,29 @@ struct LiveTailKey { animation_tick: Option, } +#[derive(Debug)] +struct AsyncTranscriptLayout { + generation: u64, + pending: Option, + measured: Option, +} + +impl AsyncTranscriptLayout { + fn new() -> Self { + Self { + generation: 0, + pending: None, + measured: None, + } + } + + fn invalidate(&mut self) { + self.generation = self.generation.saturating_add(/*rhs*/ 1); + self.pending = None; + self.measured = None; + } +} + impl TranscriptOverlay { /// Creates a transcript overlay for a fixed set of committed cells. /// @@ -589,14 +633,17 @@ impl TranscriptOverlay { state: TranscriptOverlayState, ) -> Self { let highlight_cell = Rc::new(StdCell::new(state.highlight_cell)); + let renderables = Self::render_cells( + &transcript_cells, + Rc::clone(&highlight_cell), + state.render_mode, + ); + let user_prompt_positions = Self::user_prompt_positions(&transcript_cells); + Self { view: { let mut view = PagerView::new( - Self::render_cells( - &transcript_cells, - Rc::clone(&highlight_cell), - state.render_mode, - ), + renderables, "Transcript".to_string(), state.scroll_offset, keymap, @@ -604,7 +651,7 @@ impl TranscriptOverlay { view.show_header_progress = false; view }, - user_prompt_positions: Self::user_prompt_positions(&transcript_cells), + user_prompt_positions, cells: transcript_cells, highlight_cell, render_mode: state.render_mode, @@ -614,6 +661,7 @@ impl TranscriptOverlay { scroll_selected_user_cell: None, footer_status: None, live_tail_key: None, + async_layout: AsyncTranscriptLayout::new(), is_done: false, } } @@ -687,6 +735,7 @@ impl TranscriptOverlay { let had_prior_cells = !self.cells.is_empty(); let tail_renderable = self.take_live_tail_renderable(); self.cells.push(cell); + self.invalidate_async_layout(); self.user_prompt_positions = Self::user_prompt_positions(&self.cells); self.view.renderables = Self::render_cells( &self.cells, @@ -726,6 +775,7 @@ impl TranscriptOverlay { pub(crate) fn replace_cells(&mut self, cells: Vec>) { let follow_bottom = self.view.is_scrolled_to_bottom(); self.cells = cells; + self.invalidate_async_layout(); self.user_prompt_positions = Self::user_prompt_positions(&self.cells); if self .highlight_cell @@ -771,6 +821,7 @@ impl TranscriptOverlay { } self.cells .splice(clamped_start..clamped_end, std::iter::once(consolidated)); + self.invalidate_async_layout(); self.user_prompt_positions = Self::user_prompt_positions(&self.cells); if self .highlight_cell @@ -812,6 +863,7 @@ impl TranscriptOverlay { }); if self.live_tail_key == next_key { + self.install_measured_layout_if_ready(width); return; } let follow_bottom = self.view.is_scrolled_to_bottom(); @@ -833,6 +885,82 @@ impl TranscriptOverlay { if follow_bottom { self.view.scroll_offset = usize::MAX; } + self.install_measured_layout_if_ready(width); + } + + pub(crate) fn ensure_async_layout(&mut self, width: u16, app_event_tx: AppEventSender) { + if !self.uses_async_layout() { + return; + } + let key = self.layout_key(width); + if self + .async_layout + .measured + .as_ref() + .is_some_and(|result| result.key == key) + { + self.install_measured_layout_if_ready(width); + return; + } + if self.async_layout.pending == Some(key) { + return; + } + self.async_layout.pending = Some(key); + self.async_layout.measured = None; + spawn_transcript_layout_worker(key, self.cells.clone(), app_event_tx); + } + + pub(crate) fn apply_layout_result(&mut self, result: TranscriptLayoutResult) -> bool { + if result.key != self.layout_key(result.key.width) { + return false; + } + let width = result.key.width; + self.async_layout.pending = None; + self.async_layout.measured = Some(result); + self.install_measured_layout_if_ready(width); + true + } + + fn uses_async_layout(&self) -> bool { + self.cells.len() >= ASYNC_TRANSCRIPT_LAYOUT_CELL_THRESHOLD + } + + fn layout_key(&self, width: u16) -> TranscriptLayoutKey { + TranscriptLayoutKey { + generation: self.async_layout.generation, + width, + render_mode: self.render_mode, + cell_count: self.cells.len(), + } + } + + fn is_waiting_for_async_layout(&self, width: u16) -> bool { + if !self.uses_async_layout() { + return false; + } + let Some(layout) = self.view.layout.as_ref() else { + return true; + }; + layout.width != width || layout.heights.len() != self.view.renderables.len() + } + + fn install_measured_layout_if_ready(&mut self, width: u16) -> bool { + let Some(result) = self.async_layout.measured.as_ref() else { + return false; + }; + if result.key != self.layout_key(width) { + return false; + } + let mut heights = result.heights.clone(); + for renderable in self.view.renderables.iter().skip(self.cells.len()) { + heights.push(renderable.desired_height(width) as usize); + } + self.view.install_precomputed_layout(width, heights); + true + } + + fn invalidate_async_layout(&mut self) { + self.async_layout.invalidate(); } pub(crate) fn set_highlight_cell(&mut self, cell: Option) { @@ -910,6 +1038,7 @@ impl TranscriptOverlay { HistoryRenderMode::Rich => HistoryRenderMode::Raw, HistoryRenderMode::Raw => HistoryRenderMode::Rich, }; + self.invalidate_async_layout(); self.rebuild_renderables(); } @@ -994,6 +1123,9 @@ impl TranscriptOverlay { top_h, ); let content_area = self.view.content_area(top); + if self.is_waiting_for_async_layout(content_area.width) { + return Ok(()); + } let before = self.view.clamped_scroll_offset(content_area); if !self.view.handle_key_event(tui, key_event)? { return Ok(()); @@ -1033,6 +1165,26 @@ impl TranscriptOverlay { first_fitting_right_label(width, &labels) } + fn loading_progress_label(&self, width: u16) -> String { + let total = self.user_prompt_positions.len(); + let selected = self + .highlight_cell + .get() + .and_then(|highlight_cell| { + let selected = self + .user_prompt_positions + .partition_point(|idx| *idx <= highlight_cell); + (selected > 0).then_some(selected) + }) + .unwrap_or(total); + let labels = [ + format!(" {selected} / {total} · loading... "), + format!(" {selected}/{total} · loading... "), + " loading... ".to_string(), + ]; + first_fitting_right_label(width, &labels) + } + fn rebuild_renderables(&mut self) { let tail_renderable = self.take_live_tail_renderable(); self.view.renderables = Self::render_cells( @@ -1219,6 +1371,11 @@ impl TranscriptOverlay { let bottom = Rect::new(area.x, area.y + top_h, area.width, 3); self.view.title = self.header_title(); let content_area = self.view.content_area(top); + self.install_measured_layout_if_ready(content_area.width); + if self.is_waiting_for_async_layout(content_area.width) { + self.render_loading(top, bottom, buf); + return; + } let total_len = self.view.content_height(content_area.width); self.view.resolve_pending_scroll(content_area, total_len); self.view.footer_separator_label = @@ -1226,6 +1383,26 @@ impl TranscriptOverlay { self.view.render(top, buf); self.render_hints(bottom, buf); } + + fn render_loading(&mut self, top: Rect, bottom: Rect, buf: &mut Buffer) { + Clear.render(top, buf); + let header = Rect::new(top.x, top.y, top.width, /*height*/ 1); + render_footer_separator(header, buf, String::new()); + format!(" {} ", self.header_title()) + .dim() + .render_ref(header, buf); + + let content_area = self.view.content_area(top); + if content_area.height > 0 { + Paragraph::new("Loading transcript...".dim()) + .wrap(Wrap { trim: false }) + .render(content_area, buf); + } + + let separator = Rect::new(top.x, content_area.bottom(), top.width, /*height*/ 1); + render_footer_separator(separator, buf, self.loading_progress_label(top.width)); + self.render_hints(bottom, buf); + } } impl TranscriptOverlay { @@ -1537,6 +1714,29 @@ mod tests { cells } + fn large_test_cells() -> Vec> { + (0..ASYNC_TRANSCRIPT_LAYOUT_CELL_THRESHOLD) + .map(|idx| { + Arc::new(TestCell { + lines: vec![Line::from(format!("line-{idx}"))], + }) as Arc + }) + .collect() + } + + fn measured_layout_for( + overlay: &TranscriptOverlay, + width: u16, + height: usize, + ) -> TranscriptLayoutResult { + let cell_count = overlay.committed_cell_count(); + TranscriptLayoutResult { + key: overlay.layout_key(width), + heights: vec![height; cell_count], + total_height: cell_count.saturating_mul(height), + } + } + fn session_info_cell(cwd: &str) -> Arc { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() @@ -2098,6 +2298,89 @@ mod tests { assert_eq!(second_selection, vec![false, true]); } + #[test] + fn large_transcript_renders_loading_until_layout_is_ready() { + let mut overlay = transcript_overlay(large_test_cells()); + let area = Rect::new( + /*x*/ 0, /*y*/ 0, /*width*/ 80, /*height*/ 10, + ); + let mut buf = Buffer::empty(area); + + overlay.render(area, &mut buf); + + assert!(overlay.view.layout.is_none()); + assert_snapshot!( + "large_transcript_renders_loading_until_layout_is_ready", + buffer_to_text(&buf, area) + ); + } + + #[test] + fn large_transcript_applies_matching_precomputed_layout() { + let mut overlay = transcript_overlay(large_test_cells()); + let area = Rect::new( + /*x*/ 0, /*y*/ 0, /*width*/ 80, /*height*/ 10, + ); + let result = measured_layout_for(&overlay, /*width*/ 80, /*height*/ 1); + let committed_cell_count = overlay.committed_cell_count(); + + assert!(overlay.apply_layout_result(result)); + + let layout = overlay.view.layout(/*width*/ 80); + assert_eq!(layout.heights.len(), committed_cell_count); + assert_eq!(layout.total_height, committed_cell_count); + + let mut buf = Buffer::empty(area); + overlay.render(area, &mut buf); + assert!(!buffer_to_text(&buf, area).contains("Loading transcript")); + } + + #[test] + fn large_transcript_ignores_stale_layout_result() { + let mut overlay = transcript_overlay(large_test_cells()); + let mut result = measured_layout_for(&overlay, /*width*/ 80, /*height*/ 1); + result.key.generation = result.key.generation.saturating_add(/*rhs*/ 1); + + assert!(!overlay.apply_layout_result(result)); + assert!(overlay.view.layout.is_none()); + } + + #[test] + fn prompt_selection_anchors_after_async_layout_arrives() { + let cells: Vec> = (0..ASYNC_TRANSCRIPT_LAYOUT_CELL_THRESHOLD) + .map(|idx| user_cell(&format!("prompt {idx}"))) + .collect(); + let mut overlay = transcript_overlay(cells); + let selected = overlay.committed_cell_count() / 2; + let area = Rect::new( + /*x*/ 0, /*y*/ 0, /*width*/ 80, /*height*/ 15, + ); + let top = Rect::new( + area.x, + area.y, + area.width, + area.height.saturating_sub(/*rhs*/ 3), + ); + let content_area = overlay.view.content_area(top); + + overlay.set_highlight_cell(Some(selected)); + overlay.render(area, &mut Buffer::empty(area)); + assert!(overlay.view.layout.is_none()); + + let result = measured_layout_for(&overlay, /*width*/ 80, /*height*/ 3); + assert!(overlay.apply_layout_result(result)); + overlay.render(area, &mut Buffer::empty(area)); + + let selected_row = { + let layout = overlay.view.layout(content_area.width); + layout.offsets[selected].saturating_add(overlay.prompt_first_text_row_offset(selected)) + }; + assert_eq!( + selected_row.saturating_sub(overlay.view.scroll_offset), + (content_area.height as usize) / 3, + ); + } + fn prompt_marker_reversed_states(buf: &Buffer, area: Rect) -> Vec { let mut states = Vec::new(); for y in area.y..area.bottom() { diff --git a/codex-rs/tui/src/session_log.rs b/codex-rs/tui/src/session_log.rs index 369e19d7fd..a0ce83da1a 100644 --- a/codex-rs/tui/src/session_log.rs +++ b/codex-rs/tui/src/session_log.rs @@ -151,6 +151,19 @@ pub(crate) fn log_inbound_app_event(event: &AppEvent) { }); LOGGER.write_json_line(value); } + AppEvent::TranscriptLayoutReady(result) => { + let value = json!({ + "ts": now_ts(), + "dir": "to_tui", + "kind": "transcript_layout_ready", + "generation": result.key.generation, + "width": result.key.width, + "cell_count": result.key.cell_count, + "height_count": result.heights.len(), + "total_height": result.total_height, + }); + LOGGER.write_json_line(value); + } AppEvent::StartFileSearch(query) => { let value = json!({ "ts": now_ts(), diff --git a/codex-rs/tui/src/snapshots/codex_tui__pager_overlay__tests__large_transcript_renders_loading_until_layout_is_ready.snap b/codex-rs/tui/src/snapshots/codex_tui__pager_overlay__tests__large_transcript_renders_loading_until_layout_is_ready.snap new file mode 100644 index 0000000000..1309c47b70 --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__pager_overlay__tests__large_transcript_renders_loading_until_layout_is_ready.snap @@ -0,0 +1,13 @@ +--- +source: tui/src/pager_overlay.rs +expression: "buffer_to_text(&buf, area)" +--- + Transcript ──────────────────────────────────────────────────────────────────── +Loading transcript... + + + + +─────────────────────────────────────────────────────────── 0 / 0 · loading... ─ + ↑/↓ scroll ←/→ prompts pgup/pgdn page home/end jump + q quit ctrl + o copy ⌥ + r raw esc/← prev diff --git a/codex-rs/tui/src/transcript_layout.rs b/codex-rs/tui/src/transcript_layout.rs new file mode 100644 index 0000000000..2139be34a0 --- /dev/null +++ b/codex-rs/tui/src/transcript_layout.rs @@ -0,0 +1,158 @@ +use std::fmt; +use std::sync::Arc; +use std::thread; + +use crate::app_event::AppEvent; +use crate::app_event_sender::AppEventSender; +use crate::history_cell::HistoryCell; +use crate::history_cell::HistoryRenderMode; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct TranscriptLayoutKey { + pub(crate) generation: u64, + pub(crate) width: u16, + pub(crate) render_mode: HistoryRenderMode, + pub(crate) cell_count: usize, +} + +pub(crate) struct TranscriptLayoutResult { + pub(crate) key: TranscriptLayoutKey, + pub(crate) heights: Vec, + pub(crate) total_height: usize, +} + +impl fmt::Debug for TranscriptLayoutResult { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TranscriptLayoutResult") + .field("key", &self.key) + .field("height_count", &self.heights.len()) + .field("total_height", &self.total_height) + .finish() + } +} + +pub(crate) fn spawn_transcript_layout_worker( + key: TranscriptLayoutKey, + cells: Vec>, + app_event_tx: AppEventSender, +) { + thread::spawn(move || { + let heights = measure_transcript_heights(&cells, key.width, key.render_mode); + let total_height = heights.iter().copied().sum(); + app_event_tx.send(AppEvent::TranscriptLayoutReady(TranscriptLayoutResult { + key, + heights, + total_height, + })); + }); +} + +fn measure_transcript_heights( + cells: &[Arc], + width: u16, + render_mode: HistoryRenderMode, +) -> Vec { + let worker_count = layout_worker_count(cells.len()); + if worker_count <= 1 { + return measure_transcript_height_range(cells, width, render_mode, /*start*/ 0); + } + + let chunk_size = cells.len().div_ceil(worker_count); + let mut chunks = thread::scope(|scope| { + let mut handles = Vec::with_capacity(worker_count); + for (chunk_idx, chunk) in cells.chunks(chunk_size).enumerate() { + let start = chunk_idx.saturating_mul(chunk_size); + handles.push(scope.spawn(move || { + ( + start, + measure_transcript_height_range(chunk, width, render_mode, start), + ) + })); + } + handles + .into_iter() + .map(|handle| handle.join().unwrap_or_else(|_| (0, Vec::new()))) + .collect::>() + }); + chunks.sort_by_key(|(start, _)| *start); + + let mut heights = Vec::with_capacity(cells.len()); + for (_, chunk_heights) in chunks { + heights.extend(chunk_heights); + } + heights +} + +fn measure_transcript_height_range( + cells: &[Arc], + width: u16, + render_mode: HistoryRenderMode, + start: usize, +) -> Vec { + cells + .iter() + .enumerate() + .map(|(offset, cell)| { + let idx = start.saturating_add(offset); + let spacing = usize::from(idx > 0 && !cell.is_stream_continuation()); + spacing.saturating_add(cell.desired_transcript_height_for_mode(width, render_mode) as usize) + }) + .collect() +} + +fn layout_worker_count(cell_count: usize) -> usize { + if cell_count == 0 { + return 1; + } + let parallelism = thread::available_parallelism().map_or(/*default*/ 1, usize::from); + parallelism.min(/*other*/ 4).min(cell_count) +} + +#[cfg(test)] +mod tests { + use super::*; + use ratatui::text::Line; + + #[derive(Debug)] + struct TestCell { + lines: Vec>, + stream_continuation: bool, + } + + impl HistoryCell for TestCell { + fn display_lines(&self, _width: u16) -> Vec> { + self.lines.clone() + } + + fn raw_lines(&self) -> Vec> { + self.lines.clone() + } + + fn is_stream_continuation(&self) -> bool { + self.stream_continuation + } + } + + #[test] + fn measures_transcript_heights_with_inter_cell_spacing() { + let cells: Vec> = vec![ + Arc::new(TestCell { + lines: vec![Line::from("first")], + stream_continuation: false, + }), + Arc::new(TestCell { + lines: vec![Line::from("second")], + stream_continuation: false, + }), + Arc::new(TestCell { + lines: vec![Line::from("continuation")], + stream_continuation: true, + }), + ]; + + let heights = + measure_transcript_heights(&cells, /*width*/ 80, HistoryRenderMode::Rich); + + assert_eq!(heights, vec![1, 2, 1]); + } +}