perf(web): stop re-rendering the whole chat transcript on every stream flush (#18221)

web/src/components/markdown-content/index.tsx

@@ -22,7 +22,7 @@ import { citationMarkerReg } from '@/utils/citation-utils';
 import { getExtension } from '@/utils/document-util';
 import { getDirAttribute } from '@/utils/text-direction';
 import DOMPurify from 'dompurify';
-import { useCallback, useEffect, useMemo, useState } from 'react';
+import { memo, useCallback, useEffect, useMemo } from 'react';
 import Markdown from 'react-markdown';
 import SyntaxHighlighter from 'react-syntax-highlighter';
 import rehypeKatex from 'rehype-katex';
@@ -62,6 +62,53 @@ import { SafeImg } from '@/components/safe-img';
 
 const getChunkIndex = (match: string) => parseCitationIndex(match);
 
+// Wraps every text node so citation markers can be replaced by React elements.
+// Defined at module scope: react-markdown rebuilds its whole processor whenever
+// a plugin's identity changes, and while an answer streams this component
+// re-renders many times a second.
+const rehypeWrapReference = () => {
+  return function wrapTextTransform(tree: any) {
+    visitParents(tree, 'text', (node, ancestors) => {
+      const latestAncestor = ancestors.at(-1);
+      if (
+        latestAncestor.tagName !== 'custom-typography' &&
+        latestAncestor.tagName !== 'code'
+      ) {
+        node.type = 'element';
+        node.tagName = 'custom-typography';
+        node.properties = {};
+        node.children = [{ type: 'text', value: node.value }];
+      }
+    });
+  };
+};
+
+const MarkdownRehypePlugins = [rehypeRaw, rehypeWrapReference, rehypeKatex];
+
+const MarkdownParagraph = ({ children, ...props }: any) => (
+  <p {...props}>{children}</p>
+);
+
+const MarkdownCode = (props: any) => {
+  const { children, className, ...rest } = props;
+  const restProps = omit(rest, 'node');
+  const match = /language-(\w+)/.exec(className || '');
+  return match ? (
+    <SyntaxHighlighter
+      {...restProps}
+      PreTag="div"
+      language={match[1]}
+      wrapLongLines
+    >
+      {String(children).replace(/\n$/, '')}
+    </SyntaxHighlighter>
+  ) : (
+    <code {...restProps} className={classNames(className, 'text-wrap')}>
+      {children}
+    </code>
+  );
+};
+
 const formatMetadataValue = (value: unknown) => {
   if (Array.isArray(value)) return value.join(', ');
   if (value === null || value === undefined) return '';
@@ -111,10 +158,18 @@ const MarkdownContent = ({
     );
   }, [content, loading, t]);
 
-  useEffect(() => {
+  const documentIds = useMemo(() => {
     const docAggs = reference?.doc_aggs;
-    setDocumentIds(Array.isArray(docAggs) ? docAggs.map((x) => x.doc_id) : []);
-  }, [reference, setDocumentIds]);
+    return Array.isArray(docAggs) ? docAggs.map((x) => x.doc_id) : [];
+  }, [reference?.doc_aggs]);
+
+  // Skipping the empty case matters: this component is mounted once per message,
+  // and setting a fresh empty array would re-render (and re-parse the markdown
+  // of) every message that carries no reference at all.
+  useEffect(() => {
+    if (documentIds.length === 0) return;
+    setDocumentIds(documentIds);
+  }, [documentIds, setDocumentIds]);
 
   const handleDocumentButtonClick = useCallback(
     (
@@ -137,23 +192,6 @@ const MarkdownContent = ({
     [clickDocumentButton],
   );
 
-  const rehypeWrapReference = () => {
-    return function wrapTextTransform(tree: any) {
-      visitParents(tree, 'text', (node, ancestors) => {
-        const latestAncestor = ancestors.at(-1);
-        if (
-          latestAncestor.tagName !== 'custom-typography' &&
-          latestAncestor.tagName !== 'code'
-        ) {
-          node.type = 'element';
-          node.tagName = 'custom-typography';
-          node.properties = {};
-          node.children = [{ type: 'text', value: node.value }];
-        }
-      });
-    };
-  };
-
   const getReferenceInfo = useCallback(
     (chunkIndex: number) => {
       const chunks = reference?.chunks ?? [];
@@ -295,41 +333,24 @@ const MarkdownContent = ({
   const dir = getDirAttribute(content.replace(citationMarkerReg, ''));
   const showLoadingDots = useLoadingPause(loading, content);
 
+  const markdownComponents = useMemo(
+    () =>
+      ({
+        p: MarkdownParagraph,
+        'custom-typography': ({ children }: { children: string }) =>
+          renderReference(children),
+        img: SafeImg,
+        code: MarkdownCode,
+      }) as any,
+    [renderReference],
+  );
+
   return (
     <div dir={dir} className={styles.markdownContentWrapper}>
       <Markdown
-        rehypePlugins={[rehypeRaw, rehypeWrapReference, rehypeKatex]}
+        rehypePlugins={MarkdownRehypePlugins}
         remarkPlugins={MarkdownRemarkPlugins}
-        components={
-          {
-            p: ({ children, ...props }: any) => <p {...props}>{children}</p>,
-            'custom-typography': ({ children }: { children: string }) =>
-              renderReference(children),
-            img: SafeImg,
-            code(props: any) {
-              const { children, className, ...rest } = props;
-              const restProps = omit(rest, 'node');
-              const match = /language-(\w+)/.exec(className || '');
-              return match ? (
-                <SyntaxHighlighter
-                  {...restProps}
-                  PreTag="div"
-                  language={match[1]}
-                  wrapLongLines
-                >
-                  {String(children).replace(/\n$/, '')}
-                </SyntaxHighlighter>
-              ) : (
-                <code
-                  {...restProps}
-                  className={classNames(className, 'text-wrap')}
-                >
-                  {children}
-                </code>
-              );
-            },
-          } as any
-        }
+        components={markdownComponents}
       >
         {contentWithCursor}
       </Markdown>
@@ -340,4 +361,4 @@ const MarkdownContent = ({
   );
 };
 
-export default MarkdownContent;
+export default memo(MarkdownContent);

web/src/hooks/logic-hooks.ts

@@ -437,7 +437,13 @@ export const useScrollToBottom = (
     const container = containerRef.current;
 
     const handleScroll = () => {
-      setIsAtBottom(checkIfUserAtBottom());
+      // Write the ref here rather than relying on the effect that mirrors
+      // `isAtBottom`: that effect only runs after the next render, and while the
+      // main thread is busy rendering a streaming answer an already scheduled
+      // auto-scroll would still see the stale `true`.
+      const atBottom = checkIfUserAtBottom();
+      isAtBottomRef.current = atBottom;
+      setIsAtBottom(atBottom);
     };
 
     container.addEventListener('scroll', handleScroll);
@@ -456,16 +462,22 @@ export const useScrollToBottom = (
     }
   }, [containerRef]);
 
+  // Streaming replaces `messages` many times a second. The previous
+  // rAF + setTimeout(100) chain always had several scrolls queued, and they read
+  // `isAtBottomRef` long after the user had scrolled up — yanking the view back
+  // down. One cancellable frame per change, gated on the latest position, keeps
+  // auto-follow without fighting the user.
   useEffect(() => {
     if (!messages) return;
     if (!containerRef?.current) return;
-    requestAnimationFrame(() => {
-      setTimeout(() => {
-        if (isAtBottomRef.current) {
-          scrollToBottom();
-        }
-      }, 100);
+    if (!isAtBottomRef.current) return;
+
+    const frame = requestAnimationFrame(() => {
+      if (isAtBottomRef.current) {
+        scrollToBottom();
+      }
     });
+    return () => cancelAnimationFrame(frame);
   }, [messages, containerRef, scrollToBottom]);
 
   return { scrollRef: ref, isAtBottom, scrollToBottom };

web/src/pages/next-chats/chat/chat-box/next-multiple-chat-box.tsx

@@ -57,7 +57,8 @@ import {
   UseSendSingleMessageParameter,
 } from '../../hooks/use-send-single-message';
 import { useUploadFile } from '../../hooks/use-upload-file';
-import { buildMessageItemReference } from '../../utils';
+import { useMessageReferences } from '../../hooks/use-message-references';
+import { EmptyReference } from '../../utils';
 import { useAddChatBox } from '../use-add-box';
 import { useShowInternet } from '../use-show-internet';
 
@@ -123,6 +124,11 @@ const ChatCard = forwardRef(function ChatCard(
 
   const { scrollRef } = useScrollToBottom(derivedMessages, messageContainerRef);
 
+  const messageReferences = useMessageReferences(
+    derivedMessages,
+    conversation.reference,
+  );
+
   const FormSchema = z.object(LlmSettingSchema);
 
   const form = useForm<z.infer<typeof FormSchema>>({
@@ -166,10 +172,7 @@ const ChatCard = forwardRef(function ChatCard(
   // conversation switch), not when currentDialog.llm_id changes due to Apply.
   const syncedDialogIdRef = useRef<string | undefined>(undefined);
   useLayoutEffect(() => {
-    if (
-      syncedDialogIdRef.current !== dialogId &&
-      currentDialog?.llm_id
-    ) {
+    if (syncedDialogIdRef.current !== dialogId && currentDialog?.llm_id) {
       form.setValue('llm_id', currentDialog.llm_id);
       syncedDialogIdRef.current = dialogId;
     }
@@ -290,13 +293,7 @@ const ChatCard = forwardRef(function ChatCard(
                   nickname={userInfo.nickname}
                   avatar={userInfo.avatar}
                   avatarDialog={currentDialog.icon}
-                  reference={buildMessageItemReference(
-                    {
-                      messages: derivedMessages,
-                      reference: conversation.reference,
-                    },
-                    message,
-                  )}
+                  reference={messageReferences.get(message) ?? EmptyReference}
                   // clickDocumentButton={clickDocumentButton}
                   index={i}
                   removeMessageById={removeMessageById}

web/src/pages/next-chats/chat/chat-box/single-chat-box.tsx

@@ -14,7 +14,8 @@ import {
 } from '../../hooks/use-button-disabled';
 import { useCreateConversationBeforeUploadDocument } from '../../hooks/use-create-conversation';
 import { useSendMessage } from '../../hooks/use-send-chat-message';
-import { buildMessageItemReference } from '../../utils';
+import { useMessageReferences } from '../../hooks/use-message-references';
+import { EmptyReference } from '../../utils';
 import { useShowInternet } from '../use-show-internet';
 
 interface IProps {
@@ -50,6 +51,11 @@ export function SingleChatBox({ conversation }: IProps) {
 
   const showInternet = useShowInternet();
 
+  const messageReferences = useMessageReferences(
+    messages,
+    conversation.reference,
+  );
+
   useEffect(() => {
     // Skip when the conversation prop is stale — its id doesn't match the
     // URL's current conversationId. This happens during a switch (e.g.
@@ -92,13 +98,7 @@ export function SingleChatBox({ conversation }: IProps) {
               nickname={userInfo.nickname}
               avatar={userInfo.avatar}
               avatarDialog={currentDialog.icon}
-              reference={buildMessageItemReference(
-                {
-                  messages,
-                  reference: conversation.reference,
-                },
-                message,
-              )}
+              reference={messageReferences.get(message) ?? EmptyReference}
               clickDocumentButton={clickDocumentButton}
               index={i}
               removeMessageById={removeMessageById}

web/src/pages/next-chats/hooks/use-message-references.ts

@@ -0,0 +1,71 @@
+import { MessageType } from '@/constants/chat';
+import { IMessage, IReference } from '@/interfaces/database/chat';
+import { isEmpty } from 'lodash';
+import { useMemo } from 'react';
+import { EmptyReference } from '../utils';
+
+/**
+ * Resolves the reference of every message in one pass and caches the result
+ * until the message list or the conversation's reference list changes.
+ *
+ * Doing this per message inside the render loop (the previous
+ * `buildMessageItemReference` call site) was quadratic, and — worse — handed
+ * `MessageItem` a freshly allocated object on every streaming flush, which
+ * defeated its `memo` and re-parsed the markdown of the whole transcript
+ * ~16 times a second.
+ *
+ * Keyed by message object rather than by id: a question and its answer share an
+ * id, so an id-keyed map could not tell them apart.
+ */
+export function useMessageReferences(
+  messages: IMessage[] | undefined,
+  reference: IReference[] | undefined,
+) {
+  return useMemo(() => {
+    const list = messages ?? [];
+    const references = reference ?? [];
+
+    // Index of each assistant message within the answer sequence, skipping the
+    // prologue (which has no reference) and error answers. First id wins, which
+    // is what the previous `findIndex` lookup did.
+    const answerIndexById = new Map<string, number>();
+    let answerCount = 0;
+    list.forEach((message) => {
+      if (
+        message.role !== MessageType.Assistant ||
+        message.content?.startsWith('**ERROR**:')
+      ) {
+        return;
+      }
+      if (answerCount > 0 && !answerIndexById.has(message.id)) {
+        answerIndexById.set(message.id, answerCount - 1);
+      }
+      answerCount += 1;
+    });
+
+    const resolved = new Map<IMessage, IReference>();
+    list.forEach((message) => {
+      if (!isEmpty(message.reference)) {
+        resolved.set(message, message.reference as IReference);
+        return;
+      }
+      // An assistant message that has not received any content yet is still
+      // being generated. Never resolve its reference from the conversation
+      // reference list — the indices only align with completed answers, so the
+      // lookup would surface a stale reference from a previous turn (or from a
+      // previously opened conversation) while the answer is streaming.
+      if (message.role === MessageType.Assistant && !message.content) {
+        resolved.set(message, EmptyReference);
+        return;
+      }
+      const answerIndex = answerIndexById.get(message.id);
+      resolved.set(
+        message,
+        (answerIndex === undefined ? undefined : references[answerIndex]) ??
+          EmptyReference,
+      );
+    });
+
+    return resolved;
+  }, [messages, reference]);
+}

web/src/pages/next-chats/hooks/use-send-chat-message.ts

@@ -98,19 +98,17 @@ export const useSendMessage = () => {
     }: {
       message: IMessage;
       currentConversationId?: string;
-      messages?: IMessage[];
+      messages: IMessage[];
     } & NextMessageInputOnPressEnterParameter) => {
       const sessionId = currentConversationId ?? conversationId;
 
       const { ok, aborted } = await runChatCompletionStream({
         conversationId: sessionId,
         chatId,
-        // An explicitly provided list is authoritative, even when empty
-        // (e.g. regenerating the first question must truncate history).
-        messages: [
-          ...(Array.isArray(explicitMessages) ? explicitMessages : messages),
-          message,
-        ],
+        // The history is always supplied by the caller, captured *before* it
+        // appended the question and its placeholder to the store — reading it
+        // here would include both and send the question twice.
+        messages: [...explicitMessages, message],
         enableThinking,
         enableInternet,
       });
@@ -124,7 +122,7 @@ export const useSendMessage = () => {
         notification.error({ message: t('message.requestError') });
       }
     },
-    [conversationId, chatId, messages, failStream, t],
+    [conversationId, chatId, failStream, t],
   );
 
   // Hand a failed question back to the input box, but only once the box is
@@ -156,8 +154,13 @@ export const useSendMessage = () => {
         conversationId,
       };
 
+      // Snapshot the history before appendQuestion writes the question and its
+      // assistant placeholder into the store.
+      const history =
+        useChatStreamStore.getState().sessions[conversationId]?.messages ?? [];
+
       appendQuestion(conversationId, questionMessage);
-      sendMessage({ message: questionMessage });
+      sendMessage({ message: questionMessage, messages: history });
     },
     [conversationId, isStreaming, appendQuestion, sendMessage],
   );
@@ -217,6 +220,13 @@ export const useSendMessage = () => {
 
       // Route the question to the conversation it was asked in, not whichever
       // is currently displayed.
+      //
+      // Snapshot the history before appendQuestion writes the question and its
+      // assistant placeholder into the store.
+      const history =
+        useChatStreamStore.getState().sessions[targetConversationId]
+          ?.messages ?? [];
+
       appendQuestion(targetConversationId, questionMessage);
 
       setValue('');
@@ -224,7 +234,7 @@ export const useSendMessage = () => {
         currentConversationId: targetConversationId,
         // For an existing conversation currentMessages is empty; fall back to
         // the store's message list instead of sending an empty history.
-        messages: currentMessages.length > 0 ? currentMessages : undefined,
+        messages: currentMessages.length > 0 ? currentMessages : history,
         message: {
           id,
           content: value.trim(),

web/src/pages/next-chats/share/index.tsx

@@ -14,7 +14,8 @@ import {
   useGetSharedChatSearchParams,
   useSendSharedMessage,
 } from '../hooks/use-send-shared-message';
-import { buildMessageItemReference } from '../utils';
+import { useMessageReferences } from '../hooks/use-message-references';
+import { EmptyReference } from '../utils';
 
 const ChatContainer = () => {
   const {
@@ -42,6 +43,8 @@ const ChatContainer = () => {
   const sendDisabled = useSendButtonDisabled(value);
   const { data: chatInfo } = useFetchExternalChatInfo();
 
+  const messageReferences = useMessageReferences(derivedMessages, undefined);
+
   React.useEffect(() => {
     if (locale && i18n.language !== locale) {
       changeLanguageAsync(locale, { persist: false });
@@ -78,13 +81,7 @@ const ChatContainer = () => {
                     avatarDialog={avatarDialogSrc}
                     item={message}
                     nickname="You"
-                    reference={buildMessageItemReference(
-                      {
-                        messages: derivedMessages,
-                        reference: [],
-                      },
-                      message,
-                    )}
+                    reference={messageReferences.get(message) ?? EmptyReference}
                     loading={
                       message.role === MessageType.Assistant &&
                       sendLoading &&

web/src/pages/next-chats/utils.ts

@@ -1,10 +1,5 @@
-import { EmptyConversationId, MessageType } from '@/constants/chat';
-import {
-  IConversation,
-  IMessage,
-  IReference,
-} from '@/interfaces/database/chat';
-import { isEmpty } from 'lodash';
+import { EmptyConversationId } from '@/constants/chat';
+import { IConversation, IReference } from '@/interfaces/database/chat';
 
 export const isConversationIdExist = (conversationId: string) => {
   return conversationId !== EmptyConversationId && conversationId !== '';
@@ -27,34 +22,11 @@ export const getDocumentIdsFromConversionReference = (data: IConversation) => {
   return documentIds.join(',');
 };
 
-export const buildMessageItemReference = (
-  conversation: { messages: IMessage[]; reference: IReference[] },
-  message: IMessage,
-) => {
-  const assistantMessages = conversation.messages
-    ?.filter(
-      (x) =>
-        x.role === MessageType.Assistant &&
-        !x.content?.startsWith('**ERROR**:'), // Exclude error messages
-    )
-    .slice(1);
-  const referenceIndex = assistantMessages.findIndex(
-    (x) => x.id === message.id,
-  );
-  // An assistant message that has not received any content yet is still
-  // being generated. Never resolve its reference from the conversation
-  // reference list — the indices only align with completed answers, so the
-  // lookup would surface a stale reference from a previous turn (or from a
-  // previously opened conversation) while the answer is streaming.
-  const isPendingAnswer =
-    message.role === MessageType.Assistant &&
-    !message.content &&
-    isEmpty(message?.reference);
-  const reference = !isEmpty(message?.reference)
-    ? message?.reference
-    : isPendingAnswer
-      ? undefined
-      : (conversation?.reference ?? [])[referenceIndex];
-
-  return reference ?? { doc_aggs: [], chunks: [], total: 0 };
+// Shared fallback so a message without a reference keeps handing MessageItem the
+// same object across renders. A fresh literal here would break the item's memo
+// on every streaming flush. See useMessageReferences.
+export const EmptyReference: IReference = {
+  doc_aggs: [],
+  chunks: [],
+  total: 0,
 };