diff --git a/apps/desktop/src/lib/incremental-external-store-runtime.test.ts b/apps/desktop/src/lib/incremental-external-store-runtime.test.ts new file mode 100644 index 000000000000..0c50d33370ad --- /dev/null +++ b/apps/desktop/src/lib/incremental-external-store-runtime.test.ts @@ -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[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']) + }) +}) diff --git a/apps/desktop/src/lib/incremental-external-store-runtime.ts b/apps/desktop/src/lib/incremental-external-store-runtime.ts index f4a8191b5f54..0df3ed9b2e09 100644 --- a/apps/desktop/src/lib/incremental-external-store-runtime.ts +++ b/apps/desktop/src/lib/incremental-external-store-runtime.ts @@ -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 ): 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()