fix(ui): harden streaming markdown updates

packages/ui/src/components/markdown-stream.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test"
-import { project, stream } from "./markdown-stream"
+import { canReusePendingBlock, project, stream } from "./markdown-stream"
 
 describe("markdown stream", () => {
   test("heals incomplete emphasis while streaming", () => {
@@ -83,6 +83,44 @@ describe("markdown stream", () => {
     ])
   })
 
+  test("keeps multiline reference definitions with their uses", () => {
+    expect(stream("[docs][id]\n\n[id]:\n  /guide", true)).toEqual([
+      {
+        raw: "[docs][id]\n\n[id]:\n  /guide",
+        src: "[docs][id]\n\n[id]:\n  /guide",
+        mode: "live",
+      },
+    ])
+  })
+
+  test("uses only the language portion of fence metadata", () => {
+    expect(stream("```ts title=example\nconst x = 1", true)).toEqual([
+      {
+        raw: "```ts title=example\nconst x = 1",
+        src: "const x = 1",
+        mode: "code",
+        language: "ts",
+      },
+    ])
+  })
+
+  test("preserves trailing newlines in open code fences", () => {
+    expect(stream("```ts\nconst x = 1\n", true)).toEqual([
+      {
+        raw: "```ts\nconst x = 1\n",
+        src: "const x = 1\n",
+        mode: "code",
+        language: "ts",
+      },
+    ])
+  })
+
+  test("only reuses pending blocks with compatible identity and content", () => {
+    expect(canReusePendingBlock({ mode: "full", raw: "First\n\n" }, { mode: "full", raw: "# Inserted\n\n", src: "", })).toBe(false)
+    expect(canReusePendingBlock({ mode: "code", raw: "```ts\none" }, { mode: "code", raw: "```ts\none two", src: "" })).toBe(true)
+    expect(canReusePendingBlock({ mode: "code", raw: "```ts\none" }, { mode: "live", raw: "one", src: "" })).toBe(false)
+  })
+
   test("appends plain code deltas without reprojecting frozen blocks", () => {
     const previous = project(undefined, "# Plan\n\n```ts\nconst one = 1\n", true)
     const next = project(previous, `${previous.text}const two = 2\n`, true)

packages/ui/src/components/markdown-stream.ts

@@ -16,7 +16,16 @@ export type Projection = {
 
 function refs(text: string) {
   if (!text.includes("]:")) return false
-  return /^[ \t]{0,3}\[[^\]]+\]:[ \t]*\S+/m.test(text)
+  return /^[ \t]{0,3}\[[^\]]+\]:[ \t]*(?:\S+|\r?\n[ \t]+\S+)/m.test(text)
+}
+
+function language(value: string | undefined) {
+  return value?.trim().split(/\s+/, 1)[0] || undefined
+}
+
+function openCode(raw: string) {
+  const newline = raw.indexOf("\n")
+  return newline < 0 ? "" : raw.slice(newline + 1)
 }
 
 function open(raw: string) {
@@ -60,8 +69,14 @@ export function stream(text: string, live: boolean): Block[] {
 
   const code = last as Tokens.Code
   if (!open(code.raw))
-    return [...result, { raw, src: code.text, mode: "code", language: code.lang, complete: true }]
-  return [...result, { raw, src: code.text, mode: "code", language: code.lang }]
+    return [...result, { raw, src: code.text, mode: "code", language: language(code.lang), complete: true }]
+  return [...result, { raw, src: openCode(code.raw), mode: "code", language: language(code.lang) }]
+}
+
+export function canReusePendingBlock(current: Pick<Block, "mode" | "raw"> | undefined, next: Block) {
+  if (!current || current.mode !== next.mode) return false
+  if (next.mode === "code") return next.raw.startsWith(current.raw)
+  return current.raw === next.raw
 }
 
 export function project(previous: Projection | undefined, text: string, live: boolean): Projection {
@@ -70,7 +85,6 @@ export function project(previous: Projection | undefined, text: string, live: bo
   const suffix = text.slice(previous.text.length)
   if (!suffix || tail?.mode !== "code" || tail.complete || suffix.includes("```") || suffix.includes("~~~"))
     return { text, blocks: stream(text, live) }
-  const separator = tail.src.length > 0 && previous.text.endsWith("\n") && !tail.src.endsWith("\n") ? "\n" : ""
   return {
     text,
     blocks: [
@@ -78,7 +92,7 @@ export function project(previous: Projection | undefined, text: string, live: bo
       {
         ...tail,
         raw: tail.raw + suffix,
-        src: tail.src + separator + suffix,
+        src: tail.src + suffix,
       },
     ],
   }

packages/ui/src/components/markdown-worker-protocol.test.ts

@@ -2,6 +2,14 @@ import { expect, test } from "bun:test"
 import { applyMarkdownWorkerResponse } from "./markdown-worker-protocol"
 
 const token = (content: string): [string, string] => [content, ""]
+const response = (id: number, reset: boolean, stable: [string, string][], unstable: [string, string][]) => ({
+  type: "highlight" as const,
+  id,
+  key: "code",
+  reset,
+  stable,
+  unstable,
+})
 
 test("accumulates stable worker tokens and replaces the unstable tail", () => {
   const first = applyMarkdownWorkerResponse(undefined, {
@@ -25,8 +33,15 @@ test("accumulates stable worker tokens and replaces the unstable tail", () => {
   expect(second.unstable.map((item) => item[0])).toEqual(["three"])
 })
 
+test("increments generation only when the worker resets token identity", () => {
+  const first = applyMarkdownWorkerResponse(undefined, response(1, true, [["const", ""]], []))
+  const append = applyMarkdownWorkerResponse(first, response(2, false, [[" x", ""]], []))
+  const replacement = applyMarkdownWorkerResponse(append, response(3, true, [["let y", ""]], []))
+  expect([first.generation, append.generation, replacement.generation]).toEqual([1, 1, 2])
+})
+
 test("ignores stale worker responses and resets replacement streams", () => {
-  const current = { id: 2, stable: [token("current")], unstable: [] }
+  const current = { id: 2, generation: 1, stable: [token("current")], unstable: [] }
   expect(
     applyMarkdownWorkerResponse(current, {
       type: "highlight",

packages/ui/src/components/markdown-worker-protocol.ts

@@ -20,6 +20,7 @@ export type MarkdownWorkerResponse =
 
 export type MarkdownWorkerState = {
   id: number
+  generation: number
   stable: MarkdownToken[]
   unstable: MarkdownToken[]
 }
@@ -31,6 +32,7 @@ export function applyMarkdownWorkerResponse(
   if (state && response.id <= state.id) return state
   return {
     id: response.id,
+    generation: (state?.generation ?? 0) + (response.reset ? 1 : 0),
     stable: response.reset ? response.stable : [...(state?.stable ?? []), ...response.stable],
     unstable: response.unstable,
   }

packages/ui/src/components/markdown-worker.ts

@@ -19,11 +19,14 @@ let nextID = 0
 const pending = new Map<number, Pending>()
 const states = new Map<string, MarkdownWorkerState>()
 const keys = new Set<string>()
+const latest = new Map<string, number>()
 
 export function highlightStreamingCode(key: string, text: string, language: string, complete = false) {
   const instance = getWorker()
   const id = ++nextID
+  latest.set(key, id)
   keys.delete(key)
+  latest.delete(key)
   keys.add(key)
   if (keys.size > 200) disposeStreamingCode(keys.values().next().value!)
   return new Promise<MarkdownWorkerState>((resolve, reject) => {
@@ -61,21 +64,25 @@ function getWorker() {
       return
     }
     const state = applyMarkdownWorkerResponse(states.get(event.data.key), event.data)
-    if (result.complete) {
+    if (result.complete && latest.get(event.data.key) === event.data.id) {
       states.delete(event.data.key)
       keys.delete(event.data.key)
+      latest.delete(event.data.key)
     } else states.set(event.data.key, state)
     result.resolve(state)
   }
-  worker.onerror = (event) => {
-    const error = new Error(event.message || "Markdown highlighting worker failed")
+  const fail = (message: string) => {
+    const error = new Error(message)
     pending.forEach((request) => request.reject(error))
     pending.clear()
     states.clear()
     keys.clear()
+    latest.clear()
     worker?.terminate()
     worker = undefined
   }
+  worker.onerror = (event) => fail(event.message || "Markdown highlighting worker failed")
+  worker.onmessageerror = () => fail("Markdown worker response failed")
   worker.postMessage({ type: "init", theme: OpenCodeTheme } satisfies MarkdownWorkerRequest)
   return worker
 }

packages/ui/src/components/markdown.css

@@ -127,6 +127,8 @@
   }
 
   .shiki {
+    background: var(--color-background-stronger);
+    color: var(--text-base);
     font-size: 13px;
     padding: 12px;
     border-radius: 6px;

packages/ui/src/components/markdown.tsx

@@ -3,18 +3,10 @@ import { useI18n } from "../context/i18n"
 import DOMPurify from "dompurify"
 import morphdom from "morphdom"
 import { checksum } from "@opencode-ai/core/util/encode"
-import { getSharedHighlighter } from "@pierre/diffs"
-import { ShikiStreamTokenizer } from "@shikijs/stream"
-import { ComponentProps, createEffect, createMemo, createResource, createSignal, onCleanup, splitProps } from "solid-js"
+import { ComponentProps, createEffect, createMemo, createResource, createSignal, createUniqueId, onCleanup, splitProps } from "solid-js"
 import { isServer } from "solid-js/web"
-import {
-  bundledLanguages,
-  getTokenStyleObject,
-  stringifyTokenStyle,
-  type BundledLanguage,
-  type ThemedToken,
-} from "shiki"
-import { project, type Block, type Projection } from "./markdown-stream"
+import { bundledLanguages } from "shiki"
+import { canReusePendingBlock, project, type Block, type Projection } from "./markdown-stream"
 import { disposeStreamingCode, highlightStreamingCode, MarkdownWorkerDisposedError } from "./markdown-worker"
 import type { MarkdownToken } from "./markdown-worker-protocol"
 
@@ -33,6 +25,7 @@ type RenderedBlock =
       hash: string
       language: string
       complete: boolean
+      generation: number
       stable: MarkdownToken[]
       unstable: MarkdownToken[]
     }
@@ -46,18 +39,7 @@ const max = 200
 const cache = new Map<string, Entry>()
 const renderedCodeTokens = new WeakMap<
   HTMLDivElement,
-  { language: string; stableCount: number; unstable: MarkdownToken[] }
->()
-const codeCache = new Map<
-  string,
-  {
-    language: string
-    source: string
-    tokenizer: ShikiStreamTokenizer
-    stable: ThemedToken[]
-    unstable: ThemedToken[]
-    pending: Promise<void>
-  }
+  { language: string; generation: number; stableCount: number; unstable: MarkdownToken[] }
 >()
 
 if (typeof window !== "undefined" && DOMPurify.isSupported) {
@@ -107,62 +89,13 @@ function fallback(markdown: string) {
 
 async function code(text: string, language: string | undefined, key: string, complete = false) {
   const name = language && language in bundledLanguages ? language : "text"
-  if (typeof Worker !== "undefined") {
-    try {
-      const result = await highlightStreamingCode(key, text, name, complete)
-      return { language: name, stable: result.stable, unstable: result.unstable }
-    } catch (error) {
-      if (error instanceof MarkdownWorkerDisposedError)
-        return { language: name, stable: [], unstable: [[text, ""] as MarkdownToken] }
-      // Keep highlighting available in runtimes where the worker cannot start.
-    }
+  try {
+    const result = await highlightStreamingCode(key, text, name, complete)
+    return { language: name, generation: result.generation, stable: result.stable, unstable: result.unstable }
+  } catch (error) {
+    if (!(error instanceof MarkdownWorkerDisposedError)) console.error("Markdown highlighting worker failed", error)
+    return { language: name, generation: 0, stable: [], unstable: [[text, ""] as MarkdownToken] }
   }
-  return codeOnMainThread(text, name, key)
-}
-
-async function codeOnMainThread(text: string, name: string, key: string) {
-  const existing = codeCache.get(key)
-  const entry =
-    existing && existing.language === name && text.startsWith(existing.source)
-      ? existing
-      : {
-          language: name,
-          source: "",
-          tokenizer: new ShikiStreamTokenizer({
-            highlighter: await getSharedHighlighter({
-              themes: ["OpenCode"],
-              langs: [],
-              preferredHighlighter: "shiki-wasm",
-            }),
-            lang: name,
-            theme: "OpenCode",
-          }),
-          stable: [],
-          unstable: [],
-          pending: Promise.resolve(),
-        }
-  if (!entry.tokenizer.options.highlighter.getLoadedLanguages().includes(name))
-    await entry.tokenizer.options.highlighter.loadLanguage(bundledLanguages[name as BundledLanguage])
-  const suffix = text.slice(entry.source.length)
-  entry.source = text
-  entry.pending = entry.pending.then(async () => {
-    const result = await entry.tokenizer.enqueue(suffix)
-    entry.stable.push(...result.stable.filter((token) => token.content.length > 0))
-    entry.unstable = result.unstable.filter((token) => token.content.length > 0)
-  })
-  codeCache.delete(key)
-  codeCache.set(key, entry)
-  if (codeCache.size > max) codeCache.delete(codeCache.keys().next().value!)
-  await entry.pending
-  return {
-    language: name,
-    stable: entry.stable.map(token),
-    unstable: entry.unstable.map(token),
-  }
-}
-
-function token(value: ThemedToken): MarkdownToken {
-  return [value.content, stringifyTokenStyle(value.htmlStyle ?? getTokenStyleObject(value))]
 }
 
 type CopyLabels = {
@@ -198,7 +131,7 @@ function createIcon(path: string, slot: string) {
   return icon
 }
 
-function createCopyButton(labels: CopyLabels, icons = true) {
+function createCopyButton(labels: CopyLabels) {
   const button = document.createElement("button")
   button.type = "button"
   button.setAttribute("data-component", "icon-button")
@@ -207,10 +140,8 @@ function createCopyButton(labels: CopyLabels, icons = true) {
   button.setAttribute("data-slot", "markdown-copy-button")
   button.setAttribute("aria-label", labels.copy)
   button.setAttribute("data-tooltip", labels.copy)
-  if (icons) {
-    button.appendChild(createIcon(iconPaths.copy, "copy-icon"))
-    button.appendChild(createIcon(iconPaths.check, "check-icon"))
-  }
+  button.appendChild(createIcon(iconPaths.copy, "copy-icon"))
+  button.appendChild(createIcon(iconPaths.check, "check-icon"))
   return button
 }
 
@@ -226,7 +157,7 @@ function setCopyState(button: HTMLButtonElement, labels: CopyLabels, copied: boo
   button.setAttribute("data-tooltip", labels.copy)
 }
 
-function ensureCodeWrapper(block: HTMLPreElement, labels: CopyLabels, shallow: boolean) {
+function ensureCodeWrapper(block: HTMLPreElement, labels: CopyLabels) {
   const parent = block.parentElement
   if (!parent) return
   const wrapped = parent.getAttribute("data-component") === "markdown-code"
@@ -235,7 +166,7 @@ function ensureCodeWrapper(block: HTMLPreElement, labels: CopyLabels, shallow: b
     wrapper.setAttribute("data-component", "markdown-code")
     parent.replaceChild(wrapper, block)
     wrapper.appendChild(block)
-    wrapper.appendChild(createCopyButton(labels, !shallow))
+    wrapper.appendChild(createCopyButton(labels))
     return
   }
 
@@ -282,10 +213,10 @@ function markCodeLinks(root: HTMLDivElement) {
   }
 }
 
-function decorate(root: HTMLDivElement, labels: CopyLabels, shallow = false) {
+function decorate(root: HTMLDivElement, labels: CopyLabels) {
   const blocks = Array.from(root.querySelectorAll("pre"))
   for (const block of blocks) {
-    ensureCodeWrapper(block, labels, shallow)
+    ensureCodeWrapper(block, labels)
   }
   markCodeLinks(root)
 }
@@ -358,6 +289,7 @@ export function Markdown(
   const marked = useMarked()
   const i18n = useI18n()
   const [root, setRoot] = createSignal<HTMLDivElement>()
+  const owner = createUniqueId()
   const activeCodeKeys = new Set<string>()
   const projection = createMemo((previous: Projection | undefined) =>
     project(previous, local.text, local.streaming ?? false),
@@ -390,7 +322,7 @@ export function Markdown(
       return Promise.all(
         src.projection.blocks.map(async (block, index) => {
           const key = base ? `${base}:${index}:${block.mode}` : undefined
-          const blockKey = key ?? `block:${index}`
+          const blockKey = `${owner}:${key ?? `block:${index}`}`
 
           if (block.mode === "code") {
             const result = await code(block.src, block.language, blockKey, block.complete)
@@ -477,6 +409,9 @@ export function Markdown(
     nextCodeKeys.forEach((key) => activeCodeKeys.add(key))
     content.forEach((block, index) => updateBlock(container, index, block, labels))
     while (container.children.length > content.length) container.lastElementChild?.remove()
+    container.querySelectorAll<HTMLButtonElement>('[data-slot="markdown-copy-button"]').forEach((button) =>
+      setCopyState(button, labels, button.dataset.copied === "true"),
+    )
 
     if (!copyCleanup)
       copyCleanup = setupCodeCopy(container, () => ({
@@ -509,7 +444,7 @@ function pendingBlocks(result: RenderResult | undefined, projection: Projection
   const initial = result.blocks.length === 1 && result.blocks[0]?.key === "initial"
   return projection.blocks.map((block, index) => {
     const current = initial ? undefined : result.blocks[index]
-    if (current) return current
+    if (current && canReusePendingBlock(current, block)) return current
     const key = cacheKey ? `${cacheKey}:${index}:${block.mode}` : `block:${index}`
     if (block.mode !== "code")
       return { key, mode: block.mode, raw: block.raw, hash: String(block.raw.length), html: fallback(block.src) }
@@ -521,13 +456,13 @@ function pendingBlocks(result: RenderResult | undefined, projection: Projection
       language: block.language ?? "text",
       complete: !!block.complete,
       stable: [],
+      generation: 0,
       unstable: [[block.src, ""] as MarkdownToken],
     }
   })
 }
 
 function disposeCode(key: string) {
-  codeCache.delete(key)
   disposeStreamingCode(key)
 }
 
@@ -550,11 +485,7 @@ function updateBlock(container: HTMLDivElement, index: number, block: RenderedBl
   next.dataset.markdownHash = block.hash
   next.style.display = "contents"
   next.innerHTML = block.html
-  decorate(
-    next,
-    labels,
-    current instanceof HTMLDivElement && current.querySelector('[data-slot="markdown-copy-button"]') !== null,
-  )
+  decorate(next, labels)
 
   if (!(current instanceof HTMLDivElement)) {
     container.appendChild(next)
@@ -596,7 +527,11 @@ function updateCodeBlock(
   if (code instanceof HTMLElement) {
     code.className = `language-${block.language}`
     const previous = renderedCodeTokens.get(next)
-    const reset = !previous || previous.language !== block.language || block.stable.length < previous.stableCount
+    const reset =
+      !previous ||
+      previous.language !== block.language ||
+      previous.generation !== block.generation ||
+      block.stable.length < previous.stableCount
     const stableCount = reset ? 0 : previous.stableCount
     const tail = [...block.stable.slice(stableCount), ...block.unstable]
     const prior = reset ? [] : previous.unstable
@@ -606,6 +541,7 @@ function updateCodeBlock(
     tail.slice(keep - stableCount).map(createTokenSpan).forEach((span) => code.appendChild(span))
     renderedCodeTokens.set(next, {
       language: block.language,
+      generation: block.generation,
       stableCount: block.stable.length,
       unstable: block.unstable,
     })
@@ -625,6 +561,7 @@ function updateCodeBlock(
   next.appendChild(wrapper)
   renderedCodeTokens.set(next, {
     language: block.language,
+    generation: block.generation,
     stableCount: block.stable.length,
     unstable: block.unstable,
   })