From fabd15d8c38657a92323d23af5dce939c29ee8c1 Mon Sep 17 00:00:00 2001 From: kun Date: Sat, 8 Aug 2026 23:17:47 +0800 Subject: [PATCH 1/2] =?UTF-8?q?perf:=20cut=20streamed-render=20CPU=20?= =?UTF-8?q?=E2=80=94=20skip=20Prism=20while=20streaming,=20memoize=20CodeB?= =?UTF-8?q?lock=20and=20toolResults=20map?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Streaming updates re-rendered the live message on every SSE message_update, re-tokenizing the whole growing code block with Prism each time. Render code fences as plain monospace text while the message is streaming and highlight once on completion. Also memoize CodeBlock (unchanged code must not re-run tokenization when the parent markdown re-renders) and hoist ChatWindow's toolResults Map into a useMemo keyed on messages so its identity stays stable across streaming updates. --- components/ChatWindow.tsx | 20 ++++++---- components/MarkdownBody.tsx | 2 +- components/MermaidBlock.test.mjs | 26 ++++++++++++- components/MermaidBlock.tsx | 64 +++++++++++++++++++++----------- 4 files changed, 82 insertions(+), 30 deletions(-) diff --git a/components/ChatWindow.tsx b/components/ChatWindow.tsx index fa3ed749b..c65f50c6e 100644 --- a/components/ChatWindow.tsx +++ b/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(); + 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(); const history: string[] = []; @@ -637,13 +650,6 @@ export function ChatWindow({ session, sessionRunning, newSessionCwd, newSessionD
{(() => { - const toolResultsMap = new Map(); - 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; } diff --git a/components/MarkdownBody.tsx b/components/MarkdownBody.tsx index cc6bd6822..2547f3393 100644 --- a/components/MarkdownBody.tsx +++ b/components/MarkdownBody.tsx @@ -27,7 +27,7 @@ export function MarkdownBody({ children, className, isStreaming, cwd, onOpenFile if (lang === "mermaid") { return ; } - return ; + return ; } return ( { 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客户端 diff --git a/components/MermaidBlock.tsx b/components/MermaidBlock.tsx index 0efd3db95..0d7e697fd 100644 --- a/components/MermaidBlock.tsx +++ b/components/MermaidBlock.tsx @@ -1,6 +1,6 @@ "use client"; -import { useEffect, useRef, useState, type ReactNode } from "react"; +import { memo, useEffect, 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"; @@ -82,7 +82,7 @@ export function MermaidBlock({ code, isStreaming, defaultPreview = false }: Merm ); if (!previewVisible) { - return ; + return ; } 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) {
- - {code} - + {isStreaming ? ( + 
+          {code}
+
+ ) : ( + + {code} + + )} ); -} +}); From 49b7657ad771ed5847c04da3998594cb12f084c1 Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Sun, 9 Aug 2026 18:32:27 +0800 Subject: [PATCH 2/2] fix: stabilize Mermaid preview action --- components/MermaidBlock.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/components/MermaidBlock.tsx b/components/MermaidBlock.tsx index 0d7e697fd..6826e29af 100644 --- a/components/MermaidBlock.tsx +++ b/components/MermaidBlock.tsx @@ -1,6 +1,6 @@ "use client"; -import { memo, 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(() => ( - ); + ), [isStreaming, previewVisible, t]); if (!previewVisible) { return ;