fix(webapp): read a gen-2 batch row from its shard in the batch results API

Fix the batch results API returning 404 for batches whose id routes to a
run-ops shard. The batch row is now read from the shard the id names,
matching how batch retrieve already finds it. Adds a regression test for
the sharded case.

Mono-RevId: 26a6a2f1c43ebcbff5fb9820259c4eea906bb32d

Author: d-cs

.server-changes/batch-results-sharded-batch-ids.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Fix the batch results API returning "Batch not found" for some batches that the batch retrieve API could see.

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

@@ -85,12 +85,20 @@ export class ApiBatchResultsPresenter extends BasePresenter {
     });
   }
 
-  // Passthrough: batch row off the replica, members via the single run store. No legacy read.
-  async #callPassthrough(
+  // A gen-2 batch id lives only on its shard: neither the passthrough replica nor the gen-1
+  // new/legacy probes hold the row, so both paths route the row read by the id first.
+  #shardReplicaFor(friendlyId: string): PrismaReplicaClient | undefined {
+    const shardKey = resolveShard(friendlyId);
+    if (shardKey === "new" || shardKey === "legacy") return undefined;
+    return (this.readThrough?.shardReplicas ?? runOpsShardReplicas).get(shardKey);
+  }
+
+  #readBatchRow(
+    client: PrismaClientOrTransaction,
     friendlyId: string,
     env: AuthenticatedEnvironment
