refactor(run-engine): move block-edge registration behind the coordinator

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

@@ -490,25 +490,19 @@ export class WaitpointSystem {
         this.$.runStore
       );
 
-      // Insert the blocking + historical connections via the run-ops store, routed by the owning
-      // run id so the edge co-resides with the run. Never pinned to the caller's control-plane tx:
-      // that joined `Waitpoint` on the wrong DB and wrote 0 edges. The pending check stays a
-      // SEPARATE store call so it gets its own READ COMMITTED snapshot (see the doc comment above).
-      await this.$.runStore.blockRunWithWaitpointEdges({
+      // Insert the blocking + historical connections and re-check the pending count. The
+      // coordinator keeps these as two separate store statements, in this order, for the READ
+      // COMMITTED reason documented on the method and in the doc comment above.
+      const { pendingCount } = await this.coordinator.registerBlocks({
         runId,
         waitpointIds: $waitpoints,
         projectId,
         spanIdToComplete,
         batchId: batch?.id,
         batchIndex: batch?.index,
+        client: prisma,
       });
 
-      // 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).
-      // 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;
 
       let newStatus: TaskRunExecutionStatus = "SUSPENDED";
@@ -606,10 +600,10 @@ export class WaitpointSystem {
   }): Promise<void> {
     const $waitpoints = typeof waitpoints === "string" ? [waitpoints] : waitpoints;
 
-    // Same routed edge write as blockRunWithWaitpoint, routed by the owning run id. No lock
-    // needed: ON CONFLICT DO NOTHING makes concurrent inserts safe, and the parent snapshot is
-    // already EXECUTING_WITH_WAITPOINTS from blockRunWithCreatedBatch.
-    await this.$.runStore.blockRunWithWaitpointEdges({
+    // Same routed edge write as blockRunWithWaitpoint. No lock needed: ON CONFLICT DO NOTHING
+    // makes concurrent inserts safe, and the parent snapshot is already
+    // EXECUTING_WITH_WAITPOINTS from blockRunWithCreatedBatch. No pending count here.
+    await this.coordinator.registerBlocksLockless({
       runId,
       waitpointIds: $waitpoints,
       projectId,

internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts

@@ -2,7 +2,13 @@ import type { RunStore } from "@internal/run-store";
 import type { Logger } from "@trigger.dev/core/logger";
 import type { PrismaClient } from "@trigger.dev/database";
 import { boundedIn } from "@trigger.dev/database";
-import type { ClearRunBlockStateParams, RunBlockEdge, WaitpointCoordinator } from "./types.js";
+import type {
+  ClearRunBlockStateParams,
+  RegisterBlocksLocklessParams,
+  RegisterBlocksParams,
+  RunBlockEdge,
+  WaitpointCoordinator,
+} from "./types.js";
 
 export type LegacyPostgresWaitpointCoordinatorOptions = {
   runStore: RunStore;
@@ -65,4 +71,55 @@ export class LegacyPostgresWaitpointCoordinator implements WaitpointCoordinator
       this.prisma
     );
   }
+
+  async registerBlocks({
+    client,
+    ...edge
+  }: RegisterBlocksParams): Promise<{ pendingCount: number }> {
+    await this.#writeBlockEdges(edge);
+
+    // Check if the run is actually blocked using a separate query. The separate statement is the
+    // point: under PostgreSQL READ COMMITTED each statement gets its own snapshot, so a
+    // concurrent completion that commits between the edge write and this check is still seen.
+    // It queries ALL requested ids, not just inserted ones: a row that already existed (ON
+    // CONFLICT skipped the insert) but is still PENDING must still block. Pass the caller's
+    // client so the re-read is read-your-writes on the owning PRIMARY, and pass the run id so
+    // the router counts on the run's store instead of fanning out to both DBs.
+    const pendingCount = await this.runStore.countPendingWaitpoints(
+      edge.waitpointIds,
+      client,
+      edge.runId
+    );
+
+    return { pendingCount };
+  }
+
+  async registerBlocksLockless(params: RegisterBlocksLocklessParams): Promise<void> {
+    await this.#writeBlockEdges(params);
+  }
+
+  /**
+   * The edge write, shared by both register paths so they cannot drift.
+   *
+   * Routed by the owning run id so the edge co-resides with the run. Never pinned to a caller
+   * transaction: that joined `Waitpoint` on the wrong DB, wrote 0 edges, and silently never
+   * suspended the parent. The write is idempotent (ON CONFLICT DO NOTHING).
+   */
+  #writeBlockEdges({
+    runId,
+    waitpointIds,
+    projectId,
+    spanIdToComplete,
+    batchId,
+    batchIndex,
+  }: RegisterBlocksLocklessParams): Promise<void> {
+    return this.runStore.blockRunWithWaitpointEdges({
+      runId,
+      waitpointIds,
+      projectId,
+      spanIdToComplete,
+      batchId,
+      batchIndex,
+    });
+  }
 }

internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts

@@ -1,3 +1,4 @@
+import type { ReadClient } from "@internal/run-store";
 import type { PrismaClientOrTransaction, Waitpoint } from "@trigger.dev/database";
 
 /**
@@ -15,6 +16,8 @@ import type { PrismaClientOrTransaction, Waitpoint } from "@trigger.dev/database
 export type WaitpointCoordinator = {
   clearRunBlockState(params: ClearRunBlockStateParams): Promise<{ count: number }>;
   readRunBlockState(runId: string): Promise<RunBlockEdge[]>;
+  registerBlocks(params: RegisterBlocksParams): Promise<{ pendingCount: number }>;
+  registerBlocksLockless(params: RegisterBlocksLocklessParams): Promise<void>;
 };
 
 export type ClearRunBlockStateParams = {
@@ -40,3 +43,25 @@ export type RunBlockEdge = {
   batchIndex: number | null;
   waitpoint: Pick<Waitpoint, "id" | "status" | "type" | "completedAfter">;
 };
+
+export type RegisterBlocksParams = {
+  runId: string;
+  waitpointIds: string[];
+  projectId: string;
+  spanIdToComplete?: string;
+  batchId?: string;
+  batchIndex?: number;
+  /**
+   * Read client for the pending count only. The caller resolves `tx ?? prisma` once
+   * and passes the result, so the writer is used when the caller is inside a
+   * transaction and the pending re-read is read-your-writes on the owning primary.
+   * Never forwarded to the edge write.
+   */
+  client: ReadClient;
+};
+
+/**
+ * The lockless variant writes the edge and does not count. Two methods rather than
+ * one method with a flag, so "the batch path issues no extra query" is structural.
+ */
+export type RegisterBlocksLocklessParams = Omit<RegisterBlocksParams, "client">;