diff --git a/src/react/useUIMessages.test.ts b/src/react/useUIMessages.test.ts index a4ec3beb..b6ef9ba2 100644 --- a/src/react/useUIMessages.test.ts +++ b/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" }, + ]); + }); +}); diff --git a/src/react/useUIMessages.ts b/src/react/useUIMessages.ts index bba3df3d..dda2fc22 100644 --- a/src/react/useUIMessages.ts +++ b/src/react/useUIMessages.ts @@ -155,17 +155,25 @@ export function useUIMessages>( ); 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; } +export function mergeUIMessages( + 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;