fix(desktop): no per-paragraph action bars on sealed interim messages

Since #65919 the live view seals each chunk of mid-turn assistant
commentary (message.interim) as its own finalized bubble. Every bubble
with visible text renders the hover action footer, so a tool-heavy turn
grew a copy/refresh bar under almost every paragraph — and the live
render didn't match rehydration, which merges the turn into one bubble.

Mark sealed interim bubbles with ChatMessage.interim, carry the flag
into the runtime message metadata (custom.interim), and skip the
AssistantFooter for them. The turn's final reply keeps the footer; a
previewed final that settles onto an interim bubble clears the mark so
the settled reply regains its actions. interim joins COMPARED_FIELDS /
chatMessagesEquivalent so flipping it repaints.

Also fix an id-collision flake this surfaced: stream/interim bubble ids
were Date.now()-only, so an interim seal and the next segment's first
delta in the same millisecond reused the id and the new segment appended
into the sealed bubble. Ids now include a monotonic sequence.
This commit is contained in:
Brooklyn Nicholson 2026-07-22 11:20:25 -05:00
parent cbc1054e23
commit 72dd01c553
7 changed files with 127 additions and 21 deletions

View file

@ -55,6 +55,13 @@ interface QueuedStreamDeltas {
reasoning: string
}
// Date.now() alone can collide when an interim seal and the next segment's
// first delta land in the same millisecond — the new segment would then find
// the sealed bubble by id and append into it instead of starting fresh.
let streamMessageSeq = 0
const nextStreamMessageId = (prefix: string) => `${prefix}-${Date.now()}-${++streamMessageSeq}`
export function useMessageStream({
activeGatewayProfile = 'default',
activeSessionIdRef,
@ -91,7 +98,7 @@ export function useMessageStream({
return state
}
const streamId = state.streamId ?? `assistant-stream-${Date.now()}`
const streamId = state.streamId ?? nextStreamMessageId('assistant-stream')
const groupId = state.pendingBranchGroup ?? undefined
const prev = state.messages
let nextMessages: ChatMessage[]
@ -397,19 +404,21 @@ export function useMessageStream({
let nextMessages = state.messages
if (streamId && nextMessages.some(m => m.id === streamId)) {
// Finalize the existing streaming bubble in place
// Seal the streaming bubble in place, marked interim so it renders
// without an action footer (see ChatMessage.interim).
nextMessages = nextMessages.map(m =>
m.id === streamId ? { ...m, parts: replaceTextPart(m.parts), pending: false } : m
m.id === streamId ? { ...m, parts: replaceTextPart(m.parts), pending: false, interim: true } : m
)
} else {
// No streaming bubble — create a standalone interim message
nextMessages = [
...nextMessages,
{
id: `assistant-interim-${Date.now()}`,
id: nextStreamMessageId('assistant-interim'),
role: 'assistant' as const,
parts: [assistantTextPart(authoritativeText)],
pending: false,
interim: true,
branchGroupId: state.pendingBranchGroup ?? undefined
}
]
@ -459,19 +468,15 @@ export function useMessageStream({
return mergeFinalAssistantText(parts, visibleFinalText)
}
const completeMessage = (message: ChatMessage): ChatMessage =>
completionError
? {
...message,
error: completionError,
parts: message.parts.filter(part => part.type !== 'text'),
pending: false
}
: {
...message,
parts: replaceTextPart(message.parts),
pending: false
}
// Settling the final response onto a bubble makes it the turn's real
// reply — clear `interim` so it regains the action footer.
const completeMessage = (message: ChatMessage): ChatMessage => {
const settled = { ...message, pending: false, interim: false }
return completionError
? { ...settled, error: completionError, parts: message.parts.filter(part => part.type !== 'text') }
: { ...settled, parts: replaceTextPart(message.parts) }
}
const newAssistantFromCompletion = (): ChatMessage => ({
id: `assistant-${Date.now()}`,

View file

@ -111,6 +111,36 @@ describe('useMessageStream interim text sealing', () => {
expect(texts).toContain('All checks passed.')
})
it('marks sealed interim bubbles interim and leaves the final reply unmarked', async () => {
await mountStream()
await start()
await delta('Let me check the files.')
await interim('Let me check the files.')
await delta('Now the second pass.')
await interim('Now the second pass.')
await complete('All done.')
const assistants = getState().messages.filter(m => m.role === 'assistant' && !m.hidden)
const byText = (text: string) => assistants.find(m => chatMessageText(m) === text)
expect(byText('Let me check the files.')?.interim).toBe(true)
expect(byText('Now the second pass.')?.interim).toBe(true)
expect(byText('All done.')?.interim).toBeFalsy()
})
it('clears the interim mark when a previewed final settles onto the interim bubble', async () => {
await mountStream()
await start()
await interim('same reply')
await completePreviewed('same reply')
const assistants = getState().messages.filter(m => m.role === 'assistant' && !m.hidden)
expect(assistants).toHaveLength(1)
expect(assistants[0].interim).toBeFalsy()
})
it('dedupes interim text when the final response includes it', async () => {
await mountStream()
await start()

View file

@ -73,7 +73,7 @@ const _chatMessageFieldsExhaustive: {
[K in Exclude<keyof ChatMessage, (typeof COMPARED_FIELDS)[number] | (typeof IGNORED_FIELDS)[number]>]: never
} = {}
const COMPARED_FIELDS = ['id', 'role', 'pending', 'error', 'hidden', 'branchGroupId'] as const
const COMPARED_FIELDS = ['id', 'role', 'pending', 'error', 'hidden', 'branchGroupId', 'interim'] as const
const IGNORED_FIELDS = ['timestamp', 'attachmentRefs', 'parts'] as const
// Compile-time check: every ChatMessagePart discriminant must be handled by
@ -154,7 +154,10 @@ export function chatMessagesEquivalent(a: ChatMessage, b: ChatMessage): boolean
a.pending !== b.pending ||
a.error !== b.error ||
a.hidden !== b.hidden ||
a.branchGroupId !== b.branchGroupId
a.branchGroupId !== b.branchGroupId ||
// Interim gates the action footer, so flipping it must repaint (e.g. a
// previewed final settling onto a sealed interim bubble restores the bar).
(a.interim ?? false) !== (b.interim ?? false)
) {
return false
}

View file

@ -65,6 +65,10 @@ export const AssistantMessage: FC<{
const isRunning = messageStatus === 'running'
const isPlaceholder = useAuiState(s => s.message.status?.type === 'running' && s.message.content.length === 0)
const hasVisibleText = useAuiState(s => contentHasVisibleText(s.message.content))
// Sealed mid-turn commentary keeps its text but not the footer, so a
// tool-heavy turn doesn't grow a copy/refresh bar per paragraph (see
// ChatMessage.interim).
const isInterim = useAuiState(s => s.message.metadata?.custom?.interim === true)
// Preview targets only materialize once the turn completes — while running
// the selector returns '' (stable), so per-token flushes skip the regex
@ -130,7 +134,7 @@ export const AssistantMessage: FC<{
</ErrorPrimitive.Root>
</MessagePrimitive.Error>
</div>
{hasVisibleText && (
{hasVisibleText && !isInterim && (
<AssistantFooter getMessageText={getMessageText} messageId={messageId} onBranchInNewChat={onBranchInNewChat} />
)}
</MessagePrimitive.Root>

View file

@ -328,6 +328,37 @@ function MessageHarness({ message }: { message: ThreadMessage }) {
)
}
function TranscriptHarness({ messages }: { messages: ThreadMessage[] }) {
const runtime = useExternalStoreRuntime<ThreadMessage>({
messages,
isRunning: false,
onNew: async () => {}
})
return (
<AssistantRuntimeProvider runtime={runtime}>
<Thread />
</AssistantRuntimeProvider>
)
}
function assistantInterimMessage(text: string, id = 'assistant-interim-1'): ThreadMessage {
return {
id,
role: 'assistant',
content: [{ type: 'text', text }],
status: { type: 'complete', reason: 'stop' },
createdAt,
metadata: {
unstable_state: null,
unstable_annotations: [],
unstable_data: [],
steps: [],
custom: { interim: true }
}
} as ThreadMessage
}
function RunningMessageHarness({ message }: { message: ThreadMessage }) {
const runtime = useExternalStoreRuntime<ThreadMessage>({
messages: [message],
@ -455,6 +486,34 @@ describe('assistant-ui streaming renderer', () => {
expect(container.querySelector('[data-slot="aui_composer-clearance"]')).toBeNull()
})
it('suppresses the action footer on sealed interim messages, keeping it on the final reply', () => {
const { container } = render(
<TranscriptHarness
messages={[
userMessage(),
assistantInterimMessage('Let me check the files.'),
assistantInterimMessage('Now applying the patch.', 'assistant-interim-2'),
assistantMessage('All done — patch applied.', false)
]}
/>
)
// Interim commentary stays visible…
expect(container.textContent).toContain('Let me check the files.')
expect(container.textContent).toContain('Now applying the patch.')
expect(container.textContent).toContain('All done — patch applied.')
// …but only the turn's final reply carries the copy/refresh action bar.
const actionBars = container.querySelectorAll('[data-slot="aui_msg-actions"]')
expect(actionBars).toHaveLength(1)
const finalRoot = [...container.querySelectorAll('[data-slot="aui_assistant-message-root"]')].find(root =>
root.textContent?.includes('All done — patch applied.')
)
expect(finalRoot?.querySelector('[data-slot="aui_msg-actions"]')).toBeTruthy()
})
it('renders assistant provider errors inline', () => {
render(<MessageHarness message={assistantErrorMessage('OpenRouter rejected the request (403).')} />)

View file

@ -17,6 +17,10 @@ export type ChatMessage = {
error?: string
branchGroupId?: string
hidden?: boolean
/** Sealed mid-turn commentary (`message.interim`) rendered without the
* action footer so only the turn's final reply carries copy/refresh, and
* the live view matches rehydration (which merges the turn into one bubble). */
interim?: boolean
/** Composer attachment ref strings (`@file:...`, `@image:...`) sent with this user message. */
attachmentRefs?: string[]
}

View file

@ -381,7 +381,8 @@ export function toRuntimeMessage(message: ChatMessage): ThreadMessage {
unstable_annotations: [],
unstable_data: [],
steps: [],
custom: {}
// Carries ChatMessage.interim to AssistantMessage's footer gate.
custom: message.interim ? { interim: true } : {}
}
} as ThreadMessage
}