Simplified, Express-like API

package.json

@@ -48,7 +48,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",

src/client/index.ts

@@ -1,4 +1,5 @@
 import {
+  mergeCapabilities,
   Protocol,
   ProtocolOptions,
   RequestOptions,
@@ -44,7 +45,7 @@ export type ClientOptions = ProtocolOptions & {
   /**
    * Capabilities to advertise as being supported by this client.
    */
-  capabilities: ClientCapabilities;
+  capabilities?: ClientCapabilities;
 };
 
 /**
@@ -90,10 +91,25 @@ export class Client<
    */
   constructor(
     private _clientInfo: Implementation,
-    options: ClientOptions,
+    options?: ClientOptions,
   ) {
     super(options);
-    this._capabilities = options.capabilities;
+    this._capabilities = options?.capabilities ?? {};
+  }
+
+  /**
+   * Registers new capabilities. This can only be called before connecting to a transport.
+   *
+   * The new capabilities will be merged with any existing capabilities previously given (e.g., at initialization).
+   */
+  public registerCapabilities(capabilities: ClientCapabilities): void {
+    if (this.transport) {
+      throw new Error(
+        "Cannot register capabilities after connecting to transport",
+      );
+    }
+
+    this._capabilities = mergeCapabilities(this._capabilities, capabilities);
   }
 
   protected assertCapability(

src/server/completable.test.ts

@@ -0,0 +1,46 @@
+import { z } from "zod";
+import { completable } from "./completable.js";
+
+describe("completable", () => {
+  it("preserves types and values of underlying schema", () => {
+    const baseSchema = z.string();
+    const schema = completable(baseSchema, () => []);
+
+    expect(schema.parse("test")).toBe("test");
+    expect(() => schema.parse(123)).toThrow();
+  });
+
+  it("provides access to completion function", async () => {
+    const completions = ["foo", "bar", "baz"];
+    const schema = completable(z.string(), () => completions);
+
+    expect(await schema._def.complete("")).toEqual(completions);
+  });
+
+  it("allows async completion functions", async () => {
+    const completions = ["foo", "bar", "baz"];
+    const schema = completable(z.string(), async () => completions);
+
+    expect(await schema._def.complete("")).toEqual(completions);
+  });
+
+  it("passes current value to completion function", async () => {
+    const schema = completable(z.string(), (value) => [value + "!"]);
+
+    expect(await schema._def.complete("test")).toEqual(["test!"]);
+  });
+
+  it("works with number schemas", async () => {
+    const schema = completable(z.number(), () => [1, 2, 3]);
+
+    expect(schema.parse(1)).toBe(1);
+    expect(await schema._def.complete(0)).toEqual([1, 2, 3]);
+  });
+
+  it("preserves schema description", () => {
+    const desc = "test description";
+    const schema = completable(z.string().describe(desc), () => []);
+
+    expect(schema.description).toBe(desc);
+  });
+});

src/server/completable.ts

@@ -0,0 +1,95 @@
+import {
+  ZodTypeAny,
+  ZodTypeDef,
+  ZodType,
+  ParseInput,
+  ParseReturnType,
+  RawCreateParams,
+  ZodErrorMap,
+  ProcessedCreateParams,
+} from "zod";
+
+export enum McpZodTypeKind {
+  Completable = "McpCompletable",
+}
+
+export type CompleteCallback<T extends ZodTypeAny = ZodTypeAny> = (
+  value: T["_input"],
+) => T["_input"][] | Promise<T["_input"][]>;
+
+export interface CompletableDef<T extends ZodTypeAny = ZodTypeAny>
+  extends ZodTypeDef {
+  type: T;
+  complete: CompleteCallback<T>;
+  typeName: McpZodTypeKind.Completable;
+}
+
+export class Completable<T extends ZodTypeAny> extends ZodType<
+  T["_output"],
+  CompletableDef<T>,
+  T["_input"]
+> {
+  _parse(input: ParseInput): ParseReturnType<this["_output"]> {
+    const { ctx } = this._processInputParams(input);
+    const data = ctx.data;
+    return this._def.type._parse({
+      data,
+      path: ctx.path,
+      parent: ctx,
+    });
+  }
+
+  unwrap() {
+    return this._def.type;
+  }
+
+  static create = <T extends ZodTypeAny>(
+    type: T,
+    params: RawCreateParams & {
+      complete: CompleteCallback<T>;
+    },
+  ): Completable<T> => {
+    return new Completable({
+      type,
+      typeName: McpZodTypeKind.Completable,
+      complete: params.complete,
+      ...processCreateParams(params),
+    });
+  };
+}
+
+/**
+ * Wraps a Zod type to provide autocompletion capabilities. Useful for, e.g., prompt arguments in MCP.
+ */
+export function completable<T extends ZodTypeAny>(
+  schema: T,
+  complete: CompleteCallback<T>,
+): Completable<T> {
+  return Completable.create(schema, { ...schema._def, complete });
+}
+
+// Not sure why this isn't exported from Zod:
+// https://github.com/colinhacks/zod/blob/f7ad26147ba291cb3fb257545972a8e00e767470/src/types.ts#L130
+function processCreateParams(params: RawCreateParams): ProcessedCreateParams {
+  if (!params) return {};
+  const { errorMap, invalid_type_error, required_error, description } = params;
+  if (errorMap && (invalid_type_error || required_error)) {
+    throw new Error(
+      `Can't use "invalid_type_error" or "required_error" in conjunction with custom error map.`,
+    );
+  }
+  if (errorMap) return { errorMap: errorMap, description };
+  const customMap: ZodErrorMap = (iss, ctx) => {
+    const { message } = params;
+
+    if (iss.code === "invalid_enum_value") {
+      return { message: message ?? ctx.defaultError };
+    }
+    if (typeof ctx.data === "undefined") {
+      return { message: message ?? required_error ?? ctx.defaultError };
+    }
+    if (iss.code !== "invalid_type") return { message: ctx.defaultError };
+    return { message: message ?? invalid_type_error ?? ctx.defaultError };
+  };
+  return { errorMap: customMap, description };
+}

src/server/index.test.ts

@@ -478,6 +478,7 @@ test("should handle server cancelling a request", async () => {
   // Request should be rejected
   await expect(createMessagePromise).rejects.toBe("Cancelled by test");
 });
+
 test("should handle request timeout", async () => {
   const server = new Server(
     {

src/server/index.ts

@@ -1,4 +1,5 @@
 import {
+  mergeCapabilities,
   Protocol,
   ProtocolOptions,
   RequestOptions,
@@ -32,7 +33,7 @@ export type ServerOptions = ProtocolOptions & {
   /**
    * Capabilities to advertise as being supported by this server.
    */
-  capabilities: ServerCapabilities;
+  capabilities?: ServerCapabilities;
 
   /**
    * Optional instructions describing how to use the server and its features.
@@ -89,11 +90,11 @@ export class Server<
    */
   constructor(
     private _serverInfo: Implementation,
-    options: ServerOptions,
+    options?: ServerOptions,
   ) {
     super(options);
-    this._capabilities = options.capabilities;
-    this._instructions = options.instructions;
+    this._capabilities = options?.capabilities ?? {};
+    this._instructions = options?.instructions;
 
     this.setRequestHandler(InitializeRequestSchema, (request) =>
       this._oninitialize(request),
@@ -103,6 +104,21 @@ export class Server<
     );
   }
 
+  /**
+   * Registers new capabilities. This can only be called before connecting to a transport.
+   *
+   * The new capabilities will be merged with any existing capabilities previously given (e.g., at initialization).
+   */
+  public registerCapabilities(capabilities: ServerCapabilities): void {
+    if (this.transport) {
+      throw new Error(
+        "Cannot register capabilities after connecting to transport",
+      );
+    }
+
+    this._capabilities = mergeCapabilities(this._capabilities, capabilities);
+  }
+
   protected assertCapabilityForMethod(method: RequestT["method"]): void {
     switch (method as ServerRequest["method"]) {
       case "sampling/createMessage":