feat(core,webapp,run-engine): stamp a shard key onto run, batch and waitpoint ids

apps/webapp/app/models/waitpointTag.server.ts

@@ -9,13 +9,17 @@ export async function createWaitpointTag({
   environmentId,
   projectId,
   residency,
+  shardKey,
 }: {
   tag: string;
   environmentId: string;
   projectId: string;
   // Residency from the env mint kind: a tag has no owning run, so a minted-new env pins it to NEW
   // instead of defaulting to the draining legacy DB.
   residency?: "NEW" | "LEGACY";
+  // The environment's gen-2 mint shard, when it has one. A tag has no id the router can read, so
+  // without this the row lands on a gen-1 store while the token it describes lands on the shard.
+  shardKey?: string;
 }) {
   if (tag.trim().length === 0) return;
 
@@ -30,7 +34,8 @@ export async function createWaitpointTag({
           projectId,
         },
         undefined,
-        residency
+        residency,
+        shardKey
       );
     } catch (error) {
       if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") {

apps/webapp/app/routes/api.v1.waitpoints.tokens.ts

@@ -16,6 +16,7 @@ import {
   type PrismaClientOrTransaction,
 } from "~/db.server";
 import { resolveRunIdMintKind } from "~/v3/engineVersion.server";
+import { resolveMintShard } from "~/v3/runOpsMigration/runOpsMintShard.server";
 import { logger } from "~/services/logger.server";
 import { generateHttpCallbackUrl } from "~/services/httpCallback.server";
 import { publicAccessTokenResponseHeaders } from "~/services/publicAccessTokenResponse.server";
@@ -69,6 +70,15 @@ const { action } = createActionApiRoute(
       });
       const residency = mintKind === "runOpsId" ? "NEW" : "LEGACY";
 
+      // No extra query: the org flags are already loaded on the authenticated env.
+      const standaloneShardKey =
+        mintKind === "runOpsId"
+          ? await resolveMintShard({
+              id: authentication.environment.id,
+              orgFeatureFlags: authentication.environment.organization.featureFlags,
+            })
+          : undefined;
+
       //upsert tags
       let tags: { id: string; name: string }[] = [];
       const bodyTags = typeof body.tags === "string" ? [body.tags] : body.tags;
@@ -86,6 +96,7 @@ const { action } = createActionApiRoute(
             environmentId: authentication.environment.id,
             projectId: authentication.environment.projectId,
             residency,
+            shardKey: standaloneShardKey,
           });
           if (tagRecord) {
             tags.push(tagRecord);
@@ -101,6 +112,7 @@ const { action } = createActionApiRoute(
         timeout,
         tags: bodyTags,
         standaloneResidency: residency,
+        standaloneShardKey,
       });
 
       const waitpointId = WaitpointId.toFriendlyId(result.waitpoint.id);

apps/webapp/app/runEngine/services/triggerFailedTask.server.ts

@@ -1,15 +1,15 @@
 import type { RunEngine } from "@internal/run-engine";
 import { TaskRunErrorCodes, type TaskRunError } from "@trigger.dev/core/v3";
-import { RunId, generateRunOpsId } from "@trigger.dev/core/v3/isomorphic";
+import { RunId } from "@trigger.dev/core/v3/isomorphic";
 import type {
   PrismaClientOrTransaction,
   RuntimeEnvironmentType,
   TaskRun,
 } from "@trigger.dev/database";
 import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
-import { resolveRunIdMintKind } from "~/v3/engineVersion.server";
-import { resolveInheritedMintKind } from "~/v3/runOpsMigration/resolveInheritedMintKind.server";
+import { mintFriendlyIdForKind } from "~/v3/runOpsMigration/mintAnchoredRunFriendlyId.server";
+import { resolveRunMintTarget } from "~/v3/runOpsMigration/resolveRunMintTarget.server";
 import { getEventRepository } from "~/v3/eventRepository/index.server";
 import { runStore as defaultRunStore } from "~/v3/runStore.server";
 import type { RunStore } from "@internal/run-store";
@@ -103,17 +103,16 @@ export class TriggerFailedTaskService {
       return args.runFriendlyId;
     }
 
