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

Support passing an AbortSignal to request()

Browse source
This commit is contained in:
Justin Spahr-Summers 2024-11-14 13:42:18 +00:00
commit 54d146bbdc
No known key found for this signature in database
1 changed files with 64 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

46
src/shared/protocol.ts
View file

@ -35,6 +35,21 @@ export type ProtocolOptions = {
enforceStrictCapabilities?: boolean;
};
/**
* Options that can be given per request.
*/
export type RequestOptions = {
/**
* If set, requests progress notifications from the remote end (if supported). When progress notifications are received, this callback will be invoked.
*/
onprogress?: ProgressCallback;
/**
* Can be used to cancel an in-flight request. This will cause an AbortError to be raised from request().
*/
signal?: AbortSignal;
};
/**
* Implements MCP protocol framing on top of a pluggable transport, including
* features like request/response linking, notifications, and progress.
@ -285,14 +300,14 @@ export abstract class Protocol<
protected abstract assertRequestHandlerCapability(method: string): void;
/**
* Sends a request and wait for a response, with optional progress notifications in the meantime (if supported by the server).
* Sends a request and wait for a response.
*
* Do not use this method to emit notifications! Use notification() instead.
*/
request<T extends ZodType<object>>(
request: SendRequestT,
resultSchema: T,
onprogress?: ProgressCallback,
options?: RequestOptions,
): Promise<z.infer<T>> {
return new Promise((resolve, reject) => {
if (!this._transport) {
@ -304,6 +319,8 @@ export abstract class Protocol<
this.assertCapabilityForMethod(request.method);
}
options?.signal?.throwIfAborted();
const messageId = this._requestMessageId++;
const jsonrpcRequest: JSONRPCRequest = {
...request,
@ -311,8 +328,8 @@ export abstract class Protocol<
id: messageId,
};
if (onprogress) {
this._progressHandlers.set(messageId, onprogress);
if (options?.onprogress) {
this._progressHandlers.set(messageId, options.onprogress);
jsonrpcRequest.params = {
...request.params,
_meta: { progressToken: messageId },
@ -320,6 +337,10 @@ export abstract class Protocol<
}
this._responseHandlers.set(messageId, (response) => {
if (options?.signal?.aborted) {
return;
}
if (response instanceof Error) {
return reject(response);
}
@ -332,6 +353,23 @@ export abstract class Protocol<
}
});
options?.signal?.addEventListener("abort", () => {
const reason = options?.signal?.reason;
this._responseHandlers.delete(messageId);
this._progressHandlers.delete(messageId);
this._transport?.send({
jsonrpc: "2.0",
method: "cancelled",
params: {
requestId: messageId,
reason: String(reason),
},
});
reject(reason);
});
this._transport.send(jsonrpcRequest).catch(reject);
});
}