perf(desktop): reconcile only the messages that actually moved

syncRepositoryIncrementally rewrote the entire transcript on every adapter
snapshot: one addOrUpdateMessage per message, a second full export to find
deletions, and an unconditional resetHead. During streaming that fires ~30x a
second, so the per-delta cost scaled with how long the conversation already
was — the reported symptom.

Now that settled messages keep reference identity, an identity check is a
sound "did this change?" test. Write only the items whose message or parentId
moved, and skip resetHead when the head did not move (it prunes the head's
descendants, so calling it needlessly is not free). Anything the fast path
cannot prove safe — a disjoint session swap, a changed message count, an id
with no repository entry — falls through to the original full rebuild.

Tests cover the perf contract (one write for one delta, zero for a no-op) and
the correctness cases the fast path must not break: appends, authoritative
deletion, disjoint session replacement, branch re-parenting, and an explicit
headId rewind. The two perf assertions fail on main.

Co-authored-by: Jakub Wolniewicz <frizikk@users.noreply.github.com>
This commit is contained in:
Brooklyn Nicholson 2026-07-26 02:56:08 -05:00
parent beb3f566ba
commit adf47ca832
2 changed files with 204 additions and 6 deletions

View file

@ -0,0 +1,144 @@
import { fromThreadMessageLike, getAutoStatus, MessageRepository } from '@assistant-ui/core/internal'
import type { ExportedMessageRepository, ThreadMessage } from '@assistant-ui/react'
import { describe, expect, it, vi } from 'vitest'
import { syncRepositoryIncrementally } from './incremental-external-store-runtime'
const STATUS = getAutoStatus(false, false, false, false, undefined)
function message(id: string, text: string): ThreadMessage {
return fromThreadMessageLike({ role: 'assistant', content: [{ type: 'text', text }] }, id, STATUS)
}
/** A real MessageRepository behind the same shape syncRepositoryIncrementally drives. */
function runtimeWith(items: { message: ThreadMessage; parentId: string | null }[]) {
const repository = new MessageRepository()
for (const { message: item, parentId } of items) {
repository.addOrUpdateMessage(parentId, item)
}
if (items.length > 0) {
repository.resetHead(items.at(-1)?.message.id ?? null)
}
return { repository } as unknown as Parameters<typeof syncRepositoryIncrementally>[0]
}
function chain(messages: ThreadMessage[]) {
return messages.map((item, index) => ({
message: item,
parentId: index === 0 ? null : messages[index - 1].id
}))
}
function exported(items: { message: ThreadMessage; parentId: string | null }[]): ExportedMessageRepository {
return { headId: items.at(-1)?.message.id ?? null, messages: items }
}
describe('syncRepositoryIncrementally', () => {
it('writes only the changed tail instead of the whole transcript', () => {
const settled = Array.from({ length: 200 }, (_, index) => message(`m-${index}`, `body ${index}`))
const items = chain(settled)
const runtime = runtimeWith(items)
const repository = (runtime as unknown as { repository: MessageRepository }).repository
const addOrUpdate = vi.spyOn(repository, 'addOrUpdateMessage')
const resetHead = vi.spyOn(repository, 'resetHead')
// One streamed delta: the tail grows, every settled message keeps identity.
const nextTail = message('m-199', 'body 199 + delta')
const nextItems = [...items.slice(0, -1), { message: nextTail, parentId: 'm-198' }]
const result = syncRepositoryIncrementally(runtime, exported(nextItems))
expect(addOrUpdate).toHaveBeenCalledTimes(1)
expect(addOrUpdate).toHaveBeenCalledWith('m-198', nextTail)
// The head did not move, so the descendant-pruning reset is skipped.
expect(resetHead).not.toHaveBeenCalled()
expect(result).toHaveLength(200)
expect(result.at(-1)).toBe(nextTail)
})
it('does nothing at all when the transcript is unchanged', () => {
const items = chain([message('a', 'one'), message('b', 'two')])
const runtime = runtimeWith(items)
const repository = (runtime as unknown as { repository: MessageRepository }).repository
const addOrUpdate = vi.spyOn(repository, 'addOrUpdateMessage')
const deleteMessage = vi.spyOn(repository, 'deleteMessage')
syncRepositoryIncrementally(runtime, exported(items))
expect(addOrUpdate).not.toHaveBeenCalled()
expect(deleteMessage).not.toHaveBeenCalled()
})
it('appends a new message through the full path', () => {
const first = message('a', 'one')
const items = chain([first])
const runtime = runtimeWith(items)
const second = message('b', 'two')
const result = syncRepositoryIncrementally(runtime, exported(chain([first, second])))
expect(result.map(item => item.id)).toEqual(['a', 'b'])
})
it('honours an authoritative deletion', () => {
const a = message('a', 'one')
const b = message('b', 'two')
const c = message('c', 'three')
const runtime = runtimeWith(chain([a, b, c]))
const result = syncRepositoryIncrementally(runtime, exported(chain([a, b])))
expect(result.map(item => item.id)).toEqual(['a', 'b'])
})
it('rebuilds cleanly when a disjoint transcript is swapped in', () => {
const runtime = runtimeWith(chain([message('old-1', 'one'), message('old-2', 'two')]))
const next = chain([message('new-1', 'alpha'), message('new-2', 'beta')])
const result = syncRepositoryIncrementally(runtime, exported(next))
expect(result.map(item => item.id)).toEqual(['new-1', 'new-2'])
})
it('re-parents a message when its branch parent changes', () => {
const root = message('root', 'root')
const a = message('a', 'a')
const b = message('b', 'b')
const runtime = runtimeWith([
{ message: root, parentId: null },
{ message: a, parentId: 'root' },
{ message: b, parentId: 'a' }
])
// Same ids and same message objects, but `b` moves onto a sibling branch.
const result = syncRepositoryIncrementally(runtime, {
headId: 'b',
messages: [
{ message: root, parentId: null },
{ message: a, parentId: 'root' },
{ message: b, parentId: 'root' }
]
})
expect(result.map(item => item.id)).toEqual(['root', 'b'])
})
it('moves the head when an explicit headId rewinds the branch', () => {
const a = message('a', 'one')
const b = message('b', 'two')
const runtime = runtimeWith(chain([a, b]))
const result = syncRepositoryIncrementally(runtime, {
headId: 'a',
messages: chain([a, b])
})
expect(result.map(item => item.id)).toEqual(['a'])
})
})

