Enforce v.function() on the work table

src/component/complete.ts

@@ -1,16 +1,10 @@
-import type { FunctionHandle } from "convex/server";
 import { getConvexSize, type Infer, v } from "convex/values";
 import type { Id } from "./_generated/dataModel.js";
 import { internal } from "./_generated/api.js";
 import { internalMutation, type MutationCtx } from "./_generated/server.js";
 import { kickMainLoop } from "./kick.js";
 import { createLogger } from "./logging.js";
-import {
-  getCurrentSegment,
-  type OnCompleteArgs,
-  type RunResult,
-  vResult,
-} from "./shared.js";
+import { getCurrentSegment, type RunResult, vResult } from "./shared.js";
 import { recordCompleted } from "./stats.js";
 import { assert } from "convex-helpers";
 
@@ -140,11 +134,7 @@ export async function completeHandler(
               }
             }
 
-            const handle = work.onComplete.fnHandle as FunctionHandle<
-              "mutation",
-              OnCompleteArgs,
-              void
-            >;
+            const handle = work.onComplete.fnHandle;
             const onCompleteArgs = {
               workId: work._id,
               context,

src/component/schema.ts

@@ -47,27 +47,19 @@ export default defineSchema({
   // Written on enqueue. Deleted by `complete` for success, failure, canceled.
   work: defineTable({
     fnType,
-    fnHandle: v.string(),
+    fnHandle: v.function(),
     fnName: v.string(),
     fnArgs: v.optional(v.any()),
     // Reference to large args/onComplete context if stored separately
     payloadId: v.optional(v.id("payload")),
     payloadSize: v.optional(v.number()),
     attempts: v.number(), // number of completed attempts
-    onComplete: v.optional(
-      v.object({ ...vOnCompleteFnContext.fields, fnHandle: v.string() }),
-    ),
+    onComplete: v.optional(vOnCompleteFnContext),
     retryBehavior: v.optional(retryBehavior),
     canceled: v.optional(v.boolean()),
     // May point to a deleted entry after starting, or be absent on older work.
     pendingStartId: v.optional(v.id("pendingStart")),
-  })
-    // Validates existing handles in the background. A later release enforces
-    // these, which would otherwise check every work document during the push.
-    .staged({
-      fnHandle: v.function(),
-      onComplete: v.optional(vOnCompleteFnContext),
-    }),
+  }),
 
   // Written on enqueue & rescheduled for retry, read & deleted by `main`.
   pendingStart: defineTable({

src/component/worker.ts

@@ -3,7 +3,7 @@
  * Should not touch any of loop's tables other than writing to `pendingCompletion`.
  * It is not responsible for handling retries.
  */
-import type { FunctionHandle } from "convex/server";
+import type { FunctionHandle_future } from "convex/server";
 import { v } from "convex/values";
 import { internal } from "./_generated/api.js";
 import type { Id } from "./_generated/dataModel.js";
@@ -21,7 +21,7 @@ import { assert } from "convex-helpers";
 
 const commonRunArgs = {
   workId: v.id("work"),
-  fnHandle: v.string(),
+  fnHandle: v.function(),
   fnArgs: v.optional(v.record(v.string(), v.any())),
   payloadId: v.optional(v.id("payload")),
   attempt: v.number(),
@@ -35,7 +35,7 @@ const actionOrQueryRunArgs = {
 export const runMutationWrapper = internalMutation({
   args: {
     workId: v.id("work"),
-    fnHandle: v.string(),
+    fnHandle: v.function(),
     payloadId: v.optional(v.id("payload")),
     fnArgs: v.optional(v.record(v.string(), v.any())),
     fnType: v.union(v.literal("query"), v.literal("mutation")),
@@ -55,8 +55,8 @@ export const runMutationWrapper = internalMutation({
 
     try {
       const returnValue = await (args.fnType === "query"
-        ? ctx.runQuery(args.fnHandle as FunctionHandle<"query">, fnArgs)
-        : ctx.runMutation(args.fnHandle as FunctionHandle<"mutation">, fnArgs));
+        ? ctx.runQuery(args.fnHandle, fnArgs)
+        : ctx.runMutation(args.fnHandle, fnArgs));
       // NOTE: we could run the `saveResult` handler here, or call `ctx.runMutation`,
       // but we want the mutation to be a separate transaction to reduce the window for OCCs.
       await ctx.scheduler.runAfter(0, internal.complete.complete, {
@@ -162,7 +162,7 @@ async function runOne(
   ctx: ActionCtx,
   console: Logger,
   args: {
-    fnHandle: string;
+    fnHandle: FunctionHandle_future<any>;
     fnArgs?: Record<string, unknown>;
     payloadId?: Id<"payload">;
     fnType: "action" | "query";
@@ -185,8 +185,8 @@ async function runOne(
     }
     const returnValue =
       args.fnType === "action"
-        ? await ctx.runAction(args.fnHandle as FunctionHandle<"action">, fnArgs)
-        : await ctx.runQuery(args.fnHandle as FunctionHandle<"query">, fnArgs);
+        ? await ctx.runAction(args.fnHandle, fnArgs)
+        : await ctx.runQuery(args.fnHandle, fnArgs);
     const runResult: RunResult = { kind: "success", returnValue };
     return { runResult, runOnCompleteInline: true };
   } catch (e: unknown) {