implement onWait/onResume

ericallam · ericallam · commit 268c9f6898c6 · 2025-03-25T15:55:42.000Z
diff --git a/packages/core/src/v3/lifecycle-hooks-api.ts b/packages/core/src/v3/lifecycle-hooks-api.ts
@@ -31,4 +31,5 @@ export type {
   OnCleanupHookFunction,
   AnyOnCleanupHookFunction,
   TaskCleanupHookParams,
+  TaskWait,
 } from "./lifecycleHooks/types.js";
diff --git a/packages/core/src/v3/lifecycleHooks/index.ts b/packages/core/src/v3/lifecycleHooks/index.ts
@@ -15,6 +15,7 @@ import {
   AnyOnWaitHookFunction,
   RegisteredHookFunction,
   RegisterHookFunctionParams,
+  TaskWait,
   type LifecycleHooksManager,
 } from "./types.js";
 
@@ -243,6 +244,22 @@ export class LifecycleHooksAPI {
     return this.#getManager().getGlobalCleanupHooks();
   }
 
+  public callOnWaitHookListeners(wait: TaskWait): Promise<void> {
+    return this.#getManager().callOnWaitHookListeners(wait);
+  }
+
+  public callOnResumeHookListeners(wait: TaskWait): Promise<void> {
+    return this.#getManager().callOnResumeHookListeners(wait);
+  }
+
+  public registerOnWaitHookListener(listener: (wait: TaskWait) => Promise<void>): void {
+    this.#getManager().registerOnWaitHookListener(listener);
+  }
+
+  public registerOnResumeHookListener(listener: (wait: TaskWait) => Promise<void>): void {
+    this.#getManager().registerOnResumeHookListener(listener);
+  }
+
   #getManager(): LifecycleHooksManager {
     return getGlobal(API_NAME) ?? NOOP_LIFECYCLE_HOOKS_MANAGER;
   }
diff --git a/packages/core/src/v3/lifecycleHooks/manager.ts b/packages/core/src/v3/lifecycleHooks/manager.ts
@@ -12,6 +12,7 @@ import {
   AnyOnCatchErrorHookFunction,
   AnyOnMiddlewareHookFunction,
   AnyOnCleanupHookFunction,
+  TaskWait,
 } from "./types.js";
 
 export class StandardLifecycleHooksManager implements LifecycleHooksManager {
@@ -58,6 +59,25 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
   private taskCleanupHooks: Map<string, RegisteredHookFunction<AnyOnCleanupHookFunction>> =
     new Map();
 
+  private onWaitHookListeners: ((wait: TaskWait) => Promise<void>)[] = [];
+  private onResumeHookListeners: ((wait: TaskWait) => Promise<void>)[] = [];
+
+  registerOnWaitHookListener(listener: (wait: TaskWait) => Promise<void>): void {
+    this.onWaitHookListeners.push(listener);
+  }
+
+  async callOnWaitHookListeners(wait: TaskWait): Promise<void> {
+    await Promise.all(this.onWaitHookListeners.map((listener) => listener(wait)));
+  }
+
+  registerOnResumeHookListener(listener: (wait: TaskWait) => Promise<void>): void {
+    this.onResumeHookListeners.push(listener);
+  }
+
+  async callOnResumeHookListeners(wait: TaskWait): Promise<void> {
+    await Promise.all(this.onResumeHookListeners.map((listener) => listener(wait)));
+  }
+
   registerGlobalStartHook(hook: RegisterHookFunctionParams<AnyOnStartHookFunction>): void {
     const id = generateHookId(hook);
 
@@ -377,6 +397,22 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
 }
 
 export class NoopLifecycleHooksManager implements LifecycleHooksManager {
+  registerOnWaitHookListener(listener: (wait: TaskWait) => Promise<void>): void {
+    // Noop
+  }
+
+  async callOnWaitHookListeners(wait: TaskWait): Promise<void> {
+    // Noop
+  }
+
+  registerOnResumeHookListener(listener: (wait: TaskWait) => Promise<void>): void {
+    // Noop
+  }
+
+  async callOnResumeHookListeners(wait: TaskWait): Promise<void> {
+    // Noop
+  }
+
   registerGlobalInitHook(hook: RegisterHookFunctionParams<AnyOnInitHookFunction>): void {
     // Noop
   }
diff --git a/packages/core/src/v3/lifecycleHooks/types.ts b/packages/core/src/v3/lifecycleHooks/types.ts
@@ -33,10 +33,30 @@ export type OnStartHookFunction<TPayload, TInitOutput extends TaskInitOutput = T
 
 export type AnyOnStartHookFunction = OnStartHookFunction<unknown, TaskInitOutput>;
 
+export type TaskWait =
+  | {
+      type: "duration";
+      date: Date;
+    }
+  | {
+      type: "token";
+      token: string;
+    }
+  | {
+      type: "task";
+      runId: string;
+    }
+  | {
+      type: "batch";
+      batchId: string;
+      runCount: number;
+    };
+
 export type TaskWaitHookParams<
   TPayload = unknown,
   TInitOutput extends TaskInitOutput = TaskInitOutput,
 > = {
+  wait: TaskWait;
   ctx: TaskRunContext;
   payload: TPayload;
   task: string;
@@ -55,6 +75,7 @@ export type TaskResumeHookParams<
   TInitOutput extends TaskInitOutput = TaskInitOutput,
 > = {
   ctx: TaskRunContext;
+  wait: TaskWait;
   payload: TPayload;
   task: string;
   signal?: AbortSignal;
@@ -280,4 +301,10 @@ export interface LifecycleHooksManager {
   ): void;
   getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined;
   getGlobalCleanupHooks(): RegisteredHookFunction<AnyOnCleanupHookFunction>[];
+
+  callOnWaitHookListeners(wait: TaskWait): Promise<void>;
+  registerOnWaitHookListener(listener: (wait: TaskWait) => Promise<void>): void;
+
+  callOnResumeHookListeners(wait: TaskWait): Promise<void>;
+  registerOnResumeHookListener(listener: (wait: TaskWait) => Promise<void>): void;
 }
diff --git a/packages/core/src/v3/runtime/managedRuntimeManager.ts b/packages/core/src/v3/runtime/managedRuntimeManager.ts
@@ -1,3 +1,4 @@
+import { lifecycleHooks } from "../lifecycle-hooks-api.js";
 import {
   BatchTaskRunExecutionResult,
   CompletedWaitpoint,
@@ -44,9 +45,19 @@ export class ManagedRuntimeManager implements RuntimeManager {
         this.resolversByWaitId.set(params.id, resolve);
       });
 
+      await lifecycleHooks.callOnWaitHookListeners({
+        type: "task",
+        runId: params.id,
+      });
+
       const waitpoint = await promise;
       const result = this.waitpointToTaskRunExecutionResult(waitpoint);
 
+      await lifecycleHooks.callOnResumeHookListeners({
+        type: "task",
+        runId: params.id,
+      });
+
       return result;
     });
   }
