perf(ui): render streamed markdown incrementally

.changeset/faster-streamed-markdown.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Speed up long streamed responses by preserving completed Markdown blocks and updating only the active tail.

packages/kilo-vscode/tests/markdown-incremental-dom.spec.ts

@@ -0,0 +1,208 @@
+import { expect, test } from "@playwright/test"
+import { build } from "esbuild"
+import { fileURLToPath } from "node:url"
+
+const source = fileURLToPath(new URL("../../ui/src/kilocode/markdown-incremental-dom.ts", import.meta.url))
+const bundle = await build({
+  stdin: {
+    contents: `
+      import { createIncrementalMarkdown } from ${JSON.stringify(source)}
+      globalThis.createIncrementalMarkdown = createIncrementalMarkdown
+    `,
+    resolveDir: fileURLToPath(new URL(".", import.meta.url)),
+  },
+  bundle: true,
+  format: "iife",
+  platform: "browser",
+  write: false,
+})
+const script = bundle.outputFiles[0]!.text
+
+test.beforeEach(async ({ page }) => {
+  await page.goto("about:blank")
+  await page.addScriptTag({ content: script })
+})
+
+test("keeps completed Markdown nodes while replacing only the live tail", async ({ page }) => {
+  const result = await page.evaluate(() => {
+    const create = (
+      globalThis as typeof globalThis & {
+        createIncrementalMarkdown: (decorate: () => void) => {
+          update: (
+            container: HTMLDivElement,
+            blocks: Array<{ key: string; hash: string; html: string; mode: "full" | "live" }>,
+            labels: { copy: string; copied: string },
+          ) => boolean
+        }
+      }
+    ).createIncrementalMarkdown
+    const labels = { copy: "Copy", copied: "Copied" }
+    const container = document.createElement("div")
+    document.body.append(container)
+    const renderer = create(() => {})
+    const heading = { key: "0", hash: "heading", html: "<h2>Heading</h2>", mode: "full" as const }
+    const tail = { key: "1", hash: "tail-a", html: "<p>Tail A</p>", mode: "live" as const }
+
+    renderer.update(container, [heading, tail], labels)
+    const stable = container.children[0]
+    const previous = container.children[1]
+    renderer.update(container, [heading, { ...tail, hash: "tail-b", html: "<p>Tail B</p>" }], labels)
+
+    return {
+      stable: container.children[0] === stable,
+      replaced: container.children[1] !== previous,
+      tags: Array.from(container.children).map((child) => child.tagName),
+      text: container.textContent,
+    }
+  })
+
+  expect(result).toEqual({ stable: true, replaced: true, tags: ["H2", "P"], text: "HeadingTail B" })
+})
+
+test("promotes an unchanged tail and appends the next block without wrappers", async ({ page }) => {
+  const result = await page.evaluate(() => {
+    const create = (
+      globalThis as typeof globalThis & {
+        createIncrementalMarkdown: (decorate: () => void) => {
+          update: (
+            container: HTMLDivElement,
+            blocks: Array<{ key: string; hash: string; html: string; mode: "full" | "live" }>,
+            labels: { copy: string; copied: string },
+          ) => boolean
+        }
+      }
+    ).createIncrementalMarkdown
+    const labels = { copy: "Copy", copied: "Copied" }
+    const container = document.createElement("div")
+    document.body.append(container)
+    const renderer = create(() => {})
+    const heading = { key: "0", hash: "heading", html: "<h2>Heading</h2>", mode: "full" as const }
+    const paragraph = { key: "1", hash: "paragraph", html: "<p>Stable</p>", mode: "live" as const }
+
+    renderer.update(container, [heading, paragraph], labels)
+    const stable = container.children[1]
+    renderer.update(
+      container,
+      [
+        heading,
+        { ...paragraph, mode: "full" },
+        { key: "2", hash: "tail", html: "<ul><li>Next</li></ul>", mode: "live" as const },
+      ],
+      labels,
+    )
+
+    return {
+      stable: container.children[1] === stable,
+      tags: Array.from(container.children).map((child) => child.tagName),
+    }
+  })
+
+  expect(result).toEqual({ stable: true, tags: ["H2", "P", "UL"] })
+})
+
+test("runs streaming hooks only when incremental rendering handles the update", async ({ page }) => {
+  const result = await page.evaluate(() => {
+    const create = (
+      globalThis as typeof globalThis & {
+        createIncrementalMarkdown: (
+          decorate: () => void,
+          hooks: {
+            cancel: () => void
+            ready: (container: HTMLDivElement, labels: { copy: string; copied: string }, context: string) => void
+          },
+        ) => {
+          render: (
+            streaming: boolean,
+            container: HTMLDivElement,
+            blocks: Array<{ key: string; hash: string; html: string; mode: "full" | "live" }>,
+            labels: { copy: string; copied: string },
+            context: string,
+          ) => boolean
+        }
+      }
+    ).createIncrementalMarkdown
+    const labels = { copy: "Copy", copied: "Copied" }
+    const container = document.createElement("div")
+    document.body.append(container)
+    const calls: string[] = []
+    const renderer = create(() => {}, {
+      cancel: () => calls.push("cancel"),
+      ready: (_container, _labels, context) => calls.push(context),
+    })
+    const stable = { key: "0", hash: "stable", html: "<p>Stable</p>", mode: "full" as const }
+    const tail = { key: "1", hash: "tail", html: "<p>Tail</p>", mode: "live" as const }
+
+    const completed = renderer.render(false, container, [stable, tail], labels, "ready")
+    const unsupported = renderer.render(true, container, [tail], labels, "unsupported")
+    const streaming = renderer.render(true, container, [stable, tail], labels, "ready")
+    return { calls, completed, streaming, unsupported }
+  })
+
+  expect(result).toEqual({ calls: ["cancel", "ready"], completed: false, streaming: true, unsupported: false })
+})
+
+test("falls back when there is no stable prefix", async ({ page }) => {
+  const result = await page.evaluate(() => {
+    const create = (
+      globalThis as typeof globalThis & {
+        createIncrementalMarkdown: (decorate: () => void) => {
+          update: (
+            container: HTMLDivElement,
+            blocks: Array<{ key: string; hash: string; html: string; mode: "full" | "live" }>,
+            labels: { copy: string; copied: string },
+          ) => boolean
+        }
+      }
+    ).createIncrementalMarkdown
+    const labels = { copy: "Copy", copied: "Copied" }
+    const container = document.createElement("div")
+    const renderer = create(() => {})
+    const first = { key: "0", hash: "first", html: "<p>First</p>", mode: "live" as const }
+    const second = { key: "1", hash: "second", html: "<p>Second</p>", mode: "live" as const }
+
+    return {
+      multiple: renderer.update(container, [first, second], labels),
+      single: renderer.update(container, [first], labels),
+    }
+  })
+
+  expect(result).toEqual({ multiple: false, single: false })
+})
+
+test("rebuilds safely when a Markdown boundary disappears", async ({ page }) => {
+  const result = await page.evaluate(() => {
+    const create = (
+      globalThis as typeof globalThis & {
+        createIncrementalMarkdown: (decorate: () => void) => {
+          update: (
+            container: HTMLDivElement,
+            blocks: Array<{ key: string; hash: string; html: string; mode: "full" | "live" }>,
+            labels: { copy: string; copied: string },
+          ) => boolean
+        }
+      }
+    ).createIncrementalMarkdown
+    const labels = { copy: "Copy", copied: "Copied" }
+    const container = document.createElement("div")
+    document.body.append(container)
+    const renderer = create(() => {})
+    const blocks = [
+      { key: "0", hash: "stable", html: "<p>Stable</p>", mode: "full" as const },
+      { key: "1", hash: "middle", html: "<p>Middle</p>", mode: "full" as const },
+      { key: "2", hash: "tail", html: "<p>Tail</p>", mode: "live" as const },
+    ]
+
+    renderer.update(container, blocks, labels)
+    const comments = Array.from(container.childNodes).filter((node) => node.nodeType === Node.COMMENT_NODE)
+    comments.at(-1)?.remove()
+    const updated = renderer.update(container, [blocks[0]!, { ...blocks[1]!, mode: "live" }], labels)
+
+    return {
+      updated,
+      tags: Array.from(container.children).map((child) => child.tagName),
+      text: container.textContent,
+    }
+  })
+
+  expect(result).toEqual({ updated: true, tags: ["P", "P"], text: "StableMiddle" })
+})

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

