@convex-dev/agent 0.5.0-alpha.1 → 0.6.0-alpha.0

package/src/client/index.ts

@@ -1,29 +1,32 @@
-import type { LanguageModelV2, EmbeddingModelV2 } from "@ai-sdk/provider";
+import type { JSONValue } from "@ai-sdk/provider";
 import type {
-  AssistantContent,
-  ModelMessage,
-  DeepPartial,
-  FilePart,
+  FlexibleSchema,
+  IdGenerator,
+  InferSchema,
+} from "@ai-sdk/provider-utils";
+import type {
+  CallSettings,
+  EmbeddingModel,
   GenerateObjectResult,
   GenerateTextResult,
-  ImagePart,
+  LanguageModel,
+  ModelMessage,
   StepResult,
-  StreamObjectResult,
-  StreamTextResult,
-  ToolSet,
-  UserContent,
   StopCondition,
-  Schema,
+  StreamTextResult,
   ToolChoice,
+  ToolSet,
 } from "ai";
-import {
-  embedMany,
-  generateObject,
-  generateText,
-  stepCountIs,
-  streamObject,
-  streamText,
-} from "ai";
+import { generateObject, generateText, stepCountIs, streamObject } from "ai";
+
+const MIGRATION_URL = "https://github.com/get-convex/agent/blob/main/MIGRATION.md";
+const warnedDeprecations = new Set<string>();
+function warnDeprecation(key: string, message: string) {
+  if (!warnedDeprecations.has(key)) {
+    warnedDeprecations.add(key);
+    console.warn(`[@convex-dev/agent] ${message}\n  See: ${MIGRATION_URL}`);
+  }
+}
 import { assert, omit, pick } from "convex-helpers";
 import {
   internalActionGeneric,
@@ -34,110 +37,159 @@ import {
   type PaginationResult,
   type WithoutSystemFields,
 } from "convex/server";
-import { v } from "convex/values";
-import type { MessageDoc, ThreadDoc } from "../component/schema.js";
+import { convexToJson, v, type Value } from "convex/values";
 import type { threadFieldsSupportingPatch } from "../component/threads.js";
+import { type VectorDimension } from "../component/vector/tables.js";
 import {
-  validateVectorDimension,
-  type VectorDimension,
-} from "../component/vector/tables.js";
-import {
-  deserializeMessage,
+  toModelMessage,
   serializeMessage,
   serializeNewMessagesInStep,
   serializeObjectResult,
 } from "../mapping.js";
-import { extractText, isTool } from "../shared.js";
+import { getModelName, getProviderName } from "../shared.js";
 import {
+  vMessageEmbeddings,
+  vMessageWithMetadata,
+  vSafeObjectArgs,
+  vTextArgs,
   type Message,
-  type MessageEmbeddings,
+  type MessageDoc,
   type MessageStatus,
   type MessageWithMetadata,
   type ProviderMetadata,
   type StreamArgs,
-  type Usage,
-  vMessageWithMetadata,
-  vSafeObjectArgs,
-  vTextArgs,
+  type ThreadDoc,
 } from "../validators.js";
-import { createTool, wrapTools, type ToolCtx } from "./createTool.js";
-import { listMessages } from "./listMessages.js";
-import { fetchContextMessages } from "./search.js";
 import {
-  DeltaStreamer,
-  mergeTransforms,
-  type StreamingOptions,
-  syncStreams,
-} from "./streaming.js";
+  listMessages,
+  saveMessages,
+  type SaveMessageArgs,
+  type SaveMessagesArgs,
+} from "./messages.js";
+import {
+  embedMany,
+  embedMessages,
+  fetchContextMessages,
+  generateAndSaveEmbeddings,
+} from "./search.js";
+import { startGeneration } from "./start.js";
+import { syncStreams, type StreamingOptions } from "./streaming.js";
+import { createThread, getThreadMetadata } from "./threads.js";
 import type {
   ActionCtx,
   AgentComponent,
+  Config,
   ContextOptions,
+  GenerateObjectArgs,
   GenerationOutputMetadata,
+  MaybeCustomCtx,
+  ObjectMode,
   Options,
-  OurObjectArgs,
-  OurStreamObjectArgs,
   RawRequestResponseHandler,
-  RunActionCtx,
-  RunMutationCtx,
-  RunQueryCtx,
+  MutationCtx,
   StorageOptions,
   StreamingTextArgs,
+  StreamObjectArgs,
   SyncStreamsReturnValue,
   TextArgs,
   Thread,
   UsageHandler,
-  UserActionCtx,
+  QueryCtx,
+  AgentPrompt,
+  Output,
 } from "./types.js";
-import type z from "zod";
+import { streamText } from "./streamText.js";
+import { errorToString, willContinue } from "./utils.js";
 
 export { stepCountIs } from "ai";
