| ... | ... | @@ -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 | 6 | import { CheckIcon, CopyIcon } from "lucide-react"; |
| 3 | 7 | import { |
| 4 | 8 | Children, |
| 9 | Suspense, |
| 5 | 10 | isValidElement, |
| 11 | use, |
| 6 | 12 | useCallback, |
| 7 | 13 | memo, |
| 8 | 14 | useEffect, |
| ... | ... | @@ -14,12 +20,12 @@ import { |
| 14 | 20 | import type { Components } from "react-markdown"; |
| 15 | 21 | import ReactMarkdown from "react-markdown"; |
| 16 | 22 | import remarkGfm from "remark-gfm"; |
| 17 | | import { LRUCache } from "../lib/lruCache"; |
| 18 | | import { readNativeApi } from "../nativeApi"; |
| 19 | 23 | import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffThemes"; |
| 20 | 24 | import { fnv1a32 } from "../lib/diffRendering"; |
| 25 | import { LRUCache } from "../lib/lruCache"; |
| 21 | 26 | import { useTheme } from "../hooks/useTheme"; |
| 22 | 27 | import { resolveMarkdownFileLinkTarget } from "../markdown-links"; |
| 28 | import { readNativeApi } from "../nativeApi"; |
| 23 | 29 | import { preferredTerminalEditor } from "../terminal-links"; |
| 24 | 30 | |
| 25 | 31 | interface ChatMarkdownProps { |
| ... | ... | @@ -35,11 +41,7 @@ const highlightedCodeCache = new LRUCache<string>( |
| 35 | 41 | MAX_HIGHLIGHT_CACHE_ENTRIES, |
| 36 | 42 | MAX_HIGHLIGHT_CACHE_MEMORY_BYTES, |
| 37 | 43 | ); |
| 38 | | interface InFlightHighlight { |
| 39 | | promise: Promise<string>; |
| 40 | | enableCache: () => void; |
| 41 | | } |
| 42 | | const inFlightHighlights = new Map<string, InFlightHighlight>(); |
| 44 | const highlighterPromiseCache = new Map<string, Promise<DiffsHighlighter>>(); |
| 43 | 45 | |
| 44 | 46 | function extractFenceLanguage(className: string | undefined): string { |
| 45 | 47 | const match = className?.match(CODE_FENCE_LANGUAGE_REGEX); |
| ... | ... | @@ -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 | 86 | function createHighlightCacheKey(code: string, language: string, themeName: DiffThemeName): string { |
| 101 | 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 | 94 | function estimateHighlightedSize(html: string, code: string): number { |
| 105 | 95 | return Math.max(html.length * 2, code.length * 3); |
| 106 | 96 | } |
| 107 | 97 | |
| 108 | | function getCachedOrCreateHighlight({ |
| 109 | | cacheKey, |
| 110 | | code, |
| 111 | | language, |
| 112 | | themeName, |
| 113 | | useCache, |
| 114 | | }: { |
| 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 | | } |
| 98 | function getHighlighterPromise( |
| 99 | language: string, |
| 100 | themeName: DiffThemeName, |
| 101 | ): Promise<DiffsHighlighter> { |
| 102 | const cacheKey = createHighlighterPromiseKey(language, themeName); |
| 103 | const cached = highlighterPromiseCache.get(cacheKey); |
| 104 | if (cached) return cached; |
| 135 | 105 | |
| 136 | | let shouldCache = useCache; |
| 137 | | const requestPromise = highlightCodeBlock(code, language, themeName) |
| 138 | | .then((html) => { |
| 139 | | if (shouldCache) { |
| 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 | | }, |
| 106 | const promise = getSharedHighlighter({ |
| 107 | themes: [themeName], |
| 108 | langs: [language as SupportedLanguages], |
| 109 | preferredHighlighter: "shiki-js", |
| 153 | 110 | }); |
| 154 | | return requestPromise; |
| 155 | | } |
| 156 | | |
| 157 | | interface ShikiCodeBlockProps { |
| 158 | | className: string | undefined; |
| 159 | | code: string; |
| 160 | | themeName: DiffThemeName; |
| 161 | | isStreaming: boolean; |
| 162 | | fallback: ReactNode; |
| 111 | highlighterPromiseCache.set(cacheKey, promise); |
| 112 | return promise; |
| 163 | 113 | } |
| 164 | 114 | |
| 165 | 115 | function MarkdownCodeBlock({ code, children }: { code: string; children: ReactNode }) { |
| ... | ... | @@ -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 | 176 | const language = extractFenceLanguage(className); |
| 215 | 177 | const cacheKey = createHighlightCacheKey(code, language, themeName); |
| 216 | | const [highlightState, setHighlightState] = useState<{ |
| 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; |
| 178 | const cachedHighlightedHtml = !isStreaming ? highlightedCodeCache.get(cacheKey) : null; |
| 229 | 179 | |
| 230 | | if (cachedHighlightedHtml) { |
| 231 | | setHighlightState((current) => { |
| 232 | | if (current.cacheKey === cacheKey && current.html === cachedHighlightedHtml) { |
| 233 | | return current; |
| 234 | | } |
| 235 | | return { cacheKey, html: cachedHighlightedHtml }; |
| 236 | | }); |
| 237 | | return () => { |
| 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); |
| 180 | if (cachedHighlightedHtml != null) { |
| 181 | return ( |
| 182 | <div |
| 183 | className="chat-markdown-shiki" |
| 184 | dangerouslySetInnerHTML={{ __html: cachedHighlightedHtml }} |
| 185 | /> |
| 186 | ); |
| 187 | } |
| 268 | 188 | |
| 269 | | return () => { |
| 270 | | cancelled = true; |
| 271 | | }; |
| 272 | | }, [cacheKey, cachedHighlightedHtml, code, isStreaming, language, themeName]); |
| 189 | const highlighter = use(getHighlighterPromise(language, themeName)); |
| 190 | const highlightedHtml = useMemo( |
| 191 | () => highlighter.codeToHtml(code, { lang: language, theme: themeName }), |
| 192 | [code, highlighter, language, themeName], |
| 193 | ); |
| 273 | 194 | |
| 274 | | if (!highlightedHtml) { |
| 275 | | return fallback; |
| 276 | | } |
| 195 | useEffect(() => { |
| 196 | if (!isStreaming) { |
| 197 | highlightedCodeCache.set( |
| 198 | cacheKey, |
| 199 | highlightedHtml, |
| 200 | estimateHighlightedSize(highlightedHtml, code), |
| 201 | ); |
| 202 | } |
| 203 | }, [cacheKey, code, highlightedHtml, isStreaming]); |
| 277 | 204 | |
| 278 | 205 | return ( |
| 279 | 206 | <div className="chat-markdown-shiki" dangerouslySetInnerHTML={{ __html: highlightedHtml }} /> |
| ... | ... | @@ -316,13 +243,14 @@ function ChatMarkdown({ text, cwd, isStreaming = false }: ChatMarkdownProps) { |
| 316 | 243 | |
| 317 | 244 | return ( |
| 318 | 245 | <MarkdownCodeBlock code={codeBlock.code}> |
| 319 | | <ShikiCodeBlock |
| 320 | | className={codeBlock.className} |
| 321 | | code={codeBlock.code} |
| 322 | | themeName={diffThemeName} |
| 323 | | isStreaming={isStreaming} |
| 324 | | fallback={<pre {...props}>{children}</pre>} |
| 325 | | /> |
| 246 | <Suspense fallback={<pre {...props}>{children}</pre>}> |
| 247 | <SuspenseShikiCodeBlock |
| 248 | className={codeBlock.className} |
| 249 | code={codeBlock.code} |
| 250 | themeName={diffThemeName} |
| 251 | isStreaming={isStreaming} |
| 252 | /> |
| 253 | </Suspense> |
| 326 | 254 | </MarkdownCodeBlock> |
| 327 | 255 | ); |
| 328 | 256 | }, |