From b5880ec18af187e28e74e1e959e55cd631af8292 Mon Sep 17 00:00:00 2001
From: Justin Spahr-Summers <justin@anthropic.com>
Date: Mon, 4 Nov 2024 21:07:53 +0000
Subject: [PATCH] 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 00000000..7a06c533
--- /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<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}`);
+  },
+);
diff --git a/src/server/index.ts b/src/server/index.ts
index c8ffaba7..71c22ccd 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<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;