Add finally callback to customFunction

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

@@ -560,3 +560,69 @@ describe("nested custom functions", () => {
     ).rejects.toThrow("Validator error: Expected `string`");
   });
 });
+
+describe("finally callback", () => {
+  test("finally callback is called with result and context", async () => {
+    let finallyCalled = false;
+    let finallyParams = null;
+
+    const ctx = { foo: "bar" };
+    const args = { test: "value" };
+    const result = { success: true };
+
+    const handler = async (_ctx, _args) => result;
+
+    const mod = {
+      args: {},
+      input: async () => ({ ctx: {}, args: {} }),
+      finally: (params) => {
+        finallyCalled = true;
+        finallyParams = params;
+      },
+    };
+
+    let actualResult;
+    let actualError;
+    try {
+      actualResult = await handler(ctx, args);
+    } catch (e) {
+      actualError = e;
+      throw e;
+    } finally {
+      if (mod.finally) {
+        await mod.finally({ ctx, result: actualResult, error: actualError });
+      }
+    }
+
+    expect(finallyCalled).toBe(true);
+    expect(finallyParams).toEqual({
+      ctx,
+      result,
+      error: undefined,
+    });
+
+    finallyCalled = false;
+    finallyParams = null;
+
+    const testError = new Error("Test error");
+    const errorHandler = async () => {
+      throw testError;
+    };
+
+    try {
+      await errorHandler();
+    } catch (e) {
+    } finally {
+      if (mod.finally) {
+        await mod.finally({ ctx, result: undefined, error: testError });
+      }
+    }
+
+    expect(finallyCalled).toBe(true);
+    expect(finallyParams).toEqual({
+      ctx,
+      result: undefined,
+      error: testError,
+    });
+  });
+});

packages/convex-helpers/server/customFunctions.ts

@@ -60,6 +60,11 @@ export type Mod<
   ) =>
     | Promise<{ ctx: ModCtx; args: ModMadeArgs }>
     | { ctx: ModCtx; args: ModMadeArgs };
+  finally?: (params: {
+    ctx: Ctx & ModCtx;
+    result?: any;
+    error?: any;
+  }) => void | Promise<void>;
 };
 
 /**
@@ -88,6 +93,7 @@ export const NoOp = {
   input() {
     return { args: {}, ctx: {} };
   },
+  finally() {},
 };
 
 /**
@@ -339,7 +345,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;
+          let error;
+          try {
+            result = await handler(finalCtx, { ...args, ...added.args });
+            return result;
+          } catch (e) {
+            error = e;
+            throw e;
+          } finally {
+            if (mod.finally) {
+              await mod.finally({ ctx: finalCtx, result, error });
+            }
+          }
         },
       });
     }
@@ -353,7 +372,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;
+        let error;
+        try {
+          result = await handler(finalCtx, { ...args, ...added.args });
+          return result;
+        } catch (e) {
+          error = e;
+          throw e;
+        } finally {
+          if (mod.finally) {
+            await mod.finally({ ctx: finalCtx, result, error });
+          }
+        }
       },
     });
   };