@@ -1,5 +1,6 @@
 import { marked, type Tokens } from "marked"
 import remend from "remend"
+import { stableBlocks } from "../kilocode/markdown-stable-blocks" // kilocode_change
 
 export type Block = {
   raw: string
@@ -31,6 +32,9 @@ export function stream(text: string, live: boolean) {
   const src = heal(text)
   if (refs(text)) return [{ raw: text, src, mode: "live" }] satisfies Block[]
   const tokens = marked.lexer(text)
+  const candidate = tokens.findLast((token) => token.type !== "space") // kilocode_change
+  const blocks = candidate && !open(candidate.raw) ? stableBlocks(tokens, heal) : undefined // kilocode_change
+  if (blocks) return blocks // kilocode_change
   const tail = tokens.findLastIndex((token) => token.type !== "space")
   if (tail < 0) return [{ raw: text, src, mode: "live" }] satisfies Block[]
   const last = tokens[tail]

packages/ui/src/components/markdown.tsx

@@ -9,12 +9,15 @@ import { stream } from "./markdown-stream"
 import { tryFastRender } from "../kilocode/markdown-fast-path" // kilocode_change
 import { hasMermaid, preserveMermaid, renderMermaid, type MermaidLabels } from "../kilocode/markdown-mermaid" // kilocode_change
 import { preserveStreamingHighlight } from "../kilocode/markdown-stream-highlight" // kilocode_change
+import { createIncrementalMarkdown, type MarkdownBlock } from "../kilocode/markdown-incremental-dom" // kilocode_change
 
 type Entry = {
   hash: string
   html: string
 }
 
+type Rendered = { content: string; blocks: MarkdownBlock[] } // kilocode_change
+
 const max = 200
 const cache = new Map<string, Entry>()
 
@@ -260,34 +263,34 @@ export function Markdown(
       key: local.cacheKey,
       streaming: local.streaming ?? false,
     }),
-    async (src) => {
-      if (isServer) return fallback(src.text)
-      if (!src.text) return ""
+    async (src): Promise<Rendered> => { // kilocode_change
+      if (isServer) return { content: fallback(src.text), blocks: [] } // kilocode_change
+      if (!src.text) return { content: "", blocks: [] } // kilocode_change
 
       const base = src.key ?? checksum(src.text)
       return Promise.all(
         stream(src.text, src.streaming).map(async (block, index) => {
-          const hash = checksum(block.raw)
+          const hash = checksum(block.raw) ?? "" // kilocode_change
           const key = base ? `${base}:${index}:${block.mode}` : hash
 
           if (key && hash) {
             const cached = cache.get(key)
             if (cached && cached.hash === hash) {
               touch(key, cached)
-              return cached.html
+              return { key: `${base}:${index}`, hash, html: cached.html, mode: block.mode } // kilocode_change
             }
           }
 
           const next = await Promise.resolve(marked.parse(block.src))
           const safe = sanitize(next)
           if (key && hash) touch(key, { hash, html: safe })
-          return safe
+          return { key: `${base}:${index}`, hash, html: safe, mode: block.mode } // kilocode_change
         }),
       )
-        .then((list) => list.join(""))
-        .catch(() => fallback(src.text))
+        .then((blocks) => ({ content: blocks.map((block) => block.html).join(""), blocks })) // kilocode_change
+        .catch(() => ({ content: fallback(src.text), blocks: [] })) // kilocode_change
     },
-    { initialValue: fallback(local.text) },
+    { initialValue: { content: fallback(local.text), blocks: [] } }, // kilocode_change
   )
 
   let copyCleanup: (() => void) | undefined
@@ -313,10 +316,27 @@ export function Markdown(
   let pendingContent: string | undefined
   let pendingLabels: { copy: string; copied: string } | undefined
   // kilocode_change end
+  // kilocode_change start
+  const incremental = createIncrementalMarkdown<MermaidLabels>(decorate, {
+    cancel: () => {
+      if (pendingFrame === undefined) return
+      cancelAnimationFrame(pendingFrame)
+      pendingFrame = undefined
+      pendingContent = undefined
+      pendingLabels = undefined
+    },
+    ready: (container, labels, mermaid) => {
+      copyCleanup ??= setupCodeCopy(container, () => labels)
+      kickMermaid(container, true, mermaid)
+      kickHighlight(container, labels)
+    },
+  })
+  // kilocode_change end
 
   createEffect(() => {
     const container = root()
-    const content = local.text ? (html.latest ?? html() ?? "") : ""
+    const rendered = html.latest ?? html() ?? { content: "", blocks: [] } // kilocode_change
+    const content = local.text ? rendered.content : "" // kilocode_change
     if (!container) return
     if (isServer) return
 
@@ -330,6 +350,7 @@ export function Markdown(
         pendingLabels = undefined
       }
       // kilocode_change end
+      incremental.reset() // kilocode_change
       container.innerHTML = ""
       // kilocode_change start: Mermaid diagram rendering
       mermaidState.signal.aborted = true
@@ -371,13 +392,17 @@ export function Markdown(
         pendingContent = undefined
         pendingLabels = undefined
       }
+      incremental.reset() // kilocode_change
       copyCleanup = fast.copyCleanup
       kickMermaid(container, local.streaming ?? false, mermaid)
       kickHighlight(container, labels)
       return
     }
     // kilocode_change end
 
+    if (incremental.render(local.streaming ?? false, container, rendered.blocks, labels, mermaid)) return // kilocode_change
+    incremental.reset() // kilocode_change
+
     // kilocode_change start: queue the latest content for a single rAF tick.
     // Further updates before the frame runs simply overwrite pendingContent,
     // so K rapid updates collapse to 1 parse instead of K.