From 4b7798a94e4848abe51afb6a4a74ed04f2234e22 Mon Sep 17 00:00:00 2001 From: bkellam Date: Thu, 27 Nov 2025 16:05:06 -0800 Subject: [PATCH] memoize all of the things --- packages/web/package.json | 1 + .../chat/components/chatBox/chatBox.tsx | 7 ++- .../chat/components/chatThread/answerCard.tsx | 3 +- .../chatThread/chatThreadListItem.tsx | 59 ++++++++++++++++--- .../components/chatThread/detailsCard.tsx | 4 +- .../chatThread/markdownRenderer.tsx | 3 +- .../referencedFileSourceListItem.tsx | 3 +- .../chatThread/referencedSourcesListView.tsx | 48 ++++----------- packages/web/src/features/chat/utils.ts | 10 ++++ yarn.lock | 1 + 10 files changed, 88 insertions(+), 51 deletions(-) diff --git a/packages/web/package.json b/packages/web/package.json index a4f58011..f08e6350 100644 --- a/packages/web/package.json +++ b/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", diff --git a/packages/web/src/features/chat/components/chatBox/chatBox.tsx b/packages/web/src/features/chat/components/chatBox/chatBox.tsx index 8876bd69..f95621aa 100644 --- a/packages/web/src/features/chat/components/chatBox/chatBox.tsx +++ b/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

{props.children}

} diff --git a/packages/web/src/features/chat/components/chatThread/answerCard.tsx b/packages/web/src/features/chat/components/chatThread/answerCard.tsx index 5feebd32..d420ce5e 100644 --- a/packages/web/src/features/chat/components/chatThread/answerCard.tsx +++ b/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(({ AnswerCardComponent.displayName = 'AnswerCard'; -export const AnswerCard = memo(AnswerCardComponent); \ No newline at end of file +export const AnswerCard = memo(AnswerCardComponent, isEqual); \ No newline at end of file diff --git a/packages/web/src/features/chat/components/chatThread/chatThreadListItem.tsx b/packages/web/src/features/chat/components/chatThread/chatThreadListItem.tsx index 66d13096..ae3d5c84 100644 --- a/packages/web/src/features/chat/components/chatThread/chatThreadListItem.tsx +++ b/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 { - console.log(`re-rendering chat thread list item`, index); const leftPanelRef = useRef(null); const [leftPanelHeight, setLeftPanelHeight] = useState(null); const answerRef = useRef(null); @@ -81,7 +81,6 @@ export const ChatThreadListItemComponent = forwardRef { + 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 (
- {references.length > 0 ? ( + {referencedFileSources.length > 0 ? ( !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[]) => { diff --git a/packages/web/src/features/chat/components/chatThread/detailsCard.tsx b/packages/web/src/features/chat/components/chatThread/detailsCard.tsx index 5cd6f335..26d16437 100644 --- a/packages/web/src/features/chat/components/chatThread/detailsCard.tsx +++ b/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 ( @@ -212,4 +212,4 @@ const DetailsCardComponent = ({ ) } -export const DetailsCard = memo(DetailsCardComponent); \ No newline at end of file +export const DetailsCard = memo(DetailsCardComponent, isEqual); \ No newline at end of file diff --git a/packages/web/src/features/chat/components/chatThread/markdownRenderer.tsx b/packages/web/src/features/chat/components/chatThread/markdownRenderer.tsx index b69145f0..8232af3c 100644 --- a/packages/web/src/features/chat/components/chatThread/markdownRenderer.tsx +++ b/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 }, ) => ReturnType; diff --git a/packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx b/packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx index 7792bb3a..d223fd37 100644 --- a/packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx +++ b/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(); - 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 (
No file references found @@ -215,7 +193,7 @@ const ReferencedSourcesListViewComponent = ({ >
{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); diff --git a/packages/web/src/features/chat/utils.ts b/packages/web/src/features/chat/utils.ts index f339b872..c64f1ed3 100644 --- a/packages/web/src/features/chat/utils.ts +++ b/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) + ); +} \ No newline at end of file diff --git a/yarn.lock b/yarn.lock index 41e3339f..57002159 100644 --- a/yarn.lock +++ b/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"