History

Fix code highlight cache behavior for streaming markdown

- Reuse in-flight highlight requests and allow later non-streaming renders to cache results - Stabilize code block highlight state updates to avoid stale or redundant rerenders - Export and reuse `fnv1a32` from `diffRendering` for shared cache key hashing

Julius Marminge committed 8 months ago

Commit 39eea0d1 parent 37f09512 files · +62 −36Browse files

Unified · Split

2 files
  1. apps/web/src/components/ChatMarkdown.tsx +57 −35
  2. apps/web/src/lib/diffRendering.ts +5 −1
apps/web/src/components/ChatMarkdown.tsx+57−35
@@ -14,6 +14,7 @@ import remarkGfm from "remark-gfm"; · 13 unmodified lines
1414 import { LRUCache } from "../lib/lruCache";
1515 import { readNativeApi } from "../nativeApi";
1616 import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffThemes";
17+import { fnv1a32 } from "../lib/diffRendering";
1718 import { useTheme } from "../hooks/useTheme";
1819 import { resolveMarkdownFileLinkTarget } from "../markdown-links";
1920 import { preferredTerminalEditor } from "../terminal-links";
@@ -31,17 +32,11 @@ const highlightedCodeCache = new LRUCache<string>( · 11 unmodified lines
3132 MAX_HIGHLIGHT_CACHE_ENTRIES,
3233 MAX_HIGHLIGHT_CACHE_MEMORY_BYTES,
3334 );
34−const inFlightHighlights = new Map<string, Promise<string>>();
35−
36−function hashString(input: string): string {
37− // FNV-1a 32-bit hash.
38− let hash = 0x811c9dc5;
39− for (let index = 0; index < input.length; index += 1) {
40− hash ^= input.charCodeAt(index);
41− hash = Math.imul(hash, 0x01000193);
42− }
43− return (hash >>> 0).toString(36);
35+interface InFlightHighlight {
36+ promise: Promise<string>;
37+ enableCache: () => void;
4438 }
39+const inFlightHighlights = new Map<string, InFlightHighlight>();
4540
4641 function extractFenceLanguage(className: string | undefined): string {
4742 const match = className?.match(CODE_FENCE_LANGUAGE_REGEX);
@@ -100,7 +95,7 @@ async function highlightCodeBlock(code: string, language: string, themeName: Dif · 52 unmodified lines
10095 }
10196
10297 function createHighlightCacheKey(code: string, language: string, themeName: DiffThemeName): string {
103− return `${hashString(code)}:${code.length}:${language}:${themeName}`;
98+ return `${fnv1a32(code).toString(36)}:${code.length}:${language}:${themeName}`;
10499 }
105100
106101 function estimateHighlightedSize(html: string, code: string): number {
@@ -129,12 +124,16 @@ function getCachedOrCreateHighlight({ · 22 unmodified lines
129124
130125 const existingInFlight = inFlightHighlights.get(cacheKey);
131126 if (existingInFlight) {
132− return existingInFlight;
127+ if (useCache) {
128+ existingInFlight.enableCache();
129+ }
130+ return existingInFlight.promise;
133131 }
134132
135− const request = highlightCodeBlock(code, language, themeName)
133+ let shouldCache = useCache;
134+ const requestPromise = highlightCodeBlock(code, language, themeName)
136135 .then((html) => {
137− if (useCache) {
136+ if (shouldCache) {
138137 highlightedCodeCache.set(cacheKey, html, estimateHighlightedSize(html, code));
139138 }
140139 return html;
@@ -143,8 +142,13 @@ function getCachedOrCreateHighlight({ · 2 unmodified lines
143142 inFlightHighlights.delete(cacheKey);
144143 });
145144
146− inFlightHighlights.set(cacheKey, request);
147− return request;
145+ inFlightHighlights.set(cacheKey, {
146+ promise: requestPromise,
147+ enableCache: () => {
148+ shouldCache = true;
149+ },
150+ });
151+ return requestPromise;
148152 }
149153
150154 interface ShikiCodeBlockProps {
@@ -156,24 +160,41 @@ interface ShikiCodeBlockProps { · 5 unmodified lines
156160 }
157161
158162 function ShikiCodeBlock({ className, code, themeName, isStreaming, fallback }: ShikiCodeBlockProps) {
159− const [highlightedHtml, setHighlightedHtml] = useState<string | null>(null);
163+ const language = extractFenceLanguage(className);
164+ const cacheKey = createHighlightCacheKey(code, language, themeName);
165+ const [highlightState, setHighlightState] = useState<{
166+ cacheKey: string;
167+ html: string | null;
168+ }>(() => ({
169+ cacheKey,
170+ html: isStreaming ? null : (highlightedCodeCache.get(cacheKey) ?? null),
171+ }));
172+ const cachedHighlightedHtml = isStreaming ? null : (highlightedCodeCache.get(cacheKey) ?? null);
173+ const highlightedHtml =
174+ cachedHighlightedHtml ?? (highlightState.cacheKey === cacheKey ? highlightState.html : null);
160175
161176 useEffect(() => {
162177 let cancelled = false;
163− const language = extractFenceLanguage(className);
164− const cacheKey = createHighlightCacheKey(code, language, themeName);
165178
166− if (!isStreaming) {
167− const cached = highlightedCodeCache.get(cacheKey);
168− if (cached) {
169− setHighlightedHtml(cached);
170− return () => {
171− cancelled = true;
172− };
173− }
179+ if (cachedHighlightedHtml) {
180+ setHighlightState((current) => {
181+ if (current.cacheKey === cacheKey && current.html === cachedHighlightedHtml) {
182+ return current;
183+ }
184+ return { cacheKey, html: cachedHighlightedHtml };
185+ });
186+ return () => {
187+ cancelled = true;
188+ };
174189 }
175190
176− setHighlightedHtml(null);
191+ setHighlightState((current) => {
192+ if (current.cacheKey === cacheKey) {
193+ return current;
194+ }
195+ return { cacheKey, html: null };
196+ });
197+
177198 void getCachedOrCreateHighlight({
178199 cacheKey,
179200 code,
@@ -185,18 +206,19 @@ function ShikiCodeBlock({ className, code, themeName, isStreaming, fallback }: S · 5 unmodified lines
185206 if (cancelled) {
186207 return;
187208 }
188− setHighlightedHtml(html);
209+ setHighlightState((current) => {
210+ if (current.cacheKey === cacheKey && current.html === html) {
211+ return current;
212+ }
213+ return { cacheKey, html };
214+ });
189215 })
190− .catch(() => {
191− if (!cancelled) {
192− setHighlightedHtml(null);
193− }
194− });
216+ .catch(() => undefined);
195217
196218 return () => {
197219 cancelled = true;
198220 };
199− }, [className, code, themeName, isStreaming]);
221+ }, [cacheKey, cachedHighlightedHtml, code, isStreaming, language, themeName]);
200222
201223 if (!highlightedHtml) {
202224 return fallback;
apps/web/src/lib/diffRendering.ts+5−1
@@ -3,7 +3,11 @@ const FNV_PRIME_32 = 0x01000193; · 2 unmodified lines
33 const SECONDARY_HASH_SEED = 0x9e3779b9;
44 const SECONDARY_HASH_MULTIPLIER = 0x85ebca6b;
55
6−function fnv1a32(input: string, seed: number, multiplier: number): number {
6+export function fnv1a32(
7+ input: string,
8+ seed = FNV_OFFSET_BASIS_32,
9+ multiplier = FNV_PRIME_32,
10+): number {
711 let hash = seed >>> 0;
812 for (let index = 0; index < input.length; index += 1) {
913 hash ^= input.charCodeAt(index);