fix(run-engine,run-store): close run-ops split read and routing gaps on the resume path

Repair a resume snapshot's completed-waitpoints from the owning primary when a
multi-reader replica serves the snapshot without its join rows. This is the
single-triggerAndWait case the order-based repair could not see (empty
completedWaitpointOrder), where the runner consumes an empty resume and the run
hangs. The presence-aware read co-reads snapshot visibility and its ids in one
statement, and only reads the primary when the reader lacks the snapshot, so a
single-reader replica never pays.

Also route batch item creation by batchTaskRunId, so an item stays visible to
the batch-completion count even if child and batch residency ever diverge, and
reject control-plane-only relation includes on the dedicated store with a clear
error instead of an opaque Prisma failure.

Author: d-cs

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+),
@@ -320,9 +344,23 @@ export async function getExecutionSnapshotsSince(
   // Step 3: Get waitpoint IDs for the LATEST snapshot only (first in desc order)
   const latestSnapshot = snapshots[0];
   let readClient = prisma;
-  let waitpointIds = await getSnapshotWaitpointIds(prisma, latestSnapshot.id, runStore);
+  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: completedWaitpointOrder is written on the snapshot row in the same statement as the
+  // 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.

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

@@ -1527,4 +1527,56 @@ describe("RunEngine getSnapshotsSince", () => {
       }
     }
   );
+
+  // A single triggerAndWait resume has an EMPTY completedWaitpointOrder (only batch-indexed waitpoints
+  // populate it) but a real join row. On a multi-reader replica the join read can hit a laggier reader
+  // that does not yet have the snapshot, returning 0 - and the order-based repair can't see it (order is
+  // empty). The presence check must detect the reader lacks the snapshot and repair from the primary,
+  // or the runner is handed a waitpoint-less continue and the run hangs.
+  containerTest(
+    "repairs a single-wait resume from the primary when the replica reader lacks the snapshot",
+    async ({ prisma }) => {
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+      const scenario = await setupTestScenario(prisma, authenticatedEnvironment, {
+        totalWaitpoints: 1,
+        outputSizeKB: 1,
+        snapshotConfigs: [
+          { status: "RUN_CREATED", completedWaitpointCount: 0 },
+          { status: "EXECUTING_WITH_WAITPOINTS", completedWaitpointCount: 0 },
+          { status: "EXECUTING", completedWaitpointCount: 1 },
+        ],
+      });
+      const waitpointId = scenario.waitpoints[0].id;
+      const latestSnapshot = scenario.snapshots[scenario.snapshots.length - 1];
+      // Single-wait: keep the join row, clear the order so it's the non-batch case.
+      await prisma.taskRunExecutionSnapshot.update({
+        where: { id: latestSnapshot.id },
+        data: { completedWaitpointOrder: [] },
+      });
+
+      // A multi-reader replica reader that has NOT yet applied the snapshot: any raw read (the
+      // presence+join query) returns empty, while Step 2's findMany still returns the snapshot.
+      const laggyReader = new Proxy(prisma, {
+        get(target, prop) {
+          if (prop === "$queryRaw") {
+            return async () => [];
+          }
+          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(
+        laggyReader,
+        scenario.run.id,
+        scenario.snapshots[0].id,
+        undefined,
+        prisma
+      );
+
+      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]);
+    }
+  );
 });

internal-packages/run-store/src/PostgresRunStore.dedicatedSelect.test.ts

@@ -405,4 +405,36 @@ describe("PostgresRunStore dedicated caller-select adapter (P2-store-bodies-2)",
       expect(blocking[0].taskRun.friendlyId).toBe(r.friendlyId);
     }
   );
+
+  // The dedicated finders forward include/select to a subset client typed as the full schema, so a
+  // control-plane-only relation would throw an opaque Prisma 500 for NEW data (invisible to tsc). The
+  // boundary guard rejects the known keys with a clear message; the legacy (full-schema) store is a no-op.
+  heteroRunOpsPostgresTest(
+    "dedicated batch/attempt finders reject control-plane-only includes with a clear error",
+    async ({ prisma14, prisma17 }) => {
+      const dedicated = makeStore(prisma17, "dedicated");
+      const legacy = makeStore(prisma14, "legacy");
+
+      await expect(
+        dedicated.findBatchTaskRunById("batch_x", { include: { runsBlocked: true } as never })
+      ).rejects.toThrow(/not available on the dedicated run-ops subset/);
+      await expect(
+        dedicated.findTaskRunAttempt({
+          where: { id: "att_x" },
+          include: { backgroundWorker: true },
+        } as never)
+      ).rejects.toThrow(/not available on the dedicated run-ops subset/);
+
+      // A subset-present include passes the guard and resolves normally (null for a missing batch).
+      expect(
+        await dedicated.findBatchTaskRunById("batch_missing", { include: { items: true } as never })
+      ).toBeNull();
+      // Legacy store: guard is a no-op; the full schema has the relation, so no guard throw.
+      expect(
+        await legacy.findBatchTaskRunById("batch_missing", {
+          include: { runsBlocked: true } as never,
+        })
+      ).toBeNull();
+    }
+  );
 });

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

