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

v2 refactor!: remove Zod schema parameter from public-facing APIs (#1606)

Browse source
This commit is contained in:
Konstantin Konstantinov 2026-03-04 20:34:47 +02:00 • committed by GitHub
commit f88c376c52
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 1483 additions and 1708 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -7,7 +7,6 @@ import type {
ClientNotification,
ClientRequest,
ClientResult,
CompatibilityCallToolResultSchema,
CompleteRequest,
GetPromptRequest,
Implementation,
@ -476,7 +475,7 @@ export class Client extends Protocol<ClientContext> {
return;
}
try {
const result = await this.request(
const result = await this._requestWithSchema(
{
method: 'initialize',
params: {
@ -709,22 +708,22 @@ export class Client extends Protocol<ClientContext> {
/** Sends a ping to the server to check connectivity. */
async ping(options?: RequestOptions) {
return this.request({ method: 'ping' }, EmptyResultSchema, options);
return this._requestWithSchema({ method: 'ping' }, EmptyResultSchema, options);
}
/** Requests argument autocompletion suggestions from the server for a prompt or resource. */
async complete(params: CompleteRequest['params'], options?: RequestOptions) {
return this.request({ method: 'completion/complete', params }, CompleteResultSchema, options);
return this._requestWithSchema({ method: 'completion/complete', params }, CompleteResultSchema, options);
}
/** Sets the minimum severity level for log messages sent by the server. */
async setLoggingLevel(level: LoggingLevel, options?: RequestOptions) {
return this.request({ method: 'logging/setLevel', params: { level } }, EmptyResultSchema, options);
return this._requestWithSchema({ method: 'logging/setLevel', params: { level } }, EmptyResultSchema, options);
}
/** Retrieves a prompt by name from the server, passing the given arguments for template substitution. */
async getPrompt(params: GetPromptRequest['params'], options?: RequestOptions) {
return this.request({ method: 'prompts/get', params }, GetPromptResultSchema, options);
return this._requestWithSchema({ method: 'prompts/get', params }, GetPromptResultSchema, options);
}
/**
@ -754,7 +753,7 @@ export class Client extends Protocol<ClientContext> {
console.debug('Client.listPrompts() called but server does not advertise prompts capability - returning empty list');
return { prompts: [] };
}
return this.request({ method: 'prompts/list', params }, ListPromptsResultSchema, options);
return this._requestWithSchema({ method: 'prompts/list', params }, ListPromptsResultSchema, options);
}
/**
@ -784,7 +783,7 @@ export class Client extends Protocol<ClientContext> {
console.debug('Client.listResources() called but server does not advertise resources capability - returning empty list');
return { resources: [] };
}
return this.request({ method: 'resources/list', params }, ListResourcesResultSchema, options);
return this._requestWithSchema({ method: 'resources/list', params }, ListResourcesResultSchema, options);
}
/**
@ -801,22 +800,22 @@ export class Client extends Protocol<ClientContext> {
);
return { resourceTemplates: [] };
}
return this.request({ method: 'resources/templates/list', params }, ListResourceTemplatesResultSchema, options);
return this._requestWithSchema({ method: 'resources/templates/list', params }, ListResourceTemplatesResultSchema, options);
}
/** Reads the contents of a resource by URI. */
async readResource(params: ReadResourceRequest['params'], options?: RequestOptions) {
return this.request({ method: 'resources/read', params }, ReadResourceResultSchema, options);
return this._requestWithSchema({ method: 'resources/read', params }, ReadResourceResultSchema, options);
}
/** Subscribes to change notifications for a resource. The server must support resource subscriptions. */
async subscribeResource(params: SubscribeRequest['params'], options?: RequestOptions) {
return this.request({ method: 'resources/subscribe', params }, EmptyResultSchema, options);
return this._requestWithSchema({ method: 'resources/subscribe', params }, EmptyResultSchema, options);
}
/** Unsubscribes from change notifications for a resource. */
async unsubscribeResource(params: UnsubscribeRequest['params'], options?: RequestOptions) {
return this.request({ method: 'resources/unsubscribe', params }, EmptyResultSchema, options);
return this._requestWithSchema({ method: 'resources/unsubscribe', params }, EmptyResultSchema, options);
}
/**
@ -858,11 +857,7 @@ export class Client extends Protocol<ClientContext> {
* }
* ```
*/
async callTool(
params: CallToolRequest['params'],
resultSchema: typeof CallToolResultSchema | typeof CompatibilityCallToolResultSchema = CallToolResultSchema,
options?: RequestOptions
) {
async callTool(params: CallToolRequest['params'], options?: RequestOptions) {
// Guard: required-task tools need experimental API
if (this.isToolTaskRequired(params.name)) {
throw new ProtocolError(
@ -871,7 +866,7 @@ export class Client extends Protocol<ClientContext> {
);
}
const result = await this.request({ method: 'tools/call', params }, resultSchema, options);
const result = await this._requestWithSchema({ method: 'tools/call', params }, CallToolResultSchema, options);
// Check if the tool has an outputSchema
const validator = this.getToolOutputValidator(params.name);
@ -988,7 +983,7 @@ export class Client extends Protocol<ClientContext> {
console.debug('Client.listTools() called but server does not advertise tools capability - returning empty list');
return { tools: [] };
}
const result = await this.request({ method: 'tools/list', params }, ListToolsResultSchema, options);
const result = await this._requestWithSchema({ method: 'tools/list', params }, ListToolsResultSchema, options);
// Cache the tools and their output schemas for future validation
this.cacheToolMetadata(result.tools);

7
packages/client/src/experimental/tasks/client.examples.ts
View file

@ -7,8 +7,7 @@
* @module
*/
import type { Request, RequestOptions } from '@modelcontextprotocol/core';
import { CallToolResultSchema } from '@modelcontextprotocol/core';
import type { RequestOptions } from '@modelcontextprotocol/core';
import type { Client } from '../../client/client.js';
@ -44,9 +43,9 @@ async function ExperimentalClientTasks_callToolStream(client: Client) {
/**
* Example: Using requestStream to consume task lifecycle events for any request type.
*/
async function ExperimentalClientTasks_requestStream(client: Client, request: Request, options: RequestOptions) {
async function ExperimentalClientTasks_requestStream(client: Client, options: RequestOptions) {
//#region ExperimentalClientTasks_requestStream
const stream = client.experimental.tasks.requestStream(request, CallToolResultSchema, options);
const stream = client.experimental.tasks.requestStream({ method: 'tools/call', params: { name: 'my-tool', arguments: {} } }, options);
for await (const message of stream) {
switch (message.type) {
case 'taskCreated': {

42
packages/client/src/experimental/tasks/client.ts
View file

@ -8,15 +8,18 @@
import type {
AnyObjectSchema,
CallToolRequest,
CallToolResult,
CancelTaskResult,
CreateTaskResult,
GetTaskResult,
ListTasksResult,
Request,
RequestMethod,
RequestOptions,
ResponseMessage,
ResultTypeMap,
SchemaOutput
} from '@modelcontextprotocol/core';
import { CallToolResultSchema, ProtocolError, ProtocolErrorCode } from '@modelcontextprotocol/core';
import { ProtocolError, ProtocolErrorCode } from '@modelcontextprotocol/core';
import type { Client } from '../../client/client.js';
@ -25,11 +28,10 @@ import type { Client } from '../../client/client.js';
* @internal
*/
interface ClientInternal {
requestStream<T extends AnyObjectSchema>(
request: Request,
resultSchema: T,
requestStream<M extends RequestMethod>(
request: { method: M; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<SchemaOutput<T>>, void, void>;
): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void>;
isToolTask(toolName: string): boolean;
getToolOutputValidator(toolName: string): ((data: unknown) => { valid: boolean; errorMessage?: string }) | undefined;
}
@ -90,7 +92,7 @@ export class ExperimentalClientTasks {
async *callToolStream(
params: CallToolRequest['params'],
options?: RequestOptions
): AsyncGenerator<ResponseMessage<SchemaOutput<typeof CallToolResultSchema>>, void, void> {
): AsyncGenerator<ResponseMessage<CallToolResult | CreateTaskResult>, void, void> {
// Access Client's internal methods
const clientInternal = this._client as unknown as ClientInternal;
@ -102,7 +104,7 @@ export class ExperimentalClientTasks {
task: options?.task ?? (clientInternal.isToolTask(params.name) ? {} : undefined)
};
const stream = clientInternal.requestStream({ method: 'tools/call', params }, CallToolResultSchema, optionsWithTask);
const stream = clientInternal.requestStream({ method: 'tools/call', params }, optionsWithTask);
// Get the validator for this tool (if it has an output schema)
const validator = clientInternal.getToolOutputValidator(params.name);
@ -110,8 +112,9 @@ export class ExperimentalClientTasks {
// Iterate through the stream and validate the final result if needed
for await (const message of stream) {
// If this is a result message and the tool has an output schema, validate it
if (message.type === 'result' && validator) {
const result = message.result;
// Only validate CallToolResult (has 'content'), not CreateTaskResult (has 'task')
if (message.type === 'result' && validator && 'content' in message.result) {
const result = message.result as CallToolResult;
// If tool has outputSchema, it MUST return structuredContent (unless it's an error)
if (!result.structuredContent && !result.isError) {
@ -245,7 +248,7 @@ export class ExperimentalClientTasks {
*
* @example
* ```ts source="./client.examples.ts#ExperimentalClientTasks_requestStream"
* const stream = client.experimental.tasks.requestStream(request, CallToolResultSchema, options);
* const stream = client.experimental.tasks.requestStream({ method: 'tools/call', params: { name: 'my-tool', arguments: {} } }, options);
* for await (const message of stream) {
* switch (message.type) {
* case 'taskCreated': {
@ -269,25 +272,16 @@ export class ExperimentalClientTasks {
* ```
*
* @param request - The request to send
* @param resultSchema - Zod schema for validating the result
* @param options - Optional request options (timeout, signal, task creation params, etc.)
* @returns AsyncGenerator that yields {@linkcode ResponseMessage} objects
*
* @experimental
*/
requestStream<T extends AnyObjectSchema>(
request: Request,
resultSchema: T,
requestStream<M extends RequestMethod>(
request: { method: M; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<SchemaOutput<T>>, void, void> {
): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void> {
// Delegate to the client's underlying Protocol method
type ClientWithRequestStream = {
requestStream<U extends AnyObjectSchema>(
request: Request,
resultSchema: U,
options?: RequestOptions
): AsyncGenerator<ResponseMessage<SchemaOutput<U>>, void, void>;
};
return (this._client as unknown as ClientWithRequestStream).requestStream(request, resultSchema, options);
return (this._client as unknown as ClientInternal).requestStream(request, options);
}
}

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

@ -34,6 +34,7 @@ import type {
RequestMethod,
RequestTypeMap,
Result,
ResultTypeMap,
ServerCapabilities,
Task,
TaskCreationParams,
@ -44,6 +45,7 @@ import {
CreateTaskResultSchema,
getNotificationSchema,
getRequestSchema,
getResultSchema,
GetTaskResultSchema,
isJSONRPCErrorResponse,
isJSONRPCNotification,
@ -297,7 +299,10 @@ export type BaseContext = {
*
* This is used by certain transports to correctly associate related messages.
*/
send: <U extends AnySchema>(request: Request, resultSchema: U, options?: TaskRequestOptions) => Promise<SchemaOutput<U>>;
send: <M extends RequestMethod>(
request: { method: M; params?: Record<string, unknown> },
options?: TaskRequestOptions
) => Promise<ResultTypeMap[M]>;
/**
* Sends a notification that relates to the current request being handled.
@ -808,7 +813,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
method: request.method,
_meta: request.params?._meta,
signal: abortController.signal,
send: async (r, resultSchema, options?) => {
send: async (r, options?) => {
const requestOptions: RequestOptions = { ...options, relatedRequestId: request.id };
if (relatedTaskId && !requestOptions.relatedTask) {
requestOptions.relatedTask = { taskId: relatedTaskId };
@ -817,7 +822,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
if (effectiveTaskId && taskStore) {
await taskStore.updateTaskStatus(effectiveTaskId, 'input_required');
}
return await this.request(r, resultSchema, requestOptions);
return await this.request(r, requestOptions);
},
notify: async notification => {
const notificationOptions: NotificationOptions = { relatedRequestId: request.id };
@ -1030,12 +1035,32 @@ export abstract class Protocol<ContextT extends BaseContext> {
protected abstract assertTaskHandlerCapability(method: string): void;
/**
* Sends a request and returns an AsyncGenerator that yields response messages.
* Sends a request and returns an AsyncGenerator that yields response messages,
* resolving the result schema automatically from the method name.
* The generator is guaranteed to end with either a `'result'` or `'error'` message.
*
* @experimental Use `client.experimental.tasks.requestStream()` to access this method.
*/
protected async *requestStream<T extends AnyObjectSchema>(
protected async *requestStream<M extends RequestMethod>(
request: { method: M; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void> {
const resultSchema = getResultSchema(request.method) as unknown as AnyObjectSchema;
yield* this._requestStreamWithSchema(request as Request, resultSchema, options) as AsyncGenerator<
ResponseMessage<ResultTypeMap[M]>,
void,
void
>;
}
/**
* Sends a request and returns an AsyncGenerator that yields response messages,
* using the provided schema for validation.
*
* This is the internal implementation used by SDK methods that need to specify
* a particular result schema.
*/
protected async *_requestStreamWithSchema<T extends AnyObjectSchema>(
request: Request,
resultSchema: T,
options?: RequestOptions
@ -1045,7 +1070,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
// For non-task requests, just yield the result
if (!task) {
try {
const result = await this.request(request, resultSchema, options);
const result = await this._requestWithSchema(request, resultSchema, options);
yield { type: 'result', result };
} catch (error) {
yield {
@ -1061,7 +1086,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
let taskId: string | undefined;
try {
// Send the request and get the CreateTaskResult
const createResult = await this.request(request, CreateTaskResultSchema, options);
const createResult = await this._requestWithSchema(request, CreateTaskResultSchema, options);
// Extract taskId from the result
if (createResult.task) {
@ -1132,11 +1157,30 @@ export abstract class Protocol<ContextT extends BaseContext> {
}
/**
* Sends a request and waits for a response.
* Sends a request and waits for a response, resolving the result schema
* automatically from the method name.
*
* Do not use this method to emit notifications! Use {@linkcode Protocol.notification | notification()} instead.
*/
request<T extends AnySchema>(request: Request, resultSchema: T, options?: RequestOptions): Promise<SchemaOutput<T>> {
request<M extends RequestMethod>(
request: { method: M; params?: Record<string, unknown> },
options?: RequestOptions
): Promise<ResultTypeMap[M]> {
const resultSchema = getResultSchema(request.method);
return this._requestWithSchema(request as Request, resultSchema, options) as Promise<ResultTypeMap[M]>;
}
/**
* Sends a request and waits for a response, using the provided schema for validation.
*
* This is the internal implementation used by SDK methods that need to specify
* a particular result schema (e.g., for compatibility or task-specific schemas).
*/
protected _requestWithSchema<T extends AnySchema>(
request: Request,
resultSchema: T,
options?: RequestOptions
): Promise<SchemaOutput<T>> {
const { relatedRequestId, resumptionToken, onresumptiontoken, task, relatedTask } = options ?? {};
// Send the request
@ -1299,7 +1343,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
* @experimental Use `client.experimental.tasks.getTask()` to access this method.
*/
protected async getTask(params: GetTaskRequest['params'], options?: RequestOptions): Promise<GetTaskResult> {
return this.request({ method: 'tasks/get', params }, GetTaskResultSchema, options);
return this._requestWithSchema({ method: 'tasks/get', params }, GetTaskResultSchema, options);
}
/**
@ -1312,7 +1356,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
resultSchema: T,
options?: RequestOptions
): Promise<SchemaOutput<T>> {
return this.request({ method: 'tasks/result', params }, resultSchema, options);
return this._requestWithSchema({ method: 'tasks/result', params }, resultSchema, options);
}
/**
@ -1321,7 +1365,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
* @experimental Use `client.experimental.tasks.listTasks()` to access this method.
*/
protected async listTasks(params?: { cursor?: string }, options?: RequestOptions): Promise<SchemaOutput<typeof ListTasksResultSchema>> {
return this.request({ method: 'tasks/list', params }, ListTasksResultSchema, options);
return this._requestWithSchema({ method: 'tasks/list', params }, ListTasksResultSchema, options);
}
/**
@ -1330,7 +1374,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
* @experimental Use `client.experimental.tasks.cancelTask()` to access this method.
*/
protected async cancelTask(params: { taskId: string }, options?: RequestOptions): Promise<SchemaOutput<typeof CancelTaskResultSchema>> {
return this.request({ method: 'tasks/cancel', params }, CancelTaskResultSchema, options);
return this._requestWithSchema({ method: 'tasks/cancel', params }, CancelTaskResultSchema, options);
}
/**

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

@ -2638,7 +2638,39 @@ export type ResultTypeMap = {
'tasks/cancel': CancelTaskResult;
};
/* Runtime schema lookup */
/* Runtime schema lookup — result schemas by method */
const resultSchemas: Record<string, z.core.$ZodType> = {
ping: EmptyResultSchema,
initialize: InitializeResultSchema,
'completion/complete': CompleteResultSchema,
'logging/setLevel': EmptyResultSchema,
'prompts/get': GetPromptResultSchema,
'prompts/list': ListPromptsResultSchema,
'resources/list': ListResourcesResultSchema,
'resources/templates/list': ListResourceTemplatesResultSchema,
'resources/read': ReadResourceResultSchema,
'resources/subscribe': EmptyResultSchema,
'resources/unsubscribe': EmptyResultSchema,
'tools/call': z.union([CallToolResultSchema, CreateTaskResultSchema]),
'tools/list': ListToolsResultSchema,
'sampling/createMessage': z.union([CreateMessageResultWithToolsSchema, CreateTaskResultSchema]),
'elicitation/create': z.union([ElicitResultSchema, CreateTaskResultSchema]),
'roots/list': ListRootsResultSchema,
'tasks/get': GetTaskResultSchema,
'tasks/result': ResultSchema,
'tasks/list': ListTasksResultSchema,
'tasks/cancel': CancelTaskResultSchema
};
/**
* Gets the Zod schema for validating results of a given request method.
* @see getRequestSchema for explanation of the internal type assertion.
*/
export function getResultSchema<M extends RequestMethod>(method: M): z.ZodType<ResultTypeMap[M]> {
return resultSchemas[method] as unknown as z.ZodType<ResultTypeMap[M]>;
}
/* Runtime schema lookup — request schemas by method */
type RequestSchemaType = (typeof ClientRequestSchema.options)[number] | (typeof ServerRequestSchema.options)[number];
type NotificationSchemaType = (typeof ClientNotificationSchema.options)[number] | (typeof ServerNotificationSchema.options)[number];

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

@ -41,10 +41,15 @@ interface TestProtocol {
_taskProgressTokens: Map<string, number>;
_clearTaskQueue: (taskId: string, sessionId?: string) => Promise<void>;
requestTaskStore: (request: Request, authInfo: unknown) => TaskStore;
// Protected task methods (exposed for testing)
// Protected methods (exposed for testing)
_requestWithSchema: <T extends ZodType>(request: Request, resultSchema: T, options?: unknown) => Promise<z.output<T>>;
listTasks: (params?: { cursor?: string }) => Promise<{ tasks: Task[]; nextCursor?: string }>;
cancelTask: (params: { taskId: string }) => Promise<Result>;
requestStream: <T extends Result>(request: Request, schema: ZodType<T>, options?: unknown) => AsyncGenerator<ResponseMessage<T>>;
_requestStreamWithSchema: <T extends Result>(
request: Request,
schema: ZodType<T>,
options?: unknown
) => AsyncGenerator<ResponseMessage<T>>;
}
// Mock Transport class
@ -149,6 +154,15 @@ function assertQueuedRequest(o?: QueuedMessage): asserts o is QueuedRequest {
expect(o?.type).toBe('request');
}
/**
* Helper to call the protected _requestWithSchema method from tests that
* use custom method names not present in RequestMethod.
*/
// eslint-disable-next-line @typescript-eslint/no-explicit-any
function testRequest(proto: Protocol<BaseContext>, request: Request, resultSchema: ZodType, options?: any) {
return (proto as unknown as TestProtocol)._requestWithSchema(request, resultSchema, options);
}
describe('protocol tests', () => {
let protocol: Protocol<BaseContext>;
let transport: MockTransport;
@ -176,7 +190,7 @@ describe('protocol tests', () => {
const mockSchema: ZodType<{ result: string }> = z.object({
result: z.string()
});
await protocol.request(request, mockSchema, {
await testRequest(protocol, request, mockSchema, {
timeout: 0
});
} catch (error) {
@ -230,13 +244,11 @@ describe('protocol tests', () => {
const onProgressMock = vi.fn();
// Start request but don't await - we're testing the sent message
void protocol
.request(request, mockSchema, {
onprogress: onProgressMock
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, mockSchema, {
onprogress: onProgressMock
}).catch(() => {
// May not complete, ignore error
});
expect(sendSpy).toHaveBeenCalledWith(
expect.objectContaining({
@ -270,13 +282,11 @@ describe('protocol tests', () => {
const onProgressMock = vi.fn();
// Start request but don't await - we're testing the sent message
void protocol
.request(request, mockSchema, {
onprogress: onProgressMock
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, mockSchema, {
onprogress: onProgressMock
}).catch(() => {
// May not complete, ignore error
});
expect(sendSpy).toHaveBeenCalledWith(
expect.objectContaining({
@ -310,7 +320,7 @@ describe('protocol tests', () => {
});
// Start request but don't await - we're testing the sent message
void protocol.request(request, mockSchema).catch(() => {
void testRequest(protocol, request, mockSchema).catch(() => {
// May not complete, ignore error
});
@ -341,13 +351,11 @@ describe('protocol tests', () => {
const onProgressMock = vi.fn();
// Start request but don't await - we're testing the sent message
void protocol
.request(request, mockSchema, {
onprogress: onProgressMock
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, mockSchema, {
onprogress: onProgressMock
}).catch(() => {
// May not complete, ignore error
});
expect(sendSpy).toHaveBeenCalledWith(
expect.objectContaining({
@ -380,7 +388,7 @@ describe('protocol tests', () => {
result: z.string()
});
const onProgressMock = vi.fn();
const requestPromise = protocol.request(request, mockSchema, {
const requestPromise = testRequest(protocol, request, mockSchema, {
timeout: 1000,
resetTimeoutOnProgress: false,
onprogress: onProgressMock
@ -418,7 +426,7 @@ describe('protocol tests', () => {
result: z.string()
});
const onProgressMock = vi.fn();
const requestPromise = protocol.request(request, mockSchema, {
const requestPromise = testRequest(protocol, request, mockSchema, {
timeout: 1000,
resetTimeoutOnProgress: true,
onprogress: onProgressMock
@ -459,7 +467,7 @@ describe('protocol tests', () => {
result: z.string()
});
const onProgressMock = vi.fn();
const requestPromise = protocol.request(request, mockSchema, {
const requestPromise = testRequest(protocol, request, mockSchema, {
timeout: 1000,
maxTotalTimeout: 150,
resetTimeoutOnProgress: true,
@ -506,7 +514,7 @@ describe('protocol tests', () => {
const mockSchema: ZodType<{ result: string }> = z.object({
result: z.string()
});
const requestPromise = protocol.request(request, mockSchema, {
const requestPromise = testRequest(protocol, request, mockSchema, {
timeout: 100,
resetTimeoutOnProgress: true
});
@ -521,7 +529,7 @@ describe('protocol tests', () => {
result: z.string()
});
const onProgressMock = vi.fn();
const requestPromise = protocol.request(request, mockSchema, {
const requestPromise = testRequest(protocol, request, mockSchema, {
timeout: 1000,
resetTimeoutOnProgress: true,
onprogress: onProgressMock
@ -566,7 +574,7 @@ describe('protocol tests', () => {
});
const onProgressMock = vi.fn();
const requestPromise = protocol.request(request, mockSchema, {
const requestPromise = testRequest(protocol, request, mockSchema, {
timeout: 1000,
onprogress: onProgressMock
});
@ -1060,16 +1068,14 @@ describe('Task-based execution', () => {
content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
});
void protocol
.request(request, resultSchema, {
task: {
ttl: 30000,
pollInterval: 1000
}
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, resultSchema, {
task: {
ttl: 30000,
pollInterval: 1000
}
}).catch(() => {
// May not complete, ignore error
});
expect(sendSpy).toHaveBeenCalledWith(
expect.objectContaining({
@ -1103,15 +1109,13 @@ describe('Task-based execution', () => {
content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
});
void protocol
.request(request, resultSchema, {
task: {
ttl: 60000
}
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, resultSchema, {
task: {
ttl: 60000
}
}).catch(() => {
// May not complete, ignore error
});
expect(sendSpy).toHaveBeenCalledWith(
expect.objectContaining({
@ -1141,7 +1145,7 @@ describe('Task-based execution', () => {
content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
});
const resultPromise = protocol.request(request, resultSchema, {
const resultPromise = testRequest(protocol, request, resultSchema, {
task: {
ttl: 30000
}
@ -1164,15 +1168,13 @@ describe('Task-based execution', () => {
const resultSchema = z.object({});
// Start the request (don't await completion, just let it send)
void protocol
.request(request, resultSchema, {
relatedTask: {
taskId: 'parent-task-123'
}
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, resultSchema, {
relatedTask: {
taskId: 'parent-task-123'
}
}).catch(() => {
// May not complete, ignore error
});
// Wait a bit for the request to be queued
await new Promise(resolve => setTimeout(resolve, 10));
@ -1230,20 +1232,18 @@ describe('Task-based execution', () => {
});
// Start the request (don't await completion, just let it send)
void protocol
.request(request, resultSchema, {
task: {
ttl: 60000,
pollInterval: 1000
},
relatedTask: {
taskId: 'parent-task'
},
onprogress: vi.fn()
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, resultSchema, {
task: {
ttl: 60000,
pollInterval: 1000
},
relatedTask: {
taskId: 'parent-task'
},
onprogress: vi.fn()
}).catch(() => {
// May not complete, ignore error
});
// Wait a bit for the request to be queued
await new Promise(resolve => setTimeout(resolve, 10));
@ -2482,14 +2482,12 @@ describe('Progress notification support for tasks', () => {
});
// Start a task-augmented request with progress callback
void protocol
.request(request, resultSchema, {
task: { ttl: 60000 },
onprogress: progressCallback
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, resultSchema, {
task: { ttl: 60000 },
onprogress: progressCallback
}).catch(() => {
// May not complete, ignore error
});
// Wait a bit for the request to be sent
await new Promise(resolve => setTimeout(resolve, 10));
@ -2596,14 +2594,12 @@ describe('Progress notification support for tasks', () => {
});
// Start a task-augmented request with progress callback
void protocol
.request(request, resultSchema, {
task: { ttl: 60000 },
onprogress: progressCallback
})
.catch(() => {
// May not complete, ignore error
});
void testRequest(protocol, request, resultSchema, {
task: { ttl: 60000 },
onprogress: progressCallback
}).catch(() => {
// May not complete, ignore error
});
// Wait a bit for the request to be sent
await new Promise(resolve => setTimeout(resolve, 10));
@ -2716,7 +2712,7 @@ describe('Progress notification support for tasks', () => {
})
});
void protocol.request(request, resultSchema, {
void testRequest(protocol, request, resultSchema, {
task: { ttl: 60000 },
onprogress: progressCallback
});
@ -2817,7 +2813,7 @@ describe('Progress notification support for tasks', () => {
})
});
void protocol.request(request, resultSchema, {
void testRequest(protocol, request, resultSchema, {
task: { ttl: 60000 },
onprogress: progressCallback
});
@ -2915,7 +2911,7 @@ describe('Progress notification support for tasks', () => {
})
});
void protocol.request(request, resultSchema, {
void testRequest(protocol, request, resultSchema, {
task: { ttl: 60000 },
onprogress: progressCallback
});
@ -2986,7 +2982,7 @@ describe('Progress notification support for tasks', () => {
const onProgressMock = vi.fn();
void protocol.request(request, resultSchema, {
void testRequest(protocol, request, resultSchema, {
task: {
ttl: 60000
},
@ -3011,7 +3007,7 @@ describe('Progress notification support for tasks', () => {
const onProgressMock = vi.fn();
void protocol.request(request, resultSchema, {
void testRequest(protocol, request, resultSchema, {
task: {
ttl: 30000
},
@ -3061,7 +3057,7 @@ describe('Progress notification support for tasks', () => {
const onProgressMock = vi.fn();
void protocol.request(request, resultSchema, {
void testRequest(protocol, request, resultSchema, {
task: {
ttl: 30000
},
@ -3375,7 +3371,8 @@ describe('Message interception for task-related requests', () => {
const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });
// Send a request with related task metadata (don't await - we're testing queuing)
const requestPromise = server.request(
const requestPromise = testRequest(
server,
{
method: 'ping',
params: {}
@ -3427,7 +3424,8 @@ describe('Message interception for task-related requests', () => {
await server.connect(transport);
// Send a request without related task metadata
const requestPromise = server.request(
const requestPromise = testRequest(
server,
{
method: 'ping',
params: {}
@ -3472,7 +3470,8 @@ describe('Message interception for task-related requests', () => {
const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });
// Send a request with related task metadata
const requestPromise = server.request(
const requestPromise = testRequest(
server,
{
method: 'ping',
params: {}
@ -3528,7 +3527,8 @@ describe('Message interception for task-related requests', () => {
const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });
// Send a request with related task metadata
const requestPromise = server.request(
const requestPromise = testRequest(
server,
{
method: 'ping',
params: {}
@ -3601,7 +3601,8 @@ describe('Message interception for task-related requests', () => {
const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });
// Send a request with related task metadata
void server.request(
void testRequest(
server,
{
method: 'ping',
params: {}
@ -3677,26 +3678,26 @@ describe('Message interception for task-related requests', () => {
// Fill the queue to max capacity (100 messages)
const promises: Promise<unknown>[] = [];
for (let i = 0; i < 100; i++) {
const promise = server
.request(
{
method: 'ping',
params: {}
},
z.object({}),
{
relatedTask: { taskId: task.taskId }
}
)
.catch(() => {
// Requests will remain pending until task completes or fails
});
const promise = testRequest(
server,
{
method: 'ping',
params: {}
},
z.object({}),
{
relatedTask: { taskId: task.taskId }
}
).catch(() => {
// Requests will remain pending until task completes or fails
});
promises.push(promise);
}
// Try to add one more request - should throw an error
await expect(
server.request(
testRequest(
server,
{
method: 'ping',
params: {}
@ -3766,7 +3767,8 @@ describe('Message Interception', () => {
const mockSchema = z.object({ result: z.string() });
// Send a request with relatedTask metadata
const requestPromise = protocol.request(
const requestPromise = testRequest(
protocol,
{
method: 'test/request',
params: { data: 'test' }
@ -3969,7 +3971,8 @@ describe('Message Interception', () => {
const sendSpy = vi.spyOn(transport, 'send');
// Send a request without relatedTask metadata
const requestPromise = protocol.request(
const requestPromise = testRequest(
protocol,
{
method: 'test/request',
params: { data: 'test' }
@ -4030,7 +4033,8 @@ describe('Message Interception', () => {
const mockSchema = z.object({ result: z.string() });
// Send a request with relatedTask metadata
const requestPromise = protocol.request(
const requestPromise = testRequest(
protocol,
{
method: 'test/request',
params: { data: 'test' }
@ -4174,7 +4178,8 @@ describe('Message Interception', () => {
const relatedTask = { taskId: 'task-meta-456' };
const mockSchema = z.object({ result: z.string() });
const requestPromise = protocol.request(
const requestPromise = testRequest(
protocol,
{
method: 'test/request',
params: { data: 'test' }
@ -4366,11 +4371,14 @@ describe('Queue lifecycle management', () => {
const taskId = task.taskId;
// Queue a request (catch rejection to avoid unhandled promise rejection)
const requestPromise = protocol
.request({ method: 'test/request', params: { data: 'test' } }, z.object({ result: z.string() }), {
const requestPromise = testRequest(
protocol,
{ method: 'test/request', params: { data: 'test' } },
z.object({ result: z.string() }),
{
relatedTask: { taskId }
})
.catch(err => err);
}
).catch(err => err);
// Verify request is queued
const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
@ -4391,7 +4399,7 @@ describe('Queue lifecycle management', () => {
await new Promise(resolve => setTimeout(resolve, 50));
// Verify the request promise is rejected
const result = await requestPromise;
const result = (await requestPromise) as Error;
expect(result).toBeInstanceOf(ProtocolError);
expect(result.message).toContain('Task cancelled or completed');
@ -4439,11 +4447,14 @@ describe('Queue lifecycle management', () => {
const taskId = task.taskId;
// Queue a request (catch the rejection to avoid unhandled promise rejection)
const requestPromise = protocol
.request({ method: 'test/request', params: { data: 'test' } }, z.object({ result: z.string() }), {
const requestPromise = testRequest(
protocol,
{ method: 'test/request', params: { data: 'test' } },
z.object({ result: z.string() }),
{
relatedTask: { taskId }
})
.catch(err => err);
}
).catch(err => err);
// Verify request is queued
const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
@ -4453,7 +4464,7 @@ describe('Queue lifecycle management', () => {
(protocol as unknown as TestProtocol)._clearTaskQueue(taskId);
// Verify the request promise is rejected
const result = await requestPromise;
const result = (await requestPromise) as Error;
expect(result).toBeInstanceOf(ProtocolError);
expect(result.message).toContain('Task cancelled or completed');
@ -4472,23 +4483,32 @@ describe('Queue lifecycle management', () => {
const taskId = task.taskId;
// Queue multiple requests (catch rejections to avoid unhandled promise rejections)
const request1Promise = protocol
.request({ method: 'test/request1', params: { data: 'test1' } }, z.object({ result: z.string() }), {
const request1Promise = testRequest(
protocol,
{ method: 'test/request1', params: { data: 'test1' } },
z.object({ result: z.string() }),
{
relatedTask: { taskId }
})
.catch(err => err);
}
).catch(err => err);
const request2Promise = protocol
.request({ method: 'test/request2', params: { data: 'test2' } }, z.object({ result: z.string() }), {
const request2Promise = testRequest(
protocol,
{ method: 'test/request2', params: { data: 'test2' } },
z.object({ result: z.string() }),
{
relatedTask: { taskId }
})
.catch(err => err);
}
).catch(err => err);
const request3Promise = protocol
.request({ method: 'test/request3', params: { data: 'test3' } }, z.object({ result: z.string() }), {
const request3Promise = testRequest(
protocol,
{ method: 'test/request3', params: { data: 'test3' } },
z.object({ result: z.string() }),
{
relatedTask: { taskId }
})
.catch(err => err);
}
).catch(err => err);
// Verify requests are queued
const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
@ -4498,9 +4518,9 @@ describe('Queue lifecycle management', () => {
(protocol as unknown as TestProtocol)._clearTaskQueue(taskId);
// Verify all request promises are rejected
const result1 = await request1Promise;
const result2 = await request2Promise;
const result3 = await request3Promise;
const result1 = (await request1Promise) as Error;
const result2 = (await request2Promise) as Error;
const result3 = (await request3Promise) as Error;
expect(result1).toBeInstanceOf(ProtocolError);
expect(result1.message).toContain('Task cancelled or completed');
@ -4522,11 +4542,14 @@ describe('Queue lifecycle management', () => {
const taskId = task.taskId;
// Queue a request (catch rejection to avoid unhandled promise rejection)
const requestPromise = protocol
.request({ method: 'test/request', params: { data: 'test' } }, z.object({ result: z.string() }), {
const requestPromise = testRequest(
protocol,
{ method: 'test/request', params: { data: 'test' } },
z.object({ result: z.string() }),
{
relatedTask: { taskId }
})
.catch(err => err);
}
).catch(err => err);
// Get the request ID that was sent
const requestResolvers = (protocol as unknown as TestProtocol)._requestResolvers;
@ -4541,7 +4564,7 @@ describe('Queue lifecycle management', () => {
(protocol as unknown as TestProtocol)._clearTaskQueue(taskId);
// Verify request promise is rejected
const result = await requestPromise;
const result = (await requestPromise) as Error;
expect(result).toBeInstanceOf(ProtocolError);
expect(result.message).toContain('Task cancelled or completed');
@ -4575,7 +4598,7 @@ describe('requestStream() method', () => {
// Start the request stream
const streamPromise = (async () => {
const messages = [];
const stream = (protocol as unknown as TestProtocol).requestStream(
const stream = (protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema
);
@ -4621,7 +4644,7 @@ describe('requestStream() method', () => {
// Start the request stream
const streamPromise = (async () => {
const messages = [];
const stream = (protocol as unknown as TestProtocol).requestStream(
const stream = (protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema
);
@ -4674,7 +4697,7 @@ describe('requestStream() method', () => {
// Start the request stream with already-aborted signal
const messages = [];
const stream = (protocol as unknown as TestProtocol).requestStream(
const stream = (protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema,
{
@ -4709,7 +4732,7 @@ describe('requestStream() method', () => {
await protocol.connect(transport);
const messagesPromise = toArrayAsync(
(protocol as unknown as TestProtocol).requestStream(
(protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema
)
@ -4754,7 +4777,7 @@ describe('requestStream() method', () => {
await protocol.connect(transport);
const messagesPromise = toArrayAsync(
(protocol as unknown as TestProtocol).requestStream(
(protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema,
{
@ -4800,7 +4823,7 @@ describe('requestStream() method', () => {
// Collect messages
const messages = await toArrayAsync(
(protocol as unknown as TestProtocol).requestStream(
(protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema,
{
@ -4832,7 +4855,7 @@ describe('requestStream() method', () => {
await protocol.connect(transport);
const messagesPromise = toArrayAsync(
(protocol as unknown as TestProtocol).requestStream(
(protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema
)
@ -4887,7 +4910,7 @@ describe('requestStream() method', () => {
await protocol.connect(transport);
const messagesPromise = toArrayAsync(
(protocol as unknown as TestProtocol).requestStream(
(protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema
)
@ -4953,7 +4976,7 @@ describe('requestStream() method', () => {
transport.send = vi.fn().mockRejectedValue(new Error('Network error'));
const messages = await toArrayAsync(
(protocol as unknown as TestProtocol).requestStream(
(protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema
)
@ -4981,7 +5004,7 @@ describe('requestStream() method', () => {
await protocol.connect(transport);
const messagesPromise = toArrayAsync(
(protocol as unknown as TestProtocol).requestStream(
(protocol as unknown as TestProtocol)._requestStreamWithSchema(
{ method: 'tools/call', params: { name: 'test', arguments: {} } },
CallToolResultSchema
)

32
packages/server/src/experimental/tasks/server.ts
View file

@ -15,13 +15,12 @@ import type {
ElicitResult,
GetTaskResult,
ListTasksResult,
Request,
RequestMethod,
RequestOptions,
ResponseMessage,
Result,
ResultTypeMap,
SchemaOutput
} from '@modelcontextprotocol/core';
import { CreateMessageResultSchema, ElicitResultSchema } from '@modelcontextprotocol/core';
import type { Server } from '../../server/server.js';
@ -30,7 +29,7 @@ import type { Server } from '../../server/server.js';
*
* Access via `server.experimental.tasks`:
* ```typescript
* const stream = server.experimental.tasks.requestStream(request, schema, options);
* const stream = server.experimental.tasks.requestStream(request, options);
* ```
*
* For high-level server usage with task-based tools, use {@linkcode index.McpServer | McpServer}.experimental.tasks instead.
@ -47,27 +46,24 @@ export class ExperimentalServerTasks {
* This method provides streaming access to request processing, allowing you to
* observe intermediate task status updates for task-augmented requests.
*
* @param request - The request to send
* @param resultSchema - Zod schema for validating the result
* @param request - The request to send (method name determines the result schema)
* @param options - Optional request options (timeout, signal, task creation params, etc.)
* @returns AsyncGenerator that yields {@linkcode ResponseMessage} objects
*
* @experimental
*/
requestStream<T extends AnySchema>(
request: Request,
resultSchema: T,
requestStream<M extends RequestMethod>(
request: { method: M; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<SchemaOutput<T> & Result>, void, void> {
): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void> {
// Delegate to the server's underlying Protocol method
type ServerWithRequestStream = {
requestStream<U extends AnySchema>(
request: Request,
resultSchema: U,
requestStream<N extends RequestMethod>(
request: { method: N; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<SchemaOutput<U> & Result>, void, void>;
): AsyncGenerator<ResponseMessage<ResultTypeMap[N]>, void, void>;
};
return (this._server as unknown as ServerWithRequestStream).requestStream(request, resultSchema, options);
return (this._server as unknown as ServerWithRequestStream).requestStream(request, options);
}
/**
@ -163,9 +159,8 @@ export class ExperimentalServerTasks {
method: 'sampling/createMessage',
params
},
CreateMessageResultSchema,
options
);
) as AsyncGenerator<ResponseMessage<CreateMessageResult>, void, void>;
}
/**
@ -241,9 +236,8 @@ export class ExperimentalServerTasks {
method: 'elicitation/create',
params: normalizedParams
},
ElicitResultSchema,
options
);
) as AsyncGenerator<ResponseMessage<ElicitResult>, void, void>;
}
/**

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

@ -460,7 +460,7 @@ export class Server extends Protocol<ServerContext> {
}
async ping() {
return this.request({ method: 'ping' }, EmptyResultSchema);
return this._requestWithSchema({ method: 'ping' }, EmptyResultSchema);
}
/**
@ -540,9 +540,9 @@ export class Server extends Protocol<ServerContext> {
// Use different schemas based on whether tools are provided
if (params.tools) {
return this.request({ method: 'sampling/createMessage', params }, CreateMessageResultWithToolsSchema, options);
return this._requestWithSchema({ method: 'sampling/createMessage', params }, CreateMessageResultWithToolsSchema, options);
}
return this.request({ method: 'sampling/createMessage', params }, CreateMessageResultSchema, options);
return this._requestWithSchema({ method: 'sampling/createMessage', params }, CreateMessageResultSchema, options);
}
/**
@ -562,7 +562,7 @@ export class Server extends Protocol<ServerContext> {
}
const urlParams = params as ElicitRequestURLParams;
return this.request({ method: 'elicitation/create', params: urlParams }, ElicitResultSchema, options);
return this._requestWithSchema({ method: 'elicitation/create', params: urlParams }, ElicitResultSchema, options);
}
case 'form': {
if (!this._clientCapabilities?.elicitation?.form) {
@ -572,7 +572,11 @@ export class Server extends Protocol<ServerContext> {
const formParams: ElicitRequestFormParams =
params.mode === 'form' ? (params as ElicitRequestFormParams) : { ...(params as ElicitRequestFormParams), mode: 'form' };
const result = await this.request({ method: 'elicitation/create', params: formParams }, ElicitResultSchema, options);
const result = await this._requestWithSchema(
{ method: 'elicitation/create', params: formParams },
ElicitResultSchema,
options
);
if (result.action === 'accept' && result.content && formParams.requestedSchema) {
try {
@ -629,7 +633,7 @@ export class Server extends Protocol<ServerContext> {
}
async listRoots(params?: ListRootsRequest['params'], options?: RequestOptions) {
return this.request({ method: 'roots/list', params }, ListRootsResultSchema, options);
return this._requestWithSchema({ method: 'roots/list', params }, ListRootsResultSchema, options);
}
/**