perf(cli): render streaming markdown incrementally

cli/src/components/blocks/content-with-markdown.tsx

@@ -1,11 +1,16 @@
-import { memo } from 'react'
+import { memo, useRef } from 'react'
 
 import {
   renderMarkdown,
-  renderStreamingMarkdown,
   hasMarkdown,
   type MarkdownPalette,
 } from '../../utils/markdown-renderer'
+import {
+  createStreamingMarkdownCache,
+  renderStreamingMarkdownCached,
+} from '../../utils/streaming-markdown'
+
+import type { StreamingMarkdownCache } from '../../utils/streaming-markdown'
 
 interface ContentWithMarkdownProps {
   content: string
@@ -21,13 +26,26 @@ export const ContentWithMarkdown = memo(
     codeBlockWidth,
     palette,
   }: ContentWithMarkdownProps) => {
+    // Streaming text only ever grows, so the already-rendered blocks can be
+    // kept and the tail re-rendered: parsing the whole message on every flush
+    // is what made long answers stall the render thread.
+    const cache = useRef<StreamingMarkdownCache | null>(null)
+    if (cache.current === null) {
+      cache.current = createStreamingMarkdownCache()
+    }
+
     if (!hasMarkdown(content)) {
       return content
     }
+
     const options = { codeBlockWidth, palette }
     if (isStreaming) {
-      return renderStreamingMarkdown(content, options)
+      return renderStreamingMarkdownCached(content, options, cache.current)
     }
+
+    // Done streaming: render once from scratch so the final output does not
+    // depend on how the deltas happened to be split.
+    cache.current = createStreamingMarkdownCache()
     return renderMarkdown(content, options)
   },
 )

cli/src/components/message-with-agents.tsx

@@ -1,5 +1,5 @@
 import { TextAttributes } from '@opentui/core'
-import { memo, useCallback, useMemo, type ReactNode } from 'react'
+import { memo, useCallback, useMemo, useRef, type ReactNode } from 'react'
 import React from 'react'
 import { useShallow } from 'zustand/react/shallow'
 
@@ -21,6 +21,10 @@ import {
   hasMarkdown,
   type MarkdownPalette,
 } from '../utils/markdown-renderer'
+import {
+  createStreamingMarkdownCache,
+  renderStreamingMarkdownCached,
+} from '../utils/streaming-markdown'
 
 import type { ChatMessage } from '../types/chat'
 import type { FeedbackCategory } from '@codebuff/common/constants/feedback'
