Add middleware types

Author: ericallam

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

@@ -25,6 +25,9 @@ export type {
   TaskCatchErrorHookParams,
   AnyOnCatchErrorHookFunction,
   TaskCompleteResult,
+  TaskMiddlewareHookParams,
+  AnyOnMiddlewareHookFunction,
+  OnMiddlewareHookFunction,
 } from "./lifecycleHooks/types.js";
 
 export * as lifecycleHooksAdapters from "./lifecycleHooks/adapters.js";

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

@@ -5,6 +5,7 @@ import {
   AnyOnFailureHookFunction,
   AnyOnSuccessHookFunction,
   AnyOnCatchErrorHookFunction,
+  AnyOnMiddlewareHookFunction,
 } from "./types.js";
 
 export function createInitHookAdapter<TPayload>(
@@ -76,3 +77,16 @@ export function createHandleErrorHookAdapter<TPayload>(
     return await fn(params.payload as unknown as TPayload, params.error, params);
   };
 }
+
+export function createMiddlewareHookAdapter<TPayload>(
+  fn: NonNullable<TaskOptions<string, TPayload, unknown, any>["middleware"]>
+): AnyOnMiddlewareHookFunction {
+  return async (params) => {
+    const { payload, next, ...paramsWithoutPayloadAndNext } = params;
+
+    return await fn(payload as unknown as TPayload, {
+      ...paramsWithoutPayloadAndNext,
+      next,
+    });
+  };
+}

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

@@ -10,6 +10,7 @@ import {
   AnyOnWaitHookFunction,
   AnyOnResumeHookFunction,
   AnyOnCatchErrorHookFunction,
+  AnyOnMiddlewareHookFunction,
 } from "./types.js";
 
 export class StandardLifecycleHooksManager implements LifecycleHooksManager {
@@ -46,6 +47,11 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
   private taskCatchErrorHooks: Map<string, RegisteredHookFunction<AnyOnCatchErrorHookFunction>> =
     new Map();
 
+  private globalMiddlewareHooks: Map<string, RegisteredHookFunction<AnyOnMiddlewareHookFunction>> =
+    new Map();
+  private taskMiddlewareHooks: Map<string, RegisteredHookFunction<AnyOnMiddlewareHookFunction>> =
+    new Map();
+
   registerGlobalStartHook(hook: RegisterHookFunctionParams<AnyOnStartHookFunction>): void {
     const id = generateHookId(hook);
 
@@ -298,6 +304,39 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
   getGlobalCatchErrorHooks(): RegisteredHookFunction<AnyOnCatchErrorHookFunction>[] {
     return Array.from(this.globalCatchErrorHooks.values());
   }
+
+  registerGlobalMiddlewareHook(
+    hook: RegisterHookFunctionParams<AnyOnMiddlewareHookFunction>
+  ): void {
+    const id = generateHookId(hook);
+
+    this.globalMiddlewareHooks.set(id, {
+      id,
+      name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined,
+      fn: hook.fn,
+    });
+  }
+
+  registerTaskMiddlewareHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnMiddlewareHookFunction>
+  ): void {
+    const id = generateHookId(hook);
+
+    this.taskMiddlewareHooks.set(taskId, {
+      id,
+      name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined,
+      fn: hook.fn,
+    });
+  }
+
+  getTaskMiddlewareHook(taskId: string): AnyOnMiddlewareHookFunction | undefined {
+    return this.taskMiddlewareHooks.get(taskId)?.fn;
+  }
+
+  getGlobalMiddlewareHooks(): RegisteredHookFunction<AnyOnMiddlewareHookFunction>[] {
+    return Array.from(this.globalMiddlewareHooks.values());
+  }
 }
 
 export class NoopLifecycleHooksManager implements LifecycleHooksManager {
@@ -449,6 +488,22 @@ export class NoopLifecycleHooksManager implements LifecycleHooksManager {
   getGlobalCatchErrorHooks(): [] {
     return [];
   }
+
+  registerGlobalMiddlewareHook(): void {
+    // Noop
+  }
+
+  registerTaskMiddlewareHook(): void {
+    // Noop
+  }
+
+  getTaskMiddlewareHook(): undefined {
+    return undefined;
+  }
+
+  getGlobalMiddlewareHooks(): [] {
+    return [];
+  }
 }
 
 function generateHookId(hook: RegisterHookFunctionParams<any>): string {

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

@@ -137,6 +137,20 @@ export type OnCatchErrorHookFunction<TPayload> = (
 
 export type AnyOnCatchErrorHookFunction = OnCatchErrorHookFunction<unknown>;
 
+export type TaskMiddlewareHookParams<TPayload = unknown> = {
+  ctx: TaskRunContext;
+  payload: TPayload;
+  task: string;
+  signal?: AbortSignal;
+  next: () => Promise<void>;
+};
+
+export type OnMiddlewareHookFunction<TPayload> = (
+  params: TaskMiddlewareHookParams<TPayload>
+) => Promise<void>;
+
+export type AnyOnMiddlewareHookFunction = OnMiddlewareHookFunction<unknown>;
+
 export interface LifecycleHooksManager {
   registerGlobalInitHook(hook: RegisterHookFunctionParams<AnyOnInitHookFunction>): void;
   registerTaskInitHook(
@@ -194,4 +208,11 @@ export interface LifecycleHooksManager {
   ): void;
   getTaskCatchErrorHook(taskId: string): AnyOnCatchErrorHookFunction | undefined;
   getGlobalCatchErrorHooks(): RegisteredHookFunction<AnyOnCatchErrorHookFunction>[];
+  registerGlobalMiddlewareHook(hook: RegisterHookFunctionParams<AnyOnMiddlewareHookFunction>): void;
+  registerTaskMiddlewareHook(
+    taskId: string,
+    hook: RegisterHookFunctionParams<AnyOnMiddlewareHookFunction>
+  ): void;
+  getTaskMiddlewareHook(taskId: string): AnyOnMiddlewareHookFunction | undefined;
+  getGlobalMiddlewareHooks(): RegisteredHookFunction<AnyOnMiddlewareHookFunction>[];
 }