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

apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/mothership-chat.tsx

@@ -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,39 +229,88 @@ 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
-  const precedingUserContentByIndex = useMemo(() => {
-    const out: Array<string | undefined> = []
-    let lastUserContent: string | undefined
-    for (const [index, message] of messages.entries()) {
-      out[index] = lastUserContent
-      if (message.role === 'user') lastUserContent = message.content
-    }
-    return out
-  }, [messages])
-  const assistantTurnKeyByIndex = useMemo(() => {
+
+  /**
+   * 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 initialScrollDoneRef = useRef(false)
+
+  const precedingUserContentByIndex = useMemo(() => {
+    const out: Array<string | undefined> = []
+    let lastUserContent: string | undefined
+    for (const [index, message] of messages.entries()) {
+      out[index] = lastUserContent
+      if (message.role === 'user') lastUserContent = message.content
+    }
+    return out
+  }, [messages])
+
+  /**
+   * 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 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>

apps/sim/app/workspace/[workspaceId]/home/components/user-input/user-input.tsx

@@ -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)

apps/sim/app/workspace/[workspaceId]/home/home.tsx

@@ -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(() => {})
-  }
-
-  function handleSubmit(
-    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,
-    })
+  }, [workspaceId, getCurrentRequestId, stopGeneration])
+
+  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,
+      })
 
-    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) => {