fix(session-ui): cache code highlights

packages/session-ui/src/components/markdown-cache.tsx

@@ -1,5 +1,6 @@
 import { checksum } from "@opencode-ai/core/util/encode"
 import DOMPurify from "dompurify"
+import type { MarkdownToken } from "./markdown-worker-protocol"
 import { project } from "./markdown-stream"
 
 export type MarkdownCacheEntry = {
@@ -8,8 +9,18 @@ export type MarkdownCacheEntry = {
   html: string
 }
 
+export type MarkdownCodeCacheEntry = {
+  raw: string
+  hash: string
+  language: string
+  generation: number
+  stable: MarkdownToken[]
+  unstable: MarkdownToken[]
+}
+
 const max = 200
 const cache = new Map<string, MarkdownCacheEntry>()
+const codeCache = new Map<string, MarkdownCodeCacheEntry>()
 const config = {
   USE_PROFILES: { html: true, mathMl: true },
   SANITIZE_NAMED_PROPS: true,
@@ -52,6 +63,21 @@ export function touchCachedMarkdown(key: string, value: MarkdownCacheEntry) {
   cache.delete(first)
 }
 
+export function getCachedMarkdownCode(key: string) {
+  return codeCache.get(key)
+}
+
+export function touchCachedMarkdownCode(key: string, value: MarkdownCodeCacheEntry) {
+  codeCache.delete(key)
+  codeCache.set(key, value)
+
+  if (codeCache.size <= max) return
+
+  const first = codeCache.keys().next().value
+  if (!first) return
+  codeCache.delete(first)
+}
+
 export async function preloadMarkdown(
   text: string,
   cacheKey: string,

packages/session-ui/src/components/markdown-preload.test.ts

@@ -1,5 +1,5 @@
 import { expect, test } from "bun:test"
-import { preloadMarkdown } from "./markdown-cache"
+import { getCachedMarkdownCode, preloadMarkdown, touchCachedMarkdownCode } from "./markdown-cache"
 
 test("preloads completed markdown into the render cache", async () => {
   const parsed: string[] = []
@@ -16,3 +16,24 @@ test("preloads completed markdown into the render cache", async () => {
 
   expect(parsed).toEqual(["prepared response"])
 })
+
+test("keeps completed code highlights by stable block key", () => {
+  const key = `markdown-code-${crypto.randomUUID()}:0:code`
+  touchCachedMarkdownCode(key, {
+    raw: "```ts\nconst value = 1\n```",
+    hash: "23",
+    language: "ts",
+    generation: 1,
+    stable: [["const", "color: red"]],
+    unstable: [],
+  })
+
+  expect(getCachedMarkdownCode(key)).toEqual({
+    raw: "```ts\nconst value = 1\n```",
+    hash: "23",
+    language: "ts",
+    generation: 1,
+    stable: [["const", "color: red"]],
+    unstable: [],
+  })
+})

packages/session-ui/src/components/markdown.tsx

@@ -24,7 +24,14 @@ import {
 } from "./markdown-worker"
 import { markdownBlockKey, type MarkdownToken } from "./markdown-worker-protocol"
 import { shouldResetCodeTokens, type RenderedCodeState } from "./markdown-code-state"
-import { getCachedMarkdown, sanitizeMarkdown, touchCachedMarkdown, type MarkdownCacheEntry } from "./markdown-cache"
+import {
+  getCachedMarkdown,
+  getCachedMarkdownCode,
+  sanitizeMarkdown,
+  touchCachedMarkdown,
+  touchCachedMarkdownCode,
+  type MarkdownCacheEntry,
+} from "./markdown-cache"
 import { inlineCodeKind } from "./markdown-inline-code-kind"
 
 type RenderedBlock =
@@ -66,8 +73,12 @@ function fallback(markdown: string) {
   return escape(markdown).replace(/\r\n?/g, "\n").replace(/\n/g, "<br>")
 }
 
+function highlightLanguage(language: string | undefined) {
+  return language && language in bundledLanguages ? language : "text"
+}
+
 async function code(text: string, language: string | undefined, key: string, complete = false) {
-  const name = language && language in bundledLanguages ? language : "text"
+  const name = highlightLanguage(language)
   try {
     const result = await highlightStreamingCode(key, text, name, complete)
     return { language: name, generation: result.generation, stable: result.stable, unstable: result.unstable }
@@ -297,9 +308,15 @@ function initialResult(text: string, key: string | undefined, projection: Projec
   if (!text) return { text, blocks: [] }
   const base = key ?? checksum(text)
   if (base) {
-    const blocks = projection.blocks.flatMap((block, index) => {
-      if (block.mode === "code") return []
+    const blocks = projection.blocks.flatMap((block, index): RenderedBlock[] => {
       const cacheKey = `${base}:${index}:${block.mode}`
+      if (block.mode === "code") {
+        if (!block.complete) return []
+        const cached = getCachedMarkdownCode(cacheKey)
+        if (cached?.raw !== block.raw) return []
+        touchCachedMarkdownCode(cacheKey, cached)
+        return [{ key: `${owner}:${cacheKey}`, mode: block.mode, complete: true, ...cached }]
+      }
       const cached = getCachedMarkdown(cacheKey)
       if (cached?.raw !== block.raw) return []
       return [{ key: `${owner}:${cacheKey}`, mode: block.mode, ...cached }]
@@ -370,9 +387,17 @@ export function Markdown(
           const blockKey = markdownBlockKey(owner, src.key, index, block.mode)
 
           if (block.mode === "code") {
+            const language = highlightLanguage(block.language)
+            if (block.complete && key) {
+              const shared = getCachedMarkdownCode(key)
+              if (shared?.raw === block.raw && shared.language === language) {
+                touchCachedMarkdownCode(key, shared)
+                return { key: blockKey, mode: block.mode, complete: true, ...shared }
+              }
+            }
             const cached = completedCode.get(blockKey)
             if (block.complete && cached?.raw === block.raw) return cached
-            const result = await code(block.src, block.language, blockKey, block.complete)
+            const result = await code(block.src, language, blockKey, block.complete)
             const rendered = {
               key: blockKey,
               mode: block.mode,
@@ -381,7 +406,18 @@ export function Markdown(
               complete: !!block.complete,
               ...result,
             }
-            if (block.complete) completedCode.set(blockKey, rendered)
+            if (block.complete) {
+              completedCode.set(blockKey, rendered)
+              if (key && rendered.generation > 0)
+                touchCachedMarkdownCode(key, {
+                  raw: rendered.raw,
+                  hash: rendered.hash,
+                  language: rendered.language,
+                  generation: rendered.generation,
+                  stable: rendered.stable,
+                  unstable: rendered.unstable,
+                })
+            }
             return rendered
           }