From a8c63233bfcd63ca439faa94e891cc4913aa516c Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Mon, 6 Jul 2026 01:55:04 +0100 Subject: [PATCH 1/7] fix(run-store): commit run snapshots and their completed-waitpoint links atomically createExecutionSnapshot and lockRunToWorker wrote the execution snapshot and its completed-waitpoint join rows as two separate statements, and createRun and createFailedRun (dedicated schema) wrote the run and its associated waitpoint separately. Served back from a read replica that applied the snapshot but not yet the join, the runner gets a waitpoint-less continue and never resumes, so the run hangs (or partial state persists on a crash between the writes). Wrap each primary write and its dependent write in one transaction, reusing the caller's transaction when it already has a real one, so a replica can never observe the partial state. --- .../run-store/src/PostgresRunStore.test.ts | 118 +++++++++++++++ .../run-store/src/PostgresRunStore.ts | 76 ++++++++-- .../PostgresRunStore.writeAtomicity.test.ts | 137 ++++++++++++++++++ 3 files changed, 315 insertions(+), 16 deletions(-) diff --git a/internal-packages/run-store/src/PostgresRunStore.test.ts b/internal-packages/run-store/src/PostgresRunStore.test.ts index c0712bc3e3..b0423f24bd 100644 --- a/internal-packages/run-store/src/PostgresRunStore.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.test.ts @@ -888,6 +888,124 @@ describe("PostgresRunStore", () => { } ); + // lockRunToWorker creates a PENDING_EXECUTING snapshot (which can inherit completed waitpoints) and + // then connects the _completedWaitpoints join as a separate statement. These must commit together, or + // a replica-served /snapshots/since read can return the snapshot waitpoint-less and drop the resume. + postgresTest( + "lockRunToWorker writes the snapshot and its completed-waitpoint links atomically", + async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + const runId = "run_lock_atomic"; + + await store.createRun( + buildCreateRunInput({ + runId, + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const backgroundWorker = await prisma.backgroundWorker.create({ + data: { + friendlyId: "worker_atomic", + version: "20260601.1", + runtimeEnvironmentId: environment.id, + projectId: project.id, + contentHash: "abc", + sdkVersion: "3.0.0", + cliVersion: "3.0.0", + metadata: {}, + }, + }); + const workerTask = await prisma.backgroundWorkerTask.create({ + data: { + friendlyId: "task_atomic", + slug: "my-task", + filePath: "src/my-task.ts", + exportName: "myTask", + workerId: backgroundWorker.id, + runtimeEnvironmentId: environment.id, + projectId: project.id, + }, + }); + const queue = await prisma.taskQueue.create({ + data: { + friendlyId: "queue_atomic", + name: "task/my-task", + runtimeEnvironmentId: environment.id, + projectId: project.id, + }, + }); + const priorSnapshot = await prisma.taskRunExecutionSnapshot.create({ + data: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "prior", + runStatus: "PENDING", + environmentId: environment.id, + environmentType: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + runId, + }, + }); + const waitpoint = await prisma.waitpoint.create({ + data: { + friendlyId: "wp_lock_atomic", + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: "idem-lock-atomic", + userProvidedIdempotencyKey: false, + projectId: project.id, + environmentId: environment.id, + }, + }); + + // Force the completed-waitpoint join insert to fail mid-write. + await prisma.$executeRawUnsafe('DROP TABLE "_completedWaitpoints"'); + + const snapshotId = "snap_lock_atomic"; + await expect( + // Base client as `tx` = how the dequeue path calls it; the store must still open its own tx. + store.lockRunToWorker( + runId, + { + lockedAt: new Date(), + lockedById: workerTask.id, + lockedToVersionId: backgroundWorker.id, + lockedQueueId: queue.id, + startedAt: new Date(), + baseCostInCents: 5, + machinePreset: "small-1x", + taskVersion: "20260601.1", + sdkVersion: "3.0.0", + cliVersion: "3.0.0", + maxDurationInSeconds: null, + snapshot: { + id: snapshotId, + previousSnapshotId: priorSnapshot.id, + environmentId: environment.id, + environmentType: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + completedWaitpointIds: [waitpoint.id], + completedWaitpointOrder: [waitpoint.id], + }, + }, + prisma + ) + ).rejects.toThrow(); + + // Atomic: neither the snapshot nor the run lock persists without the links. + const snap = await prisma.taskRunExecutionSnapshot.findUnique({ where: { id: snapshotId } }); + expect(snap).toBeNull(); + const run = await prisma.taskRun.findUniqueOrThrow({ where: { id: runId } }); + expect(run.status).not.toBe("DEQUEUED"); + } + ); + postgresTest( "parkPendingVersion sets status to PENDING_VERSION and stores statusReason", async ({ prisma }) => { diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 94630cdb2c..40a0a7e1e8 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -527,6 +527,25 @@ export class PostgresRunStore implements RunStore { ); } + // Run `fn` atomically: reuse the caller's interactive transaction if it gave us a real one, else open + // our own on this store's writer. A real interactive tx has no `$transaction` method; a base client + // (which callers, e.g. the engine's dequeue/resume paths, thread through for routing) DOES - so a base + // client still gets a fresh transaction. Used by write methods that create a row plus dependent rows + // (snapshot + completed-waitpoints, run + associated-waitpoint) which must commit together. + #withOptionalTransaction( + tx: PrismaClientOrTransaction | undefined, + fn: (client: PrismaClientOrTransaction) => Promise + ): Promise { + const alreadyInTransaction = + tx !== undefined && typeof (tx as { $transaction?: unknown }).$transaction !== "function"; + if (alreadyInTransaction) { + return fn(tx); + } + return (this.prisma as RunOpsTransactionalClient).$transaction((t) => + fn(t as unknown as PrismaClientOrTransaction) + ); + } + async createRun( params: CreateRunInput, tx?: PrismaClientOrTransaction @@ -547,15 +566,19 @@ export class PostgresRunStore implements RunStore { }; if (this.schemaVariant === "dedicated") { - const run = (await client.taskRun.create({ - data: { - ...params.data, - executionSnapshots: { create: snapshotCreate }, - }, - })) as TaskRun; + // The run + its associated RUN-type waitpoint are two writes here (the legacy branch below nests + // them). Commit them together so a crash / lagging read never leaves a run without its waitpoint. + return this.#withOptionalTransaction(tx, async (c) => { + const run = (await c.taskRun.create({ + data: { + ...params.data, + executionSnapshots: { create: snapshotCreate }, + }, + })) as TaskRun; - const associatedWaitpoint = await this.#createAssociatedWaitpoint(client, run.id, params); - return { ...run, associatedWaitpoint }; + const associatedWaitpoint = await this.#createAssociatedWaitpoint(c, run.id, params); + return { ...run, associatedWaitpoint }; + }); } return client.taskRun.create({ @@ -633,12 +656,15 @@ export class PostgresRunStore implements RunStore { const client = tx ?? this.prisma; if (this.schemaVariant === "dedicated") { - const run = (await client.taskRun.create({ - data: { ...params.data }, - })) as TaskRun; - - const associatedWaitpoint = await this.#createAssociatedWaitpoint(client, run.id, params); - return { ...run, associatedWaitpoint }; + // Run + associated RUN-type waitpoint are two writes here; commit them together (see createRun). + return this.#withOptionalTransaction(tx, async (c) => { + const run = (await c.taskRun.create({ + data: { ...params.data }, + })) as TaskRun; + + const associatedWaitpoint = await this.#createAssociatedWaitpoint(c, run.id, params); + return { ...run, associatedWaitpoint }; + }); } return client.taskRun.create({ @@ -954,8 +980,17 @@ export class PostgresRunStore implements RunStore { data: LockRunData, tx?: PrismaClientOrTransaction ): Promise> { - const prisma = tx ?? this.prisma; + // The run-lock update (with its nested PENDING_EXECUTING snapshot) and the completed-waitpoint + // connect must commit together, or a replica-served resume read can see the snapshot without its + // links and drop the resume. + return this.#withOptionalTransaction(tx, (c) => this.#lockRunToWorker(runId, data, c)); + } + async #lockRunToWorker( + runId: string, + data: LockRunData, + prisma: PrismaClientOrTransaction + ): Promise> { const dedicated = this.schemaVariant === "dedicated"; const result = await prisma.taskRun.update({ @@ -1533,8 +1568,17 @@ export class PostgresRunStore implements RunStore { input: CreateExecutionSnapshotInput, tx?: PrismaClientOrTransaction ): Promise> { - const prisma = tx ?? this.prisma; + // The snapshot row and its completed-waitpoint join rows MUST commit together. `/snapshots/since` + // can be served from a lagging read replica, so a snapshot that commits before its links can be + // read back waitpoint-less and the runner's resume is lost (the run hangs). This is the warm-continue + // path: the engine threads its base prisma through as `tx`, which is not a real transaction. + return this.#withOptionalTransaction(tx, (c) => this.#createExecutionSnapshot(input, c)); + } + async #createExecutionSnapshot( + input: CreateExecutionSnapshotInput, + prisma: PrismaClientOrTransaction + ): Promise> { const { run, snapshot, diff --git a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts index fb81eb90c1..1b8a086b4c 100644 --- a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts @@ -396,6 +396,71 @@ describe("fan-out deleteManyTaskRunWaitpoints honors the caller's tx on the #leg ); }); +// createExecutionSnapshot writes the snapshot row and its completed-waitpoint join rows. These MUST +// commit together: with the flag off, `/snapshots/since` is served from a lagging read replica, so a +// snapshot that commits before its `_completedWaitpoints` rows can be read waitpoint-less, and the +// runner's EXECUTING branch no-ops on an empty completedWaitpoints -> the resume is lost -> hang. +describe("createExecutionSnapshot writes the snapshot and its completed-waitpoint links atomically", () => { + heteroRunOpsPostgresTest( + "rolls the snapshot back if the completed-waitpoint insert fails (no waitpoint-less snapshot persists)", + async ({ prisma14 }) => { + const legacy = makeLegacyStore(prisma14); + const env = await seedEnvironment(prisma14, "legacy", "ces_atomic"); + const runId = `run_${CUID_25}`; + await legacy.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_ces_atomic", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + const waitpoint = await prisma14.waitpoint.create({ + data: { + friendlyId: "wp_ces_atomic", + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: "idem-ces_atomic", + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + // Force the completed-waitpoint join insert to fail mid-write. + await prisma14.$executeRawUnsafe('DROP TABLE "_completedWaitpoints"'); + + await expect( + // Pass the base client as `tx` - exactly how the engine threads its base prisma through + // (continueRunIfUnblocked -> executionSnapshotSystem.createExecutionSnapshot(prisma, ...)). + // It is NOT an interactive transaction, so the store must still open its own to stay atomic. + legacy.createExecutionSnapshot( + { + run: { id: runId, status: "EXECUTING", attemptNumber: 1 }, + snapshot: { + executionStatus: "EXECUTING_WITH_WAITPOINTS", + description: "Run was blocked by a waitpoint.", + }, + environmentId: env.environment.id, + environmentType: "DEVELOPMENT", + projectId: env.project.id, + organizationId: env.project.id, + completedWaitpoints: [{ id: waitpoint.id, index: 0 }], + }, + prisma14 + ) + ).rejects.toThrow(); + + // The snapshot must NOT persist without its links, or a replica can serve it waitpoint-less. + const snap = await prisma14.taskRunExecutionSnapshot.findFirst({ + where: { runId, executionStatus: "EXECUTING_WITH_WAITPOINTS" }, + }); + expect(snap).toBeNull(); + } + ); +}); + // RoutingRunStore.createExecutionSnapshot accepts a caller tx but must forward it to the OWNING store // only when that store is #legacy: a control-plane tx can't wrap a #new (cross-DB) write, but it can // (and should) wrap a legacy-resident snapshot so it stays atomic with the caller's operation. @@ -457,3 +522,75 @@ describe("createExecutionSnapshot honors the caller's tx on the #legacy owning s } ); }); + +// On the dedicated subset schema the associated (RUN-type) waitpoint is created as a SEPARATE +// waitpoint.create after taskRun.create (the legacy schema nests it atomically). The pair must commit +// together, or a crash / lagging read leaves a run with no completion waitpoint and its parent never resumes. +function assocWaitpoint( + env: { project: { id: string }; environment: { id: string } }, + suffix: string +) { + return { + id: `wp_${suffix}`, + friendlyId: `waitpoint_${suffix}`, + type: "RUN" as const, + status: "PENDING" as const, + idempotencyKey: `idem_${suffix}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }; +} + +describe("createRun / createFailedRun write the run and its associated waitpoint atomically (dedicated)", () => { + heteroRunOpsPostgresTest( + "createRun rolls the run back if the associated-waitpoint create fails", + async ({ prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const env = await seedEnvironment(prisma17, "dedicated", "cr_atomic"); + const runId = `run_${NEW_ID_26}`; + const input = { + ...buildCreateRunInput({ + runId, + friendlyId: "run_cr_atomic", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }), + associatedWaitpoint: assocWaitpoint(env, "cr_atomic"), + }; + + // Force #createAssociatedWaitpoint (waitpoint.create) to fail after taskRun.create. + await prisma17.$executeRawUnsafe('DROP TABLE "Waitpoint"'); + + await expect(newStore.createRun(input)).rejects.toThrow(); + + const run = await prisma17.taskRun.findFirst({ where: { id: runId } }); + expect(run).toBeNull(); + } + ); + + heteroRunOpsPostgresTest( + "createFailedRun rolls the run back if the associated-waitpoint create fails", + async ({ prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const env = await seedEnvironment(prisma17, "dedicated", "cf_atomic"); + const runId = `run_${NEW_ID_26}`; + const base = buildCreateRunInput({ + runId, + friendlyId: "run_cf_atomic", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }); + const input = { data: base.data, associatedWaitpoint: assocWaitpoint(env, "cf_atomic") }; + + await prisma17.$executeRawUnsafe('DROP TABLE "Waitpoint"'); + + await expect(newStore.createFailedRun(input)).rejects.toThrow(); + + const run = await prisma17.taskRun.findFirst({ where: { id: runId } }); + expect(run).toBeNull(); + } + ); +}); From 339ed715611cc63efde8e82e9e5323fbe4178a2d Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Mon, 6 Jul 2026 01:55:04 +0100 Subject: [PATCH 2/7] 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. --- .../run-engine/src/engine/index.ts | 19 ++- .../engine/systems/executionSnapshotSystem.ts | 26 +++- .../engine/tests/getSnapshotsSince.test.ts | 145 ++++++++++++++++++ 3 files changed, 183 insertions(+), 7 deletions(-) diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index f3091c93b8..c1eb419a62 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/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) { diff --git a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts b/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts index 52417c1687..f555d4c0fa 100644 --- a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts +++ b/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 { // 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 diff --git a/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts b/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts index 5a21bd784a..66922b7bc3 100644 --- a/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts +++ b/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)[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(); + } + } + ); }); From 3c00f145d2fea722e3ba5cce1cc92e0553ca07d7 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Mon, 6 Jul 2026 02:16:40 +0100 Subject: [PATCH 3/7] test(run-store,run-engine): cover dedicated snapshot-write atomicity and drop an unused engine Add atomicity coverage for the dedicated (#new) leg of createExecutionSnapshot and lockRunToWorker: both must roll the snapshot back when the CompletedWaitpoint join insert fails, so a lagging replica can never serve a waitpoint-less resume. The existing coverage only exercised the legacy _completedWaitpoints path. Also drop an unused RunEngine from the duplicate-order repair test; it drives getExecutionSnapshotsSince directly and needs no Redis or worker resources. --- .../engine/tests/getSnapshotsSince.test.ts | 22 +--- .../PostgresRunStore.writeAtomicity.test.ts | 112 ++++++++++++++++++ 2 files changed, 115 insertions(+), 19 deletions(-) diff --git a/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts b/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts index 66922b7bc3..c98b2300d8 100644 --- a/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts +++ b/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts @@ -1473,24 +1473,10 @@ describe("RunEngine getSnapshotsSince", () => { // 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 }) => { + async ({ prisma }) => { 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 { + // Drives getExecutionSnapshotsSince directly, so it needs no RunEngine (Redis/worker) at all. + { const scenario = await setupTestScenario(prisma, authenticatedEnvironment, { totalWaitpoints: 1, outputSizeKB: 1, @@ -1538,8 +1524,6 @@ describe("RunEngine getSnapshotsSince", () => { 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(); } } ); diff --git a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts index 1b8a086b4c..ee07f14205 100644 --- a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts @@ -594,3 +594,115 @@ describe("createRun / createFailedRun write the run and its associated waitpoint } ); }); + +// The dedicated (#new) leg connects completed waitpoints through the `CompletedWaitpoint` join table +// (createMany), where the legacy leg uses the implicit `_completedWaitpoints` M2M. Both must commit the +// snapshot and its links together: a snapshot that commits before its links can be read waitpoint-less +// from a lagging replica, and the runner's EXECUTING branch no-ops on an empty set -> the resume hangs. +describe("createExecutionSnapshot / lockRunToWorker write the snapshot and its links atomically (dedicated)", () => { + heteroRunOpsPostgresTest( + "createExecutionSnapshot rolls the snapshot back if the CompletedWaitpoint insert fails", + async ({ prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const env = await seedEnvironment(prisma17, "dedicated", "ces_ded"); + const runId = `run_${NEW_ID_26}`; + await newStore.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_ces_ded", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + + // Force the dedicated join insert (completedWaitpoint.createMany) to fail mid-write. + await prisma17.$executeRawUnsafe('DROP TABLE "CompletedWaitpoint"'); + + await expect( + // Base client as `tx` = how the engine threads its base prisma through + // (continueRunIfUnblocked -> executionSnapshotSystem.createExecutionSnapshot(prisma, ...)). + // It is NOT an interactive transaction, so the store must still open its own to stay atomic. + newStore.createExecutionSnapshot( + { + run: { id: runId, status: "EXECUTING", attemptNumber: 1 }, + snapshot: { + executionStatus: "EXECUTING_WITH_WAITPOINTS", + description: "Run was blocked by a waitpoint.", + }, + environmentId: env.environment.id, + environmentType: "DEVELOPMENT", + projectId: env.project.id, + organizationId: env.project.id, + completedWaitpoints: [{ id: `wp_${NEW_ID_26}`, index: 0 }], + }, + prisma17 as never + ) + ).rejects.toThrow(); + + const snap = await prisma17.taskRunExecutionSnapshot.findFirst({ + where: { runId, executionStatus: "EXECUTING_WITH_WAITPOINTS" }, + }); + expect(snap).toBeNull(); + } + ); + + heteroRunOpsPostgresTest( + "lockRunToWorker rolls the snapshot and run lock back if the CompletedWaitpoint insert fails", + async ({ prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const env = await seedEnvironment(prisma17, "dedicated", "lock_ded"); + const runId = `run_${NEW_ID_26}`; + await newStore.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_lock_ded", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + const prior = await prisma17.taskRunExecutionSnapshot.findFirstOrThrow({ where: { runId } }); + + await prisma17.$executeRawUnsafe('DROP TABLE "CompletedWaitpoint"'); + + const snapshotId = `snap_${NEW_ID_26}`; + await expect( + // lockedById/lockedToVersionId/lockedQueueId are FK-free scalars on the dedicated subset, so + // synthetic ids are fine; the base client as `tx` mirrors the dequeue path (no interactive tx). + newStore.lockRunToWorker( + runId, + { + lockedAt: new Date(), + lockedById: `bwt_${NEW_ID_26}`, + lockedToVersionId: `bw_${NEW_ID_26}`, + lockedQueueId: `queue_${NEW_ID_26}`, + startedAt: new Date(), + baseCostInCents: 5, + machinePreset: "small-1x", + taskVersion: "20260601.1", + sdkVersion: "3.0.0", + cliVersion: "3.0.0", + maxDurationInSeconds: null, + snapshot: { + id: snapshotId, + previousSnapshotId: prior.id, + environmentId: env.environment.id, + environmentType: "DEVELOPMENT", + projectId: env.project.id, + organizationId: env.project.id, + completedWaitpointIds: [`wp_${NEW_ID_26}`], + completedWaitpointOrder: [`wp_${NEW_ID_26}`], + }, + }, + prisma17 as never + ) + ).rejects.toThrow(); + + const snap = await prisma17.taskRunExecutionSnapshot.findUnique({ where: { id: snapshotId } }); + expect(snap).toBeNull(); + const run = await prisma17.taskRun.findUniqueOrThrow({ where: { id: runId } }); + expect(run.status).not.toBe("DEQUEUED"); + } + ); +}); From c2e29ab8dd24f7bf34c60ffa00469ba07274792a Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Mon, 6 Jul 2026 02:25:06 +0100 Subject: [PATCH 4/7] test(run-store): assert every lock column rolls back, not just the run status The lockRunToWorker atomicity test proxied "fully rolled back" through the run status alone, which would miss a partial rollback that leaves lockedAt/lockedById and friends populated. Assert those columns are null too. --- .../run-store/src/PostgresRunStore.writeAtomicity.test.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts index ee07f14205..e4576d872f 100644 --- a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts @@ -701,8 +701,13 @@ describe("createExecutionSnapshot / lockRunToWorker write the snapshot and its l const snap = await prisma17.taskRunExecutionSnapshot.findUnique({ where: { id: snapshotId } }); expect(snap).toBeNull(); + // The whole lock write must roll back, not just the status: no lock columns may leak through. const run = await prisma17.taskRun.findUniqueOrThrow({ where: { id: runId } }); expect(run.status).not.toBe("DEQUEUED"); + expect(run.lockedAt).toBeNull(); + expect(run.lockedById).toBeNull(); + expect(run.lockedToVersionId).toBeNull(); + expect(run.lockedQueueId).toBeNull(); } ); }); From ee99f156bf61a4084324b891b46d98c15a8289bb Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Mon, 6 Jul 2026 02:36:35 +0100 Subject: [PATCH 5/7] chore(run-store): satisfy oxfmt on the new atomicity test --- .../run-store/src/PostgresRunStore.writeAtomicity.test.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts index e4576d872f..37eca9e9df 100644 --- a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts @@ -699,7 +699,9 @@ describe("createExecutionSnapshot / lockRunToWorker write the snapshot and its l ) ).rejects.toThrow(); - const snap = await prisma17.taskRunExecutionSnapshot.findUnique({ where: { id: snapshotId } }); + const snap = await prisma17.taskRunExecutionSnapshot.findUnique({ + where: { id: snapshotId }, + }); expect(snap).toBeNull(); // The whole lock write must roll back, not just the status: no lock columns may leak through. const run = await prisma17.taskRun.findUniqueOrThrow({ where: { id: runId } }); From c2593efd94d31e4dce90dc3f914cd343ca8d4c2a Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Mon, 6 Jul 2026 09:29:36 +0100 Subject: [PATCH 6/7] test(testcontainers): add a lagging-replica client for read-your-writes coverage Wraps a Prisma client to serve stale reads for chosen models (missing rows or frozen snapshots) while writes pass through, so tests can reproduce the replica lag the single-database harness never exhibits. --- internal-packages/testcontainers/src/index.ts | 1 + .../testcontainers/src/laggingReplica.ts | 80 +++++++++++++++++++ 2 files changed, 81 insertions(+) create mode 100644 internal-packages/testcontainers/src/laggingReplica.ts diff --git a/internal-packages/testcontainers/src/index.ts b/internal-packages/testcontainers/src/index.ts index 1bb9b3706b..ffa0b987a4 100644 --- a/internal-packages/testcontainers/src/index.ts +++ b/internal-packages/testcontainers/src/index.ts @@ -29,6 +29,7 @@ import { } from "./utils"; export { assertNonNullable, createPostgresContainer } from "./utils"; +export { laggingReplica, type LaggingModel } from "./laggingReplica"; export { logCleanup }; export type { MinIOConnectionConfig }; diff --git a/internal-packages/testcontainers/src/laggingReplica.ts b/internal-packages/testcontainers/src/laggingReplica.ts new file mode 100644 index 0000000000..ca09470922 --- /dev/null +++ b/internal-packages/testcontainers/src/laggingReplica.ts @@ -0,0 +1,80 @@ +// A read-replica that has NOT caught up to its primary, for reproducing read-your-writes hazards the +// zero-lag single-DB test harness can't surface. Wraps a real Prisma client; for the configured +// models it serves STALE reads instead of live ones. Other models and all writes forward untouched. +// Modes: "missing" (row not replicated yet -> null/[]/0, *OrThrow throws); "frozen" (row out of date +// -> returns the provided pre-write snapshot rows, matched on top-level scalar `where` equality). +// `wasHit` asserts the stale replica was actually consulted. + +const READ_METHODS = new Set([ + "findFirst", + "findUnique", + "findFirstOrThrow", + "findUniqueOrThrow", + "findMany", + "count", +]); + +export type LaggingModel = + | { model: string; mode: "missing" } + | { model: string; mode: "frozen"; rows: readonly Record[] }; + +// Match a frozen row against a Prisma `where` using top-level scalar equality only (enough for the +// friendlyId / id / environmentId lookups these reads use); nested filters are treated as "matches". +function whereMatches(row: Record, where: unknown): boolean { + if (where == null || typeof where !== "object") return true; + return Object.entries(where as Record).every(([key, val]) => { + if (val !== null && typeof val === "object") return true; + return row[key] === val; + }); +} + +export function laggingReplica( + real: C, + configs: readonly LaggingModel[] +): { client: C; wasHit: (model?: string) => boolean } { + const byModel = new Map(configs.map((c) => [c.model, c])); + const hits = new Set(); + + const makeModelProxy = (modelName: string, realModel: object, cfg: LaggingModel) => + new Proxy(realModel, { + get(target, prop) { + if (typeof prop === "string" && READ_METHODS.has(prop)) { + return async (args?: { where?: unknown }) => { + hits.add(modelName); + if (cfg.mode === "missing") { + if (prop === "findMany") return []; + if (prop === "count") return 0; + if (prop === "findFirstOrThrow" || prop === "findUniqueOrThrow") { + throw new Error(`laggingReplica: ${modelName}.${prop} - row not visible on replica yet`); + } + return null; + } + const matched = cfg.rows.filter((r) => whereMatches(r, args?.where)); + if (prop === "findMany") return matched; + if (prop === "count") return matched.length; + const first = matched[0] ?? null; + if ((prop === "findFirstOrThrow" || prop === "findUniqueOrThrow") && !first) { + throw new Error(`laggingReplica: ${modelName}.${prop} - no frozen row matched`); + } + return first; + }; + } + const value = (target as Record)[prop]; + return typeof value === "function" ? (value as (...a: unknown[]) => unknown).bind(target) : value; + }, + }); + + const client = new Proxy(real, { + get(target, prop) { + const cfg = typeof prop === "string" ? byModel.get(prop) : undefined; + const delegate = cfg ? (target as Record)[prop as string] : undefined; + if (cfg && delegate && typeof delegate === "object") { + return makeModelProxy(prop as string, delegate, cfg); + } + const value = (target as Record)[prop]; + return typeof value === "function" ? (value as (...a: unknown[]) => unknown).bind(target) : value; + }, + }) as C; + + return { client, wasHit: (model) => (model ? hits.has(model) : hits.size > 0) }; +} From ff7f4617015da70a681e5d15dcf340bda03a0056 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Mon, 6 Jul 2026 09:29:36 +0100 Subject: [PATCH 7/7] fix(webapp): read the batch from the primary in the checkpoint WAIT_FOR_BATCH check The pre-check decides whether to suspend a run on batchRun.resumedAt but read the batch from a replica. A batch that had just resumed the parent still looked unresumed on a lagging replica, so the service checkpointed (suspended) an already-resumed run and it stalled until a sweep. Read the primary, matching the sibling WAIT_FOR_TASK arm which already does. --- .../v3/services/createCheckpoint.server.ts | 16 ++- .../createCheckpoint.batchReplicaLag.test.ts | 65 +++++++++++ .../src/checkpointBatchReplicaLag.test.ts | 106 ++++++++++++++++++ 3 files changed, 182 insertions(+), 5 deletions(-) create mode 100644 apps/webapp/test/createCheckpoint.batchReplicaLag.test.ts create mode 100644 internal-packages/run-store/src/checkpointBatchReplicaLag.test.ts diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index 088001b826..f7359cb1d5 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -147,10 +147,14 @@ export class CreateCheckpointService extends BaseService { } case "WAIT_FOR_BATCH": { // Routed by friendlyId so a run-ops id (NEW-resident) batch is found on the owning DB; - // env-scoped to the dependent attempt's run (a batch shares its dependent's env). + // env-scoped to the dependent attempt's run (a batch shares its dependent's env). Read the + // primary: a batch that just resumed the parent may lag the replica, and a stale resumedAt + // (null) would checkpoint (suspend) an already-resumed run -> it stalls until a sweep. const batchRun = await this.runStore.findBatchTaskRunByFriendlyId( reason.batchFriendlyId, - attempt.taskRun.runtimeEnvironmentId + attempt.taskRun.runtimeEnvironmentId, + undefined, + this._prisma ); if (!batchRun) { @@ -361,11 +365,13 @@ export class CreateCheckpointService extends BaseService { }); await marqs?.cancelHeartbeat(attempt.taskRunId); - // Routed by friendlyId so a run-ops id (NEW-resident) batch is found on the owning DB; - // env-scoped to the dependent attempt's run (a batch shares its dependent's env). + // Routed by friendlyId; read the primary (this._prisma) so a just-resumed batch that still + // lags the replica doesn't leave a stale resumedAt and suspend an already-resumed run. const batchRun = await this.runStore.findBatchTaskRunByFriendlyId( reason.batchFriendlyId, - attempt.taskRun.runtimeEnvironmentId + attempt.taskRun.runtimeEnvironmentId, + undefined, + this._prisma ); if (!batchRun) { diff --git a/apps/webapp/test/createCheckpoint.batchReplicaLag.test.ts b/apps/webapp/test/createCheckpoint.batchReplicaLag.test.ts new file mode 100644 index 0000000000..39dbe7b3a0 --- /dev/null +++ b/apps/webapp/test/createCheckpoint.batchReplicaLag.test.ts @@ -0,0 +1,65 @@ +// Unit red-green for the checkpoint WAIT_FOR_BATCH replica-lag fix (createCheckpoint.server.ts). +// The service decides whether to suspend a run on `batchRun.resumedAt`; reading it from a lagging +// replica makes a just-resumed batch look unresumed -> it suspends an already-resumed run -> stall. +// The fix threads the primary (`this._prisma`) into `runStore.findBatchTaskRunByFriendlyId`. Here a +// spy runStore records which client the service passed and simulates the lag (only the primary read +// sees the fresh resumedAt): RED = no client -> stale null -> no early return; GREEN = primary -> kept alive. + +import { describe, expect, it, vi } from "vitest"; + +vi.mock("~/services/logger.server", () => ({ + logger: { debug: vi.fn(), info: vi.fn(), log: vi.fn(), error: vi.fn(), warn: vi.fn() }, +})); +vi.mock("~/v3/marqs/index.server", () => ({ + marqs: { replaceMessage: vi.fn(), cancelHeartbeat: vi.fn() }, +})); + +import { CreateCheckpointService } from "~/v3/services/createCheckpoint.server"; + +describe("checkpoint WAIT_FOR_BATCH reads the primary, not a lagging replica", () => { + it("threads the primary so an already-resumed batch keeps the run alive", async () => { + // A freezable attempt so control reaches the WAIT_FOR_BATCH arm. This object IS the primary the + // fix must thread into the batch read. + const prisma = { + taskRunAttempt: { + findFirst: async () => ({ + id: "attempt_1", + status: "EXECUTING", + taskRunId: "run_1", + taskRun: { id: "run_1", status: "EXECUTING", runtimeEnvironmentId: "env_1" }, + backgroundWorker: { id: "bw_1", deployment: { imageReference: "img:1" } }, + }), + }, + }; + + let seenClient: unknown = "NOT_CALLED"; + const runStore = { + findBatchTaskRunByFriendlyId: async ( + _friendlyId: string, + _environmentId: string, + _args: unknown, + client?: unknown + ) => { + seenClient = client; + // Lagging replica: only a read handed the primary sees the just-committed resumedAt. + return { resumedAt: client === prisma ? new Date() : null }; + }, + }; + + const service = new CreateCheckpointService(prisma as never, {} as never, runStore as never); + + let result: unknown; + try { + result = await service.call({ + attemptFriendlyId: "attempt_1", + reason: { type: "WAIT_FOR_BATCH", batchFriendlyId: "batch_1" }, + } as never); + } catch { + // Buggy path falls through the pre-check into checkpoint creation (unstubbed) and throws; the + // recorded client below is what distinguishes RED from GREEN. + } + + expect(seenClient).toBe(prisma); // the fix: primary threaded into the batch read + expect(result).toEqual({ success: false, keepRunAlive: true }); // early-return, run kept alive + }); +}); diff --git a/internal-packages/run-store/src/checkpointBatchReplicaLag.test.ts b/internal-packages/run-store/src/checkpointBatchReplicaLag.test.ts new file mode 100644 index 0000000000..753a702bb2 --- /dev/null +++ b/internal-packages/run-store/src/checkpointBatchReplicaLag.test.ts @@ -0,0 +1,106 @@ +// Repro for the checkpoint WAIT_FOR_BATCH replica-lag stall (createCheckpoint.server.ts:148-181). +// +// The service resolves the batch via `runStore.findBatchTaskRunByFriendlyId(friendlyId, envId)` with +// NO client, so the read is served from the REPLICA. Its decision hinges on `batchRun.resumedAt`: +// resumedAt set -> return keepRunAlive:true (batch already resumed; the run must keep executing) +// resumedAt null -> fall through -> create the checkpoint -> SUSPEND the run +// If the batch just resumed (primary has resumedAt), but the replica still lags (resumedAt null), the +// service suspends a run whose batch already completed -> it stalls until a sweep. The sibling +// WAIT_FOR_TASK arm reads the primary (this._prisma); only the batch arm defaults to the replica. +// +// This is invisible to the normal single-DB harness (no lag). We reintroduce the lag with the shared +// `laggingReplica` primitive: the store's replica is frozen at the pre-resume snapshot while the +// primary advances. RED = the current (client-less, replica) read -> SUSPEND. GREEN = threading the +// primary (the one-line fix) -> KEEP_ALIVE. + +import { laggingReplica, postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import type { ReadClient } from "./types.js"; + +type BatchRow = { resumedAt: Date | null } | null; + +// A line-for-line mirror of the service's WAIT_FOR_BATCH pre-check. `readClient` models the fix: the +// buggy code passes nothing (replica default); the fix threads the primary. +async function precheckWaitForBatch( + store: PostgresRunStore, + batchFriendlyId: string, + environmentId: string, + readClient?: ReadClient +): Promise<"DROP_RUN" | "KEEP_ALIVE" | "SUSPEND"> { + const batchRun = (await store.findBatchTaskRunByFriendlyId( + batchFriendlyId, + environmentId, + undefined, + readClient + )) as BatchRow; + if (!batchRun) return "DROP_RUN"; // keepRunAlive:false + if (batchRun.resumedAt) return "KEEP_ALIVE"; // batch already resumed -> run continues + return "SUSPEND"; // falls through -> checkpoint created -> run suspended +} + +async function seedEnvironment(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { organization, project, environment }; +} + +describe("checkpoint WAIT_FOR_BATCH under replica lag", () => { + postgresTest( + "a batch resumed on the primary but stale on the replica suspends an already-resumed run", + async ({ prisma }) => { + const { environment } = await seedEnvironment(prisma, "ckpt_lag"); + const friendlyId = "batch_ckpt_lag"; + const batch = await prisma.batchTaskRun.create({ + data: { friendlyId, runtimeEnvironmentId: environment.id }, + }); + + // Snapshot the batch as the replica still sees it (pre-resume: resumedAt = null). + const staleBatch = await prisma.batchTaskRun.findFirstOrThrow({ where: { id: batch.id } }); + expect(staleBatch.resumedAt).toBeNull(); + + // The batch completes and resumes the parent: primary now has resumedAt set... + await prisma.batchTaskRun.update({ where: { id: batch.id }, data: { resumedAt: new Date() } }); + + // ...but the replica lags, frozen at the pre-resume snapshot. + const replica = laggingReplica(prisma, [ + { model: "batchTaskRun", mode: "frozen", rows: [staleBatch] }, + ]); + const store = new PostgresRunStore({ + prisma, + readOnlyPrisma: replica.client, + schemaVariant: "legacy", + }); + + // RED - the current service call (no client -> replica): stale null -> SUSPEND an already-resumed run. + const buggy = await precheckWaitForBatch(store, friendlyId, environment.id); + expect(replica.wasHit("batchTaskRun")).toBe(true); // proves it read the (stale) replica + expect(buggy).toBe("SUSPEND"); // the stall bug + + // GREEN - the fix (thread the primary): sees resumedAt -> keep the run alive. + const fixed = await precheckWaitForBatch(store, friendlyId, environment.id, prisma); + expect(fixed).toBe("KEEP_ALIVE"); + } + ); +});