View file

@ -35,25 +35,81 @@ const shallowEqual = (a: object, b: object): boolean => {
const getThreadListAdapter = (store: ExternalStoreAdapter) => store.adapters?.threadList ?? {}
function syncRepositoryIncrementally(
/**
* Write only the items whose (message, parentId) pair actually moved.
*
* `useRuntimeMessageRepository` caches normalized ThreadMessages by source
* identity, so a settled turn keeps the SAME object across renders. That makes
* an identity check a sound "did this change?" test: during streaming exactly
* one item the growing tail differs, and the other N-1 writes were pure
* overhead that grew with transcript length.
*
* Returns false when the export is stale (an id in `existing` is gone, or an
* incoming message has no repository entry yet), so the caller falls back to
* the full rebuild rather than guessing.
*/
function applyChangedMessages(
repository: ExternalStoreThreadRuntimeCore['repository'],
existing: readonly { message: ThreadMessage; parentId: string | null }[],
incoming: readonly { message: ThreadMessage; parentId: string | null }[]
): boolean {
if (existing.length !== incoming.length) {
return false
}
const existingById = new Map(existing.map(item => [item.message.id, item]))
for (const item of incoming) {
const current = existingById.get(item.message.id)
if (!current) {
return false
}
// Reference identity, not deep equality: the conversion cache guarantees a
// stable object for an unchanged turn, and a changed turn is a new object.
if (current.message !== item.message || current.parentId !== item.parentId) {
repository.addOrUpdateMessage(item.parentId, item.message)
}
}
return true
}
export function syncRepositoryIncrementally(
runtime: ExternalStoreThreadRuntimeCore,
messageRepository: NonNullable<ExternalStoreAdapter['messageRepository']>
): readonly ThreadMessage[] {
const repository = (runtime as unknown as { repository: ExternalStoreThreadRuntimeCore['repository'] }).repository
const incomingIds = new Set(messageRepository.messages.map(({ message }) => message.id))
const incoming = messageRepository.messages
const existing = repository.export().messages
const headId = messageRepository.headId ?? incoming.at(-1)?.message.id ?? null
// A thread switch swaps in a fully-DISJOINT transcript (no id carries over).
// Reconciling two unrelated trees in place — grafting the new chain onto the
// old one, then pruning — can strand a stale head/branch, so there's nothing
// to preserve: clear the tree first (leaves→root), then rebuild clean.
if (existing.length > 0 && !existing.some(({ message }) => incomingIds.has(message.id))) {
const incomingIds = new Set(incoming.map(({ message }) => message.id))
const disjoint = existing.length > 0 && !existing.some(({ message }) => incomingIds.has(message.id))
// Steady-state streaming: same message set, one item changed. Skip the
// whole-transcript rewrite, the prune scan, and the second export. resetHead
// deletes the head's descendants, so it only runs when the head really moved.
if (!disjoint && applyChangedMessages(repository, existing, incoming)) {
if (repository.headId !== headId) {
repository.resetHead(headId)
}
return repository.getMessages()
}
if (disjoint) {
for (const { message } of [...existing].reverse()) {
repository.deleteMessage(message.id)
}
}
for (const { message, parentId } of messageRepository.messages) {
for (const { message, parentId } of incoming) {
repository.addOrUpdateMessage(parentId, message)
}
@ -63,8 +119,6 @@ function syncRepositoryIncrementally(
}
}
const headId = messageRepository.headId ?? messageRepository.messages.at(-1)?.message.id ?? null
repository.resetHead(headId)
return repository.getMessages()