fix(run-engine,run-store,webapp): stop split-mode waits hanging on resume

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

@@ -1,6 +1,7 @@
 import type { PrismaReplicaClient } from "~/db.server";
 import {
   $replica as defaultLegacyReplica,
+  runOpsNewPrisma as defaultNewPrimary,
   runOpsNewReplica as defaultNewClient,
   runOpsSplitReadEnabled as defaultSplitReadEnabled,
 } from "~/db.server";
@@ -9,6 +10,7 @@ import { readThroughRun } from "~/v3/runOpsMigration/readThrough.server";
 type ResolveWaitpointDeps = {
   newClient?: PrismaReplicaClient;
   legacyReplica?: PrismaReplicaClient;
+  newPrimary?: PrismaReplicaClient;
   splitEnabled?: boolean;
   isPastRetention?: (id: string) => boolean;
 };
@@ -18,12 +20,14 @@ type ResolveWaitpointDeps = {
 export type ResolveWaitpointReadThroughDefaults = {
   newClient: PrismaReplicaClient;
   legacyReplica: PrismaReplicaClient;
+  newPrimary: PrismaReplicaClient;
   splitEnabled: boolean;
 };
 
 const productionDefaults: ResolveWaitpointReadThroughDefaults = {
   newClient: defaultNewClient,
   legacyReplica: defaultLegacyReplica,
+  newPrimary: defaultNewPrimary as unknown as PrismaReplicaClient,
   splitEnabled: defaultSplitReadEnabled,
 };
 
@@ -36,18 +40,36 @@ export async function resolveWaitpointThroughReadThrough<T>(opts: {
 }): Promise<T | null> {
   const defaults = opts.defaults ?? productionDefaults;
 
+  const splitEnabled = opts.deps?.splitEnabled ?? defaults.splitEnabled;
+
   const result = await readThroughRun({
     runId: opts.waitpointId,
     environmentId: opts.environmentId,
     readNew: (client) => opts.read(client),
     readLegacy: (replica) => opts.read(replica),
     deps: {
-      splitEnabled: opts.deps?.splitEnabled ?? defaults.splitEnabled,
+      splitEnabled,
       newClient: opts.deps?.newClient ?? defaults.newClient,
       legacyReplica: opts.deps?.legacyReplica ?? defaults.legacyReplica,
       isPastRetention: opts.deps?.isPastRetention,
     },
   });
 
-  return result.source === "new" || result.source === "legacy-replica" ? result.value : null;
+  if (result.source === "new" || result.source === "legacy-replica") {
+    return result.value;
+  }
+  // past-retention is an intentional not-found: the token is gone.
+  if (result.source === "past-retention") {
+    return null;
+  }
+
+  // Read-your-writes fallback for a token completed immediately after mint, before it replicated:
+  // re-read from the run-ops PRIMARY only. We deliberately never read the control-plane/legacy
+  // primary here (that is the load the replica-only read-through exists to shed), so a legacy-resident
+  // token that misses its replica stays a miss and the caller retries, rather than adding primary load.
+  const fromNewPrimary = await opts.read(opts.deps?.newPrimary ?? defaults.newPrimary);
+  if (fromNewPrimary != null) {
+    return fromNewPrimary;
+  }
+  return null;
 }

apps/webapp/app/v3/runOpsMigration/waitpointTokenResolve.server.test.ts

@@ -102,6 +102,9 @@ describe("public wait-token resolution across the split boundary", () => {
         deps: {
           newClient: prisma17 as unknown as PrismaReplicaClient,
           legacyReplica: prisma17 as unknown as PrismaReplicaClient,
+          // The run-ops-primary fallback is also pinned at run-ops, which does not hold this
+          // control-plane token, so the miss stays a miss (fallback never reads the control-plane).
+          newPrimary: prisma17 as unknown as PrismaReplicaClient,
         },
       });
 

apps/webapp/app/v3/services/createCheckpoint.server.ts

@@ -147,10 +147,14 @@ export class CreateCheckpointService extends BaseService {
       }
       case "WAIT_FOR_BATCH": {
         // Routed by friendlyId so a run-ops id (NEW-resident) batch is found on the owning DB;
-        // env-scoped to the dependent attempt's run (a batch shares its dependent's env).
+        // env-scoped to the dependent attempt's run (a batch shares its dependent's env). Read the
+        // primary: a batch that just resumed the parent may lag the replica, and a stale resumedAt
+        // (null) would checkpoint (suspend) an already-resumed run -> it stalls until a sweep.
         const batchRun = await this.runStore.findBatchTaskRunByFriendlyId(
           reason.batchFriendlyId,
-          attempt.taskRun.runtimeEnvironmentId
+          attempt.taskRun.runtimeEnvironmentId,
+          undefined,
+          this._prisma
         );
 
         if (!batchRun) {
@@ -361,11 +365,13 @@ export class CreateCheckpointService extends BaseService {
           });
           await marqs?.cancelHeartbeat(attempt.taskRunId);
 
-          // Routed by friendlyId so a run-ops id (NEW-resident) batch is found on the owning DB;
-          // env-scoped to the dependent attempt's run (a batch shares its dependent's env).
+          // Routed by friendlyId; read the primary (this._prisma) so a just-resumed batch that still
+          // lags the replica doesn't leave a stale resumedAt and suspend an already-resumed run.
           const batchRun = await this.runStore.findBatchTaskRunByFriendlyId(
             reason.batchFriendlyId,
-            attempt.taskRun.runtimeEnvironmentId
+            attempt.taskRun.runtimeEnvironmentId,
+            undefined,
+            this._prisma
           );
 
           if (!batchRun) {

apps/webapp/test/createCheckpoint.batchReplicaLag.test.ts

@@ -0,0 +1,65 @@
+// Unit red-green for the checkpoint WAIT_FOR_BATCH replica-lag fix (createCheckpoint.server.ts).
+// The service decides whether to suspend a run on `batchRun.resumedAt`; reading it from a lagging
+// replica makes a just-resumed batch look unresumed -> it suspends an already-resumed run -> stall.
+// The fix threads the primary (`this._prisma`) into `runStore.findBatchTaskRunByFriendlyId`. Here a
+// spy runStore records which client the service passed and simulates the lag (only the primary read
+// sees the fresh resumedAt): RED = no client -> stale null -> no early return; GREEN = primary -> kept alive.
+
+import { describe, expect, it, vi } from "vitest";
+
+vi.mock("~/services/logger.server", () => ({
+  logger: { debug: vi.fn(), info: vi.fn(), log: vi.fn(), error: vi.fn(), warn: vi.fn() },
+}));
+vi.mock("~/v3/marqs/index.server", () => ({
+  marqs: { replaceMessage: vi.fn(), cancelHeartbeat: vi.fn() },
+}));
+
+import { CreateCheckpointService } from "~/v3/services/createCheckpoint.server";
+
+describe("checkpoint WAIT_FOR_BATCH reads the primary, not a lagging replica", () => {
+  it("threads the primary so an already-resumed batch keeps the run alive", async () => {
+    // A freezable attempt so control reaches the WAIT_FOR_BATCH arm. This object IS the primary the
+    // fix must thread into the batch read.
+    const prisma = {
+      taskRunAttempt: {
+        findFirst: async () => ({
+          id: "attempt_1",
+          status: "EXECUTING",
+          taskRunId: "run_1",
+          taskRun: { id: "run_1", status: "EXECUTING", runtimeEnvironmentId: "env_1" },
+          backgroundWorker: { id: "bw_1", deployment: { imageReference: "img:1" } },
+        }),
+      },
+    };
+
+    let seenClient: unknown = "NOT_CALLED";
+    const runStore = {
+      findBatchTaskRunByFriendlyId: async (
+        _friendlyId: string,
+        _environmentId: string,
+        _args: unknown,
+        client?: unknown
+      ) => {
+        seenClient = client;
+        // Lagging replica: only a read handed the primary sees the just-committed resumedAt.
+        return { resumedAt: client === prisma ? new Date() : null };
+      },
+    };
+
+    const service = new CreateCheckpointService(prisma as never, {} as never, runStore as never);
+
+    let result: unknown;
+    try {
+      result = await service.call({
+        attemptFriendlyId: "attempt_1",
+        reason: { type: "WAIT_FOR_BATCH", batchFriendlyId: "batch_1" },
+      } as never);
+    } catch {
+      // Buggy path falls through the pre-check into checkpoint creation (unstubbed) and throws; the
+      // recorded client below is what distinguishes RED from GREEN.
+    }
+
+    expect(seenClient).toBe(prisma); // the fix: primary threaded into the batch read
+    expect(result).toEqual({ success: false, keepRunAlive: true }); // early-return, run kept alive
+  });
+});

apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts

@@ -151,6 +151,44 @@ describe("resolveWaitpointThroughReadThrough (hetero PG14 legacy + dedicated run
     }
   );
 
+  // Read-your-writes: a token completed immediately after mint may not be on the replicas yet. The
+  // replica-only read-through misses, and without a primary fallback we 404 a valid token (the
+  // authoritative completeWaitpoint never runs). The primary fallback must find it.
+  heteroRunOpsPostgresTest(
+    "a NEW-resident token missing on both replicas is found on the run-ops primary (no spurious 404)",
+    async ({ prisma17, prisma14 }) => {
+      // A NEW-resident token lives on the run-ops DB (prisma17); its cuid id fans new-then-legacy.
+      // Both replicas here lack it (pointed at prisma14), so only the run-ops primary can serve it.
+      const id = generateLegacyCuid();
+      const environmentId = generateRunOpsId();
+      const projectId = generateRunOpsId();
+      const seeded = await seedWaitpoint(prisma17, id, { id: environmentId, projectId });
+
+      const newClient = recording(prisma14);
+      const legacyReplica = recording(prisma14);
+      const newPrimary = recording(prisma17);
+
+      const result = await resolveWaitpointThroughReadThrough({
+        waitpointId: id,
+        environmentId,
+        read: read(id, environmentId),
+        deps: {
+          splitEnabled: true,
+          newClient: newClient.handle,
+          legacyReplica: legacyReplica.handle,
+          newPrimary: newPrimary.handle,
+        },
+      });
+
+      expect(result).not.toBeNull();
+      expect(result!.id).toBe(seeded.id);
+      // Both replicas probed and missed; the fallback read the run-ops primary (never control-plane).
+      expect(newClient.calls.length).toBe(1);
+      expect(legacyReplica.calls.length).toBe(1);
+      expect(newPrimary.calls.length).toBe(1);
+    }
+  );
+
   heteroRunOpsPostgresTest(
     "bare caller (no deps) resolves a NEW-resident waitpoint via the safe run-ops defaults",
     async ({ prisma17, prisma14 }) => {
@@ -170,6 +208,8 @@ describe("resolveWaitpointThroughReadThrough (hetero PG14 legacy + dedicated run
         defaults: {
           newClient: prisma14 as unknown as PrismaReplicaClient,
           legacyReplica: prisma14 as unknown as PrismaReplicaClient,
+          // Primaries also miss the NEW-resident waitpoint, so the miss stays a miss.
+          newPrimary: prisma14 as unknown as PrismaReplicaClient,
           splitEnabled: true,
         },
       });
@@ -183,6 +223,7 @@ describe("resolveWaitpointThroughReadThrough (hetero PG14 legacy + dedicated run
         defaults: {
           newClient: prisma17 as unknown as PrismaReplicaClient,
           legacyReplica: prisma14 as unknown as PrismaReplicaClient,
+          newPrimary: prisma17 as unknown as PrismaReplicaClient,
           splitEnabled: true,
         },
       });
@@ -206,6 +247,9 @@ describe("resolveWaitpointThroughReadThrough (hetero PG14 legacy + dedicated run
         splitEnabled: true,
         newClient: recording(prisma17).handle,
         legacyReplica: recording(prisma14).handle,
+        // The run-ops-primary fallback also misses this never-seeded token; inject a container client
+        // so it does not reach for the unconnectable production singleton.
+        newPrimary: recording(prisma17).handle,
       },
     });
 

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

