Add finally callback to customFunction

packages/convex-helpers/README.md

@@ -37,6 +37,8 @@ define custom behavior, allowing you to:
 - Consume arguments from the client that are not passed to the action, such
   as taking in an authentication parameter like an API key or session ID.
   These arguments must be sent up by the client along with each request.
+- Execute cleanup or logging logic after function execution using the `finally`
+  callback, which has access to the function's result or error.
 
 See the associated [Stack Post](https://stack.convex.dev/custom-functions)
 
@@ -50,7 +52,19 @@ const myQueryBuilder = customQuery(query, {
   input: async (ctx, args) => {
     const apiUser = await getApiUser(args.apiToken);
     const db = wrapDatabaseReader({ apiUser }, ctx.db, rlsRules);
-    return { ctx: { db, apiUser }, args: {} };
+    return {
+      ctx: { db, apiUser },
+      args: {},
+      finally: ({ result, error }) => {
+        // Optional callback that runs after the function executes
+        // Has access to resources created during input processing
+        if (error) {
+          console.error("Error in query:", error);
+        } else {
+          console.log("Query completed with result:", result);
+        }
+      },
+    };
   },
 });
 

packages/convex-helpers/package.json

@@ -1,6 +1,6 @@
 {
   "name": "convex-helpers",
-  "version": "0.1.80",
+  "version": "0.1.81-alpha.0",
   "description": "A collection of useful code to complement the official convex package.",
   "type": "module",
   "bin": {

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

@@ -36,6 +36,7 @@ import {
   describe,
   expect,
   test,
+  vi,
 } from "vitest";
 import { modules } from "./setup.test.js";
 
@@ -569,3 +570,126 @@ 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: (params) => {
+          finallyMock(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({
+        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({
+        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: (params) => {
+          finallyMock(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({
+        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: (params) => {
+          finallyMock(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({
+        result: { executed: true, foo: "bar" },
+      });
+    });
+  });
+});

packages/convex-helpers/server/customFunctions.ts

@@ -45,6 +45,12 @@ import { omit, pick } from "../index.js";
  * provided for the modified function. All returned ctx and args will show up
  * in the type signature for the modified function.
  * To remove something from `ctx`, you can return it as `undefined`.
+ *
+ * The `input` function can also return a `finally` callback that will be called
+ * after the function executes with either the result or error. This is useful for
+ * cleanup operations or logging that should happen regardless of whether the
+ * function succeeds or fails. The `finally` callback has access to resources
+ * created during input processing.
  */
 export type Mod<
   Ctx extends Record<string, any>,
@@ -57,8 +63,22 @@ export type Mod<
     ctx: Ctx,
     args: ObjectType<ModArgsValidator>,
   ) =>
-    | Promise<{ ctx: ModCtx; args: ModMadeArgs }>
-    | { ctx: ModCtx; args: ModMadeArgs };
+    | Promise<{
+        ctx: ModCtx;
+        args: ModMadeArgs;
+        finally?: (params: {
+          result?: unknown;
+          error?: unknown;
+        }) => void | Promise<void>;
+      }>
+    | {
+        ctx: ModCtx;
+        args: ModMadeArgs;
+        finally?: (params: {
+          result?: unknown;
+          error?: unknown;
+        }) => void | Promise<void>;
+      };
 };
 
 /**
@@ -101,7 +121,17 @@ export const NoOp = {
  *     const user = await getUserOrNull(ctx);
  *     const session = await db.get(sessionId);
  *     const db = wrapDatabaseReader({ user }, ctx.db, rlsRules);
- *     return { ctx: { db, user, session }, args: {} };
+ *     return {
+ *       ctx: { db, user, session },
+ *       args: {},
+ *       finally: ({ result, error }) => {
+ *         // Optional callback that runs after the function executes
+ *         // Has access to resources created during input processing
+ *         if (error) {
+ *           console.error("Error in query:", error);
+ *         }
+ *       }
+ *     };
  *   },
  * });
  *
@@ -173,7 +203,17 @@ export function customQuery<
  *     const user = await getUserOrNull(ctx);
  *     const session = await db.get(sessionId);
  *     const db = wrapDatabaseReader({ user }, ctx.db, rlsRules);
- *     return { ctx: { db, user, session }, args: {} };
+ *     return {
+ *       ctx: { db, user, session },
+ *       args: {},
+ *       finally: ({ result, error }) => {
+ *         // Optional callback that runs after the function executes
+ *         // Has access to resources created during input processing
+ *         if (error) {
+ *           console.error("Error in mutation:", error);
+ *         }
+ *       }
+ *     };
  *   },
  * });
  *
@@ -252,7 +292,21 @@ export function customMutation<
  *       throw new Error("Invalid secret key");
  *     }
  *     const user = await ctx.runQuery(internal.users.getUser, {});
- *     return { ctx: { user }, args: {} };
+ *     // Create resources that can be used in the finally callback
+ *     const logger = createLogger();
+ *     return {
+ *       ctx: { user },
+ *       args: {},
+ *       finally: ({ result, error }) => {
+ *         // Optional callback that runs after the function executes
+ *         // Has access to resources created during input processing
+ *         if (error) {
+ *           logger.error("Action failed:", error);
+ *         } else {
+ *           logger.info("Action completed successfully");
+ *         }
+ *       }
+ *     };
  *   },
  * });
  *
@@ -338,7 +392,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 (added.finally) {
+              await added.finally({ result });
+            }
+            return result;
+          } catch (e) {
+            if (added.finally) {
+              await added.finally({ error: e });
+            }
+            throw e;
+          }
         },
       });
     }
@@ -352,7 +419,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 (added.finally) {
+            await added.finally({ result });
+          }
+          return result;
+        } catch (e) {
+          if (added.finally) {
+            await added.finally({ error: e });
+          }
+          throw e;
+        }
       },
     });
   };