Add completable wrapper for Zod schemas

Author: jspahrsummers

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