Add finally callback to customFunction

packages/convex-helpers/server/customFunctions.test.ts

@@ -27,7 +27,7 @@ import {
   type Auth,
 } from "convex/server";
 import { v } from "convex/values";
-import { afterEach, beforeEach, describe, expect, test } from "vitest";
+import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
 import { modules } from "./setup.test.js";
 
 const schema = defineSchema({
@@ -560,3 +560,119 @@ describe("nested custom functions", () => {
     ).rejects.toThrow("Validator error: Expected `string`");
   });
 });
+
+describe("finally callback", () => {
+  test("finally callback is called with result and context", async () => {
+    const t = convexTest(schema, modules);
+    const finallyMock = vi.fn();
+
+    const withFinally = customQuery(query, {
+      args: {},
+      input: async () => ({ ctx: { foo: "bar" }, args: {} }),
+      finally: (ctx, params) => {
+        finallyMock(ctx, params);
+      }
+    });
+
+    const successFn = withFinally({
+      args: {},
+      handler: async (ctx) => {
+        return { success: true, foo: ctx.foo };
+      },
+    });
+
+    await t.run(async (ctx) => {
+      const result = await (successFn as any)._handler(ctx, {});
+      expect(result).toEqual({ success: true, foo: "bar" });
+
+      expect(finallyMock).toHaveBeenCalledWith(
+        expect.objectContaining({ foo: "bar" }),
+        { result: { success: true, foo: "bar" } }
+      );
+    });
+
+    finallyMock.mockClear();
+
+    const errorFn = withFinally({
+      args: {},
+      handler: async () => {
+        throw new Error("Test error");
+      },
+    });
+
+    await t.run(async (ctx) => {
+      try {
+        await (errorFn as any)._handler(ctx, {});
+        expect.fail("Should have thrown an error");
+      } catch (e: unknown) {
+        const error = e as Error;
+        expect(error.message).toContain("Test error");
+      }
+
+      expect(finallyMock).toHaveBeenCalledWith(
+        expect.objectContaining({ foo: "bar" }),
+        { error: expect.objectContaining({ message: expect.stringContaining("Test error") }) }
+      );
+    });
+  });
+
+  test("finally callback with mutation", async () => {
+    const t = convexTest(schema, modules);
+    const finallyMock = vi.fn();
+
+    const withFinally = customMutation(mutation, {
+      args: {},
+      input: async () => ({ ctx: { foo: "bar" }, args: {} }),
+      finally: (ctx, params) => {
+        finallyMock(ctx, params);
+      }
+    });
+
+    const mutationFn = withFinally({
+      args: {},
+      handler: async (ctx) => {
+        return { updated: true, foo: ctx.foo };
+      },
+    });
+
+    await t.run(async (ctx) => {
+      const result = await (mutationFn as any)._handler(ctx, {});
+      expect(result).toEqual({ updated: true, foo: "bar" });
+
+      expect(finallyMock).toHaveBeenCalledWith(
+        expect.objectContaining({ foo: "bar" }),
+        { result: { updated: true, foo: "bar" } }
+      );
+    });
+  });
+
+  test("finally callback with action", async () => {
+    const t = convexTest(schema, modules);
+    const finallyMock = vi.fn();
+
+    const withFinally = customAction(action, {
+      args: {},
+      input: async () => ({ ctx: { foo: "bar" }, args: {} }),
+      finally: (ctx, params) => {
+        finallyMock(ctx, params);
+      }
+    });
+
+    const actionFn = withFinally({
+      args: {},
+      handler: async (ctx) => {
+        return { executed: true, foo: ctx.foo };
+      },
+    });
+
+    await t.run(async (ctx) => {
+      const result = await (actionFn as any)._handler(ctx, {});
+      expect(result).toEqual({ executed: true, foo: "bar" });
+
+      expect(finallyMock).toHaveBeenCalledWith(
+        expect.objectContaining({ foo: "bar" }),
+        { result: { executed: true, foo: "bar" } }
+      );
+    });
+  });
+});

packages/convex-helpers/server/customFunctions.ts

@@ -60,6 +60,10 @@ export type Mod<
   ) =>
     | Promise<{ ctx: ModCtx; args: ModMadeArgs }>
     | { ctx: ModCtx; args: ModMadeArgs };
+  finally?: (ctx: Ctx & ModCtx, params: {
+    result?: unknown;
+    error?: unknown;
+  }) => void | Promise<void>;
 };
 
 /**
@@ -88,6 +92,7 @@ export const NoOp = {
   input() {
     return { args: {}, ctx: {} };
   },
+
 };
 
 /**
@@ -339,7 +344,20 @@ function customFnBuilder(
             pick(allArgs, Object.keys(inputArgs)) as any,
           );
           const args = omit(allArgs, Object.keys(inputArgs));
-          return handler({ ...ctx, ...added.ctx }, { ...args, ...added.args });
+          const finalCtx = { ...ctx, ...added.ctx };
+          let result;
+          try {
+            result = await handler(finalCtx, { ...args, ...added.args });
+            if (mod.finally) {
+              await mod.finally(finalCtx, { result });
+            }
+            return result;
+          } catch (e) {
+            if (mod.finally) {
+              await mod.finally(finalCtx, { error: e });
+            }
+            throw e;
+          }
         },
       });
     }
@@ -353,7 +371,20 @@ function customFnBuilder(
       returns: fn.returns,
       handler: async (ctx: any, args: any) => {
         const added = await inputMod(ctx, args);
-        return handler({ ...ctx, ...added.ctx }, { ...args, ...added.args });
+        const finalCtx = { ...ctx, ...added.ctx };
+        let result;
+        try {
+          result = await handler(finalCtx, { ...args, ...added.args });
+          if (mod.finally) {
+            await mod.finally(finalCtx, { result });
+          }
+          return result;
+        } catch (e) {
+          if (mod.finally) {
+            await mod.finally(finalCtx, { error: e });
+          }
+          throw e;
+        }
       },
     });
   };