pref: message render performence (#81)

* fix: message card always unmount when messages change

* pref: add useShallow for complex store selector
This commit is contained in:
JeffJiang
2025-05-12 20:21:54 +08:00
committed by GitHub
parent 9266201fe5
commit 229b59ab88
3 changed files with 144 additions and 117 deletions

View File

@@ -27,8 +27,11 @@ import type { Message, Option } from "~/core/messages";
import { import {
closeResearch, closeResearch,
openResearch, openResearch,
useLastFeedbackMessageId,
useLastInterruptMessage,
useMessage, useMessage,
useResearchTitle, useMessageIds,
useResearchMessage,
useStore, useStore,
} from "~/core/store"; } from "~/core/store";
import { parseJSON } from "~/core/utils"; import { parseJSON } from "~/core/utils";
@@ -47,27 +50,9 @@ export function MessageListView({
) => void; ) => void;
}) { }) {
const scrollContainerRef = useRef<ScrollContainerRef>(null); const scrollContainerRef = useRef<ScrollContainerRef>(null);
const messageIds = useStore((state) => state.messageIds); const messageIds = useMessageIds();
const interruptMessage = useStore((state) => { const interruptMessage = useLastInterruptMessage();
if (messageIds.length >= 2) { const waitingForFeedbackMessageId = useLastFeedbackMessageId();
const lastMessage = state.messages.get(
messageIds[messageIds.length - 1]!,
);
return lastMessage?.finishReason === "interrupt" ? lastMessage : null;
}
return null;
});
const waitingForFeedbackMessageId = useStore((state) => {
if (messageIds.length >= 2) {
const lastMessage = state.messages.get(
messageIds[messageIds.length - 1]!,
);
if (lastMessage && lastMessage.finishReason === "interrupt") {
return state.messageIds[state.messageIds.length - 2];
}
}
return null;
});
const responding = useStore((state) => state.responding); const responding = useStore((state) => state.responding);
const noOngoingResearch = useStore( const noOngoingResearch = useStore(
(state) => state.ongoingResearchId === null, (state) => state.ongoingResearchId === null,
@@ -138,9 +123,10 @@ function MessageListItem({
onToggleResearch?: () => void; onToggleResearch?: () => void;
}) { }) {
const message = useMessage(messageId); const message = useMessage(messageId);
const startOfResearch = useStore((state) => const researchIds = useStore((state) => state.researchIds);
state.researchIds.includes(messageId), const startOfResearch = useMemo(() => {
); return researchIds.includes(messageId);
}, [researchIds, messageId]);
if (message) { if (message) {
if ( if (
message.role === "user" || message.role === "user" ||
@@ -214,90 +200,92 @@ function MessageListItem({
} }
return null; return null;
} }
}
function MessageBubble({ function MessageBubble({
className, className,
message, message,
children, children,
}: { }: {
className?: string; className?: string;
message: Message; message: Message;
children: React.ReactNode; children: React.ReactNode;
}) { }) {
return ( return (
<div <div
className={cn( className={cn(
`flex w-fit max-w-[85%] flex-col rounded-2xl px-4 py-3 shadow`, `flex w-fit max-w-[85%] flex-col rounded-2xl px-4 py-3 shadow`,
message.role === "user" && message.role === "user" &&
"text-primary-foreground bg-brand rounded-ee-none", "text-primary-foreground bg-brand rounded-ee-none",
message.role === "assistant" && "bg-card rounded-es-none", message.role === "assistant" && "bg-card rounded-es-none",
className, className,
)} )}
> >
{children} {children}
</div> </div>
); );
} }
function ResearchCard({ function ResearchCard({
className, className,
researchId, researchId,
onToggleResearch, onToggleResearch,
}: { }: {
className?: string; className?: string;
researchId: string; researchId: string;
onToggleResearch?: () => void; onToggleResearch?: () => void;
}) { }) {
const reportId = useStore((state) => const reportId = useStore((state) => state.researchReportIds.get(researchId));
state.researchReportIds.get(researchId), const hasReport = reportId !== undefined;
); const reportGenerating = useStore(
const hasReport = useStore((state) => (state) => hasReport && state.messages.get(reportId)!.isStreaming,
state.researchReportIds.has(researchId), );
); const openResearchId = useStore((state) => state.openResearchId);
const reportGenerating = useStore( const state = useMemo(() => {
(state) => hasReport && state.messages.get(reportId!)!.isStreaming, if (hasReport) {
); return reportGenerating ? "Generating report..." : "Report generated";
const openResearchId = useStore((state) => state.openResearchId); }
const state = useMemo(() => { return "Researching...";
if (hasReport) { }, [hasReport, reportGenerating]);
return reportGenerating ? "Generating report..." : "Report generated"; const msg = useResearchMessage(researchId);
} const title = useMemo(() => {
return "Researching..."; if (msg) {
}, [hasReport, reportGenerating]); return parseJSON(msg.content ?? "", { title: "" }).title;
const title = useResearchTitle(researchId); }
const handleOpen = useCallback(() => { return undefined;
if (openResearchId === researchId) { }, [msg]);
closeResearch(); const handleOpen = useCallback(() => {
} else { if (openResearchId === researchId) {
openResearch(researchId); closeResearch();
} } else {
onToggleResearch?.(); openResearch(researchId);
}, [openResearchId, researchId, onToggleResearch]); }
return ( onToggleResearch?.();
<Card className={cn("w-full", className)}> }, [openResearchId, researchId, onToggleResearch]);
<CardHeader> return (
<CardTitle> <Card className={cn("w-full", className)}>
<RainbowText animated={state !== "Report generated"}> <CardHeader>
{title !== undefined && title !== "" ? title : "Deep Research"} <CardTitle>
</RainbowText> <RainbowText animated={state !== "Report generated"}>
</CardTitle> {title !== undefined && title !== "" ? title : "Deep Research"}
</CardHeader> </RainbowText>
<CardFooter> </CardTitle>
<div className="flex w-full"> </CardHeader>
<RollingText className="text-muted-foreground flex-grow text-sm"> <CardFooter>
{state} <div className="flex w-full">
</RollingText> <RollingText className="text-muted-foreground flex-grow text-sm">
<Button {state}
variant={!openResearchId ? "default" : "outline"} </RollingText>
onClick={handleOpen} <Button
> variant={!openResearchId ? "default" : "outline"}
{researchId !== openResearchId ? "Open" : "Close"} onClick={handleOpen}
</Button> >
</div> {researchId !== openResearchId ? "Open" : "Close"}
</CardFooter> </Button>
</Card> </div>
); </CardFooter>
} </Card>
);
} }
const GREETINGS = ["Cool", "Sounds great", "Looks good", "Great", "Awesome"]; const GREETINGS = ["Cool", "Sounds great", "Looks good", "Great", "Awesome"];

View File

@@ -17,7 +17,7 @@ import { fastForwardReplay } from "~/core/api";
import { useReplayMetadata } from "~/core/api/hooks"; import { useReplayMetadata } from "~/core/api/hooks";
import type { Option } from "~/core/messages"; import type { Option } from "~/core/messages";
import { useReplay } from "~/core/replay"; import { useReplay } from "~/core/replay";
import { sendMessage, useStore } from "~/core/store"; import { sendMessage, useMessageIds, useStore } from "~/core/store";
import { env } from "~/env"; import { env } from "~/env";
import { cn } from "~/lib/utils"; import { cn } from "~/lib/utils";
@@ -27,7 +27,8 @@ import { MessageListView } from "./message-list-view";
import { Welcome } from "./welcome"; import { Welcome } from "./welcome";
export function MessagesBlock({ className }: { className?: string }) { export function MessagesBlock({ className }: { className?: string }) {
const messageCount = useStore((state) => state.messageIds.length); const messageIds = useMessageIds();
const messageCount = messageIds.length;
const responding = useStore((state) => state.responding); const responding = useStore((state) => state.responding);
const { isReplay } = useReplay(); const { isReplay } = useReplay();
const { title: replayTitle, hasError: replayHasError } = useReplayMetadata(); const { title: replayTitle, hasError: replayHasError } = useReplayMetadata();

View File

@@ -4,6 +4,7 @@
import { nanoid } from "nanoid"; import { nanoid } from "nanoid";
import { toast } from "sonner"; import { toast } from "sonner";
import { create } from "zustand"; import { create } from "zustand";
import { useShallow } from "zustand/react/shallow";
import { chatStream, generatePodcast } from "../api"; import { chatStream, generatePodcast } from "../api";
import type { Message } from "../messages"; import type { Message } from "../messages";
@@ -305,17 +306,54 @@ export async function listenToPodcast(researchId: string) {
} }
} }
export function useResearchTitle(researchId: string) { export function useResearchMessage(researchId: string) {
const planMessage = useMessage( return useStore(
useStore.getState().researchPlanIds.get(researchId), useShallow((state) => {
const messageId = state.researchPlanIds.get(researchId);
return messageId ? state.messages.get(messageId) : undefined;
}),
); );
return planMessage
? parseJSON(planMessage.content, { title: "" }).title
: undefined;
} }
export function useMessage(messageId: string | null | undefined) { export function useMessage(messageId: string | null | undefined) {
return useStore((state) => return useStore(
messageId ? state.messages.get(messageId) : undefined, useShallow((state) =>
messageId ? state.messages.get(messageId) : undefined,
),
); );
} }
export function useMessageIds() {
return useStore(useShallow((state) => state.messageIds));
}
export function useLastInterruptMessage() {
return useStore(
useShallow((state) => {
if (state.messageIds.length >= 2) {
const lastMessage = state.messages.get(
state.messageIds[state.messageIds.length - 1]!,
);
return lastMessage?.finishReason === "interrupt" ? lastMessage : null;
}
return null;
}),
);
}
export function useLastFeedbackMessageId() {
const waitingForFeedbackMessageId = useStore(
useShallow((state) => {
if (state.messageIds.length >= 2) {
const lastMessage = state.messages.get(
state.messageIds[state.messageIds.length - 1]!,
);
if (lastMessage && lastMessage.finishReason === "interrupt") {
return state.messageIds[state.messageIds.length - 2];
}
}
return null;
}),
);
return waitingForFeedbackMessageId;
}