Merge pull request #2 from callumfrederiksen/codex/modify-message-parsing-logic-in-zodmessagehandler

Author: callumfrederiksen

packages/core/src/v3/zodMessageHandler.test.ts

@@ -0,0 +1,36 @@
+import { describe, it, expect, vi } from "vitest";
+import { EventEmitter } from "node:events";
+import { ZodMessageHandler } from "./zodMessageHandler.js";
+import { z } from "zod";
+
+describe("ZodMessageHandler.registerHandlers", () => {
+  const schema = { TEST: z.object({ foo: z.string() }) } as const;
+
+  it("handles messages with an explicit payload field", async () => {
+    const handler = vi.fn(async (payload: { foo: string }) => payload);
+    const messageHandler = new ZodMessageHandler({ schema, messages: { TEST: handler } });
+    const emitter = new EventEmitter();
+    messageHandler.registerHandlers(emitter);
+
+    const ack = await new Promise((resolve) => {
+      emitter.emit("TEST", { payload: { foo: "bar" }, version: "v1" }, resolve);
+    });
+
+    expect(handler).toHaveBeenCalledWith({ foo: "bar" });
+    expect(ack).toEqual({ foo: "bar" });
+  });
+
+  it("handles messages without a payload field", async () => {
+    const handler = vi.fn(async (payload: { foo: string }) => payload);
+    const messageHandler = new ZodMessageHandler({ schema, messages: { TEST: handler } });
+    const emitter = new EventEmitter();
+    messageHandler.registerHandlers(emitter);
+
+    const ack = await new Promise((resolve) => {
+      emitter.emit("TEST", { foo: "baz", version: "v1" }, resolve);
+    });
+
+    expect(handler).toHaveBeenCalledWith({ foo: "baz" });
+    expect(ack).toEqual({ foo: "baz" });
+  });
+});

packages/core/src/v3/zodMessageHandler.ts

@@ -179,12 +179,19 @@ export class ZodMessageHandler<TMessageCatalog extends ZodMessageCatalogSchema>
 
         let ack: Awaited<ReturnType<ZodMessageHandler<TMessageCatalog>["handleMessage"]>>;
 
-        // FIXME: this only works if the message doesn't have genuine payload prop
-        if ("payload" in message) {
-          ack = await this.handleMessage({ type: eventName, ...message });
+        // Use runtime validation to detect payload presence
+        const hasPayload =
+          typeof message === "object" &&
+          message !== null &&
+          z.object({ payload: z.unknown() }).passthrough().safeParse(message).success;
+
+        if (hasPayload) {
+          ack = await this.handleMessage({ type: eventName, ...(message as any) });
         } else {
           // Handle messages not sent by ZodMessageSender
-          const { version, ...payload } = message;
+          const messageObj =
+            typeof message === "object" && message !== null ? (message as any) : {};
+          const { version, ...payload } = messageObj;
           ack = await this.handleMessage({ type: eventName, version, payload });
         }