diff --git a/src/cli.ts b/src/cli.ts index 5d640877..d5444972 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -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: {}, + }, }, ); diff --git a/src/client/index.test.ts b/src/client/index.test.ts index 62e30b6b..7cbafd17 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -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: {}, + }, }, ); diff --git a/src/client/index.ts b/src/client/index.ts index 83d7bf73..402c3133 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -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 { @@ -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, diff --git a/src/server/index.test.ts b/src/server/index.test.ts index f5d93117..274bb07c 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -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: {}, + }, }, ); diff --git a/src/server/index.ts b/src/server/index.ts index c8be4053..4558b348 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -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 { @@ -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, diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 22d2503f..ae305413 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -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; - 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, diff --git a/src/types.ts b/src/types.ts index 0ba6fa2f..a0d2d804 100644 --- a/src/types.ts +++ b/src/types.ts @@ -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; export type ListToolsRequest = z.infer; export type ListToolsResult = z.infer; export type CallToolResult = z.infer; -export type CompatibilityCallToolResult = z.infer; +export type CompatibilityCallToolResult = z.infer< + typeof CompatibilityCallToolResultSchema +>; export type CallToolRequest = z.infer; export type ToolListChangedNotification = z.infer< typeof ToolListChangedNotificationSchema