feat(webapp,run-store): gen-2 shard arms in read-through and idempotency

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

@@ -1,5 +1,5 @@
 import type { BatchTaskRunExecutionResult } from "@trigger.dev/core/v3";
-import { ownerEngine } from "@trigger.dev/core/v3/isomorphic";
+import { resolveShard, type ShardKey } from "@trigger.dev/core/v3/isomorphic";
 import {
   $replica,
   type PrismaClientOrTransaction,
@@ -13,6 +13,8 @@ import { runStore as defaultRunStore } from "~/v3/runStore.server";
 import { BasePresenter } from "./basePresenter.server";
 
 import { boundedIn } from "@trigger.dev/database";
+import { runOpsShardReplicas } from "~/v3/runOpsMigration/shardHandles.server";
+import { logger } from "~/services/logger.server";
 /**
  * Run-ops read-through wiring. All optional; absent (or `splitEnabled` falsy) collapses `call` to
  * passthrough. `legacyReplica` is a READ REPLICA handle only — there is NO legacy-primary field.
@@ -21,6 +23,8 @@ type ApiBatchResultsReadThroughDeps = {
   splitEnabled?: boolean;
   newClient?: PrismaReplicaClient;
   legacyReplica?: PrismaReplicaClient;
+  /** Gen-2 shard replicas by shard char; empty (RUN_OPS_SHARDS unset) keeps today's behaviour. */
+  shardReplicas?: ReadonlyMap<ShardKey, PrismaReplicaClient>;
   isPastRetention?: (runId: string) => boolean;
 };
 
