diff --git a/.gitignore b/.gitignore index 6d6bb279b7..a85ff1fd8c 100644 --- a/.gitignore +++ b/.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/ diff --git a/docs/protocol-validation.md b/docs/protocol-validation.md index 5fc31a5fdf..4694f9e208 100644 --- a/docs/protocol-validation.md +++ b/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. diff --git a/packages/protocol/codegen/README.md b/packages/protocol/codegen/README.md index 7b2b0bcee5..69d7ac79b6 100644 --- a/packages/protocol/codegen/README.md +++ b/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. diff --git a/packages/protocol/codegen/ws-outbound.compile.ts b/packages/protocol/codegen/ws-outbound.compile.ts deleted file mode 100644 index 589986cf6a..0000000000 --- a/packages/protocol/codegen/ws-outbound.compile.ts +++ /dev/null @@ -1,4 +0,0 @@ -import { compile } from "zod-aot"; -import { WSOutboundMessageSchema as SourceWSOutboundMessageSchema } from "../src/messages.js"; - -export const WSOutboundMessageSchema = compile(SourceWSOutboundMessageSchema); diff --git a/packages/protocol/scripts/generate-validation-aot.mjs b/packages/protocol/scripts/generate-validation-aot.mjs index c3fd0b72ed..34c393c19e 100644 --- a/packages/protocol/scripts/generate-validation-aot.mjs +++ b/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`, ); diff --git a/packages/protocol/src/generated/validation/README.md b/packages/protocol/src/generated/validation/README.md index 7d650faf4b..10ebc31e14 100644 --- a/packages/protocol/src/generated/validation/README.md +++ b/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: diff --git a/packages/protocol/src/validation/ws-outbound-envelope.ts b/packages/protocol/src/validation/ws-outbound-envelope.ts new file mode 100644 index 0000000000..0ced88e8a9 --- /dev/null +++ b/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, +]); diff --git a/packages/protocol/src/validation/ws-outbound-schema-metadata.ts b/packages/protocol/src/validation/ws-outbound-schema-metadata.ts deleted file mode 100644 index 23c9b5e563..0000000000 --- a/packages/protocol/src/validation/ws-outbound-schema-metadata.ts +++ /dev/null @@ -1,3 +0,0 @@ -import { WSOutboundMessageSchema as SourceWSOutboundMessageSchema } from "../messages.js"; - -export const WSOutboundMessageSchema = { schema: SourceWSOutboundMessageSchema }; diff --git a/packages/protocol/src/validation/ws-outbound.ts b/packages/protocol/src/validation/ws-outbound.ts index bcb568cd8e..6cc7ad1268 100644 --- a/packages/protocol/src/validation/ws-outbound.ts +++ b/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 = { success: true; data: T } | { success: false; error: z.ZodError }; -interface WSOutboundGeneratedValidator { - safeParse(input: unknown): WSOutboundValidationResult; +interface GeneratedValidator { + safeParse(input: unknown): ValidationResult; } -// 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; +const sessionMessageValidators = SessionOutboundMessageValidators as Record< + SessionOutboundMessage["type"], + GeneratedValidator +>; -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 { + 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 } }; } diff --git a/packages/protocol/tests/validation/ws-outbound.test.ts b/packages/protocol/tests/validation/ws-outbound.test.ts index e2e3e3170a..ed147d89ad 100644 --- a/packages/protocol/tests/validation/ws-outbound.test.ts +++ b/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 } } } }, + }); + }); });