memoize all of the things

Author: brendan-kellam

packages/web/package.json

@@ -137,6 +137,7 @@
     "embla-carousel-auto-scroll": "^8.3.0",
     "embla-carousel-react": "^8.3.0",
     "escape-string-regexp": "^5.0.0",
+    "fast-deep-equal": "^3.1.3",
     "fuse.js": "^7.0.0",
     "google-auth-library": "^10.1.0",
     "graphql": "^16.9.0",

packages/web/src/features/chat/components/chatBox/chatBox.tsx

@@ -8,7 +8,7 @@ import { insertMention, slateContentToString } from "@/features/chat/utils";
 import { cn, IS_MAC } from "@/lib/utils";
 import { computePosition, flip, offset, shift, VirtualElement } from "@floating-ui/react";
 import { ArrowUp, Loader2, StopCircleIcon, TriangleAlertIcon } from "lucide-react";
-import { Fragment, KeyboardEvent, useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { Fragment, KeyboardEvent, memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { useHotkeys } from "react-hotkeys-hook";
 import { Descendant, insertText } from "slate";
 import { Editable, ReactEditor, RenderElementProps, RenderLeafProps, useFocused, useSelected, useSlate } from "slate-react";
@@ -19,6 +19,7 @@ import { useSuggestionModeAndQuery } from "./useSuggestionModeAndQuery";
 import { useSuggestionsData } from "./useSuggestionsData";
 import { useToast } from "@/components/hooks/use-toast";
 import { SearchContextQuery } from "@/lib/types";
+import isEqual from "fast-deep-equal/react";
 
 interface ChatBoxProps {
     onSubmit: (children: Descendant[], editor: CustomEditor) => void;
@@ -34,7 +35,7 @@ interface ChatBoxProps {
     onContextSelectorOpenChanged: (isOpen: boolean) => void;
 }
 
-export const ChatBox = ({
+const ChatBoxComponent = ({
     onSubmit: _onSubmit,
     onStop,
     preferredSuggestionsBoxPlacement = "bottom-start",
@@ -368,6 +369,8 @@ export const ChatBox = ({
     )
 }
 
+export const ChatBox = memo(ChatBoxComponent, isEqual);
+
 const DefaultElement = (props: RenderElementProps) => {
     return <p {...props.attributes}>{props.children}</p>
 }

packages/web/src/features/chat/components/chatThread/answerCard.tsx

@@ -17,6 +17,7 @@ import { isServiceError } from "@/lib/utils";
 import useCaptureEvent from "@/hooks/useCaptureEvent";
 import { LangfuseWeb } from "langfuse";
 import { env } from "@sourcebot/shared/client";
+import isEqual from "fast-deep-equal/react";
 
 interface AnswerCardProps {
     answerText: string;
@@ -178,4 +179,4 @@ const AnswerCardComponent = forwardRef<HTMLDivElement, AnswerCardProps>(({
 
 AnswerCardComponent.displayName = 'AnswerCard';
 
-export const AnswerCard = memo(AnswerCardComponent);
+export const AnswerCard = memo(AnswerCardComponent, isEqual);

packages/web/src/features/chat/components/chatThread/chatThreadListItem.tsx

@@ -8,12 +8,13 @@ import { CSSProperties, forwardRef, memo, useCallback, useEffect, useMemo, useRe
 import scrollIntoView from 'scroll-into-view-if-needed';
 import { Reference, referenceSchema, SBChatMessage, Source } from "../../types";
 import { useExtractReferences } from '../../useExtractReferences';
-import { getAnswerPartFromAssistantMessage, groupMessageIntoSteps, repairReferences } from '../../utils';
+import { getAnswerPartFromAssistantMessage, groupMessageIntoSteps, repairReferences, tryResolveFileReference } from '../../utils';
 import { AnswerCard } from './answerCard';
 import { DetailsCard } from './detailsCard';
 import { MarkdownRenderer, REFERENCE_PAYLOAD_ATTRIBUTE } from './markdownRenderer';
 import { ReferencedSourcesListView } from './referencedSourcesListView';
 import { uiVisiblePartTypes } from '../../constants';
+import isEqual from "fast-deep-equal/react";
 
 interface ChatThreadListItemProps {
     userMessage: SBChatMessage;
@@ -32,7 +33,6 @@ export const ChatThreadListItemComponent = forwardRef<HTMLDivElement, ChatThread
     chatId,
     index,
 }, ref) => {
-    console.log(`re-rendering chat thread list item`, index);
     const leftPanelRef = useRef<HTMLDivElement>(null);
     const [leftPanelHeight, setLeftPanelHeight] = useState<number | null>(null);
     const answerRef = useRef<HTMLDivElement>(null);
@@ -81,7 +81,6 @@ export const ChatThreadListItemComponent = forwardRef<HTMLDivElement, ChatThread
         return getAnswerPartFromAssistantMessage(assistantMessage, isStreaming);
     }, [assistantMessage, isStreaming]);
 
-    const references = useExtractReferences(answerPart);
 
     // Groups parts into steps that are associated with thinking steps that
     // should be visible to the user. By "steps", we mean parts that originated
@@ -280,6 +279,26 @@ export const ChatThreadListItemComponent = forwardRef<HTMLDivElement, ChatThread
         };
     }, [hoveredReference]);
 
+    const references = useExtractReferences(answerPart);
+
+    // Extract the file sources that are referenced by the answer part.
+    const referencedFileSources = useMemo(() => {
+        const fileSources = sources.filter((source) => source.type === 'file');
+
+        return references
+            .filter((reference) => reference.type === 'file')
+            .map((reference) => tryResolveFileReference(reference, fileSources))
+            .filter((file) => file !== undefined)
+            // de-duplicate files
+            .filter((file, index, self) =>
+                index === self.findIndex((t) =>
+                    t?.path === file?.path
+                    && t?.repo === file?.repo
+                    && t?.revision === file?.revision
+                )
+            );
+    }, [references, sources]);
+
 
     return (
         <div
@@ -365,11 +384,11 @@ export const ChatThreadListItemComponent = forwardRef<HTMLDivElement, ChatThread
                     <div
                         className="sticky top-0"
                     >
-                        {references.length > 0 ? (
+                        {referencedFileSources.length > 0 ? (
                             <ReferencedSourcesListView
                                 index={index}
                                 references={references}
-                                sources={sources}
+                                sources={referencedFileSources}
                                 hoveredReference={hoveredReference}
                                 selectedReference={selectedReference}
                                 onSelectedReferenceChanged={setSelectedReference}
@@ -396,10 +415,32 @@ export const ChatThreadListItemComponent = forwardRef<HTMLDivElement, ChatThread
 
 ChatThreadListItemComponent.displayName = 'ChatThreadListItem';
 
-// Only allow re-rendering when the message _is_ streaming.
-// This is a performance optimizations to prevent unnecessary
-// re-renders for chatThreadListItems that are not streaming.
-export const ChatThreadListItem = memo(ChatThreadListItemComponent, (_, nextProps) => !nextProps.isStreaming);
+// Custom comparison function that handles the known issue where useChat mutates
+// message objects in place during streaming, causing fast-deep-equal to return
+// true even when content changes (because it checks reference equality first).
+// See: https://github.com/vercel/ai/issues/6466
+const arePropsEqual = (
+    prevProps: ChatThreadListItemProps,
+    nextProps: ChatThreadListItemProps
+): boolean => {
+    // Always re-render if streaming status changes
+    if (prevProps.isStreaming !== nextProps.isStreaming) {
+        return false;
+    }
+
+    // If currently streaming, always allow re-render
+    // This bypasses the fast-deep-equal reference check issue when useChat
+    // mutates message objects in place during token streaming
+    if (nextProps.isStreaming) {
+        return false;
+    }
+
+    // For non-streaming messages, use deep equality
+    // At this point, useChat should have finished and created final objects
+    return isEqual(prevProps, nextProps);
+};
+
+export const ChatThreadListItem = memo(ChatThreadListItemComponent, arePropsEqual);
 
 // Finds the nearest reference element to the viewport center.
 const getNearestReferenceElement = (referenceElements: Element[]) => {

packages/web/src/features/chat/components/chatThread/detailsCard.tsx

@@ -17,6 +17,7 @@ import { SearchReposToolComponent } from './tools/searchReposToolComponent';
 import { ListAllReposToolComponent } from './tools/listAllReposToolComponent';
 import { SBChatMessageMetadata, SBChatMessagePart } from '../../types';
 import { SearchScopeIcon } from '../searchScopeIcon';
+import isEqual from "fast-deep-equal/react";
 
 
 interface DetailsCardProps {
@@ -36,7 +37,6 @@ const DetailsCardComponent = ({
     metadata,
     thinkingSteps,
 }: DetailsCardProps) => {
-
     return (
         <Card className="mb-4">
             <Collapsible open={isExpanded} onOpenChange={onExpandedChanged}>
@@ -212,4 +212,4 @@ const DetailsCardComponent = ({
     )
 }
 
-export const DetailsCard = memo(DetailsCardComponent);
+export const DetailsCard = memo(DetailsCardComponent, isEqual);

packages/web/src/features/chat/components/chatThread/markdownRenderer.tsx

@@ -20,6 +20,7 @@ import { CodeBlock } from './codeBlock';
 import { FILE_REFERENCE_REGEX } from '@/features/chat/constants';
 import { createFileReference } from '@/features/chat/utils';
 import { SINGLE_TENANT_ORG_DOMAIN } from '@/lib/constants';
+import isEqual from "fast-deep-equal/react";
 
 export const REFERENCE_PAYLOAD_ATTRIBUTE = 'data-reference-payload';
 
@@ -221,4 +222,4 @@ const MarkdownRendererComponent = forwardRef<HTMLDivElement, MarkdownRendererPro
 
 MarkdownRendererComponent.displayName = 'MarkdownRenderer';
 
-export const MarkdownRenderer = memo(MarkdownRendererComponent);
+export const MarkdownRenderer = memo(MarkdownRendererComponent, isEqual);

packages/web/src/features/chat/components/chatThread/referencedFileSourceListItem.tsx

@@ -20,6 +20,7 @@ import { createCodeFoldingExtension } from "./codeFoldingExtension";
 import useCaptureEvent from "@/hooks/useCaptureEvent";
 import { CodeHostType } from "@sourcebot/db";
 import { createAuditAction } from "@/ee/features/audit/actions";
+import isEqual from "fast-deep-equal/react";
 
 const lineDecoration = Decoration.line({
     attributes: { class: "cm-range-border-radius chat-lineHighlight" },
@@ -355,6 +356,6 @@ const ReferencedFileSourceListItem = ({
     )
 }
 
-export default memo(forwardRef(ReferencedFileSourceListItem)) as (
+export default memo(forwardRef(ReferencedFileSourceListItem), isEqual) as (
     props: ReferencedFileSourceListItemProps & { ref?: Ref<ReactCodeMirrorRef> },
 ) => ReturnType<typeof ReferencedFileSourceListItem>;

packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx

@@ -9,12 +9,14 @@ import { useQueries } from "@tanstack/react-query";
 import { ReactCodeMirrorRef } from '@uiw/react-codemirror';
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import scrollIntoView from 'scroll-into-view-if-needed';
-import { FileReference, FileSource, Reference, Source } from "../../types";
+import { FileReference, FileSource, Reference } from "../../types";
+import { tryResolveFileReference } from '../../utils';
 import ReferencedFileSourceListItem from "./referencedFileSourceListItem";
+import isEqual from 'fast-deep-equal/react';
 
 interface ReferencedSourcesListViewProps {
     references: FileReference[];
-    sources: Source[];
+    sources: FileSource[];
     index: number;
     hoveredReference?: Reference;
     onHoveredReferenceChanged: (reference?: Reference) => void;
@@ -23,13 +25,6 @@ interface ReferencedSourcesListViewProps {
     style: React.CSSProperties;
 }
 
-const resolveFileReference = (reference: FileReference, sources: FileSource[]): FileSource | undefined => {
-    return sources.find(
-        (source) => source.repo.endsWith(reference.repo) &&
-            source.path.endsWith(reference.path)
-    );
-}
-
 const ReferencedSourcesListViewComponent = ({
     references,
     sources,
@@ -59,43 +54,26 @@ const ReferencedSourcesListViewComponent = ({
         }
     }, []);
 
-    const referencedFileSources = useMemo((): FileSource[] => {
-        const fileSources = sources.filter((source) => source.type === 'file');
-
-        return references
-            .filter((reference) => reference.type === 'file')
-            .map((reference) => resolveFileReference(reference, fileSources))
-            .filter((file) => file !== undefined)
-            // de-duplicate files
-            .filter((file, index, self) =>
-                index === self.findIndex((t) =>
-                    t?.path === file?.path
-                    && t?.repo === file?.repo
-                    && t?.revision === file?.revision
-                )
-            );
-    }, [references, sources]);
-
     // Memoize the computation of references grouped by file source
     const referencesGroupedByFile = useMemo(() => {
         const groupedReferences = new Map<string, FileReference[]>();
 
-        for (const fileSource of referencedFileSources) {
+        for (const fileSource of sources) {
             const fileKey = getFileId(fileSource);
             const referencesInFile = references.filter((reference) => {
                 if (reference.type !== 'file') {
                     return false;
                 }
-                return resolveFileReference(reference, [fileSource]) !== undefined;
+                return tryResolveFileReference(reference, [fileSource]) !== undefined;
             });
             groupedReferences.set(fileKey, referencesInFile);
         }
 
         return groupedReferences;
-    }, [references, referencedFileSources, getFileId]);
+    }, [references, sources, getFileId]);
 
     const fileSourceQueries = useQueries({
-        queries: referencedFileSources.map((file) => ({
+        queries: sources.map((file) => ({
             queryKey: ['fileSource', file.path, file.repo, file.revision],
             queryFn: () => unwrapServiceError(getFileSource({
                 fileName: file.path,
@@ -112,7 +90,7 @@ const ReferencedSourcesListViewComponent = ({
             return;
         }
 
-        const fileSource = resolveFileReference(selectedReference, referencedFileSources);
+        const fileSource = tryResolveFileReference(selectedReference, sources);
         if (!fileSource) {
             return;
         }
@@ -179,7 +157,7 @@ const ReferencedSourcesListViewComponent = ({
                 behavior: 'smooth',
             });
         }
-    }, [getFileId, referencedFileSources, selectedReference]);
+    }, [getFileId, sources, selectedReference]);
 
     const onExpandedChanged = useCallback((fileId: string, isExpanded: boolean) => {
         if (isExpanded) {
@@ -200,7 +178,7 @@ const ReferencedSourcesListViewComponent = ({
         }
     }, []);
 
-    if (referencedFileSources.length === 0) {
+    if (sources.length === 0) {
         return (
             <div className="p-4 text-center text-muted-foreground text-sm">
                 No file references found
@@ -215,7 +193,7 @@ const ReferencedSourcesListViewComponent = ({
         >
             <div className="space-y-4 pr-2">
                 {fileSourceQueries.map((query, index) => {
-                    const fileSource = referencedFileSources[index];
+                    const fileSource = sources[index];
                     const fileName = fileSource.path.split('/').pop() ?? fileSource.path;
 
                     if (query.isLoading) {
@@ -280,4 +258,4 @@ const ReferencedSourcesListViewComponent = ({
 }
 
 // Memoize to prevent unnecessary re-renders
-export const ReferencedSourcesListView = memo(ReferencedSourcesListViewComponent);
+export const ReferencedSourcesListView = memo(ReferencedSourcesListViewComponent, isEqual);

packages/web/src/features/chat/utils.ts

@@ -374,3 +374,13 @@ export const buildSearchQuery = (options: {
 export const getLanguageModelKey = (model: LanguageModelInfo) => {
     return `${model.provider}-${model.model}-${model.displayName}`;
 }
+
+/**
+ * Given a file reference and a list of file sources, attempts to resolve the file source that the reference points to.
+ */
+export const tryResolveFileReference = (reference: FileReference, sources: FileSource[]): FileSource | undefined => {
+    return sources.find(
+        (source) => source.repo.endsWith(reference.repo) &&
+            source.path.endsWith(reference.path)
+    );
+}

yarn.lock

@@ -8197,6 +8197,7 @@ __metadata:
     eslint-config-next: "npm:15.5.0"
     eslint-plugin-react: "npm:^7.37.5"
     eslint-plugin-react-hooks: "npm:^5.2.0"
+    fast-deep-equal: "npm:^3.1.3"
     fuse.js: "npm:^7.0.0"
     google-auth-library: "npm:^10.1.0"
     graphql: "npm:^16.9.0"