feat: add runtime-aware restore routing and supervisor subscriptions

Adds opt-in versioned worker-queue assignment and routes checkpoint
resumes to queues matching the saved runtime, preserving the run's
dispatch class, region and channel. Legacy routing remains available,
and task retries that start fresh return to the original assignment.

Mono-RevId: ffc26bd998a3409f27f5011bb11a8bc2b2919db6

Author: carderne

apps/supervisor/src/env.test.ts

@@ -21,6 +21,104 @@ const base = {
   OTEL_EXPORTER_OTLP_ENDPOINT: "http://localhost:4318",
 };
 
+describe("worker queue selection", () => {
+  const ondemand = {
+    class: "ondemand",
+    phase: "fresh",
+    compat: "container",
+    channel: "stable",
+  };
+  const restore = {
+    class: "ondemand",
+    phase: "restore",
+    compat: "container",
+    channel: "canary",
+  };
+
+  it("keeps legacy default and scheduled selection when subscriptions are absent", () => {
+    expect(Env.parse(base).TRIGGER_WORKER_QUEUE_CLASS).toBe("default");
+    expect(
+      Env.parse({ ...base, TRIGGER_WORKER_QUEUE_CLASS: "scheduled" }).TRIGGER_WORKER_QUEUE_CLASS
+    ).toBe("scheduled");
+  });
+
+  it("parses multiple subscriptions without adding a legacy queue class", () => {
+    const subscriptions = [{ ...ondemand, weight: 0.25 }, restore];
+    const parsed = Env.parse({
+      ...base,
+      TRIGGER_CHECKPOINT_URL: "http://localhost:8089",
+      TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS: JSON.stringify(subscriptions),
+    });
+    expect(parsed.TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS).toEqual(subscriptions);
+    expect(parsed.TRIGGER_WORKER_QUEUE_CLASS).toBeUndefined();
+  });
+
+  it.each(["default", "scheduled"])(
+    "rejects subscriptions with explicit %s selection",
+    (queueClass) => {
+      expect(() =>
+        Env.parse({
+          ...base,
+          TRIGGER_WORKER_QUEUE_CLASS: queueClass,
+          TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS: JSON.stringify([ondemand]),
+        })
+      ).toThrow("mutually exclusive");
+    }
+  );
+
+  it.each(["not-json", "[]", JSON.stringify([{ ...restore, phase: "unknown" }])])(
+    "rejects invalid subscriptions: %s",
+    (subscriptions) => {
+      expect(() =>
+        Env.parse({ ...base, TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS: subscriptions })
+      ).toThrow();
+    }
+  );
+
+  it.each([
+    { config: {}, subscription: { ...ondemand, compat: "compute" }, error: "compatibility" },
+    {
+      config: { COMPUTE_GATEWAY_URL: "http://localhost:8080" },
+      subscription: ondemand,
+      error: "compatibility",
+    },
+    { config: {}, subscription: restore, error: "TRIGGER_CHECKPOINT_URL" },
+    {
+      config: {
+        KUBERNETES_FORCE_ENABLED: "true",
+        KUBERNETES_RUN_CRD_ENABLED: "true",
+        KUBERNETES_RUNNER_RUNTIME: "microvm",
+      },
+      subscription: { ...restore, compat: "compute" },
+      error: "COMPUTE_GATEWAY_URL",
+    },
+  ])(
+    "rejects incompatible or unsupported subscriptions: $error",
+    ({ config, subscription, error }) => {
+      expect(() =>
+        Env.parse({
+          ...base,
+          ...config,
+          TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS: JSON.stringify([subscription]),
+        })
+      ).toThrow(error);
+    }
+  );
+
+  it("accepts shared fresh work and compute restores on the compute backend", () => {
+    expect(() =>
+      Env.parse({
+        ...base,
+        COMPUTE_GATEWAY_URL: "http://localhost:8080",
+        TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS: JSON.stringify([
+          { ...ondemand, compat: "any" },
+          { ...restore, compat: "compute" },
+        ]),
+      })
+    ).not.toThrow();
+  });
+});
+
 describe("Env superRefine - backpressure source awareness", () => {
   it("pod-count source can be enabled without a Redis host", () => {
     expect(() =>

apps/supervisor/src/env.ts

@@ -1,6 +1,7 @@
 import { randomUUID } from "crypto";
 import { env as stdEnv } from "std-env";
 import { z } from "zod";
+import { WeightedWorkerQueueSubscriptions } from "@trigger.dev/core/v3/runEngineWorker";
 import {
   AdditionalEnvVars,
   BoolEnv,
@@ -56,7 +57,23 @@ export const Env = z
     // Which worker-queue class this supervisor fleet serves. "default" pulls the
     // region queue (standard/agent runs); "scheduled" pulls the dedicated
     // scheduled-lineage queue. Run a separate fleet per class for isolation.
-    TRIGGER_WORKER_QUEUE_CLASS: z.enum(["default", "scheduled"]).default("default"),
+    TRIGGER_WORKER_QUEUE_CLASS: z.enum(["default", "scheduled"]).optional(),
+    // JSON array of weighted v2 subscriptions; omit to retain legacy queue selection.
+    TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS: z
+      .string()
+      .transform((value, ctx) => {
+        try {
+          return JSON.parse(value) as unknown;
+        } catch {
+          ctx.addIssue({
+            code: z.ZodIssueCode.custom,
+            message: "Worker queue subscriptions must be valid JSON",
+          });
+          return z.NEVER;
+        }
+      })
+      .pipe(WeightedWorkerQueueSubscriptions)
+      .optional(),
     TRIGGER_DEQUEUE_INTERVAL_MS: z.coerce.number().int().default(250),
     TRIGGER_DEQUEUE_IDLE_INTERVAL_MS: z.coerce.number().int().default(1000),
     TRIGGER_DEQUEUE_MAX_RUN_COUNT: z.coerce.number().int().default(1),
@@ -326,6 +343,44 @@ export const Env = z
     TRIGGER_WIDE_EVENTS_NOISY_ROUTES: BoolEnv.default(false),
   })
   .superRefine((data, ctx) => {
+    if (data.TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS) {
+      if (data.TRIGGER_WORKER_QUEUE_CLASS !== undefined) {
+        ctx.addIssue({
+          code: z.ZodIssueCode.custom,
+          message:
+            "TRIGGER_WORKER_QUEUE_CLASS and TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS are mutually exclusive",
+          path: ["TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS"],
+        });
+      }
+
+      const compatibility =
+        data.COMPUTE_GATEWAY_URL ||
+        (data.KUBERNETES_FORCE_ENABLED &&
+          data.KUBERNETES_RUN_CRD_ENABLED &&
+          data.KUBERNETES_RUNNER_RUNTIME === "microvm")
+          ? "compute"
+          : "container";
+
+      for (const [index, subscription] of data.TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS.entries()) {
+        if (subscription.compat !== "any" && subscription.compat !== compatibility) {
+          ctx.addIssue({
+            code: z.ZodIssueCode.custom,
+            message: `Subscription compatibility must be any or ${compatibility} for this backend`,
+            path: ["TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS", index, "compat"],
+          });
+        } else if (
+          subscription.phase === "restore" &&
+          !(compatibility === "compute" ? data.COMPUTE_GATEWAY_URL : data.TRIGGER_CHECKPOINT_URL)
+        ) {
+          ctx.addIssue({
+            code: z.ZodIssueCode.custom,
+            message: `Restore subscriptions require ${compatibility === "compute" ? "COMPUTE_GATEWAY_URL" : "TRIGGER_CHECKPOINT_URL"}`,
+            path: ["TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS", index],
+          });
+        }
+      }
+    }
+
     if (
       data.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENABLED &&
       data.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE >=
@@ -404,6 +459,9 @@ export const Env = z
   })
   .transform((data) => ({
     ...data,
+    TRIGGER_WORKER_QUEUE_CLASS:
+      data.TRIGGER_WORKER_QUEUE_CLASS ??
+      (data.TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS === undefined ? "default" : undefined),
     COMPUTE_TRACE_OTLP_ENDPOINT: data.COMPUTE_TRACE_OTLP_ENDPOINT ?? `${data.TRIGGER_API_URL}/otel`,
   }));
 

apps/supervisor/src/index.ts

@@ -349,6 +349,7 @@ class ManagedSupervisor {
       queueConsumerEnabled: env.TRIGGER_DEQUEUE_ENABLED,
       maxRunCount: env.TRIGGER_DEQUEUE_MAX_RUN_COUNT,
       queueClass: env.TRIGGER_WORKER_QUEUE_CLASS,
+      subscriptions: env.TRIGGER_WORKER_QUEUE_SUBSCRIPTIONS,
       metricsRegistry: register,
       scaling: {
         strategy: env.TRIGGER_DEQUEUE_SCALING_STRATEGY,

apps/webapp/app/env.server.ts

@@ -1562,6 +1562,8 @@ const EnvironmentSchema = z
     // this is "1"; an org set to true splits even when this is "0"). Never
     // applies to DEVELOPMENT environments.
     TRIGGER_WORKER_QUEUE_SCHEDULED_SPLIT_ENABLED: z.string().default("0"),
+    // Controls new assignments only; existing v2 runs retain their routing.
+    TRIGGER_WORKER_QUEUE_V2_ENABLED: z.string().default("0"),
 
     TRIGGER_MOLLIFIER_ENABLED: z.string().default("0"),
     // Separate switch for the drainer (consumer side) so it can be split

apps/webapp/app/presenters/v3/SpanPresenter.server.ts

@@ -11,6 +11,7 @@ import {
 } from "@trigger.dev/core/v3";
 
 import { AttemptId, getMaxDuration, parseTraceparent } from "@trigger.dev/core/v3/isomorphic";
+import { isV2WorkerQueue } from "@trigger.dev/core/v3/workers";
 import {
   runOpsLegacyReplica,
   runOpsNewReplica,
@@ -40,6 +41,7 @@ import { buildSyntheticSpanRun } from "~/v3/mollifier/syntheticSpanRun.server";
 import { engine } from "~/v3/runEngine.server";
 import { runStore } from "~/v3/runStore.server";
 import { runTriggeredAt } from "~/v3/runTimestamps";
+import { workerRegionRegistry } from "~/v3/workerRegions.server";
 import { getTaskEventStoreTableForRun, type TaskEventStoreTable } from "~/v3/taskEventStore.server";
 import { isFailedRunStatus, isFinalRunStatus } from "~/v3/taskStatus";
 import { BasePresenter } from "./basePresenter.server";
@@ -330,23 +332,30 @@ export class SpanPresenter extends BasePresenter {
     let region: { name: string; location: string | null } | null = null;
 
     if (environment.type !== "DEVELOPMENT" && run.engine !== "V1") {
+      const v2Region = isV2WorkerQueue(run.workerQueue)
+        ? (run.region ?? baseWorkerQueue(run.workerQueue))
+        : undefined;
+      // V2 identifies geography, not the executing group. Reuse the registry to
+      // retain the indexed masterQueue lookup for the region's location metadata.
+      const masterQueue = v2Region
+        ? (workerRegionRegistry.current()?.find((group) => group.region === v2Region)
+            ?.masterQueue ?? v2Region)
+        : baseWorkerQueue(run.workerQueue);
       const workerGroup = await this._replica.workerInstanceGroup.findFirst({
         select: {
           name: true,
           location: true,
         },
-        where: {
-          // masterQueue is unique and IS the run's backing queue, so this finds
-          // the group the run actually ran on.
-          masterQueue: baseWorkerQueue(run.workerQueue),
-        },
+        where: { masterQueue },
       });
 
       // Show the stamped geo region as the name so a migrated run never reveals
       // its compute backing; fall back to the group name for unstamped runs.
       region = workerGroup
-        ? { name: run.region ?? workerGroup.name, location: workerGroup.location }
-        : null;
+        ? { name: v2Region ?? run.region ?? workerGroup.name, location: workerGroup.location }
+        : v2Region
+          ? { name: v2Region, location: null }
+          : null;
     }
 
     // Only AGENT-tagged runs can be session-bound, so skip the SessionRun lookup

apps/webapp/app/runEngine/concerns/workerQueueAssignment.server.test.ts

@@ -0,0 +1,85 @@
+import { describe, expect, it } from "vitest";
+import { workerQueueForBirth } from "~/runEngine/concerns/workerQueueAssignment.server";
+import { FEATURE_FLAG } from "~/v3/featureFlags";
+
+const assignment = {
+  workerQueue: "us-east-1-next",
+  region: "us-east-1",
+  envType: "PRODUCTION",
+  orgFeatureFlags: null,
+  globalDefault: false,
+  workerGroups: [{ masterQueue: "us-east-1-next", workloadType: "MICROVM" as const }],
+  rootTriggerSource: "api",
+  splitEnabled: true,
+};
+
+describe("workerQueueForBirth", () => {
+  it("keeps legacy assignment and scheduled lineage when v2 is off", () => {
+    expect(workerQueueForBirth(assignment)).toBe("us-east-1-next");
+    expect(workerQueueForBirth({ ...assignment, rootTriggerSource: "schedule" })).toBe(
+      "us-east-1-next:scheduled"
+    );
+  });
+
+  it("keeps legacy routing when the selected queue is absent from a loaded registry", () => {
+    const input = {
+      ...assignment,
+      region: assignment.workerQueue,
+      globalDefault: true,
+      workerGroups: [{ masterQueue: "us-west-2", workloadType: "CONTAINER" as const }],
+    };
+    expect(workerQueueForBirth(input)).toBe("us-east-1-next");
+    expect(workerQueueForBirth({ ...input, rootTriggerSource: "schedule" })).toBe(
+      "us-east-1-next:scheduled"
+    );
+  });
+
+  it("uses the geographic region and selected runtime when v2 is enabled", () => {
+    expect(workerQueueForBirth({ ...assignment, globalDefault: true })).toBe(
+      "us-east-1:v2:ondemand:fresh:compute:stable"
+    );
+    expect(
+      workerQueueForBirth({
+        ...assignment,
+        globalDefault: true,
+        workerGroups: [{ masterQueue: "us-east-1-next", workloadType: "CONTAINER" }],
+      })
+    ).toBe("us-east-1:v2:ondemand:fresh:container:stable");
+  });
+
+  it("honors org opt-in, channel and explicit shared-runtime assignment", () => {
+    expect(
+      workerQueueForBirth({
+        ...assignment,
+        orgFeatureFlags: {
+          [FEATURE_FLAG.workerQueueV2Enabled]: true,
+          [FEATURE_FLAG.workerQueueCompatibility]: "any",
+          [FEATURE_FLAG.workerQueueChannel]: "canary",
+        },
+        rootTriggerSource: "schedule",
+      })
+    ).toBe("us-east-1:v2:scheduled:fresh:any:canary");
+  });
+
+  it("honors org opt-out even when the global default is enabled", () => {
+    expect(
+      workerQueueForBirth({
+        ...assignment,
+        globalDefault: true,
+        orgFeatureFlags: { [FEATURE_FLAG.workerQueueV2Enabled]: false },
+      })
+    ).toBe("us-east-1-next");
+  });
+
+  it("leaves development on the legacy path even with an org opt-in", () => {
+    expect(
+      workerQueueForBirth({
+        ...assignment,
+        envType: "DEVELOPMENT",
+        splitEnabled: false,
+        globalDefault: true,
+        orgFeatureFlags: { [FEATURE_FLAG.workerQueueV2Enabled]: true },
+      })
+    ).toBe("us-east-1-next");
+  });
+});