@@ -70,8 +81,20 @@ export class ManagedRuntimeManager implements RuntimeManager {
         })
       );
 
+      await lifecycleHooks.callOnWaitHookListeners({
+        type: "batch",
+        batchId: params.id,
+        runCount: params.runCount,
+      });
+
       const waitpoints = await promise;
 
+      await lifecycleHooks.callOnResumeHookListeners({
+        type: "batch",
+        batchId: params.id,
+        runCount: params.runCount,
+      });
+
       return {
         id: params.id,
         items: waitpoints.map(this.waitpointToTaskRunExecutionResult),
@@ -91,8 +114,32 @@ export class ManagedRuntimeManager implements RuntimeManager {
         this.resolversByWaitId.set(waitpointFriendlyId, resolve);
       });
 
+      if (finishDate) {
+        await lifecycleHooks.callOnWaitHookListeners({
+          type: "duration",
+          date: finishDate,
+        });
+      } else {
+        await lifecycleHooks.callOnWaitHookListeners({
+          type: "token",
+          token: waitpointFriendlyId,
+        });
+      }
+
       const waitpoint = await promise;
 
+      if (finishDate) {
+        await lifecycleHooks.callOnResumeHookListeners({
+          type: "duration",
+          date: finishDate,
+        });
+      } else {
+        await lifecycleHooks.callOnResumeHookListeners({
+          type: "token",
+          token: waitpointFriendlyId,
+        });
+      }
+
       return {
         ok: !waitpoint.outputIsError,
         output: waitpoint.output,
diff --git a/packages/core/src/v3/workers/taskExecutor.ts b/packages/core/src/v3/workers/taskExecutor.ts
@@ -15,6 +15,7 @@ import {
   RegisteredHookFunction,
   TaskCompleteResult,
   TaskInitOutput,
+  TaskWait,
 } from "../lifecycleHooks/types.js";
 import { recordSpanException, TracingSDK } from "../otel/index.js";
 import { runTimelineMetrics } from "../run-timeline-metrics-api.js";
@@ -126,6 +127,14 @@ export class TaskExecutor {
 
           parsedPayload = await this.#parsePayload(payloadResult);
 
+          lifecycleHooks.registerOnWaitHookListener(async (wait) => {
+            await this.#callOnWaitFunctions(wait, parsedPayload, ctx, initOutput, signal);
+          });
+
+          lifecycleHooks.registerOnResumeHookListener(async (wait) => {
+            await this.#callOnResumeFunctions(wait, parsedPayload, ctx, initOutput, signal);
+          });
+
           const executeTask = async (payload: any) => {
             const [runError, output] = await tryCatch(
               (async () => {
@@ -383,6 +392,146 @@ export class TaskExecutor {
     });
   }
 
+  async #callOnWaitFunctions(
+    wait: TaskWait,
+    payload: unknown,
+    ctx: TaskRunContext,
+    initOutput: TaskInitOutput,
+    signal?: AbortSignal
+  ) {
+    const globalWaitHooks = lifecycleHooks.getGlobalWaitHooks();
+    const taskWaitHook = lifecycleHooks.getTaskWaitHook(this.task.id);
+
+    if (globalWaitHooks.length === 0 && !taskWaitHook) {
+      return;
+    }
+
+    const result = await runTimelineMetrics.measureMetric(
+      "trigger.dev/execution",
+      "onWait",
+      async () => {
+        for (const hook of globalWaitHooks) {
+          const [hookError] = await tryCatch(
+            this._tracer.startActiveSpan(
+              hook.name ? `onWait/${hook.name}` : "onWait/global",
+              async (span) => {
+                await hook.fn({ payload, ctx, signal, task: this.task.id, wait, init: initOutput });
+              },
+              {
+                attributes: {
+                  [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onWait",
+                  [SemanticInternalAttributes.COLLAPSED]: true,
+                },
+              }
+            )
+          );
+
+          if (hookError) {
+            throw hookError;
+          }
+        }
+
+        if (taskWaitHook) {
+          const [hookError] = await tryCatch(
+            this._tracer.startActiveSpan(
+              "onWait/task",
+              async (span) => {
+                await taskWaitHook({
+                  payload,
+                  ctx,
+                  signal,
+                  task: this.task.id,
+                  wait,
+                  init: initOutput,
+                });
+              },
+              {
+                attributes: {
+                  [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onWait",
+                  [SemanticInternalAttributes.COLLAPSED]: true,
+                },
+              }
+            )
+          );
+
+          if (hookError) {
+            throw hookError;
+          }
+        }
+      }
+    );
+  }
+
+  async #callOnResumeFunctions(
+    wait: TaskWait,
+    payload: unknown,
+    ctx: TaskRunContext,
+    initOutput: TaskInitOutput,
+    signal?: AbortSignal
+  ) {
+    const globalResumeHooks = lifecycleHooks.getGlobalResumeHooks();
+    const taskResumeHook = lifecycleHooks.getTaskResumeHook(this.task.id);
+
+    if (globalResumeHooks.length === 0 && !taskResumeHook) {
+      return;
+    }
+
+    const result = await runTimelineMetrics.measureMetric(
+      "trigger.dev/execution",
+      "onResume",
+      async () => {
+        for (const hook of globalResumeHooks) {
+          const [hookError] = await tryCatch(
+            this._tracer.startActiveSpan(
+              hook.name ? `onResume/${hook.name}` : "onResume/global",
+              async (span) => {
+                await hook.fn({ payload, ctx, signal, task: this.task.id, wait, init: initOutput });
+              },
+              {
+                attributes: {
+                  [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onResume",
+                  [SemanticInternalAttributes.COLLAPSED]: true,
+                },
+              }
+            )
+          );
+
+          if (hookError) {
+            throw hookError;
+          }
+        }
+
+        if (taskResumeHook) {
+          const [hookError] = await tryCatch(
+            this._tracer.startActiveSpan(
+              "onResume/task",
+              async (span) => {
+                await taskResumeHook({
+                  payload,
+                  ctx,
+                  signal,
+                  task: this.task.id,
+                  wait,
+                  init: initOutput,
+                });
+              },
+              {
+                attributes: {
+                  [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onResume",
+                  [SemanticInternalAttributes.COLLAPSED]: true,
+                },
+              }
+            )
+          );
+
+          if (hookError) {
+            throw hookError;
+          }
+        }
+      }
+    );
+  }
+
   async #callInitFunctions(payload: unknown, ctx: TaskRunContext, signal?: AbortSignal) {
     const globalInitHooks = lifecycleHooks.getGlobalInitHooks();
     const taskInitHook = lifecycleHooks.getTaskInitHook(this.task.id);
diff --git a/packages/core/test/taskExecutor.test.ts b/packages/core/test/taskExecutor.test.ts
diff --git a/references/hello-world/src/trigger/example.ts b/references/hello-world/src/trigger/example.ts
