|
| 1 | +// v1 compat: subclass wrappers that restore the v1 schema-first-arg |
| 2 | +// `setRequestHandler(ZodRequestSchema, handler)` form. v2's `Protocol.setRequestHandler` |
| 3 | +// is method-string-keyed; these shims extract the method literal from the schema and forward. |
| 4 | +// |
| 5 | +// Only the meta-package provides this overload, by design — the underlying |
| 6 | +// `@modelcontextprotocol/{server,client}` packages stay on the v2 API. |
| 7 | + |
| 8 | +import { |
| 9 | + Client as BaseClient, |
| 10 | + type ClientContext, |
| 11 | + type NotificationMethod, |
| 12 | + type RequestMethod, |
| 13 | + type Result, |
| 14 | + type StandardSchemaV1 |
| 15 | +} from '@modelcontextprotocol/client'; |
| 16 | +import { |
| 17 | + McpServer as BaseMcpServer, |
| 18 | + type RequestHandlerSchemas, |
| 19 | + type ResultTypeMap, |
| 20 | + Server as BaseServer, |
| 21 | + type ServerContext, |
| 22 | + type ServerOptions |
| 23 | +} from '@modelcontextprotocol/server'; |
| 24 | +import type * as z from 'zod'; |
| 25 | + |
| 26 | +type ZodRequestSchema<M extends string = string> = z.ZodObject<{ method: z.ZodLiteral<M> } & z.ZodRawShape>; |
| 27 | +type ZodNotificationSchema<M extends string = string> = ZodRequestSchema<M>; |
| 28 | + |
| 29 | +function methodFromZodSchema(schema: ZodRequestSchema): string { |
| 30 | + const shape = (schema as z.ZodObject<z.ZodRawShape>).shape; |
| 31 | + const methodField = shape?.method as z.ZodLiteral<string> | undefined; |
| 32 | + const method = methodField && 'value' in methodField ? methodField.value : undefined; |
| 33 | + if (typeof method !== 'string') { |
| 34 | + throw new TypeError( |
| 35 | + 'setRequestHandler: first argument must be a method string, a {params, result?} schema bundle, or a v1-style Zod request schema with a `.shape.method` literal.' |
| 36 | + ); |
| 37 | + } |
| 38 | + return method; |
| 39 | +} |
| 40 | + |
| 41 | +function isZodRequestSchema(arg: unknown): arg is ZodRequestSchema { |
| 42 | + return typeof arg === 'object' && arg !== null && 'shape' in arg && typeof (arg as { shape: unknown }).shape === 'object'; |
| 43 | +} |
| 44 | + |
| 45 | +type LegacyRequestHandler<S extends ZodRequestSchema, Ctx> = (request: z.infer<S>, ctx: Ctx) => Result | Promise<Result>; |
| 46 | +type LegacyNotificationHandler<S extends ZodNotificationSchema> = (notification: z.infer<S>) => void | Promise<void>; |
| 47 | + |
| 48 | +/** |
| 49 | + * Adds the v1 schema-first-arg overloads to `setRequestHandler` / `setNotificationHandler`. |
| 50 | + * Returned class is a drop-in replacement for the base; instances satisfy `instanceof Base`. |
| 51 | + */ |
| 52 | +function withV1SchemaOverloads< |
| 53 | + Ctx, |
| 54 | + // eslint-disable-next-line @typescript-eslint/no-explicit-any |
| 55 | + TBase extends abstract new (...args: any[]) => { |
| 56 | + setRequestHandler(method: string, ...rest: unknown[]): void; |
| 57 | + setNotificationHandler(method: string, ...rest: unknown[]): void; |
| 58 | + } |
| 59 | +>(Base: TBase) { |
| 60 | + abstract class WithV1 extends Base { |
| 61 | + /** v1 compat: accepts a Zod request schema with a `method` literal as the first arg. */ |
| 62 | + override setRequestHandler<S extends ZodRequestSchema>(schema: S, handler: LegacyRequestHandler<S, Ctx>): void; |
| 63 | + override setRequestHandler<M extends RequestMethod>( |
| 64 | + method: M, |
| 65 | + handler: (request: { method: M; params?: Record<string, unknown> }, ctx: Ctx) => ResultTypeMap[M] | Promise<ResultTypeMap[M]> |
| 66 | + ): void; |
| 67 | + override setRequestHandler<P extends StandardSchemaV1, R extends StandardSchemaV1 | undefined = undefined>( |
| 68 | + method: string, |
| 69 | + schemas: RequestHandlerSchemas<P, R>, |
| 70 | + handler: ( |
| 71 | + params: StandardSchemaV1.InferOutput<P>, |
| 72 | + ctx: Ctx |
| 73 | + ) => R extends StandardSchemaV1 ? StandardSchemaV1.InferOutput<R> | Promise<StandardSchemaV1.InferOutput<R>> : Result | Promise<Result> |
| 74 | + ): void; |
| 75 | + override setRequestHandler(arg1: string | ZodRequestSchema, arg2: unknown, arg3?: unknown): void { |
| 76 | + if (typeof arg1 === 'string') { |
| 77 | + return arg3 === undefined ? super.setRequestHandler(arg1, arg2) : super.setRequestHandler(arg1, arg2, arg3); |
| 78 | + } |
| 79 | + const method = methodFromZodSchema(arg1); |
| 80 | + return super.setRequestHandler(method, arg2); |
| 81 | + } |
| 82 | + |
| 83 | + /** v1 compat: accepts a Zod notification schema with a `method` literal as the first arg. */ |
| 84 | + override setNotificationHandler<S extends ZodNotificationSchema>(schema: S, handler: LegacyNotificationHandler<S>): void; |
| 85 | + override setNotificationHandler<M extends NotificationMethod>(method: M, handler: (notification: unknown) => void | Promise<void>): void; |
| 86 | + override setNotificationHandler<P extends StandardSchemaV1>( |
| 87 | + method: string, |
| 88 | + schemas: { params: P }, |
| 89 | + handler: (params: StandardSchemaV1.InferOutput<P>, notification: unknown) => void | Promise<void> |
| 90 | + ): void; |
| 91 | + override setNotificationHandler(arg1: string | ZodNotificationSchema, arg2: unknown, arg3?: unknown): void { |
| 92 | + if (typeof arg1 === 'string') { |
| 93 | + return arg3 === undefined ? super.setNotificationHandler(arg1, arg2) : super.setNotificationHandler(arg1, arg2, arg3); |
| 94 | + } |
| 95 | + const method = methodFromZodSchema(arg1); |
| 96 | + return super.setNotificationHandler(method, arg2); |
| 97 | + } |
| 98 | + } |
| 99 | + return WithV1; |
| 100 | +} |
| 101 | + |
| 102 | +/** v1-compat `Server`: adds the schema-first-arg `setRequestHandler` overload. */ |
| 103 | +export class Server extends withV1SchemaOverloads<ServerContext, typeof BaseServer>(BaseServer) {} |
| 104 | + |
| 105 | +/** v1-compat `Client`: adds the schema-first-arg `setRequestHandler` overload. */ |
| 106 | +export class Client extends withV1SchemaOverloads<ClientContext, typeof BaseClient>(BaseClient) {} |
| 107 | + |
| 108 | +/** |
| 109 | + * v1-compat `McpServer`: `.server` is the compat-wrapped {@link Server} so |
| 110 | + * `mcp.server.setRequestHandler(ZodSchema, h)` works. |
| 111 | + */ |
| 112 | +export class McpServer extends BaseMcpServer { |
| 113 | + declare readonly server: Server; |
| 114 | + constructor(serverInfo: ConstructorParameters<typeof BaseMcpServer>[0], options?: ServerOptions) { |
| 115 | + super(serverInfo, options); |
| 116 | + // Base constructor's only side effect is `this.server = new Server(...)`; replace |
| 117 | + // with the compat-wrapped subclass so `mcp.server.setRequestHandler(ZodSchema, h)` works. |
| 118 | + (this as { server: Server }).server = new Server(serverInfo, options); |
| 119 | + } |
| 120 | +} |
| 121 | + |
| 122 | +export { isZodRequestSchema as _isZodRequestSchema }; |
0 commit comments