Speed up inbound WebSocket validation (#1895)

* perf(protocol): generate inbound ws validators * perf(client): use generated ws validation * docs(protocol): document generated validation * fix(protocol): make validator generation source-only * test(protocol): cover explicit provider model normalization * fix(protocol): preserve inbound compat defaults * fix(protocol): make validator import rewrite portable * fix(protocol): harden validation safety checks * fix(protocol): guard generated validator boundaries * fix(protocol): normalize legacy inbound defaults * fix(protocol): simplify generated validation safety net * fix(protocol): keep cold typecheck source-only * fix(protocol): encapsulate validator codegen * fix(protocol): bootstrap source-alias typechecks * fix(app): keep source aliases out of bundler config * fix(protocol): harden validator codegen packaging * fix(protocol): own zod-aot patches in generator * fix(protocol): trim validator safety net * fix(client): normalize provider updates before dispatch * fix(protocol): keep validator generation out of install
2026-07-06 00:07:54 +08:00
parent 0ecf2fb8da
commit 9968f2ba50
25 changed files with 642 additions and 112 deletions
@@ -0,0 +1,106 @@
+import { mkdir, readFile, writeFile } from "node:fs/promises";
+import { createRequire } from "node:module";
+import { dirname, relative, resolve, sep } from "node:path";
+import { fileURLToPath, pathToFileURL } from "node:url";
+
+const packageRoot = resolve(dirname(fileURLToPath(import.meta.url)), "..");
+const source = resolve(packageRoot, "codegen/ws-outbound.compile.ts");
+const runtimeSchemaMetadata = resolve(packageRoot, "src/validation/ws-outbound-schema-metadata.ts");
+const output = resolve(packageRoot, "src/generated/validation/ws-outbound.aot.ts");
+
+const require = createRequire(import.meta.url);
+const zodAotEntry = require.resolve("zod-aot");
+const zodAotRoot = resolve(dirname(zodAotEntry), "..");
+const emitterPath = resolve(zodAotRoot, "dist/cli/emitter.js");
+const discriminatedUnionPath = resolve(
+  zodAotRoot,
+  "dist/core/codegen/schemas/discriminated-union.js",
+);
+
+async function ensureZodAotRuntimeImportExtensionPatch() {
+  const emitter = await readFile(emitterPath, "utf8");
+  if (emitter.includes('sourceRelPath.endsWith(".js")')) {
+    return;
+  }
+
+  const before = 'let importPath = sourceRelPath.replace(/\\.[cm]?[jt]sx?$/, "");';
+  const after =
+    'let importPath = sourceRelPath.endsWith(".js")\n        ? sourceRelPath\n        : sourceRelPath.replace(/\\.[cm]?[jt]sx?$/, "");';
+  if (!emitter.includes(before)) {
+    throw new Error("zod-aot emitter shape changed; update the runtime import extension patch");
+  }
+  await writeFile(emitterPath, emitter.replace(before, after));
+}
+
+async function ensureZodAotDiscriminatedUnionOutputPatch() {
+  let discriminatedUnionEmitter = await readFile(discriminatedUnionPath, "utf8");
+  if (
+    discriminatedUnionEmitter.includes(
+      "const needsOutputPropagation = ir.options.some(hasMutation);",
+    )
+  ) {
+    return;
+  }
+
+  const importBefore = 'import { escapeString } from "../context.js";';
+  const importAfter = 'import { escapeString, hasMutation } from "../context.js";';
+  const outputFlagBefore = "const discKey = escapeString(ir.discriminator);\n    let code = emit `";
+  const outputFlagAfter =
+    "const discKey = escapeString(ir.discriminator);\n    const needsOutputPropagation = ir.options.some(hasMutation);\n    let code = emit `";
+  const propagationBefore =
+    "        ${g.visit(option, { input: objVar, output: objVar })}\n        break;`;";
+  const propagationAfter =
+    '        ${g.visit(option, { input: objVar, output: objVar })}\n        ${needsOutputPropagation ? `${g.output}=${objVar};` : ""}\n        break;`;';
+
+  if (
+    !discriminatedUnionEmitter.includes(importBefore) ||
+    !discriminatedUnionEmitter.includes(outputFlagBefore) ||
+    !discriminatedUnionEmitter.includes(propagationBefore)
+  ) {
+    throw new Error("zod-aot discriminated-union emitter shape changed; update the output patch");
+  }
+
+  discriminatedUnionEmitter = discriminatedUnionEmitter
+    .replace(importBefore, importAfter)
+    .replace(outputFlagBefore, outputFlagAfter)
+    .replace(propagationBefore, propagationAfter);
+  await writeFile(discriminatedUnionPath, discriminatedUnionEmitter);
+}
+
+await Promise.all([
+  ensureZodAotRuntimeImportExtensionPatch(),
+  ensureZodAotDiscriminatedUnionOutputPatch(),
+]);
+
+const [{ discoverSchemas }, { compileSchemas }, { generateCompiledFileContent }] =
+  await Promise.all([
+    import(pathToFileURL(resolve(zodAotRoot, "dist/discovery.js")).href),
+    import(pathToFileURL(resolve(zodAotRoot, "dist/core/pipeline.js")).href),
+    import(pathToFileURL(resolve(zodAotRoot, "dist/cli/emitter.js")).href),
+  ]);
+
+const schemas = await discoverSchemas(source, { cacheBust: true });
+if (schemas.length === 0) {
+  throw new Error(`No zod-aot compile() exports found in ${relative(packageRoot, source)}`);
+}
+
+const compiled = compileSchemas(schemas, { mode: "inline" });
+const runtimeImportPath = relative(dirname(output), runtimeSchemaMetadata)
+  .replace(/\.[cm]?[jt]sx?$/, ".js")
+  .split(sep)
+  .join("/");
+const content = generateCompiledFileContent(compiled, runtimeImportPath, {
+  zodCompat: false,
+}).replace(
+  "// AUTO-GENERATED by zod-aot — DO NOT EDIT",
+  "// @ts-nocheck\n// AUTO-GENERATED by zod-aot — DO NOT EDIT",
+);
+
+await mkdir(dirname(output), { recursive: true });
+await writeFile(output, content);
+
+console.info(
+  `generated ${relative(packageRoot, output)} from ${relative(packageRoot, source)} (${schemas
+    .map((schema) => schema.exportName)
+    .join(", ")})`,
+);
@@ -0,0 +1,96 @@
+import { readdir, stat } from "node:fs/promises";
+import { dirname, join, relative, sep } from "node:path";
+import { fileURLToPath } from "node:url";
+import { spawn } from "node:child_process";
+
+const packageRoot = dirname(dirname(fileURLToPath(import.meta.url)));
+const protocolSrcDir = join(packageRoot, "src");
+const pollIntervalMs = 1000;
+
+let currentFingerprint = "";
+let generateInFlight = false;
+let generateAgain = false;
+
+async function collectSourceFiles(dir) {
+  const entries = await readdir(dir, { withFileTypes: true });
+  const files = [];
+
+  for (const entry of entries) {
+    const path = join(dir, entry.name);
+    const relativePath = relative(protocolSrcDir, path);
+    const parts = relativePath.split(sep);
+
+    if (parts[0] === "generated") {
+      continue;
+    }
+
+    if (entry.isDirectory()) {
+      files.push(...(await collectSourceFiles(path)));
+      continue;
+    }
+
+    if (entry.isFile() && entry.name.endsWith(".ts")) {
+      files.push(path);
+    }
+  }
+
+  return files;
+}
+
+async function fingerprintSources() {
+  const files = await collectSourceFiles(protocolSrcDir);
+  const stats = await Promise.all(
+    files.sort().map(async (file) => {
+      const fileStat = await stat(file);
+      return `${relative(packageRoot, file)}:${fileStat.mtimeMs}:${fileStat.size}`;
+    }),
+  );
+
+  return stats.join("\n");
+}
+
+function runGenerate() {
+  if (generateInFlight) {
+    generateAgain = true;
+    return;
+  }
+
+  generateInFlight = true;
+  const npmCommand = process.platform === "win32" ? "npm.cmd" : "npm";
+  const child = spawn(npmCommand, ["run", "generate:validators"], {
+    cwd: packageRoot,
+    stdio: "inherit",
+  });
+
+  child.on("exit", (code, signal) => {
+    generateInFlight = false;
+    if (code !== 0) {
+      console.error(
+        signal
+          ? `validation AOT generation stopped by ${signal}`
+          : `validation AOT generation failed with exit code ${code}`,
+      );
+    }
+
+    if (generateAgain) {
+      generateAgain = false;
+      runGenerate();
+    }
+  });
+}
+
+async function poll() {
+  try {
+    const nextFingerprint = await fingerprintSources();
+    if (currentFingerprint && nextFingerprint !== currentFingerprint) {
+      runGenerate();
+    }
+    currentFingerprint = nextFingerprint;
+  } catch (error) {
+    console.error("validation AOT watcher failed to scan protocol sources", error);
+  }
+}
+
+currentFingerprint = await fingerprintSources();
+console.log("Watching protocol schema sources for validation AOT changes...");
+setInterval(poll, pollIntervalMs);