diff --git a/apps/desktop/src/components/assistant-ui/markdown-text.tsx b/apps/desktop/src/components/assistant-ui/markdown-text.tsx index 8f17685108ff..a199b5aa1e7d 100644 --- a/apps/desktop/src/components/assistant-ui/markdown-text.tsx +++ b/apps/desktop/src/components/assistant-ui/markdown-text.tsx @@ -2,7 +2,6 @@ import { TextMessagePartProvider, useMessagePartText } from '@assistant-ui/react' import { - parseMarkdownIntoBlocks, type StreamdownTextComponents, StreamdownTextPrimitive, type SyntaxHighlighterProps, @@ -17,6 +16,7 @@ import { chunkByLines, SyntaxHighlighter } from '@/components/chat/shiki-highlig import { ZoomableImage } from '@/components/chat/zoomable-image' import { normalizeExternalUrl, openExternalLink, PrettyLink } from '@/lib/external-link' import { createMemoizedMathPlugin } from '@/lib/katex-memo' +import { parseMarkdownIntoBlocksCached } from '@/lib/markdown-blocks' import { preprocessMarkdown } from '@/lib/markdown-preprocess' import { downloadGatewayMediaFile, @@ -59,43 +59,6 @@ function preprocessWithTailRepair(text: string): string { } } -// Memoized block splitter. Streamdown calls `parseMarkdownIntoBlocks` (a full -// `marked` lex of the entire message, ~1.6ms per 28KB) inside a useMemo keyed -// on the text — but the same text is re-lexed every time a message REMOUNTS -// (virtualizer scroll, session switch). A small module-level -// LRU keyed by the exact source string removes all of those repeat parses -// with zero correctness risk (same input → same output). Streaming tail -// growth misses the cache by design (every flush is a new string) — that -// single lex is the irreducible cost. -const BLOCK_CACHE_MAX = 64 -const BLOCK_CACHE_MIN_LENGTH = 1024 -const blockCache = new Map() - -function parseMarkdownIntoBlocksCached(markdown: string): string[] { - if (markdown.length < BLOCK_CACHE_MIN_LENGTH) { - return parseMarkdownIntoBlocks(markdown) - } - - const hit = blockCache.get(markdown) - - if (hit) { - // Refresh recency (Map iteration order is insertion order). - blockCache.delete(markdown) - blockCache.set(markdown, hit) - - return hit - } - - const blocks = parseMarkdownIntoBlocks(markdown) - blockCache.set(markdown, blocks) - - if (blockCache.size > BLOCK_CACHE_MAX) { - blockCache.delete(blockCache.keys().next().value as string) - } - - return blocks -} - async function mediaSrc(path: string): Promise { if (/^(?:https?|data):/i.test(path)) { return path diff --git a/apps/desktop/src/lib/markdown-blocks.test.ts b/apps/desktop/src/lib/markdown-blocks.test.ts new file mode 100644 index 000000000000..e3f491275657 --- /dev/null +++ b/apps/desktop/src/lib/markdown-blocks.test.ts @@ -0,0 +1,111 @@ +import { parseMarkdownIntoBlocks } from '@assistant-ui/react-streamdown' +import { describe, expect, it } from 'vitest' + +import { parseMarkdownIntoBlocksCached } from './markdown-blocks' + +// The contract: streaming through the cached splitter (one call per growing +// prefix, exactly how Streamdown calls it per flush) must produce, at every +// step, the same blocks as a fresh full lex of that prefix. Byte equality — +// a divergence would change what the memoized block renderer paints. + +const CORPUS = `# Heading + +Intro paragraph with **bold**, [a link](https://example.com), \`inline\` and $x^2$ math. + +- list item one +- list item two + - nested item + +1. ordered a + +2. loose ordered b + +\`\`\`python +def f(x): + return x * 2 # comment with \`\`\` inside string? no — fence chars below +\`\`\` + +A paragraph that will be followed by a setext underline +=== + +| col a | col b | +|---|---| +| 1 | 2 | +| 3 | 4 | + +> blockquote line one +> blockquote line two +with a lazy continuation line + +
+html block content +
+ +$$ +\\int_0^1 x\\,dx = \\tfrac12 +$$ + +Final paragraph after everything, long enough to stream in pieces so the tail +block keeps getting reinterpreted while earlier blocks stay settled. +` + +// Deterministic PRNG so failures reproduce. +function mulberry32(seed: number) { + let a = seed + + return () => { + a |= 0 + a = (a + 0x6d2b79f5) | 0 + let t = Math.imul(a ^ (a >>> 15), 1 | a) + t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t + + return ((t ^ (t >>> 14)) >>> 0) / 4294967296 + } +} + +// Push the text past the cache MIN_LENGTH thresholds so the incremental +// path actually engages. +const LONG_CORPUS = Array.from({ length: 6 }, () => CORPUS).join('\n') + +describe('parseMarkdownIntoBlocksCached', () => { + it('matches a full lex at every random streaming cut (property)', () => { + for (let seed = 1; seed <= 5; seed++) { + const rand = mulberry32(seed) + let cursor = 0 + + while (cursor < LONG_CORPUS.length) { + cursor = Math.min(LONG_CORPUS.length, cursor + 1 + Math.floor(rand() * 120)) + const prefix = LONG_CORPUS.slice(0, cursor) + + expect(parseMarkdownIntoBlocksCached(prefix)).toEqual(parseMarkdownIntoBlocks(prefix)) + } + } + }) + + it('matches a full lex when streaming token-by-token through a fence boundary', () => { + const base = `${'settled paragraph one.\n\n'.repeat(100)}opening a fence now:\n` + const tail = '```js\nconst a = 1\nconst b = 2\n```\n\nafter the fence\n' + + for (let i = 1; i <= tail.length; i++) { + const text = base + tail.slice(0, i) + + expect(parseMarkdownIntoBlocksCached(text)).toEqual(parseMarkdownIntoBlocks(text)) + } + }) + + it('reconstructs the input exactly (join property the offsets rely on)', () => { + const blocks = parseMarkdownIntoBlocksCached(LONG_CORPUS) + + expect(blocks.join('')).toBe(LONG_CORPUS) + }) + + it('falls back to a full lex for non-append rewrites (edit / branch swap)', () => { + const grown = `${LONG_CORPUS}\n\nappended tail paragraph` + parseMarkdownIntoBlocksCached(grown) + + // A REWRITE that shares no prefix lineage must still be correct. + const rewritten = `completely different start\n\n${LONG_CORPUS.slice(500)}` + + expect(parseMarkdownIntoBlocksCached(rewritten)).toEqual(parseMarkdownIntoBlocks(rewritten)) + }) +}) diff --git a/apps/desktop/src/lib/markdown-blocks.ts b/apps/desktop/src/lib/markdown-blocks.ts new file mode 100644 index 000000000000..4235301b6a25 --- /dev/null +++ b/apps/desktop/src/lib/markdown-blocks.ts @@ -0,0 +1,125 @@ +import { parseMarkdownIntoBlocks } from '@assistant-ui/react-streamdown' + +/** + * Block splitting for the streaming markdown pipeline, without re-lexing the + * whole message on every token flush. + * + * `parseMarkdownIntoBlocks` is a full `marked` lex of the entire text — + * measured 3.4–9.6ms per call at 64–192KB. During streaming every flush is a + * new string, so the stock splitter pays that O(full-text) cost ~30×/s on + * long replies. Two caches remove it: + * + * 1. Exact-string LRU — a message that REMOUNTS with unchanged text + * (virtualizer scroll, session switch) reuses its parse outright. + * 2. Streaming-append cache — when the new text starts with a recently parsed + * text (the token-append case), the previous parse's blocks are reused up + * to a settled boundary and only the suffix is lexed. The boundary drops + * the previous parse's trailing whitespace-only blocks AND its last content + * block, because appended text can retroactively change how that last + * block parses (open fence, list/table continuation, setext underline, a + * lazy blockquote line). Blocks before it are separated by settled blank + * lines and cannot be affected. Cross-block reference links can't regress: + * Streamdown renders each block as an independent markdown document + * already. Verified property: `blocks.join('') === text`, and incremental + * output is asserted byte-identical to a full lex in tests across fences, + * lists, tables, setext headings, blockquotes, and HTML blocks. + * + * Any doubt — no prefix match, reconstruction mismatch — falls back to the + * full lex, i.e. exactly the previous behavior. + */ + +const EXACT_CACHE_MAX = 64 +const EXACT_CACHE_MIN_LENGTH = 1024 +const exactCache = new Map() + +// Streaming messages grow monotonically, and only a handful stream at once +// (main reply + reasoning part, maybe a tile). A tiny ring is enough; each +// entry holds the last parse for one growing text lineage. +const APPEND_CACHE_MAX = 4 +const APPEND_CACHE_MIN_LENGTH = 2048 +const appendCache: { blocks: string[]; text: string }[] = [] + +function rememberAppend(text: string, blocks: string[]): void { + if (text.length < APPEND_CACHE_MIN_LENGTH) { + return + } + + // Replace the lineage this text grew from (its cached prefix), else push. + const index = appendCache.findIndex(entry => text.startsWith(entry.text)) + + if (index !== -1) { + appendCache.splice(index, 1) + } + + appendCache.push({ blocks, text }) + + if (appendCache.length > APPEND_CACHE_MAX) { + appendCache.shift() + } +} + +function lexIncrementally(text: string): null | string[] { + const entry = appendCache.find(cached => text.length > cached.text.length && text.startsWith(cached.text)) + + if (!entry) { + return null + } + + // Settled boundary: drop trailing whitespace-only blocks, then the last + // content block (the only one appended text can reinterpret). + let keep = entry.blocks.length + + while (keep > 0 && !entry.blocks[keep - 1].trim()) { + keep -= 1 + } + + if (keep > 0) { + keep -= 1 + } + + if (keep === 0) { + return null + } + + const settled = entry.blocks.slice(0, keep) + let settledLength = 0 + + for (const block of settled) { + settledLength += block.length + } + + // Defensive reconstruction check — the splitter's join(blocks) === text + // property is what makes offsets exact. If it ever doesn't hold, full lex. + if (settledLength > entry.text.length || !text.startsWith(entry.text.slice(0, settledLength), 0)) { + return null + } + + return [...settled, ...parseMarkdownIntoBlocks(text.slice(settledLength))] +} + +export function parseMarkdownIntoBlocksCached(markdown: string): string[] { + if (markdown.length < EXACT_CACHE_MIN_LENGTH) { + return parseMarkdownIntoBlocks(markdown) + } + + const hit = exactCache.get(markdown) + + if (hit) { + // Refresh recency (Map iteration order is insertion order). + exactCache.delete(markdown) + exactCache.set(markdown, hit) + + return hit + } + + const blocks = lexIncrementally(markdown) ?? parseMarkdownIntoBlocks(markdown) + + rememberAppend(markdown, blocks) + exactCache.set(markdown, blocks) + + if (exactCache.size > EXACT_CACHE_MAX) { + exactCache.delete(exactCache.keys().next().value as string) + } + + return blocks +}