onComplete

Author: ericallam

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

@@ -6,6 +6,7 @@ import {
   RegisterHookFunctionParams,
   AnyOnFailureHookFunction,
   AnyOnSuccessHookFunction,
+  AnyOnCompleteHookFunction,
 } from "./types.js";
 
 export class StandardLifecycleHooksManager implements LifecycleHooksManager {
@@ -25,6 +26,11 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
   private taskSuccessHooks: Map<string, RegisteredHookFunction<AnyOnSuccessHookFunction>> =
     new Map();
 
+  private globalCompleteHooks: Map<string, RegisteredHookFunction<AnyOnCompleteHookFunction>> =
+    new Map();
+  private taskCompleteHooks: Map<string, RegisteredHookFunction<AnyOnCompleteHookFunction>> =
+    new Map();
+
   registerGlobalStartHook(hook: RegisterHookFunctionParams<AnyOnStartHookFunction>): void {
     const id = generateHookId(hook);
 
@@ -151,6 +157,37 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
   getGlobalSuccessHooks(): RegisteredHookFunction<AnyOnSuccessHookFunction>[] {
     return Array.from(this.globalSuccessHooks.values());
   }
+
+  registerGlobalCompleteHook(hook: RegisterHookFunctionParams<AnyOnCompleteHookFunction>): void {
+    const id = generateHookId(hook);
+
+    this.globalCompleteHooks.set(id, {
+      id,
+      name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined,
+      fn: hook.fn,
+    });
+  }
+
+  registerTaskCompleteHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnCompleteHookFunction>
+  ): void {
+    const id = generateHookId(hook);
+
+    this.taskCompleteHooks.set(taskId, {
+      id,
+      name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined,
+      fn: hook.fn,
+    });
+  }
+
+  getTaskCompleteHook(taskId: string): AnyOnCompleteHookFunction | undefined {
+    return this.taskCompleteHooks.get(taskId)?.fn;
+  }
+
+  getGlobalCompleteHooks(): RegisteredHookFunction<AnyOnCompleteHookFunction>[] {
+    return Array.from(this.globalCompleteHooks.values());
+  }
 }
 
 export class NoopLifecycleHooksManager implements LifecycleHooksManager {
@@ -229,6 +266,25 @@ export class NoopLifecycleHooksManager implements LifecycleHooksManager {
   getGlobalSuccessHooks(): RegisteredHookFunction<AnyOnSuccessHookFunction>[] {
     return [];
   }
+
+  registerGlobalCompleteHook(hook: RegisterHookFunctionParams<AnyOnCompleteHookFunction>): void {
+    // Noop
+  }
+
+  registerTaskCompleteHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnCompleteHookFunction>
+  ): void {
+    // Noop
+  }
+
+  getTaskCompleteHook(taskId: string): AnyOnCompleteHookFunction | undefined {
+    return undefined;
+  }
+
+  getGlobalCompleteHooks(): RegisteredHookFunction<AnyOnCompleteHookFunction>[] {
+    return [];
+  }
 }
 
 function generateHookId(hook: RegisterHookFunctionParams<any>): string {

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

@@ -54,6 +54,34 @@ export type OnSuccessHookFunction<TPayload, TOutput> = (
 
 export type AnyOnSuccessHookFunction = OnSuccessHookFunction<unknown, unknown>;
 
+export type TaskCompleteSuccessResult<TOutput> = {
+  ok: true;
+  data: TOutput;
+};
+
+export type TaskCompleteErrorResult = {
+  ok: false;
+  error: unknown;
+};
+
+export type TaskCompleteResult<TOutput> =
+  | TaskCompleteSuccessResult<TOutput>
+  | TaskCompleteErrorResult;
+
+export type TaskCompleteHookParams<TPayload = unknown, TOutput = unknown> = {
+  ctx: TaskRunContext;
+  payload: TPayload;
+  task: string;
+  result: TaskCompleteResult<TOutput>;
+  signal?: AbortSignal;
+};
+
+export type OnCompleteHookFunction<TPayload, TOutput> = (
+  params: TaskCompleteHookParams<TPayload, TOutput>
+) => undefined | void | Promise<undefined | void>;
+
+export type AnyOnCompleteHookFunction = OnCompleteHookFunction<unknown, unknown>;
+
 export type RegisterHookFunctionParams<THookFunction extends (params: any) => any> = {
   id?: string;
   fn: THookFunction;
@@ -94,4 +122,11 @@ export interface LifecycleHooksManager {
   ): void;
   getTaskSuccessHook(taskId: string): AnyOnSuccessHookFunction | undefined;
   getGlobalSuccessHooks(): RegisteredHookFunction<AnyOnSuccessHookFunction>[];
+  registerGlobalCompleteHook(hook: RegisterHookFunctionParams<AnyOnCompleteHookFunction>): void;
+  registerTaskCompleteHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnCompleteHookFunction>
+  ): void;
+  getTaskCompleteHook(taskId: string): AnyOnCompleteHookFunction | undefined;
+  getGlobalCompleteHooks(): RegisteredHookFunction<AnyOnCompleteHookFunction>[];
 }

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

@@ -8,6 +8,8 @@ import {
   type OnStartHookFunction,
   type AnyOnFailureHookFunction,
   type AnyOnSuccessHookFunction,
+  type AnyOnCompleteHookFunction,
+  type TaskCompleteResult,
 } from "@trigger.dev/core/v3";
 
 export type {
@@ -19,6 +21,8 @@ export type {
   OnStartHookFunction,
   AnyOnFailureHookFunction,
   AnyOnSuccessHookFunction,
+  AnyOnCompleteHookFunction,
+  TaskCompleteResult,
 };
 
 export function onInit(name: string, fn: AnyOnInitHookFunction): void;
@@ -65,3 +69,15 @@ export function onSuccess(
     fn: typeof fnOrName === "function" ? fnOrName : fn!,
   });
 }
+
+export function onComplete(name: string, fn: AnyOnCompleteHookFunction): void;
+export function onComplete(fn: AnyOnCompleteHookFunction): void;
+export function onComplete(
+  fnOrName: string | AnyOnCompleteHookFunction,
+  fn?: AnyOnCompleteHookFunction
+): void {
+  lifecycleHooks.registerGlobalCompleteHook({
+    id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined,
+    fn: typeof fnOrName === "function" ? fnOrName : fn!,
+  });
+}

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

@@ -1,4 +1,4 @@
-import { onInit, onStart, onFailure, onSuccess } from "./hooks.js";
+import { onInit, onStart, onFailure, onSuccess, onComplete } from "./hooks.js";
 import {
   batchTrigger,
   batchTriggerAndWait,
@@ -83,4 +83,5 @@ export const tasks = {
   onStart,
   onFailure,
   onSuccess,
+  onComplete,
 };