Fix typing for custom requests, notification, and result types

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<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
+  false &&
+    weatherClient.request(
+      {
+        method: "weather/get",
+        params: {
+          city: "Seattle",
+        },
+      },
+      WeatherResultSchema,
+    );
+
+  false &&
+    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;

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

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<typeof CustomRequestSchema>
+ * type CustomNotification = z.infer<typeof CustomNotificationSchema>
+ * type CustomResult = z.infer<typeof CustomResultSchema>
+ *
+ * // Create typed server
+ * const server = new Server<CustomRequest, CustomNotification, CustomResult>({
+ *   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;