Precalc viewport window for virtualization

packages/ui/src/components/message-stream-v2.tsx

@@ -28,6 +28,9 @@ const USER_BORDER_COLOR = "var(--message-user-border)"
 const ASSISTANT_BORDER_COLOR = "var(--message-assistant-border)"
 const TOOL_BORDER_COLOR = "var(--message-tool-border)"
 const VIRTUAL_ITEM_MARGIN_PX = 800
+const ESTIMATED_MESSAGE_HEIGHT = 320
+const INITIAL_FORCE_MIN_ITEMS = 12
+const INITIAL_FORCE_OVERSCAN = 6
 
 type ToolCallPart = Extract<ClientPart, { type: "tool" }>
 
@@ -299,11 +302,38 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
   })
  
   const [scrollElement, setScrollElement] = createSignal<HTMLDivElement | undefined>()
+  const [initialForceActive, setInitialForceActive] = createSignal(true)
+  const [initialForceInitialized, setInitialForceInitialized] = createSignal(false)
+  const [initialForceStartIndex, setInitialForceStartIndex] = createSignal(0)
+  const [initialForceRemaining, setInitialForceRemaining] = createSignal(0)
   const [autoScroll, setAutoScroll] = createSignal(true)
 
-  const [showScrollTopButton, setShowScrollTopButton] = createSignal(false)
+  createEffect(() => {
+    props.instanceId
+    props.sessionId
+    setInitialForceActive(true)
+    setInitialForceInitialized(false)
+    setInitialForceStartIndex(0)
+    setInitialForceRemaining(0)
+  })
 
+  createEffect(() => {
+    if (!initialForceActive() || initialForceInitialized()) return
+    const ids = messageIds()
+    if (ids.length === 0) return
+    const viewportHeight = scrollElement()?.clientHeight ?? (typeof window !== "undefined" ? window.innerHeight : 800)
+    const estimatedCount = Math.min(
+      ids.length,
+      Math.max(INITIAL_FORCE_MIN_ITEMS, Math.ceil(viewportHeight / ESTIMATED_MESSAGE_HEIGHT) + INITIAL_FORCE_OVERSCAN),
+    )
+    setInitialForceStartIndex(Math.max(0, ids.length - estimatedCount))
+    setInitialForceRemaining(estimatedCount)
+    setInitialForceInitialized(true)
+  })
+
+  const [showScrollTopButton, setShowScrollTopButton] = createSignal(false)
   const [showScrollBottomButton, setShowScrollBottomButton] = createSignal(false)
+
   let containerRef: HTMLDivElement | undefined
   let lastKnownScrollTop = 0
   let lastMeasuredScrollHeight = 0
@@ -649,13 +679,28 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
         </Show>
 
         <Index each={messageIds()}>
-          {(messageId) => (
+          {(messageId) => {
+            const messageIndex = () => messageIndexMap().get(messageId()) ?? 0
+            const forceVisible = () => initialForceActive() && messageIndex() >= initialForceStartIndex()
+            const handleMeasured = () => {
+              if (!forceVisible()) return
+              setInitialForceRemaining((value) => {
+                const next = value > 0 ? value - 1 : 0
+                if (next === 0) {
+                  setInitialForceActive(false)
+                }
+                return next
+              })
+            }
+            return (
               <VirtualItem
                 cacheKey={messageId()}
                 scrollContainer={scrollElement}
                 threshold={VIRTUAL_ITEM_MARGIN_PX}
                 placeholderClass="message-stream-placeholder"
                 virtualizationEnabled={() => !props.loading}
+                forceVisible={forceVisible}
+                onMeasured={handleMeasured}
               >
                 <MessageBlock
                   messageId={messageId()}
@@ -672,7 +717,8 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
                   onContentRendered={handleContentRendered}
                 />
               </VirtualItem>
-          )}
+            )
+          }}
         </Index>
       </div>
 

packages/ui/src/components/virtual-item.tsx

@@ -97,12 +97,17 @@ interface VirtualItemProps {
   contentClass?: string
   placeholderClass?: string
   virtualizationEnabled?: Accessor<boolean>
+  forceVisible?: Accessor<boolean>
+  onMeasured?: () => void
 }
 
 export default function VirtualItem(props: VirtualItemProps) {
   const resolved = resolveChildren(() => props.children)
+  const cachedHeight = sizeCache.get(props.cacheKey)
   const [isIntersecting, setIsIntersecting] = createSignal(true)
-  const [measuredHeight, setMeasuredHeight] = createSignal(sizeCache.get(props.cacheKey) ?? 0)
+  const [measuredHeight, setMeasuredHeight] = createSignal(cachedHeight ?? 0)
+  const [hasMeasured, setHasMeasured] = createSignal(cachedHeight !== undefined)
+  let hasReportedMeasurement = Boolean(cachedHeight && cachedHeight > 0)
   let pendingVisibility: boolean | null = null
   let visibilityFrame: number | null = null
   const flushVisibility = () => {
@@ -126,7 +131,6 @@ export default function VirtualItem(props: VirtualItemProps) {
       }
     })
   }
-  const [hasMeasured, setHasMeasured] = createSignal(sizeCache.has(props.cacheKey))
   const virtualizationEnabled = () => (props.virtualizationEnabled ? props.virtualizationEnabled() : true)
 
   let wrapperRef: HTMLDivElement | undefined
@@ -156,6 +160,10 @@ export default function VirtualItem(props: VirtualItemProps) {
     if (normalized > 0) {
       sizeCache.set(props.cacheKey, normalized)
       setHasMeasured(true)
+      if (!hasReportedMeasurement) {
+        hasReportedMeasurement = true
+        props.onMeasured?.()
+      }
     }
     setMeasuredHeight(normalized)
   }
@@ -230,6 +238,7 @@ export default function VirtualItem(props: VirtualItemProps) {
   })
 
   const shouldHideContent = createMemo(() => {
+    if (props.forceVisible?.()) return false
     if (!virtualizationEnabled()) return false
     if (!hasMeasured()) return false
     return !isIntersecting()