diff --git a/packages/session-ui/src/components/markdown-cache.tsx b/packages/session-ui/src/components/markdown-cache.tsx index 3f8f0259f3..9d4d03e940 100644 --- a/packages/session-ui/src/components/markdown-cache.tsx +++ b/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() +const codeCache = new Map() 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, diff --git a/packages/session-ui/src/components/markdown-preload.test.ts b/packages/session-ui/src/components/markdown-preload.test.ts index cba107def2..5b02853a2a 100644 --- a/packages/session-ui/src/components/markdown-preload.test.ts +++ b/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: [], + }) +}) diff --git a/packages/session-ui/src/components/markdown.tsx b/packages/session-ui/src/components/markdown.tsx index f0c3129918..c1e3141953 100644 --- a/packages/session-ui/src/components/markdown.tsx +++ b/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, "
") } +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 }