perf: zod compile reusable server validation schemas

Mono-RevId: 749cbfed4729e03d99a710f4ce62acdc0beba8ca

Author: carderne

apps/supervisor/src/clients/responseSchemas.test.ts

@@ -0,0 +1,35 @@
+import { WorkerApiRunLatestSnapshotResponseBody } from "@trigger.dev/core/v3/workers";
+import { describe, expect, it } from "vitest";
+import { z } from "zod";
+import { resolveResponseSchema } from "./responseSchemas.js";
+
+describe("supervisor response schemas", () => {
+  it("retains date normalization and validation errors for snapshots", () => {
+    const schema = resolveResponseSchema(WorkerApiRunLatestSnapshotResponseBody);
+    const result = schema.parse({
+      execution: {
+        version: "1",
+        snapshot: {
+          id: "snapshot_1",
+          friendlyId: "snap_1",
+          executionStatus: "EXECUTING",
+          description: "Run is executing",
+          createdAt: "2026-01-01T00:00:00.000Z",
+        },
+        run: { id: "run_1", friendlyId: "run_1", status: "EXECUTING" },
+        completedWaitpoints: [],
+      },
+    });
+
+    expect(result.execution.snapshot.createdAt).toEqual(new Date("2026-01-01T00:00:00.000Z"));
+    expect(schema.safeParse({ execution: {} }).success).toBe(false);
+  });
+
+  it("does not compile unfamiliar schemas or repeat their callbacks", () => {
+    let calls = 0;
+    const schema = z.string().refine(() => ++calls > 1);
+
+    expect(resolveResponseSchema(schema).safeParse("value").success).toBe(false);
+    expect(calls).toBe(1);
+  });
+});

apps/supervisor/src/clients/responseSchemas.ts

@@ -0,0 +1,33 @@
+import type { AnyZodSchema } from "@trigger.dev/core/v3";
+import {
+  WorkerApiConnectResponseBody,
+  WorkerApiContinueRunExecutionRequestBody,
+  WorkerApiDequeueResponseBody,
+  WorkerApiHeartbeatResponseBody,
+  WorkerApiRunAttemptCompleteResponseBody,
+  WorkerApiRunHeartbeatResponseBody,
+  WorkerApiRunLatestSnapshotResponseBody,
+  WorkerApiRunSnapshotsSinceResponseBody,
+  WorkerApiSuspendRunResponseBody,
+} from "@trigger.dev/core/v3/workers";
+import { z } from "zod";
+
+const responseSchemas = [
+  WorkerApiConnectResponseBody,
+  WorkerApiContinueRunExecutionRequestBody,
+  WorkerApiDequeueResponseBody,
+  WorkerApiHeartbeatResponseBody,
+  WorkerApiRunAttemptCompleteResponseBody,
+  WorkerApiRunHeartbeatResponseBody,
+  WorkerApiRunLatestSnapshotResponseBody,
+  WorkerApiRunSnapshotsSinceResponseBody,
+  WorkerApiSuspendRunResponseBody,
+];
+
+const compiledSchemas = new Map<AnyZodSchema, AnyZodSchema>(
+  responseSchemas.map((schema) => [schema, z.compile(schema)])
+);
+
+export function resolveResponseSchema<T extends AnyZodSchema>(schema: T): T {
+  return (compiledSchemas.get(schema) as T | undefined) ?? schema;
+}

apps/supervisor/src/index.ts

@@ -23,6 +23,7 @@ import {
   isKubernetesEnvironment,
 } from "@trigger.dev/core/v3/serverOnly";
 import { createK8sApi, createPodCountFetcher } from "./clients/kubernetes.js";
+import { resolveResponseSchema } from "./clients/responseSchemas.js";
 import { collectDefaultMetrics, Counter, Gauge, Histogram } from "prom-client";
 import { register } from "./metrics.js";
 import { PodCleaner } from "./services/podCleaner.js";