@@ -1662,6 +1662,41 @@ export class PostgresRunStore implements RunStore {
     return result.map((r) => r.B);
   }
 
+  // One query: LEFT JOIN the snapshot to its completed-waitpoint links so `present` (snapshot visible
+  // on this reader) and `ids` come from the SAME point-in-time. A multi-reader replica can otherwise
+  // return the snapshot (via a separate read) while a different, laggier reader returns 0 links.
+  async findSnapshotCompletedWaitpointIdsWithPresence(
+    snapshotId: string,
+    client?: ReadClient
+  ): Promise<{ present: boolean; ids: string[] }> {
+    const prisma = client ?? this.readOnlyPrisma;
+
+    const joinDelegate = (prisma as RunOpsCapableClient).completedWaitpoint;
+    if (this.schemaVariant === "dedicated" && joinDelegate) {
+      const rows = await prisma.$queryRaw<{ id: string; waitpointId: string | null }[]>`
+        SELECT s."id", cw."waitpointId"
+        FROM "TaskRunExecutionSnapshot" s
+        LEFT JOIN "CompletedWaitpoint" cw ON cw."snapshotId" = s."id"
+        WHERE s."id" = ${snapshotId}
+      `;
+      return {
+        present: rows.length > 0,
+        ids: rows.filter((r) => r.waitpointId !== null).map((r) => r.waitpointId as string),
+      };
+    }
+
+    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),
+    };
+  }
+
   // Reverse of `connectedRuns`: the run ids linked to a waitpoint. Co-resident with the RUN (the join
   // is written on the run's DB in blockRunWithWaitpointEdges), so the waitpoint's own store can MISS a
   // cross-DB run — the router fans this across BOTH DBs.
@@ -1994,12 +2029,44 @@ export class PostgresRunStore implements RunStore {
     return prisma.taskRunWaitpoint.deleteMany(args);
   }
 
+  // The dedicated subset schema lacks control-plane relations; a pass-through include/select of one
+  // throws an opaque Prisma "Unknown field" 500 for NEW-resident data - invisible to tsc, since the
+  // run-ops client is typed as the full schema. Reject the known control-plane-only keys with a clear
+  // message at the boundary. No-op on the legacy (full-schema) store.
+  #assertSubsetSelectable(
+    fields: Record<string, unknown> | null | undefined,
+    forbidden: readonly string[],
+    method: string
+  ): void {
+    if (this.schemaVariant !== "dedicated" || !fields) return;
+    for (const key of forbidden) {
+      if (fields[key]) {
+        throw new Error(
+          `${method}: "${key}" is not available on the dedicated run-ops subset schema ` +
+            `(control-plane-only); resolve it via the control-plane instead of selecting it here.`
+        );
+      }
+    }
+  }
+
   async findTaskRunAttempt<T extends Prisma.TaskRunAttemptFindFirstArgs>(
     args: Prisma.SelectSubset<T, Prisma.TaskRunAttemptFindFirstArgs>,
     client?: ReadClient
   ): Promise<Prisma.TaskRunAttemptGetPayload<T> | null> {
     const prisma = client ?? this.readOnlyPrisma;
 
+    const forbidden = ["backgroundWorker", "backgroundWorkerTask", "runtimeEnvironment", "queue"];
+    this.#assertSubsetSelectable(
+      (args as { include?: Record<string, unknown> }).include,
+      forbidden,
+      "findTaskRunAttempt"
+    );
+    this.#assertSubsetSelectable(
+      (args as { select?: Record<string, unknown> }).select,
+      forbidden,
+      "findTaskRunAttempt"
+    );
+
     return prisma.taskRunAttempt.findFirst(
       args
     ) as Promise<Prisma.TaskRunAttemptGetPayload<T> | null>;
