| ... | @@ -1,8 +1,14 @@ | ... | @@ -1,8 +1,14 @@ |
| 1 | import { getSharedHighlighter } from "@pierre/diffs"; | 1 | import { |
| | 2 | getSharedHighlighter, |
| | 3 | type DiffsHighlighter, |
| | 4 | type SupportedLanguages, |
| | 5 | } from "@pierre/diffs"; |
| 2 | import { CheckIcon, CopyIcon } from "lucide-react"; | 6 | import { CheckIcon, CopyIcon } from "lucide-react"; |
| 3 | import { | 7 | import { |
| 4 | Children, | 8 | Children, |
| | 9 | Suspense, |
| 5 | isValidElement, | 10 | isValidElement, |
| | 11 | use, |
| 6 | useCallback, | 12 | useCallback, |
| 7 | memo, | 13 | memo, |
| 8 | useEffect, | 14 | useEffect, |
| ... | @@ -14,12 +20,12 @@ import { | ... | @@ -14,12 +20,12 @@ import { |
| 14 | import type { Components } from "react-markdown"; | 20 | import type { Components } from "react-markdown"; |
| 15 | import ReactMarkdown from "react-markdown"; | 21 | import ReactMarkdown from "react-markdown"; |
| 16 | import remarkGfm from "remark-gfm"; | 22 | import remarkGfm from "remark-gfm"; |
| 17 | import { LRUCache } from "../lib/lruCache"; | | |
| 18 | import { readNativeApi } from "../nativeApi"; | | |
| 19 | import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffThemes"; | 23 | import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffThemes"; |
| 20 | import { fnv1a32 } from "../lib/diffRendering"; | 24 | import { fnv1a32 } from "../lib/diffRendering"; |
| | 25 | import { LRUCache } from "../lib/lruCache"; |
| 21 | import { useTheme } from "../hooks/useTheme"; | 26 | import { useTheme } from "../hooks/useTheme"; |
| 22 | import { resolveMarkdownFileLinkTarget } from "../markdown-links"; | 27 | import { resolveMarkdownFileLinkTarget } from "../markdown-links"; |
| | 28 | import { readNativeApi } from "../nativeApi"; |
| 23 | import { preferredTerminalEditor } from "../terminal-links"; | 29 | import { preferredTerminalEditor } from "../terminal-links"; |
| 24 | | 30 | |
| 25 | interface ChatMarkdownProps { | 31 | interface ChatMarkdownProps { |
| ... | @@ -35,11 +41,7 @@ const highlightedCodeCache = new LRUCache<string>( | ... | @@ -35,11 +41,7 @@ const highlightedCodeCache = new LRUCache<string>( |
| 35 | MAX_HIGHLIGHT_CACHE_ENTRIES, | 41 | MAX_HIGHLIGHT_CACHE_ENTRIES, |
| 36 | MAX_HIGHLIGHT_CACHE_MEMORY_BYTES, | 42 | MAX_HIGHLIGHT_CACHE_MEMORY_BYTES, |
| 37 | ); | 43 | ); |
| 38 | interface InFlightHighlight { | 44 | const highlighterPromiseCache = new Map<string, Promise<DiffsHighlighter>>(); |
| 39 | promise: Promise<string>; | | |
| 40 | enableCache: () => void; | | |
| 41 | } | | |
| 42 | const inFlightHighlights = new Map<string, InFlightHighlight>(); | | |
| 43 | | 45 | |
| 44 | function extractFenceLanguage(className: string | undefined): string { | 46 | function extractFenceLanguage(className: string | undefined): string { |
| 45 | const match = className?.match(CODE_FENCE_LANGUAGE_REGEX); | 47 | const match = className?.match(CODE_FENCE_LANGUAGE_REGEX); |
| ... | @@ -81,85 +83,33 @@ function extractCodeBlock( | ... | @@ -81,85 +83,33 @@ function extractCodeBlock( |
| 81 | }; | 83 | }; |
| 82 | } | 84 | } |
| 83 | | 85 | |
| 84 | async function highlightCodeBlock(code: string, language: string, themeName: DiffThemeName) { | | |
| 85 | try { | | |
| 86 | const highlighter = await getSharedHighlighter({ | | |
| 87 | themes: [themeName], | | |
| 88 | langs: [language], | | |
| 89 | }); | | |
| 90 | return highlighter.codeToHtml(code, { lang: language, theme: themeName }); | | |
| 91 | } catch { | | |
| 92 | const highlighter = await getSharedHighlighter({ | | |
| 93 | themes: [themeName], | | |
| 94 | langs: ["text"], | | |
| 95 | }); | | |
| 96 | return highlighter.codeToHtml(code, { lang: "text", theme: themeName }); | | |
| 97 | } | | |
| 98 | } | | |
| 99 | | | |
| 100 | function createHighlightCacheKey(code: string, language: string, themeName: DiffThemeName): string { | 86 | function createHighlightCacheKey(code: string, language: string, themeName: DiffThemeName): string { |
| 101 | return `${fnv1a32(code).toString(36)}:${code.length}:${language}:${themeName}`; | 87 | return `${fnv1a32(code).toString(36)}:${code.length}:${language}:${themeName}`; |
| 102 | } | 88 | } |
| 103 | | 89 | |
| | 90 | function createHighlighterPromiseKey(language: string, themeName: DiffThemeName): string { |
| | 91 | return `${language}:${themeName}`; |
| | 92 | } |
| | 93 | |
| 104 | function estimateHighlightedSize(html: string, code: string): number { | 94 | function estimateHighlightedSize(html: string, code: string): number { |
| 105 | return Math.max(html.length * 2, code.length * 3); | 95 | return Math.max(html.length * 2, code.length * 3); |
| 106 | } | 96 | } |
| 107 | | 97 | |
| 108 | function getCachedOrCreateHighlight({ | 98 | function getHighlighterPromise( |
| 109 | cacheKey, | 99 | language: string, |
| 110 | code, | 100 | themeName: DiffThemeName, |
| 111 | language, | 101 | ): Promise<DiffsHighlighter> { |
| 112 | themeName, | 102 | const cacheKey = createHighlighterPromiseKey(language, themeName); |
| 113 | useCache, | 103 | const cached = highlighterPromiseCache.get(cacheKey); |
| 114 | }: { | 104 | if (cached) return cached; |
| 115 | cacheKey: string; | | |
| 116 | code: string; | | |
| 117 | language: string; | | |
| 118 | themeName: DiffThemeName; | | |
| 119 | useCache: boolean; | | |
| 120 | }): Promise<string> { | | |
| 121 | if (useCache) { | | |
| 122 | const cached = highlightedCodeCache.get(cacheKey); | | |
| 123 | if (cached) { | | |
| 124 | return Promise.resolve(cached); | | |
| 125 | } | | |
| 126 | } | | |
| 127 | | | |
| 128 | const existingInFlight = inFlightHighlights.get(cacheKey); | | |
| 129 | if (existingInFlight) { | | |
| 130 | if (useCache) { | | |
| 131 | existingInFlight.enableCache(); | | |
| 132 | } | | |
| 133 | return existingInFlight.promise; | | |
| 134 | } | | |
| 135 | | 105 | |
| 136 | let shouldCache = useCache; | 106 | const promise = getSharedHighlighter({ |
| 137 | const requestPromise = highlightCodeBlock(code, language, themeName) | 107 | themes: [themeName], |
| 138 | .then((html) => { | 108 | langs: [language as SupportedLanguages], |
| 139 | if (shouldCache) { | 109 | preferredHighlighter: "shiki-js", |
| 140 | highlightedCodeCache.set(cacheKey, html, estimateHighlightedSize(html, code)); | | |
| 141 | } | | |
| 142 | return html; | | |
| 143 | }) | | |
| 144 | .finally(() => { | | |
| 145 | inFlightHighlights.delete(cacheKey); | | |
| 146 | }); | | |
| 147 | | | |
| 148 | inFlightHighlights.set(cacheKey, { | | |
| 149 | promise: requestPromise, | | |
| 150 | enableCache: () => { | | |
| 151 | shouldCache = true; | | |
| 152 | }, | | |
| 153 | }); | 110 | }); |
| 154 | return requestPromise; | 111 | highlighterPromiseCache.set(cacheKey, promise); |
| 155 | } | 112 | return promise; |
| 156 | | | |
| 157 | interface ShikiCodeBlockProps { | | |
| 158 | className: string | undefined; | | |
| 159 | code: string; | | |
| 160 | themeName: DiffThemeName; | | |
| 161 | isStreaming: boolean; | | |
| 162 | fallback: ReactNode; | | |
| 163 | } | 113 | } |
| 164 | | 114 | |
| 165 | function MarkdownCodeBlock({ code, children }: { code: string; children: ReactNode }) { | 115 | function MarkdownCodeBlock({ code, children }: { code: string; children: ReactNode }) { |
| ... | @@ -210,70 +160,47 @@ function MarkdownCodeBlock({ code, children }: { code: string; children: ReactNo | ... | @@ -210,70 +160,47 @@ function MarkdownCodeBlock({ code, children }: { code: string; children: ReactNo |
| 210 | ); | 160 | ); |
| 211 | } | 161 | } |
| 212 | | 162 | |
| 213 | function ShikiCodeBlock({ className, code, themeName, isStreaming, fallback }: ShikiCodeBlockProps) { | 163 | interface SuspenseShikiCodeBlockProps { |
| | 164 | className: string | undefined; |
| | 165 | code: string; |
| | 166 | themeName: DiffThemeName; |
| | 167 | isStreaming: boolean; |
| | 168 | } |
| | 169 | |
| | 170 | function SuspenseShikiCodeBlock({ |
| | 171 | className, |
| | 172 | code, |
| | 173 | themeName, |
| | 174 | isStreaming, |
| | 175 | }: SuspenseShikiCodeBlockProps) { |
| 214 | const language = extractFenceLanguage(className); | 176 | const language = extractFenceLanguage(className); |
| 215 | const cacheKey = createHighlightCacheKey(code, language, themeName); | 177 | const cacheKey = createHighlightCacheKey(code, language, themeName); |
| 216 | const [highlightState, setHighlightState] = useState<{ | 178 | const cachedHighlightedHtml = !isStreaming ? highlightedCodeCache.get(cacheKey) : null; |
| 217 | cacheKey: string; | | |
| 218 | html: string | null; | | |
| 219 | }>(() => ({ | | |
| 220 | cacheKey, | | |
| 221 | html: isStreaming ? null : (highlightedCodeCache.get(cacheKey) ?? null), | | |
| 222 | })); | | |
| 223 | const cachedHighlightedHtml = isStreaming ? null : (highlightedCodeCache.get(cacheKey) ?? null); | | |
| 224 | const highlightedHtml = | | |
| 225 | cachedHighlightedHtml ?? (highlightState.cacheKey === cacheKey ? highlightState.html : null); | | |
| 226 | | | |
| 227 | useEffect(() => { | | |
| 228 | let cancelled = false; | | |
| 229 | | 179 | |
| 230 | if (cachedHighlightedHtml) { | 180 | if (cachedHighlightedHtml != null) { |
| 231 | setHighlightState((current) => { | 181 | return ( |
| 232 | if (current.cacheKey === cacheKey && current.html === cachedHighlightedHtml) { | 182 | <div |
| 233 | return current; | 183 | className="chat-markdown-shiki" |
| 234 | } | 184 | dangerouslySetInnerHTML={{ __html: cachedHighlightedHtml }} |
| 235 | return { cacheKey, html: cachedHighlightedHtml }; | 185 | /> |
| 236 | }); | 186 | ); |
| 237 | return () => { | 187 | } |
| 238 | cancelled = true; | | |
| 239 | }; | | |
| 240 | } | | |
| 241 | | | |
| 242 | setHighlightState((current) => { | | |
| 243 | if (current.cacheKey === cacheKey) { | | |
| 244 | return current; | | |
| 245 | } | | |
| 246 | return { cacheKey, html: null }; | | |
| 247 | }); | | |
| 248 | | | |
| 249 | void getCachedOrCreateHighlight({ | | |
| 250 | cacheKey, | | |
| 251 | code, | | |
| 252 | language, | | |
| 253 | themeName, | | |
| 254 | useCache: !isStreaming, | | |
| 255 | }) | | |
| 256 | .then((html) => { | | |
| 257 | if (cancelled) { | | |
| 258 | return; | | |
| 259 | } | | |
| 260 | setHighlightState((current) => { | | |
| 261 | if (current.cacheKey === cacheKey && current.html === html) { | | |
| 262 | return current; | | |
| 263 | } | | |
| 264 | return { cacheKey, html }; | | |
| 265 | }); | | |
| 266 | }) | | |
| 267 | .catch(() => undefined); | | |
| 268 | | 188 | |
| 269 | return () => { | 189 | const highlighter = use(getHighlighterPromise(language, themeName)); |
| 270 | cancelled = true; | 190 | const highlightedHtml = useMemo( |
| 271 | }; | 191 | () => highlighter.codeToHtml(code, { lang: language, theme: themeName }), |
| 272 | }, [cacheKey, cachedHighlightedHtml, code, isStreaming, language, themeName]); | 192 | [code, highlighter, language, themeName], |
| | 193 | ); |
| 273 | | 194 | |
| 274 | if (!highlightedHtml) { | 195 | useEffect(() => { |
| 275 | return fallback; | 196 | if (!isStreaming) { |
| 276 | } | 197 | highlightedCodeCache.set( |
| | 198 | cacheKey, |
| | 199 | highlightedHtml, |
| | 200 | estimateHighlightedSize(highlightedHtml, code), |
| | 201 | ); |
| | 202 | } |
| | 203 | }, [cacheKey, code, highlightedHtml, isStreaming]); |
| 277 | | 204 | |
| 278 | return ( | 205 | return ( |
| 279 | <div className="chat-markdown-shiki" dangerouslySetInnerHTML={{ __html: highlightedHtml }} /> | 206 | <div className="chat-markdown-shiki" dangerouslySetInnerHTML={{ __html: highlightedHtml }} /> |
| ... | @@ -316,13 +243,14 @@ function ChatMarkdown({ text, cwd, isStreaming = false }: ChatMarkdownProps) { | ... | @@ -316,13 +243,14 @@ function ChatMarkdown({ text, cwd, isStreaming = false }: ChatMarkdownProps) { |
| 316 | | 243 | |
| 317 | return ( | 244 | return ( |
| 318 | <MarkdownCodeBlock code={codeBlock.code}> | 245 | <MarkdownCodeBlock code={codeBlock.code}> |
| 319 | <ShikiCodeBlock | 246 | <Suspense fallback={<pre {...props}>{children}</pre>}> |
| 320 | className={codeBlock.className} | 247 | <SuspenseShikiCodeBlock |
| 321 | code={codeBlock.code} | 248 | className={codeBlock.className} |
| 322 | themeName={diffThemeName} | 249 | code={codeBlock.code} |
| 323 | isStreaming={isStreaming} | 250 | themeName={diffThemeName} |
| 324 | fallback={<pre {...props}>{children}</pre>} | 251 | isStreaming={isStreaming} |
| 325 | /> | 252 | /> |
| | 253 | </Suspense> |
| 326 | </MarkdownCodeBlock> | 254 | </MarkdownCodeBlock> |
| 327 | ); | 255 | ); |
| 328 | }, | 256 | }, |