fix(memory): share bounded highlighting with cancellable Markdown retries

package.json

@@ -37,7 +37,6 @@
     "@plantuml/core": "1.2026.8",
     "@radix-ui/react-use-controllable-state": "^1.2.2",
     "@streamdown/cjk": "^1.0.2",
-    "@streamdown/code": "^1.0.2",
     "@streamdown/math": "^1.0.2",
     "@streamdown/mermaid": "^1.0.3",
     "@tauri-apps/api": "^2.12.4",
@@ -94,7 +93,7 @@
     "shiki": "^3.22.0",
     "smol-toml": "^1.9.1",
     "sonner": "^2.0.7",
-    "streamdown": "^2.7.0",
+    "streamdown": "^2.7.1",
     "tailwind-merge": "^3.4.0",
     "tailwindcss": "^4.1.18",
     "use-stick-to-bottom": "^1.1.3",

patches/README.md

@@ -71,3 +71,27 @@ scrollbar ScrollTimeline animations and continued moving the handles. A separate
 Playwright WebKit passed all eight horizontal/vertical LTR/RTL drag and track
 click cases. These are focused reproducer results, not a claim about all of
 Codeg's long-running memory usage.
+
+## Streamdown 2.7.1 highlighting lifecycle
+
+The `HighlightedCodeBlockBody` effect now supplies an optional `AbortSignal`
+to its highlighter plugin and aborts it on a code/language/theme change or
+unmount. Plugins that ignore the extra option continue to work unchanged.
+Both the executable ESM chunk and the exported TypeScript interface are patched.
+Streamdown has no CommonJS entrypoint.
+
+Codeg uses this lifecycle to remove a superseded Markdown request from the shared
+highlighter's callbacks and retry waiters. Work is capped at 16 distinct requests;
+there is at most one waiter per live code view, removed at the next streaming
+update or unmount. Waiters do not retain token trees or a second cache key. This
+avoids accumulating all historical prefixes while still eventually highlighting
+every mounted block after a cold-engine burst. Callback-only consumers without
+an abort signal have a separate 128-waiter safety limit and fall back to plaintext
+beyond it. Codeg's Markdown and file/tool views supply a lifecycle.
+
+`src/lib/code-highlighter-streamdown.test.tsx` renders the installed Streamdown
+package with the real Shiki service, including 24 concurrent blocks. It also
+checks that actual component replacement and unmount abort every plugin request.
+Service tests exercise saturation, cancelled streaming prefixes, unmounts,
+coalescing, and the cache budgets. Remove or revisit this patch when upgrading
+Streamdown or after upstream provides an equivalent cancellation lifecycle.

patches/streamdown@2.7.1.patch

