From e89b05326911562c8183d223df311ed0bebd3c25 Mon Sep 17 00:00:00 2001
From: Julius Marminge <julius0216@outlook.com>
Date: Sat, 28 Feb 2026 11:32:19 -0800
Subject: [PATCH] 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
---
 apps/web/src/components/ChatMarkdown.tsx | 92 +++++++++++++++---------
 apps/web/src/lib/diffRendering.ts        |  6 +-
 2 files changed, 62 insertions(+), 36 deletions(-)

diff --git a/apps/web/src/components/ChatMarkdown.tsx b/apps/web/src/components/ChatMarkdown.tsx
index 9b84386624523307a62eb7345b6832b10c4a5c5e..55ee4483843e36b4ae7ad4a5272975791171a0d4 100644
--- a/apps/web/src/components/ChatMarkdown.tsx
+++ b/apps/web/src/components/ChatMarkdown.tsx
@@ -14,6 +14,7 @@ import remarkGfm from "remark-gfm";
 import { LRUCache } from "../lib/lruCache";
 import { readNativeApi } from "../nativeApi";
 import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffThemes";
+import { fnv1a32 } from "../lib/diffRendering";
 import { useTheme } from "../hooks/useTheme";
 import { resolveMarkdownFileLinkTarget } from "../markdown-links";
 import { preferredTerminalEditor } from "../terminal-links";
@@ -31,17 +32,11 @@ const highlightedCodeCache = new LRUCache<string>(
   MAX_HIGHLIGHT_CACHE_ENTRIES,
   MAX_HIGHLIGHT_CACHE_MEMORY_BYTES,
 );
-const inFlightHighlights = new Map<string, Promise<string>>();
-
-function hashString(input: string): string {
-  // FNV-1a 32-bit hash.
-  let hash = 0x811c9dc5;
-  for (let index = 0; index < input.length; index += 1) {
-    hash ^= input.charCodeAt(index);
-    hash = Math.imul(hash, 0x01000193);
-  }
-  return (hash >>> 0).toString(36);
+interface InFlightHighlight {
+  promise: Promise<string>;
+  enableCache: () => void;
 }
+const inFlightHighlights = new Map<string, InFlightHighlight>();
 
 function extractFenceLanguage(className: string | undefined): string {
   const match = className?.match(CODE_FENCE_LANGUAGE_REGEX);
@@ -100,7 +95,7 @@ async function highlightCodeBlock(code: string, language: string, themeName: Dif
 }
 
 function createHighlightCacheKey(code: string, language: string, themeName: DiffThemeName): string {
-  return `${hashString(code)}:${code.length}:${language}:${themeName}`;
+  return `${fnv1a32(code).toString(36)}:${code.length}:${language}:${themeName}`;
 }
 
 function estimateHighlightedSize(html: string, code: string): number {
@@ -129,12 +124,16 @@ function getCachedOrCreateHighlight({
 
   const existingInFlight = inFlightHighlights.get(cacheKey);
   if (existingInFlight) {
-    return existingInFlight;
+    if (useCache) {
+      existingInFlight.enableCache();
+    }
+    return existingInFlight.promise;
   }
 
-  const request = highlightCodeBlock(code, language, themeName)
+  let shouldCache = useCache;
+  const requestPromise = highlightCodeBlock(code, language, themeName)
     .then((html) => {
-      if (useCache) {
+      if (shouldCache) {
         highlightedCodeCache.set(cacheKey, html, estimateHighlightedSize(html, code));
       }
       return html;
@@ -143,8 +142,13 @@ function getCachedOrCreateHighlight({
       inFlightHighlights.delete(cacheKey);
     });
 
-  inFlightHighlights.set(cacheKey, request);
-  return request;
+  inFlightHighlights.set(cacheKey, {
+    promise: requestPromise,
+    enableCache: () => {
+      shouldCache = true;
+    },
+  });
+  return requestPromise;
 }
 
 interface ShikiCodeBlockProps {
@@ -156,24 +160,41 @@ interface ShikiCodeBlockProps {
 }
 
 function ShikiCodeBlock({ className, code, themeName, isStreaming, fallback }: ShikiCodeBlockProps) {
-  const [highlightedHtml, setHighlightedHtml] = useState<string | null>(null);
+  const language = extractFenceLanguage(className);
+  const cacheKey = createHighlightCacheKey(code, language, themeName);
+  const [highlightState, setHighlightState] = useState<{
+    cacheKey: string;
+    html: string | null;
+  }>(() => ({
+    cacheKey,
+    html: isStreaming ? null : (highlightedCodeCache.get(cacheKey) ?? null),
+  }));
+  const cachedHighlightedHtml = isStreaming ? null : (highlightedCodeCache.get(cacheKey) ?? null);
+  const highlightedHtml =
+    cachedHighlightedHtml ?? (highlightState.cacheKey === cacheKey ? highlightState.html : null);
 
   useEffect(() => {
     let cancelled = false;
-    const language = extractFenceLanguage(className);
-    const cacheKey = createHighlightCacheKey(code, language, themeName);
 
-    if (!isStreaming) {
-      const cached = highlightedCodeCache.get(cacheKey);
-      if (cached) {
-        setHighlightedHtml(cached);
-        return () => {
-          cancelled = true;
-        };
-      }
+    if (cachedHighlightedHtml) {
+      setHighlightState((current) => {
+        if (current.cacheKey === cacheKey && current.html === cachedHighlightedHtml) {
+          return current;
+        }
+        return { cacheKey, html: cachedHighlightedHtml };
+      });
+      return () => {
+        cancelled = true;
+      };
     }
 
-    setHighlightedHtml(null);
+    setHighlightState((current) => {
+      if (current.cacheKey === cacheKey) {
+        return current;
+      }
+      return { cacheKey, html: null };
+    });
+
     void getCachedOrCreateHighlight({
       cacheKey,
       code,
@@ -185,18 +206,19 @@ function ShikiCodeBlock({ className, code, themeName, isStreaming, fallback }: S
         if (cancelled) {
           return;
         }
-        setHighlightedHtml(html);
+        setHighlightState((current) => {
+          if (current.cacheKey === cacheKey && current.html === html) {
+            return current;
+          }
+          return { cacheKey, html };
+        });
       })
-      .catch(() => {
-        if (!cancelled) {
-          setHighlightedHtml(null);
-        }
-      });
+      .catch(() => undefined);
 
     return () => {
       cancelled = true;
     };
-  }, [className, code, themeName, isStreaming]);
+  }, [cacheKey, cachedHighlightedHtml, code, isStreaming, language, themeName]);
 
   if (!highlightedHtml) {
     return fallback;
diff --git a/apps/web/src/lib/diffRendering.ts b/apps/web/src/lib/diffRendering.ts
index a5631ef8583b0c578c7250bbeb3dcb7760cf4f59..c61ad69f535bdae2c1770362d2c49788606dc63a 100644
--- a/apps/web/src/lib/diffRendering.ts
+++ b/apps/web/src/lib/diffRendering.ts
@@ -3,7 +3,11 @@ const FNV_PRIME_32 = 0x01000193;
 const SECONDARY_HASH_SEED = 0x9e3779b9;
 const SECONDARY_HASH_MULTIPLIER = 0x85ebca6b;
 
-function fnv1a32(input: string, seed: number, multiplier: number): number {
+export function fnv1a32(
+  input: string,
+  seed = FNV_OFFSET_BASIS_32,
+  multiplier = FNV_PRIME_32,
+): number {
   let hash = seed >>> 0;
   for (let index = 0; index < input.length; index += 1) {
     hash ^= input.charCodeAt(index);
-- 
2.54.0