Split the inbound WebSocket validator per message type (#6847)

* Split the inbound WebSocket validator per message type

zod-aot compiled WSOutboundMessageSchema into one 11 MB function. Every
inbound frame, heartbeats included, ran through it, so V8 hosts (daemon
clients, plugin hosts, the CLI, Electron) eventually Maglev-compiled the
whole function: about 800 MB of compiler memory and multi-second stalls,
again after each deoptimization.

Generate a tag-only envelope validator plus one validator per session
message, and dispatch on the message type. Issue paths, defaults and
passthrough behave as before.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Use forward slashes in generated validator imports

path.relative() returns backslash paths on Windows, which break the
import specifiers written into the generated files.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Declare the session export helper as a function

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

.gitignore

@@ -85,6 +85,8 @@ valknut-report.json/
 packages/server/src/server/fixtures/dictation/dictation-debug-largest.wav
 packages/server/src/server/fixtures/dictation/dictation-debug-largest.transcript.txt
 packages/protocol/src/generated/validation/*.aot.ts
+packages/protocol/src/generated/validation/ws-outbound-schema-metadata.ts
+packages/protocol/codegen/generated/
 
 /artifacts
 packages/desktop/.cache/

docs/protocol-validation.md

@@ -6,7 +6,9 @@ The reason is mobile performance. A captured 353 KB provider snapshot cost about
 
 ## Runtime Path
 
-`packages/protocol/src/validation/ws-outbound.ts` is the shipped boundary. It calls the generated `WSOutboundMessageSchema.safeParse` and returns the validated data. It does not normalize, repair, or re-validate the generated result.
+`packages/protocol/src/validation/ws-outbound.ts` is the shipped boundary. It validates the envelope with a generated validator that checks only each session message's `type` tag, then runs the generated validator for that one message type and returns the validated data. Issue paths from the message validator, including nested union branches, are prefixed with `message`, so callers see the same issues as when one validator covered the whole envelope. It does not normalize, repair, or re-validate the generated result.
+
+Keep one generated function per message type. zod-aot inlines every nested schema into the function it compiles, so compiling `WSOutboundMessageSchema` as a whole produced one 11 MB function. On V8 (daemon clients, plugin hosts, the CLI, and Electron), the Maglev tier compiled that function once it was hot, allocating about 800 MB and stalling for seconds, and compiled it again after each deoptimization. Hermes has no JIT, so mobile never showed it. `packages/protocol/tests/validation/ws-outbound.test.ts` fails if a generated function grows past 1 MB.
 
 Generated validators preserve unknown keys where Zod object parsing strips them. The client dispatch path uses known `type` and payload fields, so this passthrough behavior is accepted for inbound messages. The wire format is unchanged.
 
@@ -16,11 +18,10 @@ Provider model normalization is a parser-side compatibility shim in the client c
 
 The protocol package owns generation.
 
-- `packages/protocol/codegen/ws-outbound.compile.ts` is the build-time zod-aot discovery entry.
-- `packages/protocol/scripts/generate-validation-aot.mjs` runs the exact-pinned compiler and applies the small local compiler patches before generation.
+- `packages/protocol/src/validation/ws-outbound-envelope.ts` defines the tag-only envelope schema and exposes the session message schemas that get their own validators.
+- `packages/protocol/scripts/generate-validation-aot.mjs` writes the zod-aot discovery entry and the runtime schema metadata, runs the exact-pinned compiler, and applies the small local compiler patches before generation. It also emits the `type` → validator map the boundary dispatches through.
 - `packages/protocol/scripts/watch-validation-aot.mjs` reruns generation while editing protocol sources.
-- `packages/protocol/src/generated/validation/ws-outbound.aot.ts` is generated runtime code and is gitignored.
-- `packages/protocol/src/validation/ws-outbound-schema-metadata.ts` is runtime schema metadata for zod-aot fallback/default references.
+- `packages/protocol/codegen/generated/ws-outbound.compile.ts`, `packages/protocol/src/generated/validation/ws-outbound.aot.ts`, and `packages/protocol/src/generated/validation/ws-outbound-schema-metadata.ts` are generated and gitignored. The metadata module lets generated code reach schema fallback/default values.
 
 Generation runs from protocol-owned lifecycle hooks: `prebuild`, `pretypecheck`, `pretest`, and `watch`. Installs do not run generation: published packages consume protocol from prebuilt `dist`, and local build/typecheck/test flows generate the source file at the point it is actually needed.
 

packages/protocol/codegen/README.md

@@ -1,7 +1,7 @@
 # Protocol Validator Codegen
 
-This directory is build-time only. `ws-outbound.compile.ts` is the zod-aot discovery entry for the inbound WebSocket validator.
+This directory is build-time only. The generator writes `generated/ws-outbound.compile.ts`, the zod-aot discovery entry for the inbound WebSocket validators: one for the envelope and one per session message type.
 
-The generated runtime file is written to `../src/generated/validation/ws-outbound.aot.ts` and is not committed. The protocol package owns every generation trigger through its npm lifecycle scripts.
+The discovery entry and the generated runtime files under `../src/generated/validation/` are not committed. The protocol package owns every generation trigger through its npm lifecycle scripts.
 
 `zod-aot` is exact-pinned, and the protocol generator applies the small compiler patches it requires before generation. Treat changes to those patches like compiler changes: regenerate, inspect the output, and run the protocol validation regression tests before shipping.

packages/protocol/codegen/ws-outbound.compile.ts

@@ -1,4 +0,0 @@
-import { compile } from "zod-aot";
-import { WSOutboundMessageSchema as SourceWSOutboundMessageSchema } from "../src/messages.js";
-
-export const WSOutboundMessageSchema = compile(SourceWSOutboundMessageSchema);

packages/protocol/scripts/generate-validation-aot.mjs

@@ -4,8 +4,12 @@ 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 envelopeSource = resolve(packageRoot, "src/validation/ws-outbound-envelope.ts");
+const source = resolve(packageRoot, "codegen/generated/ws-outbound.compile.ts");
+const runtimeSchemaMetadata = resolve(
+  packageRoot,
+  "src/generated/validation/ws-outbound-schema-metadata.ts",
+);
 const output = resolve(packageRoot, "src/generated/validation/ws-outbound.aot.ts");
 
 const require = createRequire(import.meta.url);
@@ -72,16 +76,69 @@ await Promise.all([
   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 [
+  { loadSourceFile },
+  { discoverSchemas },
+  { compileSchemas },
+  { generateCompiledFileContent },
+] = await Promise.all([
+  import(pathToFileURL(resolve(zodAotRoot, "dist/loader.js")).href),
+  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),
+]);
+
+// One compiled validator per session message plus a tag-only envelope. zod-aot emits one function
+// per compiled schema, so this keeps every generated function small enough for V8 to optimize cheaply.
+const { sessionOutboundMessageOptions } = await loadSourceFile(envelopeSource, { cacheBust: true });
+const sessionTags = sessionOutboundMessageOptions.map((option) => option.shape.type.value);
+function sessionExport(index) {
+  return `SessionOutboundMessage${index}`;
+}
+const envelopeImport = relative(dirname(source), envelopeSource)
+  .replace(/\.ts$/, ".js")
+  .split(sep)
+  .join("/");
+const metadataImport = relative(dirname(runtimeSchemaMetadata), envelopeSource)
+  .replace(/\.ts$/, ".js")
+  .split(sep)
+  .join("/");
+
+await mkdir(dirname(source), { recursive: true });
+await writeFile(
+  source,
+  [
+    "// AUTO-GENERATED by generate-validation-aot.mjs — DO NOT EDIT",
+    'import { compile } from "zod-aot";',
+    `import { WSOutboundEnvelopeSchema, sessionOutboundMessageOptions } from "${envelopeImport}";`,
+    "export const WSOutboundEnvelope = compile(WSOutboundEnvelopeSchema);",
+    ...sessionTags.map(
+      (_, index) =>
+        `export const ${sessionExport(index)} = compile(sessionOutboundMessageOptions[${index}]);`,
+    ),
+    "",
+  ].join("\n"),
+);
+await mkdir(dirname(runtimeSchemaMetadata), { recursive: true });
+await writeFile(
+  runtimeSchemaMetadata,
+  [
+    "// AUTO-GENERATED by generate-validation-aot.mjs — DO NOT EDIT",
+    `import { WSOutboundEnvelopeSchema, sessionOutboundMessageOptions } from "${metadataImport}";`,
+    "export const WSOutboundEnvelope = { schema: WSOutboundEnvelopeSchema };",
+    ...sessionTags.map(
+      (_, index) =>
+        `export const ${sessionExport(index)} = { schema: sessionOutboundMessageOptions[${index}] };`,
+    ),
+    "",
+  ].join("\n"),
+);
 
 const schemas = await discoverSchemas(source, { cacheBust: true });
-if (schemas.length === 0) {
-  throw new Error(`No zod-aot compile() exports found in ${relative(packageRoot, source)}`);
+if (schemas.length !== sessionTags.length + 1) {
+  throw new Error(
+    `Expected ${sessionTags.length + 1} zod-aot compile() exports in ${relative(packageRoot, source)}, found ${schemas.length}`,
+  );
 }
 
 const compiled = compileSchemas(schemas, { mode: "inline" });
@@ -95,12 +152,16 @@ const content = generateCompiledFileContent(compiled, runtimeImportPath, {
   "// AUTO-GENERATED by zod-aot — DO NOT EDIT",
   "// @ts-nocheck\n// AUTO-GENERATED by zod-aot — DO NOT EDIT",
 );
+const sessionMessageValidators = [
+  "export const SessionOutboundMessageValidators = {",
+  ...sessionTags.map((tag, index) => `  ${JSON.stringify(tag)}: ${sessionExport(index)},`),
+  "};",
+  "",
+].join("\n");
 
 await mkdir(dirname(output), { recursive: true });
-await writeFile(output, content);
+await writeFile(output, `${content}\n${sessionMessageValidators}`);
 
 console.info(
-  `generated ${relative(packageRoot, output)} from ${relative(packageRoot, source)} (${schemas
-    .map((schema) => schema.exportName)
-    .join(", ")})`,
+  `generated ${relative(packageRoot, output)}: envelope + ${sessionTags.length} session message validators`,
 );

packages/protocol/src/generated/validation/README.md

@@ -1,6 +1,6 @@
 # Generated protocol validators
 
-`ws-outbound.aot.ts` is generated by zod-aot from `packages/protocol/codegen/ws-outbound.compile.ts`.
+`ws-outbound.aot.ts` is generated by zod-aot from `packages/protocol/codegen/generated/ws-outbound.compile.ts`, which the generator writes along with `ws-outbound-schema-metadata.ts`.
 
 Run:
 

packages/protocol/src/validation/ws-outbound-envelope.ts

@@ -0,0 +1,25 @@
+import { z } from "zod";
+import {
+  SessionOutboundMessageSchema,
+  WSHelloRejectedMessageSchema,
+  WSPongMessageSchema,
+} from "../messages.js";
+
+export const sessionOutboundMessageOptions = SessionOutboundMessageSchema.options;
+
+const [firstSessionTag, ...otherSessionTags] = sessionOutboundMessageOptions.map((option) =>
+  z.object({ type: option.shape.type }),
+);
+
+// WSOutboundMessageSchema with every session message reduced to its tag. It reports the same
+// envelope and discriminator issues, while each session message gets its own generated validator:
+// one function for the whole protocol is large enough that V8's Maglev tier spends hundreds of MB
+// and seconds compiling it.
+export const WSOutboundEnvelopeSchema = z.discriminatedUnion("type", [
+  WSPongMessageSchema,
+  z.object({
+    type: z.literal("session"),
+    message: z.discriminatedUnion("type", [firstSessionTag, ...otherSessionTags]),
+  }),
+  WSHelloRejectedMessageSchema,
+]);

packages/protocol/src/validation/ws-outbound-schema-metadata.ts

@@ -1,3 +0,0 @@
-import { WSOutboundMessageSchema as SourceWSOutboundMessageSchema } from "../messages.js";
-
-export const WSOutboundMessageSchema = { schema: SourceWSOutboundMessageSchema };

packages/protocol/src/validation/ws-outbound.ts

@@ -1,19 +1,54 @@
+import { ZodRealError } from "zod";
 import type { z } from "zod";
-import { WSOutboundMessageSchema } from "../generated/validation/ws-outbound.aot.js";
-import type { WSOutboundMessage } from "../messages.js";
+import {
+  SessionOutboundMessageValidators,
+  WSOutboundEnvelope,
+} from "../generated/validation/ws-outbound.aot.js";
+import type { SessionOutboundMessage, WSOutboundMessage } from "../messages.js";
 
-type WSOutboundValidationResult =
-  | { success: true; data: WSOutboundMessage }
-  | { success: false; error: z.ZodError };
+type ValidationResult<T> = { success: true; data: T } | { success: false; error: z.ZodError };
 
-interface WSOutboundGeneratedValidator {
-  safeParse(input: unknown): WSOutboundValidationResult;
+interface GeneratedValidator<T> {
+  safeParse(input: unknown): ValidationResult<T>;
 }
 
-// zod-aot emits runtime JavaScript from WSOutboundMessageSchema but not its TypeScript surface.
+// zod-aot emits runtime JavaScript from the envelope and message schemas but not its TypeScript surface.
 // Protocol regression tests cover the patched compiler cases behind this boundary.
-const wsOutboundValidator = WSOutboundMessageSchema as WSOutboundGeneratedValidator;
+const envelopeValidator = WSOutboundEnvelope as GeneratedValidator<WSOutboundMessage>;
+const sessionMessageValidators = SessionOutboundMessageValidators as Record<
+  SessionOutboundMessage["type"],
+  GeneratedValidator<SessionOutboundMessage>
+>;
 
-export function validateWSOutboundMessage(input: unknown): WSOutboundValidationResult {
-  return wsOutboundValidator.safeParse(input);
+// The message validator reports paths relative to the message, nested issues included.
+function underMessage(issue: z.core.$ZodIssue): z.core.$ZodIssue {
+  const path = ["message", ...issue.path];
+  switch (issue.code) {
+    case "invalid_union":
+      return issue.inclusive === false
+        ? { ...issue, path }
+        : { ...issue, path, errors: issue.errors.map((branch) => branch.map(underMessage)) };
+    case "invalid_key":
+    case "invalid_element":
+      return { ...issue, path, issues: issue.issues.map(underMessage) };
+    default:
+      return { ...issue, path };
+  }
+}
+
+export function validateWSOutboundMessage(input: unknown): ValidationResult<WSOutboundMessage> {
+  const envelope = envelopeValidator.safeParse(input);
+  if (!envelope.success || envelope.data.type !== "session") {
+    return envelope;
+  }
+  const message = envelope.data.message;
+  const result = sessionMessageValidators[message.type].safeParse(message);
+  if (!result.success) {
+    return { success: false, error: new ZodRealError(result.error.issues.map(underMessage)) };
+  }
+  // Message validators return a new object when they apply defaults.
+  if (result.data === message) {
+    return envelope;
+  }
+  return { success: true, data: { ...envelope.data, message: result.data } };
 }

packages/protocol/tests/validation/ws-outbound.test.ts

@@ -3,8 +3,10 @@ import { createRequire } from "node:module";
 import { dirname, join, resolve } from "node:path";
 import { fileURLToPath, pathToFileURL } from "node:url";
 import { createJiti } from "jiti";
+import type { z } from "zod";
 import { describe, expect, it } from "vitest";
-import { WSOutboundMessageSchema as GeneratedWSOutboundMessageSchema } from "../../src/generated/validation/ws-outbound.aot.js";
+import { WSOutboundMessageSchema } from "../../src/messages.js";
+import { validateWSOutboundMessage } from "../../src/validation/ws-outbound.js";
 
 interface GeneratedSchema {
   safeParse(input: unknown): { success: boolean; data?: unknown };
@@ -132,10 +134,8 @@ const SourceSchema = z.object({
   });
 
   it("accepts a minimal valid envelope and rejects a corrupted envelope", () => {
-    expect(GeneratedWSOutboundMessageSchema.safeParse({ type: "pong" }).success).toBe(true);
-    expect(GeneratedWSOutboundMessageSchema.safeParse({ type: "not_a_message" }).success).toBe(
-      false,
-    );
+    expect(validateWSOutboundMessage({ type: "pong" }).success).toBe(true);
+    expect(validateWSOutboundMessage({ type: "not_a_message" }).success).toBe(false);
   });
 
   it("accepts project config responses with and without setup commit status", () => {
@@ -158,11 +158,10 @@ const SourceSchema = z.object({
       },
     });
 
-    expect(GeneratedWSOutboundMessageSchema.safeParse(envelope(payload)).success).toBe(true);
+    expect(validateWSOutboundMessage(envelope(payload)).success).toBe(true);
     expect(
-      GeneratedWSOutboundMessageSchema.safeParse(
-        envelope({ ...payload, hasUncommittedWorktreeSetupChanges: true }),
-      ).success,
+      validateWSOutboundMessage(envelope({ ...payload, hasUncommittedWorktreeSetupChanges: true }))
+        .success,
     ).toBe(true);
   });
 
@@ -196,7 +195,7 @@ const SourceSchema = z.object({
       },
     };
 
-    expect(GeneratedWSOutboundMessageSchema.safeParse(envelope)).toEqual({
+    expect(validateWSOutboundMessage(envelope)).toEqual({
       success: true,
       data: envelope,
     });
@@ -255,7 +254,7 @@ const SourceSchema = z.object({
   ])("preserves workspaceId in a $name", ({ message }) => {
     const envelope = { type: "session", message };
 
-    expect(GeneratedWSOutboundMessageSchema.safeParse(envelope)).toEqual({
+    expect(validateWSOutboundMessage(envelope)).toEqual({
       success: true,
       data: envelope,
     });
@@ -263,11 +262,11 @@ const SourceSchema = z.object({
 
   it("emits runtime imports with .js extensions", async () => {
     const generated = await readFile(generatedWSOutboundPath, "utf8");
-    expect(generated).toContain('from "../../validation/ws-outbound-schema-metadata.js"');
+    expect(generated).toContain('from "./ws-outbound-schema-metadata.js"');
   });
 
   it("accepts a forge.search.response envelope", () => {
-    const result = GeneratedWSOutboundMessageSchema.safeParse({
+    const result = validateWSOutboundMessage({
       type: "session",
       message: {
         type: "forge.search.response",
@@ -293,7 +292,7 @@ const SourceSchema = z.object({
   });
 
   it("accepts a legacy github_search_response envelope", () => {
-    const result = GeneratedWSOutboundMessageSchema.safeParse({
+    const result = validateWSOutboundMessage({
       type: "session",
       message: {
         type: "github_search_response",
@@ -319,4 +318,86 @@ const SourceSchema = z.object({
     });
     expect(result.success).toBe(true);
   });
+
+  it("keeps each generated validator function small", async () => {
+    // A single function for the whole protocol made V8's Maglev tier spend ~800 MB and seconds compiling it.
+    const generated = await readFile(generatedWSOutboundPath, "utf8");
+    const starts = [...generated.matchAll(/function safeParse_\w+\(/g)].map((match) => match.index);
+    const sizes = starts.map((start, index) => (starts[index + 1] ?? generated.length) - start);
+
+    expect(Math.max(...sizes)).toBeLessThan(1_000_000);
+    expect(sizes.length).toBeGreaterThan(1);
+  });
+
+  it("keeps nested union issue paths rooted at the envelope", () => {
+    // Generated validators root branch issues at the input; runtime Zod reports them relative to the union.
+    const envelope = {
+      type: "session",
+      message: {
+        type: "read_project_config_response",
+        payload: { repoRoot: "/repo", ok: true, config: {}, revision: { mtimeMs: 1, size: 1 } },
+      },
+    };
+    const paths = (issues: z.core.$ZodIssue[]): PropertyKey[][] =>
+      issues.flatMap((issue) => [
+        issue.path,
+        ...(issue.code === "invalid_union" ? paths(issue.errors.flat()) : []),
+      ]);
+
+    const result = validateWSOutboundMessage(envelope);
+
+    expect(result.success).toBe(false);
+    if (result.success) return;
+    expect(paths(result.error.issues)).toEqual([
+      ["message", "payload"],
+      ["message", "payload", "requestId"],
+      ["message", "payload", "requestId"],
+      ["message", "payload", "ok"],
+      ["message", "payload", "error"],
+    ]);
+  });
+
+  it("applies defaults from the matching session message schema", () => {
+    const envelope = {
+      type: "session",
+      message: {
+        type: "agent_update",
+        payload: {
+          kind: "upsert",
+          agent: {
+            id: "agent-1",
+            provider: "claude",
+            cwd: "/repo",
+            model: "model",
+            createdAt: "2026-01-01T00:00:00.000Z",
+            updatedAt: "2026-01-01T00:00:00.000Z",
+            lastUserMessageAt: "2026-01-01T00:00:00.000Z",
+            status: "initializing",
+            capabilities: {
+              supportsStreaming: false,
+              supportsSessionPersistence: false,
+              supportsDynamicModes: false,
+              supportsMcpServers: false,
+              supportsReasoningStream: false,
+              supportsToolInvocations: false,
+            },
+            currentModeId: "default",
+            availableModes: [],
+            pendingPermissions: [],
+            persistence: { provider: "claude", sessionId: "session-1" },
+            title: "Agent",
+          },
+        },
+      },
+    };
+
+    const result = validateWSOutboundMessage(envelope);
+    const expected = WSOutboundMessageSchema.safeParse(envelope);
+
+    expect(expected.success).toBe(true);
+    expect(result).toEqual(expected);
+    expect(result.success && result.data).toMatchObject({
+      message: { payload: { agent: { capabilities: { supportsRewindConversation: false } } } },
+    });
+  });
 });