ti.sil/sourcebot
1
0
Fork 0
mirror of https://github.com/sourcebot-dev/sourcebot.git synced 2025-12-11 20:05:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

memoize all of the things

Browse source
This commit is contained in:
bkellam 2025-11-27 16:05:06 -08:00
parent da3c93e05a
commit 4b7798a94e
10 changed files with 88 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
packages/web/package.json
View file

@ -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",

7
packages/web/src/features/chat/components/chatBox/chatBox.tsx
View file

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

3
packages/web/src/features/chat/components/chatThread/answerCard.tsx
View file

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

59
packages/web/src/features/chat/components/chatThread/chatThreadListItem.tsx
View file

@ -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[]) => {

4
packages/web/src/features/chat/components/chatThread/detailsCard.tsx
View file

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

3
packages/web/src/features/chat/components/chatThread/markdownRenderer.tsx
View file

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

3
packages/web/src/features/chat/components/chatThread/referencedFileSourceListItem.tsx
View file

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

48
packages/web/src/features/chat/components/chatThread/referencedSourcesListView.tsx
View file

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

10
packages/web/src/features/chat/utils.ts
View file

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

1
yarn.lock
View file

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