@convex-dev/agent 0.2.0-alpha.2 → 0.2.1

package/src/client/index.ts

@@ -1,23 +1,21 @@
-import type { LanguageModelV2 } from "@ai-sdk/provider";
 import type {
-  AssistantContent,
-  ModelMessage,
-  DeepPartial,
-  FilePart,
+  FlexibleSchema,
+  IdGenerator,
+  InferSchema,
+  ProviderOptions,
+} from "@ai-sdk/provider-utils";
+import type {
+  CallSettings,
+  EmbeddingModel,
   GenerateObjectResult,
   GenerateTextResult,
-  ImagePart,
+  LanguageModel,
+  ModelMessage,
   StepResult,
-  StreamObjectResult,
-  StreamTextResult,
-  ToolSet,
-  UserContent,
   StopCondition,
-  Schema,
+  StreamTextResult,
   ToolChoice,
-  LanguageModel,
-  EmbeddingModel,
-  CallSettings,
+  ToolSet,
 } from "ai";
 import {
   embedMany,
@@ -31,13 +29,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 {
@@ -52,16 +49,16 @@ import {
 } from "../mapping.js";
 import { extractText, isTool } from "../shared.js";
 import {
+  vMessageEmbeddings,
+  vMessageWithMetadata,
+  vSafeObjectArgs,
+  vTextArgs,
   type Message,
   type MessageStatus,
   type MessageWithMetadata,
   type ProviderMetadata,
   type StreamArgs,
   type Usage,
-  vMessageEmbeddings,
-  vMessageWithMetadata,
-  vSafeObjectArgs,
-  vTextArgs,
 } from "../validators.js";
 import { createTool, wrapTools, type ToolCtx } from "./createTool.js";
 import {
@@ -70,7 +67,6 @@ import {
   type SaveMessageArgs,
   type SaveMessagesArgs,
 } from "./messages.js";
-import { createThread, getThreadMetadata } from "./threads.js";
 import {
   fetchContextMessages,
   getModelName,
@@ -79,36 +75,42 @@ import {
 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 type z from "zod/v3";
+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,
   deserializeMessage,
+  serializeDataOrUrl,
   serializeMessage,
 } from "../mapping.js";
 // NOTE: these are also exported via @convex-dev/agent/validators
@@ -127,12 +129,12 @@ export {
   vUserMessage,
 } from "../validators.js";
 export type { ToolCtx } from "./createTool.js";
-export { getFile, storeFile } from "./files.js";
 export {
-  filterOutOrphanedToolMessages,
-  fetchContextMessages,
-} from "./search.js";
-export { abortStream, listStreams, syncStreams } from "./streaming.js";
+  definePlaygroundAPI,
+  type AgentsFn,
+  type PlaygroundAPI,
+} from "./definePlaygroundAPI.js";
+export { getFile, storeFile } from "./files.js";
 export {
   listMessages,
   saveMessage,
@@ -140,13 +142,18 @@ export {
   type SaveMessageArgs,
   type SaveMessagesArgs,
 } from "./messages.js";
-export { createThread, getThreadMetadata } from "./threads.js";
-export { createTool, extractText, isTool };
 export {
-  definePlaygroundAPI,
-  type PlaygroundAPI,
-  type AgentsFn,
-} from "./definePlaygroundAPI.js";
+  fetchContextMessages,
+  filterOutOrphanedToolMessages,
+} from "./search.js";
+export { abortStream, listStreams, syncStreams } from "./streaming.js";
+export {
+  createThread,
+  getThreadMetadata,
+  updateThreadMetadata,
+  searchThreadTitles,
+} from "./threads.js";
+export { createTool, extractText, isTool };
 export type {
   AgentComponent,
   ContextOptions,
@@ -162,6 +169,61 @@ 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 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;
+  /**
+   * Default provider options to pass 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`.
+   */
+  providerOptions?: ProviderOptions;
+  /**
+   * 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;
+};
+
 export class Agent<
   /**
    * You can require that all `ctx` args to generateText & streamText
@@ -188,31 +250,20 @@ 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.
        */
-      name?: string;
+      name: string;
       /**
        * 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, {
-       *   chat: openai.chat("gpt-4o-mini"),
-       */
-      chat: 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, {
-       *   textEmbedding: openai.embedding("text-embedding-3-small")
+       *   languageModel: openai.chat("gpt-4o-mini"),
        */
-      textEmbedding?: EmbeddingModel<string>;
+      languageModel: LanguageModel;
       /**
        * The default system prompt to put in each request.
        * Override per-prompt by passing the "system" parameter.
@@ -225,37 +276,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 when to stop. Defaults to stepCountIs(1).
        */
       stopWhen?: StopCondition<AgentTools> | Array<StopCondition<AgentTools>>;
       /**
-       * 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.
+       * @deprecated Use `languageEmbeddingModel` instead.
        */
-      rawRequestResponseHandler?: RawRequestResponseHandler;
+      chat?: LanguageModel;
     },
   ) {}
 
@@ -269,7 +298,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?: {
       /**
@@ -285,21 +314,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
@@ -309,7 +325,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?: {
       /**
@@ -325,31 +341,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>(
+  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 };
@@ -357,8 +354,6 @@ export class Agent<
     const { thread } = await this.continueThread(ctx, {
       threadId,
       userId: args?.userId,
-      usageHandler: args?.usageHandler,
-      tools: args?.tools,
     });
     return { threadId, thread };
   }
@@ -371,7 +366,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: {
       /**
@@ -383,20 +378,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,
@@ -412,7 +395,250 @@ 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>,
+    };
+  }
+
+  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 };
+    },
+    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,
+          });
+        }
+      },
     };
   }
 
@@ -424,7 +650,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.
    */
@@ -434,102 +660,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 messages = context.savedMessages ?? [];
-    const toolCtx = {
-      ...(ctx as UserActionCtx & CustomCtx),
-      userId,
-      threadId,
-      messageId,
-      agent: this,
-    } satisfies ToolCtx;
+    const { args, promptMessageId, order, ...call } = await this.start(
+      ctx,
+      generateTextArgs,
+      { ...threadOpts, ...options },
+    );
+
     type Tools = TOOLS extends undefined ? AgentTools : TOOLS;
-    const tools = wrapTools(
-      toolCtx,
-      args.tools ?? threadTools ?? this.options.tools,
-    ) as Tools;
-    const saveOutput = opts.storageOptions?.saveMessages !== "none";
+    const steps: StepResult<Tools>[] = [];
     try {
       const result = (await generateText<Tools, OUTPUT, OUTPUT_PARTIAL>({
-        // Can be overridden
-        stopWhen: this.options.stopWhen as StopCondition<Tools> | undefined,
-        ...aiArgs,
-        tools,
+        ...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) {
-            const saved = await this.saveStep(ctx, {
-              userId,
-              threadId,
-              promptMessageId: messageId,
-              model: aiArgs.model.modelId,
-              provider: aiArgs.model.provider,
-              step,
-            });
-            messages.push(...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,
-            });
-          }
-          return args.onStepFinish?.(step);
+          steps.push(step);
+          await call.save({ step }, await willContinue(steps, args.stopWhen));
+          return generateTextArgs.onStepFinish?.(step);
         },
-      })) as GenerateTextResult<Tools, OUTPUT> & GenerationOutputMetadata;
-      result.messageId = messageId;
-      result.order = order;
-      result.messages = messages;
-      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;
     }
   }
