Add middleware hooks

Author: ericallam

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

@@ -7,6 +7,7 @@ import {
   AnyOnCompleteHookFunction,
   AnyOnFailureHookFunction,
   AnyOnInitHookFunction,
+  AnyOnMiddlewareHookFunction,
   AnyOnResumeHookFunction,
   AnyOnStartHookFunction,
   AnyOnSuccessHookFunction,
@@ -199,6 +200,27 @@ export class LifecycleHooksAPI {
     return this.#getManager().getGlobalCatchErrorHooks();
   }
 
+  public registerGlobalMiddlewareHook(
+    hook: RegisterHookFunctionParams<AnyOnMiddlewareHookFunction>
+  ): void {
+    this.#getManager().registerGlobalMiddlewareHook(hook);
+  }
+
+  public registerTaskMiddlewareHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnMiddlewareHookFunction>
+  ): void {
+    this.#getManager().registerTaskMiddlewareHook(taskId, hook);
+  }
+
+  public getTaskMiddlewareHook(taskId: string): AnyOnMiddlewareHookFunction | undefined {
+    return this.#getManager().getTaskMiddlewareHook(taskId);
+  }
+
+  public getGlobalMiddlewareHooks(): RegisteredHookFunction<AnyOnMiddlewareHookFunction>[] {
+    return this.#getManager().getGlobalMiddlewareHooks();
+  }
+
   #getManager(): LifecycleHooksManager {
     return getGlobal(API_NAME) ?? NOOP_LIFECYCLE_HOOKS_MANAGER;
   }

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

@@ -10,6 +10,7 @@ import {
   type AnyOnWaitHookFunction,
   type AnyOnResumeHookFunction,
   type AnyOnCatchErrorHookFunction,
+  type AnyOnMiddlewareHookFunction,
 } from "@trigger.dev/core/v3";
 
 export type {
@@ -23,6 +24,7 @@ export type {
   AnyOnWaitHookFunction,
   AnyOnResumeHookFunction,
   AnyOnCatchErrorHookFunction,
+  AnyOnMiddlewareHookFunction,
 };
 
 export function onStart(name: string, fn: AnyOnStartHookFunction): void;
@@ -117,3 +119,15 @@ export function onCatchError(
     fn: typeof fnOrName === "function" ? fnOrName : fn!,
   });
 }
+
+export function middleware(name: string, fn: AnyOnMiddlewareHookFunction): void;
+export function middleware(fn: AnyOnMiddlewareHookFunction): void;
+export function middleware(
+  fnOrName: string | AnyOnMiddlewareHookFunction,
+  fn?: AnyOnMiddlewareHookFunction
+): void {
+  lifecycleHooks.registerGlobalMiddlewareHook({
+    id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined,
+    fn: typeof fnOrName === "function" ? fnOrName : fn!,
+  });
+}

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

@@ -212,13 +212,6 @@ export function createTask<
     maxDuration: params.maxDuration,
     fns: {
       run: params.run,
-      init: params.init,
-      cleanup: params.cleanup,
-      middleware: params.middleware,
-      handleError: params.handleError,
-      onSuccess: params.onSuccess,
-      onFailure: params.onFailure,
-      onStart: params.onStart,
     },
   });
 
@@ -389,6 +382,12 @@ export function createSchemaTask<
     });
   }
 
+  if (params.middleware) {
+    lifecycleHooks.registerTaskMiddlewareHook(params.id, {
+      fn: lifecycleHooksAdapters.createMiddlewareHookAdapter(params.middleware),
+    });
+  }
+
   resourceCatalog.registerTaskMetadata({
     id: params.id,
     description: params.description,
@@ -398,13 +397,6 @@ export function createSchemaTask<
     maxDuration: params.maxDuration,
     fns: {
       run: params.run,
-      init: params.init,
-      cleanup: params.cleanup,
-      middleware: params.middleware,
-      handleError: params.handleError,
-      onSuccess: params.onSuccess,
-      onFailure: params.onFailure,
-      onStart: params.onStart,
       parsePayload,
     },
   });

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

@@ -7,6 +7,7 @@ import {
   onResume,
   onHandleError,
   onCatchError,
+  middleware,
 } from "./hooks.js";
 import {
   batchTrigger,
@@ -97,4 +98,5 @@ export const tasks = {
   /** @deprecated Use catchError instead */
   handleError: onHandleError,
   catchError: onCatchError,
+  middleware,
 };

references/hello-world/src/trigger/init.ts

@@ -1,6 +1,11 @@
 import { logger, tasks } from "@trigger.dev/sdk";
 // import { setDb } from "../db.js";
 
+tasks.middleware("db", ({ ctx, payload, next }) => {
+  logger.info("Hello, world from the middleware", { ctx, payload });
+  return next();
+});
+
 // tasks.onSuccess(({ ctx, payload, output }) => {
 //   logger.info("Hello, world from the success", { ctx, payload });
 // });