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

package/src/deltas.ts

@@ -18,57 +18,6 @@ import {
   type StreamMessage,
 } from "./validators.js";
 
-/**
- * Compressing parts when streaming to save bandwidth in deltas.
- */
-
-export function compressUIMessageChunks(
-  parts: UIMessageChunk[],
-): UIMessageChunk[] {
-  const compressed: UIMessageChunk[] = [];
-  for (const part of parts) {
-    const last = compressed.at(-1);
-    if (part.type === "text-delta" || part.type === "reasoning-delta") {
-      if (last?.type === part.type && part.id === last.id) {
-        last.delta += part.delta;
-      } else {
-        compressed.push(part);
-      }
-    } else {
-      compressed.push(part);
-    }
-  }
-  return compressed;
-}
-
-export function compressTextStreamParts(
-  parts: TextStreamPart<ToolSet>[],
-): TextStreamPart<ToolSet>[] {
-  const compressed: TextStreamPart<ToolSet>[] = [];
-  for (const part of parts) {
-    const last = compressed.at(-1);
-    if (part.type === "text-delta" || part.type === "reasoning-delta") {
-      if (last?.type === part.type && part.id === last.id) {
-        last.text += part.text;
-      } else {
-        compressed.push(part);
-      }
-    } else {
-      if (part.type === "file") {
-        compressed.push({
-          type: "file",
-          file: {
-            ...part.file,
-            uint8Array: undefined as unknown as Uint8Array,
-          },
-        });
-      }
-      compressed.push(part);
-    }
-  }
-  return compressed;
-}
-
 export function blankUIMessage<METADATA = unknown>(
   streamMessage: StreamMessage & { metadata?: METADATA },
   threadId: string,
@@ -238,9 +187,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(
@@ -591,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;
+}

package/src/react/useDeltaStreams.ts

@@ -0,0 +1,155 @@
+"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 { sorted } from "../shared.js";
+import { useQuery } from "convex/react";
+import { useState } from "react";
+import { assert } from "convex-helpers";
+
+export function useDeltaStreams<
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  Query extends StreamQuery<any> = StreamQuery<object>,
+>(
+  query: Query,
+  args: StreamQueryArgs<Query> | "skip",
+  options?: {
+    startOrder?: number;
+    skipStreamIds?: string[];
+  },
+): { streamMessage: StreamMessage; deltas: StreamDelta[] }[] | undefined {
+  // We hold onto and modify state directly to avoid re-running unnecessarily.
+  const [state] = useState<{
+    startOrder: number;
+    threadId: string | undefined;
+    deltaStreams:
+      | Array<{
+          streamMessage: StreamMessage;
+          deltas: StreamDelta[];
+        }>
+      | undefined;
+  }>({
+    startOrder: options?.startOrder ?? 0,
+    deltaStreams: undefined,
+    threadId: args === "skip" ? undefined : args.threadId,
+  });
+  const [cursors, setCursors] = useState<Record<string, number>>({});
+  if (args !== "skip" && state.threadId !== args.threadId) {
+    state.threadId = args.threadId;
+    state.deltaStreams = undefined;
+    state.startOrder = options?.startOrder ?? 0;
+    setCursors({});
+  }
+  if (
+    state.deltaStreams?.length ||
+    (options?.startOrder && options.startOrder < state.startOrder)
+  ) {
+    const cacheFriendlyStartOrder = options?.startOrder
+      ? // round down to the nearest 10 for some cache benefits
+        options.startOrder - (options.startOrder % 10)
+      : 0;
+    if (cacheFriendlyStartOrder !== state.startOrder) {
+      state.startOrder = cacheFriendlyStartOrder;
+    }
+  }
+
+  // Get all the active streams
+  const streamList = useQuery(
+    query,
+    args === "skip"
+      ? args
+      : ({
+          ...args,
+          streamArgs: {
+            kind: "list",
+            startOrder: state.startOrder,
+          } as StreamArgs,
+        } as FunctionArgs<Query>),
+  ) as
+    | { streams: Extract<SyncStreamsReturnValue, { kind: "list" }> }
+    | undefined;
+
+  const streamMessages =
+    args === "skip"
+      ? undefined
+      : !streamList
+        ? state.deltaStreams?.map(({ streamMessage }) => streamMessage)
+        : 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.
+  const cursorQuery = useQuery(
+    query,
+    args === "skip" || !streamMessages?.length
+      ? ("skip" as const)
+      : ({
+          ...args,
+          streamArgs: {
+            kind: "deltas",
+            cursors: streamMessages.map(({ streamId }) => ({
+              streamId,
+              cursor: cursors[streamId] ?? 0,
+            })),
+          } as StreamArgs,
+        } as FunctionArgs<Query>),
+  ) as
+    | { streams: Extract<SyncStreamsReturnValue, { kind: "deltas" }> }
+    | undefined;
+
+  const newDeltas = cursorQuery?.streams.deltas;
+  if (newDeltas?.length && streamMessages) {
+    const newDeltasByStreamId = new Map<string, StreamDelta[]>();
+    for (const delta of newDeltas) {
+      const oldCursor = cursors[delta.streamId];
+      if (oldCursor && delta.start < oldCursor) continue;
+      const existing = newDeltasByStreamId.get(delta.streamId);
+      if (existing) {
+        const previousEnd = existing.at(-1)!.end;
+        assert(
+          previousEnd === delta.start,
+          `Gap found in deltas for ${delta.streamId} jumping to ${delta.start} from ${previousEnd}`,
+        );
+        existing.push(delta);
+      } else {
+        assert(
+          !oldCursor || oldCursor === delta.start,
+          `Gap found - first delta after ${oldCursor} is ${delta.start} for stream ${delta.streamId}`,
+        );
+        newDeltasByStreamId.set(delta.streamId, [delta]);
+      }
+    }
+    const newCursors: Record<string, number> = {};
+    for (const { streamId } of streamMessages) {
+      const cursor =
+        newDeltasByStreamId.get(streamId)?.at(-1)?.end ?? cursors[streamId];
+      if (cursor !== undefined) {
+        newCursors[streamId] = cursor;
+      }
+    }
+    setCursors(newCursors);
+
+    // we defensively create a new object so object identity matches contents
+    state.deltaStreams = streamMessages.map((streamMessage) => {
+      const streamId = streamMessage.streamId;
+      const old = state.deltaStreams?.find(
+        (ds) => ds.streamMessage.streamId === streamId,
+      );
+      const newDeltas = newDeltasByStreamId.get(streamId);
+      if (!newDeltas && streamMessage === old?.streamMessage) {
+        return old;
+      }
+      return {
+        streamMessage,
+        deltas: [...(old?.deltas ?? []), ...(newDeltas ?? [])],
+      };
+    });
+  }
+  return state.deltaStreams;
+}

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"
-}