prevent uncaught websocket exceptions

nicktrn · nicktrn · commit a974e8da39bb · 2024-07-16T09:52:55.000+01:00
diff --git a/apps/webapp/app/v3/authenticatedSocketConnection.server.ts b/apps/webapp/app/v3/authenticatedSocketConnection.server.ts
@@ -52,6 +52,7 @@ export class AuthenticatedSocketConnection {
 
     this._messageHandler = new ZodMessageHandler({
       schema: clientWebsocketMessages,
+      logger,
       messages: {
         READY_FOR_TASKS: async (payload) => {
           await this._consumer.registerBackgroundWorker(
@@ -99,9 +100,22 @@ export class AuthenticatedSocketConnection {
   }
 
   async #handleMessage(ev: MessageEvent) {
-    const data = JSON.parse(ev.data.toString());
-
-    await this._messageHandler.handleMessage(data);
+    try {
+      const data = JSON.parse(ev.data.toString());
+
+      await this._messageHandler.handleMessage(data);
+    } catch (error) {
+      logger.error("[AuthenticatedSocketConnection] Failed to handle message", {
+        error:
+          error instanceof Error
+            ? {
+                message: error.message,
+                stack: error.stack,
+              }
+            : error,
+        message: ev.data.toString(),
+      });
+    }
   }
 
   async #handleClose(ev: CloseEvent) {
diff --git a/apps/webapp/app/v3/sharedSocketConnection.ts b/apps/webapp/app/v3/sharedSocketConnection.ts
@@ -96,6 +96,7 @@ export class SharedSocketConnection {
 
     this._messageHandler = new ZodMessageHandler({
       schema: clientWebsocketMessages,
+      logger,
       messages: {
         READY_FOR_TASKS: async (payload) => {
           this._sharedQueueConsumerPool.start();
diff --git a/apps/webapp/app/v3/utils/zodPubSub.server.ts b/apps/webapp/app/v3/utils/zodPubSub.server.ts
@@ -39,6 +39,7 @@ class RedisZodSubscriber<TMessageCatalog extends ZodMessageCatalogSchema>
     this._subscriber = new Redis(_options.redis);
     this._messageHandler = new ZodMessageHandler({
       schema: _options.schema,
+      logger: this._logger,
     });
   }
 
@@ -76,20 +77,26 @@ class RedisZodSubscriber<TMessageCatalog extends ZodMessageCatalogSchema>
 
     const message = this._messageHandler.parseMessage(parsedMessage);
 
-    if (typeof message.type !== "string") {
+    if (!message.success) {
+      this._logger.error(`Failed to parse message: ${message.error}`, { parsedMessage });
       return;
     }
 
-    const listener = this._listeners.get(message.type);
+    if (typeof message.data.type !== "string") {
+      this._logger.error(`Failed to parse message: invalid type`, { parsedMessage });
+      return;
+    }
+
+    const listener = this._listeners.get(message.data.type);
 
     if (!listener) {
-      this._logger.debug(`No listener for message type: ${message.type}`, { parsedMessage });
+      this._logger.debug(`No listener for message type: ${message.data.type}`, { parsedMessage });
 
       return;
     }
 
     try {
-      await listener(message.payload);
+      await listener(message.data.payload);
     } catch (error) {
       this._logger.error("Error handling message", { error, message });
     }
diff --git a/packages/cli-v3/src/workers/dev/backgroundWorker.ts b/packages/cli-v3/src/workers/dev/backgroundWorker.ts
@@ -379,20 +379,32 @@ export class BackgroundWorker {
       child.on("message", async (msg: any) => {
         const message = this._handler.parseMessage(msg);
 
-        if (message.type === "TASKS_READY" && !resolved) {
+        if (!message.success) {
           clearTimeout(timeout);
           resolved = true;
-          resolve(message.payload.tasks);
+          reject(new Error(`Failed to parse message: ${message.error}`));
           child.kill();
-        } else if (message.type === "UNCAUGHT_EXCEPTION") {
+          return;
+        }
+
+        if (message.data.type === "TASKS_READY" && !resolved) {
+          clearTimeout(timeout);
+          resolved = true;
+          resolve(message.data.payload.tasks);
+          child.kill();
+        } else if (message.data.type === "UNCAUGHT_EXCEPTION") {
           clearTimeout(timeout);
           resolved = true;
-          reject(new UncaughtExceptionError(message.payload.error, message.payload.origin));
+          reject(
+            new UncaughtExceptionError(message.data.payload.error, message.data.payload.origin)
+          );
           child.kill();
-        } else if (message.type === "TASKS_FAILED_TO_PARSE") {
+        } else if (message.data.type === "TASKS_FAILED_TO_PARSE") {
           clearTimeout(timeout);
           resolved = true;
-          reject(new TaskMetadataParseError(message.payload.zodIssues, message.payload.tasks));
+          reject(
+            new TaskMetadataParseError(message.data.payload.zodIssues, message.data.payload.tasks)
+          );
           child.kill();
         }
       });
@@ -942,9 +954,14 @@ class TaskRunProcess {
   async #handleMessage(msg: any) {
     const message = this._handler.parseMessage(msg);
 
-    switch (message.type) {
+    if (!message.success) {
+      logger.error(`Dropping message: ${message.error}`, { message });
+      return;
+    }
+
+    switch (message.data.type) {
       case "TASK_RUN_COMPLETED": {
-        const { result, execution } = message.payload;
+        const { result, execution } = message.data.payload;
 
         logger.debug(`[${this.runId}] task run completed`, {
           result,
@@ -981,7 +998,7 @@ class TaskRunProcess {
         if (this.messageId) {
           this.onTaskRunHeartbeat.post(this.messageId);
         } else {
-          this.onTaskHeartbeat.post(message.payload.id);
+          this.onTaskHeartbeat.post(message.data.payload.id);
         }
 
         break;
diff --git a/packages/core/src/v3/zodMessageHandler.ts b/packages/core/src/v3/zodMessageHandler.ts
@@ -25,6 +25,7 @@ export type ZodMessageHandlers<TCatalogSchema extends ZodMessageCatalogSchema> =
 export type ZodMessageHandlerOptions<TMessageCatalog extends ZodMessageCatalogSchema> = {
   schema: TMessageCatalog;
   messages?: ZodMessageHandlers<TMessageCatalog>;
+  logger?: StructuredLogger;
 };
 
 export type MessageFromSchema<
@@ -52,53 +53,107 @@ export interface EventEmitterLike {
 export class ZodMessageHandler<TMessageCatalog extends ZodMessageCatalogSchema> {
   #schema: TMessageCatalog;
   #handlers: ZodMessageHandlers<TMessageCatalog> | undefined;
+  #logger: StructuredLogger | Console;
 
   constructor(options: ZodMessageHandlerOptions<TMessageCatalog>) {
     this.#schema = options.schema;
     this.#handlers = options.messages;
+    this.#logger = options.logger ?? console;
   }
 
-  public async handleMessage(message: unknown) {
+  public async handleMessage(message: unknown): Promise<
+    | {
+        success: true;
+        data: unknown;
+      }
+    | {
+        success: false;
+        error: string;
+      }
+  > {
     const parsedMessage = this.parseMessage(message);
 
+    if (!parsedMessage.success) {
+      this.#logger.error(parsedMessage.error, { message });
+
+      return {
+        success: false,
+        error: parsedMessage.error,
+      };
+    }
+
     if (!this.#handlers) {
-      throw new Error("No handlers provided");
+      this.#logger.error("No handlers provided", { message });
+
+      return {
+        success: false,
+        error: "No handlers provided",
+      };
     }
 
-    const handler = this.#handlers[parsedMessage.type];
+    const handler = this.#handlers[parsedMessage.data.type];
 
     if (!handler) {
-      console.error(`No handler for message type: ${String(parsedMessage.type)}`);
-      return;
+      const error = `No handler for message type: ${String(parsedMessage.data.type)}`;
+
+      this.#logger.error(error, { message });
+
+      return {
+        success: false,
+        error,
+      };
     }
 
-    const ack = await handler(parsedMessage.payload);
+    const ack = await handler(parsedMessage.data.payload);
 
-    return ack;
+    return {
+      success: true,
+      data: ack,
+    };
   }
 
-  public parseMessage(message: unknown): MessageFromCatalog<TMessageCatalog> {
+  public parseMessage(message: unknown):
+    | {
+        success: true;
+        data: MessageFromCatalog<TMessageCatalog>;
+      }
+    | {
+        success: false;
+        error: string;
+      } {
     const parsedMessage = ZodMessageSchema.safeParse(message);
 
     if (!parsedMessage.success) {
-      throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`);
+      return {
+        success: false,
+        error: `Failed to parse message: ${JSON.stringify(parsedMessage.error)}`,
+      };
     }
 
     const schema = this.#schema[parsedMessage.data.type];
 
     if (!schema) {
-      throw new Error(`Unknown message type: ${parsedMessage.data.type}`);
+      return {
+        success: false,
+        error: `Unknown message type: ${parsedMessage.data.type}`,
+      };
     }
 
     const parsedPayload = schema.safeParse(parsedMessage.data.payload);
 
     if (!parsedPayload.success) {
-      throw new Error(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
+      return {
+        success: false,
+        error: `Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`,
+      };
     }
 
     return {
-      type: parsedMessage.data.type,
-      payload: parsedPayload.data,
+      success: true,
+      data: {
+        type: parsedMessage.data.type,
+        payload: parsedPayload.data,
+      },
     };
   }
 
@@ -117,7 +172,7 @@ export class ZodMessageHandler<TMessageCatalog extends ZodMessageCatalogSchema>
           hasCallback: !!callback,
         });
 
-        let ack;
+        let ack: Awaited<ReturnType<ZodMessageHandler<TMessageCatalog>["handleMessage"]>>;
 
         // FIXME: this only works if the message doesn't have genuine payload prop
         if ("payload" in message) {
@@ -129,7 +184,13 @@ export class ZodMessageHandler<TMessageCatalog extends ZodMessageCatalogSchema>
         }
 
         if (callback && typeof callback === "function") {
-          callback(ack);
+          if (!ack.success) {
+            // We don't know the callback type, so we can't do anything else - not all callbacks may accept a success prop
+            log.error("Failed to handle message, skipping callback", { message, error: ack.error });
+            return;
+          }
+
+          callback(ack.data);
         }
       });
     }
