diff --git a/apps/webapp/app/models/waitpointTag.server.ts b/apps/webapp/app/models/waitpointTag.server.ts index 98b4e07a17..0d521a5c83 100644 --- a/apps/webapp/app/models/waitpointTag.server.ts +++ b/apps/webapp/app/models/waitpointTag.server.ts @@ -8,10 +8,14 @@ export async function createWaitpointTag({ tag, environmentId, projectId, + residency, }: { tag: string; environmentId: string; projectId: string; + // Residency from the env mint kind: a tag has no owning run, so a minted-new env pins it to NEW + // instead of defaulting to the draining legacy DB. + residency?: "NEW" | "LEGACY"; }) { if (tag.trim().length === 0) return; @@ -19,11 +23,15 @@ export async function createWaitpointTag({ while (attempts < MAX_RETRIES) { try { - return await runStore.upsertWaitpointTag({ - environmentId, - name: tag, - projectId, - }); + return await runStore.upsertWaitpointTag( + { + environmentId, + name: tag, + projectId, + }, + undefined, + residency + ); } catch (error) { if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { // Handle unique constraint violation (conflict) diff --git a/apps/webapp/app/presenters/v3/BatchListPresenter.server.ts b/apps/webapp/app/presenters/v3/BatchListPresenter.server.ts index 44a2e7c291..6d7f60316c 100644 --- a/apps/webapp/app/presenters/v3/BatchListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/BatchListPresenter.server.ts @@ -43,6 +43,28 @@ type BatchRow = { batchVersion: string; }; +// Composite keyset cursor "_". Ordering is by createdAt then id: a batch id is +// a cuid (legacy) OR a run-ops id (new), and the two schemes occupy different lexical ranges, so `id` +// alone is not a valid chronological order across the residency split. `id` is the stable tiebreak. +// Old plain-id cursors (no "_") decode to undefined and restart from page 1 (self-healing). +type BatchCursor = { createdAt: Date; id: string }; +function encodeBatchCursor(row: BatchCursor): string { + return `${row.createdAt.getTime()}_${row.id}`; +} +function decodeBatchCursor(cursor: string | undefined): BatchCursor | undefined { + if (!cursor) return undefined; + const sep = cursor.indexOf("_"); + if (sep === -1) return undefined; + const ms = Number(cursor.slice(0, sep)); + const id = cursor.slice(sep + 1); + // Number.isFinite accepts e.g. 1e20, but new Date(1e20) is Invalid Date — reject it so a malformed + // URL cursor self-heals to page 1 instead of reaching Prisma with an invalid date. + const createdAt = new Date(ms); + if (!Number.isFinite(ms) || Number.isNaN(createdAt.getTime()) || id.length === 0) + return undefined; + return { createdAt, id }; +} + export class BatchListPresenter extends BasePresenter { // Optional run-ops read-routing. Omitted (single-DB / self-host) => everything // reads from `_replica` exactly as today (passthrough). Field names are local to @@ -86,17 +108,16 @@ export class BatchListPresenter extends BasePresenter { return scan(passthrough); } - const newRows = await scan(this.readRoute.runOpsNew ?? passthrough); - - // New DB filled the page — skip the legacy read entirely; older rows fall on a later page. - if (newRows.length >= pageSize + 1) { - return newRows; - } - - const legacyRows = await scan(this.readRoute.runOpsLegacyReplica ?? passthrough); + // Always read BOTH stores and merge. The old "skip legacy when new fills the page" shortcut is + // unsound across the residency split: legacy cuid ids ("c…") sort ABOVE new run-ops ids ("0…") + // under id order, so a new-only page can hide pre-flip legacy batches that belong ahead of it. + // Ordering is by createdAt (id tiebreak), which is chronologically correct across both schemes. + const [newRows, legacyRows] = await Promise.all([ + scan(this.readRoute.runOpsNew ?? passthrough), + scan(this.readRoute.runOpsLegacyReplica ?? passthrough), + ]); - // De-dupe by id (new wins), re-sort under the page's keyset order, re-apply the over-fetch - // LIMIT — reproduces the pageSize+1 window a single union scan would return. + // De-dupe by id (new wins), re-sort under the page's keyset order, re-apply the over-fetch LIMIT. const byId = new Map(); for (const row of newRows) { byId.set(row.id, row); @@ -107,10 +128,16 @@ export class BatchListPresenter extends BasePresenter { } } - // codepoint comparator (NEVER localeCompare): BatchTaskRun.id is ASCII (cuid or run-ops id). - const sign = direction === "forward" ? 1 : -1; // forward => DESC; backward => ASC + // forward => newest-first (createdAt DESC), backward => oldest-first (ASC); id is the stable + // tiebreak (ASCII codepoint, NEVER localeCompare). + const sign = direction === "forward" ? 1 : -1; return Array.from(byId.values()) - .sort((a, b) => (a.id < b.id ? sign : a.id > b.id ? -sign : 0)) + .sort((a, b) => { + const at = a.createdAt.getTime(); + const bt = b.createdAt.getTime(); + if (at !== bt) return at < bt ? sign : -sign; + return a.id < b.id ? sign : a.id > b.id ? -sign : 0; + }) .slice(0, pageSize + 1); } @@ -212,11 +239,28 @@ export class BatchListPresenter extends BasePresenter { } const createdAtLte: Date | undefined = time.to; + // Composite (createdAt, id) keyset — see encodeBatchCursor. An old plain-id cursor decodes to + // undefined and restarts from page 1. + const keyCursor = decodeBatchCursor(cursor); + const batches = await this.#scanBatchTaskRun(pageSize, direction, (client) => client.batchTaskRun.findMany({ where: { runtimeEnvironmentId: environmentId, - ...(cursor ? { id: direction === "forward" ? { lt: cursor } : { gt: cursor } } : {}), + ...(keyCursor + ? { + OR: + direction === "forward" + ? [ + { createdAt: { lt: keyCursor.createdAt } }, + { createdAt: keyCursor.createdAt, id: { lt: keyCursor.id } }, + ] + : [ + { createdAt: { gt: keyCursor.createdAt } }, + { createdAt: keyCursor.createdAt, id: { gt: keyCursor.id } }, + ], + } + : {}), ...(friendlyId ? { friendlyId } : {}), ...(statuses && statuses.length > 0 ? { status: { in: statuses }, batchVersion: { not: "v1" } } @@ -230,7 +274,10 @@ export class BatchListPresenter extends BasePresenter { } : {}), }, - orderBy: { id: direction === "forward" ? "desc" : "asc" }, + orderBy: [ + { createdAt: direction === "forward" ? "desc" : "asc" }, + { id: direction === "forward" ? "desc" : "asc" }, + ], take: pageSize + 1, select: { id: true, @@ -248,23 +295,24 @@ export class BatchListPresenter extends BasePresenter { const hasMore = batches.length > pageSize; - //get cursors for next and previous pages + //get cursors for next and previous pages (composite (createdAt, id) keyset) + const cur = (row?: BatchRow) => (row ? encodeBatchCursor(row) : undefined); let next: string | undefined; let previous: string | undefined; switch (direction) { case "forward": - previous = cursor ? batches.at(0)?.id : undefined; + previous = cursor ? cur(batches.at(0)) : undefined; if (hasMore) { - next = batches[pageSize - 1]?.id; + next = cur(batches[pageSize - 1]); } break; case "backward": batches.reverse(); if (hasMore) { - previous = batches[1]?.id; - next = batches[pageSize]?.id; + previous = cur(batches[1]); + next = cur(batches[pageSize]); } else { - next = batches[pageSize - 1]?.id; + next = cur(batches[pageSize - 1]); } break; } diff --git a/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts b/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts index 3acb4b64cb..ace6128ac7 100644 --- a/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts +++ b/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts @@ -16,6 +16,7 @@ import { type PrismaClientOrTransaction, } from "~/db.server"; import { type AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { resolveRunIdMintKind } from "~/v3/engineVersion.server"; import { logger } from "~/services/logger.server"; import { generateHttpCallbackUrl } from "~/services/httpCallback.server"; import { @@ -58,6 +59,16 @@ const { action } = createActionApiRoute( const timeout = await parseDelay(body.timeout); + // A token (and its tags) has no owning run, so it can't co-locate. Resolve the env mint kind so a + // minted-new env creates them on the run-ops DB (NEW) instead of defaulting to the draining LEGACY + // DB by their cuid id-shape. + const mintKind = await resolveRunIdMintKind({ + organizationId: authentication.environment.organizationId, + id: authentication.environment.id, + orgFeatureFlags: authentication.environment.organization.featureFlags, + }); + const residency = mintKind === "runOpsId" ? "NEW" : "LEGACY"; + //upsert tags let tags: { id: string; name: string }[] = []; const bodyTags = typeof body.tags === "string" ? [body.tags] : body.tags; @@ -74,6 +85,7 @@ const { action } = createActionApiRoute( tag, environmentId: authentication.environment.id, projectId: authentication.environment.projectId, + residency, }); if (tagRecord) { tags.push(tagRecord); @@ -88,6 +100,7 @@ const { action } = createActionApiRoute( idempotencyKeyExpiresAt, timeout, tags: bodyTags, + standaloneResidency: residency, }); const $responseHeaders = await responseHeaders(authentication.environment); diff --git a/apps/webapp/app/v3/services/resetIdempotencyKey.server.ts b/apps/webapp/app/v3/services/resetIdempotencyKey.server.ts index 0aa44e9466..44abaefda8 100644 --- a/apps/webapp/app/v3/services/resetIdempotencyKey.server.ts +++ b/apps/webapp/app/v3/services/resetIdempotencyKey.server.ts @@ -2,6 +2,7 @@ import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { BaseService, ServiceValidationError } from "./baseService.server"; import { logger } from "~/services/logger.server"; import { getMollifierBuffer } from "~/v3/mollifier/mollifierBuffer.server"; +import { resolveRunIdMintKind } from "~/v3/engineVersion.server"; export class ResetIdempotencyKeyService extends BaseService { public async call( @@ -9,12 +10,27 @@ export class ResetIdempotencyKeyService extends BaseService { taskIdentifier: string, authenticatedEnv: AuthenticatedEnvironment ): Promise<{ id: string }> { + // The predicate has no run id to route by. When the env mints run-ops ids its runs live on NEW, + // so pin the reset to NEW and skip the wrong-DB (0-row) write to the draining legacy DB. Resolve + // this only when the org (and its flags) is loaded on the env — which the authenticated API path + // always provides; otherwise fall back to the two-store reset (correct, just not optimized). + let residency: "NEW" | "LEGACY" = "LEGACY"; + if (authenticatedEnv.organization) { + const mintKind = await resolveRunIdMintKind({ + organizationId: authenticatedEnv.organizationId, + id: authenticatedEnv.id, + orgFeatureFlags: authenticatedEnv.organization.featureFlags, + }); + residency = mintKind === "runOpsId" ? "NEW" : "LEGACY"; + } + const { count: pgCount } = await this.runStore.clearIdempotencyKey( { byPredicate: { idempotencyKey, taskIdentifier, runtimeEnvironmentId: authenticatedEnv.id, + residency, }, }, this._prisma @@ -80,6 +96,7 @@ export class ResetIdempotencyKeyService extends BaseService { idempotencyKey, taskIdentifier, runtimeEnvironmentId: authenticatedEnv.id, + residency, }, }, this._prisma diff --git a/apps/webapp/test/batchListPresenter.readroute.test.ts b/apps/webapp/test/batchListPresenter.readroute.test.ts index ee0db66dde..9b1d55c959 100644 --- a/apps/webapp/test/batchListPresenter.readroute.test.ts +++ b/apps/webapp/test/batchListPresenter.readroute.test.ts @@ -309,9 +309,9 @@ describe("BatchListPresenter run-ops read routing (PG14 control-plane/legacy + P } ); - // Split scan merge serves new + legacy in one keyset-ordered page. + // Split scan merge serves new + legacy in one createdAt-ordered page; legacy is always read. heteroPostgresTest( - "split scan merges new (PG17) + legacy (PG14) rows under the keyset order; legacy read only when new does not fill the page", + "split scan merges new (PG17) + legacy (PG14) rows under the createdAt keyset order; legacy always read", async ({ prisma14, prisma17 }) => { const ctx14 = await seedParents(prisma14, "merge"); await mirrorEnvParents(prisma17, ctx14, "merge"); @@ -323,7 +323,8 @@ describe("BatchListPresenter run-ops read routing (PG14 control-plane/legacy + P await createBatch(prisma14, ctx14, { id: "batch_d", friendlyId: "fr_d", runCount: 4 }); await createBatch(prisma17, ctx14, { id: "batch_e", friendlyId: "fr_e", runCount: 5 }); - // Case A: small page fully served by new alone => legacy NOT read. + // Case A: always-merge — legacy is read even when new could fill the page (the old skip was + // unsound across the residency split). Page is the createdAt-ordered union of both DBs. const legacySpyA = spyClient(prisma14); const presenterA = new BatchListPresenter(prisma17, prisma17, { runOpsNew: prisma17, @@ -332,9 +333,9 @@ describe("BatchListPresenter run-ops read routing (PG14 control-plane/legacy + P splitEnabled: true, }); const pageA = await presenterA.call(baseCall(ctx14, { pageSize: 2 })); - // new ids are e, c, a -> DESC: e, c (pageSize 2). pageSize+1 = 3 rows from new fills the page. - expect(pageA.batches.map((b) => b.id)).toEqual(["batch_e", "batch_c"]); - expect(legacySpyA.counts.findMany).toBe(0); + // union newest-first (createdAt, insertion order a page of 2 = e, d. + expect(pageA.batches.map((b) => b.id)).toEqual(["batch_e", "batch_d"]); + expect(legacySpyA.counts.findMany).toBeGreaterThan(0); // Case B: page needs legacy rows => legacy IS read and the merge is keyset-ordered union. const legacySpyB = spyClient(prisma14); @@ -348,8 +349,10 @@ describe("BatchListPresenter run-ops read routing (PG14 control-plane/legacy + P // union DESC of all 5: e, d, c, b, a -> first 4. expect(pageB.batches.map((b) => b.id)).toEqual(["batch_e", "batch_d", "batch_c", "batch_b"]); expect(legacySpyB.counts.findMany).toBeGreaterThan(0); - // cursor parity: next is the 4th id (pageSize-th), previous undefined (no input cursor). - expect(pageB.pagination.next).toBe("batch_b"); + // cursor parity: next is the FULL composite (createdAt, id) cursor of the 4th row (batch_b), + // previous undefined. Assert the complete value so a bad timestamp prefix can't pass. + const bRow = pageB.batches.find((b) => b.id === "batch_b")!; + expect(pageB.pagination.next).toBe(`${new Date(bRow.createdAt).getTime()}_batch_b`); expect(pageB.pagination.previous).toBeUndefined(); expect(pageB.hasAnyBatches).toBe(true); } @@ -436,7 +439,11 @@ describe("BatchListPresenter run-ops read routing (PG14 control-plane/legacy + P const hasMore = direct.length > 2; const expectedPage = direct.slice(0, 2); expect(page.batches.map((b) => b.id)).toEqual(expectedPage.map((r) => r.id)); - expect(page.pagination.next).toBe(hasMore ? expectedPage[1].id : undefined); + expect( + hasMore + ? page.pagination.next === `${expectedPage[1].createdAt.getTime()}_${expectedPage[1].id}` + : !page.pagination.next + ).toBe(true); expect(page.pagination.previous).toBeUndefined(); expect(page.hasAnyBatches).toBe(true); @@ -453,6 +460,138 @@ describe("BatchListPresenter run-ops read routing (PG14 control-plane/legacy + P } ); + // REGRESSION: a flipped org's real id mix. A cuid ("c"=0x63) sorts ABOVE a run-ops id ("0"=0x30) + // under `id DESC`, so pre-flip legacy batches belong at the top — but #scanBatchTaskRun reads new + // first, skips legacy once the page is full, and `id < cursor` can never reach a "c…" from a "0…" + // cursor. Net: pre-flip legacy batches become unreachable. + heteroPostgresTest( + "flipped org: pre-flip legacy (cuid) batches remain reachable alongside post-flip run-ops batches", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "flip"); + await mirrorEnvParents(prisma17, ctx, "flip"); + + // Pre-flip cuid batch on legacy (sorts highest); post-flip run-ops batches on new (sort below). + const LEGACY_CUID = "cm0preflipbatch0000000001"; + await createBatch(prisma14, ctx, { id: LEGACY_CUID, friendlyId: "fr_preflip", runCount: 9 }); + + const NEW_RUNOPS = [ + "06fnewbatch00000000000000a", + "06fnewbatch00000000000000b", + "06fnewbatch00000000000000c", + ]; + for (const id of NEW_RUNOPS) { + await createBatch(prisma17, ctx, { id, friendlyId: `fr_${id.slice(-1)}`, runCount: 1 }); + } + + const presenter = new BatchListPresenter(prisma17, prisma17, { + runOpsNew: prisma17, + runOpsLegacyReplica: prisma14, + controlPlaneReplica: prisma14, + splitEnabled: true, + }); + + // The pre-flip cuid batch is the oldest, so under newest-first it lands on a later page — but it + // must be REACHABLE by paging forward, not stranded behind the run-ops ids (the skip + id-order + // bug dropped it entirely: `id < ` never matches a "c…" id). + const seen = new Set(); + let cursor: string | undefined = undefined; + for (let i = 0; i < 10; i++) { + const page = await presenter.call( + baseCall(ctx, { pageSize: 2, cursor, direction: "forward" }) + ); + page.batches.forEach((b) => seen.add(b.id)); + if (!page.pagination.next) break; + cursor = page.pagination.next; + } + expect([...seen]).toContain(LEGACY_CUID); + } + ); + + // REGRESSION (ordering): even with an always-merge fix, keyset-by-id is chronologically wrong across + // the flip — a cuid ("c") sorts above a run-ops id ("0"), so an OLDER pre-flip batch outranks a NEWER + // post-flip batch. The list is "newest first", so the later-created run-ops batch must come first. + heteroPostgresTest( + "flipped org: batches list is newest-first across the flip boundary (by createdAt, not id)", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "order"); + await mirrorEnvParents(prisma17, ctx, "order"); + + const OLD_LEGACY = "cm0oldbatch00000000000001"; // cuid, created EARLIER + const NEW_RUNOPS = "06fnewbatch000000000000001"; // run-ops, created LATER + await createBatch(prisma14, ctx, { + id: OLD_LEGACY, + friendlyId: "fr_old", + createdAt: new Date(Date.now() - 3_600_000), + }); + await createBatch(prisma17, ctx, { + id: NEW_RUNOPS, + friendlyId: "fr_new", + createdAt: new Date(), + }); + + const presenter = new BatchListPresenter(prisma17, prisma17, { + runOpsNew: prisma17, + runOpsLegacyReplica: prisma14, + controlPlaneReplica: prisma14, + splitEnabled: true, + }); + const page = await presenter.call(baseCall(ctx, { pageSize: 10 })); + // Newest-first: the later-created run-ops batch outranks the older legacy one. + expect(page.batches.map((b) => b.id)).toEqual([NEW_RUNOPS, OLD_LEGACY]); + } + ); + + // Overlap regression: batches duplicated on BOTH stores (mid-migration copies) must de-dupe to one + // each without dropping rows or underfilling pages. Proves the merge needs no post-dedup refill: + // the union of each store's top-(pageSize+1) always contains the global top, and the next page + // re-queries both stores from the cursor. + heteroPostgresTest( + "flipped org: batches duplicated across both stores de-dupe without dropping rows across pagination", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "ovl"); + await mirrorEnvParents(prisma17, ctx, "ovl"); + + const at = (secondsAgo: number) => new Date(Date.now() - secondsAgo * 1000); + // a..e newest->oldest. a,b,c live on BOTH DBs (dup); d,e only on legacy. + const rows = [ + { id: "batch_ov_a", both: true, s: 1 }, + { id: "batch_ov_b", both: true, s: 2 }, + { id: "batch_ov_c", both: true, s: 3 }, + { id: "batch_ov_d", both: false, s: 4 }, + { id: "batch_ov_e", both: false, s: 5 }, + ]; + for (const r of rows) { + // A dup is a row COPY: identical (createdAt, id) on both DBs. Compute createdAt ONCE so both + // copies match exactly (two at(r.s) calls would drift by ms and the older copy would re-surface + // on the next page under the createdAt keyset). + const createdAt = at(r.s); + await createBatch(prisma14, ctx, { id: r.id, friendlyId: `fr_${r.id}`, createdAt }); + if (r.both) { + await createBatch(prisma17, ctx, { id: r.id, friendlyId: `fr_${r.id}`, createdAt }); + } + } + + const seen: string[] = []; + let cursor: string | undefined = undefined; + for (let i = 0; i < 10; i++) { + const presenter = new BatchListPresenter(prisma17, prisma17, { + runOpsNew: prisma17, + runOpsLegacyReplica: prisma14, + controlPlaneReplica: prisma14, + splitEnabled: true, + }); + const page = await presenter.call( + baseCall(ctx, { pageSize: 2, cursor, direction: "forward" }) + ); + seen.push(...page.batches.map((b) => b.id)); + if (!page.pagination.next) break; + cursor = page.pagination.next; + } + // Every batch exactly once, newest-first; the three dups collapsed to one each. + expect(seen).toEqual(["batch_ov_a", "batch_ov_b", "batch_ov_c", "batch_ov_d", "batch_ov_e"]); + } + ); + heteroRunOpsPostgresTest( "scan against dedicated RunOpsPrismaClient (splitEnabled): returns batches from new DB", async ({ prisma14, prisma17 }) => { diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 00edf64405..bf9612ba19 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -1685,6 +1685,7 @@ export class RunEngine { idempotencyKeyExpiresAt, timeout, tags, + standaloneResidency, }: { /** The run that will block on this waitpoint. Co-locates the waitpoint with the run's DB. */ runId?: string; @@ -1694,6 +1695,8 @@ export class RunEngine { idempotencyKeyExpiresAt?: Date; timeout?: Date; tags?: string[]; + /** Standalone-token residency (no owning run) from the env mint kind; ignored when `runId` is set. */ + standaloneResidency?: "NEW" | "LEGACY"; }): Promise<{ waitpoint: Waitpoint; isCached: boolean }> { return this.waitpointSystem.createManualWaitpoint({ runId, @@ -1703,6 +1706,7 @@ export class RunEngine { idempotencyKeyExpiresAt, timeout, tags, + standaloneResidency, }); } diff --git a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts b/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts index a6e6199a49..b98f2920b3 100644 --- a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts +++ b/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 { 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,7 +175,10 @@ 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 { if (waitpointIds.length === 0) return []; @@ -178,7 +186,7 @@ async function fetchWaitpointsInChunks( 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,7 +378,12 @@ 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; @@ -377,7 +391,7 @@ export async function getExecutionSnapshotsSince( } // 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 diff --git a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts index a0d5b73349..9cf372b7b1 100644 --- a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts +++ b/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 @@ -307,6 +308,7 @@ export class WaitpointSystem { idempotencyKeyExpiresAt, timeout, tags, + standaloneResidency, }: { runId?: string; environmentId: string; @@ -315,13 +317,22 @@ export class WaitpointSystem { idempotencyKeyExpiresAt?: Date; timeout?: Date; tags?: string[]; + // For a STANDALONE token (no owning `runId`): the residency the env's mint kind resolves to, so + // the token lands on the run-ops DB (NEW) in a fully-minted-new deployment instead of defaulting + // to LEGACY by its cuid id-shape. Ignored when `runId` is set (co-location wins). + standaloneResidency?: "NEW" | "LEGACY"; }): Promise<{ waitpoint: Waitpoint; isCached: boolean }> { // Co-location invariant (see createDateTimeWaitpoint): when a `runId` is supplied the waitpoint // co-locates with that run's DB and the (env,idempotencyKey) dedup is per-run (co-resident). A // standalone token (api.v1.waitpoints.tokens.ts) passes no run id — it is created without an // owner, blocked later by whichever run waits on it (possibly cross-DB, resolved by the - // run-co-resident block edge + completion fan-out), so it routes by id-shape and dedups cross-DB. No tx here. - const colocate = runId ? { coLocateWithRunId: runId } : undefined; + // run-co-resident block edge + completion fan-out). With no owner it reads the env mint kind via + // `standaloneResidency` so a minted-new env keeps its tokens on NEW; unset, it routes by id-shape. No tx here. + const colocate = runId + ? { coLocateWithRunId: runId } + : standaloneResidency + ? { residency: standaloneResidency } + : undefined; const existingWaitpoint = idempotencyKey ? await this.$.runStore.findWaitpoint( { @@ -493,7 +504,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; diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index c206f4605f..5f636ada98 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/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 { 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 { + async countPendingWaitpoints( + waitpointIds: string[], + client?: ReadClient, + // `runId` selects residency at the router; a single store has one client and ignores it. + _runId?: string + ): Promise { 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( args: Prisma.SelectSubset, tx?: PrismaClientOrTransaction @@ -2189,7 +2227,9 @@ export class PostgresRunStore implements RunStore { async findManyWaitpoints( args: Prisma.SelectSubset, - client?: ReadClient + client?: ReadClient, + // `runId` selects residency at the router; a single store has one client and ignores it. + _runId?: string ): Promise[]> { const prisma = client ?? this.readOnlyPrisma; @@ -2495,7 +2535,9 @@ export class PostgresRunStore implements RunStore { async upsertWaitpointTag( data: { environmentId: string; name: string; projectId: string; id?: string }, - tx?: PrismaClientOrTransaction + tx?: PrismaClientOrTransaction, + // `residency` selects the store at the router; a single store has one client and ignores it. + _residency?: "NEW" | "LEGACY" ): Promise { const prisma = tx ?? this.prisma; diff --git a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts index ac042d544a..a362f8228b 100644 --- a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts +++ b/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(); } ); }); diff --git a/internal-packages/run-store/src/runOpsStore.envScopedResidency.test.ts b/internal-packages/run-store/src/runOpsStore.envScopedResidency.test.ts new file mode 100644 index 0000000000..263ccbc9d1 --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.envScopedResidency.test.ts @@ -0,0 +1,109 @@ +import { describe, expect, it } from "vitest"; +import { RoutingRunStore } from "./runOpsStore.js"; +import type { RunStore } from "./types.js"; + +// Env-scoped writes with no owning run (waitpoint tags; idempotency-key reset by predicate) must +// route to NEW when the env mints run-ops ids, instead of defaulting to LEGACY / fanning a wrong-DB +// write. Pure routing: fake RunStore slots record which store the router dispatches to. + +type Call = { method: string; args: unknown[] }; +type FakeStore = RunStore & { slot: "new" | "legacy"; calls: Call[] }; + +// `clearCount` lets a test say "this store matched N rows for the reset", so the NEW-first-then-fallback +// path can be exercised (NEW matches 0 → fall back to LEGACY). +function fakeStore(slot: "new" | "legacy", clearCount = slot === "new" ? 1 : 0): FakeStore { + const calls: Call[] = []; + const rec = + (method: string, result: unknown) => + (...args: unknown[]) => { + calls.push({ method, args }); + return Promise.resolve(result); + }; + return { + slot, + calls, + upsertWaitpointTag: rec("upsertWaitpointTag", { id: slot, slot }), + clearIdempotencyKey: rec("clearIdempotencyKey", { count: clearCount }), + } as unknown as FakeStore; +} + +function buildRouter(newClearCount?: number, legacyClearCount?: number) { + const newStore = fakeStore("new", newClearCount); + const legacyStore = fakeStore("legacy", legacyClearCount); + const router = new RoutingRunStore({ + new: newStore, + legacy: legacyStore, + classify: (id: string) => (id.startsWith("new") ? "NEW" : "LEGACY"), + }); + return { router, newStore, legacyStore }; +} + +describe("RoutingRunStore.upsertWaitpointTag — residency hint for a tag with no minted id", () => { + it("routes to NEW when residency is NEW and no id is supplied", async () => { + const { router, newStore, legacyStore } = buildRouter(); + await router.upsertWaitpointTag( + { environmentId: "env", name: "t", projectId: "p" }, + undefined, + "NEW" + ); + expect(newStore.calls.map((c) => c.method)).toEqual(["upsertWaitpointTag"]); + expect(legacyStore.calls).toHaveLength(0); + }); + + it("still falls back to LEGACY when no id and no residency are supplied", async () => { + const { router, newStore, legacyStore } = buildRouter(); + await router.upsertWaitpointTag({ environmentId: "env", name: "t", projectId: "p" }); + expect(legacyStore.calls.map((c) => c.method)).toEqual(["upsertWaitpointTag"]); + expect(newStore.calls).toHaveLength(0); + }); +}); + +describe("RoutingRunStore.clearIdempotencyKey — predicate routes NEW-first when the env mints new", () => { + it("clears on NEW and does NOT touch legacy when NEW matches (post-flip key)", async () => { + const { router, newStore, legacyStore } = buildRouter(1, 0); + const result = await router.clearIdempotencyKey({ + byPredicate: { + idempotencyKey: "k", + taskIdentifier: "task", + runtimeEnvironmentId: "env", + residency: "NEW", + }, + }); + expect(newStore.calls.map((c) => c.method)).toEqual(["clearIdempotencyKey"]); + expect(legacyStore.calls).toHaveLength(0); + expect(result.count).toBe(1); + }); + + it("falls back to LEGACY when NEW matches 0 (a key held on a pre-flip legacy run)", async () => { + // The env mints new now, but this key was created before the flip → its run lives on LEGACY. + const { router, newStore, legacyStore } = buildRouter(0, 1); + const result = await router.clearIdempotencyKey({ + byPredicate: { + idempotencyKey: "k", + taskIdentifier: "task", + runtimeEnvironmentId: "env", + residency: "NEW", + }, + }); + // NEW checked first (0 rows), then LEGACY cleared the stale key — so the reset actually works. + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + expect(result.count).toBe(1); + }); + + it("still fans out a byPredicate reset with no residency (mixed residency)", async () => { + const { router, newStore, legacyStore } = buildRouter(); + await router.clearIdempotencyKey({ + byPredicate: { idempotencyKey: "k", taskIdentifier: "task", runtimeEnvironmentId: "env" }, + }); + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + }); + + it("routes byId to the owning store (unchanged)", async () => { + const { router, newStore, legacyStore } = buildRouter(); + await router.clearIdempotencyKey({ byId: { runId: "new_run", idempotencyKey: "k" } }); + expect(newStore.calls.map((c) => c.method)).toEqual(["clearIdempotencyKey"]); + expect(legacyStore.calls).toHaveLength(0); + }); +}); diff --git a/internal-packages/run-store/src/runOpsStore.mixedResidency.test.ts b/internal-packages/run-store/src/runOpsStore.mixedResidency.test.ts index ca2027d572..ba87c75ff5 100644 --- a/internal-packages/run-store/src/runOpsStore.mixedResidency.test.ts +++ b/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); } ); }); diff --git a/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts b/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts new file mode 100644 index 0000000000..e29ecb11cb --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts @@ -0,0 +1,348 @@ +import { describe, expect, it } from "vitest"; +import { RoutingRunStore } from "./runOpsStore.js"; +import type { ReadClient, RunStore } from "./types.js"; + +// Pure routing unit tests: run-keyed waitpoint/snapshot reads must route by the run id in scope +// instead of fanning out to BOTH run-ops DBs. No DB: each slot is a fake RunStore backed by a +// per-slot set of waitpoint rows / snapshot-join ids, so the assertions are purely about WHICH store +// the router queries (the co-located run's store, never the other) and about the route-then-fallback +// that keeps a rare cross-tree token visible. Correctness against real two-DB topology is covered by +// the heteroRunOpsPostgresTest suites (crossDbTokenBlock, snapshotCompletedWaitpoints, …). + +type Call = { method: string; args: unknown[] }; + +type WaitpointRow = { id: string; status: "PENDING" | "COMPLETED" }; + +type FakeConfig = { + // Waitpoint rows resident on this store, keyed by id → status (for findManyWaitpoints / + // countPendingWaitpoints[WithPresence]). + waitpoints?: WaitpointRow[]; + // Snapshot-join waitpoint ids resident on this store (for findSnapshotCompletedWaitpointIds). + snapshotWaitpointIds?: string[]; + // Whether this store has the snapshot at all (for the WithPresence variant). + snapshotPresent?: boolean; + // Edge rows to return from findManyTaskRunWaitpoints, regardless of filter (routing-only). + edges?: Array>; +}; + +type FakeStore = RunStore & { + slot: "new" | "legacy"; + calls: Call[]; + primaryReadClient: { __primary: "new" | "legacy" }; +}; + +function idsFromWhere(where: unknown): string[] | undefined { + const id = (where as { id?: unknown } | undefined)?.id; + if (typeof id === "string") return [id]; + if (id && typeof id === "object") { + const inArr = (id as { in?: unknown }).in; + if (Array.isArray(inArr)) return inArr.filter((x): x is string => typeof x === "string"); + } + return undefined; +} + +function fakeStore(slot: "new" | "legacy", config: FakeConfig = {}): FakeStore { + const calls: Call[] = []; + const rows = config.waitpoints ?? []; + const byId = new Map(rows.map((r) => [r.id, r])); + + const record = (method: string) => (args: unknown[]) => calls.push({ method, args }); + + const store: Partial = { + slot, + calls, + primaryReadClient: { __primary: slot }, + + findManyTaskRunWaitpoints: ((args: unknown, client?: ReadClient) => { + record("findManyTaskRunWaitpoints")([args, client]); + return Promise.resolve((config.edges ?? []) as never); + }) as FakeStore["findManyTaskRunWaitpoints"], + + deleteManyTaskRunWaitpoints: ((args: unknown, tx?: unknown) => { + record("deleteManyTaskRunWaitpoints")([args, tx]); + return Promise.resolve({ count: rows.length } as never); + }) as FakeStore["deleteManyTaskRunWaitpoints"], + + findSnapshotCompletedWaitpointIds: ((snapshotId: string, client?: ReadClient) => { + record("findSnapshotCompletedWaitpointIds")([snapshotId, client]); + return Promise.resolve(config.snapshotWaitpointIds ?? []); + }) as FakeStore["findSnapshotCompletedWaitpointIds"], + + findSnapshotCompletedWaitpointIdsWithPresence: ((snapshotId: string, client?: ReadClient) => { + record("findSnapshotCompletedWaitpointIdsWithPresence")([snapshotId, client]); + return Promise.resolve({ + present: config.snapshotPresent ?? false, + ids: config.snapshotWaitpointIds ?? [], + }); + }) as FakeStore["findSnapshotCompletedWaitpointIdsWithPresence"], + + findManyWaitpoints: ((args: { where?: unknown }, client?: ReadClient) => { + record("findManyWaitpoints")([args, client]); + const requested = idsFromWhere(args.where); + const result = + requested === undefined + ? rows + : requested.map((id) => byId.get(id)).filter((r): r is WaitpointRow => r != null); + return Promise.resolve(result as never); + }) as FakeStore["findManyWaitpoints"], + + countPendingWaitpoints: ((waitpointIds: string[], client?: ReadClient) => { + record("countPendingWaitpoints")([waitpointIds, client]); + const count = waitpointIds.filter((id) => byId.get(id)?.status === "PENDING").length; + return Promise.resolve(count); + }) as FakeStore["countPendingWaitpoints"], + + countPendingWaitpointsWithPresence: ((waitpointIds: string[], client?: ReadClient) => { + record("countPendingWaitpointsWithPresence")([waitpointIds, client]); + const presentIds = waitpointIds.filter((id) => byId.has(id)); + const pendingIds = presentIds.filter((id) => byId.get(id)?.status === "PENDING"); + return Promise.resolve({ pendingIds, presentIds }); + }) as FakeStore["countPendingWaitpointsWithPresence"], + }; + + return store as unknown as FakeStore; +} + +// Deterministic residency by id prefix via the classify seam (no dependence on id-shape rules). +function buildRouter(newConfig: FakeConfig = {}, legacyConfig: FakeConfig = {}) { + const newStore = fakeStore("new", newConfig); + const legacyStore = fakeStore("legacy", legacyConfig); + const router = new RoutingRunStore({ + new: newStore, + legacy: legacyStore, + classify: (id: string) => (id.startsWith("new") ? "NEW" : "LEGACY"), + }); + return { router, newStore, legacyStore }; +} + +const WRITER = { __writer: true } as unknown as ReadClient; // non-replica → escalates to own primary + +describe("RoutingRunStore.findManyTaskRunWaitpoints — route by taskRunId (no fan-out)", () => { + it("routes an edge read keyed by a NEW run id to the new store only", async () => { + const { router, newStore, legacyStore } = buildRouter(); + await router.findManyTaskRunWaitpoints({ + where: { taskRunId: "new_run" }, + select: { taskRunId: true }, + }); + expect(newStore.calls.map((c) => c.method)).toEqual(["findManyTaskRunWaitpoints"]); + expect(legacyStore.calls).toHaveLength(0); + }); + + it("routes an edge read keyed by a LEGACY run id to the legacy store only", async () => { + const { router, newStore, legacyStore } = buildRouter(); + await router.findManyTaskRunWaitpoints({ + where: { taskRunId: "legacy_run" }, + select: { taskRunId: true }, + }); + expect(legacyStore.calls.map((c) => c.method)).toEqual(["findManyTaskRunWaitpoints"]); + expect(newStore.calls).toHaveLength(0); + }); + + it("escalates a caller writer client to the owning store's own primary", async () => { + const { router, newStore } = buildRouter(); + await router.findManyTaskRunWaitpoints( + { where: { taskRunId: "new_run" }, select: { taskRunId: true } }, + WRITER + ); + expect(newStore.calls[0]?.args[1]).toEqual({ __primary: "new" }); + }); + + it("still fans out when keyed by waitpointId (no run id in scope)", async () => { + const { router, newStore, legacyStore } = buildRouter(); + await router.findManyTaskRunWaitpoints({ + where: { waitpointId: "waitpoint_x" }, + select: { taskRunId: true }, + }); + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + }); +}); + +describe("RoutingRunStore.deleteManyTaskRunWaitpoints — route by taskRunId (no fan-out)", () => { + it("deletes only on the owning store for a classifiable taskRunId", async () => { + const { router, newStore, legacyStore } = buildRouter({ waitpoints: [] }); + const result = await router.deleteManyTaskRunWaitpoints({ + where: { taskRunId: "legacy_run", id: { in: ["waitpoint_a"] } }, + }); + expect(legacyStore.calls.map((c) => c.method)).toEqual(["deleteManyTaskRunWaitpoints"]); + expect(newStore.calls).toHaveLength(0); + expect(result.count).toBe(0); + }); + + it("still fans out and sums when there is no taskRunId in the where", async () => { + const { router, newStore, legacyStore } = buildRouter(); + await router.deleteManyTaskRunWaitpoints({ where: { waitpointId: "waitpoint_x" } }); + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + }); + + it("never threads a caller tx into the routed delete", async () => { + const { router, legacyStore } = buildRouter(); + await router.deleteManyTaskRunWaitpoints({ where: { taskRunId: "legacy_run" } }, { + $fake: "cp-tx", + } as never); + expect(legacyStore.calls[0]?.args[1]).toBeUndefined(); + }); +}); + +describe("RoutingRunStore.findSnapshotCompletedWaitpointIds — route by runId", () => { + it("routes to the run's store when a runId is threaded through", async () => { + const { router, newStore, legacyStore } = buildRouter( + { snapshotWaitpointIds: ["waitpoint_n"] }, + { snapshotWaitpointIds: ["waitpoint_l"] } + ); + const ids = await router.findSnapshotCompletedWaitpointIds( + "c".repeat(25), + undefined, + "new_run" + ); + expect(ids).toEqual(["waitpoint_n"]); + expect(legacyStore.calls).toHaveLength(0); + expect(newStore.calls.map((c) => c.method)).toEqual(["findSnapshotCompletedWaitpointIds"]); + }); + + it("still fans out and merges when no runId is supplied", async () => { + const { router, newStore, legacyStore } = buildRouter( + { snapshotWaitpointIds: ["waitpoint_n"] }, + { snapshotWaitpointIds: ["waitpoint_l"] } + ); + const ids = await router.findSnapshotCompletedWaitpointIds("c".repeat(25)); + expect(ids.sort()).toEqual(["waitpoint_l", "waitpoint_n"]); + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + }); +}); + +describe("RoutingRunStore.findSnapshotCompletedWaitpointIdsWithPresence — route by runId", () => { + it("routes to the run's store when a runId is threaded through", async () => { + const { router, newStore } = buildRouter( + { snapshotWaitpointIds: ["waitpoint_n"], snapshotPresent: true }, + { snapshotWaitpointIds: ["waitpoint_l"], snapshotPresent: true } + ); + const res = await router.findSnapshotCompletedWaitpointIdsWithPresence( + "c".repeat(25), + undefined, + "legacy_run" + ); + expect(res).toEqual({ present: true, ids: ["waitpoint_l"] }); + expect(newStore.calls).toHaveLength(0); + }); + + it("still fans out (present is the OR) when no runId is supplied", async () => { + const { router } = buildRouter( + { snapshotWaitpointIds: [], snapshotPresent: false }, + { snapshotWaitpointIds: ["waitpoint_l"], snapshotPresent: true } + ); + const res = await router.findSnapshotCompletedWaitpointIdsWithPresence("c".repeat(25)); + expect(res).toEqual({ present: true, ids: ["waitpoint_l"] }); + }); +}); + +describe("RoutingRunStore.findManyWaitpoints — route by runId then fall back for missing ids", () => { + it("queries only the run's store when every requested token co-locates with the run", async () => { + const { router, newStore, legacyStore } = buildRouter({ + waitpoints: [ + { id: "waitpoint_a", status: "COMPLETED" }, + { id: "waitpoint_b", status: "COMPLETED" }, + ], + }); + const rows = (await router.findManyWaitpoints( + { where: { id: { in: ["waitpoint_a", "waitpoint_b"] } } }, + undefined, + "new_run" + )) as WaitpointRow[]; + expect(rows.map((r) => r.id).sort()).toEqual(["waitpoint_a", "waitpoint_b"]); + expect(legacyStore.calls).toHaveLength(0); + expect(newStore.calls).toHaveLength(1); + }); + + it("falls back to the other store for ONLY the ids missing on the run's store (cross-tree token)", async () => { + const { router, legacyStore } = buildRouter( + { waitpoints: [{ id: "waitpoint_local", status: "COMPLETED" }] }, + { waitpoints: [{ id: "waitpoint_crosstree", status: "COMPLETED" }] } + ); + const rows = (await router.findManyWaitpoints( + { where: { id: { in: ["waitpoint_local", "waitpoint_crosstree"] } } }, + undefined, + "new_run" + )) as WaitpointRow[]; + expect(rows.map((r) => r.id).sort()).toEqual(["waitpoint_crosstree", "waitpoint_local"]); + // The fallback leg is queried with ONLY the missing id, never the whole set. + const fallbackCall = legacyStore.calls[0]; + expect(fallbackCall?.method).toBe("findManyWaitpoints"); + const fallbackWhere = (fallbackCall!.args[0] as { where?: unknown }).where; + expect(idsFromWhere(fallbackWhere)).toEqual(["waitpoint_crosstree"]); + }); + + it("still fans out (NEW-wins dedup) when no runId is supplied", async () => { + const { router, newStore, legacyStore } = buildRouter( + { waitpoints: [{ id: "waitpoint_a", status: "COMPLETED" }] }, + { waitpoints: [{ id: "waitpoint_a", status: "PENDING" }] } + ); + const rows = (await router.findManyWaitpoints({ + where: { id: { in: ["waitpoint_a"] } }, + })) as WaitpointRow[]; + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + // NEW-wins: the deduped row is the NEW copy (COMPLETED), not the stale legacy PENDING one. + expect(rows).toEqual([{ id: "waitpoint_a", status: "COMPLETED" }]); + }); +}); + +describe("RoutingRunStore.countPendingWaitpoints — route by runId then partition-fallback", () => { + it("counts on the run's store only when every waitpoint co-locates with the run", async () => { + const { router, newStore, legacyStore } = buildRouter({ + waitpoints: [ + { id: "waitpoint_a", status: "PENDING" }, + { id: "waitpoint_b", status: "COMPLETED" }, + ], + }); + const count = await router.countPendingWaitpoints( + ["waitpoint_a", "waitpoint_b"], + undefined, + "new_run" + ); + expect(count).toBe(1); + expect(legacyStore.calls).toHaveLength(0); + expect(newStore.calls.map((c) => c.method)).toEqual(["countPendingWaitpointsWithPresence"]); + }); + + it("counts a cross-tree pending token via the fallback so a blocked run is not prematurely unblocked", async () => { + // The classic crossDbTokenBlock shape: a LEGACY run blocks on a token resident on the NEW DB. + const { router, newStore } = buildRouter( + { waitpoints: [{ id: "waitpoint_crosstree", status: "PENDING" }] }, + { waitpoints: [] } + ); + const count = await router.countPendingWaitpoints( + ["waitpoint_crosstree"], + undefined, + "legacy_run" + ); + expect(count).toBe(1); + // Fallback queried the other store with ONLY the id missing on the run's store. + expect(newStore.calls.map((c) => c.method)).toEqual(["countPendingWaitpoints"]); + expect(newStore.calls[0]?.args[0]).toEqual(["waitpoint_crosstree"]); + }); + + it("trusts the run's store for an id present there (COMPLETED) even if a stale mirror is PENDING elsewhere", async () => { + const { router, legacyStore } = buildRouter( + { waitpoints: [{ id: "waitpoint_a", status: "COMPLETED" }] }, + { waitpoints: [{ id: "waitpoint_a", status: "PENDING" }] } + ); + const count = await router.countPendingWaitpoints(["waitpoint_a"], undefined, "new_run"); + // Present on the run's store → not in the missing set → the other store is never consulted. + expect(count).toBe(0); + expect(legacyStore.calls).toHaveLength(0); + }); + + it("still fans out and sums when no runId is supplied", async () => { + const { router, newStore, legacyStore } = buildRouter( + { waitpoints: [{ id: "waitpoint_a", status: "PENDING" }] }, + { waitpoints: [{ id: "waitpoint_b", status: "PENDING" }] } + ); + const count = await router.countPendingWaitpoints(["waitpoint_a", "waitpoint_b"]); + expect(count).toBe(2); + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + }); +}); diff --git a/internal-packages/run-store/src/runOpsStore.standaloneWaitpointResidency.test.ts b/internal-packages/run-store/src/runOpsStore.standaloneWaitpointResidency.test.ts new file mode 100644 index 0000000000..1584536790 --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.standaloneWaitpointResidency.test.ts @@ -0,0 +1,148 @@ +// A STANDALONE waitpoint/token (wait.createToken with no owning run) has an always-cuid id, so +// id-shape routing sends it to LEGACY. In a fully-minted-new deployment (env mints run-ops ids, +// legacy draining) that strands a new run's token on the draining DB and, under a read-only legacy +// connection, fails the write outright. A standalone token must instead read the env mint kind and +// land on the run's DB (NEW) via the `residency` co-location hint. Real two-DB topology; never mocked. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; + +const CUID_WAITPOINT = "c".repeat(25); // cuid id → classifies LEGACY (the always-wrong key for a new token) + +function makeRouter(prisma14: PrismaClient, prisma17: RunOpsPrismaClient) { + const newStore = new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); + const legacyStore = new PostgresRunStore({ + prisma: prisma14, + readOnlyPrisma: prisma14, + schemaVariant: "legacy", + }); + return new RoutingRunStore({ new: newStore, legacy: legacyStore }); +} + +// Seed real org/project/env on #legacy so its FKs are satisfied whichever store the waitpoint lands on +// (the #new dedicated subset is FK-free). Returns the ids the waitpoint rows reference. +async function seedLegacyEnv(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: "PRODUCTION", + slug: `prod-${suffix}`, + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_prod_${suffix}`, + pkApiKey: `pk_prod_${suffix}`, + shortcode: `short_${suffix}`, + maximumConcurrencyLimit: 10, + }, + }); + return { projectId: project.id, environmentId: environment.id }; +} + +describe("run-ops split — a standalone token pins to NEW via the residency hint (not legacy by id-shape)", () => { + heteroRunOpsPostgresTest( + "createWaitpoint with residency NEW lands a cuid token on #new, absent from #legacy", + async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => { + const router = makeRouter(prisma14, prisma17); + const { projectId, environmentId } = await seedLegacyEnv(prisma14, "swr_create"); + + await router.createWaitpoint( + { + data: { + id: CUID_WAITPOINT, + friendlyId: "wp_standalone", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${CUID_WAITPOINT}`, + userProvidedIdempotencyKey: false, + projectId, + environmentId, + }, + }, + undefined, + { residency: "NEW" } + ); + + // GREEN: the residency hint routes the cuid token to #new. RED: id-shape sends it to #legacy. + expect(await prisma17.waitpoint.count({ where: { id: CUID_WAITPOINT } })).toBe(1); + expect(await prisma14.waitpoint.count({ where: { id: CUID_WAITPOINT } })).toBe(0); + } + ); + + heteroRunOpsPostgresTest( + "upsertWaitpoint with residency NEW lands a cuid token on #new, absent from #legacy", + async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => { + const router = makeRouter(prisma14, prisma17); + const { projectId, environmentId } = await seedLegacyEnv(prisma14, "swr_upsert"); + + await router.upsertWaitpoint( + { + where: { id: CUID_WAITPOINT }, + create: { + id: CUID_WAITPOINT, + friendlyId: "wp_standalone_upsert", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${CUID_WAITPOINT}`, + userProvidedIdempotencyKey: false, + projectId, + environmentId, + }, + update: {}, + }, + undefined, + { residency: "NEW" } + ); + + expect(await prisma17.waitpoint.count({ where: { id: CUID_WAITPOINT } })).toBe(1); + expect(await prisma14.waitpoint.count({ where: { id: CUID_WAITPOINT } })).toBe(0); + } + ); + + heteroRunOpsPostgresTest( + "coLocateWithRunId wins over residency (a co-located waitpoint inherits its run, ignores the flag)", + async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => { + const router = makeRouter(prisma14, prisma17); + const { projectId, environmentId } = await seedLegacyEnv(prisma14, "swr_colocate"); + const legacyRunId = `run_${"c".repeat(25)}`; // cuid run → #legacy + + await router.createWaitpoint( + { + data: { + id: CUID_WAITPOINT, + friendlyId: "wp_colocated", + type: "DATETIME", + status: "PENDING", + idempotencyKey: `idem_${CUID_WAITPOINT}`, + userProvidedIdempotencyKey: false, + projectId, + environmentId, + }, + }, + undefined, + { coLocateWithRunId: legacyRunId, residency: "NEW" } + ); + + // Co-location must win → the waitpoint lands on the run's store (#legacy), NOT the residency hint. + expect(await prisma14.waitpoint.count({ where: { id: CUID_WAITPOINT } })).toBe(1); + expect(await prisma17.waitpoint.count({ where: { id: CUID_WAITPOINT } })).toBe(0); + } + ); +}); diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 4a98b8f003..f6b613a393 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -503,7 +503,7 @@ export class RoutingRunStore implements RunStore { return (await this.#routeOrNewForWrite(runId)).updateMetadata(runId, data, options); } - clearIdempotencyKey( + async clearIdempotencyKey( params: ClearIdempotencyKeyInput, tx?: PrismaClientOrTransaction ): Promise<{ count: number }> { @@ -513,7 +513,18 @@ export class RoutingRunStore implements RunStore { const store = this.#route(params.byId.runId); return store.clearIdempotencyKey(params, undefined); } - // `byFriendlyIds` / `byPredicate` can span mixed residency — fan out and sum. + // A `byPredicate` whose env mints run-ops ids has NEW-born runs, so check NEW first. But a key + // minted BEFORE the org flipped still lives on a LEGACY-resident run (idempotency TTL up to 30d), + // so fall back to LEGACY when NEW matched nothing — otherwise the reset 404s and the stale legacy + // key keeps deduping. In the steady (fully-drained) state NEW matches and legacy is never touched. + if ("byPredicate" in params && params.byPredicate?.residency === "NEW") { + const fromNew = await this.#new.clearIdempotencyKey(params, undefined); + if (fromNew.count > 0) { + return fromNew; + } + const fromLegacy = await this.#legacy.clearIdempotencyKey(params); + return { count: fromNew.count + fromLegacy.count }; + } return Promise.all([ this.#new.clearIdempotencyKey(params), this.#legacy.clearIdempotencyKey(params), @@ -977,13 +988,22 @@ export class RoutingRunStore implements RunStore { return store.createExecutionSnapshot(input, undefined); } - // Snapshot ids are cuids (they always classify LEGACY), and a snapshot's CompletedWaitpoint join - // co-locates with its run, which may live on either store, so fan out to BOTH and merge (like - // findWaitpointCompletedSnapshotIds) rather than route by the un-classifiable snapshot id. + // The CompletedWaitpoint join co-locates with the snapshot, which co-locates with its run. When the + // caller threads the run id (executionSnapshotSystem has it in scope), route to the run's store — no + // fan-out. Snapshot ids are cuids (they always classify LEGACY), so absent a run id we can't route + // by the snapshot id and must fan out to BOTH stores and merge (like findWaitpointCompletedSnapshotIds). async findSnapshotCompletedWaitpointIds( snapshotId: string, - client?: ReadClient + client?: ReadClient, + runId?: string ): Promise { + if (runId !== undefined) { + const store = this.#routeOrNew(runId); + return store.findSnapshotCompletedWaitpointIds( + snapshotId, + RoutingRunStore.#ownPrimary(store, client) + ); + } const [fromNew, fromLegacy] = await Promise.all([ this.#new.findSnapshotCompletedWaitpointIds( snapshotId, @@ -997,12 +1017,20 @@ export class RoutingRunStore implements RunStore { return uniqueStrings([...fromNew, ...fromLegacy]); } - // Snapshot-id has no residency to route on, so fan out; the snapshot lives on exactly one store, so - // `present` is the OR and `ids` the union. + // As above: route to the run's store when the run id is threaded through, else fan out (the snapshot + // lives on exactly one store, so `present` is the OR and `ids` the union). async findSnapshotCompletedWaitpointIdsWithPresence( snapshotId: string, - client?: ReadClient + client?: ReadClient, + runId?: string ): Promise<{ present: boolean; ids: string[] }> { + if (runId !== undefined) { + const store = this.#routeOrNew(runId); + return store.findSnapshotCompletedWaitpointIdsWithPresence( + snapshotId, + RoutingRunStore.#ownPrimary(store, client) + ); + } const [fromNew, fromLegacy] = await Promise.all([ this.#new.findSnapshotCompletedWaitpointIdsWithPresence( snapshotId, @@ -1070,20 +1098,74 @@ export class RoutingRunStore implements RunStore { return (await this.#routeOrNewForWrite(params.runId)).blockRunWithWaitpointEdges(edges); } - // A run's waitpoints can be scattered across both stores (drain in flight), so count on - // each and sum rather than assume one home. - async countPendingWaitpoints(waitpointIds: string[], client?: ReadClient): Promise { + // A run's blocking waitpoints mostly co-locate with the run; only a cross-tree token (a standalone + // MANUAL token, or waitForRun across trees) lives on the other DB. When the run id is threaded + // through, route to the run's store and fall back to the other DB for ONLY the ids absent there — + // partitioning by found-ness so a present-but-completed id is trusted from the run's store and a + // cross-tree pending token is still counted (never undercounted, which would prematurely unblock). + // Absent a run id, count on each and sum (a caller with no run id in scope). + async countPendingWaitpoints( + waitpointIds: string[], + client?: ReadClient, + runId?: string + ): Promise { + if (runId === undefined) { + const [fromNew, fromLegacy] = await Promise.all([ + this.#new.countPendingWaitpoints( + waitpointIds, + RoutingRunStore.#ownPrimary(this.#new, client) + ), + this.#legacy.countPendingWaitpoints( + waitpointIds, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), + ]); + return fromNew + fromLegacy; + } + + if (waitpointIds.length === 0) { + return 0; + } + const runStore = this.#routeOrNew(runId); + const otherStore = runStore === this.#new ? this.#legacy : this.#new; + const { pendingIds, presentIds } = await runStore.countPendingWaitpointsWithPresence( + waitpointIds, + RoutingRunStore.#ownPrimary(runStore, client) + ); + const present = new Set(presentIds); + const missing = waitpointIds.filter((id) => !present.has(id)); + if (missing.length === 0) { + return pendingIds.length; + } + const otherPending = await otherStore.countPendingWaitpoints( + missing, + RoutingRunStore.#ownPrimary(otherStore, client) + ); + return pendingIds.length + otherPending; + } + + // Fan out and union: an id lives on exactly one store in steady state (a drain-mirror can put it on + // both), so the union of pending/present ids dedups a mirror correctly. Not on any hot path — the + // router's countPendingWaitpoints routes to a sub-store's variant directly — but required by the + // interface and correct for any defensive caller. + async countPendingWaitpointsWithPresence( + waitpointIds: string[], + client?: ReadClient + ): Promise<{ pendingIds: string[]; presentIds: string[] }> { const [fromNew, fromLegacy] = await Promise.all([ - this.#new.countPendingWaitpoints( + this.#new.countPendingWaitpointsWithPresence( waitpointIds, RoutingRunStore.#ownPrimary(this.#new, client) ), - this.#legacy.countPendingWaitpoints( + this.#legacy.countPendingWaitpointsWithPresence( waitpointIds, RoutingRunStore.#ownPrimary(this.#legacy, client) ), ]); - return fromNew + fromLegacy; + return { + pendingIds: uniqueStrings([...fromNew.pendingIds, ...fromLegacy.pendingIds]), + presentIds: uniqueStrings([...fromNew.presentIds, ...fromLegacy.presentIds]), + }; } // A waitpoint co-locates with the OWNER it points at, in priority order: an explicit @@ -1100,10 +1182,34 @@ export class RoutingRunStore implements RunStore { const data = (args as { data?: unknown }).data; const ownerRunId = scalarStringField(data, "completedByTaskRunId"); const ownerBatchId = scalarStringField(data, "completedByBatchId"); - const routeId = - opts?.coLocateWithRunId ?? ownerRunId ?? ownerBatchId ?? RoutingRunStore.#waitpointId(data); - const { store, tx: routedTx } = this.#routeWaitpointWrite(routeId, tx); - return store.createWaitpoint(args, routedTx); + const store = this.#waitpointWriteStore( + opts?.coLocateWithRunId ?? ownerRunId ?? ownerBatchId, + opts?.residency, + RoutingRunStore.#waitpointId(data) + ); + // Never forward the caller's tx into a routed write (matches #routeWaitpointWrite). + return store.createWaitpoint(args, undefined); + } + + // Resolve the store a waitpoint WRITE lands on, in precedence order: an explicit OWNER id + // (coLocateWithRunId / completedBy run|batch) always wins — a co-located waitpoint inherits its + // owner's residency by id-shape. With no owner, a STANDALONE token reads the env mint kind via the + // `residency` hint (NEW when the env mints run-ops ids). Else fall back to the waitpoint's own + // id-shape (always cuid → LEGACY). + #waitpointWriteStore( + ownerId: string | undefined, + residency: Residency | undefined, + waitpointId: string | undefined + ): RunStore { + if (ownerId !== undefined) { + return this.#classifySafe(ownerId) === "NEW" ? this.#new : this.#legacy; + } + if (residency !== undefined) { + return residency === "NEW" ? this.#new : this.#legacy; + } + return typeof waitpointId === "string" && this.#classifySafe(waitpointId) === "NEW" + ? this.#new + : this.#legacy; } upsertWaitpoint( @@ -1112,13 +1218,13 @@ export class RoutingRunStore implements RunStore { opts?: WaitpointColocationOptions ): Promise> { // `coLocateWithRunId` (the owning run) wins so a DATETIME/MANUAL wait waitpoint lands on its - // run's DB; otherwise key by create.id (always the minted waitpoint id), then where. - const routeId = - opts?.coLocateWithRunId ?? + // run's DB; else a standalone token reads the `residency` hint; else key by create.id (always the + // minted waitpoint id), then where. + const waitpointId = RoutingRunStore.#waitpointId((args as { create?: unknown }).create) ?? RoutingRunStore.#waitpointId((args as { where?: unknown }).where); - const { store, tx: routedTx } = this.#routeWaitpointWrite(routeId, tx); - return store.upsertWaitpoint(args, routedTx); + const store = this.#waitpointWriteStore(opts?.coLocateWithRunId, opts?.residency, waitpointId); + return store.upsertWaitpoint(args, undefined); } // Probe by id (drain may have relocated it); an idempotency-key lookup with no id routes by @@ -1177,41 +1283,83 @@ export class RoutingRunStore implements RunStore { async findManyWaitpoints( args: Prisma.SelectSubset, - client?: ReadClient + client?: ReadClient, + runId?: string ): Promise[]> { const { scalarArgs, relations } = splitWaitpointRelationProjection( args as Record ); + const rows = (await this.#collectManyWaitpoints( + scalarArgs, + client, + runId + )) as Prisma.WaitpointGetPayload[]; + for (const row of rows) { + await this.#reresolveWaitpointRelationsCrossDb( + row as Record, + relations, + client + ); + } + return rows; + } + + // Collect the scalar waitpoint rows (relation re-resolution happens in the caller). With a run id in + // scope and a bounded id set to partition on, route to the run's store and fall back to the other DB + // for ONLY the ids missing there (a rare cross-tree token) — the two legs are disjoint by + // construction, so no dedup is needed. Otherwise fan out to BOTH and dedup by id NEW-wins. + async #collectManyWaitpoints( + scalarArgs: Record, + client: ReadClient | undefined, + runId: string | undefined + ): Promise[]> { + if (runId !== undefined) { + const requestedIds = idListFromWhere((scalarArgs.where ?? {}) as Prisma.TaskRunWhereInput); + if (requestedIds !== undefined) { + const runStore = this.#routeOrNew(runId); + const fromRun = (await runStore.findManyWaitpoints( + scalarArgs as Prisma.WaitpointFindManyArgs, + RoutingRunStore.#ownPrimary(runStore, client) + )) as Record[]; + const foundIds = new Set( + fromRun.map((w) => w.id).filter((id): id is string => typeof id === "string") + ); + const missing = requestedIds.filter((id) => !foundIds.has(id)); + if (missing.length === 0) { + return fromRun; + } + const otherStore = runStore === this.#new ? this.#legacy : this.#new; + const fromOther = (await otherStore.findManyWaitpoints( + narrowArgsToIds(scalarArgs, missing) as Prisma.WaitpointFindManyArgs, + RoutingRunStore.#ownPrimary(otherStore, client) + )) as Record[]; + return [...fromRun, ...fromOther]; + } + // No bounded id set to partition on → fall through to the fan-out path. + } + const [fromNew, fromLegacy] = await Promise.all([ this.#new.findManyWaitpoints( - scalarArgs as typeof args, + scalarArgs as Prisma.WaitpointFindManyArgs, RoutingRunStore.#ownPrimary(this.#new, client) - ), + ) as Promise[]>, this.#legacy.findManyWaitpoints( - scalarArgs as typeof args, + scalarArgs as Prisma.WaitpointFindManyArgs, RoutingRunStore.#ownPrimary(this.#legacy, client) - ), + ) as Promise[]>, ]); // A token mirrored onto both DBs during drain appears in BOTH legs; dedup by id with NEW-wins // (the NEW copy is authoritative once a run migrates), matching the router's NEW-wins invariant // (#findRunsOpen). Without this, edge-waitpoint hydration could read a stale LEGACY status and // strand the run. Rows whose projection omits `id` can't be deduped and pass through. - const byId = new Map>(); - const passthrough: Prisma.WaitpointGetPayload[] = []; + const byId = new Map>(); + const passthrough: Record[] = []; for (const w of [...fromLegacy, ...fromNew]) { - const id = (w as { id?: unknown }).id; + const id = w.id; if (typeof id === "string") byId.set(id, w); else passthrough.push(w); } - const rows = [...byId.values(), ...passthrough]; - for (const row of rows) { - await this.#reresolveWaitpointRelationsCrossDb( - row as Record, - relations, - client - ); - } - return rows; + return [...byId.values(), ...passthrough]; } // Re-resolve a waitpoint's group-A relations across BOTH DBs and attach them to `row`. Each target @@ -1391,17 +1539,32 @@ export class RoutingRunStore implements RunStore { args as Record ); - const [fromNew, fromLegacy] = await Promise.all([ - this.#new.findManyTaskRunWaitpoints( - scalarArgs as typeof args, - RoutingRunStore.#ownPrimary(this.#new, client) - ), - this.#legacy.findManyTaskRunWaitpoints( + // An edge always co-locates with its RUN (blockRunWithWaitpointEdges routes the write by runId), + // so a read keyed by a classifiable `taskRunId` routes to that run's store — no fan-out, no dedup. + // Only a `waitpointId`/`batchId` predicate (no run id) still fans across both stores. + const taskRunId = whereFieldString( + (args.where as { taskRunId?: Prisma.TaskRunWhereInput["id"] } | undefined)?.taskRunId + ); + let edges: Record[]; + if (taskRunId !== undefined) { + const store = this.#routeOrNew(taskRunId); + edges = (await store.findManyTaskRunWaitpoints( scalarArgs as typeof args, - RoutingRunStore.#ownPrimary(this.#legacy, client) - ), - ]); - const edges = dedupeEdgesById([...fromNew, ...fromLegacy]) as Record[]; + RoutingRunStore.#ownPrimary(store, client) + )) as Record[]; + } else { + const [fromNew, fromLegacy] = await Promise.all([ + this.#new.findManyTaskRunWaitpoints( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(this.#new, client) + ), + this.#legacy.findManyTaskRunWaitpoints( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), + ]); + edges = dedupeEdgesById([...fromNew, ...fromLegacy]) as Record[]; + } if (waitpoint) { await this.#hydrateEdgeWaitpointsCrossDb(edges, waitpoint, client); @@ -1470,9 +1633,15 @@ export class RoutingRunStore implements RunStore { args: Prisma.TaskRunWaitpointDeleteManyArgs, tx?: PrismaClientOrTransaction ): Promise { - // Edges co-locate with their RUN, not their waitpoint, so a waitpointId/taskRunId predicate may - // match edges on either store; delete from both so no cross-DB edge keeps a run blocked. The - // caller's `tx` is never forwarded — each leg deletes on its own store's client. + // An edge always co-locates with its RUN, so a delete keyed by a classifiable `taskRunId` routes + // to that run's store — no fan-out. Only a `waitpointId`-keyed delete (no run id) still deletes + // from both stores. The caller's `tx` is never forwarded — each leg deletes on its own client. + const taskRunId = whereFieldString( + (args.where as { taskRunId?: Prisma.TaskRunWhereInput["id"] } | undefined)?.taskRunId + ); + if (taskRunId !== undefined) { + return (await this.#routeOrNewForWrite(taskRunId)).deleteManyTaskRunWaitpoints(args); + } const [fromNew, fromLegacy] = await Promise.all([ this.#new.deleteManyTaskRunWaitpoints(args), this.#legacy.deleteManyTaskRunWaitpoints(args), @@ -1716,10 +1885,13 @@ export class RoutingRunStore implements RunStore { // residency-aware, findManyWaitpointTags must de-dupe by (environmentId, name) or names will duplicate. upsertWaitpointTag( data: { environmentId: string; name: string; projectId: string; id?: string }, - tx?: PrismaClientOrTransaction + tx?: PrismaClientOrTransaction, + residency?: Residency ): Promise { - const { store, tx: routedTx } = this.#routeWaitpointWrite(data.id, tx); - return store.upsertWaitpointTag(data, routedTx); + // No owning run; route by the env's residency hint when present, else a minted id-shape, else + // fall back to LEGACY (same precedence as a standalone waitpoint). Caller tx is never forwarded. + const store = this.#waitpointWriteStore(undefined, residency, data.id); + return store.upsertWaitpointTag(data, undefined); } // A tag keyed by (environmentId, name) can exist on BOTH DBs for one env (dual-resident, no @@ -1872,6 +2044,15 @@ function narrowToIds(args: FindRunsArgs, ids: string[]): FindRunsArgs { return { ...args, where: { ...args.where, id: { in: ids } } }; } +// Clone find-many args, replacing the `id` filter with `{ in: ids }` while keeping any other `where` +// conditions and the projection/ordering intact. Used to re-query only the ids missing on the first leg. +function narrowArgsToIds(args: Record, ids: string[]): Record { + return { + ...args, + where: { ...((args.where as Record) ?? {}), id: { in: ids } }, + }; +} + // Merge edge rows from both stores, keeping one per edge `id` (NEW seen last wins). Rows whose // projection omits `id` can't be deduped, so they pass through unchanged. function dedupeEdgesById(rows: R[]): R[] { diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index 399aab5b07..b55e8f968c 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -257,7 +257,15 @@ export type FinalizeRunData = { export type ClearIdempotencyKeyInput = | { byId: { runId: string; idempotencyKey: string }; byPredicate?: never; byFriendlyIds?: never } | { - byPredicate: { idempotencyKey: string; taskIdentifier: string; runtimeEnvironmentId: string }; + byPredicate: { + idempotencyKey: string; + taskIdentifier: string; + runtimeEnvironmentId: string; + // A predicate has no run id to route by, so it fans out to both stores. When the env mints + // run-ops ids its matching runs live on NEW, so `residency: "NEW"` routes to NEW only and + // avoids a wrong-DB (0-row) write to the draining legacy DB. Omit to fan out (mixed residency). + residency?: Residency; + }; byId?: never; byFriendlyIds?: never; } @@ -325,6 +333,15 @@ export interface ForWaitpointCompletionContext { */ export interface WaitpointColocationOptions { coLocateWithRunId?: string; + /** + * Residency for a STANDALONE waitpoint that has no owning run to co-locate with (e.g. a + * `wait.createToken()` token created via the env-scoped API). Its minted id is always a cuid, so + * id-shape routing would always send it to LEGACY; a standalone token instead reads the env mint + * kind and pins here (NEW when the env mints run-ops ids), so a fully-minted-new deployment keeps + * its tokens off the draining legacy DB. Ignored when `coLocateWithRunId` (or an owner id) is set — + * a co-located waitpoint always inherits its run/batch residency, never the flag. + */ + residency?: Residency; } export interface RunStore { @@ -674,12 +691,19 @@ export interface RunStore { ): Promise>; // Implicit-join group - findSnapshotCompletedWaitpointIds(snapshotId: string, client?: ReadClient): Promise; + /** `runId` (when known) routes to the run's store — the snapshot + its join co-locate with the run; + * omit it and the router fans out (the cuid snapshot id alone can't say which store holds the join). */ + findSnapshotCompletedWaitpointIds( + snapshotId: string, + client?: ReadClient, + runId?: string + ): Promise; /** As above, but reports in the SAME read whether the snapshot is visible on the reader: `present=false` * means this reader lacks the snapshot, so its empty id list is not authoritative (repair from primary). */ findSnapshotCompletedWaitpointIdsWithPresence( snapshotId: string, - client?: ReadClient + client?: ReadClient, + runId?: string ): Promise<{ present: boolean; ids: string[] }>; /** Run ids connected to a waitpoint (WaitpointRunConnection / `_WaitpointRunConnections`), this DB only. */ findWaitpointConnectedRunIds(waitpointId: string, client?: ReadClient): Promise; @@ -694,7 +718,20 @@ export interface RunStore { batchIndex?: number; tx?: PrismaClientOrTransaction; }): Promise; - countPendingWaitpoints(waitpointIds: string[], client?: ReadClient): Promise; + /** `runId` (when known) routes to the run's store and falls back to the other DB only for ids absent + * there (a rare cross-tree token), instead of fanning the count out to both DBs on every call. */ + countPendingWaitpoints( + waitpointIds: string[], + client?: ReadClient, + runId?: string + ): Promise; + /** Which of the given 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. */ + countPendingWaitpointsWithPresence( + waitpointIds: string[], + client?: ReadClient + ): Promise<{ pendingIds: string[]; presentIds: string[] }>; // Waitpoint group createWaitpoint( @@ -717,9 +754,12 @@ export interface RunStore { findWaitpointOnPrimary( args: Prisma.SelectSubset ): Promise | null>; + /** `runId` (when known) routes to the run's store and falls back to the other DB only for the ids + * missing there, instead of fanning every token read out to both DBs. */ findManyWaitpoints( args: Prisma.SelectSubset, - client?: ReadClient + client?: ReadClient, + runId?: string ): Promise[]>; updateWaitpoint( args: Prisma.SelectSubset, @@ -842,7 +882,11 @@ export interface RunStore { // de-dupes by id in case tag ids ever become residency-aware. upsertWaitpointTag( data: { environmentId: string; name: string; projectId: string; id?: string }, - tx?: PrismaClientOrTransaction + tx?: PrismaClientOrTransaction, + // A tag has no owning run to co-locate with; when no minted `id` pins it by id-shape, a + // minted-new env's tags read this residency (NEW) so they land with the env's tokens/runs + // instead of defaulting to LEGACY. Single-store impls ignore it. + residency?: Residency ): Promise; findManyWaitpointTags( args: {