Skip to content

chore: upgrade @assistant-ui to 0.14 and replace all custom streaming-rendering code with built-in APIs #51653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
okisdev wants to merge 9 commits into from
Closed

chore: upgrade @assistant-ui to 0.14 and replace all custom streaming-rendering code with built-in APIs #51653

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
6376c69
chore: upgrade @assistant-ui/react 0.12 to 0.14 and react-streamdown …
okisdev Jun 24, 2026
6a88a08
feat: replace custom streaming wrappers with built-in defer and smoot…
okisdev Jun 24, 2026
0d92a99
refactor: replace custom lib/remend-tail.ts with built-in tailBounded…
okisdev Jun 24, 2026
cab0f69
refactor: replace custom math delimiter helpers with built-in normali…
okisdev Jun 24, 2026
24a940f
fix(desktop): match renamed useClientLookup out-of-bounds throw in Me…
okisdev Jun 30, 2026
fbe98ca
docs(desktop): correct streaming-repair comments after the defer/smoo…
okisdev Jun 30, 2026
4b12864
Merge branch 'main' into chore/upgrade-assistant-ui
okisdev Jun 30, 2026
94756bd
test(desktop): add math-delimiter fixtures for the react-streamdown p…
okisdev Jun 30, 2026
27f8d78
docs(desktop): note the load-bearing isOptimistic invariant on the op…
okisdev Jun 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions apps/desktop/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@
"preview": "node scripts/assert-root-install.cjs && vite preview --host 127.0.0.1 --port 4174"
},
"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",
Expand Down
25 changes: 25 additions & 0 deletions apps/desktop/src/components/assistant-ui/markdown-text.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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')
})
})
204 changes: 25 additions & 179 deletions apps/desktop/src/components/assistant-ui/markdown-text.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,23 +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,
useRef,
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'
Expand All @@ -37,7 +29,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'
Expand All @@ -57,8 +48,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))
Expand All @@ -70,8 +61,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
Expand Down Expand Up @@ -306,150 +296,11 @@ function MarkdownImage({ className, src, alt, ...props }: ComponentProps<'img'>)
)
}

// Steady character-reveal for streaming text: decouples visible cadence from
// bursty arrival so text flows instead of popping (cf. assistant-ui's useSmooth,
// reimplemented for a tunable rate). Proportional drain — each frame reveals a
// slice of the backlog so the reveal converges within ~REVEAL_DRAIN_MS whatever
// the size; the per-frame cap stops a huge dump rendering as one slab. The loop
// is gated on backlog, not isRunning, so a stream that completes mid-reveal
// keeps draining its tail instead of snapping.
const REVEAL_DRAIN_MS = 500
const REVEAL_MAX_CHARS_PER_FRAME = 30
// Floor between reveal commits. Each commit republishes the text context and
// re-runs the whole Streamdown pipeline (preprocess → remend → lex → micromark
// on the open block) over the full accumulated text — at raw rAF cadence
// that's 60 full parses/second and was the dominant streaming cost for
// reasoning text. ~33ms keeps the reveal visually fluid (2 frames) while
// halving the parse work.
const REVEAL_MIN_COMMIT_MS = 33

