perf(run-store,run-engine): route run-keyed waitpoint and snapshot reads by run id

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

@@ -126,10 +126,13 @@ function enhanceExecutionSnapshotWithWaitpoints(
 async function getSnapshotWaitpointIds(
   prisma: PrismaClientOrTransaction,
   snapshotId: string,
-  runStore?: RunStore
+  runStore?: RunStore,
+  // The owning run id, so the router can route to the run's store (the completed-waitpoint join
+  // co-locates with the snapshot/run) instead of fanning out to both run-ops DBs.
+  runId?: string
 ): Promise<string[]> {
   if (runStore) {
-    return runStore.findSnapshotCompletedWaitpointIds(snapshotId, prisma);
+    return runStore.findSnapshotCompletedWaitpointIds(snapshotId, prisma, runId);
   }
 
   const result = await prisma.$queryRaw<{ B: string }[]>`
@@ -144,10 +147,12 @@ async function getSnapshotWaitpointIds(
 async function getSnapshotWaitpointIdsWithPresence(
   prisma: PrismaClientOrTransaction,
   snapshotId: string,
-  runStore?: RunStore
+  runStore?: RunStore,
+  // The owning run id, so the router can route to the run's store instead of fanning out.
+  runId?: string
 ): Promise<{ present: boolean; ids: string[] }> {
   if (runStore) {
-    return runStore.findSnapshotCompletedWaitpointIdsWithPresence(snapshotId, prisma);
+    return runStore.findSnapshotCompletedWaitpointIdsWithPresence(snapshotId, prisma, runId);
   }
 
   const rows = await prisma.$queryRaw<{ id: string; B: string | null }[]>`
@@ -170,15 +175,18 @@ async function getSnapshotWaitpointIdsWithPresence(
 async function fetchWaitpointsInChunks(
   prisma: PrismaClientOrTransaction,
   waitpointIds: string[],
-  runStore?: RunStore
+  runStore?: RunStore,
+  // The owning run id, so the router routes to the run's store and only falls back to the other DB
+  // for the rare cross-tree token, instead of fanning every chunk out to both run-ops DBs.
+  runId?: string
 ): Promise<Waitpoint[]> {
   if (waitpointIds.length === 0) return [];
 
   const allWaitpoints: Waitpoint[] = [];
   for (let i = 0; i < waitpointIds.length; i += WAITPOINT_CHUNK_SIZE) {
     const chunk = waitpointIds.slice(i, i + WAITPOINT_CHUNK_SIZE);
     const waitpoints = runStore
-      ? await runStore.findManyWaitpoints({ where: { id: { in: chunk } } }, prisma)
+      ? await runStore.findManyWaitpoints({ where: { id: { in: chunk } } }, prisma, runId)
       : await prisma.waitpoint.findMany({
           where: { id: { in: chunk } },
         });
@@ -347,7 +355,8 @@ export async function getExecutionSnapshotsSince(
   const { present, ids } = await getSnapshotWaitpointIdsWithPresence(
     prisma,
     latestSnapshot.id,
-    runStore
+    runStore,
+    runId
   );
   let waitpointIds = ids;
 
@@ -356,7 +365,7 @@ export async function getExecutionSnapshotsSince(
   // 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);
+    waitpointIds = await getSnapshotWaitpointIds(repairClient, latestSnapshot.id, runStore, runId);
     readClient = repairClient;
   }
 
@@ -369,15 +378,20 @@ export async function getExecutionSnapshotsSince(
   // 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);
+    const repaired = await getSnapshotWaitpointIds(
+      repairClient,
+      latestSnapshot.id,
+      runStore,
+      runId
+    );
     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(readClient, waitpointIds, runStore);
+  const waitpoints = await fetchWaitpointsInChunks(readClient, waitpointIds, runStore, runId);
 
   // 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/systems/waitpointSystem.ts

@@ -59,9 +59,10 @@ export class WaitpointSystem {
     runId: string;
     tx?: PrismaClientOrTransaction;
   }) {
-    // Route the delete: a run's edges may live on #new and/or #legacy (mid-drain), so it must fan
-    // across both stores. The caller's `tx` is not forwarded into either leg — each store's delete
-    // runs on its own client (the router never threads a control-plane tx into a routed write).
+    // A run's edges co-locate with the run (the edge write routes by runId), so the router routes this
+    // taskRunId-keyed delete to the run's store rather than fanning out. The caller's `tx` is not
+    // forwarded — the delete runs on the owning store's own client (the router never threads a
+    // control-plane tx into a routed write).
     const deleted = await this.$.runStore.deleteManyTaskRunWaitpoints(
       { where: { taskRunId: runId } },
       tx
@@ -493,7 +494,9 @@ export class WaitpointSystem {
 
       // Check if the run is actually blocked using a separate query (see above). Pass the writer so the
       // pending re-read is read-your-writes on the owning PRIMARY (a lagging replica can strand the run).
-      const pendingCount = await this.$.runStore.countPendingWaitpoints($waitpoints, prisma);
+      // Route by the blocked run id: its blocking waitpoints co-locate with the run, so the router
+      // counts on the run's store and only falls back to the other DB for a cross-tree token.
+      const pendingCount = await this.$.runStore.countPendingWaitpoints($waitpoints, prisma, runId);
 
       const isRunBlocked = pendingCount > 0;
 

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

@@ -1897,7 +1897,9 @@ export class PostgresRunStore implements RunStore {
 
   async findSnapshotCompletedWaitpointIds(
     snapshotId: string,
-    client?: ReadClient
+    client?: ReadClient,
+    // `runId` selects residency at the router; a single store has one client and ignores it.
+    _runId?: string
   ): Promise<string[]> {
     const prisma = client ?? this.readOnlyPrisma;
 
@@ -1924,7 +1926,9 @@ export class PostgresRunStore implements RunStore {
   // return the snapshot (via a separate read) while a different, laggier reader returns 0 links.
   async findSnapshotCompletedWaitpointIdsWithPresence(
     snapshotId: string,
-    client?: ReadClient
+    client?: ReadClient,
+    // `runId` selects residency at the router; a single store has one client and ignores it.
+    _runId?: string
   ): Promise<{ present: boolean; ids: string[] }> {
     const prisma = client ?? this.readOnlyPrisma;
 
@@ -2097,7 +2101,12 @@ export class PostgresRunStore implements RunStore {
       SELECT COUNT(*) FROM inserted`;
   }
 
-  async countPendingWaitpoints(waitpointIds: string[], client?: ReadClient): Promise<number> {
+  async countPendingWaitpoints(
+    waitpointIds: string[],
+    client?: ReadClient,
+    // `runId` selects residency at the router; a single store has one client and ignores it.
+    _runId?: string
+  ): Promise<number> {
     const prisma = client ?? this.readOnlyPrisma;
 
     if (waitpointIds.length === 0) {
@@ -2128,6 +2137,35 @@ export class PostgresRunStore implements RunStore {
     return Number(pendingCheck[0]?.pending_count ?? 0);
   }
 
+  // One `SELECT id, status` returns which ids are PENDING and which exist on this store (any status),
+  // so the router can route by run id and only re-count the ids ABSENT here on the other DB — without
+  // undercounting pending (which would prematurely unblock a run) or double-counting a drain-mirrored
+  // id. Reads only id+status, so it stays cheap for a run's small blocking set.
+  async countPendingWaitpointsWithPresence(
+    waitpointIds: string[],
+    client?: ReadClient
+  ): Promise<{ pendingIds: string[]; presentIds: string[] }> {
+    const prisma = client ?? this.readOnlyPrisma;
+
+    if (waitpointIds.length === 0) {
+      return { pendingIds: [], presentIds: [] };
+    }
+
+    const rows =
+      this.schemaVariant === "dedicated"
+        ? await prisma.$queryRaw<{ id: string; status: string }[]>`
+            SELECT id, status FROM "Waitpoint" WHERE id = ANY(${waitpointIds}::text[])
+          `
+        : await prisma.$queryRaw<{ id: string; status: string }[]>`
+            SELECT id, status FROM "Waitpoint" WHERE id IN (${Prisma.join(waitpointIds)})
+          `;
+
+    return {
+      pendingIds: rows.filter((r) => r.status === "PENDING").map((r) => r.id),
+      presentIds: rows.map((r) => r.id),
+    };
+  }
+
   async createWaitpoint<T extends Prisma.WaitpointCreateArgs>(
     args: Prisma.SelectSubset<T, Prisma.WaitpointCreateArgs>,
     tx?: PrismaClientOrTransaction
@@ -2189,7 +2227,9 @@ export class PostgresRunStore implements RunStore {
 
   async findManyWaitpoints<T extends Prisma.WaitpointFindManyArgs>(
     args: Prisma.SelectSubset<T, Prisma.WaitpointFindManyArgs>,
-    client?: ReadClient
+    client?: ReadClient,
+    // `runId` selects residency at the router; a single store has one client and ignores it.
+    _runId?: string
   ): Promise<Prisma.WaitpointGetPayload<T>[]> {
     const prisma = client ?? this.readOnlyPrisma;
 

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

@@ -748,7 +748,7 @@ describe("RoutingRunStore never threads a caller tx into either sub-store (recor
   );
 
   heteroRunOpsPostgresTest(
-    "deleteManyTaskRunWaitpoints fan-out hands BOTH sub-stores an undefined tx",
+    "deleteManyTaskRunWaitpoints routes by taskRunId to the owning store with an undefined tx",
     async ({ prisma14, prisma17 }) => {
       const legacyCalls: RecordedCall[] = [];
       const newCalls: RecordedCall[] = [];
@@ -769,14 +769,15 @@ describe("RoutingRunStore never threads a caller tx into either sub-store (recor
       );
       await seedLegacyBlockingEdge(prisma14, env, runId, "spy_del");
 
-      // Keyed by taskRunId → the both-stores fan-out branch. Pass the base control-plane client as tx.
+      // Keyed by a classifiable taskRunId (a cuid run → #legacy): routes to the owning store, no
+      // fan-out. Pass the base control-plane client as tx — it must NOT be threaded into the routed leg.
       await router.deleteManyTaskRunWaitpoints({ where: { taskRunId: runId } }, prisma14);
 
       const legacyTx = txArgsFor(legacyCalls, "deleteManyTaskRunWaitpoints");
       const newTx = txArgsFor(newCalls, "deleteManyTaskRunWaitpoints");
       expect(legacyTx.length).toBeGreaterThan(0);
-      expect(newTx.length).toBeGreaterThan(0);
-      for (const arg of [...legacyTx, ...newTx]) expect(arg).toBeUndefined();
+      expect(newTx.length).toBe(0); // routed to #legacy only; the non-owning store is never touched
+      for (const arg of legacyTx) expect(arg).toBeUndefined();
     }
   );
 });

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

@@ -834,18 +834,20 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 11b: deleteManyTaskRunWaitpoints by taskRunId fans out to both and sums (:944) ──
-  // A run's edges can straddle DBs mid-drain; a delete keyed by taskRunId (not waitpointId) must
-  // delete from BOTH DBs and sum the count.
+  // ── Case 11b: deleteManyTaskRunWaitpoints by taskRunId routes to the run's store (no fan-out) ──
+  // An edge always co-locates with its run (blockRunWithWaitpointEdges routes the write by runId), so
+  // a run's edges only ever live on ONE store — a straddle can't arise via the write path. A delete
+  // keyed by a classifiable taskRunId therefore routes to the run's store and must NOT touch the other
+  // DB. To prove the routing (not a fan-out), we manufacture an edge on the non-owning DB too and
+  // assert it survives the delete.
   heteroRunOpsPostgresTest(
-    "case 11b: deleteManyTaskRunWaitpoints by taskRunId deletes edges on both DBs and sums",
+    "case 11b: deleteManyTaskRunWaitpoints by taskRunId routes to the run's store and leaves the other DB untouched",
     async ({ prisma14, prisma17 }) => {
       const { router } = makeSplitRouter(prisma14, prisma17);
       const env = await seedSharedEnv(prisma14, "m11b");
 
-      // ONE logical run id whose edges happen to exist on BOTH DBs (the straddle the fan-out guards).
-      // The edge is FK-free on #new (unnest path) and FK-bound on #legacy, so seed a co-resident
-      // waitpoint + run on #legacy for its edge, and write the #new edge directly.
+      // A run-ops run (→ #new). Its real edge lives on #new; we also plant an edge on #legacy that the
+      // write path could never create, purely to prove the routed delete does not fan out to #legacy.
       const runId = runOpsNew("m11br");
       const legacyToken = cuidLegacy("m11bt");
       await router.createRun(buildCreateRunInput({ runId, friendlyId: "run_m11b", ...env }));
@@ -881,7 +883,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
       await prisma14.$executeRawUnsafe(
         `INSERT INTO "TaskRunWaitpoint" ("id","taskRunId","waitpointId","projectId","createdAt","updatedAt") VALUES (gen_random_uuid(),'${runId}','${legacyToken}','${env.projectId}',NOW(),NOW())`
       );
-      // #new edge (FK-free) pointing at a run-ops token absent locally — drain straddle.
+      // The run's real edge, on its own DB (#new, FK-free unnest path), pointing at a run-ops token.
       const newToken = runOpsNew("m11bn");
       await prisma17.$executeRawUnsafe(
         `INSERT INTO "TaskRunWaitpoint" ("id","taskRunId","waitpointId","projectId","createdAt","updatedAt") VALUES (gen_random_uuid(),'${runId}','${newToken}','${env.projectId}',NOW(),NOW())`
@@ -891,10 +893,10 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
       expect(await prisma17.taskRunWaitpoint.count({ where: { taskRunId: runId } })).toBe(1);
 
       const { count } = await router.deleteManyTaskRunWaitpoints({ where: { taskRunId: runId } });
-      expect(count).toBe(2); // one edge deleted on each DB, summed
-
-      expect(await prisma14.taskRunWaitpoint.count({ where: { taskRunId: runId } })).toBe(0);
+      // Routed to #new (the run's store): only its edge is deleted; #legacy is never touched.
+      expect(count).toBe(1);
       expect(await prisma17.taskRunWaitpoint.count({ where: { taskRunId: runId } })).toBe(0);
+      expect(await prisma14.taskRunWaitpoint.count({ where: { taskRunId: runId } })).toBe(1);
     }
   );
 });