Chris Estreich 8 месяцев назад
Родитель
Сommit
65958d81b9

+ 5 - 0
.changeset/five-pumpkins-carry.md

@@ -0,0 +1,5 @@
+---
+"roo-cline": patch
+---
+
+Overhaul CodeBlock rendering

+ 1 - 1
src/core/webview/ClineProvider.ts

@@ -748,7 +748,7 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
             <meta charset="utf-8">
             <meta name="viewport" content="width=device-width,initial-scale=1,shrink-to-fit=no">
             <meta name="theme-color" content="#000000">
-            <meta http-equiv="Content-Security-Policy" content="default-src 'none'; font-src ${webview.cspSource}; style-src ${webview.cspSource} 'unsafe-inline'; img-src ${webview.cspSource} data:; script-src 'nonce-${nonce}' https://us-assets.i.posthog.com; connect-src https://openrouter.ai https://api.requesty.ai https://us.i.posthog.com https://us-assets.i.posthog.com;">
+            <meta http-equiv="Content-Security-Policy" content="default-src 'none'; font-src ${webview.cspSource}; style-src ${webview.cspSource} 'unsafe-inline'; img-src ${webview.cspSource} data:; script-src ${webview.cspSource} 'wasm-unsafe-eval' 'nonce-${nonce}' https://us-assets.i.posthog.com 'strict-dynamic'; connect-src https://openrouter.ai https://api.requesty.ai https://us.i.posthog.com https://us-assets.i.posthog.com;">
             <link rel="stylesheet" type="text/css" href="${stylesUri}">
 			<link href="${codiconsUri}" rel="stylesheet" />
 			<script nonce="${nonce}">

+ 8 - 1
src/core/webview/__tests__/ClineProvider.test.ts

