From f8f5686966c28dd8543490b0b3cc4dcdb4ffb08f Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Mon, 6 Jan 2025 22:47:35 +0000 Subject: [PATCH] Server.tool() convenience API --- package-lock.json | 19 +++++-- package.json | 3 +- src/server/index.ts | 117 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 134 insertions(+), 5 deletions(-) diff --git a/package-lock.json b/package-lock.json index 5e68c8cd..44a876ac 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11,7 +11,8 @@ "dependencies": { "content-type": "^1.0.5", "raw-body": "^3.0.0", - "zod": "^3.23.8" + "zod": "^3.23.8", + "zod-to-json-schema": "^3.24.1" }, "devDependencies": { "@eslint/js": "^9.8.0", @@ -6146,12 +6147,22 @@ } }, "node_modules/zod": { - "version": "3.23.8", - "resolved": "https://registry.npmjs.org/zod/-/zod-3.23.8.tgz", - "integrity": "sha512-XBx9AXhXktjUqnepgTiE5flcKIYWi/rme0Eaj+5Y0lftuGBq+jyRu/md4WnuxqgP1ubdpNCsYEYPxrzVHD8d6g==", + "version": "3.24.1", + "resolved": "https://registry.npmjs.org/zod/-/zod-3.24.1.tgz", + "integrity": "sha512-muH7gBL9sI1nciMZV67X5fTKKBLtwpZ5VBp1vsOQzj1MhrBZ4wlVCm3gedKZWLp0Oyel8sIGfeiz54Su+OVT+A==", + "license": "MIT", "funding": { "url": "https://github.com/sponsors/colinhacks" } + }, + "node_modules/zod-to-json-schema": { + "version": "3.24.1", + "resolved": "https://registry.npmjs.org/zod-to-json-schema/-/zod-to-json-schema-3.24.1.tgz", + "integrity": "sha512-3h08nf3Vw3Wl3PK+q3ow/lIil81IT2Oa7YpQyUUDsEWbXveMesdfK1xBd2RhCkynwZndAxixji/7SYJJowr62w==", + "license": "ISC", + "peerDependencies": { + "zod": "^3.24.1" + } } } } diff --git a/package.json b/package.json index abc9e5cc..893e0100 100644 --- a/package.json +++ b/package.json @@ -43,7 +43,8 @@ "dependencies": { "content-type": "^1.0.5", "raw-body": "^3.0.0", - "zod": "^3.23.8" + "zod": "^3.23.8", + "zod-to-json-schema": "^3.24.1" }, "devDependencies": { "@eslint/js": "^9.8.0", diff --git a/src/server/index.ts b/src/server/index.ts index d15ad3c0..77c5f5c7 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -1,13 +1,19 @@ +import z, { AnyZodObject, ZodRawShape, ZodTypeAny } from "zod"; +import { zodToJsonSchema } from "zod-to-json-schema"; import { Protocol, ProtocolOptions, + RequestHandlerExtra, RequestOptions, } from "../shared/protocol.js"; import { + CallToolRequestSchema, + CallToolResult, ClientCapabilities, CreateMessageRequest, CreateMessageResultSchema, EmptyResultSchema, + ErrorCode, Implementation, InitializedNotificationSchema, InitializeRequest, @@ -16,7 +22,10 @@ import { LATEST_PROTOCOL_VERSION, ListRootsRequest, ListRootsResultSchema, + ListToolsRequestSchema, + ListToolsResult, LoggingMessageNotification, + McpError, Notification, Request, ResourceUpdatedNotification, @@ -26,6 +35,7 @@ import { ServerRequest, ServerResult, SUPPORTED_PROTOCOL_VERSIONS, + Tool } from "../types.js"; export type ServerOptions = ProtocolOptions & { @@ -35,6 +45,31 @@ export type ServerOptions = ProtocolOptions & { 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 = + Args extends ZodRawShape + ? ( + args: z.objectOutputType, + extra: RequestHandlerExtra, + ) => CallToolResult | Promise + : ( + extra: RequestHandlerExtra, + ) => CallToolResult | Promise; + +type RegisteredTool = { + description?: string; + inputSchema?: AnyZodObject; + callback: ToolCallback; +}; + +const EMPTY_OBJECT_JSON_SCHEMA = { + type: "object" as const +}; + /** * An MCP server on top of a pluggable transport. * @@ -72,6 +107,7 @@ export class Server< private _clientCapabilities?: ClientCapabilities; private _clientVersion?: Implementation; private _capabilities: ServerCapabilities; + private _registeredTools: { [name: string]: RegisteredTool } = {}; /** * Callback for when initialization has fully completed (i.e., the client has sent an `initialized` notification). @@ -305,4 +341,85 @@ export class Server< async sendPromptListChanged() { return this.notification({ method: "notifications/prompts/list_changed" }); } + + private setToolRequestHandlers() { + // TODO: Check that these handlers do not already exist + // 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(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}`); + } + + 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)); + } + }); + } + + /** + * Registers a zero-argument tool `name`, which will run the given function when the client calls it. + */ + tool(name: string, cb: ToolCallback): void; + + /** + * Registers a zero-argument tool `name` (with a description) which will run the given function when the client calls it. + */ + tool(name: string, description: string, cb: ToolCallback): void; + + /** + * 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; + + /** + * 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, ...rest: unknown[]): void { + if (this._registeredTools[name]) { + throw new Error(`Tool ${name} is already registered`); + } + + let description: string | undefined; + if (rest[0] instanceof String) { + description = rest.shift() as string; + } + + let paramsSchema: ZodRawShape | undefined; + if (rest.length > 1) { + paramsSchema = rest.shift() as ZodRawShape; + } + + const cb = rest[0] as ToolCallback; + this._registeredTools[name] = { + description, + inputSchema: paramsSchema === undefined ? undefined : z.object(paramsSchema), + callback: cb, + }; + + this.setToolRequestHandlers(); + } }