onSuccess with different schema

onFailure and onCancel (new)

Author: reeceyang

src/client/index.ts

@@ -69,7 +69,10 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
           fnHandle: string;
           fnName: string;
           fnType: "action" | "mutation" | "query";
+          onCancel?: { context?: any; fnHandle: string };
           onComplete?: { context?: any; fnHandle: string };
+          onFailure?: { context?: any; fnHandle: string };
+          onSuccess?: { context?: any; fnHandle: string };
           retryBehavior?: {
             base: number;
             initialBackoffMs: number;
@@ -93,7 +96,10 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             fnHandle: string;
             fnName: string;
             fnType: "action" | "mutation" | "query";
+            onCancel?: { context?: any; fnHandle: string };
             onComplete?: { context?: any; fnHandle: string };
+            onFailure?: { context?: any; fnHandle: string };
+            onSuccess?: { context?: any; fnHandle: string };
             retryBehavior?: {
               base: number;
               initialBackoffMs: number;

src/component/_generated/component.ts

@@ -288,6 +288,82 @@ describe("complete", () => {
       });
     });
 
+    const handlerTypes = [
+      {
+        handler: "onSuccess" as const,
+        runResult: { kind: "success", returnValue: "test result" },
+      } as const,
+      {
+        handler: "onFailure" as const,
+        runResult: {
+          kind: "failed",
+          error: "error result",
+        },
+      } as const,
+      {
+        handler: "onCancel" as const,
+        runResult: {
+          kind: "canceled",
+        },
+      } as const,
+    ];
+
+    it.each(handlerTypes)(
+      "should call the correct onComplete handler",
+      async ({ handler, runResult }) => {
+        // Create a spy on runMutation
+        const runMutationSpy = vi.fn();
+
+        const testHandle = `${runResult.kind}Handle`;
+
+        // Enqueue a work item with a onComplete handler
+        const workId = await t.mutation(api.lib.enqueue, {
+          fnHandle: "testHandle",
+          fnName: "testFunction",
+          fnArgs: { test: "data" },
+          fnType: "mutation",
+          runAt: Date.now(),
+          config: {
+            maxParallelism: 10,
+            logLevel: "WARN",
+          },
+          [handler]: {
+            fnHandle: testHandle,
+            context: { someContext: "value" },
+          },
+        });
+
+        // Simulate a job completion with a spy on runMutation
+        await t.run(async (ctx) => {
+          // Create a modified context with a spy on runMutation
+          const spyCtx = {
+            ...ctx,
+            runMutation: runMutationSpy,
+          };
+
+          await completeHandler(spyCtx, {
+            jobs: [
+              {
+                workId,
+                runResult,
+                attempt: 0,
+              },
+            ],
+          });
+
+          // Verify the handler was called with the right arguments
+          expect(runMutationSpy).toHaveBeenCalledWith(
+            testHandle,
+            expect.objectContaining({
+              workId,
+              context: { someContext: "value" },
+              result: runResult,
+            }),
+          );
+        });
+      },
+    );
+
     it("should handle multiple jobs in a single call", async () => {
       // Enqueue multiple work items
       const workId1 = await t.mutation(api.lib.enqueue, {

src/component/complete.test.ts

@@ -1,6 +1,6 @@
 import type { FunctionHandle } from "convex/server";
 import { getConvexSize, type Infer, v } from "convex/values";
-import type { Id } from "./_generated/dataModel.js";
+import type { Doc, Id } from "./_generated/dataModel.js";
 import { internal } from "./_generated/api.js";
 import { internalMutation, type MutationCtx } from "./_generated/server.js";
 import { kickMainLoop } from "./kick.js";
@@ -116,18 +116,19 @@ export async function completeHandler(
         !!maxAttempts &&
         work.attempts < maxAttempts;
       if (!retry) {
-        if (work.onComplete) {
+        const onCompleteHandler = getOnCompleteHandler(work, job.runResult);
+        if (onCompleteHandler) {
           try {
             // Retrieve large context if stored separately
-            let context = work.onComplete.context;
+            let context = onCompleteHandler.context;
             if (context === undefined && work.payloadId) {
               const payload = await ctx.db.get(work.payloadId);
               if (payload) {
                 context = payload.context;
               }
             }
 
-            const handle = work.onComplete.fnHandle as FunctionHandle<
+            const handle = onCompleteHandler.fnHandle as FunctionHandle<
               "mutation",
               OnCompleteArgs,
               void
@@ -179,6 +180,21 @@ export async function completeHandler(
   }
 }
 
+function getOnCompleteHandler(work: Doc<"work">, jobResult: RunResult) {
+  if (work.onComplete) {
+    return work.onComplete;
+  } else {
+    switch (jobResult.kind) {
+      case "success":
+        return work.onSuccess;
+      case "failed":
+        return work.onFailure;
+      case "canceled":
+        return work.onCancel;
+    }
+  }
+}
+
 function stripResult(result: RunResult): RunResult {
   if (result.kind === "success") {
     return { kind: "success", returnValue: null };

src/component/complete.ts

@@ -1,4 +1,4 @@
-import { type ObjectType, v, getConvexSize } from "convex/values";
+import { type ObjectType, v, getConvexSize, type Value } from "convex/values";
 import type { WithoutSystemFields } from "convex/server";
 import { api } from "./_generated/api.js";
 import { type Doc, type Id } from "./_generated/dataModel.js";
@@ -33,6 +33,13 @@ const INLINE_METADATA_THRESHOLD = 8_000; // 8KB threshold
 const MAX_DOC_SIZE = 1_000_000; // Some buffer for 1MiB actual limit
 const PAYLOAD_DOC_OVERHEAD = 78; // Size of { args: null, context: null }
 
+const onCompleteHandlers = [
+  "onComplete",
+  "onSuccess",
+  "onFailure",
+  "onCancel",
+] as const;
+
 const itemArgs = {
   fnHandle: v.string(),
   fnName: v.string(),
@@ -41,8 +48,12 @@ const itemArgs = {
   runAt: v.number(),
   // TODO: annotation?
   onComplete: v.optional(vOnCompleteFnContext),
+  onSuccess: v.optional(vOnCompleteFnContext),
+  onFailure: v.optional(vOnCompleteFnContext),
+  onCancel: v.optional(vOnCompleteFnContext),
   retryBehavior: v.optional(retryBehavior),
 };
+type ItemArgs = ObjectType<typeof itemArgs>;
 const enqueueArgs = {
   ...itemArgs,
   config: vConfig.partial(),
@@ -61,21 +72,31 @@ async function enqueueHandler(
   ctx: MutationCtx,
   console: Logger,
   kickSegment: bigint,
-  { runAt, ...workArgs }: ObjectType<typeof itemArgs>,
+  { runAt, ...workArgs }: ItemArgs,
 ) {
   runAt = boundScheduledTime(runAt, console);
 
+  if (
+    workArgs.onComplete &&
+    (workArgs.onSuccess || workArgs.onFailure || workArgs.onCancel)
+  ) {
+    throw new Error(
+      `Function has an onComplete handler and a onSuccess, onFailure, or onCancel handler. onComplete handler cannot be defined when a onSuccess, onFailure, or onCancel handler exists`,
+    );
+  }
+
   const fnArgsSize = getConvexSize(workArgs.fnArgs);
   if (fnArgsSize > MAX_DOC_SIZE) {
     throw new Error(
       `Function arguments for function ${workArgs.fnName} too large: ${fnArgsSize} bytes (max: ${MAX_DOC_SIZE} bytes)`,
     );
   }
 
+  // Consider all the context from onSuccess, onFailure, and onCancel all together
   let contextSize = 0;
-  const context = workArgs.onComplete?.context;
-  if (context !== undefined) {
-    contextSize = getConvexSize(context);
+  const context = getContext(workArgs);
+  if (context.length > 0) {
+    contextSize = context.map(getConvexSize).reduce((a, b) => a + b, 0);
     if (contextSize > MAX_DOC_SIZE) {
       throw new Error(
         `OnComplete context for function ${workArgs.fnName} too large: ${contextSize} bytes (max: ${MAX_DOC_SIZE} bytes)`,
@@ -88,6 +109,12 @@ async function enqueueHandler(
     attempts: 0,
   };
 
+  const deleteContext = () => {
+    for (const handler of onCompleteHandlers) {
+      delete workItem[handler]?.context;
+    }
+  };
+
   if (fnArgsSize >= INLINE_METADATA_THRESHOLD) {
     // Args are large, store separately
     const payloadDoc: { args: Record<string, any>; context?: unknown } = {
@@ -99,14 +126,14 @@ async function enqueueHandler(
       // Context is also too big to inline
       payloadDoc.context = context;
       workItem.payloadSize += contextSize;
-      delete workItem.onComplete!.context;
+      deleteContext();
     }
     workItem.payloadId = await ctx.db.insert("payload", payloadDoc);
   } else if (fnArgsSize + contextSize >= INLINE_METADATA_THRESHOLD) {
     // Args are small enough, but combined with context it's too big.
     // Store just context in this case.
     workItem.payloadId = await ctx.db.insert("payload", { context });
-    delete workItem.onComplete!.context;
+    deleteContext();
     workItem.payloadSize = contextSize + PAYLOAD_DOC_OVERHEAD;
   }
 
@@ -261,5 +288,11 @@ async function shouldCancelWorkItem(
   return true;
 }
 
+function getContext(workArgs: Partial<ItemArgs>): Value[] {
+  return onCompleteHandlers
+    .map((handler) => workArgs[handler]?.context)
+    .filter((context) => context !== undefined);
+}
+
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 const console = "THIS IS A REMINDER TO USE createLogger";

src/component/lib.ts

@@ -71,6 +71,7 @@ describe("loop", () => {
       fnArgs: {},
       logLevel: "WARN",
       attempt: 0,
+      hasOnSuccess: false,
     });
   }
 

src/component/loop.test.ts

@@ -563,7 +563,15 @@ async function beginWork(
     throw new Error("work not found");
   }
   const { attempts: attempt, fnHandle, fnArgs, payloadId } = work;
-  const args = { workId, fnHandle, fnArgs, payloadId, logLevel, attempt };
+  const args = {
+    workId,
+    fnHandle,
+    fnArgs,
+    payloadId,
+    logLevel,
+    attempt,
+    hasOnSuccess: Boolean(work.onSuccess),
+  };
   let scheduleId;
   if (work.fnType === "action") {
     scheduleId = await ctx.scheduler.runAfter(

src/component/loop.ts

@@ -57,6 +57,7 @@ describe("recovery", () => {
       fnArgs: {},
       logLevel: "WARN",
       attempt: 0,
+      hasOnSuccess: false,
     });
   }
 

src/component/recovery.test.ts

@@ -68,6 +68,10 @@ export default defineSchema({
     payloadSize: v.optional(v.number()),
     attempts: v.number(), // number of completed attempts
     onComplete: v.optional(vOnCompleteFnContext),
+    // onComplete should be undefined when any of onSuccess, onFailure, or onCancel are defined. This is enforced at the client level, but not in the schema for backwards compatibility.
+    onSuccess: v.optional(vOnCompleteFnContext),
+    onFailure: v.optional(vOnCompleteFnContext),
+    onCancel: v.optional(vOnCompleteFnContext),
     retryBehavior: v.optional(retryBehavior),
     canceled: v.optional(v.boolean()),
   }),

src/component/schema.ts

@@ -70,20 +70,25 @@ export const DEFAULT_RETRY_BEHAVIOR: RetryBehavior = {
 // This ensures that the type satisfies the schema.
 const _ = {} as RetryBehavior satisfies Infer<typeof retryBehavior>;
 
-export const vResult = v.union(
-  v.object({
-    kind: v.literal("success"),
-    returnValue: v.any(),
-  }),
-  v.object({
-    kind: v.literal("failed"),
-    error: v.string(),
-  }),
-  v.object({
-    kind: v.literal("canceled"),
-  }),
-);
+export const vSuccessResult = v.object({
+  kind: v.literal("success"),
+  returnValue: v.any(),
+});
+
+export const vFailureResult = v.object({
+  kind: v.literal("failed"),
+  error: v.string(),
+});
+
+export const vCancelResult = v.object({
+  kind: v.literal("canceled"),
+});
+
+export const vResult = v.union(vSuccessResult, vFailureResult, vCancelResult);
 export type RunResult = Infer<typeof vResult>;
+export type SuccessResult = Infer<typeof vSuccessResult>;
+export type FailureResult = Infer<typeof vFailureResult>;
+export type CancelResult = Infer<typeof vCancelResult>;
 
 export const vOnCompleteFnContext = v.object({
   fnHandle: v.string(), // mutation