implement cleanup

Author: ericallam

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() {

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<any>) => {
+          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) {