perf: skip syntax highlighting while a message is streaming

components/ChatWindow.tsx

@@ -317,6 +317,19 @@ export function ChatWindow({ session, sessionRunning, newSessionCwd, newSessionD
   const { isDragOver, handleDragEnter, handleDragOver, handleDragLeave, handleDrop } = useDragDrop(onDrop);
 
   const visibleMessages = messages.filter((m) => m.role === "user" || m.role === "assistant");
+  // Stable Map identity: `messages` doesn't change during streaming updates
+  // (the streaming message lives in streamState), so memoized MessageViews
+  // skip re-rendering on every message_update event. An inline `new Map()`
+  // here used to defeat MessageView's memo() on each streamed chunk.
+  const toolResultsMap = useMemo(() => {
+    const map = new Map<string, ToolResultMessage>();
+    for (const msg of messages) {
+      if (msg.role === "toolResult") {
+        map.set((msg as ToolResultMessage).toolCallId, msg as ToolResultMessage);
+      }
+    }
+    return map;
+  }, [messages]);
   const inputHistory = useMemo(() => {
     const seen = new Set<string>();
     const history: string[] = [];
@@ -637,13 +650,6 @@ export function ChatWindow({ session, sessionRunning, newSessionCwd, newSessionD
           <div style={{ minWidth: 0, padding: `0 ${CHAT_COLUMN_PADDING}px` }}>
             <div style={{ width: "100%", minWidth: 0, maxWidth: 820, margin: "0 auto" }}>
             {(() => {
-              const toolResultsMap = new Map<string, ToolResultMessage>();
-              for (const msg of messages) {
-                if (msg.role === "toolResult") {
-                  toolResultsMap.set((msg as ToolResultMessage).toolCallId, msg as ToolResultMessage);
-                }
-              }
-
               let lastUserIdx = -1;
               for (let i = messages.length - 1; i >= 0; i--) {
                 if (messages[i].role === "user") { lastUserIdx = i; break; }

components/MarkdownBody.tsx

@@ -27,7 +27,7 @@ export function MarkdownBody({ children, className, isStreaming, cwd, onOpenFile
         if (lang === "mermaid") {
           return <MermaidBlock code={raw.replace(/\n$/, "")} isStreaming={isStreaming} />;
         }
-        return <CodeBlock code={raw.replace(/\n$/, "")} lang={lang} />;
+        return <CodeBlock code={raw.replace(/\n$/, "")} lang={lang} isStreaming={isStreaming} />;
       }
       return (
         <code

components/MermaidBlock.test.mjs

@@ -8,7 +8,7 @@ const jiti = createJiti(import.meta.url, {
   jsx: { runtime: "automatic" },
   tsconfigPaths: true,
 });
-const { MermaidBlock } = await jiti.import("./MermaidBlock.tsx");
+const { MermaidBlock, CodeBlock } = await jiti.import("./MermaidBlock.tsx");
 const { I18nProvider } = await jiti.import("../hooks/useI18n.tsx");
 
 // Simple sequenceDiagram for testing
@@ -58,6 +58,30 @@ test("MermaidBlock renders empty graph without error", () => {
   assert.match(html, /mermaid-block-loading/);
 });
 
+function renderCode(props) {
+  return renderToStaticMarkup(
+    React.createElement(
+      I18nProvider,
+      null,
+      React.createElement(CodeBlock, props),
+    ),
+  );
+}
+
+test("CodeBlock highlights code when not streaming", () => {
+  const html = renderCode({ code: "const x = 1;", lang: "javascript" });
+
+  assert.match(html, /class="token/);
+  assert.match(html, /const/);
+});
+
+test("CodeBlock renders plain text without tokenization while streaming", () => {
+  const html = renderCode({ code: "const x = 1;", lang: "javascript", isStreaming: true });
+
+  assert.doesNotMatch(html, /class="token/);
+  assert.match(html, /const x = 1;/);
+});
+
 test("MermaidBlock handles Chinese characters in diagram", () => {
   const chineseMermaid = `sequenceDiagram
     participant PC as PC客户端

components/MermaidBlock.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { useEffect, useRef, useState, type ReactNode } from "react";
+import { memo, useEffect, useMemo, useRef, useState, type ReactNode } from "react";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { vs } from "react-syntax-highlighter/dist/cjs/styles/prism";
 import { vscDarkPlus } from "react-syntax-highlighter/dist/cjs/styles/prism";
@@ -69,7 +69,7 @@ export function MermaidBlock({ code, isStreaming, defaultPreview = false }: Merm
     };
   }, [code, currentKey, isDark, previewVisible]);
 
-  const previewButton = (
+  const previewButton = useMemo(() => (
     <button
       type="button"
       onClick={() => setShowPreview((v) => !v)}
@@ -79,10 +79,10 @@ export function MermaidBlock({ code, isStreaming, defaultPreview = false }: Merm
     >
       {previewVisible ? t("i18n.source") : t("i18n.preview")}
     </button>
-  );
+  ), [isStreaming, previewVisible, t]);
 
   if (!previewVisible) {
-    return <CodeBlock code={code} lang="mermaid" headerAction={previewButton} />;
+    return <CodeBlock code={code} lang="mermaid" headerAction={previewButton} isStreaming={isStreaming} />;
   }
 
   const body = renderState?.key === currentKey && renderState.status === "error" ? (
@@ -224,13 +224,20 @@ interface CodeBlockProps {
   code: string;
   lang: string;
   headerAction?: ReactNode;
+  isStreaming?: boolean;
 }
 
 /**
  * Syntax-highlighted code block with copy button.
  * Used as the "source" view for mermaid blocks and for all non-mermaid code fences.
+ *
+ * Memoized: parent markdown re-renders (e.g. streaming updates elsewhere in
+ * the message list) must not re-run Prism tokenization on unchanged code.
+ * While the owning message is still streaming, the block renders as plain
+ * monospace text — highlighting a growing block re-tokenizes all of it on
+ * every chunk, which is the single most expensive part of streamed rendering.
  */
-export function CodeBlock({ code, lang, headerAction }: CodeBlockProps) {
+export const CodeBlock = memo(function CodeBlock({ code, lang, headerAction, isStreaming }: CodeBlockProps) {
   const { isDark } = useTheme();
   const { t } = useI18n();
   const [copied, setCopied] = useState(false);
@@ -256,23 +263,38 @@ export function CodeBlock({ code, lang, headerAction }: CodeBlockProps) {
           </button>
         </div>
       </div>
-      <SyntaxHighlighter
-        language={lang || "text"}
-        style={isDark ? vscDarkPlus : vs}
-        showLineNumbers
-        lineNumberStyle={{ color: "var(--text-dim)", fontStyle: "normal" }}
-        customStyle={{
-          margin: 0,
-          padding: "11px 13px",
-          fontSize: 12.5,
-          lineHeight: 1.62,
-          borderRadius: 0,
-          background: "color-mix(in srgb, var(--bg) 92%, var(--bg-panel))",
-        }}
-        codeTagProps={{ style: { fontFamily: "var(--font-mono)" } }}
-      >
-        {code}
-      </SyntaxHighlighter>
+      {isStreaming ? (
+        <pre
+          style={{
+            margin: 0,
+            padding: "11px 13px",
+            fontSize: 12.5,
+            lineHeight: 1.62,
+            overflowX: "auto",
+            background: "color-mix(in srgb, var(--bg) 92%, var(--bg-panel))",
+          }}
+        >
+          <code style={{ fontFamily: "var(--font-mono)" }}>{code}</code>
+        </pre>
+      ) : (
+        <SyntaxHighlighter
+          language={lang || "text"}
+          style={isDark ? vscDarkPlus : vs}
+          showLineNumbers
+          lineNumberStyle={{ color: "var(--text-dim)", fontStyle: "normal" }}
+          customStyle={{
+            margin: 0,
+            padding: "11px 13px",
+            fontSize: 12.5,
+            lineHeight: 1.62,
+            borderRadius: 0,
+            background: "color-mix(in srgb, var(--bg) 92%, var(--bg-panel))",
+          }}
+          codeTagProps={{ style: { fontFamily: "var(--font-mono)" } }}
+        >
+          {code}
+        </SyntaxHighlighter>
+      )}
     </div>
   );
-}
+});