@convex-dev/agent 0.2.8-alpha.5 → 0.2.8-alpha.6

package/src/react/useStreamingUIMessages.ts

@@ -1,26 +1,24 @@
 "use client";
-import { omit } from "convex-helpers";
-import { useQuery } from "convex/react";
-import type { FunctionArgs } from "convex/server";
-import { useMemo, useState, useEffect, useRef } from "react";
-import {
-  readUIMessageStream,
-  type TextUIPart,
-  type UIDataTypes,
-  type UIMessageChunk,
-  type UITools,
-} from "ai";
-import type { SyncStreamsReturnValue } from "../client/types.js";
-import type { StreamArgs } from "../validators.js";
+import { useMemo, useState, useEffect } from "react";
+import { type UIDataTypes, type UIMessageChunk, type UITools } from "ai";
 import type { StreamQuery, StreamQueryArgs } from "./types.js";
 import { type UIMessage } from "../UIMessages.js";
 import {
   blankUIMessage,
   getParts,
+  updateFromUIMessageChunks,
   deriveUIMessagesFromTextStreamParts,
-  statusFromStreamStatus,
 } from "../deltas.js";
-import { sorted } from "../shared.js";
+import { useDeltaStreams } from "./useDeltaStreams.js";
+
+// Polyfill structuredClone to support readUIMessageStream on ReactNative
+if (!("structuredClone" in globalThis)) {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  void import("@ungap/structured-clone" as any).then(
+    ({ default: structuredClone }) =>
+      (globalThis.structuredClone = structuredClone),
+  );
+}
 
 /**
  * A hook that fetches streaming messages from a thread and converts them to UIMessages
@@ -51,252 +49,97 @@ export function useStreamingUIMessages<
   },
   // TODO: make generic on metadata, etc.
 ): UIMessage<METADATA, DATA_PARTS, TOOLS>[] | undefined {
-  const [uiMessages, setUIMessages] = useState<
-    Record<string, UIMessage<METADATA, DATA_PARTS, TOOLS>>
+  const [messageState, setMessageState] = useState<
+    Record<
+      string,
+      {
+        uiMessage: UIMessage<METADATA, DATA_PARTS, TOOLS>;
+        cursor: number;
+      }
+    >
   >({});
-  const [streamCursors, setStreamCursors] = useState<Record<string, number>>(
-    {},
-  );
-  const uiMessageStreamControllers = useRef<
-    Map<string, ReadableStreamDefaultController<UIMessageChunk>>
-  >(new Map());
-
-  const startOrder = options?.startOrder
-    ? // round down to the nearest 10 for some cache benefits
-      options.startOrder - (options.startOrder % 10)
-    : 0;
 
-  // Get all the active streams
-  const streamList = useQuery(
-    query,
-    args === "skip"
-      ? args
-      : ({
-          ...args,
-          streamArgs: {
-            kind: "list",
-            startOrder,
-          } as StreamArgs,
-        } as FunctionArgs<Query>),
-  ) as
-    | { streams: Extract<SyncStreamsReturnValue, { kind: "list" }> }
-    | undefined;
-
-  const minOrder = Math.max(
-    options?.startOrder ?? 0,
-    streamList?.streams.messages.length
-      ? Math.min(...streamList.streams.messages.map(({ order }) => order))
-      : 0,
-  );
-
-  // Get the cursors for all the active streams
-  const cursors = useMemo(() => {
-    if (!streamList?.streams) return [];
-    if (streamList.streams.kind !== "list") {
-      throw new Error("Expected list streams");
-    }
-    return streamList.streams.messages
-      .filter(({ streamId }) => !options?.skipStreamIds?.includes(streamId))
-      .filter(({ order }) => order >= minOrder)
-      .map(({ streamId }) => {
-        const cursor = streamCursors[streamId] ?? 0;
-        return { streamId, cursor };
-      });
-  }, [streamList, streamCursors, options?.skipStreamIds, minOrder]);
-
-  // Get the deltas for all the active streams, if any.
-  const cursorQuery = useQuery(
-    query,
-    args === "skip" || !streamList
-      ? ("skip" as const)
-      : ({
-          ...args,
-          streamArgs: { kind: "deltas", cursors } as StreamArgs,
-        } as FunctionArgs<Query>),
-  ) as
-    | { streams: Extract<SyncStreamsReturnValue, { kind: "deltas" }> }
-    | undefined;
+  const streams = useDeltaStreams(query, args, options);
 
   const threadId = args === "skip" ? undefined : args.threadId;
 
   useEffect(() => {
-    const activeStreamIds = new Set(
-      streamList?.streams.messages.map((m) => m.streamId) ?? [],
-    );
-    // Clean up finished streams
-    setUIMessages((prev) => {
-      const toRemove = Object.keys(prev).filter(
-        (streamId) => !activeStreamIds.has(streamId),
-      );
-      return toRemove.length ? omit(prev, toRemove) : prev;
-    });
-    for (const [
-      streamId,
-      controller,
-    ] of uiMessageStreamControllers.current.entries()) {
-      if (!activeStreamIds.has(streamId)) {
-        uiMessageStreamControllers.current.delete(streamId);
-        controller.close();
+    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 (!cursorQuery?.streams?.deltas || !threadId) return;
-
-    const deltasByStream = new Map<string, typeof cursorQuery.streams.deltas>();
-
-    // Group deltas by streamId
-    for (const delta of cursorQuery.streams.deltas) {
-      if (!deltasByStream.has(delta.streamId)) {
-        deltasByStream.set(delta.streamId, []);
+      if (lastDelta && lastDelta.start >= cursor) {
+        noNewDeltas = false;
+        break;
       }
-      deltasByStream.get(delta.streamId)!.push(delta);
     }
-
-    // Process each stream's deltas
-    for (const [streamId, deltas] of deltasByStream.entries()) {
-      const streamMessage = streamList?.streams?.messages.find(
-        (m) => m.streamId === streamId,
-      );
-      if (!streamMessage) continue;
-      if (streamMessage.format === "UIMessageChunk") {
-        setStreamCursors((streamCursors) => {
-          const currentCursor = streamCursors[streamId] ?? 0;
-
-          const { parts, cursor } = getParts<UIMessageChunk>(
-            deltas,
-            currentCursor,
-          );
-
-          const newStreamCursors =
-            cursor === currentCursor
-              ? streamCursors
-              : { ...streamCursors, [streamId]: cursor };
-
-          if (parts.length === 0) return newStreamCursors;
-
-          // Get existing message for this stream as starting point
-          const existingStream =
-            uiMessageStreamControllers.current.get(streamId);
-          if (existingStream) {
-            for (const part of parts) {
-              existingStream.enqueue(part);
-            }
-          } else {
-            const stream = new ReadableStream<UIMessageChunk>({
-              start(controller) {
-                uiMessageStreamControllers.current.set(streamId, controller);
-                for (const part of parts) {
-                  controller.enqueue(part);
-                }
-              },
-            });
-            const initialMessage = blankUIMessage(streamMessage, threadId);
-            setUIMessages((prev) => ({
-              ...prev,
-              [streamId]: initialMessage as UIMessage<
-                METADATA,
-                DATA_PARTS,
-                TOOLS
-              >,
-            }));
-
-            const messageStream = readUIMessageStream({
-              message: initialMessage,
-              stream,
-              onError: (error) => {
-                setUIMessages((prev) => ({
-                  ...prev,
-                  [streamId]: {
-                    ...prev[streamId],
-                    status: "failed",
-                  },
-                }));
-                console.error(`Error in stream ${streamId}:`, error);
-              },
-              terminateOnError: false,
-            });
-
-            // Process the async iterator
-            void (async () => {
-              for await (const message of messageStream) {
-                message.text = message.parts
-                  .filter((p) => p.type === "text")
-                  .map((p) => p.text)
-                  .join("");
-                setUIMessages((prev) =>
-                  // If we don't have a ui message assume we've aborted.
-                  prev[streamId]
-                    ? {
-                        ...prev,
-                        [streamId]: message as UIMessage<
-                          METADATA,
-                          DATA_PARTS,
-                          TOOLS
-                        >,
-                      }
-                    : prev,
-                );
-              }
-            })().catch((error) => {
-              console.error(`Error processing stream ${streamId}:`, error);
-            });
-          }
-          return newStreamCursors;
-        });
-      } else {
-        setUIMessages((uiMessages) => {
-          const existingUIMessage = uiMessages[streamId];
-          const [[uiMessage], [streamMetadata], changed] =
-            deriveUIMessagesFromTextStreamParts(
-              threadId,
-              [streamMessage],
-              existingUIMessage
-                ? [
-                    {
-                      streamId,
-                      cursor: streamCursors[streamId] ?? 0,
-                      message: existingUIMessage,
-                    },
-                  ]
-                : [],
-              deltas,
-            );
-          if (changed) {
-            setStreamCursors((prev) => ({
-              ...prev,
-              [streamId]: streamMetadata.cursor,
-            }));
-            return {
-              ...uiMessages,
-              [streamId]: uiMessage as UIMessage<METADATA, DATA_PARTS, TOOLS>,
-            };
-          }
-          return uiMessages;
-        });
-      }
+    if (noNewDeltas) {
+      return;
     }
-  }, [cursorQuery, streamCursors, threadId, streamList?.streams?.messages]);
+    const abortController = new AbortController();
+    void (async () => {
+      const newMessageState: Record<
+        string,
+        {
+          uiMessage: UIMessage<METADATA, DATA_PARTS, TOOLS>;
+          cursor: number;
+        }
+      > = 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,
+                },
+              ];
+            }
+          }),
+        ),
+      );
+      if (abortController.signal.aborted) return;
+      setMessageState(newMessageState);
+    })();
+    return () => {
+      abortController.abort();
+    };
+  }, [messageState, streams, threadId]);
 
   return useMemo(() => {
-    if (!streamList) return undefined;
-    return sorted(
-      Array.from(
-        Object.entries(uiMessages)
-          .map(([streamId, uiMessage]) => {
-            const streamMessage = streamList?.streams.messages.find(
-              (m) => m.streamId === streamId,
-            );
-            if (!streamMessage) return undefined;
-            uiMessage.status = statusFromStreamStatus(
-              streamMessage?.status ?? "finished",
-            );
-            uiMessage.text = uiMessage.parts
-              .filter((p): p is TextUIPart => p.type === "text")
-              .map((p) => p.text)
-              .join("");
-            return uiMessage;
-          })
-          .filter((uiMessage) => uiMessage !== undefined),
-      ),
-    );
-  }, [uiMessages, streamList]);
+    if (!streams) return undefined;
+    return streams
+      .map(
+        ({ streamMessage }) => messageState[streamMessage.streamId]?.uiMessage,
+      )
+      .filter((uiMessage) => uiMessage !== undefined);
+  }, [messageState, streams]);
 }

package/src/react/useUIMessages.ts

@@ -17,11 +17,12 @@ import type {
 } from "convex/server";
 import { useMemo } from "react";
 import type { SyncStreamsReturnValue } from "../client/types.js";
-import type { MessageStatus, StreamArgs } from "../validators.js";
+import type { StreamArgs } from "../validators.js";
 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]);
 

package/src/validators.ts

@@ -93,7 +93,7 @@ export const vSourcePart = v.union(
     sourceType: v.literal("url"),
     id: v.string(),
     url: v.string(),
-    title: v.string(),
+    title: v.optional(v.string()),
     providerOptions,
     providerMetadata,
   }),

package/dist/package.json

@@ -1,3 +0,0 @@
-{
-  "type": "module"
-}