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

packages/protocol/codegen/README.md

@@ -0,0 +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.
+
+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.
+
+`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

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

packages/protocol/package.json

@@ -23,7 +23,12 @@
     "build:clean": "npm run clean && npm run build",
     "prepack": "npm run build:clean",
     "typecheck": "tsgo --noEmit",
-    "test": "vitest run"
+    "test": "vitest run",
+    "generate:validators": "node scripts/generate-validation-aot.mjs",
+    "watch": "concurrently --kill-others --names validation,tsc --prefix-colors green,yellow \"node scripts/watch-validation-aot.mjs\" \"tsc -p tsconfig.json --watch --preserveWatchOutput\"",
+    "prebuild": "npm run generate:validators",
+    "pretypecheck": "npm run generate:validators",
+    "pretest": "npm run generate:validators"
   },
   "dependencies": {
     "zod": "^4.4.3"
@@ -31,6 +36,7 @@
   "devDependencies": {
     "@types/node": "^20.9.0",
     "typescript": "^5.2.2",
-    "vitest": "^4.1.6"
+    "vitest": "^4.1.6",
+    "zod-aot": "0.20.4"
   }
 }

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

@@ -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(", ")})`,
+);

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

@@ -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);

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

@@ -0,0 +1,11 @@
+# Generated protocol validators
+
+`ws-outbound.aot.ts` is generated by zod-aot from `packages/protocol/codegen/ws-outbound.compile.ts`.
+
+Run:
+
+```bash
+npm run generate:validators --workspace=@getpaseo/protocol
+```
+
+The generated TypeScript is intentionally gitignored. zod-aot is exact-pinned because this is a young solo-maintainer project; keep the generated validator behind protocol regression tests when upgrading it.

packages/protocol/src/loop/rpc-schemas.ts

@@ -1,5 +1,5 @@
 import { z } from "zod";
