Add cleanup hook

Author: ericallam

packages/core/src/v3/lifecycle-hooks-api.ts

@@ -28,6 +28,9 @@ export type {
   TaskMiddlewareHookParams,
   AnyOnMiddlewareHookFunction,
   OnMiddlewareHookFunction,
+  OnCleanupHookFunction,
+  AnyOnCleanupHookFunction,
+  TaskCleanupHookParams,
 } from "./lifecycleHooks/types.js";
 
 export * as lifecycleHooksAdapters from "./lifecycleHooks/adapters.js";

packages/core/src/v3/lifecycleHooks/adapters.ts

@@ -11,6 +11,8 @@ import {
   TaskFailureHookParams,
   TaskStartHookParams,
   TaskCatchErrorHookParams,
+  TaskCleanupHookParams,
+  AnyOnCleanupHookFunction,
 } from "./types.js";
 
 export function createInitHookAdapter<TPayload>(
@@ -95,3 +97,17 @@ export function createMiddlewareHookAdapter<TPayload>(
     });
   };
 }
+
+export function createCleanupHookAdapter<
+  TPayload,
+  TInitOutput extends TaskInitOutput = TaskInitOutput,
+>(
+  fn: NonNullable<TaskOptions<string, TPayload, unknown, TInitOutput>["cleanup"]>
+): AnyOnCleanupHookFunction {
+  return async (params) => {
+    return await fn(
+      params.payload as unknown as TPayload,
+      params as TaskCleanupHookParams<TPayload, TInitOutput>
+    );
+  };
+}

packages/core/src/v3/lifecycleHooks/index.ts

@@ -4,6 +4,7 @@ import { getGlobal, registerGlobal, unregisterGlobal } from "../utils/globals.js
 import { NoopLifecycleHooksManager } from "./manager.js";
 import {
   AnyOnCatchErrorHookFunction,
+  AnyOnCleanupHookFunction,
   AnyOnCompleteHookFunction,
   AnyOnFailureHookFunction,
   AnyOnInitHookFunction,
@@ -221,6 +222,27 @@ export class LifecycleHooksAPI {
     return this.#getManager().getGlobalMiddlewareHooks();
   }
 
+  public registerGlobalCleanupHook(
+    hook: RegisterHookFunctionParams<AnyOnCleanupHookFunction>
+  ): void {
+    this.#getManager().registerGlobalCleanupHook(hook);
+  }
+
+  public registerTaskCleanupHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnCleanupHookFunction>
+  ): void {
+    this.#getManager().registerTaskCleanupHook(taskId, hook);
+  }
+
+  public getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined {
+    return this.#getManager().getTaskCleanupHook(taskId);
+  }
+
+  public getGlobalCleanupHooks(): RegisteredHookFunction<AnyOnCleanupHookFunction>[] {
+    return this.#getManager().getGlobalCleanupHooks();
+  }
+
   #getManager(): LifecycleHooksManager {
     return getGlobal(API_NAME) ?? NOOP_LIFECYCLE_HOOKS_MANAGER;
   }

packages/core/src/v3/lifecycleHooks/manager.ts

