Merge pull request #63970 from NousResearch/bb/salvage-51653-assistant-ui

chore(desktop): upgrade @assistant-ui to 0.14 + use built-in streaming APIs (supersedes #51653)
This commit is contained in:
brooklyn! 2026-07-15 02:10:37 -04:00 committed by GitHub
commit 9baa7d4673
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 671 additions and 1029 deletions

View file

@ -50,8 +50,8 @@
"check": "npm run typecheck && npm run test && npm run test:desktop:all && npm run build"
},
"dependencies": {
"@assistant-ui/react": "^0.12.28",
"@assistant-ui/react-streamdown": "^0.1.11",
"@assistant-ui/react": "^0.14.23",
"@assistant-ui/react-streamdown": "^0.3.4",
"@audiowave/react": "^0.6.2",
"@chenglou/pretext": "^0.0.6",
"@codemirror/commands": "^6.10.4",

View file

@ -29,6 +29,7 @@ function settledClarifyProps(
args,
argsText: JSON.stringify(args),
isError: false,
respondToApproval: vi.fn(),
result,
resume: vi.fn(),
status: { type: 'complete' },

View file

@ -210,4 +210,29 @@ describe('preprocessMarkdown', () => {
expect(() => preprocessMarkdown(input)).not.toThrow()
})
it('keeps $$<digit>$$ display math intact instead of escaping it as currency', () => {
const output = preprocessMarkdown('$$5x = 10$$')
expect(output).toContain('$$5x = 10$$')
expect(output).not.toContain('\\$')
})
it('rewrites double-backslash bracket math to dollar delimiters', () => {
const output = preprocessMarkdown('\\\\(x^2\\\\)')
expect(output).toContain('$x^2$')
})
it('rewrites [/math] and [/inline] tag pairs to dollar delimiters', () => {
expect(preprocessMarkdown('[/math]a+b[/math]')).toContain('$$a+b$$')
expect(preprocessMarkdown('[/inline]x[/inline]')).toContain('$x$')
})
it('escapes currency dollars in prose so they are not parsed as math', () => {
const output = preprocessMarkdown('$5 and $10')
expect(output).toContain('\\$5')
expect(output).toContain('\\$10')
})
})

View file

