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

Handle capabilities checks at the request() level, make non-strict by default

Browse source
This commit is contained in:
Justin Spahr-Summers 2024-11-12 10:57:59 +00:00
commit dc45f5d148
No known key found for this signature in database
7 changed files with 254 additions and 83 deletions
Download patch file Download diff file Expand all files Collapse all files

18
src/cli.ts
View file

@ -23,7 +23,9 @@ async function runClient(url_or_command: string, args: string[]) {
version: "0.1.0",
},
{
sampling: {},
capabilities: {
sampling: {},
},
},
);
@ -73,7 +75,9 @@ async function runServer(port: number | null) {
name: "mcp-typescript test server",
version: "0.1.0",
},
{},
{
capabilities: {},
},
);
servers.push(server);
@ -111,10 +115,12 @@ async function runServer(port: number | null) {
version: "0.1.0",
},
{
prompts: {},
resources: {},
tools: {},
logging: {},
capabilities: {
prompts: {},
resources: {},
tools: {},
logging: {},
},
},
);

27
src/client/index.test.ts
View file

@ -46,7 +46,9 @@ test("should initialize with matching protocol version", async () => {
version: "1.0",
},
{
sampling: {},
capabilities: {
sampling: {},
},
},
);
@ -93,7 +95,9 @@ test("should initialize with supported older protocol version", async () => {
version: "1.0",
},
{
sampling: {},
capabilities: {
sampling: {},
},
},
);
@ -135,7 +139,9 @@ test("should reject unsupported protocol version", async () => {
version: "1.0",
},
{
sampling: {},
capabilities: {
sampling: {},
},
},
);
@ -153,8 +159,10 @@ test("should respect server capabilities", async () => {
version: "1.0",
},
{
resources: {},
tools: {},
capabilities: {
resources: {},
tools: {},
},
},
);
@ -187,7 +195,10 @@ test("should respect server capabilities", async () => {
version: "1.0",
},
{
sampling: {},
capabilities: {
sampling: {},
},
enforceStrictCapabilities: true,
},
);
@ -263,7 +274,9 @@ test("should typecheck", () => {
version: "1.0.0",
},
{
sampling: {},
capabilities: {
sampling: {},
},
},
);

112
src/client/index.ts
View file

