From ba520fcdd8fa83a13274f3bd7eb3d5405c123d87 Mon Sep 17 00:00:00 2001
From: Eric Allam <eallam@icloud.com>
Date: Tue, 25 Mar 2025 11:55:54 +0000
Subject: [PATCH] Make sure and register all the config hooks

---
 .../cli-v3/src/entryPoints/dev-run-worker.ts  | 29 +++++++++++++++----
 packages/core/src/v3/types/tasks.ts           |  2 +-
 packages/core/src/v3/workers/taskExecutor.ts  |  3 --
 packages/core/test/taskExecutor.test.ts       |  1 -
 4 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/packages/cli-v3/src/entryPoints/dev-run-worker.ts b/packages/cli-v3/src/entryPoints/dev-run-worker.ts
index a720354ae..2dc0cdc45 100644
--- a/packages/cli-v3/src/entryPoints/dev-run-worker.ts
+++ b/packages/cli-v3/src/entryPoints/dev-run-worker.ts
@@ -183,24 +183,44 @@ async function bootstrap() {
 
   if (config.init) {
     lifecycleHooks.registerGlobalInitHook({
-      id: "trigger-dev-worker",
+      id: "config",
       fn: lifecycleHooksAdapters.createInitHookAdapter(config.init),
     });
   }
 
   if (config.onStart) {
     lifecycleHooks.registerGlobalStartHook({
-      id: "trigger-dev-worker",
+      id: "config",
       fn: lifecycleHooksAdapters.createStartHookAdapter(config.onStart),
     });
   }
 
+  if (config.onSuccess) {
+    lifecycleHooks.registerGlobalSuccessHook({
+      id: "config",
+      fn: lifecycleHooksAdapters.createSuccessHookAdapter(config.onSuccess),
+    });
+  }
+
+  if (config.onFailure) {
+    lifecycleHooks.registerGlobalFailureHook({
+      id: "config",
+      fn: lifecycleHooksAdapters.createFailureHookAdapter(config.onFailure),
+    });
+  }
+
+  if (handleError) {
+    lifecycleHooks.registerGlobalCatchErrorHook({
+      id: "config",
+      fn: lifecycleHooksAdapters.createHandleErrorHookAdapter(handleError),
+    });
+  }
+
   return {
     tracer,
     tracingSDK,
     consoleInterceptor,
     config,
-    handleErrorFn: handleError,
     workerManifest,
   };
 }
@@ -242,7 +262,7 @@ const zodIpc = new ZodIpcConnection({
       }
 
       try {
-        const { tracer, tracingSDK, consoleInterceptor, config, handleErrorFn, workerManifest } =
+        const { tracer, tracingSDK, consoleInterceptor, config, workerManifest } =
           await bootstrap();
 
         _tracingSDK = tracingSDK;
@@ -359,7 +379,6 @@ const zodIpc = new ZodIpcConnection({
           tracingSDK,
           consoleInterceptor,
           retries: config.retries,
-          handleErrorFn,
         });
 
         try {
diff --git a/packages/core/src/v3/types/tasks.ts b/packages/core/src/v3/types/tasks.ts
index fa32eef40..0c425d427 100644
--- a/packages/core/src/v3/types/tasks.ts
+++ b/packages/core/src/v3/types/tasks.ts
@@ -133,7 +133,7 @@ export type HandleErrorResult =
 
 export type HandleErrorArgs = {
   ctx: Context;
-  init: unknown;
+  init?: Record<string, unknown>;
   retry?: RetryOptions;
   retryAt?: Date;
   retryDelayInMs?: number;
diff --git a/packages/core/src/v3/workers/taskExecutor.ts b/packages/core/src/v3/workers/taskExecutor.ts
index 24c2f0068..4dc5cd000 100644
--- a/packages/core/src/v3/workers/taskExecutor.ts
+++ b/packages/core/src/v3/workers/taskExecutor.ts
@@ -53,7 +53,6 @@ export type TaskExecutorOptions = {
     enabledInDev?: boolean;
     default?: RetryOptions;
   };
-  handleErrorFn: HandleErrorFunction | undefined;
 };
 
 export class TaskExecutor {
@@ -66,7 +65,6 @@ export class TaskExecutor {
         default?: RetryOptions;
       }
     | undefined;
-  private _handleErrorFn: HandleErrorFunction | undefined;
 
   constructor(
     public task: TaskMetadataWithFunctions,
@@ -76,7 +74,6 @@ export class TaskExecutor {
     this._tracer = options.tracer;
     this._consoleInterceptor = options.consoleInterceptor;
     this._retries = options.retries;
-    this._handleErrorFn = options.handleErrorFn;
   }
 
   async execute(
diff --git a/packages/core/test/taskExecutor.test.ts b/packages/core/test/taskExecutor.test.ts
index f8d436a3b..83aea93ec 100644
--- a/packages/core/test/taskExecutor.test.ts
+++ b/packages/core/test/taskExecutor.test.ts
@@ -1531,7 +1531,6 @@ function executeTask(task: TaskMetadataWithFunctions, payload: any, signal?: Abo
         maxAttempts: 1,
       },
     },
-    handleErrorFn: undefined,
   });
 
   const execution: TaskRunExecution = {