@@ -324,6 +325,7 @@ class ManagedSupervisor {
     const workerToken = getWorkerToken();
 
     this.workerSession = new SupervisorSession({
+      resolveResponseSchema,
       workerToken,
       apiUrl: env.TRIGGER_API_URL,
       instanceName: env.TRIGGER_WORKER_INSTANCE_NAME,

apps/supervisor/src/workloadServer/index.ts

@@ -66,10 +66,16 @@ const checkpointCancelRequests = new Counter({
   registers: [register],
 });
 
-const WorkloadActionParams = z.object({
-  runFriendlyId: z.string(),
-  snapshotFriendlyId: z.string(),
-});
+const WorkloadActionParams = z.compile(
+  z.object({
+    runFriendlyId: z.string(),
+    snapshotFriendlyId: z.string(),
+  })
+);
+const WorkloadDebugParams = z.compile(WorkloadActionParams.pick({ runFriendlyId: true }));
+const CompiledWorkloadRunAttemptStartRequestBody = z.compile(WorkloadRunAttemptStartRequestBody);
+const CompiledWorkloadHeartbeatRequestBody = z.compile(WorkloadHeartbeatRequestBody);
+const CompiledWorkloadDebugLogRequestBody = z.compile(WorkloadDebugLogRequestBody);
 
 // Workloads bundled into customer task images before CLI v4.4.4 use a strict
 // zod enum for checkpoint type that only allows DOCKER and KUBERNETES. The
@@ -385,7 +391,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
         "POST",
         {
           paramsSchema: WorkloadActionParams,
-          bodySchema: WorkloadRunAttemptStartRequestBody,
+          bodySchema: CompiledWorkloadRunAttemptStartRequestBody,
           handler: async (ctx) =>
             this.wideRoute(
               ctx,
@@ -490,7 +496,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
         "POST",
         {
           paramsSchema: WorkloadActionParams,
-          bodySchema: WorkloadHeartbeatRequestBody,
+          bodySchema: CompiledWorkloadHeartbeatRequestBody,
           handler: async (ctx) =>
             this.wideRoute(
               ctx,
@@ -724,8 +730,8 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
 
     if (env.SEND_RUN_DEBUG_LOGS) {
       httpServer.route("/api/v1/workload-actions/runs/:runFriendlyId/logs/debug", "POST", {
-        paramsSchema: WorkloadActionParams.pick({ runFriendlyId: true }),
-        bodySchema: WorkloadDebugLogRequestBody,
+        paramsSchema: WorkloadDebugParams,
+        bodySchema: CompiledWorkloadDebugLogRequestBody,
         handler: async (ctx) =>
           this.wideRoute(
             ctx,

apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts

@@ -53,7 +53,7 @@ const { action, loader } = createActionApiRoute(
   {
     headers: HeadersSchema,
     params: ParamsSchema,
-    body: TriggerTaskRequestBody,
+    body: z.compile(TriggerTaskRequestBody),
     allowJWT: true,
     maxContentLength: env.TASK_PAYLOAD_MAXIMUM_SIZE,
     authorization: {

apps/webapp/app/routes/api.v2.tasks.batch.ts

@@ -1,4 +1,5 @@
 import { json } from "@remix-run/server-runtime";
+import { z } from "zod";
 import type { BatchTriggerTaskV3Response } from "@trigger.dev/core/v3";
 import { BatchTriggerTaskV3RequestBody, generateJWT } from "@trigger.dev/core/v3";
 import { env } from "~/env.server";
@@ -28,7 +29,7 @@ const { action, loader } = createActionApiRoute(
     headers: HeadersSchema.extend({
       "batch-processing-strategy": BatchProcessingStrategy.nullish(),
     }),
-    body: BatchTriggerTaskV3RequestBody,
+    body: z.compile(BatchTriggerTaskV3RequestBody),
     allowJWT: true,
     maxContentLength: env.BATCH_TASK_PAYLOAD_MAXIMUM_SIZE,
     authorization: {

apps/webapp/app/routes/engine.v1.worker-actions.connect.ts

@@ -2,11 +2,12 @@ import type { TypedResponse } from "@remix-run/server-runtime";
 import { json } from "@remix-run/server-runtime";
 import type { WorkerApiConnectResponseBody } from "@trigger.dev/core/v3/workers";
 import { WorkerApiConnectRequestBody } from "@trigger.dev/core/v3/workers";
+import { z } from "zod";
 import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 export const action = createActionWorkerApiRoute(
   {
-    body: WorkerApiConnectRequestBody,
+    body: z.compile(WorkerApiConnectRequestBody),
   },
   async ({ authenticatedWorker, body }): Promise<TypedResponse<WorkerApiConnectResponseBody>> => {
     await authenticatedWorker.connect(body.metadata);

apps/webapp/app/routes/engine.v1.worker-actions.dequeue.ts

@@ -2,11 +2,12 @@ import type { TypedResponse } from "@remix-run/server-runtime";
 import { json } from "@remix-run/server-runtime";
 import type { WorkerApiDequeueResponseBody } from "@trigger.dev/core/v3/workers";
 import { WorkerApiDequeueRequestBody } from "@trigger.dev/core/v3/workers";
+import { z } from "zod";
 import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 export const action = createActionWorkerApiRoute(
   {
-    body: WorkerApiDequeueRequestBody,
+    body: z.compile(WorkerApiDequeueRequestBody),
   },
   async ({
     authenticatedWorker,

apps/webapp/app/routes/engine.v1.worker-actions.heartbeat.ts

@@ -2,11 +2,12 @@ import type { TypedResponse } from "@remix-run/server-runtime";
 import { json } from "@remix-run/server-runtime";
 import type { WorkerApiHeartbeatResponseBody } from "@trigger.dev/core/v3/workers";
 import { WorkerApiHeartbeatRequestBody } from "@trigger.dev/core/v3/workers";
+import { z } from "zod";
 import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 export const action = createActionWorkerApiRoute(
   {
-    body: WorkerApiHeartbeatRequestBody,
+    body: z.compile(WorkerApiHeartbeatRequestBody),
   },
   async ({ authenticatedWorker }): Promise<TypedResponse<WorkerApiHeartbeatResponseBody>> => {
     await authenticatedWorker.heartbeatWorkerInstance();

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.logs.debug.ts

@@ -10,7 +10,7 @@ export const action = createActionWorkerApiRoute(
     params: z.object({
       runFriendlyId: z.string(),
     }),
-    body: WorkerApiDebugLogBody,
+    body: z.compile(WorkerApiDebugLogBody),
   },
   async ({ body, params }): Promise<Response> => {
     const { runFriendlyId } = params;