fix(sync): reconcile completed turns after reconnect

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")

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