perf(ui): defer markdown syntax highlighting to prevent main thread blocking (#7102)

Switching sessions with 40+ messages caused a 2.3s UI freeze due to
synchronous Shiki/Oniguruma WASM regex tokenization running on the
main thread for every code block.

Implements two-pass rendering: first pass renders code blocks as plain
<pre><code data-lang> tags (instant), second pass highlights progressively
via setTimeout(0) to keep the UI responsive.

Key implementation details:
- Morphdom guard preserves Shiki-highlighted <pre> blocks during streaming
- replaceWithHighlighted() is wrapper-aware - preserves copy button structure
- LRU highlight cache (500 entries) with null-byte key delimiter
- Generation counter + abort signal prevent stale callbacks and race conditions
- data-source-hash stores FNV-1a hash for O(1) content-change detection
- Removed marked-shiki dependency

Original work by @daem0ndev in #6474.
Supersedes #6474 to run CI from an internal branch.

Fixes #6221

bun.lock

@@ -501,7 +501,6 @@
         "luxon": "catalog:",
         "marked": "catalog:",
         "marked-katex-extension": "5.1.6",
-        "marked-shiki": "catalog:",
         "morphdom": "2.7.8",
         "motion": "12.34.5",
         "motion-dom": "12.34.3",

packages/ui/package.json

@@ -57,7 +57,6 @@
     "luxon": "catalog:",
     "marked": "catalog:",
     "marked-katex-extension": "5.1.6",
-    "marked-shiki": "catalog:",
     "morphdom": "2.7.8",
     "motion": "12.34.5",
     "motion-dom": "12.34.3",

packages/ui/src/components/markdown.tsx

@@ -1,4 +1,4 @@
-import { useMarked } from "../context/marked"
+import { useMarked, deferredHighlight, fnv1a } from "../context/marked"
 import { useI18n } from "../context/i18n"
 import DOMPurify from "dompurify"
 import morphdom from "morphdom"
@@ -258,8 +258,13 @@ export function Markdown(
     { initialValue: "" },
   )
 
-  let copySetupTimer: ReturnType<typeof setTimeout> | undefined
   let copyCleanup: (() => void) | undefined
+  // kilocode_change start: generation counter prevents stale deferredHighlight
+  // callbacks from overwriting copyCleanup set by a newer render (issue #6221).
+  // The abort signal cancels the previous in-flight highlight pass so rapid
+  // streaming tokens don't spawn concurrent passes racing on the same DOM nodes.
+  const highlightState = { gen: 0, signal: { aborted: false } }
+  // kilocode_change end
 
   createEffect(() => {
     const container = root()
@@ -274,31 +279,72 @@ export function Markdown(
 
     const temp = document.createElement("div")
     temp.innerHTML = content
-    decorate(temp, {
+    const labels = {
       copy: i18n.t("ui.message.copy"),
       copied: i18n.t("ui.message.copied"),
-    })
+    }
+    decorate(temp, labels)
 
+    // kilocode_change start: morphdom guard for highlighted blocks (issue #6221)
+    // During streaming, morphdom re-runs on every token. Without this guard,
+    // it would revert already-highlighted <pre> blocks back to plain code.
     morphdom(container, temp, {
       childrenOnly: true,
       onBeforeElUpdated: (fromEl, toEl) => {
         if (fromEl.isEqualNode(toEl)) return false
+        // Preserve Shiki-highlighted blocks — don't let morphdom revert them
+        // to plain <pre><code> during streaming re-renders.
+        // Note: "shiki" class is on <pre> (set by Shiki's codeToHtml output).
+        // We compare data-source-hash (a lightweight FNV-1a hash stored by
+        // deferredHighlight on the highlighted <pre>) against a hash of the
+        // incoming code text to detect mid-stream content changes: if the code
+        // changed, we let morphdom update so the block can be re-queued for
+        // highlighting with the new content.
+        if (
+          fromEl instanceof HTMLElement &&
+          fromEl.tagName === "PRE" &&
+          fromEl.classList.contains("shiki") &&
+          toEl instanceof HTMLElement &&
+          toEl.tagName === "PRE" &&
+          !toEl.classList.contains("shiki")
+        ) {
+          const fromHash = fromEl.getAttribute("data-source-hash")
+          const toCode = toEl.querySelector("code")?.textContent ?? ""
+          if (fromHash === fnv1a(toCode)) return false
+          // Source changed during streaming — fall through so morphdom replaces
+          // the stale highlighted block with the updated plain block, which will
+          // be re-highlighted on the next deferredHighlight pass.
+        }
         return true
       },
     })
+    // kilocode_change end
 
-    if (copySetupTimer) clearTimeout(copySetupTimer)
-    copySetupTimer = setTimeout(() => {
+    // kilocode_change start: deferred syntax highlighting (issue #6221)
+    // DOM is now painted with plain <pre><code> blocks.
+    // Progressively highlight via setTimeout(0) to avoid blocking.
+    // onComplete re-runs setupCodeCopy since highlighting replaces DOM nodes.
+    // The generation counter ensures a stale in-flight highlight run (from a
+    // previous streaming token) doesn't overwrite the cleanup set by a newer run.
+    // Cancel any in-flight highlight pass before starting a new one — prevents
+    // concurrent passes from racing to replace the same DOM nodes.
+    highlightState.signal.aborted = true
+    const gen = ++highlightState.gen
+    const signal = { aborted: false }
+    highlightState.signal = signal
+    deferredHighlight(container, () => {
+      if (gen !== highlightState.gen) return
       if (copyCleanup) copyCleanup()
-      copyCleanup = setupCodeCopy(container, {
-        copy: i18n.t("ui.message.copy"),
-        copied: i18n.t("ui.message.copied"),
-      })
-    }, 150)
+      copyCleanup = setupCodeCopy(container, labels)
+    }, signal)
+    // kilocode_change end
   })
 
   onCleanup(() => {
-    if (copySetupTimer) clearTimeout(copySetupTimer)
+    // Invalidate any in-flight deferredHighlight callbacks so they don't call
+    // setupCodeCopy on a disconnected container after the component unmounts.
+    highlightState.signal.aborted = true
+    highlightState.gen++
     if (copyCleanup) copyCleanup()
   })
 

packages/ui/src/context/marked.tsx

@@ -1,6 +1,6 @@
 import { marked } from "marked"
 import markedKatex from "marked-katex-extension"
-import markedShiki from "marked-shiki"
+
 import katex from "katex"
 import { bundledLanguages, type BundledLanguage } from "shiki"
 import { parseFilePath } from "../file-path" // kilocode_change
@@ -424,6 +424,8 @@ function renderMathExpressions(html: string): string {
     .join("")
 }
 
+// Used only by the native parser path (props.nativeParser) — not the JS parser.
+// The JS parser uses deferredHighlight() instead for non-blocking rendering.
 async function highlightCodeBlocks(html: string): Promise<string> {
   const codeBlockRegex = /<pre><code(?:\s+class="language-([^"]*)")?>([\s\S]*?)<\/code><\/pre>/g
   const matches = [...html.matchAll(codeBlockRegex)]
@@ -464,10 +466,162 @@ export type NativeMarkdownParser = (markdown: string) => Promise<string>
 
 // kilocode_change: parseFilePath imported from ../file-path
 
+// kilocode_change start: highlight cache for deferred highlighting
+
+/** FNV-1a hash — lightweight alternative to storing full source code in DOM attributes. */
+export function fnv1a(s: string): string {
+  let h = 0x811c9dc5
+  for (let i = 0; i < s.length; i++) {
+    h ^= s.charCodeAt(i)
+    h = Math.imul(h, 0x01000193)
+  }
+  return (h >>> 0).toString(36)
+}
+
+const cache = new Map<string, string>()
+const CACHE_LIMIT = 500
+
+// Normalize common language aliases before sanitizing, so e.g. "c++" → "cpp"
+// rather than stripping to "c" (which would highlight as the wrong language).
+const LANG_ALIASES: Record<string, string> = {
+  "c++": "cpp",
+  "c#": "csharp",
+  "f#": "fsharp",
+  "objective-c++": "objective-cpp",
+}
+
+function touchHighlightCache(key: string, value: string) {
+  cache.delete(key)
+  cache.set(key, value)
+  if (cache.size <= CACHE_LIMIT) return
+  const first = cache.keys().next().value
+  if (!first) return
+  cache.delete(first)
+}
+
+function replaceWithHighlighted(block: Element, html: string, sourceHash: string) {
+  const pre = block.parentElement
+  if (!pre || !pre.isConnected) return
+  const temp = document.createElement("div")
+  temp.innerHTML = html
+  const highlighted = temp.firstElementChild
+  if (!highlighted) return
+  // Store a hash of the source code so the morphdom guard in Markdown can detect
+  // mid-stream content changes without keeping the full source in the DOM.
+  highlighted.setAttribute("data-source-hash", sourceHash)
+  // Preserve any wrapper structure (e.g., markdown-code wrapper with copy button)
+  const wrapper = pre.parentElement
+  if (wrapper?.getAttribute("data-component") === "markdown-code") {
+    wrapper.replaceChild(highlighted, pre)
+    return
+  }
+  pre.replaceWith(highlighted)
+}
+
+/**
+ * Progressively highlight unhighlighted <pre><code> blocks inside a container.
+ * Each block is highlighted via setTimeout(0) to yield back to the main thread
+ * between blocks, keeping the UI responsive.
+ *
+ * Blocks marked with data-highlighted are skipped (already processed).
+ * After highlighting, blocks are marked to survive morphdom re-runs during streaming.
+ *
+ * Returns a callback to re-run setupCodeCopy after highlighting completes,
+ * since highlight replaces DOM nodes that may have copy button wrappers.
+ */
+export async function deferredHighlight(
+  container: HTMLElement,
+  onComplete?: () => void,
+  signal?: { aborted: boolean },
+): Promise<void> {
+  const blocks = Array.from(
+    container.querySelectorAll("pre > code[data-lang]:not([data-highlighted])"),
+  )
+  if (blocks.length === 0) {
+    onComplete?.()
+    return
+  }
+
+  const highlighter = await getSharedHighlighter({ themes: ["Kilo"], langs: [] })
+
+  for (const block of blocks) {
+    // Short-circuit if the container is unmounted or the caller cancelled this run
+    // (e.g., a newer streaming token triggered a fresh deferredHighlight call).
+    if (!container.isConnected || signal?.aborted) break
+
+    const lang = block.getAttribute("data-lang") || "text"
+    const code = block.textContent ?? ""
+    if (!code) continue
+
+    const cacheKey = `${lang}\0${code}`
+    const codeHash = fnv1a(code)
+    const cached = cache.get(cacheKey)
+    if (cached) {
+      touchHighlightCache(cacheKey, cached) // refresh LRU position on hit
+      replaceWithHighlighted(block, cached, codeHash)
+      continue
+    }
+
+    await new Promise<void>((resolve) => {
+      setTimeout(() => {
+        // Re-check inside the timer callback — signal may have been aborted while
+        // this task was queued in the event loop waiting to run.
+        if (!block.isConnected || signal?.aborted) {
+          resolve()
+          return
+        }
+        try {
+          const language = lang in bundledLanguages ? lang : "text"
+          const highlight = () => {
+            // Re-check after async loadLanguage — block may have been replaced
+            // or signal aborted while the language bundle was being fetched.
+            if (!block.isConnected || signal?.aborted) {
+              resolve()
+              return
+            }
+            const html = highlighter.codeToHtml(code, { lang: language, theme: "Kilo", tabindex: false })
+            touchHighlightCache(cacheKey, html)
+            // Note: data-highlighted is NOT set on `block` here because
+            // replaceWithHighlighted replaces the parent <pre> entirely — the
+            // original <code> element leaves the DOM immediately after this call.
+            // The new highlighted <pre class="shiki"> from Shiki has no
+            // code[data-lang] child, so the querySelectorAll selector won't
+            // pick it up on subsequent deferredHighlight passes.
+            replaceWithHighlighted(block, html, codeHash)
+            resolve()
+          }
+          if (!highlighter.getLoadedLanguages().includes(language)) {
+            highlighter.loadLanguage(language as BundledLanguage).then(highlight).catch((err) => {
+              console.warn("Failed to load language for highlighting", language, err)
+              resolve()
+            })
+            return
+          }
+          highlight()
+        } catch (err) {
+          console.warn("Deferred highlight failed", lang, err)
+          resolve()
+        }
+      }, 0)
+    })
+  }
+
+  // Only fire onComplete if the container is still mounted and the run wasn't
+  // cancelled — avoids calling setupCodeCopy on a disconnected DOM tree.
+  if (container.isConnected && !signal?.aborted) {
+    onComplete?.()
+  }
+}
+// kilocode_change end
+
 export const { use: useMarked, provider: MarkedProvider } = createSimpleContext({
   name: "Marked",
   init: (props: { nativeParser?: NativeMarkdownParser }) => {
-    const jsParser = marked.use(
+    // kilocode_change start: two-pass parser — first pass skips Shiki highlighting
+    // to avoid blocking the main thread with Oniguruma WASM regex (issue #6221).
+    // Code blocks render as plain <pre><code data-lang="..."> immediately.
+    // The Markdown component calls deferredHighlight() after DOM paint.
+    const parser = marked.use(
       {
         renderer: {
           link({ href, title, text }) {
@@ -484,6 +638,20 @@ export const { use: useMarked, provider: MarkedProvider } = createSimpleContext(
             }
             return `<code>${text}</code>`
           },
+          code({ text, lang }) {
+            const escaped = text
+              .replace(/&/g, "&amp;")
+              .replace(/</g, "&lt;")
+              .replace(/>/g, "&gt;")
+              .replace(/"/g, "&quot;")
+              .replace(/'/g, "&#39;")
+            // Normalize aliases (e.g. "c++" → "cpp") before stripping special
+            // chars, so "c++" doesn't become "c" (wrong language highlight).
+            const normalized = lang ? (LANG_ALIASES[lang] ?? lang) : ""
+            const safe = normalized ? normalized.replace(/[^a-zA-Z0-9_-]/g, "") : ""
+            const attr = safe ? ` class="language-${safe}" data-lang="${safe}"` : ' data-lang="text"'
+            return `<pre><code${attr}>${escaped}</code></pre>`
+          },
           // kilocode_change end
         },
       },
@@ -491,23 +659,8 @@ export const { use: useMarked, provider: MarkedProvider } = createSimpleContext(
         throwOnError: false,
         nonStandard: true,
       }),
-      markedShiki({
-        async highlight(code, lang) {
-          const highlighter = await getSharedHighlighter({ themes: ["Kilo"], langs: [] })
-          if (!(lang in bundledLanguages)) {
-            lang = "text"
-          }
-          if (!highlighter.getLoadedLanguages().includes(lang)) {
-            await highlighter.loadLanguage(lang as BundledLanguage)
-          }
-          return highlighter.codeToHtml(code, {
-            lang: lang || "text",
-            theme: "Kilo",
-            tabindex: false,
-          })
-        },
-      }),
     )
+    // kilocode_change end
 
     if (props.nativeParser) {
       const nativeParser = props.nativeParser
@@ -520,6 +673,6 @@ export const { use: useMarked, provider: MarkedProvider } = createSimpleContext(
       }
     }
 
-    return jsParser
+    return parser
   },
 })