improvement(mothership-chat): memoize message rows for long-transcript load speed (#4520)

* improvement(mothership-chat): memoize message rows for long-transcript load speed

* improvement(mothership-chat): hoist EMPTY_BLOCKS constant to keep MessageContent memoized for block-free messages

Author: waleedlatif1

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

@@ -1,6 +1,6 @@
 'use client'
 
-import { type ComponentPropsWithoutRef, useEffect, useMemo, useRef } from 'react'
+import { type ComponentPropsWithoutRef, memo, useEffect, useMemo, useRef } from 'react'
 import { Streamdown } from 'streamdown'
 import 'streamdown/styles.css'
 import 'prismjs/components/prism-typescript'
@@ -237,7 +237,7 @@ interface ChatContentProps {
   onWorkspaceResourceSelect?: (resource: MothershipResource) => void
 }
 
-export function ChatContent({
+function ChatContentInner({
   content,
   isStreaming = false,
   onOptionSelect,
@@ -335,3 +335,5 @@ export function ChatContent({
     </div>
   )
 }
+
+export const ChatContent = memo(ChatContentInner)

apps/sim/app/workspace/[workspaceId]/home/components/message-content/message-content.tsx

@@ -1,5 +1,6 @@
 'use client'
 
+import { memo, useMemo } from 'react'
 import {
   Read as ReadTool,
   ToolSearchToolRegex,
@@ -407,14 +408,14 @@ interface MessageContentProps {
   onWorkspaceResourceSelect?: (resource: MothershipResource) => void
 }
 
-export function MessageContent({
+function MessageContentInner({
   blocks,
   fallbackContent,
   isStreaming = false,
   onOptionSelect,
   onWorkspaceResourceSelect,
 }: MessageContentProps) {
-  const parsed = blocks.length > 0 ? parseBlocks(blocks) : []
+  const parsed = useMemo(() => (blocks.length > 0 ? parseBlocks(blocks) : []), [blocks])
 
   const segments: MessageSegment[] =
     parsed.length > 0
@@ -537,3 +538,5 @@ export function MessageContent({
     </div>
   )
 }
+
+export const MessageContent = memo(MessageContentInner)

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

@@ -1,6 +1,6 @@
 'use client'
 
-import { useLayoutEffect, useRef } from 'react'
+import { memo, useCallback, useEffect, useLayoutEffect, useMemo, useRef } from 'react'
 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'
@@ -17,6 +17,9 @@ import {
 import { UserMessageContent } from '@/app/workspace/[workspaceId]/home/components/user-message-content'
 import type {
   ChatMessage,
+  ChatMessageAttachment,
+  ChatMessageContext,
+  ContentBlock,
   FileAttachmentForApi,
   MothershipResource,
   QueuedMessage,
@@ -78,6 +81,100 @@ const LAYOUT_STYLES = {
   },
 } as const
 
+const EMPTY_BLOCKS: ContentBlock[] = []
+
+interface UserMessageRowProps {
+  content: string
+  contexts?: ChatMessageContext[]
+  attachments?: ChatMessageAttachment[]
+  rowClassName: string
+  bubbleClassName: string
+  attachmentWidthClassName: string
+}
+
+const UserMessageRow = memo(function UserMessageRow({
+  content,
+  contexts,
+  attachments,
+  rowClassName,
+  bubbleClassName,
+  attachmentWidthClassName,
+}: UserMessageRowProps) {
+  const hasAttachments = Boolean(attachments?.length)
+  return (
+    <div className={rowClassName}>
+      {hasAttachments && (
+        <ChatMessageAttachments
+          attachments={attachments ?? []}
+          align='end'
+          className={attachmentWidthClassName}
+        />
+      )}
+      <div className={bubbleClassName}>
+        <UserMessageContent content={content} contexts={contexts} />
+      </div>
+    </div>
+  )
+})
+
+interface AssistantMessageRowProps {
+  message: ChatMessage
+  isStreaming: boolean
+  precedingUserContent?: string
+  chatId?: string
+  rowClassName: string
+  onOptionSelect?: (id: string) => void
+  onWorkspaceResourceSelect?: (resource: MothershipResource) => void
+}
+
+const AssistantMessageRow = memo(function AssistantMessageRow({
+  message,
+  isStreaming,
+  precedingUserContent,
+  chatId,
+  rowClassName,
+  onOptionSelect,
+  onWorkspaceResourceSelect,
+}: AssistantMessageRowProps) {
+  const blocks = message.contentBlocks ?? EMPTY_BLOCKS
+  const hasAnyBlocks = blocks.length > 0
+  const trimmedContent = message.content?.trim() ?? ''
+
+  if (!hasAnyBlocks && !trimmedContent && isStreaming) {
+    return <PendingTagIndicator />
+  }
+
+  const hasRenderableAssistant = assistantMessageHasRenderableContent(blocks, message.content ?? '')
+  if (!hasRenderableAssistant && !trimmedContent && !isStreaming) {
+    return null
+  }
+
+  const showActions = !isStreaming && (message.content || hasAnyBlocks)
+
+  return (
+    <div className={rowClassName}>
+      <MessageContent
+        blocks={blocks}
+        fallbackContent={message.content}
+        isStreaming={isStreaming}
+        onOptionSelect={onOptionSelect}
+        onWorkspaceResourceSelect={onWorkspaceResourceSelect}
+      />
+      {showActions && (
+        <div className='mt-2.5'>
+          <MessageActions
+            content={message.content}
+            chatId={chatId}
+            userQuery={precedingUserContent}
+            requestId={message.requestId}
+            messageId={message.id}
+          />
+        </div>
+      )}
+    </div>
+  )
+})
+
 export function MothershipChat({
   messages,
   isSending,
@@ -111,17 +208,31 @@ export function MothershipChat({
   const { staged: stagedMessages, isStaging } = useProgressiveList(messages, stagingKey)
   const stagedMessageCount = stagedMessages.length
   const stagedOffset = messages.length - stagedMessages.length
-  const precedingUserContentByIndex: Array<string | undefined> = []
-  let lastUserContent: string | undefined
-  for (const [index, message] of messages.entries()) {
-    precedingUserContentByIndex[index] = lastUserContent
-    if (message.role === 'user') {
-      lastUserContent = message.content
+  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 initialScrollDoneRef = useRef(false)
   const userInputRef = useRef<UserInputHandle>(null)
 
+  const onSubmitRef = useRef(onSubmit)
+  const onWorkspaceResourceSelectRef = useRef(onWorkspaceResourceSelect)
+  useEffect(() => {
+    onSubmitRef.current = onSubmit
+    onWorkspaceResourceSelectRef.current = onWorkspaceResourceSelect
+  }, [onSubmit, onWorkspaceResourceSelect])
+  const stableOnOptionSelect = useCallback((id: string) => {
+    onSubmitRef.current(id)
+  }, [])
+  const stableOnWorkspaceResourceSelect = useCallback((resource: MothershipResource) => {
+    onWorkspaceResourceSelectRef.current?.(resource)
+  }, [])
+
   function handleSendQueuedHead() {
     const topMessage = messageQueue[0]
     if (!topMessage) return
@@ -164,63 +275,31 @@ export function MothershipChat({
             {stagedMessages.map((msg, localIndex) => {
               const index = stagedOffset + localIndex
               if (msg.role === 'user') {
-                const hasAttachments = Boolean(msg.attachments?.length)
                 return (
-                  <div key={msg.id} className={styles.userRow}>
-                    {hasAttachments && (
-                      <ChatMessageAttachments
-                        attachments={msg.attachments ?? []}
-                        align='end'
-                        className={styles.attachmentWidth}
-                      />
-                    )}
-                    <div className={styles.userBubble}>
-                      <UserMessageContent content={msg.content} contexts={msg.contexts} />
-                    </div>
-                  </div>
+                  <UserMessageRow
+                    key={msg.id}
+                    content={msg.content}
+                    contexts={msg.contexts}
+                    attachments={msg.attachments}
+                    rowClassName={styles.userRow}
+                    bubbleClassName={styles.userBubble}
+                    attachmentWidthClassName={styles.attachmentWidth}
+                  />
                 )
               }
 
-              const hasAnyBlocks = Boolean(msg.contentBlocks?.length)
-              const hasRenderableAssistant = assistantMessageHasRenderableContent(
-                msg.contentBlocks ?? [],
-                msg.content ?? ''
-              )
-              const isLastAssistant = index === messages.length - 1
-              const isThisStreaming = isStreamActive && isLastAssistant
-
-              if (!hasAnyBlocks && !msg.content?.trim() && isThisStreaming) {
-                return <PendingTagIndicator key={msg.id} />
-              }
-
-              if (!hasRenderableAssistant && !msg.content?.trim() && !isThisStreaming) {
-                return null
-              }
-
-              const isLastMessage = index === messages.length - 1
-              const precedingUserContent = precedingUserContentByIndex[index]
-
+              const isLast = index === messages.length - 1
               return (
-                <div key={msg.id} className={styles.assistantRow}>
-                  <MessageContent
-                    blocks={msg.contentBlocks || []}
-                    fallbackContent={msg.content}
-                    isStreaming={isThisStreaming}
-                    onOptionSelect={isLastMessage ? onSubmit : undefined}
-                    onWorkspaceResourceSelect={onWorkspaceResourceSelect}
-                  />
-                  {!isThisStreaming && (msg.content || msg.contentBlocks?.length) && (
-                    <div className='mt-2.5'>
-                      <MessageActions
-                        content={msg.content}
-                        chatId={chatId}
-                        userQuery={precedingUserContent}
-                        requestId={msg.requestId}
-                        messageId={msg.id}
-                      />
-                    </div>
-                  )}
-                </div>
+                <AssistantMessageRow
+                  key={msg.id}
+                  message={msg}
+                  isStreaming={isStreamActive && isLast}
+                  precedingUserContent={precedingUserContentByIndex[index]}
+                  chatId={chatId}
+                  rowClassName={styles.assistantRow}
+                  onOptionSelect={isLast ? stableOnOptionSelect : undefined}
+                  onWorkspaceResourceSelect={stableOnWorkspaceResourceSelect}
+                />
               )
             })}
           </div>

apps/sim/lib/copilot/chat/display-message.ts

@@ -112,7 +112,18 @@ function toDisplayContexts(
   }))
 }
 
+const displayMessageCache = new WeakMap<PersistedMessage, ChatMessage>()
+
+/**
+ * Maps a `PersistedMessage` (server wire shape) to a `ChatMessage` (UI shape).
+ * Reference-stable: returns the same object for a given `PersistedMessage`
+ * instance so `React.memo` boundaries downstream of React Query's structural
+ * sharing can short-circuit on identity.
+ */
 export function toDisplayMessage(msg: PersistedMessage): ChatMessage {
+  const cached = displayMessageCache.get(msg)
+  if (cached) return cached
+
   const display: ChatMessage = {
     id: msg.id,
     role: msg.role,
@@ -136,5 +147,6 @@ export function toDisplayMessage(msg: PersistedMessage): ChatMessage {
 
   display.contexts = toDisplayContexts(msg.contexts)
 
+  displayMessageCache.set(msg, display)
   return display
 }