-    const mintKind = args.parentRunFriendlyId
-      ? resolveInheritedMintKind(args.parentRunFriendlyId)
-      : await resolveRunIdMintKind({
+    return mintFriendlyIdForKind(
+      await resolveRunMintTarget({
+        environment: {
           organizationId: args.organizationId,
           id: args.environmentId,
           orgFeatureFlags: args.orgFeatureFlags,
-        });
-
-    return mintKind === "runOpsId"
-      ? RunId.toFriendlyId(generateRunOpsId())
-      : RunId.generate().friendlyId;
+        },
+        parentRunFriendlyId: args.parentRunFriendlyId,
+      })
+    );
   }
 
   async call(request: TriggerFailedTaskRequest): Promise<string | null> {

apps/webapp/app/runEngine/services/triggerTask.server.ts

@@ -28,9 +28,8 @@ import { parseDelay } from "~/utils/delays";
 import { removeNullBytesFromKey } from "~/utils/nullBytes";
 import { handleMetadataPacket } from "~/utils/packets";
 import { startSpan } from "~/v3/tracing.server";
-import { resolveRunIdMintKind } from "~/v3/engineVersion.server";
-import { resolveInheritedMintKind } from "~/v3/runOpsMigration/resolveInheritedMintKind.server";
 import { mintFriendlyIdForKind } from "~/v3/runOpsMigration/mintAnchoredRunFriendlyId.server";
+import { resolveRunMintTarget } from "~/v3/runOpsMigration/resolveRunMintTarget.server";
 import type {
   TriggerTaskServiceOptions,
   TriggerTaskServiceResult,
@@ -218,15 +217,17 @@ export class RunEngineTriggerTaskService {
     parentRunFriendlyId?: string,
     region?: string
   ): Promise<string> {
-    const mintKind = parentRunFriendlyId
-      ? resolveInheritedMintKind(parentRunFriendlyId)
-      : await resolveRunIdMintKind({
+    return mintFriendlyIdForKind(
+      await resolveRunMintTarget({
+        environment: {
           organizationId: environment.organizationId,
           id: environment.id,
           orgFeatureFlags: environment.organization.featureFlags,
-        });
-
-    return mintFriendlyIdForKind(mintKind, region);
+        },
+        parentRunFriendlyId,
+        region,
+      })
+    );
   }
 
   public async call({

apps/webapp/app/v3/runEngineHandlers.server.ts

@@ -11,6 +11,7 @@ import {
   runOpsNewPrismaClient,
   runOpsNewReplicaClient,
   runOpsLegacyPrismaClient,
+  runOpsShardHandles,
 } from "~/db.server";
 import { env } from "~/env.server";
 import { findEnvironmentById, findEnvironmentFromRun } from "~/models/runtimeEnvironment.server";
@@ -1060,6 +1061,7 @@ export function setupBatchQueueCallbacks() {
       newReplica: runOpsNewReplicaClient,
       newWriter: runOpsNewPrismaClient,
       legacyWriter: runOpsLegacyPrismaClient,
+      shards: runOpsShardHandles,
       tryCompleteBatch: (batchId) => engine.tryCompleteBatch({ batchId }),
     });
   });

apps/webapp/app/v3/runEngineHandlersShared.server.ts

@@ -4,6 +4,7 @@
  * whole webapp service graph). The handlers wire the production defaults; tests
  * inject per-container stores/replicas, so these helpers never import db.server.
  */
+import { resolveShard } from "@trigger.dev/core/v3/isomorphic";
 import type { CompleteBatchResult } from "@internal/run-engine";
 import type { RunOpsPrismaClient } from "@internal/run-ops-database";
 import type { RunStore } from "@internal/run-store";
