From b171fde483f75769aa942e76cd7bf33c57dffd8b Mon Sep 17 00:00:00 2001
From: nicktrn <55853254+nicktrn@users.noreply.github.com>
Date: Sat, 29 Jun 2024 19:23:32 +0100
Subject: [PATCH] Fix index failure completion override

---
 apps/kubernetes-provider/src/index.ts         |  1 +
 apps/webapp/app/v3/handleSocketIo.server.ts   |  2 +-
 .../cli-v3/src/workers/prod/entry-point.ts    | 65 +++++++++----------
 packages/core/src/v3/schemas/messages.ts      |  2 +-
 4 files changed, 32 insertions(+), 38 deletions(-)

diff --git a/apps/kubernetes-provider/src/index.ts b/apps/kubernetes-provider/src/index.ts
index 0648bc5a8..9c3117eef 100644
--- a/apps/kubernetes-provider/src/index.ts
+++ b/apps/kubernetes-provider/src/index.ts
@@ -540,6 +540,7 @@ const taskMonitor = new TaskMonitor({
           message: details.reason,
           stack: details.logs,
         },
+        overrideCompletion: details.overrideCompletion,
       });
     } catch (error) {
       logger.error(error);
diff --git a/apps/webapp/app/v3/handleSocketIo.server.ts b/apps/webapp/app/v3/handleSocketIo.server.ts
index ede5eb9b4..99cf0b6e4 100644
--- a/apps/webapp/app/v3/handleSocketIo.server.ts
+++ b/apps/webapp/app/v3/handleSocketIo.server.ts
@@ -267,7 +267,7 @@ function createProviderNamespace(io: Server) {
         try {
           const service = new DeploymentIndexFailed();
 
-          await service.call(message.deploymentId, message.error);
+          await service.call(message.deploymentId, message.error, message.overrideCompletion);
         } catch (e) {
           logger.error("Error while indexing", { error: e });
         }
diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts
index 53b726d4f..e30388491 100644
--- a/packages/cli-v3/src/workers/prod/entry-point.ts
+++ b/packages/cli-v3/src/workers/prod/entry-point.ts
@@ -9,7 +9,7 @@ import {
   TaskRunFailedExecutionResult,
   WaitReason,
 } from "@trigger.dev/core/v3";
-import { ZodSocketConnection } from "@trigger.dev/core/v3/zodSocket";
+import { InferSocketMessageSchema, ZodSocketConnection } from "@trigger.dev/core/v3/zodSocket";
 import { HttpReply, getRandomPortNumber } from "@trigger.dev/core-apps/http";
 import { SimpleLogger } from "@trigger.dev/core-apps/logger";
 import { EXIT_CODE_ALREADY_HANDLED, EXIT_CODE_CHILD_NONZERO } from "@trigger.dev/core-apps/process";
@@ -629,6 +629,19 @@ class ProdWorker {
         }
 
         if (process.env.INDEX_TASKS === "true") {
+          const failIndex = (
+            error: InferSocketMessageSchema<
+              typeof ProdWorkerToCoordinatorMessages,
+              "INDEXING_FAILED"
+            >["error"]
+          ) => {
+            socket.emit("INDEXING_FAILED", {
+              version: "v1",
+              deploymentId: this.deploymentId,
+              error,
+            });
+          };
+
           try {
             const taskResources = await this.#initializeWorker();
 
@@ -655,15 +668,11 @@ class ProdWorker {
                 tasks: e.tasks,
               });
 
-              socket.emit("INDEXING_FAILED", {
-                version: "v1",
-                deploymentId: this.deploymentId,
-                error: {
-                  name: "TaskMetadataParseError",
-                  message: "There was an error parsing the task metadata",
-                  stack: JSON.stringify({ zodIssues: e.zodIssues, tasks: e.tasks }),
-                  stderr,
-                },
+              failIndex({
+                name: "TaskMetadataParseError",
+                message: "There was an error parsing the task metadata",
+                stack: JSON.stringify({ zodIssues: e.zodIssues, tasks: e.tasks }),
+                stderr,
               });
             } else if (e instanceof UncaughtExceptionError) {
               const error = {
@@ -675,11 +684,7 @@ class ProdWorker {
 
               logger.error("uncaught exception", { originalError: error });
 
-              socket.emit("INDEXING_FAILED", {
-                version: "v1",
-                deploymentId: this.deploymentId,
-                error,
-              });
+              failIndex(error);
             } else if (e instanceof Error) {
               const error = {
                 name: e.name,
@@ -690,34 +695,22 @@ class ProdWorker {
 
               logger.error("error", { error });
 
-              socket.emit("INDEXING_FAILED", {
-                version: "v1",
-                deploymentId: this.deploymentId,
-                error,
-              });
+              failIndex(error);
             } else if (typeof e === "string") {
               logger.error("string error", { error: { message: e } });
 
-              socket.emit("INDEXING_FAILED", {
-                version: "v1",
-                deploymentId: this.deploymentId,
-                error: {
-                  name: "Error",
-                  message: e,
-                  stderr,
-                },
+              failIndex({
+                name: "Error",
+                message: e,
+                stderr,
               });
             } else {
               logger.error("unknown error", { error: e });
 
-              socket.emit("INDEXING_FAILED", {
-                version: "v1",
-                deploymentId: this.deploymentId,
-                error: {
-                  name: "Error",
-                  message: "Unknown error",
-                  stderr,
-                },
+              failIndex({
+                name: "Error",
+                message: "Unknown error",
+                stderr,
               });
             }
 
diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts
index f2327540e..aceb1001d 100644
--- a/packages/core/src/v3/schemas/messages.ts
+++ b/packages/core/src/v3/schemas/messages.ts
@@ -347,7 +347,7 @@ export const ProviderToPlatformMessages = {
         stack: z.string().optional(),
         stderr: z.string().optional(),
       }),
-      overrideCompletion: z.string().optional(),
+      overrideCompletion: z.boolean().optional(),
     }),
   },
 };