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
14import { LRUCache } from "../lib/lruCache";14import { LRUCache } from "../lib/lruCache";
15import { readNativeApi } from "../nativeApi";15import { readNativeApi } from "../nativeApi";
16import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffThemes";16import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffThemes";
17import { fnv1a32 } from "../lib/diffRendering";
17import { useTheme } from "../hooks/useTheme";18import { useTheme } from "../hooks/useTheme";
18import { resolveMarkdownFileLinkTarget } from "../markdown-links";19import { resolveMarkdownFileLinkTarget } from "../markdown-links";
19import { preferredTerminalEditor } from "../terminal-links";20import { preferredTerminalEditor } from "../terminal-links";
@@ -31,17 +32,11 @@ const highlightedCodeCache = new LRUCache<string>( · 11 unmodified lines
31 MAX_HIGHLIGHT_CACHE_ENTRIES,32 MAX_HIGHLIGHT_CACHE_ENTRIES,
32 MAX_HIGHLIGHT_CACHE_MEMORY_BYTES,33 MAX_HIGHLIGHT_CACHE_MEMORY_BYTES,
33);34);
34const inFlightHighlights = new Map<string, Promise<string>>();35interface InFlightHighlight {
35 36 promise: Promise<string>;
36function hashString(input: string): string {37 enableCache: () => void;
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);
44}38}
39const inFlightHighlights = new Map<string, InFlightHighlight>();
45 40
46function extractFenceLanguage(className: string | undefined): string {41function extractFenceLanguage(className: string | undefined): string {
47 const match = className?.match(CODE_FENCE_LANGUAGE_REGEX);42 const match = className?.match(CODE_FENCE_LANGUAGE_REGEX);
@@ -100,7 +95,7 @@ async function highlightCodeBlock(code: string, language: string, themeName: Dif · 52 unmodified lines
100}95}
101 96
102function createHighlightCacheKey(code: string, language: string, themeName: DiffThemeName): string {97function 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}`;
104}99}
105 100
106function estimateHighlightedSize(html: string, code: string): number {101function estimateHighlightedSize(html: string, code: string): number {
@@ -129,12 +124,16 @@ function getCachedOrCreateHighlight({ · 22 unmodified lines
129 124
130 const existingInFlight = inFlightHighlights.get(cacheKey);125 const existingInFlight = inFlightHighlights.get(cacheKey);
131 if (existingInFlight) {126 if (existingInFlight) {
132 return existingInFlight;127 if (useCache) {
128 existingInFlight.enableCache();
129 }
130 return existingInFlight.promise;
133 }131 }
134 132
135 const request = highlightCodeBlock(code, language, themeName)133 let shouldCache = useCache;
134 const requestPromise = highlightCodeBlock(code, language, themeName)
136 .then((html) => {135 .then((html) => {
137 if (useCache) {136 if (shouldCache) {
138 highlightedCodeCache.set(cacheKey, html, estimateHighlightedSize(html, code));137 highlightedCodeCache.set(cacheKey, html, estimateHighlightedSize(html, code));
139 }138 }
140 return html;139 return html;
@@ -143,8 +142,13 @@ function getCachedOrCreateHighlight({ · 2 unmodified lines
143 inFlightHighlights.delete(cacheKey);142 inFlightHighlights.delete(cacheKey);
144 });143 });
145 144
146 inFlightHighlights.set(cacheKey, request);145 inFlightHighlights.set(cacheKey, {
147 return request;146 promise: requestPromise,
147 enableCache: () => {
148 shouldCache = true;
149 },
150 });
151 return requestPromise;
148}152}
149 153
150interface ShikiCodeBlockProps {154interface ShikiCodeBlockProps {
@@ -156,24 +160,41 @@ interface ShikiCodeBlockProps { · 5 unmodified lines
156}160}
157 161
158function ShikiCodeBlock({ className, code, themeName, isStreaming, fallback }: ShikiCodeBlockProps) {162function 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);
160 175
161 useEffect(() => {176 useEffect(() => {
162 let cancelled = false;177 let cancelled = false;
163 const language = extractFenceLanguage(className);
164 const cacheKey = createHighlightCacheKey(code, language, themeName);
165 178
166 if (!isStreaming) {179 if (cachedHighlightedHtml) {
167 const cached = highlightedCodeCache.get(cacheKey);180 setHighlightState((current) => {
168 if (cached) {181 if (current.cacheKey === cacheKey && current.html === cachedHighlightedHtml) {
169 setHighlightedHtml(cached);182 return current;
170 return () => {183 }
171 cancelled = true;184 return { cacheKey, html: cachedHighlightedHtml };
172 };185 });
173 }186 return () => {
187 cancelled = true;
188 };
174 }189 }
175 190
176 setHighlightedHtml(null);191 setHighlightState((current) => {
192 if (current.cacheKey === cacheKey) {
193 return current;
194 }
195 return { cacheKey, html: null };
196 });
197
177 void getCachedOrCreateHighlight({198 void getCachedOrCreateHighlight({
178 cacheKey,199 cacheKey,
179 code,200 code,
@@ -185,18 +206,19 @@ function ShikiCodeBlock({ className, code, themeName, isStreaming, fallback }: S · 5 unmodified lines
185 if (cancelled) {206 if (cancelled) {
186 return;207 return;
187 }208 }
188 setHighlightedHtml(html);209 setHighlightState((current) => {
210 if (current.cacheKey === cacheKey && current.html === html) {
211 return current;
212 }
213 return { cacheKey, html };
214 });
189 })215 })
190 .catch(() => {216 .catch(() => undefined);
191 if (!cancelled) {
192 setHighlightedHtml(null);
193 }
194 });
195 217
196 return () => {218 return () => {
197 cancelled = true;219 cancelled = true;
198 };220 };
199 }, [className, code, themeName, isStreaming]);221 }, [cacheKey, cachedHighlightedHtml, code, isStreaming, language, themeName]);
200 222
201 if (!highlightedHtml) {223 if (!highlightedHtml) {
202 return fallback;224 return fallback;
apps/web/src/lib/diffRendering.ts+5−1
@@ -3,7 +3,11 @@ const FNV_PRIME_32 = 0x01000193; · 2 unmodified lines
3const SECONDARY_HASH_SEED = 0x9e3779b9;3const SECONDARY_HASH_SEED = 0x9e3779b9;
4const SECONDARY_HASH_MULTIPLIER = 0x85ebca6b;4const SECONDARY_HASH_MULTIPLIER = 0x85ebca6b;
5 5
6function fnv1a32(input: string, seed: number, multiplier: number): number {6export function fnv1a32(
7 input: string,
8 seed = FNV_OFFSET_BASIS_32,
9 multiplier = FNV_PRIME_32,
10): number {
7 let hash = seed >>> 0;11 let hash = seed >>> 0;
8 for (let index = 0; index < input.length; index += 1) {12 for (let index = 0; index < input.length; index += 1) {
9 hash ^= input.charCodeAt(index);13 hash ^= input.charCodeAt(index);