diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index cabe0159a6..291a074a2f 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -1746,6 +1746,10 @@ const EnvironmentSchema = z RUN_OPS_MINT_ENABLED: BoolEnv.default(false), RUN_OPS_MINT_FLAG_CACHE_TTL_MS: z.coerce.number().int().default(30_000), RUN_OPS_MINT_FLAG_CACHE_MAX_ENTRIES: z.coerce.number().int().default(10_000), + // Deterministic wall-clock cutover after a runOpsMintKind flip. Must exceed the sum + // of RUN_OPS_MINT_FLAG_CACHE_TTL_MS and the control-plane cache TTL so every process + // (stale or fresh) resolves to the same kind for the whole window. See mintFlipGrace.ts. + RUN_OPS_MINT_FLIP_GRACE_MS: z.coerce.number().int().default(90_000), // Session replication (Postgres → ClickHouse sessions_v1). Shares Redis // with the runs replicator for leader locking but has its own slot and diff --git a/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts index 0e04e909eb..c9179d5912 100644 --- a/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts @@ -1,4 +1,5 @@ import type { BatchTaskRunExecutionResult } from "@trigger.dev/core/v3"; +import { ownerEngine } from "@trigger.dev/core/v3/isomorphic"; import { $replica, type PrismaClientOrTransaction, @@ -8,7 +9,6 @@ import { import type { TaskRunWithAttempts } from "~/models/taskRun.server"; import { executionResultForTaskRun } from "~/models/taskRun.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; -import { readThroughRun } from "~/v3/runOpsMigration/readThrough.server"; import { runStore as defaultRunStore } from "~/v3/runStore.server"; import { BasePresenter } from "./basePresenter.server"; @@ -43,11 +43,14 @@ const memberRunSelect = { } as const; /** - * Split on: the batch row + its item rows resolve new-run-ops first, then the LEGACY RUN-OPS - * READ REPLICA ONLY (never the legacy primary — there is no such handle); each member run is - * hydrated independently via readThroughRun keyed on the member runId, so a batch whose members - * span migrated + abandoned runs returns the complete reachable set (the batch-spanning-the-line - * read; the dangling-reference termination gate is a separate, adjacent unit). + * Split on: the batch row + its members resolve new-run-ops first, then the LEGACY RUN-OPS READ + * REPLICA ONLY (never the legacy primary — there is no such handle). Members hydrate via ONE + * grouped read against `newClient` for the whole id set, then ONE grouped read against + * `legacyReplica` for just the misses that could still be legacy-resident — the same + * residency-partitioned shape as the old per-member read-through, but batched instead of fanned + * out one query per member. A batch whose members span migrated + abandoned runs returns the + * complete reachable set (the batch-spanning-the-line read; the dangling-reference termination + * gate is a separate, adjacent unit). * * Split off (single-DB / self-host): one passthrough read for the batch row + a single store * id-set hydrate for the members — no legacy read, no known-migrated probe, no second connection. @@ -133,7 +136,7 @@ export class ApiBatchResultsPresenter extends BasePresenter { // Split: resolve the batch row new-first then off the legacy READ REPLICA only (a batch id may // be cuid or run-ops id, and a cuid-shaped id can still have been backfilled onto NEW, so id-shape // residency is not authoritative for the row — the new-first-then-legacy probe is), then - // hydrate every member run independently via the per-run read-through primitive. + // hydrate every member run in ONE grouped new-then-legacy read. async #callSplit( friendlyId: string, env: AuthenticatedEnvironment @@ -175,41 +178,35 @@ export class ApiBatchResultsPresenter extends BasePresenter { }; } - const readMemberRun = (client: PrismaClientOrTransaction, taskRunId: string) => - client.taskRun.findFirst({ - where: { id: taskRunId }, - select: memberRunSelect, - }) as Promise; - - // Per-member fan-out: each member may live on a different DB, so a single nested include cannot - // cross the seam. Promise.all preserves batchRun.items order, unchanged from today. - const memberResults = await Promise.all( - batchRun.items.map(async (item) => { - const result = await readThroughRun({ - runId: item.taskRunId, - environmentId: env.id, - readNew: (client) => readMemberRun(client, item.taskRunId), - readLegacy: (replica) => readMemberRun(replica, item.taskRunId), - deps: { - splitEnabled: true, - // Pass the SAME resolved handles the batch row used, so the batch row and its members - // never resolve against different DBs. (Letting these fall through to readThroughRun's - // own module-level defaults would diverge from the batch read's `?? this._replica`.) - newClient, - legacyReplica, - isPastRetention: this.readThrough?.isPastRetention, - }, - }); + const taskRunIds = batchRun.items.map((item) => item.taskRunId); - // not-found / past-retention members are omitted (matches today's drop-undefined behavior); - // the dangling-reference termination gate (separate unit) governs whether that's permitted. - if (result.source === "not-found" || result.source === "past-retention") { - return undefined; - } + const newRows = (await newClient.taskRun.findMany({ + where: { id: { in: taskRunIds } }, + select: memberRunSelect, + })) as TaskRunWithAttempts[]; + const runsById = new Map(newRows.map((run) => [run.id, run])); - return executionResultForTaskRun(result.value); - }) + // A run-ops id can only live on NEW, so only misses that AREN'T run-ops-shaped are candidates + // for the legacy probe — mirrors readThroughRun's per-id "NEW residency skips legacy" rule. + const legacyCandidateIds = taskRunIds.filter( + (id) => !runsById.has(id) && ownerEngine(id) !== "NEW" ); + if (legacyCandidateIds.length > 0) { + const legacyRows = (await legacyReplica.taskRun.findMany({ + where: { id: { in: legacyCandidateIds } }, + select: memberRunSelect, + })) as TaskRunWithAttempts[]; + for (const run of legacyRows) { + runsById.set(run.id, run); + } + } + + // not-found members are omitted (matches today's drop-undefined behavior); the + // dangling-reference termination gate (separate unit) governs whether that's permitted. + const memberResults = batchRun.items.map((item) => { + const run = runsById.get(item.taskRunId); + return run ? executionResultForTaskRun(run) : undefined; + }); return { id: batchRun.friendlyId, diff --git a/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts index a5c71a5555..c6a14a4492 100644 --- a/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts @@ -164,17 +164,13 @@ export class ApiRetrieveRunPresenter { collect(child.lockedToVersionId); } + const lockedWorkersByVersionId = + await controlPlaneResolver.resolveRunLockedWorkersByVersionIds([...distinctVersionIds]); const versionById = new Map( - await Promise.all( - [...distinctVersionIds].map( - async (id) => - [ - id, - (await controlPlaneResolver.resolveRunLockedWorker({ lockedToVersionId: id })) - ?.lockedToVersion?.version ?? null, - ] as const - ) - ) + [...distinctVersionIds].map((id) => [ + id, + lockedWorkersByVersionId.get(id)?.lockedToVersion?.version ?? null, + ]) ); const resolveVersion = (id: string | null): { version: string } | null => { diff --git a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts index 74f0c06fdd..29ef066572 100644 --- a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts @@ -11,6 +11,12 @@ import { waitpointStatusToApiStatus } from "./WaitpointListPresenter.server"; export type WaitpointDetail = NonNullable>>; +// Single-sourced display bound for a waitpoint's connected run friendlyIds. +export const CONNECTED_RUNS_DISPLAY_LIMIT = 5; + +// Over-read connection rows on the FK-free dedicated join so danglers don't cost a display slot. +export const CONNECTED_RUNS_CONNECTION_SCAN_LIMIT = CONNECTED_RUNS_DISPLAY_LIMIT * 5; + export class WaitpointPresenter extends BasePresenter { constructor( prisma?: PrismaClientOrTransaction, @@ -76,8 +82,7 @@ export class WaitpointPresenter extends BasePresenter { // Connected-run friendlyIds gathered across BOTH stores. The run<->waitpoint join co-locates with // the RUN (written on the run's DB), so the waitpoint's own store misses a cross-DB connection; we - // read the join on each client and resolve the run's friendlyId on that same client, then union. - // We never relation-select `connectedRuns`: it is not a field on the dedicated subset `Waitpoint`. + // read the join on each client, resolve the run's friendlyId on that same client, and union. async #connectedRunFriendlyIds(waitpointId: string): Promise { const replica = this._replica as unknown as PrismaReplicaClient; const rawClients: PrismaReplicaClient[] = @@ -91,47 +96,66 @@ export class WaitpointPresenter extends BasePresenter { const friendlyIds = new Set(); for (const client of clients) { - const runIds = await this.#connectedRunIdsOn(client, waitpointId); - if (runIds.length === 0) { - continue; + for (const friendlyId of await this.#connectedRunFriendlyIdsOn(client, waitpointId)) { + friendlyIds.add(friendlyId); } - const runs = await client.taskRun.findMany({ - where: { id: { in: runIds } }, - select: { friendlyId: true }, - take: 5, - }); - for (const run of runs) { - friendlyIds.add(run.friendlyId); - } - if (friendlyIds.size >= 5) { + if (friendlyIds.size >= CONNECTED_RUNS_DISPLAY_LIMIT) { break; } } - return Array.from(friendlyIds).slice(0, 5); + return Array.from(friendlyIds).slice(0, CONNECTED_RUNS_DISPLAY_LIMIT); } - // Schema-aware read of the run ids linked to a waitpoint: the dedicated subset uses the explicit - // `WaitpointRunConnection` model, the control-plane full schema the implicit `_WaitpointRunConnections` - // M2M (A = TaskRun.id, B = Waitpoint.id). The dedicated join delegate is absent on the full client. - async #connectedRunIdsOn(client: PrismaReplicaClient, waitpointId: string): Promise { - const joinDelegate = ( + // Connected-run friendlyIds for one store, via the ORM. Two indexed reads joined in memory instead + // of a SQL JOIN onto the (very large) TaskRun table: `id IN (...)` can only plan as a PK lookup, so + // the planner can never scan TaskRun. + // + // Dedicated subset: the explicit `WaitpointRunConnection` is scalar (`taskRunId`, no FK), so a + // connection can outlive a deleted run. We over-read connection ids, resolve runs by id (a missing + // id just drops out -- danglers cost no display slot), and cap at the display limit. + // + // Control-plane full schema: no queryable join delegate (implicit M2M), so we traverse the + // `connectedRuns` relation; it cascade-deletes with the run, so no dangler can exist. + async #connectedRunFriendlyIdsOn( + client: PrismaReplicaClient, + waitpointId: string + ): Promise { + const dedicated = ( client as unknown as { waitpointRunConnection?: { findMany: (args: unknown) => Promise<{ taskRunId: string }[]>; }; } ).waitpointRunConnection; - if (joinDelegate && typeof joinDelegate.findMany === "function") { - const links = await joinDelegate.findMany({ + + if (dedicated) { + const connections = await dedicated.findMany({ where: { waitpointId }, select: { taskRunId: true }, + take: CONNECTED_RUNS_CONNECTION_SCAN_LIMIT, }); - return links.map((link) => link.taskRunId); + if (connections.length === 0) { + return []; + } + const runs = await client.taskRun.findMany({ + where: { id: { in: connections.map((connection) => connection.taskRunId) } }, + select: { friendlyId: true }, + take: CONNECTED_RUNS_DISPLAY_LIMIT, + }); + return runs.map((run) => run.friendlyId); } - const rows = await client.$queryRaw<{ A: string }[]>` - SELECT "A" FROM "_WaitpointRunConnections" WHERE "B" = ${waitpointId} - `; - return rows.map((row) => row.A); + + const waitpoint = (await ( + client.waitpoint.findFirst as ( + args: unknown + ) => Promise<{ connectedRuns: { friendlyId: string }[] } | null> + )({ + where: { id: waitpointId }, + select: { + connectedRuns: { select: { friendlyId: true }, take: CONNECTED_RUNS_DISPLAY_LIMIT }, + }, + })) as { connectedRuns: { friendlyId: string }[] } | null; + return (waitpoint?.connectedRuns ?? []).map((run) => run.friendlyId); } public async call({ diff --git a/apps/webapp/app/routes/admin.api.v1.feature-flags.ts b/apps/webapp/app/routes/admin.api.v1.feature-flags.ts index 0c668c848c..b67a9dde02 100644 --- a/apps/webapp/app/routes/admin.api.v1.feature-flags.ts +++ b/apps/webapp/app/routes/admin.api.v1.feature-flags.ts @@ -1,9 +1,15 @@ import type { ActionFunctionArgs } from "@remix-run/server-runtime"; import { json } from "@remix-run/server-runtime"; import { prisma } from "~/db.server"; +import { env } from "~/env.server"; import { requireAdminApiRequest } from "~/services/personalAccessToken.server"; import { makeSetMultipleFlags } from "~/v3/featureFlags.server"; -import { validatePartialFeatureFlags } from "~/v3/featureFlags"; +import { + FEATURE_FLAG, + type FeatureFlagCatalog, + validatePartialFeatureFlags, +} from "~/v3/featureFlags"; +import { stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; export async function action({ request }: ActionFunctionArgs) { await requireAdminApiRequest(request); @@ -24,9 +30,51 @@ export async function action({ request }: ActionFunctionArgs) { ); } - const featureFlags = validationResult.data; + // Derived grace-stamp fields are computed server-side; never trust them from the body. + const { + runOpsMintKindPrev: _ignoredPrev, + runOpsMintKindFlippedAt: _ignoredFlippedAt, + ...requestedFlags + } = validationResult.data; + + let flagsToWrite: Partial = requestedFlags; + + if (requestedFlags.runOpsMintKind !== undefined) { + // Read the current GLOBAL mint flags so the stamp is computed against the authoritative + // stored state, mirroring the per-org route. stampMintKindFlip writes prev/flippedAt only + // on a genuine global flip, and carries an in-flight stamp forward on a same-target save. + const existingRows = await prisma.featureFlag.findMany({ + where: { + key: { + in: [ + FEATURE_FLAG.runOpsMintKind, + FEATURE_FLAG.runOpsMintKindPrev, + FEATURE_FLAG.runOpsMintKindFlippedAt, + ], + }, + }, + select: { key: true, value: true }, + }); + const existingGlobal: Record = {}; + for (const row of existingRows) { + existingGlobal[row.key] = row.value; + } + + // Anchor the cutover to the control-plane DB clock, not this process's wall clock. + const [{ now: controlPlaneNow }] = await prisma.$queryRaw< + { now: Date }[] + >`SELECT now() AS now`; + + flagsToWrite = stampMintKindFlip( + existingGlobal, + { ...requestedFlags }, + controlPlaneNow.getTime(), + env.RUN_OPS_MINT_FLIP_GRACE_MS + ) as Partial; + } + const setMultipleFlags = makeSetMultipleFlags(prisma); - const updatedFlags = await setMultipleFlags(featureFlags); + const updatedFlags = await setMultipleFlags(flagsToWrite); return json({ success: true, diff --git a/apps/webapp/app/routes/admin.api.v1.orgs.$organizationId.feature-flags.ts b/apps/webapp/app/routes/admin.api.v1.orgs.$organizationId.feature-flags.ts index e5fd7f7963..db1524f855 100644 --- a/apps/webapp/app/routes/admin.api.v1.orgs.$organizationId.feature-flags.ts +++ b/apps/webapp/app/routes/admin.api.v1.orgs.$organizationId.feature-flags.ts @@ -1,10 +1,14 @@ import type { ActionFunctionArgs, LoaderFunctionArgs } from "@remix-run/server-runtime"; import { json } from "@remix-run/server-runtime"; +import type { Prisma } from "@trigger.dev/database"; import { z } from "zod"; +import { env } from "~/env.server"; import { prisma } from "~/db.server"; import { requireAdminApiRequest } from "~/services/personalAccessToken.server"; import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; +import { selectMintBaselineSource, stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; import { validatePartialFeatureFlags } from "~/v3/featureFlags"; +import { flags as getGlobalFlags } from "~/v3/featureFlags.server"; const ParamsSchema = z.object({ organizationId: z.string(), @@ -48,20 +52,6 @@ export async function action({ request, params }: ActionFunctionArgs) { const { organizationId } = ParamsSchema.parse(params); - const organization = await prisma.organization.findFirst({ - where: { - id: organizationId, - }, - select: { - id: true, - featureFlags: true, - }, - }); - - if (!organization) { - return json({ error: "Organization not found" }, { status: 404 }); - } - try { const body = await request.json(); @@ -77,31 +67,66 @@ export async function action({ request, params }: ActionFunctionArgs) { ); } - // Merge new flags with existing flags - const existingFlags = organization.featureFlags - ? validatePartialFeatureFlags(organization.featureFlags as Record) - : { success: false as const }; - - const mergedFlags = { - ...(existingFlags.success ? existingFlags.data : {}), - ...validationResult.data, - }; + // Derived grace-stamp fields are computed server-side; never trust them from the body. + const { + runOpsMintKindPrev: _ignoredPrev, + runOpsMintKindFlippedAt: _ignoredFlippedAt, + ...requestedFlags + } = validationResult.data; + + // Seed the flip baseline from the current GLOBAL mint flags so an org's FIRST per-org override + // is graced from the currently-effective global kind, not the hardcoded default "cuid". + const globalFlags = (await getGlobalFlags()) as Record; + + // Lock the org row for the whole read -> merge -> stamp -> write so a concurrent flag save + // can't clobber the grace metadata (read-then-write race). PK lookup, one row, held to commit. + const updatedOrganization = await prisma.$transaction(async (tx) => { + const rows = await tx.$queryRaw<{ featureFlags: unknown }[]>` + SELECT "featureFlags" FROM "Organization" WHERE "id" = ${organizationId} FOR UPDATE`; + + if (rows.length === 0) { + return null; + } + + const existingRaw = rows[0].featureFlags as Record | null; + const existingResult = existingRaw + ? validatePartialFeatureFlags(existingRaw) + : ({ success: false } as const); + const existingData = existingResult.success ? existingResult.data : {}; + + // Stamp the flip from the control-plane DB clock so the grace-window cutover is anchored to + // one authoritative time source, not whichever webapp process handled this request. + const [{ now: controlPlaneNow }] = await tx.$queryRaw<{ now: Date }[]>`SELECT now() AS now`; + + const mergedFlags = stampMintKindFlip( + selectMintBaselineSource(existingRaw, globalFlags), + { + ...existingData, + ...requestedFlags, + }, + controlPlaneNow.getTime(), + env.RUN_OPS_MINT_FLIP_GRACE_MS + ); - // Update the organization's feature flags - const updatedOrganization = await prisma.organization.update({ - where: { - id: organizationId, - }, - data: { - featureFlags: mergedFlags, - }, - select: { - id: true, - slug: true, - featureFlags: true, - }, + return tx.organization.update({ + where: { + id: organizationId, + }, + data: { + featureFlags: mergedFlags as Prisma.InputJsonValue, + }, + select: { + id: true, + slug: true, + featureFlags: true, + }, + }); }); + if (!updatedOrganization) { + return json({ error: "Organization not found" }, { status: 404 }); + } + // Org feature flags are embedded in every env of the org; drop all its cached env rows. controlPlaneResolver.invalidateOrganization(organizationId); diff --git a/apps/webapp/app/routes/admin.api.v2.orgs.$organizationId.feature-flags.ts b/apps/webapp/app/routes/admin.api.v2.orgs.$organizationId.feature-flags.ts index 3c62d9c7a5..0071054bd3 100644 --- a/apps/webapp/app/routes/admin.api.v2.orgs.$organizationId.feature-flags.ts +++ b/apps/webapp/app/routes/admin.api.v2.orgs.$organizationId.feature-flags.ts @@ -2,9 +2,11 @@ import type { ActionFunctionArgs, LoaderFunctionArgs } from "@remix-run/server-r import { json } from "@remix-run/server-runtime"; import { Prisma } from "@trigger.dev/database"; import { z } from "zod"; +import { env } from "~/env.server"; import { prisma } from "~/db.server"; import { requireUser } from "~/services/session.server"; import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; +import { selectMintBaselineSource, stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; import { flags as getGlobalFlags } from "~/v3/featureFlags.server"; import { FEATURE_FLAG, @@ -103,34 +105,78 @@ export async function action({ request, params }: ActionFunctionArgs) { return json({ error: "Invalid JSON body" }, { status: 400 }); } - let featureFlags: typeof Prisma.JsonNull | Record; - if ( body === null || (typeof body === "object" && !Array.isArray(body) && Object.keys(body).length === 0) ) { - featureFlags = Prisma.JsonNull; - } else { - const validationResult = validatePartialFeatureFlags(body as Record); - if (!validationResult.success) { - return json( - { error: "Invalid feature flags", details: validationResult.error.issues }, - { status: 400 } - ); + // Clear all flags. No grace stamp (nothing to flip) and no read-then-write race. + try { + await prisma.organization.update({ + where: { id: organizationId }, + data: { featureFlags: Prisma.JsonNull }, + }); + } catch (e) { + if (e instanceof Prisma.PrismaClientKnownRequestError && e.code === "P2025") { + throw new Response("Organization not found", { status: 404 }); + } + throw e; } - featureFlags = validationResult.data; + + controlPlaneResolver.invalidateOrganization(organizationId); + return json({ success: true }); } - try { - await prisma.organization.update({ + const validationResult = validatePartialFeatureFlags(body as Record); + if (!validationResult.success) { + return json( + { error: "Invalid feature flags", details: validationResult.error.issues }, + { status: 400 } + ); + } + + // Derived grace-stamp fields are computed server-side; never trust them from the body. + const { + runOpsMintKindPrev: _ignoredPrev, + runOpsMintKindFlippedAt: _ignoredFlippedAt, + ...requestedFlags + } = validationResult.data; + + // Seed the flip baseline from the current GLOBAL mint flags so an org's FIRST per-org override + // is graced from the currently-effective global kind, not the hardcoded default "cuid". + const globalFlags = (await getGlobalFlags()) as Record; + + // Lock the org row for the whole read -> stamp -> write so a concurrent flag save can't clobber + // the grace metadata (read-then-write race). PK lookup, one row, held to commit. + const updated = await prisma.$transaction(async (tx) => { + const rows = await tx.$queryRaw<{ featureFlags: unknown }[]>` + SELECT "featureFlags" FROM "Organization" WHERE "id" = ${organizationId} FOR UPDATE`; + + if (rows.length === 0) { + return false; + } + + const existingRaw = rows[0].featureFlags as Record | null; + + // Anchor the flip stamp to the control-plane DB clock (see the v1 route), not this process's. + const [{ now: controlPlaneNow }] = await tx.$queryRaw<{ now: Date }[]>`SELECT now() AS now`; + + const stamped = stampMintKindFlip( + selectMintBaselineSource(existingRaw, globalFlags), + requestedFlags, + controlPlaneNow.getTime(), + env.RUN_OPS_MINT_FLIP_GRACE_MS + ); + + await tx.organization.update({ where: { id: organizationId }, - data: { featureFlags: featureFlags as Prisma.InputJsonValue }, + data: { featureFlags: stamped as Prisma.InputJsonValue }, }); - } catch (e) { - if (e instanceof Prisma.PrismaClientKnownRequestError && e.code === "P2025") { - throw new Response("Organization not found", { status: 404 }); - } - throw e; + + return true; + }); + + if (!updated) { + throw new Response("Organization not found", { status: 404 }); } // Org feature flags are embedded in every env of the org; drop all its cached env rows. diff --git a/apps/webapp/app/v3/featureFlags.ts b/apps/webapp/app/v3/featureFlags.ts index 637830aef0..49a18eaf35 100644 --- a/apps/webapp/app/v3/featureFlags.ts +++ b/apps/webapp/app/v3/featureFlags.ts @@ -19,6 +19,9 @@ export const FEATURE_FLAG = { computeMigrationRequireTemplate: "computeMigrationRequireTemplate", devBranchesEnabled: "devBranchesEnabled", runOpsMintKind: "runOpsMintKind", + // Grace-linger stamp carried alongside runOpsMintKind on flip. See mintFlipGrace.ts. + runOpsMintKindPrev: "runOpsMintKindPrev", + runOpsMintKindFlippedAt: "runOpsMintKindFlippedAt", } as const; export const FeatureFlagCatalog = { @@ -54,6 +57,10 @@ export const FeatureFlagCatalog = { // Per-org run-ops-id mint cutover. Defaults to "cuid"; only honored when // RUN_OPS_MINT_ENABLED is on AND isSplitEnabled() is true. [FEATURE_FLAG.runOpsMintKind]: z.enum(["cuid", "runOpsId"]), + // Grace-linger stamp: the previously-effective kind and the flip timestamp, written + // by stampMintKindFlip on a genuine flip. Display-only (see ORG_LOCKED_FLAGS). + [FEATURE_FLAG.runOpsMintKindPrev]: z.enum(["cuid", "runOpsId"]), + [FEATURE_FLAG.runOpsMintKindFlippedAt]: z.string().datetime(), }; export type FeatureFlagKey = keyof typeof FeatureFlagCatalog; @@ -70,6 +77,8 @@ export const GLOBAL_LOCKED_FLAGS: FeatureFlagKey[] = [ export const ORG_LOCKED_FLAGS: FeatureFlagKey[] = [ FEATURE_FLAG.defaultWorkerInstanceGroupId, FEATURE_FLAG.taskEventRepository, + FEATURE_FLAG.runOpsMintKindPrev, + FEATURE_FLAG.runOpsMintKindFlippedAt, ]; // Create a Zod schema from the existing catalog diff --git a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.test.ts b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.test.ts new file mode 100644 index 0000000000..cd08c8a077 --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.test.ts @@ -0,0 +1,168 @@ +// Real control-plane (single Postgres) proof that resolving the locked-worker version for many +// runs' distinct `lockedToVersionId`s is a GROUPED query, not one `backgroundWorker.findFirst` +// per id. The DB is never mocked: the call-counting proxy below delegates every call to the +// real Prisma client (the real query still runs against the real container) - it only tallies +// how many times each model.method pair was invoked. +import { postgresTest } from "@internal/testcontainers"; +import { describe, expect } from "vitest"; +import type { PrismaClient, PrismaReplicaClient } from "@trigger.dev/database"; +import { ControlPlaneCache } from "./controlPlaneCache.server"; +import { ControlPlaneResolver } from "./controlPlaneResolver.server"; + +// Wraps a real Prisma client so every `client..(...)` call is tallied by +// `"."` before being forwarded, unmodified, to the real delegate. No behavior is +// faked or short-circuited - this is instrumentation, not a mock. +function createCallCountingProxy( + client: T +): { client: T; counts: Map } { + const counts = new Map(); + const bump = (key: string) => counts.set(key, (counts.get(key) ?? 0) + 1); + const wrappedModels = new Map(); + + const wrapModel = (model: string, delegate: object) => + new Proxy(delegate, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if (typeof prop === "string" && typeof value === "function") { + return (...args: unknown[]) => { + bump(`${model}.${prop}`); + return (value as (...a: unknown[]) => unknown).apply(target, args); + }; + } + return value; + }, + }); + + const proxy = new Proxy(client, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if ( + typeof prop === "string" && + value && + typeof value === "object" && + typeof (value as { findMany?: unknown }).findMany === "function" + ) { + if (!wrappedModels.has(prop)) { + wrappedModels.set(prop, wrapModel(prop, value as object)); + } + return wrappedModels.get(prop); + } + return value; + }, + }); + + return { client: proxy as T, counts }; +} + +let n = 0; + +async function seedEnv(prisma: PrismaClient) { + const s = n++; + const organization = await prisma.organization.create({ + data: { title: `Org ${s}`, slug: `org-${s}` }, + }); + const project = await prisma.project.create({ + data: { + name: `P ${s}`, + slug: `p-${s}`, + externalRef: `proj_${s}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "PRODUCTION", + slug: `env-${s}`, + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_${s}`, + pkApiKey: `pk_${s}`, + shortcode: `sc_${s}`, + }, + }); + return { organization, project, environment }; +} + +async function seedBackgroundWorker( + prisma: PrismaClient, + ctx: { projectId: string; runtimeEnvironmentId: string }, + version: string +) { + const s = n++; + return prisma.backgroundWorker.create({ + data: { + friendlyId: `worker_${s}`, + version, + contentHash: `hash_${s}`, + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.runtimeEnvironmentId, + metadata: {}, + }, + }); +} + +describe("ControlPlaneResolver.resolveRunLockedWorkersByVersionIds", () => { + postgresTest( + "issues ONE grouped findMany and ZERO per-id findFirst for N distinct ids", + async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma); + const workers = await Promise.all( + [0, 1, 2].map((i) => + seedBackgroundWorker( + prisma, + { projectId: project.id, runtimeEnvironmentId: environment.id }, + `2024010${i}.0` + ) + ) + ); + + const { client: countedReplica, counts } = createCallCountingProxy(prisma); + + const resolver = new ControlPlaneResolver({ + controlPlanePrimary: prisma, + controlPlaneReplica: countedReplica as unknown as PrismaReplicaClient, + cache: new ControlPlaneCache(), + splitEnabled: () => true, + }); + + const ids = workers.map((w) => w.id); + const result = await resolver.resolveRunLockedWorkersByVersionIds(ids); + + expect(counts.get("backgroundWorker.findMany") ?? 0).toBe(1); + expect(counts.get("backgroundWorker.findFirst") ?? 0).toBe(0); + + for (const worker of workers) { + expect(result.get(worker.id)?.lockedToVersion?.version).toBe(worker.version); + } + } + ); + + postgresTest("cache-hit ids issue no query at all", async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma); + const worker = await seedBackgroundWorker( + prisma, + { projectId: project.id, runtimeEnvironmentId: environment.id }, + "20240101.0" + ); + + const cache = new ControlPlaneCache(); + const { client: countedReplica, counts } = createCallCountingProxy(prisma); + + const resolver = new ControlPlaneResolver({ + controlPlanePrimary: prisma, + controlPlaneReplica: countedReplica as unknown as PrismaReplicaClient, + cache, + splitEnabled: () => true, + }); + + // Warm the cache. + await resolver.resolveRunLockedWorkersByVersionIds([worker.id]); + expect(counts.get("backgroundWorker.findMany") ?? 0).toBe(1); + + // Second call for the same id is served entirely from cache - no query at all. + const result = await resolver.resolveRunLockedWorkersByVersionIds([worker.id]); + expect(counts.get("backgroundWorker.findMany") ?? 0).toBe(1); + expect(counts.get("backgroundWorker.findFirst") ?? 0).toBe(0); + expect(result.get(worker.id)?.lockedToVersion?.version).toBe(worker.version); + }); +}); diff --git a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts index fb903bc8ea..d6a706c966 100644 --- a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts @@ -70,6 +70,15 @@ function lockedWorkerKey(lockedById?: string | null, lockedToVersionId?: string return `${lockedById ?? "_"}:${lockedToVersionId ?? "_"}`; } +type LockedToVersionRow = { + id: string; + version: string; + sdkVersion: string; + runtime: string | null; + runtimeVersion: string | null; + supportsLazyAttempts: boolean; +}; + export class ControlPlaneResolver { private readonly controlPlanePrimary: PrismaClient; private readonly controlPlaneReplica: PrismaReplicaClient; @@ -243,6 +252,87 @@ export class ControlPlaneResolver { }; } + /** + * Grouped counterpart to `resolveRunLockedWorker({ lockedToVersionId })`: resolves the + * `lockedToVersion` for many distinct BackgroundWorker ids with ONE `findMany` for the + * cache misses (BackgroundWorker is control-plane only, no run-ops residency split), instead + * of one `findFirst` per id. Cache keys/shape match `resolveRunLockedWorker`'s + * `lockedById=undefined` slot, so entries populated here are reused by that method and + * vice versa. + */ + async resolveRunLockedWorkersByVersionIds( + ids: string[] + ): Promise> { + const result = new Map(); + const uniqueIds = [...new Set(ids)]; + if (uniqueIds.length === 0) { + return result; + } + + if (!this.splitEnabled()) { + const rows = await this.#queryLockedToVersionRows(this.controlPlanePrimary, uniqueIds); + for (const id of uniqueIds) { + result.set(id, this.#toResolvedRunLockedWorker(rows.get(id))); + } + return result; + } + + const misses: string[] = []; + for (const id of uniqueIds) { + const cached = this.cache.getLockedWorker(lockedWorkerKey(undefined, id)); + if (cached !== undefined) { + result.set(id, cached); + } else { + misses.push(id); + } + } + + if (misses.length > 0) { + const rows = await this.#queryLockedToVersionRows(this.controlPlaneReplica, misses); + for (const id of misses) { + const value = this.#toResolvedRunLockedWorker(rows.get(id)); + this.cache.setLockedWorker(lockedWorkerKey(undefined, id), value); + result.set(id, value); + } + } + + return result; + } + + async #queryLockedToVersionRows( + client: CpClient, + ids: string[] + ): Promise> { + const rows = await client.backgroundWorker.findMany({ + where: { id: { in: ids } }, + select: { + id: true, + version: true, + sdkVersion: true, + runtime: true, + runtimeVersion: true, + supportsLazyAttempts: true, + }, + }); + return new Map(rows.map((row) => [row.id, row])); + } + + #toResolvedRunLockedWorker(row: LockedToVersionRow | undefined): ResolvedRunLockedWorker | null { + if (!row) { + return null; + } + return { + lockedBy: null, + lockedToVersion: { + version: row.version, + sdkVersion: row.sdkVersion, + runtime: row.runtime, + runtimeVersion: row.runtimeVersion, + supportsLazyAttempts: row.supportsLazyAttempts, + }, + }; + } + async resolveWorkerVersion(args: { environmentId: string; backgroundWorkerId?: string; diff --git a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts new file mode 100644 index 0000000000..667068842f --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts @@ -0,0 +1,322 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + effectiveMintKind, + readMintResolution, + resolveMintFlag, + selectMintBaselineSource, + stampMintKindFlip, + type MintFlagResolution, +} from "./mintFlipGrace"; + +// GRACE-LINGER: during [flippedAt, flippedAt + GRACE) every process — stale or fresh — +// must resolve to the SAME (old) kind; at/after the cutover every process resolves to +// the SAME (new) kind. This collapses the cross-process divergence window. +const GRACE_MS = 90_000; +const T = 1_000_000; + +describe("effectiveMintKind", () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(() => vi.useRealTimers()); + + it("returns r.kind directly when prev is missing", () => { + const r: MintFlagResolution = { kind: "cuid" }; + expect(effectiveMintKind(r, T, GRACE_MS)).toBe("cuid"); + }); + + it("returns r.kind directly when flippedAtMs is missing", () => { + const r: MintFlagResolution = { kind: "runOpsId", prev: "cuid" }; + expect(effectiveMintKind(r, T, GRACE_MS)).toBe("runOpsId"); + }); + + it("CORE: stale and fresh resolutions agree at every instant during grace, then both flip to the new kind at cutover", () => { + const stale: MintFlagResolution = { kind: "cuid" }; + const fresh: MintFlagResolution = { kind: "runOpsId", prev: "cuid", flippedAtMs: T }; + + for (const now of [T, T + 1, T + 1_000, T + 45_000, T + GRACE_MS - 1]) { + const staleResolved = effectiveMintKind(stale, now, GRACE_MS); + const freshResolved = effectiveMintKind(fresh, now, GRACE_MS); + expect(staleResolved).toBe("cuid"); + expect(freshResolved).toBe("cuid"); + } + + for (const now of [T + GRACE_MS, T + GRACE_MS + 1, T + GRACE_MS + 60_000]) { + expect(effectiveMintKind(fresh, now, GRACE_MS)).toBe("runOpsId"); + } + }); + + it("boundary: exactly at flippedAt + GRACE resolves to the NEW kind", () => { + const fresh: MintFlagResolution = { kind: "runOpsId", prev: "cuid", flippedAtMs: T }; + expect(effectiveMintKind(fresh, T + GRACE_MS - 1, GRACE_MS)).toBe("cuid"); + expect(effectiveMintKind(fresh, T + GRACE_MS, GRACE_MS)).toBe("runOpsId"); + }); +}); + +describe("stampMintKindFlip", () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(() => vi.useRealTimers()); + + it("a genuine flip (cuid -> runOpsId) stamps prev and flippedAt", () => { + const existing = { runOpsMintKind: "cuid" }; + const outgoing = { runOpsMintKind: "runOpsId" }; + const result = stampMintKindFlip(existing, outgoing, T, GRACE_MS); + + expect(result.runOpsMintKind).toBe("runOpsId"); + expect(result.runOpsMintKindPrev).toBe("cuid"); + expect(result.runOpsMintKindFlippedAt).toBe(new Date(T).toISOString()); + }); + + it("defaults the existing effective kind to cuid when existing flags are null", () => { + const outgoing = { runOpsMintKind: "runOpsId" }; + const result = stampMintKindFlip(null, outgoing, T, GRACE_MS); + + expect(result.runOpsMintKind).toBe("runOpsId"); + expect(result.runOpsMintKindPrev).toBe("cuid"); + expect(result.runOpsMintKindFlippedAt).toBe(new Date(T).toISOString()); + }); + + it("resubmitting the same target kind mid-grace carries the stamp forward untouched (does not reset the cutover clock)", () => { + const existing = { + runOpsMintKind: "runOpsId", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: new Date(T).toISOString(), + }; + const now = T + 10_000; + const outgoing = { runOpsMintKind: "runOpsId", someOtherFlag: true }; + const result = stampMintKindFlip(existing, outgoing, now, GRACE_MS); + + expect(result.runOpsMintKind).toBe("runOpsId"); + expect(result.runOpsMintKindPrev).toBe("cuid"); + // Unrelated re-save: the cutover time must NOT slide forward. + expect(result.runOpsMintKindFlippedAt).toBe(new Date(T).toISOString()); + expect(result.someOtherFlag).toBe(true); + }); + + it("an unchanged save after grace has elapsed carries the settled stamp forward and preserves unrelated flags", () => { + const existing = { + runOpsMintKind: "runOpsId", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: new Date(T).toISOString(), + }; + const outgoing = { runOpsMintKind: "runOpsId", someOtherFlag: true }; + const result = stampMintKindFlip(existing, outgoing, T + GRACE_MS + 5_000, GRACE_MS); + + expect(result.runOpsMintKind).toBe("runOpsId"); + expect(result.someOtherFlag).toBe(true); + expect(result.runOpsMintKindPrev).toBe("cuid"); + expect(result.runOpsMintKindFlippedAt).toBe(new Date(T).toISOString()); + }); + + it("a flip-back requested after the original grace has elapsed stamps prev := the new settled (now-effective) kind, timestamped now", () => { + const existing = { + runOpsMintKind: "runOpsId", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: new Date(T).toISOString(), + }; + const now = T + GRACE_MS + 1_000; + const outgoing = { runOpsMintKind: "cuid" }; + const result = stampMintKindFlip(existing, outgoing, now, GRACE_MS); + + expect(result.runOpsMintKind).toBe("cuid"); + expect(result.runOpsMintKindPrev).toBe("runOpsId"); + expect(result.runOpsMintKindFlippedAt).toBe(new Date(now).toISOString()); + }); + + it("a flip-back mid-grace re-stamps prev to the still-effective old kind, so it keeps serving that kind (no divergence)", () => { + const existing = { + runOpsMintKind: "runOpsId", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: new Date(T).toISOString(), + }; + const now = T + 20_000; + const outgoing = { runOpsMintKind: "cuid" }; + const result = stampMintKindFlip(existing, outgoing, now, GRACE_MS); + + expect(result.runOpsMintKind).toBe("cuid"); + expect(result.runOpsMintKindPrev).toBe("cuid"); + expect(result.runOpsMintKindFlippedAt).toBe(new Date(now).toISOString()); + }); + + it("leaves runOpsMintKind untouched when the save omits it (unrelated flag change: no inject, no spurious flip)", () => { + const existing = { runOpsMintKind: "runOpsId" }; + const outgoing: Record = { someOtherFlag: true }; + const result = stampMintKindFlip(existing, outgoing, T, GRACE_MS); + + // Must not inject a default kind or stamp a flip: doing so would pin the org to an explicit + // per-org override and make a later global flip silently skip it. + expect(result.runOpsMintKind).toBeUndefined(); + expect(result.runOpsMintKindPrev).toBeUndefined(); + expect(result.runOpsMintKindFlippedAt).toBeUndefined(); + expect(result.someOtherFlag).toBe(true); + }); + + it("treats a malformed existing flippedAt as no stamp", () => { + const existing = { + runOpsMintKind: "runOpsId", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: "not-a-date", + }; + const outgoing = { runOpsMintKind: "runOpsId" }; + const result = stampMintKindFlip(existing, outgoing, T, GRACE_MS); + + expect(result.runOpsMintKind).toBe("runOpsId"); + // The malformed flippedAt is carried forward verbatim, but an unparseable timestamp is + // inert when resolved (Date.parse -> NaN -> effectiveMintKind returns the target kind). + expect(result.runOpsMintKindFlippedAt).toBe("not-a-date"); + expect( + effectiveMintKind({ kind: "runOpsId", prev: "cuid", flippedAtMs: NaN }, T, GRACE_MS) + ).toBe("runOpsId"); + }); +}); + +// SOURCE-CONSISTENCY: the kind and its grace stamp must come from the SAME source. A per-org +// runOpsMintKind override wins both the kind and the stamp; with no per-org override, BOTH the +// kind and the stamp come from the global FeatureFlag rows. Never mix (e.g. a per-org kind with +// the global stamp), which would date a grace window against the wrong flip. +describe("resolveMintFlag", () => { + it("a per-org override wins the kind AND owns the stamp, ignoring the global stamp entirely", () => { + const perOrg = { + runOpsMintKind: "runOpsId", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: new Date(T).toISOString(), + }; + const global = { + runOpsMintKind: "cuid", + runOpsMintKindPrev: "runOpsId", + runOpsMintKindFlippedAt: new Date(T + 500_000).toISOString(), + }; + expect(resolveMintFlag(perOrg, global)).toEqual({ + kind: "runOpsId", + prev: "cuid", + flippedAtMs: T, + }); + }); + + it("with NO per-org override, the kind AND the stamp come from the global rows (global flip is graced)", () => { + const global = { + runOpsMintKind: "runOpsId", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: new Date(T).toISOString(), + }; + const resolution = resolveMintFlag({}, global); + expect(resolution).toEqual({ kind: "runOpsId", prev: "cuid", flippedAtMs: T }); + // Mid-grace: a global flip resolves to the OLD kind for the whole window. + expect(effectiveMintKind(resolution, T + GRACE_MS - 1, GRACE_MS)).toBe("cuid"); + expect(effectiveMintKind(resolution, T + GRACE_MS, GRACE_MS)).toBe("runOpsId"); + }); + + it("a per-org override with NO per-org stamp does NOT borrow the global stamp (kind stays ungraced)", () => { + const perOrg = { runOpsMintKind: "runOpsId" }; + const global = { + runOpsMintKind: "cuid", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: new Date(T).toISOString(), + }; + expect(resolveMintFlag(perOrg, global)).toEqual({ + kind: "runOpsId", + prev: undefined, + flippedAtMs: undefined, + }); + }); + + it("defaults to cuid with no stamp when neither source has a kind", () => { + expect(resolveMintFlag({}, {})).toEqual({ + kind: "cuid", + prev: undefined, + flippedAtMs: undefined, + }); + expect(resolveMintFlag(null, null)).toEqual({ + kind: "cuid", + prev: undefined, + flippedAtMs: undefined, + }); + }); +}); + +// #3b: an org's FIRST per-org runOpsMintKind override must be stamped against the currently +// EFFECTIVE kind — the global FeatureFlag resolution when the org has no override yet — not the +// hardcoded default "cuid". selectMintBaselineSource picks that same source (per-org blob if it +// sets runOpsMintKind, else the global rows) so stampMintKindFlip's baseline (storedKind + +// prev + carry-forward) is correct. +describe("selectMintBaselineSource", () => { + it("returns the per-org blob when it sets runOpsMintKind (override owns the baseline)", () => { + const perOrg = { runOpsMintKind: "cuid" }; + const global = { runOpsMintKind: "runOpsId" }; + expect(selectMintBaselineSource(perOrg, global)).toBe(perOrg); + }); + + it("falls back to the global rows when the org has no runOpsMintKind override", () => { + const perOrg = { someOtherFlag: true }; + const global = { runOpsMintKind: "runOpsId" }; + expect(selectMintBaselineSource(perOrg, global)).toBe(global); + }); + + it("returns an empty record when neither source sets a kind", () => { + expect(selectMintBaselineSource(null, null)).toEqual({}); + expect(selectMintBaselineSource({ someOtherFlag: true }, null)).toEqual({}); + }); +}); + +describe("first per-org override stamps prev against the effective GLOBAL kind (#3b)", () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(() => vi.useRealTimers()); + + it("global=runOpsId, org's FIRST override -> cuid: genuine flip, prev=runOpsId, graced", () => { + const globalFlags = { runOpsMintKind: "runOpsId" }; + const orgExisting = {}; // no per-org override yet + const outgoing = { runOpsMintKind: "cuid" }; + const result = stampMintKindFlip( + selectMintBaselineSource(orgExisting, globalFlags), + outgoing, + T, + GRACE_MS + ); + + expect(result.runOpsMintKind).toBe("cuid"); + // Previously stamped prev="cuid" (the hardcoded default) OR skipped the flip entirely; must + // now be "runOpsId" (the effective global kind) so the org serves it through the grace window. + expect(result.runOpsMintKindPrev).toBe("runOpsId"); + expect(result.runOpsMintKindFlippedAt).toBe(new Date(T).toISOString()); + + const resolution = readMintResolution(result); + expect(effectiveMintKind(resolution, T + GRACE_MS - 1, GRACE_MS)).toBe("runOpsId"); + expect(effectiveMintKind(resolution, T + GRACE_MS, GRACE_MS)).toBe("cuid"); + }); + + it("global=runOpsId, org's FIRST override -> runOpsId (redundant): NOT a spurious flip, no phantom regression to cuid", () => { + const globalFlags = { runOpsMintKind: "runOpsId" }; + const orgExisting = {}; + const outgoing = { runOpsMintKind: "runOpsId" }; + const result = stampMintKindFlip( + selectMintBaselineSource(orgExisting, globalFlags), + outgoing, + T, + GRACE_MS + ); + + expect(result.runOpsMintKind).toBe("runOpsId"); + // Previously: storedKind defaulted to "cuid", so this looked like a genuine flip and stamped + // prev="cuid" — making the org serve cuid during a phantom grace window (a regression). + expect(result.runOpsMintKindPrev).toBeUndefined(); + expect(result.runOpsMintKindFlippedAt).toBeUndefined(); + }); + + it("mid-grace global flip carried into a redundant org override keeps the global stamp", () => { + const globalFlags = { + runOpsMintKind: "runOpsId", + runOpsMintKindPrev: "cuid", + runOpsMintKindFlippedAt: new Date(T).toISOString(), + }; + const orgExisting = {}; + const outgoing = { runOpsMintKind: "runOpsId" }; + const result = stampMintKindFlip( + selectMintBaselineSource(orgExisting, globalFlags), + outgoing, + T + 10_000, + GRACE_MS + ); + + expect(result.runOpsMintKind).toBe("runOpsId"); + expect(result.runOpsMintKindPrev).toBe("cuid"); + expect(result.runOpsMintKindFlippedAt).toBe(new Date(T).toISOString()); + }); +}); diff --git a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts new file mode 100644 index 0000000000..7e2366edcd --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts @@ -0,0 +1,118 @@ +import type { RunIdMintKind } from "./runOpsMintKind.server"; + +export type { RunIdMintKind }; + +export type MintFlagResolution = { + kind: RunIdMintKind; + prev?: RunIdMintKind; + flippedAtMs?: number; +}; + +const DEFAULT_MINT_KIND: RunIdMintKind = "cuid"; + +// Cutover boundary. `nowMs` is the reader's wall clock but `flippedAtMs` (in `r`) is DB-clock +// (admin routes) — so this assumes NTP-synced hosts with skew << graceMs, letting every process +// cross [flippedAtMs, flippedAtMs + graceMs) together (OLD then NEW). Accepted residual: a badly +// mis-synced host can cross early/late and briefly reopen a skew-wide cross-DB duplicate window. +export function effectiveMintKind( + r: MintFlagResolution, + nowMs: number, + graceMs: number +): RunIdMintKind { + if (r.prev === undefined || r.flippedAtMs === undefined) { + return r.kind; + } + return nowMs < r.flippedAtMs + graceMs ? r.prev : r.kind; +} + +function readMintKind(flags: Record, key: string): RunIdMintKind | undefined { + const value = flags[key]; + return value === "cuid" || value === "runOpsId" ? value : undefined; +} + +// Reads the { kind, prev, flippedAtMs } trio out of one flag record — either an org's +// featureFlags override blob or the global FeatureFlag rows projected into a record. Pure. +export function readMintResolution( + flags: Record | null | undefined +): MintFlagResolution { + const source = flags ?? {}; + const kind = readMintKind(source, "runOpsMintKind") ?? DEFAULT_MINT_KIND; + const prev = readMintKind(source, "runOpsMintKindPrev"); + const flippedAtRaw = source.runOpsMintKindFlippedAt; + const parsed = typeof flippedAtRaw === "string" ? Date.parse(flippedAtRaw) : NaN; + const flippedAtMs = Number.isNaN(parsed) ? undefined : parsed; + return { kind, prev, flippedAtMs }; +} + +// SOURCE-CONSISTENT resolution: a per-org runOpsMintKind override wins the kind AND owns the +// grace stamp; with no per-org override, the kind AND the stamp both come from the global rows. +// The stamp is never read from a different source than the kind, which would date a grace +// window against the wrong flip. +export function resolveMintFlag( + perOrgOverrides: Record | null | undefined, + globalFlags: Record | null | undefined +): MintFlagResolution { + if (readMintKind(perOrgOverrides ?? {}, "runOpsMintKind") !== undefined) { + return readMintResolution(perOrgOverrides); + } + return readMintResolution(globalFlags); +} + +// Picks the flag record that currently determines an org's effective mint kind: the per-org +// override blob when it sets runOpsMintKind, otherwise the global FeatureFlag rows. Same source +// resolveMintFlag() reads, but returned as a record so it can seed stampMintKindFlip's baseline +// (storedKind for flip-detection, prev on a genuine flip, and stamp carry-forward). This makes an +// org's first per-org override stamp against the effective GLOBAL kind, not the default "cuid". +export function selectMintBaselineSource( + perOrgOverrides: Record | null | undefined, + globalFlags: Record | null | undefined +): Record { + if (readMintKind(perOrgOverrides ?? {}, "runOpsMintKind") !== undefined) { + return perOrgOverrides ?? {}; + } + return globalFlags ?? {}; +} + +function resolveEffectiveFromFlags( + flags: Record | null | undefined, + nowMs: number, + graceMs: number +): RunIdMintKind { + return effectiveMintKind(readMintResolution(flags), nowMs, graceMs); +} + +// Stamps a grace window only when the outgoing TARGET kind differs from the stored one (a +// genuine flip); prev := the currently-effective kind. A save that leaves the target kind +// unchanged carries any in-flight stamp forward, so it can't reset the cutover clock. +export function stampMintKindFlip( + existingFlags: Record | null | undefined, + outgoingFlags: Record, + nowMs: number, + graceMs: number +): Record { + // Only act when the save actually SETS runOpsMintKind. Omitting it (an unrelated flag change) + // must not inject the default kind, which would pin the org and make a later global flip skip it. + const outgoingKind = readMintKind(outgoingFlags, "runOpsMintKind"); + if (outgoingKind === undefined) { + return outgoingFlags; + } + const storedKind = readMintKind(existingFlags ?? {}, "runOpsMintKind") ?? DEFAULT_MINT_KIND; + + if (outgoingKind !== storedKind) { + // Genuine target change: serve the currently-effective kind through the new grace window. + outgoingFlags.runOpsMintKindPrev = resolveEffectiveFromFlags(existingFlags, nowMs, graceMs); + outgoingFlags.runOpsMintKindFlippedAt = new Date(nowMs).toISOString(); + return outgoingFlags; + } + + const existing = existingFlags ?? {}; + const existingPrev = existing.runOpsMintKindPrev; + const existingFlippedAt = existing.runOpsMintKindFlippedAt; + if (existingPrev !== undefined) { + outgoingFlags.runOpsMintKindPrev = existingPrev; + } + if (existingFlippedAt !== undefined) { + outgoingFlags.runOpsMintKindFlippedAt = existingFlippedAt; + } + return outgoingFlags; +} diff --git a/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.flipLatency.test.ts b/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.flipLatency.test.ts index 0465c8d1c8..11ce521b31 100644 --- a/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.flipLatency.test.ts +++ b/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.flipLatency.test.ts @@ -2,15 +2,15 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { BoundedTtlCache } from "~/services/realtime/boundedTtlCache"; import { computeRunIdMintKind, type RunIdMintKind } from "./runOpsMintKind.server"; -// LOCK of the CURRENT (intentional) flip-latency behavior, NOT a change request. -// resolveRunIdMintKind caches the per-org mint kind in a process-singleton -// BoundedTtlCache (TTL RUN_OPS_MINT_FLAG_CACHE_TTL_MS, 30000ms default) with get/set -// and NO invalidation hook (runOpsMintKind.server.ts:38-45,56-81). So after a flag -// flip a process keeps minting the stale kind until its cached entry expires; in -// multi-instance prod each process expires independently. This suite reconstructs the -// same flag fn over a real cache and pins both edges of that window. +// LOCK of the raw per-process cache's flip-latency behavior in isolation, NOT a change +// request. Production resolveRunIdMintKind now wraps this same staleness in a deterministic +// wall-clock grace window (mintFlipGrace.ts) so every process resolves to the SAME effective +// kind for the whole window, then all cross together. This raw staleness is now an +// intentional, safe input to that resolution. computeRunIdMintKind is unaffected, so this +// suite's assertions stand as-is. -// Mirror of resolveRunIdMintKind's flag fn (runOpsMintKind.server.ts:56-81). +// Bare cached-flag closure — deliberately NOT the production flag fn, which now layers +// grace resolution on top (see runOpsMintKind.server.ts). function makeCachedFlag( cache: BoundedTtlCache, liveFlag: () => RunIdMintKind diff --git a/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts b/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts index 305b851bdc..90b96217e1 100644 --- a/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts @@ -4,7 +4,8 @@ import { logger } from "~/services/logger.server"; import { BoundedTtlCache } from "~/services/realtime/boundedTtlCache"; import { singleton } from "~/utils/singleton"; import { FEATURE_FLAG } from "~/v3/featureFlags"; -import { makeFlag } from "~/v3/featureFlags.server"; +import { DEFAULT_CP_CACHE_TTL_MS } from "./controlPlaneCache.server"; +import { effectiveMintKind, resolveMintFlag, type MintFlagResolution } from "./mintFlipGrace"; import { isSplitEnabled } from "./splitMode.server"; export type RunIdMintKind = "cuid" | "runOpsId"; @@ -34,16 +35,33 @@ export async function computeRunIdMintKind( } // ENV-BOUND wrapper — the only place env/$replica/isSplitEnabled are read. -const flagFn = singleton("runOpsMintFlag", () => makeFlag($replica)); const mintCache = singleton( "runOpsMintCache", () => - new BoundedTtlCache( + new BoundedTtlCache( env.RUN_OPS_MINT_FLAG_CACHE_TTL_MS, env.RUN_OPS_MINT_FLAG_CACHE_MAX_ENTRIES ) ); +// BOOT-TIME SAFETY CHECK (warning only, never throws): the grace window only collapses +// the cross-process divergence window if it outlasts BOTH caches a flag flip has to drain +// through — this process's own mint-flag cache AND the org-flags control-plane cache a +// stale process might still be reading through. If it doesn't, warn loudly but keep booting. +const controlPlaneCacheTtlMs = env.CONTROL_PLANE_CACHE_TTL_MS ?? DEFAULT_CP_CACHE_TTL_MS; +if (env.RUN_OPS_MINT_FLIP_GRACE_MS <= env.RUN_OPS_MINT_FLAG_CACHE_TTL_MS + controlPlaneCacheTtlMs) { + logger.warn( + "[runOpsMintKind] RUN_OPS_MINT_FLIP_GRACE_MS does not exceed the sum of " + + "RUN_OPS_MINT_FLAG_CACHE_TTL_MS and the control-plane cache TTL; a flag flip can still " + + "cross-DB-duplicate a concurrent root trigger during the divergence window", + { + RUN_OPS_MINT_FLIP_GRACE_MS: env.RUN_OPS_MINT_FLIP_GRACE_MS, + RUN_OPS_MINT_FLAG_CACHE_TTL_MS: env.RUN_OPS_MINT_FLAG_CACHE_TTL_MS, + controlPlaneCacheTtlMs, + } + ); +} + export async function resolveRunIdMintKind(environment: { organizationId: string; id: string; @@ -54,10 +72,14 @@ export async function resolveRunIdMintKind(environment: { masterEnabled: env.RUN_OPS_MINT_ENABLED, splitEnabled: isSplitEnabled, flag: async (orgId, orgFeatureFlags) => { - // The cache stores only "cuid"|"runOpsId" (never undefined), so the cache's - // "stored-undefined == miss" caveat never applies here. + // The cache stores the full { kind, prev, flippedAtMs } trio (never undefined), so the + // cache's "stored-undefined == miss" caveat never applies here. A cache HIT still passes + // back through effectiveMintKind so a cached-but-stale entry crosses the grace boundary + // on schedule, without needing an invalidation hook. const cached = mintCache.get(orgId); - if (cached !== undefined) return cached; + if (cached !== undefined) { + return effectiveMintKind(cached, Date.now(), env.RUN_OPS_MINT_FLIP_GRACE_MS); + } // Hot-path pass-through: use the org flags the authenticated environment already // carries; only fall back to a DB read when the caller did NOT pass them (non-trigger @@ -72,13 +94,34 @@ export async function resolveRunIdMintKind(environment: { }) )?.featureFlags; - const kind = await flagFn({ - key: FEATURE_FLAG.runOpsMintKind, - defaultValue: "cuid", - overrides: (overrides as Record) ?? {}, + const overridesRecord = (overrides as Record) ?? {}; + + // One global read over the three mint-flag keys (kind + grace stamp), folded into the + // single cache-miss round-trip. This replaces the former single-key flag read, so a + // GLOBAL flip is now grace-stamped WITHOUT adding any new per-mint/per-resolve query. + // (The cache-hit branch above never touches the DB.) + const globalRows = await $replica.featureFlag.findMany({ + where: { + key: { + in: [ + FEATURE_FLAG.runOpsMintKind, + FEATURE_FLAG.runOpsMintKindPrev, + FEATURE_FLAG.runOpsMintKindFlippedAt, + ], + }, + }, + select: { key: true, value: true }, }); - mintCache.set(orgId, kind); - return kind; + const globalFlags: Record = {}; + for (const row of globalRows) { + globalFlags[row.key] = row.value; + } + + // Source-consistent: a per-org override wins the kind AND its stamp; otherwise the + // global row wins the kind AND its stamp. + const resolution: MintFlagResolution = resolveMintFlag(overridesRecord, globalFlags); + mintCache.set(orgId, resolution); + return effectiveMintKind(resolution, Date.now(), env.RUN_OPS_MINT_FLIP_GRACE_MS); }, }); } diff --git a/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts b/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts new file mode 100644 index 0000000000..eb322c48a1 --- /dev/null +++ b/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts @@ -0,0 +1,337 @@ +// Flow-level seam proof for ApiBatchResultsPresenter, ABOVE the run-store-level coverage in +// internal-packages/run-store (batchItemMisroute, batchCompletionResidency). Proves the batch +// RESULTS READ assembles correctly when one batch's members are genuinely split across the real +// dedicated run-ops subset schema (prisma17 / RunOpsPrismaClient) and the full control-plane +// schema (prisma14) — not a mirrored full schema on both sides. No mocks. +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import type { PrismaClient } from "@trigger.dev/database"; +import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; +import { describe, expect, vi } from "vitest"; +import type { PrismaReplicaClient } from "~/db.server"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { ApiBatchResultsPresenter } from "~/presenters/v3/ApiBatchResultsPresenter.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +// A prisma handle that throws on any access — proves the passthrough constructor args are never +// touched on the split path. +const throwingPrisma = new Proxy( + {}, + { + get(_t, prop) { + throw new Error( + `passthrough handle must not be touched on the split path (got .${String(prop)})` + ); + }, + } +) as unknown as PrismaReplicaClient; + +// 25-char cuid-shaped id, no v1 version marker at index 25 -> classifies LEGACY. +function generateLegacyCuid(): string { + const alphabet = "0123456789abcdefghijklmnopqrstuvwxyz"; + const suffix = Array.from({ length: 24 }, () => alphabet[Math.floor(Math.random() * 36)]).join( + "" + ); + return `c${suffix}`; +} + +let seedCounter = 0; + +// TaskRun on the full control-plane schema has real FKs into RuntimeEnvironment/Project. +async function seedLegacyEnv(prisma14: PrismaClient, slug: string) { + const n = seedCounter++; + const organization = await prisma14.organization.create({ + data: { title: `Org ${slug}`, slug: `org-${slug}-${n}` }, + }); + const project = await prisma14.project.create({ + data: { + name: `Proj ${slug}`, + slug: `proj-${slug}-${n}`, + organizationId: organization.id, + externalRef: `ext-${slug}-${n}`, + }, + }); + const environment = await prisma14.runtimeEnvironment.create({ + data: { + slug: `env-${slug}-${n}`, + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: `api-${slug}-${n}`, + pkApiKey: `pk-${slug}-${n}`, + shortcode: `sc-${slug}-${n}`, + }, + }); + return { organization, project, environment }; +} + +type SeedCtx = Awaited>; + +type MemberSeed = { + id: string; + friendlyId: string; + status: "COMPLETED_SUCCESSFULLY" | "COMPLETED_WITH_ERRORS"; + output?: string; + error?: unknown; +}; + +// Drop the TaskRunAttempt worker/queue FKs so attempts can be seeded without standing up +// BackgroundWorker/TaskQueue parents — incidental to this read path. +async function relaxLegacyAttemptFk(prisma14: PrismaClient) { + for (const sql of [ + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_backgroundWorkerId_fkey"`, + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_backgroundWorkerTaskId_fkey"`, + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_queueId_fkey"`, + ]) { + await prisma14.$executeRawUnsafe(sql); + } +} + +async function seedLegacyMember(prisma14: PrismaClient, ctx: SeedCtx, m: MemberSeed) { + const run = await prisma14.taskRun.create({ + data: { + id: m.id, + friendlyId: m.friendlyId, + taskIdentifier: "my-task", + status: m.status, + payload: JSON.stringify({}), + payloadType: "application/json", + traceId: m.id, + spanId: m.id, + queue: "main", + runtimeEnvironmentId: ctx.environment.id, + projectId: ctx.project.id, + organizationId: ctx.organization.id, + environmentType: "PRODUCTION", + engine: "V2", + }, + }); + + await prisma14.taskRunAttempt.create({ + data: { + friendlyId: `attempt_${m.id}`, + number: 1, + taskRunId: run.id, + backgroundWorkerId: "bw", + backgroundWorkerTaskId: "bwt", + runtimeEnvironmentId: ctx.environment.id, + queueId: "q", + status: m.status === "COMPLETED_SUCCESSFULLY" ? "COMPLETED" : "FAILED", + output: m.output, + outputType: "application/json", + error: m.error as any, + }, + }); + + return run; +} + +// TaskRun/TaskRunAttempt on the DEDICATED (run-ops subset) schema: runtimeEnvironmentId, +// organizationId, projectId, backgroundWorkerId/backgroundWorkerTaskId/queueId are all +// scalar-only there (no real relation) — no parent rows to seed, unlike the legacy side. +async function seedNewMember( + prisma17: RunOpsPrismaClient, + ctx: { envId: string; orgId: string; projectId: string }, + m: MemberSeed +) { + await prisma17.taskRun.create({ + data: { + id: m.id, + engine: "V2", + status: m.status, + friendlyId: m.friendlyId, + runtimeEnvironmentId: ctx.envId, + environmentType: "PRODUCTION", + organizationId: ctx.orgId, + projectId: ctx.projectId, + taskIdentifier: "my-task", + payload: JSON.stringify({}), + payloadType: "application/json", + traceContext: {}, + traceId: m.id, + spanId: m.id, + queue: "main", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }, + }); + + await prisma17.taskRunAttempt.create({ + data: { + friendlyId: `attempt_${m.id}`, + number: 1, + taskRunId: m.id, + backgroundWorkerId: "bw", + backgroundWorkerTaskId: "bwt", + runtimeEnvironmentId: ctx.envId, + queueId: "q", + status: m.status === "COMPLETED_SUCCESSFULLY" ? "COMPLETED" : "FAILED", + output: m.output, + outputType: "application/json", + error: m.error as any, + }, + }); +} + +// BatchTaskRunItem.taskRunId is a real FK to TaskRun on the dedicated schema too. A batch seeded +// on the dedicated DB whose items reference a LEGACY-resident (or wholly missing) member has no +// local row for that member — drop the FK so the cross-seam item row can exist, matching the +// physical reality of a split batch. +async function relaxNewBatchItemFk(prisma17: RunOpsPrismaClient) { + await prisma17.$executeRawUnsafe( + `ALTER TABLE "BatchTaskRunItem" DROP CONSTRAINT IF EXISTS "BatchTaskRunItem_taskRunId_fkey"` + ); +} + +async function seedBatchOnNew( + prisma17: RunOpsPrismaClient, + envId: string, + friendlyId: string, + memberIds: string[] +) { + const batch = await prisma17.batchTaskRun.create({ + data: { + friendlyId, + runtimeEnvironmentId: envId, + runCount: memberIds.length, + runIds: [], + batchVersion: "runengine:v2", + }, + }); + // Items in a deterministic order so the result `items` order is assertable. + for (const taskRunId of memberIds) { + await prisma17.batchTaskRunItem.create({ + data: { batchTaskRunId: batch.id, taskRunId, status: "COMPLETED" }, + }); + } + return batch; +} + +const env = (ctx: SeedCtx) => + ({ + id: ctx.environment.id, + type: ctx.environment.type, + slug: ctx.environment.slug, + organizationId: ctx.organization.id, + organization: { slug: ctx.organization.slug, title: ctx.organization.title }, + projectId: ctx.project.id, + project: { name: ctx.project.name }, + }) as unknown as AuthenticatedEnvironment; + +describe("ApiBatchResultsPresenter split mode — real run-ops dedicated schema seam", () => { + // The core assembly proof: one batch, members genuinely split across the two physically + // distinct, differently-shaped DBs (dedicated subset vs full control-plane). + heteroRunOpsPostgresTest( + "a batch with members split across the dedicated NEW schema and the legacy schema returns the complete union", + async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => { + const ctx = await seedLegacyEnv(prisma14, "split-new-legacy"); + await relaxLegacyAttemptFk(prisma14); + await relaxNewBatchItemFk(prisma17); + + const newMemberId = generateRunOpsId(); + expect(newMemberId.length).toBe(26); + const legacyMemberId = generateLegacyCuid(); + expect(legacyMemberId.length).toBe(25); + + // NEW-resident member: lives ONLY on the dedicated run-ops schema (prisma17). + await seedNewMember( + prisma17, + { envId: ctx.environment.id, orgId: ctx.organization.id, projectId: ctx.project.id }, + { + id: newMemberId, + friendlyId: "run_new_member", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ from: "new" }), + } + ); + + // LEGACY-resident member: lives ONLY on the full control-plane schema (prisma14). + await seedLegacyMember(prisma14, ctx, { + id: legacyMemberId, + friendlyId: "run_legacy_member", + status: "COMPLETED_WITH_ERRORS", + error: { type: "BUILT_IN_ERROR", name: "Err", message: "boom", stackTrace: "" }, + }); + + // The batch row + items live on the NEW dedicated DB; items reference both members. + const batchFriendlyId = "batch_split_seam"; + await seedBatchOnNew(prisma17, ctx.environment.id, batchFriendlyId, [ + newMemberId, + legacyMemberId, + ]); + + const presenter = new ApiBatchResultsPresenter(throwingPrisma, throwingPrisma, { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + }); + + const result = await presenter.call(batchFriendlyId, env(ctx)); + + expect(result).toBeDefined(); + expect(result!.id).toBe(batchFriendlyId); + expect(result!.items).toHaveLength(2); + + // Order follows item order: the NEW-resident member first, the legacy-resident second. + const [first, second] = result!.items; + expect(first).toEqual({ + ok: true, + id: "run_new_member", + taskIdentifier: "my-task", + output: JSON.stringify({ from: "new" }), + outputType: "application/json", + }); + expect(second).toMatchObject({ + ok: false, + id: "run_legacy_member", + taskIdentifier: "my-task", + }); + } + ); + + // Dangling member: an item id that resolves on NEITHER the dedicated NEW schema nor the legacy + // schema must be dropped from the result, not thrown — the presenter degrades to the reachable + // set. Give the missing id a legacy shape so it also exercises the "not NEW-shaped -> legacy + // candidate" branch, proving the legacy probe finding nothing doesn't throw either. + heteroRunOpsPostgresTest( + "a member id that resolves on neither the dedicated NEW schema nor the legacy schema is dropped, not thrown", + async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => { + const ctx = await seedLegacyEnv(prisma14, "dangling"); + await relaxNewBatchItemFk(prisma17); + + const presentId = generateRunOpsId(); + const missingId = generateLegacyCuid(); // referenced by an item but seeded on NEITHER DB + + await seedNewMember( + prisma17, + { envId: ctx.environment.id, orgId: ctx.organization.id, projectId: ctx.project.id }, + { + id: presentId, + friendlyId: "run_present", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ present: true }), + } + ); + + const batchFriendlyId = "batch_dangling_seam"; + await seedBatchOnNew(prisma17, ctx.environment.id, batchFriendlyId, [presentId, missingId]); + + const presenter = new ApiBatchResultsPresenter(throwingPrisma, throwingPrisma, { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + }); + + const result = await presenter.call(batchFriendlyId, env(ctx)); + + expect(result).toBeDefined(); + expect(result!.id).toBe(batchFriendlyId); + // The dangling member is silently dropped; the reachable member still returns — the call + // must not throw despite the item referencing a run absent from both physical DBs. + expect(result!.items).toHaveLength(1); + expect(result!.items[0]).toMatchObject({ ok: true, id: "run_present" }); + } + ); +}); diff --git a/apps/webapp/test/apiBatchResultsPresenter.splitNPlus1.test.ts b/apps/webapp/test/apiBatchResultsPresenter.splitNPlus1.test.ts new file mode 100644 index 0000000000..614e112992 --- /dev/null +++ b/apps/webapp/test/apiBatchResultsPresenter.splitNPlus1.test.ts @@ -0,0 +1,251 @@ +// RED→GREEN: kill the #callSplit per-member N+1 in ApiBatchResultsPresenter. +// +// Today, #callSplit hydrates every batch member independently via `readThroughRun` inside +// `Promise.all(batchRun.items.map(...))` — that is one (up to two, new-then-legacy) `taskRun` +// query PER member. `#callPassthrough` already does the grouped one-query form via +// `this.runStore.findRuns({ where: { id: { in: taskRunIds } } })`. +// +// The fix replaces the per-member fan-out with ONE grouped call to the RunStore's +// `findRunsByIds` method, mirroring `#callPassthrough`. This test proves the query-count +// reduction with a call-counting Proxy over a REAL testcontainer Postgres client: every call is +// delegated unchanged to the real client (the DB still runs the query) — this is instrumentation, +// not a mock. +import { PostgresRunStore } from "@internal/run-store"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import type { PrismaReplicaClient } from "~/db.server"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { ApiBatchResultsPresenter } from "~/presenters/v3/ApiBatchResultsPresenter.server"; + +// 26-char run-ops v1 body: 24-char base32hex core + region char + version "1" at index 25. +// ownerEngine classifies residency by the VERSION CHAR AT INDEX 25, not by length — a naive id +// generator would misclassify this as LEGACY unless the last two chars are a valid region+version. +function newRunId(c: string) { + return c.repeat(24) + "01"; +} + +type CallCounts = { findMany: number; findFirst: number }; + +// Wrap a REAL client's `taskRun` delegate to tally findMany/findFirst calls, delegating every +// call unchanged to the real client (the DB still runs the query — pure instrumentation). +function countingClient(real: PrismaClient): { client: PrismaClient; counts: CallCounts } { + const counts: CallCounts = { findMany: 0, findFirst: 0 }; + const countingTaskRun = new Proxy((real as any).taskRun, { + get(target, prop) { + if (prop === "findMany" || prop === "findFirst") { + counts[prop as "findMany" | "findFirst"]++; + } + return (target as any)[prop]; + }, + }); + const client = new Proxy(real, { + get(target, prop) { + if (prop === "taskRun") { + return countingTaskRun; + } + return (target as any)[prop]; + }, + }) as PrismaClient; + return { client, counts }; +} + +let seedCounter = 0; + +async function seedEnv(prisma: PrismaClient, slug: string) { + const n = seedCounter++; + const organization = await prisma.organization.create({ + data: { title: `Org ${slug}`, slug: `org-${slug}-${n}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Proj ${slug}`, + slug: `proj-${slug}-${n}`, + organizationId: organization.id, + externalRef: `ext-${slug}-${n}`, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + slug: `env-${slug}-${n}`, + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: `api-${slug}-${n}`, + pkApiKey: `pk-${slug}-${n}`, + shortcode: `sc-${slug}-${n}`, + }, + }); + return { organization, project, environment }; +} + +type SeedCtx = Awaited>; + +type MemberSeed = { + id: string; + friendlyId: string; + status: "COMPLETED_SUCCESSFULLY" | "COMPLETED_WITH_ERRORS"; + output?: string; + error?: unknown; +}; + +async function seedMember(prisma: PrismaClient, ctx: SeedCtx, m: MemberSeed) { + const run = await prisma.taskRun.create({ + data: { + id: m.id, + friendlyId: m.friendlyId, + taskIdentifier: "my-task", + status: m.status, + payload: JSON.stringify({}), + payloadType: "application/json", + traceId: m.id, + spanId: m.id, + queue: "main", + runtimeEnvironmentId: ctx.environment.id, + projectId: ctx.project.id, + organizationId: ctx.organization.id, + environmentType: "PRODUCTION", + engine: "V2", + }, + }); + + await prisma.taskRunAttempt.create({ + data: { + friendlyId: `attempt_${m.id}`, + number: 1, + taskRunId: run.id, + backgroundWorkerId: "bw", + backgroundWorkerTaskId: "bwt", + runtimeEnvironmentId: ctx.environment.id, + queueId: "q", + status: m.status === "COMPLETED_SUCCESSFULLY" ? "COMPLETED" : "FAILED", + output: m.output, + outputType: "application/json", + error: m.error as any, + }, + }); + + return run; +} + +// Drop the TaskRunAttempt worker/queue FKs so attempts can be seeded (their output/error is what's +// under test) without standing up BackgroundWorker/TaskQueue parents — incidental to this read path. +async function relaxFks(prisma: PrismaClient) { + for (const sql of [ + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_backgroundWorkerId_fkey"`, + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_backgroundWorkerTaskId_fkey"`, + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_queueId_fkey"`, + ]) { + await prisma.$executeRawUnsafe(sql); + } +} + +async function seedBatch( + prisma: PrismaClient, + ctx: SeedCtx, + friendlyId: string, + memberIds: string[] +) { + const batch = await prisma.batchTaskRun.create({ + data: { + friendlyId, + runtimeEnvironmentId: ctx.environment.id, + runCount: memberIds.length, + runIds: [], + batchVersion: "runengine:v2", + }, + }); + // Items in a deterministic order so the result `items` order is assertable. + for (const taskRunId of memberIds) { + await prisma.batchTaskRunItem.create({ + data: { batchTaskRunId: batch.id, taskRunId, status: "COMPLETED" }, + }); + } + return batch; +} + +const env = (ctx: SeedCtx) => + ({ + id: ctx.environment.id, + type: ctx.environment.type, + slug: ctx.environment.slug, + organizationId: ctx.organization.id, + organization: { slug: ctx.organization.slug, title: ctx.organization.title }, + projectId: ctx.project.id, + project: { name: ctx.project.name }, + }) as unknown as AuthenticatedEnvironment; + +describe("ApiBatchResultsPresenter split mode — member hydration is grouped, not per-member", () => { + postgresTest( + "a batch of N members costs ONE findMany (never N findFirst) and returns the same result", + async ({ prisma }) => { + const ctx = await seedEnv(prisma, "n1"); + await relaxFks(prisma); + + const { client: countingPrisma, counts } = countingClient(prisma); + + const memberIds = [newRunId("a"), newRunId("b"), newRunId("c")]; + await seedMember(prisma, ctx, { + id: memberIds[0], + friendlyId: "run_a", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ from: "a" }), + }); + await seedMember(prisma, ctx, { + id: memberIds[1], + friendlyId: "run_b", + status: "COMPLETED_WITH_ERRORS", + error: { type: "BUILT_IN_ERROR", name: "Err", message: "boom", stackTrace: "" }, + }); + await seedMember(prisma, ctx, { + id: memberIds[2], + friendlyId: "run_c", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ from: "c" }), + }); + + await seedBatch(prisma, ctx, "batch_n1", memberIds); + + const runStore = new PostgresRunStore({ + prisma: countingPrisma, + readOnlyPrisma: countingPrisma, + }); + + const presenter = new ApiBatchResultsPresenter( + countingPrisma, + countingPrisma, + { + splitEnabled: true, + newClient: countingPrisma as unknown as PrismaReplicaClient, + legacyReplica: countingPrisma as unknown as PrismaReplicaClient, + }, + runStore + ); + + const result = await presenter.call("batch_n1", env(ctx)); + + expect(result).toBeDefined(); + expect(result!.id).toBe("batch_n1"); + expect(result!.items).toHaveLength(3); + expect(result!.items[0]).toEqual({ + ok: true, + id: "run_a", + taskIdentifier: "my-task", + output: JSON.stringify({ from: "a" }), + outputType: "application/json", + }); + expect(result!.items[1]).toMatchObject({ ok: false, id: "run_b" }); + expect(result!.items[2]).toEqual({ + ok: true, + id: "run_c", + taskIdentifier: "my-task", + output: JSON.stringify({ from: "c" }), + outputType: "application/json", + }); + + // The grouped-read proof: ONE findMany for the whole member set, never a findFirst per member. + expect(counts.findMany).toBe(1); + expect(counts.findFirst).toBe(0); + } + ); +}); diff --git a/apps/webapp/test/apiRetrieveRunPresenter.groupedLockedWorker.test.ts b/apps/webapp/test/apiRetrieveRunPresenter.groupedLockedWorker.test.ts new file mode 100644 index 0000000000..1b4a200751 --- /dev/null +++ b/apps/webapp/test/apiRetrieveRunPresenter.groupedLockedWorker.test.ts @@ -0,0 +1,266 @@ +import { containerTest } from "@internal/testcontainers"; +import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import { beforeEach, describe, expect, vi } from "vitest"; + +// `findRun` reads the module-level `prisma`/`$replica` (control-plane handles) both directly +// and, transitively, via the `runStore`/`controlPlaneResolver` singletons that also import them. +// A lazy Proxy - not a plain value - is captured by those singletons at their own (one-time) +// construction, so pointing it at a real container client after the fact still routes every +// call there. Not a DB mock: every call forwards to a real Postgres container. +const dbHolder = vi.hoisted(() => ({ client: undefined as PrismaClient | undefined })); + +vi.mock("~/db.server", () => { + const proxy = new Proxy( + {}, + { + get(_t, prop) { + if (!dbHolder.client) throw new Error("dbHolder.client not set for this test"); + return (dbHolder.client as any)[prop]; + }, + } + ); + return { prisma: proxy, $replica: proxy }; +}); + +vi.mock("~/v3/objectStore.server", () => ({ + generatePresignedUrl: vi.fn(async () => ({ success: false, error: "not-used" })), +})); + +import { ApiRetrieveRunPresenter } from "~/presenters/v3/ApiRetrieveRunPresenter.server"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +// Wraps a real Prisma client so every `client..(...)` call is tallied by +// `"."` before being forwarded, unmodified, to the real delegate. +function createCallCountingProxy(client: PrismaClient): { + client: PrismaClient; + counts: Map; +} { + const counts = new Map(); + const wrappedModels = new Map(); + + const wrapModel = (model: string, delegate: object) => + new Proxy(delegate, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if (typeof prop === "string" && typeof value === "function") { + return (...args: unknown[]) => { + counts.set(`${model}.${prop}`, (counts.get(`${model}.${prop}`) ?? 0) + 1); + return (value as (...a: unknown[]) => unknown).apply(target, args); + }; + } + return value; + }, + }); + + const proxy = new Proxy(client, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if ( + typeof prop === "string" && + value && + typeof value === "object" && + typeof (value as { findMany?: unknown }).findMany === "function" + ) { + if (!wrappedModels.has(prop)) { + wrappedModels.set(prop, wrapModel(prop, value as object)); + } + return wrappedModels.get(prop); + } + return value; + }, + }); + + return { client: proxy as PrismaClient, counts }; +} + +async function seedOrgProjectEnv(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `test-${suffix}`, slug: `test-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `test-${suffix}`, + slug: `test-${suffix}`, + organizationId: organization.id, + externalRef: `test-${suffix}`, + }, + }); + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: `test-${suffix}`, + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short-${suffix}`, + }, + }); + return { organization, project, runtimeEnvironment }; +} + +function authEnv( + organization: { id: string }, + project: { id: string; externalRef: string }, + runtimeEnvironment: { id: string; slug: string } +): AuthenticatedEnvironment { + return { + id: runtimeEnvironment.id, + slug: runtimeEnvironment.slug, + organizationId: organization.id, + organization: { id: organization.id }, + project: { id: project.id, externalRef: project.externalRef }, + } as unknown as AuthenticatedEnvironment; +} + +async function seedBackgroundWorker( + prisma: PrismaClient, + ctx: { projectId: string; runtimeEnvironmentId: string }, + version: string +) { + return prisma.backgroundWorker.create({ + data: { + friendlyId: `worker_${generateRunOpsId()}`, + version, + contentHash: `hash_${generateRunOpsId()}`, + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.runtimeEnvironmentId, + metadata: {}, + }, + }); +} + +interface SeedRunOpts { + id: string; + friendlyId: string; + runtimeEnvironmentId: string; + projectId: string; + organizationId: string; + lockedToVersionId?: string; + parentTaskRunId?: string; + rootTaskRunId?: string; +} + +async function seedRun(prisma: PrismaClient, opts: SeedRunOpts) { + return prisma.taskRun.create({ + data: { + id: opts.id, + friendlyId: opts.friendlyId, + taskIdentifier: "my-task", + payload: JSON.stringify({ hello: "world" }), + payloadType: "application/json", + traceId: `trace_${opts.id}`, + spanId: `span_${opts.id}`, + queue: "task/my-task", + runtimeEnvironmentId: opts.runtimeEnvironmentId, + projectId: opts.projectId, + organizationId: opts.organizationId, + environmentType: "DEVELOPMENT", + engine: "V2", + lockedToVersionId: opts.lockedToVersionId, + parentTaskRunId: opts.parentTaskRunId, + rootTaskRunId: opts.rootTaskRunId, + }, + }); +} + +beforeEach(() => { + dbHolder.client = undefined; +}); + +describe("ApiRetrieveRunPresenter.findRun locked-worker version resolution", () => { + containerTest( + "resolves run+parent+root+children lockedToVersion with ONE grouped query, not one per id", + async ({ prisma }) => { + const proxied = createCallCountingProxy(prisma); + dbHolder.client = proxied.client; + + const { organization, project, runtimeEnvironment } = await seedOrgProjectEnv( + prisma, + "grouped" + ); + const workerCtx = { projectId: project.id, runtimeEnvironmentId: runtimeEnvironment.id }; + + // Two distinct versions - `workerA` is deliberately reused across run/root/one child to + // prove the pre-existing dedup-by-Set still collapses to the same distinct-id count. + const workerA = await seedBackgroundWorker(prisma, workerCtx, "2024.1.0"); + const workerB = await seedBackgroundWorker(prisma, workerCtx, "2024.2.0"); + const workerC = await seedBackgroundWorker(prisma, workerCtx, "2024.3.0"); + + const rootId = generateRunOpsId(); + const parentId = generateRunOpsId(); + const runId = generateRunOpsId(); + const childId1 = generateRunOpsId(); + const childId2 = generateRunOpsId(); + + await seedRun(prisma, { + id: rootId, + friendlyId: `run_${rootId}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + lockedToVersionId: workerA.id, + }); + await seedRun(prisma, { + id: parentId, + friendlyId: `run_${parentId}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + rootTaskRunId: rootId, + lockedToVersionId: workerB.id, + }); + await seedRun(prisma, { + id: runId, + friendlyId: `run_${runId}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + parentTaskRunId: parentId, + rootTaskRunId: rootId, + lockedToVersionId: workerA.id, + }); + await seedRun(prisma, { + id: childId1, + friendlyId: `run_${childId1}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + parentTaskRunId: runId, + rootTaskRunId: rootId, + lockedToVersionId: workerC.id, + }); + await seedRun(prisma, { + id: childId2, + friendlyId: `run_${childId2}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + parentTaskRunId: runId, + rootTaskRunId: rootId, + lockedToVersionId: workerA.id, + }); + + const env = authEnv(organization, project, runtimeEnvironment); + const found = await ApiRetrieveRunPresenter.findRun(`run_${runId}`, env); + + expect(found).not.toBeNull(); + expect(found!.lockedToVersion?.version).toBe(workerA.version); + expect(found!.parentTaskRun?.lockedToVersion?.version).toBe(workerB.version); + expect(found!.rootTaskRun?.lockedToVersion?.version).toBe(workerA.version); + const versionByChildFriendlyId = new Map( + found!.childRuns.map((c) => [c.friendlyId, c.lockedToVersion?.version]) + ); + expect(versionByChildFriendlyId.get(`run_${childId1}`)).toBe(workerC.version); + expect(versionByChildFriendlyId.get(`run_${childId2}`)).toBe(workerA.version); + + // 3 distinct lockedToVersionIds (workerA, workerB, workerC) across 5 rows -> ONE grouped + // findMany, ZERO per-id findFirst. + expect(proxied.counts.get("backgroundWorker.findMany") ?? 0).toBe(1); + expect(proxied.counts.get("backgroundWorker.findFirst") ?? 0).toBe(0); + } + ); +}); diff --git a/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts new file mode 100644 index 0000000000..a8263b80e3 --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts @@ -0,0 +1,236 @@ +// RED->GREEN guard: WaitpointPresenter connected-run gather must BOUND the fetch of a waitpoint's +// connected-run ids, not just the number displayed. A "displayed count <= 5" assertion would +// false-green (the take:5 on the run resolve already bounds the DISPLAY). Instead this captures the +// real `taskRun.findMany` call args via a Proxy over a REAL testcontainer Postgres client (no mocks) +// and asserts the IN-list is capped at the scan limit (danglers over-read), never unbounded. +// +// Exercises the dedicated-schema (Prisma `waitpointRunConnection`) branch: waitpoint + more than the +// scan-limit connected runs seeded on the NEW dedicated run-ops client (RunOpsPrismaClient, prisma17). +import { describe, expect, vi } from "vitest"; + +const legacyReplicaHolder = vi.hoisted(() => ({ client: undefined as any })); +const newClientHolder = vi.hoisted(() => ({ client: undefined as any })); + +vi.mock("~/db.server", async () => { + const { Prisma } = await import("@trigger.dev/database"); + const lazyProxy = (holder: { client: any }, label: string) => + new Proxy( + {}, + { + get(_t, prop) { + if (!holder.client) { + throw new Error(`${label} not set for this test`); + } + return holder.client[prop]; + }, + } + ); + const replicaProxy = lazyProxy(legacyReplicaHolder, "legacyReplicaHolder.client"); + return { + prisma: replicaProxy, + $replica: replicaProxy, + runOpsNewPrisma: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsNewReplica: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsLegacyPrisma: replicaProxy, + runOpsLegacyReplica: replicaProxy, + sqlDatabaseSchema: Prisma.sql([`public`]), + DATABASE_SCHEMA: "public", + }; +}); + +vi.mock("~/services/clickhouse/clickhouseFactoryInstance.server", () => ({ + clickhouseFactory: { + getClickhouseForOrganization: async () => ({}), + }, +})); + +// Echo the runId set back as runs so the presenter's final CH hydrate never runs for real -- the +// thing under test is the connected-run-id GATHER (the taskRun.findMany args), not this step. +vi.mock("~/presenters/v3/NextRunListPresenter.server", () => ({ + NextRunListPresenter: class { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + constructor(..._args: unknown[]) {} + async call(_organizationId: string, _environmentId: string, opts: { runId?: string[] }) { + return { + runs: (opts.runId ?? []).map((friendlyId) => ({ + friendlyId, + taskIdentifier: "echoed", + })), + }; + } + }, +})); + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { + CONNECTED_RUNS_CONNECTION_SCAN_LIMIT, + WaitpointPresenter, +} from "~/presenters/v3/WaitpointPresenter.server"; + +vi.setConfig({ testTimeout: 90_000 }); + +type SeedContext = { + organizationId: string; + projectId: string; + environmentId: string; +}; + +async function seedParents(prisma: PrismaClient, slug: string): Promise { + const organization = await prisma.organization.create({ + data: { title: `org-${slug}`, slug: `org-${slug}` }, + }); + const project = await prisma.project.create({ + data: { + name: `proj-${slug}`, + slug: `proj-${slug}`, + organizationId: organization.id, + externalRef: `proj-${slug}`, + }, + }); + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: `env-${slug}`, + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${slug}`, + pkApiKey: `pk_dev_${slug}`, + shortcode: `sc-${slug}`, + }, + }); + return { + organizationId: organization.id, + projectId: project.id, + environmentId: runtimeEnvironment.id, + }; +} + +async function seedWaitpoint( + prisma: PrismaClient | RunOpsPrismaClient, + ctx: SeedContext, + friendlyId: string +) { + return (prisma as PrismaClient).waitpoint.create({ + data: { + friendlyId, + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: `idem-${friendlyId}`, + userProvidedIdempotencyKey: false, + outputType: "application/json", + outputIsError: false, + completedAt: new Date(), + tags: [], + projectId: ctx.projectId, + environmentId: ctx.environmentId, + }, + }); +} + +async function seedRun( + prisma: PrismaClient | RunOpsPrismaClient, + ctx: SeedContext, + friendlyId: string +) { + return (prisma as PrismaClient).taskRun.create({ + data: { + friendlyId, + taskIdentifier: "my-task", + status: "PENDING", + payload: JSON.stringify({ foo: friendlyId }), + payloadType: "application/json", + traceId: friendlyId, + spanId: friendlyId, + queue: "test", + runtimeEnvironmentId: ctx.environmentId, + projectId: ctx.projectId, + organizationId: ctx.organizationId, + environmentType: "DEVELOPMENT", + engine: "V2", + }, + }); +} + +// Wrap the REAL dedicated run-ops client's `taskRun.findMany` to capture the args of every call, +// delegating unchanged to the real client (the DB still runs the query -- pure instrumentation, +// never a mock). +function capturingTaskRunFindMany(real: RunOpsPrismaClient): { + client: RunOpsPrismaClient; + calls: { where?: { id?: { in?: string[] } } }[]; +} { + const calls: { where?: { id?: { in?: string[] } } }[] = []; + const wrappedTaskRun = new Proxy((real as any).taskRun, { + get(target, prop) { + if (prop === "findMany") { + return (...args: any[]) => { + calls.push(args[0]); + return (target as any)[prop](...args); + }; + } + return (target as any)[prop]; + }, + }); + const client = new Proxy(real as object, { + get(target, prop) { + if (prop === "taskRun") { + return wrappedTaskRun; + } + return (target as any)[prop]; + }, + }) as RunOpsPrismaClient; + return { client, calls }; +} + +// Seed MORE than the scan limit so the assertion bites: an unbounded gather IN-lists all of them, +// a correctly bounded one caps at CONNECTED_RUNS_CONNECTION_SCAN_LIMIT. +const CONNECTED_RUN_COUNT = CONNECTED_RUNS_CONNECTION_SCAN_LIMIT + 5; + +describe("WaitpointPresenter bounds the connected-run-id FETCH", () => { + heteroRunOpsPostgresTest( + "a waitpoint with more connections than the scan limit caps the IN-list at the scan limit", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "bounded"); + const waitpoint = await seedWaitpoint(prisma17, ctx, "waitpoint_bounded"); + + const runs = await Promise.all( + Array.from({ length: CONNECTED_RUN_COUNT }, (_, i) => seedRun(prisma17, ctx, `run_b${i}`)) + ); + for (const run of runs) { + await prisma17.waitpointRunConnection.create({ + data: { taskRunId: run.id, waitpointId: waitpoint.id }, + }); + } + + legacyReplicaHolder.client = prisma14; + newClientHolder.client = prisma17; + + const { client: countingPrisma17, calls } = capturingTaskRunFindMany(prisma17); + + const presenter = new WaitpointPresenter(undefined, undefined, { + splitEnabled: true, + newClient: countingPrisma17 as unknown as PrismaClient, + legacyReplica: prisma14, + }); + + await presenter.call({ + friendlyId: waitpoint.friendlyId, + environmentId: ctx.environmentId, + projectId: ctx.projectId, + }); + + // The guard: assert the FETCH (the IN-list built from the connected-run-id gather) is + // bounded, not just the eventual displayed count. An unbounded gather would IN-list every + // connection row; the dedicated branch over-reads up to CONNECTED_RUNS_CONNECTION_SCAN_LIMIT + // (25) so a display slot is never lost to a dangler, so the IN-list must be capped at the + // scan limit -- above the display limit (5), but never unbounded. + expect(calls.length).toBeGreaterThan(0); + for (const call of calls) { + expect(call.where?.id?.in?.length ?? 0).toBeLessThanOrEqual( + CONNECTED_RUNS_CONNECTION_SCAN_LIMIT + ); + } + } + ); +}); diff --git a/apps/webapp/test/waitpointPresenter.controlPlane.test.ts b/apps/webapp/test/waitpointPresenter.controlPlane.test.ts index f3513252ed..dfb8187e89 100644 --- a/apps/webapp/test/waitpointPresenter.controlPlane.test.ts +++ b/apps/webapp/test/waitpointPresenter.controlPlane.test.ts @@ -32,6 +32,7 @@ vi.mock("~/db.server", async () => { $replica: proxy, runOpsNewPrisma: proxy, sqlDatabaseSchema: Prisma.sql([`public`]), + DATABASE_SCHEMA: "public", }; }); diff --git a/apps/webapp/test/waitpointPresenter.danglingConnectedRuns.test.ts b/apps/webapp/test/waitpointPresenter.danglingConnectedRuns.test.ts new file mode 100644 index 0000000000..0c2565c545 --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.danglingConnectedRuns.test.ts @@ -0,0 +1,199 @@ +// RED->GREEN guard: #connectedRunIdsOn used to cap the connection-id fetch at +// CONNECTED_RUNS_DISPLAY_LIMIT BEFORE checking the runs exist. The dedicated store's +// WaitpointRunConnection is FK-free, so dangling rows can starve out real ones within the cap. +// The fix existence-filters AT THE QUERY (JOIN to TaskRun) so the LIMIT only ever lands on rows +// whose run exists. +import { describe, expect, vi } from "vitest"; + +const legacyReplicaHolder = vi.hoisted(() => ({ client: undefined as any })); +const newClientHolder = vi.hoisted(() => ({ client: undefined as any })); + +vi.mock("~/db.server", async () => { + const { Prisma } = await import("@trigger.dev/database"); + const lazyProxy = (holder: { client: any }, label: string) => + new Proxy( + {}, + { + get(_t, prop) { + if (!holder.client) { + throw new Error(`${label} not set for this test`); + } + return holder.client[prop]; + }, + } + ); + const replicaProxy = lazyProxy(legacyReplicaHolder, "legacyReplicaHolder.client"); + return { + prisma: replicaProxy, + $replica: replicaProxy, + runOpsNewPrisma: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsNewReplica: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsLegacyPrisma: replicaProxy, + runOpsLegacyReplica: replicaProxy, + sqlDatabaseSchema: Prisma.sql([`public`]), + DATABASE_SCHEMA: "public", + }; +}); + +vi.mock("~/services/clickhouse/clickhouseFactoryInstance.server", () => ({ + clickhouseFactory: { + getClickhouseForOrganization: async () => ({}), + }, +})); + +vi.mock("~/presenters/v3/NextRunListPresenter.server", () => ({ + NextRunListPresenter: class { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + constructor(..._args: unknown[]) {} + async call(_organizationId: string, _environmentId: string, opts: { runId?: string[] }) { + return { + runs: (opts.runId ?? []).map((friendlyId) => ({ + friendlyId, + taskIdentifier: "echoed", + })), + }; + } + }, +})); + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { + CONNECTED_RUNS_DISPLAY_LIMIT, + WaitpointPresenter, +} from "~/presenters/v3/WaitpointPresenter.server"; + +vi.setConfig({ testTimeout: 90_000 }); + +type SeedContext = { + organizationId: string; + projectId: string; + environmentId: string; +}; + +async function seedParents(prisma: PrismaClient, slug: string): Promise { + const organization = await prisma.organization.create({ + data: { title: `org-${slug}`, slug: `org-${slug}` }, + }); + const project = await prisma.project.create({ + data: { + name: `proj-${slug}`, + slug: `proj-${slug}`, + organizationId: organization.id, + externalRef: `proj-${slug}`, + }, + }); + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: `env-${slug}`, + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${slug}`, + pkApiKey: `pk_dev_${slug}`, + shortcode: `sc-${slug}`, + }, + }); + return { + organizationId: organization.id, + projectId: project.id, + environmentId: runtimeEnvironment.id, + }; +} + +async function seedWaitpoint(prisma: RunOpsPrismaClient, ctx: SeedContext, friendlyId: string) { + return (prisma as unknown as PrismaClient).waitpoint.create({ + data: { + friendlyId, + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: `idem-${friendlyId}`, + userProvidedIdempotencyKey: false, + outputType: "application/json", + outputIsError: false, + completedAt: new Date(), + tags: [], + projectId: ctx.projectId, + environmentId: ctx.environmentId, + }, + }); +} + +async function seedRun(prisma: RunOpsPrismaClient, ctx: SeedContext, friendlyId: string) { + return (prisma as unknown as PrismaClient).taskRun.create({ + data: { + friendlyId, + taskIdentifier: "my-task", + status: "PENDING", + payload: JSON.stringify({ foo: friendlyId }), + payloadType: "application/json", + traceId: friendlyId, + spanId: friendlyId, + queue: "test", + runtimeEnvironmentId: ctx.environmentId, + projectId: ctx.projectId, + organizationId: ctx.organizationId, + environmentType: "DEVELOPMENT", + engine: "V2", + }, + }); +} + +// Inserted before the real rows so the pre-fix cap-then-check code is starved of real ids. +const DANGLING_CONNECTION_COUNT = 10; +const REAL_CONNECTED_RUN_COUNT = 3; // <= CONNECTED_RUNS_DISPLAY_LIMIT (5): all must show up + +describe("WaitpointPresenter#connectedRunIdsOn is robust to dangling (FK-free) connection rows", () => { + heteroRunOpsPostgresTest( + "returns the existing connected runs even when dangling connections precede them", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "dangling"); + const waitpoint = await seedWaitpoint(prisma17, ctx, "waitpoint_dangling"); + + // Dangling connections: taskRunId points at a run that was NEVER created. The dedicated + // `WaitpointRunConnection` model is scalar/FK-free, so this insert is legal. + for (let i = 0; i < DANGLING_CONNECTION_COUNT; i++) { + await prisma17.waitpointRunConnection.create({ + data: { taskRunId: `nonexistent_run_${i}`, waitpointId: waitpoint.id }, + }); + } + + // Real, existing connected runs -- created (and connected) AFTER the dangling rows. + const realRuns = await Promise.all( + Array.from({ length: REAL_CONNECTED_RUN_COUNT }, (_, i) => + seedRun(prisma17, ctx, `run_dangling_real_${i}`) + ) + ); + for (const run of realRuns) { + await prisma17.waitpointRunConnection.create({ + data: { taskRunId: run.id, waitpointId: waitpoint.id }, + }); + } + + legacyReplicaHolder.client = prisma14; + newClientHolder.client = prisma17; + + const presenter = new WaitpointPresenter(undefined, undefined, { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaClient, + legacyReplica: prisma14, + }); + + const result = await presenter.call({ + friendlyId: waitpoint.friendlyId, + environmentId: ctx.environmentId, + projectId: ctx.projectId, + }); + + expect(result).not.toBeNull(); + const friendlyIds = (result?.connectedRuns ?? []).map((r) => r.friendlyId).sort(); + expect(friendlyIds.length).toBeLessThanOrEqual(CONNECTED_RUNS_DISPLAY_LIMIT); + // The bug: cap-before-existence-check lets dangling rows starve out the real ones, so this + // fails (returns fewer than REAL_CONNECTED_RUN_COUNT, often zero) on unfixed code. + expect(friendlyIds).toEqual( + realRuns.map((r) => r.friendlyId).sort((a, b) => a.localeCompare(b)) + ); + } + ); +}); diff --git a/apps/webapp/test/waitpointPresenter.dedicatedConnectedRuns.readthrough.test.ts b/apps/webapp/test/waitpointPresenter.dedicatedConnectedRuns.readthrough.test.ts index a71a47d53a..f8df63d7cf 100644 --- a/apps/webapp/test/waitpointPresenter.dedicatedConnectedRuns.readthrough.test.ts +++ b/apps/webapp/test/waitpointPresenter.dedicatedConnectedRuns.readthrough.test.ts @@ -31,6 +31,7 @@ vi.mock("~/db.server", async () => { runOpsLegacyPrisma: replicaProxy, runOpsLegacyReplica: replicaProxy, sqlDatabaseSchema: Prisma.sql([`public`]), + DATABASE_SCHEMA: "public", }; }); diff --git a/apps/webapp/test/waitpointPresenter.readthrough.test.ts b/apps/webapp/test/waitpointPresenter.readthrough.test.ts index a60562a950..72b794b30d 100644 --- a/apps/webapp/test/waitpointPresenter.readthrough.test.ts +++ b/apps/webapp/test/waitpointPresenter.readthrough.test.ts @@ -39,6 +39,7 @@ vi.mock("~/db.server", async () => { runOpsLegacyPrisma: replicaProxy, runOpsLegacyReplica: replicaProxy, sqlDatabaseSchema: Prisma.sql([`public`]), + DATABASE_SCHEMA: "public", }; }); @@ -67,11 +68,17 @@ import { setupClickhouseReplication } from "./utils/replicationUtils"; vi.setConfig({ testTimeout: 90_000 }); -// A read client whose waitpoint.findFirst is recorded; throws if used after being marked -// forbidden, so we can prove a store was NEVER read. Every other access forwards to the real -// client (so the inlined `environment` join + connectedRuns relation still resolve). +// A read client whose waitpoint.findFirst calls are recorded, split by kind. The presenter issues +// two shapes of waitpoint.findFirst: a HYDRATE (loads the waitpoint scalar row -- no `connectedRuns` +// in its select) and a connectedRuns-RELATION read (control-plane branch of the connected-runs +// gather -- select is exactly `{ connectedRuns }`). The `forbidden` guard throws ONLY on a HYDRATE, +// so a store can be proven to never be HYDRATED while still allowing the connectedRuns cross-DB +// union to legitimately read it (that union always touched both stores; it was invisible when the +// old code did it via raw SQL). Every other access forwards to the real client. function recording(client: PrismaClient, opts: { forbidden?: boolean } = {}) { const calls: unknown[] = []; + const hydrateCalls: unknown[] = []; + const connectedRunsCalls: unknown[] = []; return { handle: new Proxy(client, { get(target, prop) { @@ -81,8 +88,16 @@ function recording(client: PrismaClient, opts: { forbidden?: boolean } = {}) { if (wpProp === "findFirst") { return (args: unknown) => { calls.push(args); - if (opts.forbidden) { - throw new Error("this store must never be read"); + const select = (args as { select?: Record })?.select ?? {}; + const keys = Object.keys(select); + const isConnectedRunsRead = keys.length === 1 && keys[0] === "connectedRuns"; + if (isConnectedRunsRead) { + connectedRunsCalls.push(args); + } else { + hydrateCalls.push(args); + if (opts.forbidden) { + throw new Error("this store must never be hydrated"); + } } return (wpTarget as any).findFirst(args); }; @@ -95,6 +110,8 @@ function recording(client: PrismaClient, opts: { forbidden?: boolean } = {}) { }, }) as unknown as PrismaReplicaClient, calls, + hydrateCalls, + connectedRunsCalls, }; } @@ -254,9 +271,12 @@ describe("WaitpointPresenter dual-DB read-through (hetero PG14 + PG17, no connec const result = await presenter.call(callArgs(ctx, seeded.friendlyId)); expect(result?.id).toBe(seeded.friendlyId); - // New-first short-circuit: legacy never probed (the throwing handle proves it). - expect(newClient.calls.length).toBe(1); - expect(legacy.calls.length).toBe(0); + // New-first short-circuit: the HYDRATE answered from new and never fell through to a legacy + // hydrate (the throwing handle proves it). The connectedRuns cross-DB union still reads legacy + // legitimately -- that read is allowed and must NOT count as a hydrate. + expect(newClient.hydrateCalls.length).toBe(1); + expect(legacy.hydrateCalls.length).toBe(0); + expect(legacy.connectedRunsCalls.length).toBeGreaterThan(0); } ); @@ -285,8 +305,11 @@ describe("WaitpointPresenter dual-DB read-through (hetero PG14 + PG17, no connec expect(result?.id).toBe(seeded.friendlyId); expect(result?.tags).toEqual(["one"]); - // Exactly one read on the single client; the second handle is never touched. - expect(single.calls.length).toBe(1); + // Two reads on the single client, both on the one handle: the first findFirst hydrates the + // waitpoint, the second loads the `connectedRuns` relation (the implicit M2M has no queryable + // join delegate, so the ORM must traverse the relation with a second findFirst). The second + // handle is never touched -- passthrough is structural, the split branch never fires. + expect(single.calls.length).toBe(2); expect(second.calls.length).toBe(0); } ); diff --git a/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts b/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts new file mode 100644 index 0000000000..ac321e210a --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts @@ -0,0 +1,208 @@ +// FLOW-level gap: a waitpoint's connected runs SPLIT across BOTH physical DBs at once (some legacy, +// some new-resident) rather than the existing single-direction cross-DB cases in +// waitpointPresenter.dedicatedConnectedRuns.readthrough.test.ts. Asserts #connectedRunFriendlyIds +// unions friendlyIds gathered from BOTH stores in one read AND stays bounded to +// CONNECTED_RUNS_DISPLAY_LIMIT even when the combined total exceeds it. Real two-physical-DB +// topology (heteroRunOpsPostgresTest); never mocked beyond the same db.server/clickhouse/ +// NextRunListPresenter echo seams the sibling presenter tests use. +import { describe, expect, vi } from "vitest"; + +const legacyReplicaHolder = vi.hoisted(() => ({ client: undefined as any })); +const newClientHolder = vi.hoisted(() => ({ client: undefined as any })); + +vi.mock("~/db.server", async () => { + const { Prisma } = await import("@trigger.dev/database"); + const lazyProxy = (holder: { client: any }, label: string) => + new Proxy( + {}, + { + get(_t, prop) { + if (!holder.client) { + throw new Error(`${label} not set for this test`); + } + return holder.client[prop]; + }, + } + ); + const replicaProxy = lazyProxy(legacyReplicaHolder, "legacyReplicaHolder.client"); + return { + prisma: replicaProxy, + $replica: replicaProxy, + runOpsNewPrisma: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsNewReplica: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsLegacyPrisma: replicaProxy, + runOpsLegacyReplica: replicaProxy, + sqlDatabaseSchema: Prisma.sql([`public`]), + DATABASE_SCHEMA: "public", + }; +}); + +vi.mock("~/services/clickhouse/clickhouseFactoryInstance.server", () => ({ + clickhouseFactory: { + getClickhouseForOrganization: async () => ({}), + }, +})); + +// Echo the runId set back as runs so `result.connectedRuns` == the friendlyIds the presenter +// gathered cross-DB (isolates the gather from the CH hydrate, as the sibling tests do). +vi.mock("~/presenters/v3/NextRunListPresenter.server", () => ({ + NextRunListPresenter: class { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + constructor(..._args: unknown[]) {} + async call(_organizationId: string, _environmentId: string, opts: { runId?: string[] }) { + return { + runs: (opts.runId ?? []).map((friendlyId) => ({ + friendlyId, + taskIdentifier: "echoed", + })), + }; + } + }, +})); + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { + CONNECTED_RUNS_DISPLAY_LIMIT, + WaitpointPresenter, +} from "~/presenters/v3/WaitpointPresenter.server"; + +vi.setConfig({ testTimeout: 90_000 }); + +type SeedContext = { + organizationId: string; + projectId: string; + environmentId: string; +}; + +async function seedParents(prisma: PrismaClient, slug: string): Promise { + const organization = await prisma.organization.create({ + data: { title: `org-${slug}`, slug: `org-${slug}` }, + }); + const project = await prisma.project.create({ + data: { + name: `proj-${slug}`, + slug: `proj-${slug}`, + organizationId: organization.id, + externalRef: `proj-${slug}`, + }, + }); + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: `env-${slug}`, + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${slug}`, + pkApiKey: `pk_dev_${slug}`, + shortcode: `sc-${slug}`, + }, + }); + return { + organizationId: organization.id, + projectId: project.id, + environmentId: runtimeEnvironment.id, + }; +} + +async function seedRun( + prisma: PrismaClient | RunOpsPrismaClient, + ctx: SeedContext, + friendlyId: string +) { + return (prisma as PrismaClient).taskRun.create({ + data: { + friendlyId, + taskIdentifier: "my-task", + status: "PENDING", + payload: JSON.stringify({ foo: friendlyId }), + payloadType: "application/json", + traceId: friendlyId, + spanId: friendlyId, + queue: "test", + runtimeEnvironmentId: ctx.environmentId, + projectId: ctx.projectId, + organizationId: ctx.organizationId, + environmentType: "DEVELOPMENT", + engine: "V2", + }, + }); +} + +describe("WaitpointPresenter — connected runs SPLIT across both physical DBs", () => { + heteroRunOpsPostgresTest( + "unions connected-run friendlyIds from both stores and stays bounded to the display limit", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "split"); + + // Waitpoint resident on LEGACY. + const waitpoint = await prisma14.waitpoint.create({ + data: { + friendlyId: "waitpoint_split", + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: "idem-waitpoint_split", + userProvidedIdempotencyKey: false, + outputType: "application/json", + outputIsError: false, + completedAt: new Date(), + tags: [], + projectId: ctx.projectId, + environmentId: ctx.environmentId, + }, + }); + + // 2 connected runs resident + joined on NEW (below the limit on its own). + const NEW_RUN_FRIENDLY_IDS = ["run_split_new0", "run_split_new1"]; + for (const friendlyId of NEW_RUN_FRIENDLY_IDS) { + const run = await seedRun(prisma17, ctx, friendlyId); + await prisma17.waitpointRunConnection.create({ + data: { taskRunId: run.id, waitpointId: waitpoint.id }, + }); + } + + // 6 connected runs resident + joined on LEGACY via the implicit M2M — more than enough, on its + // own, to push the combined total past CONNECTED_RUNS_DISPLAY_LIMIT (5). + const LEGACY_RUN_FRIENDLY_IDS = Array.from({ length: 6 }, (_, i) => `run_split_leg${i}`); + for (const friendlyId of LEGACY_RUN_FRIENDLY_IDS) { + const run = await seedRun(prisma14, ctx, friendlyId); + await prisma14.waitpoint.update({ + where: { id: waitpoint.id }, + data: { connectedRuns: { connect: [{ id: run.id }] } }, + }); + } + + legacyReplicaHolder.client = prisma14; + newClientHolder.client = prisma17; + + const presenter = new WaitpointPresenter(undefined, undefined, { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaClient, + legacyReplica: prisma14, + }); + + const result = await presenter.call({ + friendlyId: waitpoint.friendlyId, + environmentId: ctx.environmentId, + projectId: ctx.projectId, + }); + + const returnedIds = result?.connectedRuns.map((r) => r.friendlyId) ?? []; + + // Bounded: the combined total across both DBs (8) must never surface more than the display + // limit, proving the fetch is capped globally, not just per-DB. + expect(returnedIds.length).toBeLessThanOrEqual(CONNECTED_RUNS_DISPLAY_LIMIT); + expect(returnedIds.length).toBeGreaterThan(0); + + // BOTH stores contributed: at least one NEW-resident and one LEGACY-resident connected run + // friendlyId made it into the result. A single-store gather (the bug this guards against) + // would surface only one side. + const fromNew = returnedIds.filter((id) => NEW_RUN_FRIENDLY_IDS.includes(id)); + const fromLegacy = returnedIds.filter((id) => LEGACY_RUN_FRIENDLY_IDS.includes(id)); + expect(fromNew.length).toBeGreaterThan(0); + expect(fromLegacy.length).toBeGreaterThan(0); + expect(fromNew.length + fromLegacy.length).toBe(returnedIds.length); + } + ); +}); diff --git a/internal-packages/run-engine/src/engine/tests/legacyRunNewTokenResumeFlow.test.ts b/internal-packages/run-engine/src/engine/tests/legacyRunNewTokenResumeFlow.test.ts new file mode 100644 index 0000000000..20f68f13d2 --- /dev/null +++ b/internal-packages/run-engine/src/engine/tests/legacyRunNewTokenResumeFlow.test.ts @@ -0,0 +1,302 @@ +// Completion -> resume FLOW test for the direction crossDbTokenBlock.test.ts's writer fix enables but +// never itself exercises end-to-end: a LEGACY (cuid) run blocks on a standalone MANUAL token that is +// NEW-resident (run-ops id) — e.g. the token was minted co-located with a run-ops sibling. This test +// drives the real engine path (blockRunWithWaitpoint -> completeWaitpoint -> continueRunIfUnblocked) +// across the two PHYSICAL DBs to prove the blocking edge is discovered and the run actually resumes, +// not just that the edge write succeeds (that unit-level assertion already lives in +// runOpsStore.crossDbTokenBlock.test.ts). Mirrors completeWaitpointReadResidency.test.ts's own +// cross-DB case, but with the run/token residencies swapped. Real two-physical-DB topology; never +// mocked. + +import { + heteroRunOpsPostgresTest, + network, + redisContainer, + redisOptions, +} from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { PostgresRunStore, RoutingRunStore, type CreateRunInput } from "@internal/run-store"; +import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { expect, vi } from "vitest"; +import { RunEngine } from "../index.js"; + +const twoDbEngineTest = heteroRunOpsPostgresTest.extend<{ + redisContainer: any; + redisOptions: any; +}>({ + network, + redisContainer, + redisOptions, +}); + +// 25-char cuid body (no v1 version marker at index 25) -> classifies LEGACY -> #legacy (prisma14). +const CUID_25 = "c".repeat(25); + +function baseEngineOptions(redisOptions: any, prisma: any) { + return { + prisma, + worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 }, + queue: { + redis: redisOptions, + masterQueueConsumersDisabled: true, + processWorkerQueueDebounceMs: 50, + }, + runLock: { redis: redisOptions }, + machines: { + defaultMachine: "small-1x" as const, + machines: { + "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("test", "0.0.0"), + }; +} + +async function seedControlPlaneEnv(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 { organization, project, environment }; +} + +function buildCreateRunInput(p: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: p.runId, + engine: "V2", + status: "EXECUTING", + friendlyId: p.friendlyId, + runtimeEnvironmentId: p.runtimeEnvironmentId, + environmentType: "PRODUCTION", + organizationId: p.organizationId, + projectId: p.projectId, + taskIdentifier: "parent-task", + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: `trace_${p.runId}`, + spanId: `span_${p.runId}`, + runTags: [], + queue: "task/parent-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: new Date("2024-01-01T00:00:00.000Z"), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: p.runtimeEnvironmentId, + environmentType: "PRODUCTION", + projectId: p.projectId, + organizationId: p.organizationId, + }, + }; +} + +// Seed an EXECUTING LEGACY (cuid) parent run on #legacy (prisma14), plus a standalone MANUAL token +// that is NEW-resident (run-ops id, minted co-located with some other run-ops sibling) on #new +// (prisma17) ONLY — the token is never created on #legacy. +async function seedExecutingLegacyParentAndNewToken( + prisma14: PrismaClient, + prisma17: RunOpsPrismaClient, + router: RoutingRunStore, + parentRunId: string, + waitpointId: string, + suffix: string +) { + const env = await seedControlPlaneEnv(prisma14, suffix); + + await router.createRun( + buildCreateRunInput({ + runId: parentRunId, + friendlyId: `run_${suffix}_parent`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + + const created = await router.findLatestExecutionSnapshot(parentRunId); + await router.createExecutionSnapshot( + { + run: { id: parentRunId, status: "EXECUTING", attemptNumber: 1 }, + snapshot: { executionStatus: "EXECUTING", description: "parent executing" }, + previousSnapshotId: created!.id, + environmentId: env.environment.id, + environmentType: "PRODUCTION", + projectId: env.project.id, + organizationId: env.organization.id, + }, + prisma14 + ); + + // Standalone MANUAL token lives on #new ONLY (run-ops id) — NOT on #legacy. + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: `wp_${suffix}`, + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + return env; +} + +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 }); +} + +describe("RunEngine completion->resume flow — a LEGACY run blocked on a NEW-resident token resumes across the seam", () => { + twoDbEngineTest( + "completeWaitpoint on the NEW token finds the #legacy edge and unblocks the LEGACY run", + async ({ prisma14, prisma17, redisOptions }) => { + const router = makeRouter(prisma14 as unknown as PrismaClient, prisma17); + + // The harness builds the schema with `prisma db push`, which re-creates the FKs that the + // run-ops split migrations drop in prod (see runOpsStore.crossDbTokenBlock.test.ts). Mirror + // those drops on this clone so the cross-DB edge write exercises real prod state. + await (prisma14 as unknown as PrismaClient).$executeRawUnsafe( + `ALTER TABLE "TaskRunWaitpoint" DROP CONSTRAINT IF EXISTS "TaskRunWaitpoint_waitpointId_fkey"` + ); + await (prisma14 as unknown as PrismaClient).$executeRawUnsafe( + `ALTER TABLE "_WaitpointRunConnections" DROP CONSTRAINT IF EXISTS "_WaitpointRunConnections_B_fkey"` + ); + // continueRunIfUnblocked's resume snapshot connects the now-COMPLETED (NEW-resident) token via + // _completedWaitpoints (migration 20260705230000 drops this FK in prod; see + // runOpsStore.crossDbCompletedWaitpoint.test.ts). + await (prisma14 as unknown as PrismaClient).$executeRawUnsafe( + `ALTER TABLE "_completedWaitpoints" DROP CONSTRAINT IF EXISTS "_completedWaitpoints_B_fkey"` + ); + + const engine = new RunEngine({ + store: router, + ...baseEngineOptions(redisOptions, prisma14), + }); + + try { + const parentRunId = `run_${CUID_25}`; // LEGACY run -> #legacy + const tokenId = `waitpoint_${generateRunOpsId()}`; // NEW-resident standalone token -> #new + const env = await seedExecutingLegacyParentAndNewToken( + prisma14 as unknown as PrismaClient, + prisma17, + router, + parentRunId, + tokenId, + "xdbresume" + ); + + // Block the LEGACY run on the NEW token. The edge must land on #legacy (the run's own DB, + // FK-free now that the migration-dropped constraints are mirrored above). + await engine.blockRunWithWaitpoint({ + runId: parentRunId, + waitpoints: tokenId, + projectId: env.project.id, + organizationId: env.organization.id, + tx: prisma14 as unknown as PrismaClient, + }); + + expect( + await (prisma14 as unknown as PrismaClient).taskRunWaitpoint.count({ + where: { taskRunId: parentRunId }, + }) + ).toBe(1); + expect(await prisma17.taskRunWaitpoint.count({ where: { taskRunId: parentRunId } })).toBe( + 0 + ); + + const blocked = await engine.getRunExecutionData({ runId: parentRunId }); + expect(blocked?.snapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + const enqueueSpy = vi.spyOn((engine as any).worker, "enqueue"); + + // Complete the NEW token via the engine path. forWaitpointCompletion resolves #new (run-ops + // id, no pins); completeWaitpoint marks it COMPLETED there, then fans the waitpointId edge + // read across BOTH DBs (findManyTaskRunWaitpoints) to discover the #legacy-resident edge. + const completed = await engine.completeWaitpoint({ + id: tokenId, + output: { value: '{"resumed":"legacy-run-new-token"}', isError: false }, + }); + expect(completed.status).toBe("COMPLETED"); + + // Token COMPLETED on #new only. + expect((await prisma17.waitpoint.findFirst({ where: { id: tokenId } }))?.status).toBe( + "COMPLETED" + ); + expect( + await (prisma14 as unknown as PrismaClient).waitpoint.findFirst({ + where: { id: tokenId }, + }) + ).toBeNull(); + + // The fan-out found the #legacy edge and enqueued the unblock for the LEGACY run. + const continueEnqueued = enqueueSpy.mock.calls.some( + ([arg]) => + (arg as any)?.job === "continueRunIfUnblocked" && + (arg as any)?.payload?.runId === parentRunId + ); + expect(continueEnqueued).toBe(true); + + // Driving the unblock body re-resolves the NEW token's COMPLETED status across the seam + // (continueRunIfUnblocked's blockingWaitpoints read rehydrates `waitpoint` cross-DB) and + // actually resumes the LEGACY run — the higher-level outcome, not just the edge write. + const result = await (engine as any).waitpointSystem.continueRunIfUnblocked({ + runId: parentRunId, + }); + expect(result.status).toBe("unblocked"); + + const after = await engine.getRunExecutionData({ runId: parentRunId }); + expect(after?.snapshot.executionStatus).not.toBe("EXECUTING_WITH_WAITPOINTS"); + } finally { + await engine.quit(); + } + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.connectedRunsBounded.test.ts b/internal-packages/run-store/src/PostgresRunStore.connectedRunsBounded.test.ts new file mode 100644 index 0000000000..e20ad22805 --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.connectedRunsBounded.test.ts @@ -0,0 +1,286 @@ +// RED→GREEN for bounding the connected-runs read path so it can never emit an unbounded id list. +// +// Mirrors the fix already shipped in `apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts` +// `#connectedRunIdsOn`: existence-JOIN to TaskRun (so a dangling connection row can never occupy a +// LIMIT slot ahead of a real run), then LIMIT to `CONNECTED_RUNS_LIMIT` on BOTH schema branches. +// +// `PostgresRunStore.findWaitpointConnectedRunIds` currently has neither: the dedicated branch does +// a bare `waitpointRunConnection.findMany` (no LIMIT, no existence check) and the legacy branch does +// a bare `SELECT "A" FROM "_WaitpointRunConnections"` (same). `RoutingRunStore.findWaitpointConnectedRunIds` +// unions the two stores' results with no final cap, so even if each store were individually bounded +// the cross-DB union would not be. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import { describe, expect } from "vitest"; +import { CONNECTED_RUNS_LIMIT, PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; + +function seedEnvironmentDedicated(suffix: string) { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; +} + +async function seedEnvironmentLegacy(prisma: any, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { organization, project, environment }; +} + +function taskRunData(opts: { + id: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}) { + return { + id: opts.id, + engine: "V2" as const, + status: "PENDING" as const, + friendlyId: opts.friendlyId, + runtimeEnvironmentId: opts.runtimeEnvironmentId, + environmentType: "DEVELOPMENT" as const, + organizationId: opts.organizationId, + projectId: opts.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${opts.id}`, + spanId: `span_${opts.id}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }; +} + +describe("PostgresRunStore.findWaitpointConnectedRunIds — bounded via existence-JOIN + LIMIT", () => { + heteroRunOpsPostgresTest( + "dedicated schema: caps at CONNECTED_RUNS_LIMIT and never returns a dangling (run-less) connection", + async ({ prisma17 }) => { + const store = new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); + + const newEnv = seedEnvironmentDedicated("bound_new"); + const waitpointId = "waitpoint_bound_new"; + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: "wp_bound_new", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: newEnv.project.id, + environmentId: newEnv.environment.id, + }, + }); + + // Dangling connections FIRST: taskRunId points at runs that were NEVER created. Legal on the + // dedicated schema — `WaitpointRunConnection.taskRunId` is a scalar, no FK (see schema.prisma). + const danglingRunIds = Array.from({ length: 20 }, (_, i) => `run_dangling_${i}`); + for (const runId of danglingRunIds) { + await store.blockRunWithWaitpointEdges({ + runId, + waitpointIds: [waitpointId], + projectId: newEnv.project.id, + }); + } + + // A few REAL connected runs, inserted AFTER the dangling ones. + const realRunIds = ["run_real_0", "run_real_1", "run_real_2"]; + for (const [i, id] of realRunIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_bound_new_${i}`, + organizationId: newEnv.organization.id, + projectId: newEnv.project.id, + runtimeEnvironmentId: newEnv.environment.id, + }), + }); + await store.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: newEnv.project.id, + }); + } + + const result = await store.findWaitpointConnectedRunIds(waitpointId); + + // Bounded. + expect(result.length).toBeLessThanOrEqual(CONNECTED_RUNS_LIMIT); + // Never a dangling id — every returned id must be one of the REAL runs. + for (const id of result) { + expect(realRunIds).toContain(id); + } + // Not starved out: all 3 real runs (well under the limit) must be present. + for (const id of realRunIds) { + expect(result).toContain(id); + } + } + ); + + heteroRunOpsPostgresTest( + "legacy schema: caps at CONNECTED_RUNS_LIMIT instead of returning every connected run", + async ({ prisma14 }) => { + const store = new PostgresRunStore({ + prisma: prisma14 as never, + readOnlyPrisma: prisma14 as never, + schemaVariant: "legacy", + }); + + const legEnv = await seedEnvironmentLegacy(prisma14, "bound_leg"); + const waitpointId = "waitpoint_bound_leg"; + await prisma14.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: "wp_bound_leg", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: legEnv.project.id, + environmentId: legEnv.environment.id, + }, + }); + + // More real connected runs than CONNECTED_RUNS_LIMIT — the legacy `_WaitpointRunConnections` + // table still enforces an FK from "A" to TaskRun, so every row here is a real run. + const realRunIds = Array.from({ length: CONNECTED_RUNS_LIMIT + 3 }, (_, i) => `run_leg_${i}`); + for (const [i, id] of realRunIds.entries()) { + await prisma14.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_bound_leg_${i}`, + organizationId: legEnv.organization.id, + projectId: legEnv.project.id, + runtimeEnvironmentId: legEnv.environment.id, + }), + }); + await store.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: legEnv.project.id, + }); + } + + const result = await store.findWaitpointConnectedRunIds(waitpointId); + + expect(result.length).toBeLessThanOrEqual(CONNECTED_RUNS_LIMIT); + } + ); +}); + +describe("RoutingRunStore.findWaitpointConnectedRunIds — the cross-DB union is also bounded", () => { + heteroRunOpsPostgresTest( + "slices the merged NEW+LEGACY result to CONNECTED_RUNS_LIMIT", + async ({ prisma14, prisma17 }) => { + const legacyStore = new PostgresRunStore({ + prisma: prisma14 as never, + readOnlyPrisma: prisma14 as never, + schemaVariant: "legacy", + }); + const newStore = new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + + const legEnv = await seedEnvironmentLegacy(prisma14, "bound_union_leg"); + const newEnv = seedEnvironmentDedicated("bound_union_new"); + const waitpointId = "waitpoint_bound_union"; + // The legacy `TaskRunWaitpoint.waitpointId` FK is still live in this test DB (built via + // `prisma db push` from schema.prisma, which still declares the relation even though a later + // migration drops the constraint in real deployments) — the legacy leg needs a real row here. + await prisma14.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: "wp_bound_union", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: legEnv.project.id, + environmentId: legEnv.environment.id, + }, + }); + + // Each store individually caps at CONNECTED_RUNS_LIMIT, but a distinct set of runs on EACH side + // means the union of the two capped results can still exceed the limit without a final slice. + const legacyRunIds = Array.from( + { length: CONNECTED_RUNS_LIMIT }, + (_, i) => `run_union_leg_${i}` + ); + for (const [i, id] of legacyRunIds.entries()) { + await prisma14.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_union_leg_${i}`, + organizationId: legEnv.organization.id, + projectId: legEnv.project.id, + runtimeEnvironmentId: legEnv.environment.id, + }), + }); + await legacyStore.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: legEnv.project.id, + }); + } + + const newRunIds = Array.from( + { length: CONNECTED_RUNS_LIMIT }, + (_, i) => `run_union_new_${i}` + ); + for (const [i, id] of newRunIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_union_new_${i}`, + organizationId: newEnv.organization.id, + projectId: newEnv.project.id, + runtimeEnvironmentId: newEnv.environment.id, + }), + }); + await newStore.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: newEnv.project.id, + }); + } + + const result = await router.findWaitpointConnectedRunIds(waitpointId); + + expect(result.length).toBeLessThanOrEqual(CONNECTED_RUNS_LIMIT); + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts b/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts new file mode 100644 index 0000000000..a2ce43103c --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts @@ -0,0 +1,130 @@ +// RED→GREEN for bounding the DISPLAY `connectedRuns` RELATION hydration (findWaitpoint / +// findManyWaitpoints with `include`/`select: { connectedRuns }`), NOT the id-list helper +// `findWaitpointConnectedRunIds` (covered separately in connectedRunsBounded.test.ts). +// +// The dedicated-schema hydrator `hydrateConnectedRuns` goes through the generic +// `batchHydrateJoinRelation`, which fetched EVERY WaitpointRunConnection link with no per-parent +// cap — so a heavily-fanned-in waitpoint hydrated an unbounded connectedRuns list, bypassing the +// CONNECTED_RUNS_LIMIT the presenter/`findWaitpointConnectedRunIds` already enforce. +// +// The fix bounds it per parent via a window function + existence-JOIN to TaskRun, mirroring the +// bounded helper: a dangling connection row never occupies a LIMIT slot, and the returned relation +// is capped at CONNECTED_RUNS_LIMIT. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import { describe, expect } from "vitest"; +import { CONNECTED_RUNS_LIMIT, PostgresRunStore } from "./PostgresRunStore.js"; + +function seedEnvironmentDedicated(suffix: string) { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; +} + +function taskRunData(opts: { + id: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}) { + return { + id: opts.id, + engine: "V2" as const, + status: "PENDING" as const, + friendlyId: opts.friendlyId, + runtimeEnvironmentId: opts.runtimeEnvironmentId, + environmentType: "DEVELOPMENT" as const, + organizationId: opts.organizationId, + projectId: opts.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${opts.id}`, + spanId: `span_${opts.id}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }; +} + +describe("PostgresRunStore.findWaitpoint — connectedRuns relation is bounded", () => { + heteroRunOpsPostgresTest( + "dedicated schema: include connectedRuns caps at CONNECTED_RUNS_LIMIT and skips dangling connections", + async ({ prisma17 }) => { + const store = new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); + + const env = seedEnvironmentDedicated("cr_incl_new"); + const waitpointId = "waitpoint_cr_incl_new"; + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: "wp_cr_incl_new", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + // Dangling connections (taskRunId points at runs that were never created). Legal on the + // dedicated schema — WaitpointRunConnection.taskRunId is a scalar with no FK. + const danglingRunIds = Array.from({ length: 20 }, (_, i) => `run_cr_dangling_${i}`); + for (const runId of danglingRunIds) { + await store.blockRunWithWaitpointEdges({ + runId, + waitpointIds: [waitpointId], + projectId: env.project.id, + }); + } + + // More REAL connected runs than the limit. + const realRunIds = Array.from( + { length: CONNECTED_RUNS_LIMIT + 3 }, + (_, i) => `run_cr_real_${i}` + ); + for (const [i, id] of realRunIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_cr_incl_new_${i}`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }), + }); + await store.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: env.project.id, + }); + } + + const waitpoint = (await store.findWaitpoint({ + where: { id: waitpointId }, + include: { connectedRuns: true }, + })) as unknown as { connectedRuns: { id: string }[] } | null; + + expect(waitpoint).not.toBeNull(); + const connectedRuns = waitpoint!.connectedRuns; + + // Bounded to exactly the cap (fixture connects more real runs than the limit). + expect(connectedRuns).toHaveLength(CONNECTED_RUNS_LIMIT); + // Never a dangling id — every returned run must be a REAL run. + for (const run of connectedRuns) { + expect(realRunIds).toContain(run.id); + expect(danglingRunIds).not.toContain(run.id); + } + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts new file mode 100644 index 0000000000..10e9aa8086 --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts @@ -0,0 +1,175 @@ +// RED→GREEN: `findRunsByIds` is the grouped replacement for `Promise.all(ids.map(id => +// findRun(id)))` — it must issue ONE `findMany` for the whole id batch, never a `findFirst` +// per id. `postgresTest` is a REAL PrismaClient over the standard (non-dedicated) schema. The +// counting proxy wraps its `taskRun` delegate to tally `findFirst`/`findMany` calls while +// delegating every call to the real client — the DB still runs the query; this is +// instrumentation, not a mock. + +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import type { CreateRunInput } from "./types.js"; + +type CallCounts = { findMany: number; findFirst: number }; + +// Wraps the named delegates on a REAL client to tally `findMany`/`findFirst` calls PER delegate, +// delegating every call unchanged to the real client (the DB still runs the query). +function countingClient( + real: PrismaClient, + delegateNames: string[] +): { client: PrismaClient; counts: Record } { + const counts: Record = {}; + for (const name of delegateNames) { + counts[name] = { findMany: 0, findFirst: 0 }; + } + const wrapped = new Map( + delegateNames.map((name) => [ + name, + new Proxy((real as any)[name], { + get(target, prop) { + if (prop === "findMany" || prop === "findFirst") { + counts[name][prop as "findMany" | "findFirst"]++; + } + return (target as any)[prop]; + }, + }), + ]) + ); + const client = new Proxy(real, { + get(target, prop) { + if (typeof prop === "string" && wrapped.has(prop)) { + return wrapped.get(prop); + } + return (target as any)[prop]; + }, + }) as PrismaClient; + return { client, counts }; +} + +async function seedEnvironment(prisma: PrismaClient) { + const organization = await prisma.organization.create({ + data: { + title: "Test Organization", + slug: "test-organization-findrunsbyids", + }, + }); + + const project = await prisma.project.create({ + data: { + name: "Test Project", + slug: "test-project-findrunsbyids", + externalRef: "proj_findrunsbyids", + organizationId: organization.id, + }, + }); + + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: "tr_dev_apikey_findrunsbyids", + pkApiKey: "pk_dev_apikey_findrunsbyids", + shortcode: "short_code_findrunsbyids", + }, + }); + + return { organization, project, environment }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +describe("PostgresRunStore.findRunsByIds", () => { + postgresTest( + "resolves N ids with ONE grouped findMany, never a findFirst per id", + async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + const seedStore = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + + const runId1 = "run_findbyids_1"; + const runId2 = "run_findbyids_2"; + await seedStore.createRun( + buildCreateRunInput({ + runId: runId1, + friendlyId: "run_findbyids_1_friendly", + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + await seedStore.createRun( + buildCreateRunInput({ + runId: runId2, + friendlyId: "run_findbyids_2_friendly", + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const counting = countingClient(prisma, ["taskRun"]); + const readStore = new PostgresRunStore({ prisma, readOnlyPrisma: counting.client }); + + const byId = await readStore.findRunsByIds([runId1, runId2], { + select: { friendlyId: true, status: true }, + }); + + expect(byId.size).toBe(2); + expect(byId.get(runId1)?.friendlyId).toBe("run_findbyids_1_friendly"); + expect(byId.get(runId1)?.status).toBe("PENDING"); + expect(byId.get(runId2)?.friendlyId).toBe("run_findbyids_2_friendly"); + expect(byId.get(runId2)?.status).toBe("PENDING"); + + // The select omitted `id`; the id we force-inject for keying must not leak into the value + // (it would otherwise violate the declared payload type and expose an unrequested id). + expect("id" in (byId.get(runId1) as object)).toBe(false); + expect("id" in (byId.get(runId2) as object)).toBe(false); + + // GROUPED: one findMany for the WHOLE batch, never one findFirst per id. + expect(counting.counts.taskRun.findMany).toBe(1); + expect(counting.counts.taskRun.findFirst).toBe(0); + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts b/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts new file mode 100644 index 0000000000..6323ca4858 --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts @@ -0,0 +1,239 @@ +// RED→GREEN: the dedicated-schema relation hydrators (`hydrateBlockingTaskRuns` / +// `#hydrateDedicatedRelations`) resolve a relation PER PARENT ROW, so for N parents each requesting +// the same relation, the store issues N round trips (a `findFirst`/`findMany` per row) instead of one +// grouped query for the whole batch. `heteroRunOpsPostgresTest.prisma17` is a REAL RunOpsPrismaClient +// over the dedicated subset schema. The counting proxy wraps its `taskRun` delegate to tally +// `findFirst`/`findMany` calls while delegating every call to the real client — the DB still runs the +// query; this is instrumentation, not a mock. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import type { CreateRunInput } from "./types.js"; + +type CallCounts = { findMany: number; findFirst: number }; + +// Wraps the named delegates on a REAL client to tally `findMany`/`findFirst` calls PER delegate, +// delegating every call unchanged to the real client (the DB still runs the query). +function countingClient( + real: RunOpsPrismaClient, + delegateNames: string[] +): { client: RunOpsPrismaClient; counts: Record; queryRaw: { count: number } } { + const counts: Record = {}; + for (const name of delegateNames) { + counts[name] = { findMany: 0, findFirst: 0 }; + } + const queryRaw = { count: 0 }; + const wrapped = new Map( + delegateNames.map((name) => [ + name, + new Proxy((real as any)[name], { + get(target, prop) { + if (prop === "findMany" || prop === "findFirst") { + counts[name][prop as "findMany" | "findFirst"]++; + } + return (target as any)[prop]; + }, + }), + ]) + ); + const client = new Proxy(real, { + get(target, prop) { + if (typeof prop === "string" && wrapped.has(prop)) { + return wrapped.get(prop); + } + // Tally the grouped raw query the bounded connectedRuns hydrator issues, delegating unchanged. + if (prop === "$queryRaw") { + return (...args: unknown[]) => { + queryRaw.count++; + return (target as any).$queryRaw(...args); + }; + } + return (target as any)[prop]; + }, + }) as RunOpsPrismaClient; + return { client, counts, queryRaw }; +} + +function seedEnvironmentDedicated(suffix: string) { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: '{"hello":"world"}', + payloadType: "application/json", + traceContext: { trace: "ctx" }, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + runTags: [], + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: new Date("2024-01-01T00:00:00.000Z"), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +function makeStore(prisma: RunOpsPrismaClient) { + return new PostgresRunStore({ + prisma: prisma as never, + readOnlyPrisma: prisma as never, + schemaVariant: "dedicated", + }); +} + +describe("PostgresRunStore dedicated relation hydrators — grouped batch reads", () => { + heteroRunOpsPostgresTest( + "hydrateBlockingTaskRuns resolves N edges' `taskRun` with ONE grouped findMany, not N findFirst", + async ({ prisma17 }) => { + const seedStore = makeStore(prisma17); + const env = seedEnvironmentDedicated("hbt"); + const waitpointId = "wp_hbt_shared"; + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: "wp_hbt_friendly", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + const runIds = ["run_hbt_1", "run_hbt_2", "run_hbt_3"]; + for (const runId of runIds) { + await seedStore.createRun( + buildCreateRunInput({ + runId, + friendlyId: `${runId}_friendly`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + await seedStore.blockRunWithWaitpointEdges({ + runId, + waitpointIds: [waitpointId], + projectId: env.project.id, + }); + } + + const counting = countingClient(prisma17, ["taskRun"]); + const readStore = makeStore(counting.client); + + const waitpoint = (await readStore.findWaitpoint({ + where: { id: waitpointId }, + include: { + blockingTaskRuns: { select: { taskRun: { select: { id: true, friendlyId: true } } } }, + }, + })) as { blockingTaskRuns: { taskRun: { id: string; friendlyId: string } | null }[] } | null; + + const blocking = waitpoint?.blockingTaskRuns ?? []; + expect(blocking).toHaveLength(3); + expect(blocking.map((b) => b.taskRun?.id).sort()).toEqual([...runIds].sort()); + + // GROUPED: one findMany for the WHOLE batch, never one findFirst per edge. + expect(counting.counts.taskRun.findMany).toBe(1); + expect(counting.counts.taskRun.findFirst).toBe(0); + } + ); + + heteroRunOpsPostgresTest( + "findManyWaitpoints hydrates N parents' `connectedRuns` with ONE grouped join query + ONE grouped target query", + async ({ prisma17 }) => { + const seedStore = makeStore(prisma17); + const env = seedEnvironmentDedicated("cr"); + + const waitpointIds = ["wp_cr_1", "wp_cr_2", "wp_cr_3"]; + const runIds = ["run_cr_1", "run_cr_2", "run_cr_3"]; + for (const [i, waitpointId] of waitpointIds.entries()) { + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: `${waitpointId}_friendly`, + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + await seedStore.createRun( + buildCreateRunInput({ + runId: runIds[i], + friendlyId: `${runIds[i]}_friendly`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + // One connection per waitpoint, to its own distinct run. + await seedStore.blockRunWithWaitpointEdges({ + runId: runIds[i], + waitpointIds: [waitpointId], + projectId: env.project.id, + }); + } + + const counting = countingClient(prisma17, ["waitpointRunConnection", "taskRun"]); + const readStore = makeStore(counting.client); + + const rows = (await readStore.findManyWaitpoints({ + where: { id: { in: waitpointIds } }, + include: { connectedRuns: { select: { id: true, friendlyId: true } } }, + })) as { id: string; connectedRuns: { id: string; friendlyId: string }[] }[]; + + expect(rows).toHaveLength(3); + const byWaitpointId = new Map(rows.map((r) => [r.id, r.connectedRuns])); + for (const [i, waitpointId] of waitpointIds.entries()) { + const connected = byWaitpointId.get(waitpointId) ?? []; + expect(connected).toHaveLength(1); + expect(connected[0].id).toBe(runIds[i]); + } + + // GROUPED: one bounded join query (raw, ROW_NUMBER-per-parent) + one target query for the + // WHOLE batch, never one pair per parent. The bounded hydrator replaces the delegate + // `waitpointRunConnection.findMany` with a single `$queryRaw`. + expect(counting.queryRaw.count).toBe(1); + expect(counting.counts.waitpointRunConnection.findMany).toBe(0); + expect(counting.counts.taskRun.findMany).toBe(1); + expect(counting.counts.taskRun.findFirst).toBe(0); + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts b/internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts new file mode 100644 index 0000000000..0db044bd94 --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts @@ -0,0 +1,226 @@ +// RED→GREEN: the dedicated-schema relation hydrators (`batchHydrateJoinRelation` / +// `batchHydrateEdgeTarget`) fetch the target row's FULL columns (no `select`) and only strip fields +// in JS afterwards via `applyProjection`. At prod scale (a 6.68B-row `TaskRun`) that over-fetches +// every wide TOASTed column (`payload`/`output`/`context`) even when the caller only asked for +// `friendlyId`. The RESULT is identical with/without the fix (`applyProjection` trims either way), +// so this asserts the QUERY SHAPE Prisma actually receives, via a `$extends` query hook on a REAL +// `heteroRunOpsPostgresTest.prisma17` client — the DB still runs the query, this only observes it. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import type { CreateRunInput } from "./types.js"; + +function seedEnvironmentDedicated(suffix: string) { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: '{"hello":"world"}', + payloadType: "application/json", + traceContext: { trace: "ctx" }, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + runTags: [], + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: new Date("2024-01-01T00:00:00.000Z"), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +function makeStore(prisma: RunOpsPrismaClient) { + return new PostgresRunStore({ + prisma: prisma as never, + readOnlyPrisma: prisma as never, + schemaVariant: "dedicated", + }); +} + +// Wraps `taskRun.findMany` with a REAL Prisma Client Extension query hook: every call's `args` is +// captured, then delegated unchanged to `query(args)` so the DB still executes it. This is +// observation, not a mock — the extension can't fabricate a result, only see what's sent. +function withCapturedTaskRunFindManyArgs(real: RunOpsPrismaClient) { + const capturedArgs: Record[] = []; + const extended = (real as unknown as { $extends: (config: unknown) => unknown }).$extends({ + query: { + taskRun: { + findMany({ + args, + query, + }: { + args: Record; + query: (args: unknown) => Promise; + }) { + capturedArgs.push(args); + return query(args); + }, + }, + }, + }); + return { client: extended as RunOpsPrismaClient, capturedArgs }; +} + +describe("PostgresRunStore dedicated relation hydrators — target findMany select pushdown", () => { + heteroRunOpsPostgresTest( + "hydrateConnectedRuns narrows the target taskRun.findMany to the caller's select instead of fetching the full row", + async ({ prisma17 }) => { + const seedStore = makeStore(prisma17); + const env = seedEnvironmentDedicated("select_pushdown"); + + const runId = "run_select_pushdown"; + const waitpointId = "wp_select_pushdown"; + + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: `${waitpointId}_friendly`, + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + await seedStore.createRun( + buildCreateRunInput({ + runId, + friendlyId: `${runId}_friendly`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + + // Seeds BOTH the TaskRunWaitpoint edge and the WaitpointRunConnection join row that + // `connectedRuns` reads from. + await seedStore.blockRunWithWaitpointEdges({ + runId, + waitpointIds: [waitpointId], + projectId: env.project.id, + }); + + const { client: capturingClient, capturedArgs } = withCapturedTaskRunFindManyArgs(prisma17); + const readStore = makeStore(capturingClient); + + const waitpoint = (await readStore.findWaitpoint({ + where: { id: waitpointId }, + include: { connectedRuns: { select: { friendlyId: true } } }, + })) as { connectedRuns: { friendlyId: string }[] } | null; + + expect(capturedArgs).toHaveLength(1); + const targetArgs = capturedArgs[0] as { where: unknown; select?: Record }; + + // THE FIX: the target `findMany` carries a `select` narrowed to the caller's projection (plus + // the `id` the hydrator keys off), never a bare where-only (= full row) query. + expect(targetArgs.select).toBeDefined(); + expect(targetArgs.select).toMatchObject({ friendlyId: true, id: true }); + expect(targetArgs.select?.payload).toBeUndefined(); + expect(targetArgs.select?.output).toBeUndefined(); + expect(targetArgs.select?.context).toBeUndefined(); + + // `applyProjection` still trims correctly: exactly `{ friendlyId }`, no `id` leak, no wide + // columns — proving the pushdown didn't change the caller-visible result. + expect(waitpoint?.connectedRuns).toHaveLength(1); + expect(waitpoint?.connectedRuns[0]).toEqual({ friendlyId: `${runId}_friendly` }); + } + ); + + heteroRunOpsPostgresTest( + "hydrateEdgeTaskRun (batchHydrateEdgeTarget) narrows the target taskRun.findMany to the caller's select instead of fetching the full row", + async ({ prisma17 }) => { + const seedStore = makeStore(prisma17); + const env = seedEnvironmentDedicated("edge_pushdown"); + + const runId = "run_edge_pushdown"; + const waitpointId = "wp_edge_pushdown"; + + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: `${waitpointId}_friendly`, + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + await seedStore.createRun( + buildCreateRunInput({ + runId, + friendlyId: `${runId}_friendly`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + + await seedStore.blockRunWithWaitpointEdges({ + runId, + waitpointIds: [waitpointId], + projectId: env.project.id, + }); + + const { client: capturingClient, capturedArgs } = withCapturedTaskRunFindManyArgs(prisma17); + const readStore = makeStore(capturingClient); + + // TaskRunWaitpoint's `taskRun` relation is the dedicated `hydrateEdgeTaskRun` -> + // `batchHydrateEdgeTarget` path (the scalar-FK-on-the-parent variant). + const edges = (await readStore.findManyTaskRunWaitpoints({ + where: { waitpointId }, + select: { taskRun: { select: { friendlyId: true } } }, + })) as { taskRun: { friendlyId: string } | null }[]; + + expect(capturedArgs).toHaveLength(1); + const targetArgs = capturedArgs[0] as { where: unknown; select?: Record }; + + expect(targetArgs.select).toBeDefined(); + expect(targetArgs.select).toMatchObject({ friendlyId: true, id: true }); + expect(targetArgs.select?.payload).toBeUndefined(); + expect(targetArgs.select?.output).toBeUndefined(); + expect(targetArgs.select?.context).toBeUndefined(); + + expect(edges).toHaveLength(1); + expect(edges[0].taskRun).toEqual({ friendlyId: `${runId}_friendly` }); + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.sharedHydratedTarget.test.ts b/internal-packages/run-store/src/PostgresRunStore.sharedHydratedTarget.test.ts new file mode 100644 index 0000000000..482f532f5d --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.sharedHydratedTarget.test.ts @@ -0,0 +1,139 @@ +// RED→GREEN: bare-projection dedicated-relation hydration (e.g. `connectedRuns: true`) shares ONE +// target object reference across every parent bucket that links it, so an in-place mutation on one +// parent's hydrated target aliases into another's. Two waitpoints connecting to the SAME run should +// come back with DISTINCT `TaskRun` objects, not the same instance. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import type { CreateRunInput } from "./types.js"; + +function seedEnvironmentDedicated(suffix: string) { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: '{"hello":"world"}', + payloadType: "application/json", + traceContext: { trace: "ctx" }, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + runTags: [], + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: new Date("2024-01-01T00:00:00.000Z"), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +function makeStore(prisma: RunOpsPrismaClient) { + return new PostgresRunStore({ + prisma: prisma as never, + readOnlyPrisma: prisma as never, + schemaVariant: "dedicated", + }); +} + +describe("PostgresRunStore dedicated relation hydrators — shared target identity (bare projection)", () => { + heteroRunOpsPostgresTest( + "findManyWaitpoints hydrates bare `connectedRuns: true` for two waitpoints sharing one run as DISTINCT objects", + async ({ prisma17 }) => { + const store = makeStore(prisma17); + const env = seedEnvironmentDedicated("shared_target"); + + const runId = "run_shared_target"; + const waitpointIds = ["wp_shared_1", "wp_shared_2"]; + + await store.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_shared_target_friendly", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + + for (const waitpointId of waitpointIds) { + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: `${waitpointId}_friendly`, + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + } + + // Both waitpoints connect to the SAME run: the target-row `findMany` in + // `batchHydrateJoinRelation` returns ONE row for the run, shared across both parents' buckets. + await store.blockRunWithWaitpointEdges({ + runId, + waitpointIds, + projectId: env.project.id, + }); + + // Bare projection (`connectedRuns: true`, no sub-select) — the buggy path: `applyProjection` + // returns the row UNCHANGED instead of a fresh object. + const rows = (await store.findManyWaitpoints({ + where: { id: { in: waitpointIds } }, + include: { connectedRuns: true }, + })) as { id: string; connectedRuns: { id: string; friendlyId: string }[] }[]; + + expect(rows).toHaveLength(2); + const byWaitpointId = new Map(rows.map((r) => [r.id, r.connectedRuns])); + const target1 = byWaitpointId.get("wp_shared_1")?.[0]; + const target2 = byWaitpointId.get("wp_shared_2")?.[0]; + + expect(target1).toBeDefined(); + expect(target2).toBeDefined(); + expect(target1!.id).toBe(runId); + expect(target2!.id).toBe(runId); + + // THE BUG: both hydrated targets are the SAME object reference. + expect(target1).not.toBe(target2); + + // Mutating a top-level field on one parent's hydrated target must NOT change the other's. + (target1 as { friendlyId: string }).friendlyId = "mutated"; + expect(target2!.friendlyId).not.toBe("mutated"); + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 5d8ea6c8df..fa8df51f28 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -89,6 +89,10 @@ export interface RunOpsTransactionalClient extends RunOpsCapableClient { */ export type RunStoreSchemaVariant = "legacy" | "dedicated"; +// Mirrors the webapp's `CONNECTED_RUNS_DISPLAY_LIMIT` +// (apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts) — keep the values in sync. +export const CONNECTED_RUNS_LIMIT = 5; + export type PostgresRunStoreOptions = { prisma: RunOpsCapableClient; readOnlyPrisma: RunOpsCapableClient; @@ -99,13 +103,15 @@ export type PostgresRunStoreOptions = { // A caller sub-select for a relation: `{ select?, include? }` or `true` for a bare `key: true`. type SubProjection = { select?: any; include?: any } | true | undefined; -// Hydrates one dedicated-schema relation for a single parent row, honoring the caller's sub-projection. +// Hydrates one dedicated-schema relation for a WHOLE batch of parent rows in one grouped pass: +// one query for the join/target rows spanning every parent id, never one per parent. Returns a +// Map keyed by parent `id` to the already-defaulted (null / []) hydrated value. type DedicatedRelationHydrator = ( client: RunOpsCapableClient, - parent: Record, + parents: Record[], projection: { select?: any; include?: any } | undefined, store: PostgresRunStore -) => Promise; +) => Promise>; // The dedicated-schema relation keys (with hydrators) for a single Prisma model. type DedicatedRelationSpec = Record; @@ -119,13 +125,23 @@ function projectionOf(sub: SubProjection): { select?: any; include?: any } | und } // Apply a caller sub-projection to a hydrated row (or array) so only requested fields remain. +// +// Bare-projection path (no `select`): return a SHALLOW CLONE, not the row itself, so every parent +// bucket that links the same target gets a distinct top-level object — two parents sharing one +// target (e.g. two waitpoints connected to the same run) must not alias through a shared reference. +// This only protects top-level mutation; a deep in-place mutation of a nested field would still +// alias, which matches the realistic redaction/patch cases and avoids a costly deep clone on hot +// reads. function applyProjection | null>( row: T, projection: { select?: any; include?: any } | undefined ): T { - if (!row || !projection?.select) { + if (!row) { return row; } + if (!projection?.select) { + return { ...row } as T; + } const keys = Object.keys(projection.select).filter((k) => projection.select[k]); const out: Record = {}; for (const k of keys) { @@ -175,153 +191,247 @@ function stripDedicatedRelations( return { stripped: args, requested }; } -// --- per-model dedicated-schema relation hydrators --- +// --- per-model dedicated-schema relation hydrators (batched across the WHOLE parent array) --- -// Waitpoint where completedByTaskRunId = run.id (the @unique scalar back-pointer); at most one. -const hydrateAssociatedWaitpoint: DedicatedRelationHydrator = async ( - client, - parent, - projection -) => { - const wp = (await client.waitpoint.findFirst({ - where: { completedByTaskRunId: parent.id as string }, - })) as Record | null; - return applyProjection(wp, projection); -}; +// Narrows a hydrator's target `findMany` to the caller's `select` (avoids fetching the wide +// TOASTed columns just to strip them in `applyProjection`); a bare/`include` projection stays a +// full-row fetch. `keepKeys` are the column(s) the hydrator's Map is keyed on. +function targetFindManyArgs( + where: unknown, + projection: { select?: any; include?: any } | undefined, + keepKeys: string[] +): { where: unknown; select?: Record } { + if (projection?.select) { + const select: Record = { ...projection.select }; + for (const key of keepKeys) { + select[key] = true; + } + return { where, select }; + } + return { where }; +} -// Display connections for a run: WaitpointRunConnection → Waitpoint rows. -const hydrateConnectedWaitpoints: DedicatedRelationHydrator = async ( - client, - parent, - projection -) => { - const join = client.waitpointRunConnection; - if (!join) { - return []; +// Generic to-many relation reached via an explicit join model: one grouped query for the join rows +// spanning every parent id, then one grouped query for the distinct target rows, then an in-memory +// (DB-free) assembly per parent. `joinParentField`/`joinTargetField` name the join row's two FK +// columns; `targetDelegate` is the model the join points at. +async function batchHydrateJoinRelation( + join: RunOpsDelegate<"findMany"> | undefined, + targetDelegate: RunOpsDelegate<"findMany">, + parentIds: string[], + joinParentField: string, + joinTargetField: string, + projection: { select?: any; include?: any } | undefined +): Promise> { + const byParent = new Map(parentIds.map((id) => [id, []])); + if (!join || parentIds.length === 0) { + return byParent; } const links = (await join.findMany({ - where: { taskRunId: parent.id as string }, - select: { waitpointId: true }, - })) as { waitpointId: string }[]; + where: { [joinParentField]: { in: parentIds } }, + select: { [joinParentField]: true, [joinTargetField]: true }, + })) as Record[]; if (links.length === 0) { - return []; + return byParent; + } + const targetIds = [...new Set(links.map((l) => l[joinTargetField]))]; + const rows = (await targetDelegate.findMany( + targetFindManyArgs({ id: { in: targetIds } }, projection, ["id"]) + )) as Record[]; + const byTargetId = new Map(rows.map((r) => [r.id as string, r])); + for (const link of links) { + const target = byTargetId.get(link[joinTargetField]); + const bucket = byParent.get(link[joinParentField]); + if (target && bucket) { + bucket.push(applyProjection(target, projection)); + } } - const rows = (await client.waitpoint.findMany({ - where: { id: { in: links.map((l) => l.waitpointId) } }, - })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); -}; + return byParent; +} -// Completed waitpoints for a snapshot: CompletedWaitpoint join → Waitpoint rows. -const hydrateCompletedWaitpoints: DedicatedRelationHydrator = async ( +// Waitpoint where completedByTaskRunId = run.id (the @unique scalar back-pointer); at most one. +const hydrateAssociatedWaitpoint: DedicatedRelationHydrator = async ( client, - parent, + parents, projection ) => { - const join = client.completedWaitpoint; - if (!join) { - return []; - } - const links = (await join.findMany({ - where: { snapshotId: parent.id as string }, - select: { waitpointId: true }, - })) as { waitpointId: string }[]; - if (links.length === 0) { - return []; + const parentIds = parents.map((p) => p.id as string); + const byParent = new Map(parentIds.map((id) => [id, null])); + if (parentIds.length === 0) { + return byParent; + } + const rows = (await client.waitpoint.findMany( + targetFindManyArgs({ completedByTaskRunId: { in: parentIds } }, projection, [ + "completedByTaskRunId", + ]) + )) as Record[]; + for (const row of rows) { + const runId = row.completedByTaskRunId as string | undefined; + if (runId && byParent.has(runId)) { + byParent.set(runId, applyProjection(row, projection)); + } } - const rows = (await client.waitpoint.findMany({ - where: { id: { in: links.map((l) => l.waitpointId) } }, - })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); + return byParent; }; +// Display connections for a run: WaitpointRunConnection → Waitpoint rows. +const hydrateConnectedWaitpoints: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateJoinRelation( + client.waitpointRunConnection, + client.waitpoint, + parents.map((p) => p.id as string), + "taskRunId", + "waitpointId", + projection + ); + +// Completed waitpoints for a snapshot: CompletedWaitpoint join → Waitpoint rows. +const hydrateCompletedWaitpoints: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateJoinRelation( + client.completedWaitpoint, + client.waitpoint, + parents.map((p) => p.id as string), + "snapshotId", + "waitpointId", + projection + ); + // Runs a waitpoint is blocking: TaskRunWaitpoint rows keyed by waitpointId. A nested `taskRun` // select (the run-engine's getWaitpoint shape) is resolved from the scalar TaskRunWaitpoint.taskRunId. -const hydrateBlockingTaskRuns: DedicatedRelationHydrator = async (client, parent, projection) => { +const hydrateBlockingTaskRuns: DedicatedRelationHydrator = async (client, parents, projection) => { + const parentIds = parents.map((p) => p.id as string); + const byParent = new Map(parentIds.map((id) => [id, []])); + if (parentIds.length === 0) { + return byParent; + } const edges = (await client.taskRunWaitpoint.findMany({ - where: { waitpointId: parent.id as string }, + where: { waitpointId: { in: parentIds } }, })) as Record[]; const nestedTaskRun = projection?.select?.taskRun; - if (!nestedTaskRun) { - return edges; - } - const runProjection = projectionOf(nestedTaskRun as SubProjection); - return Promise.all( - edges.map(async (edge) => { - const run = (await client.taskRun.findFirst({ - where: { id: edge.taskRunId as string }, - })) as Record | null; - return { ...edge, taskRun: applyProjection(run, runProjection) }; - }) - ); + const runProjection = nestedTaskRun ? projectionOf(nestedTaskRun as SubProjection) : undefined; + let byRunId = new Map>(); + if (nestedTaskRun) { + const runIds = [...new Set(edges.map((e) => e.taskRunId as string))]; + const runs = ( + runIds.length > 0 + ? await client.taskRun.findMany( + targetFindManyArgs({ id: { in: runIds } }, runProjection, ["id"]) + ) + : [] + ) as Record[]; + byRunId = new Map(runs.map((r) => [r.id as string, r])); + } + for (const edge of edges) { + const bucket = byParent.get(edge.waitpointId as string); + if (!bucket) continue; + bucket.push( + nestedTaskRun + ? { + ...edge, + taskRun: applyProjection(byRunId.get(edge.taskRunId as string) ?? null, runProjection), + } + : edge + ); + } + return byParent; }; -// Display connections for a waitpoint: WaitpointRunConnection → TaskRun rows. -const hydrateConnectedRuns: DedicatedRelationHydrator = async (client, parent, projection) => { - const join = client.waitpointRunConnection; - if (!join) { - return []; - } - const links = (await join.findMany({ - where: { waitpointId: parent.id as string }, - select: { taskRunId: true }, - })) as { taskRunId: string }[]; +// Display connections for a waitpoint: WaitpointRunConnection → TaskRun rows. Bounded per parent to +// CONNECTED_RUNS_LIMIT via a window function + existence-JOIN to TaskRun, mirroring the id-list +// helper findWaitpointConnectedRunIds: a dangling (run-less) connection row never occupies a LIMIT +// slot, and a heavily-fanned-in waitpoint never hydrates an unbounded connectedRuns list. This is +// the DISPLAY relation only — functional blocking reads (hydrateBlockingTaskRuns) stay uncapped. +const hydrateConnectedRuns: DedicatedRelationHydrator = async (client, parents, projection) => { + const parentIds = parents.map((p) => p.id as string); + const byParent = new Map(parentIds.map((id) => [id, []])); + if (parentIds.length === 0) { + return byParent; + } + // One grouped query for the bounded edges across every parent: ROW_NUMBER partitioned per + // waitpoint keeps at most CONNECTED_RUNS_LIMIT rows per parent (uses @@index([waitpointId])). + const links = (await client.$queryRaw` + SELECT ranked."waitpointId" AS "waitpointId", ranked."taskRunId" AS "taskRunId" + FROM ( + SELECT c."waitpointId", c."taskRunId", + ROW_NUMBER() OVER (PARTITION BY c."waitpointId" ORDER BY c."id") AS rn + FROM "WaitpointRunConnection" c + JOIN "TaskRun" t ON t."id" = c."taskRunId" + WHERE c."waitpointId" = ANY(${parentIds}::text[]) + ) ranked + WHERE ranked.rn <= ${CONNECTED_RUNS_LIMIT} + `) as { waitpointId: string; taskRunId: string }[]; if (links.length === 0) { - return []; + return byParent; + } + const targetIds = [...new Set(links.map((l) => l.taskRunId))]; + const rows = (await client.taskRun.findMany( + targetFindManyArgs({ id: { in: targetIds } }, projection, ["id"]) + )) as Record[]; + const byTargetId = new Map(rows.map((r) => [r.id as string, r])); + for (const link of links) { + const target = byTargetId.get(link.taskRunId); + const bucket = byParent.get(link.waitpointId); + if (target && bucket) { + bucket.push(applyProjection(target, projection)); + } } - const rows = (await client.taskRun.findMany({ - where: { id: { in: links.map((l) => l.taskRunId) } }, - })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); + return byParent; }; // Snapshots that completed a waitpoint: CompletedWaitpoint join → TaskRunExecutionSnapshot rows. const hydrateCompletedExecutionSnapshots: DedicatedRelationHydrator = async ( client, - parent, + parents, projection -) => { - const join = client.completedWaitpoint; - if (!join) { - return []; - } - const links = (await join.findMany({ - where: { waitpointId: parent.id as string }, - select: { snapshotId: true }, - })) as { snapshotId: string }[]; - if (links.length === 0) { - return []; - } - const rows = (await client.taskRunExecutionSnapshot.findMany({ - where: { id: { in: links.map((l) => l.snapshotId) } }, - })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); -}; - -// The waitpoint a block edge points at, resolved from the edge's scalar `waitpointId`. The edge's -// own client only finds a co-resident token; the router re-resolves cross-DB. -const hydrateEdgeWaitpoint: DedicatedRelationHydrator = async (client, parent, projection) => { - const waitpointId = parent.waitpointId as string | undefined; - if (!waitpointId) { - return null; - } - const wp = (await client.waitpoint.findFirst({ - where: { id: waitpointId }, - })) as Record | null; - return applyProjection(wp, projection); -}; +) => + batchHydrateJoinRelation( + client.completedWaitpoint, + client.taskRunExecutionSnapshot, + parents.map((p) => p.id as string), + "waitpointId", + "snapshotId", + projection + ); -// The run a block edge belongs to, resolved from the edge's scalar `taskRunId`. -const hydrateEdgeTaskRun: DedicatedRelationHydrator = async (client, parent, projection) => { - const taskRunId = parent.taskRunId as string | undefined; - if (!taskRunId) { - return null; +// The waitpoint each block edge points at, resolved from its scalar `waitpointId`. The edge's own +// client only finds a co-resident token; the router re-resolves cross-DB. +const hydrateEdgeWaitpoint: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateEdgeTarget(client.waitpoint, parents, "waitpointId", projection); + +// The run each block edge belongs to, resolved from its scalar `taskRunId`. +const hydrateEdgeTaskRun: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateEdgeTarget(client.taskRun, parents, "taskRunId", projection); + +// Generic to-one relation reached via a scalar FK ON the parent itself (not a join model): one +// grouped target query for every distinct FK value across the batch. +async function batchHydrateEdgeTarget( + targetDelegate: RunOpsDelegate<"findMany">, + parents: Record[], + fkField: string, + projection: { select?: any; include?: any } | undefined +): Promise> { + const byParent = new Map(); + const targetIds: string[] = []; + for (const p of parents) { + byParent.set(p.id as string, null); + const fk = p[fkField] as string | undefined; + if (fk) targetIds.push(fk); + } + if (targetIds.length === 0) { + return byParent; + } + const rows = (await targetDelegate.findMany( + targetFindManyArgs({ id: { in: [...new Set(targetIds)] } }, projection, ["id"]) + )) as Record[]; + const byTargetId = new Map(rows.map((r) => [r.id as string, r])); + for (const p of parents) { + const fk = p[fkField] as string | undefined; + if (fk) { + byParent.set(p.id as string, applyProjection(byTargetId.get(fk) ?? null, projection)); + } } - const run = (await client.taskRun.findFirst({ - where: { id: taskRunId }, - })) as Record | null; - return applyProjection(run, projection); -}; + return byParent; +} const TASK_RUN_DEDICATED: DedicatedRelationSpec = { associatedWaitpoint: hydrateAssociatedWaitpoint, @@ -1444,15 +1554,72 @@ export class PostgresRunStore implements RunStore { cursor, ...stripped, })) as Record[]; + await this.#hydrateDedicatedRelations( + prisma as RunOpsCapableClient, + rows, + requested, + TASK_RUN_DEDICATED + ); + return rows; + } + + // Grouped replacement for `Promise.all(ids.map(id => findRun(id)))`: a thin wrapper over + // `findRuns` (so it inherits dedicated-relation hydration + read-your-writes client routing), + // bounding the whole id batch into one round trip instead of one per id. + findRunsByIds( + ids: string[], + args: { select: S }, + client?: ReadClient + ): Promise>>; + findRunsByIds( + ids: string[], + args: { include: I }, + client?: ReadClient + ): Promise>>; + findRunsByIds(ids: string[], client?: ReadClient): Promise>; + async findRunsByIds( + ids: string[], + argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | ReadClient, + _client?: ReadClient + ): Promise> { + if (ids.length === 0) { + return new Map(); + } + const hasSelectOrInclude = + argsOrClient != null && + typeof argsOrClient === "object" && + ("select" in argsOrClient || "include" in argsOrClient); + const args = hasSelectOrInclude + ? (argsOrClient as { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }) + : undefined; + // Slot recovery mirrors `findRuns`'s overloads: when `argsOrClient` isn't a + // `{ select | include }` object it may itself BE the client (2-arg call) or be undefined + // with the client in the 3rd slot (an explicit `(ids, undefined, client)` call). + const client = + args === undefined ? ((argsOrClient as ReadClient | undefined) ?? _client) : _client; + // Force `id` into the projection so the map can key off it, even when the caller's select + // omits it — `findRuns` would otherwise strip it back out as an added-for-merge-only field. + const projected = args?.select + ? { select: { ...args.select, id: true } } + : args?.include + ? { include: args.include } + : {}; + const rows = (await this.findRuns( + { where: { id: { in: ids } }, ...projected } as Parameters[0], + client + )) as Record[]; + const byId = new Map(); + // Strip the id we force-injected for map keying when the caller's select did not ask for it, + // so returned values match the declared payload type and never leak an unrequested id. + const stripInjectedId = args?.select != null && !("id" in args.select); for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - TASK_RUN_DEDICATED - ); + const key = row.id as string; + if (stripInjectedId) { + delete row.id; + } + byId.set(key, row); } - return rows; + return byId; } // --- run-ops persistence --- @@ -1563,14 +1730,12 @@ export class PostgresRunStore implements RunStore { cursor, ...stripped, })) as Record[]; - for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - SNAPSHOT_DEDICATED - ); - } + await this.#hydrateDedicatedRelations( + prisma as RunOpsCapableClient, + rows, + requested, + SNAPSHOT_DEDICATED + ); return rows as Prisma.TaskRunExecutionSnapshotGetPayload[]; } @@ -1710,23 +1875,34 @@ export class PostgresRunStore implements RunStore { // Reverse of `connectedRuns`: the run ids linked to a waitpoint. Co-resident with the RUN (the join // is written on the run's DB in blockRunWithWaitpointEdges), so the waitpoint's own store can MISS a // cross-DB run — the router fans this across BOTH DBs. + // Bounded to CONNECTED_RUNS_LIMIT via an existence-JOIN to TaskRun, mirroring the webapp's + // `#connectedRunIdsOn`: a dangling connection row (dedicated schema: FK-free `taskRunId`) can + // never occupy a LIMIT slot ahead of a real run, and a heavily-fanned-in waitpoint can never emit + // an unbounded id list. async findWaitpointConnectedRunIds(waitpointId: string, client?: ReadClient): Promise { const prisma = client ?? this.readOnlyPrisma; const joinDelegate = (prisma as RunOpsCapableClient).waitpointRunConnection; if (this.schemaVariant === "dedicated" && joinDelegate) { - const links = (await joinDelegate.findMany({ - where: { waitpointId }, - select: { taskRunId: true }, - })) as { taskRunId: string }[]; - return links.map((l) => l.taskRunId); + const rows = await prisma.$queryRaw<{ taskRunId: string }[]>` + SELECT c."taskRunId" AS "taskRunId" + FROM "WaitpointRunConnection" c + JOIN "TaskRun" t ON t."id" = c."taskRunId" + WHERE c."waitpointId" = ${waitpointId} + LIMIT ${CONNECTED_RUNS_LIMIT} + `; + return rows.map((row) => row.taskRunId); } // Legacy implicit M2M `_WaitpointRunConnections`: A = TaskRun.id, B = Waitpoint.id (alphabetical). - const result = await prisma.$queryRaw<{ A: string }[]>` - SELECT "A" FROM "_WaitpointRunConnections" WHERE "B" = ${waitpointId} + const rows = await prisma.$queryRaw<{ A: string }[]>` + SELECT c."A" AS "A" + FROM "_WaitpointRunConnections" c + JOIN "TaskRun" t ON t."id" = c."A" + WHERE c."B" = ${waitpointId} + LIMIT ${CONNECTED_RUNS_LIMIT} `; - return result.map((r) => r.A); + return rows.map((row) => row.A); } // Reverse of `completedExecutionSnapshots`: the snapshot ids that completed a waitpoint. The join is @@ -1949,14 +2125,12 @@ export class PostgresRunStore implements RunStore { cursor, ...stripped, })) as Record[]; - for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - WAITPOINT_DEDICATED - ); - } + await this.#hydrateDedicatedRelations( + prisma as RunOpsCapableClient, + rows, + requested, + WAITPOINT_DEDICATED + ); return rows as Prisma.WaitpointGetPayload[]; } @@ -2019,14 +2193,12 @@ export class PostgresRunStore implements RunStore { cursor, ...stripped, })) as Record[]; - for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - TASK_RUN_WAITPOINT_DEDICATED - ); - } + await this.#hydrateDedicatedRelations( + prisma as RunOpsCapableClient, + rows, + requested, + TASK_RUN_WAITPOINT_DEDICATED + ); return rows as Prisma.TaskRunWaitpointGetPayload[]; } @@ -2267,14 +2439,15 @@ export class PostgresRunStore implements RunStore { if (!row) { return row; } - await this.#hydrateDedicatedRelations(client, row, requested, spec); + await this.#hydrateDedicatedRelations(client, [row], requested, spec); return row; } - // Hydrate each requested dedicated-schema relation key onto `row` in place, honoring the caller's sub-select. + // Hydrate each requested dedicated-schema relation key onto EVERY row in `rows` in ONE grouped + // pass per key (never one query per row), honoring the caller's sub-select. async #hydrateDedicatedRelations( client: RunOpsCapableClient, - row: Record, + rows: Record[], requested: Record, spec: DedicatedRelationSpec ): Promise { @@ -2284,7 +2457,10 @@ export class PostgresRunStore implements RunStore { continue; } const subArgs = requested[key]; - row[key] = await hydrator(client, row, projectionOf(subArgs), this); + const byParentId = await hydrator(client, rows, projectionOf(subArgs), this); + for (const row of rows) { + row[key] = byParentId.get(row.id as string); + } } } diff --git a/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts new file mode 100644 index 0000000000..4b2230bec0 --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts @@ -0,0 +1,392 @@ +// RED→GREEN for the grouped-read primitive `RoutingRunStore.findRunsByIds`. +// +// Today, callers that need N runs by id do `Promise.all(ids.map(id => readThroughRun(id)))`, which +// fans out to `findRun` PER ID — each one either a direct `taskRun.findFirst` on the owning store, or +// (for an unclassifiable id) a NEW-then-LEGACY probe. For N ids that is O(N) round trips. +// +// `findRunsByIds` is the grouped replacement: it reuses the router's existing bounded id-set +// machinery (`#findRunsByIdSet` via `findRuns`), which queries NEW for the WHOLE id set in one +// `findMany`, then probes LEGACY in one more `findMany` only for the ids NEW missed. So for any N +// (spread across both stores), the call count is CONSTANT (one grouped `findMany` per store queried), +// never N `findFirst` calls. +// +// `heteroRunOpsPostgresTest` gives the REAL split topology: prisma17 = dedicated subset schema +// (#new), prisma14 = full legacy schema on a SEPARATE physical PG container (#legacy). The counting +// proxy wraps each REAL client's `taskRun` delegate to tally `findFirst`/`findMany` calls while +// delegating every call to the real client — the DB still runs the query; this is instrumentation, +// not a mock. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; + +type AnyClient = PrismaClient | RunOpsPrismaClient; + +// ownerEngine classifies a 26-char body with version "1" at index 25 (and a valid base32hex/region +// char elsewhere) as NEW; anything else (including any 25-char id, regardless of content) as LEGACY. +// The distinguishing prefix varies by `n` so multiple ids in one test stay unique. +const CUID_25 = (n: number) => `c${String(n).padStart(2, "0")}`.padEnd(25, "c"); +const NEW_ID_26 = (n: number) => `${String(n).padStart(2, "0")}${"k".repeat(22)}01`; + +type CallCounts = { findMany: number; findFirst: number }; + +// Wrap a REAL client's `taskRun` delegate to tally `findFirst`/`findMany` invocations, delegating +// every call unchanged to the real client (the DB still runs the query — pure instrumentation). +function countingClient(real: C): { client: C; counts: CallCounts } { + const counts: CallCounts = { findMany: 0, findFirst: 0 }; + const countingTaskRun = new Proxy((real as any).taskRun, { + get(target, prop) { + if (prop === "findMany" || prop === "findFirst") { + counts[prop as "findMany" | "findFirst"]++; + } + return (target as any)[prop]; + }, + }); + const client = new Proxy(real, { + get(target, prop) { + if (prop === "taskRun") { + return countingTaskRun; + } + return (target as any)[prop]; + }, + }) as C; + return { client, counts }; +} + +async function seedEnvironmentLegacy(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { organization, project, environment }; +} + +function seedEnvironmentDedicated(suffix: string) { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; +} + +function taskRunData(opts: { + id: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}) { + return { + id: opts.id, + engine: "V2" as const, + status: "PENDING" as const, + friendlyId: opts.friendlyId, + runtimeEnvironmentId: opts.runtimeEnvironmentId, + environmentType: "DEVELOPMENT" as const, + organizationId: opts.organizationId, + projectId: opts.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${opts.id}`, + spanId: `span_${opts.id}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }; +} + +describe("RoutingRunStore.findRunsByIds — grouped, residency-partitioned read of a run id set", () => { + heteroRunOpsPostgresTest( + "returns all N rows with a CONSTANT number of grouped findMany calls (never per-id findFirst)", + async ({ prisma14, prisma17 }) => { + const legacyCounting = countingClient(prisma14); + const newCounting = countingClient(prisma17); + + const legacyStore = new PostgresRunStore({ + prisma: legacyCounting.client, + readOnlyPrisma: legacyCounting.client, + schemaVariant: "legacy", + }); + const newStore = new PostgresRunStore({ + prisma: newCounting.client as never, + readOnlyPrisma: newCounting.client as never, + schemaVariant: "dedicated", + }); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + + // Seed 2 LEGACY-resident (cuid) runs + 2 NEW-resident (run-ops id) runs — ids spread across + // both stores. + const legEnv = await seedEnvironmentLegacy(prisma14, "grp_leg"); + const newEnv = seedEnvironmentDedicated("grp_new"); + + const legacyIds = [`run_${CUID_25(1)}`, `run_${CUID_25(2)}`]; + const newIds = [`run_${NEW_ID_26(1)}`, `run_${NEW_ID_26(2)}`]; + + for (const [i, id] of legacyIds.entries()) { + await prisma14.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_grp_leg_${i}`, + organizationId: legEnv.organization.id, + projectId: legEnv.project.id, + runtimeEnvironmentId: legEnv.environment.id, + }), + }); + } + for (const [i, id] of newIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_grp_new_${i}`, + organizationId: newEnv.organization.id, + projectId: newEnv.project.id, + runtimeEnvironmentId: newEnv.environment.id, + }), + }); + } + + const allIds = [...legacyIds, ...newIds]; + + // Reset counters after seeding (seeding uses `prisma.taskRun.create`, not find*). + legacyCounting.counts.findMany = 0; + legacyCounting.counts.findFirst = 0; + newCounting.counts.findMany = 0; + newCounting.counts.findFirst = 0; + + const result = await router.findRunsByIds(allIds, { + select: { friendlyId: true }, + }); + + // All N rows returned, keyed by id. + expect(result.size).toBe(4); + for (const id of allIds) { + expect(result.has(id)).toBe(true); + } + expect(result.get(legacyIds[0])?.friendlyId).toBe("run_grp_leg_0"); + expect(result.get(newIds[0])?.friendlyId).toBe("run_grp_new_0"); + + // The select omitted `id`; the id we force-inject for keying must not leak into the value. + expect("id" in (result.get(legacyIds[0]) as object)).toBe(false); + expect("id" in (result.get(newIds[0]) as object)).toBe(false); + + // GROUPED: exactly one findMany call per store queried (NEW always queried for the whole + // set; LEGACY queried once more for the misses) — never N per-id findFirst calls. + expect(newCounting.counts.findMany).toBe(1); + expect(legacyCounting.counts.findMany).toBe(1); + expect(newCounting.counts.findFirst).toBe(0); + expect(legacyCounting.counts.findFirst).toBe(0); + } + ); +}); + +describe("RoutingRunStore.findManyTaskRunWaitpoints — grouped taskRun hydration", () => { + heteroRunOpsPostgresTest( + "hydrates N edges' `taskRun` with a CONSTANT number of grouped findMany calls (never per-edge findFirst)", + async ({ prisma14, prisma17 }) => { + const legacyCounting = countingClient(prisma14); + const newCounting = countingClient(prisma17); + + const legacyStore = new PostgresRunStore({ + prisma: legacyCounting.client, + readOnlyPrisma: legacyCounting.client, + schemaVariant: "legacy", + }); + const newStore = new PostgresRunStore({ + prisma: newCounting.client as never, + readOnlyPrisma: newCounting.client as never, + schemaVariant: "dedicated", + }); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + + const legEnv = await seedEnvironmentLegacy(prisma14, "edge_leg"); + const newEnv = seedEnvironmentDedicated("edge_new"); + + const legacyRunIds = [`run_${CUID_25(11)}`, `run_${CUID_25(12)}`]; + const newRunIds = [`run_${NEW_ID_26(11)}`, `run_${NEW_ID_26(12)}`]; + for (const [i, id] of legacyRunIds.entries()) { + await prisma14.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_edge_leg_${i}`, + organizationId: legEnv.organization.id, + projectId: legEnv.project.id, + runtimeEnvironmentId: legEnv.environment.id, + }), + }); + } + for (const [i, id] of newRunIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_edge_new_${i}`, + organizationId: newEnv.organization.id, + projectId: newEnv.project.id, + runtimeEnvironmentId: newEnv.environment.id, + }), + }); + } + + // Block edges are FK-free on the dedicated (#new) subset schema, so they can point at a run on + // EITHER DB — create one edge per run, all on #new. + const allRunIds = [...legacyRunIds, ...newRunIds]; + const edgeIds: string[] = []; + for (const runId of allRunIds) { + const edge = await prisma17.taskRunWaitpoint.create({ + data: { + taskRunId: runId, + waitpointId: `wp_${runId}`, + projectId: newEnv.project.id, + }, + }); + edgeIds.push(edge.id); + } + + legacyCounting.counts.findMany = 0; + legacyCounting.counts.findFirst = 0; + newCounting.counts.findMany = 0; + newCounting.counts.findFirst = 0; + + const rows = await router.findManyTaskRunWaitpoints({ + where: { id: { in: edgeIds } }, + select: { + id: true, + taskRunId: true, + taskRun: { select: { id: true, friendlyId: true } }, + }, + }); + + expect(rows).toHaveLength(4); + const byRunId = new Map(rows.map((r) => [r.taskRunId, r.taskRun])); + expect((byRunId.get(legacyRunIds[0]) as any)?.friendlyId).toBe("run_edge_leg_0"); + expect((byRunId.get(newRunIds[0]) as any)?.friendlyId).toBe("run_edge_new_0"); + + // GROUPED: one findMany per store queried for the WHOLE edge set, never one findFirst per edge. + expect(newCounting.counts.findMany).toBe(1); + expect(legacyCounting.counts.findMany).toBe(1); + expect(newCounting.counts.findFirst).toBe(0); + expect(legacyCounting.counts.findFirst).toBe(0); + } + ); +}); + +describe("RoutingRunStore.findWaitpoint connectedRuns — grouped, order-preserving hydration", () => { + heteroRunOpsPostgresTest( + "hydrates N connectedRuns with a CONSTANT number of grouped findMany calls, in the join's own order", + async ({ prisma14, prisma17 }) => { + const legacyCounting = countingClient(prisma14); + const newCounting = countingClient(prisma17); + + const legacyStore = new PostgresRunStore({ + prisma: legacyCounting.client, + readOnlyPrisma: legacyCounting.client, + schemaVariant: "legacy", + }); + const newStore = new PostgresRunStore({ + prisma: newCounting.client as never, + readOnlyPrisma: newCounting.client as never, + schemaVariant: "dedicated", + }); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + + const legEnv = await seedEnvironmentLegacy(prisma14, "conn_leg"); + const newEnv = seedEnvironmentDedicated("conn_new"); + const waitpointId = `waitpoint_${CUID_25(20)}`; + await prisma14.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: "wp_conn", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: legEnv.project.id, + environmentId: legEnv.environment.id, + }, + }); + + const legacyRunIds = [`run_${CUID_25(21)}`, `run_${CUID_25(22)}`]; + const newRunIds = [`run_${NEW_ID_26(21)}`, `run_${NEW_ID_26(22)}`]; + for (const [i, id] of legacyRunIds.entries()) { + await prisma14.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_conn_leg_${i}`, + organizationId: legEnv.organization.id, + projectId: legEnv.project.id, + runtimeEnvironmentId: legEnv.environment.id, + }), + }); + await router.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: legEnv.project.id, + }); + } + for (const [i, id] of newRunIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_conn_new_${i}`, + organizationId: newEnv.organization.id, + projectId: newEnv.project.id, + runtimeEnvironmentId: newEnv.environment.id, + }), + }); + await router.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: newEnv.project.id, + }); + } + + // The order the connection join itself returns, independent of any assumption about Postgres's + // internal row order — the fix must preserve exactly THIS order, not resort by residency/leg. + const expectedOrder = await router.findWaitpointConnectedRunIds(waitpointId); + expect(expectedOrder).toHaveLength(4); + + legacyCounting.counts.findMany = 0; + legacyCounting.counts.findFirst = 0; + newCounting.counts.findMany = 0; + newCounting.counts.findFirst = 0; + + const waitpoint = (await router.findWaitpoint({ + where: { id: waitpointId }, + include: { connectedRuns: { select: { id: true, friendlyId: true } } }, + })) as { connectedRuns: { id: string; friendlyId: string }[] } | null; + + const connected = waitpoint?.connectedRuns ?? []; + expect(connected).toHaveLength(4); + expect(connected.map((r) => r.id)).toEqual(expectedOrder); + + // GROUPED: one findMany per store for the WHOLE connected-run set, never one findFirst per run. + expect(newCounting.counts.findMany).toBe(1); + expect(legacyCounting.counts.findMany).toBe(1); + expect(newCounting.counts.findFirst).toBe(0); + expect(legacyCounting.counts.findFirst).toBe(0); + } + ); +}); diff --git a/internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts b/internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts new file mode 100644 index 0000000000..8897c29566 --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts @@ -0,0 +1,296 @@ +// RED→GREEN lock for RoutingRunStore.findRun's ON-MISS FAN-OUT for a CLASSIFIABLE id. +// +// THE BUG: findRun for a classifiable id routed to the SINGLE owning store (by id-shape +// classification) and returned whatever it gave — no on-miss fallback. When a run's PHYSICAL +// residency does not match its id-shape classification (e.g. a pre-#4154 27-char base62 run that +// lives on the NEW store but now classifies LEGACY), findRun routed to the wrong store, missed, +// and returned null → a spurious 404 — even though the run is physically present on the OTHER DB +// (and runs.list surfaces it). The unclassifiable path already fanned out NEW→LEGACY; this makes +// the classifiable path equally robust. +// +// Uses the REAL two-physical-DB split (heteroRunOpsPostgresTest: prisma14 = full/legacy on PG14, +// prisma17 = dedicated run-ops subset on PG17). NEVER mocked. The residency/classification MISMATCH +// is simulated deterministically by injecting a custom `classify` fn into the RoutingRunStore +// constructor — the physical row is written to the NEW store while `classify` reports its id LEGACY. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import { ownerEngine, type Residency } from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; +import type { CreateRunInput, RunStore } from "./types.js"; + +// ownerEngine classifies by internal-id LENGTH/version char: 25 chars → cuid → LEGACY, +// a v1 body (version "1" at index 25) → run-ops id → NEW. +function cuidLegacy(seed: string): string { + return (seed + "c".repeat(25)).slice(0, 25); +} +function runOpsNew(seed: string): string { + return (seed.replace(/[^0-9a-v]/g, "0") + "k".repeat(24)).slice(0, 24) + "01"; +} + +function makeDedicatedStore(prisma17: RunOpsPrismaClient) { + return new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); +} + +function makeLegacyStore(prisma14: PrismaClient) { + return new PostgresRunStore({ + prisma: prisma14, + readOnlyPrisma: prisma14, + schemaVariant: "legacy", + }); +} + +// Wrap a real store so findRun/findRunOnPrimary calls are COUNTED while every method still delegates +// to the REAL PostgresRunStore (this is instrumentation, not a behavior mock — the underlying reads, +// writes, getters all run for real). Lets us assert the FAST PATH does not touch the other store. +function countingReads( + inner: RunStore, + counts: { findRun: number; findRunOnPrimary: number } +): RunStore { + return new Proxy(inner, { + get(target, prop) { + // Read via target[prop] so getters (e.g. primaryReadClient) run with `this` = the real store. + const value = (target as unknown as Record)[prop]; + if (typeof value !== "function") return value; + if (prop === "findRun" || prop === "findRunOnPrimary") { + return (...args: unknown[]) => { + counts[prop as "findRun" | "findRunOnPrimary"] += 1; + return (value as (...a: unknown[]) => unknown).apply(target, args); + }; + } + return (value as (...a: unknown[]) => unknown).bind(target); + }, + }) as unknown as RunStore; +} + +async function seedLegacyEnvironment(prisma14: PrismaClient, suffix: string) { + const organization = await prisma14.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma14.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma14.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + environmentId: environment.id, + }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +// Insert a TaskRun row DIRECTLY onto the NEW (dedicated) store, bypassing routing, so we can force a +// residency/classification MISMATCH: the row is physically on #new while `classify` calls its id LEGACY. +async function insertRunOnNewStore( + prisma17: RunOpsPrismaClient, + params: { + runId: string; + friendlyId: string; + environmentId: string; + organizationId: string; + projectId: string; + } +) { + await prisma17.taskRun.create({ + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.environmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }, + }); +} + +describe("RoutingRunStore.findRun — on-miss fan-out for a classifiable id (residency ≠ classification)", () => { + // ── THE BUG: a run physically on #new whose id classifies LEGACY must still be found ── + // Without the on-miss fallback, findRun routes to #legacy (per classification), misses, returns null. + heteroRunOpsPostgresTest( + "returns a #new-resident run whose id classifies LEGACY (owning-store miss → other-store fallback)", + async ({ prisma14, prisma17 }) => { + const env = await seedLegacyEnvironment(prisma14, "mm1"); + const newStore = makeDedicatedStore(prisma17); + const legacyStore = makeLegacyStore(prisma14); + + // A run-ops-shaped id (so ownerEngine would say NEW), but we FORCE classify → LEGACY to model + // a residency/classification mismatch: physically on #new, classified LEGACY. + const mismatchId = runOpsNew("mm1"); + const classify = (id: string): Residency => (id === mismatchId ? "LEGACY" : ownerEngine(id)); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore, classify }); + + await insertRunOnNewStore(prisma17, { + runId: mismatchId, + friendlyId: "run_mm1", + environmentId: env.environmentId, + organizationId: env.organizationId, + projectId: env.projectId, + }); + + // Physical residency sanity: on #new only. + expect(await prisma17.taskRun.findUnique({ where: { id: mismatchId } })).not.toBeNull(); + expect(await prisma14.taskRun.findUnique({ where: { id: mismatchId } })).toBeNull(); + + // classify → LEGACY routes to #legacy (miss); the fix falls back to #new and finds the run. + const byId = (await router.findRun({ id: mismatchId }, { select: { id: true } })) as Record< + string, + unknown + > | null; + expect(byId?.id).toBe(mismatchId); + + // Same on the read-your-writes primary variant (a caller-passed writer → findRunOnPrimary). + const byIdPrimary = (await router.findRun( + { id: mismatchId }, + { select: { id: true } }, + prisma14 + )) as Record | null; + expect(byIdPrimary?.id).toBe(mismatchId); + } + ); + + // ── FAST PATH: a run found in its CLASSIFIED store is a SINGLE read (no second-store probe) ── + heteroRunOpsPostgresTest( + "does NOT read the other store when the classified (owning) store hits", + async ({ prisma14, prisma17 }) => { + const env = await seedLegacyEnvironment(prisma14, "mm2"); + + // NEW-resident run-ops-id run: owning store = #new. Wrap #legacy to catch any stray probe. + const newCounts = { findRun: 0, findRunOnPrimary: 0 }; + const legacyCounts = { findRun: 0, findRunOnPrimary: 0 }; + const newStore = countingReads(makeDedicatedStore(prisma17), newCounts); + const legacyStore = countingReads(makeLegacyStore(prisma14), legacyCounts); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + + const newId = runOpsNew("mm2n"); // classifies NEW + await router.createRun( + buildCreateRunInput({ + runId: newId, + friendlyId: "run_mm2_new", + organizationId: env.organizationId, + projectId: env.projectId, + runtimeEnvironmentId: env.runtimeEnvironmentId, + }) + ); + + const hit = (await router.findRun({ id: newId }, { select: { id: true } })) as Record< + string, + unknown + > | null; + expect(hit?.id).toBe(newId); + // Owning store read exactly once; the other store NOT touched (fast path preserved). + expect(newCounts.findRun).toBe(1); + expect(legacyCounts.findRun).toBe(0); + expect(legacyCounts.findRunOnPrimary).toBe(0); + + // Symmetric: a cuid run whose owning store is #legacy must not probe #new on a hit. + const legacyId = cuidLegacy("mm2l"); // classifies LEGACY + await router.createRun( + buildCreateRunInput({ + runId: legacyId, + friendlyId: "run_mm2_legacy", + organizationId: env.organizationId, + projectId: env.projectId, + runtimeEnvironmentId: env.runtimeEnvironmentId, + }) + ); + newCounts.findRun = 0; + legacyCounts.findRun = 0; + + const hitLegacy = (await router.findRun( + { id: legacyId }, + { select: { id: true } } + )) as Record | null; + expect(hitLegacy?.id).toBe(legacyId); + expect(legacyCounts.findRun).toBe(1); + expect(newCounts.findRun).toBe(0); + } + ); + + // ── A genuine miss on BOTH stores still returns null (fan-out exhausted) ── + heteroRunOpsPostgresTest( + "returns null when the run is on neither store", + async ({ prisma14, prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const legacyStore = makeLegacyStore(prisma14); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + expect( + await router.findRun({ id: cuidLegacy("ghost") }, { select: { id: true } }) + ).toBeNull(); + } + ); +}); diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 27cab57d1a..ec0decce15 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -27,6 +27,7 @@ import type { WaitpointColocationOptions, } from "./types.js"; import { isReadReplicaClient } from "./readReplicaClient.js"; +import { CONNECTED_RUNS_LIMIT } from "./PostgresRunStore.js"; /** * Run-ops routing substrate for the TaskRun-core method group. Implements {@link RunStore} @@ -238,10 +239,9 @@ export class RoutingRunStore implements RunStore { const onPrimary = readYourWrites(argsOrClient, _client); const id = idFromWhere(where); if (id !== undefined) { - // Residency-classifiable (id/friendlyId): route to the owning store. - const store = this.#routeOrNew(id); - const method = onPrimary ? "findRunOnPrimary" : "findRun"; - return (store[method] as (...rest: unknown[]) => Promise)(where, args); + // Residency-classifiable (id/friendlyId): route to the owning store, then fall back to the + // OTHER store on a miss. + return this.#findRunRouted(id, where, args, onPrimary); } // Unclassifiable where (e.g. spanId, idempotencyKey): the run may live on either DB, // so fan out NEW-first then LEGACY rather than defaulting to NEW — defaulting silently @@ -249,6 +249,31 @@ export class RoutingRunStore implements RunStore { return this.#findRunUnrouted(where, args, onPrimary); } + // A classifiable id names its OWNING store by id-shape, but physical residency can diverge from + // classification (a pre-#4154 base62 run lives on #new yet classifies LEGACY). Read the owning + // store first — a hit is a SINGLE read (the fast path) — then, ONLY on a miss, probe the OTHER + // store before returning null, so a run whose residency ≠ its id-shape is found rather than 404'd. + // Both legs run the SAME index-covered TaskRun lookup; the fan-out cost is paid only on the (rare) + // miss. Mirrors #findRunUnrouted's shape but keyed on the classified owner. + async #findRunRouted( + id: string, + where: Prisma.TaskRunWhereInput, + args: unknown, + onPrimary: boolean + ): Promise { + const method = onPrimary ? "findRunOnPrimary" : "findRun"; + const owning = this.#routeOrNew(id); + const fromOwning = await (owning[method] as (...rest: unknown[]) => Promise)( + where, + args + ); + if (fromOwning != null) { + return fromOwning; + } + const other = owning === this.#new ? this.#legacy : this.#new; + return (other[method] as (...rest: unknown[]) => Promise)(where, args); + } + async #findRunUnrouted( where: Prisma.TaskRunWhereInput, args: unknown, @@ -376,6 +401,62 @@ export class RoutingRunStore implements RunStore { return finalizeRows([...byId.values()], args, addedFields); } + // Canonical grouped replacement for `Promise.all(ids.map(id => readThroughRun(id)))`: reuses + // `findRuns`'s bounded id-set path (`#findRunsByIdSet`), so NEW is queried once for the whole set + // and LEGACY once more only for the misses — never one round trip per id. Returns an id-keyed Map; + // missing/duplicate ids are simply absent/collapsed. + findRunsByIds( + ids: string[], + args: { select: S }, + client?: ReadClient + ): Promise>>; + findRunsByIds( + ids: string[], + args: { include: I }, + client?: ReadClient + ): Promise>>; + findRunsByIds(ids: string[], client?: ReadClient): Promise>; + async findRunsByIds( + ids: string[], + argsOrClient?: + | { select?: Record; include?: Record } + | ReadClient, + _client?: ReadClient + ): Promise> { + if (ids.length === 0) { + return new Map(); + } + const args = selectOrIncludeArgs(argsOrClient); + // Mirrors `readYourWrites`'s slot recovery: when `argsOrClient` isn't a `{select|include}` + // object it may itself BE the client (2-arg call) or be undefined with the client in the + // 3rd slot (an explicit `(ids, undefined, client)` call, e.g. from a relation hydrator). + const client = + args === undefined ? ((argsOrClient as ReadClient | undefined) ?? _client) : _client; + // Force `id` into the projection so the map can key off it, even when the caller's select + // omits it — `findRuns` would otherwise strip it back out as an added-for-merge-only field. + const projected = args?.select + ? { select: { ...args.select, id: true } } + : args?.include + ? { include: args.include } + : {}; + const rows = (await this.findRuns( + { where: { id: { in: ids } }, ...projected } as FindRunsArgs, + client + )) as Record[]; + const byId = new Map(); + // Strip the id we force-injected for map keying when the caller's select did not ask for it, + // so returned values match the declared payload type and never leak an unrequested id. + const stripInjectedId = args?.select != null && !("id" in (args.select as object)); + for (const row of rows) { + const key = row.id as string; + if (stripInjectedId) { + delete row.id; + } + byId.set(key, row); + } + return byId; + } + // --------------------------------------------------------------------------- // TaskRun-core: update-family — route by run id in params // --------------------------------------------------------------------------- @@ -880,7 +961,8 @@ export class RoutingRunStore implements RunStore { // Keyed by waitpointId, but the WaitpointRunConnection / CompletedWaitpoint join co-locates with the // RUN/snapshot — which can be on the OTHER DB from a cross-DB token — so fan out to BOTH stores and // merge. Dedup by value: a token mirrored onto both DBs during drain can carry the same join - // row on each leg. + // row on each leg. Each sub-store already caps at CONNECTED_RUNS_LIMIT, but a disjoint run set on + // each side can still make the union exceed it, so slice again after the merge. async findWaitpointConnectedRunIds(waitpointId: string, client?: ReadClient): Promise { const [fromNew, fromLegacy] = await Promise.all([ this.#new.findWaitpointConnectedRunIds( @@ -892,7 +974,7 @@ export class RoutingRunStore implements RunStore { RoutingRunStore.#ownPrimary(this.#legacy, client) ), ]); - return uniqueStrings([...fromNew, ...fromLegacy]); + return uniqueStrings([...fromNew, ...fromLegacy]).slice(0, CONNECTED_RUNS_LIMIT); } async findWaitpointCompletedSnapshotIds( @@ -1123,23 +1205,20 @@ export class RoutingRunStore implements RunStore { } // connectedRuns: the WaitpointRunConnection join co-locates with the run, so read the connected run - // ids from EACH store, then resolve the TaskRun rows across BOTH DBs (findRun routes by id). + // ids from EACH store, then resolve the TaskRun rows in ONE grouped, residency-partitioned read + // (`findRunsByIds`) and reorder to match `runIds` (the join's own order). async #reresolveConnectedRunsCrossDb( waitpointId: string, projection: SubProjection, client?: ReadClient ): Promise { const runIds = await this.findWaitpointConnectedRunIds(waitpointId, client); - const findRun = (this.findRun as (...rest: unknown[]) => Promise).bind(this); const args = projectionAsArgs(projection); - const runs: unknown[] = []; - for (const runId of runIds) { - const run = await findRun({ id: runId }, args, client); - if (run != null) { - runs.push(run); - } - } - return runs; + const findRunsByIds = ( + this.findRunsByIds as (...rest: unknown[]) => Promise> + ).bind(this); + const byId = runIds.length > 0 ? await findRunsByIds(runIds, args, client) : new Map(); + return runIds.map((id) => byId.get(id)).filter((run) => run != null); } // completedExecutionSnapshots: the CompletedWaitpoint join co-locates with the snapshot/run, so read @@ -1300,21 +1379,25 @@ export class RoutingRunStore implements RunStore { } } - // Resolve each edge's `taskRun` from its scalar `taskRunId` across BOTH stores (findRun routes by - // id and falls back NEW→LEGACY). A missing run is left null (display-only callers tolerate it; the - // blocked-run resume path keys off `waitpoint`). + // Resolve every edge's `taskRun` from its scalar `taskRunId` in ONE grouped, residency-partitioned + // read (`findRunsByIds`) rather than one `findRun` per edge. A missing run is left null + // (display-only callers tolerate it; the blocked-run resume path keys off `waitpoint`). async #hydrateEdgeTaskRunsCrossDb( edges: Record[], projection: SubProjection, client?: ReadClient ): Promise { - // Bind to `this`: findRun reaches the private #routeOrNew/#findRunUnrouted members, so an unbound - // reference loses `this` and throws on the first private access. - const findRun = (this.findRun as (...rest: unknown[]) => Promise).bind(this); + const ids = uniqueStrings(edges.map((e) => e.taskRunId)); const args = projectionAsArgs(projection); + // Bind to `this`: findRunsByIds reaches private members, so an unbound reference throws. + const findRunsByIds = ( + this.findRunsByIds as (...rest: unknown[]) => Promise> + ).bind(this); + const byId = + ids.length > 0 ? await findRunsByIds(ids, args, client) : new Map(); for (const edge of edges) { const id = edge.taskRunId as string | undefined; - const run = id ? await findRun({ id }, args, client) : null; + const run = id ? byId.get(id) : undefined; edge.taskRun = applyEdgeProjection((run as Record) ?? null, projection); } } diff --git a/internal-packages/run-store/src/scheduleSeam.test.ts b/internal-packages/run-store/src/scheduleSeam.test.ts new file mode 100644 index 0000000000..43c4ce63bf --- /dev/null +++ b/internal-packages/run-store/src/scheduleSeam.test.ts @@ -0,0 +1,306 @@ +// SCHEDULES over the run-ops cutover SEAM. +// +// schedule-engine (internal-packages/schedule-engine) has NO residency logic of its own: a cron +// schedule fires and calls back into `onTriggerScheduledTask` (apps/webapp/app/v3/scheduleEngine.server.ts), +// which mints the run via the ORDINARY TriggerTaskService -> RunEngine.trigger -> RoutingRunStore.createRun +// path — the same mint path already exercised by internal-packages/run-engine/src/engine/tests/triggerCreateRouting.test.ts +// and the `case 1` / `case 1b` findRuns-by-id-set tests in runOpsStore.mixedResidency.test.ts. There is no +// schedule-specific residency branch to test at the mint step; schedules simply inherit whatever the trigger +// path already does. Case A below adds exactly ONE small confirming test (using the previously-untested +// `scheduleId`/`scheduleInstanceId` scalar fields specifically) rather than re-deriving that generic coverage. +// +// The one piece of schedule-adjacent logic that genuinely spans the seam and had NO hetero coverage is +// `rescheduleRun` (RoutingRunStore.rescheduleRun, runOpsStore.ts:617) — the write delegate behind the +// delayed-run "reschedule" API (apps/webapp/app/v3/services/rescheduleTaskRun.server.ts) and +// `DelayedRunSystem.rescheduleDelayedRun`. It is a pure `#routeForWrite(runId)` mechanical delegate with +// no dedicated mixed-residency test anywhere in the existing matrix (runOpsStore.mixedResidency.test.ts +// covers batch/waitpoint/find methods but not rescheduleRun). Case B below closes that gap. +// +// Real two-physical-DB fixture, NO MOCKS: heteroRunOpsPostgresTest (prisma14 = full control-plane/LEGACY +// schema, prisma17 = RunOpsPrismaClient dedicated NEW-subset schema). + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; +import type { CreateRunInput, RunStoreSchemaVariant } from "./types.js"; + +type AnyClient = PrismaClient | RunOpsPrismaClient; + +// ownerEngine classifies by internal-id LENGTH after stripping a leading `_` +// (runOpsResidency.ts): 25 chars → cuid → LEGACY, a v1 body (version "1" at index 25 of a 26-char +// body) → run-ops id → NEW. Mirrors the helpers in runOpsStore.mixedResidency.test.ts so ids classify +// the same way here (the documented ID pitfall: a naive generator misclassifies NEW ids as LEGACY). +function cuidLegacy(seed: string): string { + return (seed + "c".repeat(25)).slice(0, 25); // 25 chars → LEGACY (#legacy / prisma14) +} +function runOpsNew(seed: string): string { + return (seed.replace(/[^0-9a-v]/g, "0") + "k".repeat(24)).slice(0, 24) + "01"; // 26 chars, version "1" at index 25 → NEW (#new / prisma17) +} + +async function seedEnvironment( + prisma: AnyClient, + schemaVariant: RunStoreSchemaVariant, + suffix: string +) { + if (schemaVariant === "dedicated") { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; + } + const organization = await (prisma as PrismaClient).organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await (prisma as PrismaClient).project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await (prisma as PrismaClient).runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { organization, project, environment }; +} + +async function seedSharedEnv(prisma14: PrismaClient, suffix: string) { + const legacy = await seedEnvironment(prisma14, "legacy", suffix); + return { + organizationId: legacy.organization.id, + projectId: legacy.project.id, + runtimeEnvironmentId: legacy.environment.id, + environmentId: legacy.environment.id, + }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; + status?: "PENDING" | "DELAYED"; + scheduleId?: string; + scheduleInstanceId?: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: params.status ?? "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-scheduled-task", + payload: '{"hello":"world"}', + payloadType: "application/json", + traceContext: { trace: "ctx" }, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + runTags: [], + queue: "task/my-scheduled-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: new Date("2024-01-01T00:00:00.000Z"), + ...(params.scheduleId && { scheduleId: params.scheduleId }), + ...(params.scheduleInstanceId && { scheduleInstanceId: params.scheduleInstanceId }), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: params.status ?? "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +function makeDedicatedStore(prisma17: RunOpsPrismaClient) { + return new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); +} + +function makeLegacyStore(prisma14: PrismaClient) { + return new PostgresRunStore({ + prisma: prisma14, + readOnlyPrisma: prisma14, + schemaVariant: "legacy", + }); +} + +// The REAL production split topology: #new = dedicated subset on prisma17, #legacy = full schema on +// prisma14. +function makeSplitRouter(prisma14: PrismaClient, prisma17: RunOpsPrismaClient) { + const legacyStore = makeLegacyStore(prisma14); + const newStore = makeDedicatedStore(prisma17); + return { + router: new RoutingRunStore({ new: newStore, legacy: legacyStore }), + legacyStore, + newStore, + }; +} + +describe("Schedules over the cutover seam", () => { + // ── Case A: a schedule-minted run (scheduleId/scheduleInstanceId set) lands on the correct + // physical store under mixed residency, and is found by plain id lookup (the ONLY lookup that + // exists — there is no scheduleInstanceId-keyed TaskRun query anywhere in the codebase; the field + // is write-only metadata stamped at creation). This confirms schedules have no distinct mint-routing + // seam beyond the already-covered generic createRun routing. ── + heteroRunOpsPostgresTest( + "case A: a schedule-minted run routes to the owning store and resolves by id, for both residencies", + async ({ prisma14, prisma17 }) => { + const { router } = makeSplitRouter(prisma14, prisma17); + const env = await seedSharedEnv(prisma14, "schedA"); + + const legacyRun = cuidLegacy("schAl"); // pre-cutover mint shape → #legacy + const newRun = runOpsNew("schAn"); // run-ops mint shape → #new + + await router.createRun( + buildCreateRunInput({ + runId: legacyRun, + friendlyId: "run_schedA_legacy", + scheduleId: "sched_A", + scheduleInstanceId: "schedinst_A_legacy", + ...env, + }) + ); + await router.createRun( + buildCreateRunInput({ + runId: newRun, + friendlyId: "run_schedA_new", + scheduleId: "sched_A", + scheduleInstanceId: "schedinst_A_new", + ...env, + }) + ); + + // Physical residency: each landed on its OWN DB only, with the scheduleInstanceId intact. + const onLegacy = await prisma14.taskRun.findUnique({ where: { id: legacyRun } }); + expect(onLegacy?.scheduleInstanceId).toBe("schedinst_A_legacy"); + expect(await prisma17.taskRun.findUnique({ where: { id: legacyRun } })).toBeNull(); + + const onNew = await prisma17.taskRun.findUnique({ where: { id: newRun } }); + expect(onNew?.scheduleInstanceId).toBe("schedinst_A_new"); + expect(await prisma14.taskRun.findUnique({ where: { id: newRun } })).toBeNull(); + + // The only lookup a schedule-minted run ever needs (by id) resolves through the router on + // BOTH residencies. + const foundLegacy = (await router.findRun( + { id: legacyRun }, + { select: { id: true, scheduleInstanceId: true } } + )) as Record | null; + expect(foundLegacy?.id).toBe(legacyRun); + expect(foundLegacy?.scheduleInstanceId).toBe("schedinst_A_legacy"); + + const foundNew = (await router.findRun( + { id: newRun }, + { select: { id: true, scheduleInstanceId: true } } + )) as Record | null; + expect(foundNew?.id).toBe(newRun); + expect(foundNew?.scheduleInstanceId).toBe("schedinst_A_new"); + } + ); + + // ── Case B: rescheduleRun (the write delegate behind the delayed-run "reschedule" API and + // DelayedRunSystem.rescheduleDelayedRun) routes to the OWNING store for a mixed-residency + // population, and does NOT touch the other physical DB. No existing hetero test covers this + // write path. ── + heteroRunOpsPostgresTest( + "case B: rescheduleRun routes the write to the owning store only, for a mixed-residency population", + async ({ prisma14, prisma17 }) => { + const { router } = makeSplitRouter(prisma14, prisma17); + const env = await seedSharedEnv(prisma14, "schedB"); + + const legacyRun = cuidLegacy("schBl"); + const newRun = runOpsNew("schBn"); + + await router.createRun( + buildCreateRunInput({ + runId: legacyRun, + friendlyId: "run_schedB_legacy", + status: "DELAYED", + ...env, + }) + ); + await router.createRun( + buildCreateRunInput({ + runId: newRun, + friendlyId: "run_schedB_new", + status: "DELAYED", + ...env, + }) + ); + + const legacyDelayUntil = new Date("2027-05-01T00:00:00.000Z"); + const newDelayUntil = new Date("2027-06-01T00:00:00.000Z"); + + const updatedLegacy = await router.rescheduleRun(legacyRun, { + delayUntil: legacyDelayUntil, + snapshot: { + environmentId: env.environmentId, + environmentType: "DEVELOPMENT", + projectId: env.projectId, + organizationId: env.organizationId, + }, + }); + const updatedNew = await router.rescheduleRun(newRun, { + delayUntil: newDelayUntil, + snapshot: { + environmentId: env.environmentId, + environmentType: "DEVELOPMENT", + projectId: env.projectId, + organizationId: env.organizationId, + }, + }); + + expect(updatedLegacy.id).toBe(legacyRun); + expect(updatedLegacy.delayUntil).toEqual(legacyDelayUntil); + expect(updatedNew.id).toBe(newRun); + expect(updatedNew.delayUntil).toEqual(newDelayUntil); + + // The write landed on the OWNING physical DB only. + const legacyRow = await prisma14.taskRun.findUnique({ where: { id: legacyRun } }); + expect(legacyRow?.delayUntil).toEqual(legacyDelayUntil); + const legacySnapshots = await prisma14.taskRunExecutionSnapshot.findMany({ + where: { runId: legacyRun, executionStatus: "DELAYED" }, + }); + expect(legacySnapshots).toHaveLength(1); + + const newRow = await prisma17.taskRun.findUnique({ where: { id: newRun } }); + expect(newRow?.delayUntil).toEqual(newDelayUntil); + const newSnapshots = await prisma17.taskRunExecutionSnapshot.findMany({ + where: { runId: newRun, executionStatus: "DELAYED" }, + }); + expect(newSnapshots).toHaveLength(1); + + // Cross-DB isolation: the OTHER physical DB's row for each run is untouched by the other + // reschedule call — neither run exists on the non-owning DB at all, so there is nothing there + // to have been (mis)updated. + expect(await prisma17.taskRun.findUnique({ where: { id: legacyRun } })).toBeNull(); + expect(await prisma14.taskRun.findUnique({ where: { id: newRun } })).toBeNull(); + } + ); +}); diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index 1d39498f39..26ccb52aea 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -573,6 +573,21 @@ export interface RunStore { client?: ReadClient ): Promise; + // Grouped replacement for `Promise.all(ids.map(id => findRun(id)))`: one round trip for the + // whole id batch instead of one per id. Returns an id-keyed Map; missing/duplicate ids are + // simply absent/collapsed. + findRunsByIds( + ids: string[], + args: { select: S }, + client?: ReadClient + ): Promise>>; + findRunsByIds( + ids: string[], + args: { include: I }, + client?: ReadClient + ): Promise>>; + findRunsByIds(ids: string[], client?: ReadClient): Promise>; + // --- run-ops persistence --- // Snapshots, waitpoints, implicit M:N joins, dependents, attempts and checkpoints. The // generic model wrappers are thin generics over the Prisma `*Args` types so include/select