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

package/src/UIMessages.ts

@@ -27,6 +27,8 @@ import type {
 } from "./validators.js";
 import { omit, pick } from "convex-helpers";
 
+export type UIStatus = "streaming" | MessageStatus;
+
 export type UIMessage<
   METADATA = unknown,
   DATA_PARTS extends UIDataTypes = UIDataTypes,
@@ -35,7 +37,7 @@ export type UIMessage<
   key: string;
   order: number;
   stepOrder: number;
-  status: "streaming" | MessageStatus;
+  status: UIStatus;
   agentName?: string;
   text: string;
   _creationTime: number;
@@ -160,7 +162,7 @@ export function toUIMessages<
   messages: (MessageDoc & ExtraFields<METADATA>)[],
 ): UIMessage<METADATA, DATA_PARTS, TOOLS>[] {
   // Group assistant and tool messages together
-  const assistantGroups = groupAssistantMessages(messages);
+  const assistantGroups = groupAssistantMessages(sorted(messages));
 
   const uiMessages: UIMessage<METADATA, DATA_PARTS, TOOLS>[] = [];
   for (const group of assistantGroups) {
@@ -196,13 +198,10 @@ function groupAssistantMessages<METADATA = unknown>(
 ): Group<METADATA>[] {
   const groups: Group<METADATA>[] = [];
 
-  // Sort messages by order and stepOrder first to handle out-of-order arrivals
-  const sortedMessages = sorted(messages);
-
   let currentAssistantGroup: (MessageDoc & ExtraFields<METADATA>)[] = [];
   let currentOrder: number | undefined;
 
-  for (const message of sortedMessages) {
+  for (const message of messages) {
     const coreMessage = message.message && deserializeMessage(message.message);
     if (!coreMessage) continue;
 
@@ -239,7 +238,6 @@ function groupAssistantMessages<METADATA = unknown>(
       currentAssistantGroup.push(message);
 
       // End group if this is an assistant message without tool calls
-      // But only if we're processing messages in order (which we are now due to sorting)
       if (coreMessage.role === "assistant" && !message.tool) {
         groups.push({
           role: "assistant",
@@ -294,7 +292,7 @@ function createSystemUIMessage<
 
 function extractTextFromMessageDoc(message: MessageDoc): string {
   return (
-    message.text || (message.message && extractText(message.message)) || ""
+    (message.message && extractText(message.message)) || message.text || ""
   );
 }
 
@@ -373,9 +371,9 @@ function createAssistantUIMessage<
 
   // Concatenate text from all messages in group
   const allText = group
-    .map((msg) => msg.text || ((msg.message && extractText(msg.message)) ?? ""))
+    .map((msg) => extractTextFromMessageDoc(msg))
     .filter(Boolean)
-    .join("");
+    .join(" ");
 
   // Get status from last message
   const lastMessage = group[group.length - 1];
@@ -393,8 +391,7 @@ function createAssistantUIMessage<
     const content = coreMessage.content;
     const nonStringContent =
       content && typeof content !== "string" ? content : [];
-    const text =
-      message.text || ((message.message && extractText(message.message)) ?? "");
+    const text = extractTextFromMessageDoc(message);
 
     const partCommon = {
       state: message.streaming ? ("streaming" as const) : ("done" as const),
@@ -465,6 +462,10 @@ function createAssistantUIMessage<
           break;
         }
         case "tool-result": {
+          const output =
+            contentPart.output?.type === "json"
+              ? contentPart.output.value
+              : contentPart.output;
           const call = allParts.find(
             (part) =>
               part.type === `tool-${contentPart.toolName}` &&
@@ -475,13 +476,10 @@ function createAssistantUIMessage<
             if (message.error) {
               call.state = "output-error";
               call.errorText = message.error;
-              call.output = contentPart.output;
+              call.output = output;
             } else {
               call.state = "output-available";
-              call.output =
-                contentPart.output?.type === "json"
-                  ? contentPart.output.value
-                  : contentPart.output;
+              call.output = output;
             }
           } else {
             console.warn(
@@ -503,10 +501,7 @@ function createAssistantUIMessage<
                 toolCallId: contentPart.toolCallId,
                 state: "output-available",
                 input: undefined,
-                output:
-                  contentPart.output?.type === "json"
-                    ? contentPart.output.value
-                    : contentPart.output,
+                output,
                 callProviderMetadata: message.providerMetadata,
               } satisfies ToolUIPart<TOOLS>);
             }

package/src/client/createTool.ts

@@ -44,7 +44,7 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(def: {
     ctx: Ctx,
     args: INPUT,
     options: ToolCallOptions,
-  ) => PromiseLike<OUTPUT>;
+  ) => PromiseLike<OUTPUT> | AsyncIterable<OUTPUT>;
   /**
    * Provide the context to use, e.g. when defining the tool at runtime.
    */
@@ -85,7 +85,7 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(def: {
     ctx: def.ctx,
     description: def.description,
     inputSchema: def.args,
-    async execute(args: INPUT, options: ToolCallOptions) {
+    execute(args: INPUT, options: ToolCallOptions) {
       if (!getCtx(this)) {
         throw new Error(
           "To use a Convex tool, you must either provide the ctx" +

package/src/client/mockModel.ts

@@ -4,7 +4,7 @@ import type {
   LanguageModelV2StreamPart,
 } from "@ai-sdk/provider";
 import { simulateReadableStream, type ProviderMetadata } from "ai";
-import { pick } from "convex-helpers";
+import { assert, pick } from "convex-helpers";
 
 export const DEFAULT_TEXT = `
 A A A A A A A A A A A A A A A
@@ -22,8 +22,18 @@ export type MockModelArgs = {
     | (() => LanguageModelV2["supportedUrls"]);
   chunkDelayInMs?: number;
   initialDelayInMs?: number;
-  // provide either content or doGenerate & doStream
+  /** A list of the responses for multiple steps.
+   * For tool calls, the first list would include a tool call part,
+   * then the next list would be after the tool response or another tool call.
+   * Tool responses come from actual tool calls!
+   */
+  contentSteps?: LanguageModelV2Content[][];
+  /** A single list of content responded from each step.
+   * Provide contentSteps instead if you want to do multi-step responses with
+   * tool calls.
+   */
   content?: LanguageModelV2Content[];
+  // provide either content, contentResponses or doGenerate & doStream
   doGenerate?: LanguageModelV2["doGenerate"];
   doStream?: LanguageModelV2["doStream"];
   providerMetadata?: ProviderMetadata;
@@ -35,6 +45,10 @@ export type MockModelArgs = {
       };
 };
 
+function atMostOneOf(...args: unknown[]) {
+  return args.filter(Boolean).length <= 1;
+}
+
 export function mockModel(args?: MockModelArgs): LanguageModelV2 {
   return new MockLanguageModel(args ?? {});
 }
@@ -54,10 +68,19 @@ export class MockLanguageModel implements LanguageModelV2 {
   doStreamCalls: Parameters<LanguageModelV2["doStream"]>[0][] = [];
 
   constructor(args: MockModelArgs) {
+    assert(
+      atMostOneOf(
+        args.content,
+        args.contentSteps,
+        args.doGenerate && args.doStream,
+      ),
+      "Expected only one of content, contentSteps, or doGenerate and doStream",
+    );
     this.provider = args.provider || "mock-provider";
     this.modelId = args.modelId || "mock-model-id";
     const {
       content = [{ type: "text", text: DEFAULT_TEXT }],
+      contentSteps = [content],
       chunkDelayInMs = 0,
       initialDelayInMs = 0,
       supportedUrls = {},
@@ -72,77 +95,83 @@ export class MockLanguageModel implements LanguageModelV2 {
       "Mock error message";
     const metadata = pick(args, ["providerMetadata"]);
 
-    const chunks: LanguageModelV2StreamPart[] = [
-      { type: "stream-start", warnings: [] },
-    ];
-    chunks.push(
-      ...content.flatMap((c, ci): LanguageModelV2StreamPart[] => {
-        if (c.type !== "text" && c.type !== "reasoning") {
-          return [c];
-        }
-        const metadata = pick(c, ["providerMetadata"]);
-        const deltas = c.text.split(" ");
-        const parts: LanguageModelV2StreamPart[] = [];
-        if (c.type === "reasoning") {
-          parts.push({
-            type: "reasoning-start",
-            id: `reasoning-${ci}`,
-            ...metadata,
-          });
-          parts.push(
-            ...deltas.map(
-              (delta, di) =>
-                ({
-                  type: "reasoning-delta",
-                  delta: (di ? " " : "") + delta,
-                  id: `reasoning-${ci}`,
-                  ...metadata,
-                }) satisfies LanguageModelV2StreamPart,
-            ),
-          );
-          parts.push({
-            type: "reasoning-end",
-            id: `reasoning-${ci}`,
-            ...metadata,
-          });
-        } else if (c.type === "text") {
-          parts.push({
-            type: "text-start",
-            id: `txt-${ci}`,
-            ...metadata,
-          });
-          parts.push(
-            ...deltas.map(
-              (delta, di) =>
-                ({
-                  type: "text-delta",
-                  delta: (di ? " " : "") + delta,
-                  id: `txt-${ci}`,
-                        ...metadata,
-                }) satisfies LanguageModelV2StreamPart,
-            ),
-          );
-          parts.push({
-            type: "text-end",
-            id: `txt-${ci}`,
-            ...metadata,
+    const chunkResponses: LanguageModelV2StreamPart[][] = contentSteps.map(
+      (content) => {
+        const chunks: LanguageModelV2StreamPart[] = [
+          { type: "stream-start", warnings: [] },
+        ];
+        chunks.push(
+          ...content.flatMap((c, ci): LanguageModelV2StreamPart[] => {
+            if (c.type !== "text" && c.type !== "reasoning") {
+              return [c];
+            }
+            const metadata = pick(c, ["providerMetadata"]);
+            const deltas = c.text.split(" ");
+            const parts: LanguageModelV2StreamPart[] = [];
+            if (c.type === "reasoning") {
+              parts.push({
+                type: "reasoning-start",
+                id: `reasoning-${ci}`,
+                ...metadata,
+              });
+              parts.push(
+                ...deltas.map(
+                  (delta, di) =>
+                    ({
+                      type: "reasoning-delta",
+                      delta: (di ? " " : "") + delta,
+                      id: `reasoning-${ci}`,
+                      ...metadata,
+                    }) satisfies LanguageModelV2StreamPart,
+                ),
+              );
+              parts.push({
+                type: "reasoning-end",
+                id: `reasoning-${ci}`,
+                ...metadata,
+              });
+            } else if (c.type === "text") {
+              parts.push({
+                type: "text-start",
+                id: `txt-${ci}`,
+                ...metadata,
+              });
+              parts.push(
+                ...deltas.map(
+                  (delta, di) =>
+                    ({
+                      type: "text-delta",
+                      delta: (di ? " " : "") + delta,
+                      id: `txt-${ci}`,
+                      ...metadata,
+                    }) satisfies LanguageModelV2StreamPart,
+                ),
+              );
+              parts.push({
+                type: "text-end",
+                id: `txt-${ci}`,
+                ...metadata,
+              });
+            }
+            return parts;
+          }),
+        );
+        if (fail) {
+          chunks.push({
+            type: "error",
+            error,
           });
         }
-        return parts;
-      }),
+        chunks.push({
+          type: "finish",
+          finishReason: fail ? "error" : "stop",
+          usage: DEFAULT_USAGE,
+          ...metadata,
+        });
+        return chunks;
+      },
     );
-    if (fail) {
-      chunks.push({
-        type: "error",
-        error,
-      });
-    }
-    chunks.push({
-      type: "finish",
-      finishReason: fail ? "error" : "stop",
-      usage: DEFAULT_USAGE,
-      ...metadata,
-    });
+    let callIndex = 0;
     this.doGenerate = async (options) => {
       this.doGenerateCalls.push(options);
 
@@ -153,14 +182,16 @@ export class MockLanguageModel implements LanguageModelV2 {
         return args.doGenerate(options);
       } else if (Array.isArray(args.doGenerate)) {
         return args.doGenerate[this.doGenerateCalls.length];
-      } else if (content) {
-        return {
-          content,
-          finishReason: "stop",
+      } else if (contentSteps.length) {
+        const result = {
+          content: contentSteps[callIndex % contentSteps.length],
+          finishReason: "stop" as const,
           usage: DEFAULT_USAGE,
           ...metadata,
           warnings: [],
         };
+        callIndex++;
+        return result;
       } else {
         throw new Error("Unexpected: no content or doGenerate");
       }
@@ -176,12 +207,13 @@ export class MockLanguageModel implements LanguageModelV2 {
         return args.doStream(options);
       } else if (Array.isArray(args.doStream)) {
         return args.doStream[this.doStreamCalls.length];
-      } else if (content) {
+      } else if (contentSteps) {
         const stream = simulateReadableStream({
-          chunks,
+          chunks: chunkResponses[callIndex % chunkResponses.length],
           initialDelayInMs,
           chunkDelayInMs,
         });
+        callIndex++;
 
         if (options.abortSignal) {
           options.abortSignal.addEventListener("abort", () => {

package/src/component/streams.ts

@@ -3,6 +3,7 @@ import { v } from "convex/values";
 import {
   type MessageWithMetadataInternal,
   type StreamDelta,
+  type StreamMessage,
   vStreamDelta,
   vStreamMessage,
 } from "../validators.js";
@@ -136,23 +137,27 @@ export const list = query({
       ["order", "stepOrder"],
     ).take(100);
 
-    return messages.map((m) => ({
-      streamId: m._id,
-      status: m.state.kind,
-      ...pick(m, [
-        "format",
-        "order",
-        "stepOrder",
-        "userId",
-        "agentName",
-        "model",
-        "provider",
-        "providerOptions",
-      ]),
-    }));
+    return messages.map((m) => publicStreamMessage(m));
   },
 });
 
+function publicStreamMessage(m: Doc<"streamingMessages">): StreamMessage {
+  return {
+    streamId: m._id,
+    status: m.state.kind,
+    ...pick(m, [
+      "format",
+      "order",
+      "stepOrder",
+      "userId",
+      "agentName",
+      "model",
+      "provider",
+      "providerOptions",
+    ]),
+  };
+}
+
 export const abortByOrder = mutation({
   args: { threadId: v.id("threads"), order: v.number(), reason: v.string() },
   returns: v.boolean(),
@@ -526,13 +531,7 @@ export async function getStreamingMessagesWithMetadata(
           .take(1000);
         const uiMessages = await deriveUIMessagesFromDeltas(
           threadId,
-          [
-            {
-              ...streamingMessage,
-              status: "streaming",
-              streamId: streamingMessage._id,
-            },
-          ],
+          [publicStreamMessage(streamingMessage)],
           deltas,
         );
         // We don't save messages that have already been saved