From fe932c8307ad083187ad30f56fb16765685b1150 Mon Sep 17 00:00:00 2001
From: Shantur <shantur@Shanturs-MacBook-Pro-M5.local>
Date: Tue, 31 Mar 2026 15:18:44 +0100
Subject: [PATCH] fix(ui): avoid caching incomplete code highlighting

Only cache markdown HTML after Shiki has the required fence languages loaded so virtualized assistant messages can re-render with syntax highlighting when remounted.
---
 packages/ui/src/components/markdown.tsx | 13 +++++--
 packages/ui/src/lib/markdown.ts         | 45 ++++++++++++++++++++-----
 2 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/packages/ui/src/components/markdown.tsx b/packages/ui/src/components/markdown.tsx
index 69e3f11a..56889ca3 100644
--- a/packages/ui/src/components/markdown.tsx
+++ b/packages/ui/src/components/markdown.tsx
@@ -123,7 +123,11 @@ export function Markdown(props: MarkdownProps) {
     version: () => resolved().version,
   })
 
-  const commitCacheEntry = (snapshot: ReturnType<typeof resolved>, renderedHtml: string) => {
+  const commitCacheEntry = (
+    snapshot: ReturnType<typeof resolved>,
+    renderedHtml: string,
+    options?: { cache?: boolean },
+  ) => {
     const cacheEntry: RenderCache = {
       text: snapshot.text,
       html: renderedHtml,
@@ -131,7 +135,9 @@ export function Markdown(props: MarkdownProps) {
       mode: `${snapshot.version}:${snapshot.escapeRawHtml ? "escaped" : "raw"}`,
     }
     setHtml(renderedHtml)
-    cacheHandle.set(cacheEntry)
+    if (options?.cache ?? true) {
+      cacheHandle.set(cacheEntry)
+    }
     notifyRendered()
   }
 
@@ -142,9 +148,10 @@ export function Markdown(props: MarkdownProps) {
       suppressHighlight: !snapshot.highlightEnabled,
       escapeRawHtml: snapshot.escapeRawHtml,
     })
+    const shouldCache = !snapshot.highlightEnabled || !markdown.hasPendingCodeHighlight(snapshot.text)
 
     if (latestRequestKey === snapshot.requestKey) {
-      commitCacheEntry(snapshot, rendered)
+      commitCacheEntry(snapshot, rendered, { cache: shouldCache })
     }
   }
 
diff --git a/packages/ui/src/lib/markdown.ts b/packages/ui/src/lib/markdown.ts
index f544405f..c83c6894 100644
--- a/packages/ui/src/lib/markdown.ts
+++ b/packages/ui/src/lib/markdown.ts
@@ -120,14 +120,7 @@ function resolveLanguage(token: string): { canonical: string | null; raw: string
   return { canonical: null, raw: normalized }
 }
 
-async function ensureLanguages(content: string) {
-  if (highlightSuppressed) {
-    return
-  }
-
-  // Extract code-fence language tokens via `marked` so we correctly handle code blocks
-  // that contain backticks (e.g. JS template literals). Regex-based fence scans tend
-  // to miss these and prevent languages from loading.
+function collectCodeFenceLanguages(content: string): string[] {
   const foundLanguages = new Set<string>()
   try {
     const tokens = marked.lexer(content) as any
@@ -139,10 +132,44 @@ async function ensureLanguages(content: string) {
       }
     })
   } catch {
-    // If tokenization fails for any reason, skip language preloading.
+    return []
+  }
+
+  return [...foundLanguages]
+}
+
+export function hasPendingCodeHighlight(content: string): boolean {
+  const languages = collectCodeFenceLanguages(content)
+  for (const token of languages) {
+    const rawToken = normalizeLanguageToken(token)
+    if (!rawToken || rawToken === "text") {
+      continue
+    }
+
+    const { canonical, raw } = resolveLanguage(token)
+    const langKey = canonical || raw
+    if (langKey === "text" || raw === "text") {
+      continue
+    }
+
+    if (!highlighter || !loadedLanguages.has(langKey)) {
+      return true
+    }
+  }
+
+  return false
+}
+
+async function ensureLanguages(content: string) {
+  if (highlightSuppressed) {
     return
   }
 
+  // Extract code-fence language tokens via `marked` so we correctly handle code blocks
+  // that contain backticks (e.g. JS template literals). Regex-based fence scans tend
+  // to miss these and prevent languages from loading.
+  const foundLanguages = collectCodeFenceLanguages(content)
+
   // Queue language loading tasks
   for (const token of foundLanguages) {
     const rawToken = normalizeLanguageToken(token)