diff --git a/src/server/completable.test.ts b/src/server/completable.test.ts new file mode 100644 index 00000000..6040ff3f --- /dev/null +++ b/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); + }); +}); diff --git a/src/server/completable.ts b/src/server/completable.ts new file mode 100644 index 00000000..3b5bc164 --- /dev/null +++ b/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 = ( + value: T["_input"], +) => T["_input"][] | Promise; + +export interface CompletableDef + extends ZodTypeDef { + type: T; + complete: CompleteCallback; + typeName: McpZodTypeKind.Completable; +} + +export class Completable extends ZodType< + T["_output"], + CompletableDef, + T["_input"] +> { + _parse(input: ParseInput): ParseReturnType { + 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 = ( + type: T, + params: RawCreateParams & { + complete: CompleteCallback; + }, + ): Completable => { + 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( + schema: T, + complete: CompleteCallback, +): Completable { + 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 }; +}