@@ -1655,7 +1655,6 @@ export class RunEngine {
     completedAfter,
     idempotencyKey,
     idempotencyKeyExpiresAt,
-    tx,
   }: {
     /** The run that will block on this waitpoint. Co-locates the waitpoint with the run's DB. */
     runId?: string;
@@ -1664,7 +1663,6 @@ export class RunEngine {
     completedAfter: Date;
     idempotencyKey?: string;
     idempotencyKeyExpiresAt?: Date;
-    tx?: PrismaClientOrTransaction;
   }) {
     return this.waitpointSystem.createDateTimeWaitpoint({
       runId,
@@ -1673,7 +1671,6 @@ export class RunEngine {
       completedAfter,
       idempotencyKey,
       idempotencyKeyExpiresAt,
-      tx,
     });
   }
 
@@ -2097,13 +2094,24 @@ export class RunEngine {
       this.readOnlyPrisma !== this.prisma;
     const prisma = tx ?? (useReplica ? this.readOnlyPrisma : this.prisma);
 
-    const query = async (client: PrismaClientOrTransaction) => {
-      const snapshots = await getExecutionSnapshotsSince(client, runId, snapshotId, this.runStore);
+    const query = async (
+      client: PrismaClientOrTransaction,
+      repairClient?: PrismaClientOrTransaction
+    ) => {
+      const snapshots = await getExecutionSnapshotsSince(
+        client,
+        runId,
+        snapshotId,
+        this.runStore,
+        repairClient
+      );
       return snapshots.map(executionDataFromSnapshot);
     };
 
     try {
-      return await query(prisma);
+      // When reading the replica, pass the primary so a snapshot whose completed-waitpoint join rows
+      // have not replicated yet is repaired from the primary instead of resuming the run waitpoint-less.
+      return await query(prisma, useReplica ? this.prisma : undefined);
     } catch (e) {
       if (useReplica && e instanceof ExecutionSnapshotNotFoundError) {
         // Replica lag: the runner learned this snapshot id from the writer before the
@@ -2115,7 +2123,7 @@ export class RunEngine {
         if (maxMs > 0) {
           await setTimeout(minMs + Math.random() * Math.max(0, maxMs - minMs));
           try {
-            const result = await query(this.readOnlyPrisma);
+            const result = await query(this.readOnlyPrisma, this.prisma);
             this.snapshotsSinceReplicaMissCounter.add(1, { outcome: "replica_retry" });
             return result;
           } catch (replicaRetryError) {

internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts

@@ -138,6 +138,30 @@ async function getSnapshotWaitpointIds(
   return result.map((r) => r.B);
 }
 
+// Reads the snapshot's completed-waitpoint ids AND whether the snapshot is visible on this reader, in
+// one query, so a multi-reader replica can't return the ids from a different point-in-time than the
+// snapshot. `present=false` -> this reader lacks the snapshot; its empty id list is not authoritative.
+async function getSnapshotWaitpointIdsWithPresence(
+  prisma: PrismaClientOrTransaction,
+  snapshotId: string,
+  runStore?: RunStore
+): Promise<{ present: boolean; ids: string[] }> {
+  if (runStore) {
+    return runStore.findSnapshotCompletedWaitpointIdsWithPresence(snapshotId, prisma);
+  }
+
+  const rows = await prisma.$queryRaw<{ id: string; B: string | null }[]>`
+    SELECT s."id", cw."B"
+    FROM "TaskRunExecutionSnapshot" s
+    LEFT JOIN "_completedWaitpoints" cw ON cw."A" = s."id"
+    WHERE s."id" = ${snapshotId}
+  `;
+  return {
+    present: rows.length > 0,
+    ids: rows.filter((r) => r.B !== null).map((r) => r.B as string),
+  };
+}
+
 /**
  * Fetches waitpoints in chunks to avoid NAPI string conversion limits.
  * This is necessary because waitpoints can have large outputs (100KB+),
@@ -260,7 +284,10 @@ export async function getExecutionSnapshotsSince(
   prisma: PrismaClientOrTransaction,
   runId: string,
   sinceSnapshotId: string,
-  runStore?: RunStore
+  runStore?: RunStore,
+  // The primary, for read-repair when `prisma` is a lagging read replica (see Step 3). Omit when
+  // `prisma` is already the primary.
+  repairClient?: PrismaClientOrTransaction
 ): Promise<EnhancedExecutionSnapshot[]> {
   // Step 1: Find the createdAt of the sinceSnapshotId
   const sinceSnapshot = runStore
@@ -316,10 +343,41 @@ export async function getExecutionSnapshotsSince(
 
   // Step 3: Get waitpoint IDs for the LATEST snapshot only (first in desc order)
   const latestSnapshot = snapshots[0];
-  const waitpointIds = await getSnapshotWaitpointIds(prisma, latestSnapshot.id, runStore);
+  let readClient = prisma;
+  const { present, ids } = await getSnapshotWaitpointIdsWithPresence(
+    prisma,
+    latestSnapshot.id,
+    runStore
+  );
+  let waitpointIds = ids;
+
+  // Read-repair (multi-reader): Step 2 saw this snapshot, but on a multi-reader replica the join read
+  // can hit a laggier reader that does not have it yet. present=false means its empty id list is not
+  // authoritative - re-read from the primary so the runner is not handed a waitpoint-less continue
+  // (which it silently drops, hanging the run). Single-reader replicas never hit this (present stays true).
+  if (repairClient && repairClient !== prisma && !present) {
+    waitpointIds = await getSnapshotWaitpointIds(repairClient, latestSnapshot.id, runStore);
+    readClient = repairClient;
+  }
+
+  // Read-repair (batch): completedWaitpointOrder is written on the snapshot row in the same statement as the
+  // snapshot, so it is authoritative. If it lists more waitpoints than the join read returned, the
+  // _completedWaitpoints rows are stale on this client (a lagging read replica) - re-read them from the
+  // primary so the runner is not handed a waitpoint-less continue, which it drops and the run hangs.
+  // Distinct: completedWaitpointOrder can list the same id twice (a run batched under one idempotency
+  // key) while the _completedWaitpoints join is deduped, so compare unique ids or the repair fires every
+  // poll even against a caught-up replica.
+  const expectedCount = new Set(latestSnapshot.completedWaitpointOrder ?? []).size;
+  if (repairClient && repairClient !== prisma && waitpointIds.length < expectedCount) {
+    const repaired = await getSnapshotWaitpointIds(repairClient, latestSnapshot.id, runStore);
+    if (repaired.length > waitpointIds.length) {
+      waitpointIds = repaired;
+      readClient = repairClient;
+    }
+  }
 
   // Step 4: Fetch waitpoints in chunks to avoid NAPI string conversion limits
-  const waitpoints = await fetchWaitpointsInChunks(prisma, waitpointIds, runStore);
+  const waitpoints = await fetchWaitpointsInChunks(readClient, waitpointIds, runStore);
 
   // Step 5: Build enhanced snapshots - only latest gets waitpoints, others get empty arrays
   // The runner only uses completedWaitpoints from the latest snapshot anyway