import { DiffsHighlighter, getSharedHighlighter, SupportedLanguages } from "@pierre/diffs"; import { CheckIcon, CopyIcon } from "lucide-react"; import React, { Children, Suspense, isValidElement, use, useCallback, memo, useEffect, useMemo, useRef, useState, type ReactNode, } from "react"; import type { Components } from "react-markdown"; import ReactMarkdown from "react-markdown"; import remarkGfm from "remark-gfm"; import { openInPreferredEditor } from "../editorPreferences"; import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffRendering"; import { fnv1a32 } from "../lib/diffRendering"; import { LRUCache } from "../lib/lruCache"; import { useTheme } from "../hooks/useTheme"; import { resolveMarkdownFileLinkTarget } from "../markdown-links"; import { readNativeApi } from "../nativeApi"; class CodeHighlightErrorBoundary extends React.Component< { fallback: ReactNode; children: ReactNode }, { hasError: boolean } > { constructor(props: { fallback: ReactNode; children: ReactNode }) { super(props); this.state = { hasError: false }; } static getDerivedStateFromError() { return { hasError: true }; } override render() { if (this.state.hasError) { return this.props.fallback; } return this.props.children; } } interface ChatMarkdownProps { text: string; cwd: string | undefined; isStreaming?: boolean; } const CODE_FENCE_LANGUAGE_REGEX = /(?:^|\s)language-([^\s]+)/; const MAX_HIGHLIGHT_CACHE_ENTRIES = 500; const MAX_HIGHLIGHT_CACHE_MEMORY_BYTES = 50 * 1024 * 1024; const highlightedCodeCache = new LRUCache( MAX_HIGHLIGHT_CACHE_ENTRIES, MAX_HIGHLIGHT_CACHE_MEMORY_BYTES, ); const highlighterPromiseCache = new Map>(); function extractFenceLanguage(className: string | undefined): string { const match = className?.match(CODE_FENCE_LANGUAGE_REGEX); const raw = match?.[1] ?? "text"; // Shiki doesn't bundle a gitignore grammar; ini is a close match (#685) return raw === "gitignore" ? "ini" : raw; } function nodeToPlainText(node: ReactNode): string { if (typeof node === "string" || typeof node === "number") { return String(node); } if (Array.isArray(node)) { return node.map((child) => nodeToPlainText(child)).join(""); } if (isValidElement<{ children?: ReactNode }>(node)) { return nodeToPlainText(node.props.children); } return ""; } function extractCodeBlock( children: ReactNode, ): { className: string | undefined; code: string } | null { const childNodes = Children.toArray(children); if (childNodes.length !== 1) { return null; } const onlyChild = childNodes[0]; if ( !isValidElement<{ className?: string; children?: ReactNode }>(onlyChild) || onlyChild.type !== "code" ) { return null; } return { className: onlyChild.props.className, code: nodeToPlainText(onlyChild.props.children), }; } function createHighlightCacheKey(code: string, language: string, themeName: DiffThemeName): string { return `${fnv1a32(code).toString(36)}:${code.length}:${language}:${themeName}`; } function estimateHighlightedSize(html: string, code: string): number { return Math.max(html.length * 2, code.length * 3); } function getHighlighterPromise(language: string): Promise { const cached = highlighterPromiseCache.get(language); if (cached) return cached; const promise = getSharedHighlighter({ themes: [resolveDiffThemeName("dark"), resolveDiffThemeName("light")], langs: [language as SupportedLanguages], preferredHighlighter: "shiki-js", }).catch((err) => { highlighterPromiseCache.delete(language); if (language === "text") { // "text" itself failed — Shiki cannot initialize at all, surface the error throw err; } // Language not supported by Shiki — fall back to "text" return getHighlighterPromise("text"); }); highlighterPromiseCache.set(language, promise); return promise; } function MarkdownCodeBlock({ code, children }: { code: string; children: ReactNode }) { const [copied, setCopied] = useState(false); const copiedTimerRef = useRef | null>(null); const handleCopy = useCallback(() => { if (typeof navigator === "undefined" || navigator.clipboard == null) { return; } void navigator.clipboard .writeText(code) .then(() => { if (copiedTimerRef.current != null) { clearTimeout(copiedTimerRef.current); } setCopied(true); copiedTimerRef.current = setTimeout(() => { setCopied(false); copiedTimerRef.current = null; }, 1200); }) .catch(() => undefined); }, [code]); useEffect( () => () => { if (copiedTimerRef.current != null) { clearTimeout(copiedTimerRef.current); copiedTimerRef.current = null; } }, [], ); return (
{children}
); } interface SuspenseShikiCodeBlockProps { className: string | undefined; code: string; themeName: DiffThemeName; isStreaming: boolean; } function SuspenseShikiCodeBlock({ className, code, themeName, isStreaming, }: SuspenseShikiCodeBlockProps) { const language = extractFenceLanguage(className); const cacheKey = createHighlightCacheKey(code, language, themeName); const cachedHighlightedHtml = !isStreaming ? highlightedCodeCache.get(cacheKey) : null; if (cachedHighlightedHtml != null) { return (
); } const highlighter = use(getHighlighterPromise(language)); const highlightedHtml = useMemo(() => { try { return highlighter.codeToHtml(code, { lang: language, theme: themeName }); } catch (error) { // Log highlighting failures for debugging while falling back to plain text console.warn( `Code highlighting failed for language "${language}", falling back to plain text.`, error instanceof Error ? error.message : error, ); // If highlighting fails for this language, render as plain text return highlighter.codeToHtml(code, { lang: "text", theme: themeName }); } }, [code, highlighter, language, themeName]); useEffect(() => { if (!isStreaming) { highlightedCodeCache.set( cacheKey, highlightedHtml, estimateHighlightedSize(highlightedHtml, code), ); } }, [cacheKey, code, highlightedHtml, isStreaming]); return (
); } function ChatMarkdown({ text, cwd, isStreaming = false }: ChatMarkdownProps) { const { resolvedTheme } = useTheme(); const diffThemeName = resolveDiffThemeName(resolvedTheme); const markdownComponents = useMemo( () => ({ a({ node: _node, href, ...props }) { const targetPath = resolveMarkdownFileLinkTarget(href, cwd); if (!targetPath) { return ; } return ( { event.preventDefault(); event.stopPropagation(); const api = readNativeApi(); if (api) { void openInPreferredEditor(api, targetPath); } else { console.warn("Native API not found. Unable to open file in editor."); } }} /> ); }, pre({ node: _node, children, ...props }) { const codeBlock = extractCodeBlock(children); if (!codeBlock) { return
{children}
; } return ( {children}}> {children}}> ); }, }), [cwd, diffThemeName, isStreaming], ); return (
{text}
); } export default memo(ChatMarkdown);