@convex-dev/agent 0.2.5-alpha.0 → 0.2.5

package/src/client/index.ts

@@ -59,7 +59,7 @@ import {
   type ProviderMetadata,
   type StreamArgs,
 } from "../validators.js";
-import { createTool, wrapTools, type ToolCtx } from "./createTool.js";
+import { wrapTools, type ToolCtx } from "./createTool.js";
 import {
   listMessages,
   saveMessages,
@@ -164,7 +164,8 @@ export {
   updateThreadMetadata,
   searchThreadTitles,
 } from "./threads.js";
-export { createTool, extractText, isTool };
+export { extractText, isTool, sorted } from "../shared.js";
+export { createTool } from "./createTool.js";
 export type {
   AgentComponent,
   ContextOptions,
@@ -178,9 +179,10 @@ export type {
   ThreadDoc,
   UsageHandler,
 };
+export { mockModel } from "./mockModel.js";
 
- // 10k characters should be more than enough for most cases, and stays under
- // the 8k token limit for some models.
+// 10k characters should be more than enough for most cases, and stays under
+// the 8k token limit for some models.
 const MAX_EMBEDDING_TEXT_LENGTH = 10_000;
 
 export type Config = {
@@ -429,7 +431,12 @@ export class Agent<
     };
   }
 
-  async start<TOOLS extends ToolSet | undefined, T>(
+  async start<
+    TOOLS extends ToolSet | undefined,
+    T extends {
+      _internal?: { generateId?: IdGenerator };
+    },
+  >(
     ctx: ActionCtx & CustomCtx,
     /**
      * These are the arguments you'll pass to the LLM call such as
@@ -482,9 +489,6 @@ export class Agent<
        * aborted, it will trigger this signal when detected.
        */
       abortSignal?: AbortSignal;
-      // We optimistically override the generateId function to use the pending
-      // message id.
-      _internal?: { generateId?: IdGenerator };
       stopWhen?:
         | StopCondition<TOOLS extends undefined ? AgentTools : TOOLS>
         | Array<StopCondition<TOOLS extends undefined ? AgentTools : TOOLS>>;
@@ -495,7 +499,6 @@ export class Agent<
       system?: string;
       model: LanguageModel;
       messages: ModelMessage[];
-      // abortSignal?: AbortSignal;
       tools?: TOOLS extends undefined ? AgentTools : TOOLS;
     } & CallSettings;
     order: number;
@@ -513,23 +516,17 @@ export class Agent<
     getSavedMessages: () => MessageDoc[];
   }> {
     const { threadId, ...opts } = { ...this.options, ...options };
-    const context = await this._saveMessagesAndFetchContext(ctx, args, {
+    const context = await this._saveMessagesAndFetchContext(ctx, {
       userId: options?.userId,
       threadId: options?.threadId,
+      messages: args.messages,
+      prompt: args.prompt,
+      promptMessageId: args.promptMessageId,
       ...opts,
     });
     let pendingMessageId = context.pendingMessageId;
-    // TODO: extract pending message if one exists
-    const { args: aiArgs, promptMessageId, order, stepOrder, userId } = context;
-    const messages = context.savedMessages ?? [];
-    if (pendingMessageId) {
-      if (!aiArgs._internal?.generateId) {
-        aiArgs._internal = {
-          ...aiArgs._internal,
-          generateId: () => pendingMessageId ?? crypto.randomUUID(),
-        };
-      }
-    }
+    const { messages, promptMessageId, order, stepOrder, userId } = context;
+    const savedMessages = context.savedMessages ?? [];
     const toolCtx = {
       ...(ctx as UserActionCtx & CustomCtx),
       userId,
@@ -553,10 +550,9 @@ export class Agent<
         });
       }
     };
-    let activeModel = aiArgs.model;
-    if (aiArgs.abortSignal) {
-      const abortSignal = aiArgs.abortSignal;
-      aiArgs.abortSignal.addEventListener(
+    if (args.abortSignal) {
+      const abortSignal = args.abortSignal;
+      abortSignal.addEventListener(
         "abort",
         async () => {
           await fail(abortSignal.reason ?? "Aborted");
@@ -564,21 +560,39 @@ export class Agent<
         { once: true },
       );
     }
+    const aiArgs = {
+      ...this.options.callSettings,
+      ...this.options.providerOptions,
+      ...omit(args, ["messages", "prompt", "promptMessageId"]),
+      model: args.model ?? this.options.languageModel,
+      system: args.system ?? this.options.instructions,
+      messages,
+      stopWhen:
+        args.stopWhen ??
+        this.options.stopWhen ??
+        stepCountIs(this.options.maxSteps ?? 1),
+      tools,
+    } as T & {
+      model: LanguageModel;
+      messages: ModelMessage[];
+      tools?: TOOLS extends undefined ? AgentTools : TOOLS;
+    } & CallSettings;
+    if (pendingMessageId) {
+      if (!aiArgs._internal?.generateId) {
+        aiArgs._internal = {
+          ...aiArgs._internal,
+          generateId: () => pendingMessageId ?? crypto.randomUUID(),
+        };
+      }
+    }
+    let activeModel = aiArgs.model;
     return {
-      args: {
-        stopWhen:
-          args.stopWhen ??
-          this.options.stopWhen ??
-          stepCountIs(this.options.maxSteps ?? 1),
-        ...aiArgs,
-        tools,
-        // abortSignal: abortController.signal,
-      },
+      args: aiArgs,
       order: order ?? 0,
       stepOrder: stepOrder ?? 0,
       userId,
       promptMessageId,
-      getSavedMessages: () => messages,
+      getSavedMessages: () => savedMessages,
       updateModel: (model: LanguageModel | undefined) => {
         if (model) {
           activeModel = model;
@@ -640,18 +654,18 @@ export class Agent<
           if (createPendingMessage) {
             if (lastMessage.status === "failed") {
               pendingMessageId = undefined;
-              messages.push(...saved.messages);
+              savedMessages.push(...saved.messages);
               await fail(
                 lastMessage.error ??
                   "Aborting - the pending message was marked as failed",
               );
             } else {
               pendingMessageId = lastMessage._id;
-              messages.push(...saved.messages.slice(0, -1));
+              savedMessages.push(...saved.messages.slice(0, -1));
             }
           } else {
             pendingMessageId = undefined;
-            messages.push(...saved.messages);
+            savedMessages.push(...saved.messages);
           }
         }
         const output = "object" in toSave ? toSave.object : toSave.step;
@@ -796,17 +810,23 @@ export class Agent<
     const opts = { ...this.options, ...options };
     const streamer =
       threadId && opts.saveStreamDeltas
-        ? new DeltaStreamer(this.component, ctx, opts.saveStreamDeltas, {
-            threadId,
-            userId,
-            agentName: this.options.name,
-            model: getModelName(args.model),
-            provider: getProviderName(args.model),
-            providerOptions: args.providerOptions,
-            order,
-            stepOrder,
-            abortSignal: args.abortSignal,
-          })
+        ? new DeltaStreamer(
+            this.component,
+            ctx,
+            opts.saveStreamDeltas,
+            call.fail,
+            {
+              threadId,
+              userId,
+              agentName: this.options.name,
+              model: getModelName(args.model),
+              provider: getProviderName(args.model),
+              providerOptions: args.providerOptions,
+              order,
+              stepOrder,
+              abortSignal: args.abortSignal,
+            },
+          )
         : undefined;
 
     const result = streamText({
@@ -1652,30 +1672,23 @@ export class Agent<
     });
   }
 
-  async _saveMessagesAndFetchContext<
-    T extends {
-      prompt?: string | (ModelMessage | Message)[];
-      messages?: (ModelMessage | Message)[];
-      system?: string;
-      promptMessageId?: string;
-      pendingMessageId?: string;
-      model?: LanguageModel;
-    },
-  >(
+  async _saveMessagesAndFetchContext(
     ctx: RunActionCtx,
-    args: T,
     {
       userId: argsUserId,
       threadId,
       contextOptions,
       storageOptions,
+      ...args
     }: {
+      prompt: string | (ModelMessage | Message)[] | undefined;
+      messages: (ModelMessage | Message)[] | undefined;
+      promptMessageId: string | undefined;
       userId: string | null | undefined;
       threadId: string | undefined;
     } & Options,
   ): Promise<{
-    args: Extract<T, { model: LanguageModel; messages: ModelMessage[] }> &
-      CallSettings;
+    messages: ModelMessage[];
     userId: string | undefined;
     promptMessageId: string | undefined;
     pendingMessageId: string | undefined;
@@ -1685,7 +1698,7 @@ export class Agent<
   }> {
     // If only a promptMessageId is provided, this will be empty.
     const messages: (ModelMessage | Message)[] = args.messages ?? [];
-    const prompt: ModelMessage[] = !args.prompt
+    const promptArray: ModelMessage[] = !args.prompt
       ? []
       : Array.isArray(args.prompt)
         ? args.prompt.map((p) => deserializeMessage(p))
@@ -1722,14 +1735,14 @@ export class Agent<
     if (threadId && storageOptions?.saveMessages !== "none") {
       let saved: { messages: MessageDoc[] };
       if (
-        messages.length + prompt.length &&
+        messages.length + promptArray.length &&
         // If it was a promptMessageId, we don't want to save it again.
         (!args.promptMessageId || storageOptions?.saveMessages === "all")
       ) {
         const saveAll = storageOptions?.saveMessages === "all";
         const coreMessages: (ModelMessage | Message)[] = [
           ...messages,
-          ...prompt,
+          ...promptArray,
         ];
         const toSave = saveAll ? coreMessages : coreMessages.slice(-1);
         const metadata = Array.from({ length: toSave.length }, () => ({}));
@@ -1738,8 +1751,8 @@ export class Agent<
           userId,
           messages: [...toSave, { role: "assistant", content: [] }],
           metadata: [...metadata, { status: "pending" }],
-          failPendingSteps: !!args.pendingMessageId,
-          pendingMessageId: args.pendingMessageId,
+          // TODO: sanity check
+          failPendingSteps: !!args.promptMessageId,
         });
         promptMessageId = saved.messages.at(-2)!._id;
       } else {
@@ -1748,8 +1761,7 @@ export class Agent<
           userId,
           messages: [{ role: "assistant", content: [] }],
           metadata: [{ status: "pending" }],
-          failPendingSteps: !!args.pendingMessageId,
-          pendingMessageId: args.pendingMessageId,
+          failPendingSteps: !!args.promptMessageId,
         });
       }
       pendingMessageId = saved.messages.at(-1)!._id;
@@ -1783,7 +1795,7 @@ export class Agent<
     let processedMessages: ModelMessage[] = [
       ...prePrompt,
       ...messages,
-      ...prompt,
+      ...promptArray,
       ...existingResponses,
     ].map((m) => deserializeMessage(m));
 
@@ -1792,17 +1804,8 @@ export class Agent<
       processedMessages = await inlineMessagesFiles(processedMessages);
     }
 
-    const { prompt: _, model, ...rest } = args;
     return {
-      args: {
-        ...this.options.callSettings,
-        ...this.options.providerOptions,
-        ...rest,
-        model: model ?? this.options.languageModel,
-        system: args.system ?? this.options.instructions,
-        messages: processedMessages,
-      } as Extract<T, { model: LanguageModel; messages: ModelMessage[] }> &
-        CallSettings,
+      messages: processedMessages,
       userId,
       promptMessageId,
       pendingMessageId,

package/src/client/mockModel.ts

@@ -0,0 +1,195 @@
+import type {
+  LanguageModelV2,
+  LanguageModelV2StreamPart,
+} from "@ai-sdk/provider";
+import type { ReasoningPart, TextPart } from "@ai-sdk/provider-utils";
+import { simulateReadableStream } from "ai";
+
+const longDefaultText = `
+A A A A A A A A A A A A A A A
+B B B B B B B B B B B B B B B
+C C C C C C C C C C C C C C C
+D D D D D D D D D D D D D D D
+`;
+const defaultUsage = { outputTokens: 10, inputTokens: 3, totalTokens: 13 };
+
+export type MockModelArgs = {
+  provider?: LanguageModelV2["provider"];
+  modelId?: LanguageModelV2["modelId"];
+  supportedUrls?:
+    | LanguageModelV2["supportedUrls"]
+    | (() => LanguageModelV2["supportedUrls"]);
+  chunkDelayInMs?: number;
+  initialDelayInMs?: number;
+  // provide either content or doGenerate & doStream
+  content?: (TextPart | ReasoningPart)[];
+  doGenerate?: LanguageModelV2["doGenerate"];
+  doStream?: LanguageModelV2["doStream"];
+  fail?:
+    | boolean
+    | {
+        probability?: number;
+        error?: string;
+      };
+};
+
+export function mockModel(args?: MockModelArgs): LanguageModelV2 {
+  return new MockLanguageModel(args ?? {});
+}
+
+export class MockLanguageModel implements LanguageModelV2 {
+  readonly specificationVersion = "v2";
+
+  private _supportedUrls: () => LanguageModelV2["supportedUrls"];
+
+  readonly provider: LanguageModelV2["provider"];
+  readonly modelId: LanguageModelV2["modelId"];
+
+  doGenerate: LanguageModelV2["doGenerate"];
+  doStream: LanguageModelV2["doStream"];
+
+  doGenerateCalls: Parameters<LanguageModelV2["doGenerate"]>[0][] = [];
+  doStreamCalls: Parameters<LanguageModelV2["doStream"]>[0][] = [];
+
+  constructor(args: MockModelArgs) {
+    this.provider = args.provider || "mock-provider";
+    this.modelId = args.modelId || "mock-model-id";
+    const {
+      content = [{ type: "text", text: longDefaultText }],
+      chunkDelayInMs = 200,
+      initialDelayInMs = 1000,
+      supportedUrls = {},
+    } = args;
+    const fail =
+      args.fail &&
+      (args.fail === true ||
+        !args.fail.probability ||
+        Math.random() < args.fail.probability);
+    const error =
+      (typeof args.fail === "object" && args.fail.error) ||
+      "Mock error message";
+
+    const chunks: LanguageModelV2StreamPart[] = [
+      { type: "stream-start", warnings: [] },
+    ];
+    chunks.push(
+      ...content.flatMap((c, ci) => {
+        const deltas = c.text.split(" ");
+        let parts: LanguageModelV2StreamPart[] = [];
+        if (c.type === "reasoning") {
+          parts.push({
+            type: "reasoning-start",
+            id: `${ci}-reasoning-start`,
+          });
+          parts.push(
+            ...deltas.map(
+              (delta, di) =>
+                ({
+                  type: "reasoning-delta",
+                  delta,
+                  id: `${ci}-reasoning-${di}`,
+                  providerMetadata: {
+                    mockProvider: { mock: { reasoningDetails: null } },
+                  },
+                }) satisfies LanguageModelV2StreamPart,
+            ),
+          );
+          parts.push({
+            type: "reasoning-end",
+            id: `${ci}-reasoning-end`,
+          });
+        } else if (c.type === "text") {
+          parts.push({
+            type: "text-start",
+            id: `${ci}-text-start`,
+          });
+          parts = deltas.map((delta, di) => ({
+            type: "text-delta",
+            delta,
+            id: `${ci}-text-${di}`,
+          }));
+          parts.push({
+            type: "text-end",
+            id: `${ci}-text-end`,
+          });
+        }
+        return parts;
+      }),
+    );
+    if (fail) {
+      chunks.push({
+        type: "error",
+        error,
+      });
+    }
+    chunks.push({
+      type: "finish",
+      finishReason: fail ? "error" : "stop",
+      usage: defaultUsage,
+      providerMetadata: {
+        mockProvider: { mock: "mock metadata" },
+      },
+    });
+    this.doGenerate = async (options) => {
+      this.doGenerateCalls.push(options);
+
+      if (fail) {
+        throw new Error(error);
+      }
+      if (typeof args.doGenerate === "function") {
+        return args.doGenerate(options);
+      } else if (Array.isArray(args.doGenerate)) {
+        return args.doGenerate[this.doGenerateCalls.length];
+      } else if (content) {
+        return {
+          content,
+          finishReason: "stop",
+          usage: defaultUsage,
+          providerMetadata: { mockProvider: { mock: "mock metadata" } },
+          warnings: [],
+        };
+      } else {
+        throw new Error("Unexpected: no content or doGenerate");
+      }
+    };
+    this._supportedUrls =
+      typeof supportedUrls === "function"
+        ? supportedUrls
+        : async () => supportedUrls;
+    this.doStream = async (options) => {
+      this.doStreamCalls.push(options);
+
+      if (typeof args.doStream === "function") {
+        return args.doStream(options);
+      } else if (Array.isArray(args.doStream)) {
+        return args.doStream[this.doStreamCalls.length];
+      } else if (content) {
+        const stream = simulateReadableStream({
+          chunks,
+          initialDelayInMs,
+          chunkDelayInMs,
+        });
+        if (options.abortSignal) {
+          throw new Error("abortSignal in mock model");
+        }
+        return {
+          stream,
+          request: { body: {} },
+          response: { headers: {} },
+        };
+      } else if (args.doStream) {
+        return args.doStream;
+      } else {
+        throw new Error("Provide either content or doStream");
+      }
+    };
+    this._supportedUrls =
+      typeof supportedUrls === "function"
+        ? supportedUrls
+        : async () => supportedUrls;
+  }
+
+  get supportedUrls() {
+    return this._supportedUrls();
+  }
+}

package/src/client/streaming.ts

@@ -197,6 +197,7 @@ export class DeltaStreamer {
     public readonly component: AgentComponent,
     public readonly ctx: RunActionCtx,
     options: true | StreamingOptions,
+    private onAsyncAbort: (reason: string) => Promise<void>,
     public readonly metadata: {
       threadId: string;
       userId?: string;
@@ -210,21 +211,22 @@ export class DeltaStreamer {
     },
   ) {
     this.options =
-      typeof options === "boolean"
+      options === true
         ? DEFAULT_STREAMING_OPTIONS
         : { ...DEFAULT_STREAMING_OPTIONS, ...options };
     this.#nextParts = [];
     this.abortController = new AbortController();
     if (metadata.abortSignal) {
       metadata.abortSignal.addEventListener("abort", async () => {
+        if (this.abortController.signal.aborted) {
+          return;
+        }
         if (this.streamId) {
           this.abortController.abort();
-          const finalDelta = this.#createDelta();
           await this.#ongoingWrite;
           await this.ctx.runMutation(this.component.streams.abort, {
             streamId: this.streamId,
             reason: "abortSignal",
-            finalDelta,
           });
         }
       });
@@ -265,9 +267,12 @@ export class DeltaStreamer {
         delta,
       );
       if (!success) {
+        await this.onAsyncAbort("async abort");
         this.abortController.abort();
+        return;
       }
     } catch (e) {
+      await this.onAsyncAbort(e instanceof Error ? e.message : "unknown error");
       this.abortController.abort();
       throw e;
     }
@@ -302,11 +307,9 @@ export class DeltaStreamer {
     if (!this.streamId) {
       return;
     }
-    const finalDelta = this.#createDelta();
     await this.#ongoingWrite;
     await this.ctx.runMutation(this.component.streams.finish, {
       streamId: this.streamId,
-      finalDelta,
     });
   }
 
@@ -318,12 +321,10 @@ export class DeltaStreamer {
     if (!this.streamId) {
       return;
     }
-    const finalDelta = this.#createDelta();
     await this.#ongoingWrite;
     await this.ctx.runMutation(this.component.streams.abort, {
       streamId: this.streamId,
       reason,
-      finalDelta,
     });
   }
 }

package/src/component/messages.ts

@@ -351,7 +351,7 @@ export const finalizeMessage = mutation({
       return;
     }
     // See if we can add any in-progress data
-    if (message.message === undefined) {
+    if (!message.message?.content.length) {
       const messages = await getStreamingMessagesWithMetadata(
         ctx,
         message,