From bafd9e74f3b681c0bff38664c4399d55a4d258c0 Mon Sep 17 00:00:00 2001 From: ihrpr Date: Sat, 5 Apr 2025 19:52:29 +0100 Subject: [PATCH] Add ways to associate related requests and notifications --- src/server/mcp.test.ts | 8 ++-- src/server/mcp.ts | 16 ++++---- src/shared/protocol.ts | 82 ++++++++++++++++++++++++++++++----------- src/shared/transport.ts | 6 ++- 4 files changed, 78 insertions(+), 34 deletions(-) diff --git a/src/server/mcp.test.ts b/src/server/mcp.test.ts index 2e91a568..ae11279b 100644 --- a/src/server/mcp.test.ts +++ b/src/server/mcp.test.ts @@ -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: [ diff --git a/src/server/mcp.ts b/src/server/mcp.ts index 8f4a909c..484084fc 100644 --- a/src/server/mcp.ts +++ b/src/server/mcp.ts @@ -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 ZodRawShape ? ( args: z.objectOutputType, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => CallToolResult | Promise - : (extra: RequestHandlerExtra) => CallToolResult | Promise; + : (extra: RequestHandlerExtra) => CallToolResult | Promise; type RegisteredTool = { description?: string; @@ -717,7 +719,7 @@ export type ResourceMetadata = Omit; * Callback to list all resources matching a given template. */ export type ListResourcesCallback = ( - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => ListResourcesResult | Promise; /** @@ -725,7 +727,7 @@ export type ListResourcesCallback = ( */ export type ReadResourceCallback = ( uri: URL, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => ReadResourceResult | Promise; type RegisteredResource = { @@ -740,7 +742,7 @@ type RegisteredResource = { export type ReadResourceTemplateCallback = ( uri: URL, variables: Variables, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => ReadResourceResult | Promise; type RegisteredResourceTemplate = { @@ -760,9 +762,9 @@ export type PromptCallback< > = Args extends PromptArgsRawShape ? ( args: z.objectOutputType, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => GetPromptResult | Promise - : (extra: RequestHandlerExtra) => GetPromptResult | Promise; + : (extra: RequestHandlerExtra) => GetPromptResult | Promise; type RegisteredPrompt = { description?: string; diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index a6e47184..b072e578 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -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 = { + /** + * 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; + + /** + * Sends a request that relates to the current request being handled. + * + * This is used by certain transports to correctly associate related messages. + */ + sendRequest: >(request: SendRequestT, resultSchema: U, options?: RequestOptions) => Promise>; + }; /** * Information about a request's timeout state @@ -122,7 +152,7 @@ export abstract class Protocol< string, ( request: JSONRPCRequest, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => Promise > = new Map(); private _requestHandlerAbortControllers: Map = @@ -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 = { 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> { + 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 { + async notification(notification: SendNotificationT, options?: NotificationOptions): Promise { 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, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => SendResultT | Promise, ): 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)); + }); } /** diff --git a/src/shared/transport.ts b/src/shared/transport.ts index b80e2a51..84d2c829 100644 --- a/src/shared/transport.ts +++ b/src/shared/transport.ts @@ -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; + send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise; /** * Closes the connection.