-export { vMessageDoc, vThreadDoc } from "../component/schema.js";
 export {
+  docsToModelMessages,
+  toModelMessage,
+  //** @deprecated use toModelMessage instead */
+  toModelMessage as deserializeMessage,
+  guessMimeType,
   serializeDataOrUrl,
-  deserializeMessage,
   serializeMessage,
+  toUIFilePart,
 } 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 { extractText, isTool, sorted } from "../shared.js";
 export {
   vAssistantMessage,
+  vContent,
   vContextOptions,
   vMessage,
+  vMessageDoc,
   vPaginationResult,
   vProviderMetadata,
+  vSource,
   vStorageOptions,
   vStreamArgs,
   vSystemMessage,
+  vThreadDoc,
   vToolMessage,
   vUsage,
   vUserMessage,
+  type Message,
+  type MessageDoc,
+  type SourcePart,
+  type ThreadDoc,
+  type Usage,
 } from "../validators.js";
-export type { ToolCtx } from "./createTool.js";
+export { createTool, type ToolCtx } from "./createTool.js";
+export {
+  definePlaygroundAPI,
+  type AgentsFn,
+  type PlaygroundAPI,
+} from "./definePlaygroundAPI.js";
 export { getFile, storeFile } from "./files.js";
 export {
-  filterOutOrphanedToolMessages,
+  listMessages,
+  listUIMessages,
+  saveMessage,
+  saveMessages,
+  type SaveMessageArgs,
+  type SaveMessagesArgs,
+} from "./messages.js";
+export { mockModel } from "./mockModel.js";
+export {
   fetchContextMessages,
+  filterOutOrphanedToolMessages,
+  fetchContextWithPrompt,
+  generateAndSaveEmbeddings,
+  embedMessages,
+  embedMany,
 } from "./search.js";
-export { abortStream, listStreams, syncStreams } from "./streaming.js";
-export { createTool, extractText, isTool, listMessages };
+export { startGeneration } from "./start.js";
 export {
-  definePlaygroundAPI,
-  type PlaygroundAPI,
-  type AgentsFn,
-} from "./definePlaygroundAPI.js";
+  DEFAULT_STREAMING_OPTIONS,
+  DeltaStreamer,
+  abortStream,
+  compressUIMessageChunks,
+  listStreams,
+  syncStreams,
+  vStreamMessagesReturnValue,
+} from "./streaming.js";
+export {
+  createThread,
+  getThreadMetadata,
+  searchThreadTitles,
+  updateThreadMetadata,
+} from "./threads.js";
+export type { ContextHandler } from "./types.js";
+export { toUIMessages, fromUIMessages, type UIMessage } from "../UIMessages.js";
+
 export type {
   AgentComponent,
+  Config,
   ContextOptions,
-  MessageDoc,
   ProviderMetadata,
   RawRequestResponseHandler,
   StorageOptions,
   StreamArgs,
   SyncStreamsReturnValue,
   Thread,
-  ThreadDoc,
-  Usage,
   UsageHandler,
 };
 
@@ -162,36 +214,24 @@ export class Agent<
    * });
    */
   CustomCtx extends object = object,
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
   AgentTools extends ToolSet = any,
 > {
   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"),
+       *   languageModel: openai.chat("gpt-4o-mini"),
        */
-      chat: LanguageModelV2;
-      /**
-       * 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")
-       */
-      textEmbedding?: EmbeddingModelV2<string>;
+      languageModel: LanguageModel;
       /**
        * The default system prompt to put in each request.
        * Override per-prompt by passing the "system" parameter.
@@ -204,38 +244,30 @@ 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 maximum number of calls to make to an LLM in case it fails.
-       * This can be overridden at each generate/stream callsite.
-       */
-      maxRetries?: number;
-      /**
-       * The usage handler to use for this agent.
+       * determines when to stop. Defaults to the AI SDK default.
        */
-      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;
+      stopWhen?:
+        | StopCondition<NoInfer<AgentTools>>
+        | Array<StopCondition<NoInfer<AgentTools>>>;
     },
-  ) {}
+  ) {
+    if (this.options.textEmbeddingModel && !this.options.embeddingModel) {
+      warnDeprecation(
+        "textEmbeddingModel",
+        "textEmbeddingModel is deprecated. Use embeddingModel instead.",
+      );
+    }
+  }
+
+  /**
+   * Get the embedding model, prioritizing embeddingModel over textEmbeddingModel.
+   * @private
+   */
+  private getEmbeddingModel(): EmbeddingModel | undefined {
+    return this.options.embeddingModel ?? this.options.textEmbeddingModel;
+  }
 
   /**
    * Start a new thread with the agent. This will have a fresh history, though if
@@ -247,8 +279,8 @@ 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>(
-    ctx: RunActionCtx & CustomCtx,
+  async createThread(
+    ctx: ActionCtx & CustomCtx,
     args?: {
       /**
        * The userId to associate with the thread. If not provided, the thread will be
@@ -263,21 +295,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
@@ -287,8 +306,8 @@ export class Agent<
    * @param args The thread metadata.
    * @returns The threadId of the new thread.
    */
