@convex-dev/agent 0.1.6 → 0.1.7-alpha.1

package/src/client/index.ts

@@ -31,7 +31,6 @@ import {
   serializeMessage,
   serializeNewMessagesInStep,
   serializeObjectResult,
-  serializeStep,
 } from "../mapping.js";
 import {
   DEFAULT_MESSAGE_RANGE,
@@ -40,7 +39,7 @@ import {
   isTool,
 } from "../shared.js";
 import {
-  type MessageWithMetadata as InnerMessageWithMetadata,
+  type MessageWithMetadata,
   type MessageStatus,
   type ProviderMetadata,
   type SearchOptions,
@@ -60,13 +59,14 @@ import type {
   AgentComponent,
   ContextOptions,
   GenerationOutputMetadata,
-  OpaqueIds,
   Options,
   OurObjectArgs,
   OurStreamObjectArgs,
+  RawRequestResponseHandler,
   RunActionCtx,
   RunMutationCtx,
   RunQueryCtx,
+  ActionCtx,
   StorageOptions,
   StreamingTextArgs,
   SyncStreamsReturnValue,
@@ -75,6 +75,8 @@ import type {
   UsageHandler,
 } from "./types.js";
 
+export { storeFile, getFile } from "./files.js";
+export { serializeDataOrUrl } from "../mapping.js";
 export { vMessageDoc, vThreadDoc } from "../component/schema.js";
 export {
   vAssistantMessage,
@@ -169,6 +171,11 @@ export class Agent<AgentTools extends ToolSet> {
        * 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;
     }
   ) {}
 
@@ -253,7 +260,7 @@ export class Agent<AgentTools extends ToolSet> {
     threadId: string;
   }>;
   async createThread<ThreadTools extends ToolSet | undefined = undefined>(
-    ctx: RunActionCtx | RunMutationCtx,
+    ctx: ActionCtx | RunMutationCtx,
     args?: {
       userId: string;
       title?: string;
@@ -297,7 +304,7 @@ export class Agent<AgentTools extends ToolSet> {
    * @returns Functions bound to the userId and threadId on a `{thread}` object.
    */
   async continueThread<ThreadTools extends ToolSet | undefined = undefined>(
-    ctx: RunActionCtx,
+    ctx: ActionCtx,
     args: {
       /**
        * The associated thread created by {@link createThread}
@@ -342,317 +349,451 @@ export class Agent<AgentTools extends ToolSet> {
   }
 
   /**
-   * 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.
+   * This behaves like {@link generateText} from the "ai" package except that
+   * it add context based on the userId and threadId and saves the input and
+   * resulting messages to the thread, if specified.
+   * 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}.
+   * @returns The result of the generateText function.
    */
-  async getThreadMetadata(
-    ctx: RunQueryCtx,
-    args: { threadId: string }
-  ): Promise<ThreadDoc> {
-    const thread = await ctx.runQuery(this.component.threads.getThread, {
-      threadId: args.threadId,
+  async generateText<
+    TOOLS extends ToolSet | undefined = undefined,
+    OUTPUT = never,
+    OUTPUT_PARTIAL = never,
+  >(
+    ctx: ActionCtx,
+    {
+      userId: argsUserId,
+      threadId,
+      usageHandler,
+      tools: threadTools,
+    }: {
+      userId?: string;
+      threadId?: string;
+      /**
+       * The usage handler to use for this thread. Overrides any handler
+       * set in the agent constructor.
+       */
+      usageHandler?: UsageHandler;
+      /** @deprecated Pass `tools` in the next parameter instead. This is only intended to pass through thread-default tools.  */
+      tools?: ToolSet;
+    },
+    args: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL>,
+    options?: Options
+  ): Promise<
+    GenerateTextResult<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT> &
+      GenerationOutputMetadata
+  > {
+    const context = await this._saveMessagesAndFetchContext(ctx, args, {
+      userId: argsUserId,
+      threadId,
+      ...options,
     });
-    if (!thread) {
-      throw new Error("Thread not found");
-    }
-    return thread;
-  }
-
-  /**
-   * Update the metadata for a thread.
-   * @param ctx A ctx object from a mutation or action.
-   * @param args.threadId The thread to update the metadata for.
-   * @param args.patch The patch to apply to the thread.
-   * @returns The updated thread metadata.
-   */
-  async updateThreadMetadata(
-    ctx: RunMutationCtx,
-    args: { threadId: string; patch: Partial<WithoutSystemFields<ThreadDoc>> }
-  ): Promise<ThreadDoc> {
-    const thread = await ctx.runMutation(
-      this.component.threads.updateThread,
-      args
+    const { args: aiArgs, messageId, userId } = context;
+    const toolCtx = { ...ctx, userId, threadId, messageId, agent: this };
+    const tools = wrapTools(
+      toolCtx,
+      args.tools ?? threadTools ?? this.options.tools
+    ) as TOOLS extends undefined ? AgentTools : TOOLS;
+    const saveOutputMessages = this._shouldSaveOutputMessages(
+      options?.storageOptions
     );
-    return thread;
-  }
-
-  /**
-   * List messages from a thread.
-   * @param ctx A ctx object from a query, mutation, or action.
-   * @param args.threadId The thread to list messages from.
-   * @param args.paginationOpts Pagination options (e.g. via usePaginatedQuery).
-   * @param args.excludeToolMessages Whether to exclude tool messages.
-   *   False by default.
-   * @param args.statuses What statuses to include. All by default.
-   * @returns The MessageDoc's in a format compatible with usePaginatedQuery.
-   */
-  async listMessages(
-    ctx: RunQueryCtx,
-    args: {
-      threadId: string;
-      paginationOpts: PaginationOptions;
-      excludeToolMessages?: boolean;
-      statuses?: MessageStatus[];
-    }
-  ): Promise<PaginationResult<MessageDoc>> {
-    if (args.paginationOpts.numItems === 0) {
-      return {
-        page: [],
-        isDone: true,
-        continueCursor: args.paginationOpts.cursor ?? "",
-      };
-    }
-    return ctx.runQuery(this.component.messages.listMessagesByThreadId, {
-      order: "desc",
-      ...args,
-    });
-  }
-
-  /**
-   * A function that handles fetching stream deltas, used with the React hooks
-   * `useThreadMessages` or `useStreamingThreadMessages`.
-   * @param ctx A ctx object from a query, mutation, or action.
-   * @param args.threadId The thread to sync streams for.
-   * @param args.streamArgs The stream arguments with per-stream cursors.
-   * @returns The deltas for each stream from their existing cursor.
-   */
-  async syncStreams(
-    ctx: RunQueryCtx,
-    args: {
-      threadId: string;
-      streamArgs: StreamArgs | undefined;
-    }
-  ): Promise<SyncStreamsReturnValue | undefined> {
-    if (!args.streamArgs) return undefined;
-    if (args.streamArgs.kind === "list") {
-      return {
-        kind: "list",
-        messages: await ctx.runQuery(this.component.streams.list, {
-          threadId: args.threadId,
-        }),
-      };
-    } else {
-      return {
-        kind: "deltas",
-        deltas: await ctx.runQuery(this.component.streams.listDeltas, {
-          threadId: args.threadId,
-          cursors: args.streamArgs.cursors,
-        }),
-      };
+    const trackUsage = usageHandler ?? this.options.usageHandler;
+    try {
+      const result = (await generateText({
+        // Can be overridden
+        maxSteps: this.options.maxSteps,
+        ...aiArgs,
+        tools,
+        onStepFinish: async (step) => {
+          if (threadId && messageId && saveOutputMessages) {
+            await this.saveStep(ctx, {
+              userId,
+              threadId,
+              promptMessageId: messageId,
+              step,
+            });
+          }
+          if (this.options.rawRequestResponseHandler) {
+            await this.options.rawRequestResponseHandler(ctx, {
+              userId,
+              threadId,
+              agentName: this.options.name,
+              request: step.request,
+              response: step.response,
+            });
+          }
+          if (trackUsage && step.usage) {
+            await trackUsage(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);
+        },
+      })) as GenerateTextResult<
+        TOOLS extends undefined ? AgentTools : TOOLS,
+        OUTPUT
+      > &
+        GenerationOutputMetadata;
+      result.messageId = messageId;
+      return result;
+    } catch (error) {
+      if (threadId && messageId) {
+        console.error("RollbackMessage", messageId);
+        await ctx.runMutation(this.component.messages.rollbackMessage, {
+          messageId,
+          error: (error as Error).message,
+        });
+      }
+      throw error;
     }
   }
 
   /**
-   * Fetch the context messages for a thread.
-   * @param ctx Either a query, mutation, or action ctx.
-   *   If it is not an action context, you can't do text or
-   *   vector search.
-   * @param args The associated thread, user, message
-   * @returns
+   * This behaves like {@link streamText} from the "ai" package except that
+   * it add context based on the userId and threadId and saves the input and
+   * resulting messages to the thread, if specified.
+   * Use {@link continueThread} to get a version of this function already scoped
+   * to a thread (and optionally userId).
    */
-  async fetchContextMessages(
-    ctx: RunQueryCtx | RunActionCtx,
-    args: {
-      userId: string | undefined;
-      threadId: string | undefined;
-      messages: CoreMessage[];
+  async streamText<
+    TOOLS extends ToolSet | undefined = undefined,
+    OUTPUT = never,
+    PARTIAL_OUTPUT = never,
+  >(
+    ctx: ActionCtx,
+    {
+      userId: argsUserId,
+      threadId,
+      usageHandler,
       /**
-       * 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.
+       * @deprecated Pass `tools` in the next parameter instead.
+       * This is only intended to pass through thread-default tools.
        */
-      upToAndIncludingMessageId?: string;
-      contextOptions: ContextOptions | undefined;
+      tools: threadTools,
+    }: {
+      userId?: string;
+      threadId?: string;
+      usageHandler?: UsageHandler;
+      tools?: ToolSet;
+    },
+    /**
+     * The arguments to the streamText function, similar to the ai `streamText` function.
+     */
+    args: StreamingTextArgs<AgentTools, TOOLS, OUTPUT, PARTIAL_OUTPUT>,
+    /**
+     * The {@link ContextOptions} and {@link StorageOptions}
+     * options to use for fetching contextual messages and saving input/output messages.
+     */
+    options?: Options & {
+      /**
+       * Whether to save incremental data (deltas) from streaming responses.
+       * Defaults to false.
+       * If false, it will not save any deltas to the database.
+       * If true, it will save deltas with {@link DEFAULT_STREAMING_OPTIONS}.
+       *
+       * Regardless of this option, when streaming you are able to use this
+       * `streamText` function as you would with the "ai" package's version:
+       * iterating over the text, streaming it over HTTP, etc.
+       */
+      saveStreamDeltas?: boolean | StreamingOptions;
     }
-  ): Promise<MessageDoc[]> {
-    assert(args.userId || args.threadId, "Specify userId or threadId");
-    // Fetch the latest messages from the thread
-    let included: Set<string> | undefined;
-    const opts = this._mergedContextOptions(args.contextOptions);
-    const contextMessages: MessageDoc[] = [];
-    if (
-      args.threadId &&
-      (opts.recentMessages !== 0 || args.upToAndIncludingMessageId)
-    ) {
-      const { page } = await ctx.runQuery(
-        this.component.messages.listMessagesByThreadId,
-        {
-          threadId: args.threadId,
-          excludeToolMessages:
-            opts.includeToolCalls === true ? false : opts.excludeToolMessages,
-          paginationOpts: {
-            numItems: opts.recentMessages ?? DEFAULT_RECENT_MESSAGES,
-            cursor: null,
-          },
-          upToAndIncludingMessageId: args.upToAndIncludingMessageId,
-          order: "desc",
-          statuses: ["success"],
+  ): Promise<
+    StreamTextResult<
+      TOOLS extends undefined ? AgentTools : TOOLS,
+      PARTIAL_OUTPUT
+    > &
+      GenerationOutputMetadata
+  > {
+    const context = await this._saveMessagesAndFetchContext(ctx, args, {
+      userId: argsUserId,
+      threadId,
+      ...options,
+    });
+    const { args: aiArgs, messageId, order, stepOrder, userId } = context;
+    const toolCtx = { ...ctx, userId, threadId, messageId, agent: this };
+    const tools = wrapTools(
+      toolCtx,
+      args.tools ?? threadTools ?? this.options.tools
+    ) as TOOLS extends undefined ? AgentTools : TOOLS;
+    const saveOutputMessages = this._shouldSaveOutputMessages(
+      options?.storageOptions
+    );
+    const trackUsage = usageHandler ?? this.options.usageHandler;
+    const streamer =
+      threadId && options?.saveStreamDeltas
+        ? new DeltaStreamer(this.component, ctx, options.saveStreamDeltas, {
+            threadId,
+            userId,
+            agentName: this.options.name,
+            model: aiArgs.model.modelId,
+            provider: aiArgs.model.provider,
+            providerOptions: aiArgs.providerOptions,
+            order,
+            stepOrder,
+          })
+        : undefined;
+
+    const result = streamText({
+      // Can be overridden
+      maxSteps: this.options.maxSteps,
+      ...aiArgs,
+      tools,
+      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 && saveOutputMessages) {
+          await ctx.runMutation(this.component.messages.rollbackMessage, {
+            messageId,
+            error: (error.error as Error).message,
+          });
         }
-      );
-      included = new Set(page.map((m) => m._id));
-      contextMessages.push(
-        // Reverse since we fetched in descending order
-        ...page.reverse()
-      );
-    }
-    if (opts.searchOptions?.textSearch || opts.searchOptions?.vectorSearch) {
-      const targetMessage = contextMessages.find(
-        (m) => m._id === args.upToAndIncludingMessageId
-      )?.message;
-      const messagesToSearch = targetMessage
-        ? [targetMessage, ...args.messages]
-        : args.messages;
-      if (!("runAction" in ctx)) {
-        throw new Error("searchUserMessages only works in an action");
-      }
-      const searchMessages = await ctx.runAction(
-        this.component.messages.searchMessages,
-        {
-          searchAllMessagesForUserId: opts?.searchOtherThreads
-            ? args.userId ??
-              (args.threadId &&
-                (
-                  await ctx.runQuery(this.component.threads.getThread, {
-                    threadId: args.threadId,
-                  })
-                )?.userId)
-            : undefined,
-          threadId: args.threadId,
-          beforeMessageId: args.upToAndIncludingMessageId,
-          ...(await this._searchOptionsWithDefaults(opts, messagesToSearch)),
+        return args.onError?.(error);
+      },
+      onStepFinish: async (step) => {
+        // console.log("onStepFinish", step);
+        // TODO: compare delta to the output. internally drop the deltas when committing
+        if (threadId && messageId) {
+          const saved = await this.saveStep(ctx, {
+            userId,
+            threadId,
+            promptMessageId: messageId,
+            step,
+          });
+          // TODO: figure out pending/not
+          await streamer?.finish(saved.messages);
         }
-      );
-      // TODO: track what messages we used for context
-      contextMessages.unshift(
-        ...searchMessages.filter((m) => !included?.has(m._id))
-      );
-    }
-    // Ensure we don't include tool messages without a corresponding tool call
-    return filterOutOrphanedToolMessages(
-      contextMessages.sort((a, b) =>
-        // Sort the raw MessageDocs by order and stepOrder
-        a.order === b.order ? a.stepOrder - b.stepOrder : a.order - b.order
-      )
-    );
+        if (this.options.rawRequestResponseHandler) {
+          await this.options.rawRequestResponseHandler(ctx, {
+            userId,
+            threadId,
+            agentName: this.options.name,
+            request: step.request,
+            response: step.response,
+          });
+        }
+        if (trackUsage && step.usage) {
+          await trackUsage(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);
+      },
+    }) as StreamTextResult<
+      TOOLS extends undefined ? AgentTools : TOOLS,
+      PARTIAL_OUTPUT
+    > &
+      GenerationOutputMetadata;
+    result.messageId = messageId;
+    return result;
   }
 
   /**
-   * Get the embeddings for a set of messages.
-   * @param messages The messages to get the embeddings for.
-   * @returns The embeddings for the messages.
+   * This behaves like {@link generateObject} from the "ai" package except that
+   * it add context based on the userId and threadId and saves the input and
+   * resulting messages to the thread, if specified.
+   * Use {@link continueThread} to get a version of this function already scoped
+   * to a thread (and optionally userId).
    */
-  async generateEmbeddings(messages: CoreMessage[]) {
-    let embeddings:
-      | {
-          vectors: (number[] | null)[];
-          dimension: VectorDimension;
-          model: string;
-        }
-      | undefined;
-    if (this.options.textEmbedding) {
-      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;
+  async generateObject<T>(
+    ctx: RunActionCtx,
+    {
+      userId: argsUserId,
+      threadId,
+      usageHandler,
+    }: { userId?: string; threadId?: string; usageHandler?: UsageHandler },
+    /**
+     * The arguments to the generateObject function, similar to the ai.generateObject function.
+     */
+    args: OurObjectArgs<T>,
+    /**
+     * 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 context = await this._saveMessagesAndFetchContext(ctx, args, {
+      userId: argsUserId,
+      threadId,
+      ...options,
+    });
+    const { args: aiArgs, messageId, userId } = context;
+    const trackUsage = usageHandler ?? this.options.usageHandler;
+    const saveOutputMessages = this._shouldSaveOutputMessages(
+      options?.storageOptions
+    );
+    try {
+      const result = (await generateObject(
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        aiArgs as any
+      )) as GenerateObjectResult<T> & GenerationOutputMetadata;
+
+      if (threadId && messageId && saveOutputMessages) {
+        await this.saveObject(ctx, {
+          threadId,
+          promptMessageId: messageId,
+          result,
+          userId,
+        });
       }
-      // Then embed those messages.
-      const textEmbeddings = await this.options.textEmbedding.doEmbed({
-        values: messageTexts.filter((t): t is string => !!t),
-      });
-      // TODO: record usage of embeddings
-      // Then assemble the embeddings into a single array with nulls for the messages without text.
-      const embeddingsOrNull = Array(messages.length).fill(null);
-      textIndexes.forEach((i, j) => {
-        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,
-        };
+      result.messageId = messageId;
+      if (this.options.rawRequestResponseHandler) {
+        await this.options.rawRequestResponseHandler(ctx, {
+          userId,
+          threadId,
+          agentName: this.options.name,
+          request: result.request,
+          response: result.response,
+        });
+      }
+      if (trackUsage && result.usage) {
+        await trackUsage(ctx, {
+          userId,
+          threadId,
+          agentName: this.options.name,
+          model: aiArgs.model.modelId,
+          provider: aiArgs.model.provider,
+          usage: result.usage,
+          providerMetadata: result.providerMetadata,
+        });
+      }
+      return result;
+    } catch (error) {
+      if (threadId && messageId) {
+        await ctx.runMutation(this.component.messages.rollbackMessage, {
+          messageId,
+          error: (error as Error).message,
+        });
       }
+      throw error;
     }
-    return embeddings;
   }
 
   /**
-   * Generate embeddings for a set of messages, and save them to the database.
-   * It will not generate or save embeddings for messages that already have an
-   * embedding.
-   * @param ctx The ctx parameter to an action.
-   * @param args The messageIds to generate embeddings for.
+   * This behaves like `streamObject` from the "ai" package except that
+   * it add context based on the userId and threadId and saves the input and
+   * resulting messages to the thread, if specified.
+   * Use {@link continueThread} to get a version of this function already scoped
+   * to a thread (and optionally userId).
    */
-  async generateAndSaveEmbeddings(
+  async streamObject<T>(
     ctx: RunActionCtx,
-    args: {
-      messageIds: string[];
-    }
-  ) {
-    const messages = (
-      await ctx.runQuery(this.component.messages.getMessagesByIds, {
-        messageIds: args.messageIds,
-      })
-    ).filter((m): m is NonNullable<typeof m> => m !== null);
-    if (messages.length !== args.messageIds.length) {
-      throw new Error(
-        "Some messages were not found: " +
-          args.messageIds
-            .filter((id) => !messages.some((m) => m?._id === id))
-            .join(", ")
-      );
-    }
-    if (messages.some((m) => !m.message)) {
-      throw new Error(
-        "Some messages don't have a message: " +
-          args.messageIds
-            .map((id, i) => (!messages[i].message ? id : undefined))
-            .filter((id): id is string => id !== undefined)
-            .join(", ")
-      );
-    }
-    const messagesMissingEmbeddings = messages.filter((m) => !m.embeddingId);
-    if (messagesMissingEmbeddings.length === 0) {
-      return;
-    }
-    const embeddings = await this.generateEmbeddings(
-      messagesMissingEmbeddings.map((m) => m!.message!)
-    );
-    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
-        ),
+    {
+      userId: argsUserId,
+      threadId,
+      usageHandler,
+    }: { userId?: string; threadId?: string; usageHandler?: UsageHandler },
+    /**
+     * The arguments to the streamObject function, similar to the ai `streamObject` function.
+     */
+    args: OurStreamObjectArgs<T>,
+    /**
+     * 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
+  > {
+    // TODO: unify all this shared code between all the generate* and stream* functions
+    const context = await this._saveMessagesAndFetchContext(ctx, args, {
+      userId: argsUserId,
+      threadId,
+      ...options,
     });
+    const { args: aiArgs, messageId, userId } = context;
+    const trackUsage = usageHandler ?? this.options.usageHandler;
+    const saveOutputMessages = this._shouldSaveOutputMessages(
+      options?.storageOptions
+    );
+    const stream = streamObject<T>({
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      ...(aiArgs as any),
+      onError: async (error) => {
+        console.error("onError", error);
+        return args.onError?.(error);
+      },
+      onFinish: async (result) => {
+        if (threadId && messageId && saveOutputMessages) {
+          await this.saveObject(ctx, {
+            userId,
+            threadId,
+            promptMessageId: messageId,
+            result: {
+              object: result.object,
+              finishReason: "stop",
+              usage: result.usage,
+              warnings: result.warnings,
+              request: await stream.request,
+              response: result.response,
+              providerMetadata: result.providerMetadata,
+              experimental_providerMetadata:
+                result.experimental_providerMetadata,
+              logprobs: undefined,
+              toJsonResponse: stream.toTextStreamResponse,
+            },
+          });
+        }
+        if (trackUsage && result.usage) {
+          await trackUsage(ctx, {
+            userId,
+            threadId,
+            agentName: this.options.name,
+            model: aiArgs.model.modelId,
+            provider: aiArgs.model.provider,
+            usage: result.usage,
+            providerMetadata: result.providerMetadata,
+          });
+        }
+        if (this.options.rawRequestResponseHandler) {
+          await this.options.rawRequestResponseHandler(ctx, {
+            userId,
+            threadId,
+            agentName: this.options.name,
+            request: await stream.request,
+            response: result.response,
+          });
+        }
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        return args.onFinish?.(result as any);
+      },
+    }) as StreamObjectResult<DeepPartial<T>, T, never> &
+      GenerationOutputMetadata;
+    stream.messageId = messageId;
+    return stream;
   }
 
+  /**
+   * 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.
+   */
   async saveMessage(
     ctx: RunMutationCtx,
     args: {
@@ -687,7 +828,7 @@ export class Agent<AgentTools extends ToolSet> {
         }
     )
   ) {
-    const { lastMessageId } = await this.saveMessages(ctx, {
+    const { lastMessageId, messages } = await this.saveMessages(ctx, {
       threadId: args.threadId,
       userId: args.userId,
       messages:
@@ -697,7 +838,7 @@ export class Agent<AgentTools extends ToolSet> {
       metadata: args.metadata ? [args.metadata] : undefined,
       skipEmbeddings: args.skipEmbeddings,
     });
-    return { messageId: lastMessageId };
+    return { messageId: lastMessageId, message: messages.at(-1)! };
   }
 
   /**
@@ -707,7 +848,7 @@ export class Agent<AgentTools extends ToolSet> {
    * @returns
    */
   async saveMessages(
-    ctx: RunMutationCtx,
+    ctx: RunMutationCtx | RunActionCtx,
     args: {
       threadId: string;
       userId?: string;
@@ -766,7 +907,14 @@ export class Agent<AgentTools extends ToolSet> {
         );
       }
     } else {
-      embeddings = await this.generateEmbeddings(args.messages);
+      embeddings = await this.generateEmbeddings(
+        ctx,
+        {
+          userId: args.userId,
+          threadId: args.threadId,
+        },
+        args.messages
+      );
     }
     const result = await ctx.runMutation(this.component.messages.addMessages, {
       threadId: args.threadId,
@@ -774,12 +922,19 @@ export class Agent<AgentTools extends ToolSet> {
       agentName: this.options.name,
       promptMessageId: args.promptMessageId,
       embeddings,
-      messages: args.messages.map(
-        (m, i) =>
-          ({
+      messages: await Promise.all(
+        args.messages.map(async (m, i) => {
+          const { message, fileIds } = await serializeMessage(
+            ctx,
+            this.component,
+            m
+          );
+          return {
             ...args.metadata?.[i],
-            message: serializeMessage(m),
-          }) as MessageWithMetadata
+            message,
+            fileIds,
+          } as MessageWithMetadata;
+        })
       ),
       failPendingSteps: args.failPendingSteps ?? false,
       pending: args.pending ?? false,
@@ -791,361 +946,462 @@ export class Agent<AgentTools extends ToolSet> {
   }
 
   /**
-   * Save messages to the thread.
-   * Useful as a step in Workflows, e.g.
-   * ```ts
-   * const saveMessages = agent.asSaveMessagesMutation();
-   *
-   * const myWorkflow = workflow.define()
-   * ```
-   * @returns A mutation that can be used to save messages to the thread.
-   */
-  asSaveMessagesMutation() {
-    return internalMutationGeneric({
-      args: {
-        threadId: v.string(),
-        userId: v.optional(v.string()),
-        promptMessageId: v.optional(v.string()),
-        messages: v.array(vMessageWithMetadata),
-        pending: v.optional(v.boolean()),
-        failPendingSteps: v.optional(v.boolean()),
-      },
-      handler: async (ctx, args) => {
-        const { lastMessageId, messages } = await this.saveMessages(ctx, {
-          ...args,
-          messages: args.messages.map((m) => m.message),
-          metadata: args.messages.map(({ message: _, ...m }) => m),
-        });
-        return {
-          lastMessageId,
-          messageIds: messages.map((m) => m._id),
-        };
-      },
-    });
-  }
-
-  /**
-   * Explicitly save a "step" created by the AI SDK.
-   * @param ctx The ctx argument to a mutation or action.
-   * @param args The Step generated by the AI SDK.
+   * List messages from a thread.
+   * @param ctx A ctx object from a query, mutation, or action.
+   * @param args.threadId The thread to list messages from.
+   * @param args.paginationOpts Pagination options (e.g. via usePaginatedQuery).
+   * @param args.excludeToolMessages Whether to exclude tool messages.
+   *   False by default.
+   * @param args.statuses What statuses to include. All by default.
+   * @returns The MessageDoc's in a format compatible with usePaginatedQuery.
    */
-  async saveStep<TOOLS extends ToolSet>(
-    ctx: RunMutationCtx,
+  async listMessages(
+    ctx: RunQueryCtx,
     args: {
-      userId?: string;
       threadId: string;
-      /**
-       * The message this step is in response to.
-       */
-      promptMessageId: string;
-      /**
-       * The step to save, possibly including multiple tool calls.
-       */
-      step: StepResult<TOOLS>;
-      /**
-       * The model used to generate the step.
-       * Defaults to the chat model for the Agent.
-       */
-      model?: string;
-      /**
-       * The provider of the model used to generate the step.
-       * Defaults to the chat provider for the Agent.
-       */
-      provider?: string;
+      paginationOpts: PaginationOptions;
+      excludeToolMessages?: boolean;
+      statuses?: MessageStatus[];
     }
-  ): Promise<MessageDoc[]> {
-    const step = serializeStep(args.step as StepResult<ToolSet>);
-    const messages = serializeNewMessagesInStep(args.step, {
-      provider: args.provider ?? this.options.chat.provider,
-      model: args.model ?? this.options.chat.modelId,
-    });
-    const embeddings = await this.generateEmbeddings(
-      messages.map((m) => m.message)
-    );
-    const saved = await ctx.runMutation(this.component.messages.addStep, {
-      userId: args.userId,
-      threadId: args.threadId,
-      promptMessageId: args.promptMessageId,
-      step: { step, messages, embeddings },
-      failPendingSteps: false,
+  ): Promise<PaginationResult<MessageDoc>> {
+    if (args.paginationOpts.numItems === 0) {
+      return {
+        page: [],
+        isDone: true,
+        continueCursor: args.paginationOpts.cursor ?? "",
+      };
+    }
+    return ctx.runQuery(this.component.messages.listMessagesByThreadId, {
+      order: "desc",
+      ...args,
     });
-    return saved;
   }
 
   /**
-   * Commit or rollback a message that was pending.
-   * This is done automatically when saving messages by default.
-   * If creating pending messages, you can call this when the full "transaction" is done.
-   * @param ctx The ctx argument to your mutation or action.
-   * @param args What message to save. Generally the parent message sent into
-   *   the generateText call.
+   * A function that handles fetching stream deltas, used with the React hooks
+   * `useThreadMessages` or `useStreamingThreadMessages`.
+   * @param ctx A ctx object from a query, mutation, or action.
+   * @param args.threadId The thread to sync streams for.
+   * @param args.streamArgs The stream arguments with per-stream cursors.
+   * @returns The deltas for each stream from their existing cursor.
    */
-  async completeMessage(
-    ctx: RunMutationCtx,
+  async syncStreams(
+    ctx: RunQueryCtx,
     args: {
       threadId: string;
-      messageId: string;
-      result: { kind: "error"; error: string } | { kind: "success" };
+      streamArgs: StreamArgs | undefined;
     }
-  ): 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,
-      });
+  ): Promise<SyncStreamsReturnValue | undefined> {
+    if (!args.streamArgs) return undefined;
+    if (args.streamArgs.kind === "list") {
+      return {
+        kind: "list",
+        messages: await ctx.runQuery(this.component.streams.list, {
+          threadId: args.threadId,
+        }),
+      };
+    } else {
+      return {
+        kind: "deltas",
+        deltas: await ctx.runQuery(this.component.streams.listDeltas, {
+          threadId: args.threadId,
+          cursors: args.streamArgs.cursors,
+        }),
+      };
     }
   }
 
   /**
-   * This behaves like {@link generateText} from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified.
-   * 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}.
-   * @returns The result of the generateText function.
+   * Fetch the context messages for a thread.
+   * @param ctx Either a query, mutation, or action ctx.
+   *   If it is not an action context, you can't do text or
+   *   vector search.
+   * @param args The associated thread, user, message
+   * @returns
    */
-  async generateText<
-    TOOLS extends ToolSet | undefined = undefined,
-    OUTPUT = never,
-    OUTPUT_PARTIAL = never,
-  >(
-    ctx: RunActionCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      usageHandler,
-      tools: threadTools,
-    }: {
-      userId?: string;
-      threadId?: string;
+  async fetchContextMessages(
+    ctx: RunQueryCtx | RunActionCtx,
+    args: {
+      userId: string | undefined;
+      threadId: string | undefined;
+      messages: CoreMessage[];
       /**
-       * The usage handler to use for this thread. Overrides any handler
-       * set in the agent constructor.
+       * 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.
        */
-      usageHandler?: UsageHandler;
-      /** @deprecated Pass `tools` in the next parameter instead. This is only intended to pass through thread-default tools.  */
-      tools?: ToolSet;
-    },
-    args: TextArgs<AgentTools, TOOLS, OUTPUT, OUTPUT_PARTIAL>,
-    options?: Options
-  ): Promise<
-    GenerateTextResult<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT> &
-      GenerationOutputMetadata
-  > {
-    const context = await this._saveMessagesAndFetchContext(ctx, args, {
-      userId: argsUserId,
-      threadId,
-      ...options,
-    });
-    const { args: aiArgs, messageId, userId } = context;
-    const toolCtx = { ...ctx, userId, threadId, messageId, agent: this };
-    const tools = wrapTools(
-      toolCtx,
-      args.tools ?? threadTools ?? this.options.tools
-    ) as TOOLS extends undefined ? AgentTools : TOOLS;
-    const saveOutputMessages =
-      options?.storageOptions?.saveOutputMessages ??
-      this.options.storageOptions?.saveOutputMessages;
-    const trackUsage = usageHandler ?? this.options.usageHandler;
-    try {
-      const result = (await generateText({
-        // Can be overridden
-        maxSteps: this.options.maxSteps,
-        ...aiArgs,
-        tools,
-        onStepFinish: async (step) => {
-          if (threadId && messageId && saveOutputMessages !== false) {
-            await this.saveStep(ctx, {
-              userId,
-              threadId,
-              promptMessageId: messageId,
-              step,
-            });
-          }
-          if (trackUsage && step.usage) {
-            await trackUsage(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);
-        },
-      })) as GenerateTextResult<
-        TOOLS extends undefined ? AgentTools : TOOLS,
-        OUTPUT
-      > &
-        GenerationOutputMetadata;
-      result.messageId = messageId;
-      return result;
-    } catch (error) {
-      if (threadId && messageId) {
-        console.error("RollbackMessage", messageId);
-        await ctx.runMutation(this.component.messages.rollbackMessage, {
-          messageId,
-          error: (error as Error).message,
-        });
+      upToAndIncludingMessageId?: string;
+      contextOptions: ContextOptions | undefined;
+    }
+  ): Promise<MessageDoc[]> {
+    assert(args.userId || args.threadId, "Specify userId or threadId");
+    // Fetch the latest messages from the thread
+    let included: Set<string> | undefined;
+    const opts = this._mergedContextOptions(args.contextOptions);
+    const contextMessages: MessageDoc[] = [];
+    if (
+      args.threadId &&
+      (opts.recentMessages !== 0 || args.upToAndIncludingMessageId)
+    ) {
+      const { page } = await ctx.runQuery(
+        this.component.messages.listMessagesByThreadId,
+        {
+          threadId: args.threadId,
+          excludeToolMessages:
+            opts.includeToolCalls === true ? false : opts.excludeToolMessages,
+          paginationOpts: {
+            numItems: opts.recentMessages ?? DEFAULT_RECENT_MESSAGES,
+            cursor: null,
+          },
+          upToAndIncludingMessageId: args.upToAndIncludingMessageId,
+          order: "desc",
+          statuses: ["success"],
+        }
+      );
+      included = new Set(page.map((m) => m._id));
+      contextMessages.push(
+        // Reverse since we fetched in descending order
+        ...page.reverse()
+      );
+    }
+    if (opts.searchOptions?.textSearch || opts.searchOptions?.vectorSearch) {
+      const targetMessage = contextMessages.find(
+        (m) => m._id === args.upToAndIncludingMessageId
+      )?.message;
+      const messagesToSearch = targetMessage
+        ? [targetMessage, ...args.messages]
+        : args.messages;
+      if (!("runAction" in ctx)) {
+        throw new Error("searchUserMessages only works in an action");
       }
-      throw error;
+      const searchMessages = await ctx.runAction(
+        this.component.messages.searchMessages,
+        {
+          searchAllMessagesForUserId: opts?.searchOtherThreads
+            ? args.userId ??
+              (args.threadId &&
+                (
+                  await ctx.runQuery(this.component.threads.getThread, {
+                    threadId: args.threadId,
+                  })
+                )?.userId)
+            : undefined,
+          threadId: args.threadId,
+          beforeMessageId: args.upToAndIncludingMessageId,
+          ...(await this._searchOptionsWithEmbeddingAndDefaults(
+            ctx,
+            { userId: args.userId, threadId: args.threadId },
+            opts,
+            messagesToSearch
+          )),
+        }
+      );
+      // TODO: track what messages we used for context
+      contextMessages.unshift(
+        ...searchMessages.filter((m) => !included?.has(m._id))
+      );
+    }
+    // Ensure we don't include tool messages without a corresponding tool call
+    return filterOutOrphanedToolMessages(
+      contextMessages.sort((a, b) =>
+        // Sort the raw MessageDocs by order and stepOrder
+        a.order === b.order ? a.stepOrder - b.stepOrder : a.order - b.order
+      )
+    );
+  }
+
+  /**
+   * 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.
+   */
+  async getThreadMetadata(
+    ctx: RunQueryCtx,
+    args: { threadId: string }
+  ): Promise<ThreadDoc> {
+    const thread = await ctx.runQuery(this.component.threads.getThread, {
+      threadId: args.threadId,
+    });
+    if (!thread) {
+      throw new Error("Thread not found");
     }
+    return thread;
   }
 
   /**
-   * This behaves like {@link streamText} from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified.
-   * Use {@link continueThread} to get a version of this function already scoped
-   * to a thread (and optionally userId).
+   * Update the metadata for a thread.
+   * @param ctx A ctx object from a mutation or action.
+   * @param args.threadId The thread to update the metadata for.
+   * @param args.patch The patch to apply to the thread.
+   * @returns The updated thread metadata.
    */
-  async streamText<
-    TOOLS extends ToolSet | undefined = undefined,
-    OUTPUT = never,
-    PARTIAL_OUTPUT = never,
-  >(
+  async updateThreadMetadata(
+    ctx: RunMutationCtx,
+    args: { threadId: string; patch: Partial<WithoutSystemFields<ThreadDoc>> }
+  ): Promise<ThreadDoc> {
+    const thread = await ctx.runMutation(
+      this.component.threads.updateThread,
+      args
+    );
+    return thread;
+  }
+
+  /**
+   * Get the embeddings for a set of messages.
+   * @param messages The messages to get the embeddings for.
+   * @returns The embeddings for the messages.
+   */
+  async generateEmbeddings(
     ctx: RunActionCtx,
     {
-      userId: argsUserId,
+      userId,
       threadId,
-      usageHandler,
-      /**
-       * @deprecated Pass `tools` in the next parameter instead.
-       * This is only intended to pass through thread-default tools.
-       */
-      tools: threadTools,
     }: {
-      userId?: string;
-      threadId?: string;
-      usageHandler?: UsageHandler;
-      tools?: ToolSet;
+      userId: string | undefined;
+      threadId: string | undefined;
     },
-    /**
-     * The arguments to the streamText function, similar to the ai `streamText` function.
-     */
-    args: StreamingTextArgs<AgentTools, TOOLS, OUTPUT, PARTIAL_OUTPUT>,
-    /**
-     * The {@link ContextOptions} and {@link StorageOptions}
-     * options to use for fetching contextual messages and saving input/output messages.
-     */
-    options?: Options & {
-      /**
-       * Whether to save incremental data (deltas) from streaming responses.
-       * Defaults to false.
-       * If false, it will not save any deltas to the database.
-       * If true, it will save deltas with {@link DEFAULT_STREAMING_OPTIONS}.
-       *
-       * Regardless of this option, when streaming you are able to use this
-       * `streamText` function as you would with the "ai" package's version:
-       * iterating over the text, streaming it over HTTP, etc.
-       */
-      saveStreamDeltas?: boolean | StreamingOptions;
+    messages: CoreMessage[]
+  ) {
+    if (!this.options.textEmbedding) {
+      return undefined;
     }
-  ): Promise<
-    StreamTextResult<
-      TOOLS extends undefined ? AgentTools : TOOLS,
-      PARTIAL_OUTPUT
-    > &
-      GenerationOutputMetadata
-  > {
-    const context = await this._saveMessagesAndFetchContext(ctx, args, {
-      userId: argsUserId,
+    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,
-      ...options,
+      values: messageTexts.filter((t): t is string => !!t),
     });
-    const { args: aiArgs, messageId, order, stepOrder, userId } = context;
-    const toolCtx = { ...ctx, userId, threadId, messageId, agent: this };
-    const tools = wrapTools(
-      toolCtx,
-      args.tools ?? threadTools ?? this.options.tools
-    ) as TOOLS extends undefined ? AgentTools : TOOLS;
-    const storageOptions = {
-      ...this.options.storageOptions,
-      ...options?.storageOptions,
-    };
-    const saveOutputMessages = storageOptions.saveOutputMessages;
-    const trackUsage = usageHandler ?? this.options.usageHandler;
-    const streamer =
-      threadId && options?.saveStreamDeltas
-        ? new DeltaStreamer(this.component, ctx, options.saveStreamDeltas, {
-            threadId,
-            userId,
-            agentName: this.options.name,
-            model: aiArgs.model.modelId,
-            provider: aiArgs.model.provider,
-            providerOptions: aiArgs.providerOptions,
-            order,
-            stepOrder,
-          })
-        : undefined;
+    // TODO: record usage of embeddings
+    // Then assemble the embeddings into a single array with nulls for the messages without text.
+    const embeddingsOrNull = Array(messages.length).fill(null);
+    textIndexes.forEach((i, j) => {
+      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;
+  }
 
-    const result = streamText({
-      // Can be overridden
-      maxSteps: this.options.maxSteps,
-      ...aiArgs,
-      tools,
-      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 && saveOutputMessages !== false) {
-          await ctx.runMutation(this.component.messages.rollbackMessage, {
-            messageId,
-            error: (error.error as Error).message,
-          });
-        }
-        return args.onError?.(error);
-      },
-      onStepFinish: async (step) => {
-        // console.log("onStepFinish", step);
-        // TODO: compare delta to the output. internally drop the deltas when committing
-        if (threadId && messageId) {
-          const saved = await this.saveStep(ctx, {
-            userId,
-            threadId,
-            promptMessageId: messageId,
-            step,
-          });
-          await streamer?.finish(saved);
-        }
-        if (trackUsage && step.usage) {
-          await trackUsage(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);
+  /**
+   * Generate embeddings for a set of messages, and save them to the database.
+   * It will not generate or save embeddings for messages that already have an
+   * embedding.
+   * @param ctx The ctx parameter to an action.
+   * @param args The messageIds to generate embeddings for.
+   */
+  async generateAndSaveEmbeddings(
+    ctx: RunActionCtx,
+    args: {
+      messageIds: string[];
+    }
+  ) {
+    const messages = (
+      await ctx.runQuery(this.component.messages.getMessagesByIds, {
+        messageIds: args.messageIds,
+      })
+    ).filter((m): m is NonNullable<typeof m> => m !== null);
+    if (messages.length !== args.messageIds.length) {
+      throw new Error(
+        "Some messages were not found: " +
+          args.messageIds
+            .filter((id) => !messages.some((m) => m?._id === id))
+            .join(", ")
+      );
+    }
+    if (messages.some((m) => !m.message)) {
+      throw new Error(
+        "Some messages don't have a message: " +
+          args.messageIds
+            .map((id, i) => (!messages[i].message ? id : undefined))
+            .filter((id): id is string => id !== undefined)
+            .join(", ")
+      );
+    }
+    const messagesMissingEmbeddings = messages.filter((m) => !m.embeddingId);
+    if (messagesMissingEmbeddings.length === 0) {
+      return;
+    }
+    const embeddings = await this.generateEmbeddings(
+      ctx,
+      {
+        userId: messagesMissingEmbeddings[0]!.userId,
+        threadId: messagesMissingEmbeddings[0]!.threadId,
       },
-    }) as StreamTextResult<
-      TOOLS extends undefined ? AgentTools : TOOLS,
-      PARTIAL_OUTPUT
-    > &
-      GenerationOutputMetadata;
-    result.messageId = messageId;
-    return result;
+      messagesMissingEmbeddings.map((m) => m!.message!)
+    );
+    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
+        ),
+    });
+  }
+
+  /**
+   * Explicitly save a "step" created by the AI SDK.
+   * @param ctx The ctx argument to a mutation or action.
+   * @param args The Step generated by the AI SDK.
+   */
+  async saveStep<TOOLS extends ToolSet>(
+    ctx: ActionCtx,
+    args: {
+      userId?: string;
+      threadId: string;
+      /**
+       * The message this step is in response to.
+       */
+      promptMessageId: string;
+      /**
+       * The step to save, possibly including multiple tool calls.
+       */
+      step: StepResult<TOOLS>;
+      /**
+       * The model used to generate the step.
+       * Defaults to the chat model for the Agent.
+       */
+      model?: string;
+      /**
+       * The provider of the model used to generate the step.
+       * Defaults to the chat provider for the Agent.
+       */
+      provider?: string;
+    }
+  ): Promise<{ messages: MessageDoc[]; pending?: MessageDoc }> {
+    const messages = await serializeNewMessagesInStep(
+      ctx,
+      this.component,
+      args.step,
+      {
+        provider: args.provider ?? this.options.chat.provider,
+        model: args.model ?? this.options.chat.modelId,
+      }
+    );
+    const embeddings = await this.generateEmbeddings(
+      ctx,
+      { userId: args.userId, threadId: args.threadId },
+      messages.map((m) => m.message)
+    );
+    const saved = await ctx.runMutation(this.component.messages.addMessages, {
+      userId: args.userId,
+      threadId: args.threadId,
+      agentName: this.options.name,
+      promptMessageId: args.promptMessageId,
+      messages,
+      embeddings,
+      failPendingSteps: false,
+    });
+    return saved;
+  }
+
+  /**
+   * Manually save the result of a generateObject call to the thread.
+   * This happens automatically when using {@link generateObject} or {@link streamObject}
+   * from the `thread` object created by {@link continueThread} or {@link createThread}.
+   * @param ctx The context passed from the mutation or action function calling this.
+   * @param args The arguments to the saveObject function.
+   */
+  async saveObject(
+    ctx: RunActionCtx,
+    args: {
+      userId: string | undefined;
+      threadId: string;
+      promptMessageId: string;
+      result: GenerateObjectResult<unknown>;
+      metadata?: Omit<MessageWithMetadata, "message">;
+    }
+  ): Promise<void> {
+    const { messages } = serializeObjectResult(args.result, {
+      model: this.options.chat.modelId,
+      provider: this.options.chat.provider,
+    });
+    const embeddings = await this.generateEmbeddings(
+      ctx,
+      { userId: args.userId, threadId: args.threadId },
+      messages.map((m) => m.message)
+    );
+
+    await ctx.runMutation(this.component.messages.addMessages, {
+      userId: args.userId,
+      threadId: args.threadId,
+      promptMessageId: args.promptMessageId,
+      failPendingSteps: false,
+      messages,
+      embeddings,
+      agentName: this.options.name,
+      pending: false,
+    });
+  }
+
+  /**
+   * Commit or rollback a message that was pending.
+   * This is done automatically when saving messages by default.
+   * If creating pending messages, you can call this when the full "transaction" is done.
+   * @param ctx The ctx argument to your mutation or action.
+   * @param args What message to save. Generally the parent message sent into
+   *   the generateText call.
+   */
+  async completeMessage(
+    ctx: RunMutationCtx,
+    args: {
+      threadId: string;
+      messageId: string;
+      result: { kind: "error"; error: string } | { kind: "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,
+      });
+    }
   }
 
   async _saveMessagesAndFetchContext<
@@ -1222,9 +1478,10 @@ export class Agent<AgentTools extends ToolSet> {
     if (
       threadId &&
       messages.length &&
+      storageOptions?.saveMessages !== "none" &&
       storageOptions?.saveAnyInputMessages !== false
     ) {
-      const saveAll = storageOptions?.saveAllInputMessages;
+      const saveAll = storageOptions?.saveMessages === "all";
       const coreMessages = saveAll ? messages : messages.slice(-1);
       const saved = await this.saveMessages(ctx, {
         threadId,
@@ -1257,195 +1514,9 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
-  /**
-   * This behaves like {@link generateObject} from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified.
-   * Use {@link continueThread} to get a version of this function already scoped
-   * to a thread (and optionally userId).
-   */
-  async generateObject<T>(
-    ctx: RunActionCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      usageHandler,
-    }: { userId?: string; threadId?: string; usageHandler?: UsageHandler },
-    /**
-     * The arguments to the generateObject function, similar to the ai.generateObject function.
-     */
-    args: OurObjectArgs<T>,
-    /**
-     * 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 context = await this._saveMessagesAndFetchContext(ctx, args, {
-      userId: argsUserId,
-      threadId,
-      ...options,
-    });
-    const { args: aiArgs, messageId, userId } = context;
-    const trackUsage = usageHandler ?? this.options.usageHandler;
-    const saveOutputMessages =
-      options?.storageOptions?.saveOutputMessages ??
-      this.options.storageOptions?.saveOutputMessages;
-    try {
-      const result = (await generateObject(
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        aiArgs as any
-      )) as GenerateObjectResult<T> & GenerationOutputMetadata;
-
-      if (threadId && messageId && saveOutputMessages !== false) {
-        await this.saveObject(ctx, {
-          threadId,
-          promptMessageId: messageId,
-          result,
-          userId,
-        });
-      }
-      result.messageId = messageId;
-      if (trackUsage && result.usage) {
-        await trackUsage(ctx, {
-          userId,
-          threadId,
-          agentName: this.options.name,
-          model: aiArgs.model.modelId,
-          provider: aiArgs.model.provider,
-          usage: result.usage,
-          providerMetadata: result.providerMetadata,
-        });
-      }
-      return result;
-    } catch (error) {
-      if (threadId && messageId) {
-        await ctx.runMutation(this.component.messages.rollbackMessage, {
-          messageId,
-          error: (error as Error).message,
-        });
-      }
-      throw error;
-    }
-  }
-
-  /**
-   * This behaves like `streamObject` from the "ai" package except that
-   * it add context based on the userId and threadId and saves the input and
-   * resulting messages to the thread, if specified.
-   * Use {@link continueThread} to get a version of this function already scoped
-   * to a thread (and optionally userId).
-   */
-  async streamObject<T>(
-    ctx: RunActionCtx,
-    {
-      userId: argsUserId,
-      threadId,
-      usageHandler,
-    }: { userId?: string; threadId?: string; usageHandler?: UsageHandler },
-    /**
-     * The arguments to the streamObject function, similar to the ai `streamObject` function.
-     */
-    args: OurStreamObjectArgs<T>,
-    /**
-     * 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
-  > {
-    // TODO: unify all this shared code between all the generate* and stream* functions
-    const context = await this._saveMessagesAndFetchContext(ctx, args, {
-      userId: argsUserId,
-      threadId,
-      ...options,
-    });
-    const { args: aiArgs, messageId, userId } = context;
-    const trackUsage = usageHandler ?? this.options.usageHandler;
-    const saveOutputMessages =
-      options?.storageOptions?.saveOutputMessages ??
-      this.options.storageOptions?.saveOutputMessages;
-    const stream = streamObject<T>({
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      ...(aiArgs as any),
-      onError: async (error) => {
-        console.error("onError", error);
-        return args.onError?.(error);
-      },
-      onFinish: async (result) => {
-        if (threadId && messageId && saveOutputMessages !== false) {
-          await this.saveObject(ctx, {
-            userId,
-            threadId,
-            promptMessageId: messageId,
-            result: {
-              object: result.object,
-              finishReason: "stop",
-              usage: result.usage,
-              warnings: result.warnings,
-              request: await stream.request,
-              response: result.response,
-              providerMetadata: result.providerMetadata,
-              experimental_providerMetadata:
-                result.experimental_providerMetadata,
-              logprobs: undefined,
-              toJsonResponse: stream.toTextStreamResponse,
-            },
-          });
-        }
-        if (trackUsage && result.usage) {
-          await trackUsage(ctx, {
-            userId,
-            threadId,
-            agentName: this.options.name,
-            model: aiArgs.model.modelId,
-            provider: aiArgs.model.provider,
-            usage: result.usage,
-            providerMetadata: result.providerMetadata,
-          });
-        }
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        return args.onFinish?.(result as any);
-      },
-    }) as StreamObjectResult<DeepPartial<T>, T, never> &
-      GenerationOutputMetadata;
-    stream.messageId = messageId;
-    return stream;
-  }
-
-  /**
-   * Manually save the result of a generateObject call to the thread.
-   * This happens automatically when using {@link generateObject} or {@link streamObject}
-   * from the `thread` object created by {@link continueThread} or {@link createThread}.
-   * @param ctx The context passed from the mutation or action function calling this.
-   * @param args The arguments to the saveObject function.
-   */
-  async saveObject(
-    ctx: RunMutationCtx,
-    args: {
-      userId: string | undefined;
-      threadId: string;
-      promptMessageId: string;
-      result: GenerateObjectResult<unknown>;
-      metadata?: Omit<MessageWithMetadata, "message">;
-    }
-  ): Promise<void> {
-    const { step, messages } = serializeObjectResult(args.result, {
-      model: this.options.chat.modelId,
-      provider: this.options.chat.provider,
-    });
-    const embeddings = await this.generateEmbeddings(
-      messages.map((m) => m.message)
-    );
-
-    await ctx.runMutation(this.component.messages.addStep, {
-      userId: args.userId,
-      threadId: args.threadId,
-      promptMessageId: args.promptMessageId,
-      failPendingSteps: false,
-      step: { step, messages, embeddings },
-    });
+  _shouldSaveOutputMessages(storageOpts?: StorageOptions): boolean {
+    const opts = storageOpts ?? this.options.storageOptions;
+    return opts?.saveOutputMessages !== false && opts?.saveMessages !== "none";
   }
 
   _mergedContextOptions(opts: ContextOptions | undefined): ContextOptions {
@@ -1462,7 +1533,9 @@ export class Agent<AgentTools extends ToolSet> {
     };
   }
 
-  async _searchOptionsWithDefaults(
+  async _searchOptionsWithEmbeddingAndDefaults(
+    ctx: RunActionCtx,
+    { userId, threadId }: { userId?: string; threadId?: string },
     contextOptions: ContextOptions,
     messages: CoreMessage[]
   ): Promise<SearchOptions> {
@@ -1487,16 +1560,58 @@ export class Agent<AgentTools extends ToolSet> {
       this.options.textEmbedding
     ) {
       search.vector = (
-        await this.options.textEmbedding.doEmbed({
+        await this.doEmbed(ctx, {
+          threadId,
+          userId,
           values: [text],
         })
       ).embeddings[0];
-      // TODO: record usage of embeddings
       search.vectorModel = this.options.textEmbedding.modelId;
     }
     return search;
   }
 
+  async doEmbed(
+    ctx: RunActionCtx,
+    options: {
+      userId: string | undefined;
+      threadId: string | undefined;
+      values: string[];
+      abortSignal?: AbortSignal;
+      headers?: Record<string, string | undefined>;
+    }
+  ): Promise<{ embeddings: number[][] }> {
+    const embedding = this.options.textEmbedding;
+    assert(embedding, "textEmbedding is required");
+    const result = await embedding.doEmbed({
+      values: options.values,
+      abortSignal: options.abortSignal,
+      headers: options.headers,
+    });
+    if (this.options.usageHandler && result.usage) {
+      await this.options.usageHandler(ctx, {
+        userId: options.userId,
+        threadId: options.threadId,
+        agentName: this.options.name,
+        model: embedding.modelId,
+        provider: embedding.provider,
+        providerMetadata: result.rawResponse
+          ? { [embedding.provider]: result.rawResponse }
+          : undefined,
+        usage: {
+          promptTokens: result.usage.tokens,
+          completionTokens: 0,
+          totalTokens: result.usage.tokens,
+        },
+      });
+    }
+    return { embeddings: result.embeddings };
+  }
+
+  /**
+   * WORKFLOW UTILITIES
+   */
+
   /**
    * Create a mutation that creates a thread so you can call it from a Workflow.
    * e.g.
@@ -1522,7 +1637,7 @@ export class Agent<AgentTools extends ToolSet> {
         title: v.optional(v.string()),
         summary: v.optional(v.string()),
       },
-      handler: async (ctx, args) => {
+      handler: async (ctx, args): Promise<{ threadId: string }> => {
         const { threadId } = await this.createThread(ctx, args);
         return { threadId };
       },
@@ -1643,6 +1758,50 @@ export class Agent<AgentTools extends ToolSet> {
       },
     });
   }
+
+  /**
+   * 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.
+   */
+  asSaveMessagesMutation() {
+    return internalMutationGeneric({
+      args: {
+        threadId: v.string(),
+        userId: v.optional(v.string()),
+        promptMessageId: v.optional(v.string()),
+        messages: v.array(vMessageWithMetadata),
+        pending: v.optional(v.boolean()),
+        failPendingSteps: v.optional(v.boolean()),
+      },
+      handler: async (ctx, args) => {
+        const { lastMessageId, messages } = await this.saveMessages(ctx, {
+          ...args,
+          messages: args.messages.map((m) => m.message),
+          metadata: args.messages.map(({ message: _, ...m }) => m),
+        });
+        return {
+          lastMessageId,
+          messageIds: messages.map((m) => m._id),
+        };
+      },
+    });
+  }
 }
 
 export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
@@ -1672,5 +1831,4 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
   return result;
 }
 
-type MessageWithMetadata = OpaqueIds<InnerMessageWithMetadata>;
 type CoreMessageMaybeWithId = CoreMessage & { id?: string | undefined };