@ -1,22 +1,15 @@
'use client'
import { TextMessagePartProvider, useMessagePartText } from '@assistant-ui/react'
import { type SmoothOptions, TextMessagePartProvider, useMessagePartText } from '@assistant-ui/react'
import {
parseMarkdownIntoBlocks,
type StreamdownTextComponents,
StreamdownTextPrimitive,
type SyntaxHighlighterProps
type SyntaxHighlighterProps,
tailBoundedRemend
} from '@assistant-ui/react-streamdown'
import { code } from '@streamdown/code'
import {
type ComponentProps,
memo,
type ReactNode,
useDeferredValue,
useEffect,
useMemo,
useState
} from 'react'
import { type ComponentProps, memo, useEffect, useMemo, useState } from 'react'
import { ExpandableBlock } from '@/components/chat/expandable-block'
import { PreviewAttachment } from '@/components/chat/preview-attachment'
@ -37,7 +30,6 @@ import {
mediaStreamUrl
} from '@/lib/media'
import { previewTargetFromMarkdownHref } from '@/lib/preview-targets'
import { tailBoundedRemend } from '@/lib/remend-tail'
import { cn } from '@/lib/utils'
import { detectEmbed, extractAlert, MarkdownAlert, RichCodeBlock, UrlEmbed } from './embeds'
@ -57,8 +49,8 @@ import { detectEmbed, extractAlert, MarkdownAlert, RichCodeBlock, UrlEmbed } fro
const mathPlugin = createMemoizedMathPlugin({ singleDollarTextMath: true })
// Replaces Streamdown's `parseIncompleteMarkdown` (full-text remend per
// flush) with a tail-bounded repair — see lib/remend-tail.ts. Must stay
// module-scope so the prop identity is stable across renders.
// flush) with a tail-bounded repair. Must stay module-scope so the prop
// identity is stable across renders.
function preprocessWithTailRepair(text: string): string {
try {
return tailBoundedRemend(preprocessMarkdown(text))
@ -70,8 +62,7 @@ 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) and whenever multiple surfaces render
// the same content (deferred + smooth reveal republish). A small module-level
// (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
@ -347,44 +338,11 @@ function MarkdownImage({ className, src, alt, ...props }: ComponentProps<'img'>)
)
}
/**
* Re-publish the active message-part context with React's `useDeferredValue`
* applied to the streaming text and status. The outer wrapper still re-renders
* on every token, but the work it does is trivial (one hook, one provider).
*
* The expensive subtree (Streamdown micromark mdast hast React) lives
* inside `<TextMessagePartProvider>` and reads the deferred text via the
* normal `useMessagePartText` hook. React's concurrent scheduler then has
* permission to:
* - skip intermediate token states when the next token arrives mid-render
* (it abandons the in-flight deferred render and starts over)
* - deprioritize the markdown render when the main thread is busy with an
* urgent task (typing, scrolling, layout work elsewhere)
*
* Net effect: per-token CPU is unchanged but the *blocking* part of that work
* goes away typing-while-streaming stays a single-frame paint, scroll
* stutter disappears, and the longtask histogram tightens because long
* commits can be interrupted and discarded.
*
* Industry standard (Streamdown's own block-array setState already uses
* `useTransition`); this just lifts the deferral up to the consumer text
* boundary so it covers the whole pipeline, not just the inner setState.
*/
function DeferStreamingText({ children }: { children: ReactNode }) {
const { text, status } = useMessagePartText()
const deferredText = useDeferredValue(text)
const isRunning = status.type === 'running'
return (
<TextMessagePartProvider isRunning={isRunning} text={deferredText}>
{children}
</TextMessagePartProvider>
)
}
interface MarkdownTextSurfaceProps {
containerClassName?: string
containerProps?: ComponentProps<'div'>
defer?: boolean
smooth?: boolean | SmoothOptions
}
// Headings shrink to chat scale rather than the prose default (h1≈xl). Kept
@ -433,7 +391,7 @@ function HugeTextFallback({ containerClassName, text }: { containerClassName?: s
)
}
function MarkdownTextSurface({ containerClassName, containerProps }: MarkdownTextSurfaceProps) {
function MarkdownTextSurface({ containerClassName, containerProps, defer, smooth }: MarkdownTextSurfaceProps) {
const { status, text } = useMessagePartText()
const isStreaming = status.type === 'running'
@ -561,26 +519,29 @@ function MarkdownTextSurface({ containerClassName, containerProps }: MarkdownTex
components={components}
containerClassName={cn(MARKDOWN_CONTAINER_CLASS_NAME, containerClassName)}
containerProps={containerProps}
defer={defer}
lineNumbers={false}
mode="streaming"
// Incomplete-markdown repair is handled by `preprocessWithTailRepair`
// below (tail-bounded remend) instead of Streamdown's built-in pass,
// which re-runs remend over the ENTIRE message on every flush — ~18%
// of streaming script time on 50KB+ messages. The repair itself stays
// always-on (even between flushes / for completed messages): an
// unclosed ```python ... ``` whose body contains `$` (shell snippets,
// JS template strings, dollar amounts) would otherwise leak those
// dollars to the math parser and render broken inline math. Shiki is
// independently deferred via `defer={isStreaming}` on the
// SyntaxHighlighter component.
// Incomplete-markdown repair runs in preprocessWithTailRepair on the
// full accumulated text; the built-in tail-bounded remend is disabled
// because a custom parseMarkdownIntoBlocksFn is supplied, and
// parseIncompleteMarkdown stays false to avoid a second full-text
// remend pass.
parseIncompleteMarkdown={false}
parseMarkdownIntoBlocksFn={parseMarkdownIntoBlocksCached}
plugins={plugins}
preprocess={preprocessWithTailRepair}
smooth={smooth}
/>
)
}
const SMOOTH_OPTIONS: SmoothOptions = {
drainMs: 500,
maxCharsPerFrame: 30,
minCommitMs: 33
}
interface MarkdownTextContentProps extends MarkdownTextSurfaceProps {
isRunning: boolean
text: string
@ -592,19 +553,13 @@ export function MarkdownTextContent({ isRunning, text, ...surfaceProps }: Markdo
// the whole message), blanking the Thinking widget.
return (
<TextMessagePartProvider isRunning={isRunning} text={text}>
<DeferStreamingText>
<MarkdownTextSurface {...surfaceProps} />
</DeferStreamingText>
<MarkdownTextSurface defer smooth={SMOOTH_OPTIONS} {...surfaceProps} />
</TextMessagePartProvider>
)
}
const MarkdownTextImpl = () => {
return (
<DeferStreamingText>
<MarkdownTextSurface />
</DeferStreamingText>
)
return <MarkdownTextSurface defer />
}
export const MarkdownText = memo(MarkdownTextImpl)

View file

@ -13,7 +13,7 @@ function Boom({ error }: { error: Error | null }): null {
return null
}
const lookupError = new Error('tapClientLookup: Index 2 out of bounds (length: 2)')
const lookupError = new Error('useClientLookup: Index 2 out of bounds (length: 2)')
describe('MessageRenderBoundary', () => {
it('renders children when nothing throws', () => {
@ -26,7 +26,7 @@ describe('MessageRenderBoundary', () => {
expect(screen.getByText('content')).toBeTruthy()
})
it('swallows the transient tapClientLookup out-of-bounds store race', () => {
it('swallows the transient useClientLookup out-of-bounds store race', () => {
const spy = vi.spyOn(console, 'error').mockImplementation(() => undefined)
const { container } = render(

View file

@ -1,6 +1,6 @@
import { Component, type ReactNode } from 'react'
// `@assistant-ui/store`'s index-keyed child-scope lookup (`tapClientLookup`)
// `@assistant-ui/store`'s index-keyed child-scope lookup (`useClientLookup`)
// throws — rather than returning undefined — when a subscriber reads an index
// that the message/parts list no longer has. This races during high-frequency
// store replacement (session switch mid-stream, gateway reconnect replay): a
@ -10,7 +10,7 @@ import { Component, type ReactNode } from 'react'
// without a local boundary it unwinds to the root and blanks the whole app.
// Upstream-tracked: assistant-ui/assistant-ui#4051, #3652.
const isTransientLookupError = (error: unknown): boolean =>
error instanceof Error && /tapClient(Lookup|Resource).*out of bounds/.test(error.message)
error instanceof Error && /(useClientLookup|tapClient(Lookup|Resource)).*out of bounds/.test(error.message)
interface Props {
// Changes whenever the message list mutates; remounting clears the caught

View file

@ -1,4 +1,4 @@
import { ExportedMessageRepository } from '@assistant-ui/core/internal'
import { ExportedMessageRepository } from '@assistant-ui/react'
// Clicking a user bubble must open the inline edit composer — through the
// app's incremental external-store runtime (which reimplements capability
// resolution, incl. `edit: onEdit !== undefined`) and the stock runtime.

View file

@ -8,6 +8,8 @@ import {
import {
type AssistantRuntime,
type ExternalStoreAdapter,
fromThreadMessageLike,
generateId,
type ThreadMessage,
useRuntimeAdapters
} from '@assistant-ui/react'
@ -134,11 +136,19 @@ class IncrementalExternalStoreThreadRuntimeCore extends ExternalStoreThreadRunti
self._notifyEventSubscribers(store.isRunning ? 'runStart' : 'runEnd', {})
}
// metadata.isOptimistic keeps this placeholder ephemeral: core evicts
// off-branch optimistic messages on head moves and omits them from export().
if (hasUpcomingMessage(isRunning, messages)) {
self._assistantOptimisticId = this.repository.appendOptimisticMessage(messages.at(-1)?.id ?? null, {
role: 'assistant',
content: []
})
const optimisticId = generateId()
this.repository.addOrUpdateMessage(
messages.at(-1)?.id ?? null,
fromThreadMessageLike(
{ role: 'assistant', content: [], metadata: { isOptimistic: true } },
optimisticId,
{ type: 'running' }
)
)
self._assistantOptimisticId = optimisticId
}
this.repository.resetHead(self._assistantOptimisticId ?? messages.at(-1)?.id ?? null)

View file

@ -1,3 +1,5 @@
import { escapeCurrencyDollars, normalizeMathDelimiters } from '@assistant-ui/react-streamdown'
import { isLikelyProseFence, sanitizeLanguageTag } from '@/lib/markdown-code'
import { stripPreviewTargets } from '@/lib/preview-targets'
@ -310,41 +312,6 @@ function normalizeFenceBlocks(text: string): string {
return out.join('\n')
}
// Convert LaTeX bracket delimiters to remark-math's dollar-sign syntax.
// Models often emit `\(...\)` for inline math and `\[...\]` for display
// math (the standard LaTeX convention) instead of `$...$` / `$$...$$`.
// remark-math only natively recognizes the dollar form, so we rewrite at
// preprocess time. Done with simple non-greedy matches keyed on the
// escaped-bracket sequences — these are rare enough in non-math content
// (you'd have to write a literal `\(` followed eventually by a literal
// `\)` with NO interleaving newline-paragraph-break) that false positives
// are extremely unlikely.
const LATEX_INLINE_RE = /\\\(([^\n]+?)\\\)/g
const LATEX_DISPLAY_RE = /\\\[([\s\S]+?)\\\]/g
function rewriteLatexBracketDelimiters(text: string): string {
return text
.replace(LATEX_INLINE_RE, (_, body: string) => `$${body}$`)
.replace(LATEX_DISPLAY_RE, (_, body: string) => `$$${body}$$`)
}
// Escape `$<digit>` patterns so they don't get eaten as math delimiters.
// Models commonly write currency amounts ($5, $19.99, $1,299) in prose.
// With `singleDollarTextMath: true`, remark-math is greedy and matches
// EVERY pair of `$`s — including the open of `$5` to the next `$10`,
// rendering "5 in my pocket and you have " as italicized math text.
// The de-facto convention across math-supporting LLM UIs is to treat
// `$` followed by a digit as currency rather than math, since math
// expressions almost always start with a letter or `\command`. Trade-
// off: a math expression like `$5x = 10$` would have its leading 5
// escaped — annoying but rare. The escape `\$` survives to render as
// a literal `$` in the final output.
const CURRENCY_DOLLAR_RE = /(^|[^\\])\$(?=\d)/g
function escapeCurrencyDollars(text: string): string {
return text.replace(CURRENCY_DOLLAR_RE, '$1\\$')
}
export function preprocessMarkdown(text: string): string {
const cleaned = text.replace(REASONING_BLOCK_RE, '').replace(PREVIEW_MARKER_RE, '')
const scrubbed = scrubBacktickNoise(cleaned)
@ -377,12 +344,10 @@ export function preprocessMarkdown(text: string): string {
const leading = part.match(/^\s*/)?.[0] ?? ''
const trailing = part.match(/\s*$/)?.[0] ?? ''
// rewriteLatexBracketDelimiters runs only on prose segments so
// we don't accidentally touch `\(` inside a code block.
// escapeCurrencyDollars likewise only runs on prose, so legit
// `$5` literals inside fenced code stay intact.
// Run only on prose segments so `$5` literals and `\(` inside code
// blocks stay intact.
const transformed = normalizeVisibleProse(
stripPreviewTargets(rewriteLatexBracketDelimiters(escapeCurrencyDollars(part)))
stripPreviewTargets(normalizeMathDelimiters(escapeCurrencyDollars(part)))
)
return leading + transformed + trailing

View file

@ -1,105 +0,0 @@
import { parseMarkdownIntoBlocks } from '@assistant-ui/react-streamdown'
import remend from 'remend'
import { describe, expect, it } from 'vitest'
import { findRemendWindowStart, tailBoundedRemend } from './remend-tail'
const CORPUS = `# Heading one
Intro paragraph with **bold**, *italic*, \`inline code\`, and a [link](https://example.com).
## Code
\`\`\`python
def main():
cost = "$5"
print(f"total: $\{cost}")
\`\`\`
Some text after the fence with $x^2 + y^2$ inline math.
$$
\\int_0^1 f(x) dx
$$
- list item one with **bold**
- list item two
| col a | col b |
| ----- | ----- |
| 1 | 2 |
~~~js
const s = \`template \${value}\`
~~~
Final paragraph with ~~strike~~ and unfinished [link text](https://exa
`
/**
* Render-equivalence oracle: full-text remend and tail-bounded remend may
* differ in raw string output ONLY in ways that cannot affect rendering
* i.e. after block splitting, every block must be identical. (Streamdown
* renders blocks independently, so block-level equality IS render equality.)
*/
function blocksOf(text: string): string[] {
return parseMarkdownIntoBlocks(text)
}
describe('tailBoundedRemend', () => {
it('matches full remend block output at every streaming prefix', () => {
for (let end = 1; end <= CORPUS.length; end++) {
const prefix = CORPUS.slice(0, end)
const full = blocksOf(remend(prefix))
const tail = blocksOf(tailBoundedRemend(prefix))
expect(tail, `prefix length ${end}: ${JSON.stringify(prefix.slice(-60))}`).toEqual(full)
}
})
it('repairs an unclosed fence opened early in a long message', () => {
const text = `intro\n\n\`\`\`python\n${'x = 1\n'.repeat(500)}print("$dollar")`
const repaired = tailBoundedRemend(text)
expect(blocksOf(repaired)).toEqual(blocksOf(remend(text)))
// the window must reach back to the fence opener
expect(findRemendWindowStart(text)).toBe(text.indexOf('```python'))
})
it('bounds the window to the tail paragraph when no fence is open', () => {
const text = `para one\n\npara two\n\npara three with **bold`
const start = findRemendWindowStart(text)
expect(start).toBe(text.indexOf('para three'))
expect(tailBoundedRemend(text)).toBe(remend(text))
})
it('widens the window across an open $$ math block', () => {
const text = `before\n\n$$\n\\frac{a}{b}`
const start = findRemendWindowStart(text)
expect(start).toBeLessThanOrEqual(text.indexOf('$$'))
expect(blocksOf(tailBoundedRemend(text))).toEqual(blocksOf(remend(text)))
})
it('handles closed constructs without modification', () => {
const text = `done **bold** and \`code\`\n\n\`\`\`js\nconst a = 1\n\`\`\`\n\nlast line.`
expect(tailBoundedRemend(text)).toBe(text)
})
it('intentionally diverges from full remend on cross-block dangling openers', () => {
// Full remend scans the whole document and appends `**` for an opener
// left dangling in an EARLIER block, dumping stray asterisks into the
// unrelated tail block ("|**"). Because Streamdown splits into blocks
// after the repair, that opener never renders as bold either way — the
// tail-bounded result is the cleaner of the two. This test documents
// the divergence so a future remend upgrade that changes the behavior
// gets noticed.
const text = `- item with **dangling\n- item two\n\n|`
expect(remend(text).endsWith('|**')).toBe(true)
expect(tailBoundedRemend(text).endsWith('|')).toBe(true)
expect(tailBoundedRemend(text).endsWith('|**')).toBe(false)
})
})

View file

@ -1,108 +0,0 @@
import remend from 'remend'
// Tail-bounded incomplete-markdown repair.
//
// Streamdown's built-in `parseIncompleteMarkdown` runs `remend` over the whole
// accumulated message on every streaming flush (~18% of script time on 50KB+
// messages). But repairs only ever matter in the trailing block: inline
// constructs can't cross a blank line, and Streamdown splits into blocks AFTER
// the repair, so a dangling opener in an earlier block can't reach the tail.
// We run `remend` on just that block instead.
const BACKTICK = 96 // `
const TILDE = 126 // ~
const SPACE = 32
const TAB = 9
const BACKSLASH = 92
const isSpace = (c: number) => c === SPACE || c === TAB
/**
* Index of the last top-level block start the char after the most recent
* blank line that sits outside any open code fence or `$$` math block. An
* unclosed fence/math always begins after that blank, so it stays wholly
* inside the window without separate tracking. One cheap char pass, no regex.
*/
export function findRemendWindowStart(text: string): number {
const n = text.length
let inFence = false
let fenceChar = 0
let fenceRun = 0
let inMath = false
let boundary = 0
let pending = -1 // a blank line, committed to `boundary` once content follows
for (let lineStart = 0; lineStart <= n; ) {
let lineEnd = text.indexOf('\n', lineStart)
if (lineEnd === -1) {
lineEnd = n
}
let i = lineStart
while (i < lineEnd && isSpace(text.charCodeAt(i))) {
i += 1
}
const first = i < lineEnd ? text.charCodeAt(i) : -1
let marker = false
// Fence open/close (``` or ~~~, ≤3 spaces indent).
if ((first === BACKTICK || first === TILDE) && i - lineStart <= 3) {
let run = i
while (run < lineEnd && text.charCodeAt(run) === first) {
run += 1
}
if (run - i >= 3) {
marker = true
if (!inFence) {
inFence = true
fenceChar = first
fenceRun = run - i
} else if (first === fenceChar && run - i >= fenceRun && onlyWhitespace(text, run, lineEnd)) {
inFence = false
}
}
}
// Toggle `$$` math state on plain lines ($$ inside a fence is literal).
if (!inFence && !marker) {
for (let s = text.indexOf('$$', lineStart); s !== -1 && s < lineEnd - 1; s = text.indexOf('$$', s + 2)) {
if (s === 0 || text.charCodeAt(s - 1) !== BACKSLASH) {
inMath = !inMath
}
}
}
if (first === -1 && !inFence && !inMath) {
pending = lineEnd + 1
} else if (pending !== -1) {
boundary = pending
pending = -1
}
lineStart = lineEnd + 1
}
return boundary
}
function onlyWhitespace(text: string, from: number, to: number): boolean {
for (let i = from; i < to; i += 1) {
if (!isSpace(text.charCodeAt(i))) {
return false
}
}
return true
}
export function tailBoundedRemend(text: string): string {
const start = findRemendWindowStart(text)
return start <= 0 ? remend(text) : text.slice(0, start) + remend(text.slice(start))
}

1288
package-lock.json generated

File diff suppressed because it is too large Load diff

View file

@ -38,7 +38,6 @@
},
"overrides": {
"lodash": "4.18.1",
"@assistant-ui/store": "0.2.13",
"yauzl": "^3.3.1"
},
"engines": {