-  async createThread<ThreadTools extends ToolSet | undefined = undefined>(
-    ctx: RunMutationCtx,
+  async createThread(
+    ctx: MutationCtx,
     args?: {
       /**
        * The userId to associate with the thread. If not provided, the thread will be
@@ -303,33 +322,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>(
-    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>;
-  }> {
+  ): Promise<{ threadId: string }>;
+  async createThread(
+    ctx: (ActionCtx & CustomCtx) | MutationCtx,
+    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 };
@@ -337,13 +335,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 };
   }
 
   /**
@@ -354,7 +347,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: {
       /**
@@ -366,20 +359,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,
@@ -395,32 +376,81 @@ 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 extends {
+      _internal?: { generateId?: IdGenerator };
     },
-  ): Promise<ThreadDoc[]> {
-    return ctx.runQuery(this.component.threads.searchThreadTitles, {
-      userId,
-      query,
-      limit: limit ?? 10,
-    });
+  >(
+    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 &
+      AgentPrompt & {
+        /**
+         * The tools to use for the tool calls. This will override tools specified
+         * in the Agent constructor or createThread / continueThread.
+         */
+        tools?: TOOLS;
+        /**
+         * 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;
+        stopWhen?:
+          | StopCondition<TOOLS extends undefined ? AgentTools : TOOLS>
+          | Array<StopCondition<TOOLS extends undefined ? AgentTools : TOOLS>>;
+      },
+    options?: Options & { userId?: string | null; threadId?: string },
+  ): Promise<{
+    args: T & {
+      system?: string;
+      model: LanguageModel;
+      prompt?: never;
+      messages: ModelMessage[];
+      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[];
+  }> {
+    type Tools = TOOLS extends undefined ? AgentTools : TOOLS;
+    return startGeneration<T, Tools, CustomCtx>(
+      ctx,
+      this.component,
+      {
+        ...args,
+        tools: (args.tools ?? this.options.tools) as Tools,
+        system: args.system ?? this.options.instructions,
+        stopWhen: (args.stopWhen ?? this.options.stopWhen) as any,
+      },
+      {
+        ...this.options,
+        ...options,
+        agentName: this.options.name,
+        agentForToolCtx: this,
+      },
+    );
   }
 
   /**
@@ -430,110 +460,60 @@ export class Agent<
    * Use {@link continueThread} to get a version of this function already scoped
    * 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
-   * for the {@link ContextOptions} and {@link StorageOptions}.
+   * @param scope: The user and thread to associate the message with
+   * @param generateTextArgs The arguments to the generateText function, along
+   * with {@link AgentPrompt} options, such as promptMessageId.
+   * @param options Extra controls for the {@link ContextOptions} and {@link StorageOptions}.
    * @returns The result of the generateText function.
    */
   async generateText<
     TOOLS extends ToolSet | undefined = undefined,
-    OUTPUT = never,
-    OUTPUT_PARTIAL = never,
+    OUTPUT extends Output<any, any, any> = never,
   >(
     ctx: ActionCtx & CustomCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      usageHandler,
-      tools: threadTools,
-    }: {
-      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 },
+    /**
+     * The arguments to the generateText function, similar to the ai sdk's
+     * {@link generateText} function, along with Agent prompt options.
+     */
+    generateTextArgs: AgentPrompt &
+      TextArgs<AgentTools, TOOLS, OUTPUT>,
     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,
-    } 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,
+      const result = (await generateText<Tools, OUTPUT>({
+        ...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, 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;
     }
   }
@@ -547,26 +527,16 @@ export class Agent<
    */
   async streamText<
     TOOLS extends ToolSet | undefined = undefined,
-    OUTPUT = never,
-    PARTIAL_OUTPUT = never,
+    OUTPUT extends Output<any, any, any> = never,
   >(
     ctx: ActionCtx & CustomCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      usageHandler,
-      /** Note: to get better type inference, pass tools in the next arg */
-      tools: threadTools,
-    }: {
-      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.
+     * The arguments to the streamText function, similar to the ai sdk's
+     * {@link streamText} function, along with Agent prompt options.
      */
-    args: StreamingTextArgs<AgentTools, TOOLS, OUTPUT, PARTIAL_OUTPUT>,
+    streamTextArgs: AgentPrompt &
+      StreamingTextArgs<AgentTools, TOOLS, OUTPUT>,
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
@@ -587,113 +557,29 @@ export class Agent<
   ): Promise<
     StreamTextResult<
       TOOLS extends undefined ? AgentTools : TOOLS,
-      PARTIAL_OUTPUT
+      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, 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 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,
-            order,
-            stepOrder,
-            abortSignal: aiArgs.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,
-      experimental_transform: mergeTransforms(
-        options?.saveStreamDeltas,
-        args.experimental_transform,
-      ),
-      onChunk: async (event) => {
-        await streamer?.addParts([event.chunk]);
-        // console.log("onChunk", chunk);
-        return args.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,
-          });
-        }
-        // TODO: update the streamer to error state
-        return args.onError?.(error);
+    type Tools = TOOLS extends undefined ? AgentTools : TOOLS;
+    return streamText<Tools, OUTPUT>(
+      ctx,
+      this.component,
+      {
+        ...streamTextArgs,
+        model: streamTextArgs.model ?? this.options.languageModel,
+        tools: (streamTextArgs.tools ?? this.options.tools) as Tools,
+        system: streamTextArgs.system ?? this.options.instructions,
+        stopWhen: (streamTextArgs.stopWhen ?? this.options.stopWhen) as any,
       },
-      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,
-          });
-        }
-        return args.onStepFinish?.(step);
+      {
+        ...threadOpts,
+        ...this.options,
+        agentName: this.options.name,
+        agentForToolCtx: this,
+        ...options,
       },
