diff --git a/.changeset/faster-streamed-markdown.md b/.changeset/faster-streamed-markdown.md new file mode 100644 index 00000000000..f1146b9ec6b --- /dev/null +++ b/.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. diff --git a/packages/kilo-vscode/tests/markdown-incremental-dom.spec.ts b/packages/kilo-vscode/tests/markdown-incremental-dom.spec.ts new file mode 100644 index 00000000000..afc75f77b45 --- /dev/null +++ b/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: "

Heading

", mode: "full" as const } + const tail = { key: "1", hash: "tail-a", html: "

Tail A

", 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: "

Tail B

" }], 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: "

Heading

", mode: "full" as const } + const paragraph = { key: "1", hash: "paragraph", html: "

Stable

", 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: "", 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: "

Stable

", mode: "full" as const } + const tail = { key: "1", hash: "tail", html: "

Tail

", 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: "

First

", mode: "live" as const } + const second = { key: "1", hash: "second", html: "

Second

", 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: "

Stable

", mode: "full" as const }, + { key: "1", hash: "middle", html: "

Middle

", mode: "full" as const }, + { key: "2", hash: "tail", html: "

Tail

", 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" }) +}) diff --git a/packages/ui/src/components/markdown-stream.ts b/packages/ui/src/components/markdown-stream.ts index ea35b0c140d..ae034275c70 100644 --- a/packages/ui/src/components/markdown-stream.ts +++ b/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] diff --git a/packages/ui/src/components/markdown.tsx b/packages/ui/src/components/markdown.tsx index cee1ef3bee4..813bc924108 100644 --- a/packages/ui/src/components/markdown.tsx +++ b/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() @@ -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 => { // 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(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,6 +392,7 @@ export function Markdown( pendingContent = undefined pendingLabels = undefined } + incremental.reset() // kilocode_change copyCleanup = fast.copyCleanup kickMermaid(container, local.streaming ?? false, mermaid) kickHighlight(container, labels) @@ -378,6 +400,9 @@ export function Markdown( } // 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. diff --git a/packages/ui/src/kilocode/markdown-incremental-dom.ts b/packages/ui/src/kilocode/markdown-incremental-dom.ts new file mode 100644 index 00000000000..bd44cbf2763 --- /dev/null +++ b/packages/ui/src/kilocode/markdown-incremental-dom.ts @@ -0,0 +1,128 @@ +type Labels = { + copy: string + copied: string +} + +export type MarkdownBlock = { + key: string + hash: string + html: string + mode: "full" | "live" +} + +type Record = { + key: string + hash: string + start: Comment + end: Comment +} + +type Decorate = (root: HTMLDivElement, labels: Labels) => void + +type Hooks = { + cancel: () => void + ready: (container: HTMLDivElement, labels: Labels, context: Context) => void +} + +export function createIncrementalMarkdown(decorate: Decorate, hooks?: Hooks) { + let records: Record[] = [] + + const reset = () => { + records = [] + } + + const parse = (html: string, labels: Labels) => { + const root = document.createElement("div") + root.innerHTML = html + decorate(root, labels) + const fragment = document.createDocumentFragment() + while (root.firstChild) fragment.appendChild(root.firstChild) + return fragment + } + + const remove = (record: Record) => { + const parent = record.start.parentNode + if (!parent || record.end.parentNode !== parent) return false + + const nodes: ChildNode[] = [] + let node: ChildNode | null = record.start + while (node && node.parentNode === parent) { + nodes.push(node) + if (node === record.end) { + for (const item of nodes) parent.removeChild(item) + return true + } + node = node.nextSibling + } + return false + } + + const replace = (record: Record, block: MarkdownBlock, labels: Labels) => { + let node = record.start.nextSibling + while (node && node !== record.end) { + const next: ChildNode | null = node.nextSibling + node.parentNode?.removeChild(node) + node = next + } + record.end.parentNode?.insertBefore(parse(block.html, labels), record.end) + record.hash = block.hash + } + + const append = (container: HTMLDivElement, block: MarkdownBlock, labels: Labels) => { + // Comment boundaries preserve the exact Markdown child structure, so existing + // direct-child CSS and copy/highlight behavior do not need wrapper exceptions. + const start = document.createComment(`markdown:${block.key}:start`) + const end = document.createComment(`markdown:${block.key}:end`) + container.appendChild(start) + container.appendChild(parse(block.html, labels)) + container.appendChild(end) + records.push({ key: block.key, hash: block.hash, start, end }) + } + + const update = (container: HTMLDivElement, blocks: MarkdownBlock[], labels: Labels) => { + if (blocks.length < 2) return false + if (blocks.slice(0, -1).some((block) => block.mode !== "full")) return false + + if (records.some((record) => !record.start.isConnected || !record.end.isConnected)) reset() + const shared = Math.min(records.length, blocks.length) + if (records.slice(0, shared).some((record, index) => record.key !== blocks[index]?.key)) reset() + + if (records.length === 0) container.replaceChildren() + + while (records.length > blocks.length) { + const record = records.at(-1) + if (!record || !remove(record)) { + reset() + return false + } + records.pop() + } + + for (let index = 0; index < blocks.length; index++) { + const block = blocks[index]! + const record = records[index] + if (!record) { + append(container, block, labels) + continue + } + if (record.hash === block.hash) continue + replace(record, block, labels) + } + return true + } + + const render = ( + streaming: boolean, + container: HTMLDivElement, + blocks: MarkdownBlock[], + labels: Labels, + context: Context, + ) => { + if (!streaming || !update(container, blocks, labels)) return false + hooks?.cancel() + hooks?.ready(container, labels, context) + return true + } + + return { reset, render, update } +} diff --git a/packages/ui/src/kilocode/markdown-stable-blocks.test.ts b/packages/ui/src/kilocode/markdown-stable-blocks.test.ts new file mode 100644 index 00000000000..3dcf7592477 --- /dev/null +++ b/packages/ui/src/kilocode/markdown-stable-blocks.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, test } from "bun:test" +import { marked } from "marked" +import remend from "remend" +import { stream } from "../components/markdown-stream" +import { stableBlocks } from "./markdown-stable-blocks" + +async function render(text: string) { + const html = await Promise.all(stream(text, true).map((block) => Promise.resolve(marked.parse(block.src)))) + return html.join("") +} + +describe("stable markdown blocks", () => { + test("keeps completed top-level tokens stable and only heals the tail", () => { + expect( + stableBlocks( + [ + { type: "heading", raw: "# Title\n\n" }, + { type: "paragraph", raw: "First" }, + { type: "space", raw: "\n\n" }, + { type: "paragraph", raw: "Second **open" }, + ], + (raw) => `${raw}**`, + ), + ).toEqual([ + { raw: "# Title\n\n", src: "# Title\n\n", mode: "full" }, + { raw: "First\n\n", src: "First\n\n", mode: "full" }, + { raw: "Second **open", src: "Second **open**", mode: "live" }, + ]) + }) + + test("leaves a single mutable token on the existing streaming path", () => { + expect(stableBlocks([{ type: "paragraph", raw: "Still streaming" }], (raw) => raw)).toBeUndefined() + }) + + test("matches canonical streaming HTML for mixed completed blocks", async () => { + const text = [ + "# Report", + "", + "A completed paragraph with **emphasis**.", + "", + "- first item", + "- second item", + "", + "```ts", + "export const value = 1", + "```", + "", + "The final paragraph is *still streaming", + ].join("\n") + + expect(await render(text)).toBe(await marked.parse(remend(text, { linkMode: "text-only" }))) + expect(stream(text, true).map((block) => block.mode)).toEqual(["full", "full", "full", "full", "live"]) + }) +}) diff --git a/packages/ui/src/kilocode/markdown-stable-blocks.ts b/packages/ui/src/kilocode/markdown-stable-blocks.ts new file mode 100644 index 00000000000..1878f9787de --- /dev/null +++ b/packages/ui/src/kilocode/markdown-stable-blocks.ts @@ -0,0 +1,25 @@ +import type { Block } from "../components/markdown-stream" + +type Token = { + type: string + raw: string +} + +export function stableBlocks(tokens: Token[], live: (raw: string) => string): Block[] | undefined { + const indexes = tokens.flatMap((token, index) => (token.type === "space" ? [] : [index])) + if (indexes.length < 2) return + + const raw = (start: number, end = tokens.length) => + tokens + .slice(start, end) + .map((token) => token.raw) + .join("") + // Completed top-level tokens keep stable hashes across stream updates. The + // existing parse and sanitize cache can then reuse them while only the tail changes. + const stable = indexes.slice(0, -1).map((start, index) => { + const value = raw(start, indexes[index + 1]) + return { raw: value, src: value, mode: "full" as const } + }) + const tail = raw(indexes.at(-1)!) + return [...stable, { raw: tail, src: live(tail), mode: "live" }] +}