@@ -0,0 +1,23 @@
+diff --git a/dist/highlighted-body-JLMTFACH.js b/dist/highlighted-body-JLMTFACH.js
+index 4aa2fedc6052d65b7cb966ce40786674ec9c9633..edec7aad4b6141c89627b96eedc6d2f2a9aa98fb 100644
+--- a/dist/highlighted-body-JLMTFACH.js
++++ b/dist/highlighted-body-JLMTFACH.js
+@@ -1,2 +1,2 @@
+ "use client";
+-import {B,f,c}from'./chunk-7TTQLYRW.js';import {useContext,useState,useEffect}from'react';import {jsx}from'react/jsx-runtime';var T=({code:s,isIncomplete:l=false,language:e,maxHeight:d,raw:t,className:a,startLine:c$1,lineNumbers:p,...f$1})=>{let{shikiTheme:r}=useContext(B),o=f(),[H,i]=useState(t);return useEffect(()=>{if(!o){i(t);return}let n=true,u=o.highlight({code:s,isIncomplete:l,language:e,themes:r},B=>{n&&i(B);});return u&&i(u),()=>{n=false;}},[s,l,e,r,o,t]),jsx(c,{className:a,language:e,lineNumbers:p,maxHeight:d,result:H,startLine:c$1,...f$1})};export{T as HighlightedCodeBlockBody};
+\ No newline at end of file
++import {B,f,c}from'./chunk-7TTQLYRW.js';import {useContext,useState,useEffect}from'react';import {jsx}from'react/jsx-runtime';var T=({code:s,isIncomplete:l=false,language:e,maxHeight:d,raw:t,className:a,startLine:c$1,lineNumbers:p,...f$1})=>{let{shikiTheme:r}=useContext(B),o=f(),[H,i]=useState(t);return useEffect(()=>{if(!o){i(t);return}let n=true,abort=new AbortController(),u=o.highlight({code:s,isIncomplete:l,language:e,themes:r,signal:abort.signal},B=>{n&&i(B);});return u&&i(u),()=>{n=false;abort.abort();}},[s,l,e,r,o,t]),jsx(c,{className:a,language:e,lineNumbers:p,maxHeight:d,result:H,startLine:c$1,...f$1})};export{T as HighlightedCodeBlockBody};
+\ No newline at end of file
+diff --git a/dist/index.d.ts b/dist/index.d.ts
+index d201fee607ae5238f102ed23463967ae6c18ac2c..89731cacb33f834a39947ddbdd39dce5ded7dd7e 100644
+--- a/dist/index.d.ts
++++ b/dist/index.d.ts
+@@ -182,6 +182,8 @@ interface HighlightResult {
+  * Options for highlighting code
+  */
+ interface HighlightOptions {
++    /** Aborted when the code view changes or unmounts. */
++    signal?: AbortSignal;
+     code: string;
+     /**
+      * Whether the code block is still streaming. Its result is superseded by

pnpm-workspace.yaml

@@ -37,3 +37,4 @@ minimumReleaseAgeExclude:
 patchedDependencies:
   '@tauri-apps/api@2.12.4': patches/@tauri-apps__api@2.12.4.patch
   overlayscrollbars@2.15.1: patches/overlayscrollbars@2.15.1.patch
+  streamdown@2.7.1: patches/streamdown@2.7.1.patch

src/components/ai-elements/code-block.tsx

@@ -1,12 +1,7 @@
 "use client"
 
 import type { ComponentProps, CSSProperties, HTMLAttributes } from "react"
-import type {
-  BundledLanguage,
-  BundledTheme,
-  HighlighterGeneric,
-  ThemedToken,
-} from "shiki"
+import type { BundledLanguage, ThemedToken, TokensResult } from "shiki"
 
 import { Button } from "@/components/ui/button"
 import { ScrollArea } from "@/components/ui/scroll-area"
@@ -18,6 +13,13 @@ import {
   SelectValue,
 } from "@/components/ui/select"
 import { cn, copyTextToClipboard } from "@/lib/utils"
+import {
+  DEFAULT_CODE_THEMES,
+  getCachedHighlightedCode,
+  highlightCode as highlightSharedCode,
+  rawCodeTokens,
+  subscribeHighlightedCode,
+} from "@/lib/code-highlighter"
 import { CheckIcon, CopyIcon } from "lucide-react"
 import {
   createContext,
@@ -118,127 +120,22 @@ const CodeBlockContext = createContext<CodeBlockContextType>({
   code: "",
 })
 
-// Highlighter cache (singleton per language)
-const highlighterCache = new Map<
-  string,
-  Promise<HighlighterGeneric<BundledLanguage, BundledTheme>>
->()
-
-// Token cache
-const tokensCache = new Map<string, TokenizedCode>()
-
-// Subscribers for async token updates
-const subscribers = new Map<string, Set<(result: TokenizedCode) => void>>()
-
-const getTokensCacheKey = (code: string, language: BundledLanguage) => {
-  const start = code.slice(0, 100)
-  const end = code.length > 100 ? code.slice(-100) : ""
-  return `${language}:${code.length}:${start}:${end}`
-}
-
-const getHighlighter = (
-  language: BundledLanguage
-): Promise<HighlighterGeneric<BundledLanguage, BundledTheme>> => {
-  const cached = highlighterCache.get(language)
-  if (cached) {
-    return cached
-  }
-
-  // Load Shiki's engine lazily: a static `import` would pull it (and its
-  // textmate machinery) into the first-paint bundle for every message list,
-  // even a text-only conversation. Highlighting is already async with an
-  // immediate raw-token fallback (see `createRawTokens`), so deferring the
-  // engine import only extends that existing pre-highlight window slightly.
-  const highlighterPromise = import("shiki").then(({ createHighlighter }) =>
-    createHighlighter({
-      langs: [language],
-      themes: ["github-light", "github-dark"],
-    })
-  )
-
-  highlighterCache.set(language, highlighterPromise)
-  return highlighterPromise
-}
-
-// Create raw tokens for immediate display while highlighting loads
-const createRawTokens = (code: string): TokenizedCode => ({
-  bg: "transparent",
-  fg: "inherit",
-  tokens: code.split("\n").map((line) =>
-    line === ""
-      ? []
-      : [
-          {
-            color: "inherit",
-            content: line,
-          } as ThemedToken,
-        ]
-  ),
+const normalizeTokens = (result: TokensResult): TokenizedCode => ({
+  tokens: result.tokens,
+  fg: result.fg ?? "inherit",
+  bg: result.bg ?? "transparent",
 })
 
-// Synchronous highlight with callback for async results
 export const highlightCode = (
   code: string,
   language: BundledLanguage,
-  // oxlint-disable-next-line eslint-plugin-promise(prefer-await-to-callbacks)
   callback?: (result: TokenizedCode) => void
 ): TokenizedCode | null => {
-  const tokensCacheKey = getTokensCacheKey(code, language)
-
-  // Return cached result if available
-  const cached = tokensCache.get(tokensCacheKey)
-  if (cached) {
-    return cached
-  }
-
-  // Subscribe callback if provided
-  if (callback) {
-    if (!subscribers.has(tokensCacheKey)) {
-      subscribers.set(tokensCacheKey, new Set())
-    }
-    subscribers.get(tokensCacheKey)?.add(callback)
-  }
-
-  // Start highlighting in background - fire-and-forget async pattern
-  getHighlighter(language)
-    // oxlint-disable-next-line eslint-plugin-promise(prefer-await-to-then)
-    .then((highlighter) => {
-      const availableLangs = highlighter.getLoadedLanguages()
-      const langToUse = availableLangs.includes(language) ? language : "text"
-
-      const result = highlighter.codeToTokens(code, {
-        lang: langToUse,
-        themes: {
-          dark: "github-dark",
-          light: "github-light",
-        },
-      })
-
-      const tokenized: TokenizedCode = {
-        bg: result.bg ?? "transparent",
-        fg: result.fg ?? "inherit",
-        tokens: result.tokens,
-      }
-
-      // Cache the result
-      tokensCache.set(tokensCacheKey, tokenized)
-
-      // Notify all subscribers
-      const subs = subscribers.get(tokensCacheKey)
-      if (subs) {
-        for (const sub of subs) {
-          sub(tokenized)
-        }
-        subscribers.delete(tokensCacheKey)
-      }
-    })
-    // oxlint-disable-next-line eslint-plugin-promise(prefer-await-to-then), eslint-plugin-promise(prefer-await-to-callbacks)
-    .catch((error) => {
-      console.error("Failed to highlight code:", error)
-      subscribers.delete(tokensCacheKey)
-    })
-
-  return null
+  const result = highlightSharedCode(
+    { code, language, themes: DEFAULT_CODE_THEMES },
+    callback ? (tokens) => callback(normalizeTokens(tokens)) : undefined
+  )
+  return result ? normalizeTokens(result) : null
 }
 
 // Line number styles using CSS counters
@@ -391,13 +288,17 @@ export function useHighlightedTokens(
   language: BundledLanguage
 ): TokenizedCode {
   // Memoized raw tokens for immediate display
-  const rawTokens = useMemo(() => createRawTokens(code), [code])
+  const rawTokens = useMemo(() => normalizeTokens(rawCodeTokens(code)), [code])
 
   // Synchronous cached-or-raw value, recomputed when code/language changes
-  const syncTokenized = useMemo(
-    () => highlightCode(code, language) ?? rawTokens,
-    [code, language, rawTokens]
-  )
+  const syncTokenized = useMemo(() => {
+    const cached = getCachedHighlightedCode({
+      code,
+      language,
+      themes: DEFAULT_CODE_THEMES,
+    })
+    return cached ? normalizeTokens(cached) : rawTokens
+  }, [code, language, rawTokens])
 
   // Async highlighted result, tagged with its source code/language
   const [asyncState, setAsyncState] = useState<{
@@ -406,20 +307,16 @@ export function useHighlightedTokens(
     tokenized: TokenizedCode
   } | null>(null)
 
-  useEffect(() => {
-    let cancelled = false
-
-    // Subscribe to async highlighting result
-    highlightCode(code, language, (result) => {
-      if (!cancelled) {
-        setAsyncState({ code, language, tokenized: result })
-      }
-    })
-
-    return () => {
-      cancelled = true
-    }
-  }, [code, language])
+  useEffect(
+    () =>
+      subscribeHighlightedCode(
+        { code, language, themes: DEFAULT_CODE_THEMES },
+        (result) => {
+          setAsyncState({ code, language, tokenized: normalizeTokens(result) })
+        }
+      ),
+    [code, language]
+  )
 
   // Use async result only if it matches current code/language
   return asyncState?.code === code && asyncState?.language === language

src/components/ai-elements/message.test.tsx

@@ -35,11 +35,16 @@ vi.mock("@streamdown/math", () => ({
   createMathPlugin: () => ({}),
 }))
 vi.mock("@streamdown/mermaid", () => ({ mermaid: {} }))
-vi.mock("@streamdown/code", () => ({
-  code: {
+vi.mock("@/lib/code-highlighter", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/lib/code-highlighter")>()),
+  createCodePlugin: async () => ({
+    name: "shiki",
+    type: "code-highlighter",
     highlight: vi.fn(),
     supportsLanguage: vi.fn(() => true),
-  },
+    getSupportedLanguages: () => [],
+    getThemes: () => ["github-light", "github-dark"],
+  }),
 }))
 
 vi.mock("@/components/ai-elements/link-safety", () => ({