From 189cc84113ec5ce9cb132a8d24f808ec4be3fa92 Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Wed, 8 Jan 2025 15:46:43 +0000 Subject: [PATCH] Autocomplete on resource template variables --- src/server/index.test.ts | 183 +++++++++++++++++++++++++++++++++----- src/server/index.ts | 185 ++++++++++++++++++++++++++++----------- 2 files changed, 294 insertions(+), 74 deletions(-) diff --git a/src/server/index.test.ts b/src/server/index.test.ts index cb686d68..2957b872 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -560,32 +560,34 @@ test("should handle request timeout", async () => { describe("ResourceTemplate", () => { test("should create ResourceTemplate with string pattern", () => { - const template = new ResourceTemplate("test://{category}/{id}", undefined); + const template = new ResourceTemplate("test://{category}/{id}", { + list: undefined, + }); expect(template.uriTemplate.toString()).toBe("test://{category}/{id}"); expect(template.listCallback).toBeUndefined(); }); test("should create ResourceTemplate with UriTemplate", () => { const uriTemplate = new UriTemplate("test://{category}/{id}"); - const template = new ResourceTemplate(uriTemplate, undefined); + const template = new ResourceTemplate(uriTemplate, { list: undefined }); expect(template.uriTemplate).toBe(uriTemplate); expect(template.listCallback).toBeUndefined(); }); test("should create ResourceTemplate with list callback", async () => { - const listCallback = jest.fn().mockResolvedValue({ + const list = jest.fn().mockResolvedValue({ resources: [{ name: "Test", uri: "test://example" }], }); - const template = new ResourceTemplate("test://{id}", listCallback); - expect(template.listCallback).toBe(listCallback); + const template = new ResourceTemplate("test://{id}", { list }); + expect(template.listCallback).toBe(list); const abortController = new AbortController(); const result = await template.listCallback?.({ signal: abortController.signal, }); expect(result?.resources).toHaveLength(1); - expect(listCallback).toHaveBeenCalled(); + expect(list).toHaveBeenCalled(); }); }); @@ -1068,7 +1070,7 @@ describe("Server.resource", () => { server.resource( "test", - new ResourceTemplate("test://resource/{id}", undefined), + new ResourceTemplate("test://resource/{id}", { list: undefined }), async () => ({ contents: [ { @@ -1113,18 +1115,20 @@ describe("Server.resource", () => { server.resource( "test", - new ResourceTemplate("test://resource/{id}", async () => ({ - resources: [ - { - name: "Resource 1", - uri: "test://resource/1", - }, - { - name: "Resource 2", - uri: "test://resource/2", - }, - ], - })), + new ResourceTemplate("test://resource/{id}", { + list: async () => ({ + resources: [ + { + name: "Resource 1", + uri: "test://resource/1", + }, + { + name: "Resource 2", + uri: "test://resource/2", + }, + ], + }), + }), async (uri) => ({ contents: [ { @@ -1169,7 +1173,9 @@ describe("Server.resource", () => { server.resource( "test", - new ResourceTemplate("test://resource/{category}/{id}", undefined), + new ResourceTemplate("test://resource/{category}/{id}", { + list: undefined, + }), async (uri, { category, id }) => ({ contents: [ { @@ -1236,7 +1242,7 @@ describe("Server.resource", () => { server.resource( "test", - new ResourceTemplate("test://resource/{id}", undefined), + new ResourceTemplate("test://resource/{id}", { list: undefined }), async () => ({ contents: [ { @@ -1250,7 +1256,7 @@ describe("Server.resource", () => { expect(() => { server.resource( "test", - new ResourceTemplate("test://resource/{id}", undefined), + new ResourceTemplate("test://resource/{id}", { list: undefined }), async () => ({ contents: [ { @@ -1337,6 +1343,139 @@ describe("Server.resource", () => { ), ).rejects.toThrow(/Resource test:\/\/nonexistent not found/); }); + + test("should support completion of resource template parameters", async () => { + const server = new Server({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + resources: {}, + }, + }, + ); + + server.resource( + "test", + new ResourceTemplate("test://resource/{category}", { + list: undefined, + complete: { + category: () => ["books", "movies", "music"], + }, + }), + async () => ({ + contents: [ + { + uri: "test://resource/test", + text: "Test content", + }, + ], + }), + ); + + 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/resource", + uri: "test://resource/{category}", + }, + argument: { + name: "category", + value: "", + }, + }, + }, + CompleteResultSchema, + ); + + expect(result.completion.values).toEqual(["books", "movies", "music"]); + expect(result.completion.total).toBe(3); + }); + + test("should support filtered completion of resource template parameters", async () => { + const server = new Server({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + resources: {}, + }, + }, + ); + + server.resource( + "test", + new ResourceTemplate("test://resource/{category}", { + list: undefined, + complete: { + category: (test) => + ["books", "movies", "music"].filter((value) => + value.startsWith(test), + ), + }, + }), + async () => ({ + contents: [ + { + uri: "test://resource/test", + text: "Test content", + }, + ], + }), + ); + + 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/resource", + uri: "test://resource/{category}", + }, + argument: { + name: "category", + value: "m", + }, + }, + }, + CompleteResultSchema, + ); + + expect(result.completion.values).toEqual(["movies", "music"]); + expect(result.completion.total).toBe(2); + }); }); describe("Server.prompt", () => { diff --git a/src/server/index.ts b/src/server/index.ts index 2816ee8f..1810883a 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -20,6 +20,7 @@ import { CallToolRequestSchema, CallToolResult, ClientCapabilities, + CompleteRequest, CompleteRequestSchema, CompleteResult, CreateMessageRequest, @@ -48,10 +49,12 @@ import { Notification, Prompt, PromptArgument, + PromptReference, ReadResourceRequestSchema, ReadResourceResult, Request, Resource, + ResourceReference, ResourceUpdatedNotification, Result, ServerCapabilities, @@ -500,6 +503,86 @@ export class Server< this.setToolRequestHandlers(); } + private setCompletionRequestHandler() { + this.assertCanSetRequestHandler(CompleteRequestSchema.shape.method.value); + + this.setRequestHandler( + CompleteRequestSchema, + async (request): Promise => { + switch (request.params.ref.type) { + case "ref/prompt": + return this.handlePromptCompletion(request, request.params.ref); + + case "ref/resource": + return this.handleResourceCompletion(request, request.params.ref); + + default: + throw new McpError( + ErrorCode.InvalidParams, + `Invalid completion reference: ${request.params.ref}`, + ); + } + }, + ); + } + + private async handlePromptCompletion( + request: CompleteRequest, + ref: PromptReference, + ): Promise { + const prompt = this._registeredPrompts[ref.name]; + if (!prompt) { + throw new McpError( + ErrorCode.InvalidParams, + `Prompt ${request.params.ref.name} not found`, + ); + } + + if (!prompt.argsSchema) { + return EMPTY_COMPLETION_RESULT; + } + + const field = prompt.argsSchema.shape[request.params.argument.name]; + if (!(field instanceof Completable)) { + return EMPTY_COMPLETION_RESULT; + } + + const def: CompletableDef = field._def; + const suggestions = await def.complete(request.params.argument.value); + return createCompletionResult(suggestions); + } + + private async handleResourceCompletion( + request: CompleteRequest, + ref: ResourceReference, + ): Promise { + const template = Object.values(this._registeredResourceTemplates).find( + (t) => t.resourceTemplate.uriTemplate.toString() === ref.uri, + ); + + if (!template) { + if (this._registeredResources[ref.uri]) { + // Attempting to autocomplete a fixed resource URI is not an error in the spec (but probably should be). + return EMPTY_COMPLETION_RESULT; + } + + throw new McpError( + ErrorCode.InvalidParams, + `Resource template ${request.params.ref.uri} not found`, + ); + } + + const completer = template.resourceTemplate.completeCallback( + request.params.argument.name, + ); + if (!completer) { + return EMPTY_COMPLETION_RESULT; + } + + const suggestions = await completer(request.params.argument.value); + return createCompletionResult(suggestions); + } + private setResourceRequestHandlers() { this.assertCanSetRequestHandler( ListResourcesRequestSchema.shape.method.value, @@ -588,6 +671,8 @@ export class Server< ); }, ); + + this.setCompletionRequestHandler(); } /** @@ -708,7 +793,6 @@ export class Server< ListPromptsRequestSchema.shape.method.value, ); this.assertCanSetRequestHandler(GetPromptRequestSchema.shape.method.value); - this.assertCanSetRequestHandler(CompleteRequestSchema.shape.method.value); this.registerCapabilities({ prompts: {}, @@ -763,55 +847,7 @@ 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, - }, - }; - }, - ); + this.setCompletionRequestHandler(); } prompt(name: string, cb: PromptCallback): void; @@ -854,6 +890,13 @@ export class Server< } } +/** + * A callback to complete one variable within a resource template's URI template. + */ +export type CompleteResourceTemplateCallback = ( + value: string, +) => string[] | Promise; + /** * A resource template combines a URI pattern with optional functionality to enumerate * all resources matching that pattern. @@ -863,7 +906,19 @@ export class ResourceTemplate { constructor( uriTemplate: string | UriTemplate, - private _listCallback: ListResourcesCallback | undefined, + private _callbacks: { + /** + * A callback to list all resources matching this template. This is required to specified, even if `undefined`, to avoid accidentally forgetting resource listing. + */ + list: ListResourcesCallback | undefined; + + /** + * An optional callback to autocomplete variables within the URI template. Useful for clients and users to discover possible values. + */ + complete?: { + [variable: string]: CompleteResourceTemplateCallback; + }; + }, ) { this._uriTemplate = typeof uriTemplate === "string" @@ -882,7 +937,16 @@ export class ResourceTemplate { * Gets the list callback, if one was provided. */ get listCallback(): ListResourcesCallback | undefined { - return this._listCallback; + return this._callbacks.list; + } + + /** + * Gets the callback for completing a specific URI template variable, if one was provided. + */ + completeCallback( + variable: string, + ): CompleteResourceTemplateCallback | undefined { + return this._callbacks.complete?.[variable]; } } @@ -982,3 +1046,20 @@ function promptArgumentsFromSchema( }), ); } + +function createCompletionResult(suggestions: string[]): CompleteResult { + return { + completion: { + values: suggestions.slice(0, 100), + total: suggestions.length, + hasMore: suggestions.length > 100, + }, + }; +} + +const EMPTY_COMPLETION_RESULT: CompleteResult = { + completion: { + values: [], + hasMore: false, + }, +};