@convex-dev/agent 0.2.8-alpha.2 → 0.2.8-alpha.4

package/src/deltas.test.ts

@@ -1,41 +1,175 @@
 import { describe, it, expect } from "vitest";
 import {
+  blankUIMessage,
   deriveUIMessagesFromTextStreamParts,
   updateFromTextStreamParts,
+  updateFromUIMessageChunks,
 } from "./deltas.js";
 import type { StreamMessage, StreamDelta } from "./validators.js";
 import { omit } from "convex-helpers";
-import type { TextStreamPart, ToolSet, ToolUIPart } from "ai";
+import type { Tool, ToolUIPart, TypedToolResult } from "ai";
 
-function makeStreamMessage(
-  streamId: string,
-  order: number,
-  stepOrder: number,
-): StreamMessage {
-  return { streamId, order, stepOrder } as StreamMessage;
-}
+describe("UIMessageChunks", () => {
+  it("updates a UIMessage with a tool call and follow up", async () => {
+    const uiMessage = blankUIMessage(
+      {
+        streamId: "s1",
+        status: "streaming",
+        order: 0,
+        stepOrder: 1,
+        format: "UIMessageChunk",
+        agentName: "agent1",
+      },
+      "thread1",
+    );
+    expect(uiMessage.text).toBe("");
+    expect(uiMessage.parts).toEqual([]);
+    const updatedMessage = await updateFromUIMessageChunks(uiMessage, [
+      { type: "start" },
+      { type: "start-step" },
+      { type: "reasoning-start", id: "reasoning-0" },
+      { type: "reasoning-delta", id: "reasoning-0", delta: "Okay" },
+      {
+        type: "reasoning-delta",
+        id: "reasoning-0",
+        delta: ", the user is asking...",
+      },
+      { type: "text-start", id: "txt-1" },
+      {
+        type: "text-delta",
+        id: "txt-1",
+        delta: "Hey ho.",
+      },
+      { type: "reasoning-end", id: "reasoning-0" },
+      { type: "text-end", id: "txt-1" },
+      { type: "tool-input-start", toolCallId: "0ychh9k6f", toolName: "say" },
+      {
+        type: "tool-input-delta",
+        toolCallId: "0ychh9k6f",
+        inputTextDelta:
+          '{"question":"What is your favorite flavor of ice cream?"}',
+      },
+      {
+        type: "tool-input-available",
+        toolCallId: "0ychh9k6f",
+        toolName: "say",
+        input: { question: "What is your favorite flavor of ice cream?" },
+        providerMetadata: { openai: { itemId: "123" } },
+      },
+      {
+        type: "tool-output-available",
+        toolCallId: "0ychh9k6f",
+        output: "I'm sorry I can't help you. Stop asking me questions.",
+      },
+      { type: "finish-step" },
+      { type: "start-step" },
+      { type: "tool-input-start", toolCallId: "1ychh9k6f", toolName: "say" },
+      {
+        type: "tool-input-delta",
+        toolCallId: "1ychh9k6f",
+        inputTextDelta:
+          '{"question":"What is your favorite flavor of ice cream?"}',
+      },
+      {
+        type: "tool-input-available",
+        toolCallId: "1ychh9k6f",
+        toolName: "say",
+        input: { question: "What is your favorite flavor of ice cream?" },
+      },
+      {
+        type: "tool-output-available",
+        toolCallId: "1ychh9k6f",
+        output: "I'm serious.",
+      },
+      { type: "finish-step" },
+      { type: "start-step" },
+      { type: "text-start", id: "msg_0" },
+      {
+        type: "text-delta",
+        id: "msg_0",
+        delta: "The best ice cream flavor is vanilla",
+      },
+      {
+        type: "text-delta",
+        id: "msg_0",
+        delta: ".",
+      },
+      { type: "text-end", id: "msg_0" },
+      { type: "finish-step" },
+      { type: "finish" },
+    ]);
+    expect(updatedMessage.text).toBe(
+      "Hey ho.The best ice cream flavor is vanilla.",
+    );
+    const expectedParts = [
+      {
+        type: "step-start",
+      },
+      {
+        state: "done",
+        text: "Okay, the user is asking...",
+        type: "reasoning",
+      },
+      {
+        state: "done",
+        text: "Hey ho.",
+        type: "text",
+      },
+      {
+        callProviderMetadata: {
+          openai: {
+            itemId: "123",
+          },
+        },
+        input: {
+          question: "What is your favorite flavor of ice cream?",
+        },
+        output: "I'm sorry I can't help you. Stop asking me questions.",
 
-function makeDelta(
-  streamId: string,
-  start: number,
-  end: number,
-  parts: TextStreamPart<ToolSet>[],
-): StreamDelta {
-  return { streamId, start, end, parts };
-}
+        state: "output-available",
+        toolCallId: "0ychh9k6f",
+        type: "tool-say",
+      },
+      {
+        type: "step-start",
+      },
+      {
+        input: {
+          question: "What is your favorite flavor of ice cream?",
+        },
+        output: "I'm serious.",
+        state: "output-available",
+        toolCallId: "1ychh9k6f",
+        type: "tool-say",
+      },
+      {
+        type: "step-start",
+      },
+      {
+        state: "done",
+        text: "The best ice cream flavor is vanilla.",
+        type: "text",
+      },
+    ];
+    expect(updatedMessage.parts).toEqual(expectedParts);
+    expect(updatedMessage.parts).toHaveLength(8);
+  });
+});
 
 describe("mergeDeltas", () => {
   it("merges a single text-delta into a message", () => {
     const streamId = "s1";
-    const streamMessages = [makeStreamMessage(streamId, 1, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 5, [
-        { type: "text-delta", id: "1", text: "Hello" },
-      ]),
+      {
+        streamId,
+        start: 0,
+        end: 5,
+        parts: [{ type: "text-delta", id: "1", text: "Hello" }],
+      } satisfies StreamDelta,
     ];
     const [messages, newStreams, changed] = deriveUIMessagesFromTextStreamParts(
       "thread1",
-      streamMessages,
+      [{ streamId, order: 1, stepOrder: 0, status: "streaming" }],
       [],
       deltas,
     );
@@ -48,18 +182,23 @@ describe("mergeDeltas", () => {
 
   it("merges multiple deltas for the same stream", () => {
     const streamId = "s1";
-    const streamMessages = [makeStreamMessage(streamId, 1, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 5, [
-        { type: "text-delta", id: "1", text: "Hello" },
-      ]),
-      makeDelta(streamId, 5, 11, [
-        { type: "text-delta", id: "2", text: " World!" },
-      ]),
+      {
+        streamId,
+        start: 0,
+        end: 5,
+        parts: [{ type: "text-delta", id: "1", text: "Hello" }],
+      },
+      {
+        streamId,
+        start: 5,
+        end: 11,
+        parts: [{ type: "text-delta", id: "2", text: " World!" }],
+      },
     ];
     const [messages, newStreams, changed] = deriveUIMessagesFromTextStreamParts(
       "thread1",
-      streamMessages,
+      [{ streamId, order: 1, stepOrder: 0, status: "streaming" }],
       [],
       deltas,
     );
@@ -71,29 +210,38 @@ describe("mergeDeltas", () => {
 
   it("handles tool-call and tool-result parts", () => {
     const streamId = "s2";
-    const streamMessages = [makeStreamMessage(streamId, 2, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 1, [
-        {
-          type: "tool-call",
-          toolCallId: "call1",
-          toolName: "myTool",
-          input: "What's the meaning of life?",
-        },
-      ]),
-      makeDelta(streamId, 1, 2, [
-        {
-          type: "tool-result",
-          toolCallId: "call1",
-          toolName: "myTool",
-          input: undefined,
-          output: "42",
-        },
-      ]),
+      {
+        streamId,
+        start: 0,
+        end: 1,
+        parts: [
+          {
+            type: "tool-call",
+            toolCallId: "call1",
+            toolName: "myTool",
+            input: "What's the meaning of life?",
+          },
+        ],
+      } satisfies StreamDelta,
+      {
+        streamId,
+        start: 1,
+        end: 2,
+        parts: [
+          {
+            type: "tool-result",
+            toolCallId: "call1",
+            toolName: "myTool",
+            input: undefined,
+            output: "42",
+          } satisfies TypedToolResult<{ myTool: Tool }>,
+        ],
+      } satisfies StreamDelta,
     ];
     const [[message], _, changed] = deriveUIMessagesFromTextStreamParts(
       "thread1",
-      streamMessages,
+      [{ streamId, order: 2, stepOrder: 0, status: "streaming" }],
       [],
       deltas,
     );
@@ -114,11 +262,10 @@ describe("mergeDeltas", () => {
 
   it("returns changed=false if no new deltas", () => {
     const streamId = "s3";
-    const streamMessages = [makeStreamMessage(streamId, 3, 0)];
     const deltas: StreamDelta[] = [];
     const [, newStreams, changed] = deriveUIMessagesFromTextStreamParts(
       "thread1",
-      streamMessages,
+      [{ streamId, order: 3, stepOrder: 0, status: "streaming" }],
       [],
       deltas,
     );
@@ -127,15 +274,26 @@ describe("mergeDeltas", () => {
   });
 
   it("handles multiple streams and sorts by order/stepOrder", () => {
-    const s1 = makeStreamMessage("s1", 1, 0);
-    const s2 = makeStreamMessage("s2", 2, 0);
     const deltas = [
-      makeDelta("s2", 0, 3, [{ type: "text-delta", id: "1", text: "B" }]),
-      makeDelta("s1", 0, 3, [{ type: "text-delta", id: "2", text: "A" }]),
+      {
+        streamId: "s2",
+        start: 0,
+        end: 3,
+        parts: [{ type: "text-delta", id: "1", text: "B" }],
+      } satisfies StreamDelta,
+      {
+        streamId: "s1",
+        start: 0,
+        end: 3,
+        parts: [{ type: "text-delta", id: "2", text: "A" }],
+      } satisfies StreamDelta,
     ];
     const [messages, _, changed] = deriveUIMessagesFromTextStreamParts(
       "thread1",
-      [s2, s1],
+      [
+        { streamId: "s1", order: 1, stepOrder: 0, status: "streaming" },
+        { streamId: "s2", order: 2, stepOrder: 0, status: "streaming" },
+      ],
       [],
       deltas,
     );
@@ -150,19 +308,29 @@ describe("mergeDeltas", () => {
 
   it("does not duplicate text content when merging sequential text-deltas", () => {
     const streamId = "s4";
-    const streamMessages = [makeStreamMessage(streamId, 4, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 5, [
-        { type: "text-delta", id: "1", text: "Hello" },
-      ]),
-      makeDelta(streamId, 5, 11, [
-        { type: "text-delta", id: "2", text: " World!" },
-      ]),
-      makeDelta(streamId, 11, 12, [{ type: "text-delta", id: "3", text: "!" }]),
-    ];
+      {
+        streamId,
+        start: 0,
+        end: 5,
+        parts: [{ type: "text-delta", id: "1", text: "Hello" }],
+      },
+      {
+        streamId,
+        start: 5,
+        end: 11,
+        parts: [{ type: "text-delta", id: "2", text: " World!" }],
+      },
+      {
+        streamId,
+        start: 11,
+        end: 12,
+        parts: [{ type: "text-delta", id: "3", text: "!" }],
+      },
+    ] satisfies StreamDelta[];
     const [messages] = deriveUIMessagesFromTextStreamParts(
       "thread1",
-      streamMessages,
+      [{ streamId, order: 4, stepOrder: 0, status: "streaming" }],
       [],
       deltas,
     );
@@ -179,20 +347,34 @@ describe("mergeDeltas", () => {
 
   it("does not duplicate reasoning parts", () => {
     const streamId = "s6";
-    const streamMessages = [makeStreamMessage(streamId, 6, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 1, [
-        { type: "reasoning-start", id: "1" },
-        { type: "reasoning-delta", id: "1", text: "I'm thinking..." },
-      ]),
-      makeDelta(streamId, 1, 2, [
-        { type: "reasoning-delta", id: "1", text: " Still thinking..." },
-      ]),
-      makeDelta(streamId, 2, 3, [{ type: "reasoning-end", id: "1" }]),
+      {
+        streamId,
+        start: 0,
+        end: 1,
+        parts: [
+          { type: "reasoning-start", id: "1" },
+          { type: "reasoning-delta", id: "1", text: "I'm thinking..." },
+        ],
+      },
+      {
+        streamId,
+        start: 1,
+        end: 2,
+        parts: [
+          { type: "reasoning-delta", id: "1", text: " Still thinking..." },
+        ],
+      },
+      {
+        streamId,
+        start: 2,
+        end: 3,
+        parts: [{ type: "reasoning-end", id: "1" }],
+      },
     ];
     const [messages] = deriveUIMessagesFromTextStreamParts(
       "thread1",
-      streamMessages,
+      [{ streamId, order: 6, stepOrder: 0, status: "streaming" }],
       [],
       deltas,
     );
@@ -209,14 +391,25 @@ describe("mergeDeltas", () => {
 
   it("applyDeltasToStreamMessage is idempotent and does not duplicate content", () => {
     const streamId = "s7";
-    const streamMessage = makeStreamMessage(streamId, 7, 0);
+    const streamMessage = {
+      streamId,
+      order: 7,
+      stepOrder: 0,
+      status: "streaming",
+    } satisfies StreamMessage;
     const deltas = [
-      makeDelta(streamId, 0, 5, [
-        { type: "text-delta", id: "1", text: "Hello" },
-      ]),
-      makeDelta(streamId, 5, 11, [
-        { type: "text-delta", id: "2", text: " World!" },
-      ]),
+      {
+        streamId,
+        start: 0,
+        end: 5,
+        parts: [{ type: "text-delta", id: "1", text: "Hello" }],
+      },
+      {
+        streamId,
+        start: 5,
+        end: 11,
+        parts: [{ type: "text-delta", id: "2", text: " World!" }],
+      },
     ];
     // First call: apply both deltas
     let [result, changed] = updateFromTextStreamParts(
@@ -237,7 +430,12 @@ describe("mergeDeltas", () => {
     // Third call: add a new delta
     const moreDeltas = [
       ...deltas,
-      makeDelta(streamId, 11, 12, [{ type: "text-delta", id: "3", text: "!" }]),
+      {
+        streamId,
+        start: 11,
+        end: 12,
+        parts: [{ type: "text-delta", id: "3", text: "!" }],
+      },
     ];
     [result, changed] = updateFromTextStreamParts(
       "thread1",
@@ -260,14 +458,22 @@ describe("mergeDeltas", () => {
 
   it("mergeDeltas is pure and does not mutate inputs", () => {
     const streamId = "s8";
-    const streamMessages = [makeStreamMessage(streamId, 8, 0)];
+    const streamMessages = [
+      { streamId, order: 8, stepOrder: 0, status: "streaming" },
+    ] satisfies StreamMessage[];
     const deltas = [
-      makeDelta(streamId, 0, 5, [
-        { type: "text-delta", id: "1", text: "Hello" },
-      ]),
-      makeDelta(streamId, 5, 11, [
-        { type: "text-delta", id: "2", text: " World!" },
-      ]),
+      {
+        streamId,
+        start: 0,
+        end: 5,
+        parts: [{ type: "text-delta", id: "1", text: "Hello" }],
+      },
+      {
+        streamId,
+        start: 5,
+        end: 11,
+        parts: [{ type: "text-delta", id: "2", text: " World!" }],
+      },
     ];
     // Deep freeze inputs to catch mutation
     function deepFreeze(obj: unknown): unknown {
@@ -309,14 +515,22 @@ describe("mergeDeltas", () => {
     );
     expect(changed1).toBe(changed2);
     // Inputs should remain unchanged
-    expect(streamMessages).toEqual([makeStreamMessage(streamId, 8, 0)]);
+    expect(streamMessages).toMatchObject([
+      { streamId, order: 8, stepOrder: 0, status: "streaming" },
+    ]);
     expect(deltas).toEqual([
-      makeDelta(streamId, 0, 5, [
-        { type: "text-delta", id: "1", text: "Hello" },
-      ]),
-      makeDelta(streamId, 5, 11, [
-        { type: "text-delta", id: "2", text: " World!" },
-      ]),
+      {
+        streamId,
+        start: 0,
+        end: 5,
+        parts: [{ type: "text-delta", id: "1", text: "Hello" }],
+      },
+      {
+        streamId,
+        start: 5,
+        end: 11,
+        parts: [{ type: "text-delta", id: "2", text: " World!" }],
+      },
     ]);
   });
 });

package/src/deltas.ts

@@ -127,6 +127,10 @@ export async function updateFromUIMessageChunks(
   });
   let message = uiMessage;
   for await (const messagePart of messageStream) {
+    assert(
+      messagePart.id === message.id,
+      `Expecting to only make one UIMessage in a stream, but have ${JSON.stringify(message)} and created ${JSON.stringify(messagePart)}`,
+    );
     message = messagePart;
   }
   if (failed) {

package/src/react/types.ts

@@ -13,7 +13,7 @@ export type StreamQuery<Args = Record<string, unknown>> = FunctionReference<
   { streams: SyncStreamsReturnValue }
 >;
 
-export type StreamMessagesArgs<Query extends StreamQuery<unknown>> =
+export type StreamQueryArgs<Query extends StreamQuery<unknown>> =
   Query extends StreamQuery<unknown>
     ? Expand<BetterOmit<FunctionArgs<Query>, "streamArgs">>
     : never;

package/src/react/useStreamingUIMessages.ts

@@ -12,7 +12,7 @@ import {
 } from "ai";
 import type { SyncStreamsReturnValue } from "../client/types.js";
 import type { StreamArgs } from "../validators.js";
-import type { StreamQuery, StreamMessagesArgs } from "./types.js";
+import type { StreamQuery, StreamQueryArgs } from "./types.js";
 import { type UIMessage } from "../UIMessages.js";
 import {
   blankUIMessage,
@@ -44,7 +44,7 @@ export function useStreamingUIMessages<
   Query extends StreamQuery<any> = StreamQuery<object>,
 >(
   query: Query,
-  args: StreamMessagesArgs<Query> | "skip",
+  args: StreamQueryArgs<Query> | "skip",
   options?: {
     startOrder?: number;
     skipStreamIds?: string[];
@@ -61,6 +61,11 @@ export function useStreamingUIMessages<
     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,
@@ -70,13 +75,20 @@ export function useStreamingUIMessages<
           ...args,
           streamArgs: {
             kind: "list",
-            startOrder: options?.startOrder ?? 0,
+            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 [];
@@ -85,11 +97,12 @@ export function useStreamingUIMessages<
     }
     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]);
+  }, [streamList, streamCursors, options?.skipStreamIds, minOrder]);
 
   // Get the deltas for all the active streams, if any.
   const cursorQuery = useQuery(
@@ -98,7 +111,6 @@ export function useStreamingUIMessages<
       ? ("skip" as const)
       : ({
           ...args,
-          paginationOpts: { cursor: null, numItems: 0 },
           streamArgs: { kind: "deltas", cursors } as StreamArgs,
         } as FunctionArgs<Query>),
   ) as
@@ -107,8 +119,26 @@ export function useStreamingUIMessages<
 
   const threadId = args === "skip" ? undefined : args.threadId;
 
-  // Process new deltas and convert to UIMessageChunks, then use readUIMessageStream
   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 (!cursorQuery?.streams?.deltas || !threadId) return;
 
     const deltasByStream = new Map<string, typeof cursorQuery.streams.deltas>();
@@ -246,28 +276,6 @@ export function useStreamingUIMessages<
     }
   }, [cursorQuery, streamCursors, threadId, streamList?.streams?.messages]);
 
-  // Clean up finished streams
-  useEffect(() => {
-    const activeStreamIds = new Set(
-      streamList?.streams.messages.map((m) => m.streamId) ?? [],
-    );
-    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)) {
-        controller.close();
-        uiMessageStreamControllers.current.delete(streamId);
-      }
-    }
-  }, [streamList]);
-
   return useMemo(() => {
     if (!streamList) return undefined;
     return sorted(