-  ): Promise<BatchTaskRunExecutionResult | undefined> {
-    const batchRun = await this._replica.batchTaskRun.findFirst({
+  ) {
+    return client.batchTaskRun.findFirst({
       where: {
         friendlyId,
         runtimeEnvironmentId: env.id,
@@ -103,6 +111,18 @@ export class ApiBatchResultsPresenter extends BasePresenter {
         },
       },
     });
+  }
+
+  // Passthrough: batch row off the replica, members via the single run store. No legacy read.
+  async #callPassthrough(
+    friendlyId: string,
+    env: AuthenticatedEnvironment
+  ): Promise<BatchTaskRunExecutionResult | undefined> {
+    const batchRun = await this.#readBatchRow(
+      this.#shardReplicaFor(friendlyId) ?? this._replica,
+      friendlyId,
+      env
+    );
 
     if (!batchRun) {
       return undefined;
@@ -150,26 +170,12 @@ export class ApiBatchResultsPresenter extends BasePresenter {
     const newClient = (this.readThrough?.newClient ?? this._replica) as PrismaReplicaClient;
     const legacyReplica = (this.readThrough?.legacyReplica ?? this._replica) as PrismaReplicaClient;
 
-    const readBatch = (client: PrismaClientOrTransaction) =>
-      client.batchTaskRun.findFirst({
-        where: {
-          friendlyId,
-          runtimeEnvironmentId: env.id,
-        },
-        include: {
-          items: {
-            select: {
-              taskRunId: true,
-            },
-          },
-        },
-      });
-
-    let batchRun = await readBatch(newClient);
+    const batchShard = this.#shardReplicaFor(friendlyId);
+    let batchRun = await this.#readBatchRow(batchShard ?? newClient, friendlyId, env);
 
-    // Legacy READ REPLICA probe, only on a new-probe miss; skipped when past retention.
-    if (!batchRun && !this.readThrough?.isPastRetention?.(friendlyId)) {
-      batchRun = await readBatch(legacyReplica);
+    // Legacy READ REPLICA probe, only on a gen-1 new-probe miss; skipped when past retention.
+    if (!batchRun && !batchShard && !this.readThrough?.isPastRetention?.(friendlyId)) {
+      batchRun = await this.#readBatchRow(legacyReplica, friendlyId, env);
     }
 
     if (!batchRun) {

apps/webapp/app/v3/runOpsMigration/track1-baseline.json

@@ -68,37 +68,28 @@
     "WaitpointTag.project"
   ],
   "totals": {
-    "violations": 5,
-    "detectorI": 5,
+    "violations": 4,
+    "detectorI": 4,
     "detectorII": 0,
     "detectorIII": 0,
     "write": 0,
-    "read": 5,
+    "read": 4,
     "files": 1,
     "legacyAnnotations": 0
   },
   "violations": [
     {
       "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts",
-      "line": 93,
+      "line": 101,
       "model": "BatchTaskRun",
       "delegate": "batchTaskRun",
       "callKind": "read",
       "detector": "i",
-      "snippet": "const batchRun = await this._replica.batchTaskRun.findFirst({"
+      "snippet": "return client.batchTaskRun.findFirst({"
     },
     {
       "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts",
-      "line": 154,
-      "model": "BatchTaskRun",
-      "delegate": "batchTaskRun",
-      "callKind": "read",
-      "detector": "i",
-      "snippet": "client.batchTaskRun.findFirst({"
-    },
-    {
-      "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts",
-      "line": 214,
+      "line": 220,
       "model": "TaskRun",
       "delegate": "taskRun",
       "callKind": "read",
@@ -107,7 +98,7 @@
     },
     {
       "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts",
-      "line": 224,
+      "line": 230,
       "model": "TaskRun",
       "delegate": "taskRun",
       "callKind": "read",
@@ -116,7 +107,7 @@
     },
     {
       "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts",
-      "line": 241,
+      "line": 247,
       "model": "TaskRun",
       "delegate": "taskRun",
       "callKind": "read",

apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts

@@ -7,6 +7,7 @@ import { heteroRunOpsPostgresTest, makeNShardRunOpsPostgresTest } from "@interna
 import type { RunOpsPrismaClient } from "@internal/run-ops-database";
 import type { PrismaClient } from "@trigger.dev/database";
 import { generateRunOpsId, generateRunOpsIdV2 } from "@trigger.dev/core/v3/isomorphic";
+import { PostgresRunStore } from "@internal/run-store";
 import { describe, expect, vi } from "vitest";
 import type { PrismaReplicaClient } from "~/db.server";
 import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
@@ -475,4 +476,66 @@ describe("ApiBatchResultsPresenter split mode — real run-ops dedicated schema
     },
     180_000
   );
+
+  // A gen-2 BATCH id is minted onto its shard, so the batch row (not just its members) lives
+  // only there. Both read paths must route the row by its own id; before this, the gen-1
+  // probes missed and the endpoint 404'd for a batch that `batch.retrieve` found.
+  for (const splitEnabled of [true, false]) {
+    oneShardTest(
+      `a gen-2 batch whose row lives only on its shard returns its members (split ${
+        splitEnabled ? "on" : "off"
+      })`,
+      async ({ legacyPrisma, newPrisma, shardPrismas }) => {
+        const shardPrisma = shardPrismas[0]!;
+        const ctx = await seedLegacyEnv(legacyPrisma, `gen2-batch-row-${splitEnabled}`);
+
+        const memberIds = [generateRunOpsIdV2("a"), generateRunOpsIdV2("a")];
+        for (const [i, id] of memberIds.entries()) {
+          await seedNewMember(
+            shardPrisma,
+            { envId: ctx.environment.id, orgId: ctx.organization.id, projectId: ctx.project.id },
+            {
+              id,
+              friendlyId: `run_shard_${i}`,
+              status: "COMPLETED_SUCCESSFULLY",
+              output: JSON.stringify({ i }),
+            }
+          );
+        }
+
+        const batchFriendlyId = `batch_${generateRunOpsIdV2("a")}`;
+        await seedBatchOnNew(shardPrisma, ctx.environment.id, batchFriendlyId, memberIds);
+
+        const shardReplicas = new Map([["a", shardPrisma as unknown as PrismaReplicaClient]]);
+        const presenter = splitEnabled
+          ? new ApiBatchResultsPresenter(throwingPrisma, throwingPrisma, {
+              splitEnabled: true,
+              newClient: newPrisma as unknown as PrismaReplicaClient,
+              legacyReplica: legacyPrisma as unknown as PrismaReplicaClient,
+              shardReplicas,
+            })
+          : // Passthrough hydrates members through the primary handle; the REPLICA (where the row
+            // is read without the fix) is the gen-1 store, which does not hold this batch.
+            new ApiBatchResultsPresenter(
+              shardPrisma as unknown as PrismaReplicaClient,
+              newPrisma as unknown as PrismaReplicaClient,
+              { splitEnabled: false, shardReplicas },
+              new PostgresRunStore({
+                prisma: shardPrisma as any,
+                readOnlyPrisma: shardPrisma as any,
+              })
+            );
+
+        const result = await presenter.call(batchFriendlyId, env(ctx));
+
+        expect(result).toBeDefined();
+        expect(result!.id).toBe(batchFriendlyId);
+        expect(result!.items.map((item) => (item as { output?: string }).output)).toEqual([
+          JSON.stringify({ i: 0 }),
+          JSON.stringify({ i: 1 }),
+        ]);
+      },
+      180_000
+    );
+  }
 });