Use Zod for message validation at transport layer

Author: jspahrsummers

src/client/sse.ts

@@ -1,6 +1,5 @@
-import { validateMessage } from "../shared/message.js";
 import { Transport } from "../shared/transport.js";
-import { JSONRPCMessage } from "../types.js";
+import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
 
 /**
  * Client transport for SSE: this will connect to a server using Server-Sent Events for receiving
@@ -57,8 +56,7 @@ export class SSEClientTransport implements Transport {
         const messageEvent = event as MessageEvent;
         let message: JSONRPCMessage;
         try {
-          message = JSON.parse(messageEvent.data);
-          validateMessage(message);
+          message = JSONRPCMessageSchema.parse(JSON.parse(messageEvent.data));
         } catch (error) {
           this.onerror?.(error as Error);
           return;

src/client/websocket.ts

@@ -1,6 +1,5 @@
-import { validateMessage } from "../shared/message.js";
 import { Transport } from "../shared/transport.js";
-import { JSONRPCMessage } from "../types.js";
+import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
 
 const SUBPROTOCOL = "mcp";
 
@@ -38,8 +37,7 @@ export class WebSocketClientTransport implements Transport {
       this._socket.onmessage = (event: MessageEvent) => {
         let message: JSONRPCMessage;
         try {
-          message = JSON.parse(event.data);
-          validateMessage(message);
+          message = JSONRPCMessageSchema.parse(JSON.parse(event.data));
         } catch (error) {
           this.onerror?.(error as Error);
           return;

src/server/sse.ts

@@ -1,8 +1,7 @@
 import { randomUUID } from "node:crypto";
 import { IncomingMessage, ServerResponse } from "node:http";
-import { validateMessage } from "../shared/message.js";
 import { Transport } from "../shared/transport.js";
-import { JSONRPCMessage } from "../types.js";
+import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
 import getRawBody from "raw-body";
 import contentType from "content-type";
 
@@ -101,15 +100,16 @@ export class SSEServerTransport implements Transport {
   /**
    * Handle a client message, regardless of how it arrived. This can be used to inform the server of messages that arrive via a means different than HTTP POST.
    */
-  async handleMessage(message: JSONRPCMessage): Promise<void> {
+  async handleMessage(message: unknown): Promise<void> {
+    let parsedMessage: JSONRPCMessage;
     try {
-      validateMessage(message);
+      parsedMessage = JSONRPCMessageSchema.parse(message);
     } catch (error) {
       this.onerror?.(error as Error);
       throw error;
     }
 
-    this.onmessage?.(message);
+    this.onmessage?.(parsedMessage);
   }
 
   async close(): Promise<void> {

src/shared/message.ts

@@ -1,5 +1,4 @@
-import { JSONRPCMessage } from "../types.js";
-import { validateMessage } from "./message.js";
+import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
 
 /**
  * Buffers a continuous stdio stream into discrete JSON-RPC messages.
@@ -32,9 +31,7 @@ export class ReadBuffer {
 }
 
 export function deserializeMessage(line: string): JSONRPCMessage {
-  const message = JSON.parse(line) as JSONRPCMessage;
-  validateMessage(message);
-  return message;
+  return JSONRPCMessageSchema.parse(JSON.parse(line));
 }
 
 export function serializeMessage(message: JSONRPCMessage): string {