Better virtualization (#2653)

web/src/app/chat/ChatPage.tsx

@@ -759,7 +759,15 @@ export function ChatPage({
     setAboveHorizon(scrollDist.current > 500);
   };
 
-  scrollableDivRef?.current?.addEventListener("scroll", updateScrollTracking);
+  useEffect(() => {
+    const scrollableDiv = scrollableDivRef.current;
+    if (scrollableDiv) {
+      scrollableDiv.addEventListener("scroll", updateScrollTracking);
+      return () => {
+        scrollableDiv.removeEventListener("scroll", updateScrollTracking);
+      };
+    }
+  }, []);
 
   const handleInputResize = () => {
     setTimeout(() => {
@@ -1137,7 +1145,9 @@ export function ChatPage({
 
       await delay(50);
       while (!stack.isComplete || !stack.isEmpty()) {
-        await delay(0.5);
+        if (stack.isEmpty()) {
+          await delay(0.5);
+        }
 
         if (!stack.isEmpty() && !controller.signal.aborted) {
           const packet = stack.nextPacket();

web/src/app/chat/message/CodeBlock.tsx

@@ -1,20 +1,22 @@
 import React, { useState, ReactNode, useCallback, useMemo, memo } from "react";
 import { FiCheck, FiCopy } from "react-icons/fi";
 
-const CODE_BLOCK_PADDING_TYPE = { padding: "1rem" };
+const CODE_BLOCK_PADDING = { padding: "1rem" };
 
 interface CodeBlockProps {
-  className?: string | undefined;
+  className?: string;
   children?: ReactNode;
-  content: string;
+  codeText: string;
-  [key: string]: any;
 }
 
+const MemoizedCodeLine = memo(({ content }: { content: ReactNode }) => (
+  <>{content}</>
+));
+
 export const CodeBlock = memo(function CodeBlock({
   className = "",
   children,
-  content,
+  codeText,
-  ...props
 }: CodeBlockProps) {
   const [copied, setCopied] = useState(false);
 
@@ -26,132 +28,99 @@ export const CodeBlock = memo(function CodeBlock({
       .join(" ");
   }, [className]);
 
-  const codeText = useMemo(() => {
+  const handleCopy = useCallback(() => {
-    let codeText: string | null = null;
+    if (!codeText) return;
-    if (
+    navigator.clipboard.writeText(codeText).then(() => {
-      props.node?.position?.start?.offset &&
+      setCopied(true);
-      props.node?.position?.end?.offset
+      setTimeout(() => setCopied(false), 2000);
-    ) {
+    });
-      codeText = content.slice(
+  }, [codeText]);
-        props.node.position.start.offset,
-        props.node.position.end.offset
-      );
-      codeText = codeText.trim();
 
-      // Find the last occurrence of closing backticks
+  const CopyButton = memo(() => (
-      const lastBackticksIndex = codeText.lastIndexOf("```");
+    <div
-      if (lastBackticksIndex !== -1) {
+      className="ml-auto cursor-pointer select-none"
-        codeText = codeText.slice(0, lastBackticksIndex + 3);
+      onMouseDown={handleCopy}
-      }
+    >
+      {copied ? (
+        <div className="flex items-center space-x-2">
+          <FiCheck size={16} />
+          <span>Copied!</span>
+        </div>
+      ) : (
+        <div className="flex items-center space-x-2">
+          <FiCopy size={16} />
+          <span>Copy code</span>
+        </div>
+      )}
+    </div>
+  ));
+  CopyButton.displayName = "CopyButton";
 
-      // Remove the language declaration and trailing backticks
+  const CodeContent = memo(() => {
-      const codeLines = codeText.split("\n");
+    if (!language) {
-      if (
+      if (typeof children === "string") {
-        codeLines.length > 1 &&
+        return (
-        (codeLines[0].startsWith("```") ||
+          <code
-          codeLines[0].trim().startsWith("```"))
+            className={`
-      ) {
+            font-mono 
-        codeLines.shift(); // Remove the first line with the language declaration
+            text-gray-800 
-        if (
+            bg-gray-50 
-          codeLines[codeLines.length - 1] === "```" ||
+            border 
-          codeLines[codeLines.length - 1]?.trim() === "```"
+            border-gray-300 
-        ) {
+            rounded 
-          codeLines.pop(); // Remove the last line with the trailing backticks
+            px-1
-        }
+            py-[3px]
-
+            text-xs 
-        const minIndent = codeLines
+            whitespace-pre-wrap 
-          .filter((line) => line.trim().length > 0)
+            break-words 
-          .reduce((min, line) => {
+            overflow-hidden
-            const match = line.match(/^\s*/);
+            mb-1
-            return Math.min(min, match ? match[0].length : 0);
+            ${className}
-          }, Infinity);
+          `}
-
+          >
-        const formattedCodeLines = codeLines.map((line) =>
+            {children}
-          line.slice(minIndent)
+          </code>
         );
-        codeText = formattedCodeLines.join("\n");
       }
-    }
+      return (
-
+        <pre style={CODE_BLOCK_PADDING}>
-    // handle unknown languages. They won't have a `node.position.start.offset`
+          <code className={`text-sm ${className}`}>
-    if (!codeText) {
+            {Array.isArray(children)
-      const findTextNode = (node: any): string | null => {
+              ? children.map((child, index) => (
-        if (node.type === "text") {
+                  <MemoizedCodeLine key={index} content={child} />
-          return node.value;
+                ))
-        }
+              : children}
-        let finalResult = "";
+          </code>
-        if (node.children) {
+        </pre>
-          for (const child of node.children) {
+      );
-            const result = findTextNode(child);
-            if (result) {
-              finalResult += result;
-            }
-          }
-        }
-        return finalResult;
-      };
-
-      codeText = findTextNode(props.node);
-    }
-
-    return codeText;
-  }, [content, props.node]);
-
-  const handleCopy = useCallback(
-    (event: React.MouseEvent) => {
-      event.preventDefault();
-      if (!codeText) {
-        return;
-      }
-
-      navigator.clipboard.writeText(codeText).then(() => {
-        setCopied(true);
-        setTimeout(() => setCopied(false), 2000);
-      });
-    },
-    [codeText]
-  );
-
-  if (!language) {
-    if (typeof children === "string") {
-      return <code className={className}>{children}</code>;
     }
 
     return (
-      <pre style={CODE_BLOCK_PADDING_TYPE}>
+      <pre className="overflow-x-scroll" style={CODE_BLOCK_PADDING}>
-        <code {...props} className={`text-sm ${className}`}>
+        <code className="text-xs overflow-x-auto">
-          {children}
+          {Array.isArray(children)
+            ? children.map((child, index) => (
+                <MemoizedCodeLine key={index} content={child} />
+              ))
+            : children}
         </code>
       </pre>
     );
-  }
+  });
+  CodeContent.displayName = "CodeContent";
 
   return (
     <div className="overflow-x-hidden">
-      <div className="flex mx-3 py-2 text-xs">
+      {language && (
-        {language}
+        <div className="flex mx-3 py-2 text-xs">
-        {codeText && (
+          {language}
-          <div
+          {codeText && <CopyButton />}
-            className="ml-auto cursor-pointer select-none"
+        </div>
-            onMouseDown={handleCopy}
+      )}
-          >
+      <CodeContent />
-            {copied ? (
-              <div className="flex items-center space-x-2">
-                <FiCheck size={16} />
-                <span>Copied!</span>
-              </div>
-            ) : (
-              <div className="flex items-center space-x-2">
-                <FiCopy size={16} />
-                <span>Copy code</span>
-              </div>
-            )}
-          </div>
-        )}
-      </div>
-      <pre {...props} className="overflow-x-scroll" style={{ padding: "1rem" }}>
-        <code className={`text-xs overflow-x-auto `}>{children}</code>
-      </pre>
     </div>
   );
 });
+
+CodeBlock.displayName = "CodeBlock";
+MemoizedCodeLine.displayName = "MemoizedCodeLine";

web/src/app/chat/message/MemoizedTextComponents.tsx

@@ -25,9 +25,9 @@ export const MemoizedLink = memo((props: any) => {
   }
 });
 
-export const MemoizedParagraph = memo(({ node, ...props }: any) => (
+export const MemoizedParagraph = memo(({ ...props }: any) => {
-  <p {...props} className="text-default" />
+  return <p {...props} className="text-default" />;
-));
+});
 
 MemoizedLink.displayName = "MemoizedLink";
 MemoizedParagraph.displayName = "MemoizedParagraph";

web/src/app/chat/message/Messages.tsx

@@ -54,6 +54,7 @@ import RegenerateOption from "../RegenerateOption";
 import { LlmOverride } from "@/lib/hooks";
 import { ContinueGenerating } from "./ContinueMessage";
 import { MemoizedLink, MemoizedParagraph } from "./MemoizedTextComponents";
+import { extractCodeText } from "./codeUtils";
 
 const TOOLS_WITH_CUSTOM_HANDLING = [
   SEARCH_TOOL_NAME,
@@ -253,6 +254,40 @@ export const AIMessage = ({
     new Set((docs || []).map((doc) => doc.source_type))
   ).slice(0, 3);
 
+  const markdownComponents = useMemo(
+    () => ({
+      a: MemoizedLink,
+      p: MemoizedParagraph,
+      code: ({ node, inline, className, children, ...props }: any) => {
+        const codeText = extractCodeText(
+          node,
+          finalContent as string,
+          children
+        );
+
+        return (
+          <CodeBlock className={className} codeText={codeText}>
+            {children}
+          </CodeBlock>
+        );
+      },
+    }),
+    [messageId, content]
+  );
+
+  const renderedMarkdown = useMemo(() => {
+    return (
+      <ReactMarkdown
+        className="prose max-w-full text-base"
+        components={markdownComponents}
+        remarkPlugins={[remarkGfm]}
+        rehypePlugins={[[rehypePrism, { ignoreMissing: true }]]}
+      >
+        {finalContent as string}
+      </ReactMarkdown>
+    );
+  }, [finalContent]);
+
   const includeMessageSwitcher =
     currentMessageInd !== undefined &&
     onMessageSelection &&
@@ -352,27 +387,7 @@ export const AIMessage = ({
 
                         {typeof content === "string" ? (
                           <div className="overflow-x-visible max-w-content-max">
-                            <ReactMarkdown
+                            {renderedMarkdown}
-                              key={messageId}
-                              className="prose max-w-full text-base"
-                              components={{
-                                a: MemoizedLink,
-                                p: MemoizedParagraph,
-                                code: (props) => (
-                                  <CodeBlock
-                                    className="w-full"
-                                    {...props}
-                                    content={content as string}
-                                  />
-                                ),
-                              }}
-                              remarkPlugins={[remarkGfm]}
-                              rehypePlugins={[
-                                [rehypePrism, { ignoreMissing: true }],
-                              ]}
-                            >
-                              {finalContent as string}
-                            </ReactMarkdown>
                           </div>
                         ) : (
                           content

web/src/app/chat/message/codeUtils.ts

@@ -0,0 +1,47 @@
+export function extractCodeText(
+  node: any,
+  content: string,
+  children: React.ReactNode
+): string {
+  let codeText: string | null = null;
+  if (
+    node?.position?.start?.offset != null &&
+    node?.position?.end?.offset != null
+  ) {
+    codeText = content.slice(
+      node.position.start.offset,
+      node.position.end.offset
+    );
+    codeText = codeText.trim();
+
+    // Find the last occurrence of closing backticks
+    const lastBackticksIndex = codeText.lastIndexOf("```");
+    if (lastBackticksIndex !== -1) {
+      codeText = codeText.slice(0, lastBackticksIndex + 3);
+    }
+
+    // Remove the language declaration and trailing backticks
+    const codeLines = codeText.split("\n");
+    if (codeLines.length > 1 && codeLines[0].trim().startsWith("```")) {
+      codeLines.shift(); // Remove the first line with the language declaration
+      if (codeLines[codeLines.length - 1]?.trim() === "```") {
+        codeLines.pop(); // Remove the last line with the trailing backticks
+      }
+
+      const minIndent = codeLines
+        .filter((line) => line.trim().length > 0)
+        .reduce((min, line) => {
+          const match = line.match(/^\s*/);
+          return Math.min(min, match ? match[0].length : 0);
+        }, Infinity);
+
+      const formattedCodeLines = codeLines.map((line) => line.slice(minIndent));
+      codeText = formattedCodeLines.join("\n");
+    }
+  } else {
+    // Fallback if position offsets are not available
+    codeText = children?.toString() || null;
+  }
+
+  return codeText || "";
+}

web/src/components/chat_search/MinimalMarkdown.tsx

@@ -1,4 +1,5 @@
 import { CodeBlock } from "@/app/chat/message/CodeBlock";
+import { extractCodeText } from "@/app/chat/message/codeUtils";
 import {
   MemoizedLink,
   MemoizedParagraph,
@@ -10,13 +11,11 @@ import remarkGfm from "remark-gfm";
 interface MinimalMarkdownProps {
   content: string;
   className?: string;
-  useCodeBlock?: boolean;
 }
 
 export const MinimalMarkdown: React.FC<MinimalMarkdownProps> = ({
   content,
   className = "",
-  useCodeBlock = false,
 }) => {
   return (
     <ReactMarkdown
@@ -24,11 +23,15 @@ export const MinimalMarkdown: React.FC<MinimalMarkdownProps> = ({
       components={{
         a: MemoizedLink,
         p: MemoizedParagraph,
-        code: useCodeBlock
+        code: ({ node, inline, className, children, ...props }: any) => {
-          ? (props) => (
+          const codeText = extractCodeText(node, content, children);
-              <CodeBlock className="w-full" {...props} content={content} />
+
-            )
+          return (
-          : (props) => <code {...props} />,
+            <CodeBlock className={className} codeText={codeText}>
+              {children}
+            </CodeBlock>
+          );
+        },
       }}
       remarkPlugins={[remarkGfm]}
     >

web/src/components/search/results/AnswerSection.tsx

@@ -1,7 +1,5 @@
 import { Quote } from "@/lib/search/interfaces";
 import { ResponseSection, StatusOptions } from "./ResponseSection";
-import ReactMarkdown from "react-markdown";
-import remarkGfm from "remark-gfm";
 import { MinimalMarkdown } from "@/components/chat_search/MinimalMarkdown";
 
 const TEMP_STRING = "__$%^TEMP$%^__";
@@ -40,12 +38,7 @@ export const AnswerSection = (props: AnswerSectionProps) => {
     status = "success";
     header = <></>;
 
-    body = (
+    body = <MinimalMarkdown content={replaceNewlines(props.answer || "")} />;
-      <MinimalMarkdown
-        useCodeBlock
-        content={replaceNewlines(props.answer || "")}
-      />
-    );
 
     // error while building answer (NOTE: if error occurs during quote generation
     // the above if statement will hit and the error will not be displayed)
@@ -61,9 +54,7 @@ export const AnswerSection = (props: AnswerSectionProps) => {
   } else if (props.answer) {
     status = "success";
     header = <></>;
-    body = (
+    body = <MinimalMarkdown content={replaceNewlines(props.answer)} />;
-      <MinimalMarkdown useCodeBlock content={replaceNewlines(props.answer)} />
-    );
   }
 
   return (