@ -1,4 +1,8 @@
import { ProgressCallback, Protocol } from "../shared/protocol.js";
import {
ProgressCallback,
Protocol,
ProtocolOptions,
} from "../shared/protocol.js";
import { Transport } from "../shared/transport.js";
import {
CallToolRequest,
@ -36,6 +40,13 @@ import {
UnsubscribeRequest,
} from "../types.js";
export type ClientOptions = ProtocolOptions & {
/**
* Capabilities to advertise as being supported by this client.
*/
capabilities: ClientCapabilities;
};
/**
* An MCP client on top of a pluggable transport.
*
@ -72,15 +83,28 @@ export class Client<
> {
private _serverCapabilities?: ServerCapabilities;
private _serverVersion?: Implementation;
private _capabilities: ClientCapabilities;
/**
* Initializes this client with the given name and version information.
*/
constructor(
private _clientInfo: Implementation,
private _capabilities: ClientCapabilities,
options: ClientOptions,
) {
super();
super(options);
this._capabilities = options.capabilities;
}
protected assertCapability(
capability: keyof ServerCapabilities,
method: string,
): void {
if (!this._serverCapabilities?.[capability]) {
throw new Error(
`Server does not support ${capability} (required for ${method})`,
);
}
}
override async connect(transport: Transport): Promise<void> {
@ -136,14 +160,69 @@ export class Client<
return this._serverVersion;
}
private assertCapability(
capability: keyof ServerCapabilities,
method: string,
) {
if (!this._serverCapabilities?.[capability]) {
throw new Error(
`Server does not support ${capability} (required for ${method})`,
);
protected assertCapabilityForMethod(method: RequestT["method"]): void {
switch (method as ClientRequest["method"]) {
case "logging/setLevel":
if (!this._serverCapabilities?.logging) {
throw new Error(
"Server does not support logging (required for logging/setLevel)",
);
}
break;
case "prompts/get":
case "prompts/list":
if (!this._serverCapabilities?.prompts) {
throw new Error(
`Server does not support prompts (required for ${method})`,
);
}
break;
case "resources/list":
case "resources/templates/list":
case "resources/read":
case "resources/subscribe":
case "resources/unsubscribe":
if (!this._serverCapabilities?.resources) {
throw new Error(
`Server does not support resources (required for ${method})`,
);
}
if (
method === "resources/subscribe" &&
!this._serverCapabilities.resources.subscribe
) {
throw new Error("Server does not support resource subscriptions");
}
break;
case "tools/call":
case "tools/list":
if (!this._serverCapabilities?.tools) {
throw new Error(
`Server does not support tools (required for ${method})`,
);
}
break;
case "completion/complete":
if (!this._serverCapabilities?.prompts) {
throw new Error(
"Server does not support prompts (required for completion/complete)",
);
}
break;
case "initialize":
// No specific capability required for initialize
break;
case "ping":
// No specific capability required for ping
break;
}
}
@ -155,7 +234,6 @@ export class Client<
params: CompleteRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("prompts", "completion/complete");
return this.request(
{ method: "completion/complete", params },
CompleteResultSchema,
@ -164,7 +242,6 @@ export class Client<
}
async setLoggingLevel(level: LoggingLevel) {
this.assertCapability("logging", "logging/setLevel");
return this.request(
{ method: "logging/setLevel", params: { level } },
EmptyResultSchema,
@ -175,7 +252,6 @@ export class Client<
params: GetPromptRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("prompts", "prompts/get");
return this.request(
{ method: "prompts/get", params },
GetPromptResultSchema,
@ -187,7 +263,6 @@ export class Client<
params?: ListPromptsRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("prompts", "prompts/list");
return this.request(
{ method: "prompts/list", params },
ListPromptsResultSchema,
@ -199,7 +274,6 @@ export class Client<
params?: ListResourcesRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("resources", "resources/list");
return this.request(
{ method: "resources/list", params },
ListResourcesResultSchema,
@ -211,7 +285,6 @@ export class Client<
params?: ListResourceTemplatesRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("resources", "resources/templates/list");
return this.request(
{ method: "resources/templates/list", params },
ListResourceTemplatesResultSchema,
@ -223,7 +296,6 @@ export class Client<
params: ReadResourceRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("resources", "resources/read");
return this.request(
{ method: "resources/read", params },
ReadResourceResultSchema,
@ -232,7 +304,6 @@ export class Client<
}
async subscribeResource(params: SubscribeRequest["params"]) {
this.assertCapability("resources", "resources/subscribe");
return this.request(
{ method: "resources/subscribe", params },
EmptyResultSchema,
@ -240,7 +311,6 @@ export class Client<
}
async unsubscribeResource(params: UnsubscribeRequest["params"]) {
this.assertCapability("resources", "resources/unsubscribe");
return this.request(
{ method: "resources/unsubscribe", params },
EmptyResultSchema,
@ -254,7 +324,6 @@ export class Client<
| typeof CompatibilityCallToolResultSchema = CallToolResultSchema,
onprogress?: ProgressCallback,
) {
this.assertCapability("tools", "tools/call");
return this.request(
{ method: "tools/call", params },
resultSchema,
@ -266,7 +335,6 @@ export class Client<
params?: ListToolsRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("tools", "tools/list");
return this.request(
{ method: "tools/list", params },
ListToolsResultSchema,

60
src/server/index.test.ts
View file

@ -46,10 +46,12 @@ test("should accept latest protocol version", async () => {
version: "1.0",
},
{
prompts: {},
resources: {},
tools: {},
logging: {},
capabilities: {
prompts: {},
resources: {},
tools: {},
logging: {},
},
},
);
@ -105,10 +107,12 @@ test("should accept supported older protocol version", async () => {
version: "1.0",
},
{
prompts: {},
resources: {},
tools: {},
logging: {},
capabilities: {
prompts: {},
resources: {},
tools: {},
logging: {},
},
},
);
@ -163,10 +167,12 @@ test("should handle unsupported protocol version", async () => {
version: "1.0",
},
{
prompts: {},
resources: {},
tools: {},
logging: {},
capabilities: {
prompts: {},
resources: {},
tools: {},
logging: {},
},
},
);
@ -197,19 +203,25 @@ test("should respect client capabilities", async () => {
version: "1.0",
},
{
prompts: {},
resources: {},
tools: {},
logging: {},
capabilities: {
prompts: {},
resources: {},
tools: {},
logging: {},
},
enforceStrictCapabilities: true,
},
);
const client = new Client(
{
name: "test client",
version: "1.0",
},
{
sampling: {},
capabilities: {
sampling: {},
},
},
);
@ -245,9 +257,7 @@ test("should respect client capabilities", async () => {
).resolves.not.toThrow();
// This should still throw because roots are not supported by the client
await expect(server.listRoots()).rejects.toThrow(
"Client does not support roots",
);
await expect(server.listRoots()).rejects.toThrow(/^Client does not support/);
});
/*
@ -301,10 +311,12 @@ test("should typecheck", () => {
version: "1.0.0",
},
{
prompts: {},
resources: {},
tools: {},
logging: {},
capabilities: {
prompts: {},
resources: {},
tools: {},
logging: {},
},
},
);

45
src/server/index.ts
View file

@ -1,4 +1,8 @@
import { ProgressCallback, Protocol } from "../shared/protocol.js";
import {
ProgressCallback,
Protocol,
ProtocolOptions,
} from "../shared/protocol.js";
import {
ClientCapabilities,
CreateMessageRequest,
@ -24,6 +28,13 @@ import {
SUPPORTED_PROTOCOL_VERSIONS,
} from "../types.js";
export type ServerOptions = ProtocolOptions & {
/**
* Capabilities to advertise as being supported by this server.
*/
capabilities: ClientCapabilities;
};
/**
* An MCP server on top of a pluggable transport.
*
@ -60,6 +71,7 @@ export class Server<
> {
private _clientCapabilities?: ClientCapabilities;
private _clientVersion?: Implementation;
private _capabilities: ServerCapabilities;
/**
* Callback for when initialization has fully completed (i.e., the client has sent an `initialized` notification).
@ -71,9 +83,10 @@ export class Server<
*/
constructor(
private _serverInfo: Implementation,
private _capabilities: ServerCapabilities,
options: ServerOptions,
) {
super();
super(options);
this._capabilities = options.capabilities;
this.setRequestHandler(InitializeRequestSchema, (request) =>
this._oninitialize(request),
@ -83,6 +96,30 @@ export class Server<
);
}
protected assertCapabilityForMethod(method: RequestT["method"]): void {
switch (method as ServerRequest["method"]) {
case "sampling/createMessage":
if (!this._clientCapabilities?.sampling) {
throw new Error(
`Client does not support sampling (required for ${method})`,
);
}
break;
case "roots/list":
if (!this._clientCapabilities?.roots) {
throw new Error(
`Client does not support listing roots (required for ${method})`,
);
}
break;
case "ping":
// No specific capability required for ping
break;
}
}
private async _oninitialize(
request: InitializeRequest,
): Promise<InitializeResult> {
@ -137,7 +174,6 @@ export class Server<
params: CreateMessageRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("sampling", "sampling/createMessage");
return this.request(
{ method: "sampling/createMessage", params },
CreateMessageResultSchema,
@ -149,7 +185,6 @@ export class Server<
params?: ListRootsRequest["params"],
onprogress?: ProgressCallback,
) {
this.assertCapability("roots", "roots/list");
return this.request(
{ method: "roots/list", params },
ListRootsResultSchema,

29
src/shared/protocol.ts
View file

@ -21,11 +21,23 @@ import { Transport } from "./transport.js";
*/
export type ProgressCallback = (progress: Progress) => void;
/**
* Additional initialization options.
*/
export type ProtocolOptions = {
/**
* Whether to restrict emitted requests to only those that the remote side has indicated that they can handle, through their advertised capabilities.
*
* Currently this defaults to false, for backwards compatibility with SDK versions that did not advertise capabilities correctly. In future, this will default to true.
*/
enforceStrictCapabilities?: boolean;
};
/**
* Implements MCP protocol framing on top of a pluggable transport, including
* features like request/response linking, notifications, and progress.
*/
export class Protocol<
export abstract class Protocol<
SendRequestT extends Request,
SendNotificationT extends Notification,
SendResultT extends Result,
@ -70,7 +82,7 @@ export class Protocol<
*/
fallbackNotificationHandler?: (notification: Notification) => Promise<void>;
constructor() {
constructor(private _options?: ProtocolOptions) {
this.setNotificationHandler(ProgressNotificationSchema, (notification) => {
this._onprogress(notification as unknown as ProgressNotification);
});
@ -245,6 +257,15 @@ export class Protocol<
await this._transport?.close();
}
/**
* A method to check if a capability is supported by the remote side, for the given method to be called.
*
* This should be implemented by subclasses.
*/
protected abstract assertCapabilityForMethod(
method: SendRequestT["method"],
): void;
/**
* Sends a request and wait for a response, with optional progress notifications in the meantime (if supported by the server).
*
@ -261,6 +282,10 @@ export class Protocol<
return;
}
if (this._options?.enforceStrictCapabilities === true) {
this.assertCapabilityForMethod(request.method);
}
const messageId = this._requestMessageId++;
const jsonrpcRequest: JSONRPCRequest = {
...request,

46
src/types.ts
View file

@ -1,7 +1,10 @@
import { z } from "zod";
export const LATEST_PROTOCOL_VERSION = "2024-11-05";
export const SUPPORTED_PROTOCOL_VERSIONS = [LATEST_PROTOCOL_VERSION, "2024-10-07"];
export const SUPPORTED_PROTOCOL_VERSIONS = [
LATEST_PROTOCOL_VERSION,
"2024-10-07",
];
/* JSON-RPC types */
export const JSONRPC_VERSION = "2.0";
@ -179,7 +182,7 @@ export const ClientCapabilitiesSchema = z
z
.object({
/**
* Whether the client supports notifications for changes to the roots list.
* Whether the client supports issuing notifications for changes to the roots list.
*/
listChanged: z.optional(z.boolean()),
})
@ -223,7 +226,7 @@ export const ServerCapabilitiesSchema = z
z
.object({
/**
* Whether this server supports notifications for changes to the prompt list.
* Whether this server supports issuing notifications for changes to the prompt list.
*/
listChanged: z.optional(z.boolean()),
})
@ -236,11 +239,12 @@ export const ServerCapabilitiesSchema = z
z
.object({
/**
* Whether this server supports subscribing to resource updates.
* Whether this server supports clients subscribing to resource updates.
*/
subscribe: z.optional(z.boolean()),
/**
* Whether this server supports notifications for changes to the resource list.
* Whether this server supports issuing notifications for changes to the resource list.
*/
listChanged: z.optional(z.boolean()),
})
@ -253,7 +257,7 @@ export const ServerCapabilitiesSchema = z
z
.object({
/**
* Whether this server supports notifications for changes to the tool list.
* Whether this server supports issuing notifications for changes to the tool list.
*/
listChanged: z.optional(z.boolean()),
})
@ -725,9 +729,11 @@ export const CallToolResultSchema = ResultSchema.extend({
/**
* CallToolResultSchema extended with backwards compatibility to protocol version 2024-10-07.
*/
export const CompatibilityCallToolResultSchema = CallToolResultSchema.or(ResultSchema.extend({
toolResult: z.unknown(),
}));
export const CompatibilityCallToolResultSchema = CallToolResultSchema.or(
ResultSchema.extend({
toolResult: z.unknown(),
}),
);
/**
* Used by the client to invoke a tool provided by the server.
@ -802,12 +808,14 @@ export const LoggingMessageNotificationSchema = NotificationSchema.extend({
/**
* Hints to use for model selection.
*/
export const ModelHintSchema = z.object({
/**
* A hint for a model name.
*/
name: z.string().optional(),
}).passthrough();
export const ModelHintSchema = z
.object({
/**
* A hint for a model name.
*/
name: z.string().optional(),
})
.passthrough();
/**
* The server's preferences for model selection, requested of the client during sampling.
@ -886,7 +894,9 @@ export const CreateMessageResultSchema = ResultSchema.extend({
/**
* The reason why sampling stopped.
*/
stopReason: z.optional(z.enum(["endTurn", "stopSequence", "maxTokens"]).or(z.string())),
stopReason: z.optional(
z.enum(["endTurn", "stopSequence", "maxTokens"]).or(z.string()),
),
role: z.enum(["user", "assistant"]),
content: z.discriminatedUnion("type", [
TextContentSchema,
@ -1156,7 +1166,9 @@ export type Tool = z.infer<typeof ToolSchema>;
export type ListToolsRequest = z.infer<typeof ListToolsRequestSchema>;
export type ListToolsResult = z.infer<typeof ListToolsResultSchema>;
export type CallToolResult = z.infer<typeof CallToolResultSchema>;
export type CompatibilityCallToolResult = z.infer<typeof CompatibilityCallToolResultSchema>;
export type CompatibilityCallToolResult = z.infer<
typeof CompatibilityCallToolResultSchema
>;
export type CallToolRequest = z.infer<typeof CallToolRequestSchema>;
export type ToolListChangedNotification = z.infer<
typeof ToolListChangedNotificationSchema