Fix streaming UI message dedupe

src/react/useUIMessages.test.ts

@@ -1,5 +1,9 @@
 import { describe, it, expect } from "vitest";
-import { dedupeMessages } from "./useUIMessages.js";
+import {
+  dedupeMessages,
+  mergeUIMessages,
+  type UIMessageLike,
+} from "./useUIMessages.js";
 
 type TestMessage = {
   order: number;
@@ -8,6 +12,39 @@ type TestMessage = {
   id: string;
 };
 
+type TestUIMessage = UIMessageLike & {
+  id: string;
+  key: string;
+  text: string;
+  _creationTime: number;
+};
+
+function testUIMessage({
+  id,
+  order,
+  stepOrder,
+  status,
+  text,
+}: {
+  id: string;
+  order: number;
+  stepOrder: number;
+  status: TestUIMessage["status"];
+  text: string;
+}): TestUIMessage {
+  return {
+    id,
+    key: `thread-${order}-${stepOrder}`,
+    order,
+    stepOrder,
+    status,
+    role: "assistant",
+    parts: [{ type: "text", text }],
+    text,
+    _creationTime: 0,
+  };
+}
+
 describe("dedupeMessages", () => {
   it("should prefer messages from messages list when streaming messages are absent", () => {
     const messages: TestMessage[] = [
@@ -253,3 +290,45 @@ describe("dedupeMessages", () => {
     expect(result[0].id).toBe("messages-success");
   });
 });
+
+describe("mergeUIMessages", () => {
+  it("dedupes streaming steps before combining assistant messages", () => {
+    const messages = [
+      testUIMessage({
+        id: "persisted-step-1",
+        order: 1,
+        stepOrder: 1,
+        status: "success",
+        text: "7",
+      }),
+      testUIMessage({
+        id: "persisted-step-2",
+        order: 1,
+        stepOrder: 2,
+        status: "success",
+        text: "8",
+      }),
+    ];
+    const streamMessages = [
+      testUIMessage({
+        id: "stream-step-2",
+        order: 1,
+        stepOrder: 2,
+        status: "streaming",
+        text: "8",
+      }),
+    ];
+
+    const result = mergeUIMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(1);
+    expect(result[0].id).toBe("persisted-step-1");
+    expect(result[0].stepOrder).toBe(1);
+    expect(result[0].status).toBe("success");
+    expect(result[0].text).toBe("7 8");
+    expect(result[0].parts).toEqual([
+      { type: "text", text: "7" },
+      { type: "text", text: "8" },
+    ]);
+  });
+});

src/react/useUIMessages.ts

@@ -155,17 +155,25 @@ export function useUIMessages<Query extends UIMessagesQuery<any, any>>(
   );
 
   const merged = useMemo(() => {
-    // Messages may have been split by pagination. Re-combine them here.
-    const combined = combineUIMessages(sorted(paginated.results));
     return {
       ...paginated,
-      results: dedupeMessages(combined, streamMessages ?? []),
+      results: mergeUIMessages(paginated.results, streamMessages ?? []),
     };
   }, [paginated, streamMessages]);
 
   return merged as UIMessagesQueryResult<Query>;
 }
 
+export function mergeUIMessages<M extends UIMessageLike>(
+  messages: M[],
+  streamMessages: M[],
+): M[] {
+  const deduped = dedupeMessages(sorted(messages), streamMessages);
+  // Messages may have been split by pagination. Re-combine them here
+  // after dedupe has had access to each message's original stepOrder.
+  return combineUIMessages(deduped as unknown as UIMessage[]) as unknown as M[];
+}
+
 export function dedupeMessages<
   M extends {
     order: number;