Autocomplete on prompt arguments

Author: jspahrsummers

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);
+  });
 });

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<CompleteResult> => {
+        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<ZodString> = 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<Args>,
   ): void;
-
   prompt<Args extends PromptArgsRawShape>(
     name: string,
     description: string,
@@ -896,7 +951,9 @@ type RegisteredResourceTemplate = {
 };
 
 type PromptArgsRawShape = {
-  [k: string]: ZodString | ZodOptional<ZodString>;
+  [k: string]:
+    | ZodType<string, ZodTypeDef, string>
+    | ZodOptional<ZodType<string, ZodTypeDef, string>>;
 };
 
 export type PromptCallback<