From fc8db537d5877ffba01e242b36c62830a40d9800 Mon Sep 17 00:00:00 2001 From: LukeParkerDev <10430890+Hona@users.noreply.github.com> Date: Mon, 15 Jun 2026 14:41:16 +0200 Subject: [PATCH] fix(ui): handle split streaming fences --- .../ui/src/components/markdown-stream.test.ts | 35 +++++++++++++++++++ packages/ui/src/components/markdown-stream.ts | 8 ++++- .../markdown-worker-protocol.test.ts | 8 ++++- .../components/markdown-worker-protocol.ts | 4 +++ packages/ui/src/components/markdown-worker.ts | 5 +-- 5 files changed, 56 insertions(+), 4 deletions(-) diff --git a/packages/ui/src/components/markdown-stream.test.ts b/packages/ui/src/components/markdown-stream.test.ts index b2d2e5110c..10933feb1e 100644 --- a/packages/ui/src/components/markdown-stream.test.ts +++ b/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, + }) + }) }) diff --git a/packages/ui/src/components/markdown-stream.ts b/packages/ui/src/components/markdown-stream.ts index 0351f6f5f4..072de44e4b 100644 --- a/packages/ui/src/components/markdown-stream.ts +++ b/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, diff --git a/packages/ui/src/components/markdown-worker-protocol.test.ts b/packages/ui/src/components/markdown-worker-protocol.test.ts index 18dcef8eac..4ec9ca69d5 100644 --- a/packages/ui/src/components/markdown-worker-protocol.test.ts +++ b/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) +}) diff --git a/packages/ui/src/components/markdown-worker-protocol.ts b/packages/ui/src/components/markdown-worker-protocol.ts index 36d05d2aa3..392863575b 100644 --- a/packages/ui/src/components/markdown-worker-protocol.ts +++ b/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, diff --git a/packages/ui/src/components/markdown-worker.ts b/packages/ui/src/components/markdown-worker.ts index e80f0c782f..d68e4e550d 100644 --- a/packages/ui/src/components/markdown-worker.ts +++ b/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((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)