diff --git a/cli/src/components/blocks/content-with-markdown.tsx b/cli/src/components/blocks/content-with-markdown.tsx index ea25e61651..bb492cb93f 100644 --- a/cli/src/components/blocks/content-with-markdown.tsx +++ b/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(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) }, ) diff --git a/cli/src/components/message-with-agents.tsx b/cli/src/components/message-with-agents.tsx index cdc0440298..97d9e3ee3c 100644 --- a/cli/src/components/message-with-agents.tsx +++ b/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 | 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) diff --git a/cli/src/utils/__tests__/streaming-markdown.test.tsx b/cli/src/utils/__tests__/streaming-markdown.test.tsx new file mode 100644 index 0000000000..e1bad462a0 --- /dev/null +++ b/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 => { + const setup = await createTestRenderer({ width: WIDTH, height: HEIGHT }) + const root = createRoot(setup.renderer) + try { + flushSync(() => + root.render( + + {node as React.ReactNode} + , + ), + ) + 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') + }) +}) diff --git a/cli/src/utils/markdown-renderer.tsx b/cli/src/utils/markdown-renderer.tsx index 662602cc25..d8afa77bc7 100644 --- a/cli/src/utils/markdown-renderer.tsx +++ b/cli/src/utils/markdown-renderer.tsx @@ -124,6 +124,7 @@ interface RenderState { const createRenderState = ( palette: MarkdownPalette, codeBlockWidth: number, + keyPrefix = '', ): RenderState => { let counter = 0 return { @@ -131,7 +132,7 @@ const createRenderState = ( codeBlockWidth, nextKey: () => { counter += 1 - return `markdown-${counter}` + return `${keyPrefix}markdown-${counter}` }, } } @@ -989,6 +990,119 @@ const renderNode = ( } } +/** + * Merges adjacent styled runs that render identically, including runs separated + * only by whitespace. + * + * A frame costs the terminal roughly linearly in the number of nodes it has to + * walk - ~60ms for 1,500 nodes regardless of whether anything changed - so a + * code block emitted as one span per line is the expensive shape, not the + * characters. Merging `line'\n'line` into a single + * span removes those per-line nodes without changing a rendered character. + */ +const coalesceNodes = (nodes: ReactNode[]): ReactNode[] => { + const merged: ReactNode[] = [] + + const sameStyledStrings = ( + a: ReactNode, + b: ReactNode, + ): [Record, string, string] | null => { + if (!React.isValidElement(a) || !React.isValidElement(b)) return null + if (a.type !== b.type) return null + const aProps = a.props as Record + const bProps = b.props as Record + const aKeys = Object.keys(aProps).filter((k) => k !== 'children') + const bKeys = Object.keys(bProps).filter((k) => k !== 'children') + if (aKeys.length !== bKeys.length) return null + for (const key of aKeys) { + if (aProps[key] !== bProps[key]) return null + } + if ( + typeof aProps.children !== 'string' || + typeof bProps.children !== 'string' + ) { + return null + } + return [aProps, aProps.children, bProps.children] + } + + for (let i = 0; i < nodes.length; i++) { + const node = nodes[i] + const prev = merged[merged.length - 1] + + if ( + typeof prev === 'string' && + typeof node === 'string' && + prev.trim().length > 0 && + node.trim().length > 0 + ) { + merged[merged.length - 1] = prev + node + continue + } + + const direct = sameStyledStrings(prev, node) + if (direct) { + const [props, aChildren, bChildren] = direct + merged[merged.length - 1] = React.createElement( + (prev as React.ReactElement).type as React.ElementType, + { ...props, children: aChildren + bChildren }, + ) + continue + } + + if (!React.isValidElement(node)) { + merged.push(node) + continue + } + + // Absorb a following run of same-styled nodes separated only by + // whitespace. Whitespace-only strings after the last node are left alone: + // they are the block separator the document-level trim drops. + let end = i + 1 + let separator = '' + let absorbed: string | null = null + let props: Record | null = null + + while (end < nodes.length) { + const next = nodes[end] + if (typeof next === 'string' && next.trim().length === 0) { + separator += next + end += 1 + continue + } + const match = sameStyledStrings(node, next) + if (!match) break + const [nextProps, , nextChildren] = match + props = nextProps + absorbed = (absorbed ?? (node.props as { children: string }).children) + + separator + + nextChildren + separator = '' + end += 1 + } + + if (absorbed !== null && props !== null) { + merged.push( + React.createElement( + (node as React.ReactElement).type as React.ElementType, + { ...props, children: absorbed }, + ), + ) + // The run ended at a node that could not be merged; the whitespace after + // the last merged node was buffered, not consumed, so re-emit it. + if (separator) { + merged.push(separator) + } + i = end - 1 + continue + } + + merged.push(node) + } + + return merged +} + const normalizeOutput = (nodes: ReactNode[]): ReactNode => { const trimmed = trimTrailingWhitespaceNodes(nodes) if (trimmed.length === 0) { @@ -1006,6 +1120,35 @@ const normalizeOutput = (nodes: ReactNode[]): ReactNode => { ) } +/** + * Parses and renders markdown without trimming the trailing block separator. + * + * `renderMarkdown` drops that separator, which is correct for a whole document + * but not when several rendered fragments are stitched back together: the + * incremental streaming renderer needs the separators to line up exactly as a + * single render would produce them. + */ +export function renderMarkdownNodes( + markdown: string, + options: MarkdownRenderOptions = {}, + keyPrefix = '', +): ReactNode[] { + try { + const palette = resolvePalette(options.palette) + const codeBlockWidth = options.codeBlockWidth ?? 80 + const state = createRenderState(palette, codeBlockWidth, keyPrefix) + const ast = processor.parse(markdown) as Root + applyInlineFallbackFormatting(ast) + return coalesceNodes(renderNode(ast, state, ast.type, undefined)) + } catch (error) { + logger.error(error, 'Failed to parse markdown') + return [markdown] + } +} + +/** Drops the whitespace-only nodes a rendered document ends with. */ +export const trimTrailingWhitespace = trimTrailingWhitespaceNodes + export function renderMarkdown( markdown: string, options: MarkdownRenderOptions = {}, @@ -1016,7 +1159,7 @@ export function renderMarkdown( const state = createRenderState(palette, codeBlockWidth) const ast = processor.parse(markdown) as Root applyInlineFallbackFormatting(ast) - const nodes = renderNode(ast, state, ast.type, undefined) + const nodes = coalesceNodes(renderNode(ast, state, ast.type, undefined)) return normalizeOutput(nodes) } catch (error) { logger.error(error, 'Failed to parse markdown') diff --git a/cli/src/utils/streaming-markdown.ts b/cli/src/utils/streaming-markdown.ts new file mode 100644 index 0000000000..d84a04a567 --- /dev/null +++ b/cli/src/utils/streaming-markdown.ts @@ -0,0 +1,204 @@ +import { Fragment, createElement, memo } from 'react' + +import { + hasIncompleteCodeFence, + renderMarkdownNodes, + trimTrailingWhitespace, +} from './markdown-renderer' + +import type { MarkdownRenderOptions } from './markdown-renderer' +import type { ReactNode } from 'react' + +/** + * Incremental renderer for the block that is still streaming. + * + * `renderMarkdown` costs roughly 3ms per kilobyte, and the live message is + * re-rendered on every flush while it grows - so a long answer re-parsed its + * whole transcript about ten times a second: a 200KB answer cost ~480ms of + * parsing per flush, on the render thread. + * + * Streaming text is append-only, so every complete block before the last + * boundary can be parsed once and reused. Only the tail - the block currently + * being written - is re-rendered per flush. + */ + +/** + * Rendered once per committed segment. `memo` keeps React from re-walking the + * already-rendered prefix on every flush: the element objects are reused, so + * the reconcile cost stays proportional to the tail rather than the message. + */ +const RenderedSegment = memo(function RenderedSegment({ + nodes, +}: { + nodes: ReactNode[] +}) { + return createElement(Fragment, null, ...nodes) +}) + +export type StreamingMarkdownCache = { + /** Already parsed and reusable. `content` must start with this. */ + prefix: string + /** Committed segments, each a memoized element. */ + segmentElements: ReactNode[] + /** How many segments have been committed; keys segment node keys apart. */ + segments: number + /** Whether a fenced code block is still open at the end of `prefix`. */ + fenceOpen: boolean + /** Options the nodes were rendered with; a change invalidates them. */ + codeBlockWidth: number + palette: unknown +} + +export const createStreamingMarkdownCache = (): StreamingMarkdownCache => ({ + prefix: '', + segmentElements: [], + segments: 0, + fenceOpen: false, + codeBlockWidth: -1, + palette: undefined, +}) + +/** A blank line only separates blocks cleanly outside lists and quotes. */ +const LIST_OR_QUOTE_LINE = /^[ \t]*(?:[-*+]|\d+[.)]|>)/ + +const previousLine = (content: string, index: number): string => { + let start = index - 1 + while (start >= 0 && content[start] !== '\n') start-- + return content.slice(start + 1, index) +} + +const nextLine = (content: string, index: number): string => { + let start = index + while ( + start < content.length && + (content[start] === '\n' || content[start] === ' ' || content[start] === '\t') + ) { + start++ + } + let end = start + while (end < content.length && content[end] !== '\n') end++ + return content.slice(start, end) +} + +const isSafeBoundary = (content: string, blankIndex: number): boolean => { + const next = nextLine(content, blankIndex) + // Nothing rendered after the boundary yet: no stable content to commit. + if (!next) return false + + // A list or blockquote can continue across a blank line (loose lists), and + // parsing the halves separately would restart ordered numbering or split the + // quote. Leave those boundaries to the tail. Excluding them also means the + // last committed block is never a paragraph followed by a list, which is the + // only case where the renderer's trailing separator depends on the sibling + // on the other side of the boundary. + const prev = previousLine(content, blankIndex) + if (LIST_OR_QUOTE_LINE.test(prev) || LIST_OR_QUOTE_LINE.test(next)) { + return false + } + + return true +} + +/** + * Renders streaming markdown by reusing the parse of the stable prefix. + * + * Produces the same output as `renderStreamingMarkdown` for content that + * extends what was rendered before; falls back to a full parse when the + * content is replaced (chat reload, edit) or has no safe boundary yet. + */ +export function renderStreamingMarkdownCached( + content: string, + options: MarkdownRenderOptions, + cache: StreamingMarkdownCache, +): ReactNode { + const codeBlockWidth = options.codeBlockWidth ?? 80 + + if ( + codeBlockWidth !== cache.codeBlockWidth || + options.palette !== cache.palette || + !content.startsWith(cache.prefix) + ) { + cache.prefix = '' + cache.segmentElements = [] + cache.segments = 0 + cache.fenceOpen = false + cache.codeBlockWidth = codeBlockWidth + cache.palette = options.palette + } + + let fenceOpen = cache.fenceOpen + let split = -1 + let fenceAtSplit = false + + for (let i = cache.prefix.length; i < content.length; i++) { + if (content.charCodeAt(i) === 96 /* ` */ && content.startsWith('```', i)) { + fenceOpen = !fenceOpen + i += 2 + continue + } + if ( + !fenceOpen && + content.charCodeAt(i) === 10 /* \n */ && + content[i + 1] === '\n' && + isSafeBoundary(content, i) + ) { + split = i + fenceAtSplit = fenceOpen + i += 1 + } + } + + if (split > cache.prefix.length) { + const segment = content.slice(cache.prefix.length, split) + const key = `md-seg-${cache.segments++}` + // The renderer ends a block with its separator (`\n\n`); hold that at the + // top level instead of inside the memoized element, so the final trim can + // still drop it when nothing follows. Every boundary this renderer picks + // sits between two top-level blocks and never before a list or quote, + // which is the only case where the separator would be a single `\n`. + cache.segmentElements.push( + createElement(RenderedSegment, { + key, + nodes: trimTrailingWhitespace( + renderMarkdownNodes(segment, options, `${key}-`), + ), + }), + '\n\n', + ) + cache.prefix = content.slice(0, split) + cache.fenceOpen = fenceAtSplit + } + + const tail = content.slice(cache.prefix.length) + + if (tail && hasIncompleteCodeFence(tail)) { + // Same split as renderStreamingMarkdown: rendered blocks, then the + // unterminated fence as literal text so it is not half-highlighted. That + // function trims the trailing separator before appending the literal text + // (so the pending fence continues the previous line); trimming here too + // keeps the stitched output identical. + const lastFence = tail.lastIndexOf('```') + const trimmedBefore = trimTrailingWhitespace([ + ...cache.segmentElements, + ...renderMarkdownNodes(tail.slice(0, lastFence), options, 'md-tail-'), + ]) + return createElement( + Fragment, + null, + ...trimmedBefore, + tail.slice(lastFence), + ) + } + + const tailNodes = tail + ? renderMarkdownNodes(tail, options, 'md-tail-') + : [] + const trimmed = trimTrailingWhitespace([...cache.segmentElements, ...tailNodes]) + if (trimmed.length === 0) { + return '' + } + if (trimmed.length === 1) { + return trimmed[0] + } + return createElement(Fragment, null, ...trimmed) +}