diff --git a/packages/session-ui/src/components/markdown-worker-protocol.test.ts b/packages/session-ui/src/components/markdown-worker-protocol.test.ts index 1971b69af4cc..9d6726553a60 100644 --- a/packages/session-ui/src/components/markdown-worker-protocol.test.ts +++ b/packages/session-ui/src/components/markdown-worker-protocol.test.ts @@ -1,9 +1,12 @@ import { expect, test } from "bun:test" import { applyMarkdownWorkerResponse, + applyProjectionBlocks, + diffProjectionBlocks, markdownBlockKey, shouldReleaseMarkdownWorkerState, } from "./markdown-worker-protocol" +import { project, type Block, type Projection } from "./markdown-stream" const token = (content: string): [string, string] => [content, ""] @@ -83,6 +86,33 @@ test("releases only the latest completed worker state", () => { expect(shouldReleaseMarkdownWorkerState(false, 4, 4)).toBe(false) }) +test("rebuilds each streamed projection from the blocks that changed", () => { + const text = + "# Plan\n\nFirst paragraph.\n\n```ts\nconst a = 1\n```\n\n- item one\n- item two\n\n[docs]: https://example.com\n\nTail" + + const steps = Array.from({ length: Math.ceil(text.length / 3) }, (_, index) => [index * 3 + 3, true] as const) + const sent: number[] = [] + + const final = [...steps, [text.length, false] as const].reduce<{ worker?: Projection; main?: Block[] }>( + (state, [end, live]) => { + const worker = project(state.worker, text.slice(0, end), live) + const diff = diffProjectionBlocks(state.worker?.blocks, worker.blocks) + // Blocks cross to the main thread by structured clone, so only its own earlier objects can be reused. + const main = applyProjectionBlocks(state.main, diff.from, structuredClone(diff.blocks)) + sent.push(diff.blocks.length) + + expect(main).toEqual(worker.blocks) + main.slice(0, diff.from).forEach((block, index) => expect(block).toBe(state.main![index]!)) + + return { worker, main } + }, + {}, + ) + + expect(final.worker!.blocks.length).toBeGreaterThan(3) + expect(sent.slice(-6, -1).every((count) => count < final.worker!.blocks.length)).toBe(true) +}) + test("prefixes pending and dispatched block keys with the component owner", () => { expect(markdownBlockKey("owner", "message", 2, "code")).toBe("owner:message:2:code") expect(markdownBlockKey("owner", undefined, 2, "code")).toBe("owner:block:2") diff --git a/packages/session-ui/src/components/markdown-worker-protocol.ts b/packages/session-ui/src/components/markdown-worker-protocol.ts index 03031009f091..8def43897886 100644 --- a/packages/session-ui/src/components/markdown-worker-protocol.ts +++ b/packages/session-ui/src/components/markdown-worker-protocol.ts @@ -1,4 +1,4 @@ -import type { Projection } from "./markdown-stream" +import type { Block } from "./markdown-stream" export type MarkdownToken = [content: string, style: string] @@ -10,7 +10,7 @@ export type MarkdownWorkerRequest = export type MarkdownWorkerResponse = | { type: "parse"; id: number; html: string } - | { type: "project"; id: number; key: string; projection: Projection } + | { type: "project"; id: number; key: string; from: number; blocks: Block[] } | { type: "highlight" id: number @@ -35,6 +35,19 @@ export function shouldReleaseMarkdownWorkerState(complete: boolean, latestID: nu return complete && latestID === responseID } +// A streamed answer keeps its settled blocks, so each projection response carries only the blocks from the first one +// that changed since the previous response for its key. Resending them all copied the whole answer on every delta. +export function diffProjectionBlocks(previous: Block[] | undefined, next: Block[]) { + const same = previous ? next.findIndex((block, index) => block !== previous[index]) : 0 + const from = same === -1 ? next.length : same + + return { from, blocks: next.slice(from) } +} + +export function applyProjectionBlocks(previous: Block[] | undefined, from: number, blocks: Block[]) { + return [...(previous ?? []).slice(0, from), ...blocks] +} + export function markdownBlockKey(owner: string, cacheKey: string | undefined, index: number, mode: string) { return `${owner}:${cacheKey ? `${cacheKey}:${index}:${mode}` : `block:${index}`}` } diff --git a/packages/session-ui/src/components/markdown-worker.ts b/packages/session-ui/src/components/markdown-worker.ts index 5330bdd79084..87dbc5ca76ff 100644 --- a/packages/session-ui/src/components/markdown-worker.ts +++ b/packages/session-ui/src/components/markdown-worker.ts @@ -1,12 +1,13 @@ import { applyMarkdownWorkerResponse, + applyProjectionBlocks, shouldReleaseMarkdownWorkerState, type MarkdownWorkerRequest, type MarkdownWorkerResponse, type MarkdownWorkerState, } from "./markdown-worker-protocol" import { createWorkerTransport } from "./markdown-worker-transport" -import type { Projection } from "./markdown-stream" +import type { Block, Projection } from "./markdown-stream" type HighlightPending = { key: string @@ -17,6 +18,7 @@ type HighlightPending = { type ProjectPending = { key: string + text: string resolve: (projection: Projection) => void reject: (error: Error) => void } @@ -38,6 +40,9 @@ const projects = new Map() const parses = new Map() +// The blocks of the last projection the worker sent for each key, which its next response builds on. +const projected = new Map() + const states = new Map() const keys = new Set() @@ -91,13 +96,14 @@ export function projectMarkdown(key: string, text: string, live: boolean) { const id = ++nextID return new Promise((resolve, reject) => { - projects.set(id, { key, resolve, reject }) + projects.set(id, { key, text, resolve, reject }) projectTransport.send({ type: "project", id, key, text, live }) }) } export function disposeMarkdownProjection(key: string) { projectTransport.dispose(key) + projected.delete(key) projects.forEach((request, id) => { if (request.key !== key) return projects.delete(id) @@ -170,6 +176,7 @@ function getWorker() { if (event.data.type === "project") { const result = projects.get(event.data.id) + // Only disposal drops a sent request, and the worker forgets that key's projection too. if (!result) { projectTransport.complete(event.data.key, event.data.id) @@ -177,7 +184,9 @@ function getWorker() { } projects.delete(event.data.id) - result.resolve(event.data.projection) + const blocks = applyProjectionBlocks(projected.get(event.data.key), event.data.from, event.data.blocks) + projected.set(event.data.key, blocks) + result.resolve({ text: result.text, blocks }) projectTransport.complete(event.data.key, event.data.id) return @@ -272,6 +281,7 @@ function getWorker() { pending.clear() projects.clear() parses.clear() + projected.clear() states.clear() keys.clear() latest.clear() diff --git a/packages/session-ui/src/components/markdown.tsx b/packages/session-ui/src/components/markdown.tsx index b14492f3e095..a48fd44a0f14 100644 --- a/packages/session-ui/src/components/markdown.tsx +++ b/packages/session-ui/src/components/markdown.tsx @@ -642,6 +642,7 @@ export function Markdown( const completedCode = new Map>() const settledMarkdown = new Map>() let renderedBlocks: RenderedBlock[] = [] + let lastRender: { base: string | undefined; blocks: Block[]; rendered: RenderedBlock[] } | undefined let appliedLabels: CopyLabels | undefined let streamed = false @@ -719,68 +720,86 @@ export function Markdown( const base = src.key ?? (!streamed ? checksum(src.text) : undefined) - return Promise.all( - src.projection.blocks.map(async (block, index) => { - const key = base ? `${base}:${index}:${block.mode}` : undefined - const blockKey = markdownBlockKey(owner, src.key, index, block.mode) - - if (block.mode === "code") { - 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 rendered = { - key: blockKey, - mode: block.mode, - raw: block.raw, - hash: String(block.raw.length), - complete: !!block.complete, - ...result, - } + const render = async (block: Block, index: number): Promise => { + const key = base ? `${base}:${index}:${block.mode}` : undefined + const blockKey = markdownBlockKey(owner, src.key, index, block.mode) - if (block.complete) completedCode.set(blockKey, rendered) + if (block.mode === "code") { + const cached = completedCode.get(blockKey) - return rendered + if (block.complete && cached?.raw === block.raw) return cached + const result = await code(block.src, block.language, blockKey, block.complete) + + const rendered = { + key: blockKey, + mode: block.mode, + raw: block.raw, + hash: String(block.raw.length), + complete: !!block.complete, + ...result, } - if (block.mode === "full") { - const settled = settledMarkdown.get(blockKey) + if (block.complete) completedCode.set(blockKey, rendered) - if (settled && (settled.src ?? settled.raw) === block.src) { - if (key) { - const cached = getCachedMarkdown(key) + return rendered + } - if (cached) touchCachedMarkdown(key, cached) - } + if (block.mode === "full") { + const settled = settledMarkdown.get(blockKey) - return settled - } - } - - const ready = block.mode === "full" ? getReadyMarkdown(block, key) : undefined - const entry = ready ?? (await renderCachedMarkdown(block, key, lifetime.signal)) + if (settled && (settled.src ?? settled.raw) === block.src) { + if (key) { + const cached = getCachedMarkdown(key) - if (block.mode === "full") { - const settled = { - key: blockKey, - mode: block.mode, - ...entry, + if (cached) touchCachedMarkdown(key, cached) } - settledMarkdown.set(blockKey, settled) - return settled } + } + + const ready = block.mode === "full" ? getReadyMarkdown(block, key) : undefined + const entry = ready ?? (await renderCachedMarkdown(block, key, lifetime.signal)) - return { + if (block.mode === "full") { + const settled = { key: blockKey, mode: block.mode, ...entry, } - }), + + settledMarkdown.set(blockKey, settled) + + return settled + } + + return { + key: blockKey, + mode: block.mode, + ...entry, + } + } + + // The worker keeps an unchanged block as the same object, so the block the last render made from it is reused. + // Rendering every block of the answer again on each delta cost a promise and two key strings per block. + const previous = lastRender?.base === base ? lastRender : undefined + + const queued = src.projection.blocks.map((block, index) => + previous?.blocks[index] === block ? previous.rendered[index]! : render(block, index), + ) + + return Promise.all( + queued.flatMap((item, index) => + item instanceof Promise ? [item.then((block) => [index, block] as const)] : [], + ), ) - .then((blocks) => ({ text: src.text, blocks, ready: true }) satisfies RenderResult) + .then((done) => { + const resolved = new Map(done) + const blocks = queued.map((item, index) => (item instanceof Promise ? resolved.get(index)! : item)) + lastRender = { base, blocks: src.projection.blocks, rendered: blocks } + + return { text: src.text, blocks, ready: true } satisfies RenderResult + }) .catch( () => (lifetime.signal.aborted diff --git a/packages/session-ui/src/components/markdown.worker.ts b/packages/session-ui/src/components/markdown.worker.ts index 1d537dc262e6..cc288dad5e04 100644 --- a/packages/session-ui/src/components/markdown.worker.ts +++ b/packages/session-ui/src/components/markdown.worker.ts @@ -3,7 +3,12 @@ import { codeToTokens, isSupportedLanguage, LiveTokenizer, type ThemedToken } from "@pierre/highlights" import { createMarkdownParser } from "@opencode/ui/context/marked-parser" import { OpenCodeTheme } from "@opencode/ui/context/highlight-theme" -import type { MarkdownToken, MarkdownWorkerRequest, MarkdownWorkerResponse } from "./markdown-worker-protocol" +import { + diffProjectionBlocks, + type MarkdownToken, + type MarkdownWorkerRequest, + type MarkdownWorkerResponse, +} from "./markdown-worker-protocol" import { createLatestWorkerQueue } from "./markdown-worker-queue" import { project, type Projection } from "./markdown-stream" @@ -64,9 +69,15 @@ async function parse(request: Extract) async function runProject(request: Extract) { try { - const projection = project(projections.get(request.key), request.text, request.live) + const previous = projections.get(request.key) + const projection = project(previous, request.text, request.live) projections.set(request.key, projection) - post({ type: "project", id: request.id, key: request.key, projection }) + post({ + type: "project", + id: request.id, + key: request.key, + ...diffProjectionBlocks(previous?.blocks, projection.blocks), + }) } catch (error) { post({ type: "error", @@ -164,11 +175,7 @@ function token(value: ThemedToken): MarkdownToken { return [ value.content, - [ - value.color && `color:${value.color}`, - fontStyle & 1 && "font-style:italic", - fontStyle & 2 && "font-weight:bold", - ] + [value.color && `color:${value.color}`, fontStyle & 1 && "font-style:italic", fontStyle & 2 && "font-weight:bold"] .filter(Boolean) .join(";"), ]