diff --git a/README.md b/README.md index 05fec59e..f30d5179 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,26 @@ export const emailSent = pool.defineOnComplete({ }); ``` +If you only need to react to terminal failures, use `onFailure` instead of +`onComplete`. It accepts the same handler and context. The enqueue call above +becomes: + +```ts +await pool.enqueueAction(ctx, internal.email.send, args, { + onFailure: internal.email.emailSent, + context: { emailType: args.emailType, userId: args.userId }, + retry: false, // don't retry this action, as we can't guarantee idempotency. +}); +``` + +`onFailure` runs after retries are exhausted, or immediately after a +`NonRetryableError`. It is skipped for successful jobs and jobs canceled through +the Workpool API. Canceling a scheduled function directly in the dashboard is +treated as a failure and can trigger retries and `onFailure`. Like `onComplete`, +it runs in a separate transaction from the work. A callback error does not retry +the original work. Use `onComplete` when you also need success or cancellation +notifications; specifying both options is an error. + ### Idempotency Idempotent actions are actions that can be run multiple times safely. This @@ -310,7 +330,11 @@ options include: set to `true`, it will use the `defaultRetryBehavior`. If it's set to a custom config, it will use that (and do retries). - `onComplete`: A mutation to run after the function finishes. -- `context`: Any data you want to pass to the `onComplete` mutation. +- `onFailure`: A mutation to run after a terminal failure, using the same + arguments as `onComplete`. Not called on success or cancellation. Cannot be + combined with `onComplete`. +- `context`: Any data you want to pass to the `onComplete` or `onFailure` + mutation. - `runAt` and `runAfter`: Similar to `ctx.scheduler.run*`, allows you to schedule the work to run later. By default it's immediate. diff --git a/src/client/index.test.ts b/src/client/index.test.ts index c6bda3c3..b4a519e2 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -1,18 +1,458 @@ import { + anyApi, + componentsGeneric, + defineSchema, + defineTable, + internalActionGeneric, + internalMutationGeneric, + internalQueryGeneric, makeFunctionReference, + type ApiFromModules, + type DataModelFromSchemaDefinition, + type FunctionArgs, type GenericDataModel, type GenericMutationCtx, } from "convex/server"; -import { type Infer, v } from "convex/values"; -import { expectTypeOf, test } from "vitest"; +import { type Infer, type ObjectType, v } from "convex/values"; +import { convexTest } from "convex-test"; +import { + afterEach, + beforeEach, + describe, + expect, + expectTypeOf, + test, + vi, +} from "vitest"; +import workpool from "../test.js"; +import type { api } from "../component/_generated/api.js"; import { type EnqueueOptions, + NonRetryableError, type OnCompleteArgs, + type RetryOption, + Workpool, type WorkId, - type Workpool, + type WorkpoolComponent, vOnCompleteArgs, + vResult, + vWorkId, } from "./index.js"; +const schema = defineSchema({ + events: defineTable({ + key: v.string(), + kind: v.union( + v.literal("attempt"), + v.literal("work"), + v.literal("callback"), + ), + workId: v.optional(vWorkId), + result: v.optional(vResult), + context: v.optional(v.any()), + }).index("key", ["key"]), +}); +const runArgs = { + key: v.string(), + fail: v.optional(v.boolean()), + failures: v.optional(v.number()), + nonRetryable: v.optional(v.boolean()), + payload: v.optional(v.string()), +}; +type RunArgs = ObjectType; +type Context = { key: string; padding?: string; throw?: boolean }; + +// Observe real callback execution while retaining Convex validation and rollback. +const callback = vi.fn( + async ( + ctx: GenericMutationCtx>, + args: OnCompleteArgs & { workId: WorkId }, + ) => { + await ctx.db.insert("events", { + key: args.context?.key ?? "default", + kind: "callback", + workId: args.workId, + context: args.context, + result: args.result, + }); + if (args.context?.throw) { + throw new Error("callback failed"); + } + return null; + }, +); +const fixtures = { + attempt: internalMutationGeneric({ + args: { key: v.string() }, + returns: v.number(), + handler: async (ctx, { key }) => { + await ctx.db.insert("events", { key, kind: "attempt" }); + const events = await ctx.db + .query("events") + .withIndex("key", (q) => q.eq("key", key)) + .collect(); + return events.filter((e) => e.kind === "attempt").length; + }, + }), + action: internalActionGeneric({ + args: runArgs, + returns: v.string(), + handler: async (ctx, args): Promise => { + const attempt = await ctx.runMutation(refs.attempt, { key: args.key }); + if (args.fail || attempt <= (args.failures ?? 0)) { + if (args.nonRetryable) throw new NonRetryableError("work failed"); + throw new Error("work failed"); + } + return "action result"; + }, + }), + mutation: internalMutationGeneric({ + args: runArgs, + returns: v.string(), + handler: async (ctx, args) => { + await ctx.db.insert("events", { key: args.key, kind: "work" }); + if (args.fail) throw new Error("work failed"); + return "mutation result"; + }, + }), + query: internalQueryGeneric({ + args: runArgs, + returns: v.string(), + handler: async (_ctx, args) => { + if (args.fail) throw new Error("work failed"); + return "query result"; + }, + }), + complete: internalMutationGeneric({ + args: vOnCompleteArgs( + v.object({ + key: v.string(), + padding: v.optional(v.string()), + throw: v.optional(v.boolean()), + }), + ), + returns: v.null(), + handler: callback, + }), + optionalComplete: internalMutationGeneric({ + args: vOnCompleteArgs(), + returns: v.null(), + handler: callback, + }), +}; +const refs = ( + anyApi as unknown as ApiFromModules<{ callbacks: typeof fixtures }> +).callbacks; +// An application with real callback functions and a nested workpool component. +// The generated path provides the module root expected by convex-test. +const modules = { + "./_generated/server.ts": async () => ({}), + "./callbacks.ts": async () => fixtures, +}; +const component = componentsGeneric().workpool as unknown as WorkpoolComponent; +const pool = new Workpool(component, { maxParallelism: 3, logLevel: "ERROR" }); +const retries = { maxAttempts: 3, initialBackoffMs: 1, base: 2 }; +type Kind = "action" | "mutation" | "query"; +type Mode = "onComplete" | "onFailure"; + +describe("completion callbacks through the client and scheduler", () => { + let t: ReturnType; + function setup() { + const t = convexTest(schema, modules); + workpool.register(t); + return t; + } + beforeEach(() => { + vi.useFakeTimers(); + callback.mockClear(); + t = setup(); + }); + afterEach(async () => { + await t.finishAllScheduledFunctions(vi.runAllTimers); + vi.useRealTimers(); + }); + const drain = () => t.finishAllScheduledFunctions(vi.runAllTimers); + const events = (key: string) => + t.query((ctx) => + ctx.db + .query("events") + .withIndex("key", (q) => q.eq("key", key)) + .collect(), + ); + function enqueue( + kind: Kind, + args: RunArgs, + options: { + mode?: Mode; + context?: Context; + retry?: RetryOption["retry"]; + runAt?: number; + } = {}, + ) { + const { + mode = "onFailure", + context = { key: args.key }, + ...rest + } = options; + const opts = + mode === "onFailure" + ? { onFailure: refs.complete, context, ...rest } + : { onComplete: refs.complete, context, ...rest }; + return t.mutation((ctx) => { + switch (kind) { + case "action": + return pool.enqueueAction(ctx, refs.action, args, opts); + case "mutation": + return pool.enqueueMutation(ctx, refs.mutation, args, opts); + case "query": + return pool.enqueueQuery(ctx, refs.query, args, opts); + } + }); + } + + test.each(["action", "mutation", "query"] as const)( + "onFailure skips successful %s work", + async (kind) => { + const id = await enqueue(kind, { key: kind }); + await drain(); + expect(callback).not.toHaveBeenCalled(); + expect(await t.query((ctx) => pool.status(ctx, id))).toEqual({ + state: "finished", + }); + }, + ); + + test.each(["action", "mutation", "query"] as const)( + "onFailure reuses an existing completion handler for failed %s work", + async (kind) => { + const id = await enqueue(kind, { key: kind, fail: true }); + await drain(); + expect(callback).toHaveBeenCalledTimes(1); + expect( + (await events(kind)).filter((e) => e.kind === "callback"), + ).toMatchObject([ + { + workId: id, + context: { key: kind }, + result: { kind: "failed", error: "work failed" }, + }, + ]); + expect(await t.query((ctx) => pool.status(ctx, id))).toEqual({ + state: "finished", + }); + if (kind === "mutation") { + expect((await events(kind)).some((e) => e.kind === "work")).toBe(false); + } + }, + ); + + test("onFailure runs once after retry exhaustion", async () => { + await enqueue("action", { key: "retry", fail: true }, { retry: retries }); + await drain(); + expect(callback).toHaveBeenCalledTimes(1); + expect((await events("retry")).map((e) => e.kind)).toEqual([ + "attempt", + "attempt", + "attempt", + "callback", + ]); + }); + + test("onFailure is skipped if a retry succeeds", async () => { + await enqueue("action", { key: "retry", failures: 2 }, { retry: retries }); + await drain(); + expect(callback).not.toHaveBeenCalled(); + expect((await events("retry")).map((e) => e.kind)).toEqual([ + "attempt", + "attempt", + "attempt", + ]); + }); + + test("NonRetryableError skips remaining attempts and invokes onFailure", async () => { + await enqueue( + "action", + { key: "terminal", fail: true, nonRetryable: true }, + { retry: retries }, + ); + await drain(); + expect(callback).toHaveBeenCalledTimes(1); + expect((await events("terminal")).map((e) => e.kind)).toEqual([ + "attempt", + "callback", + ]); + expect((await events("terminal"))[1].result).toEqual({ + kind: "failed", + error: "work failed", + }); + }); + + test.each(["onComplete", "onFailure"] as const)( + "%s preserves cancellation semantics", + async (mode) => { + const id = await enqueue( + "action", + { key: "cancel", fail: true }, + { mode, runAt: Date.now() + 60_000 }, + ); + await t.mutation((ctx) => pool.cancel(ctx, id)); + await drain(); + const recorded = await events("cancel"); + if (mode === "onFailure") { + expect(callback).not.toHaveBeenCalled(); + expect(recorded).toEqual([]); + } else { + expect(callback).toHaveBeenCalledTimes(1); + expect(recorded).toMatchObject([ + { workId: id, kind: "callback", result: { kind: "canceled" } }, + ]); + } + expect(await t.query((ctx) => pool.status(ctx, id))).toEqual({ + state: "finished", + }); + }, + ); + + test.each(["action", "mutation", "query"] as const)( + "onComplete still receives successful %s results", + async (kind) => { + const id = await enqueue(kind, { key: kind }, { mode: "onComplete" }); + await drain(); + expect(callback).toHaveBeenCalledTimes(1); + expect( + (await events(kind)).filter((e) => e.kind === "callback"), + ).toMatchObject([ + { + workId: id, + result: { kind: "success", returnValue: `${kind} result` }, + }, + ]); + }, + ); + + test.each([false, true])( + "onFailure restores large context (large arguments: %s)", + async (largeArgs) => { + const context = { key: "large", padding: "x".repeat(12_000) }; + const id = await enqueue( + "action", + { + key: "large", + fail: true, + ...(largeArgs ? { payload: "x".repeat(12_000) } : {}), + }, + { context }, + ); + await drain(); + expect(callback).toHaveBeenCalledTimes(1); + const recorded = await events("large"); + expect(recorded.map((e) => e.kind)).toEqual(["attempt", "callback"]); + expect(recorded[1]).toMatchObject({ + context, + result: { kind: "failed", error: "work failed" }, + }); + expect(await t.query((ctx) => pool.status(ctx, id))).toEqual({ + state: "finished", + }); + }, + ); + + test("onFailure supports omitted context", async () => { + const id = await t.mutation((ctx) => + pool.enqueueMutation( + ctx, + refs.mutation, + { key: "default", fail: true }, + { onFailure: refs.optionalComplete }, + ), + ); + await drain(); + expect(callback).toHaveBeenCalledTimes(1); + expect((await events("default"))[0]).toMatchObject({ + workId: id, + result: { kind: "failed" }, + }); + expect((await events("default"))[0].context).toBeUndefined(); + }); + + test.each(["action", "mutation", "query"] as const)( + "batch %s enqueue only calls back for failed items", + async (kind) => { + const args = [{ key: "success" }, { key: "failure", fail: true }]; + const opts = { onFailure: refs.complete, context: { key: "batch" } }; + const ids = await t.mutation((ctx) => { + switch (kind) { + case "action": + return pool.enqueueActionBatch(ctx, refs.action, args, opts); + case "mutation": + return pool.enqueueMutationBatch(ctx, refs.mutation, args, opts); + case "query": + return pool.enqueueQueryBatch(ctx, refs.query, args, opts); + } + }); + await drain(); + expect(callback).toHaveBeenCalledTimes(1); + expect(await events("batch")).toMatchObject([ + { workId: ids[1], result: { kind: "failed" } }, + ]); + expect(await t.query((ctx) => pool.statusBatch(ctx, ids))).toEqual([ + { state: "finished" }, + { state: "finished" }, + ]); + }, + ); + + test("a failing onFailure callback rolls back without retrying the original work", async () => { + const id = await enqueue( + "action", + { key: "broken", fail: true }, + { retry: retries, context: { key: "broken", throw: true } }, + ); + await drain(); + expect(callback).toHaveBeenCalledTimes(1); + expect((await events("broken")).map((e) => e.kind)).toEqual([ + "attempt", + "attempt", + "attempt", + ]); + expect(await t.query((ctx) => pool.status(ctx, id))).toEqual({ + state: "finished", + }); + }); + + test.each([false, true])( + "rejects mixed completion options at runtime (batch: %s)", + async (batch) => { + // JavaScript callers can bypass TypeScript's mutually exclusive options. + const opts = { + onComplete: refs.complete, + onFailure: refs.complete, + context: { key: "invalid" }, + } as unknown as EnqueueOptions; + await expect( + t.mutation((ctx) => + batch + ? pool.enqueueMutationBatch( + ctx, + refs.mutation, + [{ key: "invalid" }], + opts, + ) + : pool.enqueueMutation( + ctx, + refs.mutation, + { key: "invalid" }, + opts, + ), + ), + ).rejects.toThrow("Cannot define both onComplete and onFailure"); + await drain(); + expect(await events("invalid")).toEqual([]); + expect(callback).not.toHaveBeenCalled(); + }, + ); +}); + type MutationCtx = GenericMutationCtx; const action = makeFunctionReference<"action", { value: number }, number>( "work:action", @@ -67,3 +507,66 @@ test("runAt and runAfter are mutually exclusive", () => { const options: EnqueueOptions = { runAt: 1, runAfter: 2 }; expectTypeOf(options).toMatchTypeOf(); }); + +test("onFailure accepts existing onComplete handlers on all enqueue methods", () => { + const options = { onFailure: onComplete, context: { label: "job" } }; + const mutation = makeFunctionReference<"mutation", { value: number }, number>( + "work:mutation", + ); + const query = makeFunctionReference<"query", { value: number }, number>( + "work:query", + ); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => + pool.enqueueAction(ctx, action, { value: 1 }, options), + ).returns.toEqualTypeOf>(); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => + pool.enqueueMutation(ctx, mutation, { value: 1 }, options), + ).returns.toEqualTypeOf>(); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => + pool.enqueueQuery(ctx, query, { value: 1 }, options), + ).returns.toEqualTypeOf>(); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => + pool.enqueueActionBatch(ctx, action, [{ value: 1 }], options), + ).returns.toEqualTypeOf>(); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => + pool.enqueueMutationBatch(ctx, mutation, [{ value: 1 }], options), + ).returns.toEqualTypeOf>(); + expectTypeOf((pool: Workpool, ctx: MutationCtx) => + pool.enqueueQueryBatch(ctx, query, [{ value: 1 }], options), + ).returns.toEqualTypeOf>(); +}); + +test("onFailure preserves context typing", () => { + expectTypeOf((pool: Workpool, ctx: MutationCtx) => { + void pool.enqueueAction( + ctx, + action, + { value: 1 }, + { + onFailure: onComplete, + // @ts-expect-error The reused completion handler requires a string label. + context: { label: 1 }, + }, + ); + }).returns.toBeVoid(); +}); + +test("onComplete and onFailure are mutually exclusive", () => { + // @ts-expect-error The completion options cannot be combined. + const options: EnqueueOptions = { onComplete, onFailure: onComplete }; + expectTypeOf(options).toMatchTypeOf(); +}); + +test("callback options can each be disabled with null", () => { + expectTypeOf({ onComplete: null }).toMatchTypeOf(); + expectTypeOf({ onFailure: null }).toMatchTypeOf(); +}); + +test("generated enqueue types match the component validators", () => { + expectTypeOf< + FunctionArgs + >().toEqualTypeOf>(); + expectTypeOf< + FunctionArgs + >().toEqualTypeOf>(); +}); diff --git a/src/client/index.ts b/src/client/index.ts index 989a67df..44f397c9 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -462,61 +462,83 @@ export type EnqueueOptions = { */ name?: string; /** - * A mutation to run after the function succeeds, fails, or is canceled. - * The context type is for your use, feel free to provide a validator for it. - * e.g. - * ```ts - * export const completion = workpool.defineOnComplete({ - * context: v.string(), - * handler: async (ctx, {workId, context, result}) => { - * // context has been validated as a string - * // ... do something with the result - * }, - * }); - * ``` - * or more manually: - * ```ts - * export const completion = internalMutation({ - * args: vOnCompleteArgs(v.string()), - * handler: async (ctx, args) => { - * console.log(args.result, "Got Context back -> ", args.context, Date.now() - args.context); - * }, - * }); - * ``` - */ - onComplete?: FunctionReference< - "mutation", - FunctionVisibility, - OnCompleteArgs - > | null; - - /** - * A context object to pass to the `onComplete` mutation. + * A context object to pass to the `onComplete` or `onFailure` mutation. * Useful for passing data from the enqueue site to the onComplete site. */ context?: Context; } & ( | { /** - * The time (ms since epoch) to run the action at. - * If not provided, the action will be run as soon as possible. - * Note: this is advisory only. It may run later. + * A mutation to run after the function succeeds, fails, or is canceled. + * Cannot be used together with `onFailure`. + * The context type is for your use, feel free to provide a validator for it. + * e.g. + * ```ts + * export const completion = workpool.defineOnComplete({ + * context: v.string(), + * handler: async (ctx, {workId, context, result}) => { + * // context has been validated as a string + * // ... do something with the result + * }, + * }); + * ``` + * or more manually: + * ```ts + * export const completion = internalMutation({ + * args: vOnCompleteArgs(v.string()), + * handler: async (ctx, args) => { + * console.log(args.result, "Got Context back -> ", args.context, Date.now() - args.context); + * }, + * }); + * ``` */ - runAt?: number; - - runAfter?: never; + onComplete?: FunctionReference< + "mutation", + FunctionVisibility, + OnCompleteArgs + > | null; + onFailure?: never; } | { /** - * The number of milliseconds to run the action after. - * If not provided, the action will be run as soon as possible. - * Note: this is advisory only. It may run later. + * A mutation to run when the function fails with no retries remaining, + * or throws a NonRetryableError. Not called on success or cancellation + * through the Workpool API. Recovery treats direct scheduler cancellation + * as a failure, which can trigger retries and this callback. + * Uses the same arguments and separate transaction as `onComplete`, + * so an existing `onComplete` handler can be passed here unchanged. + * Cannot be used together with `onComplete`. */ - runAfter?: number; - - runAt?: never; + onFailure?: FunctionReference< + "mutation", + FunctionVisibility, + OnCompleteArgs + > | null; + onComplete?: never; } -); +) & + ( + | { + /** + * The time (ms since epoch) to run the action at. + * If not provided, the action will be run as soon as possible. + * Note: this is advisory only. It may run later. + */ + runAt?: number; + + runAfter?: never; + } + | { + /** + * The number of milliseconds to run the action after. + * If not provided, the action will be run as soon as possible. + * Note: this is advisory only. It may run later. + */ + runAfter?: number; + + runAt?: never; + } + ); export type OnCompleteArgs = { /** @@ -566,6 +588,9 @@ async function enqueueArgs( Partial & { retryBehavior?: RetryBehavior }) | undefined, ) { + if (opts?.onComplete && opts.onFailure) { + throw new Error("Cannot define both onComplete and onFailure handlers."); + } const [fnHandle, fnName] = typeof fn === "string" && fn.startsWith("function://") ? [fn, opts?.name ?? fn] @@ -578,7 +603,14 @@ async function enqueueArgs( fnHandle: await createFunctionHandle(opts.onComplete), context: opts.context, } - : undefined, + : opts?.onFailure + ? { + onStatusHandle: { + failed: await createFunctionHandle(opts.onFailure), + }, + context: opts.context, + } + : undefined, runAt: getRunAt(opts), retryBehavior: opts?.retryBehavior, config: { diff --git a/src/component/_generated/component.ts b/src/component/_generated/component.ts index 2ea8267f..76f1c7ef 100644 --- a/src/component/_generated/component.ts +++ b/src/component/_generated/component.ts @@ -69,7 +69,9 @@ export type ComponentApi = fnHandle: string; fnName: string; fnType: "action" | "mutation" | "query"; - onComplete?: { context?: any; fnHandle: string }; + onComplete?: + | { context?: any; fnHandle: string } + | { context?: any; onStatusHandle: { failed: string } }; retryBehavior?: { base: number; initialBackoffMs: number; @@ -93,7 +95,9 @@ export type ComponentApi = fnHandle: string; fnName: string; fnType: "action" | "mutation" | "query"; - onComplete?: { context?: any; fnHandle: string }; + onComplete?: + | { context?: any; fnHandle: string } + | { context?: any; onStatusHandle: { failed: string } }; retryBehavior?: { base: number; initialBackoffMs: number; diff --git a/src/component/complete.ts b/src/component/complete.ts index da03f26e..a4fe2759 100644 --- a/src/component/complete.ts +++ b/src/component/complete.ts @@ -126,7 +126,15 @@ export async function completeHandler( work.attempts < maxAttempts; if (!retry) { let scheduledId = undefined; - if (work.onComplete) { + // Workers and recovery both dispatch terminal callbacks here. + const fnHandle = + work.onComplete && + ("fnHandle" in work.onComplete + ? work.onComplete.fnHandle + : job.runResult.kind === "failed" + ? work.onComplete.onStatusHandle.failed + : undefined); + if (work.onComplete && fnHandle) { try { // Retrieve large context if stored separately let context = work.onComplete.context; @@ -137,7 +145,7 @@ export async function completeHandler( } } - const handle = work.onComplete.fnHandle as FunctionHandle< + const handle = fnHandle as FunctionHandle< "mutation", OnCompleteArgs, void diff --git a/src/component/lib.test.ts b/src/component/lib.test.ts index d09aa10a..459b4391 100644 --- a/src/component/lib.test.ts +++ b/src/component/lib.test.ts @@ -37,6 +37,67 @@ describe("lib", () => { }); describe("enqueue", () => { + it.each([false, true])( + "accepts legacy and failure callback handles (batch: %s)", + async (batch) => { + const callbacks = [ + { fnHandle: "completeHandle", context: { key: "complete" } }, + { + onStatusHandle: { failed: "failureHandle" }, + context: { key: "failure" }, + }, + ]; + const items = callbacks.map((onComplete) => ({ + fnHandle: "workHandle", + fnName: "work", + fnArgs: {}, + fnType: "mutation" as const, + runAt: Date.now(), + onComplete, + })); + const config = { logLevel: "WARN" as const }; + const ids = batch + ? await t.mutation(api.lib.enqueueBatch, { items, config }) + : await Promise.all( + items.map((item) => + t.mutation(api.lib.enqueue, { ...item, config }), + ), + ); + const stored = await t.run(async (ctx) => + Promise.all(ids.map((id) => ctx.db.get("work", id))), + ); + expect(stored.map((work) => work?.onComplete)).toEqual(callbacks); + }, + ); + + it.each([false, true])( + "rejects mixed callback handles at the component boundary (batch: %s)", + async (batch) => { + // Bypass the client guard to exercise the component's own validator. + const item = { + fnHandle: "workHandle", + fnName: "work", + fnArgs: {}, + fnType: "mutation" as const, + runAt: Date.now(), + onComplete: { + fnHandle: "completeHandle", + onStatusHandle: { failed: "failureHandle" }, + }, + }; + await expect( + batch + ? t.mutation(api.lib.enqueueBatch, { items: [item], config: {} }) + : t.mutation(api.lib.enqueue, { ...item, config: {} }), + ).rejects.toThrow(/Validator error/); + await t.run(async (ctx) => { + expect(await ctx.db.query("work").collect()).toEqual([]); + expect(await ctx.db.query("pendingStart").collect()).toEqual([]); + expect(await ctx.db.query("globals").collect()).toEqual([]); + }); + }, + ); + it("should successfully enqueue a work item", async () => { const id = await t.mutation(api.lib.enqueue, { fnHandle: "testHandle", diff --git a/src/component/recovery.test.ts b/src/component/recovery.test.ts index 32f97005..fb033be3 100644 --- a/src/component/recovery.test.ts +++ b/src/component/recovery.test.ts @@ -1,4 +1,6 @@ import { convexTest } from "convex-test"; +import { createFunctionHandle, makeFunctionReference } from "convex/server"; +import { v } from "convex/values"; import type { DocumentByName, GenericDatabaseReader, @@ -18,15 +20,32 @@ import { } from "vitest"; import { internal } from "./_generated/api.js"; import type { Doc, Id } from "./_generated/dataModel.js"; -import type { MutationCtx } from "./_generated/server.js"; +import { internalMutation, type MutationCtx } from "./_generated/server.js"; import batchWorker from "@convex-dev/batch-worker/test"; import { recoveryHandler } from "./recovery.js"; import schema from "./schema.js"; import { modules } from "./setup.test.js"; +import { type OnCompleteArgs, vResult } from "./shared.js"; describe("recovery", () => { + const callback = vi.fn<(args: OnCompleteArgs) => void>(); + const callbackRef = makeFunctionReference<"mutation", OnCompleteArgs>( + "recoveryCallbacks:complete", + ); async function setupTest() { - const t = convexTest(schema, modules); + const t = convexTest(schema, { + ...modules, + "./recoveryCallbacks.ts": async () => ({ + complete: internalMutation({ + args: { + workId: v.id("work"), + context: v.any(), + result: vResult, + }, + handler: async (_ctx, args) => callback(args), + }), + }), + }); batchWorker.register(t); return t; } @@ -64,6 +83,7 @@ describe("recovery", () => { beforeEach(async () => { vi.useFakeTimers(); + callback.mockClear(); t = await setupTest(); // Set up globals for logging @@ -79,6 +99,65 @@ describe("recovery", () => { vi.useRealTimers(); }); + describe.each(["onComplete", "onFailure"] as const)( + "%s callbacks", + (mode) => { + it.each([ + ["missing", "Scheduled job not found"], + ["failed", "Function execution failed"], + ["canceled", "Canceled via scheduler"], + ] as const)("runs once for a %s scheduled job", async (state, error) => { + const job = await t.run(async (ctx) => { + const handle = await createFunctionHandle(callbackRef); + const workId = await makeDummyWork(ctx, { + onComplete: + mode === "onComplete" + ? { fnHandle: handle, context: { key: state } } + : { + onStatusHandle: { failed: handle }, + context: { key: state }, + }, + }); + const scheduledId = await makeDummyScheduledFunction(ctx, workId); + // Prevent the placeholder worker from running when callbacks are drained. + await ctx.scheduler.cancel(scheduledId); + return { workId, scheduledId, attempt: 0, started: Date.now() }; + }); + await t.run(async (ctx) => { + const scheduled = await ctx.db.system.get( + "_scheduled_functions", + job.scheduledId, + ); + assert(scheduled); + // Only emulate the scheduler state; recovery and callback dispatch run + // normally, including scheduling and executing the callback mutation. + ctx.db.system.get = patchedSystemGet(ctx.db, { + [job.scheduledId]: + state === "missing" + ? null + : { + ...scheduled, + state: + state === "failed" + ? { kind: state, error } + : { kind: state }, + }, + }); + await recoveryHandler(ctx, { jobs: [job] }); + }); + // A repeated recovery scan must not dispatch the callback again. + await t.mutation(internal.recovery.recover, { jobs: [job] }); + await t.finishAllScheduledFunctions(vi.runAllTimers); + expect(callback).toHaveBeenCalledExactlyOnceWith({ + workId: job.workId, + context: { key: state }, + result: { kind: "failed", error }, + }); + expect(await t.run((ctx) => ctx.db.get("work", job.workId))).toBeNull(); + }); + }, + ); + describe("recover", () => { it("should skip jobs that already have a pendingCompletion", async () => { // Create work and scheduled function diff --git a/src/component/shared.ts b/src/component/shared.ts index 371db8bb..834a0f65 100644 --- a/src/component/shared.ts +++ b/src/component/shared.ts @@ -109,10 +109,18 @@ export type RunResult = kind: "canceled"; }; -export const vOnCompleteFnContext = v.object({ - fnHandle: v.string(), // mutation - context: v.optional(v.any()), -}); +// Keep the legacy handle shape for existing clients and queued work. The union +// rejects combining an all-outcomes handler with handlers for specific outcomes. +export const vOnCompleteFnContext = v.union( + v.object({ + fnHandle: v.string(), // mutation + context: v.optional(v.any()), + }), + v.object({ + onStatusHandle: v.object({ failed: v.string() }), // mutations + context: v.optional(v.any()), + }), +); export type OnCompleteArgs = { /**