fix(tui): keep write previews from flashing unstyled

# ------------------------ >8 ------------------------
# Do not modify or remove the line above.
# Everything below it will be ignored.
#
# Conflicts:
#	packages/tui/src/routes/session/index.tsx

packages/tui/src/plugin/context.tsx

@@ -1,6 +1,12 @@
 import type { PluginInfo } from "@opencode-ai/client"
 import type { Plugin } from "@opencode-ai/plugin/tui"
-import { createMarkdownCodeBlockRenderer, type MarkdownCodeBlockRenderer, type MarkdownOptions } from "@opentui/core"
+import {
+  CodeRenderable,
+  createMarkdownCodeBlockRenderer,
+  type MarkdownCodeBlockRenderer,
+  type MarkdownOptions,
+  type Renderable,
+} from "@opentui/core"
 import {
   batch,
   createContext,
@@ -88,8 +94,26 @@ export function combineMarkdownRenderers(
   for (const source of sources) {
     for (const [language, render] of Object.entries(source)) renderers.set(language, render)
   }
-  if (renderers.size === 0) return undefined
-  return createMarkdownCodeBlockRenderer(renderers)
+  const render = createMarkdownCodeBlockRenderer(renderers)
+  return (token, context) => {
+    if (token.type !== "code" && token.type !== "list") return render?.(token, context)
+    let fallback: Renderable | null | undefined
+    const custom = render?.(token, {
+      ...context,
+      defaultRender: () => (fallback ??= context.defaultRender()),
+    })
+    if (custom && custom !== fallback) return custom
+    const node = fallback ?? context.defaultRender()
+    if (!node) return undefined
+    deferCodePaint(node)
+    return node
+  }
+}
+
+function deferCodePaint(node: Renderable) {
+  // Streaming nodes already retain styled text; only completed fences opt into raw first paint.
+  if (node instanceof CodeRenderable && !node.streaming) node.drawUnstyledText = false
+  node.getChildren().forEach(deferCodePaint)
 }
 
 export function PluginProvider(props: ParentProps<{ packages: PackageResolver; directories: string[] }>) {

packages/tui/src/routes/session/index.tsx

@@ -114,6 +114,7 @@ import { SessionLocationMissing } from "./location-missing"
 import { isRecord } from "../../util/record"
 import { createHistoryPrepend } from "./history"
 import { useSessionTerminals } from "../../context/session-terminals"
+import { installSyntaxHighlightCache } from "../../util/syntax-highlight-cache"
 
 addDefaultParsers(parsers.parsers)
 
@@ -161,6 +162,7 @@ export function Session(props: {
   visibleTerminalID?: string
   width?: number
 }) {
+  installSyntaxHighlightCache()
   const setEpilogue = useEpilogue()
   const clipboard = useClipboard()
   const writeExport = async (file: string, content: string) => {
@@ -3349,6 +3351,7 @@ function Write(props: ToolProps) {
         >
           <line_number fg={theme.text.subdued} minWidth={3} paddingRight={1}>
             <code
+              drawUnstyledText={false}
               conceal={false}
               fg={theme.text.default}
               filetype={filetype(stringValue(props.input.path))}

packages/tui/src/util/syntax-highlight-cache.ts

@@ -0,0 +1,38 @@
+import { getTreeSitterClient, type TreeSitterClient } from "@opentui/core"
+
+const CACHE_SIZE = 500
+const installed = new WeakSet<TreeSitterClient>()
+
+export function installSyntaxHighlightCache() {
+  const client = getTreeSitterClient()
+  if (installed.has(client)) return
+  installed.add(client)
+  client.highlightOnce = cacheHighlights(client.highlightOnce.bind(client))
+}
+
+export function cacheHighlights(highlight: TreeSitterClient["highlightOnce"], capacity = CACHE_SIZE) {
+  const cache = new Map<string, ReturnType<TreeSitterClient["highlightOnce"]>>()
+
+  return (content: string, filetype: string) => {
+    const key = `${filetype}\0${content}`
+    const cached = cache.get(key)
+    if (cached) {
+      cache.delete(key)
+      cache.set(key, cached)
+      return cached
+    }
+
+    const result = highlight(content, filetype)
+    cache.set(key, result)
+    if (cache.size > capacity) cache.delete(cache.keys().next().value!)
+
+    void result
+      .then((value) => {
+        if (value.error && cache.get(key) === result) cache.delete(key)
+      })
+      .catch(() => {
+        if (cache.get(key) === result) cache.delete(key)
+      })
+    return result
+  }
+}

packages/tui/test/app-lifecycle.test.tsx

@@ -1193,6 +1193,85 @@ test.each(["manual", "select"] as const)(
   },
 )
 
+test.each([80, 120].flatMap((width) => ["fence", "list", "write"].map((kind) => ({ width, kind }))))(
+  "session code is highlighted on open and tab return: %j",
+  async (input) => {
+    await using state = await tmpdir()
+    const session = {
+      id: "ses_highlight",
+      title: "Highlight fixture",
+      projectID: "project",
+      location: { directory },
+      agent: "build",
+      model: { providerID: "fixture", id: "model" },
+      cost: 0,
+      tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+      time: { created: 1, updated: 2 },
+    }
+    await using setup = await createAppFixture({
+      width: input.width,
+      state: state.path,
+      config: { animations: false, tabs: { enabled: true }, session: { sidebar: "hide" } },
+      args: { sessionID: session.id },
+      fetch: (url) => {
+        if (url.pathname === `/api/session/${session.id}`) return json({ data: session })
+        if (url.pathname === `/api/session/${session.id}/message`)
+          return json({
+            data: [
+              {
+                id: "msg_highlight",
+                type: "assistant",
+                agent: session.agent,
+                model: session.model,
+                time: { created: 2, completed: 3 },
+                content: [
+                  input.kind === "write"
+                    ? {
+                        type: "tool",
+                        id: "call_highlight",
+                        name: "write",
+                        time: { created: 2, completed: 3 },
+                        state: {
+                          status: "completed",
+                          input: { path: `${directory}/snippet.ts`, content: "const flashFixture: number = 42" },
+                          content: [{ type: "text", text: "Wrote file" }],
+                        },
+                      }
+                    : {
+                        type: "text",
+                        text:
+                          input.kind === "list"
+                            ? "- Example:\n\n  ```typescript\n  const flashFixture: number = 42\n  ```"
+                            : "```typescript\nconst flashFixture: number = 42\n```",
+                      },
+                ],
+              },
+            ],
+            cursor: {},
+          })
+        if (url.pathname === `/api/session/${session.id}/inbox`) return json({ data: [] })
+        if (url.pathname === `/api/session/${session.id}/permission`) return json({ data: [] })
+      },
+    })
+    await setup.ready
+    const line = () =>
+      setup.captureSpans().lines.find((line) => line.spans.some((span) => span.text.includes("flashFixture")))
+    await setup.waitForFrame((frame) => frame.includes("flashFixture"))
+    const first = line()
+    await setup.waitForFrame(() => (line()?.spans.filter((span) => span.text.trim()).length ?? 0) > 4)
+    const highlighted = line()
+    setup.mockInput.pressKey("x", { ctrl: true })
+    setup.mockInput.pressKey("n")
+    await setup.waitForFrame((frame) => !frame.includes("flashFixture"))
+    setup.mockInput.pressKey("1", { ctrl: true })
+    await setup.waitForFrame((frame) => frame.includes("flashFixture"))
+    const returned = line()
+    await setup.waitForFrame(() => (line()?.spans.filter((span) => span.text.trim()).length ?? 0) > 4)
+    expect(returned).toEqual(highlighted)
+    expect(first).toEqual(line())
+  },
+)
+
 async function createAppFixture(
   input: {
     width?: number

packages/tui/test/plugin-markdown.test.ts

@@ -1,6 +1,13 @@
 import { expect, test } from "bun:test"
-import { SyntaxStyle, TextRenderable, type MarkdownOptions, type RenderNodeContext } from "@opentui/core"
-import { createTestRenderer } from "@opentui/core/testing"
+import {
+  CodeRenderable,
+  MarkdownRenderable,
+  SyntaxStyle,
+  TextRenderable,
+  type MarkdownOptions,
+  type RenderNodeContext,
+} from "@opentui/core"
+import { createTestRenderer, MockTreeSitterClient } from "@opentui/core/testing"
 import { combineMarkdownRenderers } from "../src/plugin/context"
 
 const code = (language: string) => ({ type: "code" as const, lang: language, text: "content", raw: "" })
@@ -29,6 +36,132 @@ test("later Markdown renderer registrations take precedence", async () => {
   const combined = combineMarkdownRenderers([{ mermaid: () => first }, { mermaid: () => second }])!
 
   expect(combined(code("mermaid"), context)).toBe(second)
-  expect(combineMarkdownRenderers([])).toBeUndefined()
+  expect(combineMarkdownRenderers([])).toBeFunction()
   renderer.destroy()
 })
+
+test.each(["returned", "declined"])("reuses a plugin's %s default code block", async (fallback) => {
+  const output = await createTestRenderer({ width: 80, height: 10 })
+  const defaults: Array<ReturnType<RenderNodeContext["defaultRender"]>> = []
+  const markdown = new MarkdownRenderable(output.renderer, {
+    content: "```typescript\nconst fixture = 42\n```",
+    syntaxStyle: context.syntaxStyle,
+    renderNode: combineMarkdownRenderers([
+      {
+        typescript: (_token, context) => {
+          defaults.push(context.defaultRender())
+          return fallback === "returned" ? defaults.at(-1) : undefined
+        },
+      },
+    ]),
+  })
+  output.renderer.root.add(markdown)
+  try {
+    const code = markdown.getChildren()[0]
+    expect(defaults).toHaveLength(1)
+    expect(defaults[0]).toBe(code)
+    if (!(code instanceof CodeRenderable)) throw new Error("Expected fenced code")
+    expect(code.drawUnstyledText).toBe(false)
+    await output.renderOnce()
+    await code.highlightingDone
+    await output.renderOnce()
+    expect(output.captureCharFrame()).toContain("const fixture = 42")
+  } finally {
+    output.renderer.destroy()
+  }
+  expect(defaults[0]?.isDestroyed).toBe(true)
+})
+
+test.each(["typescript", "text", "unknown-fixture-language", ""])(
+  "renders %s fences without losing content",
+  async (language) => {
+    const output = await createTestRenderer({ width: 80, height: 10 })
+    const markdown = new MarkdownRenderable(output.renderer, {
+      content: `\`\`\`${language}\nconst fixture = 42\n\`\`\``,
+      syntaxStyle: context.syntaxStyle,
+      renderNode: combineMarkdownRenderers([]),
+    })
+    output.renderer.root.add(markdown)
+    try {
+      await output.renderOnce()
+      const code = markdown.getChildren()[0]
+      expect(code).toBeInstanceOf(CodeRenderable)
+      if (!(code instanceof CodeRenderable)) throw new Error("Expected fenced code")
+      await code.highlightingDone
+      await output.renderOnce()
+      expect(output.captureCharFrame()).toContain("const fixture = 42")
+    } finally {
+      output.renderer.destroy()
+    }
+  },
+)
+
+test("keeps the fenced code node while streaming", async () => {
+  const output = await createTestRenderer({ width: 80, height: 10 })
+  const markdown = new MarkdownRenderable(output.renderer, {
+    content: "```typescript\nconst fixture = 4",
+    internalBlockMode: "top-level",
+    streaming: true,
+    syntaxStyle: context.syntaxStyle,
+    renderNode: combineMarkdownRenderers([]),
+  })
+  output.renderer.root.add(markdown)
+  try {
+    await output.renderOnce()
+    const code = markdown.getChildren()[0]
+    if (!(code instanceof CodeRenderable)) throw new Error("Expected fenced code")
+    await code.highlightingDone
+    markdown.content = "```typescript\nconst fixture = 42"
+    await output.renderOnce()
+    expect(markdown.getChildren()[0] === code).toBe(true)
+    await code.highlightingDone
+    await output.renderOnce()
+    expect(output.captureCharFrame()).toContain("const fixture = 42")
+    markdown.content += "\n```"
+    markdown.streaming = false
+    await output.renderOnce()
+    const completed = markdown.getChildren()[0]
+    if (!(completed instanceof CodeRenderable)) throw new Error("Expected completed fenced code")
+    await completed.highlightingDone
+    await output.renderOnce()
+    expect(output.captureCharFrame()).toContain("const fixture = 42")
+  } finally {
+    output.renderer.destroy()
+  }
+})
+
+test.each(["error", "rejection"])(
+  "reserves code layout and shows a readable fallback on parser %s",
+  async (failure) => {
+    const output = await createTestRenderer({ width: 80, height: 10 })
+    const pending = Promise.withResolvers<{ error: string }>()
+    const client = new MockTreeSitterClient()
+    client.highlightOnce = () => pending.promise
+    const markdown = new MarkdownRenderable(output.renderer, {
+      content: "```typescript\nconst fixture = 42\n```",
+      syntaxStyle: context.syntaxStyle,
+      treeSitterClient: client,
+      renderNode: combineMarkdownRenderers([]),
+    })
+    output.renderer.root.add(markdown)
+    try {
+      await output.renderOnce()
+      const code = markdown.getChildren()[0]
+      if (!(code instanceof CodeRenderable)) throw new Error("Expected fenced code")
+      expect(code.isHighlighting).toBe(true)
+      const height = code.height
+      expect(height).toBeGreaterThan(0)
+      expect(output.captureCharFrame()).not.toContain("const fixture")
+      if (failure === "error") pending.resolve({ error: "Parser unavailable" })
+      if (failure === "rejection") pending.reject(new Error("Worker unavailable"))
+      await code.highlightingDone
+      await output.renderOnce()
+      expect(output.captureCharFrame()).toContain("const fixture = 42")
+      expect(code.height).toBe(height)
+    } finally {
+      pending.resolve({ error: "Parser unavailable" })
+      output.renderer.destroy()
+      await client.destroy()
+    }
+  },
+)

packages/tui/test/util/syntax-highlight-cache.test.ts

@@ -0,0 +1,69 @@
+import { describe, expect, test } from "bun:test"
+import { cacheHighlights } from "../../src/util/syntax-highlight-cache"
+
+describe("syntax highlight cache", () => {
+  test("reuses completed and in-flight highlights", async () => {
+    let calls = 0
+    const highlight = cacheHighlights(async () => {
+      calls++
+      return { highlights: [[0, 5, "keyword"]] }
+    })
+
+    const first = highlight("const", "typescript")
+    const second = highlight("const", "typescript")
+
+    expect(second).toBe(first)
+    expect(await second).toEqual({ highlights: [[0, 5, "keyword"]] })
+    expect(await highlight("const", "typescript")).toEqual({ highlights: [[0, 5, "keyword"]] })
+    expect(calls).toBe(1)
+  })
+
+  test("evicts least recently used highlights", async () => {
+    let calls = 0
+    const highlight = cacheHighlights(async () => {
+      calls++
+      return { highlights: [] }
+    }, 2)
+
+    await highlight("one", "text")
+    await highlight("two", "text")
+    await highlight("one", "text")
+    await highlight("three", "text")
+    await highlight("two", "text")
+
+    expect(calls).toBe(4)
+  })
+
+  test("retries failed highlights", async () => {
+    let calls = 0
+    const highlight = cacheHighlights(async () => {
+      calls++
+      if (calls === 1) return { error: "parser unavailable" }
+      return { highlights: [] }
+    })
+
+    await highlight("const", "typescript")
+    await highlight("const", "typescript")
+
+    expect(calls).toBe(2)
+  })
+
+  test("an evicted failure does not delete its replacement", async () => {
+    const pending = Promise.withResolvers<{ highlights: [] }>()
+    let calls = 0
+    const highlight = cacheHighlights(() => {
+      calls++
+      if (calls === 1) return pending.promise
+      return Promise.resolve({ highlights: [] })
+    }, 1)
+
+    const stale = highlight("one", "text")
+    await highlight("two", "text")
+    const current = highlight("one", "text")
+    pending.reject(new Error("parser unavailable"))
+
+    await expect(stale).rejects.toThrow("parser unavailable")
+    expect(highlight("one", "text")).toBe(current)
+    expect(calls).toBe(3)
+  })
+})