@@ -83,8 +84,23 @@ export async function resolveBatchRunOpsWriter(
     newReplica: RunOpsPrismaClient;
     newWriter: RunOpsPrismaClient;
     legacyWriter: RunOpsPrismaClient;
+    shards?: ReadonlyArray<{ key: string; writer: RunOpsPrismaClient }>;
   }
 ): Promise<RunOpsPrismaClient> {
+  // The probe below is binary, so without this a gen-2 batch resolves to a store holding no such
+  // row, and the update throws before the batch waitpoint completes.
+  const shardKey = resolveShard(batchId);
+  if (shardKey !== "new" && shardKey !== "legacy") {
+    const shard = deps.shards?.find((s) => s.key === shardKey);
+    if (!shard) {
+      // Writing to a guessed store is what strands a run. Fail loud instead.
+      throw new Error(
+        `resolveBatchRunOpsWriter: batch "${batchId}" names shard "${shardKey}", which is not configured`
+      );
+    }
+    return shard.writer;
+  }
+
   const onNew = await deps.newReplica.batchTaskRun.findFirst({
     where: { id: batchId },
     select: { id: true },
@@ -106,6 +122,7 @@ export type BatchCompletionDeps = {
   newReplica: RunOpsPrismaClient;
   newWriter: RunOpsPrismaClient;
   legacyWriter: RunOpsPrismaClient;
+  shards?: ReadonlyArray<{ key: string; writer: RunOpsPrismaClient }>;
   tryCompleteBatch: (batchId: string) => Promise<unknown>;
 };
 
@@ -136,6 +153,7 @@ export async function handleBatchCompletion(
     newReplica: deps.newReplica,
     newWriter: deps.newWriter,
     legacyWriter: deps.legacyWriter,
+    shards: deps.shards,
   });
 
   try {

apps/webapp/app/v3/runOpsMigration/gen2MintInertness.test.ts

@@ -0,0 +1,112 @@
+import { describe, expect, it, vi } from "vitest";
+import { classifyKind, mintWaitpointIdFor, resolveShard } from "@trigger.dev/core/v3/isomorphic";
+import { resolveInheritedMintKind } from "./resolveInheritedMintKind.server";
+import {
+  mintAnchoredRunFriendlyId,
+  mintFriendlyIdForKind,
+} from "./mintAnchoredRunFriendlyId.server";
+import { batchIdForMintKind } from "./mintBatchFriendlyId.server";
+import { resolveRunMintTarget } from "./resolveRunMintTarget.server";
+
+// Gate off means resolveMintShard answers "new". Every assertion is "the id is what it was".
+const offShard = vi.fn().mockResolvedValue("new" as const);
+const environment = { organizationId: "org_1", id: "env_1", orgFeatureFlags: {} };
+
+describe("gate off — run mint paths", () => {
+  it("a root run on the run-ops path mints a gen-1 v1 id", async () => {
+    const target = await resolveRunMintTarget({
+      environment,
+      region: "us-east-1",
+      deps: {
+        resolveRunIdMintKind: vi.fn().mockResolvedValue("runOpsId"),
+        resolveMintShard: offShard,
+      },
+    });
+    const body = mintFriendlyIdForKind(target).slice(4);
+    expect(body.length).toBe(26);
+    expect(body[24]).toBe("e"); // the region char, as today
+    expect(body[25]).toBe("1");
+  });
+
+  it("a root run on a non-cut-over org mints a cuid", async () => {
+    const target = await resolveRunMintTarget({
+      environment,
+      deps: {
+        resolveRunIdMintKind: vi.fn().mockResolvedValue("cuid"),
+        resolveMintShard: offShard,
+      },
+    });
+    expect(mintFriendlyIdForKind(target).slice(4).length).toBe(25);
+  });
+
+  it("a child of a gen-1 parent keeps the caller's region char", async () => {
+    // The pre-split code passed the region on both arms; dropping it on the inherited arm would
+    // silently stamp the default.
+    const target = await resolveRunMintTarget({
+      environment,
+      parentRunFriendlyId: `run_${"a".repeat(24)}01`,
+      region: "us-east-1",
+      deps: {
+        resolveRunIdMintKind: vi.fn().mockResolvedValue("runOpsId"),
+        resolveMintShard: offShard,
+      },
+    });
+    const body = mintFriendlyIdForKind(target).slice(4);
+    expect(body[24]).toBe("e");
+    expect(body[25]).toBe("1");
+  });
+
+  it("a gen-2 parent's shard still outranks the caller's region", async () => {
+    const target = await resolveRunMintTarget({
+      environment,
+      parentRunFriendlyId: `run_${"a".repeat(24)}a2`,
+      region: "us-east-1",
+      deps: {
+        resolveRunIdMintKind: vi.fn().mockResolvedValue("runOpsId"),
+        resolveMintShard: offShard,
+      },
+    });
+    expect(mintFriendlyIdForKind(target).slice(4)[24]).toBe("a");
+  });
+
+  it("a child of a gen-1 parent mints a gen-1 v1 id", () => {
+    const body = mintFriendlyIdForKind(resolveInheritedMintKind(`run_${"a".repeat(24)}01`)).slice(
+      4
+    );
+    expect(body[25]).toBe("1");
+  });
+
+  it("a child of a cuid parent mints a cuid", () => {
+    expect(
+      mintFriendlyIdForKind(resolveInheritedMintKind(`run_${"b".repeat(25)}`)).slice(4).length
+    ).toBe(25);
+  });
+});
+
+describe("gate off — batch and item paths", () => {
+  it("a batch with no shard char mints a gen-1 v1 id", () => {
+    const r = batchIdForMintKind({ kind: "runOpsId" });
+    expect(r.id.length).toBe(26);
+    expect(r.id[25]).toBe("1");
+    expect(classifyKind(r.id)).toBe("runOpsId");
+  });
+
+  it("a batch on a non-cut-over org mints a cuid", () => {
+    expect(batchIdForMintKind({ kind: "cuid" }).id.length).toBe(25);
+  });
+
+  it("a batch item anchored on a gen-1 batch mints a gen-1 v1 id", () => {
+    const body = mintAnchoredRunFriendlyId(`batch_${"a".repeat(24)}01`).slice(4);
+    expect(body[25]).toBe("1");
+  });
+});
+
+describe("gate off — waitpoint paths", () => {
+  it("every gen-1 or legacy anchor yields a cuid waitpoint id", () => {
+    for (const anchor of [`${"a".repeat(24)}01`, "c".repeat(25), undefined]) {
+      const r = mintWaitpointIdFor(anchor);
+      expect(r.id.length).toBe(25);
+      expect(resolveShard(r.id)).toBe("legacy");
+    }
+  });
+});

apps/webapp/app/v3/runOpsMigration/mintAnchoredRunFriendlyId.server.test.ts

@@ -28,4 +28,16 @@ describe("mintAnchoredRunFriendlyId", () => {
     expect(parsed.format).toBe("b32hex");
     expect(parsed.format === "b32hex" && parsed.region).toBe(REGION_CODES["us-east-1"]);
   });
+
+  it("a gen-2 batch anchor mints an item on the batch's shard", () => {
+    const body = mintAnchoredRunFriendlyId(`batch_${"a".repeat(24)}a2`).slice("run_".length);
+    expect(body).toHaveLength(26);
+    expect(body[24]).toBe("a");
+    expect(body[25]).toBe("2");
+  });
+
+  it("a gen-2 batch anchor ignores a caller region: the shard owns index 24", () => {
+    const body = mintAnchoredRunFriendlyId(`batch_${"a".repeat(24)}a2`, "us-east-1").slice(4);
+    expect(body[24]).toBe("a");
+  });
 });

apps/webapp/app/v3/runOpsMigration/mintAnchoredRunFriendlyId.server.ts

@@ -1,15 +1,20 @@
-import { generateRunOpsId, RunId, type ResidencyKind } from "@trigger.dev/core/v3/isomorphic";
+import { generateRunOpsId, generateRunOpsIdV2, RunId } from "@trigger.dev/core/v3/isomorphic";
+import type { MintTarget } from "./mintTarget";
 import { resolveInheritedMintKind } from "./resolveInheritedMintKind.server";
 
-// Shared id-generation branch for every run-mint path: "runOpsId" -> NEW store, "cuid" -> LEGACY.
-export function mintFriendlyIdForKind(mintKind: ResidencyKind, region?: string): string {
-  return mintKind === "runOpsId"
-    ? RunId.toFriendlyId(generateRunOpsId(region))
-    : RunId.generate().friendlyId;
+// A shardChar selects one gen-2 shard and takes index 24; without one the region takes that slot.
+export function mintFriendlyIdForKind(target: MintTarget): string {
+  if (target.kind !== "runOpsId") {
+    return RunId.generate().friendlyId;
+  }
+
+  return RunId.toFriendlyId(
+    target.shardChar ? generateRunOpsIdV2(target.shardChar) : generateRunOpsId(target.region)
+  );
 }
 
 // Anchor a batch item's mint on the BATCH's friendlyId (id-shape, zero I/O), never the per-org
 // flag, so the item and its BatchTaskRun stay co-resident across a mid-batch flag flip.
 export function mintAnchoredRunFriendlyId(batchFriendlyId: string, region?: string): string {
-  return mintFriendlyIdForKind(resolveInheritedMintKind(batchFriendlyId), region);
+  return mintFriendlyIdForKind({ ...resolveInheritedMintKind(batchFriendlyId), region });
 }