Add transactional success completion for mutations

README.md

@@ -130,10 +130,11 @@ export const emailSent = internalMutation({
 });
 ```
 
-Note: the `onComplete` handler runs in a different transaction than the job
-enqueued. If you want to run it in the same transaction, you can do that work at
-the end of the enqueued function, before returning. This is generally faster and
-more typesafe when handling the "success" case.
+By default, the `onComplete` handler runs in a different transaction than the
+enqueued job. For mutations, `completeTransactionally: true` makes successful
+work and its selected callback commit together, as described below. You can also
+do that work at the end of the enqueued function, before returning, which avoids
+another function call and preserves the result's type directly.
 
 You can also use this equivalent helper to define an `onComplete` mutation. Note
 the `DataModel` type parameter, if you want ctx.db to be type safe.
@@ -169,6 +170,39 @@ with `cancel` can still report `"success"` or a terminal `"failed"`. Direct
 scheduler cancelation (for example, from the dashboard) is treated as a failure
 and can trigger retries instead.
 
+### Transactional success for mutations
+
+Use `completeTransactionally: true` with `enqueueMutation` or
+`enqueueMutationBatch` to commit each mutation, its selected success callback,
+and workpool completion bookkeeping together:
+
+```ts
+await pool.enqueueMutation(ctx, internal.tasks.process, args, {
+  onComplete: internal.tasks.handleResult,
+  onCompleteExcludeKinds: ["canceled"],
+  completeTransactionally: true,
+});
+```
+
+The callback can read the mutation's writes. If the success callback or
+completion bookkeeping throws, the scheduled mutation fails and all those writes
+roll back. Workpool's periodic recovery then marks the work as failed and
+invokes the callback for `"failed"`, unless excluded. Until recovery runs, the
+job remains running and occupies a concurrency slot. There is no separate
+fallback execution of the success callback.
+
+Ordinary work errors still use the existing failure-completion path. Failure and
+cancellation callbacks keep their usual transaction behavior. Excluding every
+kind, or omitting the callback, still allows transactional bookkeeping, avoiding
+the separately scheduled completion mutation. Jobs in a batch remain independent
+transactions. Actions and queries do not support this option.
+
+This option requires Convex 1.41 or newer. The work and callback share one
+transaction budget. Workpool uses nested transaction limits to reserve space for
+completion bookkeeping; work that fits by itself may exceed the combined budget.
+Limit overruns or execution failures can still fail the transaction, in which
+case recovery handles it as above.
+
 ### Idempotency
 
 Idempotent actions are actions that can be run multiple times safely. This
@@ -336,6 +370,8 @@ options include:
 - `onCompleteExcludeKinds`: Optional list of result kinds to skip: `"success"`,
   `"failed"`, and `"canceled"`. Defaults to excluding none, so every outcome
   invokes the callback.
+- `completeTransactionally`: Mutation-only option to commit successful work, its
+  selected callback, and completion bookkeeping together. Defaults to false.
 - `context`: Any data you want to pass to the `onComplete` mutation.
 - `runAt` and `runAfter`: Similar to `ctx.scheduler.run*`, allows you to
   schedule the work to run later. By default it's immediate.

src/client/index.test.ts

@@ -22,6 +22,7 @@ import workpool from "../test.js";
 import type { api } from "../component/_generated/api.js";
 import {
   type EnqueueOptions,
+  type TransactionalEnqueueOptions,
   NonRetryableError,
   type OnCompleteArgs,
   type RetryOption,
@@ -31,6 +32,8 @@ import {
   vOnCompleteArgs,
   vResult,
   vWorkId,
+  enqueue,
+  enqueueBatch,
 } from "./index.js";
 
 const schema = defineSchema({
@@ -190,6 +193,109 @@ describe("completion callbacks through the client and scheduler", () => {
         .collect(),
     );
 
+  test.each([false, true])(
+    "transactional success works through mutation enqueue (batch: %s)",
+    async (batch) => {
+      const args = [
+        { key: "atomic-success" },
+        { key: "atomic-failure", fail: true },
+      ];
+      const options = {
+        completeTransactionally: true,
+        onComplete: refs.complete,
+        context: { key: "atomic-callbacks" },
+      };
+      const ids = await t.mutation(async (ctx) =>
+        batch
+          ? pool.enqueueMutationBatch(ctx, refs.mutation, args, options)
+          : Promise.all(
+              args.map((args) =>
+                pool.enqueueMutation(ctx, refs.mutation, args, options),
+              ),
+            ),
+      );
+      await drain();
+      expect((await events("atomic-success")).map((e) => e.kind)).toEqual([
+        "work",
+      ]);
+      expect(await events("atomic-failure")).toEqual([]);
+      const results = await events("atomic-callbacks");
+      expect(results).toHaveLength(2);
+      expect(results.find((e) => e.workId === ids[0])?.result).toEqual({
+        kind: "success",
+        returnValue: "mutation result",
+      });
+      expect(results.find((e) => e.workId === ids[1])?.result).toEqual({
+        kind: "failed",
+        error: "work failed",
+      });
+      expect(await t.query((ctx) => pool.statusBatch(ctx, ids))).toEqual([
+        { state: "finished" },
+        { state: "finished" },
+      ]);
+    },
+  );
+
+  test.each([false, true])(
+    "transactional mutation enqueue rolls back work when the callback fails (batch: %s)",
+    async (batch) => {
+      const args = [{ key: "rollback-first" }, { key: "rollback-second" }];
+      const options = {
+        completeTransactionally: true,
+        onComplete: refs.complete,
+        context: { key: "rollback-callback", throw: true },
+      };
+      try {
+        const ids = await t.mutation(async (ctx) =>
+          batch
+            ? pool.enqueueMutationBatch(ctx, refs.mutation, args, options)
+            : Promise.all(
+                args.map((args) =>
+                  pool.enqueueMutation(ctx, refs.mutation, args, options),
+                ),
+              ),
+        );
+        await vi.waitFor(() => expect(callback).toHaveBeenCalledTimes(2), {
+          timeout: 10_000,
+        });
+        await t.finishInProgressScheduledFunctions();
+        for (const { key } of args) expect(await events(key)).toEqual([]);
+        expect(await events("rollback-callback")).toEqual([]);
+        expect(
+          await t.query((ctx) => pool.statusBatch(ctx, ids)),
+        ).toMatchObject([{ state: "running" }, { state: "running" }]);
+      } finally {
+        // Inspect rollback before periodic recovery. worker.test.ts covers
+        // recovery with convex-test's missing scheduled-error text supplied.
+        await t.finishInProgressScheduledFunctions();
+        vi.clearAllTimers();
+      }
+    },
+  );
+
+  test("canceling pending transactional work preserves the cancellation callback", async () => {
+    const id = await t.mutation((ctx) =>
+      pool.enqueueMutation(
+        ctx,
+        refs.mutation,
+        { key: "atomic-cancel" },
+        {
+          completeTransactionally: true,
+          onComplete: refs.complete,
+          onCompleteExcludeKinds: ["success", "failed"],
+          context: { key: "atomic-cancel" },
+          runAt: Date.now() + 60_000,
+        },
+      ),
+    );
+    await t.mutation((ctx) => pool.cancel(ctx, id));
+    await drain();
+    expect(await events("atomic-cancel")).toMatchObject([
+      { kind: "callback", result: { kind: "canceled" } },
+    ]);
+    expect(callback).toHaveBeenCalledTimes(1);
+  });
+
   describe.each(excludeFilters.map((excludeKinds) => ({ excludeKinds })))(
     "exclude filter $excludeKinds",
     ({ excludeKinds }) => {
@@ -860,6 +966,77 @@ test("exclude filters only accept terminal result kinds", () => {
   }).toMatchTypeOf<EnqueueOptions>();
 });
 
+test("transactional completion is only exposed for mutations", () => {
+  const options = {
+    onComplete,
+    onCompleteExcludeKinds: ["failed", "canceled"] as const,
+    context: { label: "job" },
+    completeTransactionally: true,
+  } satisfies TransactionalEnqueueOptions<{ label: string }, number>;
+  const mutation = makeFunctionReference<"mutation", { value: number }, number>(
+    "work:mutation",
+  );
+  const query = makeFunctionReference<"query", { value: number }, number>(
+    "work:query",
+  );
+  expectTypeOf((pool: Workpool, ctx: MutationCtx) => {
+    // @ts-expect-error Actions cannot opt into transactional completion.
+    void pool.enqueueAction(ctx, action, { value: 1 }, options);
+    // @ts-expect-error Batched actions cannot opt into transactional completion.
+    void pool.enqueueActionBatch(ctx, action, [{ value: 1 }], options);
+    // @ts-expect-error Queries cannot opt into transactional completion.
+    void pool.enqueueQuery(ctx, query, { value: 1 }, options);
+    // @ts-expect-error Batched queries cannot opt into transactional completion.
+    void pool.enqueueQueryBatch(ctx, query, [{ value: 1 }], options);
+    return pool.enqueueMutation(ctx, mutation, { value: 1 }, options);
+  }).returns.toEqualTypeOf<Promise<WorkId>>();
+  expectTypeOf((pool: Workpool, ctx: MutationCtx) =>
+    pool.enqueueMutationBatch(ctx, mutation, [{ value: 1 }], options),
+  ).returns.toEqualTypeOf<Promise<WorkId[]>>();
+});
+
+test("standalone enqueue gates transactional completion on fnType", () => {
+  const options = {
+    onComplete,
+    context: { label: "job" },
+    completeTransactionally: true,
+  };
+  const mutation = makeFunctionReference<"mutation", { value: number }, number>(
+    "work:mutation",
+  );
+  const query = makeFunctionReference<"query", { value: number }, number>(
+    "work:query",
+  );
+  const mixedFn = makeFunctionReference<
+    "mutation" | "action",
+    { value: number },
+    number
+  >("work:mixed");
+  const arg = { value: 1 };
+  const batch = [arg];
+  expectTypeOf((c: WorkpoolComponent, ctx: MutationCtx) => {
+    // @ts-expect-error Actions cannot opt into transactional completion.
+    void enqueue(c, ctx, "action", action, arg, options);
+    // @ts-expect-error Queries cannot opt into transactional completion.
+    void enqueue(c, ctx, "query", query, arg, options);
+    // @ts-expect-error Batched actions cannot opt into transactional completion.
+    void enqueueBatch(c, ctx, "action", action, batch, options);
+    // @ts-expect-error Batched queries cannot opt into transactional completion.
+    void enqueueBatch(c, ctx, "query", query, batch, options);
+    // A union that includes a mutation must still be rejected as a whole: a
+    // distributive conditional would collapse this to boolean and let it through.
+    const mixed = "mutation" as "mutation" | "action";
+    // @ts-expect-error An action-capable union cannot opt in.
+    void enqueue(c, ctx, mixed, mixedFn, arg, options);
+    // @ts-expect-error A batched action-capable union cannot opt in.
+    void enqueueBatch(c, ctx, mixed, mixedFn, batch, options);
+    return enqueue(c, ctx, "mutation", mutation, arg, options);
+  }).returns.toEqualTypeOf<Promise<WorkId>>();
+  expectTypeOf((c: WorkpoolComponent, ctx: MutationCtx) =>
+    enqueueBatch(c, ctx, "mutation", mutation, batch, options),
+  ).returns.toEqualTypeOf<Promise<WorkId[]>>();
+});
+
 test("generated enqueue types match the component validators", () => {
   expectTypeOf<
     FunctionArgs<WorkpoolComponent["lib"]["enqueue"]>

src/client/index.ts

@@ -101,7 +101,8 @@ export class Workpool {
       NoInfer<ReturnType>
     >,
     fnArgs: Args,
-    options?: RetryOption & EnqueueOptions<Context, ReturnType>,
+    options?: RetryOption &
+      EnqueueOptions<Context, ReturnType> & { completeTransactionally?: never },
   ): Promise<WorkId> {
     const retryBehavior = getRetryBehavior(
       this.options.defaultRetryBehavior,
@@ -140,7 +141,8 @@ export class Workpool {
       NoInfer<ReturnType>
     >,
     argsArray: Array<Args>,
-    options?: RetryOption & EnqueueOptions<Context, ReturnType>,
+    options?: RetryOption &
+      EnqueueOptions<Context, ReturnType> & { completeTransactionally?: never },
   ): Promise<WorkId[]> {
     const retryBehavior = getRetryBehavior(
       this.options.defaultRetryBehavior,
@@ -176,7 +178,7 @@ export class Workpool {
       NoInfer<ReturnType>
     >,
     fnArgs: Args,
-    options?: EnqueueOptions<Context, ReturnType>,
+    options?: TransactionalEnqueueOptions<Context, ReturnType>,
   ): Promise<WorkId> {
     return enqueue(this.component, ctx, "mutation", fn, fnArgs, {
       ...this.options,
@@ -207,7 +209,7 @@ export class Workpool {
       NoInfer<ReturnType>
     >,
     argsArray: Array<Args>,
-    options?: EnqueueOptions<Context, ReturnType>,
+    options?: TransactionalEnqueueOptions<Context, ReturnType>,
   ): Promise<WorkId[]> {
     return enqueueBatch(this.component, ctx, "mutation", fn, argsArray, {
       ...this.options,
@@ -236,7 +238,9 @@ export class Workpool {
       NoInfer<ReturnType>
     >,
     fnArgs: Args,
-    options?: EnqueueOptions<Context, ReturnType>,
+    options?: EnqueueOptions<Context, ReturnType> & {
+      completeTransactionally?: never;
+    },
   ): Promise<WorkId> {
     return enqueue(this.component, ctx, "query", fn, fnArgs, {
       ...this.options,
@@ -268,7 +272,9 @@ export class Workpool {
       NoInfer<ReturnType>
     >,
     argsArray: Array<Args>,
-    options?: EnqueueOptions<Context, ReturnType>,
+    options?: EnqueueOptions<Context, ReturnType> & {
+      completeTransactionally?: never;
+    },
   ): Promise<WorkId[]> {
     return enqueueBatch(this.component, ctx, "query", fn, argsArray, {
       ...this.options,
@@ -525,6 +531,22 @@ export type EnqueueOptions<Context = unknown, ReturnValue = unknown> = {
     }
 );
 
+/** Options for work that can complete in the enqueuing transaction. */
+export type TransactionalEnqueueOptions<
+  Context = unknown,
+  ReturnValue = unknown,
+> = EnqueueOptions<Context, ReturnValue> & {
+  /**
+   * Commit successful work, its selected onComplete callback, and completion
+   * bookkeeping in one transaction. Defaults to false. Callback or bookkeeping
+   * errors roll back the work; recovery later reports the job as failed.
+   * The work and callback share transaction limits, with headroom reserved for
+   * bookkeeping. Failure and cancellation callbacks keep their usual behavior.
+   * Only supported for mutations.
+   */
+  completeTransactionally?: boolean;
+};
+
 export type OnCompleteArgs<Context = unknown, Returns = unknown> = {
   /**
    * The ID of the work that completed.
@@ -569,7 +591,7 @@ async function enqueueArgs<Context, ReturnType>(
     | FunctionReference<FunctionType, FunctionVisibility>
     | FunctionHandle<FunctionType, DefaultFunctionArgs>,
   opts:
-    | (EnqueueOptions<Context, ReturnType> &
+    | (TransactionalEnqueueOptions<Context, ReturnType> &
         Partial<Config> & { retryBehavior?: RetryBehavior })
     | undefined,
 ) {
@@ -580,6 +602,7 @@ async function enqueueArgs<Context, ReturnType>(
   return {
     fnHandle,
     fnName,
+    completeTransactionally: opts?.completeTransactionally,
     onComplete: opts?.onComplete
       ? {
           fnHandle: await createFunctionHandle(opts.onComplete),
@@ -636,6 +659,10 @@ export async function enqueueBatch<
     retryBehavior?: RetryBehavior;
     maxParallelism?: number;
     logLevel?: LogLevel;
+    /** Mutations only; see {@link TransactionalEnqueueOptions}. */
+    // Wrapped in a tuple so a union FnType is checked as a whole rather than
+    // distributed, which would admit action-capable calls.
+    completeTransactionally?: [FnType] extends ["mutation"] ? boolean : never;
   },
 ): Promise<WorkId[]> {
   const { config, ...defaults } = await enqueueArgs(fn, options);
@@ -699,6 +726,10 @@ export async function enqueue<
     retryBehavior?: RetryBehavior;
     maxParallelism?: number;
     logLevel?: LogLevel;
+    /** Mutations only; see {@link TransactionalEnqueueOptions}. */
+    // Wrapped in a tuple so a union FnType is checked as a whole rather than
+    // distributed, which would admit action-capable calls.
+    completeTransactionally?: [FnType] extends ["mutation"] ? boolean : never;
   },
 ): Promise<WorkId> {
   const id = await ctx.runMutation(component.lib.enqueue, {