fix: schema-arg setRequestHandler bypassed per-method wrapping in Client/Server

Normalize schema-arg to method string + parse-wrapped handler, then fall through to
the per-method dispatch (tools/call task validation, elicitation/create capability checks).
Previously the schema form short-circuited via _registerCompatRequestHandler, so e.g.
setRequestHandler(CallToolRequestSchema, h) and setRequestHandler('tools/call', h) had
different runtime behavior.

Also reword @deprecated message: schema form is not deprecated for non-spec methods (the
method-string overload is constrained to RequestMethod), so the advice now reads
'For spec methods, pass the method string instead.'

Author: felixweinberger

packages/client/src/client/client.ts

@@ -51,6 +51,7 @@ import {
     ElicitRequestSchema,
     ElicitResultSchema,
     EmptyResultSchema,
+    extractMethodLiteral,
     extractTaskManagerOptions,
     GetPromptResultSchema,
     InitializeResultSchema,
@@ -342,19 +343,23 @@ export class Client extends Protocol<ClientContext> {
         method: M,
         handler: (request: RequestTypeMap[M], ctx: ClientContext) => ResultTypeMap[M] | Promise<ResultTypeMap[M]>
     ): void;
-    /** @deprecated Pass the method string instead. */
+    /** @deprecated For spec methods, pass the method string instead. */
     public override setRequestHandler<T extends ZodLikeRequestSchema>(
         requestSchema: T,
         handler: (request: ReturnType<T['parse']>, ctx: ClientContext) => Result | Promise<Result>
     ): void;
-    public override setRequestHandler(method: string | ZodLikeRequestSchema, schemaHandler: unknown): void {
-        if (isZodLikeSchema(method)) {
-            return this._registerCompatRequestHandler(
-                method,
-                schemaHandler as (request: unknown, ctx: ClientContext) => Result | Promise<Result>
-            );
+    public override setRequestHandler(methodOrSchema: string | ZodLikeRequestSchema, schemaHandler: unknown): void {
+        let method: string;
+        let handler: (request: Request, ctx: ClientContext) => ClientResult | Promise<ClientResult>;
+        if (isZodLikeSchema(methodOrSchema)) {
+            const schema = methodOrSchema;
+            const userHandler = schemaHandler as (request: unknown, ctx: ClientContext) => Result | Promise<Result>;
+            method = extractMethodLiteral(schema);
+            handler = (req, ctx) => userHandler(schema.parse(req), ctx);
+        } else {
+            method = methodOrSchema;
+            handler = schemaHandler as (request: Request, ctx: ClientContext) => ClientResult | Promise<ClientResult>;
         }
-        const handler = schemaHandler as (request: Request, ctx: ClientContext) => ClientResult | Promise<ClientResult>;
         if (method === 'elicitation/create') {
             const wrappedHandler = async (request: Request, ctx: ClientContext): Promise<ClientResult> => {
                 const validatedRequest = parseSchema(ElicitRequestSchema, request);

packages/core/src/shared/protocol.ts

@@ -1030,7 +1030,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
         method: M,
         handler: (request: RequestTypeMap[M], ctx: ContextT) => Result | Promise<Result>
     ): void;
-    /** @deprecated Pass the method string instead. */
+    /** @deprecated For spec methods, pass the method string instead. */
     setRequestHandler<T extends ZodLikeRequestSchema>(
         requestSchema: T,
         handler: (request: ReturnType<T['parse']>, ctx: ContextT) => Result | Promise<Result>
@@ -1096,7 +1096,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
         method: M,
         handler: (notification: NotificationTypeMap[M]) => void | Promise<void>
     ): void;
-    /** @deprecated Pass the method string instead. */
+    /** @deprecated For spec methods, pass the method string instead. */
     setNotificationHandler<T extends ZodLikeRequestSchema>(
         notificationSchema: T,
         handler: (notification: ReturnType<T['parse']>) => void | Promise<void>

packages/server/src/server/server.ts

@@ -46,6 +46,7 @@ import {
     CreateTaskResultSchema,
     ElicitResultSchema,
     EmptyResultSchema,
+    extractMethodLiteral,
     extractTaskManagerOptions,
     isZodLikeSchema,
     LATEST_PROTOCOL_VERSION,
@@ -230,19 +231,23 @@ export class Server extends Protocol<ServerContext> {
         method: M,
         handler: (request: RequestTypeMap[M], ctx: ServerContext) => ResultTypeMap[M] | Promise<ResultTypeMap[M]>
     ): void;
-    /** @deprecated Pass the method string instead. */
+    /** @deprecated For spec methods, pass the method string instead. */
     public override setRequestHandler<T extends ZodLikeRequestSchema>(
         requestSchema: T,
         handler: (request: ReturnType<T['parse']>, ctx: ServerContext) => Result | Promise<Result>
     ): void;
-    public override setRequestHandler(method: string | ZodLikeRequestSchema, schemaHandler: unknown): void {
-        if (isZodLikeSchema(method)) {
-            return this._registerCompatRequestHandler(
-                method,
-                schemaHandler as (request: unknown, ctx: ServerContext) => Result | Promise<Result>
-            );
+    public override setRequestHandler(methodOrSchema: string | ZodLikeRequestSchema, schemaHandler: unknown): void {
+        let method: string;
+        let handler: (request: Request, ctx: ServerContext) => ServerResult | Promise<ServerResult>;
+        if (isZodLikeSchema(methodOrSchema)) {
+            const schema = methodOrSchema;
+            const userHandler = schemaHandler as (request: unknown, ctx: ServerContext) => Result | Promise<Result>;
+            method = extractMethodLiteral(schema);
+            handler = (req, ctx) => userHandler(schema.parse(req), ctx);
+        } else {
+            method = methodOrSchema;
+            handler = schemaHandler as (request: Request, ctx: ServerContext) => ServerResult | Promise<ServerResult>;
         }
-        const handler = schemaHandler as (request: Request, ctx: ServerContext) => ServerResult | Promise<ServerResult>;
         if (method === 'tools/call') {
             const wrappedHandler = async (request: Request, ctx: ServerContext): Promise<ServerResult> => {
                 const validatedRequest = parseSchema(CallToolRequestSchema, request);

packages/server/test/server/setRequestHandlerSchemaParity.test.ts

@@ -0,0 +1,49 @@
+import { describe, expect, it } from 'vitest';
+
+import { CallToolRequestSchema, InMemoryTransport } from '@modelcontextprotocol/core';
+
+import { Server } from '../../src/server/server.js';
+
+/**
+ * Regression test: setRequestHandler(CallToolRequestSchema, h) and
+ * setRequestHandler('tools/call', h) must apply the same per-method
+ * wrapping (task-result validation when params.task is set).
+ */
+describe('Server.setRequestHandler — Zod-schema form parity', () => {
+    async function setup(register: (s: Server) => void) {
+        const server = new Server({ name: 't', version: '1.0' }, { capabilities: { tools: {} } });
+        register(server);
+        const [ct, st] = InMemoryTransport.createLinkedPair();
+        await server.connect(st);
+        await ct.start();
+        return { ct };
+    }
+
+    async function callToolWithTask(ct: InMemoryTransport): Promise<{ result?: unknown; error?: unknown }> {
+        return await new Promise(resolve => {
+            ct.onmessage = m => {
+                const msg = m as { result?: unknown; error?: unknown };
+                if ('result' in msg || 'error' in msg) resolve(msg);
+            };
+            ct.send({
+                jsonrpc: '2.0',
+                id: 1,
+                method: 'tools/call',
+                params: { name: 'x', arguments: {}, task: { ttl: 1000 } }
+            });
+        });
+    }
+
+    it('schema form gets the same task-result validation as string form', async () => {
+        const invalidTaskResult = { content: [{ type: 'text' as const, text: 'not a task result' }] };
+
+        const viaString = await setup(s => s.setRequestHandler('tools/call', () => invalidTaskResult));
+        const viaSchema = await setup(s => s.setRequestHandler(CallToolRequestSchema, () => invalidTaskResult));
+
+        const stringRes = await callToolWithTask(viaString.ct);
+        const schemaRes = await callToolWithTask(viaSchema.ct);
+
+        expect(stringRes.error).toBeDefined();
+        expect(schemaRes.error).toEqual(stringRes.error);
+    });
+});