Merge remote-tracking branch 'upstream/main' into request-auth-info

This commit is contained in:
Allen Zhou 2025-04-12 16:41:22 -07:00
commit 6f852a1117
No known key found for this signature in database
GPG key ID: 5817C145DE5C5556
2 changed files with 5074 additions and 102 deletions

View file

@ -3,6 +3,10 @@ import {
CancelledNotificationSchema,
ClientCapabilities,
ErrorCode,
isJSONRPCError,
isJSONRPCRequest,
isJSONRPCResponse,
isJSONRPCNotification,
JSONRPCError,
JSONRPCNotification,
JSONRPCRequest,
@ -79,27 +83,57 @@ 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;
/**
* Information about a validated access token, provided to request handlers.
*/
authInfo?: AuthInfo;
/**
* Information about a validated access token, provided to request handlers.
*/
authInfo?: AuthInfo;
/**
* 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
@ -128,7 +162,7 @@ export abstract class Protocol<
string,
(
request: JSONRPCRequest,
extra: RequestHandlerExtra,
extra: RequestHandlerExtra<SendRequestT, SendNotificationT>,
) => Promise<SendResultT>
> = new Map();
private _requestHandlerAbortControllers: Map<RequestId, AbortController> =
@ -247,12 +281,14 @@ export abstract class Protocol<
};
this._transport.onmessage = (message, extra) => {
if (!("method" in message)) {
if (isJSONRPCResponse(message) || isJSONRPCError(message)) {
this._onresponse(message);
} else if ("id" in message) {
} else if (isJSONRPCRequest(message)) {
this._onrequest(message, extra);
} else {
} else if (isJSONRPCNotification(message)) {
this._onnotification(message);
} else {
this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`));
}
};
@ -321,11 +357,16 @@ export abstract class Protocol<
const abortController = new AbortController();
this._requestHandlerAbortControllers.set(request.id, abortController);
// Create fullExtra object with both abort signal and sessionId from transport
// Create `fullExtra` object with both abort signal and sessionId from transport
// in addition to any authInfo provided
const fullExtra: RequestHandlerExtra = {
const fullExtra: 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 }),
authInfo: extra?.authInfo,
};
@ -372,7 +413,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)}`));
@ -381,7 +422,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);
@ -410,7 +451,7 @@ export abstract class Protocol<
this._progressHandlers.delete(messageId);
this._cleanupTimeout(messageId);
if ("result" in response) {
if (isJSONRPCResponse(response)) {
handler(response);
} else {
const error = new McpError(
@ -468,6 +509,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"));
@ -508,7 +551,7 @@ export abstract class Protocol<
requestId: messageId,
reason: String(reason),
},
})
}, { relatedRequestId })
.catch((error) =>
this._onerror(new Error(`Failed to send cancellation: ${error}`)),
);
@ -546,7 +589,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);
});
@ -556,7 +599,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");
}
@ -568,7 +611,7 @@ export abstract class Protocol<
jsonrpc: "2.0",
};
await this._transport.send(jsonrpcNotification);
await this._transport.send(jsonrpcNotification, options);
}
/**
@ -584,14 +627,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));
});
}
/**

View file

@ -1,5 +1,5 @@
import { AuthInfo } from "../server/auth/types.js";
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.
@ -16,8 +16,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.