From 7cc7e81387b56bda5db23b82ef31721172ae9909 Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Mon, 4 Nov 2024 21:03:57 +0000 Subject: [PATCH 1/6] Support custom client request, notification, and result types --- src/client/index.test.ts | 71 ++++++++++++++++++++++++++++++++++++++++ src/client/index.ts | 35 +++++++++++++++++--- 2 files changed, 102 insertions(+), 4 deletions(-) create mode 100644 src/client/index.test.ts diff --git a/src/client/index.test.ts b/src/client/index.test.ts new file mode 100644 index 0000000..edb0057 --- /dev/null +++ b/src/client/index.test.ts @@ -0,0 +1,71 @@ +import { Client } from "./index.js"; +import { z } from "zod"; +import { RequestSchema, NotificationSchema, ResultSchema } from "../types.js"; + +/* +Test that custom request/notification/result schemas can be used with the Client class. +*/ +const GetWeatherRequestSchema = RequestSchema.extend({ + method: z.literal("weather/get"), + params: z.object({ + city: z.string(), + }), +}); + +const GetForecastRequestSchema = RequestSchema.extend({ + method: z.literal("weather/forecast"), + params: z.object({ + city: z.string(), + days: z.number(), + }), +}); + +const WeatherForecastNotificationSchema = NotificationSchema.extend({ + method: z.literal("weather/alert"), + params: z.object({ + severity: z.enum(["warning", "watch"]), + message: z.string(), + }), +}); + +const WeatherRequestSchema = GetWeatherRequestSchema.or( + GetForecastRequestSchema, +); +const WeatherNotificationSchema = WeatherForecastNotificationSchema; +const WeatherResultSchema = ResultSchema.extend({ + temperature: z.number(), + conditions: z.string(), +}); + +type WeatherRequest = z.infer; +type WeatherNotification = z.infer; +type WeatherResult = z.infer; + +// Create a typed Client for weather data +const weatherClient = new Client< + WeatherRequest, + WeatherNotification, + WeatherResult +>({ + name: "WeatherClient", + version: "1.0.0", +}); + +// Typecheck that only valid weather requests/notifications/results are allowed +weatherClient.request( + { + method: "weather/get", + params: { + city: "Seattle", + }, + }, + WeatherResultSchema, +); + +weatherClient.notification({ + method: "weather/alert", + params: { + severity: "warning", + message: "Storm approaching", + }, +}); diff --git a/src/client/index.ts b/src/client/index.ts index 9b9aa81..e4ae4da 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -6,7 +6,10 @@ import { ClientResult, Implementation, InitializeResultSchema, + Notification, PROTOCOL_VERSION, + Request, + Result, ServerCapabilities, } from "../types.js"; @@ -14,11 +17,35 @@ import { * An MCP client on top of a pluggable transport. * * The client will automatically begin the initialization flow with the server when connect() is called. + * + * To use with custom types, extend the base Request/Notification/Result types and pass them as type parameters: + * + * ```typescript + * // Custom schemas + * const CustomRequestSchema = RequestSchema.extend({...}) + * const CustomNotificationSchema = NotificationSchema.extend({...}) + * const CustomResultSchema = ResultSchema.extend({...}) + * + * // Type aliases + * type CustomRequest = z.infer + * type CustomNotification = z.infer + * type CustomResult = z.infer + * + * // Create typed client + * const client = new Client({ + * name: "CustomClient", + * version: "1.0.0" + * }) + * ``` */ -export class Client extends Protocol< - ClientRequest, - ClientNotification, - ClientResult +export class Client< + RequestT extends Request = Request, + NotificationT extends Notification = Notification, + ResultT extends Result = Result, +> extends Protocol< + ClientRequest | RequestT, + ClientNotification | NotificationT, + ClientResult | ResultT > { private _serverCapabilities?: ServerCapabilities; private _serverVersion?: Implementation; From b5880ec18af187e28e74e1e959e55cd631af8292 Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Mon, 4 Nov 2024 21:07:53 +0000 Subject: [PATCH 2/6] Support custom server request, notification, and result types --- src/server/index.test.ts | 67 ++++++++++++++++++++++++++++++++++++++++ src/server/index.ts | 35 ++++++++++++++++++--- 2 files changed, 98 insertions(+), 4 deletions(-) create mode 100644 src/server/index.test.ts diff --git a/src/server/index.test.ts b/src/server/index.test.ts new file mode 100644 index 0000000..7a06c53 --- /dev/null +++ b/src/server/index.test.ts @@ -0,0 +1,67 @@ +import { Server } from "./index.js"; +import { z } from "zod"; +import { RequestSchema, NotificationSchema, ResultSchema } from "../types.js"; + +/* +Test that custom request/notification/result schemas can be used with the Server class. +*/ +const GetWeatherRequestSchema = RequestSchema.extend({ + method: z.literal("weather/get"), + params: z.object({ + city: z.string(), + }), +}); + +const GetForecastRequestSchema = RequestSchema.extend({ + method: z.literal("weather/forecast"), + params: z.object({ + city: z.string(), + days: z.number(), + }), +}); + +const WeatherForecastNotificationSchema = NotificationSchema.extend({ + method: z.literal("weather/alert"), + params: z.object({ + severity: z.enum(["warning", "watch"]), + message: z.string(), + }), +}); + +const WeatherRequestSchema = GetWeatherRequestSchema.or( + GetForecastRequestSchema, +); +const WeatherNotificationSchema = WeatherForecastNotificationSchema; +const WeatherResultSchema = ResultSchema.extend({ + temperature: z.number(), + conditions: z.string(), +}); + +type WeatherRequest = z.infer; +type WeatherNotification = z.infer; +type WeatherResult = z.infer; + +// Create a typed Server for weather data +const weatherServer = new Server< + WeatherRequest, + WeatherNotification, + WeatherResult +>({ + name: "WeatherServer", + version: "1.0.0", +}); + +// Typecheck that only valid weather requests/notifications/results are allowed +weatherServer.setRequestHandler(GetWeatherRequestSchema, (request) => { + return { + temperature: 72, + conditions: "sunny", + }; +}); + +weatherServer.setNotificationHandler( + WeatherForecastNotificationSchema, + (notification) => { + console.log(`Weather alert: ${notification.params.message}`); + }, +); diff --git a/src/server/index.ts b/src/server/index.ts index c8ffaba..71c22cc 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -6,7 +6,10 @@ import { InitializeRequest, InitializeRequestSchema, InitializeResult, + Notification, PROTOCOL_VERSION, + Request, + Result, ServerNotification, ServerRequest, ServerResult, @@ -21,11 +24,35 @@ import { * An MCP server on top of a pluggable transport. * * This server will automatically respond to the initialization flow as initiated from the client. + * + * To use with custom types, extend the base Request/Notification/Result types and pass them as type parameters: + * + * ```typescript + * // Custom schemas + * const CustomRequestSchema = RequestSchema.extend({...}) + * const CustomNotificationSchema = NotificationSchema.extend({...}) + * const CustomResultSchema = ResultSchema.extend({...}) + * + * // Type aliases + * type CustomRequest = z.infer + * type CustomNotification = z.infer + * type CustomResult = z.infer + * + * // Create typed server + * const server = new Server({ + * name: "CustomServer", + * version: "1.0.0" + * }) + * ``` */ -export class Server extends Protocol< - ServerRequest, - ServerNotification, - ServerResult +export class Server< + RequestT extends Request = Request, + NotificationT extends Notification = Notification, + ResultT extends Result = Result, +> extends Protocol< + ServerRequest | RequestT, + ServerNotification | NotificationT, + ServerResult | ResultT > { private _clientCapabilities?: ClientCapabilities; private _clientVersion?: Implementation; From 9774e2fdfe74b44490fe31bc8f36e59d9cccfbd0 Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Mon, 4 Nov 2024 21:10:35 +0000 Subject: [PATCH 3/6] Disable unused vars lint --- src/client/index.test.ts | 1 + src/server/index.test.ts | 1 + 2 files changed, 2 insertions(+) diff --git a/src/client/index.test.ts b/src/client/index.test.ts index edb0057..c4de6d7 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -1,3 +1,4 @@ +/* eslint-disable @typescript-eslint/no-unused-vars */ import { Client } from "./index.js"; import { z } from "zod"; import { RequestSchema, NotificationSchema, ResultSchema } from "../types.js"; diff --git a/src/server/index.test.ts b/src/server/index.test.ts index 7a06c53..4a9ac59 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -1,3 +1,4 @@ +/* eslint-disable @typescript-eslint/no-unused-vars */ import { Server } from "./index.js"; import { z } from "zod"; import { RequestSchema, NotificationSchema, ResultSchema } from "../types.js"; From 202626066868f34f4c935ae9e315286a0b834e75 Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Mon, 4 Nov 2024 21:12:17 +0000 Subject: [PATCH 4/6] Wrap into one test (required) --- src/client/index.test.ts | 112 ++++++++++++++++++++------------------- src/server/index.test.ts | 108 +++++++++++++++++++------------------ 2 files changed, 112 insertions(+), 108 deletions(-) diff --git a/src/client/index.test.ts b/src/client/index.test.ts index c4de6d7..1dc5ead 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -6,67 +6,69 @@ import { RequestSchema, NotificationSchema, ResultSchema } from "../types.js"; /* Test that custom request/notification/result schemas can be used with the Client class. */ -const GetWeatherRequestSchema = RequestSchema.extend({ - method: z.literal("weather/get"), - params: z.object({ - city: z.string(), - }), -}); +test("should typecheck", () => { + const GetWeatherRequestSchema = RequestSchema.extend({ + method: z.literal("weather/get"), + params: z.object({ + city: z.string(), + }), + }); -const GetForecastRequestSchema = RequestSchema.extend({ - method: z.literal("weather/forecast"), - params: z.object({ - city: z.string(), - days: z.number(), - }), -}); + const GetForecastRequestSchema = RequestSchema.extend({ + method: z.literal("weather/forecast"), + params: z.object({ + city: z.string(), + days: z.number(), + }), + }); -const WeatherForecastNotificationSchema = NotificationSchema.extend({ - method: z.literal("weather/alert"), - params: z.object({ - severity: z.enum(["warning", "watch"]), - message: z.string(), - }), -}); + const WeatherForecastNotificationSchema = NotificationSchema.extend({ + method: z.literal("weather/alert"), + params: z.object({ + severity: z.enum(["warning", "watch"]), + message: z.string(), + }), + }); -const WeatherRequestSchema = GetWeatherRequestSchema.or( - GetForecastRequestSchema, -); -const WeatherNotificationSchema = WeatherForecastNotificationSchema; -const WeatherResultSchema = ResultSchema.extend({ - temperature: z.number(), - conditions: z.string(), -}); + const WeatherRequestSchema = GetWeatherRequestSchema.or( + GetForecastRequestSchema, + ); + const WeatherNotificationSchema = WeatherForecastNotificationSchema; + const WeatherResultSchema = ResultSchema.extend({ + temperature: z.number(), + conditions: z.string(), + }); -type WeatherRequest = z.infer; -type WeatherNotification = z.infer; -type WeatherResult = z.infer; + type WeatherRequest = z.infer; + type WeatherNotification = z.infer; + type WeatherResult = z.infer; -// Create a typed Client for weather data -const weatherClient = new Client< - WeatherRequest, - WeatherNotification, - WeatherResult ->({ - name: "WeatherClient", - version: "1.0.0", -}); + // Create a typed Client for weather data + const weatherClient = new Client< + WeatherRequest, + WeatherNotification, + WeatherResult + >({ + name: "WeatherClient", + version: "1.0.0", + }); -// Typecheck that only valid weather requests/notifications/results are allowed -weatherClient.request( - { - method: "weather/get", - params: { - city: "Seattle", + // Typecheck that only valid weather requests/notifications/results are allowed + weatherClient.request( + { + method: "weather/get", + params: { + city: "Seattle", + }, }, - }, - WeatherResultSchema, -); + WeatherResultSchema, + ); -weatherClient.notification({ - method: "weather/alert", - params: { - severity: "warning", - message: "Storm approaching", - }, + weatherClient.notification({ + method: "weather/alert", + params: { + severity: "warning", + message: "Storm approaching", + }, + }); }); diff --git a/src/server/index.test.ts b/src/server/index.test.ts index 4a9ac59..e75edbe 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -6,63 +6,65 @@ import { RequestSchema, NotificationSchema, ResultSchema } from "../types.js"; /* Test that custom request/notification/result schemas can be used with the Server class. */ -const GetWeatherRequestSchema = RequestSchema.extend({ - method: z.literal("weather/get"), - params: z.object({ - city: z.string(), - }), -}); +test("should typecheck", () => { + const GetWeatherRequestSchema = RequestSchema.extend({ + method: z.literal("weather/get"), + params: z.object({ + city: z.string(), + }), + }); -const GetForecastRequestSchema = RequestSchema.extend({ - method: z.literal("weather/forecast"), - params: z.object({ - city: z.string(), - days: z.number(), - }), -}); + const GetForecastRequestSchema = RequestSchema.extend({ + method: z.literal("weather/forecast"), + params: z.object({ + city: z.string(), + days: z.number(), + }), + }); -const WeatherForecastNotificationSchema = NotificationSchema.extend({ - method: z.literal("weather/alert"), - params: z.object({ - severity: z.enum(["warning", "watch"]), - message: z.string(), - }), -}); + const WeatherForecastNotificationSchema = NotificationSchema.extend({ + method: z.literal("weather/alert"), + params: z.object({ + severity: z.enum(["warning", "watch"]), + message: z.string(), + }), + }); -const WeatherRequestSchema = GetWeatherRequestSchema.or( - GetForecastRequestSchema, -); -const WeatherNotificationSchema = WeatherForecastNotificationSchema; -const WeatherResultSchema = ResultSchema.extend({ - temperature: z.number(), - conditions: z.string(), -}); + const WeatherRequestSchema = GetWeatherRequestSchema.or( + GetForecastRequestSchema, + ); + const WeatherNotificationSchema = WeatherForecastNotificationSchema; + const WeatherResultSchema = ResultSchema.extend({ + temperature: z.number(), + conditions: z.string(), + }); -type WeatherRequest = z.infer; -type WeatherNotification = z.infer; -type WeatherResult = z.infer; + type WeatherRequest = z.infer; + type WeatherNotification = z.infer; + type WeatherResult = z.infer; -// Create a typed Server for weather data -const weatherServer = new Server< - WeatherRequest, - WeatherNotification, - WeatherResult ->({ - name: "WeatherServer", - version: "1.0.0", -}); + // Create a typed Server for weather data + const weatherServer = new Server< + WeatherRequest, + WeatherNotification, + WeatherResult + >({ + name: "WeatherServer", + version: "1.0.0", + }); -// Typecheck that only valid weather requests/notifications/results are allowed -weatherServer.setRequestHandler(GetWeatherRequestSchema, (request) => { - return { - temperature: 72, - conditions: "sunny", - }; -}); + // Typecheck that only valid weather requests/notifications/results are allowed + weatherServer.setRequestHandler(GetWeatherRequestSchema, (request) => { + return { + temperature: 72, + conditions: "sunny", + }; + }); -weatherServer.setNotificationHandler( - WeatherForecastNotificationSchema, - (notification) => { - console.log(`Weather alert: ${notification.params.message}`); - }, -); + weatherServer.setNotificationHandler( + WeatherForecastNotificationSchema, + (notification) => { + console.log(`Weather alert: ${notification.params.message}`); + }, + ); +}); From af9df1ff1494dcf3cae3d505682ff9dc53c99d27 Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Mon, 4 Nov 2024 21:13:00 +0000 Subject: [PATCH 5/6] Don't actually run request() and notification() --- src/client/index.test.ts | 32 +++++++++++++++++--------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/src/client/index.test.ts b/src/client/index.test.ts index 1dc5ead..58d7a22 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -54,21 +54,23 @@ test("should typecheck", () => { }); // Typecheck that only valid weather requests/notifications/results are allowed - weatherClient.request( - { - method: "weather/get", - params: { - city: "Seattle", + false && + weatherClient.request( + { + method: "weather/get", + params: { + city: "Seattle", + }, }, - }, - WeatherResultSchema, - ); + WeatherResultSchema, + ); - weatherClient.notification({ - method: "weather/alert", - params: { - severity: "warning", - message: "Storm approaching", - }, - }); + false && + weatherClient.notification({ + method: "weather/alert", + params: { + severity: "warning", + message: "Storm approaching", + }, + }); }); From f5fa3cf8e390c3f7619e2c2aeb952883c6801fec Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Mon, 4 Nov 2024 21:22:05 +0000 Subject: [PATCH 6/6] Oops, ignore more lints --- src/client/index.test.ts | 2 ++ src/server/index.test.ts | 2 ++ 2 files changed, 4 insertions(+) diff --git a/src/client/index.test.ts b/src/client/index.test.ts index 58d7a22..d93ca39 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -1,4 +1,6 @@ /* eslint-disable @typescript-eslint/no-unused-vars */ +/* eslint-disable no-constant-binary-expression */ +/* eslint-disable @typescript-eslint/no-unused-expressions */ import { Client } from "./index.js"; import { z } from "zod"; import { RequestSchema, NotificationSchema, ResultSchema } from "../types.js"; diff --git a/src/server/index.test.ts b/src/server/index.test.ts index e75edbe..be33c58 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -1,4 +1,6 @@ /* eslint-disable @typescript-eslint/no-unused-vars */ +/* eslint-disable no-constant-binary-expression */ +/* eslint-disable @typescript-eslint/no-unused-expressions */ import { Server } from "./index.js"; import { z } from "zod"; import { RequestSchema, NotificationSchema, ResultSchema } from "../types.js";