From aaf2ed8a200f21e45f74763ff199eaef009f59bf Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sat, 22 Mar 2025 15:47:59 +0000 Subject: [PATCH] onWait and onResume --- .../core/src/v3/lifecycleHooks/manager.ts | 109 ++++++++++++++++++ packages/core/src/v3/lifecycleHooks/types.ts | 40 +++++++ packages/trigger-sdk/src/v3/hooks.ts | 25 ++++ packages/trigger-sdk/src/v3/tasks.ts | 4 +- 4 files changed, 177 insertions(+), 1 deletion(-) diff --git a/packages/core/src/v3/lifecycleHooks/manager.ts b/packages/core/src/v3/lifecycleHooks/manager.ts index 11dab9182..86b2f1abb 100644 --- a/packages/core/src/v3/lifecycleHooks/manager.ts +++ b/packages/core/src/v3/lifecycleHooks/manager.ts @@ -7,6 +7,8 @@ import { AnyOnFailureHookFunction, AnyOnSuccessHookFunction, AnyOnCompleteHookFunction, + AnyOnWaitHookFunction, + AnyOnResumeHookFunction, } from "./types.js"; export class StandardLifecycleHooksManager implements LifecycleHooksManager { @@ -31,6 +33,13 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager { private taskCompleteHooks: Map> = new Map(); + private globalWaitHooks: Map> = new Map(); + private taskWaitHooks: Map> = new Map(); + + private globalResumeHooks: Map> = + new Map(); + private taskResumeHooks: Map> = new Map(); + registerGlobalStartHook(hook: RegisterHookFunctionParams): void { const id = generateHookId(hook); @@ -188,6 +197,68 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager { getGlobalCompleteHooks(): RegisteredHookFunction[] { return Array.from(this.globalCompleteHooks.values()); } + + registerGlobalWaitHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalWaitHooks.set(id, { + id, + name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined, + fn: hook.fn, + }); + } + + registerTaskWaitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskWaitHooks.set(taskId, { + id, + name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined, + fn: hook.fn, + }); + } + + getTaskWaitHook(taskId: string): AnyOnWaitHookFunction | undefined { + return this.taskWaitHooks.get(taskId)?.fn; + } + + getGlobalWaitHooks(): RegisteredHookFunction[] { + return Array.from(this.globalWaitHooks.values()); + } + + registerGlobalResumeHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalResumeHooks.set(id, { + id, + name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined, + fn: hook.fn, + }); + } + + registerTaskResumeHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskResumeHooks.set(taskId, { + id, + name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined, + fn: hook.fn, + }); + } + + getTaskResumeHook(taskId: string): AnyOnResumeHookFunction | undefined { + return this.taskResumeHooks.get(taskId)?.fn; + } + + getGlobalResumeHooks(): RegisteredHookFunction[] { + return Array.from(this.globalResumeHooks.values()); + } } export class NoopLifecycleHooksManager implements LifecycleHooksManager { @@ -285,6 +356,44 @@ export class NoopLifecycleHooksManager implements LifecycleHooksManager { getGlobalCompleteHooks(): RegisteredHookFunction[] { return []; } + + registerGlobalWaitHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskWaitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskWaitHook(taskId: string): AnyOnWaitHookFunction | undefined { + return undefined; + } + + getGlobalWaitHooks(): RegisteredHookFunction[] { + return []; + } + + registerGlobalResumeHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskResumeHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskResumeHook(taskId: string): AnyOnResumeHookFunction | undefined { + return undefined; + } + + getGlobalResumeHooks(): RegisteredHookFunction[] { + return []; + } } function generateHookId(hook: RegisterHookFunctionParams): string { diff --git a/packages/core/src/v3/lifecycleHooks/types.ts b/packages/core/src/v3/lifecycleHooks/types.ts index bcf7c36db..831e39040 100644 --- a/packages/core/src/v3/lifecycleHooks/types.ts +++ b/packages/core/src/v3/lifecycleHooks/types.ts @@ -26,6 +26,32 @@ export type OnStartHookFunction = ( export type AnyOnStartHookFunction = OnStartHookFunction; +export type TaskWaitHookParams = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + signal?: AbortSignal; +}; + +export type OnWaitHookFunction = ( + params: TaskWaitHookParams +) => undefined | void | Promise; + +export type AnyOnWaitHookFunction = OnWaitHookFunction; + +export type TaskResumeHookParams = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + signal?: AbortSignal; +}; + +export type OnResumeHookFunction = ( + params: TaskResumeHookParams +) => undefined | void | Promise; + +export type AnyOnResumeHookFunction = OnResumeHookFunction; + export type TaskFailureHookParams = { ctx: TaskRunContext; payload: TPayload; @@ -129,4 +155,18 @@ export interface LifecycleHooksManager { ): void; getTaskCompleteHook(taskId: string): AnyOnCompleteHookFunction | undefined; getGlobalCompleteHooks(): RegisteredHookFunction[]; + registerGlobalWaitHook(hook: RegisterHookFunctionParams): void; + registerTaskWaitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskWaitHook(taskId: string): AnyOnWaitHookFunction | undefined; + getGlobalWaitHooks(): RegisteredHookFunction[]; + registerGlobalResumeHook(hook: RegisterHookFunctionParams): void; + registerTaskResumeHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskResumeHook(taskId: string): AnyOnResumeHookFunction | undefined; + getGlobalResumeHooks(): RegisteredHookFunction[]; } diff --git a/packages/trigger-sdk/src/v3/hooks.ts b/packages/trigger-sdk/src/v3/hooks.ts index 2ff538089..119566fed 100644 --- a/packages/trigger-sdk/src/v3/hooks.ts +++ b/packages/trigger-sdk/src/v3/hooks.ts @@ -10,6 +10,8 @@ import { type AnyOnSuccessHookFunction, type AnyOnCompleteHookFunction, type TaskCompleteResult, + type AnyOnWaitHookFunction, + type AnyOnResumeHookFunction, } from "@trigger.dev/core/v3"; export type { @@ -23,6 +25,8 @@ export type { AnyOnSuccessHookFunction, AnyOnCompleteHookFunction, TaskCompleteResult, + AnyOnWaitHookFunction, + AnyOnResumeHookFunction, }; export function onInit(name: string, fn: AnyOnInitHookFunction): void; @@ -81,3 +85,24 @@ export function onComplete( fn: typeof fnOrName === "function" ? fnOrName : fn!, }); } + +export function onWait(name: string, fn: AnyOnWaitHookFunction): void; +export function onWait(fn: AnyOnWaitHookFunction): void; +export function onWait(fnOrName: string | AnyOnWaitHookFunction, fn?: AnyOnWaitHookFunction): void { + lifecycleHooks.registerGlobalWaitHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} + +export function onResume(name: string, fn: AnyOnResumeHookFunction): void; +export function onResume(fn: AnyOnResumeHookFunction): void; +export function onResume( + fnOrName: string | AnyOnResumeHookFunction, + fn?: AnyOnResumeHookFunction +): void { + lifecycleHooks.registerGlobalResumeHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} diff --git a/packages/trigger-sdk/src/v3/tasks.ts b/packages/trigger-sdk/src/v3/tasks.ts index e0ebe10c7..c9b1b4be4 100644 --- a/packages/trigger-sdk/src/v3/tasks.ts +++ b/packages/trigger-sdk/src/v3/tasks.ts @@ -1,4 +1,4 @@ -import { onInit, onStart, onFailure, onSuccess, onComplete } from "./hooks.js"; +import { onInit, onStart, onFailure, onSuccess, onComplete, onWait, onResume } from "./hooks.js"; import { batchTrigger, batchTriggerAndWait, @@ -84,4 +84,6 @@ export const tasks = { onFailure, onSuccess, onComplete, + onWait, + onResume, };