Support custom client request, notification, and result types

Author: jspahrsummers

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<typeof WeatherRequestSchema>;
+type WeatherNotification = z.infer<typeof WeatherNotificationSchema>;
+type WeatherResult = z.infer<typeof WeatherResultSchema>;
+
+// 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",
+  },
+});

src/client/index.ts

@@ -6,19 +6,46 @@ import {
   ClientResult,
   Implementation,
   InitializeResultSchema,
+  Notification,
   PROTOCOL_VERSION,
+  Request,
+  Result,
   ServerCapabilities,
 } from "../types.js";
 
 /**
  * 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<typeof CustomRequestSchema>
+ * type CustomNotification = z.infer<typeof CustomNotificationSchema>
+ * type CustomResult = z.infer<typeof CustomResultSchema>
+ *
+ * // Create typed client
+ * const client = new Client<CustomRequest, CustomNotification, CustomResult>({
+ *   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;