Replace react markdown with react remark for better performance when streaming

2025-12-20 12:21:13 -05:00 · 2024-10-01 01:11:25 -04:00
parent 5fa3794610
commit c664403d4e
4 changed files with 133 additions and 1210 deletions
--- a/webview-ui/src/components/chat/ChatRow.tsx
+++ b/webview-ui/src/components/chat/ChatRow.tsx
@@ -1,14 +1,14 @@
 import { VSCodeBadge, VSCodeProgressRing } from "@vscode/webview-ui-toolkit/react"
 import deepEqual from "fast-deep-equal"
 import React, { memo, useMemo } from "react"
-import ReactMarkdown from "react-markdown"
 import { ClaudeApiReqInfo, ClaudeMessage, ClaudeSayTool } from "../../../../src/shared/ExtensionMessage"
 import { COMMAND_OUTPUT_STRING } from "../../../../src/shared/combineCommandSequences"
 import { vscode } from "../../utils/vscode"
 import CodeAccordian, { removeLeadingNonAlphanumeric } from "../common/CodeAccordian"
 import CodeBlock, { CODE_BLOCK_BG_COLOR } from "../common/CodeBlock"
-import { highlightMentions } from "./TaskHeader"
+import MarkdownBlock from "../common/MarkdownBlock"
 import Thumbnails from "../common/Thumbnails"
