diff --git a/src/client/sse.ts b/src/client/sse.ts index 7b3a128a..dd06b77f 100644 --- a/src/client/sse.ts +++ b/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; diff --git a/src/client/websocket.ts b/src/client/websocket.ts index bd6039b4..47856cc7 100644 --- a/src/client/websocket.ts +++ b/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; diff --git a/src/server/sse.ts b/src/server/sse.ts index e4c08a2c..86d5fe59 100644 --- a/src/server/sse.ts +++ b/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 { + async handleMessage(message: unknown): Promise { + 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 { diff --git a/src/shared/message.ts b/src/shared/message.ts deleted file mode 100644 index 25162b1e..00000000 --- a/src/shared/message.ts +++ /dev/null @@ -1,10 +0,0 @@ -import { JSONRPCMessage } from "../types.js"; - -// Throws an exception if the message is invalid or unsupported. -export function validateMessage(message: JSONRPCMessage): void { - if (message.jsonrpc !== "2.0") { - throw new Error( - `Message is not in a supported JSON-RPC version: ${message}`, - ); - } -} diff --git a/src/shared/stdio.ts b/src/shared/stdio.ts index c9a51d40..a30e90bc 100644 --- a/src/shared/stdio.ts +++ b/src/shared/stdio.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 {