@@ -2051,6 +2118,12 @@ export class PostgresRunStore implements RunStore {
   ): Promise<Prisma.BatchTaskRunGetPayload<{ include: T }> | null> {
     const prisma = client ?? this.prisma;
 
+    this.#assertSubsetSelectable(
+      args?.include as Record<string, unknown> | undefined,
+      ["runsBlocked", "waitpoints", "runtimeEnvironment"],
+      "findBatchTaskRunById"
+    );
+
     return prisma.batchTaskRun.findFirst({
       where: { id },
       ...(args?.include ? { include: args.include } : {}),
@@ -2065,6 +2138,12 @@ export class PostgresRunStore implements RunStore {
   ): Promise<Prisma.BatchTaskRunGetPayload<{ include: T }> | null> {
     const prisma = client ?? this.readOnlyPrisma;
 
+    this.#assertSubsetSelectable(
+      args?.include as Record<string, unknown> | undefined,
+      ["runsBlocked", "waitpoints", "runtimeEnvironment"],
+      "findBatchTaskRunByFriendlyId"
+    );
+
     return prisma.batchTaskRun.findFirst({
       where: { friendlyId, runtimeEnvironmentId: environmentId },
       ...(args?.include ? { include: args.include } : {}),
@@ -2081,6 +2160,12 @@ export class PostgresRunStore implements RunStore {
   ): Promise<Prisma.BatchTaskRunGetPayload<{ include: T }> | null> {
     const prisma = client ?? this.prisma;
 
+    this.#assertSubsetSelectable(
+      args?.include as Record<string, unknown> | undefined,
+      ["runsBlocked", "waitpoints", "runtimeEnvironment"],
+      "findBatchTaskRunByIdempotencyKey"
+    );
+
     return prisma.batchTaskRun.findFirst({
       where: { runtimeEnvironmentId: environmentId, idempotencyKey },
       ...(args?.include ? { include: args.include } : {}),

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

@@ -92,4 +92,40 @@ describe("run-ops split — completing a batch item routes by the batch id, not
       expect(await prisma14.batchTaskRunItem.count({ where: { batchTaskRunId: batchId } })).toBe(0);
     }
   );
+
+  // createBatchTaskRunItem must co-locate the item with its BATCH (so the completion count, routed by
+  // batchTaskRunId, finds it), not with the child run. Routing by taskRunId would place a divergent-
+  // residency item on the child's DB -> invisible to the count -> the batch never completes -> parent hangs.
+  heteroRunOpsPostgresTest(
+    "createBatchTaskRunItem places the item on the batch's DB, routing by batchTaskRunId not taskRunId",
+    async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => {
+      const router = makeSplitRouter(prisma14, prisma17);
+      const envId = "env_batchitem";
+      const batchId = `batch_${NEW_ID_26}`; // run-ops id -> #new
+      const runId = "c".repeat(25); // cuid -> classifies LEGACY (the divergent routing key)
+
+      await prisma17.batchTaskRun.create({
+        data: {
+          id: batchId,
+          friendlyId: "batch_create_new",
+          runtimeEnvironmentId: envId,
+          runCount: 1,
+          status: "PROCESSING",
+        },
+      });
+      // The run physically lives on #new (the batch's DB) so the item's FKs resolve there.
+      await seedDedicatedRun(prisma17, envId, runId);
+
+      await router.createBatchTaskRunItem({
+        batchTaskRunId: batchId,
+        taskRunId: runId,
+        status: "PENDING",
+      });
+
+      // GREEN: routed by batchTaskRunId (NEW) -> item on #new, visible to the batch-completion count.
+      // RED: routed by the cuid taskRunId -> #legacy -> the create FK-fails / the count never sees it.
+      expect(await prisma17.batchTaskRunItem.count({ where: { batchTaskRunId: batchId } })).toBe(1);
+      expect(await prisma14.batchTaskRunItem.count({ where: { batchTaskRunId: batchId } })).toBe(0);
+    }
+  );
 });

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

@@ -44,4 +44,60 @@ describe("run-ops split — completed waitpoints for a cuid snapshot are found o
       expect(ids).toEqual([WAITPOINT_ID]);
     }
   );
+
+  // WithPresence reports, in one read, whether the snapshot is visible on the reader (so a multi-reader
+  // replica can distinguish "no waitpoints" from "reader has not applied the snapshot yet").
+  heteroRunOpsPostgresTest(
+    "findSnapshotCompletedWaitpointIdsWithPresence reports present+ids for a dedicated snapshot, absent otherwise",
+    async ({ prisma17 }: { prisma17: RunOpsPrismaClient }) => {
+      const newStore = new PostgresRunStore({
+        prisma: prisma17 as never,
+        readOnlyPrisma: prisma17 as never,
+        schemaVariant: "dedicated",
+      });
+      const runId = "run_" + "k".repeat(24) + "01";
+      await prisma17.taskRun.create({
+        data: {
+          id: runId,
+          friendlyId: "run_pres",
+          engine: "V2",
+          status: "PENDING",
+          taskIdentifier: "t",
+          payload: "{}",
+          payloadType: "application/json",
+          traceId: "tr",
+          spanId: "sp",
+          queue: "q",
+          runtimeEnvironmentId: "env_pres",
+          projectId: "proj_pres",
+        },
+      });
+      const snap = await prisma17.taskRunExecutionSnapshot.create({
+        data: {
+          id: "c".repeat(25),
+          engine: "V2",
+          executionStatus: "EXECUTING",
+          description: "continue",
+          runStatus: "PENDING",
+          runId,
+          environmentId: "env_pres",
+          environmentType: "DEVELOPMENT",
+          projectId: "proj_pres",
+          organizationId: "org_pres",
+        },
+      });
+      await prisma17.completedWaitpoint.create({
+        data: { snapshotId: snap.id, waitpointId: WAITPOINT_ID },
+      });
+
+      expect(await newStore.findSnapshotCompletedWaitpointIdsWithPresence(snap.id)).toEqual({
+        present: true,
+        ids: [WAITPOINT_ID],
+      });
+      // A snapshot the reader does not have -> present:false, so its empty ids are not trusted as authoritative.
+      expect(
+        await newStore.findSnapshotCompletedWaitpointIdsWithPresence("c".repeat(24) + "zz")
+      ).toEqual({ present: false, ids: [] });
+    }
+  );
 });