-import { AgentProviderSchema } from "@getpaseo/protocol/provider-manifest";
+import { AgentProviderSchema } from "../provider-manifest.js";
 
 export const LoopLogEntrySchema = z.object({
   seq: z.number().int().positive(),

packages/protocol/src/messages.providers-snapshot.test.ts

@@ -38,41 +38,6 @@ describe("provider snapshot message schemas", () => {
     expect(parsed.enabled).toBe(true);
   });
 
-  test("normalizes thinking option defaults on provider snapshot models", () => {
-    const parsed = ProviderSnapshotEntrySchema.parse({
-      provider: "claude",
-      status: "ready",
-      models: [
-        {
-          provider: "claude",
-          id: "MiniMax-M2.7",
-          label: "MiniMax-M2.7",
-          isDefault: true,
-          contextWindowMaxTokens: 1_000_000,
-          thinkingOptions: [
-            { id: "off", label: "Off" },
-            { id: "max", label: "Max", isDefault: true },
-          ],
-        },
-      ],
-    });
-
-    expect(parsed.models).toEqual([
-      {
-        provider: "claude",
-        id: "MiniMax-M2.7",
-        label: "MiniMax-M2.7",
-        isDefault: true,
-        contextWindowMaxTokens: 1_000_000,
-        thinkingOptions: [
-          { id: "off", label: "Off" },
-          { id: "max", label: "Max", isDefault: true },
-        ],
-        defaultThinkingOptionId: "max",
-      },
-    ]);
-  });
-
   test("defaults missing enabled state in providers snapshot response entries", () => {
     const parsed = GetProvidersSnapshotResponseMessageSchema.parse({
       type: "get_providers_snapshot_response",

packages/protocol/src/messages.ts

@@ -2,9 +2,9 @@ import { z } from "zod";
 import { TerminalActivitySchema } from "./terminal-activity.js";
 import { CLIENT_CAPS } from "./client-capabilities.js";
 import { AGENT_LIFECYCLE_STATUSES } from "./agent-lifecycle.js";
-import { MAX_EXPLICIT_AGENT_TITLE_CHARS } from "@getpaseo/protocol/agent-title-limits";
-import { AgentProviderSchema } from "@getpaseo/protocol/provider-manifest";
-import { normalizeAgentModelDefinition, TOOL_CALL_ICON_NAMES } from "./agent-types.js";
+import { MAX_EXPLICIT_AGENT_TITLE_CHARS } from "./agent-title-limits.js";
+import { AgentProviderSchema } from "./provider-manifest.js";
+import { TOOL_CALL_ICON_NAMES } from "./agent-types.js";
 import {
   ChatCreateRequestSchema,
   ChatListRequestSchema,
@@ -40,7 +40,7 @@ import {
   ScheduleDeleteResponseSchema,
   ScheduleRunOnceResponseSchema,
   ScheduleUpdateResponseSchema,
-} from "@getpaseo/protocol/schedule/rpc-schemas";
+} from "./schedule/rpc-schemas.js";
 import {
   LoopRunRequestSchema,
   LoopListRequestSchema,
@@ -52,7 +52,7 @@ import {
   LoopInspectResponseSchema,
   LoopLogsResponseSchema,
   LoopStopResponseSchema,
-} from "@getpaseo/protocol/loop/rpc-schemas";
+} from "./loop/rpc-schemas.js";
 import {
   BrowserAutomationExecuteRequestSchema,
   BrowserAutomationExecuteResponseSchema,
@@ -73,7 +73,7 @@ import {
   type PaseoMetadataGenerationEntry,
   type PaseoScriptEntryRaw,
   type ProjectConfigRpcError,
-} from "@getpaseo/protocol/paseo-config-schema";
+} from "./paseo-config-schema.js";
 export {
   PaseoConfigRawSchema,
   PaseoLifecycleCommandRawSchema,
@@ -247,19 +247,17 @@ export const AgentFeatureSchema = z.discriminatedUnion("type", [
   AgentFeatureSelectSchema,
 ]);
 
-const AgentModelDefinitionSchema: z.ZodType<AgentModelDefinition> = z
-  .object({
-    provider: AgentProviderSchema,
-    id: z.string(),
-    label: z.string(),
-    description: z.string().optional(),
-    isDefault: z.boolean().optional(),
-    metadata: z.record(z.string(), z.unknown()).optional(),
-    contextWindowMaxTokens: z.number().optional(),
-    thinkingOptions: z.array(AgentSelectOptionSchema).optional(),
-    defaultThinkingOptionId: z.string().optional(),
-  })
-  .transform(normalizeAgentModelDefinition);
+const AgentModelDefinitionSchema: z.ZodType<AgentModelDefinition> = z.object({
+  provider: AgentProviderSchema,
+  id: z.string(),
+  label: z.string(),
+  description: z.string().optional(),
+  isDefault: z.boolean().optional(),
+  metadata: z.record(z.string(), z.unknown()).optional(),
+  contextWindowMaxTokens: z.number().optional(),
+  thinkingOptions: z.array(AgentSelectOptionSchema).optional(),
+  defaultThinkingOptionId: z.string().optional(),
+});
 
 export const ProviderSnapshotEntrySchema = z.object({
   provider: AgentProviderSchema,
@@ -361,20 +359,21 @@ const AgentPermissionActionSchema = z.object({
   intent: z.enum(["implement", "implement_resume", "dismiss"]).optional(),
 });
 
-export const AgentPermissionResponseSchema: z.ZodType<AgentPermissionResponse> = z.union([
-  z.object({
-    behavior: z.literal("allow"),
-    selectedActionId: z.string().optional(),
-    updatedInput: z.record(z.string(), z.unknown()).optional(),
-    updatedPermissions: z.array(AgentPermissionUpdateSchema).optional(),
-  }),
-  z.object({
-    behavior: z.literal("deny"),
-    selectedActionId: z.string().optional(),
-    message: z.string().optional(),
-    interrupt: z.boolean().optional(),
-  }),
-]);
+export const AgentPermissionResponseSchema: z.ZodType<AgentPermissionResponse> =
+  z.discriminatedUnion("behavior", [
+    z.object({
+      behavior: z.literal("allow"),
+      selectedActionId: z.string().optional(),
+      updatedInput: z.record(z.string(), z.unknown()).optional(),
+      updatedPermissions: z.array(AgentPermissionUpdateSchema).optional(),
+    }),
+    z.object({
+      behavior: z.literal("deny"),
+      selectedActionId: z.string().optional(),
+      message: z.string().optional(),
+      interrupt: z.boolean().optional(),
+    }),
+  ]);
 
 export const AgentPermissionRequestPayloadSchema: z.ZodType<AgentPermissionRequest, unknown> =
   z.object({
@@ -554,13 +553,16 @@ const ToolCallCanceledPayloadSchema = ToolCallBasePayloadSchema.extend({
   error: z.null(),
 });
 
-const ToolCallTimelineItemPayloadSchema: z.ZodType<ToolCallTimelineItem, unknown> = z.union([
-  ToolCallRunningPayloadSchema,
-  ToolCallCompletedPayloadSchema,
-  ToolCallFailedPayloadSchema,
-  ToolCallCanceledPayloadSchema,
-]);
+const ToolCallTimelineItemPayloadSchema: z.ZodType<ToolCallTimelineItem, unknown> =
+  z.discriminatedUnion("status", [
+    ToolCallRunningPayloadSchema,
+    ToolCallCompletedPayloadSchema,
+    ToolCallFailedPayloadSchema,
+    ToolCallCanceledPayloadSchema,
+  ]);
 
+// zod-aot 0.20.4 miscompiles this as a nested discriminated union by omitting
+// the inner tool_call branch from the generated outer dispatch.
 export const AgentTimelineItemPayloadSchema: z.ZodType<AgentTimelineItem, unknown> = z.union([
   z.object({
     type: z.literal("user_message"),
@@ -3113,7 +3115,9 @@ export const SetDaemonConfigResponseMessageSchema = z.object({
 
 export const ReadProjectConfigResponseMessageSchema = z.object({
   type: z.literal("read_project_config_response"),
-  payload: z.discriminatedUnion("ok", [
+  // zod-aot 0.2.0 miscompiles boolean discriminators as string options
+  // (`"true"`/`"false"`), so keep this sequential until upstream fixes it.
+  payload: z.union([
     z.object({
       requestId: z.string(),
       repoRoot: z.string(),
@@ -3132,7 +3136,9 @@ export const ReadProjectConfigResponseMessageSchema = z.object({
 
 export const WriteProjectConfigResponseMessageSchema = z.object({
   type: z.literal("write_project_config_response"),
-  payload: z.discriminatedUnion("ok", [
+  // zod-aot 0.2.0 miscompiles boolean discriminators as string options
+  // (`"true"`/`"false"`), so keep this sequential until upstream fixes it.
+  payload: z.union([
     z.object({
       requestId: z.string(),
       repoRoot: z.string(),

packages/protocol/src/schedule/types.ts

@@ -1,5 +1,5 @@
 import { z } from "zod";
-import { AgentProviderSchema } from "@getpaseo/protocol/provider-manifest";
+import { AgentProviderSchema } from "../provider-manifest.js";
 
 export const ScheduleStatusSchema = z.enum(["active", "paused", "completed"]);
 export type ScheduleStatus = z.infer<typeof ScheduleStatusSchema>;

packages/protocol/src/tool-call-display.ts

@@ -1,5 +1,5 @@
 import type { ToolCallTimelineItem } from "./agent-types.js";
-import { getPaseoToolLeafName, isPaseoToolName } from "@getpaseo/protocol/tool-name-normalization";
+import { getPaseoToolLeafName, isPaseoToolName } from "./tool-name-normalization.js";
 import { stripCwdPrefix } from "./path-utils.js";
 
 export type ToolCallDisplayInput = Pick<

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

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

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

@@ -0,0 +1,19 @@
+import type { z } from "zod";
+import { WSOutboundMessageSchema } from "../generated/validation/ws-outbound.aot.js";
+import type { WSOutboundMessage } from "../messages.js";
+
+type WSOutboundValidationResult =
+  | { success: true; data: WSOutboundMessage }
+  | { success: false; error: z.ZodError };
+
+interface WSOutboundGeneratedValidator {
+  safeParse(input: unknown): WSOutboundValidationResult;
+}
+
+// zod-aot emits runtime JavaScript from WSOutboundMessageSchema but not its TypeScript surface.
+// Protocol regression tests cover the patched compiler cases behind this boundary.
+const wsOutboundValidator = WSOutboundMessageSchema as WSOutboundGeneratedValidator;
+
+export function validateWSOutboundMessage(input: unknown): WSOutboundValidationResult {
+  return wsOutboundValidator.safeParse(input);
+}

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

@@ -0,0 +1,145 @@
+import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
+import { createRequire } from "node:module";
+import { dirname, join, resolve } from "node:path";
+import { fileURLToPath, pathToFileURL } from "node:url";
+import { createJiti } from "jiti";
+import { describe, expect, it } from "vitest";
+import { WSOutboundMessageSchema as GeneratedWSOutboundMessageSchema } from "../../src/generated/validation/ws-outbound.aot.js";
+
+interface GeneratedSchema {
+  safeParse(input: unknown): { success: boolean; data?: unknown };
+}
+
+const protocolRoot = resolve(dirname(fileURLToPath(import.meta.url)), "../..");
+const generatedWSOutboundPath = resolve(
+  protocolRoot,
+  "src/generated/validation/ws-outbound.aot.ts",
+);
+const require = createRequire(import.meta.url);
+
+async function compileInlineSchema(sourceSchema: string): Promise<GeneratedSchema> {
+  const scratchRoot = resolve(protocolRoot, "../../.tmp");
+  await mkdir(scratchRoot, { recursive: true });
+  const tempDir = await mkdtemp(join(scratchRoot, "paseo-zod-aot-"));
+
+  try {
+    const sourcePath = join(tempDir, "schema.source.js");
+    const outputPath = join(tempDir, "schema.generated.ts");
+    await writeFile(join(tempDir, "package.json"), '{"type":"module"}\n');
+    await writeFile(
+      sourcePath,
+      [
+        'import { z } from "zod";',
+        'import { compile } from "zod-aot";',
+        sourceSchema,
+        "export const Schema = compile(SourceSchema);",
+        "",
+      ].join("\n"),
+    );
+
+    const zodAotEntry = require.resolve("zod-aot");
+    const zodAotRoot = resolve(dirname(zodAotEntry), "..");
+    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(sourcePath, { cacheBust: true });
+    const compiled = compileSchemas(schemas, { mode: "inline" });
+    const content = generateCompiledFileContent(compiled, "./schema.source.js", {
+      zodCompat: false,
+    });
+    await writeFile(outputPath, content);
+
+    const jiti = createJiti(import.meta.url, { moduleCache: false });
+    const generated = await jiti.import(outputPath);
+    return generated.Schema as GeneratedSchema;
+  } finally {
+    await rm(tempDir, { recursive: true, force: true });
+  }
+}
+
+describe("WS outbound zod-aot validation", () => {
+  it("applies defaults inside discriminated-union branches", async () => {
+    const schema = await compileInlineSchema(`
+const SourceSchema = z.discriminatedUnion("type", [
+  z.object({
+    type: z.literal("with_default"),
+    enabled: z.boolean().default(true),
+  }),
+  z.object({
+    type: z.literal("without_default"),
+    label: z.string(),
+  }),
+]);
+`);
+
+    expect(schema.safeParse({ type: "with_default" })).toMatchObject({
+      success: true,
+      data: { type: "with_default", enabled: true },
+    });
+  });
+
+  it("routes tool-call-like status unions through the current sequential item union", async () => {
+    const schema = await compileInlineSchema(`
+const ToolCallItemSchema = z.discriminatedUnion("status", [
+  z.object({ type: z.literal("tool_call"), status: z.literal("running"), callId: z.string() }),
+  z.object({ type: z.literal("tool_call"), status: z.literal("completed"), callId: z.string(), output: z.string() }),
+  z.object({ type: z.literal("tool_call"), status: z.literal("failed"), callId: z.string(), error: z.string() }),
+  z.object({ type: z.literal("tool_call"), status: z.literal("canceled"), callId: z.string() }),
+]);
+
+const TimelineItemSchema = z.union([
+  z.object({ type: z.literal("assistant_message"), text: z.string() }),
+  ToolCallItemSchema,
+]);
+
+const SourceSchema = z.object({
+  item: TimelineItemSchema,
+});
+`);
+
+    expect(
+      schema.safeParse({ item: { type: "tool_call", status: "running", callId: "run" } }),
+    ).toMatchObject({
+      success: true,
+      data: { item: { type: "tool_call", status: "running", callId: "run" } },
+    });
+    expect(
+      schema.safeParse({
+        item: { type: "tool_call", status: "completed", callId: "done", output: "ok" },
+      }),
+    ).toMatchObject({
+      success: true,
+      data: { item: { type: "tool_call", status: "completed", callId: "done", output: "ok" } },
+    });
+    expect(
+      schema.safeParse({
+        item: { type: "tool_call", status: "failed", callId: "fail", error: "boom" },
+      }),
+    ).toMatchObject({
+      success: true,
+      data: { item: { type: "tool_call", status: "failed", callId: "fail", error: "boom" } },
+    });
+    expect(
+      schema.safeParse({ item: { type: "tool_call", status: "canceled", callId: "stop" } }),
+    ).toMatchObject({
+      success: true,
+      data: { item: { type: "tool_call", status: "canceled", callId: "stop" } },
+    });
+  });
+
+  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,
+    );
+  });
+
+  it("emits runtime imports with .js extensions", async () => {
+    const generated = await readFile(generatedWSOutboundPath, "utf8");
+    expect(generated).toContain('from "../../validation/ws-outbound-schema-metadata.js"');
+  });
+});