perf(mothership): virtualize chat transcript and isolate input from stream re-renders (#5019)

* perf(mothership): virtualize chat transcript and isolate input from stream re-renders Long chats rendered every message into the DOM with no windowing — a custom rAF "progressive list" only smeared the mount cost across frames without capping it. At ~1000 messages this was 52k DOM nodes and a 21s main-thread block on open, and the input toolbar re-rendered on every streamed token. - Virtualize the message list with @tanstack/react-virtual using dynamic measureElement, stable per-row keys, and a tuned size estimate. Only the visible window mounts, so load cost is now flat regardless of transcript length. Remove the now-redundant useProgressiveList hook. - Memoize UserInput and stabilize its callbacks (useCallback in MothershipChat and home) so streaming ticks no longer re-render the entire input toolbar. - Keep the existing useAutoScroll for streaming stick-to-bottom (it reads the virtualizer's real scrollHeight) and add a per-chat scrollToIndex for initial positioning before paint. Measured on a cloned 1032-message chat: time-to-rendered 26.3s -> 1.7s, main-thread blocked 21.4s -> 0.8s, DOM nodes 52k -> 1.4k, typing-while- streaming p-max 104ms -> 26ms. Adds scripts/perf/ harness used to validate. * address review: pending-chat scroll, flash on tail unmount, empty-row gap, per-role estimate - Pending-chat initial scroll (Cursor, high): seed the scrolled-chat guard with a unique sentinel so a not-yet-persisted chat (undefined chatId) with messages still scrolls to bottom instead of being treated as already-scrolled. - Streaming-row flash (review of virtualization): pin the last row in the rendered window via rangeExtractor so scrolling it out of the overscan window and back mid-stream can't unmount/remount it and re-fire the reveal fade. - Empty assistant row gap (Greptile): move the row gap from the virtual-item wrapper into the row content so a null-rendering (finalised, empty) assistant turn collapses to zero height instead of leaving a pb-6 blank slot. - Per-role row-height estimate instead of a single blended constant, so the scrollbar drifts less as off-screen rows resolve. - Drop the scripts/perf harness from the PR. * fix(mothership): preserve user-row top spacing in virtualized layout Restoring pt-3/pt-2 on the user row keeps the exact inter-row rhythm from the old space-y-6 layout: assistant→user gaps stay 24+12px and user→assistant stay 24px, instead of becoming uniform when the gap moved to per-row pb. * fix(mothership): don't re-scroll when a pending chat persists its id The per-chat initial-scroll guard treated undefined→persisted-id as a chat switch and scrolled to the bottom again, yanking the viewport down if the user had scrolled up mid-stream. Treat that transition as the same conversation: adopt the id without re-scrolling. Genuine chat switches (id→different id) still re-scroll.
2026-09-24 15:45:35 +08:00 · 2026-06-12 23:05:57 -07:00
parent 26224b8edb
commit 005fa1027b
4 changed files with 194 additions and 227 deletions
@@ -1,6 +1,7 @@
 'use client'

 import { memo, useCallback, useEffect, useLayoutEffect, useMemo, useRef } from 'react'
+import { defaultRangeExtractor, type Range, useVirtualizer } from '@tanstack/react-virtual'
 import { cn } from '@/lib/core/utils/cn'
 import { MessageActions } from '@/app/workspace/[workspaceId]/components'
 import { ChatMessageAttachments } from '@/app/workspace/[workspaceId]/home/components/chat-message-attachments'
@@ -26,7 +27,6 @@ import type {
  QueuedMessage,
 } from '@/app/workspace/[workspaceId]/home/types'
 import { useAutoScroll } from '@/hooks/use-auto-scroll'
-import { useProgressiveList } from '@/hooks/use-progressive-list'
 import type { ChatContext } from '@/stores/panel'
 import { MothershipChatSkeleton } from './components/mothership-chat-skeleton'

@@ -61,11 +61,38 @@ interface MothershipChatProps {
  className?: string
 }

+/**
+ * Per-role row-height estimates seed the virtualizer before each row is measured.
+ * They only size the scrollbar for not-yet-rendered rows — every visible row is
+ * measured precisely via `measureElement` — so approximate values suffice. Split
+ * by role because user bubbles are short and assistant turns are tall; a single
+ * blended number would over/under-shoot both and drift the scrollbar more.
+ */
+const ROW_HEIGHT_ESTIMATE = {
+  'mothership-view': { user: 64, assistant: 280 },
+  'copilot-view': { user: 48, assistant: 180 },
+} as const
+
+/**
+ * Rows render farther beyond the viewport edges than the default so fast scroll
+ * and the streaming tail stay painted without a blank flash before measurement.
+ */
+const OVERSCAN = 6
+
+/**
+ * Initial-scroll sentinel. Distinct from every real `chatId` value — including
+ * `undefined` (a not-yet-persisted chat) — so the first scroll-to-bottom fires
+ * even before a chat has an id, instead of treating `undefined` as "already
+ * scrolled this chat".
+ */
+const UNSCROLLED = Symbol('unscrolled')
+
 const LAYOUT_STYLES = {
  'mothership-view': {
    scrollContainer:
      'min-h-0 flex-1 overflow-y-auto overflow-x-hidden px-6 pt-4 pb-8 [scrollbar-gutter:stable_both-edges]',
-    content: 'mx-auto max-w-[48rem] space-y-6',
+    sizer: 'relative mx-auto w-full max-w-[48rem]',
+    rowGap: 'pb-6',
    userRow: 'flex flex-col items-end gap-[6px] pt-3',
    attachmentWidth: 'max-w-[70%]',
    userBubble: 'max-w-[70%] overflow-hidden rounded-[16px] bg-[var(--surface-5)] px-3.5 py-2',
@@ -75,7 +102,8 @@ const LAYOUT_STYLES = {
  },
  'copilot-view': {
    scrollContainer: 'min-h-0 flex-1 overflow-y-auto overflow-x-hidden px-3 pt-2 pb-4',
-    content: 'space-y-4',
+    sizer: 'relative w-full',
+    rowGap: 'pb-4',
    userRow: 'flex flex-col items-end gap-[6px] pt-2',
    attachmentWidth: 'max-w-[85%]',
    userBubble: 'max-w-[85%] overflow-hidden rounded-[16px] bg-[var(--surface-5)] px-3 py-2',
@@ -201,14 +229,41 @@ export function MothershipChat({
 }: MothershipChatProps) {
  const styles = LAYOUT_STYLES[layout]
  const isStreamActive = isSending || isReconnecting
-  const { ref: scrollContainerRef, scrollToBottom } = useAutoScroll(isStreamActive, {
-    scrollOnMount: true,
-  })
+  const scrollElementRef = useRef<HTMLDivElement | null>(null)
+  const { ref: autoScrollRef } = useAutoScroll(isStreamActive)
+  const setScrollElement = useCallback(
+    (el: HTMLDivElement | null) => {
+      scrollElementRef.current = el
+      autoScrollRef(el)
+    },
+    [autoScrollRef]
+  )
+
  const hasMessages = messages.length > 0
-  const stagingKey = chatId ?? 'pending-chat'
-  const { staged: stagedMessages, isStaging } = useProgressiveList(messages, stagingKey)
-  const stagedMessageCount = stagedMessages.length
-  const stagedOffset = messages.length - stagedMessages.length
+
+  /**
+   * Stable per-row identity for virtualizer measurement caching and React
+   * reconciliation. User rows key on their message id; assistant rows key on
+   * their turn position (`assistant:<userId>:<ordinal>`) so a streaming
+   * placeholder keeps the same element — and its smooth-text state — when the
+   * persisted message arrives with a new id.
+   */
+  const rowKeyByIndex = useMemo(() => {
+    const out: string[] = []
+    let lastUserId: string | undefined
+    let ordinal = 0
+    for (const [index, message] of messages.entries()) {
+      if (message.role === 'user') {
+        lastUserId = message.id
+        ordinal = 0
+        out[index] = message.id
+      } else {
+        out[index] = lastUserId ? `assistant:${lastUserId}:${ordinal++}` : message.id
+      }
+    }
+    return out
+  }, [messages])
+
  const precedingUserContentByIndex = useMemo(() => {
    const out: Array<string | undefined> = []
    let lastUserContent: string | undefined
@@ -218,22 +273,44 @@ export function MothershipChat({
    }
    return out
  }, [messages])
-  const assistantTurnKeyByIndex = useMemo(() => {
-    const out: string[] = []
-    let lastUserId: string | undefined
-    let ordinal = 0
-    for (const [index, message] of messages.entries()) {
-      if (message.role === 'user') {
-        lastUserId = message.id
-        ordinal = 0
-      } else {
-        out[index] = lastUserId ? `assistant:${lastUserId}:${ordinal++}` : message.id
+
+  /**
+   * Always keep the last row in the rendered window. It is the live/streaming
+   * row; unmounting it (by scrolling far enough up that it leaves the overscan
+   * window) and remounting it mid-stream would reset its smooth-text reveal
+   * state and re-fire the fade-in animation — a visible flash. Pinning it costs
+   * one extra always-mounted row.
+   */
+  const lastIndex = messages.length - 1
+  const rangeExtractor = useCallback(
+    (range: Range) => {
+      const indexes = defaultRangeExtractor(range)
+      if (lastIndex >= 0 && !indexes.includes(lastIndex)) {
+        indexes.push(lastIndex)
      }
-    }
-    return out
-  }, [messages])
-  const initialScrollDoneRef = useRef(false)
+      return indexes
+    },
+    [lastIndex]
+  )
+
+  const virtualizer = useVirtualizer({
+    count: messages.length,
+    getScrollElement: () => scrollElementRef.current,
+    estimateSize: (index) => {
+      const estimate = ROW_HEIGHT_ESTIMATE[layout]
+      return messages[index]?.role === 'user' ? estimate.user : estimate.assistant
+    },
+    overscan: OVERSCAN,
+    getItemKey: (index) => rowKeyByIndex[index] ?? index,
+    rangeExtractor,
+  })
+
+  const scrolledChatRef = useRef<string | undefined | typeof UNSCROLLED>(UNSCROLLED)
  const userInputRef = useRef<UserInputHandle>(null)
+  const messageQueueRef = useRef(messageQueue)
+  useEffect(() => {
+    messageQueueRef.current = messageQueue
+  }, [messageQueue])

  const onSubmitRef = useRef(onSubmit)
  useEffect(() => {
@@ -243,37 +320,48 @@ export function MothershipChat({
    onSubmitRef.current(id)
  }, [])

-  function handleSendQueuedHead() {
-    const topMessage = messageQueue[0]
+  const handleSendQueuedHead = useCallback(() => {
+    const topMessage = messageQueueRef.current[0]
    if (!topMessage) return
    void onSendQueuedMessage(topMessage.id)
-  }
+  }, [onSendQueuedMessage])

-  function handleEditQueued(id: string) {
-    const msg = onEditQueuedMessage(id)
-    if (msg) userInputRef.current?.loadQueuedMessage(msg)
-  }
+  const handleEditQueued = useCallback(
+    (id: string) => {
+      const msg = onEditQueuedMessage(id)
+      if (msg) userInputRef.current?.loadQueuedMessage(msg)
+    },
+    [onEditQueuedMessage]
+  )

-  function handleEditQueuedTail() {
-    const tail = messageQueue[messageQueue.length - 1]
+  const handleEditQueuedTail = useCallback(() => {
+    const tail = messageQueueRef.current[messageQueueRef.current.length - 1]
    if (!tail) return
    handleEditQueued(tail.id)
-  }
+  }, [handleEditQueued])

+  /**
+   * Land at the most recent message once per chat — on open and when switching
+   * chats. The ref tracks which `chatId` we last scrolled for (seeded with
+   * {@link UNSCROLLED} so a pending, id-less chat still scrolls on first mount),
+   * so it re-fires on a genuine chat switch, including between chats of equal
+   * length. A pending chat persisting its id (`undefined` → string) is the SAME
+   * conversation, so adopt the id without re-scrolling — otherwise the viewport
+   * would snap back to the bottom after the user scrolled up mid-stream. Runs
+   * before paint so a long transcript never flashes at the top. Subsequent
+   * growth within the same chat is handled by {@link useAutoScroll}'s streaming
+   * sticky-scroll, not here.
+   */
  useLayoutEffect(() => {
-    if (!hasMessages) {
-      initialScrollDoneRef.current = false
-      return
-    }
-    if (initialScrollDoneRef.current || initialScrollBlocked) return
-    initialScrollDoneRef.current = true
-    scrollToBottom()
-  }, [hasMessages, initialScrollBlocked, scrollToBottom])
+    const scrolledFor = scrolledChatRef.current
+    if (!hasMessages || initialScrollBlocked || scrolledFor === chatId) return
+    const isPendingPersist = scrolledFor === undefined && chatId !== undefined
+    scrolledChatRef.current = chatId
+    if (isPendingPersist) return
+    virtualizer.scrollToIndex(lastIndex, { align: 'end' })
+  }, [chatId, hasMessages, initialScrollBlocked, lastIndex, virtualizer])

-  useLayoutEffect(() => {
-    if (!isStaging || initialScrollBlocked || !initialScrollDoneRef.current) return
-    scrollToBottom()
-  }, [isStaging, stagedMessageCount, initialScrollBlocked, scrollToBottom])
+  const virtualItems = virtualizer.getVirtualItems()

  return (
    <ChatSurfaceProvider
@@ -284,37 +372,42 @@ export function MothershipChat({
      onWorkspaceResourceSelect={onWorkspaceResourceSelect}
    >
      <div className={cn('flex h-full min-h-0 flex-col', className)}>
-        <div ref={scrollContainerRef} className={styles.scrollContainer}>
+        <div ref={setScrollElement} className={styles.scrollContainer}>
          {isLoading && !hasMessages ? (
            <MothershipChatSkeleton layout={layout} />
          ) : (
-            <div className={styles.content}>
-              {stagedMessages.map((msg, localIndex) => {
-                const index = stagedOffset + localIndex
-                if (msg.role === 'user') {
-                  return (
-                    <UserMessageRow
-                      key={msg.id}
-                      content={msg.content}
-                      contexts={msg.contexts}
-                      attachments={msg.attachments}
-                      rowClassName={styles.userRow}
-                      bubbleClassName={styles.userBubble}
-                      attachmentWidthClassName={styles.attachmentWidth}
-                    />
-                  )
-                }
-
-                const isLast = index === messages.length - 1
+            <div className={styles.sizer} style={{ height: virtualizer.getTotalSize() }}>
+              {virtualItems.map((virtualItem) => {
+                const index = virtualItem.index
+                const msg = messages[index]
+                const isLast = index === lastIndex
                return (
-                  <AssistantMessageRow
-                    key={assistantTurnKeyByIndex[index] ?? msg.id}
-                    message={msg}
-                    isStreaming={isStreamActive && isLast}
-                    precedingUserContent={precedingUserContentByIndex[index]}
-                    rowClassName={styles.assistantRow}
-                    onOptionSelect={isLast ? stableOnOptionSelect : undefined}
-                  />
+                  <div
+                    key={virtualItem.key}
+                    data-index={index}
+                    ref={virtualizer.measureElement}
+                    className='absolute top-0 left-0 w-full'
+                    style={{ transform: `translateY(${virtualItem.start}px)` }}
+                  >
+                    {msg.role === 'user' ? (
+                      <UserMessageRow
+                        content={msg.content}
+                        contexts={msg.contexts}
+                        attachments={msg.attachments}
+                        rowClassName={cn(styles.userRow, styles.rowGap)}
+                        bubbleClassName={styles.userBubble}
+                        attachmentWidthClassName={styles.attachmentWidth}
+                      />
+                    ) : (
+                      <AssistantMessageRow
+                        message={msg}
+                        isStreaming={isStreamActive && isLast}
+                        precedingUserContent={precedingUserContentByIndex[index]}
+                        rowClassName={cn(styles.assistantRow, styles.rowGap)}
+                        onOptionSelect={isLast ? stableOnOptionSelect : undefined}
+                      />
+                    )}
+                  </div>
                )
              })}
            </div>
@@ -3,6 +3,7 @@
 import type React from 'react'
 import {
  forwardRef,
+  memo,
  useCallback,
  useEffect,
  useImperativeHandle,
@@ -145,7 +146,7 @@ export interface UserInputHandle {
  populatePrompt: (text: string) => void
 }

-export const UserInput = forwardRef<UserInputHandle, UserInputProps>(function UserInput(
+const UserInputImpl = forwardRef<UserInputHandle, UserInputProps>(function UserInput(
  {
    defaultValue = '',
    draftScopeKey,
@@ -1445,3 +1446,10 @@ export const UserInput = forwardRef<UserInputHandle, UserInputProps>(function Us
    </div>
  )
 })
+
+/**
+ * Memoized so streaming ticks in the parent transcript — which re-render
+ * {@link MothershipChat} on every chunk — do not re-render the entire input
+ * toolbar. Relies on callers passing stable callbacks (see `MothershipChat`).
+ */
+export const UserInput = memo(UserInputImpl)
@@ -233,36 +233,35 @@ export function Home({ chatId, userName, userId, initialResourceId = null }: Hom
    }
  }, [resources, collapseResource])

-  function handleStopGeneration() {
+  const handleStopGeneration = useCallback(() => {
    captureEvent(posthogRef.current, 'task_generation_aborted', {
      workspace_id: workspaceId,
      view: 'mothership',
      request_id: getCurrentRequestId(),
    })
    void stopGeneration().catch(() => {})
-  }
+  }, [workspaceId, getCurrentRequestId, stopGeneration])

-  function handleSubmit(
-    text: string,
-    fileAttachments?: FileAttachmentForApi[],
-    contexts?: ChatContext[]
-  ) {
-    const trimmed = text.trim()
-    if (!trimmed && !(fileAttachments && fileAttachments.length > 0)) return
+  const handleSubmit = useCallback(
+    (text: string, fileAttachments?: FileAttachmentForApi[], contexts?: ChatContext[]) => {
+      const trimmed = text.trim()
+      if (!trimmed && !(fileAttachments && fileAttachments.length > 0)) return

-    captureEvent(posthogRef.current, 'task_message_sent', {
-      workspace_id: workspaceId,
-      has_attachments: !!(fileAttachments && fileAttachments.length > 0),
-      has_contexts: !!(contexts && contexts.length > 0),
-      is_new_task: !chatId,
-    })
+      captureEvent(posthogRef.current, 'task_message_sent', {
+        workspace_id: workspaceId,
+        has_attachments: !!(fileAttachments && fileAttachments.length > 0),
+        has_contexts: !!(contexts && contexts.length > 0),
+        is_new_task: !chatId,
+      })

-    if (initialViewInputRef.current) {
-      setIsInputEntering(true)
-    }
+      if (initialViewInputRef.current) {
+        setIsInputEntering(true)
+      }

-    sendMessage(trimmed || 'Analyze the attached file(s).', fileAttachments, contexts)
-  }
+      sendMessage(trimmed || 'Analyze the attached file(s).', fileAttachments, contexts)
+    },
+    [workspaceId, chatId, sendMessage]
+  )

  useEffect(() => {
    const handler = (e: Event) => {
@@ -1,133 +0,0 @@
-'use client'
-
-import { useEffect, useLayoutEffect, useRef, useState } from 'react'
-
-interface ProgressiveListOptions {
-  /** Number of items to render in the initial batch (most recent items) */
-  initialBatch?: number
-  /** Number of items to add per animation frame */
-  batchSize?: number
-}
-
-const DEFAULTS = {
-  initialBatch: 10,
-  batchSize: 5,
-} satisfies Required<ProgressiveListOptions>
-
-interface ProgressiveListState {
-  key: string
-  count: number
-  caughtUp: boolean
-}
-
-function createInitialState(
-  key: string,
-  itemCount: number,
-  initialBatch: number
-): ProgressiveListState {
-  const count = Math.min(itemCount, initialBatch)
-  return {
-    key,
-    count,
-    caughtUp: itemCount > 0 && count >= itemCount,
-  }
-}
-
-/**
- * Progressively renders a list of items so that first paint is fast.
- *
- * On mount (or when `key` changes), only the most recent `initialBatch`
- * items are rendered. The rest are added in `batchSize` increments via
- * `requestAnimationFrame`.
- *
- * @param items    Full list of items to render.
- * @param key      A session/conversation identifier. When it changes,
- *                 staging restarts for the new list.
- * @param options  Tuning knobs for batch sizes.
- * @returns        The currently staged (visible) subset of items.
- */
-export function useProgressiveList<T>(
-  items: T[],
-  key: string,
-  options?: ProgressiveListOptions
-): { staged: T[]; isStaging: boolean } {
-  const initialBatch = Math.max(0, options?.initialBatch ?? DEFAULTS.initialBatch)
-  const batchSize = Math.max(1, options?.batchSize ?? DEFAULTS.batchSize)
-  const [state, setState] = useState(() => createInitialState(key, items.length, initialBatch))
-  const latestItemCountRef = useRef(items.length)
-
-  useLayoutEffect(() => {
-    latestItemCountRef.current = items.length
-  }, [items.length])
-
-  const renderState =
-    state.key === key && (state.count > 0 || items.length === 0 || state.caughtUp)
-      ? state
-      : createInitialState(key, items.length, initialBatch)
-
-  useEffect(() => {
-    setState((prev) => {
-      if (prev.key !== key) {
-        return createInitialState(key, items.length, initialBatch)
-      }
-
-      if (items.length === 0) {
-        if (prev.count === 0 && !prev.caughtUp) {
-          return prev
-        }
-        return { key, count: 0, caughtUp: false }
-      }
-
-      if (prev.caughtUp) {
-        if (prev.count === items.length) {
-          return prev
-        }
-        return { key, count: items.length, caughtUp: true }
-      }
-
-      const minimumCount = Math.min(items.length, initialBatch)
-      if (prev.count >= minimumCount && prev.count <= items.length) {
-        return prev
-      }
-
-      const count = Math.min(items.length, Math.max(prev.count, minimumCount))
-      return {
-        key,
-        count,
-        caughtUp: count >= items.length,
-      }
-    })
-  }, [key, items.length, initialBatch])
-
-  useEffect(() => {
-    if (state.key !== key || state.caughtUp || state.count >= items.length) {
-      return
-    }
-
-    const frame = requestAnimationFrame(() => {
-      setState((prev) => {
-        if (prev.key !== key || prev.caughtUp) {
-          return prev
-        }
-
-        const itemCount = latestItemCountRef.current
-        const count = Math.min(itemCount, prev.count + batchSize)
-        return {
-          key,
-          count,
-          caughtUp: count >= itemCount,
-        }
-      })
-    })
-
-    return () => cancelAnimationFrame(frame)
-  }, [state.key, state.count, state.caughtUp, key, items.length, batchSize])
-
-  const effectiveCount = renderState.caughtUp
-    ? items.length
-    : Math.min(renderState.count, items.length)
-  const staged = items.slice(Math.max(0, items.length - effectiveCount))
-  const isStaging = effectiveCount < items.length
-
-  return { staged, isStaging }
-}