@@ -181,16 +185,57 @@ export class ApiBatchResultsPresenter extends BasePresenter {
 
     const taskRunIds = batchRun.items.map((item) => item.taskRunId);
 
-    const newRows = (await newClient.taskRun.findMany({
-      where: { id: { in: boundedIn(taskRunIds) } },
-      select: memberRunSelect,
-    })) as TaskRunWithAttempts[];
+    // A gen-2 id is directly routable to its own shard, so it must not join the gen-1 read:
+    // it would miss there, and (being dedicated-family) never reach the legacy probe either.
+    const shardReplicas = this.readThrough?.shardReplicas ?? runOpsShardReplicas;
+    const genOneIds: string[] = [];
+    const idsByShard = new Map<ShardKey, string[]>();
+    for (const id of taskRunIds) {
+      const shardKey = resolveShard(id);
+      if (shardKey === "new" || shardKey === "legacy") {
+        genOneIds.push(id);
+      } else if (shardReplicas.has(shardKey)) {
+        const group = idsByShard.get(shardKey);
+        group ? group.push(id) : idsByShard.set(shardKey, [id]);
+      } else {
+        // Not routable and not a gen-1 shape. A gen-1 store is the wrong database, and a
+        // dedicated-family id never reaches the legacy probe, so falling back there would
+        // drop the member silently. Drop it loudly instead.
+        logger.error("ApiBatchResultsPresenter: gen-2 member on an unconfigured shard key", {
+          runId: id,
+          shardKey,
+          configured: [...shardReplicas.keys()],
+        });
+      }
+    }
+
+    const newRows = (
+      genOneIds.length > 0
+        ? ((await newClient.taskRun.findMany({
+            where: { id: { in: boundedIn(genOneIds) } },
+            select: memberRunSelect,
+          })) as TaskRunWithAttempts[])
+        : []
+    ).concat(
+      (
+        await Promise.all(
+          [...idsByShard.entries()].map(
+            async ([shardKey, ids]) =>
+              (await shardReplicas.get(shardKey)!.taskRun.findMany({
+                where: { id: { in: boundedIn(ids) } },
+                select: memberRunSelect,
+              })) as TaskRunWithAttempts[]
+          )
+        )
+      ).flat()
+    );
     const runsById = new Map(newRows.map((run) => [run.id, run]));
 
-    // A run-ops id can only live on NEW, so only misses that AREN'T run-ops-shaped are candidates
-    // for the legacy probe — mirrors readThroughRun's per-id "NEW residency skips legacy" rule.
-    const legacyCandidateIds = taskRunIds.filter(
-      (id) => !runsById.has(id) && ownerEngine(id) !== "NEW"
+    // A dedicated-family id (gen-1 v1 or gen-2) can only live on its own store, so only
+    // misses that AREN'T dedicated-shaped are candidates for the legacy probe — mirrors
+    // readThroughRun's per-id "dedicated residency skips legacy" rule.
+    const legacyCandidateIds = genOneIds.filter(
+      (id) => !runsById.has(id) && resolveShard(id) === "legacy"
     );
     if (legacyCandidateIds.length > 0) {
       const legacyRows = (await legacyReplica.taskRun.findMany({

apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts

@@ -1,4 +1,4 @@
-import { ownerEngine, RunId } from "@trigger.dev/core/v3/isomorphic";
+import { resolveShard, RunId, type ShardKey } from "@trigger.dev/core/v3/isomorphic";
 import type { PrismaClientOrTransaction, TaskRun, Waitpoint } from "@trigger.dev/database";
 import { env } from "~/env.server";
 import { logger } from "~/services/logger.server";
@@ -13,9 +13,10 @@ import { computeClaimTtlSeconds } from "~/v3/mollifier/claimTtl";
 import { makeResolveMollifierFlag } from "~/v3/mollifier/mollifierGate.server";
 import { runStore } from "~/v3/runStore.server";
 import { runOpsLegacyPrisma, runOpsNewPrisma } from "~/db.server";
+import { runOpsShardWriters } from "~/v3/runOpsMigration/shardHandles.server";
 import { isSplitEnabled } from "~/v3/runOpsMigration/splitMode.server";
 import { resolveRunIdMintKind } from "~/v3/engineVersion.server";
-import { resolveIdempotencyDedupClient } from "./idempotencyResidency.server";
+import { clientForShardKey, resolveIdempotencyDedupClient } from "./idempotencyResidency.server";
 import type { TraceEventConcern, TriggerTaskRequest } from "../types";
 
 // In-memory per-org mollifier-enabled check, shared with `evaluateGate`
@@ -32,6 +33,24 @@ const resolveOrgMollifierFlag = makeResolveMollifierFlag();
 // PG's unique index as the backstop.
 const MAX_CLEARED_WINNER_REACQUIRES = 5;
 
+// Every run-ops store keyed by shard key. Both idempotency call sites resolve through this
+// one map, so they cannot disagree about which store owns an id.
+//
+// Built on first use, not at import: dereferencing the db.server handles at module scope
+// breaks any test that mocks `~/db.server` without them, and this module is imported by
+// triggerTask. Memoised because the trigger path is the hottest in the system.
+let cachedShardClients: ReadonlyMap<ShardKey, PrismaClientOrTransaction> | undefined;
+
+function idempotencyShardClients(): ReadonlyMap<ShardKey, PrismaClientOrTransaction> {
+  return (cachedShardClients ??= new Map<ShardKey, PrismaClientOrTransaction>([
+    ["legacy", runOpsLegacyPrisma],
+    ["new", runOpsNewPrisma],
+    ...[...runOpsShardWriters.entries()].map(
+      ([key, writer]) => [key, writer as PrismaClientOrTransaction] as const
+    ),
+  ]));
+}
+
 // Claim ownership context returned to the caller when the
 // IdempotencyKeyConcern won a pre-gate claim. Caller MUST publish the
 // winning runId on pipeline success (`publishClaim`) or release the
@@ -172,12 +191,9 @@ export class IdempotencyKeyConcern {
       {
         isSplitEnabled,
         fallbackClient: this.prisma,
-        newClient: runOpsNewPrisma,
-        legacyClient: runOpsLegacyPrisma,
+        clients: idempotencyShardClients(),
         resolveMintKind: resolveRunIdMintKind,
-        // `isMigrated` is intentionally omitted: until a child of a swept
-        // legacy-id parent can be born on the new DB, the swept-marker override
-        // would never change the answer, so a child routes by parent id-shape.
+        logger,
       }
     );
 
@@ -640,12 +656,15 @@ export class IdempotencyKeyConcern {
     } catch {
       return null;
     }
-    let client: PrismaClientOrTransaction;
-    try {
-      client = ownerEngine(internalId) === "NEW" ? runOpsNewPrisma : runOpsLegacyPrisma;
-    } catch {
-      client = this.prisma;
-    }
+    // The routing store routes by id and never forwards this object, so its identity only
+    // signals read-your-writes. Resolving it through the shard map keeps the two idempotency
+    // call sites in agreement and stops this reading as gen-2-unaware.
+    const client = clientForShardKey(
+      resolveShard(internalId),
+      idempotencyShardClients(),
+      this.prisma,
+      logger
+    );
     return runStore.findRun(
       { id: internalId, runtimeEnvironmentId: environmentId },
       { include: { associatedWaitpoint: true } },

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

@@ -1,6 +1,7 @@
 import { describe, expect, it } from "vitest";
 import { RunId } from "@trigger.dev/core/v3/isomorphic";
 import {
+  clientForShardKey,
   resolveIdempotencyDedupClient,
   type ResolveIdempotencyClientDeps,
 } from "./idempotencyResidency.server";
@@ -9,20 +10,30 @@ import {
 const FALLBACK = { __tag: "fallback" } as never;
 const NEW_CLIENT = { __tag: "new" } as never;
 const LEGACY_CLIENT = { __tag: "legacy" } as never;
+const SHARD_A_CLIENT = { __tag: "shard-a" } as never;
+
+function clientMap() {
+  return new Map([
+    ["new", NEW_CLIENT],
+    ["legacy", LEGACY_CLIENT],
+    ["a", SHARD_A_CLIENT],
+  ]);
+}
 
 function makeDeps(over: Partial<ResolveIdempotencyClientDeps>): ResolveIdempotencyClientDeps {
   return {
     isSplitEnabled: async () => true,
     fallbackClient: FALLBACK,
-    newClient: NEW_CLIENT,
-    legacyClient: LEGACY_CLIENT,
+    clients: clientMap(),
     resolveMintKind: async () => "runOpsId",
+    // Kept as an injected seam: the real resolveShard is total, so only an injected
+    // classifier can exercise the throw-to-fallback arm below.
     classify: (id) => {
-      if (id.length === 26 && id[25] === "1") return "NEW";
-      if (id.length === 25) return "LEGACY";
+      if (id.length === 26 && id[25] === "2") return id[24]!;
+      if (id.length === 26 && id[25] === "1") return "new";
+      if (id.length === 25) return "legacy";
       throw new Error(`unclassifiable: ${id.length}`);
     },
-    isMigrated: undefined,
     ...over,
   };
 }
@@ -72,29 +83,49 @@ describe("resolveIdempotencyDedupClient", () => {
     expect(client).toBe(LEGACY_CLIENT);
   });
 
-  it("routes a swept (migrated) cuid-parent child to the NEW client", async () => {
-    const cuidParent = RunId.toFriendlyId("c".repeat(25));
+  it("falls back to the fallback client when a present parent id is unclassifiable", async () => {
     const client = await resolveIdempotencyDedupClient(
-      { environmentForMint: env, parentRunFriendlyId: cuidParent },
-      makeDeps({ isMigrated: async () => true })
+      { environmentForMint: env, parentRunFriendlyId: "run_not-a-valid-length" },
+      makeDeps({})
     );
-    expect(client).toBe(NEW_CLIENT);
+    expect(client).toBe(FALLBACK);
   });
 
-  it("routes a non-migrated cuid-parent child to the LEGACY client even when isMigrated is provided", async () => {
-    const cuidParent = RunId.toFriendlyId("d".repeat(25));
+  it("routes a child to its OWN SHARD client when the parent is a gen-2 id", async () => {
+    const genTwoParent = RunId.toFriendlyId("e".repeat(24) + "a2");
     const client = await resolveIdempotencyDedupClient(
-      { environmentForMint: env, parentRunFriendlyId: cuidParent },
-      makeDeps({ isMigrated: async () => false })
+      { environmentForMint: env, parentRunFriendlyId: genTwoParent },
+      makeDeps({ resolveMintKind: async () => "cuid" }) // mint flag must NOT win for a child
     );
-    expect(client).toBe(LEGACY_CLIENT);
+    expect(client).toBe(SHARD_A_CLIENT);
   });
 
-  it("falls back to the fallback client when a present parent id is unclassifiable", async () => {
+  it("falls back and logs when a gen-2 parent names an unconfigured shard key", async () => {
+    const errors: unknown[] = [];
+    const genTwoParent = RunId.toFriendlyId("f".repeat(24) + "z2");
     const client = await resolveIdempotencyDedupClient(
-      { environmentForMint: env, parentRunFriendlyId: "run_not-a-valid-length" },
-      makeDeps({})
+      { environmentForMint: env, parentRunFriendlyId: genTwoParent },
+      makeDeps({ logger: { error: (_m, meta) => errors.push(meta) } })
     );
     expect(client).toBe(FALLBACK);
+    expect(errors).toHaveLength(1);
+  });
+});
+
+describe("clientForShardKey", () => {
+  it("selects the same client the map holds for each reserved key and shard key", () => {
+    const clients = clientMap();
+    expect(clientForShardKey("new", clients, FALLBACK)).toBe(NEW_CLIENT);
+    expect(clientForShardKey("legacy", clients, FALLBACK)).toBe(LEGACY_CLIENT);
+    expect(clientForShardKey("a", clients, FALLBACK)).toBe(SHARD_A_CLIENT);
+  });
+
+  it("returns the fallback and logs for a key the map does not hold", () => {
+    const errors: unknown[] = [];
+    const client = clientForShardKey("z", clientMap(), FALLBACK, {
+      error: (_m, meta) => errors.push(meta),
+    });
+    expect(client).toBe(FALLBACK);
+    expect(errors).toHaveLength(1);
   });
 });

apps/webapp/app/runEngine/concerns/idempotencyResidency.server.ts

@@ -1,22 +1,46 @@
-import { ownerEngine, RunId, type Residency } from "@trigger.dev/core/v3/isomorphic";
+import { resolveShard, RunId, type ShardKey } from "@trigger.dev/core/v3/isomorphic";
 import type { PrismaClientOrTransaction } from "@trigger.dev/database";
 
 type MintKind = "cuid" | "runOpsId";
 
+type Logger = { error: (message: string, meta?: Record<string, unknown>) => void };
+
 export type ResolveIdempotencyClientDeps = {
   isSplitEnabled: () => Promise<boolean>;
   fallbackClient: PrismaClientOrTransaction;
-  newClient: PrismaClientOrTransaction;
-  legacyClient: PrismaClientOrTransaction;
+  /** Every store keyed by shard key: the reserved `legacy`/`new` plus one entry per gen-2 shard. */
+  clients: ReadonlyMap<ShardKey, PrismaClientOrTransaction>;
   resolveMintKind: (environment: {
     organizationId: string;
     id: string;
     orgFeatureFlags?: unknown;
   }) => Promise<MintKind>;
-  classify?: (id: string) => Residency;
-  isMigrated?: (id: string) => Promise<boolean>;
+  classify?: (id: string) => ShardKey;
+  logger?: Logger;
 };
 
+/**
+ * The one place an id becomes a client. `ShardKey` collapses to `string`, so the compiler
+ * cannot catch a wrong key here — an absent key takes an explicit logged branch to the
+ * fallback rather than a silent `?? legacy`.
+ */
+export function clientForShardKey(
+  shardKey: ShardKey,
+  clients: ReadonlyMap<ShardKey, PrismaClientOrTransaction>,
+  fallback: PrismaClientOrTransaction,
+  logger?: Logger
+): PrismaClientOrTransaction {
+  const client = clients.get(shardKey);
+  if (client === undefined) {
+    logger?.error("idempotency: no client configured for shard key", {
+      shardKey,
+      configured: [...clients.keys()],
+    });
+    return fallback;
+  }
+  return client;
+}
+
 export async function resolveIdempotencyDedupClient(
   args: {
     environmentForMint: { organizationId: string; id: string; orgFeatureFlags?: unknown };
@@ -28,9 +52,9 @@ export async function resolveIdempotencyDedupClient(
     return deps.fallbackClient;
   }
 
-  const classify = deps.classify ?? ownerEngine;
-  const clientFor = (residency: Residency): PrismaClientOrTransaction =>
-    residency === "NEW" ? deps.newClient : deps.legacyClient;
+  const classify = deps.classify ?? resolveShard;
+  const clientFor = (shardKey: ShardKey): PrismaClientOrTransaction =>
+    clientForShardKey(shardKey, deps.clients, deps.fallbackClient, deps.logger);
 
   if (args.parentRunFriendlyId) {
     let parentInternalId: string;
@@ -39,18 +63,18 @@ export async function resolveIdempotencyDedupClient(
     } catch {
       return deps.fallbackClient;
     }
-    let residency: Residency;
+    let shardKey: ShardKey;
     try {
-      residency = classify(parentInternalId);
+      shardKey = classify(parentInternalId);
     } catch {
       return deps.fallbackClient;
     }
-    if (residency === "LEGACY" && deps.isMigrated && (await deps.isMigrated(parentInternalId))) {
-      return deps.newClient;
-    }
-    return clientFor(residency);
+    return clientFor(shardKey);
   }
 
+  // Mint kind, not an id: there is no shard to decode, so this keeps resolving to the
+  // gen-1 pair exactly as before. Which shard a gen-2 env mints into is the mint layer's
+  // decision, and this client is a read-your-writes signal rather than a correctness gate.
   const kind = await deps.resolveMintKind(args.environmentForMint);
-  return clientFor(kind === "runOpsId" ? "NEW" : "LEGACY");
+  return clientFor(kind === "runOpsId" ? "new" : "legacy");
 }