fix customCtx args type

Author: ianmacartney

packages/convex-helpers/server/customFunctions.ts

@@ -76,10 +76,15 @@ import { omit, pick } from "../index.js";
  * has access to resources created during input processing via closure.
  */
 export type Customization<
+  // The ctx object from the original function.
   Ctx extends Record<string, any>,
+  // The validators for the args the customization function consumes.
   CustomArgsValidator extends PropertyValidators,
+  // The ctx object produced: a patch applied to the original ctx.
   CustomCtx extends Record<string, any>,
+  // The args produced by the customization function.
   CustomMadeArgs extends Record<string, any>,
+  // Extra args that are passed to the input function.
   ExtraArgs extends Record<string, any> = Record<string, any>,
 > = {
   args: CustomArgsValidator;
@@ -155,31 +160,30 @@ export function customCtxAndArgs<
   CustomMadeArgs,
   ExtraArgs
 > {
+  // This is already the right type. This function just helps you define it.
   return objectWithArgsAndInput;
 }
 
 /**
  * A helper for defining a Customization when your mod doesn't need to add or remove
  * anything from args.
- * @param mod A function that defines how to modify the ctx.
+ * @param modifyCtx A function that defines how to modify the ctx.
  * @returns A ctx delta to be applied to the original ctx.
  */
 export function customCtx<
   InCtx extends Record<string, any>,
   OutCtx extends Record<string, any>,
   ExtraArgs extends Record<string, any> = Record<string, any>,
 >(
-  mod: (original: InCtx, extra: ExtraArgs) => Promise<OutCtx> | OutCtx,
-): Customization<
-  InCtx,
-  Record<string, never>,
-  OutCtx,
-  Record<string, never>,
-  ExtraArgs
-> {
+  modifyCtx: (original: InCtx, extra: ExtraArgs) => Promise<OutCtx> | OutCtx,
+  // eslint-disable-next-line @typescript-eslint/no-empty-object-type
+): Customization<InCtx, {}, OutCtx, object, ExtraArgs> {
   return {
     args: {},
-    input: async (ctx, _, extra) => ({ ctx: await mod(ctx, extra), args: {} }),
+    input: async (ctx, _, extra) => ({
+      ctx: await modifyCtx(ctx, extra),
+      args: {},
+    }),
   };
 }
 

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

@@ -4,6 +4,7 @@ import type {
   ApiFromModules,
   RegisteredQuery,
   DefaultFunctionArgs,
+  FunctionReference,
 } from "convex/server";
 import { defineTable, defineSchema, queryGeneric, anyApi } from "convex/server";
 import type { Equals } from "../index.js";
@@ -361,6 +362,7 @@ const testApi: ApiFromModules<{
     addC: typeof addC;
     addCU: typeof addCU;
     addCU2: typeof addCU2;
+    addCtxWithExistingArg: typeof addCtxWithExistingArg;
     add: typeof add;
     addUnverified: typeof addUnverified;
     addUnverified2: typeof addUnverified2;
@@ -644,6 +646,32 @@ describe("zod functions", () => {
     });
   });
 
+  test("add ctx with existing arg", async () => {
+    const t = convexTest(schema, modules);
+    expect(
+      await t.query(testApi.addCtxWithExistingArg, { b: "foo" }),
+    ).toMatchObject({
+      ctxA: "hi",
+      argB: "foo",
+    });
+    expectTypeOf(testApi.addCtxWithExistingArg).toExtend<
+      FunctionReference<
+        "query",
+        "public",
+        { b: string },
+        { ctxA: string; argB: string }
+      >
+    >();
+    expectTypeOf<
+      FunctionReference<
+        "query",
+        "public",
+        { b: string },
+        { ctxA: string; argB: string }
+      >
+    >().toExtend<typeof testApi.addCtxWithExistingArg>();
+  });
+
   test("add args", async () => {
     const t = convexTest(schema, modules);
     expect(await t.query(testApi.add, {})).toMatchObject({