fix(run-engine): repair completed-waitpoints from the primary on a stale replica read

getExecutionSnapshotsSince serves /snapshots/since from the read replica. When a
snapshot's completedWaitpointOrder lists more (distinct) waitpoints than the join
read returns, the join rows have not replicated yet; re-read them from the primary
so the runner resumes instead of hanging. Distinct ids matter because a batched run
can list the same waitpoint more than once while the join is deduped.

Author: d-cs

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

@@ -2097,13 +2097,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 +2126,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

@@ -260,7 +260,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 +319,27 @@ 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;
+  let waitpointIds = await getSnapshotWaitpointIds(prisma, latestSnapshot.id, runStore);
+
+  // Read-repair: 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

internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts

@@ -3,7 +3,9 @@ import { trace } from "@internal/tracing";
 import { generateFriendlyId } from "@trigger.dev/core/v3/isomorphic";
 import { setTimeout } from "node:timers/promises";
 import { describe, expect } from "vitest";
+import type { PrismaClient } from "@trigger.dev/database";
 import { RunEngine } from "../index.js";
+import { getExecutionSnapshotsSince } from "../systems/executionSnapshotSystem.js";
 import { copySnapshotsToReplica, createTestMetricsMeter } from "./helpers/replicaTestHelpers.js";
 import { setupTestScenario } from "./helpers/snapshotTestHelpers.js";
 import { setupAuthenticatedEnvironment, setupBackgroundWorker } from "./setup.js";
@@ -1398,4 +1400,147 @@ describe("RunEngine getSnapshotsSince", () => {
       }
     }
   );
+
+  // This models a BATCH resume: production only populates completedWaitpointOrder for waitpoints that
+  // carry a batch index, so the read-repair only ever engages for batch resumes. Single-waitpoint
+  // continues have an empty order (repair is a no-op) and are covered instead by the atomic write
+  // (snapshot + join commit in one tx, which a replica applies atomically).
+  containerTest(
+    "repairs a batch resume's completed-waitpoints from the primary when the replica lags the join rows",
+    async ({ prisma, schemaOnlyPrisma, redisOptions }) => {
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+
+      const engine = new RunEngine({
+        prisma,
+        // Replica has the snapshot rows but lags on the _completedWaitpoints join (see below).
+        readOnlyPrisma: schemaOnlyPrisma,
+        readReplicaSnapshotsSinceEnabled: true,
+        readReplicaSnapshotsSinceRetryDelay: { minMs: 1, maxMs: 2 },
+        worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 },
+        queue: { redis: redisOptions },
+        runLock: { redis: redisOptions },
+        machines: {
+          defaultMachine: "small-1x",
+          machines: {
+            "small-1x": {
+              name: "small-1x" as const,
+              cpu: 0.5,
+              memory: 0.5,
+              centsPerMs: 0.0001,
+            },
+          },
+          baseCostInCents: 0.0001,
+        },
+        tracer: trace.getTracer("test", "0.0.0"),
+      });
+
+      try {
+        // Primary: a run whose latest snapshot is a warm-continue carrying one completed waitpoint -
+        // completedWaitpointOrder is set on the row AND the _completedWaitpoints join + waitpoint exist.
+        const scenario = await setupTestScenario(prisma, authenticatedEnvironment, {
+          totalWaitpoints: 1,
+          outputSizeKB: 1,
+          snapshotConfigs: [
+            { status: "RUN_CREATED", completedWaitpointCount: 0 },
+            { status: "EXECUTING", completedWaitpointCount: 0 },
+            { status: "EXECUTING_WITH_WAITPOINTS", completedWaitpointCount: 0 },
+            { status: "EXECUTING", completedWaitpointCount: 1 },
+          ],
+        });
+        const waitpointId = scenario.waitpoints[0].id;
+
+        // Replica lag: it receives the snapshot ROWS (incl. completedWaitpointOrder) but NOT the
+        // _completedWaitpoints join rows or the waitpoint - the exact state that drops the resume.
+        await copySnapshotsToReplica(prisma, schemaOnlyPrisma, scenario.run.id);
+
+        const result = await engine.getSnapshotsSince({
+          runId: scenario.run.id,
+          snapshotId: scenario.snapshots[0].id,
+        });
+
+        expect(result).not.toBeNull();
+        const latest = result![result!.length - 1];
+        // The runner must receive the completed waitpoint (repaired from the primary), not an empty set.
+        expect(latest.completedWaitpoints.map((w) => w.id)).toEqual([waitpointId]);
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
+
+  // completedWaitpointOrder can contain the same waitpoint id twice (a run batched under one idempotency
+  // key), while the _completedWaitpoints join is deduped. The read-repair must compare DISTINCT ids, or it
+  // fires on every poll for such a snapshot even against a caught-up replica - silently defeating offload.
+  containerTest(
+    "does not repair from the primary when completedWaitpointOrder has duplicates but the join is complete",
+    async ({ prisma, redisOptions }) => {
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+      const engine = new RunEngine({
+        prisma,
+        worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 },
+        queue: { redis: redisOptions },
+        runLock: { redis: redisOptions },
+        machines: {
+          defaultMachine: "small-1x",
+          machines: {
+            "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 },
+          },
+          baseCostInCents: 0.0001,
+        },
+        tracer: trace.getTracer("test", "0.0.0"),
+      });
+
+      try {
+        const scenario = await setupTestScenario(prisma, authenticatedEnvironment, {
+          totalWaitpoints: 1,
+          outputSizeKB: 1,
+          snapshotConfigs: [
+            { status: "RUN_CREATED", completedWaitpointCount: 0 },
+            { status: "EXECUTING", completedWaitpointCount: 1 },
+          ],
+        });
+        const waitpointId = scenario.waitpoints[0].id;
+        const latestSnapshot = scenario.snapshots[scenario.snapshots.length - 1];
+        // The join stays a single row; the order column lists the same id twice (the batch-dup case).
+        await prisma.taskRunExecutionSnapshot.update({
+          where: { id: latestSnapshot.id },
+          data: { completedWaitpointOrder: [waitpointId, waitpointId] },
+        });
+
+        // Count join reads on the repair client; the replica read (`prisma`) is already complete, so a
+        // correct repair must NEVER touch the repair client for this snapshot.
+        const repairCalls = { queryRaw: 0 };
+        const repairClient = new Proxy(prisma, {
+          get(target, prop) {
+            if (prop === "$queryRaw") {
+              return (...args: unknown[]) => {
+                repairCalls.queryRaw++;
+                return (target as unknown as { $queryRaw: (...a: unknown[]) => unknown }).$queryRaw(
+                  ...args
+                );
+              };
+            }
+            const value = (target as Record<string | symbol, unknown>)[prop];
+            return typeof value === "function" ? value.bind(target) : value;
+          },
+        }) as unknown as PrismaClient;
+
+        const result = await getExecutionSnapshotsSince(
+          prisma,
+          scenario.run.id,
+          scenario.snapshots[0].id,
+          undefined,
+          repairClient
+        );
+
+        const latest = result[result.length - 1];
+        // The duplicate is intentional (batch dup) - enhance expands by completedWaitpointOrder.
+        expect(latest.completedWaitpoints.map((w) => w.id)).toEqual([waitpointId, waitpointId]);
+        // No spurious primary read: distinct(order) === join count, so the repair must not fire.
+        expect(repairCalls.queryRaw).toBe(0);
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
 });