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

Add ways to associate related requests and notifications

Browse source
This commit is contained in:
ihrpr 2025-04-05 19:52:29 +01:00
commit bafd9e74f3
No known key found for this signature in database
4 changed files with 77 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

8
src/server/mcp.test.ts
View file

@ -85,6 +85,8 @@ describe("ResourceTemplate", () => {
const abortController = new AbortController();
const result = await template.listCallback?.({
signal: abortController.signal,
sendRequest: () => { throw new Error("Not implemented") },
sendNotification: () => { throw new Error("Not implemented") }
});
expect(result?.resources).toHaveLength(1);
expect(list).toHaveBeenCalled();
@ -318,7 +320,7 @@ describe("tool()", () => {
// This should succeed
mcpServer.tool("tool1", () => ({ content: [] }));
// This should also succeed and not throw about request handlers
mcpServer.tool("tool2", () => ({ content: [] }));
});
@ -815,7 +817,7 @@ describe("resource()", () => {
},
],
}));
// This should also succeed and not throw about request handlers
mcpServer.resource("resource2", "test://resource2", async () => ({
contents: [
@ -1321,7 +1323,7 @@ describe("prompt()", () => {
},
],
}));
// This should also succeed and not throw about request handlers
mcpServer.prompt("prompt2", async () => ({
messages: [

16
src/server/mcp.ts
View file

@ -37,6 +37,8 @@ import {
PromptArgument,
GetPromptResult,
ReadResourceResult,
ServerRequest,
ServerNotification,
} from "../types.js";
import { Completable, CompletableDef } from "./completable.js";
import { UriTemplate, Variables } from "../shared/uriTemplate.js";
@ -694,9 +696,9 @@ export type ToolCallback<Args extends undefined | ZodRawShape = undefined> =
Args extends ZodRawShape
? (
args: z.objectOutputType<Args, ZodTypeAny>,
extra: RequestHandlerExtra,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
) => CallToolResult | Promise<CallToolResult>
: (extra: RequestHandlerExtra) => CallToolResult | Promise<CallToolResult>;
: (extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => CallToolResult | Promise<CallToolResult>;
type RegisteredTool = {
description?: string;
@ -717,7 +719,7 @@ export type ResourceMetadata = Omit<Resource, "uri" | "name">;
* Callback to list all resources matching a given template.
*/
export type ListResourcesCallback = (
extra: RequestHandlerExtra,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
) => ListResourcesResult | Promise<ListResourcesResult>;
/**
@ -725,7 +727,7 @@ export type ListResourcesCallback = (
*/
export type ReadResourceCallback = (
uri: URL,
extra: RequestHandlerExtra,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
) => ReadResourceResult | Promise<ReadResourceResult>;
type RegisteredResource = {
@ -740,7 +742,7 @@ type RegisteredResource = {
export type ReadResourceTemplateCallback = (
uri: URL,
variables: Variables,
extra: RequestHandlerExtra,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
) => ReadResourceResult | Promise<ReadResourceResult>;
type RegisteredResourceTemplate = {
@ -760,9 +762,9 @@ export type PromptCallback<
> = Args extends PromptArgsRawShape
? (
args: z.objectOutputType<Args, ZodTypeAny>,
extra: RequestHandlerExtra,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
) => GetPromptResult | Promise<GetPromptResult>
: (extra: RequestHandlerExtra) => GetPromptResult | Promise<GetPromptResult>;
: (extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => GetPromptResult | Promise<GetPromptResult>;
type RegisteredPrompt = {
description?: string;

82
src/shared/protocol.ts
View file

@ -78,22 +78,52 @@ export type RequestOptions = {
* If not specified, there is no maximum total timeout.
*/
maxTotalTimeout?: number;
/**
* May be used to indicate to the transport which incoming request to associate this outgoing request with.
*/
relatedRequestId?: RequestId;
};
/**
* Options that can be given per notification.
*/
export type NotificationOptions = {
/**
* May be used to indicate to the transport which incoming request to associate this outgoing notification with.
*/
relatedRequestId?: RequestId;
}
/**
* Extra data given to request handlers.
*/
export type RequestHandlerExtra = {
/**
* An abort signal used to communicate if the request was cancelled from the sender's side.
*/
signal: AbortSignal;
export type RequestHandlerExtra<SendRequestT extends Request,
SendNotificationT extends Notification> = {
/**
* An abort signal used to communicate if the request was cancelled from the sender's side.
*/
signal: AbortSignal;
/**
* The session ID from the transport, if available.
*/
sessionId?: string;
};
/**
* The session ID from the transport, if available.
*/
sessionId?: string;
/**
* Sends a notification that relates to the current request being handled.
*
* This is used by certain transports to correctly associate related messages.
*/
sendNotification: (notification: SendNotificationT) => Promise<void>;
/**
* Sends a request that relates to the current request being handled.
*
* This is used by certain transports to correctly associate related messages.
*/
sendRequest: <U extends ZodType<object>>(request: SendRequestT, resultSchema: U, options?: RequestOptions) => Promise<z.infer<U>>;
};
/**
* Information about a request's timeout state
@ -122,7 +152,7 @@ export abstract class Protocol<
string,
(
request: JSONRPCRequest,
extra: RequestHandlerExtra,
extra: RequestHandlerExtra<SendRequestT, SendNotificationT>,
) => Promise<SendResultT>
> = new Map();
private _requestHandlerAbortControllers: Map<RequestId, AbortController> =
@ -316,9 +346,14 @@ export abstract class Protocol<
this._requestHandlerAbortControllers.set(request.id, abortController);
// Create extra object with both abort signal and sessionId from transport
const extra: RequestHandlerExtra = {
const extra: RequestHandlerExtra<SendRequestT, SendNotificationT> = {
signal: abortController.signal,
sessionId: this._transport?.sessionId,
sendNotification:
(notification) =>
this.notification(notification, { relatedRequestId: request.id }),
sendRequest: (r, resultSchema, options?) =>
this.request(r, resultSchema, { ...options, relatedRequestId: request.id })
};
// Starting with Promise.resolve() puts any synchronous errors into the monad as well.
@ -364,7 +399,7 @@ export abstract class Protocol<
private _onprogress(notification: ProgressNotification): void {
const { progressToken, ...params } = notification.params;
const messageId = Number(progressToken);
const handler = this._progressHandlers.get(messageId);
if (!handler) {
this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`));
@ -373,7 +408,7 @@ export abstract class Protocol<
const responseHandler = this._responseHandlers.get(messageId);
const timeoutInfo = this._timeoutInfo.get(messageId);
if (timeoutInfo && responseHandler && timeoutInfo.resetTimeoutOnProgress) {
try {
this._resetTimeout(messageId);
@ -460,6 +495,8 @@ export abstract class Protocol<
resultSchema: T,
options?: RequestOptions,
): Promise<z.infer<T>> {
const { relatedRequestId } = options ?? {};
return new Promise((resolve, reject) => {
if (!this._transport) {
reject(new Error("Not connected"));
@ -500,7 +537,7 @@ export abstract class Protocol<
requestId: messageId,
reason: String(reason),
},
})
}, { relatedRequestId })
.catch((error) =>
this._onerror(new Error(`Failed to send cancellation: ${error}`)),
);
@ -538,7 +575,7 @@ export abstract class Protocol<
this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false);
this._transport.send(jsonrpcRequest).catch((error) => {
this._transport.send(jsonrpcRequest, { relatedRequestId }).catch((error) => {
this._cleanupTimeout(messageId);
reject(error);
});
@ -548,7 +585,7 @@ export abstract class Protocol<
/**
* Emits a notification, which is a one-way message that does not expect a response.
*/
async notification(notification: SendNotificationT): Promise<void> {
async notification(notification: SendNotificationT, options?: NotificationOptions): Promise<void> {
if (!this._transport) {
throw new Error("Not connected");
}
@ -560,7 +597,7 @@ export abstract class Protocol<
jsonrpc: "2.0",
};
await this._transport.send(jsonrpcNotification);
await this._transport.send(jsonrpcNotification, options);
}
/**
@ -576,14 +613,15 @@ export abstract class Protocol<
requestSchema: T,
handler: (
request: z.infer<T>,
extra: RequestHandlerExtra,
extra: RequestHandlerExtra<SendRequestT, SendNotificationT>,
) => SendResultT | Promise<SendResultT>,
): void {
const method = requestSchema.shape.method.value;
this.assertRequestHandlerCapability(method);
this._requestHandlers.set(method, (request, extra) =>
Promise.resolve(handler(requestSchema.parse(request), extra)),
);
this._requestHandlers.set(method, (request, extra) => {
return Promise.resolve(handler(requestSchema.parse(request), extra));
});
}
/**

6
src/shared/transport.ts
View file

@ -1,4 +1,4 @@
import { JSONRPCMessage } from "../types.js";
import { JSONRPCMessage, RequestId } from "../types.js";
/**
* Describes the minimal contract for a MCP transport that a client or server can communicate over.
@ -15,8 +15,10 @@ export interface Transport {
/**
* Sends a JSON-RPC message (request or response).
*
* If present, `relatedRequestId` is used to indicate to the transport which incoming request to associate this outgoing message with.
*/
send(message: JSONRPCMessage): Promise<void>;
send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise<void>;
/**
* Closes the connection.