Spaces:
Sleeping
Sleeping
File size: 2,661 Bytes
c16e487 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | 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 };
}
|