Merge pull request #10486 from Kilo-Org/arrow-abstract

fix: stabilize streamed markdown code blocks

.changeset/steady-code-streams.md

@@ -0,0 +1,7 @@
+---
+"kilo-code": patch
+"@kilocode/kilo-ui": patch
+"@opencode-ai/ui": patch
+---
+
+Keep streamed Markdown code blocks stable while assistant output is still arriving.

packages/kilo-ui/src/components/message-part.tsx

@@ -1377,7 +1377,12 @@ PART_MAPPING["text"] = function TextPartDisplay(props) {
     <Show when={throttledText() && showSyntheticPart()}>
       <div data-component="text-part">
         <div data-slot="text-part-body">
-          <Markdown text={throttledText()} cacheKey={part().id} onClick={handleMarkdownClick} />
+          <Markdown
+            text={throttledText()}
+            cacheKey={part().id}
+            streaming={streaming()}
+            onClick={handleMarkdownClick}
+          />
         </div>
         <Show when={showCopy()}>
           <div data-slot="assistant-copy-wrapper">
@@ -1575,7 +1580,7 @@ PART_MAPPING["reasoning"] = function ReasoningPartDisplay(props: MessagePartProp
           </Collapsible.Trigger>
           <Collapsible.Content>
             <div data-slot="reasoning-content" ref={ref} onScroll={onScroll} onWheel={onWheel}>
-              <Markdown text={view().body} cacheKey={id} />
+              <Markdown text={view().body} cacheKey={id} streaming={!done()} />
             </div>
           </Collapsible.Content>
         </Collapsible>

packages/kilo-vscode/tests/unit/kilo-ui-contract.test.ts

@@ -137,6 +137,19 @@ describe("DataProvider contract (runtime)", () => {
   })
 })
 
