diff --git a/src/client/index.ts b/src/client/index.ts index 61cd6ea6..c8c729bd 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -5,12 +5,9 @@ import { ClientRequest, ClientResult, Implementation, - InitializeResult, + InitializeResultSchema, PROTOCOL_VERSION, ServerCapabilities, - ServerNotification, - ServerRequest, - ServerResult, } from "../types.js"; /** @@ -19,9 +16,6 @@ import { * The client will automatically begin the initialization flow with the server when connect() is called. */ export class Client extends Protocol< - ServerRequest, - ServerNotification, - ServerResult, ClientRequest, ClientNotification, ClientResult @@ -39,14 +33,17 @@ export class Client extends Protocol< override async connect(transport: Transport): Promise { await super.connect(transport); - const result = (await this.request({ - method: "initialize", - params: { - protocolVersion: 1, - capabilities: {}, - clientInfo: this._clientInfo, + const result = await this.request( + { + method: "initialize", + params: { + protocolVersion: 1, + capabilities: {}, + clientInfo: this._clientInfo, + }, }, - })) as InitializeResult; + InitializeResultSchema, + ); if (result === undefined) { throw new Error(`Server sent invalid initialize result: ${result}`); diff --git a/src/server/index.ts b/src/server/index.ts index a17de0c2..01af0837 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -1,11 +1,10 @@ import { Protocol } from "../shared/protocol.js"; import { ClientCapabilities, - ClientNotification, - ClientRequest, - ClientResult, Implementation, + InitializedNotificationSchema, InitializeRequest, + InitializeRequestSchema, InitializeResult, PROTOCOL_VERSION, ServerNotification, @@ -19,9 +18,6 @@ import { * This server will automatically respond to the initialization flow as initiated from the client. */ export class Server extends Protocol< - ClientRequest, - ClientNotification, - ClientResult, ServerRequest, ServerNotification, ServerResult @@ -40,10 +36,10 @@ export class Server extends Protocol< constructor(private _serverInfo: Implementation) { super(); - this.setRequestHandler("initialize", (request) => - this._oninitialize(request as InitializeRequest), + this.setRequestHandler(InitializeRequestSchema, (request) => + this._oninitialize(request), ); - this.setNotificationHandler("notification/initialized", () => + this.setNotificationHandler(InitializedNotificationSchema, () => this.oninitialized?.(), ); } diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index ec1083d3..96c1ae95 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -1,3 +1,4 @@ +import { AnyZodObject, ZodLiteral, ZodObject, z } from "zod"; import { ErrorCode, JSONRPCError, @@ -25,9 +26,6 @@ export type ProgressCallback = (progress: Progress) => void; * features like request/response linking, notifications, and progress. */ export class Protocol< - ReceiveRequestT extends Request, - ReceiveNotificationT extends Notification, - ReceiveResultT extends Result, SendRequestT extends Request, SendNotificationT extends Notification, SendResultT extends Result, @@ -36,15 +34,15 @@ export class Protocol< private _requestMessageId = 0; private _requestHandlers: Map< string, - (request: ReceiveRequestT) => Promise + (request: JSONRPCRequest) => Promise > = new Map(); private _notificationHandlers: Map< string, - (notification: ReceiveNotificationT) => Promise + (notification: JSONRPCNotification) => Promise > = new Map(); private _responseHandlers: Map< number, - (response: ReceiveResultT | Error) => void + (response: JSONRPCResponse | Error) => void > = new Map(); private _progressHandlers: Map = new Map(); @@ -65,25 +63,20 @@ export class Protocol< /** * A handler to invoke for any request types that do not have their own handler installed. */ - fallbackRequestHandler?: (request: ReceiveRequestT) => Promise; + fallbackRequestHandler?: (request: Request) => Promise; /** * A handler to invoke for any notification types that do not have their own handler installed. */ - fallbackNotificationHandler?: ( - notification: ReceiveNotificationT, - ) => Promise; + fallbackNotificationHandler?: (notification: Notification) => Promise; constructor() { - this.setNotificationHandler( - ProgressNotificationSchema.shape.method.value, - (notification) => { - this._onprogress(notification as unknown as ProgressNotification); - }, - ); + this.setNotificationHandler(ProgressNotificationSchema, (notification) => { + this._onprogress(notification as unknown as ProgressNotification); + }); this.setRequestHandler( - PingRequestSchema.shape.method.value, + PingRequestSchema, // Automatic pong by default. (_request) => ({}) as SendResultT, ); @@ -106,11 +99,11 @@ export class Protocol< this._transport.onmessage = (message) => { if (!("method" in message)) { - this._onresponse(message as JSONRPCResponse | JSONRPCError); + this._onresponse(message); } else if ("id" in message) { - this._onrequest(message as JSONRPCRequest); + this._onrequest(message); } else { - this._onnotification(message as JSONRPCNotification); + this._onnotification(message); } }; } @@ -142,7 +135,7 @@ export class Protocol< return; } - handler(notification as unknown as ReceiveNotificationT).catch((error) => + handler(notification).catch((error) => this._onerror( new Error(`Uncaught error in notification handler: ${error}`), ), @@ -171,7 +164,7 @@ export class Protocol< return; } - handler(request as unknown as ReceiveRequestT) + handler(request) .then( (result) => { this._transport?.send({ @@ -228,7 +221,7 @@ export class Protocol< this._responseHandlers.delete(Number(messageId)); this._progressHandlers.delete(Number(messageId)); if ("result" in response) { - handler(response.result as ReceiveResultT); + handler(response); } else { const error = new McpError( response.error.code, @@ -255,11 +248,11 @@ export class Protocol< * * Do not use this method to emit notifications! Use notification() instead. */ - // TODO: This could infer a better response type based on the method - request( + request( request: SendRequestT, + resultSchema: T, onprogress?: ProgressCallback, - ): Promise { + ): Promise> { return new Promise((resolve, reject) => { if (!this._transport) { reject(new Error("Not connected")); @@ -283,9 +276,14 @@ export class Protocol< this._responseHandlers.set(messageId, (response) => { if (response instanceof Error) { - reject(response); - } else { - resolve(response); + return reject(response); + } + + try { + const result = resultSchema.parse(response.result); + resolve(result); + } catch (error) { + reject(error); } }); @@ -314,13 +312,16 @@ export class Protocol< * * Note that this will replace any previous request handler for the same method. */ - // TODO: This could infer a better request type based on the method. - setRequestHandler( - method: string, - handler: (request: ReceiveRequestT) => SendResultT | Promise, + setRequestHandler< + T extends ZodObject<{ + method: ZodLiteral; + }>, + >( + requestSchema: T, + handler: (request: z.infer) => SendResultT | Promise, ): void { - this._requestHandlers.set(method, (request) => - Promise.resolve(handler(request)), + this._requestHandlers.set(requestSchema.shape.method.value, (request) => + Promise.resolve(handler(requestSchema.parse(request))), ); } @@ -336,13 +337,18 @@ export class Protocol< * * Note that this will replace any previous notification handler for the same method. */ - // TODO: This could infer a better notification type based on the method. - setNotificationHandler( - method: string, - handler: (notification: T) => void | Promise, + setNotificationHandler< + T extends ZodObject<{ + method: ZodLiteral; + }>, + >( + notificationSchema: T, + handler: (notification: z.infer) => void | Promise, ): void { - this._notificationHandlers.set(method, (notification) => - Promise.resolve(handler(notification as T)), + this._notificationHandlers.set( + notificationSchema.shape.method.value, + (notification) => + Promise.resolve(handler(notificationSchema.parse(notification))), ); }