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/client/package.json

@@ -27,7 +27,7 @@
   },
   "scripts": {
     "clean": "node ../../scripts/clean-package-dist.mjs",
-    "build": "tsc -p tsconfig.json --incremental false",
+    "build": "npm run build --workspace=@getpaseo/protocol && tsc -p tsconfig.json --incremental false",
     "build:clean": "npm run clean && npm run build",
     "prepack": "npm run build:clean",
     "typecheck": "tsgo --noEmit",

packages/client/src/compat/normalize-provider-models.ts

@@ -0,0 +1,74 @@
+import type { AgentModelDefinition, ProviderSnapshotEntry } from "@getpaseo/protocol/agent-types";
+import { normalizeAgentModelDefinition } from "@getpaseo/protocol/agent-types";
+import type {
+  GetProvidersSnapshotResponseMessage,
+  ListProviderModelsResponseMessage,
+  SessionOutboundMessage,
+} from "@getpaseo/protocol/messages";
+
+type ListProviderModelsPayload = ListProviderModelsResponseMessage["payload"];
+type GetProvidersSnapshotPayload = GetProvidersSnapshotResponseMessage["payload"];
+type ProvidersSnapshotUpdatePayload = Extract<
+  SessionOutboundMessage,
+  { type: "providers_snapshot_update" }
+>["payload"];
+
+// COMPAT(model-normalize): daemon normalizes at source (provider-registry) — shim covers older daemons; drop when floor >= v0.1.104
+function normalizeAgentModels(
+  models: AgentModelDefinition[] | undefined,
+): AgentModelDefinition[] | undefined {
+  if (!models) {
+    return models;
+  }
+
+  let changed = false;
+  const normalized = models.map((model) => {
+    const next = normalizeAgentModelDefinition(model);
+    changed ||= next !== model;
+    return next;
+  });
+
+  return changed ? normalized : models;
+}
+
+function normalizeProviderSnapshotEntry(entry: ProviderSnapshotEntry): ProviderSnapshotEntry {
+  const models = normalizeAgentModels(entry.models);
+  return models === entry.models ? entry : { ...entry, models };
+}
+
+function normalizeProviderSnapshotEntries(
+  entries: ProviderSnapshotEntry[],
+): ProviderSnapshotEntry[] {
+  let changed = false;
+  const normalized = entries.map((entry) => {
+    const next = normalizeProviderSnapshotEntry(entry);
+    changed ||= next !== entry;
+    return next;
+  });
+
+  return changed ? normalized : entries;
+}
+
+export function normalizeListProviderModelsPayload(
+  payload: ListProviderModelsPayload,
+): ListProviderModelsPayload {
+  const models = normalizeAgentModels(payload.models);
+  return models === payload.models ? payload : { ...payload, models };
+}
+
+export function normalizeProvidersSnapshotPayload<
+  T extends GetProvidersSnapshotPayload | ProvidersSnapshotUpdatePayload,
+>(payload: T): T {
+  const entries = normalizeProviderSnapshotEntries(payload.entries);
+  return entries === payload.entries ? payload : { ...payload, entries };
+}
+
+export function normalizeProviderSnapshotUpdateMessage(
+  msg: SessionOutboundMessage,
+): SessionOutboundMessage {
+  if (msg.type !== "providers_snapshot_update") {
+    return msg;
+  }
+
+  return { ...msg, payload: normalizeProvidersSnapshotPayload(msg.payload) };
+}

packages/client/src/daemon-client.ts

@@ -13,8 +13,8 @@ import {
   DaemonUpdateResponseSchema,
   SessionInboundMessageSchema,
   type ServerInfoStatusPayload,
-  WSOutboundMessageSchema,
 } from "@getpaseo/protocol/messages";
+import { validateWSOutboundMessage } from "@getpaseo/protocol/validation/ws-outbound";
 import type {
   AgentStreamEventPayload,
   AgentSnapshotPayload,
@@ -118,6 +118,11 @@ import {
   type WebSocketFactory,
 } from "./daemon-client-transport.js";
 import { DaemonClientRuntimeMetrics } from "./daemon-client-runtime-metrics.js";
+import {
+  normalizeListProviderModelsPayload,
+  normalizeProviderSnapshotUpdateMessage,
+  normalizeProvidersSnapshotPayload,
+} from "./compat/normalize-provider-models.js";
 import { TerminalStreamRouter, type TerminalStreamEvent } from "./terminal-stream-router.js";
 import type {
   BrowserAutomationExecuteRequest,
@@ -370,6 +375,7 @@ type WriteProjectConfigPayload = Extract<
   SessionOutboundMessage,
   { type: "write_project_config_response" }
 >["payload"];
+
 type ListCommandsPayload = ListCommandsResponse["payload"];
 type ListCommandsDraftConfig = Pick<
   AgentSessionConfig,
@@ -3728,7 +3734,7 @@ export class DaemonClient {
     provider: AgentProvider,
     options?: { cwd?: string; requestId?: string },
   ): Promise<ListProviderModelsPayload> {
-    return this.sendCorrelatedSessionRequest({
+    const payload = await this.sendCorrelatedSessionRequest({
       requestId: options?.requestId,
       message: {
         type: "list_provider_models_request",
@@ -3739,6 +3745,7 @@ export class DaemonClient {
       // Provider SDK cold starts (especially model discovery) can exceed 60s.
       timeout: 90000,
     });
+    return normalizeListProviderModelsPayload(payload);
   }
 
   async listProviderModes(
@@ -3788,7 +3795,7 @@ export class DaemonClient {
     cwd?: string;
     requestId?: string;
   }): Promise<GetProvidersSnapshotPayload> {
-    return this.sendCorrelatedSessionRequest({
+    const payload = await this.sendCorrelatedSessionRequest({
       requestId: options?.requestId,
       message: {
         type: "get_providers_snapshot_request",
@@ -3796,6 +3803,7 @@ export class DaemonClient {
       },
       responseType: "get_providers_snapshot_response",
     });
+    return normalizeProvidersSnapshotPayload(payload);
   }
 
   async getDaemonConfig(
@@ -4754,7 +4762,7 @@ export class DaemonClient {
       return;
     }
 
-    const parsed = WSOutboundMessageSchema.safeParse(parsedJson);
+    const parsed = validateWSOutboundMessage(parsedJson);
     if (!parsed.success) {
       const msgType =
         parsedJson != null &&
@@ -5017,8 +5025,10 @@ export class DaemonClient {
   }
 
   private handleSessionMessage(msg: SessionOutboundMessage): void {
-    if (msg.type === "status") {
-      const serverInfo = parseServerInfoStatusPayload(msg.payload);
+    const consumerMessage = normalizeProviderSnapshotUpdateMessage(msg);
+
+    if (consumerMessage.type === "status") {
+      const serverInfo = parseServerInfoStatusPayload(consumerMessage.payload);
       if (serverInfo) {
         this.lastServerInfoMessage = serverInfo;
         if (this.connectionState.status === "connecting") {
@@ -5034,39 +5044,39 @@ export class DaemonClient {
       }
     }
 
-    if (msg.type === "terminal_stream_exit") {
-      this.terminalStreams.removeTerminal(msg.payload.terminalId);
+    if (consumerMessage.type === "terminal_stream_exit") {
+      this.terminalStreams.removeTerminal(consumerMessage.payload.terminalId);
     }
 
     if (this.rawMessageListeners.size > 0) {
       for (const handler of this.rawMessageListeners) {
         try {
-          handler(msg);
+          handler(consumerMessage);
         } catch {
           // no-op
         }
       }
     }
 
-    const handlers = this.messageHandlers.get(msg.type);
+    const handlers = this.messageHandlers.get(consumerMessage.type);
     if (handlers) {
       for (const handler of handlers) {
         try {
-          handler(msg);
+          handler(consumerMessage);
         } catch {
           // no-op
         }
       }
     }
 
-    const event = this.toEvent(msg);
+    const event = this.toEvent(consumerMessage);
     if (event) {
       for (const handler of this.eventListeners) {
         handler(event);
       }
     }
 
-    this.resolveWaiters(msg);
+    this.resolveWaiters(consumerMessage);
   }
 
   private resolveWaiters(msg: SessionOutboundMessage): void {

packages/client/src/index.test.ts

@@ -629,20 +629,37 @@ test("provider actions delegate to existing provider RPCs and local snapshot upd
   });
 
   const snapshotUpdates: string[] = [];
+  const snapshotModelDefaults: Array<string | undefined> = [];
   const unsubscribe = client.providers.subscribe((update) => {
     snapshotUpdates.push(update.generatedAt);
+    snapshotModelDefaults.push(update.entries[0]?.models?.[0]?.defaultThinkingOptionId);
   });
   ws.message(
     sessionMessage({
       type: "providers_snapshot_update",
       payload: {
         cwd: "/repo/sdk",
-        entries: [{ provider: "codex", status: "ready", enabled: true }],
+        entries: [
+          {
+            provider: "codex",
+            status: "ready",
+            enabled: true,
+            models: [
+              {
+                provider: "codex",
+                id: "gpt-5.5",
+                label: "GPT-5.5",
+                thinkingOptions: [{ id: "high", label: "High", isDefault: true }],
+              },
+            ],
+          },
+        ],
         generatedAt: "2026-05-16T01:00:00.000Z",
       },
     }),
   );
   expect(snapshotUpdates).toEqual(["2026-05-16T01:00:00.000Z"]);
+  expect(snapshotModelDefaults).toEqual(["high"]);
 
   unsubscribe();
   await client.close();

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"');
+  });
+});