perf(tui): lazily seed virtual history heights (#16523)

ui-tui/src/__tests__/useVirtualHistoryHeights.test.ts

@@ -0,0 +1,39 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import { ensureVirtualItemHeight } from '../hooks/useVirtualHistory.js'
+
+describe('ensureVirtualItemHeight', () => {
+  it('reuses cached heights without invoking the estimator', () => {
+    const heights = new Map([['a', 7]])
+    const estimateHeight = vi.fn(() => 99)
+
+    expect(ensureVirtualItemHeight(heights, 'a', 0, 4, estimateHeight)).toBe(7)
+    expect(estimateHeight).not.toHaveBeenCalled()
+    expect(heights.get('a')).toBe(7)
+  })
+
+  it('lazily seeds missing heights from the estimator', () => {
+    const heights = new Map<string, number>()
+    const estimateHeight = vi.fn((index: number) => 10 + index)
+
+    expect(ensureVirtualItemHeight(heights, 'b', 2, 4, estimateHeight)).toBe(12)
+    expect(estimateHeight).toHaveBeenCalledTimes(1)
+    expect(estimateHeight).toHaveBeenCalledWith(2, 'b')
+    expect(heights.get('b')).toBe(12)
+  })
+
+  it('falls back to the default estimate when no estimator is provided', () => {
+    const heights = new Map<string, number>()
+
+    expect(ensureVirtualItemHeight(heights, 'c', 0, 4)).toBe(4)
+    expect(heights.get('c')).toBe(4)
+  })
+
+  it('normalizes non-positive estimates to a minimum of one row', () => {
+    const heights = new Map<string, number>()
+    const estimateHeight = vi.fn(() => 0)
+
+    expect(ensureVirtualItemHeight(heights, 'd', 0, 0, estimateHeight)).toBe(1)
+    expect(heights.get('d')).toBe(1)
+  })
+})

ui-tui/src/app/useMainApp.ts

@@ -218,23 +218,15 @@ export function useMainApp(gw: GatewayClient) {
     return cache
   }, [heightCacheKey])
 
-  const initialHeights = useMemo(() => {
+  const estimateRowHeight = useCallback(
-    const out = new Map<string, number>()
+    (index: number) =>
-
+      estimatedMsgHeight(virtualRows[index]!.msg, cols, {
-    for (const row of virtualRows) {
+        compact: ui.compact,
-      out.set(
+        details: detailsVisible,
-        row.key,
+        limitHistory: index < virtualRows.length - FULL_RENDER_TAIL_ITEMS
-        heightCache.get(row.key) ??
+      }),
-          estimatedMsgHeight(row.msg, cols, {
+    [cols, detailsVisible, ui.compact, virtualRows]
-            compact: ui.compact,
+  )
-            details: detailsVisible,
-            limitHistory: row.index < virtualRows.length - FULL_RENDER_TAIL_ITEMS
-          })
-      )
-    }
-
-    return out
-  }, [cols, detailsVisible, heightCache, ui.compact, virtualRows])
 
   const syncHeightCache = useCallback(
     (heights: ReadonlyMap<string, number>) => {
@@ -250,7 +242,8 @@ export function useMainApp(gw: GatewayClient) {
   )
 
   const virtualHistory = useVirtualHistory(scrollRef, virtualRows, cols, {
-    initialHeights,
+    estimateHeight: estimateRowHeight,
+    initialHeights: heightCache,
     liveTailActive: turnLiveTailActive,
     onHeightsChange: syncHeightCache
   })

ui-tui/src/hooks/useVirtualHistory.ts

@@ -76,12 +76,32 @@ export const shouldSetVirtualClamp = ({
   viewportHeight: number
 }) => itemCount > 0 && viewportHeight > 0 && !sticky && !liveTailActive
 
+export const ensureVirtualItemHeight = (
+  heights: Map<string, number>,
+  key: string,
+  index: number,
+  estimate: number,
+  estimateHeight?: (index: number, key: string) => number
+) => {
+  const cached = heights.get(key)
+
+  if (cached !== undefined) {
+    return Math.max(1, Math.floor(cached))
+  }
+
+  const seeded = Math.max(1, Math.floor(estimateHeight?.(index, key) ?? estimate))
+  heights.set(key, seeded)
+
+  return seeded
+}
+
 export function useVirtualHistory(
   scrollRef: RefObject<ScrollBoxHandle | null>,
   items: readonly { key: string }[],
   columns: number,
   {
     estimate = ESTIMATE,
+    estimateHeight,
     initialHeights,
     liveTailActive = false,
     onHeightsChange,
@@ -208,7 +228,7 @@ export function useVirtualHistory(
     arr[0] = 0
 
     for (let i = 0; i < n; i++) {
-      arr[i + 1] = arr[i]! + Math.max(1, Math.floor(heights.current.get(items[i]!.key) ?? estimate))
+      arr[i + 1] = arr[i]! + ensureVirtualItemHeight(heights.current, items[i]!.key, i, estimate, estimateHeight)
     }
 
     offsetsCache.current = { arr, n, version: offsetVersion.current }
@@ -280,7 +300,7 @@ export function useVirtualHistory(
     let coverage = 0
 
     for (let i = start; i < end; i++) {
-      coverage += heights.current.get(items[i]!.key) ?? PESSIMISTIC
+      coverage += ensureVirtualItemHeight(heights.current, items[i]!.key, i, PESSIMISTIC, estimateHeight)
     }
 
     if (sticky) {
@@ -288,13 +308,13 @@ export function useVirtualHistory(
 
       while (start > minStart && coverage < needed) {
         start--
-        coverage += heights.current.get(items[start]!.key) ?? PESSIMISTIC
+        coverage += ensureVirtualItemHeight(heights.current, items[start]!.key, start, PESSIMISTIC, estimateHeight)
       }
     } else {
       const maxEnd = Math.min(n, start + maxMounted)
 
       while (end < maxEnd && coverage < needed) {
-        coverage += heights.current.get(items[end]!.key) ?? PESSIMISTIC
+        coverage += ensureVirtualItemHeight(heights.current, items[end]!.key, end, PESSIMISTIC, estimateHeight)
         end++
       }
     }
@@ -498,6 +518,7 @@ interface MeasuredNode {
 interface VirtualHistoryOptions {
   coldStartCount?: number
   estimate?: number
+  estimateHeight?: (index: number, key: string) => number
   initialHeights?: ReadonlyMap<string, number>
   liveTailActive?: boolean
   maxMounted?: number