diff --git a/src/client/index.ts b/src/client/index.ts index 1ad51616..da0d04c5 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -1,4 +1,5 @@ import { + mergeCapabilities, Protocol, ProtocolOptions, RequestOptions, @@ -44,7 +45,7 @@ export type ClientOptions = ProtocolOptions & { /** * Capabilities to advertise as being supported by this client. */ - capabilities: ClientCapabilities; + capabilities?: ClientCapabilities; }; /** @@ -90,10 +91,25 @@ export class Client< */ constructor( private _clientInfo: Implementation, - options: ClientOptions, + options?: ClientOptions, ) { super(options); - this._capabilities = options.capabilities; + this._capabilities = options?.capabilities ?? {}; + } + + /** + * Registers new capabilities. This can only be called before connecting to a transport. + * + * The new capabilities will be merged with any existing capabilities previously given (e.g., at initialization). + */ + public registerCapabilities(capabilities: ClientCapabilities): void { + if (this.transport) { + throw new Error( + "Cannot register capabilities after connecting to transport", + ); + } + + this._capabilities = mergeCapabilities(this._capabilities, capabilities); } protected assertCapability( diff --git a/src/server/index.ts b/src/server/index.ts index a5345dd4..ab4f4136 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -1,6 +1,7 @@ import z, { AnyZodObject, ZodRawShape, ZodTypeAny } from "zod"; import { zodToJsonSchema } from "zod-to-json-schema"; import { + mergeCapabilities, Protocol, ProtocolOptions, RequestHandlerExtra, @@ -35,19 +36,19 @@ import { ServerRequest, ServerResult, SUPPORTED_PROTOCOL_VERSIONS, - Tool + Tool, } from "../types.js"; export type ServerOptions = ProtocolOptions & { /** * Capabilities to advertise as being supported by this server. */ - capabilities: ServerCapabilities; + capabilities?: ServerCapabilities; }; /** * Callback for a tool handler registered with Server.tool(). - * + * * Parameters will include tool arguments, if applicable, as well as other request handler context. */ export type ToolCallback = @@ -56,9 +57,7 @@ export type ToolCallback = args: z.objectOutputType, extra: RequestHandlerExtra, ) => CallToolResult | Promise - : ( - extra: RequestHandlerExtra, - ) => CallToolResult | Promise; + : (extra: RequestHandlerExtra) => CallToolResult | Promise; type RegisteredTool = { description?: string; @@ -67,7 +66,7 @@ type RegisteredTool = { }; const EMPTY_OBJECT_JSON_SCHEMA = { - type: "object" as const + type: "object" as const, }; /** @@ -119,10 +118,10 @@ export class Server< */ constructor( private _serverInfo: Implementation, - options: ServerOptions, + options?: ServerOptions, ) { super(options); - this._capabilities = options.capabilities; + this._capabilities = options?.capabilities ?? {}; this.setRequestHandler(InitializeRequestSchema, (request) => this._oninitialize(request), @@ -132,6 +131,21 @@ export class Server< ); } + /** + * Registers new capabilities. This can only be called before connecting to a transport. + * + * The new capabilities will be merged with any existing capabilities previously given (e.g., at initialization). + */ + public registerCapabilities(capabilities: ServerCapabilities): void { + if (this.transport) { + throw new Error( + "Cannot register capabilities after connecting to transport", + ); + } + + this._capabilities = mergeCapabilities(this._capabilities, capabilities); + } + protected assertCapabilityForMethod(method: RequestT["method"]): void { switch (method as ServerRequest["method"]) { case "sampling/createMessage": @@ -348,36 +362,54 @@ export class Server< // TODO: Register tool capability - this.setRequestHandler(ListToolsRequestSchema, (): ListToolsResult => ({ - tools: Object.entries(this._registeredTools).map(([name, tool]): Tool => { - return { - name, - description: tool.description, - inputSchema: tool.inputSchema ? zodToJsonSchema(tool.inputSchema) as Tool["inputSchema"] : EMPTY_OBJECT_JSON_SCHEMA, - }; - }) - })); + this.setRequestHandler( + ListToolsRequestSchema, + (): ListToolsResult => ({ + tools: Object.entries(this._registeredTools).map( + ([name, tool]): Tool => { + return { + name, + description: tool.description, + inputSchema: tool.inputSchema + ? (zodToJsonSchema(tool.inputSchema) as Tool["inputSchema"]) + : EMPTY_OBJECT_JSON_SCHEMA, + }; + }, + ), + }), + ); - this.setRequestHandler(CallToolRequestSchema, async (request, extra): Promise => { - const tool = this._registeredTools[request.params.name]; - if (!tool) { - throw new McpError(ErrorCode.InvalidParams, `Tool ${request.params.name} not found`); - } - - if (tool.inputSchema) { - const parseResult = await tool.inputSchema.safeParseAsync(request.params.arguments); - if (!parseResult.success) { - throw new McpError(ErrorCode.InvalidParams, `Invalid arguments for tool ${request.params.name}: ${parseResult.error.message}`); + this.setRequestHandler( + CallToolRequestSchema, + async (request, extra): Promise => { + const tool = this._registeredTools[request.params.name]; + if (!tool) { + throw new McpError( + ErrorCode.InvalidParams, + `Tool ${request.params.name} not found`, + ); } - const args = parseResult.data; - const cb = tool.callback as ToolCallback - return await Promise.resolve(cb(args, extra)); - } else { - const cb = tool.callback as ToolCallback; - return await Promise.resolve(cb(extra)); - } - }); + if (tool.inputSchema) { + const parseResult = await tool.inputSchema.safeParseAsync( + request.params.arguments, + ); + if (!parseResult.success) { + throw new McpError( + ErrorCode.InvalidParams, + `Invalid arguments for tool ${request.params.name}: ${parseResult.error.message}`, + ); + } + + const args = parseResult.data; + const cb = tool.callback as ToolCallback; + return await Promise.resolve(cb(args, extra)); + } else { + const cb = tool.callback as ToolCallback; + return await Promise.resolve(cb(extra)); + } + }, + ); } /** @@ -393,12 +425,21 @@ export class Server< /** * Registers a tool `name` accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments. */ - tool(name: string, paramsSchema: Args, cb: ToolCallback): void; + tool( + name: string, + paramsSchema: Args, + cb: ToolCallback, + ): void; /** * Registers a tool `name` (with a description) accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments. */ - tool(name: string, description: string, paramsSchema: Args, cb: ToolCallback): void; + tool( + name: string, + description: string, + paramsSchema: Args, + cb: ToolCallback, + ): void; tool(name: string, ...rest: unknown[]): void { if (this._registeredTools[name]) { @@ -418,7 +459,8 @@ export class Server< const cb = rest[0] as ToolCallback; this._registeredTools[name] = { description, - inputSchema: paramsSchema === undefined ? undefined : z.object(paramsSchema), + inputSchema: + paramsSchema === undefined ? undefined : z.object(paramsSchema), callback: cb, }; diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 9b4ba569..cf33cd2e 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -1,6 +1,7 @@ import { ZodLiteral, ZodObject, ZodType, z } from "zod"; import { CancelledNotificationSchema, + ClientCapabilities, ErrorCode, JSONRPCError, JSONRPCNotification, @@ -15,6 +16,7 @@ import { Request, RequestId, Result, + ServerCapabilities, } from "../types.js"; import { Transport } from "./transport.js"; @@ -547,3 +549,17 @@ export abstract class Protocol< this._notificationHandlers.delete(method); } } + +export function mergeCapabilities(base: T, additional: T): T { + return Object.entries(additional).reduce( + (acc, [key, value]) => { + if (value && typeof value === "object") { + acc[key] = acc[key] ? { ...acc[key], ...value } : value; + } else { + acc[key] = value; + } + return acc; + }, + { ...base }, + ); +}