modelcontextprotocol/typescript-sdk
Watch
1
0
Fork
You've already forked typescript-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

fix: replace schema-based handlers with string method names (#1446)

Browse source
This commit is contained in:
Matt 2026-02-02 23:57:07 +00:00 • committed by GitHub
commit f6e8204b86
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 490 additions and 878 deletions
Download patch file Download diff file Expand all files Collapse all files

64
packages/client/src/client/client.ts
View file

@ -1,5 +1,4 @@
import type {
AnyObjectSchema,
CallToolRequest,
ClientCapabilities,
ClientNotification,
@ -20,20 +19,20 @@ import type {
ListToolsRequest,
LoggingLevel,
Notification,
NotificationMethod,
ProtocolOptions,
ReadResourceRequest,
Request,
RequestHandlerExtra,
RequestMethod,
RequestOptions,
RequestTypeMap,
Result,
SchemaOutput,
ServerCapabilities,
SubscribeRequest,
Tool,
Transport,
UnsubscribeRequest,
ZodV3Internal,
ZodV4Internal
UnsubscribeRequest
} from '@modelcontextprotocol/core';
import {
AjvJsonSchemaValidator,
@ -49,10 +48,8 @@ import {
ElicitResultSchema,
EmptyResultSchema,
ErrorCode,
getObjectShape,
GetPromptResultSchema,
InitializeResultSchema,
isZ4Schema,
LATEST_PROTOCOL_VERSION,
ListChangedOptionsBaseSchema,
ListPromptsResultSchema,
@ -61,13 +58,10 @@ import {
ListToolsResultSchema,
McpError,
mergeCapabilities,
PromptListChangedNotificationSchema,
Protocol,
ReadResourceResultSchema,
ResourceListChangedNotificationSchema,
safeParse,
SUPPORTED_PROTOCOL_VERSIONS,
ToolListChangedNotificationSchema
SUPPORTED_PROTOCOL_VERSIONS
} from '@modelcontextprotocol/core';
import { ExperimentalClientTasks } from '../experimental/tasks/client.js';
@ -281,21 +275,21 @@ export class Client<
*/
private _setupListChangedHandlers(config: ListChangedHandlers): void {
if (config.tools && this._serverCapabilities?.tools?.listChanged) {
this._setupListChangedHandler('tools', ToolListChangedNotificationSchema, config.tools, async () => {
this._setupListChangedHandler('tools', 'notifications/tools/list_changed', config.tools, async () => {
const result = await this.listTools();
return result.tools;
});
}
if (config.prompts && this._serverCapabilities?.prompts?.listChanged) {
this._setupListChangedHandler('prompts', PromptListChangedNotificationSchema, config.prompts, async () => {
this._setupListChangedHandler('prompts', 'notifications/prompts/list_changed', config.prompts, async () => {
const result = await this.listPrompts();
return result.prompts;
});
}
if (config.resources && this._serverCapabilities?.resources?.listChanged) {
this._setupListChangedHandler('resources', ResourceListChangedNotificationSchema, config.resources, async () => {
this._setupListChangedHandler('resources', 'notifications/resources/list_changed', config.resources, async () => {
const result = await this.listResources();
return result.resources;
});
@ -334,38 +328,16 @@ export class Client<
/**
* Override request handler registration to enforce client-side validation for elicitation.
*/
public override setRequestHandler<T extends AnyObjectSchema>(
requestSchema: T,
public override setRequestHandler<M extends RequestMethod>(
method: M,
handler: (
request: SchemaOutput<T>,
request: RequestTypeMap[M],
extra: RequestHandlerExtra<ClientRequest | RequestT, ClientNotification | NotificationT>
) => ClientResult | ResultT | Promise<ClientResult | ResultT>
): void {
const shape = getObjectShape(requestSchema);
const methodSchema = shape?.method;
if (!methodSchema) {
throw new Error('Schema is missing a method literal');
}
// Extract literal value using type-safe property access
let methodValue: unknown;
if (isZ4Schema(methodSchema)) {
const v4Schema = methodSchema as unknown as ZodV4Internal;
const v4Def = v4Schema._zod?.def;
methodValue = v4Def?.value ?? v4Schema.value;
} else {
const v3Schema = methodSchema as unknown as ZodV3Internal;
const legacyDef = v3Schema._def;
methodValue = legacyDef?.value ?? v3Schema.value;
}
if (typeof methodValue !== 'string') {
throw new TypeError('Schema method literal must be a string');
}
const method = methodValue;
if (method === 'elicitation/create') {
const wrappedHandler = async (
request: SchemaOutput<T>,
request: RequestTypeMap[M],
extra: RequestHandlerExtra<ClientRequest | RequestT, ClientNotification | NotificationT>
): Promise<ClientResult | ResultT> => {
const validatedRequest = safeParse(ElicitRequestSchema, request);
@ -433,12 +405,12 @@ export class Client<
};
// Install the wrapped handler
return super.setRequestHandler(requestSchema, wrappedHandler as unknown as typeof handler);
return super.setRequestHandler(method, wrappedHandler as unknown as typeof handler);
}
if (method === 'sampling/createMessage') {
const wrappedHandler = async (
request: SchemaOutput<T>,
request: RequestTypeMap[M],
extra: RequestHandlerExtra<ClientRequest | RequestT, ClientNotification | NotificationT>
): Promise<ClientResult | ResultT> => {
const validatedRequest = safeParse(CreateMessageRequestSchema, request);
@ -479,11 +451,11 @@ export class Client<
};
// Install the wrapped handler
return super.setRequestHandler(requestSchema, wrappedHandler as unknown as typeof handler);
return super.setRequestHandler(method, wrappedHandler as unknown as typeof handler);
}
// Other handlers use default behavior
return super.setRequestHandler(requestSchema, handler);
return super.setRequestHandler(method, handler);
}
protected assertCapability(capability: keyof ServerCapabilities, method: string): void {
@ -898,7 +870,7 @@ export class Client<
*/
private _setupListChangedHandler<T>(
listType: string,
notificationSchema: { shape: { method: { value: string } } },
notificationMethod: NotificationMethod,
options: ListChangedOptions<T>,
fetcher: () => Promise<T[]>
): void {
@ -949,7 +921,7 @@ export class Client<
};
// Register notification handler
this.setNotificationHandler(notificationSchema as AnyObjectSchema, handler);
this.setNotificationHandler(notificationMethod, handler);
}
async sendRootsListChanged() {

2
packages/client/test/client/sse.test.ts
View file

@ -111,7 +111,7 @@ describe('SSEClientTransport', () => {
// Create a server that returns 403
await resourceServer.close();
resourceServer = createServer((req, res) => {
resourceServer = createServer((_req, res) => {
res.writeHead(403);
res.end();
});

68
packages/core/src/shared/protocol.ts
View file

@ -14,6 +14,8 @@ import type {
JSONRPCResultResponse,
MessageExtraInfo,
Notification,
NotificationMethod,
NotificationTypeMap,
Progress,
ProgressNotification,
RelatedTaskMetadata,
@ -21,6 +23,8 @@ import type {
RequestId,
RequestInfo,
RequestMeta,
RequestMethod,
RequestTypeMap,
Result,
ServerCapabilities,
Task,
@ -28,30 +32,25 @@ import type {
TaskStatusNotification
} from '../types/types.js';
import {
CancelledNotificationSchema,
CancelTaskRequestSchema,
CancelTaskResultSchema,
CreateTaskResultSchema,
ErrorCode,
GetTaskPayloadRequestSchema,
GetTaskRequestSchema,
getNotificationSchema,
getRequestSchema,
GetTaskResultSchema,
isJSONRPCErrorResponse,
isJSONRPCNotification,
isJSONRPCRequest,
isJSONRPCResultResponse,
isTaskAugmentedRequestParams,
ListTasksRequestSchema,
ListTasksResultSchema,
McpError,
PingRequestSchema,
ProgressNotificationSchema,
RELATED_TASK_META_KEY,
TaskStatusNotificationSchema
} from '../types/types.js';
import type { AnyObjectSchema, AnySchema, SchemaOutput } from '../util/zodCompat.js';
import type { AnySchema, SchemaOutput } from '../util/zodCompat.js';
import { safeParse } from '../util/zodCompat.js';
import { getMethodLiteral, parseWithCompat } from '../util/zodJsonSchemaCompat.js';
import { parseWithCompat } from '../util/zodJsonSchemaCompat.js';
import type { ResponseMessage } from './responseMessage.js';
import type { Transport, TransportSendOptions } from './transport.js';
@ -368,16 +367,16 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
fallbackNotificationHandler?: (notification: Notification) => Promise<void>;
constructor(private _options?: ProtocolOptions) {
this.setNotificationHandler(CancelledNotificationSchema, notification => {
this.setNotificationHandler('notifications/cancelled', notification => {
this._oncancel(notification);
});
this.setNotificationHandler(ProgressNotificationSchema, notification => {
this._onprogress(notification as unknown as ProgressNotification);
this.setNotificationHandler('notifications/progress', notification => {
this._onprogress(notification);
});
this.setRequestHandler(
PingRequestSchema,
'ping',
// Automatic pong by default.
_request => ({}) as SendResultT
);
@ -386,7 +385,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
this._taskStore = _options?.taskStore;
this._taskMessageQueue = _options?.taskMessageQueue;
if (this._taskStore) {
this.setRequestHandler(GetTaskRequestSchema, async (request, extra) => {
this.setRequestHandler('tasks/get', async (request, extra) => {
const task = await this._taskStore!.getTask(request.params.taskId, extra.sessionId);
if (!task) {
throw new McpError(ErrorCode.InvalidParams, 'Failed to retrieve task: Task not found');
@ -394,13 +393,12 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
// Per spec: tasks/get responses SHALL NOT include related-task metadata
// as the taskId parameter is the source of truth
// @ts-expect-error SendResultT cannot contain GetTaskResult, but we include it in our derived types everywhere else
return {
...task
} as SendResultT;
} as unknown as SendResultT;
});
this.setRequestHandler(GetTaskPayloadRequestSchema, async (request, extra) => {
this.setRequestHandler('tasks/result', async (request, extra) => {
const handleTaskResult = async (): Promise<SendResultT> => {
const taskId = request.params.taskId;
@ -478,7 +476,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
taskId: taskId
}
}
} as SendResultT;
} as unknown as SendResultT;
}
return await handleTaskResult();
@ -487,15 +485,14 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
return await handleTaskResult();
});
this.setRequestHandler(ListTasksRequestSchema, async (request, extra) => {
this.setRequestHandler('tasks/list', async (request, extra) => {
try {
const { tasks, nextCursor } = await this._taskStore!.listTasks(request.params?.cursor, extra.sessionId);
// @ts-expect-error SendResultT cannot contain ListTasksResult, but we include it in our derived types everywhere else
return {
tasks,
nextCursor,
_meta: {}
} as SendResultT;
} as unknown as SendResultT;
} catch (error) {
throw new McpError(
ErrorCode.InvalidParams,
@ -504,7 +501,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
}
});
this.setRequestHandler(CancelTaskRequestSchema, async (request, extra) => {
this.setRequestHandler('tasks/cancel', async (request, extra) => {
try {
// Get the current task to check if it's in a terminal state, in case the implementation is not atomic
const task = await this._taskStore!.getTask(request.params.taskId, extra.sessionId);
@ -1401,18 +1398,18 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
*
* Note that this will replace any previous request handler for the same method.
*/
setRequestHandler<T extends AnyObjectSchema>(
requestSchema: T,
setRequestHandler<M extends RequestMethod>(
method: M,
handler: (
request: SchemaOutput<T>,
request: RequestTypeMap[M],
extra: RequestHandlerExtra<SendRequestT, SendNotificationT>
) => SendResultT | Promise<SendResultT>
): void {
const method = getMethodLiteral(requestSchema);
this.assertRequestHandlerCapability(method);
const schema = getRequestSchema(method);
this._requestHandlers.set(method, (request, extra) => {
const parsed = parseWithCompat(requestSchema, request) as SchemaOutput<T>;
const parsed = parseWithCompat(schema, request) as RequestTypeMap[M];
return Promise.resolve(handler(parsed, extra));
});
}
@ -1420,14 +1417,14 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
/**
* Removes the request handler for the given method.
*/
removeRequestHandler(method: string): void {
removeRequestHandler(method: RequestMethod): void {
this._requestHandlers.delete(method);
}
/**
* Asserts that a request handler has not already been set for the given method, in preparation for a new one being automatically installed.
*/
assertCanSetRequestHandler(method: string): void {
assertCanSetRequestHandler(method: RequestMethod): void {
if (this._requestHandlers.has(method)) {
throw new Error(`A request handler for ${method} already exists, which would be overridden`);
}
@ -1438,13 +1435,14 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
*
* Note that this will replace any previous notification handler for the same method.
*/
setNotificationHandler<T extends AnyObjectSchema>(
notificationSchema: T,
handler: (notification: SchemaOutput<T>) => void | Promise<void>
setNotificationHandler<M extends NotificationMethod>(
method: M,
handler: (notification: NotificationTypeMap[M]) => void | Promise<void>
): void {
const method = getMethodLiteral(notificationSchema);
const schema = getNotificationSchema(method);
this._notificationHandlers.set(method, notification => {
const parsed = parseWithCompat(notificationSchema, notification) as SchemaOutput<T>;
const parsed = parseWithCompat(schema, notification) as NotificationTypeMap[M];
return Promise.resolve(handler(parsed));
});
}
@ -1452,7 +1450,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
/**
* Removes the notification handler for the given method.
*/
removeNotificationHandler(method: string): void {
removeNotificationHandler(method: NotificationMethod): void {
this._notificationHandlers.delete(method);
}

38
packages/core/src/types/types.ts
View file

@ -2605,3 +2605,41 @@ export type ClientResult = Infer<typeof ClientResultSchema>;
export type ServerRequest = Infer<typeof ServerRequestSchema>;
export type ServerNotification = Infer<typeof ServerNotificationSchema>;
export type ServerResult = Infer<typeof ServerResultSchema>;
/* Protocol type maps */
type MethodToTypeMap<U> = {
[T in U as T extends { method: infer M extends string } ? M : never]: T;
};
export type RequestMethod = ClientRequest['method'] | ServerRequest['method'];
export type NotificationMethod = ClientNotification['method'] | ServerNotification['method'];
export type RequestTypeMap = MethodToTypeMap<ClientRequest | ServerRequest>;
export type NotificationTypeMap = MethodToTypeMap<ClientNotification | ServerNotification>;
/* Runtime schema lookup */
type RequestSchemaType = (typeof ClientRequestSchema.options)[number] | (typeof ServerRequestSchema.options)[number];
type NotificationSchemaType = (typeof ClientNotificationSchema.options)[number] | (typeof ServerNotificationSchema.options)[number];
function buildSchemaMap<T extends { shape: { method: { value: string } } }>(schemas: readonly T[]): Record<string, T> {
const map: Record<string, T> = {};
for (const schema of schemas) {
const method = schema.shape.method.value;
map[method] = schema;
}
return map;
}
const requestSchemas = buildSchemaMap([...ClientRequestSchema.options, ...ServerRequestSchema.options] as const) as Record<
RequestMethod,
RequestSchemaType
>;
const notificationSchemas = buildSchemaMap([...ClientNotificationSchema.options, ...ServerNotificationSchema.options] as const) as Record<
NotificationMethod,
NotificationSchemaType
>;
export function getRequestSchema<M extends RequestMethod>(method: M) {
return requestSchemas[method];
}
export function getNotificationSchema<M extends NotificationMethod>(method: M) {
return notificationSchemas[method];
}

23
packages/core/src/util/zodJsonSchemaCompat.ts
View file

@ -9,8 +9,8 @@ import type * as z4c from 'zod/v4/core';
import * as z4mini from 'zod/v4-mini';
import { zodToJsonSchema } from 'zod-to-json-schema';
import type { AnyObjectSchema, AnySchema } from './zodCompat.js';
import { getLiteralValue, getObjectShape, isZ4Schema, safeParse } from './zodCompat.js';
import type { AnyObjectSchema, AnySchema, SchemaOutput } from './zodCompat.js';
import { isZ4Schema, safeParse } from './zodCompat.js';
type JsonSchema = Record<string, unknown>;
@ -44,25 +44,10 @@ export function toJsonSchemaCompat(schema: AnyObjectSchema, opts?: CommonOpts):
}) as JsonSchema;
}
export function getMethodLiteral(schema: AnyObjectSchema): string {
const shape = getObjectShape(schema);
const methodSchema = shape?.method as AnySchema | undefined;
if (!methodSchema) {
throw new Error('Schema is missing a method literal');
}
const value = getLiteralValue(methodSchema);
if (typeof value !== 'string') {
throw new TypeError('Schema method literal must be a string');
}
return value;
}
export function parseWithCompat(schema: AnySchema, data: unknown): unknown {
export function parseWithCompat<T extends AnySchema>(schema: T, data: unknown): SchemaOutput<T> {
const result = safeParse(schema, data);
if (!result.success) {
throw result.error;
}
return result.data;
return result.data as SchemaOutput<T>;
}

83
packages/core/test/shared/protocol.test.ts
View file

@ -29,7 +29,7 @@ import type {
Task,
TaskCreationParams
} from '../../src/types/types.js';
import { CallToolRequestSchema, ErrorCode, McpError, RELATED_TASK_META_KEY } from '../../src/types/types.js';
import { ErrorCode, McpError, RELATED_TASK_META_KEY } from '../../src/types/types.js';
// Type helper for accessing private/protected Protocol properties in tests
interface TestProtocol {
@ -1266,13 +1266,13 @@ describe('Task-based execution', () => {
await protocol.connect(transport);
protocol.setRequestHandler(CallToolRequestSchema, async request => {
protocol.setRequestHandler('tools/call', async request => {
// Tool implementor can access task creation parameters from request.params.task
expect(request.params.task).toEqual({
ttl: 60000,
pollInterval: 1000
});
return { result: 'success' };
return { content: [{ type: 'text', text: 'success' }] };
});
transport.onmessage?.({
@ -2009,7 +2009,7 @@ describe('Task-based execution', () => {
await serverProtocol.connect(serverTransport);
// Set up a handler that uses sendRequest and sendNotification
serverProtocol.setRequestHandler(CallToolRequestSchema, async (_request, extra) => {
serverProtocol.setRequestHandler('tools/call', async (_request, extra) => {
// Send a notification using the extra.sendNotification
await extra.sendNotification({
method: 'notifications/message',
@ -2093,15 +2093,11 @@ describe('Request Cancellation vs Task Cancellation', () => {
// Set up a request handler that checks if it was aborted
let wasAborted = false;
const TestRequestSchema = z.object({
method: z.literal('test/longRunning'),
params: z.optional(z.record(z.string(), z.unknown()))
});
protocol.setRequestHandler(TestRequestSchema, async (_request, extra) => {
protocol.setRequestHandler('ping', async (_request, extra) => {
// Simulate a long-running operation
await new Promise(resolve => setTimeout(resolve, 100));
wasAborted = extra.signal.aborted;
return { _meta: {} } as Result;
return {};
});
// Simulate an incoming request
@ -2110,7 +2106,7 @@ describe('Request Cancellation vs Task Cancellation', () => {
transport.onmessage({
jsonrpc: '2.0',
id: requestId,
method: 'test/longRunning',
method: 'ping',
params: {}
});
}
@ -2308,28 +2304,24 @@ describe('Request Cancellation vs Task Cancellation', () => {
// Set up a request handler
let requestCompleted = false;
const TestMethodSchema = z.object({
method: z.literal('test/method'),
params: z.optional(z.record(z.string(), z.unknown()))
});
protocol.setRequestHandler(TestMethodSchema, async () => {
protocol.setRequestHandler('ping', async () => {
await new Promise(resolve => setTimeout(resolve, 50));
requestCompleted = true;
return { _meta: {} } as Result;
return {};
});
// Create a task
// Create a task (simulating a long-running tools/call)
const task = await taskStore.createTask({ ttl: 60000 }, 'req-1', {
method: 'test/method',
params: {}
method: 'tools/call',
params: { name: 'long-running-tool', arguments: {} }
});
// Start a request
// Start an unrelated ping request
if (transport.onmessage) {
transport.onmessage({
jsonrpc: '2.0',
id: 123,
method: 'test/method',
method: 'ping',
params: {}
});
}
@ -2488,7 +2480,7 @@ describe('Progress notification support for tasks', () => {
await protocol.connect(transport);
// Set up a request handler that will complete the task
protocol.setRequestHandler(CallToolRequestSchema, async (request, extra) => {
protocol.setRequestHandler('tools/call', async (_request, extra) => {
if (extra.taskStore) {
const task = await extra.taskStore.createTask({ ttl: 60000 });
@ -3686,17 +3678,8 @@ describe('Message Interception', () => {
await protocol.connect(transport);
// Set up a request handler that returns a result
const TestRequestSchema = z.object({
method: z.literal('test/taskRequest'),
params: z
.object({
_meta: z.optional(z.record(z.string(), z.unknown()))
})
.passthrough()
});
protocol.setRequestHandler(TestRequestSchema, async () => {
return { content: 'test result' } as Result;
protocol.setRequestHandler('ping', async () => {
return {};
});
// Simulate an incoming request with relatedTask metadata
@ -3705,7 +3688,7 @@ describe('Message Interception', () => {
transport.onmessage?.({
jsonrpc: '2.0',
id: requestId,
method: 'test/taskRequest',
method: 'ping',
params: {
_meta: {
'io.modelcontextprotocol/related-task': { taskId }
@ -3725,7 +3708,7 @@ describe('Message Interception', () => {
expect(queuedMessage!.type).toBe('response');
if (queuedMessage!.type === 'response') {
expect(queuedMessage!.message.id).toBe(requestId);
expect(queuedMessage!.message.result).toEqual({ content: 'test result' });
expect(queuedMessage!.message.result).toEqual({});
}
});
@ -3733,16 +3716,7 @@ describe('Message Interception', () => {
await protocol.connect(transport);
// Set up a request handler that throws an error
const TestRequestSchema = z.object({
method: z.literal('test/taskRequestError'),
params: z
.object({
_meta: z.optional(z.record(z.string(), z.unknown()))
})
.passthrough()
});
protocol.setRequestHandler(TestRequestSchema, async () => {
protocol.setRequestHandler('ping', async () => {
throw new McpError(ErrorCode.InternalError, 'Test error message');
});
@ -3752,7 +3726,7 @@ describe('Message Interception', () => {
transport.onmessage?.({
jsonrpc: '2.0',
id: requestId,
method: 'test/taskRequestError',
method: 'ping',
params: {
_meta: {
'io.modelcontextprotocol/related-task': { taskId }
@ -3815,13 +3789,8 @@ describe('Message Interception', () => {
const sendSpy = vi.spyOn(transport, 'send');
// Set up a request handler
const TestRequestSchema = z.object({
method: z.literal('test/normalRequest'),
params: z.optional(z.record(z.string(), z.unknown()))
});
protocol.setRequestHandler(TestRequestSchema, async () => {
return { content: 'normal result' } as Result;
protocol.setRequestHandler('tools/call', async () => {
return { content: [{ type: 'text', text: 'done' }] };
});
// Simulate an incoming request WITHOUT relatedTask metadata
@ -3829,8 +3798,8 @@ describe('Message Interception', () => {
transport.onmessage?.({
jsonrpc: '2.0',
id: requestId,
method: 'test/normalRequest',
params: {}
method: 'tools/call',
params: { name: 'test-tool' }
});
// Wait for the handler to complete
@ -3841,7 +3810,7 @@ describe('Message Interception', () => {
expect.objectContaining({
jsonrpc: '2.0',
id: requestId,
result: { content: 'normal result' }
result: { content: [{ type: 'text', text: 'done' }] }
})
);
});

109
packages/core/test/shared/protocolTransportHandling.test.ts
View file

@ -1,9 +1,8 @@
import { beforeEach, describe, expect, test } from 'vitest';
import * as z from 'zod/v4';
import { Protocol } from '../../src/shared/protocol.js';
import type { Transport } from '../../src/shared/transport.js';
import type { JSONRPCMessage, Notification, Request, Result } from '../../src/types/types.js';
import type { EmptyResult, JSONRPCMessage, Notification, Request, Result } from '../../src/types/types.js';
// Mock Transport class
class MockTransport implements Transport {
@ -48,54 +47,24 @@ describe('Protocol transport handling bug', () => {
test('should send response to the correct transport when multiple clients are connected', async () => {
// Set up a request handler that simulates processing time
let resolveHandler: (value: Result) => void;
const handlerPromise = new Promise<Result>(resolve => {
let resolveHandler: (value: EmptyResult) => void;
const handlerPromise = new Promise<EmptyResult>(resolve => {
resolveHandler = resolve;
});
const TestRequestSchema = z.object({
method: z.literal('test/method'),
params: z
.object({
from: z.string()
})
.optional()
});
protocol.setRequestHandler(TestRequestSchema, async request => {
console.log(`Processing request from ${request.params?.from}`);
return handlerPromise;
});
protocol.setRequestHandler('ping', async () => handlerPromise);
// Client A connects and sends a request
await protocol.connect(transportA);
const requestFromA = {
jsonrpc: '2.0' as const,
method: 'test/method',
params: { from: 'clientA' },
id: 1
};
// Simulate client A sending a request
transportA.onmessage?.(requestFromA);
transportA.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 1 });
// While A's request is being processed, client B connects
// This overwrites the transport reference in the protocol
await protocol.connect(transportB);
const requestFromB = {
jsonrpc: '2.0' as const,
method: 'test/method',
params: { from: 'clientB' },
id: 2
};
// Client B sends its own request
transportB.onmessage?.(requestFromB);
transportB.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 2 });
// Now complete A's request
resolveHandler!({ data: 'responseForA' } as Result);
resolveHandler!({});
// Wait for async operations to complete
await new Promise(resolve => setTimeout(resolve, 10));
@ -104,70 +73,33 @@ describe('Protocol transport handling bug', () => {
console.log('Transport A received:', transportA.sentMessages);
console.log('Transport B received:', transportB.sentMessages);
// FIXED: Each transport now receives its own response
// Transport A should receive response for request ID 1
expect(transportA.sentMessages.length).toBe(1);
expect(transportA.sentMessages[0]).toMatchObject({
jsonrpc: '2.0',
id: 1,
result: { data: 'responseForA' }
});
expect(transportA.sentMessages).toHaveLength(1);
expect(transportA.sentMessages[0]).toMatchObject({ jsonrpc: '2.0', id: 1, result: {} });
// Transport B should only receive its own response (when implemented)
expect(transportB.sentMessages.length).toBe(1);
expect(transportB.sentMessages[0]).toMatchObject({
jsonrpc: '2.0',
id: 2,
result: { data: 'responseForA' } // Same handler result in this test
});
// Transport B should receive response for request ID 2
expect(transportB.sentMessages).toHaveLength(1);
expect(transportB.sentMessages[0]).toMatchObject({ jsonrpc: '2.0', id: 2, result: {} });
});
test('demonstrates the timing issue with multiple rapid connections', async () => {
const delays: number[] = [];
const results: { transport: string; response: JSONRPCMessage[] }[] = [];
const DelayedRequestSchema = z.object({
method: z.literal('test/delayed'),
params: z
.object({
delay: z.number(),
client: z.string()
})
.optional()
});
// Set up handler with variable delay
protocol.setRequestHandler(DelayedRequestSchema, async (request, extra) => {
const delay = request.params?.delay || 0;
delays.push(delay);
// Set up handler with variable delay based on request id
protocol.setRequestHandler('ping', async (_request, extra) => {
const delay = extra.requestId === 1 ? 50 : 10;
await new Promise(resolve => setTimeout(resolve, delay));
return {
processedBy: `handler-${extra.requestId}`,
delay: delay
} as Result;
return {};
});
// Rapid succession of connections and requests
await protocol.connect(transportA);
transportA.onmessage?.({
jsonrpc: '2.0' as const,
method: 'test/delayed',
params: { delay: 50, client: 'A' },
id: 1
});
transportA.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 1 });
// Connect B while A is processing
setTimeout(async () => {
await protocol.connect(transportB);
transportB.onmessage?.({
jsonrpc: '2.0' as const,
method: 'test/delayed',
params: { delay: 10, client: 'B' },
id: 2
});
transportB.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 2 });
}, 10);
// Wait for all processing
@ -183,8 +115,7 @@ describe('Protocol transport handling bug', () => {
console.log('Timing test results:', results);
// FIXED: Each transport receives its own responses
expect(transportA.sentMessages.length).toBe(1);
expect(transportB.sentMessages.length).toBe(1);
expect(transportA.sentMessages).toHaveLength(1);
expect(transportB.sentMessages).toHaveLength(1);
});
});

57
packages/server/src/server/mcp.ts
View file

@ -36,23 +36,14 @@ import type {
import {
assertCompleteRequestPrompt,
assertCompleteRequestResourceTemplate,
CallToolRequestSchema,
CompleteRequestSchema,
ErrorCode,
getLiteralValue,
getObjectShape,
getParseErrorMessage,
GetPromptRequestSchema,
getSchemaDescription,
isSchemaOptional,
ListPromptsRequestSchema,
ListResourcesRequestSchema,
ListResourceTemplatesRequestSchema,
ListToolsRequestSchema,
McpError,
normalizeObjectSchema,
objectFromShape,
ReadResourceRequestSchema,
safeParseAsync,
toJsonSchemaCompat,
UriTemplate,
@ -128,8 +119,8 @@ export class McpServer {
return;
}
this.server.assertCanSetRequestHandler(getMethodValue(ListToolsRequestSchema));
this.server.assertCanSetRequestHandler(getMethodValue(CallToolRequestSchema));
this.server.assertCanSetRequestHandler('tools/list');
this.server.assertCanSetRequestHandler('tools/call');
this.server.registerCapabilities({
tools: {
@ -138,7 +129,7 @@ export class McpServer {
});
this.server.setRequestHandler(
ListToolsRequestSchema,
'tools/list',
(): ListToolsResult => ({
tools: Object.entries(this._registeredTools)
.filter(([, tool]) => tool.enabled)
@ -176,7 +167,7 @@ export class McpServer {
})
);
this.server.setRequestHandler(CallToolRequestSchema, async (request, extra): Promise<CallToolResult | CreateTaskResult> => {
this.server.setRequestHandler('tools/call', async (request, extra): Promise<CallToolResult | CreateTaskResult> => {
try {
const tool = this._registeredTools[request.params.name];
if (!tool) {
@ -404,13 +395,13 @@ export class McpServer {
return;
}
this.server.assertCanSetRequestHandler(getMethodValue(CompleteRequestSchema));
this.server.assertCanSetRequestHandler('completion/complete');
this.server.registerCapabilities({
completions: {}
});
this.server.setRequestHandler(CompleteRequestSchema, async (request): Promise<CompleteResult> => {
this.server.setRequestHandler('completion/complete', async (request): Promise<CompleteResult> => {
switch (request.params.ref.type) {
case 'ref/prompt': {
assertCompleteRequestPrompt(request);
@ -490,9 +481,9 @@ export class McpServer {
return;
}
this.server.assertCanSetRequestHandler(getMethodValue(ListResourcesRequestSchema));
this.server.assertCanSetRequestHandler(getMethodValue(ListResourceTemplatesRequestSchema));
this.server.assertCanSetRequestHandler(getMethodValue(ReadResourceRequestSchema));
this.server.assertCanSetRequestHandler('resources/list');
this.server.assertCanSetRequestHandler('resources/templates/list');
this.server.assertCanSetRequestHandler('resources/read');
this.server.registerCapabilities({
resources: {
@ -500,7 +491,7 @@ export class McpServer {
}
});
this.server.setRequestHandler(ListResourcesRequestSchema, async (request, extra) => {
this.server.setRequestHandler('resources/list', async (_request, extra) => {
const resources = Object.entries(this._registeredResources)
.filter(([_, resource]) => resource.enabled)
.map(([uri, resource]) => ({
@ -528,7 +519,7 @@ export class McpServer {
return { resources: [...resources, ...templateResources] };
});
this.server.setRequestHandler(ListResourceTemplatesRequestSchema, async () => {
this.server.setRequestHandler('resources/templates/list', async () => {
const resourceTemplates = Object.entries(this._registeredResourceTemplates).map(([name, template]) => ({
name,
uriTemplate: template.resourceTemplate.uriTemplate.toString(),
@ -538,7 +529,7 @@ export class McpServer {
return { resourceTemplates };
});
this.server.setRequestHandler(ReadResourceRequestSchema, async (request, extra) => {
this.server.setRequestHandler('resources/read', async (request, extra) => {
const uri = new URL(request.params.uri);
// First check for exact resource match
@ -571,8 +562,8 @@ export class McpServer {
return;
}
this.server.assertCanSetRequestHandler(getMethodValue(ListPromptsRequestSchema));
this.server.assertCanSetRequestHandler(getMethodValue(GetPromptRequestSchema));
this.server.assertCanSetRequestHandler('prompts/list');
this.server.assertCanSetRequestHandler('prompts/get');
this.server.registerCapabilities({
prompts: {
@ -581,7 +572,7 @@ export class McpServer {
});
this.server.setRequestHandler(
ListPromptsRequestSchema,
'prompts/list',
(): ListPromptsResult => ({
prompts: Object.entries(this._registeredPrompts)
.filter(([, prompt]) => prompt.enabled)
@ -596,7 +587,7 @@ export class McpServer {
})
);
this.server.setRequestHandler(GetPromptRequestSchema, async (request, extra): Promise<GetPromptResult> => {
this.server.setRequestHandler('prompts/get', async (request, extra): Promise<GetPromptResult> => {
const prompt = this._registeredPrompts[request.params.name];
if (!prompt) {
throw new McpError(ErrorCode.InvalidParams, `Prompt ${request.params.name} not found`);
@ -1263,22 +1254,6 @@ function promptArgumentsFromSchema(schema: AnyObjectSchema): PromptArgument[] {
});
}
function getMethodValue(schema: AnyObjectSchema): string {
const shape = getObjectShape(schema);
const methodSchema = shape?.method as AnySchema | undefined;
if (!methodSchema) {
throw new Error('Schema is missing a method literal');
}
// Extract literal value - works for both v3 and v4
const value = getLiteralValue(methodSchema);
if (typeof value === 'string') {
return value;
}
throw new Error('Schema method literal must be a string');
}
function createCompletionResult(suggestions: string[]): CompleteResult {
return {
completion: {

54
packages/server/src/server/server.ts
View file

@ -1,5 +1,4 @@
import type {
AnyObjectSchema,
ClientCapabilities,
CreateMessageRequest,
CreateMessageRequestParamsBase,
@ -22,18 +21,17 @@ import type {
ProtocolOptions,
Request,
RequestHandlerExtra,
RequestMethod,
RequestOptions,
RequestTypeMap,
ResourceUpdatedNotification,
Result,
SchemaOutput,
ServerCapabilities,
ServerNotification,
ServerRequest,
ServerResult,
ToolResultContent,
ToolUseContent,
ZodV3Internal,
ZodV4Internal
ToolUseContent
} from '@modelcontextprotocol/core';
import {
AjvJsonSchemaValidator,
@ -47,10 +45,6 @@ import {
ElicitResultSchema,
EmptyResultSchema,
ErrorCode,
getObjectShape,
InitializedNotificationSchema,
InitializeRequestSchema,
isZ4Schema,
LATEST_PROTOCOL_VERSION,
ListRootsResultSchema,
LoggingLevelSchema,
@ -58,7 +52,6 @@ import {
mergeCapabilities,
Protocol,
safeParse,
SetLevelRequestSchema,
SUPPORTED_PROTOCOL_VERSIONS
} from '@modelcontextprotocol/core';
@ -162,11 +155,11 @@ export class Server<
this._instructions = options?.instructions;
this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new AjvJsonSchemaValidator();
this.setRequestHandler(InitializeRequestSchema, request => this._oninitialize(request));
this.setNotificationHandler(InitializedNotificationSchema, () => this.oninitialized?.());
this.setRequestHandler('initialize', request => this._oninitialize(request));
this.setNotificationHandler('notifications/initialized', () => this.oninitialized?.());
if (this._capabilities.logging) {
this.setRequestHandler(SetLevelRequestSchema, async (request, extra) => {
this.setRequestHandler('logging/setLevel', async (request, extra) => {
const transportSessionId: string | undefined =
extra.sessionId || (extra.requestInfo?.headers.get('mcp-session-id') as string) || undefined;
const { level } = request.params;
@ -222,39 +215,16 @@ export class Server<
/**
* Override request handler registration to enforce server-side validation for tools/call.
*/
public override setRequestHandler<T extends AnyObjectSchema>(
requestSchema: T,
public override setRequestHandler<M extends RequestMethod>(
method: M,
handler: (
request: SchemaOutput<T>,
request: RequestTypeMap[M],
extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
) => ServerResult | ResultT | Promise<ServerResult | ResultT>
): void {
const shape = getObjectShape(requestSchema);
const methodSchema = shape?.method;
if (!methodSchema) {
throw new Error('Schema is missing a method literal');
}
// Extract literal value using type-safe property access
let methodValue: unknown;
if (isZ4Schema(methodSchema)) {
const v4Schema = methodSchema as unknown as ZodV4Internal;
const v4Def = v4Schema._zod?.def;
methodValue = v4Def?.value ?? v4Schema.value;
} else {
const v3Schema = methodSchema as unknown as ZodV3Internal;
const legacyDef = v3Schema._def;
methodValue = legacyDef?.value ?? v3Schema.value;
}
if (typeof methodValue !== 'string') {
throw new TypeError('Schema method literal must be a string');
}
const method = methodValue;
if (method === 'tools/call') {
const wrappedHandler = async (
request: SchemaOutput<T>,
request: RequestTypeMap[M],
extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
): Promise<ServerResult | ResultT> => {
const validatedRequest = safeParse(CallToolRequestSchema, request);
@ -293,11 +263,11 @@ export class Server<
};
// Install the wrapped handler
return super.setRequestHandler(requestSchema, wrappedHandler as unknown as typeof handler);
return super.setRequestHandler(method, wrappedHandler as unknown as typeof handler);
}
// Other handlers use default behavior
return super.setRequestHandler(requestSchema, handler);
return super.setRequestHandler(method, handler);
}
protected assertCapabilityForMethod(method: RequestT['method']): void {

2
packages/server/test/server/stdio.test.ts
View file

@ -17,7 +17,7 @@ beforeEach(() => {
outputBuffer = new ReadBuffer();
output = new Writable({
write(chunk, encoding, callback) {
write(chunk, _encoding, callback) {
outputBuffer.append(chunk);
callback();
}