diff --git a/src/client/index.test.ts b/src/client/index.test.ts new file mode 100644 index 00000000..edb0057a --- /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 9b9aa81e..e4ae4da9 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;