@convex-dev/agent 0.1.18 → 0.2.0

package/src/client/index.ts

@@ -1,22 +1,26 @@
-import type { EmbeddingModelV1, LanguageModelV1 } from "@ai-sdk/provider";
 import type {
-  AssistantContent,
-  CoreMessage,
-  DeepPartial,
-  FilePart,
+  FlexibleSchema,
+  IdGenerator,
+  InferSchema,
+} from "@ai-sdk/provider-utils";
+import type {
+  CallSettings,
+  EmbeddingModel,
   GenerateObjectResult,
   GenerateTextResult,
-  ImagePart,
+  LanguageModel,
+  ModelMessage,
   StepResult,
-  StreamObjectResult,
+  StopCondition,
   StreamTextResult,
+  ToolChoice,
   ToolSet,
-  UserContent,
 } from "ai";
 import {
   embedMany,
   generateObject,
   generateText,
+  stepCountIs,
   streamObject,
   streamText,
 } from "ai";
@@ -24,13 +28,12 @@ import { assert, omit, pick } from "convex-helpers";
 import {
   internalActionGeneric,
   internalMutationGeneric,
-  type GenericActionCtx,
   type GenericDataModel,
   type PaginationOptions,
   type PaginationResult,
   type WithoutSystemFields,
 } from "convex/server";
