diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index b072e578..07bfc02c 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -3,6 +3,10 @@ import { CancelledNotificationSchema, ClientCapabilities, ErrorCode, + isJSONRPCError, + isJSONRPCRequest, + isJSONRPCResponse, + isJSONRPCNotification, JSONRPCError, JSONRPCNotification, JSONRPCRequest, @@ -271,12 +275,14 @@ export abstract class Protocol< }; this._transport.onmessage = (message) => { - if (!("method" in message)) { + if (isJSONRPCResponse(message) || isJSONRPCError(message)) { this._onresponse(message); - } else if ("id" in message) { + } else if (isJSONRPCRequest(message)) { this._onrequest(message); - } else { + } else if (isJSONRPCNotification(message)) { this._onnotification(message); + } else { + this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); } }; @@ -437,7 +443,7 @@ export abstract class Protocol< this._progressHandlers.delete(messageId); this._cleanupTimeout(messageId); - if ("result" in response) { + if (isJSONRPCResponse(response)) { handler(response); } else { const error = new McpError( diff --git a/src/types.ts b/src/types.ts index da4afd2a..db6bf125 100644 --- a/src/types.ts +++ b/src/types.ts @@ -78,6 +78,9 @@ export const JSONRPCRequestSchema = z .merge(RequestSchema) .strict(); +export const isJSONRPCRequest = (value: unknown): value is JSONRPCRequest => + JSONRPCRequestSchema.safeParse(value).success; + /** * A notification which does not expect a response. */ @@ -88,6 +91,11 @@ export const JSONRPCNotificationSchema = z .merge(NotificationSchema) .strict(); +export const isJSONRPCNotification = ( + value: unknown +): value is JSONRPCNotification => + JSONRPCNotificationSchema.safeParse(value).success; + /** * A successful (non-error) response to a request. */ @@ -99,6 +107,9 @@ export const JSONRPCResponseSchema = z }) .strict(); +export const isJSONRPCResponse = (value: unknown): value is JSONRPCResponse => + JSONRPCResponseSchema.safeParse(value).success; + /** * Error codes defined by the JSON-RPC specification. */ @@ -139,6 +150,9 @@ export const JSONRPCErrorSchema = z }) .strict(); +export const isJSONRPCError = (value: unknown): value is JSONRPCError => + JSONRPCErrorSchema.safeParse(value).success; + export const JSONRPCMessageSchema = z.union([ JSONRPCRequestSchema, JSONRPCNotificationSchema,