From 9b95fa8203da0f3c60240faff1ff24dfca88f290 Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Wed, 8 Jan 2025 14:20:22 +0000 Subject: [PATCH] Server.prompt() convenience API --- src/server/index.test.ts | 286 +++++++++++++++++++++++++++++++++++++++ src/server/index.ts | 149 +++++++++++++++++++- 2 files changed, 433 insertions(+), 2 deletions(-) diff --git a/src/server/index.test.ts b/src/server/index.test.ts index b8626005..89bb578b 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -20,6 +20,8 @@ import { ListResourcesResultSchema, ListResourceTemplatesResultSchema, ReadResourceResultSchema, + ListPromptsResultSchema, + GetPromptResultSchema, } from "../types.js"; import { Transport } from "../shared/transport.js"; import { InMemoryTransport } from "../inMemory.js"; @@ -1334,3 +1336,287 @@ describe("Server.resource", () => { ).rejects.toThrow(/Resource test:\/\/nonexistent not found/); }); }); + +describe("Server.prompt", () => { + test("should register zero-argument prompt", async () => { + const server = new Server({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + server.prompt("test", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response", + }, + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "prompts/list", + }, + ListPromptsResultSchema, + ); + + expect(result.prompts).toHaveLength(1); + expect(result.prompts[0].name).toBe("test"); + expect(result.prompts[0].arguments).toBeUndefined(); + }); + + test("should register prompt with args schema", async () => { + const server = new Server({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + server.prompt( + "test", + { + name: z.string(), + value: z.string(), + }, + async ({ name, value }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `${name}: ${value}`, + }, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "prompts/list", + }, + ListPromptsResultSchema, + ); + + expect(result.prompts).toHaveLength(1); + expect(result.prompts[0].name).toBe("test"); + expect(result.prompts[0].arguments).toEqual([ + { name: "name", required: true }, + { name: "value", required: true }, + ]); + }); + + test("should register prompt with description", async () => { + const server = new Server({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + server.prompt("test", "Test description", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response", + }, + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "prompts/list", + }, + ListPromptsResultSchema, + ); + + expect(result.prompts).toHaveLength(1); + expect(result.prompts[0].name).toBe("test"); + expect(result.prompts[0].description).toBe("Test description"); + }); + + test("should validate prompt args", async () => { + const server = new Server({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + prompts: {}, + }, + }, + ); + + server.prompt( + "test", + { + name: z.string(), + value: z.string().min(3), + }, + async ({ name, value }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `${name}: ${value}`, + }, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + await expect( + client.request( + { + method: "prompts/get", + params: { + name: "test", + arguments: { + name: "test", + value: "ab", // Too short + }, + }, + }, + GetPromptResultSchema, + ), + ).rejects.toThrow(/Invalid arguments/); + }); + + test("should prevent duplicate prompt registration", () => { + const server = new Server({ + name: "test server", + version: "1.0", + }); + + server.prompt("test", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response", + }, + }, + ], + })); + + expect(() => { + server.prompt("test", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response 2", + }, + }, + ], + })); + }).toThrow(/already registered/); + }); + + test("should throw McpError for invalid prompt name", async () => { + const server = new Server({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + prompts: {}, + }, + }, + ); + + server.prompt("test-prompt", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response", + }, + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + await expect( + client.request( + { + method: "prompts/get", + params: { + name: "nonexistent-prompt", + }, + }, + GetPromptResultSchema, + ), + ).rejects.toThrow(/Prompt nonexistent-prompt not found/); + }); +}); diff --git a/src/server/index.ts b/src/server/index.ts index bc5ea61d..c1efdd85 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -1,4 +1,11 @@ -import z, { AnyZodObject, ZodRawShape, ZodTypeAny } from "zod"; +import z, { + AnyZodObject, + ZodObject, + ZodOptional, + ZodRawShape, + ZodString, + ZodTypeAny, +} from "zod"; import { zodToJsonSchema } from "zod-to-json-schema"; import { mergeCapabilities, @@ -15,12 +22,16 @@ import { CreateMessageResultSchema, EmptyResultSchema, ErrorCode, + GetPromptRequestSchema, + GetPromptResult, Implementation, InitializedNotificationSchema, InitializeRequest, InitializeRequestSchema, InitializeResult, LATEST_PROTOCOL_VERSION, + ListPromptsRequestSchema, + ListPromptsResult, ListResourcesRequestSchema, ListResourcesResult, ListResourceTemplatesRequestSchema, @@ -31,6 +42,8 @@ import { LoggingMessageNotification, McpError, Notification, + Prompt, + PromptArgument, ReadResourceRequestSchema, ReadResourceResult, Request, @@ -90,11 +103,12 @@ export class Server< private _clientCapabilities?: ClientCapabilities; private _clientVersion?: Implementation; private _capabilities: ServerCapabilities; - private _registeredTools: { [name: string]: RegisteredTool } = {}; private _registeredResources: { [uri: string]: RegisteredResource } = {}; private _registeredResourceTemplates: { [name: string]: RegisteredResourceTemplate; } = {}; + private _registeredTools: { [name: string]: RegisteredTool } = {}; + private _registeredPrompts: { [name: string]: RegisteredPrompt } = {}; /** * Callback for when initialization has fully completed (i.e., the client has sent an `initialized` notification). @@ -683,6 +697,106 @@ export class Server< this.setResourceRequestHandlers(); } + + private setPromptRequestHandlers() { + this.assertCanSetRequestHandler( + ListPromptsRequestSchema.shape.method.value, + ); + this.assertCanSetRequestHandler(GetPromptRequestSchema.shape.method.value); + + this.registerCapabilities({ + prompts: {}, + }); + + this.setRequestHandler( + ListPromptsRequestSchema, + (): ListPromptsResult => ({ + prompts: Object.entries(this._registeredPrompts).map( + ([name, prompt]): Prompt => { + return { + name, + description: prompt.description, + arguments: prompt.argsSchema + ? promptArgumentsFromSchema(prompt.argsSchema) + : undefined, + }; + }, + ), + }), + ); + + this.setRequestHandler( + GetPromptRequestSchema, + async (request, extra): Promise => { + const prompt = this._registeredPrompts[request.params.name]; + if (!prompt) { + throw new McpError( + ErrorCode.InvalidParams, + `Prompt ${request.params.name} not found`, + ); + } + + if (prompt.argsSchema) { + const parseResult = await prompt.argsSchema.safeParseAsync( + request.params.arguments, + ); + if (!parseResult.success) { + throw new McpError( + ErrorCode.InvalidParams, + `Invalid arguments for prompt ${request.params.name}: ${parseResult.error.message}`, + ); + } + + const args = parseResult.data; + const cb = prompt.callback as PromptCallback; + return await Promise.resolve(cb(args, extra)); + } else { + const cb = prompt.callback as PromptCallback; + return await Promise.resolve(cb(extra)); + } + }, + ); + } + + prompt(name: string, cb: PromptCallback): void; + prompt(name: string, description: string, cb: PromptCallback): void; + prompt( + name: string, + argsSchema: Args, + cb: PromptCallback, + ): void; + + prompt( + name: string, + description: string, + argsSchema: Args, + cb: PromptCallback, + ): void; + + prompt(name: string, ...rest: unknown[]): void { + if (this._registeredPrompts[name]) { + throw new Error(`Prompt ${name} is already registered`); + } + + let description: string | undefined; + if (typeof rest[0] === "string") { + description = rest.shift() as string; + } + + let argsSchema: PromptArgsRawShape | undefined; + if (rest.length > 1) { + argsSchema = rest.shift() as PromptArgsRawShape; + } + + const cb = rest[0] as PromptCallback; + this._registeredPrompts[name] = { + description, + argsSchema: argsSchema === undefined ? undefined : z.object(argsSchema), + callback: cb, + }; + + this.setPromptRequestHandlers(); + } } /** @@ -780,3 +894,34 @@ type RegisteredResourceTemplate = { metadata?: ResourceMetadata; readCallback: ReadResourceTemplateCallback; }; + +type PromptArgsRawShape = { + [k: string]: ZodString | ZodOptional; +}; + +export type PromptCallback< + Args extends undefined | PromptArgsRawShape = undefined, +> = Args extends PromptArgsRawShape + ? ( + args: z.objectOutputType, + extra: RequestHandlerExtra, + ) => GetPromptResult | Promise + : (extra: RequestHandlerExtra) => GetPromptResult | Promise; + +type RegisteredPrompt = { + description?: string; + argsSchema?: ZodObject; + callback: PromptCallback; +}; + +function promptArgumentsFromSchema( + schema: ZodObject, +): PromptArgument[] { + return Object.entries(schema.shape).map( + ([name, field]): PromptArgument => ({ + name, + description: field.description, + required: !field.isOptional(), + }), + ); +}