make uimessages from scratch each time

Author: ianmacartney

example/ui/main.tsx

@@ -9,6 +9,7 @@ import FilesImages from "./files/FilesImages";
 import RateLimiting from "./rate_limiting/RateLimiting";
 import { WeatherFashion } from "./workflows/WeatherFashion";
 import RagBasic from "./rag/RagBasic";
+import StreamArray from "./objects/StreamArray";
 
 const convex = new ConvexReactClient(import.meta.env.VITE_CONVEX_URL as string);
 

src/deltas.ts

@@ -188,9 +188,6 @@ export function getParts<T extends StreamDelta["parts"][number]>(
     }
     if (cursor !== delta.start) {
       if (cursor >= delta.end) {
-        console.debug(
-          `Got duplicate delta for stream ${delta.streamId} at ${delta.start}`,
-        );
         continue;
       } else if (cursor < delta.start) {
         console.warn(

src/react/useDeltaStreams.ts

@@ -1,9 +1,10 @@
 "use client";
 
 import type { StreamQuery, StreamQueryArgs } from "./types.js";
+import type { SyncStreamsReturnValue } from "../client/types.js";
 import type { FunctionArgs } from "convex/server";
 import type { StreamArgs, StreamDelta, StreamMessage } from "../validators.js";
-import type { SyncStreamsReturnValue } from "@convex-dev/agent";
+import { sorted } from "../shared.js";
 import { useQuery } from "convex/react";
 import { useState } from "react";
 import { assert } from "convex-helpers";
@@ -75,10 +76,12 @@ export function useDeltaStreams<
       ? undefined
       : !streamList
         ? state.deltaStreams?.map(({ streamMessage }) => streamMessage)
-        : streamList.streams.messages.filter(
-            ({ streamId, order }) =>
-              !options?.skipStreamIds?.includes(streamId) &&
-              (!options?.startOrder || order >= options.startOrder),
+        : sorted(
+            streamList.streams.messages.filter(
+              ({ streamId, order }) =>
+                !options?.skipStreamIds?.includes(streamId) &&
+                (!options?.startOrder || order >= options.startOrder),
+            ),
           );
 
   // Get the deltas for all the active streams, if any.

src/react/useStreamingUIMessages.ts

@@ -6,9 +6,8 @@ import { type UIMessage } from "../UIMessages.js";
 import {
   blankUIMessage,
   getParts,
-  statusFromStreamStatus,
   updateFromUIMessageChunks,
-  updateFromTextStreamParts,
+  deriveUIMessagesFromTextStreamParts,
 } from "../deltas.js";
 import { useDeltaStreams } from "./useDeltaStreams.js";
 
@@ -66,75 +65,68 @@ export function useStreamingUIMessages<
 
   useEffect(() => {
     if (!streams) return;
+    // return if there are no new deltas beyond the cursors
+    let noNewDeltas = true;
+    for (const stream of streams) {
+      const lastDelta = stream.deltas.at(-1);
+      const cursor = messageState[stream.streamMessage.streamId]?.cursor;
+      if (!cursor) {
+        noNewDeltas = false;
+        break;
+      }
+      if (lastDelta && lastDelta.start >= cursor) {
+        noNewDeltas = false;
+        break;
+      }
+    }
+    if (noNewDeltas) {
+      return;
+    }
     const abortController = new AbortController();
     void (async () => {
-      let changed = false;
       const newMessageState: Record<
         string,
         {
           uiMessage: UIMessage<METADATA, DATA_PARTS, TOOLS>;
           cursor: number;
         }
-      > = {};
-      for (const stream of streams) {
-        if (abortController.signal.aborted) return;
-        const oldState = messageState[stream.streamMessage.streamId];
-        let uiMessage = oldState?.uiMessage;
-        if (!oldState) {
-          changed = true;
-          uiMessage = blankUIMessage(
-            stream.streamMessage,
-            threadId,
-          ) as UIMessage<METADATA, DATA_PARTS, TOOLS>;
-        }
-        const { parts, cursor } = getParts<UIMessageChunk>(
-          stream.deltas,
-          oldState?.cursor,
-        );
-        if (parts.length) {
-          changed = true;
-          if (stream.streamMessage.format === "UIMessageChunk") {
-            uiMessage = (await updateFromUIMessageChunks(
-              uiMessage,
-              parts,
-            )) as UIMessage<METADATA, DATA_PARTS, TOOLS>;
-            if (
-              uiMessage.status !== "failed" &&
-              uiMessage.status !== "success"
-            ) {
-              uiMessage.status = statusFromStreamStatus(
-                stream.streamMessage.status,
+      > = Object.fromEntries(
+        await Promise.all(
+          streams.map(async ({ deltas, streamMessage }) => {
+            const { parts, cursor } = getParts<UIMessageChunk>(deltas, 0);
+            if (streamMessage.format === "UIMessageChunk") {
+              // Unfortunately this can't handle resuming from a UIMessage and
+              // adding more chunks, so we re-create it from scratch each time.
+              const uiMessage = await updateFromUIMessageChunks(
+                blankUIMessage(streamMessage, threadId),
+                parts,
               );
+              return [
+                streamMessage.streamId,
+                {
+                  uiMessage,
+                  cursor,
+                },
+              ];
+            } else {
+              const [uiMessages] = deriveUIMessagesFromTextStreamParts(
+                threadId,
+                [streamMessage],
+                [],
+                deltas,
+              );
+              return [
+                streamMessage.streamId,
+                {
+                  uiMessage: uiMessages[0],
+                  cursor,
+                },
+              ];
             }
-          } else if (
-            stream.streamMessage.format === "TextStreamPart" ||
-            !stream.streamMessage.format
-          ) {
-            const updated = updateFromTextStreamParts(
-              threadId,
-              stream.streamMessage,
-              {
-                streamId: stream.streamMessage.streamId,
-                cursor,
-                message: uiMessage,
-              },
-              stream.deltas,
-            )[0];
-            uiMessage = updated.message as UIMessage<
-              METADATA,
-              DATA_PARTS,
-              TOOLS
-            >;
-          } else {
-            console.error("Unknown format", stream.streamMessage.format);
-          }
-        }
-        newMessageState[stream.streamMessage.streamId] = {
-          uiMessage,
-          cursor,
-        };
-      }
-      if (!changed || abortController.signal.aborted) return;
+          }),
+        ),
+      );
+      if (abortController.signal.aborted) return;
       setMessageState(newMessageState);
     })();
     return () => {