+describe("Assistant Markdown streaming contract (source)", () => {
+  const src = fs.readFileSync(KILO_MESSAGE_PART_FILE, "utf-8")
+  const block =
+    src.match(
+      /PART_MAPPING\["text"\]\s*=\s*function TextPartDisplay[\s\S]*?(?=\/\/ Expanded mode|PART_MAPPING\["reasoning"\])/,
+    )?.[0] ?? ""
+
+  it("passes active text streams through Markdown's streaming mode", () => {
+    expect(block).not.toBe("")
+    expect(block).toContain("streaming={streaming()}")
+  })
+})
+
 describe("Edit tool diff-first click contract (source)", () => {
   const src = fs.readFileSync(KILO_MESSAGE_PART_FILE, "utf-8")
 

packages/kilo-vscode/tests/unit/markdown-raf-coalesce.test.ts

@@ -24,6 +24,7 @@ import { join } from "node:path"
  */
 describe("Markdown rAF-coalesced parse — regression guard", () => {
   const path = join(__dirname, "..", "..", "..", "ui", "src", "components", "markdown.tsx")
+  const helper = join(__dirname, "..", "..", "..", "ui", "src", "kilocode", "markdown-stream-highlight.ts")
 
   const stripComments = (src: string): string =>
     src
@@ -32,6 +33,7 @@ describe("Markdown rAF-coalesced parse — regression guard", () => {
       .replace(/^\s*\/\/.*$/gm, "")
 
   const src = stripComments(readFileSync(path, "utf8"))
+  const body = stripComments(readFileSync(helper, "utf8"))
 
   it("render effect uses requestAnimationFrame to coalesce parses", () => {
     // Locate the createEffect that owns the morphdom call.
@@ -50,4 +52,11 @@ describe("Markdown rAF-coalesced parse — regression guard", () => {
     // rapid updates can collapse into it.
     expect(src).toMatch(/\b(pendingFrame|pendingContent)\b/)
   })
+
+  it("delegates streamed Shiki refreshes to the Kilo-owned helper", () => {
+    expect(src).toContain("preserveStreamingHighlight(fromEl, toEl, local.streaming ?? false)")
+    expect(body).toContain("export function preserveStreamingHighlight")
+    expect(body).toMatch(/continues\(before, after\)[\s\S]*queue\(from, after, lang\)/)
+    expect(body).toMatch(/const done = \(\) => \{\s*job\.busy = false\s*if \(!pre\.isConnected\) return/)
+  })
 })

packages/ui/src/components/markdown.tsx

@@ -8,6 +8,7 @@ import { isServer } from "solid-js/web"
 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
 
 type Entry = {
   hash: string
@@ -436,6 +437,7 @@ export function Markdown(
             const fromHash = fromEl.getAttribute("data-source-hash")
             const toCode = toEl.querySelector("code")?.textContent ?? ""
             if (fromHash === fnv1a(toCode)) return false
+            if (preserveStreamingHighlight(fromEl, toEl, local.streaming ?? false)) return false // kilocode_change
             // Source changed during streaming — fall through so morphdom replaces // kilocode_change
             // the stale highlighted block with the updated plain block, which will
             // be re-highlighted on the next deferredHighlight pass.

packages/ui/src/kilocode/markdown-stream-highlight.ts

@@ -0,0 +1,93 @@
+import { getSharedHighlighter } from "@pierre/diffs"
+import { bundledLanguages, type BundledLanguage } from "shiki"
+import { fnv1a } from "../context/marked"
+
+type Job = {
+  code: string
+  lang: string
+  busy: boolean
+}
+
+const jobs = new WeakMap<HTMLPreElement, Job>()
+
+function continues(before: string, after: string) {
+  const base = before.endsWith("\n") ? before.slice(0, -1) : before
+  return !!base && after.startsWith(base)
+}
+
+async function source(lang: string, code: string) {
+  try {
+    const highlighter = await getSharedHighlighter({ themes: ["Kilo"], langs: [] })
+    const language = lang in bundledLanguages ? lang : "text"
+    if (!highlighter.getLoadedLanguages().includes(language)) {
+      await highlighter.loadLanguage(language as BundledLanguage)
+    }
+    return highlighter.codeToHtml(code, { lang: language, theme: "Kilo", tabindex: false })
+  } catch (err) {
+    console.warn("Streaming code highlight failed", lang, err)
+    return
+  }
+}
+
+function update(pre: HTMLPreElement, html: string, code: string) {
+  if (!pre.isConnected) return
+  const temp = document.createElement("div")
+  temp.innerHTML = html
+  const next = temp.firstElementChild
+  if (!(next instanceof HTMLPreElement)) return
+  const x = pre.scrollLeft
+  for (const name of pre.getAttributeNames()) {
+    pre.removeAttribute(name)
+  }
+  for (const attr of next.attributes) {
+    pre.setAttribute(attr.name, attr.value)
+  }
+  pre.setAttribute("data-source-hash", fnv1a(code))
+  pre.replaceChildren(...Array.from(next.childNodes))
+  pre.scrollLeft = x
+}
+
+async function refresh(pre: HTMLPreElement, code: string, lang: string) {
+  if (!pre.isConnected || !code) return
+  await new Promise<void>((resolve) => setTimeout(resolve, 0))
+  if (!pre.isConnected) return
+  const html = await source(lang, code)
+  if (!html || !pre.isConnected) return
+  update(pre, html, code)
+}
+
+function run(pre: HTMLPreElement, job: Job) {
+  const code = job.code
+  const lang = job.lang
+  job.busy = true
+  const done = () => {
+    job.busy = false
+    if (!pre.isConnected) return
+    if (job.code !== code || job.lang !== lang) {
+      run(pre, job)
+      return
+    }
+  }
+  void refresh(pre, code, lang).then(done, done)
+}
+
+function queue(pre: HTMLPreElement, code: string, lang: string) {
+  const job = jobs.get(pre) ?? { code, lang, busy: false }
+  job.code = code
+  job.lang = lang
+  jobs.set(pre, job)
+  if (job.busy) return
+  run(pre, job)
+}
+
+export function preserveStreamingHighlight(from: Element, to: Element, streaming: boolean) {
+  if (!streaming) return false
+  if (!(from instanceof HTMLPreElement) || !(to instanceof HTMLPreElement)) return false
+  if (!from.classList.contains("shiki") || to.classList.contains("shiki")) return false
+  const before = from.querySelector("code")?.textContent ?? ""
+  const after = to.querySelector("code")?.textContent ?? ""
+  const lang = to.querySelector("code")?.getAttribute("data-lang") || "text"
+  if (!after || lang === "mermaid" || !continues(before, after)) return false
+  queue(from, after, lang)
+  return true
+}