diff --git a/src/server/index.test.ts b/src/server/index.test.ts index 89bb578b..cb686d68 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -22,12 +22,14 @@ import { ReadResourceResultSchema, ListPromptsResultSchema, GetPromptResultSchema, + CompleteResultSchema, } from "../types.js"; import { Transport } from "../shared/transport.js"; import { InMemoryTransport } from "../inMemory.js"; import { Client } from "../client/index.js"; import { UriTemplate } from "../shared/uriTemplate.js"; import { ResourceTemplate } from "./index.js"; +import { completable } from "./completable.js"; test("should accept latest protocol version", async () => { let sendPromiseResolve: (value: unknown) => void; @@ -1619,4 +1621,135 @@ describe("Server.prompt", () => { ), ).rejects.toThrow(/Prompt nonexistent-prompt not found/); }); + test("should support completion of prompt arguments", 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", + { + name: completable(z.string(), () => ["Alice", "Bob", "Charlie"]), + }, + async ({ name }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `Hello ${name}`, + }, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "completion/complete", + params: { + ref: { + type: "ref/prompt", + name: "test-prompt", + }, + argument: { + name: "name", + value: "", + }, + }, + }, + CompleteResultSchema, + ); + + expect(result.completion.values).toEqual(["Alice", "Bob", "Charlie"]); + expect(result.completion.total).toBe(3); + }); + + test("should support filtered completion of prompt arguments", 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", + { + name: completable(z.string(), (test) => + ["Alice", "Bob", "Charlie"].filter((value) => value.startsWith(test)), + ), + }, + async ({ name }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `Hello ${name}`, + }, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "completion/complete", + params: { + ref: { + type: "ref/prompt", + name: "test-prompt", + }, + argument: { + name: "name", + value: "A", + }, + }, + }, + CompleteResultSchema, + ); + + expect(result.completion.values).toEqual(["Alice"]); + expect(result.completion.total).toBe(1); + }); }); diff --git a/src/server/index.ts b/src/server/index.ts index c1efdd85..2816ee8f 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -4,7 +4,9 @@ import z, { ZodOptional, ZodRawShape, ZodString, + ZodType, ZodTypeAny, + ZodTypeDef, } from "zod"; import { zodToJsonSchema } from "zod-to-json-schema"; import { @@ -18,6 +20,8 @@ import { CallToolRequestSchema, CallToolResult, ClientCapabilities, + CompleteRequestSchema, + CompleteResult, CreateMessageRequest, CreateMessageResultSchema, EmptyResultSchema, @@ -58,6 +62,7 @@ import { Tool, } from "../types.js"; import { UriTemplate, Variables } from "../shared/uriTemplate.js"; +import { Completable, CompletableDef } from "./completable.js"; export type ServerOptions = ProtocolOptions & { /** @@ -703,6 +708,7 @@ export class Server< ListPromptsRequestSchema.shape.method.value, ); this.assertCanSetRequestHandler(GetPromptRequestSchema.shape.method.value); + this.assertCanSetRequestHandler(CompleteRequestSchema.shape.method.value); this.registerCapabilities({ prompts: {}, @@ -756,6 +762,56 @@ export class Server< } }, ); + + this.setRequestHandler( + CompleteRequestSchema, + async (request): Promise => { + if (request.params.ref.type !== "ref/prompt") { + throw new McpError( + ErrorCode.InvalidParams, + "Only prompt completions are supported", + ); + } + + const prompt = this._registeredPrompts[request.params.ref.name]; + if (!prompt) { + throw new McpError( + ErrorCode.InvalidParams, + `Prompt ${request.params.ref.name} not found`, + ); + } + + if (!prompt.argsSchema) { + return { + completion: { + values: [], + hasMore: false, + }, + }; + } + + const field = prompt.argsSchema.shape[request.params.argument.name]; + if (!(field instanceof Completable)) { + return { + completion: { + values: [], + hasMore: false, + }, + }; + } + + const def: CompletableDef = field._def; + const completer = def.complete; + const suggestions = await completer(request.params.argument.value); + return { + completion: { + values: suggestions.slice(0, 100), + total: suggestions.length, + hasMore: suggestions.length > 100, + }, + }; + }, + ); } prompt(name: string, cb: PromptCallback): void; @@ -765,7 +821,6 @@ export class Server< argsSchema: Args, cb: PromptCallback, ): void; - prompt( name: string, description: string, @@ -896,7 +951,9 @@ type RegisteredResourceTemplate = { }; type PromptArgsRawShape = { - [k: string]: ZodString | ZodOptional; + [k: string]: + | ZodType + | ZodOptional>; }; export type PromptCallback<