function useSmoothReveal(text: string, isRunning: boolean): string {
const [displayed, setDisplayed] = useState(isRunning ? '' : text)
const targetRef = useRef(text)
const shownRef = useRef(displayed)
const frameRef = useRef<number | null>(null)
const lastTickRef = useRef(0)

shownRef.current = displayed
targetRef.current = text

useEffect(() => {
if (typeof window === 'undefined') {
return
}

// Non-extending change (regenerate / branch / history swap): restart from
// empty while streaming, else snap to the replacement.
if (!text.startsWith(shownRef.current)) {
shownRef.current = isRunning ? '' : text
setDisplayed(shownRef.current)
}

if (shownRef.current.length >= text.length || frameRef.current !== null) {
return
}

lastTickRef.current = performance.now()

const tick = () => {
const now = performance.now()
const dt = now - lastTickRef.current

// Skip this frame if the floor hasn't elapsed — the backlog math below
// is dt-proportional, so delayed commits reveal proportionally more.
if (dt < REVEAL_MIN_COMMIT_MS) {
frameRef.current = requestAnimationFrame(tick)

return
}

lastTickRef.current = now

const remaining = targetRef.current.length - shownRef.current.length

const add = Math.min(
remaining,
// dt-scaled so the per-commit cap stays equivalent to the old
// per-frame cap at any commit cadence.
Math.ceil((REVEAL_MAX_CHARS_PER_FRAME * dt) / 16.7),
Math.max(1, Math.ceil((remaining * dt) / REVEAL_DRAIN_MS))
)

shownRef.current = targetRef.current.slice(0, shownRef.current.length + add)
setDisplayed(shownRef.current)

frameRef.current = shownRef.current.length < targetRef.current.length ? requestAnimationFrame(tick) : null
}

frameRef.current = requestAnimationFrame(tick)
}, [text, isRunning])

useEffect(
() => () => {
if (frameRef.current !== null && typeof window !== 'undefined') {
cancelAnimationFrame(frameRef.current)
}
},
[]
)

return displayed
}

// Re-publish the part context with a smooth character-reveal, above
// DeferStreamingText so the reveal feeds the deferred markdown pipeline. Status
// stays running while revealing so the caret persists past the underlying part
// settling.
function SmoothStreamingText({ children }: { children: ReactNode }) {
const { text, status } = useMessagePartText()
const isRunning = status.type === 'running'
const revealed = useSmoothReveal(text, isRunning)

return (
<TextMessagePartProvider isRunning={isRunning || revealed !== text} text={revealed}>
{children}
</TextMessagePartProvider>
)
}

