diff --git a/packages/core/src/v3/workers/taskExecutor.ts b/packages/core/src/v3/workers/taskExecutor.ts index 21ec33bc2..0fe63f0a6 100644 --- a/packages/core/src/v3/workers/taskExecutor.ts +++ b/packages/core/src/v3/workers/taskExecutor.ts @@ -174,6 +174,8 @@ export class TaskExecutor { ); } + await this.#cleanupAndWaitUntil(payload, ctx, initOutput, signal); + return { id: execution.run.id, ok: false, @@ -191,6 +193,8 @@ export class TaskExecutor { if (outputError) { recordSpanException(span, outputError); + await this.#cleanupAndWaitUntil(payload, ctx, initOutput, signal); + return this.#internalErrorResult( execution, TaskRunErrorCodes.TASK_OUTPUT_ERROR, @@ -208,6 +212,8 @@ export class TaskExecutor { if (exportError) { recordSpanException(span, exportError); + await this.#cleanupAndWaitUntil(payload, ctx, initOutput, signal); + return this.#internalErrorResult( execution, TaskRunErrorCodes.TASK_OUTPUT_ERROR, @@ -236,6 +242,8 @@ export class TaskExecutor { signal ); + await this.#cleanupAndWaitUntil(payload, ctx, initOutput, signal); + return { ok: true, id: execution.run.id, @@ -690,21 +698,90 @@ export class TaskExecutor { ); } - async #callTaskCleanup( + async #cleanupAndWaitUntil( payload: unknown, ctx: TaskRunContext, - init: unknown, + initOutput: any, signal?: AbortSignal ) { - const cleanupFn = this.task.fns.cleanup; + await this.#callCleanupFunctions(payload, ctx, initOutput, signal); + await this.#blockForWaitUntil(); + } - if (!cleanupFn) { + async #callCleanupFunctions( + payload: unknown, + ctx: TaskRunContext, + initOutput: any, + signal?: AbortSignal + ) { + const globalCleanupHooks = lifecycleHooks.getGlobalCleanupHooks(); + const taskCleanupHook = lifecycleHooks.getTaskCleanupHook(this.task.id); + + if (globalCleanupHooks.length === 0 && !taskCleanupHook) { return; } - return this._tracer.startActiveSpan("cleanup", async (span) => { - return await cleanupFn(payload, { ctx, init, signal }); - }); + return this._tracer.startActiveSpan( + "hooks.cleanup", + async (span) => { + return await runTimelineMetrics.measureMetric( + "trigger.dev/execution", + "cleanup", + async () => { + for (const hook of globalCleanupHooks) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ?? "global", + async (span) => { + await hook.fn({ + payload, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "tabler-function", + }, + } + ) + ); + // Ignore errors from cleanup functions + } + + if (taskCleanupHook) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + "task", + async (span) => { + await taskCleanupHook({ + payload, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "tabler-function", + }, + } + ) + ); + // Ignore errors from cleanup functions + } + } + ); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "tabler-function", + }, + } + ); } async #blockForWaitUntil() { diff --git a/packages/core/test/taskExecutor.test.ts b/packages/core/test/taskExecutor.test.ts index ad0b70117..beac1e06a 100644 --- a/packages/core/test/taskExecutor.test.ts +++ b/packages/core/test/taskExecutor.test.ts @@ -779,7 +779,8 @@ describe("TaskExecutor", () => { }, }); - expect((result as any).result.retry.delay).toBeCloseTo(30000, -1); + expect((result as any).result.retry.delay).toBeGreaterThan(29900); + expect((result as any).result.retry.delay).toBeLessThan(30100); }); test("should execute middleware hooks in correct order around other hooks", async () => { @@ -1203,6 +1204,219 @@ describe("TaskExecutor", () => { }, }); }); + + test("should call cleanup hooks in correct order after other hooks but before middleware completion", async () => { + const executionOrder: string[] = []; + + // Register global init hook + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + executionOrder.push("init"); + return { + foo: "bar", + }; + }, + }); + + // Register global start hook + lifecycleHooks.registerGlobalStartHook({ + id: "global-start", + fn: async () => { + executionOrder.push("start"); + }, + }); + + // Register global success hook + lifecycleHooks.registerGlobalSuccessHook({ + id: "global-success", + fn: async () => { + executionOrder.push("success"); + }, + }); + + // Register global complete hook + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async () => { + executionOrder.push("complete"); + }, + }); + + // Register global cleanup hooks + lifecycleHooks.registerGlobalCleanupHook({ + id: "global-cleanup-1", + fn: async ({ init }) => { + executionOrder.push("global-cleanup-1"); + // Verify we have access to init data + expect(init).toEqual({ foo: "bar" }); + }, + }); + + lifecycleHooks.registerGlobalCleanupHook({ + id: "global-cleanup-2", + fn: async ({ init }) => { + executionOrder.push("global-cleanup-2"); + // Verify we have access to init data + expect(init).toEqual({ foo: "bar" }); + }, + }); + + // Register task-specific cleanup hook + lifecycleHooks.registerTaskCleanupHook("test-task", { + id: "task-cleanup", + fn: async ({ init }) => { + executionOrder.push("task-cleanup"); + // Verify we have access to init data + expect(init).toEqual({ foo: "bar" }); + }, + }); + + // Register middleware to verify cleanup happens before middleware completion + lifecycleHooks.registerGlobalMiddlewareHook({ + id: "global-middleware", + fn: async ({ next }) => { + executionOrder.push("middleware-before"); + await next(); + executionOrder.push("middleware-after"); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run"); + return { + output: "test-output", + init: params.init, + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }); + + // Verify the execution order: + // 1. Middleware starts + // 2. Init hook + // 3. Start hook + // 4. Run function + // 5. Success hook + // 6. Complete hook + // 7. Cleanup hooks + // 8. Middleware completes + expect(executionOrder).toEqual([ + "middleware-before", + "init", + "start", + "run", + "success", + "complete", + "global-cleanup-1", + "global-cleanup-2", + "task-cleanup", + "middleware-after", + ]); + + // Verify the final result + expect(result).toEqual({ + result: { + ok: true, + id: "test-run-id", + output: '{"json":{"output":"test-output","init":{"foo":"bar"}}}', + outputType: "application/super+json", + }, + }); + }); + + test("should call cleanup hooks even when task fails", async () => { + const executionOrder: string[] = []; + const expectedError = new Error("Task failed intentionally"); + + // Register global init hook + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + executionOrder.push("init"); + return { + foo: "bar", + }; + }, + }); + + // Register failure hook + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure", + fn: async () => { + executionOrder.push("failure"); + }, + }); + + // Register complete hook + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async () => { + executionOrder.push("complete"); + }, + }); + + // Register cleanup hooks + lifecycleHooks.registerGlobalCleanupHook({ + id: "global-cleanup", + fn: async ({ init }) => { + executionOrder.push("global-cleanup"); + // Verify we have access to init data even after failure + expect(init).toEqual({ foo: "bar" }); + }, + }); + + lifecycleHooks.registerTaskCleanupHook("test-task", { + id: "task-cleanup", + fn: async ({ init }) => { + executionOrder.push("task-cleanup"); + // Verify we have access to init data even after failure + expect(init).toEqual({ foo: "bar" }); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async () => { + executionOrder.push("run"); + throw expectedError; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }); + + // Verify cleanup hooks are called even after failure + expect(executionOrder).toEqual([ + "init", + "run", + "failure", + "complete", + "global-cleanup", + "task-cleanup", + ]); + + // Verify the error result + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Task failed intentionally", + name: "Error", + stackTrace: expect.any(String), + }, + skippedRetrying: false, + }, + }); + }); }); function executeTask(task: TaskMetadataWithFunctions, payload: any) {