fix(ui): handle split streaming fences

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

@@ -145,4 +145,39 @@ describe("markdown stream", () => {
       language: "ts",
     })
   })
+
+  test("closes code fences split across provider deltas", () => {
+    const open = project(undefined, "```ts\nconst x = 1\n", true)
+    const one = project(open, `${open.text}\``, true)
+    const two = project(one, `${one.text}\``, true)
+    const closed = project(two, `${two.text}\``, true)
+    const prose = project(closed, `${closed.text}\nafter`, true)
+
+    expect(closed.blocks.at(-1)).toEqual({
+      raw: "```ts\nconst x = 1\n```",
+      src: "const x = 1",
+      mode: "code",
+      language: "ts",
+      complete: true,
+    })
+    expect(prose.blocks).toEqual([
+      { raw: "```ts\nconst x = 1\n```\n", src: "```ts\nconst x = 1\n```\n", mode: "full" },
+      { raw: "after", src: "after", mode: "live" },
+    ])
+  })
+
+  test("closes tilde fences split across provider deltas", () => {
+    const open = project(undefined, "~~~ts\nconst x = 1\n", true)
+    const one = project(open, `${open.text}~`, true)
+    const two = project(one, `${one.text}~`, true)
+    const closed = project(two, `${two.text}~`, true)
+
+    expect(closed.blocks.at(-1)).toEqual({
+      raw: "~~~ts\nconst x = 1\n~~~",
+      src: "const x = 1",
+      mode: "code",
+      language: "ts",
+      complete: true,
+    })
+  })
 })

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

@@ -39,6 +39,12 @@ function open(raw: string) {
   return !new RegExp(`^[\\t ]{0,3}${char}{${size},}[\\t ]*$`).test(last)
 }
 
+function closesFence(raw: string, suffix: string) {
+  const mark = raw.match(/^[ \t]{0,3}(`{3,}|~{3,})/)?.[1]
+  if (!mark) return suffix.includes("```") || suffix.includes("~~~")
+  return `${raw.slice(-(mark.length - 1))}${suffix}`.includes(mark)
+}
+
 function heal(text: string) {
   return remend(text, { linkMode: "text-only" })
 }
@@ -83,7 +89,7 @@ export function project(previous: Projection | undefined, text: string, live: bo
   if (!live || !previous || !text.startsWith(previous.text)) return { text, blocks: stream(text, live) }
   const tail = previous.blocks.at(-1)
   const suffix = text.slice(previous.text.length)
-  if (!suffix || tail?.mode !== "code" || tail.complete || suffix.includes("```") || suffix.includes("~~~"))
+  if (!suffix || tail?.mode !== "code" || tail.complete || closesFence(tail.raw, suffix))
     return { text, blocks: stream(text, live) }
   return {
     text,

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

@@ -1,5 +1,5 @@
 import { expect, test } from "bun:test"
-import { applyMarkdownWorkerResponse } from "./markdown-worker-protocol"
+import { applyMarkdownWorkerResponse, shouldReleaseMarkdownWorkerState } from "./markdown-worker-protocol"
 
 const token = (content: string): [string, string] => [content, ""]
 const response = (id: number, reset: boolean, stable: [string, string][], unstable: [string, string][]) => ({
@@ -64,3 +64,9 @@ test("ignores stale worker responses and resets replacement streams", () => {
     }).stable.map((item) => item[0]),
   ).toEqual(["replacement"])
 })
+
+test("releases only the latest completed worker state", () => {
+  expect(shouldReleaseMarkdownWorkerState(true, 4, 4)).toBe(true)
+  expect(shouldReleaseMarkdownWorkerState(true, 5, 4)).toBe(false)
+  expect(shouldReleaseMarkdownWorkerState(false, 4, 4)).toBe(false)
+})

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

@@ -25,6 +25,10 @@ export type MarkdownWorkerState = {
   unstable: MarkdownToken[]
 }
 
+export function shouldReleaseMarkdownWorkerState(complete: boolean, latestID: number | undefined, responseID: number) {
+  return complete && latestID === responseID
+}
+
 export function applyMarkdownWorkerResponse(
   state: MarkdownWorkerState | undefined,
   response: Extract<MarkdownWorkerResponse, { type: "highlight" }>,

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

@@ -2,6 +2,7 @@ import MarkdownShikiWorkerUrl from "./markdown-shiki.worker.ts?worker&url"
 import { OpenCodeTheme } from "../context/marked"
 import {
   applyMarkdownWorkerResponse,
+  shouldReleaseMarkdownWorkerState,
   type MarkdownWorkerRequest,
   type MarkdownWorkerResponse,
   type MarkdownWorkerState,
@@ -26,7 +27,6 @@ export function highlightStreamingCode(key: string, text: string, language: stri
   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) => {
@@ -37,6 +37,7 @@ export function highlightStreamingCode(key: string, text: string, language: stri
 
 export function disposeStreamingCode(key: string) {
   keys.delete(key)
+  latest.delete(key)
   states.delete(key)
   pending.forEach((request, id) => {
     if (request.key !== key) return
@@ -64,7 +65,7 @@ function getWorker() {
       return
     }
     const state = applyMarkdownWorkerResponse(states.get(event.data.key), event.data)
-    if (result.complete && latest.get(event.data.key) === event.data.id) {
+    if (shouldReleaseMarkdownWorkerState(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)