-    }) as StreamTextResult<
-      TOOLS extends undefined ? AgentTools : TOOLS,
-      PARTIAL_OUTPUT
-    > &
-      GenerationOutputMetadata;
-    result.messageId = messageId;
-    result.order = order;
-    return result;
+    );
   }
 
   /**
@@ -703,81 +589,47 @@ 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 FlexibleSchema<unknown> = FlexibleSchema<JSONValue>,
+    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.
+     * The arguments to the generateObject function, similar to the ai sdk's
+     * {@link generateObject} function, along with Agent prompt options.
      */
-    args: OurObjectArgs<T>,
+    generateObjectArgs: AgentPrompt &
+      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;
     }
   }
@@ -789,94 +641,67 @@ 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 FlexibleSchema<unknown> = FlexibleSchema<JSONValue>,
+    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.
+     * The arguments to the streamObject function, similar to the ai sdk's
+     * {@link streamObject} function, along with Agent prompt options.
      */
-    args: OurStreamObjectArgs<T>,
+    streamObjectArgs: AgentPrompt & StreamObjectArgs<SCHEMA, OUTPUT, RESULT>,
     /**
      * 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";
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const stream = streamObject<any>({
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      ...(aiArgs as any),
+    const { args, promptMessageId, order, fail, save, getSavedMessages } =
+      await this.start(ctx, streamObjectArgs, { ...threadOpts, ...options });
+
+    const stream = streamObject<SCHEMA, OUTPUT, RESULT>({
+      ...(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,
-              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,
+            reasoning: undefined,
+          },
+        });
+        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);
   }
 
   /**
@@ -887,7 +712,7 @@ export class Agent<
    * @returns The messageId of the saved message.
    */
   async saveMessage(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: SaveMessageArgs & {
       /**
        * If true, it will not generate embeddings for the message.
@@ -898,14 +723,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 +735,11 @@ export class Agent<
           : [args.message],
       metadata: args.metadata ? [args.metadata] : undefined,
       skipEmbeddings: args.skipEmbeddings,
+      promptMessageId: args.promptMessageId,
+      pendingMessageId: args.pendingMessageId,
     });
-    return { messageId: lastMessageId, message: messages.at(-1)! };
+    const message = messages.at(-1)!;
+    return { messageId: message._id, message };
   }
 
   /**
@@ -926,7 +751,7 @@ export class Agent<
    * @returns
    */
   async saveMessages(
-    ctx: RunMutationCtx | RunActionCtx,
+    ctx: MutationCtx | ActionCtx,
     args: SaveMessagesArgs & {
       /**
        * Skip generating embeddings for the messages. Useful if you're
@@ -936,20 +761,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.getEmbeddingModel()) {
       if (!("runAction" in ctx)) {
         console.warn(
           "You're trying to save messages and generate embeddings, but you're in a mutation. " +
@@ -967,10 +784,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,
         );
       }
@@ -993,7 +807,7 @@ export class Agent<
    * @returns The MessageDoc's in a format compatible with usePaginatedQuery.
    */
   async listMessages(
-    ctx: RunQueryCtx,
+    ctx: QueryCtx | MutationCtx | ActionCtx,
     args: {
       threadId: string;
       paginationOpts: PaginationOptions;
@@ -1013,7 +827,7 @@ export class Agent<
    * @returns The deltas for each stream from their existing cursor.
    */
   async syncStreams(
-    ctx: RunQueryCtx,
+    ctx: QueryCtx | MutationCtx | ActionCtx,
     args: {
       threadId: string;
       streamArgs: StreamArgs | undefined;
@@ -1033,15 +847,26 @@ export class Agent<
    * @returns
    */
   async fetchContextMessages(
-    ctx: RunQueryCtx | RunActionCtx,
+    ctx: QueryCtx | MutationCtx | ActionCtx,
     args: {
       userId: string | undefined;
       threadId: string | undefined;
-      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
-       * limited, as it's post-filtering the results.
+       * If targetMessageId is not provided, this text will be used
+       * for text and vector search
+       */
+      searchText?: string;
+      /**
+       * If provided, it will use this message for text/vector search (if enabled)
+       * and will only fetch messages up to (and including) this message's "order"
+       */
+      targetMessageId?: string;
+      /**
+       * @deprecated use searchText and targetMessageId instead
+       */
+      messages?: (ModelMessage | Message)[];
+      /**
+       * @deprecated use targetMessageId instead
        */
       upToAndIncludingMessageId?: string;
       contextOptions: ContextOptions | undefined;
@@ -1057,19 +882,22 @@ export class Agent<
       contextOptions,
       getEmbedding: async (text) => {
         assert("runAction" in ctx);
+        const embeddingModel = this.getEmbeddingModel();
         assert(
-          this.options.textEmbedding,
-          "A textEmbedding model is required to be set on the Agent that you're doing vector search with",
+          embeddingModel,
+          "An embeddingModel (or textEmbeddingModel) is required to be set on the Agent that you're doing vector search with",
         );
         return {
           embedding: (
-            await this.doEmbed(ctx, {
+            await embedMany(ctx, {
+              ...this.options,
+              agentName: this.options.name,
               userId: args.userId,
               threadId: args.threadId,
               values: [text],
             })
           ).embeddings[0],
-          embeddingModel: this.options.textEmbedding.modelId,
+          embeddingModel: embeddingModel,
         };
       },
     });
@@ -1082,7 +910,7 @@ export class Agent<
    * @returns The metadata for the thread.
    */
   async getThreadMetadata(
-    ctx: RunQueryCtx,
+    ctx: QueryCtx | MutationCtx | ActionCtx,
     args: { threadId: string },
   ): Promise<ThreadDoc> {
     return getThreadMetadata(ctx, this.component, args);
@@ -1096,7 +924,7 @@ export class Agent<
    * @returns The updated thread metadata.
    */
   async updateThreadMetadata(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: {
       threadId: string;
       patch: Partial<
@@ -1117,55 +945,22 @@ export class Agent<
    * @returns The embeddings for the messages.
    */
   async generateEmbeddings(
-    ctx: RunActionCtx,
-    {
-      userId,
-      threadId,
-    }: {
-      userId: string | undefined;
-      threadId: string | undefined;
-    },
+    ctx: ActionCtx,
+    args: { userId: string | undefined; threadId: string | undefined },
     messages: (ModelMessage | Message)[],
-  ) {
-    if (!this.options.textEmbedding) {
-      return undefined;
-    }
-    let embeddings:
-      | {
-          vectors: (number[] | null)[];
-          dimension: VectorDimension;
-          model: string;
-        }
-      | undefined;
-    const messageTexts = messages.map((m) => !isTool(m) && extractText(m));
-    // Find the indexes of the messages that have text.
-    const textIndexes = messageTexts
-      .map((t, i) => (t ? i : undefined))
-      .filter((i) => i !== undefined);
-    if (textIndexes.length === 0) {
-      return undefined;
-    }
-    // Then embed those messages.
-    const textEmbeddings = await this.doEmbed(ctx, {
-      userId,
-      threadId,
-      values: messageTexts as string[],
-    });
-    // 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) => {
-      embeddingsOrNull[i] = textEmbeddings.embeddings[j];
-    });
-    if (textEmbeddings.embeddings.length > 0) {
-      const dimension = textEmbeddings.embeddings[0].length;
-      validateVectorDimension(dimension);
-      embeddings = {
-        vectors: embeddingsOrNull,
-        dimension,
-        model: this.options.textEmbedding.modelId,
-      };
-    }
-    return embeddings;
+  ): Promise<
+    | {
+        vectors: (number[] | null)[];
+        dimension: VectorDimension;
+        model: string;
+      }
+    | undefined
+  > {
+    return embedMessages(
+      ctx,
+      { ...args, ...this.options, agentName: this.options.name },
+      messages,
+    );
   }
 
   /**
@@ -1176,10 +971,8 @@ export class Agent<
    * @param args The messageIds to generate embeddings for.
    */
   async generateAndSaveEmbeddings(
-    ctx: RunActionCtx,
-    args: {
-      messageIds: string[];
-    },
+    ctx: ActionCtx,
+    args: { messageIds: string[] },
   ) {
     const messages = (
       await ctx.runQuery(this.component.messages.getMessagesByIds, {
@@ -1194,10 +987,6 @@ export class Agent<
             .join(", "),
       );
     }
-    await this._generateAndSaveEmbeddings(ctx, messages);
-  }
-
-  async _generateAndSaveEmbeddings(ctx: RunActionCtx, messages: MessageDoc[]) {
     if (messages.some((m) => !m.message)) {
       throw new Error(
         "Some messages don't have a message: " +
@@ -1207,45 +996,24 @@ export class Agent<
             .join(", "),
       );
     }
-    const messagesMissingEmbeddings = messages.filter((m) => !m.embeddingId);
-    if (messagesMissingEmbeddings.length === 0) {
-      return;
+    const embeddingModel = this.getEmbeddingModel();
+    if (!embeddingModel) {
+      throw new Error(
+        "No embeddings were generated for the messages. You must pass an embeddingModel (or textEmbeddingModel) to the agent constructor.",
+      );
     }
-    const embeddings = await this.generateEmbeddings(
+    await generateAndSaveEmbeddings(
       ctx,
+      this.component,
       {
-        userId: messagesMissingEmbeddings[0]!.userId,
-        threadId: messagesMissingEmbeddings[0]!.threadId,
+        ...this.options,
+        agentName: this.options.name,
+        threadId: messages[0].threadId,
+        userId: messages[0].userId,
+        embeddingModel,
       },
-      messagesMissingEmbeddings.map((m) => deserializeMessage(m!.message!)),
+      messages,
     );
-    if (!embeddings) {
-      if (!this.options.textEmbedding) {
-        throw new Error(
-          "No embeddings were generated for the messages. You must pass a textEmbedding model to the agent constructor.",
-        );
-      }
-      throw new Error(
-        "No embeddings were generated for these messages: " +
-          messagesMissingEmbeddings.map((m) => m!._id).join(", "),
-      );
-    }
-    await ctx.runMutation(this.component.vector.index.insertBatch, {
-      vectorDimension: embeddings.dimension,
-      vectors: messagesMissingEmbeddings
-        .map((m, i) => ({
-          messageId: m!._id,
-          model: embeddings.model,
-          table: "messages",
-          userId: m.userId,
-          threadId: m.threadId,
-          vector: embeddings.vectors[i],
-        }))
-        .filter(
-          (v): v is Extract<typeof v, { vector: number[] }> =>
-            v.vector !== null,
-        ),
-    });
   }
 
   /**
@@ -1277,22 +1045,22 @@ 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(
       ctx,
       { userId: args.userId, threadId: args.threadId },
-      messages.map((m) => deserializeMessage(m.message)),
+      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,
@@ -1301,7 +1069,6 @@ export class Agent<
       embeddings,
       failPendingSteps: false,
     });
-    return saved;
   }
 
   /**
@@ -1322,23 +1089,29 @@ 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(
       ctx,
       { userId: args.userId, threadId: args.threadId },
-      messages.map((m) => deserializeMessage(m.message)),
+      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,
@@ -1346,7 +1119,6 @@ export class Agent<
       messages,
       embeddings,
       agentName: this.options.name,
-      pending: false,
     });
   }
 
@@ -1358,25 +1130,17 @@ export class Agent<
    * @param args What message to save. Generally the parent message sent into
    *   the generateText call.
    */
-  async completeMessage(
-    ctx: RunMutationCtx,
+  async finalizeMessage(
+    ctx: MutationCtx | ActionCtx,
     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,
+    });
   }
 
   /**
@@ -1385,13 +1149,13 @@ export class Agent<
    * @param args The message fields to update.
    */
   async updateMessage(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: {
       /** The id of the message to update. */
       messageId: string;
       patch: {
         /** The message to replace the existing message. */
-        message: (ModelMessage & { id?: string }) | Message;
+        message: ModelMessage | Message;
         /** The status to set on the message. */
         status: "success" | "error";
         /** The error message to set on the message. */
@@ -1435,10 +1199,8 @@ export class Agent<
    * @param args The ids of the messages to delete.
    */
   async deleteMessages(
-    ctx: RunMutationCtx,
-    args: {
-      messageIds: string[];
-    },
+    ctx: MutationCtx | ActionCtx,
+    args: { messageIds: string[] },
   ): Promise<void> {
     await ctx.runMutation(this.component.messages.deleteByIds, args);
   }
@@ -1450,10 +1212,8 @@ export class Agent<
    * @param args The id of the message to delete.
    */
   async deleteMessage(
-    ctx: RunMutationCtx,
-    args: {
-      messageId: string;
-    },
+    ctx: MutationCtx | ActionCtx,
+    args: { messageId: string },
   ): Promise<void> {
     await ctx.runMutation(this.component.messages.deleteByIds, {
       messageIds: [args.messageId],
@@ -1498,7 +1258,7 @@ export class Agent<
    * @param args The range of messages to delete.
    */
   async deleteMessageRange(
-    ctx: RunMutationCtx,
+    ctx: MutationCtx | ActionCtx,
     args: {
       threadId: string;
       startOrder: number;
@@ -1524,11 +1284,8 @@ export class Agent<
    * @param args The id of the thread to delete and optionally the page size to use for the delete.
    */
   async deleteThreadAsync(
-    ctx: RunMutationCtx,
-    args: {
-      threadId: string;
-      pageSize?: number;
-    },
+    ctx: MutationCtx | ActionCtx,
+    args: { threadId: string; pageSize?: number },
   ): Promise<void> {
     await ctx.runMutation(this.component.threads.deleteAllForThreadIdAsync, {
       threadId: args.threadId,
@@ -1544,11 +1301,8 @@ export class Agent<
    * @param args The id of the thread to delete and optionally the page size to use for the delete.
    */
   async deleteThreadSync(
-    ctx: RunActionCtx,
-    args: {
-      threadId: string;
-      pageSize?: number;
-    },
+    ctx: ActionCtx,
+    args: { threadId: string; pageSize?: number },
   ): Promise<void> {
     await ctx.runAction(this.component.threads.deleteAllForThreadIdSync, {
       threadId: args.threadId,
@@ -1556,278 +1310,6 @@ export class Agent<
     });
   }
 
-  async _saveMessagesAndFetchContext<
-    T extends {
-      id?: string;
-      prompt?: string | (ModelMessage | Message)[];
-      messages?: (ModelMessage | Message)[];
-      system?: string;
-      promptMessageId?: string;
-      model?: LanguageModelV2;
-      maxRetries?: number;
-    },
-  >(
-    ctx: RunActionCtx,
-    args: T,
-    {
-      userId: argsUserId,
-      threadId,
-      contextOptions,
-      storageOptions,
-    }: {
-      userId: string | undefined;
-      threadId: string | undefined;
-    } & Options,
-  ): Promise<{
-    args: T & { model: LanguageModelV2 };
-    userId: string | undefined;
-    messageId: string | undefined;
-    order: number | undefined;
-    stepOrder: number | undefined;
-  }> {
-    // If only a promptMessageId is provided, this will be empty.
-    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);
-    // If only a messageId is provided, this will add that message to the end.
-    const contextMessages = await this.fetchContextMessages(ctx, {
-      userId,
-      threadId,
-      upToAndIncludingMessageId: args.promptMessageId,
-      messages,
-      contextOptions,
-    });
-    // If it was a promptMessageId, pop it off context messages
-    // and add to the end of messages.
-    const promptMessageIndex = args.promptMessageId
-      ? contextMessages.findIndex((m) => m._id === args.promptMessageId)
-      : -1;
-    const promptMessage =
-      promptMessageIndex !== -1
-        ? contextMessages.splice(promptMessageIndex, 1)[0]
-        : undefined;
-
-    let messageId = promptMessage?._id;
-    let order = promptMessage?.order;
-    let stepOrder = promptMessage?.stepOrder;
-    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 }, () => ({}));
-      metadata[metadata.length - 1] = { id: args.id };
-      const saved = await this.saveMessages(ctx, {
-        threadId,
-        userId,
-        messages: toSave,
-        metadata,
-        failPendingSteps: true,
-      });
-      messageId = saved.lastMessageId;
-      order = saved.messages.at(-1)?.order;
-      stepOrder = saved.messages.at(-1)?.stepOrder;
-    }
-
-    if (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) {
-        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 = [
-      ...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);
-    }
-
-    const { prompt: _, model, ...rest } = args;
-    return {
-      args: {
-        ...rest,
-        maxRetries: args.maxRetries ?? this.options.maxRetries,
-        model: model ?? this.options.chat,
-        system: args.system ?? this.options.instructions,
-        messages: processedMessages,
-      } as T & { model: LanguageModelV2 },
-      userId,
-      messageId,
-      order,
-      stepOrder,
-    };
-  }
-
-  async doEmbed(
-    ctx: RunActionCtx,
-    options: {
-      userId: string | undefined;
-      threadId: string | undefined;
-      values: string[];
-      abortSignal?: AbortSignal;
-      headers?: Record<string, string>;
-    },
-  ): Promise<{ embeddings: number[][] }> {
-    const embeddingModel = this.options.textEmbedding;
-    assert(
-      embeddingModel,
-      "a textEmbedding model is required to be set on the Agent that you're doing vector search with",
-    );
-    const result = await embedMany({
-      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,
-        providerMetadata: undefined,
-        usage: {
-          inputTokens: result.usage.tokens,
-          outputTokens: 0,
-          totalTokens: result.usage.tokens,
-        },
-      });
-    }
-    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
    */
@@ -1871,23 +1353,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)
@@ -1896,57 +1362,34 @@ 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) => {
         const stream =
-          args.stream === true ? spec?.stream || true : spec?.stream ?? false;
-        const targetArgs = { userId: args.userId, threadId: args.threadId };
+          args.stream === true ? spec?.stream || true : (spec?.stream ?? false);
+        const { userId, threadId, prompt, messages, maxSteps, ...rest } = args;
+        const targetArgs = { userId, threadId };
         const llmArgs = {
-          stopWhen,
-          ...omit(args, ["storageOptions", "contextOptions"]),
-          messages: args.messages?.map(deserializeMessage),
-          prompt: Array.isArray(args.prompt)
-            ? args.prompt.map(deserializeMessage)
-            : args.prompt,
+          stopWhen: spec?.stopWhen,
+          ...overrides,
+          ...omit(rest, ["storageOptions", "contextOptions", "stream"]),
+          messages: messages?.map(toModelMessage),
+          prompt: Array.isArray(prompt) ? prompt.map(toModelMessage) : prompt,
           toolChoice: args.toolChoice as ToolChoice<AgentTools>,
         } satisfies StreamingTextArgs<AgentTools>;
-        if (args.maxSteps) {
-          llmArgs.stopWhen = stepCountIs(args.maxSteps);
+        if (maxSteps) {
+          llmArgs.stopWhen = stepCountIs(maxSteps);
         }
         const opts = {
-          ...this.options,
-          ...(spec && pick(spec, ["contextOptions", "storageOptions"])),
+          ...pick(spec, ["contextOptions", "storageOptions"]),
           ...pick(args, ["contextOptions", "storageOptions"]),
           saveStreamDeltas: stream,
         };
@@ -1954,9 +1397,8 @@ export class Agent<
           spec?.customCtx
             ? { ...ctx_, ...spec.customCtx(ctx_, targetArgs, llmArgs) }
             : ctx_
-        ) as UserActionCtx & CustomCtx;
+        ) as GenericActionCtx<GenericDataModel> & CustomCtx;
         if (stream) {
-          // eslint-disable-next-line @typescript-eslint/no-explicit-any
           const result = await this.streamText<any>(
             ctx,
             targetArgs,
@@ -1966,13 +1408,13 @@ 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,
+            warnings: await result.warnings,
+            savedMessageIds: result.savedMessages?.map((m) => m._id) ?? [],
           };
         } else {
-          // eslint-disable-next-line @typescript-eslint/no-explicit-any
           const res = await this.generateText<any>(
             ctx,
             targetArgs,
@@ -1981,10 +1423,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) ?? [],
           };
         }
       },
@@ -1997,57 +1440,49 @@ export class Agent<
    * the normal parameters to {@link generateObject}, plus {@link ContextOptions}
    * and stopWhen.
    */
-  asObjectAction<T>(
-    spec: OurObjectArgs<T>,
-    options?: {
-      contextOptions?: ContextOptions;
-      storageOptions?: StorageOptions;
-    },
+  asObjectAction<T, DataModel extends GenericDataModel>(
+    objectArgs: GenerateObjectArgs<FlexibleSchema<T>> & Partial<AgentPrompt>,
+    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(toModelMessage),
+          prompt: Array.isArray(args.prompt)
+            ? args.prompt.map(toModelMessage)
+            : args.prompt,
+        } as GenerateObjectArgs<FlexibleSchema<T>>;
+        const ctx = (
+          options?.customCtx
+            ? { ...ctx_, ...options.customCtx(ctx_, targetArgs, llmArgs as any) }
+            : ctx_
+        ) as GenericActionCtx<GenericDataModel> & 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) ?? [],
         };
       },
     });
   }
 
   /**
-   * Save messages to the thread.
-   * Useful as a step in Workflows, e.g.
-   * ```ts
-   * const saveMessages = agent.asSaveMessagesMutation();
-   *
-   * const myWorkflow = workflow.define({
-   *   args: {...},
-   *   handler: async (step, args) => {
-   *     // do things to create (but not save)messages
-   *     const { messageIds } = await step.runMutation(internal.foo.saveMessages, {
-   *       threadId: args.threadId,
-   *       messages: args.messages,
-   *     });
-   *     // ...
-   *   },
-   * })
-   * ```
-   * @returns A mutation that can be used to save messages to the thread.
+   * @deprecated Use {@link saveMessages} directly instead.
    */
   asSaveMessagesMutation() {
     return internalMutationGeneric({
@@ -2056,226 +1491,21 @@ 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) => deserializeMessage(m.message)),
+          messages: args.messages.map((m) => toModelMessage(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"])),
         };
       },
     });
   }
 }
-
-/**
- * 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: ((ModelMessage & { id?: string | undefined }) | Message)[];
-  /**
-   * 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,
-      };
-    }
-  }
-  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,
-  };
-}
-
-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: ModelMessage | Message;
-    }
-  | {
-      /*
-       * 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],
-    };
-  }
-  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)! };
-}