/**
* 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
Expand Down Expand Up @@ -498,7 +349,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'

Expand Down Expand Up @@ -626,26 +477,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
Expand All @@ -654,21 +508,13 @@ interface MarkdownTextContentProps extends MarkdownTextSurfaceProps {
export function MarkdownTextContent({ isRunning, text, ...surfaceProps }: MarkdownTextContentProps) {
return (
<TextMessagePartProvider isRunning={isRunning} text={text}>
<SmoothStreamingText>
<DeferStreamingText>
<MarkdownTextSurface {...surfaceProps} />
</DeferStreamingText>
</SmoothStreamingText>
<MarkdownTextSurface defer smooth={SMOOTH_OPTIONS} {...surfaceProps} />

@OutThisLife OutThisLife Jun 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The pipeline reorder (smooth → defer → preprocess ⟶ preprocess → smooth → defer) is not intermediate-frame identical for reasoning text, even though the final render is. In the built-in StreamdownTextPrimitive, useSmooth runs on the already-remend-repaired full text, and the revealed prefix is then rendered with parseIncompleteMarkdown: false and no per-prefix repair.

Consequence (reasoning only — this MarkdownTextContent path with smooth; the body-text MarkdownText is defer-only and genuinely equivalent): the OLD code re-ran tailBoundedRemend on each revealed prefix, so an incomplete **bold showed up already-styled during the reveal. The NEW code reveals a prefix of the repaired text, so the opener is shown before its closer → **, backtick, $ briefly flash as literal syntax at the typewriter frontier until the reveal catches up.

Questions:

  1. Was this transient flicker observed during your manual reasoning-stream testing, and is it considered acceptable? (We think it's likely fine — it's cosmetic and reasoning-only — but the PR body's "end result is identical" should be narrowed to final state.)
  2. SMOOTH_OPTIONS matches the old constants, but the built-in TextStreamAnimator uses a different rate algorithm than the deleted useSmoothReveal (incl. a maxCharIntervalMs default of 5ms). Did the reveal cadence visibly match the old feel side-by-side, or just approximately?

@okisdev okisdev Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

confirmed, and you're right that this is the one place "end result is identical" overclaims. it's intermediate-frame divergent on the reasoning path only (MarkdownTextContent with smooth; body text is defer-only and genuinely equivalent). the mechanism is exactly as you describe: the built-in runs preprocess on the full text, useSmooth slices a prefix of the already-repaired string, and because we pass parseMarkdownIntoBlocksFn the built-in tail-remend is gated off (!parseMarkdownIntoBlocksFn) with parseIncompleteMarkdown false, so the reveal frontier is never re-repaired and an unclosed **, backtick, or $ shows raw until its closer is revealed.

treating it as acceptable: cosmetic and reasoning-only. narrowed the PR body to final-state equivalence and fixed the misleading inline comment in fbe98ca. if we decide the flicker isn't shippable, the fix is either dropping parseMarkdownIntoBlocksFn on the smooth surface (costs the reasoning block-parse cache) or setting parseIncompleteMarkdown non-false there (full remend per flush on reasoning).

on cadence (Q2): approximately, not exact. maxCharIntervalMs is unset so it defaults to 5ms against the old ~33ms floor, so the tail reveals faster. if the feel is off we can set maxCharIntervalMs explicitly to match.

</TextMessagePartProvider>
)
}

const MarkdownTextImpl = () => {
return (
<DeferStreamingText>
<MarkdownTextSurface />
</DeferStreamingText>
)
return <MarkdownTextSurface defer />
}

export const MarkdownText = memo(MarkdownTextImpl)
4 changes: 2 additions & 2 deletions apps/desktop/src/components/assistant-ui/message-render-boundary.test.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand All @@ -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(
Expand Down
4 changes: 2 additions & 2 deletions apps/desktop/src/components/assistant-ui/message-render-boundary.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion apps/desktop/src/components/assistant-ui/thread/user-message-edit.test.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
18 changes: 14 additions & 4 deletions apps/desktop/src/lib/incremental-external-store-runtime.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import {
import {
type AssistantRuntime,
type ExternalStoreAdapter,
fromThreadMessageLike,
generateId,
type ThreadMessage,
useRuntimeAdapters
} from '@assistant-ui/react'
Expand Down Expand Up @@ -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 } },

@OutThisLife OutThisLife Jun 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The inlining of the removed appendOptimisticMessage looks faithful (generateId + fromThreadMessageLike(..., { type: 'running' }) + addOrUpdateMessage), and adding metadata.isOptimistic: true for the new off-branch eviction is a sensible, necessary adaptation rather than a gratuitous change.

This is the highest-risk spot since it's the one place we hand-reimplement a core method against a core that jumped two minors. Questions:

  1. Did you confirm the placeholder is correctly evicted (no ghost/empty assistant bubble) across: stop mid-stream, regenerate/reload, edit-and-resend, and branch switch? Those are the paths where a stale optimistic id historically leaks.
  2. Is metadata.isOptimistic a documented/stable contract in core 0.2.x's eviction logic, or an implementation detail we're relying on? If the latter, can we add a short comment pointing at the upstream code so a future core bump doesn't silently break cleanup?

@okisdev okisdev Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

faithful, yes. on whether isOptimistic is a contract or an implementation detail: it's a documented public field, not something we're reaching past. core's published types carry the JSDoc on ThreadMessage.metadata.isOptimistic: "Marks a client-side optimistic placeholder. Such messages are evicted once off the head branch and are never persisted." it landed in #4162, and core's own external-store runtime uses the identical generateId() + fromThreadMessageLike({ ..., metadata: { isOptimistic: true } }, id, { type: 'running' }) pattern, so we're matching core's usage exactly rather than relying on an internal detail. added a short invariant note at the call site in 27f8d78 flagging that dependency (without a PR-pointer, since the contract lives in core's type) so a future core bump that touches it gets caught here.

on eviction (Q1): verified clean against the installed 0.2.18. stop bypasses core's cancelRun (our Stop wires onCancel directly to the gateway), so the placeholder is cleaned by our own sync instead: hasUpcomingMessage is false on the next snapshot, so we deleteMessage the tracked optimistic id and don't re-add it. branch-switch routes through core switchToBranch, which early-returns while running, and when idle there's no optimistic to evict; within every sync the placeholder is the on-branch head at each resetHead, so core's off-branch eviction never touches it. no double-evict or ghost bubble across stop, regenerate, edit-and-resend, or branch switch.

optimisticId,
{ type: 'running' }
)
)
self._assistantOptimisticId = optimisticId
}

this.repository.resetHead(self._assistantOptimisticId ?? messages.at(-1)?.id ?? null)
Expand Down
Loading