如何缩小泛型函数makeMethodRpcSchema的返回类型？

现有代码定义 #

通用RPC 2.0请求Schema #

import { z } from "zod";

const rpcRequestMessageSchema = z.object({
  jsonrpc: z.literal("2.0"),
  method: z.string(),
  params: z
    .union([z.array(z.unknown()), z.object({}).passthrough()])
    .optional(),
  id: z.union([z.string(), z.number(), z.null()]).optional(),
});

特定RPC方法的参数与结果Schema #

const m1name = "m1" as const;
const m1paramsschema = z.object({ m1foo: z.string() });
const m1resultschema = z.object({ m1bar: z.string() });

const m2name = "m2" as const;
const m2paramsschema = z.object({ m2foo: z.string() });
const m2resultschema = z.object({ m2bar: z.string() });

const allMethods = {
  [m1name]: { params: m1paramsschema, result: m1resultschema },
  [m2name]: { params: m2paramsschema, result: m2resultschema },
};

当前的makeMethodRpcSchema工具函数 #

function makeMethodRpcSchema<Method extends keyof typeof allMethods>(
  method: Method
) {
  return rpcRequestMessageSchema.extend({
    method: z.literal(method),
    params: allMethods[method].params,
  });
}

问题现象 #

调用makeMethodRpcSchema("m1")时，生成的类型会错误包含其他方法（如m2）的参数字段：

const rpcm1Schema = makeMethodRpcSchema("m1");
type RpcM1 = z.infer<typeof rpcm1Schema>;
const obj: RpcM1  = {
  params: {
    // 编辑器同时提示m1foo和m2foo，但只有m1foo是合法的
  }
}

解决方案 #

问题核心是TypeScript无法自动推断allMethods[method].params的精确关联类型，需要通过显式类型约束或断言，让Zod正确缩小返回值的类型范围。

优化后的工具函数 #

import { z, ZodObject } from "zod";

// 先定义allMethods的精确类型
type AllMethods = typeof allMethods;

function makeMethodRpcSchema<Method extends keyof AllMethods>(
  method: Method
): ZodObject<{
  jsonrpc: typeof rpcRequestMessageSchema.shape.jsonrpc;
  method: z.ZodLiteral<Method>;
  params: AllMethods[Method]["params"];
  id: typeof rpcRequestMessageSchema.shape.id;
}> {
  return rpcRequestMessageSchema.extend({
    method: z.literal(method),
    params: allMethods[method].params,
  }) as ZodObject<{
    jsonrpc: typeof rpcRequestMessageSchema.shape.jsonrpc;
    method: z.ZodLiteral<Method>;
    params: AllMethods[Method]["params"];
    id: typeof rpcRequestMessageSchema.shape.id;
  }>;
}

或者更简洁的泛型推导写法：

function makeMethodRpcSchema<Method extends keyof typeof allMethods>(
  method: Method
) {
  return rpcRequestMessageSchema.extend({
    method: z.literal(method),
    params: allMethods[method].params,
  }) as typeof rpcRequestMessageSchema extends z.ZodObject<infer T>
    ? z.ZodObject<Omit<T, "method" | "params"> & {
        method: z.ZodLiteral<Method>;
        params: typeof allMethods[Method]["params"];
      }>
    : never;
}

验证效果 #

现在调用makeMethodRpcSchema("m1")后，RpcM1类型的params字段只会提示m1foo，符合预期：

const rpcm1Schema = makeMethodRpcSchema("m1");
type RpcM1 = z.infer<typeof rpcm1Schema>;
const obj: RpcM1 = {
  jsonrpc: "2.0",
  method: "m1",
  params: {
    m1foo: "test" // 仅提示合法的m1foo字段
  }
};

原理说明 #

原函数中，extend()的返回类型被TypeScript宽泛推断为包含所有方法参数的联合类型。通过显式指定返回值的Zod对象结构，我们强制TypeScript将params字段与传入的method进行精确关联，从而实现类型的正确缩小。

内容的提问来源于stack exchange，提问作者aryzing