@@ -547,22 +715,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.
@@ -587,112 +749,87 @@ 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 messages = context.savedMessages ?? [];
-    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({
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      stopWhen: this.options.stopWhen as any, // Can be overridden
-      ...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);
-          messages.push(...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;
-    result.messages = messages;
-    return result;
+    >;
+    const metadata: GenerationOutputMetadata = {
+      promptMessageId,
+      order,
+      savedMessages: call.getSavedMessages(),
+      messageId: promptMessageId,
+    };
+    if (
+      (typeof options?.saveStreamDeltas === "object" &&
+        !options.saveStreamDeltas.returnImmediately) ||
+      options?.saveStreamDeltas === true
+    ) {
+      await result.consumeStream();
+    }
+    return Object.assign(result, metadata);
   }
 
   /**
@@ -702,84 +839,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 messages = context.savedMessages ?? [];
-    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) {
-        const saved = await this.saveObject(ctx, {
-          threadId,
-          promptMessageId: messageId,
-          result,
-          userId,
-          model: aiArgs.model.modelId,
-          provider: aiArgs.model.provider,
-        });
-        messages.push(...saved.messages);
-      }
-      result.messageId = messageId;
-      result.order = order;
-      result.messages = messages;
-      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;
     }
   }
@@ -791,97 +889,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 extends z.Schema | Schema>(
-    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 messages = context.savedMessages ?? [];
-    const saveOutput = opts.storageOptions?.saveMessages !== "none";
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const stream = streamObject<any>({
+    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) {
-          const saved = 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,
-              toJsonResponse: stream.toTextStreamResponse,
-            },
-            model: aiArgs.model.modelId,
-            provider: aiArgs.model.provider,
-          });
-          messages.push(...saved.messages);
-        }
-        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;
-    stream.messages = messages;
-    return stream;
+    });
+    const metadata: GenerationOutputMetadata = {
+      promptMessageId,
+      order,
+      savedMessages: getSavedMessages(),
+      messageId: promptMessageId,
+    };
+    return Object.assign(stream, metadata);
   }
 
   /**
@@ -915,6 +999,7 @@ export class Agent<
           : [args.message],
       metadata: args.metadata ? [args.metadata] : undefined,
       skipEmbeddings: args.skipEmbeddings,
+      pendingMessageId: args.pendingMessageId,
     });
     const message = messages.at(-1)!;
     return { messageId: message._id, message };
@@ -944,7 +1029,7 @@ export class Agent<
     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. " +
@@ -1050,8 +1135,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: (
@@ -1061,7 +1146,7 @@ export class Agent<
               values: [text],
             })
           ).embeddings[0],
-          embeddingModel: this.options.textEmbedding,
+          textEmbeddingModel: this.options.textEmbeddingModel,
         };
       },
     });
@@ -1116,7 +1201,7 @@ export class Agent<
     }: { userId: string | undefined; threadId: string | undefined },
     messages: (ModelMessage | Message)[],
   ) {
-    if (!this.options.textEmbedding) {
+    if (!this.options.textEmbeddingModel) {
       return undefined;
     }
     let embeddings:
@@ -1134,11 +1219,12 @@ 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,
     });
     // Then assemble the embeddings into a single array with nulls for the messages without text.
     const embeddingsOrNull = Array(messages.length).fill(null);
@@ -1148,7 +1234,7 @@ export class Agent<
     if (textEmbeddings.embeddings.length > 0) {
       const dimension = textEmbeddings.embeddings[0].length;
       validateVectorDimension(dimension);
-      const model = getModelName(this.options.textEmbedding);
+      const model = getModelName(this.options.textEmbeddingModel);
       embeddings = { vectors: embeddingsOrNull, dimension, model };
     }
     return embeddings;
@@ -1204,9 +1290,9 @@ export class Agent<
       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(
@@ -1262,13 +1348,13 @@ export class Agent<
       provider?: string;
     },
   ): Promise<{ messages: MessageDoc[] }> {
-    const messages = await serializeNewMessagesInStep(
+    const { messages } = await serializeNewMessagesInStep(
       ctx,
       this.component,
       args.step,
       {
-        provider: args.provider ?? getProviderName(this.options.chat),
-        model: args.model ?? getModelName(this.options.chat),
+        provider: args.provider ?? getProviderName(this.options.languageModel),
+        model: args.model ?? getModelName(this.options.languageModel),
       },
     );
     const embeddings = await this.generateEmbeddings(
@@ -1311,14 +1397,20 @@ export class Agent<
       this.component,
       args.result,
       {
-        model: args.model ?? getModelName(this.options.chat),
-        provider: args.provider ?? getProviderName(this.options.chat),
+        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(
       ctx,
       { userId: args.userId, threadId: args.threadId },
-      messages.map((m) => deserializeMessage(m.message)),
+      messages.map((m) => m.message),
     );
 
     return ctx.runMutation(this.component.messages.addMessages, {
@@ -1340,25 +1432,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,
+    });
   }
 
   /**
@@ -1530,12 +1614,12 @@ export class Agent<
 
   async _saveMessagesAndFetchContext<
     T extends {
-      id?: string;
       prompt?: string | (ModelMessage | Message)[];
       messages?: (ModelMessage | Message)[];
       system?: string;
       promptMessageId?: string;
-      model?: LanguageModelV2;
+      pendingMessageId?: string;
+      model?: LanguageModel;
     },
   >(
     ctx: RunActionCtx,
@@ -1545,11 +1629,16 @@ export class Agent<
       threadId,
       contextOptions,
       storageOptions,
-    }: { userId: string | undefined; threadId: string | undefined } & Options,
+    }: {
+      userId: string | null | undefined;
+      threadId: string | undefined;
+    } & Options,
   ): Promise<{
-    args: T & { model: LanguageModelV2 };
+    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;
@@ -1565,7 +1654,8 @@ export class Agent<
       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,
@@ -1584,32 +1674,46 @@ export class Agent<
         ? contextMessages.splice(promptMessageIndex, 1)[0]
         : undefined;
 
-    let messageId = promptMessage?._id;
+    let promptMessageId = promptMessage?._id;
     let order = promptMessage?.order;
     let stepOrder = promptMessage?.stepOrder;
     let savedMessages = undefined;
-    if (
-      threadId &&
-      messages.length + prompt.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 = [...messages, ...prompt];
-      const toSave = saveAll ? coreMessages : coreMessages.slice(-1);
-      const metadata = Array.from({ length: toSave.length }, () => ({}));
-      const saved = await this.saveMessages(ctx, {
-        threadId,
-        userId,
-        messages: toSave,
-        metadata,
-        failPendingSteps: true,
-      });
-      messageId = saved.messages.at(-1)!._id;
+    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;
-      savedMessages = saved.messages;
+      // Don't return the pending message
+      savedMessages = saved.messages.slice(0, -1);
     }
 
     if (promptMessage?.message) {
@@ -1620,7 +1724,7 @@ export class Agent<
       // 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]);
       }
     }
@@ -1650,20 +1754,23 @@ export class Agent<
 
     // 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,
+        ...this.options.providerOptions,
         ...rest,
-        model: model ?? this.options.chat,
+        model: model ?? this.options.languageModel,
         system: args.system ?? this.options.instructions,
         messages: processedMessages,
-      } as T & { model: LanguageModelV2 },
+      } as Extract<T, { model: LanguageModel; messages: ModelMessage[] }> &
+        CallSettings,
       userId,
-      messageId,
+      promptMessageId,
+      pendingMessageId,
       savedMessages,
       order,
       stepOrder,
@@ -1680,10 +1787,10 @@ 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,
@@ -1710,83 +1817,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: (ModelMessage | Message)[],
-  ): Promise<(ModelMessage | Message)[]> {
-    // Process each message to convert localhost URLs to base64
-    return Promise.all(
-      messages.map(async (message): Promise<ModelMessage | Message> => {
-        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
    */
@@ -1830,23 +1860,7 @@ export class Agent<
    *   {@link ContextOptions}, {@link StorageOptions}, and {@link stopWhen}.
    */
   asTextAction<DataModel extends GenericDataModel>(
-    spec?: {
-      /**
-       * When to stop generating text.
-       * Defaults to the {@link Agent["options"].stopWhen} option.
-       */
-      stopWhen?: StopCondition<AgentTools> | Array<StopCondition<AgentTools>>;
-      /**
-       * 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)
@@ -1855,36 +1869,14 @@ 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 stopWhen = spec?.stopWhen ?? this.options.stopWhen;
     return internalActionGeneric({
       args: vTextArgs,
       handler: async (ctx_, args) => {
@@ -1892,7 +1884,8 @@ export class Agent<
           args.stream === true ? spec?.stream || true : (spec?.stream ?? false);
         const targetArgs = { userId: args.userId, threadId: args.threadId };
         const llmArgs = {
-          stopWhen,
+          stopWhen: spec?.stopWhen ?? this.options.stopWhen,
+          ...overrides,
           ...omit(args, ["storageOptions", "contextOptions"]),
           messages: args.messages?.map(deserializeMessage),
           prompt: Array.isArray(args.prompt)
@@ -1905,7 +1898,7 @@ export class Agent<
         }
         const opts = {
           ...this.options,
-          ...(spec && pick(spec, ["contextOptions", "storageOptions"])),
+          ...pick(spec, ["contextOptions", "storageOptions"]),
           ...pick(args, ["contextOptions", "storageOptions"]),
           saveStreamDeltas: stream,
         };
@@ -1925,10 +1918,11 @@ export class Agent<
           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 {
           // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -1940,10 +1934,11 @@ export class Agent<
           );
           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) ?? [],
           };
         }
       },
@@ -1957,32 +1952,44 @@ export class Agent<
    * and stopWhen.
    */
   asObjectAction<T>(
-    spec: OurObjectArgs<T>,
-    options?: {
-      contextOptions?: ContextOptions;
-      storageOptions?: StorageOptions;
-    },
+    objectArgs: Omit<
+      Parameters<typeof generateObject<FlexibleSchema<T>>>[0],
+      "model"
+    >,
+    options?: Options & MaybeCustomCtx<CustomCtx, DataModel, AgentTools>,
   ) {
     return internalActionGeneric({
       args: vSafeObjectArgs,
-      handler: async (ctx, args) => {
-        const { userId, threadId, ...rest } = args;
+      handler: async (ctx_, args) => {
+        const { userId, threadId, callSettings, ...rest } = args;
         const overrides = pick(rest, ["contextOptions", "storageOptions"]);
-        const value = await this.generateObject(
-          ctx,
-          { userId, threadId },
-          {
-            ...spec,
-            ...omit(rest, ["contextOptions", "storageOptions"]),
-          } as OurObjectArgs<unknown>,
-          { ...this.options, ...options, ...overrides },
-        );
+        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) ?? [],
         };
       },
     });
@@ -2033,3 +2040,30 @@ export class Agent<
     });
   }
 }
+
+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,
+    );
+  }
+  return !(await stopWhen?.({ steps }));
+}
+
+function errorToString(error: unknown): string {
+  if (error instanceof Error) {
+    return error.message;
+  }
+  return String(error);
+}