Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Replace abortAfterTimeout with RequestOptions.timeout #60

Merged
merged 6 commits into from
Nov 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# MCP TypeScript SDK
# MCP TypeScript SDK ![NPM Version](https://img.shields.io/npm/v/%40modelcontextprotocol%2Fsdk)

TypeScript implementation of the Model Context Protocol (MCP), providing both client and server capabilities for integrating with LLM surfaces.

Expand Down
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "@modelcontextprotocol/sdk",
"version": "0.5.0",
"version": "0.6.0",
"description": "Model Context Protocol implementation for TypeScript",
"license": "MIT",
"author": "Anthropic, PBC (https://anthropic.com)",
Expand Down
56 changes: 56 additions & 0 deletions src/client/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
ListToolsRequestSchema,
CreateMessageRequestSchema,
ListRootsRequestSchema,
ErrorCode,
} from "../types.js";
import { Transport } from "../shared/transport.js";
import { Server } from "../server/index.js";
Expand Down Expand Up @@ -491,3 +492,58 @@ test("should handle client cancelling a request", async () => {
// Request should be rejected
await expect(listResourcesPromise).rejects.toBe("Cancelled by test");
});

test("should handle request timeout", async () => {
const server = new Server(
{
name: "test server",
version: "1.0",
},
{
capabilities: {
resources: {},
},
},
);

// Set up server with a delayed response
server.setRequestHandler(
ListResourcesRequestSchema,
async (_request, extra) => {
const timer = new Promise((resolve) => {
const timeout = setTimeout(resolve, 100);
extra.signal.addEventListener("abort", () => clearTimeout(timeout));
});

await timer;
return {
resources: [],
};
},
);

const [clientTransport, serverTransport] =
InMemoryTransport.createLinkedPair();

const client = new Client(
{
name: "test client",
version: "1.0",
},
{
capabilities: {},
},
);

await Promise.all([
client.connect(clientTransport),
server.connect(serverTransport),
]);

// Request with 0 msec timeout should fail immediately
await expect(
client.listResources(undefined, { timeout: 0 }),
).rejects.toMatchObject({
code: ErrorCode.RequestTimeout,
});
});
70 changes: 70 additions & 0 deletions src/server/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
ListResourcesRequestSchema,
ListToolsRequestSchema,
SetLevelRequestSchema,
ErrorCode,
} from "../types.js";
import { Transport } from "../shared/transport.js";
import { InMemoryTransport } from "../inMemory.js";
Expand Down Expand Up @@ -475,3 +476,72 @@ test("should handle server cancelling a request", async () => {
// Request should be rejected
await expect(createMessagePromise).rejects.toBe("Cancelled by test");
});
test("should handle request timeout", async () => {
const server = new Server(
{
name: "test server",
version: "1.0",
},
{
capabilities: {
sampling: {},
},
},
);

// Set up client that delays responses
const client = new Client(
{
name: "test client",
version: "1.0",
},
{
capabilities: {
sampling: {},
},
},
);

client.setRequestHandler(
CreateMessageRequestSchema,
async (_request, extra) => {
await new Promise((resolve, reject) => {
const timeout = setTimeout(resolve, 100);
extra.signal.addEventListener("abort", () => {
clearTimeout(timeout);
reject(extra.signal.reason);
});
});

return {
model: "test",
role: "assistant",
content: {
type: "text",
text: "Test response",
},
};
},
);

const [clientTransport, serverTransport] =
InMemoryTransport.createLinkedPair();

await Promise.all([
client.connect(clientTransport),
server.connect(serverTransport),
]);

// Request with 0 msec timeout should fail immediately
await expect(
server.createMessage(
{
messages: [],
maxTokens: 10,
},
{ timeout: 0 },
),
).rejects.toMatchObject({
code: ErrorCode.RequestTimeout,
});
});
70 changes: 57 additions & 13 deletions src/shared/protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,11 @@ export type ProtocolOptions = {
enforceStrictCapabilities?: boolean;
};

/**
* The default request timeout, in miliseconds.
*/
export const DEFAULT_REQUEST_TIMEOUT_MSEC = 60000;

/**
* Options that can be given per request.
*/
Expand All @@ -48,10 +53,15 @@ export type RequestOptions = {

/**
* Can be used to cancel an in-flight request. This will cause an AbortError to be raised from request().
*
* Use abortAfterTimeout() to easily implement timeouts using this signal.
*/
signal?: AbortSignal;

/**
* A timeout (in milliseconds) for this request. If exceeded, an McpError with code `RequestTimeout` will be raised from request().
*
* If not specified, `DEFAULT_REQUEST_TIMEOUT_MSEC` will be used as the timeout.
*/
timeout?: number;
};

/**
Expand Down Expand Up @@ -381,7 +391,13 @@ export abstract class Protocol<
};
}

let timeoutId: ReturnType<typeof setTimeout> | undefined = undefined;

this._responseHandlers.set(messageId, (response) => {
if (timeoutId !== undefined) {
clearTimeout(timeoutId);
}

if (options?.signal?.aborted) {
return;
}
Expand All @@ -398,24 +414,52 @@ export abstract class Protocol<
}
});

options?.signal?.addEventListener("abort", () => {
const reason = options?.signal?.reason;
const cancel = (reason: unknown) => {
this._responseHandlers.delete(messageId);
this._progressHandlers.delete(messageId);

this._transport?.send({
jsonrpc: "2.0",
method: "cancelled",
params: {
requestId: messageId,
reason: String(reason),
},
});
this._transport
?.send({
jsonrpc: "2.0",
method: "cancelled",
params: {
requestId: messageId,
reason: String(reason),
},
})
.catch((error) =>
this._onerror(new Error(`Failed to send cancellation: ${error}`)),
);

reject(reason);
};

options?.signal?.addEventListener("abort", () => {
if (timeoutId !== undefined) {
clearTimeout(timeoutId);
}

cancel(options?.signal?.reason);
});

this._transport.send(jsonrpcRequest).catch(reject);
const timeout = options?.timeout ?? DEFAULT_REQUEST_TIMEOUT_MSEC;
timeoutId = setTimeout(
() =>
cancel(
new McpError(ErrorCode.RequestTimeout, "Request timed out", {
timeout,
}),
),
timeout,
);

this._transport.send(jsonrpcRequest).catch((error) => {
if (timeoutId !== undefined) {
clearTimeout(timeoutId);
}

reject(error);
});
});
}

Expand Down
1 change: 1 addition & 0 deletions src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ export const JSONRPCResponseSchema = z
export enum ErrorCode {
// SDK error codes
ConnectionClosed = -1,
RequestTimeout = -2,

// Standard JSON-RPC error codes
ParseError = -32700,
Expand Down
15 changes: 0 additions & 15 deletions src/utils.test.ts

This file was deleted.

11 changes: 0 additions & 11 deletions src/utils.ts

This file was deleted.