+import { highlightMentions } from "./TaskHeader"

 interface ChatRowProps {
 	message: ClaudeMessage
@@ -764,116 +764,9 @@ const ProgressIndicator = () => (
 )

 const Markdown = memo(({ markdown }: { markdown?: string }) => {
-	// react-markdown lets us customize elements, so here we're using their example of replacing code blocks with SyntaxHighlighter. However when there are no language matches (` or ``` without a language specifier) then we default to a normal code element for inline code. Code blocks without a language specifier shouldn't be a common occurrence as we prompt Claude to always use a language specifier.
-	// when claude wraps text in thinking tags, he doesnt use line breaks so we need to insert those ourselves to render markdown correctly
-	// const parsed = markdown?.replace(/<thinking>([\s\S]*?)<\/thinking>/g, (match, content) => {
-	// 	return content
-	// 	// return `_<thinking>_\n\n${content}\n\n_</thinking>_`
-	// })
-	const parsed = markdown
 	return (
-		<div style={{ wordBreak: "break-word", overflowWrap: "anywhere", marginBottom: -10, marginTop: -10 }}>
-			<ReactMarkdown
-				children={parsed}
-				components={{
-					p(props) {
-						const { style, ...rest } = props
-						return (
-							<p
-								style={{
-									...style,
-									margin: 0,
-									marginTop: 10,
-									marginBottom: 10,
-									whiteSpace: "pre-wrap",
-									wordBreak: "break-word",
-									overflowWrap: "anywhere",
-								}}
-								{...rest}
-							/>
-						)
-					},
-					ol(props) {
-						const { style, ...rest } = props
-						return (
-							<ol
-								style={{
-									...style,
-									padding: "0 0 0 20px",
-									margin: "10px 0",
-									wordBreak: "break-word",
-									overflowWrap: "anywhere",
-								}}
-								{...rest}
-							/>
-						)
-					},
-					ul(props) {
-						const { style, ...rest } = props
-						return (
-							<ul
-								style={{
-									...style,
-									padding: "0 0 0 20px",
-									margin: "10px 0",
-									wordBreak: "break-word",
-									overflowWrap: "anywhere",
-								}}
-								{...rest}
-							/>
-						)
-					},
-					// pre always surrounds a code, and we custom handle code blocks below. Pre has some non-10 margin, while all other elements in markdown have a 10 top/bottom margin and the outer div has a -10 top/bottom margin to counteract this between chat rows. However we render markdown in a completion_result row so make sure to add padding as necessary when used within other rows.
-					pre(props) {
-						const { style, ...rest } = props
-						return (
-							<pre
-								style={{
-									...style,
-									marginTop: 10,
-									marginBlock: 10,
-								}}
-								{...rest}
-							/>
-						)
-					},
-					// https://github.com/remarkjs/react-markdown?tab=readme-ov-file#use-custom-components-syntax-highlight
-					code(props) {
-						const { children, className, node, ...rest } = props
-						const match = /language-(\w+)/.exec(className || "")
-						return match ? (
-							<div
-								style={{
-									borderRadius: 3,
-									border: "1px solid var(--vscode-editorGroup-border)",
-									overflow: "hidden",
-								}}>
-								<CodeBlock
-									source={`${"```"}${match[1]}\n${String(children).replace(/\n$/, "")}\n${"```"}`}
-								/>
-							</div>
-						) : (
-							<code
-								{...rest}
-								className={className}
-								style={{
-									whiteSpace: "pre-line",
-									wordBreak: "break-word",
-									overflowWrap: "anywhere",
-									backgroundColor: "var(--vscode-textCodeBlock-background)",
-									color: "var(--vscode-textPreformat-foreground)",
-									fontFamily: "var(--vscode-editor-font-family)",
-									fontSize: "var(--vscode-editor-font-size)",
-									borderRadius: "3px",
-									border: "1px solid var(--vscode-textSeparator-foreground)",
-									padding: "0px 2px",
-								}}>
-								{children}
-							</code>
-						)
-					},
-				}}
-			/>
+		<div style={{ wordBreak: "break-word", overflowWrap: "anywhere", marginBottom: -15, marginTop: -15 }}>
+			<MarkdownBlock markdown={markdown} />
 		</div>
 	)
 })
--- a/webview-ui/src/components/common/MarkdownBlock.tsx
+++ b/webview-ui/src/components/common/MarkdownBlock.tsx
@@ -0,0 +1,129 @@
+import { memo, useEffect } from "react"
+import { useRemark } from "react-remark"
+import rehypeHighlight, { Options } from "rehype-highlight"
+import styled from "styled-components"
+import { visit } from "unist-util-visit"
+import { useExtensionState } from "../../context/ExtensionStateContext"
+import { CODE_BLOCK_BG_COLOR } from "./CodeBlock"
+
+interface MarkdownBlockProps {
+	markdown?: string
+}
+
+const StyledMarkdown = styled.div`
+	pre {
+		background-color: ${CODE_BLOCK_BG_COLOR};
+		border-radius: 3px;
+		margin: 13x 0;
+		padding: 10px 10px;
+		max-width: calc(100vw - 20px);
+		overflow-x: scroll;
+		overflow-y: hidden;
+	}
+
+	pre > code {
+		.hljs-deletion {
+			background-color: var(--vscode-diffEditor-removedTextBackground);
+			display: inline-block;
+			width: 100%;
+		}
+		.hljs-addition {
+			background-color: var(--vscode-diffEditor-insertedTextBackground);
+			display: inline-block;
+			width: 100%;
+		}
+	}
+
+	code {
+		span.line:empty {
+			display: none;
+		}
+		word-wrap: break-word;
+		border-radius: 3px;
+		background-color: ${CODE_BLOCK_BG_COLOR};
+		font-size: var(--vscode-editor-font-size, var(--vscode-font-size, 12px));
+		font-family: var(--vscode-editor-font-family);
+	}
+
+	code:not(pre > code) {
+		font-family: var(--vscode-editor-font-family, monospace);
+		color: var(--vscode-textPreformat-foreground, #f78383);
+		background-color: var(--vscode-textCodeBlock-background, #1e1e1e);
+		padding: 0px 2px;
+		border-radius: 3px;
+		border: 1px solid var(--vscode-textSeparator-foreground, #424242);
+		white-space: pre-line;
+		word-break: break-word;
+		overflow-wrap: anywhere;
+	}
+
+	font-family: var(--vscode-font-family), system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Oxygen,
+		Ubuntu, Cantarell, "Open Sans", "Helvetica Neue", sans-serif;
+	font-size: var(--vscode-font-size, 13px);
+
+	p,
+	li,
+	ol,
+	ul {
+		line-height: 1.25;
+	}
+`
+
+const StyledPre = styled.pre<{ theme: any }>`
+	& .hljs {
+		color: var(--vscode-editor-foreground, #fff);
+	}
+
+	${(props) =>
+		Object.keys(props.theme)
+			.map((key, index) => {
+				return `
+      & ${key} {
+        color: ${props.theme[key]};
+      }
+    `
+			})
+			.join("")}
+`
+
+const MarkdownBlock = memo(({ markdown }: MarkdownBlockProps) => {
+	const { theme } = useExtensionState()
+	const [reactContent, setMarkdown] = useRemark({
+		remarkPlugins: [
+			() => {
+				return (tree) => {
+					visit(tree, "code", (node: any) => {
+						if (!node.lang) {
+							node.lang = "javascript"
+						} else if (node.lang.includes(".")) {
+							node.lang = node.lang.split(".").slice(-1)[0]
+						}
+					})
+				}
+			},
+		],
+		rehypePlugins: [
+			rehypeHighlight as any,
+			{
+				// languages: {},
+			} as Options,
+		],
+		rehypeReactOptions: {
+			components: {
+				pre: ({ node, ...preProps }: any) => <StyledPre {...preProps} theme={theme} />,
+			},
+		},
+	})
+
+	useEffect(() => {
+		setMarkdown(markdown || "")
+	}, [markdown, setMarkdown, theme])
+
+	return (
+		<div style={{}}>
+			<StyledMarkdown>{reactContent}</StyledMarkdown>
+		</div>
+	)
+})
+
+export default MarkdownBlock