Add ways to associate related requests and notifications

Author: ihrpr

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: [

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 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,15 +719,15 @@ 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>;
 
 /**
  * Callback to read a resource at a given URI.
  */
 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;

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;
 };
 
 /**
- * Extra data given to request handlers.
+ * Options that can be given per notification.
  */
-export type RequestHandlerExtra = {
+export type NotificationOptions = {
   /**
-   * An abort signal used to communicate if the request was cancelled from the sender's side.
+   * May be used to indicate to the transport which incoming request to associate this outgoing notification with.
    */
-  signal: AbortSignal;
+  relatedRequestId?: RequestId;
+}
 
-  /**
-   * The session ID from the transport, if available.
-   */
-  sessionId?: string;
-};
+/**
+ * Extra data given to request handlers.
+ */
+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;
+
+    /**
+     * 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));
+    });
   }
 
   /**

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<void>;
+  send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise<void>;
 
   /**
    * Closes the connection.