fix(session-ui): send and render only the markdown blocks a delta changed

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")

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}`}`
 }

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<number, ProjectPending>()
 
 const parses = new Map<number, ParsePending>()
 
+// The blocks of the last projection the worker sent for each key, which its next response builds on.
+const projected = new Map<string, Block[]>()
+
 const states = new Map<string, MarkdownWorkerState>()
 
 const keys = new Set<string>()
@@ -91,13 +96,14 @@ export function projectMarkdown(key: string, text: string, live: boolean) {
   const id = ++nextID
 
   return new Promise<Projection>((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,14 +176,17 @@ 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)
 
         return
       }
 
       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()

packages/session-ui/src/components/markdown.tsx

@@ -642,6 +642,7 @@ export function Markdown(
   const completedCode = new Map<string, Extract<RenderedBlock, { mode: "code" }>>()
   const settledMarkdown = new Map<string, Exclude<RenderedBlock, { mode: "code" }>>()
   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<RenderedBlock> => {
+        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

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<MarkdownWorkerRequest, { type: "parse" }>)
 
 async function runProject(request: Extract<MarkdownWorkerRequest, { type: "project" }>) {
   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(";"),
   ]