fix code block cannot be selected while generating

examples/server/webui/src/components/ChatMessage.tsx

@@ -2,7 +2,7 @@ import { useMemo, useState } from 'react';
 import { useAppContext } from '../utils/app.context';
 import { Message, PendingMessage } from '../utils/types';
 import { classNames } from '../utils/misc';
-import MarkdownDisplay from './MarkdownDisplay';
+import MarkdownDisplay, { CopyButton } from './MarkdownDisplay';
 
 interface SplitMessage {
   content: PendingMessage['content'];
@@ -207,20 +207,19 @@ export default function ChatMessage({
           {/* assistant message */}
           {msg.role === 'assistant' && (
             <>
-              <button
+              {!isPending && (
+                <button
+                  className="badge btn-mini show-on-hover mr-2"
+                  onClick={regenerate}
+                  disabled={msg.content === null}
+                >
+                  🔄 Regenerate
+                </button>
+              )}
+              <CopyButton
                 className="badge btn-mini show-on-hover mr-2"
-                onClick={regenerate}
-                disabled={msg.content === null}
-              >
-                🔄 Regenerate
-              </button>
-              <button
-                className="badge btn-mini show-on-hover mr-2"
-                onClick={() => navigator.clipboard.writeText(msg.content || '')}
-                disabled={msg.content === null}
-              >
-                📋 Copy
-              </button>
+                content={msg.content}
+              />
             </>
           )}
         </div>

examples/server/webui/src/components/ChatScreen.tsx

@@ -27,7 +27,7 @@ export default function ChatScreen() {
       msgListElem.scrollHeight -
       msgListElem.scrollTop -
       msgListElem.clientHeight;
-    if (!requiresNearBottom || spaceToBottom < 100) {
+    if (!requiresNearBottom || spaceToBottom < 50) {
       setTimeout(
         () => msgListElem.scrollTo({ top: msgListElem.scrollHeight }),
         1

examples/server/webui/src/components/MarkdownDisplay.tsx

@@ -6,7 +6,9 @@ import rehypeKatex from 'rehype-katex';
 import remarkMath from 'remark-math';
 import remarkBreaks from 'remark-breaks';
 import 'katex/dist/katex.min.css';
-import { copyStr } from '../utils/misc';
+import { classNames, copyStr } from '../utils/misc';
+import { ElementContent, Root } from 'hast';
+import { visit } from 'unist-util-visit';
 
 export default function MarkdownDisplay({ content }: { content: string }) {
   const preprocessedContent = useMemo(
@@ -16,9 +18,11 @@ export default function MarkdownDisplay({ content }: { content: string }) {
   return (
     <Markdown
       remarkPlugins={[remarkGfm, remarkMath, remarkBreaks]}
-      rehypePlugins={[rehypeHightlight, rehypeKatex]}
+      rehypePlugins={[rehypeHightlight, rehypeKatex, rehypeCustomCopyButton]}
       components={{
-        pre: (props) => <Pre {...props} origContent={preprocessedContent} />,
+        button: (props) => (
+          <CopyCodeButton {...props} origContent={preprocessedContent} />
+        ),
       }}
     >
       {preprocessedContent}
@@ -26,15 +30,14 @@ export default function MarkdownDisplay({ content }: { content: string }) {
   );
 }
 
-const Pre: React.ElementType<
-  React.ClassAttributes<HTMLPreElement> &
-    React.HTMLAttributes<HTMLPreElement> &
+const CopyCodeButton: React.ElementType<
+  React.ClassAttributes<HTMLButtonElement> &
+    React.HTMLAttributes<HTMLButtonElement> &
     ExtraProps & { origContent: string }
-> = ({ node, origContent, ...props }) => {
+> = ({ node, origContent }) => {
   const startOffset = node?.position?.start.offset ?? 0;
   const endOffset = node?.position?.end.offset ?? 0;
 
-  const [copied, setCopied] = useState(false);
   const copiedContent = useMemo(
     () =>
       origContent
@@ -44,29 +47,70 @@ const Pre: React.ElementType<
     [origContent, startOffset, endOffset]
   );
 
-  if (!node?.position) {
-    return <pre {...props} />;
-  }
-
   return (
-    <div className="relative my-4">
-      <div
-        className="text-right sticky top-4 mb-2 mr-2 h-0"
-        onClick={() => {
-          copyStr(copiedContent);
-          setCopied(true);
-        }}
-        onMouseLeave={() => setCopied(false)}
-      >
-        <button className="badge btn-mini">
-          {copied ? 'Copied!' : '📋 Copy'}
-        </button>
-      </div>
-      <pre {...props} />
+    <div
+      className={classNames({
+        'text-right sticky top-4 mb-2 mr-2 h-0': true,
+        'display-none': !node?.position,
+      })}
+    >
+      <CopyButton className="badge btn-mini" content={copiedContent} />
     </div>
   );
 };
 
+export const CopyButton = ({
+  content,
+  className,
+}: {
+  content: string;
+  className?: string;
+}) => {
+  const [copied, setCopied] = useState(false);
+  return (
+    <button
+      className={className}
+      onClick={() => {
+        copyStr(content);
+        setCopied(true);
+      }}
+      onMouseLeave={() => setCopied(false)}
+    >
+      {copied ? 'Copied!' : '📋 Copy'}
+    </button>
+  );
+};
+
+/**
+ * This injects the "button" element before each "pre" element.
+ * The actual button will be replaced with a react component in the MarkdownDisplay.
+ * We don't replace "pre" node directly because it will cause the node to re-render, which causes this bug: https://github.com/ggerganov/llama.cpp/issues/9608
+ */
+function rehypeCustomCopyButton() {
+  return function (tree: Root) {
+    visit(tree, 'element', function (node) {
+      if (node.tagName === 'pre' && !node.properties.visited) {
+        const preNode = { ...node };
+        // replace current node
+        preNode.properties.visited = 'true';
+        node.tagName = 'div';
+        node.properties = {
+          className: 'relative my-4',
+        };
+        // add node for button
+        const btnNode: ElementContent = {
+          type: 'element',
+          tagName: 'button',
+          properties: {},
+          children: [],
+          position: node.position,
+        };
+        node.children = [btnNode, preNode];
+      }
+    });
+  };
+}
+
 /**
  * The part below is copied and adapted from:
  * https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts