fix(run-store): keep caller replica reads on a replica; escalate only writer/tx

Review follow-up to the read-your-writes routing fix. The presence-only signal
(`client != null`) escalated EVERY caller-passed read client to the owning
store's primary, including an explicit read replica (e.g. `$replica`) passed by
span/trace/session lookups and the run-ops read-through — defeating replica
scaling.

Writer and replica Prisma clients are structurally identical at runtime (a
replica is a `new PrismaClient(...)` too, so it also exposes `$transaction`), so
shape can't tell them apart. The client builder now brands replica handles
(`markReadReplicaClient`) and the routing store reads the brand: a writer/tx
still escalates to the owning primary (read-your-writes, incl. the dequeue hot
path), while a branded replica or no client keeps the owning store's replica.

Also probe the PRIMARY in `forWaitpointCompletion`'s store-resolution (it selects
the store a subsequent write lands on), mirroring `#resolveWaitpointStore` — a
replica-lagged probe could mis-resolve the owner and strand the run.

Tests: branded-replica-stays-on-replica vs writer-escalates (routed reads); and
forWaitpointCompletion resolving the owner under replica lag. Seed the snapshot↔
waitpoint link via the Prisma relation API instead of a raw join-table insert.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: d-cs

apps/webapp/app/db.server.ts

@@ -8,6 +8,7 @@ import {
   type PrismaTransactionOptions,
 } from "@trigger.dev/database";
 import { RunOpsPrismaClient } from "@internal/run-ops-database";
+import { markReadReplicaClient } from "@internal/run-store";
 import invariant from "tiny-invariant";
 import { z } from "zod";
 import { env } from "./env.server";
@@ -170,7 +171,11 @@ export const prisma = singleton("prisma", () =>
 
 export const $replica: PrismaReplicaClient = singleton("replica", () => {
   const replica = getReplicaClient();
-  return replica ? captureInfrastructureErrors(tagDatasource("replica", replica)) : prisma;
+  // Brand ONLY a real replica so the run-store routing layer keeps replica reads off the primary.
+  // No replica configured → fall back to the writer `prisma`, which must stay UNBRANDED.
+  return replica
+    ? markReadReplicaClient(captureInfrastructureErrors(tagDatasource("replica", replica)))
+    : prisma;
 });
 
 export type RunOpsClients = { writer: PrismaClient; replica: PrismaReplicaClient };
@@ -254,9 +259,14 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => {
         captureInfraErrorsRunOps(
           tagDatasourceRunOps("writer", buildRunOpsWriterClient({ url, clientType }))
         ),
+      // Brand the run-ops replica (only built for a real replica URL) so routed replica reads stay
+      // off the primary. When no replica URL is set, selectRunOpsTopology reuses the writer here —
+      // which this callback never touches, so the writer stays unbranded.
       buildNewReplica: (url, clientType) =>
-        captureInfraErrorsRunOps(
-          tagDatasourceRunOps("replica", buildRunOpsReplicaClient({ url, clientType }))
+        markReadReplicaClient(
+          captureInfraErrorsRunOps(
+            tagDatasourceRunOps("replica", buildRunOpsReplicaClient({ url, clientType }))
+          )
         ),
     }
   );

internal-packages/run-store/src/index.ts

@@ -1,3 +1,4 @@
 export * from "./types.js";
 export * from "./PostgresRunStore.js";
 export * from "./runOpsStore.js";
+export * from "./readReplicaClient.js";

internal-packages/run-store/src/readReplicaClient.ts

@@ -0,0 +1,26 @@
+// A writer and a read-replica Prisma client are structurally identical at runtime (a replica is a
+// `new PrismaClient(...)` too, so it also exposes `$transaction`). The routing layer therefore
+// cannot tell a caller-passed replica from a writer by shape, yet it must — a writer/tx read has to
+// reach the owning store's PRIMARY (read-your-writes) while a replica read stays on a replica (read
+// scaling). So the client builder brands replica handles and the routing store reads the brand.
+export const READ_REPLICA_BRAND: unique symbol = Symbol.for("trigger.dev/run-store/read-replica");
+
+// Brand a replica client (returns the same object). MUST only be called on a genuine replica: the
+// routing layer trusts the brand to mean "do not escalate this read to the primary". An unbranded
+// replica just escalates as before — a scaling miss, never a correctness bug.
+export function markReadReplicaClient<T extends object>(client: T): T {
+  try {
+    (client as Record<symbol, unknown>)[READ_REPLICA_BRAND] = true;
+  } catch {
+    // Frozen/exotic clients may reject the assignment; only costs the optimization, not correctness.
+  }
+  return client;
+}
+
+export function isReadReplicaClient(client: unknown): boolean {
+  return (
+    !!client &&
+    typeof client === "object" &&
+    (client as Record<symbol, unknown>)[READ_REPLICA_BRAND] === true
+  );
+}

internal-packages/run-store/src/runOpsStore.routedReadPrimary.test.ts

@@ -16,6 +16,7 @@ import type { PrismaClient } from "@trigger.dev/database";
 import { describe, expect } from "vitest";
 import { PostgresRunStore } from "./PostgresRunStore.js";
 import { RoutingRunStore } from "./runOpsStore.js";
+import { markReadReplicaClient } from "./readReplicaClient.js";
 import type { CreateRunInput } from "./types.js";
 
 // ownerEngine classifies by internal-id LENGTH: 25 chars → cuid → LEGACY, 27 → ksuid → NEW.
@@ -175,9 +176,12 @@ describe("run-ops split — routed reads honor a caller-passed client via the ow
           environmentId: seed.environment.id,
         },
       });
-      await prisma14.$executeRaw`
-        INSERT INTO "_completedWaitpoints" ("A", "B") VALUES (${snapshotId}, ${waitpoint.id})
-      `;
+      // Link via the Prisma relation API (not a raw insert into the implicit join table) so a
+      // relation rename fails at compile time rather than silently seeding nothing.
+      await prisma14.taskRunExecutionSnapshot.update({
+        where: { id: snapshotId },
+        data: { completedWaitpoints: { connect: { id: waitpoint.id } } },
+      });
       expect(await router.findSnapshotCompletedWaitpointIds(snapshotId, prisma14)).toEqual([
         waitpoint.id,
       ]);
@@ -263,6 +267,51 @@ describe("run-ops split — routed reads honor a caller-passed client via the ow
     }
   );
 
+  // Read scaling: a caller that passes an explicit READ REPLICA (e.g. `$replica`) must NOT be
+  // escalated to the primary — only true read-your-writes (a writer/tx) should. A replica is a
+  // full PrismaClient at runtime (it has `$transaction` too), so shape can't distinguish it; the
+  // client builder brands it and the router honors the brand. Proven here by branding a client and
+  // showing the read stays on the owning store's (empty) replica — same as passing no client —
+  // while an unbranded writer escalates and finds the fresh row.
+  heteroPostgresTest(
+    "a branded read-replica client stays on the replica; a writer escalates to the primary",
+    async ({ prisma14, prisma17 }) => {
+      const { router } = splitTopology("LEGACY", prisma14, prisma17);
+      const seed = await seedEnvironment(prisma14, "replica_leg");
+      const runId = `run_${CUID_25}`;
+      await router.createRun(
+        buildCreateRunInput({
+          runId,
+          friendlyId: "run_replica_leg",
+          organizationId: seed.organization.id,
+          projectId: seed.project.id,
+          runtimeEnvironmentId: seed.environment.id,
+        })
+      );
+
+      // The router discards the caller's client object (it can't cross DBs) and reads the brand
+      // only, so a branded marker faithfully stands in for a passed `$replica`.
+      const replicaClient = markReadReplicaClient({} as unknown as PrismaClient);
+
+      // findRun (readYourWrites path): branded replica → owning replica (empty) → miss.
+      expect(
+        await router.findRun({ id: runId }, { select: { id: true } }, replicaClient)
+      ).toBeNull();
+      // Control: an unbranded writer escalates to the owning primary → finds the fresh row.
+      expect((await router.findRun({ id: runId }, { select: { id: true } }, prisma14))?.id).toBe(
+        runId
+      );
+
+      // findLatestExecutionSnapshot (#ownPrimary path): same replica-stays-on-replica invariant.
+      expect(await router.findLatestExecutionSnapshot(runId, replicaClient)).toBeNull();
+      expect((await router.findLatestExecutionSnapshot(runId, prisma14))?.executionStatus).toBe(
+        "RUN_CREATED"
+      );
+      // No client behaves identically to the branded replica.
+      expect(await router.findLatestExecutionSnapshot(runId)).toBeNull();
+    }
+  );
+
   heteroPostgresTest(
     "LEGACY cuid: waitpoint reads honor a caller client",
     async ({ prisma14, prisma17 }) => {

internal-packages/run-store/src/runOpsStore.test.ts

@@ -1257,6 +1257,39 @@ describe("RoutingRunStore.findRuns split-mode fan-out + drain", () => {
     }
   );
 
+  // forWaitpointCompletion selects the store a subsequent WRITE (updateManyWaitpoints) lands on,
+  // so its resolution probe must read each store's PRIMARY — not the replica. Here the owning
+  // (NEW) store's replica lags (points at an empty DB), so a replica probe would miss the fresh
+  // waitpoint and mis-resolve to the id-shape's default (LEGACY), stranding the run.
+  heteroPostgresTest(
+    "forWaitpointCompletion probes the primary, resolving the owner even under replica lag",
+    async ({ prisma14, prisma17 }) => {
+      const legacyStore = new PostgresRunStore({ prisma: prisma14, readOnlyPrisma: prisma14 });
+      // NEW store writes to prisma17 but its replica is the (empty, w.r.t. this waitpoint) prisma14.
+      const newStore = new PostgresRunStore({ prisma: prisma17, readOnlyPrisma: prisma14 });
+      const router = new RoutingRunStore({ new: newStore, legacy: legacyStore });
+      const seed17 = await seedEnvironment(prisma17, "wpc_lag17");
+
+      const cuidWaitpointId = `waitpoint_${"c".repeat(25)}`; // id-shape → LEGACY (the wrong owner)
+      await prisma17.waitpoint.create({
+        data: {
+          id: cuidWaitpointId,
+          friendlyId: "waitpoint_wpc_lag",
+          type: "MANUAL",
+          idempotencyKey: "wpc-lag-key",
+          userProvidedIdempotencyKey: false,
+          projectId: seed17.project.id,
+          environmentId: seed17.environment.id,
+        },
+      });
+
+      // Only the NEW store's PRIMARY (prisma17) has the row; a replica probe (prisma14) misses it.
+      expect(await router.forWaitpointCompletion(cuidWaitpointId, { routeKind: "MANUAL" })).toBe(
+        newStore
+      );
+    }
+  );
+
   // A waitpoint must be born on the same DB as its run (cuid → LEGACY, ksuid → NEW) so that
   // completion and the blocking edge — which already routes by run id — line up. A cuid
   // waitpoint landing on NEW is the regression that strands a non-opted org's wait forever.

internal-packages/run-store/src/runOpsStore.ts

@@ -26,6 +26,7 @@ import type {
   TaskRunWithWaitpoint,
   WaitpointColocationOptions,
 } from "./types.js";
+import { isReadReplicaClient } from "./readReplicaClient.js";
 
 /**
  * Run-ops routing substrate for the TaskRun-core method group. Implements {@link RunStore}
@@ -57,11 +58,11 @@ export class RoutingRunStore implements RunStore {
 
   // Map a caller-passed read client onto a routed store. The caller's client is bound to the
   // control-plane connection — the wrong database for a NEW-resident row — so it is never forwarded
-  // verbatim. Its PRESENCE is the read-your-writes signal (mirroring `client ?? readOnlyPrisma` in
-  // PostgresRunStore): the routed read runs on the owning store's OWN primary; no client keeps the
-  // store's default.
+  // verbatim. A WRITER/tx signals read-your-writes (the just-written row must beat replica lag), so
+  // the routed read runs on the owning store's OWN primary. A caller-passed REPLICA (branded) or no
+  // client keeps the owning store's replica, preserving read scaling.
   static #ownPrimary(store: RunStore, client: ReadClient | undefined): ReadClient | undefined {
-    return client === undefined ? undefined : store.primaryReadClient;
+    return client != null && !isReadReplicaClient(client) ? store.primaryReadClient : undefined;
   }
 
   // An unclassifiable id is treated as LEGACY (probe the control-plane DB rather than drop a
@@ -1174,12 +1175,17 @@ export class RoutingRunStore implements RunStore {
           : this.#legacy;
     // Resolve to where the waitpoint ACTUALLY lives: a migrated run's waitpoint can be on NEW
     // with a LEGACY-classified id (or vice versa), so verify and fall back rather than route
-    // by id-shape alone and miss it (which leaves the blocked run stuck forever).
-    if (await preferred.findWaitpoint({ where: { id: waitpointId } })) {
+    // by id-shape alone and miss it (which leaves the blocked run stuck forever). This guard
+    // selects the store a WRITE (updateManyWaitpoints) then lands on, so it must probe each
+    // store's PRIMARY (mirroring #resolveWaitpointStore's onPrimary): a just-created waitpoint the
+    // replica hasn't caught up on would otherwise mis-resolve the owner and strand the run.
+    if (
+      await preferred.findWaitpoint({ where: { id: waitpointId } }, preferred.primaryReadClient)
+    ) {
       return preferred;
     }
     const other = preferred === this.#new ? this.#legacy : this.#new;
-    if (await other.findWaitpoint({ where: { id: waitpointId } })) {
+    if (await other.findWaitpoint({ where: { id: waitpointId } }, other.primaryReadClient)) {
       return other;
     }
     return preferred;
@@ -1524,18 +1530,18 @@ function selectOrIncludeArgs(
   return undefined;
 }
 
-// A read-your-writes call passes a client — the writer or an ambient tx, whose just-written row
-// must be read back before the replica has it. Recover the caller's client from the overloaded
-// read args — slot two when it isn't a `{ select | include }` object, else slot three — and report
-// whether one was passed. (Its presence, not its identity, is the signal: even an explicitly
-// passed replica handle can't be forwarded across DBs, so it resolves to the owning primary.)
+// A read-your-writes call passes a WRITER or ambient tx, whose just-written row must be read back
+// before the replica has it. Recover the caller's client from the overloaded read args — slot two
+// when it isn't a `{ select | include }` object, else slot three — and report whether it warrants
+// escalation to the owning primary. A branded replica does NOT: it can't be forwarded across DBs,
+// but it signals a replica-intended read, so the owning store keeps its own replica (read scaling).
 function readYourWrites(
   argsOrClient: { select?: unknown; include?: unknown } | ReadClient | unknown,
   client: ReadClient | undefined
 ): boolean {
   const passedClient =
     selectOrIncludeArgs(argsOrClient) === undefined ? (argsOrClient ?? client) : client;
-  return passedClient != null;
+  return passedClient != null && !isReadReplicaClient(passedClient);
 }
 
 // Read a plain scalar string field off a create-data object (e.g. `data.completedByTaskRunId`).