Message list virtualization (#1993)

package-lock.json

 {
   "name": "dyad",
-  "version": "0.30.0-beta.1",
+  "version": "0.31.0-beta.1",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "dyad",
-      "version": "0.30.0-beta.1",
+      "version": "0.31.0-beta.1",
       "license": "MIT",
       "dependencies": {
         "@ai-sdk/amazon-bedrock": "^3.0.15",

src/components/chat/MessagesList.tsx

-import type React from "react";
+import React from "react";
 import type { Message } from "@/ipc/ipc_types";
-import { forwardRef, useState } from "react";
+import { forwardRef, useState, useCallback, useMemo } from "react";
+import { Virtuoso } from "react-virtuoso";
 import ChatMessage from "./ChatMessage";
 import { OpenRouterSetupBanner, SetupBanner } from "../SetupBanner";
 
@@ -26,60 +27,56 @@ interface MessagesListProps {
   messagesEndRef: React.RefObject<HTMLDivElement | null>;
 }
 
-export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
-  function MessagesList({ messages, messagesEndRef }, ref) {
-    const appId = useAtomValue(selectedAppIdAtom);
-    const { versions, revertVersion } = useVersions(appId);
-    const { streamMessage, isStreaming } = useStreamChat();
-    const { isAnyProviderSetup, isProviderSetup } = useLanguageModelProviders();
-    const { settings } = useSettings();
-    const setMessagesById = useSetAtom(chatMessagesByIdAtom);
-    const [isUndoLoading, setIsUndoLoading] = useState(false);
-    const [isRetryLoading, setIsRetryLoading] = useState(false);
-    const selectedChatId = useAtomValue(selectedChatIdAtom);
-    const { userBudget } = useUserBudgetInfo();
-    // Only fetch token count when not streaming
-    const { result: tokenCountResult } = useCountTokens(
-      !isStreaming ? selectedChatId : null,
-      "",
-    );
+// Memoize ChatMessage at module level to prevent recreation on every render
+const MemoizedChatMessage = React.memo(ChatMessage);
 
-    const renderSetupBanner = () => {
-      const selectedModel = settings?.selectedModel;
-      if (
-        selectedModel?.name === "free" &&
-        selectedModel?.provider === "auto" &&
-        !isProviderSetup("openrouter")
-      ) {
-        return <OpenRouterSetupBanner className="w-full" />;
-      }
-      if (!isAnyProviderSetup()) {
-        return <SetupBanner />;
-      }
-      return null;
-    };
+// Context type for Virtuoso
+interface FooterContext {
+  messages: Message[];
+  messagesEndRef: React.RefObject<HTMLDivElement | null>;
+  isStreaming: boolean;
+  tokenCountResult: ReturnType<typeof useCountTokens>["result"];
+  isUndoLoading: boolean;
+  isRetryLoading: boolean;
+  setIsUndoLoading: (loading: boolean) => void;
+  setIsRetryLoading: (loading: boolean) => void;
+  versions: ReturnType<typeof useVersions>["versions"];
+  revertVersion: ReturnType<typeof useVersions>["revertVersion"];
+  streamMessage: ReturnType<typeof useStreamChat>["streamMessage"];
+  selectedChatId: number | null;
+  appId: number | null;
+  setMessagesById: ReturnType<typeof useSetAtom<typeof chatMessagesByIdAtom>>;
+  settings: ReturnType<typeof useSettings>["settings"];
+  userBudget: ReturnType<typeof useUserBudgetInfo>["userBudget"];
+  renderSetupBanner: () => React.ReactNode;
+}
+
+// Footer component for Virtuoso - receives context via props
+function FooterComponent({ context }: { context?: FooterContext }) {
+  if (!context) return null;
+
+  const {
+    messages,
+    messagesEndRef,
+    isStreaming,
+    tokenCountResult,
+    isUndoLoading,
+    isRetryLoading,
+    setIsUndoLoading,
+    setIsRetryLoading,
+    versions,
+    revertVersion,
+    streamMessage,
+    selectedChatId,
+    appId,
+    setMessagesById,
+    settings,
+    userBudget,
+    renderSetupBanner,
+  } = context;
 
   return (
-      <div
-        className="absolute inset-0 overflow-y-auto p-4"
-        ref={ref}
-        data-testid="messages-list"
-      >
-        {messages.length > 0
-          ? messages.map((message, index) => (
-              <ChatMessage
-                key={index}
-                message={message}
-                isLastMessage={index === messages.length - 1}
-              />
-            ))
-          : !renderSetupBanner() && (
-              <div className="flex flex-col items-center justify-center h-full max-w-2xl mx-auto">
-                <div className="flex flex-1 items-center justify-center text-gray-500">
-                  No messages yet
-                </div>
-              </div>
-            )}
+    <>
       {/* Show context limit banner when close to token limit */}
       {!isStreaming && tokenCountResult && (
         <ContextLimitBanner
@@ -87,6 +84,7 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
           contextWindow={tokenCountResult.contextWindow}
         />
       )}
+
       {!isStreaming && (
         <div className="flex max-w-3xl mx-auto gap-2">
           {!!messages.length &&
@@ -118,9 +116,7 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
                           versionId: previousAssistantMessage.commitHash,
                         });
                         const chat =
-                            await IpcClient.getInstance().getChat(
-                              selectedChatId,
-                            );
+                          await IpcClient.getInstance().getChat(selectedChatId);
                         setMessagesById((prev) => {
                           const next = new Map(prev);
                           next.set(selectedChatId, chat.messages);
@@ -195,9 +191,7 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
                       previousAssistantMessage?.role === "assistant" &&
                       previousAssistantMessage?.commitHash
                     ) {
-                        console.debug(
-                          "Reverting to previous assistant version",
-                        );
+                      console.debug("Reverting to previous assistant version");
                       await revertVersion({
                         versionId: previousAssistantMessage.commitHash,
                       });
@@ -267,6 +261,190 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
         )}
       <div ref={messagesEndRef} />
       {renderSetupBanner()}
+    </>
+  );
+}
+
+export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
+  function MessagesList({ messages, messagesEndRef }, ref) {
+    const appId = useAtomValue(selectedAppIdAtom);
+    const { versions, revertVersion } = useVersions(appId);
+    const { streamMessage, isStreaming } = useStreamChat();
+    const { isAnyProviderSetup, isProviderSetup } = useLanguageModelProviders();
+    const { settings } = useSettings();
+    const setMessagesById = useSetAtom(chatMessagesByIdAtom);
+    const [isUndoLoading, setIsUndoLoading] = useState(false);
+    const [isRetryLoading, setIsRetryLoading] = useState(false);
+    const selectedChatId = useAtomValue(selectedChatIdAtom);
+    const { userBudget } = useUserBudgetInfo();
+
+    // Virtualization only renders visible DOM elements, which creates issues for E2E tests:
+    // 1. Off-screen logs don't exist in the DOM and can't be queried by test selectors
+    // 2. Tests would need complex scrolling logic to bring elements into view before interaction
+    // 3. Race conditions and timing issues occur when waiting for virtualized elements to render after scrolling
+    const isTestMode = settings?.isTestMode;
+    // Only fetch token count when not streaming
+    const { result: tokenCountResult } = useCountTokens(
+      !isStreaming ? selectedChatId : null,
+      "",
+    );
+
+    // Wrap state setters in useCallback to stabilize references
+    const handleSetIsUndoLoading = useCallback((loading: boolean) => {
+      setIsUndoLoading(loading);
+    }, []);
+
+    const handleSetIsRetryLoading = useCallback((loading: boolean) => {
+      setIsRetryLoading(loading);
+    }, []);
+
+    // Stabilize renderSetupBanner with proper dependencies
+    const renderSetupBanner = useCallback(() => {
+      const selectedModel = settings?.selectedModel;
+      if (
+        selectedModel?.name === "free" &&
+        selectedModel?.provider === "auto" &&
+        !isProviderSetup("openrouter")
+      ) {
+        return <OpenRouterSetupBanner className="w-full" />;
+      }
+      if (!isAnyProviderSetup()) {
+        return <SetupBanner />;
+      }
+      return null;
+    }, [
+      settings?.selectedModel?.name,
+      settings?.selectedModel?.provider,
+      isProviderSetup,
+      isAnyProviderSetup,
+    ]);
+
+    // Memoized item renderer for virtualized list
+    const itemContent = useCallback(
+      (index: number, message: Message) => {
+        const isLastMessage = index === messages.length - 1;
+        const messageKey = message.id;
+
+        return (
+          <div className="px-4" key={messageKey}>
+            <MemoizedChatMessage
+              message={message}
+              isLastMessage={isLastMessage}
+            />
+          </div>
+        );
+      },
+      [messages.length],
+    );
+
+    // Create context object for Footer component with stable references
+    const footerContext = useMemo<FooterContext>(
+      () => ({
+        messages,
+        messagesEndRef,
+        isStreaming,
+        tokenCountResult,
+        isUndoLoading,
+        isRetryLoading,
+        setIsUndoLoading: handleSetIsUndoLoading,
+        setIsRetryLoading: handleSetIsRetryLoading,
+        versions,
+        revertVersion,
+        streamMessage,
+        selectedChatId,
+        appId,
+        setMessagesById,
+        settings,
+        userBudget,
+        renderSetupBanner,
+      }),
+      [
+        messages,
+        messagesEndRef,
+        isStreaming,
+        tokenCountResult,
+        isUndoLoading,
+        isRetryLoading,
+        handleSetIsUndoLoading,
+        handleSetIsRetryLoading,
+        versions,
+        revertVersion,
+        streamMessage,
+        selectedChatId,
+        appId,
+        setMessagesById,
+        settings,
+        userBudget,
+        renderSetupBanner,
+      ],
+    );
+
+    // Render empty state or setup banner
+    if (messages.length === 0) {
+      const setupBanner = renderSetupBanner();
+      if (setupBanner) {
+        return (
+          <div
+            className="absolute inset-0 overflow-y-auto p-4"
+            ref={ref}
+            data-testid="messages-list"
+          >
+            {setupBanner}
+          </div>
+        );
+      }
+      return (
+        <div
+          className="absolute inset-0 overflow-y-auto p-4"
+          ref={ref}
+          data-testid="messages-list"
+        >
+          <div className="flex flex-col items-center justify-center h-full max-w-2xl mx-auto">
+            <div className="flex flex-1 items-center justify-center text-gray-500">
+              No messages yet
+            </div>
+          </div>
+        </div>
+      );
+    }
+
+    // In test mode, render all messages without virtualization
+    // so E2E tests can query all messages in the DOM
+    if (isTestMode) {
+      return (
+        <div
+          className="absolute inset-0 p-4 overflow-y-auto"
+          ref={ref}
+          data-testid="messages-list"
+        >
+          {messages.map((message, index) => {
+            const isLastMessage = index === messages.length - 1;
+            return (
+              <div className="px-4" key={message.id}>
+                <ChatMessage message={message} isLastMessage={isLastMessage} />
+              </div>
+            );
+          })}
+          <FooterComponent context={footerContext} />
+        </div>
+      );
+    }
+
+    return (
+      <div
+        className="absolute inset-0 p-4"
+        ref={ref}
+        data-testid="messages-list"
+      >
+        <Virtuoso
+          data={messages}
+          increaseViewportBy={{ top: 1000, bottom: 500 }}
+          initialTopMostItemIndex={messages.length - 1}
+          followOutput="smooth"
+          itemContent={itemContent}
+          components={{ Footer: FooterComponent }}
+          context={footerContext}
+        />
       </div>
     );
   },