@@ -357,6 +361,29 @@ const AgentMessage = memo(
     const isStreaming = useChatStore((state) => state.streamingAgents.has(message.id))
     const setFocusedAgentId = useChatStore((state) => state.setFocusedAgentId)
 
+    // A streaming agent re-renders on every flush, and a full markdown parse of
+    // its whole body per flush is the same quadratic cost the root message had.
+    const markdownCache = useRef<ReturnType<typeof createStreamingMarkdownCache> | null>(null)
+    if (markdownCache.current === null) {
+      markdownCache.current = createStreamingMarkdownCache()
+    }
+
+    const agentCodeBlockWidth = Math.max(
+      10,
+      availableWidth - AGENT_CONTENT_HORIZONTAL_PADDING,
+    )
+    const agentPalette: MarkdownPalette | undefined = useMemo(
+      () => markdownPalette ? {
+        ...markdownPalette,
+        codeTextFg: theme?.foreground ?? markdownPalette.codeTextFg,
+      } : undefined,
+      [markdownPalette, theme?.foreground],
+    )
+    const agentMarkdownOptions = useMemo(
+      () => ({ codeBlockWidth: agentCodeBlockWidth, palette: agentPalette! }),
+      [agentCodeBlockWidth, agentPalette],
+    )
+
     // Guard against missing agent info (should not happen for agent variant messages)
     if (!message.agent) {
       return (
@@ -389,21 +416,15 @@ const AgentMessage = memo(
         ? lastLine.replace(/[#*_`~\[\]()]/g, '').trim()
         : ''
 
-    const agentCodeBlockWidth = Math.max(
-      10,
-      availableWidth - AGENT_CONTENT_HORIZONTAL_PADDING,
-    )
-    const agentPalette: MarkdownPalette | undefined = markdownPalette ? {
-      ...markdownPalette,
-      codeTextFg: theme?.foreground ?? markdownPalette.codeTextFg,
-    } : undefined
-    const agentMarkdownOptions = {
-      codeBlockWidth: agentCodeBlockWidth,
-      palette: agentPalette!,
-    }
-    const displayContent = hasMarkdown(rawDisplayContent)
-      ? renderMarkdown(rawDisplayContent, agentMarkdownOptions)
-      : rawDisplayContent
+    const displayContent = !hasMarkdown(rawDisplayContent)
+      ? rawDisplayContent
+      : isStreaming
+        ? renderStreamingMarkdownCached(
+            rawDisplayContent,
+            agentMarkdownOptions,
+            markdownCache.current,
+          )
+        : renderMarkdown(rawDisplayContent, agentMarkdownOptions)
 
     const handleTitleClick = (): void => {
       onToggleCollapsed(message.id)

cli/src/utils/__tests__/streaming-markdown.test.tsx

@@ -0,0 +1,192 @@
+import { describe, expect, test } from 'bun:test'
+import { createTestRenderer } from '@opentui/core/testing'
+import { createRoot, flushSync } from '@opentui/react'
+import React from 'react'
+
+import { renderStreamingMarkdown } from '../markdown-renderer'
+import {
+  createStreamingMarkdownCache,
+  renderStreamingMarkdownCached,
+} from '../streaming-markdown'
+
+import type { ReactNode } from 'react'
+
+const WIDTH = 72
+const HEIGHT = 90
+
+const renderFrame = async (node: ReactNode): Promise<string> => {
+  const setup = await createTestRenderer({ width: WIDTH, height: HEIGHT })
+  const root = createRoot(setup.renderer)
+  try {
+    flushSync(() =>
+      root.render(
+        <box flexDirection="column">
+          <text>{node as React.ReactNode}</text>
+        </box>,
+      ),
+    )
+    await setup.renderOnce()
+    return setup.captureCharFrame()
+  } finally {
+    flushSync(() => root.unmount())
+    setup.renderer.destroy()
+  }
+}
+
+const DOC = `## Streaming equivalence
+
+A paragraph with **bold**, *italic*, \`inline code\` and a [link](https://example.com).
+
+- first bullet with \`code\`
+- second bullet
+- third bullet
+
+1. ordered one
+2. ordered two
+
+> a quoted line
+> with two lines
+
+| column a | column b |
+| -------- | -------- |
+| value 1  | value 2  |
+
+\`\`\`ts
+export const answer = 42
+function add(a: number, b: number) {
+  return a + b
+}
+\`\`\`
+
+A closing paragraph that mentions a < b and 5 > 3 as plain text.
+`
+
+/** Every prefix that ends at a line boundary, i.e. where a flush could land. */
+const prefixEnds = (text: string): number[] => {
+  const ends: number[] = []
+  for (let i = 1; i < text.length; i++) {
+    if (text[i] === '\n') ends.push(i + 1)
+  }
+  return ends
+}
+
+describe('renderStreamingMarkdownCached', () => {
+  test('matches a full render at every streaming step', async () => {
+    const cache = createStreamingMarkdownCache()
+    const options = { codeBlockWidth: 72 }
+
+    for (const end of prefixEnds(DOC)) {
+      const content = DOC.slice(0, end)
+      const expected = await renderFrame(renderStreamingMarkdown(content, options))
+      const actual = await renderFrame(
+        renderStreamingMarkdownCached(content, options, cache),
+      )
+
+      // Guard against a vacuous pass: an empty or error frame would compare
+      // equal to another empty frame.
+      expect(expected).toContain('Streaming equivalence')
+      expect({ at: end, frame: actual }).toEqual({ at: end, frame: expected })
+    }
+  })
+
+  test('matches a full render inside an unterminated code fence', async () => {
+    const cache = createStreamingMarkdownCache()
+    const options = { codeBlockWidth: 72 }
+    const open = `Intro text.\n\n\`\`\`ts\nconst a = 1\nconst b = 2\n`
+
+    const expected = await renderFrame(renderStreamingMarkdown(open, options))
+    const actual = await renderFrame(
+      renderStreamingMarkdownCached(open, options, cache),
+    )
+
+    expect(expected).toContain('const a = 1')
+    expect(actual).toEqual(expected)
+  })
+
+  test('a paragraph directly followed by a list keeps the same spacing', async () => {
+    const cache = createStreamingMarkdownCache()
+    const options = { codeBlockWidth: 72 }
+    const content = `Intro line.\n\n- one\n- two\n\nAfter the list.\n`
+
+    const expected = await renderFrame(renderStreamingMarkdown(content, options))
+    const actual = await renderFrame(
+      renderStreamingMarkdownCached(content, options, cache),
+    )
+
+    expect(expected).toContain('After the list.')
+    expect(actual).toEqual(expected)
+  })
+
+  test('replaced content re-renders from scratch', async () => {
+    const cache = createStreamingMarkdownCache()
+    const options = { codeBlockWidth: 72 }
+
+    renderStreamingMarkdownCached(DOC, options, cache)
+
+    const replacement = `# A different message\n\nWith **its own** body.\n`
+    const expected = await renderFrame(
+      renderStreamingMarkdown(replacement, options),
+    )
+    const actual = await renderFrame(
+      renderStreamingMarkdownCached(replacement, options, cache),
+    )
+
+    expect(expected).toContain('A different message')
+    expect(actual).toEqual(expected)
+  })
+
+  test('a width change invalidates the cached prefix', async () => {
+    const cache = createStreamingMarkdownCache()
+
+    renderStreamingMarkdownCached(DOC, { codeBlockWidth: 72 }, cache)
+
+    const expected = await renderFrame(
+      renderStreamingMarkdown(DOC, { codeBlockWidth: 40 }),
+    )
+    const actual = await renderFrame(
+      renderStreamingMarkdownCached(DOC, { codeBlockWidth: 40 }, cache),
+    )
+
+    expect(expected).toContain('Streaming equivalence')
+    expect(actual).toEqual(expected)
+  })
+})
+
+describe('merged styled runs keep block boundaries', () => {
+  // Regression: merging same-styled nodes must not swallow the blank line
+  // between a code block and the paragraph that follows it.
+  test('a paragraph after a code block starts on its own line', async () => {
+    const doc = `\`\`\`ts
+const a = 1
+const b = 2
+\`\`\`
+
+After the block.
+`
+    const frame = await renderFrame(
+      renderStreamingMarkdown(doc, { codeBlockWidth: 60 }),
+    )
+    const lines = frame.split('\n').map((line) => line.trimEnd())
+
+    expect(lines).toContain('const b = 2')
+    expect(lines).toContain('After the block.')
+    expect(lines.some((line) => line.includes('const b = 2After'))).toBe(false)
+  })
+
+  test('code block lines stay on separate lines', async () => {
+    const doc = `\`\`\`ts
+const a = 1
+const b = 2
+const c = 3
+\`\`\`
+`
+    const frame = await renderFrame(
+      renderStreamingMarkdown(doc, { codeBlockWidth: 60 }),
+    )
+    const lines = frame.split('\n').map((line) => line.trimEnd())
+
+    expect(lines).toContain('const a = 1')
+    expect(lines).toContain('const b = 2')
+    expect(lines).toContain('const c = 3')
+  })
+})