pref: message render performence (#81)

* fix: message card always unmount when messages change

* pref: add useShallow for complex store selector

web/src/app/chat/components/message-list-view.tsx

@@ -27,8 +27,11 @@ import type { Message, Option } from "~/core/messages";
 import {
   closeResearch,
   openResearch,
+  useLastFeedbackMessageId,
+  useLastInterruptMessage,
   useMessage,
-  useResearchTitle,
+  useMessageIds,
+  useResearchMessage,
   useStore,
 } from "~/core/store";
 import { parseJSON } from "~/core/utils";
@@ -47,27 +50,9 @@ export function MessageListView({
   ) => void;
 }) {
   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 noOngoingResearch = useStore(
     (state) => state.ongoingResearchId === null,
@@ -138,9 +123,10 @@ function MessageListItem({
   onToggleResearch?: () => void;
 }) {
   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.role === "user" ||
@@ -214,16 +200,17 @@ function MessageListItem({
     }
     return null;
   }
+}
 
-  function MessageBubble({
+function MessageBubble({
   className,
   message,
   children,
-  }: {
+}: {
   className?: string;
   message: Message;
   children: React.ReactNode;
-  }) {
+}) {
   return (
     <div
       className={cn(
@@ -237,25 +224,21 @@ function MessageListItem({
       {children}
     </div>
   );
-  }
+}
 
-  function ResearchCard({
+function ResearchCard({
   className,
   researchId,
   onToggleResearch,
-  }: {
+}: {
   className?: string;
   researchId: string;
   onToggleResearch?: () => void;
-  }) {
+}) {
-    const reportId = useStore((state) =>
+  const reportId = useStore((state) => state.researchReportIds.get(researchId));
-      state.researchReportIds.get(researchId),
+  const hasReport = reportId !== undefined;
-    );
-    const hasReport = useStore((state) =>
-      state.researchReportIds.has(researchId),
-    );
   const reportGenerating = useStore(
-      (state) => hasReport && state.messages.get(reportId!)!.isStreaming,
+    (state) => hasReport && state.messages.get(reportId)!.isStreaming,
   );
   const openResearchId = useStore((state) => state.openResearchId);
   const state = useMemo(() => {
@@ -264,7 +247,13 @@ function MessageListItem({
     }
     return "Researching...";
   }, [hasReport, reportGenerating]);
-    const title = useResearchTitle(researchId);
+  const msg = useResearchMessage(researchId);
+  const title = useMemo(() => {
+    if (msg) {
+      return parseJSON(msg.content ?? "", { title: "" }).title;
+    }
+    return undefined;
+  }, [msg]);
   const handleOpen = useCallback(() => {
     if (openResearchId === researchId) {
       closeResearch();
@@ -297,7 +286,6 @@ function MessageListItem({
       </CardFooter>
     </Card>
   );
-  }
 }
 
 const GREETINGS = ["Cool", "Sounds great", "Looks good", "Great", "Awesome"];

web/src/app/chat/components/messages-block.tsx

@@ -17,7 +17,7 @@ import { fastForwardReplay } from "~/core/api";
 import { useReplayMetadata } from "~/core/api/hooks";
 import type { Option } from "~/core/messages";
 import { useReplay } from "~/core/replay";
-import { sendMessage, useStore } from "~/core/store";
+import { sendMessage, useMessageIds, useStore } from "~/core/store";
 import { env } from "~/env";
 import { cn } from "~/lib/utils";
 
@@ -27,7 +27,8 @@ import { MessageListView } from "./message-list-view";
 import { Welcome } from "./welcome";
 
 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 { isReplay } = useReplay();
   const { title: replayTitle, hasError: replayHasError } = useReplayMetadata();

web/src/core/store/store.ts

@@ -4,6 +4,7 @@
 import { nanoid } from "nanoid";
 import { toast } from "sonner";
 import { create } from "zustand";
+import { useShallow } from "zustand/react/shallow";
 
 import { chatStream, generatePodcast } from "../api";
 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) {
-  return useStore((state) =>
+  return useStore(
+    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;
+}