Add onSaveMessages callback

.cursor/rules/convex_rules.mdc

@@ -159,7 +159,7 @@ export const listWithExtraArg = query({
     handler: async (ctx, args) => {
         return await ctx.db
         .query("messages")
-        .filter((q) => q.eq(q.field("author"), args.author))
+        .withIndex("by_author", (q) => q.eq("author", args.author))
         .order("desc")
         .paginate(args.paginationOpts);
     },
@@ -198,7 +198,7 @@ export const exampleQuery = query({
     handler: async (ctx, args) => {
         const idToUsername: Record<Id<"users">, string> = {};
         for (const userId of args.userIds) {
-            const user = await ctx.db.get(userId);
+            const user = await ctx.db.get("users", userId);
             if (user) {
                 idToUsername[user._id] = user.username;
             }
@@ -236,8 +236,8 @@ const messages = await ctx.db
 
 
 ## Mutation guidelines
-- Use `ctx.db.replace` to fully replace an existing document. This method will throw an error if the document does not exist.
-- Use `ctx.db.patch` to shallow merge updates into an existing document. This method will throw an error if the document does not exist.
+- Use `ctx.db.replace` to fully replace an existing document. This method will throw an error if the document does not exist. Syntax: `await ctx.db.replace('tasks', taskId, { name: 'Buy milk', completed: false })`
+- Use `ctx.db.patch` to shallow merge updates into an existing document. This method will throw an error if the document does not exist. Syntax: `await ctx.db.patch('tasks', taskId, { completed: true })`
 
 ## Action guidelines
 - Always add `"use node";` to the top of files containing actions that use Node.js built-in modules.
@@ -307,7 +307,7 @@ export const exampleQuery = query({
     args: { fileId: v.id("_storage") },
     returns: v.null(),
     handler: async (ctx, args) => {
-        const metadata: FileMetadata | null = await ctx.db.system.get(args.fileId);
+        const metadata: FileMetadata | null = await ctx.db.system.get("_storage", args.fileId);
         console.log(metadata);
         return null;
     },
@@ -434,7 +434,7 @@ Internal Functions:
   "description": "This example shows how to build a chat app without authentication.",
   "version": "1.0.0",
   "dependencies": {
-    "convex": "^1.17.4",
+    "convex": "^1.31.2",
     "openai": "^4.79.0"
   },
   "devDependencies": {
@@ -667,6 +667,35 @@ export default defineSchema({
 });
 ```
 
+#### convex/tsconfig.json
+```typescript
+{
+  /* This TypeScript project config describes the environment that
+   * Convex functions run in and is used to typecheck them.
+   * You can modify it, but some settings required to use Convex.
+   */
+  "compilerOptions": {
+    /* These settings are not required by Convex and can be modified. */
+    "allowJs": true,
+    "strict": true,
+    "moduleResolution": "Bundler",
+    "jsx": "react-jsx",
+    "skipLibCheck": true,
+    "allowSyntheticDefaultImports": true,
+
+    /* These compiler options are required by Convex */
+    "target": "ESNext",
+    "lib": ["ES2021", "dom"],
+    "forceConsistentCasingInFileNames": true,
+    "module": "ESNext",
+    "isolatedModules": true,
+    "noEmit": true
+  },
+  "include": ["./**/*"],
+  "exclude": ["./_generated"]
+}
+```
+
 #### src/App.tsx
 ```typescript
 export default function App() {

src/client/index.ts

@@ -87,6 +87,8 @@ import type {
   Options,
   RawRequestResponseHandler,
   MutationCtx,
+  SaveMessagesCallbackArgs,
+  SaveMessagesHandler,
   StorageOptions,
   StreamingTextArgs,
   StreamObjectArgs,
@@ -123,6 +125,7 @@ export {
   vMessageDoc,
   vPaginationResult,
   vProviderMetadata,
+  vSaveMessagesArgs,
   vSource,
   vStorageOptions,
   vStreamArgs,
@@ -186,6 +189,8 @@ export type {
   ContextOptions,
   ProviderMetadata,
   RawRequestResponseHandler,
+  SaveMessagesCallbackArgs,
+  SaveMessagesHandler,
   StorageOptions,
   StreamArgs,
   SyncStreamsReturnValue,
@@ -793,6 +798,7 @@ export class Agent<
       ...rest,
       agentName: this.options.name,
       embeddings,
+      onSaveMessages: this.options.onSaveMessages,
     });
   }
 

src/client/messages.ts

@@ -1,5 +1,9 @@
 import type { ModelMessage } from "ai";
-import type { PaginationOptions, PaginationResult } from "convex/server";
+import {
+  createFunctionHandle,
+  type PaginationOptions,
+  type PaginationResult,
+} from "convex/server";
 import type { MessageDoc } from "../validators.js";
 import { validateVectorDimension } from "../component/vector/tables.js";
 import {
@@ -17,6 +21,7 @@ import type {
   MutationCtx,
   QueryCtx,
   ActionCtx,
+  SaveMessagesHandler,
 } from "./types.js";
 import { parse } from "convex-helpers/validators";
 
@@ -104,6 +109,11 @@ export type SaveMessagesArgs = {
    * A pending message ID to replace when adding messages.
    */
   pendingMessageId?: string;
+  /**
+   * Optional callback mutation to invoke after messages are saved.
+   * Called within the same transaction as the message save.
+   */
+  onSaveMessages?: SaveMessagesHandler;
 };
 
 /**
@@ -131,6 +141,11 @@ export async function saveMessages(
       };
     }
   }
+  // Convert function reference to handle string for passing to component
+  const onSaveMessagesHandle = args.onSaveMessages
+    ? await createFunctionHandle(args.onSaveMessages)
+    : undefined;
+
   const result = await ctx.runMutation(component.messages.addMessages, {
     threadId: args.threadId,
     userId: args.userId ?? undefined,
@@ -153,6 +168,7 @@ export async function saveMessages(
       }),
     ),
     failPendingSteps: args.failPendingSteps ?? false,
+    onSaveMessages: onSaveMessagesHandle,
   });
   return { messages: result.messages };
 }

src/client/start.ts

@@ -14,7 +14,13 @@ import {
   serializeObjectResult,
 } from "../mapping.js";
 import { embedMessages, fetchContextWithPrompt } from "./search.js";
-import type { ActionCtx, AgentComponent, Config, Options } from "./types.js";
+import type {
+  ActionCtx,
+  AgentComponent,
+  Config,
+  Options,
+  SaveMessagesHandler,
+} from "./types.js";
 import type { Message, MessageDoc } from "../validators.js";
 import {
   getModelName,
@@ -25,7 +31,11 @@ import { wrapTools, type ToolCtx } from "./createTool.js";
 import type { Agent } from "./index.js";
 import { assert, omit } from "convex-helpers";
 import { saveInputMessages } from "./saveInputMessages.js";
-import type { GenericActionCtx, GenericDataModel } from "convex/server";
+import {
+  createFunctionHandle,
+  type GenericActionCtx,
+  type GenericDataModel,
+} from "convex/server";
 
 export async function startGeneration<
   T,
@@ -93,6 +103,7 @@ export async function startGeneration<
       languageModel?: LanguageModel;
       agentName: string;
       agentForToolCtx?: Agent;
+      onSaveMessages?: SaveMessagesHandler;
     },
 ): Promise<{
   args: T & {
@@ -124,6 +135,11 @@ export async function startGeneration<
         ?.userId) ??
     undefined;
 
+  // Convert function reference to handle string for passing to component
+  const onSaveMessagesHandle = opts.onSaveMessages
+    ? await createFunctionHandle(opts.onSaveMessages)
+    : undefined;
+
   const context = await fetchContextWithPrompt(ctx, component, {
     ...opts,
     userId,
@@ -282,6 +298,7 @@ export async function startGeneration<
           embeddings,
           failPendingSteps: false,
           finishStreamId,
+          onSaveMessages: onSaveMessagesHandle,
         });
         const lastMessage = saved.messages.at(-1)!;
         if (createPendingMessage) {

src/client/types.ts

@@ -30,6 +30,7 @@ export interface Output<_T = any, _P = any, _E = any> {
   createElementStreamTransform: any;
 }
 import type {
+  FunctionReference,
   GenericActionCtx,
   GenericDataModel,
   GenericMutationCtx,
@@ -39,6 +40,7 @@ import type {
 import type {
   MessageDoc,
   ProviderMetadata,
+  SaveMessagesCallbackArgs,
   StreamDelta,
   StreamMessage,
   ThreadDoc,
@@ -151,6 +153,12 @@ export type Config = {
    * log the raw request body or response headers to a table, or logs.
    */
   rawRequestResponseHandler?: RawRequestResponseHandler;
+  /**
+   * Called whenever messages are saved to the thread. This includes messages
+   * saved via generateText, streamText, generateObject, streamObject,
+   * saveMessage, and saveMessages.
+   */
+  onSaveMessages?: SaveMessagesHandler;
   /**
    * @deprecated Reach out if you use this. Otherwise will be removed soon.
    * Default provider options to pass for the LLM calls.
@@ -348,6 +356,50 @@ export type RawRequestResponseHandler = (
   },
 ) => void | Promise<void>;
 
+export type { SaveMessagesCallbackArgs } from "../validators.js";
+
+/**
+ * A reference to a mutation function that will be called whenever messages are
+ * saved to a thread. This callback is invoked **within the same transaction**
+ * as the message save, making it transactional.
+ *
+ * This includes messages saved via generateText, streamText, generateObject,
+ * streamObject, saveMessage, and saveMessages.
+ *
+ * Use this to trigger side effects when messages are saved, such as updating
+ * counters, creating notifications, or syncing with external systems.
+ *
+ * @example
+ * ```ts
+ * // In your convex/myModule.ts:
+ * import { vSaveMessagesArgs } from "@convex-dev/agent";
+ *
+ * export const onNewMessages = internalMutation({
+ *   args: vSaveMessagesArgs,
+ *   handler: async (ctx, args) => {
+ *     // This runs in the same transaction as the message save
+ *     await ctx.db.insert("messageEvents", {
+ *       threadId: args.threadId,
+ *       messageCount: args.messages.length,
+ *       timestamp: Date.now(),
+ *     });
+ *   },
+ * });
+ *
+ * // In your agent configuration:
+ * const agent = new Agent(components.agent, {
+ *   name: "myAgent",
+ *   languageModel: openai.chat("gpt-4o-mini"),
+ *   onSaveMessages: internal.myModule.onNewMessages,
+ * });
+ * ```
+ */
+export type SaveMessagesHandler = FunctionReference<
+  "mutation",
+  "internal" | "public",
+  SaveMessagesCallbackArgs
+>;
+
 export type AgentComponent = ComponentApi;
 
 export type TextArgs<
@@ -621,6 +673,11 @@ export type Options = {
    * ordering will not apply. This excludes the system message / instructions.
    */
   contextHandler?: ContextHandler;
+  /**
+   * Called whenever messages are saved to the thread.
+   * Overrides the onSaveMessages handler set in the agent constructor.
+   */
+  onSaveMessages?: SaveMessagesHandler;
 };
 
 export type SyncStreamsReturnValue =

src/component/_generated/component.ts

@@ -698,6 +698,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
               | { message: string; type: "other" }
             >;
           }>;
+          onSaveMessages?: string;
           pendingMessageId?: string;
           promptMessageId?: string;
           threadId: string;

src/component/messages.ts

@@ -2,6 +2,7 @@ import { assert, omit, pick } from "convex-helpers";
 import { mergedStream, stream } from "convex-helpers/server/stream";
 import {
   paginationOptsValidator,
+  type FunctionReference,
   type WithoutSystemFields,
 } from "convex/server";
 import type { ObjectType } from "convex/values";
@@ -19,6 +20,7 @@ import {
   vMessageWithMetadataInternal,
   vPaginationResult,
   type MessageDoc,
+  type SaveMessagesCallbackArgs,
 } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
 import type { Doc, Id } from "./_generated/dataModel.js";
@@ -144,6 +146,9 @@ const addMessagesArgs = {
   // If provided, finish this stream atomically with the message save.
   // This prevents UI flickering from separate mutations (issue #181).
   finishStreamId: v.optional(v.id("streamingMessages")),
+  // Optional callback mutation to invoke after messages are saved.
+  // Called within the same transaction as the message save.
+  onSaveMessages: v.optional(v.string()),
 };
 export const addMessages = mutation({
   args: addMessagesArgs,
@@ -166,6 +171,7 @@ async function addMessagesHandler(
     failPendingSteps,
     // Destructured separately to exclude from `...rest` (used in addMessages args, not message fields)
     finishStreamId,
+    onSaveMessages,
     messages,
     promptMessageId,
     pendingMessageId,
@@ -313,7 +319,24 @@ async function addMessagesHandler(
   if (finishStreamId) {
     await finishHandler(ctx, { streamId: finishStreamId });
   }
-  return { messages: toReturn.map(publicMessage) };
+  const savedMessages = toReturn.map(publicMessage);
+  // Call the onSaveMessages callback if provided, within the same transaction
+  if (onSaveMessages && savedMessages.length > 0) {
+    const callbackArgs: SaveMessagesCallbackArgs = {
+      userId,
+      threadId,
+      messages: savedMessages,
+    };
+    await ctx.runMutation(
+      onSaveMessages as unknown as FunctionReference<
+        "mutation",
+        "public" | "internal",
+        SaveMessagesCallbackArgs
+      >,
+      callbackArgs,
+    );
+  }
+  return { messages: savedMessages };
 }
 
 // exported for tests

src/validators.ts

@@ -671,6 +671,13 @@ export const vMessageDoc = v.object({
 });
 export type MessageDoc = Infer<typeof vMessageDoc>; // Public
 
+export const vSaveMessagesArgs = v.object({
+  userId: v.optional(v.string()),
+  threadId: v.string(),
+  messages: v.array(vMessageDoc),
+});
+export type SaveMessagesCallbackArgs = Infer<typeof vSaveMessagesArgs>;
+
 export const vThreadDoc = v.object({
   _id: v.string(),
   _creationTime: v.number(),