From 831b341516e04463b9167ab92e273044b5becc11 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 27 Aug 2026 15:59:26 +0100 Subject: [PATCH 1/2] test(run-store,webapp): cover the run-ops router at three shards The fake-slot router suites only ever built two stores, so several rules that only apply above two were untestable by construction. clearIdempotencyKey is the sole caller of the "every other shard" helper, and with two stores that helper returns a single entry, which hides a take-the-first bug. The absent-id partition has the same blind spot: a gen-2 id and a cuid select the same store when only one other store exists. Three suites now run at two shards and at three. Each new assertion was checked by mutating the router to the failure it guards, confirming the three-shard arm fails while the two-shard arm still passes, then reverting. Two webapp tests defined their own local RoutingRunStore. They compiled against a two-store model whatever the real class did, and one described a routing rule the code never implemented. Both now build the real router over the two Postgres stores they already create. No production code changes. --- .../performTaskRunAlertsStoreRouting.test.ts | 154 +----- .../updateMetadataStoreRoutingHetero.test.ts | 248 +-------- .../runOpsStore.envScopedResidency.test.ts | 187 ++++--- ...runOpsStore.forWaitpointCompletion.test.ts | 4 + .../src/runOpsStore.newMethods.test.ts | 319 +++++++----- .../src/runOpsStore.runKeyedRouting.test.ts | 479 ++++++++++-------- 6 files changed, 618 insertions(+), 773 deletions(-) diff --git a/apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts b/apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts index 141b8274a8a..0cd239d675a 100644 --- a/apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts +++ b/apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts @@ -1,16 +1,15 @@ // Real heterogeneous legacy + new Postgres proof for the alert-hydration TaskRun read. -// The DB is never mocked. A test-only RunStore wraps two real PostgresRunStore -// instances and routes findRun by id residency (run-ops id → NEW, cuid → LEGACY), -// mirroring the sibling routing suite. The ProjectAlertChannel read must stay control-plane. +// The DB is never mocked. The REAL RoutingRunStore wraps two real PostgresRunStore instances and +// routes findRun by id residency, mirroring the sibling routing suite. The ProjectAlertChannel +// read must stay control-plane. // // The alert env-type read (parentEnvironment?.type ?? type) is resolved via the app // ControlPlaneResolver over a control-plane client DISTINCT from the run-ops store, proving the // cross-provider inversion. The prior version co-located env + run and masked it. import { heteroPostgresTest, postgresTest } from "@internal/testcontainers"; -import { PostgresRunStore } from "@internal/run-store"; -import type { ReadClient, RunStore } from "@internal/run-store"; -import type { Prisma, PrismaClient } from "@trigger.dev/database"; -import { generateRunOpsId, ownerEngine } from "@trigger.dev/core/v3/isomorphic"; +import { PostgresRunStore, RoutingRunStore } from "@internal/run-store"; +import type { PrismaClient } from "@trigger.dev/database"; +import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; import { describe, expect } from "vitest"; import { ControlPlaneCache } from "~/v3/runOpsMigration/controlPlaneCache.server"; import { ControlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; @@ -28,145 +27,16 @@ function buildControlPlaneResolver(controlPlane: PrismaClient) { vi.setConfig({ testTimeout: 60_000 }); -// Test-only routing store: resolve findRun by id length (27 → NEW, else LEGACY), -// dropping any forwarded client so each inner store uses its OWN prisma. NOT a mock — -// real DB I/O against two PostgresRunStore instances. -class RoutingRunStore implements RunStore { - readonly #newStore: PostgresRunStore; - readonly #legacyStore: PostgresRunStore; - - constructor(newStore: PostgresRunStore, legacyStore: PostgresRunStore) { - this.#newStore = newStore; - this.#legacyStore = legacyStore; - } - - #resolveById(runId: string): PostgresRunStore { - return ownerEngine(runId) === "NEW" ? this.#newStore : this.#legacyStore; - } - - #idFromWhere(where: Prisma.TaskRunWhereInput): string | undefined { - const id = (where as { id?: unknown }).id; - return typeof id === "string" ? id : undefined; - } - - async findRun( - where: Prisma.TaskRunWhereInput, - argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | ReadClient, - _client?: ReadClient - ): Promise { - const id = this.#idFromWhere(where); - if (id !== undefined) { - return (this.#resolveById(id).findRun as any)(where, argsOrClient); - } - const fromNew = await (this.#newStore.findRun as any)(where, argsOrClient); - return fromNew ?? (this.#legacyStore.findRun as any)(where, argsOrClient); - } - - // The remaining RunStore methods are not exercised here; delegate to NEW to satisfy - // the interface. - findRunOrThrow(...a: any[]): any { - return (this.#newStore.findRunOrThrow as any)(...a); - } - findRuns(...a: any[]): any { - return (this.#newStore.findRuns as any)(...a); - } - createRun(p: any, tx?: any): any { - return this.#resolveById(p.data.id).createRun(p, tx); - } - createCancelledRun(p: any, tx?: any): any { - return this.#resolveById(p.data.id).createCancelledRun(p, tx); - } - createFailedRun(p: any, tx?: any): any { - return this.#resolveById(p.data.id).createFailedRun(p, tx); - } - updateMetadata(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).updateMetadata as any)(...[runId, ...a]); - } - startAttempt(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).startAttempt as any)(runId, ...a); - } - completeAttemptSuccess(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).completeAttemptSuccess as any)(runId, ...a); - } - recordRetryOutcome(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).recordRetryOutcome as any)(runId, ...a); - } - requeueRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).requeueRun as any)(runId, ...a); - } - recordBulkActionMembership(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).recordBulkActionMembership as any)(runId, ...a); - } - cancelRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).cancelRun as any)(runId, ...a); - } - failRunPermanently(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).failRunPermanently as any)(runId, ...a); - } - expireRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).expireRun as any)(runId, ...a); - } - expireRunsBatch(runIds: string[], ...a: any[]): any { - return (this.#resolveById(runIds[0] ?? "").expireRunsBatch as any)(runIds, ...a); - } - lockRunToWorker(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).lockRunToWorker as any)(runId, ...a); - } - parkPendingVersion(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).parkPendingVersion as any)(runId, ...a); - } - promotePendingVersionRuns(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).promotePendingVersionRuns as any)(runId, ...a); - } - expireParkedRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).expireParkedRun as any)(runId, ...a); - } - suspendForCheckpoint(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).suspendForCheckpoint as any)(runId, ...a); - } - resumeFromCheckpoint(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).resumeFromCheckpoint as any)(runId, ...a); - } - rescheduleRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).rescheduleRun as any)(runId, ...a); - } - enqueueDelayedRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).enqueueDelayedRun as any)(runId, ...a); - } - rewriteDebouncedRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).rewriteDebouncedRun as any)(runId, ...a); - } - clearIdempotencyKey(params: any, tx?: any): any { - const runId = params?.byId?.runId ?? ""; - return this.#resolveById(runId).clearIdempotencyKey(params, tx); - } - pushTags(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).pushTags as any)(runId, ...a); - } - pushRealtimeStream(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).pushRealtimeStream as any)(runId, ...a); - } - finalizeRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).finalizeRun as any)(runId, ...a); - } - findManyBatchTaskRunItems(...a: any[]): any { - return (this.#newStore.findManyBatchTaskRunItems as any)(...a); - } - findBatchTaskRunItem(...a: any[]): any { - return (this.#newStore.findBatchTaskRunItem as any)(...a); - } - upsertWaitpointTag(...a: any[]): any { - return (this.#newStore.upsertWaitpointTag as any)(...a); - } - findManyWaitpointTags(...a: any[]): any { - return (this.#newStore.findManyWaitpointTags as any)(...a); - } -} +// The alert-hydration TaskRun read runs through the REAL RoutingRunStore over two real +// PostgresRunStore instances (NEW = PG17, LEGACY = PG14). The DB is never mocked. The router +// resolves residency from the id shape — a v1 run-ops id (26 chars, version "1" at index 25) to +// NEW, a 25-char cuid to LEGACY — and never forwards a caller-passed control-plane client into a +// routed read, so each store uses its OWN prisma. function buildRoutingStore(prisma17: PrismaClient, prisma14: PrismaClient) { const newStore = new PostgresRunStore({ prisma: prisma17, readOnlyPrisma: prisma17 }); const legacyStore = new PostgresRunStore({ prisma: prisma14, readOnlyPrisma: prisma14 }); - return new RoutingRunStore(newStore, legacyStore); + return new RoutingRunStore({ new: newStore, legacy: legacyStore }); } async function seedProject(prisma: PrismaClient, suffix: string) { diff --git a/apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts b/apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts index 17d6ec9fc2f..3ce92d18106 100644 --- a/apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts +++ b/apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts @@ -1,243 +1,31 @@ import { heteroPostgresTest } from "@internal/testcontainers"; -import { PostgresRunStore } from "@internal/run-store"; -import type { ReadClient, RunStore } from "@internal/run-store"; -import type { Prisma, PrismaClient } from "@trigger.dev/database"; +import { PostgresRunStore, RoutingRunStore } from "@internal/run-store"; +import type { PrismaClient } from "@trigger.dev/database"; import { parsePacket } from "@trigger.dev/core/v3"; -import { generateRunOpsId, ownerEngine } from "@trigger.dev/core/v3/isomorphic"; +import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; import { setTimeout } from "timers/promises"; import { describe, expect } from "vitest"; import { UpdateMetadataService } from "~/services/metadata/updateMetadata.server"; vi.setConfig({ testTimeout: 60_000 }); -/** - * A test-only RunStore that routes residency-bearing operations to one of two - * inner PostgresRunStore instances (NEW = PG17, LEGACY = PG14) purely by run-id - * classification — NOT by whatever client the service forwards as `tx`. - * - * This is the load-bearing design point: the UpdateMetadataService forwards - * `this._prisma` as the tx/client to every findRun/updateMetadata call. To prove - * STORE residency routing (and not the forwarded prisma), this wrapper IGNORES - * the forwarded client for residency-bearing calls and resolves to its own inner - * store by id length, then calls the inner store WITHOUT forwarding the outer tx - * (passes undefined), so the inner PostgresRunStore uses its own prisma17/prisma14. - * - * Classification contract (version char): a v1 id (26 chars, version "1" at index 25) => NEW store; - * 25-char cuid => LEGACY store. - */ -class RoutingRunStore implements RunStore { - readonly #newStore: PostgresRunStore; - readonly #legacyStore: PostgresRunStore; - - constructor(newStore: PostgresRunStore, legacyStore: PostgresRunStore) { - this.#newStore = newStore; - this.#legacyStore = legacyStore; - } - - // Resolve by the version char: a v1 body => NEW, otherwise LEGACY (25-char cuid). - #resolveById(runId: string): PostgresRunStore { - return ownerEngine(runId) === "NEW" ? this.#newStore : this.#legacyStore; - } - - // Extract a classifiable run id from a `where`. Prefers `where.id`; if only a - // friendlyId is present the stub does not classify, so the caller falls back - // to read-through (try NEW, then LEGACY). - #idFromWhere(where: Prisma.TaskRunWhereInput): string | undefined { - const id = (where as { id?: unknown }).id; - return typeof id === "string" ? id : undefined; - } - - // ---- Reads (residency routing; drop forwarded client) ---- - - async findRun( - where: Prisma.TaskRunWhereInput, - argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | ReadClient, - _client?: ReadClient - ): Promise { - const id = this.#idFromWhere(where); - if (id !== undefined) { - // Classifiable by id shape — route to the owning store, dropping the - // forwarded client so the inner store uses its OWN prisma. - return (this.#resolveById(id).findRun as any)(where, argsOrClient); - } - // Not classifiable (friendlyId-only / other) — read-through: NEW then LEGACY. - const fromNew = await (this.#newStore.findRun as any)(where, argsOrClient); - if (fromNew) { - return fromNew; - } - return (this.#legacyStore.findRun as any)(where, argsOrClient); - } - - async findRunOrThrow( - where: Prisma.TaskRunWhereInput, - argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | ReadClient, - _client?: ReadClient - ): Promise { - const id = this.#idFromWhere(where); - if (id !== undefined) { - return (this.#resolveById(id).findRunOrThrow as any)(where, argsOrClient); - } - const fromNew = await (this.#newStore.findRun as any)(where, argsOrClient); - if (fromNew) { - return fromNew; - } - return (this.#legacyStore.findRunOrThrow as any)(where, argsOrClient); - } - - async findRunOnPrimary( - where: Prisma.TaskRunWhereInput, - args?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } - ): Promise { - const id = this.#idFromWhere(where); - if (id !== undefined) { - return (this.#resolveById(id).findRunOnPrimary as any)(where, args); - } - const fromNew = await (this.#newStore.findRunOnPrimary as any)(where, args); - if (fromNew) { - return fromNew; - } - return (this.#legacyStore.findRunOnPrimary as any)(where, args); - } - - async findRunOrThrowOnPrimary( - where: Prisma.TaskRunWhereInput, - args?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } - ): Promise { - const id = this.#idFromWhere(where); - if (id !== undefined) { - return (this.#resolveById(id).findRunOrThrowOnPrimary as any)(where, args); - } - const fromNew = await (this.#newStore.findRunOnPrimary as any)(where, args); - if (fromNew) { - return fromNew; - } - return (this.#legacyStore.findRunOrThrowOnPrimary as any)(where, args); - } - - async findRuns( - args: { where: Prisma.TaskRunWhereInput }, - _client?: ReadClient - ): Promise { - const id = this.#idFromWhere(args.where); - if (id !== undefined) { - return (this.#resolveById(id).findRuns as any)(args); - } - // Read-through across both stores, NEW first. - const fromNew = (await (this.#newStore.findRuns as any)(args)) as unknown[]; - const fromLegacy = (await (this.#legacyStore.findRuns as any)(args)) as unknown[]; - return [...fromNew, ...fromLegacy]; - } - - // ---- Field touches (residency routing; drop forwarded tx) ---- - - async updateMetadata( - runId: string, - data: Parameters[1], - options: Parameters[2], - _tx?: unknown - ): Promise<{ count: number }> { - // Route by run id, dropping the forwarded tx so the inner store writes to - // its OWN prisma — this is what proves the CAS targets the owning store. - return this.#resolveById(runId).updateMetadata(runId, data, options); - } - - // ---- Everything else: delegate by run id to satisfy the RunStore interface; - // not exercised by these tests. ---- - - createRun(params: any, _tx?: unknown): any { - return this.#resolveById(params.data.id).createRun(params); - } - createCancelledRun(params: any, _tx?: unknown): any { - return this.#resolveById(params.data.id).createCancelledRun(params); - } - createFailedRun(params: any, _tx?: unknown): any { - return this.#resolveById(params.data.id).createFailedRun(params); - } - startAttempt(runId: string, data: any, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).startAttempt as any)(runId, data, args); - } - completeAttemptSuccess(runId: string, data: any, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).completeAttemptSuccess as any)(runId, data, args); - } - recordRetryOutcome(runId: string, data: any, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).recordRetryOutcome as any)(runId, data, args); - } - requeueRun(runId: string, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).requeueRun as any)(runId, args); - } - recordBulkActionMembership(runId: string, bulkActionId: string, _tx?: unknown): any { - return this.#resolveById(runId).recordBulkActionMembership(runId, bulkActionId); - } - cancelRun(runId: string, data: any, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).cancelRun as any)(runId, data, args); - } - failRunPermanently(runId: string, data: any, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).failRunPermanently as any)(runId, data, args); - } - expireRun(runId: string, data: any, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).expireRun as any)(runId, data, args); - } - expireRunsBatch(runIds: string[], data: any, _tx?: unknown): any { - return this.#resolveById(runIds[0] ?? "").expireRunsBatch(runIds, data); - } - lockRunToWorker(runId: string, data: any, _tx?: unknown): any { - return this.#resolveById(runId).lockRunToWorker(runId, data); - } - parkPendingVersion(runId: string, data: any, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).parkPendingVersion as any)(runId, data, args); - } - promotePendingVersionRuns(runId: string, args?: any, _tx?: unknown): any { - return this.#resolveById(runId).promotePendingVersionRuns(runId, args); - } - expireParkedRun(runId: string, data: any, _tx?: unknown): any { - return (this.#resolveById(runId).expireParkedRun as any)(runId, data); - } - suspendForCheckpoint(runId: string, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).suspendForCheckpoint as any)(runId, args); - } - resumeFromCheckpoint(runId: string, args: any, _tx?: unknown): any { - return (this.#resolveById(runId).resumeFromCheckpoint as any)(runId, args); - } - rescheduleRun(runId: string, data: any, _tx?: unknown): any { - return this.#resolveById(runId).rescheduleRun(runId, data); - } - enqueueDelayedRun(runId: string, data: any, _tx?: unknown): any { - return this.#resolveById(runId).enqueueDelayedRun(runId, data); - } - rewriteDebouncedRun(runId: string, data: any, _tx?: unknown): any { - return this.#resolveById(runId).rewriteDebouncedRun(runId, data); - } - clearIdempotencyKey(params: any, _tx?: unknown): any { - const runId = params?.byId?.runId ?? ""; - return this.#resolveById(runId).clearIdempotencyKey(params); - } - pushTags(runId: string, tags: string[], where: any, _tx?: unknown): any { - return this.#resolveById(runId).pushTags(runId, tags, where); - } - pushRealtimeStream(runId: string, streamId: string, _tx?: unknown): any { - return this.#resolveById(runId).pushRealtimeStream(runId, streamId); - } - finalizeRun(runId: string, ...a: any[]): any { - return (this.#resolveById(runId).finalizeRun as any)(runId, ...a); - } - findManyBatchTaskRunItems(...a: any[]): any { - return (this.#newStore.findManyBatchTaskRunItems as any)(...a); - } - findBatchTaskRunItem(...a: any[]): any { - return (this.#newStore.findBatchTaskRunItem as any)(...a); - } - upsertWaitpointTag(...a: any[]): any { - return (this.#newStore.upsertWaitpointTag as any)(...a); - } - findManyWaitpointTags(...a: any[]): any { - return (this.#newStore.findManyWaitpointTags as any)(...a); - } -} +// Real heterogeneous NEW + LEGACY Postgres proof for UpdateMetadataService, exercising the REAL +// RoutingRunStore over two real PostgresRunStore instances (NEW = PG17, LEGACY = PG14). The DB is +// never mocked. +// +// The load-bearing design point: UpdateMetadataService forwards `this._prisma` as the tx/client to +// every findRun/updateMetadata call. That client is bound to the control plane — the wrong database +// for a run resident on either store — so the router must never forward it verbatim. It does not: +// a non-replica client escalates to the OWNING store's own primary, so residency routing is proved +// rather than the forwarded prisma. +// +// Residency comes from the id shape: a v1 run-ops id (26 chars, version "1" at index 25) resolves +// to NEW, a 25-char cuid to LEGACY. function buildRoutingStore(prisma17: PrismaClient, prisma14: PrismaClient) { const newStore = new PostgresRunStore({ prisma: prisma17, readOnlyPrisma: prisma17 }); const legacyStore = new PostgresRunStore({ prisma: prisma14, readOnlyPrisma: prisma14 }); - return new RoutingRunStore(newStore, legacyStore); + return new RoutingRunStore({ new: newStore, legacy: legacyStore }); } // 25-char cuid-format id (starts with "c"), no v1 version marker. @@ -466,8 +254,8 @@ describe("UpdateMetadataService store routing (hetero)", () => { logLevel: "error", }); - // Call WITHOUT an environment arg, so the `where` is just `{ id: runId }` and - // the router classifies by id length (25 => LEGACY). + // Call WITHOUT an environment arg, so the `where` is just `{ id: runId }` and the router + // resolves residency from the id shape (a 25-char cuid is not a v1 body => LEGACY). const result = await service.call(runId, { operations: [{ type: "set", key: "x", value: 1 }], }); diff --git a/internal-packages/run-store/src/runOpsStore.envScopedResidency.test.ts b/internal-packages/run-store/src/runOpsStore.envScopedResidency.test.ts index 263ccbc9d11..4bd4df793e3 100644 --- a/internal-packages/run-store/src/runOpsStore.envScopedResidency.test.ts +++ b/internal-packages/run-store/src/runOpsStore.envScopedResidency.test.ts @@ -5,13 +5,23 @@ 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. +// +// Every case runs at TWO shards (the compat pair) and at THREE (the pair plus one gen-2 shard). +// clearIdempotencyKey is the ONLY caller of #shardsExcept, which must fan out over every remaining +// store rather than take the first. At two shards that list holds one entry, so the fan-out is +// degenerate and a take-the-first bug is invisible. The third shard makes it observable. +// +// The seam stays `classify` (binary NEW/LEGACY) on purpose: these paths are selected by the +// residency hint and by fan-out membership, never by id shape. The gen-2 id algebra lives in +// runOpsStore.shardMap.test.ts. type Call = { method: string; args: unknown[] }; -type FakeStore = RunStore & { slot: "new" | "legacy"; calls: Call[] }; +type Slot = "new" | "legacy" | "a"; +type FakeStore = RunStore & { slot: Slot; 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 { +// path can be exercised (NEW matches 0 → fall back to the other stores). +function fakeStore(slot: Slot, clearCount = slot === "new" ? 1 : 0): FakeStore { const calls: Call[] = []; const rec = (method: string, result: unknown) => @@ -27,83 +37,128 @@ function fakeStore(slot: "new" | "legacy", clearCount = slot === "new" ? 1 : 0): } as unknown as FakeStore; } -function buildRouter(newClearCount?: number, legacyClearCount?: number) { - const newStore = fakeStore("new", newClearCount); - const legacyStore = fakeStore("legacy", legacyClearCount); +type Topology = { name: string; gen2Keys: readonly Slot[] }; + +const TOPOLOGIES: readonly Topology[] = [ + { name: "compat pair", gen2Keys: [] }, + { name: "one gen-2 shard", gen2Keys: ["a"] }, +]; + +type Counts = { new?: number; legacy?: number; a?: number }; + +function buildRouter(topology: Topology, counts: Counts = {}) { + const newStore = fakeStore("new", counts.new); + const legacyStore = fakeStore("legacy", counts.legacy); + const gen2Stores = topology.gen2Keys.map((key) => fakeStore(key, counts[key])); const router = new RoutingRunStore({ new: newStore, legacy: legacyStore, classify: (id: string) => (id.startsWith("new") ? "NEW" : "LEGACY"), + ...(gen2Stores.length > 0 + ? { shards: gen2Stores.map((store) => ({ key: store.slot, store })) } + : {}), }); - return { router, newStore, legacyStore }; + return { router, newStore, legacyStore, gen2Stores }; } -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); - }); +for (const topology of TOPOLOGIES) { + describe(`RoutingRunStore.upsertWaitpointTag — residency hint for a tag with no minted id (${topology.name})`, () => { + it("routes to NEW when residency is NEW and no id is supplied", async () => { + const { router, newStore, legacyStore, gen2Stores } = buildRouter(topology); + 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); + for (const store of gen2Stores) { + expect(store.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); + it("still falls back to LEGACY when no id and no residency are supplied", async () => { + const { router, newStore, legacyStore, gen2Stores } = buildRouter(topology); + await router.upsertWaitpointTag({ environmentId: "env", name: "t", projectId: "p" }); + expect(legacyStore.calls.map((c) => c.method)).toEqual(["upsertWaitpointTag"]); + expect(newStore.calls).toHaveLength(0); + // #idlessWaitpointShard stays legacy at any shard count. + for (const store of gen2Stores) { + expect(store.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", - }, + describe(`RoutingRunStore.clearIdempotencyKey — predicate routes NEW-first when the env mints new (${topology.name})`, () => { + it("clears on NEW and does NOT touch the other stores when NEW matches (post-flip key)", async () => { + const { router, newStore, legacyStore, gen2Stores } = buildRouter(topology, { + new: 1, + legacy: 0, + a: 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); + // The NEW short circuit must not widen with the shard count. + for (const store of gen2Stores) { + expect(store.calls).toHaveLength(0); + } + expect(result.count).toBe(1); }); - 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", - }, + it("falls back to EVERY other store when NEW matches 0 (a key held on a pre-flip run)", async () => { + // The env mints new now, but this key was created before the flip → its run lives elsewhere. + // #shardsExcept(NEW) must yield every remaining store, not just the first one. + const { router, newStore, legacyStore, gen2Stores } = buildRouter(topology, { + new: 0, + legacy: 1, + a: 1, + }); + const result = await router.clearIdempotencyKey({ + byPredicate: { + idempotencyKey: "k", + taskIdentifier: "task", + runtimeEnvironmentId: "env", + residency: "NEW", + }, + }); + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + for (const store of gen2Stores) { + expect(store.calls).toHaveLength(1); + } + // A take-the-first fan-out returns 1 here, leaving the stale key on shard "a" deduping. + expect(result.count).toBe(1 + gen2Stores.length); }); - // 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" }, + it("still fans out a byPredicate reset with no residency (mixed residency)", async () => { + const { router, newStore, legacyStore, gen2Stores } = buildRouter(topology); + await router.clearIdempotencyKey({ + byPredicate: { idempotencyKey: "k", taskIdentifier: "task", runtimeEnvironmentId: "env" }, + }); + // #sumCounts spans every distinct store, so the fan widens with the shard count. + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); + for (const store of gen2Stores) { + expect(store.calls).toHaveLength(1); + } }); - 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); + it("routes byId to the owning store (unchanged)", async () => { + const { router, newStore, legacyStore, gen2Stores } = buildRouter(topology); + await router.clearIdempotencyKey({ byId: { runId: "new_run", idempotencyKey: "k" } }); + expect(newStore.calls.map((c) => c.method)).toEqual(["clearIdempotencyKey"]); + expect(legacyStore.calls).toHaveLength(0); + for (const store of gen2Stores) { + expect(store.calls).toHaveLength(0); + } + }); }); -}); +} diff --git a/internal-packages/run-store/src/runOpsStore.forWaitpointCompletion.test.ts b/internal-packages/run-store/src/runOpsStore.forWaitpointCompletion.test.ts index eff2a71cd97..1db1f8cb1bf 100644 --- a/internal-packages/run-store/src/runOpsStore.forWaitpointCompletion.test.ts +++ b/internal-packages/run-store/src/runOpsStore.forWaitpointCompletion.test.ts @@ -3,6 +3,10 @@ import { PostgresRunStore } from "./PostgresRunStore.js"; import { RoutingRunStore } from "./runOpsStore.js"; import type { RunStore } from "./types.js"; +// This suite is deliberately TWO-STORE. runOpsStore.shardMap.test.ts runs this same method at four +// stores ("routes a gen-2 waitpoint directly, with no probe"; "lets a gen-2 waitpoint id beat the +// cross-tree legacy pin"), so a shard arm here would repeat that coverage rather than add any. +// // forWaitpointCompletion is async: it picks a preferred store from the id-shape + pins, then // PROBES findWaitpoint to resolve where the token ACTUALLY lives (drain can relocate a cuid // waitpoint onto NEW, or a run-ops token can be pinned LEGACY), falling back to the other store. diff --git a/internal-packages/run-store/src/runOpsStore.newMethods.test.ts b/internal-packages/run-store/src/runOpsStore.newMethods.test.ts index 163ddf2f3c8..8a56057a117 100644 --- a/internal-packages/run-store/src/runOpsStore.newMethods.test.ts +++ b/internal-packages/run-store/src/runOpsStore.newMethods.test.ts @@ -1,21 +1,52 @@ import { describe, expect, it } from "vitest"; import { RoutingRunStore } from "./runOpsStore.js"; +import type { RoutingStoreMetrics } from "./routingStoreMetrics.js"; import type { FinalizeRunData, RunStore } from "./types.js"; // Pure routing unit tests for the five store methods added in Track 1. No DB: each slot is a fake // RunStore that records the calls it receives, so the assertions are purely about WHICH store the // router dispatches to (by residency key) and WHAT it forwards (never a control-plane tx into a // routed write; caller client presence escalates to the owning store's own primary). +// +// Every case runs at TWO shards (the compat pair) and at THREE (the pair plus one gen-2 shard). +// The routed cases prove a write never drifts onto a shard that cannot own it. The merge case is +// topology-indexed: #precedence is [legacy, new, ...gen2] with last-write-wins, so the winner of a +// duplicate id CHANGES when a gen-2 leg holds it, and #reportDuplicateId alarms because the +// reporting set is no longer a subset of the gen-1 pair. +// +// The seam is `resolveShard`, not `classify`: a gen-1 classifier maps through a binary ternary and +// can only ever name the two reserved keys, so it can never reach a gen-2 shard. type Call = { method: string; args: unknown[] }; +// Shard key "z", because the merged tag rows already use the ids "a", "b" and "c". +type Slot = "new" | "legacy" | "z"; + type FakeStore = RunStore & { - slot: "new" | "legacy"; + slot: Slot; calls: Call[]; - primaryReadClient: { __primary: "new" | "legacy" }; + primaryReadClient: { __primary: Slot }; }; -function fakeStore(slot: "new" | "legacy"): FakeStore { +// The gen-2 leg collides on "a" ONLY, so the merged id set is unchanged and the single observable +// difference between the two topologies is which leg wins that duplicate. +function tagRows(slot: Slot) { + if (slot === "new") { + return [ + { id: "b", src: "new" }, + { id: "a", src: "new" }, + ]; + } + if (slot === "legacy") { + return [ + { id: "c", src: "legacy" }, + { id: "a", src: "legacy" }, + ]; + } + return [{ id: "a", src: slot }]; +} + +function fakeStore(slot: Slot): FakeStore { const calls: Call[] = []; const record = (method: string, result: unknown) => @@ -31,162 +62,178 @@ function fakeStore(slot: "new" | "legacy"): FakeStore { findManyBatchTaskRunItems: record("findManyBatchTaskRunItems", [{ slot }]), findBatchTaskRunItem: record("findBatchTaskRunItem", { slot }), upsertWaitpointTag: record("upsertWaitpointTag", { slot }), - // Slot-specific rows so the merge/dedupe (NEW-wins) is observable; id "a" collides across legs. - findManyWaitpointTags: record( - "findManyWaitpointTags", - slot === "new" - ? [ - { id: "b", src: "new" }, - { id: "a", src: "new" }, - ] - : [ - { id: "c", src: "legacy" }, - { id: "a", src: "legacy" }, - ] - ), + findManyWaitpointTags: record("findManyWaitpointTags", tagRows(slot)), } as unknown as FakeStore; } -// Deterministic residency by id prefix, injected via the classify seam so the tests don't depend on -// id-shape length rules. -function buildRouter() { +type Topology = { name: string; gen2Keys: readonly Slot[] }; + +const TOPOLOGIES: readonly Topology[] = [ + { name: "compat pair", gen2Keys: [] }, + { name: "one gen-2 shard", gen2Keys: ["z"] }, +]; + +// Deterministic residency by id prefix, injected via the resolveShard seam so the tests don't +// depend on id-shape length rules. +function buildRouter(topology: Topology, metrics?: RoutingStoreMetrics) { const newStore = fakeStore("new"); const legacyStore = fakeStore("legacy"); + const gen2Stores = topology.gen2Keys.map((key) => fakeStore(key)); const router = new RoutingRunStore({ new: newStore, legacy: legacyStore, - classify: (id: string) => (id.startsWith("new") ? "NEW" : "LEGACY"), + resolveShard: (id: string) => + id.startsWith("new") ? "new" : id.startsWith("z_") ? "z" : "legacy", + ...(gen2Stores.length > 0 + ? { shards: gen2Stores.map((store) => ({ key: store.slot, store })) } + : {}), + ...(metrics ? { metrics } : {}), }); - return { router, newStore, legacyStore }; + return { router, newStore, legacyStore, gen2Stores }; } const DATA: FinalizeRunData = { status: "COMPLETED_SUCCESSFULLY", completedAt: new Date() }; -describe("RoutingRunStore.finalizeRun", () => { - it("routes by runId and forwards the projection, never the tx", async () => { - const { router, newStore, legacyStore } = buildRouter(); - const projection = { select: { id: true } }; - await router.finalizeRun("new_run", DATA, projection); - expect(newStore.calls).toHaveLength(1); - expect(newStore.calls[0]?.args).toEqual(["new_run", DATA, projection]); - expect(legacyStore.calls).toHaveLength(0); - }); +for (const topology of TOPOLOGIES) { + describe(`RoutingRunStore.finalizeRun (${topology.name})`, () => { + it("routes by runId and forwards the projection, never the tx", async () => { + const { router, newStore, legacyStore } = buildRouter(topology); + const projection = { select: { id: true } }; + await router.finalizeRun("new_run", DATA, projection); + expect(newStore.calls).toHaveLength(1); + expect(newStore.calls[0]?.args).toEqual(["new_run", DATA, projection]); + expect(legacyStore.calls).toHaveLength(0); + }); - it("routes a cuid/legacy runId to the legacy store", async () => { - const { router, newStore, legacyStore } = buildRouter(); - await router.finalizeRun("legacy_run", DATA, { include: { attempts: true } }); - expect(legacyStore.calls[0]?.args).toEqual([ - "legacy_run", - DATA, - { include: { attempts: true } }, - ]); - expect(newStore.calls).toHaveLength(0); - }); + it("routes a cuid/legacy runId to the legacy store", async () => { + const { router, newStore, legacyStore } = buildRouter(topology); + await router.finalizeRun("legacy_run", DATA, { include: { attempts: true } }); + expect(legacyStore.calls[0]?.args).toEqual([ + "legacy_run", + DATA, + { include: { attempts: true } }, + ]); + expect(newStore.calls).toHaveLength(0); + }); - it("drops a caller-passed control-plane tx (never threaded into the routed write)", async () => { - const { router, legacyStore } = buildRouter(); - const controlPlaneTx = { $fake: "cp-tx" }; - await router.finalizeRun("legacy_run", DATA, controlPlaneTx as never); - // The tx is neither a select/include projection nor forwarded: the sub-store sees a 3-arg call - // whose projection slot is undefined. - expect(legacyStore.calls[0]?.args).toEqual(["legacy_run", DATA, undefined]); - }); -}); - -describe("RoutingRunStore.findManyBatchTaskRunItems", () => { - it("routes by batchTaskRunId first", async () => { - const { router, newStore, legacyStore } = buildRouter(); - await router.findManyBatchTaskRunItems({ - batchTaskRunId: "new_batch", - taskRunId: "legacy_run", + it("drops a caller-passed control-plane tx (never threaded into the routed write)", async () => { + const { router, legacyStore } = buildRouter(topology); + const controlPlaneTx = { $fake: "cp-tx" }; + await router.finalizeRun("legacy_run", DATA, controlPlaneTx as never); + // The tx is neither a select/include projection nor forwarded: the sub-store sees a 3-arg call + // whose projection slot is undefined. + expect(legacyStore.calls[0]?.args).toEqual(["legacy_run", DATA, undefined]); }); - expect(newStore.calls[0]?.method).toBe("findManyBatchTaskRunItems"); - expect(legacyStore.calls).toHaveLength(0); }); - it("falls back to taskRunId when no batchTaskRunId is present", async () => { - const { router, newStore, legacyStore } = buildRouter(); - await router.findManyBatchTaskRunItems({ taskRunId: "legacy_run" }); - expect(legacyStore.calls[0]?.method).toBe("findManyBatchTaskRunItems"); - expect(newStore.calls).toHaveLength(0); - }); + describe(`RoutingRunStore.findManyBatchTaskRunItems (${topology.name})`, () => { + it("routes by batchTaskRunId first", async () => { + const { router, newStore, legacyStore } = buildRouter(topology); + await router.findManyBatchTaskRunItems({ + batchTaskRunId: "new_batch", + taskRunId: "legacy_run", + }); + expect(newStore.calls[0]?.method).toBe("findManyBatchTaskRunItems"); + expect(legacyStore.calls).toHaveLength(0); + }); - it("escalates a caller client to the owning store's own primary (read-your-writes)", async () => { - const { router, newStore } = buildRouter(); - // A non-replica client object signals read-your-writes; it must NOT be forwarded verbatim. - await router.findManyBatchTaskRunItems({ batchTaskRunId: "new_batch" }, undefined, { - writer: true, - } as never); - expect(newStore.calls[0]?.args[2]).toEqual({ __primary: "new" }); - }); -}); - -describe("RoutingRunStore.findBatchTaskRunItem", () => { - it("routes by batchTaskRunId", async () => { - const { router, newStore, legacyStore } = buildRouter(); - await router.findBatchTaskRunItem({ batchTaskRunId: "legacy_batch", taskRunId: "new_run" }); - expect(legacyStore.calls[0]?.method).toBe("findBatchTaskRunItem"); - expect(newStore.calls).toHaveLength(0); - }); -}); - -describe("RoutingRunStore.upsertWaitpointTag", () => { - it("routes the write by the tag's minted id-shape (env mint-kind)", async () => { - const { router, newStore, legacyStore } = buildRouter(); - await router.upsertWaitpointTag({ - environmentId: "env", - name: "t", - projectId: "p", - id: "new_tag", + it("falls back to taskRunId when no batchTaskRunId is present", async () => { + const { router, newStore, legacyStore } = buildRouter(topology); + await router.findManyBatchTaskRunItems({ taskRunId: "legacy_run" }); + expect(legacyStore.calls[0]?.method).toBe("findManyBatchTaskRunItems"); + expect(newStore.calls).toHaveLength(0); + }); + + it("escalates a caller client to the owning store's own primary (read-your-writes)", async () => { + const { router, newStore } = buildRouter(topology); + // A non-replica client object signals read-your-writes; it must NOT be forwarded verbatim. + await router.findManyBatchTaskRunItems({ batchTaskRunId: "new_batch" }, undefined, { + writer: true, + } as never); + expect(newStore.calls[0]?.args[2]).toEqual({ __primary: "new" }); }); - expect(newStore.calls[0]?.method).toBe("upsertWaitpointTag"); - expect(legacyStore.calls).toHaveLength(0); }); - it("falls back to legacy when no minted id is supplied", async () => { - const { router, newStore, legacyStore } = buildRouter(); - await router.upsertWaitpointTag({ environmentId: "env", name: "t", projectId: "p" }); - expect(legacyStore.calls[0]?.method).toBe("upsertWaitpointTag"); - expect(newStore.calls).toHaveLength(0); + describe(`RoutingRunStore.findBatchTaskRunItem (${topology.name})`, () => { + it("routes by batchTaskRunId", async () => { + const { router, newStore, legacyStore } = buildRouter(topology); + await router.findBatchTaskRunItem({ batchTaskRunId: "legacy_batch", taskRunId: "new_run" }); + expect(legacyStore.calls[0]?.method).toBe("findBatchTaskRunItem"); + expect(newStore.calls).toHaveLength(0); + }); }); - it("never threads a control-plane tx into either leg", async () => { - const { router, newStore, legacyStore } = buildRouter(); - const tx = { $fake: "cp-tx" }; - await router.upsertWaitpointTag( - { environmentId: "env", name: "t", projectId: "p", id: "legacy_tag" }, - tx as never - ); - // The routed write runs on the owning store's own client, so the tx is dropped on the LEGACY leg too. - expect(legacyStore.calls[0]?.args[1]).toBeUndefined(); - - const tx2 = { $fake: "cp-tx-2" }; - await router.upsertWaitpointTag( - { environmentId: "env", name: "t", projectId: "p", id: "new_tag" }, - tx2 as never - ); - // NEW leg likewise never receives the control-plane tx. - expect(newStore.calls[0]?.args[1]).toBeUndefined(); + describe(`RoutingRunStore.upsertWaitpointTag (${topology.name})`, () => { + it("routes the write by the tag's minted id-shape (env mint-kind)", async () => { + const { router, newStore, legacyStore } = buildRouter(topology); + await router.upsertWaitpointTag({ + environmentId: "env", + name: "t", + projectId: "p", + id: "new_tag", + }); + expect(newStore.calls[0]?.method).toBe("upsertWaitpointTag"); + expect(legacyStore.calls).toHaveLength(0); + }); + + it("falls back to legacy when no minted id is supplied", async () => { + const { router, newStore, legacyStore } = buildRouter(topology); + await router.upsertWaitpointTag({ environmentId: "env", name: "t", projectId: "p" }); + expect(legacyStore.calls[0]?.method).toBe("upsertWaitpointTag"); + expect(newStore.calls).toHaveLength(0); + }); + + it("never threads a control-plane tx into either leg", async () => { + const { router, newStore, legacyStore } = buildRouter(topology); + const tx = { $fake: "cp-tx" }; + await router.upsertWaitpointTag( + { environmentId: "env", name: "t", projectId: "p", id: "legacy_tag" }, + tx as never + ); + // The routed write runs on the owning store's own client, so the tx is dropped on the LEGACY leg too. + expect(legacyStore.calls[0]?.args[1]).toBeUndefined(); + + const tx2 = { $fake: "cp-tx-2" }; + await router.upsertWaitpointTag( + { environmentId: "env", name: "t", projectId: "p", id: "new_tag" }, + tx2 as never + ); + // NEW leg likewise never receives the control-plane tx. + expect(newStore.calls[0]?.args[1]).toBeUndefined(); + }); }); -}); - -describe("RoutingRunStore.findManyWaitpointTags", () => { - it("fans out to both stores, de-dupes NEW-wins, and re-imposes orderBy/take/skip globally", async () => { - const { router, newStore, legacyStore } = buildRouter(); - const result = (await router.findManyWaitpointTags({ - where: { environmentId: "env" }, - orderBy: { id: "desc" }, - take: 2, - skip: 1, - })) as Array<{ id: string; src: string }>; - - // Union {a,b,c} sorted desc = [c,b,a]; slice(1,3) = [b,a]; "a" collides so NEW wins. - expect(result.map((r) => r.id)).toEqual(["b", "a"]); - expect(result.find((r) => r.id === "a")?.src).toBe("new"); - - // Each leg is widened: skip dropped to 0, take widened to skip+take. - expect((newStore.calls[0]!.args[0] as { take: number; skip: number }).take).toBe(3); - expect((newStore.calls[0]!.args[0] as { take: number; skip: number }).skip).toBe(0); - expect((legacyStore.calls[0]!.args[0] as { take: number; skip: number }).take).toBe(3); + + describe(`RoutingRunStore.findManyWaitpointTags (${topology.name})`, () => { + it("fans out to both stores, de-dupes NEW-wins, and re-imposes orderBy/take/skip globally", async () => { + const { router, newStore, legacyStore, gen2Stores } = buildRouter(topology); + const result = (await router.findManyWaitpointTags({ + where: { environmentId: "env" }, + orderBy: { id: "desc" }, + take: 2, + skip: 1, + })) as Array<{ id: string; src: string }>; + + // Union {a,b,c} sorted desc = [c,b,a]; slice(1,3) = [b,a]. The id set does not move with the + // topology; only the winner of the "a" collision does, because #precedence puts gen-2 last. + expect(result.map((r) => r.id)).toEqual(["b", "a"]); + expect(result.find((r) => r.id === "a")?.src).toBe(gen2Stores.at(-1)?.slot ?? "new"); + + // Each leg is widened: skip dropped to 0, take widened to skip+take. + expect((newStore.calls[0]!.args[0] as { take: number; skip: number }).take).toBe(3); + expect((newStore.calls[0]!.args[0] as { take: number; skip: number }).skip).toBe(0); + expect((legacyStore.calls[0]!.args[0] as { take: number; skip: number }).take).toBe(3); + }); + + it("alarms for a tag id held by a gen-2 shard, and stays silent across the gen-1 pair", async () => { + const seen: string[][] = []; + const { router } = buildRouter(topology, { + recordDuplicateId: (keys) => seen.push(keys), + recordWaitpointProbeFallback() {}, + }); + await router.findManyWaitpointTags({ where: { environmentId: "env" } }); + // A duplicate confined to {legacy, new} is the known drain-mirror case and stays silent. Once a + // gen-2 leg reports the same id the set is no longer a subset of the pair, so it must alarm. + expect(seen).toEqual(topology.gen2Keys.length > 0 ? [["legacy", "new", "z"]] : []); + }); }); -}); +} diff --git a/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts b/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts index e5cb229ef06..74b45f661bd 100644 --- a/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts +++ b/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts @@ -8,6 +8,14 @@ import type { ReadClient, RunStore } from "./types.js"; // 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, …). +// +// Every case runs at TWO shards (the compat pair) and at THREE (the pair plus one gen-2 shard). +// #collectManyWaitpoints and countPendingWaitpoints partition the ids MISSING from the run's store: +// a gen-2 id goes to its own shard, a cuid to both gen-1 stores. At two shards those two rules pick +// the same single store, so the partition is unobservable. The third shard separates them. +// +// The seam is `resolveShard`, not `classify`: a gen-1 classifier maps through a binary ternary and +// can never name a gen-2 shard. type Call = { method: string; args: unknown[] }; @@ -25,10 +33,12 @@ type FakeConfig = { edges?: Array>; }; +type Slot = "new" | "legacy" | "a"; + type FakeStore = RunStore & { - slot: "new" | "legacy"; + slot: Slot; calls: Call[]; - primaryReadClient: { __primary: "new" | "legacy" }; + primaryReadClient: { __primary: Slot }; }; function idsFromWhere(where: unknown): string[] | undefined { @@ -41,7 +51,7 @@ function idsFromWhere(where: unknown): string[] | undefined { return undefined; } -function fakeStore(slot: "new" | "legacy", config: FakeConfig = {}): FakeStore { +function fakeStore(slot: Slot, config: FakeConfig = {}): FakeStore { const calls: Call[] = []; const rows = config.waitpoints ?? []; const byId = new Map(rows.map((r) => [r.id, r])); @@ -103,247 +113,318 @@ function fakeStore(slot: "new" | "legacy", config: FakeConfig = {}): FakeStore { 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 = {}) { +type Topology = { name: string; gen2Keys: readonly Slot[] }; + +const TOPOLOGIES: readonly Topology[] = [ + { name: "compat pair", gen2Keys: [] }, + { name: "one gen-2 shard", gen2Keys: ["a"] }, +]; + +// Deterministic residency by id prefix via the resolveShard seam (no dependence on id-shape rules). +function buildRouterFor( + topology: Topology, + newConfig: FakeConfig = {}, + legacyConfig: FakeConfig = {}, + shardConfig: FakeConfig = {} +) { const newStore = fakeStore("new", newConfig); const legacyStore = fakeStore("legacy", legacyConfig); + const gen2Stores = topology.gen2Keys.map((key) => fakeStore(key, shardConfig)); const router = new RoutingRunStore({ new: newStore, legacy: legacyStore, - classify: (id: string) => (id.startsWith("new") ? "NEW" : "LEGACY"), + resolveShard: (id: string) => + id.startsWith("new") ? "new" : id.startsWith("a_") ? "a" : "legacy", + ...(gen2Stores.length > 0 + ? { shards: gen2Stores.map((store) => ({ key: store.slot, store })) } + : {}), }); - return { router, newStore, legacyStore }; + return { router, newStore, legacyStore, gen2Stores }; } 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 }, +for (const topology of TOPOLOGIES) { + describe(`RoutingRunStore.findManyTaskRunWaitpoints — route by taskRunId (no fan-out) (${topology.name})`, () => { + it("routes an edge read keyed by a NEW run id to the new store only", async () => { + const { router, newStore, legacyStore } = buildRouterFor(topology); + await router.findManyTaskRunWaitpoints({ + where: { taskRunId: "new_run" }, + select: { taskRunId: true }, + }); + expect(newStore.calls.map((c) => c.method)).toEqual(["findManyTaskRunWaitpoints"]); + expect(legacyStore.calls).toHaveLength(0); }); - 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 }, + it("routes an edge read keyed by a LEGACY run id to the legacy store only", async () => { + const { router, newStore, legacyStore } = buildRouterFor(topology); + await router.findManyTaskRunWaitpoints({ + where: { taskRunId: "legacy_run" }, + select: { taskRunId: true }, + }); + expect(legacyStore.calls.map((c) => c.method)).toEqual(["findManyTaskRunWaitpoints"]); + expect(newStore.calls).toHaveLength(0); }); - 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("escalates a caller writer client to the owning store's own primary", async () => { + const { router, newStore } = buildRouterFor(topology); + 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 }, + it("still fans out when keyed by waitpointId (no run id in scope)", async () => { + const { router, newStore, legacyStore } = buildRouterFor(topology); + await router.findManyTaskRunWaitpoints({ + where: { waitpointId: "waitpoint_x" }, + select: { taskRunId: true }, + }); + expect(newStore.calls).toHaveLength(1); + expect(legacyStore.calls).toHaveLength(1); }); - 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"] } }, + describe(`RoutingRunStore.deleteManyTaskRunWaitpoints — route by taskRunId (no fan-out) (${topology.name})`, () => { + it("deletes only on the owning store for a classifiable taskRunId", async () => { + const { router, newStore, legacyStore } = buildRouterFor(topology, { 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); }); - 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("still fans out and sums when there is no taskRunId in the where", async () => { + const { router, newStore, legacyStore } = buildRouterFor(topology); + 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(); + it("never threads a caller tx into the routed delete", async () => { + const { router, legacyStore } = buildRouterFor(topology); + 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"]); - }); + describe(`RoutingRunStore.findSnapshotCompletedWaitpointIds — route by runId (${topology.name})`, () => { + it("routes to the run's store when a runId is threaded through", async () => { + const { router, newStore, legacyStore } = buildRouterFor( + topology, + { 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); + it("still fans out and merges when no runId is supplied", async () => { + const { router, newStore, legacyStore } = buildRouterFor( + topology, + { 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); + describe(`RoutingRunStore.findSnapshotCompletedWaitpointIdsWithPresence — route by runId (${topology.name})`, () => { + it("routes to the run's store when a runId is threaded through", async () => { + const { router, newStore } = buildRouterFor( + topology, + { 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 } = buildRouterFor( + topology, + { snapshotWaitpointIds: [], snapshotPresent: false }, + { snapshotWaitpointIds: ["waitpoint_l"], snapshotPresent: true } + ); + const res = await router.findSnapshotCompletedWaitpointIdsWithPresence("c".repeat(25)); + expect(res).toEqual({ present: true, ids: ["waitpoint_l"] }); + }); }); - 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 (${topology.name})`, () => { + it("queries only the run's store when every requested token co-locates with the run", async () => { + const { router, newStore, legacyStore } = buildRouterFor(topology, { + 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, gen2Stores } = buildRouterFor( + topology, + { 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"]); + // A cuid absent id probes the gen-1 pair ONLY. It must never reach a gen-2 shard. + for (const store of gen2Stores) { + expect(store.calls).toHaveLength(0); + } + }); + + it("still fans out (NEW-wins dedup) when no runId is supplied", async () => { + const { router, newStore, legacyStore } = buildRouterFor( + topology, + { 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.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" }, - ], + describe(`RoutingRunStore.countPendingWaitpoints — route by runId then partition-fallback (${topology.name})`, () => { + it("counts on the run's store only when every waitpoint co-locates with the run", async () => { + const { router, newStore, legacyStore } = buildRouterFor(topology, { + 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 } = buildRouterFor( + topology, + { 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. It uses the + // presence variant so the results can be unioned by id (a drain mirror counts once at N). + expect(newStore.calls.map((c) => c.method)).toEqual(["countPendingWaitpointsWithPresence"]); + 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 } = buildRouterFor( + topology, + { 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 } = buildRouterFor( + topology, + { 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); }); - 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( +// A gen-2 shard exists only in the three-shard topology, so this case has no two-shard counterpart. +describe("RoutingRunStore — a gen-2 absent id partitions to its own shard", () => { + it("sends a missing gen-2 waitpoint id to its shard and never to the gen-1 partner", async () => { + const { router, newStore, legacyStore, gen2Stores } = buildRouterFor( + TOPOLOGIES[1]!, { waitpoints: [{ id: "waitpoint_local", status: "COMPLETED" }] }, - { waitpoints: [{ id: "waitpoint_crosstree", status: "COMPLETED" }] } + { waitpoints: [] }, + { waitpoints: [{ id: "a_waitpoint", status: "COMPLETED" }] } ); const rows = (await router.findManyWaitpoints( - { where: { id: { in: ["waitpoint_local", "waitpoint_crosstree"] } } }, + { where: { id: { in: ["waitpoint_local", "a_waitpoint"] } } }, 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(rows.map((r) => r.id).sort()).toEqual(["a_waitpoint", "waitpoint_local"]); 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); + // The id names its shard, so the partition must not widen to the gen-1 partner. expect(legacyStore.calls).toHaveLength(0); - expect(newStore.calls.map((c) => c.method)).toEqual(["countPendingWaitpointsWithPresence"]); + expect(gen2Stores[0]!.calls).toHaveLength(1); + const where = (gen2Stores[0]!.calls[0]!.args[0] as { where?: unknown }).where; + expect(idsFromWhere(where)).toEqual(["a_waitpoint"]); }); - 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" + it("counts a pending gen-2 token through the partition fallback", async () => { + const { router, legacyStore, gen2Stores } = buildRouterFor( + TOPOLOGIES[1]!, + { waitpoints: [] }, + { waitpoints: [] }, + { waitpoints: [{ id: "a_waitpoint", status: "PENDING" }] } ); + const count = await router.countPendingWaitpoints(["a_waitpoint"], undefined, "new_run"); expect(count).toBe(1); - // Fallback queried the other store with ONLY the id missing on the run's store. It uses the - // presence variant so the results can be unioned by id (a drain mirror counts once at N). - expect(newStore.calls.map((c) => c.method)).toEqual(["countPendingWaitpointsWithPresence"]); - 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); + expect(gen2Stores[0]!.calls.map((c) => c.method)).toEqual([ + "countPendingWaitpointsWithPresence", + ]); }); }); From 800e8c62b2807b40bbb750fe0e457c7dbedac441 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 28 Aug 2026 11:33:38 +0100 Subject: [PATCH 2/2] test(webapp): give the NEW store its production schema variant Production builds the NEW run-ops store with schemaVariant "dedicated", but both hetero routing tests omitted it, so the store defaulted to "legacy". The two variants take different read paths: findRun routes through the dedicated relation-stripping adapter only when the variant is set, so the tests were exercising the legacy path for a store that is dedicated in production. Verified the branch is now reached by making the dedicated arm throw, which failed 3 of the 5 tests. --- apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts | 6 +++++- apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts b/apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts index 0cd239d675a..ff54586c01a 100644 --- a/apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts +++ b/apps/webapp/test/performTaskRunAlertsStoreRouting.test.ts @@ -34,7 +34,11 @@ vi.setConfig({ testTimeout: 60_000 }); // routed read, so each store uses its OWN prisma. function buildRoutingStore(prisma17: PrismaClient, prisma14: PrismaClient) { - const newStore = new PostgresRunStore({ prisma: prisma17, readOnlyPrisma: prisma17 }); + const newStore = new PostgresRunStore({ + prisma: prisma17, + readOnlyPrisma: prisma17, + schemaVariant: "dedicated", + }); const legacyStore = new PostgresRunStore({ prisma: prisma14, readOnlyPrisma: prisma14 }); return new RoutingRunStore({ new: newStore, legacy: legacyStore }); } diff --git a/apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts b/apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts index 3ce92d18106..306af1766fb 100644 --- a/apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts +++ b/apps/webapp/test/updateMetadataStoreRoutingHetero.test.ts @@ -23,7 +23,11 @@ vi.setConfig({ testTimeout: 60_000 }); // to NEW, a 25-char cuid to LEGACY. function buildRoutingStore(prisma17: PrismaClient, prisma14: PrismaClient) { - const newStore = new PostgresRunStore({ prisma: prisma17, readOnlyPrisma: prisma17 }); + const newStore = new PostgresRunStore({ + prisma: prisma17, + readOnlyPrisma: prisma17, + schemaVariant: "dedicated", + }); const legacyStore = new PostgresRunStore({ prisma: prisma14, readOnlyPrisma: prisma14 }); return new RoutingRunStore({ new: newStore, legacy: legacyStore }); }