combine UIMessages across page boundaries

Author: ianmacartney

src/deltas.ts

@@ -18,7 +18,6 @@ import {
   type StreamMessage,
 } from "./validators.js";
 
-
 export function blankUIMessage<METADATA = unknown>(
   streamMessage: StreamMessage & { metadata?: METADATA },
   threadId: string,
@@ -538,3 +537,65 @@ function mergeProviderMetadata(
   }
   return merged;
 }
+
+export function combineUIMessages(messages: UIMessage[]): UIMessage[] {
+  const combined = messages.reduce((acc, message) => {
+    if (!acc.length) {
+      return [message];
+    }
+    const previous = acc.at(-1)!;
+    if (previous.role !== message.role) {
+      acc.push(message);
+      return acc;
+    }
+    // We will replace it with a combined message
+    acc.pop();
+    const newParts = [...previous.parts];
+    for (const part of message.parts) {
+      const toolCallId = getToolCallId(part);
+      if (!toolCallId) {
+        newParts.push(part);
+        continue;
+      }
+      const previousPartIndex = newParts.findIndex(
+        (p) => getToolCallId(p) === toolCallId,
+      );
+      const previousPart = newParts.splice(previousPartIndex, 1)[0];
+      if (!previousPart) {
+        newParts.push(part);
+        continue;
+      }
+      newParts.push(mergeParts(previousPart, part));
+    }
+    acc.push({
+      ...previous,
+      ...pick(message, ["status", "metadata", "agentName"]),
+      parts: newParts,
+      text: newParts
+        .filter((p) => p.type === "text")
+        .map((p) => p.text)
+        .join(""),
+    });
+    return acc;
+  }, [] as UIMessage[]);
+  return combined;
+}
+
+function getToolCallId(
+  part: UIMessage["parts"][number] & { toolCallId?: string },
+) {
+  return part.toolCallId;
+}
+
+function mergeParts(
+  previousPart: UIMessage["parts"][number],
+  part: UIMessage["parts"][number],
+): UIMessage["parts"][number] {
+  const merged: Record<string, unknown> = { ...previousPart };
+  for (const [key, value] of Object.entries(part)) {
+    if (value !== undefined) {
+      merged[key] = value;
+    }
+  }
+  return merged as ToolUIPart | DynamicToolUIPart;
+}

src/react/useUIMessages.ts

@@ -22,6 +22,7 @@ import type { StreamQuery } from "./types.js";
 import { type UIMessage, type UIStatus } from "../UIMessages.js";
 import { sorted } from "../shared.js";
 import { useStreamingUIMessages } from "./useStreamingUIMessages.js";
+import { combineUIMessages } from "../deltas.js";
 
 export type UIMessageLike = {
   order: number;
@@ -155,9 +156,11 @@ 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(paginated.results, streamMessages ?? []),
+      results: dedupeMessages(combined, streamMessages ?? []),
     };
   }, [paginated, streamMessages]);