@@ -11,6 +11,7 @@ import {
   AnyOnResumeHookFunction,
   AnyOnCatchErrorHookFunction,
   AnyOnMiddlewareHookFunction,
+  AnyOnCleanupHookFunction,
 } from "./types.js";
 
 export class StandardLifecycleHooksManager implements LifecycleHooksManager {
@@ -52,6 +53,11 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
   private taskMiddlewareHooks: Map<string, RegisteredHookFunction<AnyOnMiddlewareHookFunction>> =
     new Map();
 
+  private globalCleanupHooks: Map<string, RegisteredHookFunction<AnyOnCleanupHookFunction>> =
+    new Map();
+  private taskCleanupHooks: Map<string, RegisteredHookFunction<AnyOnCleanupHookFunction>> =
+    new Map();
+
   registerGlobalStartHook(hook: RegisterHookFunctionParams<AnyOnStartHookFunction>): void {
     const id = generateHookId(hook);
 
@@ -337,6 +343,37 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
   getGlobalMiddlewareHooks(): RegisteredHookFunction<AnyOnMiddlewareHookFunction>[] {
     return Array.from(this.globalMiddlewareHooks.values());
   }
+
+  registerGlobalCleanupHook(hook: RegisterHookFunctionParams<AnyOnCleanupHookFunction>): void {
+    const id = generateHookId(hook);
+
+    this.globalCleanupHooks.set(id, {
+      id,
+      name: hook.id,
+      fn: hook.fn,
+    });
+  }
+
+  registerTaskCleanupHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnCleanupHookFunction>
+  ): void {
+    const id = generateHookId(hook);
+
+    this.taskCleanupHooks.set(taskId, {
+      id,
+      name: hook.id,
+      fn: hook.fn,
+    });
+  }
+
+  getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined {
+    return this.taskCleanupHooks.get(taskId)?.fn;
+  }
+
+  getGlobalCleanupHooks(): RegisteredHookFunction<AnyOnCleanupHookFunction>[] {
+    return Array.from(this.globalCleanupHooks.values());
+  }
 }
 
 export class NoopLifecycleHooksManager implements LifecycleHooksManager {
@@ -504,6 +541,25 @@ export class NoopLifecycleHooksManager implements LifecycleHooksManager {
   getGlobalMiddlewareHooks(): [] {
     return [];
   }
+
+  registerGlobalCleanupHook(hook: RegisterHookFunctionParams<AnyOnCleanupHookFunction>): void {
+    // Noop
+  }
+
+  registerTaskCleanupHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnCleanupHookFunction>
+  ): void {
+    // Noop
+  }
+
+  getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined {
+    return undefined;
+  }
+
+  getGlobalCleanupHooks(): RegisteredHookFunction<AnyOnCleanupHookFunction>[] {
+    return [];
+  }
 }
 
 function generateHookId(hook: RegisterHookFunctionParams<any>): string {

packages/core/src/v3/lifecycleHooks/types.ts

@@ -192,6 +192,23 @@ export type OnMiddlewareHookFunction<TPayload> = (
 
 export type AnyOnMiddlewareHookFunction = OnMiddlewareHookFunction<unknown>;
 
+export type TaskCleanupHookParams<
+  TPayload = unknown,
+  TInitOutput extends TaskInitOutput = TaskInitOutput,
+> = {
+  ctx: TaskRunContext;
+  payload: TPayload;
+  task: string;
+  signal?: AbortSignal;
+  init?: TInitOutput;
+};
+
+export type OnCleanupHookFunction<TPayload, TInitOutput extends TaskInitOutput = TaskInitOutput> = (
+  params: TaskCleanupHookParams<TPayload, TInitOutput>
+) => undefined | void | Promise<undefined | void>;
+
+export type AnyOnCleanupHookFunction = OnCleanupHookFunction<unknown, TaskInitOutput>;
+
 export interface LifecycleHooksManager {
   registerGlobalInitHook(hook: RegisterHookFunctionParams<AnyOnInitHookFunction>): void;
   registerTaskInitHook(
@@ -256,4 +273,11 @@ export interface LifecycleHooksManager {
   ): void;
   getTaskMiddlewareHook(taskId: string): AnyOnMiddlewareHookFunction | undefined;
   getGlobalMiddlewareHooks(): RegisteredHookFunction<AnyOnMiddlewareHookFunction>[];
+  registerGlobalCleanupHook(hook: RegisterHookFunctionParams<AnyOnCleanupHookFunction>): void;
+  registerTaskCleanupHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnCleanupHookFunction>
+  ): void;
+  getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined;
+  getGlobalCleanupHooks(): RegisteredHookFunction<AnyOnCleanupHookFunction>[];
 }

packages/trigger-sdk/src/v3/shared.ts

@@ -1624,4 +1624,10 @@ function registerTaskLifecycleHooks<
       fn: lifecycleHooksAdapters.createMiddlewareHookAdapter(params.middleware),
     });
   }
+
+  if (params.cleanup) {
+    lifecycleHooks.registerTaskCleanupHook(taskId, {
+      fn: lifecycleHooksAdapters.createCleanupHookAdapter(params.cleanup),
+    });
+  }
 }