diff --git a/src/client/index.test.ts b/src/client/index.test.ts new file mode 100644 index 0000000..d93ca39 --- /dev/null +++ b/src/client/index.test.ts @@ -0,0 +1,78 @@ +/* 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"; + +/* +Test that custom request/notification/result schemas can be used with the Client class. +*/ +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 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 + false && + weatherClient.request( + { + method: "weather/get", + params: { + city: "Seattle", + }, + }, + WeatherResultSchema, + ); + + false && + 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; diff --git a/src/server/index.test.ts b/src/server/index.test.ts new file mode 100644 index 0000000..be33c58 --- /dev/null +++ b/src/server/index.test.ts @@ -0,0 +1,72 @@ +/* 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"; + +/* +Test that custom request/notification/result schemas can be used with the Server class. +*/ +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 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;