Better virtualization

web/src/app/chat/ChatPage.tsx

@@ -755,7 +755,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(() => {
@@ -1133,7 +1141,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;
-  [key: string]: any;
+  codeText: string;
 }
 
+const MemoizedCodeLine = memo(({ content }: { content: ReactNode }) => (
+  <>{content}</>
+));
+
 export const CodeBlock = memo(function CodeBlock({
   className = "",
   children,
-  content,
-  ...props
+  codeText,
 }: CodeBlockProps) {
   const [copied, setCopied] = useState(false);
 
@@ -26,132 +28,99 @@ export const CodeBlock = memo(function CodeBlock({
       .join(" ");
   }, [className]);
 
-  const codeText = useMemo(() => {
-    let codeText: string | null = null;
-    if (
-      props.node?.position?.start?.offset &&
-      props.node?.position?.end?.offset
-    ) {
-      codeText = content.slice(
-        props.node.position.start.offset,
-        props.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].startsWith("```") ||
-          codeLines[0].trim().startsWith("```"))
-      ) {
-        codeLines.shift(); // Remove the first line with the language declaration
-        if (
-          codeLines[codeLines.length - 1] === "```" ||
-          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)
+  const handleCopy = useCallback(() => {
+    if (!codeText) return;
+    navigator.clipboard.writeText(codeText).then(() => {
+      setCopied(true);
+      setTimeout(() => setCopied(false), 2000);
+    });
+  }, [codeText]);
+
+  const CopyButton = memo(() => (
+    <div
+      className="ml-auto cursor-pointer select-none"
+      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";
+
+  const CodeContent = memo(() => {
+    if (!language) {
+      if (typeof children === "string") {
+        return (
+          <code
+            className={`
+            font-mono 
+            text-gray-800 
+            bg-gray-50 
+            border 
+            border-gray-300 
+            rounded 
+            px-1
+            py-[3px]
+            text-xs 
+            whitespace-pre-wrap 
+            break-words 
+            overflow-hidden
+            mb-1
+            ${className}
+          `}
+          >
+            {children}
+          </code>
         );
-        codeText = formattedCodeLines.join("\n");
       }
-    }
-
-    // handle unknown languages. They won't have a `node.position.start.offset`
-    if (!codeText) {
-      const findTextNode = (node: any): string | null => {
-        if (node.type === "text") {
-          return node.value;
-        }
-        let finalResult = "";
-        if (node.children) {
-          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}>
+          <code className={`text-sm ${className}`}>
+            {Array.isArray(children)
+              ? children.map((child, index) => (
+                  <MemoizedCodeLine key={index} content={child} />
+                ))
+              : children}
+          </code>
+        </pre>
+      );
     }
 
     return (
-      <pre style={CODE_BLOCK_PADDING_TYPE}>
-        <code {...props} className={`text-sm ${className}`}>
-          {children}
+      <pre className="overflow-x-scroll" style={CODE_BLOCK_PADDING}>
+        <code className="text-xs overflow-x-auto">
+          {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}
-        {codeText && (
-          <div
-            className="ml-auto cursor-pointer select-none"
-            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>
-        )}
-      </div>
-      <pre {...props} className="overflow-x-scroll" style={{ padding: "1rem" }}>
-        <code className={`text-xs overflow-x-auto `}>{children}</code>
-      </pre>
+      {language && (
+        <div className="flex mx-3 py-2 text-xs">
+          {language}
+          {codeText && <CopyButton />}
+        </div>
+      )}
+      <CodeContent />
     </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) => (
-  <p {...props} className="text-default" />
-));
+export const MemoizedParagraph = memo(({ ...props }: any) => {
+  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
-                              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>
+                            {renderedMarkdown}
                           </div>
                         ) : (
                           content