feat(run-engine,run-store): make retried snapshot and waitpoint operations replay-safe

Internal groundwork for surviving brief database connection blips. No
user-facing change on its own.

Mono-RevId: b0eff85f8500b01f86dc17907da1b999589ebfcc

Author: d-cs

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

@@ -56,12 +56,14 @@ function createEngineOptions(redisOptions: any, prisma: any, store?: PostgresRun
 class CountingPostgresRunStore extends PostgresRunStore {
   public creates = 0;
   public latestReads = 0;
+  public snapshotCreateIds: (string | undefined)[] = [];
 
   override async createExecutionSnapshot(
     input: CreateExecutionSnapshotInput,
     tx?: any
   ): ReturnType<PostgresRunStore["createExecutionSnapshot"]> {
     this.creates++;
+    this.snapshotCreateIds.push(input.id);
     return super.createExecutionSnapshot(input, tx);
   }
 
@@ -130,6 +132,98 @@ describe("executionSnapshotSystem store routing (single-DB passthrough)", () =>
     }
   });
 
+  // A transition routed through the engine hands the store a stable id, so a blip-retry inside the
+  // store replays the same transition idempotently instead of duplicating it.
+  containerTest(
+    "snapshot transitions carry a store-supplied id",
+    async ({ prisma, redisOptions }) => {
+      const countingStore = new CountingPostgresRunStore({ prisma, readOnlyPrisma: prisma });
+      const engine = new RunEngine(createEngineOptions(redisOptions, prisma, countingStore));
+
+      try {
+        const run = await triggerRun(engine, prisma, "run_snapid1");
+        const latest = await getLatestExecutionSnapshot(prisma, run.id, countingStore);
+        countingStore.snapshotCreateIds.length = 0;
+
+        const created = await engine.executionSnapshotSystem.createExecutionSnapshot(prisma, {
+          run: { id: run.id, status: latest.runStatus, attemptNumber: latest.attemptNumber },
+          snapshot: { executionStatus: latest.executionStatus, description: "test transition" },
+          previousSnapshotId: latest.id,
+          environmentId: latest.environmentId,
+          environmentType: latest.environmentType,
+          projectId: latest.projectId,
+          organizationId: latest.organizationId,
+        });
+
+        expect(created.id).toBeDefined();
+        expect(SnapshotId.fromFriendlyId(SnapshotId.toFriendlyId(created.id))).toBe(created.id);
+        expect(countingStore.snapshotCreateIds).toEqual([created.id]);
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
+
+  // The transition-id invariant: each engine transition mints its OWN id, so two transitions never
+  // collide; but replaying ONE transition (the same supplied id, as a blip-retry does) stays a single
+  // row. This is what makes the store's conflict-ignoring insert safe: ids are per-transition, so a
+  // replay is only ever the same transition, never different data reusing an id.
+  containerTest(
+    "distinct engine transitions mint distinct ids; a replayed store write retains one id",
+    async ({ prisma, redisOptions }) => {
+      const countingStore = new CountingPostgresRunStore({ prisma, readOnlyPrisma: prisma });
+      const engine = new RunEngine(createEngineOptions(redisOptions, prisma, countingStore));
+
+      try {
+        const run = await triggerRun(engine, prisma, "run_txids1");
+        const latest = await getLatestExecutionSnapshot(prisma, run.id, countingStore);
+
+        const base = {
+          run: { id: run.id, status: latest.runStatus, attemptNumber: latest.attemptNumber },
+          environmentId: latest.environmentId,
+          environmentType: latest.environmentType,
+          projectId: latest.projectId,
+          organizationId: latest.organizationId,
+        };
+
+        // Two separate engine transitions receive two different ids (fresh mint per transition).
+        const t1 = await engine.executionSnapshotSystem.createExecutionSnapshot(prisma, {
+          ...base,
+          snapshot: { executionStatus: latest.executionStatus, description: "transition 1" },
+          previousSnapshotId: latest.id,
+        });
+        const t2 = await engine.executionSnapshotSystem.createExecutionSnapshot(prisma, {
+          ...base,
+          snapshot: { executionStatus: latest.executionStatus, description: "transition 2" },
+          previousSnapshotId: t1.id,
+        });
+        expect(t1.id).not.toBe(t2.id);
+
+        // Replaying ONE store operation (the same supplied transition id) retains one id and one row.
+        const transitionId = SnapshotId.generate().id;
+        const input = {
+          id: transitionId,
+          run: { id: run.id, status: latest.runStatus, attemptNumber: latest.attemptNumber ?? 1 },
+          snapshot: { executionStatus: latest.executionStatus, description: "replayed transition" },
+          environmentId: latest.environmentId,
+          environmentType: latest.environmentType,
+          projectId: latest.projectId,
+          organizationId: latest.organizationId,
+        };
+        const first = await countingStore.createExecutionSnapshot(input);
+        const second = await countingStore.createExecutionSnapshot(input);
+
+        expect(first.id).toBe(transitionId);
+        expect(second.id).toBe(transitionId);
+        expect(await prisma.taskRunExecutionSnapshot.count({ where: { id: transitionId } })).toBe(
+          1
+        );
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
+
   // getLatestExecutionSnapshot reads through the store, routed by run id.
   containerTest(
     "getLatestExecutionSnapshot reads through the store routed by run id",

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

@@ -437,6 +437,7 @@ export class ExecutionSnapshotSystem {
   public async createExecutionSnapshot(
     prisma: PrismaClientOrTransaction,
     {
+      snapshotId,
       run,
       snapshot,
       previousSnapshotId,
@@ -451,6 +452,12 @@ export class ExecutionSnapshotSystem {
       completedWaitpoints,
       error,
     }: {
+      /**
+       * Caller-supplied TRANSITION id: minted once per logical transition (e.g. so a publish guard can
+       * arm keyed by it before the write) and NEVER reused for different snapshot data. The store insert
+       * ignores conflicts, so a reused id keeps the first committed row and silently drops the new data.
+       */
+      snapshotId?: string;
       run: { id: string; status: TaskRunStatus; attemptNumber?: number | null };
       snapshot: {
         executionStatus: TaskRunExecutionStatus;
@@ -478,8 +485,16 @@ export class ExecutionSnapshotSystem {
     // The heartbeat/eventBus side effects below are unchanged.
     store?: RunStore
   ) {
+    // Mint the snapshot id here (above the store) unless the caller supplied one, so a connection-blip
+    // retry inside the store replays the SAME transition idempotently (the store's conflict-ignoring
+    // insert keyed by the id, then a read) instead of duplicating it. The id is a transition id: minted
+    // once per transition, never reused for different data. A caller supplies it only to arm a publish
+    // guard keyed by the id before the write.
+    const id = snapshotId ?? SnapshotId.generate().id;
+
     const newSnapshot = await (store ?? this.$.runStore).createExecutionSnapshot(
       {
+        id,
         run,
         snapshot,
         previousSnapshotId,

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

@@ -82,6 +82,15 @@ class CountingRunStore extends PostgresRunStore {
     this.calls.push("updateManyWaitpoints");
     return super.updateManyWaitpoints(args, tx);
   }
+  // The completion write goes through markWaitpointCompleted (the replay-safe completion method), so
+  // it is the DB step the guard must precede on every completion route.
+  override async markWaitpointCompleted(
+    waitpointId: string,
+    completion: Parameters<PostgresRunStore["markWaitpointCompleted"]>[1]
+  ): Promise<Prisma.BatchPayload> {
+    this.calls.push("markWaitpointCompleted");
+    return super.markWaitpointCompleted(waitpointId, completion);
+  }
   override async findManyTaskRunWaitpoints<T extends Prisma.TaskRunWaitpointFindManyArgs>(
     args: Prisma.SelectSubset<T, Prisma.TaskRunWaitpointFindManyArgs>,
     client?: any
@@ -753,10 +762,15 @@ describe("WaitpointSystem completion fan-out + residency store-selection guard",
   // than) the first completion DB write on the same store.
   function expectGuardFiredBeforeUpdate(calls: string[]) {
     const guardIdx = calls.indexOf("forWaitpointCompletion");
-    const updateIdx = calls.indexOf("updateManyWaitpoints");
+    // The completion DB write is markWaitpointCompleted (some routes may also touch
+    // updateManyWaitpoints); take whichever completion write lands first.
+    const writeIdx = ["markWaitpointCompleted", "updateManyWaitpoints"]
+      .map((m) => calls.indexOf(m))
+      .filter((i) => i >= 0)
+      .sort((a, b) => a - b)[0];
     expect(guardIdx).toBeGreaterThanOrEqual(0);
-    expect(updateIdx).toBeGreaterThanOrEqual(0);
-    expect(guardIdx).toBeLessThanOrEqual(updateIdx);
+    expect(writeIdx).toBeGreaterThanOrEqual(0);
+    expect(guardIdx).toBeLessThanOrEqual(writeIdx);
   }
 
   // ----- Group 1: EXHAUSTIVE route enumeration -----
@@ -781,10 +795,10 @@ describe("WaitpointSystem completion fan-out + residency store-selection guard",
 
         expect(store.calls).toContain("forWaitpointCompletion");
         const guardIdx = store.calls.indexOf("forWaitpointCompletion");
-        const updateIdx = store.calls.indexOf("updateManyWaitpoints");
-        // Synchronous route: guard is strictly the first DB step.
+        const writeIdx = store.calls.indexOf("markWaitpointCompleted");
+        // Synchronous route: guard is strictly the first DB step, before the completion write.
         expect(guardIdx).toBe(0);
-        expect(updateIdx).toBeGreaterThan(guardIdx);
+        expect(writeIdx).toBeGreaterThan(guardIdx);
       } finally {
         await engine.quit();
       }

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

@@ -147,17 +147,10 @@ export class LegacyPostgresWaitpointCoordinator implements WaitpointCoordinator
     }
 
     // 1. Complete the Waitpoint (if not completed)
+    // Replay-safe completion: the store builds the PENDING guard and COMPLETED values internally and
+    // retries this (and only this) waitpoint-update path on a connection blip.
     const [updateError, updateResult] = await tryCatch(
-      store.updateManyWaitpoints({
-        where: { id: waitpointId, status: "PENDING" },
-        data: {
-          status: "COMPLETED",
-          completedAt: new Date(),
-          output: output?.value,
-          outputType: output?.type,
-          outputIsError: output?.isError,
-        },
-      })
+      store.markWaitpointCompleted(waitpointId, { output })
     );
 
     if (updateError) {

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

@@ -0,0 +1,136 @@
+// #maybeInfraRetry must charge a blip-retry to the correct budget: replica reads to the replica
+// budget (so a replica retry storm can't drain the writer's), writer ops to the writer budget, and a
+// caller-transaction client is never retried. Selection is by the stored read-only object OR the
+// replica brand, so a routing-layer-forwarded replica wrapper (a fresh object, not the stored one)
+// still bills the replica budget. Deterministic one-shot fault injection, real testcontainer DB.
+import { postgresTest } from "@internal/testcontainers";
+import type { PrismaClient } from "@trigger.dev/database";
+import { generateInternalId } from "@trigger.dev/core/v3/isomorphic";
+import { expect } from "vitest";
+import { PostgresRunStore } from "./PostgresRunStore.js";
+import { setupSnapshotIdFixture } from "./testFixtures/snapshotIdFixture.js";
+import { withOneShotBlip } from "./testFixtures/oneShotBlip.js";
+import { markReadReplicaClient } from "./readReplicaClient.js";
+
+const options = { enabled: true, maxAttempts: 12, backoffMinMs: 20, backoffMaxMs: 120 };
+
+async function seedWaitpoint(prisma: PrismaClient, env: { id: string; projectId: string }) {
+  const id = generateInternalId();
+  await prisma.waitpoint.create({
+    data: {
+      id,
+      friendlyId: `wp_${id}`,
+      type: "MANUAL",
+      status: "PENDING",
+      idempotencyKey: `idem_${id}`,
+      userProvidedIdempotencyKey: false,
+      projectId: env.projectId,
+      environmentId: env.id,
+    },
+  });
+  return id;
+}
+
+postgresTest(
+  "a branded replica wrapper (not the stored one) charges its retry to the replica budget",
+  async ({ prisma }) => {
+    const { env } = await setupSnapshotIdFixture(prisma);
+    const wpId = await seedWaitpoint(prisma as PrismaClient, env);
+
+    let writer = 0;
+    let replica = 0;
+    const store = new PostgresRunStore({
+      prisma: prisma as never,
+      readOnlyPrisma: prisma as never,
+      infraRetry: { options, onRetry: () => writer++ },
+      readInfraRetry: { options, onRetry: () => replica++ },
+    });
+
+    // A freshly wrapped, branded replica handle: a different object than the stored readOnlyPrisma, so
+    // the old identity-only check would have billed the writer. The brand routes it to the replica.
+    const branded = markReadReplicaClient(withOneShotBlip(prisma, "waitpoint", "findFirst"));
+    const wp = await store.findWaitpoint({ where: { id: wpId } }, branded as never);
+
+    expect(wp?.id).toBe(wpId);
+    expect(replica).toBe(1);
+    expect(writer).toBe(0);
+  }
+);
+
+postgresTest(
+  "an op through the stored writer charges its retry to the writer budget",
+  async ({ prisma }) => {
+    const { env } = await setupSnapshotIdFixture(prisma);
+    const wpId = await seedWaitpoint(prisma as PrismaClient, env);
+
+    let writer = 0;
+    let replica = 0;
+    const faultingWriter = withOneShotBlip(prisma, "waitpoint", "updateMany");
+    const store = new PostgresRunStore({
+      prisma: faultingWriter as never,
+      readOnlyPrisma: prisma as never, // distinct stored replica so identity selection is meaningful
+      infraRetry: { options, onRetry: () => writer++ },
+      readInfraRetry: { options, onRetry: () => replica++ },
+    });
+
+    // markWaitpointCompleted is the retried writer op (updateManyWaitpoints is not retried); it runs
+    // waitpoint.updateMany on the writer, which the one-shot blip faults once, charging the writer budget.
+    await store.markWaitpointCompleted(wpId, {
+      output: { value: "{}", type: "application/json", isError: false },
+    });
+
+    expect(writer).toBe(1);
+    expect(replica).toBe(0);
+  }
+);
+
+postgresTest(
+  "a read through the stored replica charges its retry to the replica budget",
+  async ({ prisma }) => {
+    const { env } = await setupSnapshotIdFixture(prisma);
+    const wpId = await seedWaitpoint(prisma as PrismaClient, env);
+
+    let writer = 0;
+    let replica = 0;
+    const faultingReplica = withOneShotBlip(prisma, "waitpoint", "findFirst");
+    const store = new PostgresRunStore({
+      prisma: prisma as never,
+      readOnlyPrisma: faultingReplica as never,
+      infraRetry: { options, onRetry: () => writer++ },
+      readInfraRetry: { options, onRetry: () => replica++ },
+    });
+
+    const wp = await store.findWaitpoint({ where: { id: wpId } }); // no client -> stored readOnlyPrisma
+
+    expect(wp?.id).toBe(wpId);
+    expect(replica).toBe(1);
+    expect(writer).toBe(0);
+  }
+);
+
+postgresTest(
+  "a caller-transaction client is never retried (boundary preserved for both budgets)",
+  async ({ prisma }) => {
+    const { env } = await setupSnapshotIdFixture(prisma);
+    const wpId = await seedWaitpoint(prisma as PrismaClient, env);
+
+    let writer = 0;
+    let replica = 0;
+    const faulting = withOneShotBlip(prisma, "waitpoint", "findFirst");
+    const store = new PostgresRunStore({
+      prisma: prisma as never,
+      readOnlyPrisma: prisma as never,
+      infraRetry: { options, onRetry: () => writer++ },
+      readInfraRetry: { options, onRetry: () => replica++ },
+    });
+
+    // A tx client has no `$transaction`, so the op runs exactly once: the blip surfaces and neither
+    // budget's onRetry fires.
+    await expect(
+      (faulting as any).$transaction((tx: any) => store.findWaitpoint({ where: { id: wpId } }, tx))
+    ).rejects.toThrow();
+
+    expect(writer).toBe(0);
+    expect(replica).toBe(0);
+  }
+);