-import { v } from "convex/values";
+import { convexToJson, v, type Value } from "convex/values";
 import type { MessageDoc, ThreadDoc } from "../component/schema.js";
 import type { threadFieldsSupportingPatch } from "../component/threads.js";
 import {
@@ -38,57 +41,77 @@ import {
   type VectorDimension,
 } from "../component/vector/tables.js";
 import {
-  type AIMessageWithoutId,
   deserializeMessage,
-  promptOrMessagesToCoreMessages,
   serializeMessage,
   serializeNewMessagesInStep,
   serializeObjectResult,
 } from "../mapping.js";
 import { extractText, isTool } from "../shared.js";
 import {
-  type MessageEmbeddings,
+  vMessageEmbeddings,
+  vMessageWithMetadata,
+  vSafeObjectArgs,
+  vTextArgs,
+  type Message,
   type MessageStatus,
   type MessageWithMetadata,
   type ProviderMetadata,
   type StreamArgs,
   type Usage,
-  vMessageWithMetadata,
-  vSafeObjectArgs,
-  vTextArgs,
 } from "../validators.js";
-import { createTool, wrapTools } from "./createTool.js";
-import { listMessages } from "./listMessages.js";
-import { fetchContextMessages } from "./search.js";
+import { createTool, wrapTools, type ToolCtx } from "./createTool.js";
+import {
+  listMessages,
+  saveMessages,
+  type SaveMessageArgs,
+  type SaveMessagesArgs,
+} from "./messages.js";
+import {
+  fetchContextMessages,
+  getModelName,
+  getProviderName,
+} from "./search.js";
 import {
   DeltaStreamer,
   mergeTransforms,
-  type StreamingOptions,
   syncStreams,
+  type StreamingOptions,
 } from "./streaming.js";
+import { createThread, getThreadMetadata } from "./threads.js";
 import type {
   ActionCtx,
   AgentComponent,
   ContextOptions,
+  DefaultObjectSchema,
   GenerationOutputMetadata,
+  MaybeCustomCtx,
+  GenerateObjectArgs,
+  ObjectMode,
+  ObjectSchema,
   Options,
-  OurObjectArgs,
-  OurStreamObjectArgs,
   RawRequestResponseHandler,
   RunActionCtx,
   RunMutationCtx,
   RunQueryCtx,
   StorageOptions,
   StreamingTextArgs,
+  StreamObjectArgs,
   SyncStreamsReturnValue,
   TextArgs,
   Thread,
   UsageHandler,
   UserActionCtx,
 } from "./types.js";
+import { inlineMessagesFiles } from "./files.js";
+import type { DataModel } from "../component/_generated/dataModel.js";
 
+export { stepCountIs } from "ai";
 export { vMessageDoc, vThreadDoc } from "../component/schema.js";
-export { serializeDataOrUrl } from "../mapping.js";
+export {
+  deserializeMessage,
+  serializeDataOrUrl,
+  serializeMessage,
+} from "../mapping.js";
 // NOTE: these are also exported via @convex-dev/agent/validators
 // a future version may put them all here or move these over there
 export {
@@ -105,22 +128,26 @@ export {
   vUserMessage,
 } from "../validators.js";
 export type { ToolCtx } from "./createTool.js";
-export { getFile, storeFile } from "./files.js";
-export { filterOutOrphanedToolMessages } from "./search.js";
-export { abortStream, listStreams } from "./streaming.js";
-export {
-  createTool,
-  extractText,
-  fetchContextMessages,
-  isTool,
-  listMessages,
-  syncStreams,
-};
 export {
   definePlaygroundAPI,
-  type PlaygroundAPI,
   type AgentsFn,
+  type PlaygroundAPI,
 } from "./definePlaygroundAPI.js";
+export { getFile, storeFile } from "./files.js";
+export {
+  listMessages,
+  saveMessage,
+  saveMessages,
+  type SaveMessageArgs,
+  type SaveMessagesArgs,
+} from "./messages.js";
+export {
+  fetchContextMessages,
+  filterOutOrphanedToolMessages,
+} from "./search.js";
+export { abortStream, listStreams, syncStreams } from "./streaming.js";
+export { createThread, getThreadMetadata } from "./threads.js";
+export { createTool, extractText, isTool };
 export type {
   AgentComponent,
   ContextOptions,
@@ -136,6 +163,55 @@ export type {
   UsageHandler,
 };
 
+export type Config = {
+  /**
+   * The LLM model to use for generating / streaming text and objects.
+   * e.g.
+   * import { openai } from "@ai-sdk/openai"
+   * const myAgent = new Agent(components.agent, {
+   *   languageModel: openai.chat("gpt-4o-mini"),
+   */
+  languageModel?: LanguageModel;
+  /**
+   * The model to use for text embeddings. Optional.
+   * If specified, it will use this for generating vector embeddings
+   * of chats, and can opt-in to doing vector search for automatic context
+   * on generateText, etc.
+   * e.g.
+   * import { openai } from "@ai-sdk/openai"
+   * const myAgent = new Agent(components.agent, {
+   *   ...
+   *   textEmbeddingModel: openai.embedding("text-embedding-3-small")
+   */
+  textEmbeddingModel?: EmbeddingModel<string>;
+  /**
+   * Options to determine what messages are included as context in message
+   * generation. To disable any messages automatically being added, pass:
+   * { recentMessages: 0 }
+   */
+  contextOptions?: ContextOptions;
+  /**
+   * Determines whether messages are automatically stored when passed as
+   * arguments or generated.
+   */
+  storageOptions?: StorageOptions;
+  /**
+   * The default settings to use for the LLM calls.
+   * This can be overridden at each generate/stream callsite on a per-field
+   * basis. To clear a default setting, you'll need to pass `undefined`.
+   */
+  callSettings?: CallSettings;
+  /**
+   * The usage handler to use for this agent.
+   */
+  usageHandler?: UsageHandler;
+  /**
+   * Called for each LLM request/response, so you can do things like
+   * log the raw request body or response headers to a table, or logs.
+   */
+  rawRequestResponseHandler?: RawRequestResponseHandler;
+};
+
 export class Agent<
   /**
    * You can require that all `ctx` args to generateText & streamText
@@ -162,7 +238,7 @@ export class Agent<
 > {
   constructor(
     public component: AgentComponent,
-    public options: {
+    public options: Config & {
       /**
        * The name for the agent. This will be attributed on each message
        * created by this agent.
@@ -173,20 +249,9 @@ export class Agent<
        * e.g.
        * import { openai } from "@ai-sdk/openai"
        * const myAgent = new Agent(components.agent, {
-       *   chat: openai.chat("gpt-4o-mini"),
-       */
-      chat: LanguageModelV1;
-      /**
-       * The model to use for text embeddings. Optional.
-       * If specified, it will use this for generating vector embeddings
-       * of chats, and can opt-in to doing vector search for automatic context
-       * on generateText, etc.
-       * e.g.
-       * import { openai } from "@ai-sdk/openai"
-       * const myAgent = new Agent(components.agent, {
-       *   textEmbedding: openai.embedding("text-embedding-3-small")
+       *   languageModel: openai.chat("gpt-4o-mini"),
        */
-      textEmbedding?: EmbeddingModelV1<string>;
+      languageModel: LanguageModel;
       /**
        * The default system prompt to put in each request.
        * Override per-prompt by passing the "system" parameter.
@@ -199,36 +264,15 @@ export class Agent<
        * (import { createTool } from "@convex-dev/agent")
        */
       tools?: AgentTools;
-      /**
-       * Options to determine what messages are included as context in message
-       * generation. To disable any messages automatically being added, pass:
-       * { recentMessages: 0 }
-       */
-      contextOptions?: ContextOptions;
-      /**
-       * Determines whether messages are automatically stored when passed as
-       * arguments or generated.
-       */
-      storageOptions?: StorageOptions;
       /**
        * When generating or streaming text with tools available, this
-       * determines the default max number of iterations.
-       */
-      maxSteps?: number;
-      /**
-       * The maximum number of calls to make to an LLM in case it fails.
-       * This can be overridden at each generate/stream callsite.
+       * determines when to stop. Defaults to stepCountIs(1).
        */
-      maxRetries?: number;
+      stopWhen?: StopCondition<AgentTools> | Array<StopCondition<AgentTools>>;
       /**
-       * The usage handler to use for this agent.
+       * @deprecated Use `languageEmbeddingModel` instead.
        */
-      usageHandler?: UsageHandler;
-      /**
-       * Called for each LLM request/response, so you can do things like
-       * log the raw request body or response headers to a table, or logs.
-       */
-      rawRequestResponseHandler?: RawRequestResponseHandler;
+      chat?: LanguageModel;
     },
   ) {}
 
@@ -242,7 +286,7 @@ export class Agent<
    * @param args The thread metadata.
    * @returns The threadId of the new thread and the thread object.
    */
-  async createThread<ThreadTools extends ToolSet | undefined = undefined>(
+  async createThread(
     ctx: RunActionCtx & CustomCtx,
     args?: {
       /**
@@ -258,21 +302,8 @@ export class Agent<
        * The summary of the thread. Not currently used for anything.
        */
       summary?: string;
-      /**
-       * The usage handler to use for this thread. Overrides any handler
-       * set in the agent constructor.
-       */
-      usageHandler?: UsageHandler;
-      /**
-       * The tools to use for this thread.
-       * Overrides any tools passed in the agent constructor.
-       */
-      tools?: ThreadTools;
     },
-  ): Promise<{
-    threadId: string;
-    thread: Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>;
-  }>;
+  ): Promise<{ threadId: string; thread: Thread<AgentTools> }>;
   /**
    * Start a new thread with the agent. This will have a fresh history, though if
    * you pass in a userId you can have it search across other threads for relevant
@@ -282,7 +313,7 @@ export class Agent<
    * @param args The thread metadata.
    * @returns The threadId of the new thread.
    */
-  async createThread<ThreadTools extends ToolSet | undefined = undefined>(
+  async createThread(
     ctx: RunMutationCtx,
     args?: {
       /**
@@ -298,33 +329,12 @@ export class Agent<
        * The summary of the thread. Not currently used for anything.
        */
       summary?: string;
-      /**
-       * The usage handler to use for this thread. Overrides any handler
-       * set in the agent constructor.
-       */
-      usageHandler?: UsageHandler;
-      /**
-       * The tools to use for this thread.
-       * Overrides any tools passed in the agent constructor.
-       */
-      tools?: ThreadTools;
     },
-  ): Promise<{
-    threadId: string;
-  }>;
-  async createThread<ThreadTools extends ToolSet | undefined = undefined>(
+  ): Promise<{ threadId: string }>;
+  async createThread(
     ctx: (ActionCtx & CustomCtx) | RunMutationCtx,
-    args?: {
-      userId: string | null;
-      title?: string;
-      summary?: string;
-      usageHandler?: UsageHandler;
-      tools?: ThreadTools;
-    },
-  ): Promise<{
-    threadId: string;
-    thread?: Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>;
-  }> {
+    args?: { userId: string | null; title?: string; summary?: string },
+  ): Promise<{ threadId: string; thread?: Thread<AgentTools> }> {
     const threadId = await createThread(ctx, this.component, args);
     if (!("runAction" in ctx) || "workflowId" in ctx) {
       return { threadId };
@@ -332,13 +342,8 @@ export class Agent<
     const { thread } = await this.continueThread(ctx, {
       threadId,
       userId: args?.userId,
-      usageHandler: args?.usageHandler,
-      tools: args?.tools,
     });
-    return {
-      threadId,
-      thread,
-    };
+    return { threadId, thread };
   }
 
   /**
@@ -349,7 +354,7 @@ export class Agent<
    * @param { threadId, userId }: the thread and user to associate the messages with.
    * @returns Functions bound to the userId and threadId on a `{thread}` object.
    */
-  async continueThread<ThreadTools extends ToolSet | undefined = undefined>(
+  async continueThread(
     ctx: ActionCtx & CustomCtx,
     args: {
       /**
@@ -361,20 +366,8 @@ export class Agent<
        * relevant messages from the same user as context for the LLM calls.
        */
       userId?: string | null;
-      /**
-       * The usage handler to use for this thread. Overrides any handler
-       * set in the agent constructor.
-       */
-      usageHandler?: UsageHandler;
-      /**
-       * The tools to use for this thread.
-       * Overrides any tools passed in the agent constructor.
-       */
-      tools?: ThreadTools;
     },
-  ): Promise<{
-    thread: Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>;
-  }> {
+  ): Promise<{ thread: Thread<AgentTools> }> {
     return {
       thread: {
         threadId: args.threadId,
@@ -390,32 +383,251 @@ export class Agent<
         streamText: this.streamText.bind(this, ctx, args),
         generateObject: this.generateObject.bind(this, ctx, args),
         streamObject: this.streamObject.bind(this, ctx, args),
-      } as Thread<ThreadTools extends undefined ? AgentTools : ThreadTools>,
+      } as Thread<AgentTools>,
     };
   }
 
-  /**
-   * Search for threads by title, paginated.
-   * @param ctx The context passed from the query/mutation/action.
-   * @returns The threads matching the search, paginated.
-   */
-  async searchThreadTitles(
-    ctx: RunQueryCtx,
-    {
-      userId,
-      query,
-      limit,
-    }: {
-      userId?: string | undefined;
-      query: string;
-      limit?: number;
+  async start<TOOLS extends ToolSet | undefined, T>(
+    ctx: ActionCtx & CustomCtx,
+    /**
+     * These are the arguments you'll pass to the LLM call such as
+     * `generateText` or `streamText`. This function will look up the context
+     * and provide functions to save the steps, abort the generation, and more.
+     * The type of the arguments returned infers from the type of the arguments
+     * you pass here.
+     */
+    args: T & {
+      /**
+       * If provided, this message will be used as the "prompt" for the LLM call,
+       * instead of the prompt or messages.
+       * This is useful if you want to first save a user message, then use it as
+       * the prompt for the LLM call in another call.
+       */
+      promptMessageId?: string;
+      /**
+       * The model to use for the LLM calls. This will override the model specified
+       * in the Agent constructor.
+       */
+      model?: LanguageModel;
+      /**
+       * The tools to use for the tool calls. This will override tools specified
+       * in the Agent constructor or createThread / continueThread.
+       */
+      tools?: TOOLS;
+      /**
+       * The single prompt message to use for the LLM call. This will be the
+       * last message in the context. If it's a string, it will be a user role.
+       */
+      prompt?: string | (ModelMessage | Message)[];
+      /**
+       * If provided alongside prompt, the ordering will be:
+       * 1. system prompt
+       * 2. search context
+       * 3. recent messages
+       * 4. these messages
+       * 5. prompt messages, including those already on the same `order` as
+       *   the promptMessageId message, if provided.
+       */
+      messages?: (ModelMessage | Message)[];
+      /**
+       * This will be the first message in the context, and overrides the
+       * agent's instructions.
+       */
+      system?: string;
+      /**
+       * The abort signal to be passed to the LLM call. If triggered, it will
+       * mark the pending message as failed. If the generation is asynchronously
+       * 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 };
     },
-  ): Promise<ThreadDoc[]> {
-    return ctx.runQuery(this.component.threads.searchThreadTitles, {
-      userId,
-      query,
-      limit: limit ?? 10,
+    options?: Options & { userId?: string | null; threadId?: string },
+  ): Promise<{
+    args: T & {
+      system?: string;
+      model: LanguageModel;
+      messages: ModelMessage[];
+      // abortSignal?: AbortSignal;
+      tools?: TOOLS extends undefined ? AgentTools : TOOLS;
+    } & CallSettings;
+    order: number;
+    stepOrder: number;
+    userId: string | undefined;
+    promptMessageId: string | undefined;
+    updateModel: (model: LanguageModel | undefined) => void;
+    save: <TOOLS extends ToolSet>(
+      toSave:
+        | { step: StepResult<TOOLS> }
+        | { object: GenerateObjectResult<unknown> },
+      createPendingMessage?: boolean,
+    ) => Promise<void>;
+    fail: (reason: string) => Promise<void>;
+    getSavedMessages: () => MessageDoc[];
+  }> {
+    const { threadId, ...opts } = { ...this.options, ...options };
+    const context = await this._saveMessagesAndFetchContext(ctx, args, {
+      userId: options?.userId,
+      threadId: options?.threadId,
+      ...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 toolCtx = {
+      ...(ctx as UserActionCtx & CustomCtx),
+      userId,
+      threadId,
+      promptMessageId,
+      agent: this,
+    } satisfies ToolCtx;
+    type Tools = TOOLS extends undefined ? AgentTools : TOOLS;
+    const tools = wrapTools(toolCtx, args.tools ?? this.options.tools) as Tools;
+    const saveOutput = opts.storageOptions?.saveMessages !== "none";
+    const fail = async (reason: string) => {
+      if (threadId && promptMessageId) {
+        console.error("RollbackMessage", promptMessageId, reason);
+      }
+      if (pendingMessageId) {
+        await ctx.runMutation(this.component.messages.finalizeMessage, {
+          messageId: pendingMessageId,
+          result: { status: "failed", error: reason },
+        });
+      }
+    };
+    let activeModel = aiArgs.model;
+    if (aiArgs.abortSignal) {
+      const abortSignal = aiArgs.abortSignal;
+      aiArgs.abortSignal.addEventListener(
+        "abort",
+        async () => {
+          await fail(abortSignal.reason ?? "Aborted");
+        },
+        { once: true },
+      );
+    }
+    return {
+      args: {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        stopWhen: (args as any).stopWhen ?? this.options.stopWhen,
+        ...aiArgs,
+        tools,
+        // abortSignal: abortController.signal,
+      },
+      order: order ?? 0,
+      stepOrder: stepOrder ?? 0,
+      userId,
+      promptMessageId,
+      getSavedMessages: () => messages,
+      updateModel: (model: LanguageModel | undefined) => {
+        if (model) {
+          activeModel = model;
+        }
+      },
+      fail,
+      save: async <TOOLS extends ToolSet>(
+        toSave:
+          | { step: StepResult<TOOLS> }
+          | { object: GenerateObjectResult<unknown> },
+        createPendingMessage?: boolean,
+      ) => {
+        if (threadId && promptMessageId && saveOutput) {
+          const metadata = {
+            // TODO: get up to date one when user selects mid-generation
+            model: getModelName(activeModel),
+            provider: getProviderName(activeModel),
+          };
+          const serialized =
+            "object" in toSave
+              ? await serializeObjectResult(
+                  ctx,
+                  this.component,
+                  toSave.object,
+                  metadata,
+                )
+              : await serializeNewMessagesInStep(
+                  ctx,
+                  this.component,
+                  toSave.step,
+                  metadata,
+                );
+          const embeddings = await this.generateEmbeddings(
+            ctx,
+            { userId, threadId },
+            serialized.messages.map((m) => m.message),
+          );
+          if (createPendingMessage) {
+            serialized.messages.push({
+              message: { role: "assistant", content: [] },
+              status: "pending",
+            });
+            embeddings?.vectors.push(null);
+          }
+          const saved = await ctx.runMutation(
+            this.component.messages.addMessages,
+            {
+              userId,
+              threadId,
+              agentName: this.options.name,
+              promptMessageId,
+              pendingMessageId,
+              messages: serialized.messages,
+              embeddings,
+              failPendingSteps: false,
+            },
+          );
+          const lastMessage = saved.messages.at(-1)!;
+          if (createPendingMessage) {
+            if (lastMessage.status === "failed") {
+              pendingMessageId = undefined;
+              messages.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));
+            }
+          } else {
+            pendingMessageId = undefined;
+            messages.push(...saved.messages);
+          }
+        }
+        const output = "object" in toSave ? toSave.object : toSave.step;
+        if (this.options.rawRequestResponseHandler) {
+          await this.options.rawRequestResponseHandler(ctx, {
+            userId,
+            threadId,
+            agentName: this.options.name,
+            request: output.request,
+            response: output.response,
+          });
+        }
+        if (opts.usageHandler && output.usage) {
+          await opts.usageHandler(ctx, {
+            userId,
+            threadId,
+            agentName: this.options.name,
+            model: getModelName(activeModel),
+            provider: getProviderName(activeModel),
+            usage: output.usage,
+            providerMetadata: output.providerMetadata,
+          });
+        }
+      },
+    };
   }
 
   /**
@@ -426,7 +638,7 @@ export class Agent<
    * to a thread (and optionally userId).
    * @param ctx The context passed from the action function calling this.
    * @param { userId, threadId }: The user and thread to associate the message with
-   * @param args The arguments to the generateText function, along with extra controls
+   * @param generateTextArgs The arguments to the generateText function, along with extra controls
    * for the {@link ContextOptions} and {@link StorageOptions}.
    * @returns The result of the generateText function.
    */
@@ -436,102 +648,44 @@ export class Agent<
     OUTPUT_PARTIAL = never,
   >(
     ctx: ActionCtx & CustomCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      tools: threadTools,
-      ...usageHandler
-    }: {
-      userId?: string | null;
-      threadId?: string;
-      /**
-       * The usage handler to use for this thread. Overrides any handler
-       * set in the agent constructor.
-       */
-      usageHandler?: UsageHandler;
-      /** Note: to get better type inference, pass tools in the next arg */
-      tools?: ToolSet;
-    },
-    args: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL>,
+    threadOpts: { userId?: string | null; threadId?: string },
+    generateTextArgs: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL>,
     options?: Options,
   ): Promise<
     GenerateTextResult<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT> &
       GenerationOutputMetadata
   > {
-    const opts = { ...this.options, ...options, ...usageHandler };
-    const context = await this._saveMessagesAndFetchContext(ctx, args, {
-      userId: argsUserId ?? undefined,
-      threadId,
-      ...opts,
-    });
-    const { args: aiArgs, messageId, order, userId } = context;
-    const toolCtx = {
-      ...(ctx as UserActionCtx & CustomCtx),
-      userId,
-      threadId,
-      messageId,
-      agent: this,
-    };
-    const tools = wrapTools(
-      toolCtx,
-      args.tools ?? threadTools ?? this.options.tools,
-    ) as TOOLS extends undefined ? AgentTools : TOOLS;
-    const saveOutput = opts.storageOptions?.saveMessages !== "none";
+    const { args, promptMessageId, order, ...call } = await this.start(
+      ctx,
+      generateTextArgs,
+      { ...threadOpts, ...options },
+    );
+
+    type Tools = TOOLS extends undefined ? AgentTools : TOOLS;
+    const steps: StepResult<Tools>[] = [];
     try {
-      const result = (await generateText({
-        // Can be overridden
-        maxSteps: this.options.maxSteps,
-        ...aiArgs,
-        tools,
+      const result = (await generateText<Tools, OUTPUT, OUTPUT_PARTIAL>({
+        ...args,
+        prepareStep: async (options) => {
+          const result = await generateTextArgs.prepareStep?.(options);
+          call.updateModel(result?.model ?? options.model);
+          return result;
+        },
         onStepFinish: async (step) => {
-          if (threadId && messageId && saveOutput) {
-            await this.saveStep(ctx, {
-              userId,
-              threadId,
-              promptMessageId: messageId,
-              model: aiArgs.model.modelId,
-              provider: aiArgs.model.provider,
-              step,
-            });
-          }
-          if (this.options.rawRequestResponseHandler) {
-            await this.options.rawRequestResponseHandler(ctx, {
-              userId,
-              threadId,
-              agentName: this.options.name,
-              request: step.request,
-              response: step.response,
-            });
-          }
-          if (opts.usageHandler && step.usage) {
-            await opts.usageHandler(ctx, {
-              userId,
-              threadId,
-              agentName: this.options.name,
-              model: aiArgs.model.modelId,
-              provider: aiArgs.model.provider,
-              usage: step.usage,
-              providerMetadata: step.providerMetadata,
-            });
-          }
-          return args.onStepFinish?.(step);
+          steps.push(step);
+          await call.save({ step }, await willContinue(steps, args.stopWhen));
+          return generateTextArgs.onStepFinish?.(step);
         },
-      })) as GenerateTextResult<
-        TOOLS extends undefined ? AgentTools : TOOLS,
-        OUTPUT
-      > &
-        GenerationOutputMetadata;
-      result.messageId = messageId;
-      result.order = order;
-      return result;
+      })) as GenerateTextResult<Tools, OUTPUT>;
+      const metadata: GenerationOutputMetadata = {
+        promptMessageId,
+        order,
+        savedMessages: call.getSavedMessages(),
+        messageId: promptMessageId,
+      };
+      return Object.assign(result, metadata);
     } catch (error) {
-      if (threadId && messageId) {
-        console.error("RollbackMessage", messageId);
-        await ctx.runMutation(this.component.messages.rollbackMessage, {
-          messageId,
-          error: (error as Error).message,
-        });
-      }
+      await call.fail(errorToString(error));
       throw error;
     }
   }
@@ -549,22 +703,16 @@ export class Agent<
     PARTIAL_OUTPUT = never,
   >(
     ctx: ActionCtx & CustomCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      /** Note: to get better type inference, pass tools in the next arg */
-      tools: threadTools,
-      ...usageHandler
-    }: {
-      userId?: string | null;
-      threadId?: string;
-      usageHandler?: UsageHandler;
-      tools?: ToolSet;
-    },
+    threadOpts: { userId?: string | null; threadId?: string },
     /**
      * The arguments to the streamText function, similar to the ai `streamText` function.
      */
-    args: StreamingTextArgs<AgentTools, TOOLS, OUTPUT, PARTIAL_OUTPUT>,
+    streamTextArgs: StreamingTextArgs<
+      AgentTools,
+      TOOLS,
+      OUTPUT,
+      PARTIAL_OUTPUT
+    >,
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
@@ -589,109 +737,80 @@ export class Agent<
     > &
       GenerationOutputMetadata
   > {
-    const opts = { ...this.options, ...options, ...usageHandler };
-    const context = await this._saveMessagesAndFetchContext(ctx, args, {
-      userId: argsUserId ?? undefined,
-      threadId,
-      ...opts,
-    });
-    const { args: aiArgs, messageId, order, stepOrder, userId } = context;
-    const toolCtx = {
-      ...(ctx as UserActionCtx & CustomCtx),
-      userId,
-      threadId,
-      messageId,
-      agent: this,
-    };
-    const tools = wrapTools(
-      toolCtx,
-      args.tools ?? threadTools ?? this.options.tools,
-    ) as TOOLS extends undefined ? AgentTools : TOOLS;
-    const saveOutput = opts.storageOptions?.saveMessages !== "none";
+    const { threadId } = threadOpts;
+    const { args, userId, order, stepOrder, promptMessageId, ...call } =
+      await this.start(ctx, streamTextArgs, { ...threadOpts, ...options });
+
+    type Tools = TOOLS extends undefined ? AgentTools : TOOLS;
+    const steps: StepResult<Tools>[] = [];
+
+    const opts = { ...this.options, ...options };
     const streamer =
       threadId && opts.saveStreamDeltas
         ? new DeltaStreamer(this.component, ctx, opts.saveStreamDeltas, {
             threadId,
             userId,
             agentName: this.options.name,
-            model: aiArgs.model.modelId,
-            provider: aiArgs.model.provider,
-            providerOptions: aiArgs.providerOptions,
+            model: getModelName(args.model),
+            provider: getProviderName(args.model),
+            providerOptions: args.providerOptions,
             order,
             stepOrder,
-            abortSignal: aiArgs.abortSignal,
+            abortSignal: args.abortSignal,
           })
         : undefined;
 
     const result = streamText({
-      // Can be overridden
-      maxSteps: this.options.maxSteps,
-      ...aiArgs,
-      tools,
-      abortSignal: streamer?.abortController.signal ?? aiArgs.abortSignal,
+      ...args,
+      abortSignal: streamer?.abortController.signal ?? args.abortSignal,
+      // TODO: this is probably why reasoning isn't streaming
       experimental_transform: mergeTransforms(
         options?.saveStreamDeltas,
-        args.experimental_transform,
+        streamTextArgs.experimental_transform,
       ),
       onChunk: async (event) => {
         await streamer?.addParts([event.chunk]);
         // console.log("onChunk", chunk);
-        return args.onChunk?.(event);
+        return streamTextArgs.onChunk?.(event);
       },
       onError: async (error) => {
         console.error("onError", error);
-        if (threadId && messageId && saveOutput) {
-          await ctx.runMutation(this.component.messages.rollbackMessage, {
-            messageId,
-            error: (error.error as Error).message,
-          });
+        await call.fail(errorToString(error.error));
+        await streamer?.fail(errorToString(error.error));
+        return streamTextArgs.onError?.(error);
+      },
+      // onFinish: async (event) => {
+      //   return streamTextArgs.onFinish?.(event);
+      // },
+      prepareStep: async (options) => {
+        const result = await streamTextArgs.prepareStep?.(options);
+        if (result) {
+          const model = result.model ?? options.model;
+          call.updateModel(model);
+          return result;
         }
-        // TODO: update the streamer to error state
-        return args.onError?.(error);
+        return undefined;
       },
       onStepFinish: async (step) => {
-        // console.log("onStepFinish", step);
-        if (threadId && messageId && saveOutput) {
-          const saved = await this.saveStep(ctx, {
-            userId,
-            threadId,
-            model: aiArgs.model.modelId,
-            provider: aiArgs.model.provider,
-            promptMessageId: messageId,
-            step,
-          });
-          await streamer?.finish(saved.messages);
-        }
-        if (this.options.rawRequestResponseHandler) {
-          await this.options.rawRequestResponseHandler(ctx, {
-            userId,
-            threadId,
-            agentName: this.options.name,
-            request: step.request,
-            response: step.response,
-          });
-        }
-        if (opts.usageHandler && step.usage) {
-          await opts.usageHandler(ctx, {
-            userId,
-            threadId,
-            agentName: this.options.name,
-            model: aiArgs.model.modelId,
-            provider: aiArgs.model.provider,
-            usage: step.usage,
-            providerMetadata: step.providerMetadata,
-          });
+        steps.push(step);
+        const createPendingMessage = await willContinue(steps, args.stopWhen);
+        await call.save({ step }, createPendingMessage);
+        if (!createPendingMessage) {
+          await streamer?.finish();
         }
         return args.onStepFinish?.(step);
       },
     }) as StreamTextResult<
       TOOLS extends undefined ? AgentTools : TOOLS,
       PARTIAL_OUTPUT
-    > &
-      GenerationOutputMetadata;
-    result.messageId = messageId;
-    result.order = order;
-    return result;
+    >;
+    const metadata: GenerationOutputMetadata = {
+      promptMessageId,
+      order,
+      savedMessages: call.getSavedMessages(),
+      messageId: promptMessageId,
+    };
+    return Object.assign(result, metadata);
   }
 
   /**
@@ -701,81 +820,45 @@ export class Agent<
    * Use {@link continueThread} to get a version of this function already scoped
    * to a thread (and optionally userId).
    */
-  async generateObject<T>(
-    ctx: ActionCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      ...usageHandler
-    }: {
-      userId?: string | null;
-      threadId?: string;
-      usageHandler?: UsageHandler;
-    },
+  async generateObject<
+    SCHEMA extends ObjectSchema = DefaultObjectSchema,
+    OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+      ? "enum"
+      : "object",
+    RESULT = OUTPUT extends "array"
+      ? Array<InferSchema<SCHEMA>>
+      : InferSchema<SCHEMA>,
+  >(
+    ctx: ActionCtx & CustomCtx,
+    threadOpts: { userId?: string | null; threadId?: string },
     /**
      * The arguments to the generateObject function, similar to the ai.generateObject function.
      */
-    args: OurObjectArgs<T>,
+    generateObjectArgs: GenerateObjectArgs<SCHEMA, OUTPUT, RESULT>,
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
      */
     options?: Options,
-  ): Promise<GenerateObjectResult<T> & GenerationOutputMetadata> {
-    const opts = { ...this.options, ...options, ...usageHandler };
-    const context = await this._saveMessagesAndFetchContext(ctx, args, {
-      userId: argsUserId ?? undefined,
-      threadId,
-      ...opts,
-    });
-    const { args: aiArgs, messageId, order, userId } = context;
-    const saveOutput = opts.storageOptions?.saveMessages !== "none";
+  ): Promise<GenerateObjectResult<RESULT> & GenerationOutputMetadata> {
+    const { args, promptMessageId, order, fail, save, getSavedMessages } =
+      await this.start(ctx, generateObjectArgs, { ...threadOpts, ...options });
+
     try {
       const result = (await generateObject(
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        aiArgs as any,
-      )) as GenerateObjectResult<T> & GenerationOutputMetadata;
+        args,
+      )) as GenerateObjectResult<RESULT>;
 
-      if (threadId && messageId && saveOutput) {
-        await this.saveObject(ctx, {
-          threadId,
-          promptMessageId: messageId,
-          result,
-          userId,
-          model: aiArgs.model.modelId,
-          provider: aiArgs.model.provider,
-        });
-      }
-      result.messageId = messageId;
-      result.order = order;
-      if (this.options.rawRequestResponseHandler) {
-        await this.options.rawRequestResponseHandler(ctx, {
-          userId,
-          threadId,
-          agentName: this.options.name,
-          request: result.request,
-          response: result.response,
-        });
-      }
-      if (opts.usageHandler && result.usage) {
-        await opts.usageHandler(ctx, {
-          userId,
-          threadId,
-          agentName: this.options.name,
-          model: aiArgs.model.modelId,
-          provider: aiArgs.model.provider,
-          usage: result.usage,
-          providerMetadata: result.providerMetadata,
-        });
-      }
-      return result;
+      await save({ object: result });
+      const metadata: GenerationOutputMetadata = {
+        promptMessageId,
+        order,
+        savedMessages: getSavedMessages(),
+        messageId: promptMessageId,
+      };
+      return Object.assign(result, metadata);
     } catch (error) {
-      if (threadId && messageId) {
-        await ctx.runMutation(this.component.messages.rollbackMessage, {
-          messageId,
-          error: (error as Error).message,
-        });
-      }
+      await fail(errorToString(error));
       throw error;
     }
   }
@@ -787,96 +870,83 @@ export class Agent<
    * Use {@link continueThread} to get a version of this function already scoped
    * to a thread (and optionally userId).
    */
-  async streamObject<T>(
-    ctx: ActionCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      ...usageHandler
-    }: {
-      userId?: string | null;
-      threadId?: string;
-      usageHandler?: UsageHandler;
-    },
+  async streamObject<
+    SCHEMA extends ObjectSchema = DefaultObjectSchema,
+    OUTPUT extends ObjectMode = InferSchema<SCHEMA> extends string
+      ? "enum"
+      : "object",
+    RESULT = OUTPUT extends "array"
+      ? Array<InferSchema<SCHEMA>>
+      : InferSchema<SCHEMA>,
+  >(
+    ctx: ActionCtx & CustomCtx,
+    threadOpts: { userId?: string | null; threadId?: string },
     /**
      * The arguments to the streamObject function, similar to the ai `streamObject` function.
      */
-    args: OurStreamObjectArgs<T>,
+    streamObjectArgs: StreamObjectArgs<SCHEMA, OUTPUT, RESULT> & {
+      /**
+       * If provided, this message will be used as the "prompt" for the LLM call,
+       * instead of the prompt or messages.
+       * This is useful if you want to first save a user message, then use it as
+       * the prompt for the LLM call in another call.
+       */
+      promptMessageId?: string;
+      /**
+       * The model to use for the LLM calls. This will override the model specified
+       * in the Agent constructor.
+       */
+      model?: LanguageModel;
+      /**
+       * The tools to use for the tool calls. This will override tools specified
+       * in the Agent constructor or createThread / continueThread.
+       */
+    },
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
      */
     options?: Options,
   ): Promise<
-    StreamObjectResult<DeepPartial<T>, T, never> & GenerationOutputMetadata
+    ReturnType<typeof streamObject<SCHEMA, OUTPUT, RESULT>> &
+      GenerationOutputMetadata
   > {
-    // TODO: unify all this shared code between all the generate* and stream* functions
-    const opts = { ...this.options, ...options, ...usageHandler };
-    const context = await this._saveMessagesAndFetchContext(ctx, args, {
-      userId: argsUserId ?? undefined,
-      threadId,
-      ...opts,
-    });
-    const { args: aiArgs, messageId, order, userId } = context;
-    const saveOutput = opts.storageOptions?.saveMessages !== "none";
-    const stream = streamObject<T>({
+    const { args, promptMessageId, order, fail, save, getSavedMessages } =
+      await this.start(ctx, streamObjectArgs, { ...threadOpts, ...options });
+
+    const stream = streamObject<SCHEMA, OUTPUT, RESULT>({
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      ...(aiArgs as any),
+      ...(args as any),
       onError: async (error) => {
-        console.error("onError", error);
+        console.error(" streamObject onError", error);
+        // TODO: content that we have so far
+        // content: stream.fullStream.
+        await fail(errorToString(error.error));
         return args.onError?.(error);
       },
       onFinish: async (result) => {
-        if (threadId && messageId && saveOutput) {
-          await this.saveObject(ctx, {
-            userId,
-            threadId,
-            promptMessageId: messageId,
-            result: {
-              object: result.object,
-              finishReason: "stop",
-              usage: result.usage,
-              warnings: result.warnings,
-              request: await stream.request,
-              response: result.response,
-              providerMetadata: result.providerMetadata,
-              experimental_providerMetadata:
-                result.experimental_providerMetadata,
-              logprobs: undefined,
-              toJsonResponse: stream.toTextStreamResponse,
-            },
-            model: aiArgs.model.modelId,
-            provider: aiArgs.model.provider,
-          });
-        }
-        if (opts.usageHandler && result.usage) {
-          await opts.usageHandler(ctx, {
-            userId,
-            threadId,
-            agentName: this.options.name,
-            model: aiArgs.model.modelId,
-            provider: aiArgs.model.provider,
+        await save({
+          object: {
+            object: result.object,
+            finishReason: result.error ? "error" : "stop",
             usage: result.usage,
-            providerMetadata: result.providerMetadata,
-          });
-        }
-        if (this.options.rawRequestResponseHandler) {
-          await this.options.rawRequestResponseHandler(ctx, {
-            userId,
-            threadId,
-            agentName: this.options.name,
+            warnings: result.warnings,
             request: await stream.request,
             response: result.response,
-          });
-        }
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        return args.onFinish?.(result as any);
+            providerMetadata: result.providerMetadata,
+            toJsonResponse: stream.toTextStreamResponse,
+          },
+        });
+        return args.onFinish?.(result);
       },
-    }) as StreamObjectResult<DeepPartial<T>, T, never> &
-      GenerationOutputMetadata;
-    stream.messageId = messageId;
-    stream.order = order;
-    return stream;
+    });
+    const metadata: GenerationOutputMetadata = {
+      promptMessageId,
+      order,
+      savedMessages: getSavedMessages(),
+      messageId: promptMessageId,
+    };
+    return Object.assign(stream, metadata);
   }
 
   /**
@@ -898,14 +968,11 @@ export class Agent<
       skipEmbeddings?: boolean;
     },
   ) {
-    const { lastMessageId, messages } = await this.saveMessages(ctx, {
+    const { messages } = await this.saveMessages(ctx, {
       threadId: args.threadId,
       userId: args.userId,
       embeddings: args.embedding
-        ? {
-            model: args.embedding.model,
-            vectors: [args.embedding.vector],
-          }
+        ? { model: args.embedding.model, vectors: [args.embedding.vector] }
         : undefined,
       messages:
         args.prompt !== undefined
@@ -913,8 +980,10 @@ export class Agent<
           : [args.message],
       metadata: args.metadata ? [args.metadata] : undefined,
       skipEmbeddings: args.skipEmbeddings,
+      pendingMessageId: args.pendingMessageId,
     });
-    return { messageId: lastMessageId, message: messages.at(-1)! };
+    const message = messages.at(-1)!;
+    return { messageId: message._id, message };
   }
 
   /**
@@ -936,20 +1005,12 @@ export class Agent<
        */
       skipEmbeddings?: boolean;
     },
-  ): Promise<{
-    lastMessageId: string;
-    messages: MessageDoc[];
-  }> {
-    let embeddings:
-      | {
-          vectors: (number[] | null)[];
-          model: string;
-        }
-      | undefined;
+  ): Promise<{ messages: MessageDoc[] }> {
+    let embeddings: { vectors: (number[] | null)[]; model: string } | undefined;
     const { skipEmbeddings, ...rest } = args;
     if (args.embeddings) {
       embeddings = args.embeddings;
-    } else if (!skipEmbeddings && this.options.textEmbedding) {
+    } else if (!skipEmbeddings && this.options.textEmbeddingModel) {
       if (!("runAction" in ctx)) {
         console.warn(
           "You're trying to save messages and generate embeddings, but you're in a mutation. " +
@@ -967,10 +1028,7 @@ export class Agent<
       } else {
         embeddings = await this.generateEmbeddings(
           ctx,
-          {
-            userId: args.userId ?? undefined,
-            threadId: args.threadId,
-          },
+          { userId: args.userId ?? undefined, threadId: args.threadId },
           args.messages,
         );
       }
@@ -1037,7 +1095,7 @@ export class Agent<
     args: {
       userId: string | undefined;
       threadId: string | undefined;
-      messages: CoreMessage[];
+      messages: (ModelMessage | Message)[];
       /**
        * If provided, it will search for messages up to and including this message.
        * Note: if this is far in the past, text and vector search results may be more
@@ -1058,8 +1116,8 @@ export class Agent<
       getEmbedding: async (text) => {
         assert("runAction" in ctx);
         assert(
-          this.options.textEmbedding,
-          "A textEmbedding model is required to be set on the Agent that you're doing vector search with",
+          this.options.textEmbeddingModel,
+          "A textEmbeddingModel is required to be set on the Agent that you're doing vector search with",
         );
         return {
           embedding: (
@@ -1069,7 +1127,7 @@ export class Agent<
               values: [text],
             })
           ).embeddings[0],
-          embeddingModel: this.options.textEmbedding.modelId,
+          textEmbeddingModel: this.options.textEmbeddingModel,
         };
       },
     });
@@ -1121,13 +1179,10 @@ export class Agent<
     {
       userId,
       threadId,
-    }: {
-      userId: string | undefined;
-      threadId: string | undefined;
-    },
-    messages: CoreMessage[],
+    }: { userId: string | undefined; threadId: string | undefined },
+    messages: (ModelMessage | Message)[],
   ) {
-    if (!this.options.textEmbedding) {
+    if (!this.options.textEmbeddingModel) {
       return undefined;
     }
     let embeddings:
@@ -1145,13 +1200,13 @@ export class Agent<
     if (textIndexes.length === 0) {
       return undefined;
     }
+    const values = messageTexts.filter((t): t is string => !!t);
     // Then embed those messages.
     const textEmbeddings = await this.doEmbed(ctx, {
       userId,
       threadId,
-      values: messageTexts as string[],
+      values,
     });
-    // TODO: record usage of embeddings
     // Then assemble the embeddings into a single array with nulls for the messages without text.
     const embeddingsOrNull = Array(messages.length).fill(null);
     textIndexes.forEach((i, j) => {
@@ -1160,11 +1215,8 @@ export class Agent<
     if (textEmbeddings.embeddings.length > 0) {
       const dimension = textEmbeddings.embeddings[0].length;
       validateVectorDimension(dimension);
-      embeddings = {
-        vectors: embeddingsOrNull,
-        dimension,
-        model: this.options.textEmbedding.modelId,
-      };
+      const model = getModelName(this.options.textEmbeddingModel);
+      embeddings = { vectors: embeddingsOrNull, dimension, model };
     }
     return embeddings;
   }
@@ -1178,9 +1230,7 @@ export class Agent<
    */
   async generateAndSaveEmbeddings(
     ctx: RunActionCtx,
-    args: {
-      messageIds: string[];
-    },
+    args: { messageIds: string[] },
   ) {
     const messages = (
       await ctx.runQuery(this.component.messages.getMessagesByIds, {
@@ -1218,12 +1268,12 @@ export class Agent<
         userId: messagesMissingEmbeddings[0]!.userId,
         threadId: messagesMissingEmbeddings[0]!.threadId,
       },
-      messagesMissingEmbeddings.map((m) => m!.message!),
+      messagesMissingEmbeddings.map((m) => deserializeMessage(m!.message!)),
     );
     if (!embeddings) {
-      if (!this.options.textEmbedding) {
+      if (!this.options.textEmbeddingModel) {
         throw new Error(
-          "No embeddings were generated for the messages. You must pass a textEmbedding model to the agent constructor.",
+          "No embeddings were generated for the messages. You must pass a textEmbeddingModel to the agent constructor.",
         );
       }
       throw new Error(
@@ -1278,14 +1328,14 @@ export class Agent<
        */
       provider?: string;
     },
-  ): Promise<{ messages: MessageDoc[]; pending?: MessageDoc }> {
-    const messages = await serializeNewMessagesInStep(
+  ): Promise<{ messages: MessageDoc[] }> {
+    const { messages } = await serializeNewMessagesInStep(
       ctx,
       this.component,
       args.step,
       {
-        provider: args.provider ?? this.options.chat.provider,
-        model: args.model ?? this.options.chat.modelId,
+        provider: args.provider ?? getProviderName(this.options.languageModel),
+        model: args.model ?? getModelName(this.options.languageModel),
       },
     );
     const embeddings = await this.generateEmbeddings(
@@ -1293,7 +1343,7 @@ export class Agent<
       { userId: args.userId, threadId: args.threadId },
       messages.map((m) => m.message),
     );
-    const saved = await ctx.runMutation(this.component.messages.addMessages, {
+    return ctx.runMutation(this.component.messages.addMessages, {
       userId: args.userId,
       threadId: args.threadId,
       agentName: this.options.name,
@@ -1302,7 +1352,6 @@ export class Agent<
       embeddings,
       failPendingSteps: false,
     });
-    return saved;
   }
 
   /**
@@ -1323,14 +1372,20 @@ export class Agent<
       result: GenerateObjectResult<unknown>;
       metadata?: Omit<MessageWithMetadata, "message">;
     },
-  ): Promise<void> {
+  ): Promise<{ messages: MessageDoc[] }> {
     const { messages } = await serializeObjectResult(
       ctx,
       this.component,
       args.result,
       {
-        model: args.model ?? this.options.chat.modelId,
-        provider: args.provider ?? this.options.chat.provider,
+        model:
+          args.model ??
+          args.metadata?.model ??
+          getModelName(this.options.languageModel),
+        provider:
+          args.provider ??
+          args.metadata?.provider ??
+          getProviderName(this.options.languageModel),
       },
     );
     const embeddings = await this.generateEmbeddings(
@@ -1339,7 +1394,7 @@ export class Agent<
       messages.map((m) => m.message),
     );
 
-    await ctx.runMutation(this.component.messages.addMessages, {
+    return ctx.runMutation(this.component.messages.addMessages, {
       userId: args.userId,
       threadId: args.threadId,
       promptMessageId: args.promptMessageId,
@@ -1347,7 +1402,6 @@ export class Agent<
       messages,
       embeddings,
       agentName: this.options.name,
-      pending: false,
     });
   }
 
@@ -1359,25 +1413,17 @@ export class Agent<
    * @param args What message to save. Generally the parent message sent into
    *   the generateText call.
    */
-  async completeMessage(
+  async finalizeMessage(
     ctx: RunMutationCtx,
     args: {
-      threadId: string;
       messageId: string;
-      result: { kind: "error"; error: string } | { kind: "success" };
+      result: { status: "failed"; error: string } | { status: "success" };
     },
   ): Promise<void> {
-    const result = args.result;
-    if (result.kind === "success") {
-      await ctx.runMutation(this.component.messages.commitMessage, {
-        messageId: args.messageId,
-      });
-    } else {
-      await ctx.runMutation(this.component.messages.rollbackMessage, {
-        messageId: args.messageId,
-        error: result.error,
-      });
-    }
+    await ctx.runMutation(this.component.messages.finalizeMessage, {
+      messageId: args.messageId,
+      result: args.result,
+    });
   }
 
   /**
@@ -1392,7 +1438,7 @@ export class Agent<
       messageId: string;
       patch: {
         /** The message to replace the existing message. */
-        message: CoreMessage & { id?: string };
+        message: ModelMessage | Message;
         /** The status to set on the message. */
         status: "success" | "error";
         /** The error message to set on the message. */
@@ -1437,9 +1483,7 @@ export class Agent<
    */
   async deleteMessages(
     ctx: RunMutationCtx,
-    args: {
-      messageIds: string[];
-    },
+    args: { messageIds: string[] },
   ): Promise<void> {
     await ctx.runMutation(this.component.messages.deleteByIds, args);
   }
@@ -1452,9 +1496,7 @@ export class Agent<
    */
   async deleteMessage(
     ctx: RunMutationCtx,
-    args: {
-      messageId: string;
-    },
+    args: { messageId: string },
   ): Promise<void> {
     await ctx.runMutation(this.component.messages.deleteByIds, {
       messageIds: [args.messageId],
@@ -1526,10 +1568,7 @@ export class Agent<
    */
   async deleteThreadAsync(
     ctx: RunMutationCtx,
-    args: {
-      threadId: string;
-      pageSize?: number;
-    },
+    args: { threadId: string; pageSize?: number },
   ): Promise<void> {
     await ctx.runMutation(this.component.threads.deleteAllForThreadIdAsync, {
       threadId: args.threadId,
@@ -1546,10 +1585,7 @@ export class Agent<
    */
   async deleteThreadSync(
     ctx: RunActionCtx,
-    args: {
-      threadId: string;
-      pageSize?: number;
-    },
+    args: { threadId: string; pageSize?: number },
   ): Promise<void> {
     await ctx.runAction(this.component.threads.deleteAllForThreadIdSync, {
       threadId: args.threadId,
@@ -1559,13 +1595,12 @@ export class Agent<
 
   async _saveMessagesAndFetchContext<
     T extends {
-      id?: string;
-      prompt?: string;
-      messages?: CoreMessage[] | AIMessageWithoutId[];
+      prompt?: string | (ModelMessage | Message)[];
+      messages?: (ModelMessage | Message)[];
       system?: string;
       promptMessageId?: string;
-      model?: LanguageModelV1;
-      maxRetries?: number;
+      pendingMessageId?: string;
+      model?: LanguageModel;
     },
   >(
     ctx: RunActionCtx,
@@ -1576,23 +1611,32 @@ export class Agent<
       contextOptions,
       storageOptions,
     }: {
-      userId: string | undefined;
+      userId: string | null | undefined;
       threadId: string | undefined;
     } & Options,
   ): Promise<{
-    args: T & { model: LanguageModelV1 };
+    args: Extract<T, { model: LanguageModel; messages: ModelMessage[] }> &
+      CallSettings;
     userId: string | undefined;
-    messageId: string | undefined;
+    promptMessageId: string | undefined;
+    pendingMessageId: string | undefined;
     order: number | undefined;
     stepOrder: number | undefined;
+    savedMessages: MessageDoc[] | undefined;
   }> {
     // If only a promptMessageId is provided, this will be empty.
-    const messages = promptOrMessagesToCoreMessages(args);
+    const messages = args.messages ?? [];
+    const prompt: (ModelMessage | Message)[] = !args.prompt
+      ? []
+      : Array.isArray(args.prompt)
+        ? args.prompt
+        : [{ role: "user", content: args.prompt }];
     const userId =
       argsUserId ??
       (threadId &&
         (await ctx.runQuery(this.component.threads.getThread, { threadId }))
-          ?.userId);
+          ?.userId) ??
+      undefined;
     // If only a messageId is provided, this will add that message to the end.
     const contextMessages = await this.fetchContextMessages(ctx, {
       userId,
@@ -1603,74 +1647,111 @@ export class Agent<
     });
     // If it was a promptMessageId, pop it off context messages
     // and add to the end of messages.
-    // TODO: slice it from the prompt message, to append all of them
+    const promptMessageIndex = args.promptMessageId
+      ? contextMessages.findIndex((m) => m._id === args.promptMessageId)
+      : -1;
     const promptMessage =
-      !!args.promptMessageId &&
-      contextMessages.at(-1)?._id === args.promptMessageId
-        ? contextMessages.pop()
+      promptMessageIndex !== -1
+        ? contextMessages.splice(promptMessageIndex, 1)[0]
         : undefined;
-    if (promptMessage && args.prompt) {
-      // If they specify both a promptMessageId and a prompt, we prefer
-      // the prompt to stand in for the promptMessageId message.
-      promptMessage.message = { role: "user", content: args.prompt };
-    }
-    let messageId = promptMessage?._id;
+
+    let promptMessageId = promptMessage?._id;
     let order = promptMessage?.order;
     let stepOrder = promptMessage?.stepOrder;
-    if (
-      threadId &&
-      messages.length &&
-      storageOptions?.saveMessages !== "none" &&
-      // 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 = saveAll ? messages : messages.slice(-1);
-      const saved = await this.saveMessages(ctx, {
-        threadId,
-        userId,
-        messages: coreMessages,
-        metadata: coreMessages.map((_, i) =>
-          i === coreMessages.length - 1 ? { id: args.id } : {},
-        ),
-        failPendingSteps: true,
-      });
-      messageId = saved.lastMessageId;
-      order = saved.messages.at(-1)?.order;
-      stepOrder = saved.messages.at(-1)?.stepOrder;
+    let savedMessages = undefined;
+    let pendingMessageId = undefined;
+    if (threadId && storageOptions?.saveMessages !== "none") {
+      let saved: { messages: MessageDoc[] };
+      if (
+        messages.length + prompt.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 = [...messages, ...prompt];
+        const toSave = saveAll ? coreMessages : coreMessages.slice(-1);
+        const metadata = Array.from({ length: toSave.length }, () => ({}));
+        saved = await this.saveMessages(ctx, {
+          threadId,
+          userId,
+          messages: [...toSave, { role: "assistant", content: [] }],
+          metadata: [...metadata, { status: "pending" }],
+          failPendingSteps: true,
+          pendingMessageId: args.pendingMessageId,
+        });
+        promptMessageId = saved.messages.at(-2)!._id;
+      } else {
+        saved = await this.saveMessages(ctx, {
+          threadId,
+          userId,
+          messages: [{ role: "assistant", content: [] }],
+          metadata: [{ status: "pending" }],
+          failPendingSteps: true,
+          pendingMessageId: args.pendingMessageId,
+        });
+      }
+      pendingMessageId = saved.messages.at(-1)!._id;
+      order = saved.messages.at(-1)!.order;
+      stepOrder = saved.messages.at(-1)!.stepOrder;
+      // Don't return the pending message
+      savedMessages = saved.messages.slice(0, -1);
     }
+
     if (promptMessage?.message) {
-      // Add the message after saving the messages, so it's not saved again.
-      messages.push(deserializeMessage(promptMessage.message));
+      if (!args.prompt) {
+        // If they override the prompt, we skip the existing prompt message.
+        messages.push(deserializeMessage(promptMessage.message));
+      }
       // Lazily generate embeddings for the prompt message, if it doesn't have
       // embeddings yet. This can happen if the message was saved in a mutation
       // where the LLM is not available.
-      if (!promptMessage.embeddingId && this.options.textEmbedding) {
+      if (!promptMessage.embeddingId && this.options.textEmbeddingModel) {
         await this._generateAndSaveEmbeddings(ctx, [promptMessage]);
       }
     }
 
+    const prePrompt = contextMessages.map((m) =>
+      deserializeMessage(m.message!),
+    );
+    let existingResponses: ModelMessage[] = [];
+    if (promptMessageIndex !== -1) {
+      // pull any messages that already responded to the prompt off
+      // and add them after the prompt
+      existingResponses = prePrompt.splice(promptMessageIndex);
+    }
+
     let processedMessages = [
-      ...contextMessages.map((m) => deserializeMessage(m.message!)),
+      ...prePrompt,
       ...messages,
+      ...prompt,
+      ...existingResponses,
     ];
+    if (promptMessageIndex === -1) {
+      processedMessages.push(...prompt);
+    } else {
+      // We add the prompt where the prompt message was
+      processedMessages.splice(promptMessageIndex, 0, ...prompt);
+    }
 
     // Process messages to inline localhost files (if not, file urls pointing to localhost will be sent to LLM providers)
     if (process.env.CONVEX_CLOUD_URL?.startsWith("http://127.0.0.1")) {
-      processedMessages = await this._inlineMessagesFiles(processedMessages);
+      processedMessages = await inlineMessagesFiles(processedMessages);
     }
 
     const { prompt: _, model, ...rest } = args;
     return {
       args: {
+        ...this.options.callSettings,
         ...rest,
-        maxRetries: args.maxRetries ?? this.options.maxRetries,
-        model: model ?? this.options.chat,
+        model: model ?? this.options.languageModel,
         system: args.system ?? this.options.instructions,
         messages: processedMessages,
-      } as T & { model: LanguageModelV1 },
+      } as Extract<T, { model: LanguageModel; messages: ModelMessage[] }> &
+        CallSettings,
       userId,
-      messageId,
+      promptMessageId,
+      pendingMessageId,
+      savedMessages,
       order,
       stepOrder,
     };
@@ -1686,29 +1767,29 @@ export class Agent<
       headers?: Record<string, string>;
     },
   ): Promise<{ embeddings: number[][] }> {
-    const embeddingModel = this.options.textEmbedding;
+    const embeddingModel = this.options.textEmbeddingModel;
     assert(
       embeddingModel,
-      "a textEmbedding model is required to be set on the Agent that you're doing vector search with",
+      "a textEmbeddingModel is required to be set on the Agent that you're doing vector search with",
     );
     const result = await embedMany({
+      ...this.options.callSettings,
       model: embeddingModel,
       values: options.values,
       abortSignal: options.abortSignal,
       headers: options.headers,
-      maxRetries: this.options.maxRetries,
     });
     if (this.options.usageHandler && result.usage) {
       await this.options.usageHandler(ctx, {
         userId: options.userId,
         threadId: options.threadId,
         agentName: this.options.name,
-        model: embeddingModel.modelId,
-        provider: embeddingModel.provider,
+        model: getModelName(embeddingModel),
+        provider: getProviderName(embeddingModel),
         providerMetadata: undefined,
         usage: {
-          promptTokens: result.usage.tokens,
-          completionTokens: 0,
+          inputTokens: result.usage.tokens,
+          outputTokens: 0,
           totalTokens: result.usage.tokens,
         },
       });
@@ -1716,95 +1797,6 @@ export class Agent<
     return { embeddings: result.embeddings };
   }
 
-  /**
-   * Process messages to inline file and image URLs that point to localhost
-   * by converting them to base64. This solves the problem of LLMs not being
-   * able to access localhost URLs.
-   */
-  private async _inlineMessagesFiles(
-    messages: CoreMessage[],
-  ): Promise<CoreMessage[]> {
-    // Process each message to convert localhost URLs to base64
-    return Promise.all(
-      messages.map(async (message): Promise<CoreMessage> => {
-        if (
-          (message.role !== "user" && message.role !== "assistant") ||
-          typeof message.content === "string" ||
-          !Array.isArray(message.content)
-        ) {
-          return message;
-        }
-
-        const processedContent = await Promise.all(
-          message.content.map(async (part) => {
-            if (part.type === "image" && part.image instanceof URL) {
-              assert(
-                message.role === "user",
-                "Images can only be in user messages",
-              );
-              if (this._isLocalhostUrl(part.image)) {
-                const imageData = await this._downloadFile(part.image);
-                return {
-                  ...part,
-                  image: imageData,
-                } as ImagePart;
-              }
-            }
-
-            // Handle file parts
-            if (part.type === "file" && part.data instanceof URL) {
-              if (this._isLocalhostUrl(part.data)) {
-                const fileData = await this._downloadFile(part.data);
-                return {
-                  ...part,
-                  data: fileData,
-                } as FilePart;
-              }
-            }
-
-            return part;
-          }),
-        );
-        if (message.role === "user") {
-          return {
-            ...message,
-            content: processedContent as UserContent,
-          };
-        } else {
-          return {
-            ...message,
-            content: processedContent as AssistantContent,
-          };
-        }
-      }),
-    );
-  }
-
-  /**
-   * Check if a URL points to localhost
-   */
-  private _isLocalhostUrl(url: URL): boolean {
-    return (
-      url.hostname === "localhost" ||
-      url.hostname === "127.0.0.1" ||
-      url.hostname === "::1" ||
-      url.hostname === "0.0.0.0"
-    );
-  }
-
-  /**
-   * Download a file from a URL
-   */
-  private async _downloadFile(url: URL): Promise<ArrayBuffer> {
-    // Fetch the file
-    const response = await fetch(url);
-    if (!response.ok) {
-      throw new Error(`Failed to fetch ${url}: ${response.statusText}`);
-    }
-
-    return await response.arrayBuffer();
-  }
-
   /**
    * WORKFLOW UTILITIES
    */
@@ -1845,26 +1837,10 @@ export class Agent<
    * Create an action out of this agent so you can call it from workflows or other actions
    * without a wrapping function.
    * @param spec Configuration for the agent acting as an action, including
-   *   {@link ContextOptions}, {@link StorageOptions}, and maxSteps.
+   *   {@link ContextOptions}, {@link StorageOptions}, and {@link stopWhen}.
    */
   asTextAction<DataModel extends GenericDataModel>(
-    spec?: {
-      /**
-       * The maximum number of steps to take in this action.
-       * Defaults to the {@link Agent.maxSteps} option.
-       */
-      maxSteps?: number;
-      /**
-       * The {@link ContextOptions} to use for fetching contextual messages and
-       * saving input/output messages.
-       * Defaults to the {@link Agent.contextOptions} option.
-       */
-      contextOptions?: ContextOptions;
-      /**
-       * The {@link StorageOptions} to use for saving input/output messages.
-       * Defaults to the {@link Agent.storageOptions} option.
-       */
-      storageOptions?: StorageOptions;
+    spec: MaybeCustomCtx<CustomCtx, DataModel, AgentTools> & {
       /**
        * Whether to stream the text.
        * If false, it will generate the text in a single call. (default)
@@ -1873,49 +1849,36 @@ export class Agent<
        * defaults if you pass true.
        */
       stream?: boolean | StreamingOptions;
-    } & (CustomCtx extends Record<string, unknown>
-      ? {
-          /**
-           * If you have a custom ctx that you use with the Agent
-           * (e.g. new Agent<{ orgId: string }>(...))
-           * you need to provide this function to add any extra fields.
-           * e.g.
-           * ```ts
-           * const myAgent = new Agent<{ orgId: string }>(...);
-           * const myAction = myAgent.asTextAction({
-           *   customCtx: (ctx: ActionCtx, target, llmArgs) => {
-           *     const orgId = await lookupOrgId(ctx, target.threadId);
-           *     return { orgId };
-           *   },
-           * });
-           * ```
-           * Then, in your tools, you can
-           */
-          customCtx: (
-            ctx: GenericActionCtx<DataModel>,
-            target: {
-              userId?: string | undefined;
-              threadId?: string | undefined;
-            },
-            llmArgs: TextArgs<AgentTools>,
-          ) => CustomCtx;
-        }
-      : { customCtx?: never }),
+      /**
+       * When to stop generating text.
+       * Defaults to the {@link Agent["options"].stopWhen} option.
+       */
+      stopWhen?: StopCondition<AgentTools> | Array<StopCondition<AgentTools>>;
+    } & Options,
+    overrides?: CallSettings,
   ) {
-    const maxSteps = spec?.maxSteps ?? this.options.maxSteps;
     return internalActionGeneric({
       args: vTextArgs,
       handler: async (ctx_, args) => {
         const stream =
-          args.stream === true ? spec?.stream || true : spec?.stream ?? false;
+          args.stream === true ? spec?.stream || true : (spec?.stream ?? false);
         const targetArgs = { userId: args.userId, threadId: args.threadId };
         const llmArgs = {
-          maxSteps,
+          stopWhen: spec?.stopWhen ?? this.options.stopWhen,
+          ...overrides,
           ...omit(args, ["storageOptions", "contextOptions"]),
-        };
+          messages: args.messages?.map(deserializeMessage),
+          prompt: Array.isArray(args.prompt)
+            ? args.prompt.map(deserializeMessage)
+            : args.prompt,
+          toolChoice: args.toolChoice as ToolChoice<AgentTools>,
+        } satisfies StreamingTextArgs<AgentTools>;
+        if (args.maxSteps) {
+          llmArgs.stopWhen = stepCountIs(args.maxSteps);
+        }
         const opts = {
           ...this.options,
-          ...(spec && pick(spec, ["contextOptions", "storageOptions"])),
+          ...pick(spec, ["contextOptions", "storageOptions"]),
           ...pick(args, ["contextOptions", "storageOptions"]),
           saveStreamDeltas: stream,
         };
@@ -1925,23 +1888,37 @@ export class Agent<
             : ctx_
         ) as UserActionCtx & CustomCtx;
         if (stream) {
-          const result = await this.streamText(ctx, targetArgs, llmArgs, opts);
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          const result = await this.streamText<any>(
+            ctx,
+            targetArgs,
+            llmArgs,
+            opts,
+          );
           await result.consumeStream();
           return {
             text: await result.text,
-            messageId: result.messageId,
+            promptMessageId: result.promptMessageId,
             order: result.order,
             finishReason: await result.finishReason,
             warnings: result.warnings,
+            savedMessageIds: result.savedMessages?.map((m) => m._id) ?? [],
           };
         } else {
-          const res = await this.generateText(ctx, targetArgs, llmArgs, opts);
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          const res = await this.generateText<any>(
+            ctx,
+            targetArgs,
+            llmArgs,
+            opts,
+          );
           return {
             text: res.text,
-            messageId: res.messageId,
+            promptMessageId: res.promptMessageId,
             order: res.order,
             finishReason: res.finishReason,
             warnings: res.warnings,
+            savedMessageIds: res.savedMessages?.map((m) => m._id) ?? [],
           };
         }
       },
@@ -1952,36 +1929,47 @@ export class Agent<
    * it from workflows or other actions without a wrapping function.
    * @param spec Configuration for the agent acting as an action, including
    * the normal parameters to {@link generateObject}, plus {@link ContextOptions}
-   * and maxSteps.
+   * and stopWhen.
    */
   asObjectAction<T>(
-    spec: OurObjectArgs<T> & { maxSteps?: number },
-    options?: {
-      contextOptions?: ContextOptions;
-      storageOptions?: StorageOptions;
-    },
+    objectArgs: Omit<
+      Parameters<typeof generateObject<FlexibleSchema<T>>>[0],
+      "model"
+    >,
+    options?: Options & MaybeCustomCtx<CustomCtx, DataModel, AgentTools>,
   ) {
-    const maxSteps = spec?.maxSteps ?? this.options.maxSteps;
     return internalActionGeneric({
       args: vSafeObjectArgs,
-      handler: async (ctx, args) => {
-        const overrides = pick(args, ["userId", "threadId"]);
-        const value = await this.generateObject(
-          ctx,
-          { userId: args.userId, threadId: args.threadId },
-          {
-            ...spec,
-            maxSteps,
-            ...omit(args, ["userId", "threadId"]),
-          } as unknown as OurObjectArgs<unknown>,
-          { ...this.options, ...options, ...overrides },
-        );
+      handler: async (ctx_, args) => {
+        const { userId, threadId, callSettings, ...rest } = args;
+        const overrides = pick(rest, ["contextOptions", "storageOptions"]);
+        const targetArgs = { userId, threadId };
+        const llmArgs = {
+          ...objectArgs,
+          ...callSettings,
+          ...omit(rest, ["storageOptions", "contextOptions"]),
+          messages: args.messages?.map(deserializeMessage),
+          prompt: Array.isArray(args.prompt)
+            ? args.prompt.map(deserializeMessage)
+            : args.prompt,
+        } as Omit<Parameters<typeof generateObject>[0], "model">;
+        const ctx = (
+          options?.customCtx
+            ? { ...ctx_, ...options.customCtx(ctx_, targetArgs, llmArgs) }
+            : ctx_
+        ) as UserActionCtx & CustomCtx;
+        const value = await this.generateObject(ctx, targetArgs, llmArgs, {
+          ...this.options,
+          ...options,
+          ...overrides,
+        });
         return {
-          object: value.object as T,
-          messageId: value.messageId,
+          object: convexToJson(value.object as Value) as T,
+          promptMessageId: value.promptMessageId,
           order: value.order,
           finishReason: value.finishReason,
           warnings: value.warnings,
+          savedMessageIds: value.savedMessages?.map((m) => m._id) ?? [],
         };
       },
     });
@@ -2014,228 +2002,48 @@ export class Agent<
         userId: v.optional(v.string()),
         promptMessageId: v.optional(v.string()),
         messages: v.array(vMessageWithMetadata),
-        pending: v.optional(v.boolean()),
         failPendingSteps: v.optional(v.boolean()),
+        embeddings: v.optional(vMessageEmbeddings),
       },
       handler: async (ctx, args) => {
-        const { lastMessageId, messages } = await this.saveMessages(ctx, {
+        const { messages } = await this.saveMessages(ctx, {
           ...args,
-          messages: args.messages.map((m) => m.message),
+          messages: args.messages.map((m) => deserializeMessage(m.message)),
           metadata: args.messages.map(({ message: _, ...m }) => m),
+          skipEmbeddings: true,
         });
         return {
-          lastMessageId,
-          messageIds: messages.map((m) => m._id),
+          lastMessageId: messages.at(-1)!._id,
+          messages: messages.map((m) => pick(m, ["_id", "order", "stepOrder"])),
         };
       },
     });
   }
 }
 
-type CoreMessageMaybeWithId = CoreMessage & { id?: string | undefined };
-
-/**
- * Create a thread to store messages with an Agent.
- * @param ctx The context from a mutation or action.
- * @param component The Agent component, usually `components.agent`.
- * @param args The associated thread metadata.
- * @returns The id of the created thread.
- */
-export async function createThread(
-  ctx: RunMutationCtx,
-  component: AgentComponent,
-  args?: {
-    userId?: string | null;
-    title?: string;
-    summary?: string;
-  },
-) {
-  const { _id: threadId } = await ctx.runMutation(
-    component.threads.createThread,
-    {
-      userId: args?.userId ?? undefined,
-      title: args?.title,
-      summary: args?.summary,
-    },
-  );
-  return threadId;
-}
-
-/**
- * Get the metadata for a thread.
- * @param ctx A ctx object from a query, mutation, or action.
- * @param args.threadId The thread to get the metadata for.
- * @returns The metadata for the thread.
- */
-export async function getThreadMetadata(
-  ctx: RunQueryCtx,
-  component: AgentComponent,
-  args: { threadId: string },
-): Promise<ThreadDoc> {
-  const thread = await ctx.runQuery(component.threads.getThread, {
-    threadId: args.threadId,
-  });
-  if (!thread) {
-    throw new Error("Thread not found");
-  }
-  return thread;
-}
-
-type SaveMessagesArgs = {
-  threadId: string;
-  userId?: string | null;
-  /**
-   * The message that these messages are in response to. They will be
-   * the same "order" as this message, at increasing stepOrder(s).
-   */
-  promptMessageId?: string;
-  /**
-   * The messages to save.
-   */
-  messages: CoreMessageMaybeWithId[];
-  /**
-   * Metadata to save with the messages. Each element corresponds to the
-   * message at the same index.
-   */
-  metadata?: Omit<MessageWithMetadata, "message">[];
-  /**
-   * If false, it will "commit" the messages immediately.
-   * If true, it will mark them as pending until the final step has finished.
-   * Defaults to false.
-   */
-  pending?: boolean;
-  /**
-   * If true, it will fail any pending steps.
-   * Defaults to false.
-   */
-  failPendingSteps?: boolean;
-  /**
-   * The embeddings to save with the messages.
-   */
-  embeddings?: Omit<MessageEmbeddings, "dimension">;
-};
-
-/**
- * Explicitly save messages associated with the thread (& user if provided)
- */
-export async function saveMessages(
-  ctx: RunMutationCtx,
-  component: AgentComponent,
-  args: SaveMessagesArgs & {
-    /**
-     * The agent name to associate with the messages.
-     */
-    agentName?: string;
-  },
-) {
-  let embeddings: MessageEmbeddings | undefined;
-  if (args.embeddings) {
-    const dimension = args.embeddings.vectors.find((v) => v !== null)?.length;
-    if (dimension) {
-      validateVectorDimension(dimension);
-      embeddings = {
-        model: args.embeddings.model,
-        dimension,
-        vectors: args.embeddings.vectors,
-      };
-    }
+async function willContinue(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  steps: StepResult<any>[],
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  stopWhen: StopCondition<any> | Array<StopCondition<any>> | undefined,
+): Promise<boolean> {
+  const step = steps.at(-1)!;
+  // we aren't doing another round after a tool result
+  // TODO: whether to handle continuing after too much context used..
+  if (step.finishReason !== "tool-calls") return false;
+  // we don't have a tool result, so we'll wait for more
+  if (step.toolCalls.length > step.toolResults.length) return false;
+  if (Array.isArray(stopWhen)) {
+    return (await Promise.all(stopWhen.map(async (s) => s({ steps })))).every(
+      (stop) => !stop,
+    );
   }
-  const result = await ctx.runMutation(component.messages.addMessages, {
-    threadId: args.threadId,
-    userId: args.userId ?? undefined,
-    agentName: args.agentName,
-    promptMessageId: args.promptMessageId,
-    embeddings,
-    messages: await Promise.all(
-      args.messages.map(async (m, i) => {
-        const { message, fileIds } = await serializeMessage(ctx, component, m);
-        return {
-          ...args.metadata?.[i],
-          message,
-          fileIds,
-        } as MessageWithMetadata;
-      }),
-    ),
-    failPendingSteps: args.failPendingSteps ?? false,
-    pending: args.pending ?? false,
-  });
-  return {
-    lastMessageId: result.messages.at(-1)!._id,
-    messages: result.messages,
-  };
+  return !(await stopWhen?.({ steps }));
 }
 
-type SaveMessageArgs = {
-  threadId: string;
-  userId?: string | null;
-  /**
-   * Metadata to save with the messages. Each element corresponds to the
-   * message at the same index.
-   */
-  metadata?: Omit<MessageWithMetadata, "message">;
-  /**
-   * The embedding to save with the message.
-   */
-  embedding?: {
-    vector: number[];
-    model: string;
-  };
-} & (
-  | {
-      prompt?: undefined;
-      /**
-       * The message to save.
-       */
-      message: CoreMessage;
-    }
-  | {
-      /*
-       * The prompt to save with the message.
-       */
-      prompt: string;
-      message?: undefined;
-    }
-);
-
-/**
- * Save a message to the thread.
- * @param ctx A ctx object from a mutation or action.
- * @param args The message and what to associate it with (user / thread)
- * You can pass extra metadata alongside the message, e.g. associated fileIds.
- * @returns The messageId of the saved message.
- */
-export async function saveMessage(
-  ctx: RunMutationCtx,
-  component: AgentComponent,
-  args: SaveMessageArgs & {
-    /**
-     * The agent name to associate with the message.
-     */
-    agentName?: string;
-  },
-) {
-  let embeddings:
-    | {
-        vectors: number[][];
-        model: string;
-      }
-    | undefined;
-  if (args.embedding && args.embedding.vector) {
-    embeddings = {
-      model: args.embedding.model,
-      vectors: [args.embedding.vector],
-    };
+function errorToString(error: unknown): string {
+  if (error instanceof Error) {
+    return error.message;
   }
-  const { lastMessageId, messages } = await saveMessages(ctx, component, {
-    threadId: args.threadId,
-    userId: args.userId ?? undefined,
-    agentName: args.agentName,
-    messages:
-      args.prompt !== undefined
-        ? [{ role: "user", content: args.prompt }]
-        : [args.message],
-    metadata: args.metadata ? [args.metadata] : undefined,
-    embeddings,
-  });
-  return { messageId: lastMessageId, message: messages.at(-1)! };
+  return String(error);
 }