diff --git a/packages/ui/src/sync/__tests__/materialization.test.ts b/packages/ui/src/sync/__tests__/materialization.test.ts index 9cc5a67f..19251632 100644 --- a/packages/ui/src/sync/__tests__/materialization.test.ts +++ b/packages/ui/src/sync/__tests__/materialization.test.ts @@ -144,6 +144,37 @@ describe("materializeSessionSnapshots", () => { expect(result.message.ses_1[0]).not.toBe(staleMessage) }) + test("replaces a locally aborted assistant message with the authoritative completed snapshot", () => { + const unfinishedMessage = message("msg_1") + if (unfinishedMessage.role !== "assistant") throw new Error("Expected assistant fixture") + const abortedMessage: Message = { + ...unfinishedMessage, + time: { created: 1, completed: 5000 }, + error: { name: "MessageAbortedError", data: { message: "aborted" } }, + } + const completedMessage: Message = { + ...unfinishedMessage, + time: { created: 1, completed: 4000 }, + } + const state = { + message: { ses_1: [abortedMessage] }, + part: { msg_1: [] }, + } + + const result = materializeSessionSnapshots( + state, + "ses_1", + [{ info: completedMessage, parts: [] }], + ) + + const reconciled = result.message.ses_1[0] + expect(reconciled).toBe(completedMessage) + expect(reconciled?.role).toBe("assistant") + if (reconciled?.role !== "assistant") throw new Error("Expected assistant result") + expect("error" in reconciled).toBe(false) + expect(reconciled.time.completed).toBe(4000) + }) + test("does not preserve omitted optimistic user text parts beside server snapshot parts", () => { const optimisticPart = { id: "prt_optimistic", messageID: "msg_1", type: "text", text: "Hello" } as Part const serverPart = part("prt_server", "msg_1", "text", "Hello") diff --git a/packages/ui/src/sync/materialization.ts b/packages/ui/src/sync/materialization.ts index 0747dd6e..89a653d2 100644 --- a/packages/ui/src/sync/materialization.ts +++ b/packages/ui/src/sync/materialization.ts @@ -270,7 +270,21 @@ export function materializeSessionSnapshots( const snapshots = nextMessages.map((message) => recordsByMessageID.get(message.id)!) const existingMessages = state.message[sessionID] const currentMessages = existingMessages ?? [] - const messages = mergeMessages(currentMessages, nextMessages) + const incomingByID = new Map(nextMessages.map((message) => [message.id, message] as const)) + let reconciledCurrentMessages = currentMessages + for (let index = 0; index < currentMessages.length; index += 1) { + const existing = currentMessages[index] + const incoming = incomingByID.get(existing.id) + if ( + existing.role !== "assistant" + || existing.error?.name !== "MessageAbortedError" + || incoming?.role !== "assistant" + || incoming.time.completed === undefined + ) continue + if (reconciledCurrentMessages === currentMessages) reconciledCurrentMessages = [...currentMessages] + reconciledCurrentMessages[index] = incoming + } + const messages = mergeMessages(reconciledCurrentMessages, nextMessages) const messagesChanged = messages !== currentMessages || (existingMessages === undefined && snapshots.length === 0) let partsChanged = false