@@ -375,7 +375,14 @@ describe("ClineProvider", () => {
 		expect(mockWebviewView.webview.html).toContain(
 			"connect-src https://openrouter.ai https://api.requesty.ai https://us.i.posthog.com https://us-assets.i.posthog.com;",
 		)
-		expect(mockWebviewView.webview.html).toContain("script-src 'nonce-")
+
+		// Extract the script-src directive section and verify required security elements
+		const html = mockWebviewView.webview.html
+		const scriptSrcMatch = html.match(/script-src[^;]*;/)
+		expect(scriptSrcMatch).not.toBeNull()
+		expect(scriptSrcMatch![0]).toContain("'nonce-")
+		// Verify wasm-unsafe-eval is present for Shiki syntax highlighting
+		expect(scriptSrcMatch![0]).toContain("'wasm-unsafe-eval'")
 	})
 
 	test("postMessageToWebview sends message to webview", async () => {

+ 1 - 1
webview-ui/package.json

@@ -57,6 +57,7 @@
 		"remark-gfm": "^4.0.1",
 		"remove-markdown": "^0.6.0",
 		"shell-quote": "^1.8.2",
+		"shiki": "^3.2.1",
 		"styled-components": "^6.1.13",
 		"tailwind-merge": "^2.6.0",
 		"tailwindcss": "^4.0.0",
@@ -91,7 +92,6 @@
 		"jest": "^29.7.0",
 		"jest-environment-jsdom": "^29.7.0",
 		"jest-simple-dot-reporter": "^1.0.5",
-		"shiki": "^2.3.2",
 		"storybook": "^8.5.6",
 		"storybook-dark-mode": "^4.0.2",
 		"ts-jest": "^29.2.5",

+ 2 - 4
webview-ui/src/components/chat/BrowserSessionRow.tsx

@@ -357,9 +357,7 @@ const BrowserSessionRow = memo((props: BrowserSessionRowProps) => {
 						<span style={{ fontSize: "0.8em" }}>{t("chat:browser.consoleLogs")}</span>
 					</div>
 					{consoleLogsExpanded && (
-						<CodeBlock
-							source={`${"```"}shell\n${displayState.consoleLogs || t("chat:browser.noNewLogs")}\n${"```"}`}
-						/>
+						<CodeBlock source={displayState.consoleLogs || t("chat:browser.noNewLogs")} language="shell" />
 					)}
 				</div>
 			</div>
@@ -488,7 +486,7 @@ const BrowserSessionRowContent = ({
 									overflow: "hidden",
 									backgroundColor: CODE_BLOCK_BG_COLOR,
 								}}>
-								<CodeBlock source={`${"```"}shell\n${message.text}\n${"```"}`} forceWrap={true} />
+								<CodeBlock source={message.text} language="shell" />
 							</div>
 						</>
 					)

+ 2 - 6
webview-ui/src/components/chat/ChatRow.tsx

@@ -511,7 +511,7 @@ export const ChatRowContent = ({
 						<CodeAccordian
 							code={tool.content!}
 							path={tool.path! + (tool.filePattern ? `/(${tool.filePattern})` : "")}
-							language="plaintext"
+							language="log"
 							isExpanded={isExpanded}
 							onToggleExpand={onToggleExpand}
 						/>
@@ -731,10 +731,7 @@ export const ChatRowContent = ({
 											backgroundColor: "var(--vscode-editor-background)",
 											borderTop: "none",
 										}}>
-										<CodeBlock
-											source={`${"```"}plaintext\n${message.text || ""}\n${"```"}`}
-											forceWrap={true}
-										/>
+										<CodeBlock source={`${"```"}plaintext\n${message.text || ""}\n${"```"}`} />
 									</div>
 								)}
 							</div>
@@ -1111,7 +1108,6 @@ export const ChatRowContent = ({
 													language="json"
 													isExpanded={true}
 													onToggleExpand={onToggleExpand}
-													forceWrap={true}
 												/>
 											</div>
 										)}

+ 2 - 8
webview-ui/src/components/common/CodeAccordian.tsx

@@ -15,7 +15,6 @@ interface CodeAccordianProps {
 	onToggleExpand: () => void
 	isLoading?: boolean
 	progressStatus?: ToolProgressStatus
-	forceWrap?: boolean
 }
 
 /*
@@ -40,7 +39,6 @@ const CodeAccordian = ({
 	onToggleExpand,
 	isLoading,
 	progressStatus,
-	forceWrap,
 }: CodeAccordianProps) => {
 	const inferredLanguage = useMemo(
 		() => code && (language ?? (path ? getLanguageFromPath(path) : undefined)),
@@ -126,12 +124,8 @@ const CodeAccordian = ({
 						maxWidth: "100%",
 					}}>
 					<CodeBlock
-						source={`${"```"}${diff !== undefined ? "diff" : inferredLanguage}\n${(
-							code ??
-							diff ??
-							""
-						).trim()}\n${"```"}`}
-						forceWrap={forceWrap}
+						source={(code ?? diff ?? "").trim()}
+						language={diff !== undefined ? "diff" : inferredLanguage}
 					/>
 				</div>
 			)}

+ 671 - 108
webview-ui/src/components/common/CodeBlock.tsx

@@ -1,11 +1,19 @@
-import { memo, useEffect } from "react"
-import { useRemark } from "react-remark"
-import rehypeHighlight, { Options } from "rehype-highlight"
+import { memo, useEffect, useRef, useCallback, useState } from "react"
 import styled from "styled-components"
-import { visit } from "unist-util-visit"
-import { useExtensionState } from "@src/context/ExtensionStateContext"
-
+import { useCopyToClipboard } from "@src/utils/clipboard"
+import { getHighlighter, isLanguageLoaded, normalizeLanguage, ExtendedLanguage } from "@src/utils/highlighter"
+import { bundledLanguages } from "shiki"
+import type { ShikiTransformer } from "shiki"
 export const CODE_BLOCK_BG_COLOR = "var(--vscode-editor-background, --vscode-sideBar-background, rgb(30 30 30))"
+export const WRAPPER_ALPHA = "cc" // 80% opacity
+// Configuration constants
+export const WINDOW_SHADE_SETTINGS = {
+	transitionDelayS: 0.2,
+	collapsedHeight: 500, // Default collapsed height in pixels
+}
+
+// Tolerance in pixels for determining when a container is considered "at the bottom"
+export const SCROLL_SNAP_TOLERANCE = 20
 
 /*
 overflowX: auto + inner div with padding results in an issue where the top/left/bottom padding renders but the right padding inside does not count as overflow as the width of the element is not exceeded. Once the inner div is outside the boundaries of the parent it counts as overflow.
@@ -17,26 +25,122 @@ minWidth: "max-content",
 
 interface CodeBlockProps {
 	source?: string
-	forceWrap?: boolean
+	rawSource?: string // Add rawSource prop for copying raw text
+	language?: string
+	preStyle?: React.CSSProperties
+	initialWordWrap?: boolean
+	collapsedHeight?: number
+	initialWindowShade?: boolean
+	onLanguageChange?: (language: string) => void
 }
 
-const StyledMarkdown = styled.div<{ forceWrap: boolean }>`
-	${({ forceWrap }) =>
-		forceWrap &&
-		`
-    pre, code {
-      white-space: pre-wrap;
-      word-break: break-all;
-      overflow-wrap: anywhere;
-    }
-  `}
+const ButtonIcon = styled.span`
+	display: inline-block;
+	width: 1.2em;
+	text-align: center;
+	vertical-align: middle;
+`
+
+const CodeBlockButton = styled.button`
+	background: transparent;
+	border: none;
+	color: var(--vscode-foreground);
+	cursor: var(--copy-button-cursor, default);
+	padding: 4px;
+	margin: 0 0px;
+	display: flex;
+	align-items: center;
+	opacity: 0.4;
+	border-radius: 3px;
+	pointer-events: var(--copy-button-events, none);
+	margin-left: 4px;
+	height: 24px;
+
+	&:hover {
+		background: var(--vscode-toolbar-hoverBackground);
+		opacity: 1;
+	}
+`
+
+const CodeBlockButtonWrapper = styled.div`
+	position: fixed;
+	top: var(--copy-button-top);
+	right: var(--copy-button-right, 8px);
+	height: auto;
+	z-index: 100;
+	background: ${CODE_BLOCK_BG_COLOR}${WRAPPER_ALPHA};
+	overflow: visible;
+	pointer-events: none;
+	opacity: var(--copy-button-opacity, 0);
+	padding: 4px 6px;
+	border-radius: 3px;
+	display: inline-flex;
+	align-items: center;
+	justify-content: center;
+
+	&:hover {
+		background: var(--vscode-editor-background);
+		opacity: 1 !important;
+	}
+
+	${CodeBlockButton} {
+		position: relative;
+		top: 0;
+		right: 0;
+	}
+`
+
+const CodeBlockContainer = styled.div`
+	position: relative;
+	overflow: hidden;
+	border-bottom: 4px solid var(--vscode-sideBar-background);
+	background-color: ${CODE_BLOCK_BG_COLOR};
+
+	${CodeBlockButtonWrapper} {
+		opacity: 0;
+		pointer-events: none;
+		transition: opacity 0.2s; /* Keep opacity transition for buttons */
+	}
+
+	&[data-partially-visible="true"]:hover ${CodeBlockButtonWrapper} {
+		opacity: 1;
+		pointer-events: all;
+		cursor: pointer;
+	}
+`
+
+export const StyledPre = styled.div<{
+	preStyle?: React.CSSProperties
+	wordwrap?: "true" | "false" | undefined
+	windowshade?: "true" | "false"
+	collapsedHeight?: number
+}>`
+	background-color: ${CODE_BLOCK_BG_COLOR};
+	max-height: ${({ windowshade, collapsedHeight }) =>
+		windowshade === "true" ? `${collapsedHeight || WINDOW_SHADE_SETTINGS.collapsedHeight}px` : "none"};
+	overflow-y: auto;
+	padding: 10px;
+	// transition: max-height ${WINDOW_SHADE_SETTINGS.transitionDelayS} ease-out;
+	border-radius: 5px;
+	${({ preStyle }) => preStyle && { ...preStyle }}
 
 	pre {
 		background-color: ${CODE_BLOCK_BG_COLOR};
 		border-radius: 5px;
 		margin: 0;
-		min-width: ${({ forceWrap }) => (forceWrap ? "auto" : "max-content")};
-		padding: 10px 10px;
+		padding: 10px;
+		width: 100%;
+		box-sizing: border-box;
+	}
+
+	pre,
+	code {
+		/* Undefined wordwrap defaults to true (pre-wrap) behavior */
+		white-space: ${({ wordwrap }) => (wordwrap === "false" ? "pre" : "pre-wrap")};
+		word-break: ${({ wordwrap }) => (wordwrap === "false" ? "normal" : "normal")};
+		overflow-wrap: ${({ wordwrap }) => (wordwrap === "false" ? "normal" : "break-word")};
+		font-size: var(--vscode-editor-font-size, var(--vscode-font-size, 12px));
+		font-family: var(--vscode-editor-font-family);
 	}
 
 	pre > code {
@@ -52,108 +156,567 @@ const StyledMarkdown = styled.div<{ forceWrap: boolean }>`
 		}
 	}
 
-	code {
-		span.line:empty {
-			display: none;
-		}
-		word-wrap: break-word;
-		border-radius: 5px;
+	.hljs {
+		color: var(--vscode-editor-foreground, #fff);
 		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);
-		color: #f78383;
+const LanguageSelect = styled.select`
+	font-size: 12px;
+	color: var(--vscode-foreground);
+	opacity: 0.4;
+	font-family: monospace;
+	appearance: none;
+	background: transparent;
+	border: none;
+	cursor: pointer;
+	padding: 4px;
+	margin: 0;
+	vertical-align: middle;
+	height: 24px;
+
+	& option {
+		background: var(--vscode-editor-background);
+		color: var(--vscode-foreground);
+		padding: 0;
+		margin: 0;
 	}
 
-	background-color: ${CODE_BLOCK_BG_COLOR};
-	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-editor-font-size, var(--vscode-font-size, 12px));
-	color: var(--vscode-editor-foreground, #fff);
-
-	p,
-	li,
-	ol,
-	ul {
-		line-height: 1.5;
+	&::-webkit-scrollbar {
+		width: 6px;
 	}
-`
 
-const StyledPre = styled.pre<{ theme: any }>`
-	& .hljs {
-		color: var(--vscode-editor-foreground, #fff);
+	&::-webkit-scrollbar-thumb {
+		background: var(--vscode-scrollbarSlider-background);
 	}
 
-	${(props) =>
-		Object.keys(props.theme)
-			.map((key) => {
-				return `
-      & ${key} {
-        color: ${props.theme[key]};
-      }
-    `
-			})
-			.join("")}
+	&::-webkit-scrollbar-track {
+		background: var(--vscode-editor-background);
+	}
+
+	&:hover {
+		opacity: 1;
+		background: var(--vscode-toolbar-hoverBackground);
+		border-radius: 3px;
+	}
+
+	&:focus {
+		opacity: 1;
+		outline: none;
+		border-radius: 3px;
+	}
 `
 
-const CodeBlock = memo(({ source, forceWrap = false }: CodeBlockProps) => {
-	const { theme } = useExtensionState()
-	const [reactContent, setMarkdownSource] = useRemark({
-		remarkPlugins: [
-			() => {
-				return (tree) => {
-					visit(tree, "code", (node: any) => {
-						if (!node.lang) {
-							node.lang = "javascript"
-						} else if (node.lang.includes(".")) {
-							// if the language is a file, get the extension
-							node.lang = node.lang.split(".").slice(-1)[0]
+const CodeBlock = memo(
+	({
+		source,
+		rawSource,
+		language,
+		preStyle,
+		initialWordWrap = true,
+		initialWindowShade = true,
+		collapsedHeight,
+		onLanguageChange,
+	}: CodeBlockProps) => {
+		const [wordWrap, setWordWrap] = useState(initialWordWrap)
+		const [windowShade, setWindowShade] = useState(initialWindowShade)
+		const [currentLanguage, setCurrentLanguage] = useState<ExtendedLanguage>(() => normalizeLanguage(language))
+		const userChangedLanguageRef = useRef(false)
+		const [highlightedCode, setHighlightedCode] = useState<string>("")
+		const [showCollapseButton, setShowCollapseButton] = useState(true)
+		const codeBlockRef = useRef<HTMLDivElement>(null)
+		const preRef = useRef<HTMLDivElement>(null)
+		const copyButtonWrapperRef = useRef<HTMLDivElement>(null)
+		const { showCopyFeedback, copyWithFeedback } = useCopyToClipboard()
+
+		// Update current language when prop changes, but only if user hasn't made a selection
+		useEffect(() => {
+			const normalizedLang = normalizeLanguage(language)
+			if (normalizedLang !== currentLanguage && !userChangedLanguageRef.current) {
+				setCurrentLanguage(normalizedLang)
+			}
+		}, [language, currentLanguage])
+
+		// Syntax highlighting with cached Shiki instance
+		useEffect(() => {
+			const fallback = `<pre style="padding: 0; margin: 0;"><code class="hljs language-${currentLanguage || "txt"}">${source || ""}</code></pre>`
+			const highlight = async () => {
+				// Show plain text if language needs to be loaded
+				if (currentLanguage && !isLanguageLoaded(currentLanguage)) {
+					setHighlightedCode(fallback)
+				}
+
+				const highlighter = await getHighlighter(currentLanguage)
+				const html = await highlighter.codeToHtml(source || "", {
+					lang: currentLanguage || "txt",
+					theme: document.body.className.toLowerCase().includes("light") ? "github-light" : "github-dark",
+					transformers: [
+						{
+							pre(node) {
+								node.properties.style = "padding: 0; margin: 0;"
+								return node
+							},
+							code(node) {
+								// Add hljs classes for consistent styling
+								node.properties.class = `hljs language-${currentLanguage}`
+								return node
+							},
+							line(node) {
+								// Preserve existing line handling
+								node.properties.class = node.properties.class || ""
+								return node
+							},
+						},
+					] as ShikiTransformer[],
+				})
+				setHighlightedCode(html)
+			}
+
+			highlight().catch((e) => {
+				console.error("[CodeBlock] Syntax highlighting error:", e, "\nStack trace:", e.stack)
+				setHighlightedCode(fallback)
+			})
+		}, [source, currentLanguage, collapsedHeight])
+
+		// Check if content height exceeds collapsed height whenever content changes
+		useEffect(() => {
+			const codeBlock = codeBlockRef.current
+			if (codeBlock) {
+				const actualHeight = codeBlock.scrollHeight
+				setShowCollapseButton(actualHeight >= WINDOW_SHADE_SETTINGS.collapsedHeight)
+			}
+		}, [highlightedCode])
+
+		// Ref to track if user was scrolled up *before* the source update potentially changes scrollHeight
+		const wasScrolledUpRef = useRef(false)
+
+		// Ref to track if outer container was near bottom
+		const outerContainerNearBottomRef = useRef(false)
+
+		// Effect to listen to scroll events and update the ref
+		useEffect(() => {
+			const preElement = preRef.current
+			if (!preElement) return
+
+			const handleScroll = () => {
+				const isAtBottom =
+					Math.abs(preElement.scrollHeight - preElement.scrollTop - preElement.clientHeight) <
+					SCROLL_SNAP_TOLERANCE
+				wasScrolledUpRef.current = !isAtBottom
+			}
+
+			preElement.addEventListener("scroll", handleScroll, { passive: true })
+			// Initial check in case it starts scrolled up
+			handleScroll()
+
+			return () => {
+				preElement.removeEventListener("scroll", handleScroll)
+			}
+		}, []) // Empty dependency array: runs once on mount
+
+		// Effect to track outer container scroll position
+		useEffect(() => {
+			const scrollContainer = document.querySelector('[data-virtuoso-scroller="true"]')
+			if (!scrollContainer) return
+
+			const handleOuterScroll = () => {
+				const isAtBottom =
+					Math.abs(scrollContainer.scrollHeight - scrollContainer.scrollTop - scrollContainer.clientHeight) <
+					SCROLL_SNAP_TOLERANCE
+				outerContainerNearBottomRef.current = isAtBottom
+			}
+
+			scrollContainer.addEventListener("scroll", handleOuterScroll, { passive: true })
+			// Initial check
+			handleOuterScroll()
+
+			return () => {
+				scrollContainer.removeEventListener("scroll", handleOuterScroll)
+			}
+		}, []) // Empty dependency array: runs once on mount
+
+		// Store whether we should scroll after highlighting completes
+		const shouldScrollAfterHighlightRef = useRef(false)
+
+		// Check if we should scroll when source changes
+		useEffect(() => {
+			// Only set the flag if we're at the bottom when source changes
+			if (preRef.current && source && !wasScrolledUpRef.current) {
+				shouldScrollAfterHighlightRef.current = true
+			} else {
+				shouldScrollAfterHighlightRef.current = false
+			}
+		}, [source])
+
+		const updateCodeBlockButtonPosition = useCallback((forceHide = false) => {
+			const codeBlock = codeBlockRef.current
+			const copyWrapper = copyButtonWrapperRef.current
+			if (!codeBlock) return
+
+			const rectCodeBlock = codeBlock.getBoundingClientRect()
+			const scrollContainer = document.querySelector('[data-virtuoso-scroller="true"]')
+			if (!scrollContainer) return
+
+			// Get wrapper height dynamically
+			let wrapperHeight
+			if (copyWrapper) {
+				const copyRect = copyWrapper.getBoundingClientRect()
+				// If height is 0 due to styling, estimate from children
+				if (copyRect.height > 0) {
+					wrapperHeight = copyRect.height
+				} else if (copyWrapper.children.length > 0) {
+					// Try to get height from the button inside
+					const buttonRect = copyWrapper.children[0].getBoundingClientRect()
+					const buttonStyle = window.getComputedStyle(copyWrapper.children[0] as Element)
+					const buttonPadding =
+						parseInt(buttonStyle.getPropertyValue("padding-top") || "0", 10) +
+						parseInt(buttonStyle.getPropertyValue("padding-bottom") || "0", 10)
+					wrapperHeight = buttonRect.height + buttonPadding
+				}
+			}
+
+			// If we still don't have a height, calculate from font size
+			if (!wrapperHeight) {
+				const fontSize = parseInt(window.getComputedStyle(document.body).getPropertyValue("font-size"), 10)
+				wrapperHeight = fontSize * 2.5 // Approximate button height based on font size
+			}
+
+			const scrollRect = scrollContainer.getBoundingClientRect()
+			const copyButtonEdge = 48
+			const isPartiallyVisible =
+				rectCodeBlock.top < scrollRect.bottom - copyButtonEdge &&
+				rectCodeBlock.bottom >= scrollRect.top + copyButtonEdge
+
+			// Calculate margin from existing padding in the component
+			const computedStyle = window.getComputedStyle(codeBlock)
+			const paddingValue = parseInt(computedStyle.getPropertyValue("padding") || "0", 10)
+			const margin =
+				paddingValue > 0 ? paddingValue : parseInt(computedStyle.getPropertyValue("padding-top") || "0", 10)
+
+			// Update visibility state and button interactivity
+			const isVisible = !forceHide && isPartiallyVisible
+			codeBlock.setAttribute("data-partially-visible", isPartiallyVisible ? "true" : "false")
+			codeBlock.style.setProperty("--copy-button-cursor", isVisible ? "pointer" : "default")
+			codeBlock.style.setProperty("--copy-button-events", isVisible ? "all" : "none")
+			codeBlock.style.setProperty("--copy-button-opacity", isVisible ? "1" : "0")
+
+			if (isPartiallyVisible) {
+				// Keep button within code block bounds using dynamic measurements
+				const topPosition = Math.max(
+					scrollRect.top + margin,
+					Math.min(rectCodeBlock.bottom - wrapperHeight - margin, rectCodeBlock.top + margin),
+				)
+				const rightPosition = Math.max(margin, scrollRect.right - rectCodeBlock.right + margin)
+
+				codeBlock.style.setProperty("--copy-button-top", `${topPosition}px`)
+				codeBlock.style.setProperty("--copy-button-right", `${rightPosition}px`)
+			}
+		}, [])
+
+		useEffect(() => {
+			const handleScroll = () => updateCodeBlockButtonPosition()
+			const handleResize = () => updateCodeBlockButtonPosition()
+
+			const scrollContainer = document.querySelector('[data-virtuoso-scroller="true"]')
+			if (scrollContainer) {
+				scrollContainer.addEventListener("scroll", handleScroll)
+				window.addEventListener("resize", handleResize)
+				updateCodeBlockButtonPosition()
+			}
+
+			return () => {
+				if (scrollContainer) {
+					scrollContainer.removeEventListener("scroll", handleScroll)
+					window.removeEventListener("resize", handleResize)
+				}
+			}
+		}, [updateCodeBlockButtonPosition])
+
+		// Update button position and scroll when highlightedCode changes
+		useEffect(() => {
+			if (highlightedCode) {
+				// Update button position
+				setTimeout(updateCodeBlockButtonPosition, 0)
+
+				// Scroll to bottom if needed (immediately after Shiki updates)
+				if (shouldScrollAfterHighlightRef.current) {
+					// Scroll inner container
+					if (preRef.current) {
+						preRef.current.scrollTop = preRef.current.scrollHeight
+						wasScrolledUpRef.current = false
+					}
+
+					// Also scroll outer container if it was near bottom
+					if (outerContainerNearBottomRef.current) {
+						const scrollContainer = document.querySelector('[data-virtuoso-scroller="true"]')
+						if (scrollContainer) {
+							scrollContainer.scrollTop = scrollContainer.scrollHeight
+							outerContainerNearBottomRef.current = true
 						}
-					})
+					}
+
+					// Reset the flag
+					shouldScrollAfterHighlightRef.current = false
+				}
+			}
+		}, [highlightedCode, updateCodeBlockButtonPosition])
+
+		// Advanced inertial scroll chaining
+		// This effect handles the transition between scrolling the code block and the outer container.
+		// When a user scrolls to the boundary of a code block (top or bottom), this implementation:
+		// 1. Detects the boundary condition
+		// 2. Applies inertial scrolling to the outer container for a smooth transition
+		// 3. Adds physics-based momentum for natural deceleration
+		// This creates a seamless experience where scrolling flows naturally between nested scrollable areas
+		useEffect(() => {
+			if (!preRef.current) return
+
+			// Find the outer scrollable container
+			const getScrollContainer = () => {
+				return document.querySelector('[data-virtuoso-scroller="true"]') as HTMLElement
+			}
+
+			// Inertial scrolling implementation
+			let velocity = 0
+			let animationFrameId: number | null = null
+			const FRICTION = 0.85 // Friction coefficient (lower = more friction)
+			const MIN_VELOCITY = 0.5 // Minimum velocity before stopping
+
+			// Animation function for inertial scrolling
+			const animate = () => {
+				const scrollContainer = getScrollContainer()
+				if (!scrollContainer) return
+
+				// Apply current velocity
+				if (Math.abs(velocity) > MIN_VELOCITY) {
+					scrollContainer.scrollBy(0, velocity)
+					velocity *= FRICTION // Apply friction
+					animationFrameId = requestAnimationFrame(animate)
+				} else {
+					velocity = 0
+					animationFrameId = null
+				}
+			}
+
+			// Wheel event handler with inertial scrolling
+			const handleWheel = (e: WheelEvent) => {
+				// If shift is pressed, let the browser handle default horizontal scrolling
+				if (e.shiftKey) {
+					return
+				}
+				if (!preRef.current) return
+
+				// Only handle wheel events if the inner container has a scrollbar,
+				// otherwise let the browser handle the default scrolling
+				const hasScrollbar = preRef.current.scrollHeight > preRef.current.clientHeight
+
+				// Pass through events if we don't need special handling
+				if (!hasScrollbar) {
+					return
+				}
+
+				const scrollContainer = getScrollContainer()
+				if (!scrollContainer) return
+
+				// Check if we're at the top or bottom of the inner container
+				const isAtVeryTop = preRef.current.scrollTop === 0
+				const isAtVeryBottom =
+					Math.abs(preRef.current.scrollHeight - preRef.current.scrollTop - preRef.current.clientHeight) < 1
+
+				// Handle scrolling at container boundaries
+				if ((e.deltaY < 0 && isAtVeryTop) || (e.deltaY > 0 && isAtVeryBottom)) {
+					// Prevent default to stop inner container from handling
+					e.preventDefault()
+
+					const boost = 0.15
+					velocity += e.deltaY * boost
+
+					// Start animation if not already running
+					if (!animationFrameId) {
+						animationFrameId = requestAnimationFrame(animate)
+					}
+				}
+			}
+
+			// Add wheel event listener to inner container
+			const preElement = preRef.current
+			preElement.addEventListener("wheel", handleWheel, { passive: false })
+
+			// Clean up
+			return () => {
+				preElement.removeEventListener("wheel", handleWheel)
+
+				// Cancel any ongoing animation
+				if (animationFrameId) {
+					cancelAnimationFrame(animationFrameId)
+				}
+			}
+		}, [])
+
+		// Track text selection state
+		const [isSelecting, setIsSelecting] = useState(false)
+
+		useEffect(() => {
+			if (!preRef.current) return
+
+			const handleMouseDown = (e: MouseEvent) => {
+				// Only trigger if clicking the pre element directly
+				if (e.currentTarget === preRef.current) {
+					setIsSelecting(true)
+				}
+			}
+
+			const handleMouseUp = () => {
+				setIsSelecting(false)
+			}
+
+			const preElement = preRef.current
+			preElement.addEventListener("mousedown", handleMouseDown)
+			document.addEventListener("mouseup", handleMouseUp)
+
+			return () => {
+				preElement.removeEventListener("mousedown", handleMouseDown)
+				document.removeEventListener("mouseup", handleMouseUp)
+			}
+		}, [])
+
+		const handleCopy = useCallback(
+			(e: React.MouseEvent) => {
+				e.stopPropagation()
+
+				// Check if code block is partially visible before allowing copy
+				const codeBlock = codeBlockRef.current
+				if (!codeBlock || codeBlock.getAttribute("data-partially-visible") !== "true") {
+					return
+				}
+				const textToCopy = rawSource !== undefined ? rawSource : source || ""
+				if (textToCopy) {
+					copyWithFeedback(textToCopy, e)
 				}
 			},
-		],
-		rehypePlugins: [
-			rehypeHighlight as any,
-			{
-				// languages: {},
-			} as Options,
-		],
-		rehypeReactOptions: {
-			components: {
-				pre: ({ node: _, ...preProps }: any) => <StyledPre {...preProps} theme={theme} />,
-			},
-		},
-	})
-
-	useEffect(() => {
-		setMarkdownSource(source || "")
-	}, [source, setMarkdownSource, theme])
-
-	return (
-		<div
-			style={{
-				overflowY: forceWrap ? "visible" : "auto",
-				maxHeight: forceWrap ? "none" : "100%",
-				backgroundColor: CODE_BLOCK_BG_COLOR,
-			}}>
-			<StyledMarkdown forceWrap={forceWrap}>{reactContent}</StyledMarkdown>
-		</div>
-	)
-})
+			[source, rawSource, copyWithFeedback],
+		)
+
+		return (
+			<CodeBlockContainer ref={codeBlockRef}>
+				<StyledPre
+					ref={preRef}
+					preStyle={preStyle}
+					wordwrap={wordWrap ? "true" : "false"}
+					windowshade={windowShade ? "true" : "false"}
+					collapsedHeight={collapsedHeight}
+					onMouseDown={() => updateCodeBlockButtonPosition(true)}
+					onMouseUp={() => updateCodeBlockButtonPosition(false)}
+					// onScroll prop is removed - handled by the useEffect scroll listener now
+				>
+					<div dangerouslySetInnerHTML={{ __html: highlightedCode }} />
+				</StyledPre>
+				{!isSelecting && (
+					<CodeBlockButtonWrapper
+						ref={copyButtonWrapperRef}
+						onMouseOver={() => updateCodeBlockButtonPosition()}
+						style={{ gap: 0 }}>
+						{language && (
+							<LanguageSelect
+								value={currentLanguage}
+								style={{
+									alignContent: "middle",
+									width: `${Math.max(3, (currentLanguage?.length || 5) + 1)}ch`,
+									textAlign: "right",
+									marginRight: 0,
+								}}
+								onClick={(e) => {
+									e.currentTarget.focus()
+								}}
+								onChange={(e) => {
+									const newLang = normalizeLanguage(e.target.value)
+									userChangedLanguageRef.current = true
+									setCurrentLanguage(newLang)
+									if (onLanguageChange) {
+										onLanguageChange(newLang)
+									}
+								}}>
+								{
+									// Display original language at top of list for quick selection
+									language && (
+										<option
+											value={normalizeLanguage(language)}
+											style={{ fontWeight: "bold", textAlign: "left", fontSize: "1.2em" }}>
+											{normalizeLanguage(language)}
+										</option>
+									)
+								}
+								{
+									// Display all available languages in alphabetical order
+									Object.keys(bundledLanguages)
+										.sort()
+										.map((lang) => {
+											const normalizedLang = normalizeLanguage(lang)
+											return (
+												<option
+													key={normalizedLang}
+													value={normalizedLang}
+													style={{
+														fontWeight:
+															normalizedLang === currentLanguage ? "bold" : "normal",
+														textAlign: "left",
+														fontSize:
+															normalizedLang === currentLanguage ? "1.2em" : "inherit",
+													}}>
+													{normalizedLang}
+												</option>
+											)
+										})
+								}
+							</LanguageSelect>
+						)}
+						{showCollapseButton && (
+							<CodeBlockButton
+								onClick={() => {
+									// Get the current code block element and scrollable container
+									const codeBlock = codeBlockRef.current
+									const scrollContainer = document.querySelector('[data-virtuoso-scroller="true"]')
+									if (!codeBlock || !scrollContainer) return
+
+									// Toggle window shade state
+									setWindowShade(!windowShade)
+
+									// After UI updates, ensure code block is visible and update button position
+									setTimeout(
+										() => {
+											codeBlock.scrollIntoView({
+												behavior: "smooth",
+												block: "nearest",
+											})
+
+											// Wait for scroll to complete before updating button position
+											setTimeout(() => {
+												updateCodeBlockButtonPosition()
+											}, 50)
+										},
+										WINDOW_SHADE_SETTINGS.transitionDelayS * 1000 + 50,
+									)
+								}}
+								title={`${windowShade ? "Expand" : "Collapse"} code block`}>
+								<ButtonIcon style={{ fontSize: "16px" }}>{windowShade ? "⌄" : "⌃"}</ButtonIcon>
+							</CodeBlockButton>
+						)}
+						<CodeBlockButton
+							onClick={() => setWordWrap(!wordWrap)}
+							title={`${wordWrap ? "Disable" : "Enable"} word wrap`}>
+							<ButtonIcon style={{ fontSize: "16px", fontWeight: 900 }}>
+								{wordWrap ? "⟼" : "⤸"}
+							</ButtonIcon>
+						</CodeBlockButton>
+						<CodeBlockButton onClick={handleCopy} title="Copy code">
+							<ButtonIcon className={`codicon codicon-${showCopyFeedback ? "check" : "copy"}`} />
+						</CodeBlockButton>
+					</CodeBlockButtonWrapper>
+				)}
+			</CodeBlockContainer>
+		)
+	},
+)
 
 export default CodeBlock

+ 36 - 78
webview-ui/src/components/common/MarkdownBlock.tsx

@@ -1,10 +1,11 @@
 import React, { 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 "@src/context/ExtensionStateContext"
-import { CODE_BLOCK_BG_COLOR } from "./CodeBlock"
+
+import CodeBlock from "./CodeBlock"
 import MermaidBlock from "./MermaidBlock"
 
 interface MarkdownBlockProps {
@@ -25,19 +26,21 @@ const remarkUrlToLink = () => {
 		visit(tree, "text", (node: any, index, parent) => {
 			const urlRegex = /https?:\/\/[^\s<>)"]+/g
 			const matches = node.value.match(urlRegex)
-			if (!matches) return
+
+			if (!matches) {
+				return
+			}
 
 			const parts = node.value.split(urlRegex)
 			const children: any[] = []
 
 			parts.forEach((part: string, i: number) => {
-				if (part) children.push({ type: "text", value: part })
+				if (part) {
+					children.push({ type: "text", value: part })
+				}
+
 				if (matches[i]) {
-					children.push({
-						type: "link",
-						url: matches[i],
-						children: [{ type: "text", value: matches[i] }],
-					})
+					children.push({ type: "link", url: matches[i], children: [{ type: "text", value: matches[i] }] })
 				}
 			})
 
@@ -52,45 +55,6 @@ const remarkUrlToLink = () => {
 }
 
 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: auto;
-		overflow-y: hidden;
-		white-space: pre-wrap;
-	}
-
-	:where(h1, h2, h3, h4, h5, h6):has(code) code {
-		font-size: inherit;
-	}
-
-	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);
@@ -124,6 +88,7 @@ const StyledMarkdown = styled.div`
 		"Open Sans",
 		"Helvetica Neue",
 		sans-serif;
+
 	font-size: var(--vscode-font-size, 13px);
 
 	p,
@@ -153,23 +118,6 @@ const StyledMarkdown = styled.div`
 	}
 `
 
-const StyledPre = styled.pre<{ theme: any }>`
-	& .hljs {
-		color: var(--vscode-editor-foreground, #fff);
-	}
-
-	${(props) =>
-		Object.keys(props.theme)
-			.map((key) => {
-				return `
-      & ${key} {
-        color: ${props.theme[key]};
-      }
-    `
-			})
-			.join("")}
-`
-
 const MarkdownBlock = memo(({ markdown }: MarkdownBlockProps) => {
 	const { theme } = useExtensionState()
 	const [reactContent, setMarkdown] = useRemark({
@@ -179,7 +127,7 @@ const MarkdownBlock = memo(({ markdown }: MarkdownBlockProps) => {
 				return (tree) => {
 					visit(tree, "code", (node: any) => {
 						if (!node.lang) {
-							node.lang = "javascript"
+							node.lang = "text"
 						} else if (node.lang.includes(".")) {
 							node.lang = node.lang.split(".").slice(-1)[0]
 						}
@@ -187,33 +135,43 @@ const MarkdownBlock = memo(({ markdown }: MarkdownBlockProps) => {
 				}
 			},
 		],
-		rehypePlugins: [
-			rehypeHighlight as any,
-			{
-				// languages: {},
-			} as Options,
-		],
+		rehypePlugins: [],
 		rehypeReactOptions: {
 			components: {
-				pre: ({ node: _, children, ...preProps }: any) => {
+				pre: ({ node: _, children }: any) => {
+					// Check for Mermaid diagrams first
 					if (Array.isArray(children) && children.length === 1 && React.isValidElement(children[0])) {
 						const child = children[0] as React.ReactElement<{ className?: string }>
+
 						if (child.props?.className?.includes("language-mermaid")) {
 							return child
 						}
 					}
-					return (
-						<StyledPre {...preProps} theme={theme}>
-							{children}
-						</StyledPre>
-					)
+
+					// For all other code blocks, use CodeBlock with copy button
+					const codeNode = children?.[0]
+
+					if (!codeNode?.props?.children) {
+						return null
+					}
+
+					const language =
+						(Array.isArray(codeNode.props?.className)
+							? codeNode.props.className
+							: [codeNode.props?.className]
+						).map((c: string) => c?.replace("language-", ""))[0] || "javascript"
+
+					const rawText = codeNode.props.children[0] || ""
+					return <CodeBlock source={rawText} language={language} />
 				},
 				code: (props: any) => {
 					const className = props.className || ""
+
 					if (className.includes("language-mermaid")) {
 						const codeText = String(props.children || "")
 						return <MermaidBlock code={codeText} />
 					}
+
 					return <code {...props} />
 				},
 			},

+ 159 - 0
webview-ui/src/components/common/__tests__/CodeBlock.test.tsx

@@ -0,0 +1,159 @@
+import { render, screen, fireEvent, act } from "@testing-library/react"
+import "@testing-library/jest-dom"
+import CodeBlock from "../CodeBlock"
+
+// Mock shiki module
+jest.mock("shiki", () => ({
+	bundledLanguages: {
+		typescript: {},
+		javascript: {},
+		txt: {},
+	},
+}))
+
+// Mock the highlighter utility
+jest.mock("../../../utils/highlighter", () => {
+	const mockHighlighter = {
+		codeToHtml: jest.fn().mockImplementation((code, options) => {
+			const theme = options.theme === "github-light" ? "light" : "dark"
+			return `<pre><code class="hljs language-${options.lang}">${code} [${theme}-theme]</code></pre>`
+		}),
+	}
+
+	return {
+		normalizeLanguage: jest.fn((lang) => lang || "txt"),
+		isLanguageLoaded: jest.fn().mockReturnValue(true),
+		getHighlighter: jest.fn().mockResolvedValue(mockHighlighter),
+	}
+})
+
+// Mock clipboard utility
+jest.mock("../../../utils/clipboard", () => ({
+	useCopyToClipboard: () => ({
+		showCopyFeedback: false,
+		copyWithFeedback: jest.fn(),
+	}),
+}))
+
+describe("CodeBlock", () => {
+	const mockIntersectionObserver = jest.fn()
+	const originalGetComputedStyle = window.getComputedStyle
+
+	beforeEach(() => {
+		// Mock scroll container
+		const scrollContainer = document.createElement("div")
+		scrollContainer.setAttribute("data-virtuoso-scroller", "true")
+		document.body.appendChild(scrollContainer)
+
+		// Mock IntersectionObserver
+		window.IntersectionObserver = mockIntersectionObserver
+
+		// Mock getComputedStyle
+		window.getComputedStyle = jest.fn().mockImplementation((element) => ({
+			...originalGetComputedStyle(element),
+			getPropertyValue: () => "12px",
+		}))
+	})
+
+	afterEach(() => {
+		jest.clearAllMocks()
+		const scrollContainer = document.querySelector('[data-virtuoso-scroller="true"]')
+		if (scrollContainer) {
+			document.body.removeChild(scrollContainer)
+		}
+		window.getComputedStyle = originalGetComputedStyle
+	})
+
+	it("renders basic syntax highlighting", async () => {
+		const code = "const x = 1;\nconsole.log(x);"
+
+		await act(async () => {
+			render(<CodeBlock source={code} language="typescript" />)
+		})
+
+		expect(screen.getByText(/const x = 1/)).toBeInTheDocument()
+	})
+
+	it("handles theme switching", async () => {
+		const code = "const x = 1;"
+
+		await act(async () => {
+			const { rerender } = render(<CodeBlock source={code} language="typescript" />)
+
+			// Simulate light theme
+			document.body.className = "light"
+			rerender(<CodeBlock source={code} language="typescript" />)
+		})
+
+		expect(screen.getByText(/\[light-theme\]/)).toBeInTheDocument()
+
+		await act(async () => {
+			document.body.className = "dark"
+			render(<CodeBlock source={code} language="typescript" />)
+		})
+
+		expect(screen.getByText(/\[dark-theme\]/)).toBeInTheDocument()
+	})
+
+	it("handles invalid language gracefully", async () => {
+		const code = "some code"
+
+		await act(async () => {
+			render(<CodeBlock source={code} language="invalid-lang" />)
+		})
+
+		expect(screen.getByText(/some code/)).toBeInTheDocument()
+	})
+
+	it("handles WASM loading errors", async () => {
+		const mockError = new Error("WASM load failed")
+		const highlighterUtil = require("../../../utils/highlighter")
+		highlighterUtil.getHighlighter.mockRejectedValueOnce(mockError)
+
+		const code = "const x = 1;"
+		const consoleSpy = jest.spyOn(console, "error").mockImplementation()
+
+		await act(async () => {
+			render(<CodeBlock source={code} language="typescript" />)
+		})
+
+		expect(consoleSpy).toHaveBeenCalledWith(
+			"[CodeBlock] Syntax highlighting error:",
+			mockError,
+			"\nStack trace:",
+			mockError.stack,
+		)
+		expect(screen.getByText(/const x = 1;/)).toBeInTheDocument()
+
+		consoleSpy.mockRestore()
+	})
+
+	it("verifies highlighter utility is used correctly", async () => {
+		const code = "const x = 1;"
+		const highlighterUtil = require("../../../utils/highlighter")
+
+		await act(async () => {
+			render(<CodeBlock source={code} language="typescript" />)
+		})
+
+		// Verify getHighlighter was called with the right language
+		expect(highlighterUtil.getHighlighter).toHaveBeenCalledWith("typescript")
+		expect(highlighterUtil.normalizeLanguage).toHaveBeenCalledWith("typescript")
+	})
+
+	it("handles copy functionality", async () => {
+		const code = "const x = 1;"
+		const { container } = render(<CodeBlock source={code} language="typescript" />)
+
+		// Simulate code block visibility
+		const codeBlock = container.querySelector("[data-partially-visible]")
+		if (codeBlock) {
+			codeBlock.setAttribute("data-partially-visible", "true")
+		}
+
+		const copyButton = screen.getByTitle("Copy code")
+		await act(async () => {
+			fireEvent.click(copyButton)
+		})
+	})
+})

+ 5 - 0
webview-ui/src/index.tsx

@@ -5,6 +5,11 @@ import "./index.css"
 import App from "./App"
 import "../../node_modules/@vscode/codicons/dist/codicon.css"
 
+import { getHighlighter } from "./utils/highlighter"
+
+// Initialize Shiki early to hide initialization latency (async)
+getHighlighter().catch((error: Error) => console.error("Failed to initialize Shiki highlighter:", error))
+
 createRoot(document.getElementById("root")!).render(
 	<StrictMode>
 		<App />

+ 214 - 0
webview-ui/src/utils/highlighter.ts

@@ -0,0 +1,214 @@
+import {
+	createHighlighter,
+	type Highlighter,
+	type BundledTheme,
+	type BundledLanguage,
+	bundledLanguages,
+	bundledThemes,
+} from "shiki"
+
+// Extend BundledLanguage to include 'txt' because Shiki supports this but it is
+// not listed in the bundled languages
+export type ExtendedLanguage = BundledLanguage | "txt"
+
+// Map common language aliases to their Shiki BundledLanguage equivalent
+const languageAliases: Record<string, ExtendedLanguage> = {
+	// Plain text variants
+	text: "txt",
+	plaintext: "txt",
+	plain: "txt",
+
+	// Shell/Bash variants
+	sh: "shell",
+	bash: "shell",
+	zsh: "shell",
+	shellscript: "shell",
+	"shell-script": "shell",
+	console: "shell",
+	terminal: "shell",
+
+	// JavaScript variants
+	js: "javascript",
+	node: "javascript",
+	nodejs: "javascript",
+
+	// TypeScript variants
+	ts: "typescript",
+
+	// Python variants
+	py: "python",
+	python3: "python",
+	py3: "python",
+
+	// Ruby variants
+	rb: "ruby",
+
+	// Markdown variants
+	md: "markdown",
+
+	// C++ variants
+	cpp: "c++",
+	cc: "c++",
+
+	// C# variants
+	cs: "c#",
+	csharp: "c#",
+
+	// HTML variants
+	htm: "html",
+
+	// YAML variants
+	yml: "yaml",
+
+	// Docker variants
+	dockerfile: "docker",
+
+	// CSS variants
+	styles: "css",
+	style: "css",
+
+	// JSON variants
+	jsonc: "json",
+	json5: "json",
+
+	// XML variants
+	xaml: "xml",
+	xhtml: "xml",
+	svg: "xml",
+
+	// SQL variants
+	mysql: "sql",
+	postgresql: "sql",
+	postgres: "sql",
+	pgsql: "sql",
+	plsql: "sql",
+	oracle: "sql",
+}
+
+// Track which languages we've warned about to avoid duplicate warnings
+const warnedLanguages = new Set<string>()
+
+// Normalize language to a valid Shiki language
+export function normalizeLanguage(language: string | undefined): ExtendedLanguage {
+	if (language === undefined) {
+		return "txt"
+	}
+
+	// Convert to lowercase for consistent matching
+	const normalizedInput = language.toLowerCase()
+
+	// If it's already a valid bundled language, return it
+	if (normalizedInput in bundledLanguages) {
+		return normalizedInput as BundledLanguage
+	}
+
+	// Check if it's an alias
+	if (normalizedInput in languageAliases) {
+		return languageAliases[normalizedInput]
+	}
+
+	// Warn about unrecognized language and default to txt (only once per language)
+	if (language !== "txt" && !warnedLanguages.has(language)) {
+		console.warn(`[Shiki] Unrecognized language '${language}', defaulting to txt.`)
+		warnedLanguages.add(language)
+	}
+
+	return "txt"
+}
+
+// Export function to check if a language is loaded
+export const isLanguageLoaded = (language: string): boolean => {
+	return state.loadedLanguages.has(normalizeLanguage(language))
+}
+
+// Artificial delay for testing language loading (ms) - for testing
+const LANGUAGE_LOAD_DELAY = 0
+
+// Common languages for first-stage initialization
+const initialLanguages: BundledLanguage[] = ["shell", "log"]
+
+// Singleton state
+let state: {
+	instance: Highlighter | null
+	instanceInitPromise: Promise<Highlighter> | null
+	loadedLanguages: Set<ExtendedLanguage>
+	pendingLanguageLoads: Map<ExtendedLanguage, Promise<void>>
+} = {
+	instance: null,
+	instanceInitPromise: null,
+	loadedLanguages: new Set<ExtendedLanguage>(["txt"]),
+	pendingLanguageLoads: new Map(),
+}
+
+export const getHighlighter = async (language?: string): Promise<Highlighter> => {
+	try {
+		const shikilang = normalizeLanguage(language)
+
+		// Initialize highlighter if needed
+		if (!state.instanceInitPromise) {
+			state.instanceInitPromise = (async () => {
+				// const startTime = performance.now()
+				// console.debug("[Shiki] Initialization started...")
+
+				const instance = await createHighlighter({
+					themes: Object.keys(bundledThemes) as BundledTheme[],
+					langs: initialLanguages,
+				})
+
+				// const elapsed = Math.round(performance.now() - startTime)
+				// console.debug(`[Shiki] Initialization complete (${elapsed}ms)`)
+
+				state.instance = instance
+
+				// Track initially loaded languages
+				initialLanguages.forEach((lang) => state.loadedLanguages.add(lang))
+
+				return instance
+			})()
+		}
+
+		// Wait for initialization to complete
+		const instance = await state.instanceInitPromise
+
+		// Load requested language if needed (txt is already in loadedLanguages)
+		if (!state.loadedLanguages.has(shikilang)) {
+			// Check for existing pending load
+			let loadingPromise = state.pendingLanguageLoads.get(shikilang)
+
+			if (!loadingPromise) {
+				// const loadStart = performance.now()
+				// Create new loading promise
+				loadingPromise = (async () => {
+					try {
+						// Add artificial delay for testing if nonzero
+						if (LANGUAGE_LOAD_DELAY > 0) {
+							await new Promise((resolve) => setTimeout(resolve, LANGUAGE_LOAD_DELAY))
+						}
+
+						await instance.loadLanguage(shikilang as BundledLanguage)
+						state.loadedLanguages.add(shikilang)
+
+						// const loadTime = Math.round(performance.now() - loadStart)
+						// console.debug(`[Shiki] Loaded language ${shikilang} (${loadTime}ms)`)
+					} catch (error) {
+						console.error(`[Shiki] Failed to load language ${shikilang}:`, error)
+						throw error
+					} finally {
+						// Clean up pending promise after completion
+						state.pendingLanguageLoads.delete(shikilang)
+					}
+				})()
+
+				// Store the promise
+				state.pendingLanguageLoads.set(shikilang, loadingPromise)
+			}
+
+			await loadingPromise
+		}
+
+		return instance
+	} catch (error) {
+		console.error("[Shiki] Error in getHighlighter:", error)
+		throw error
+	}
+}

+ 1 - 1
webview-ui/tsconfig.json

@@ -10,7 +10,7 @@
 		"forceConsistentCasingInFileNames": true,
 		"noFallthroughCasesInSwitch": true,
 		"module": "esnext",
-		"moduleResolution": "node",
+		"moduleResolution": "bundler",
 		"resolveJsonModule": true,
 		"isolatedModules": true,
 		"noEmit": true,

+ 22 - 2
webview-ui/vite.config.ts

@@ -1,10 +1,26 @@
 import { resolve } from "path"
 import fs from "fs"
 
-import { defineConfig } from "vite"
+import { defineConfig, type Plugin } from "vite"
 import react from "@vitejs/plugin-react"
 import tailwindcss from "@tailwindcss/vite"
 
+function wasmPlugin(): Plugin {
+	return {
+		name: "wasm",
+		async load(id: string) {
+			if (id.endsWith(".wasm")) {
+				const wasmBinary = await import(id)
+
+				return `
+          			const wasmModule = new WebAssembly.Module(${wasmBinary.default});
+          			export default wasmModule;
+        		`
+			}
+		},
+	}
+}
+
 // Custom plugin to write the server port to a file
 const writePortToFile = () => {
 	return {
@@ -31,7 +47,7 @@ const writePortToFile = () => {
 
 // https://vitejs.dev/config/
 export default defineConfig({
-	plugins: [react(), tailwindcss(), writePortToFile()],
+	plugins: [react(), tailwindcss(), writePortToFile(), wasmPlugin()],
 	resolve: {
 		alias: {
 			"@": resolve(__dirname, "./src"),
@@ -65,4 +81,8 @@ export default defineConfig({
 		"process.platform": JSON.stringify(process.platform),
 		"process.env.VSCODE_TEXTMATE_DEBUG": JSON.stringify(process.env.VSCODE_TEXTMATE_DEBUG),
 	},
+	optimizeDeps: {
+		exclude: ["@vscode/codicons", "vscode-oniguruma", "shiki"],
+	},
+	assetsInclude: ["**/*.wasm"],
 })