fix(webapp): route a gen-2 batch's completion write to its own shard

Minting gen-2 batch ids broke batch waits. The batch-completion writer was
resolved by a binary probe: look for the row on the new store, otherwise
assume legacy. A gen-2 batch lives on neither, so the probe fell through to
legacy, the update found no row and threw, the callback died before
tryCompleteBatch, the batch waitpoint stayed pending, and the parent run
waited forever with nothing logged as a hang.

Found by running it: a gen-2 batchTriggerAndWait parent never resumed, while
the same task on a gen-1 batch completed in twenty seconds.

A gen-2 batch id names its own shard, so it now routes by that and never
probes. An id naming an unconfigured shard throws rather than guessing a
store, because guessing is precisely what strands the run.

Both new tests fail without this change, the first on a fake client that
throws if the new store is probed at all.

Author: d-cs

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";
@@ -82,8 +83,25 @@ export async function resolveBatchRunOpsWriter(
     newReplica: RunOpsPrismaClient;
     newWriter: RunOpsPrismaClient;
     legacyWriter: RunOpsPrismaClient;
+    shards?: ReadonlyArray<{ key: string; writer: RunOpsPrismaClient }>;
   }
 ): Promise<RunOpsPrismaClient> {
+  // A gen-2 batch names its own shard in its id, so route by that and never probe. The
+  // probe below is binary — NEW, else assume LEGACY — so a gen-2 batch would fall through
+  // to a store that holds no such row, and the completion update would throw before the
+  // batch waitpoint could complete, leaving the parent run blocked with nothing logged.
+  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 },
@@ -105,6 +123,7 @@ export type BatchCompletionDeps = {
   newReplica: RunOpsPrismaClient;
   newWriter: RunOpsPrismaClient;
   legacyWriter: RunOpsPrismaClient;
+  shards?: ReadonlyArray<{ key: string; writer: RunOpsPrismaClient }>;
   tryCompleteBatch: (batchId: string) => Promise<unknown>;
 };
 
@@ -135,6 +154,7 @@ export async function handleBatchCompletion(
     newReplica: deps.newReplica,
     newWriter: deps.newWriter,
     legacyWriter: deps.legacyWriter,
+    shards: deps.shards,
   });
 
   try {

apps/webapp/test/runEngineHandlers.test.ts

@@ -490,6 +490,42 @@ describe("runEngineHandlers batch completion", () => {
 });
 
 describe("runEngineHandlers batch residency routing", () => {
+  // A gen-2 batch lives on its own shard. The binary probe below it looks only on the
+  // NEW store and then assumes LEGACY, so without a shard arm the completion update runs
+  // on a database that has no such row: Prisma throws "no record was found for an
+  // update", the callback dies before tryCompleteBatch, the BATCH waitpoint stays
+  // PENDING and the parent run waits forever with nothing logged as a hang.
+  it("a gen-2 batch resolves to its own shard writer", async () => {
+    const shardWriter = {} as never; // identity is the whole assertion; no database is touched
+    const gen2BatchId = `${"a".repeat(24)}a2`;
+
+    const writer = await resolveBatchRunOpsWriter(gen2BatchId, {
+      newReplica: {
+        batchTaskRun: {
+          findFirst: async () => {
+            throw new Error("a gen-2 batch id must never probe the NEW store");
+          },
+        },
+      } as never,
+      newWriter: {} as never,
+      legacyWriter: {} as never,
+      shards: [{ key: "a", writer: shardWriter as never }],
+    });
+
+    expect(writer).toBe(shardWriter);
+  });
+
+  it("an unconfigured shard key fails loud rather than writing elsewhere", async () => {
+    await expect(
+      resolveBatchRunOpsWriter(`${"a".repeat(24)}z2`, {
+        newReplica: {} as never,
+        newWriter: {} as never,
+        legacyWriter: {} as never,
+        shards: [{ key: "a", writer: {} as never }],
+      })
+    ).rejects.toThrow(/shard/i);
+  });
+
   // True single-DB invariant: the topology's cpFallback makes newReplica and
   // legacyWriter the SAME control-plane client, so the probe always resolves to
   // that one client regardless of where length-classification would guess.