Merge remote-tracking branch 'origin/main' into feat/waitpoint-mint-flag-wiring-tri-13442

Author: d-cs

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/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts

@@ -38,7 +38,14 @@ const { action } = createActionApiRoute(
       });
 
       if (!waitpoint) {
-        throw json({ error: "Waitpoint not found" }, { status: 404 });
+        // Retryable: a miss here can be replica lag. resolveWaitpointThroughReadThrough
+        // deliberately does not read the legacy primary, so it relies on the caller retrying.
+        // A plain 404 is not retried by the SDK, which would turn a transient miss into a
+        // permanent failure.
+        throw json(
+          { error: "Waitpoint not found" },
+          { status: 404, headers: { "x-should-retry": "true" } }
+        );
       }
 
       const _result = await engine.blockRunWithWaitpoint({
@@ -55,6 +62,11 @@ const { action } = createActionApiRoute(
         { status: 200 }
       );
     } catch (error) {
+      // A Response thrown inside the try is a deliberate status (the 404 above), not a
+      // failure. Re-throw it untouched, or every intentional 4xx here becomes a 500.
+      if (error instanceof Response) {
+        throw error;
+      }
       logger.error("Failed to wait for waitpoint", { runId, waitpointId, error });
       throw json({ error: "Failed to wait for waitpoint token" }, { status: 500 });
     }

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,16 @@ const resolveOrgMollifierFlag = makeResolveMollifierFlag();
 // PG's unique index as the backstop.
 const MAX_CLEARED_WINNER_REACQUIRES = 5;
 
+// The store that owns a shard key. A function, not a map: the handles are module constants and
+// `runOpsShardWriters` is already keyed, so a second structure would add an allocation and, if
+// memoised, mutable module state. Reading them lazily also keeps this module importable by
+// triggerTask under a `~/db.server` mock that omits them.
+function idempotencyClientFor(shardKey: ShardKey): PrismaClientOrTransaction | undefined {
+  if (shardKey === "legacy") return runOpsLegacyPrisma;
+  if (shardKey === "new") return runOpsNewPrisma;
+  return runOpsShardWriters.get(shardKey);
+}
+
 // 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 +183,9 @@ export class IdempotencyKeyConcern {
       {
         isSplitEnabled,
         fallbackClient: this.prisma,
-        newClient: runOpsNewPrisma,
-        legacyClient: runOpsLegacyPrisma,
+        clientFor: idempotencyClientFor,
         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 +648,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),
+      idempotencyClientFor,
+      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,
+    clientFor: (key) => clientMap().get(key),
     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,51 @@ 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();
+    const clientFor = (key: string) => clients.get(key);
+    expect(clientForShardKey("new", clientFor, FALLBACK)).toBe(NEW_CLIENT);
+    expect(clientForShardKey("legacy", clientFor, FALLBACK)).toBe(LEGACY_CLIENT);
+    expect(clientForShardKey("a", clientFor, FALLBACK)).toBe(SHARD_A_CLIENT);
+  });
+
+  it("returns the fallback and logs for a key the map does not hold", () => {
+    const errors: unknown[] = [];
+    const map = clientMap();
+    const client = clientForShardKey("z", (key) => map.get(key), 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,44 @@
-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;
+  /** The store that owns a shard key: the reserved `legacy`/`new`, or a gen-2 shard. */
+  clientFor: (shardKey: ShardKey) => PrismaClientOrTransaction | undefined;
   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`. The configured set is not repeated in the log:
+ * boot already prints the shard table.
+ */
+export function clientForShardKey(
+  shardKey: ShardKey,
+  clientFor: (shardKey: ShardKey) => PrismaClientOrTransaction | undefined,
+  fallback: PrismaClientOrTransaction,
+  logger?: Logger
+): PrismaClientOrTransaction {
+  const client = clientFor(shardKey);
+  if (client === undefined) {
+    logger?.error("idempotency: no client configured for shard key", { shardKey });
+    return fallback;
+  }
+  return client;
+}
+
 export async function resolveIdempotencyDedupClient(
   args: {
     environmentForMint: { organizationId: string; id: string; orgFeatureFlags?: unknown };
@@ -28,9 +50,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.clientFor, deps.fallbackClient, deps.logger);
 
   if (args.parentRunFriendlyId) {
     let parentInternalId: string;
@@ -39,18 +61,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");
 }