From c64242702fe9ff993dd6d862105d4faef1aed23c Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 15:14:55 +0100 Subject: [PATCH 01/21] perf(webapp,run-store): group per-item run reads into batched queries Several run-ops read paths issued one database query per item where a single batched query returns the same data: batch results hydrated each member run separately, the run retrieve endpoint resolved each locked worker version separately, the dedicated-schema relation hydrator ran per row, and a few services read runs in per-id loops. These now group into batched reads. Adds a grouped findRunsByIds to the run store (one residency-partitioned read for a set of run ids) and routes the presenters and services through it. The dedicated-schema relation hydrator batches across all rows per relation. The waitpoint connected-runs read also regains its display limit at the fetch, so it no longer loads every connected run into memory before slicing to five. --- .../v3/ApiBatchResultsPresenter.server.ts | 75 ++-- .../v3/ApiRetrieveRunPresenter.server.ts | 16 +- .../v3/WaitpointPresenter.server.ts | 13 +- .../controlPlaneResolver.server.test.ts | 168 ++++++++ .../controlPlaneResolver.server.ts | 90 ++++ .../services/cancelDevSessionRuns.server.ts | 211 ++++++--- ...atchResultsPresenter.dedicatedSeam.test.ts | 337 +++++++++++++++ ...iBatchResultsPresenter.splitNPlus1.test.ts | 251 +++++++++++ ...veRunPresenter.groupedLockedWorker.test.ts | 266 ++++++++++++ .../cancelDevSessionRunsGroupedReads.test.ts | 183 ++++++++ ...ointPresenter.connectedRunsBounded.test.ts | 232 ++++++++++ ...tpointPresenter.splitConnectedRuns.test.ts | 207 +++++++++ .../tests/legacyRunNewTokenResumeFlow.test.ts | 302 +++++++++++++ .../PostgresRunStore.findRunsByIds.test.ts | 170 ++++++++ .../PostgresRunStore.groupedRelations.test.ts | 228 ++++++++++ .../run-store/src/PostgresRunStore.ts | 401 +++++++++++------- .../src/runOpsStore.groupedReads.test.ts | 388 +++++++++++++++++ .../run-store/src/runOpsStore.ts | 84 +++- .../run-store/src/scheduleSeam.test.ts | 306 +++++++++++++ internal-packages/run-store/src/types.ts | 15 + 20 files changed, 3651 insertions(+), 292 deletions(-) create mode 100644 apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.test.ts create mode 100644 apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts create mode 100644 apps/webapp/test/apiBatchResultsPresenter.splitNPlus1.test.ts create mode 100644 apps/webapp/test/apiRetrieveRunPresenter.groupedLockedWorker.test.ts create mode 100644 apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts create mode 100644 apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts create mode 100644 apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts create mode 100644 internal-packages/run-engine/src/engine/tests/legacyRunNewTokenResumeFlow.test.ts create mode 100644 internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts create mode 100644 internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts create mode 100644 internal-packages/run-store/src/runOpsStore.groupedReads.test.ts create mode 100644 internal-packages/run-store/src/scheduleSeam.test.ts 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..c7c4e67605 100644 --- a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts @@ -11,6 +11,10 @@ import { waitpointStatusToApiStatus } from "./WaitpointListPresenter.server"; export type WaitpointDetail = NonNullable>>; +// Single-sourced bound for connected run friendlyIds: applied at the FETCH in #connectedRunIdsOn, +// not just at display time. +export const CONNECTED_RUNS_DISPLAY_LIMIT = 5; + export class WaitpointPresenter extends BasePresenter { constructor( prisma?: PrismaClientOrTransaction, @@ -98,16 +102,16 @@ export class WaitpointPresenter extends BasePresenter { const runs = await client.taskRun.findMany({ where: { id: { in: runIds } }, select: { friendlyId: true }, - take: 5, + take: CONNECTED_RUNS_DISPLAY_LIMIT, }); 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 @@ -125,11 +129,14 @@ export class WaitpointPresenter extends BasePresenter { const links = await joinDelegate.findMany({ where: { waitpointId }, select: { taskRunId: true }, + take: CONNECTED_RUNS_DISPLAY_LIMIT, }); return links.map((link) => link.taskRunId); } + // CONNECTED_RUNS_DISPLAY_LIMIT is a compile-time constant int, not an interpolated value. const rows = await client.$queryRaw<{ A: string }[]>` SELECT "A" FROM "_WaitpointRunConnections" WHERE "B" = ${waitpointId} + LIMIT ${CONNECTED_RUNS_DISPLAY_LIMIT} `; return rows.map((row) => row.A); } 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/services/cancelDevSessionRuns.server.ts b/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts index 3575a75052..3ebc5df76a 100644 --- a/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts +++ b/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts @@ -1,10 +1,17 @@ import { type RunStore } from "@internal/run-store"; +import { ownerEngine } from "@trigger.dev/core/v3/isomorphic"; import { z } from "zod"; -import { type PrismaClientOrTransaction } from "~/db.server"; +import { + runOpsLegacyReplica as defaultLegacyReplica, + runOpsNewReplica as defaultNewClient, + type PrismaClientOrTransaction, + type PrismaReplicaClient, +} from "~/db.server"; import { findLatestSession } from "~/models/runtimeEnvironment.server"; import { logger } from "~/services/logger.server"; import { commonWorker } from "../commonWorker.server"; -import { type ReadThroughDeps, readThroughRun } from "../runOpsMigration/readThrough.server"; +import { type ReadThroughDeps } from "../runOpsMigration/readThrough.server"; +import { isSplitEnabled } from "../runOpsMigration/splitMode.server"; import { BaseService } from "./baseService.server"; import { type CancelableTaskRun, CancelTaskRunService } from "./cancelTaskRun.server"; @@ -17,9 +24,19 @@ export const CancelDevSessionRunsServiceOptions = z.object({ export type CancelDevSessionRunsServiceOptions = z.infer; +const RUN_SELECT = { + id: true, + engine: true, + status: true, + friendlyId: true, + taskEventStore: true, + createdAt: true, + completedAt: true, +} as const; + export class CancelDevSessionRunsService extends BaseService { // Injectable read-through deps for the run-ops TaskRun read. Undefined in production: - // readThroughRun then uses its ~/db.server singleton handles and the boot split flag, + // the grouped read then uses its ~/db.server singleton handles and the boot split flag, // so single-DB is unchanged. Tests inject the hetero new/legacy handles + splitEnabled. readonly #readThroughDeps?: ReadThroughDeps; @@ -69,80 +86,116 @@ export class CancelDevSessionRunsService extends BaseService { const cancelTaskRunService = new CancelTaskRunService(); - // readThroughRun resolves residency from the run id alone; an env scope is only - // available when a cancelled session was resolved. - const environmentId = cancelledSession?.environmentId ?? ""; + // Read every run up front, grouped by id shape and (when the split is on) by + // residency, instead of the old per-run readThroughRun probe. Cancellation stays a + // per-run mutation loop below, unchanged. + const runsById = await this.#readRunsGrouped(options.runIds); for (const runId of options.runIds) { - await this.#cancelInProgressRun( - runId, - cancelTaskRunService, - options.cancelledAt, - options.reason, - environmentId - ); + logger.debug("Cancelling in progress run", { runId }); + + const taskRun = runsById.get(runId); + + if (!taskRun) { + continue; + } + + try { + await cancelTaskRunService.call(taskRun, { + reason: options.reason, + cancelAttempts: true, + cancelledAt: options.cancelledAt, + }); + } catch (e) { + logger.error("Failed to cancel in progress run", { + runId, + error: e, + }); + } } } - async #cancelInProgressRun( - runId: string, - service: CancelTaskRunService, - cancelledAt: Date, - reason: string, - environmentId: string - ) { - logger.debug("Cancelling in progress run", { runId }); - - // Read-through: new store first, legacy read replica for an old - // in-retention run; single plain read in single-DB passthrough. - const where = runId.startsWith("run_") ? { friendlyId: runId } : { id: runId }; - - const result = await readThroughRun({ - runId, - environmentId, - readNew: (client) => - client.taskRun.findFirst({ - where, - select: { - id: true, - engine: true, - status: true, - friendlyId: true, - taskEventStore: true, - createdAt: true, - completedAt: true, - }, - }), - readLegacy: (replica) => - replica.taskRun.findFirst({ - where, - select: { - id: true, - engine: true, - status: true, - friendlyId: true, - taskEventStore: true, - createdAt: true, - completedAt: true, - }, - }), - deps: this.#readThroughDeps, - }); + // Grouped replacement for readThroughRun-per-id: a constant number of queries + // regardless of N. Same fan-out order as before (NEW-classified ids read only the + // new store; LEGACY-classified ids probe new first, then the legacy replica for + // misses; passthrough reads the single store once), just batched via findMany. + async #readRunsGrouped(runIds: string[]): Promise> { + const result = new Map(); + + if (runIds.length === 0) { + return result; + } + + const deps = this.#readThroughDeps; + const newClient = deps?.newClient ?? defaultNewClient; + const legacyReplica = deps?.legacyReplica ?? defaultLegacyReplica; + const splitEnabled = deps?.splitEnabled ?? (await isSplitEnabled()); + + const friendlyIds: string[] = []; + const internalIds: string[] = []; + + for (const runId of runIds) { + (runId.startsWith("run_") ? friendlyIds : internalIds).push(runId); + } - if (result.source === "not-found" || result.source === "past-retention") { - return; + const applyRows = ( + rows: CancelableTaskRun[], + requestedIds: string[], + key: "id" | "friendlyId" + ) => { + const byKey = new Map(rows.map((row) => [row[key], row])); + for (const id of requestedIds) { + const row = byKey.get(id); + if (row) { + result.set(id, row); + } + } + }; + + if (!splitEnabled) { + applyRows(await readGroup(newClient, internalIds, "id"), internalIds, "id"); + applyRows(await readGroup(newClient, friendlyIds, "friendlyId"), friendlyIds, "friendlyId"); + return result; } - const taskRun = result.value; + // classifyResidency is total (never throws): NEW ids read only the new store, + // LEGACY ids fan out new-then-legacy-replica. + const newIds: string[] = []; + const newFriendlyIds: string[] = []; + const legacyIds: string[] = []; + const legacyFriendlyIds: string[] = []; - try { - await service.call(taskRun, { reason, cancelAttempts: true, cancelledAt }); - } catch (e) { - logger.error("Failed to cancel in progress run", { - runId, - error: e, - }); + for (const id of internalIds) { + (ownerEngine(id) === "NEW" ? newIds : legacyIds).push(id); } + for (const id of friendlyIds) { + (ownerEngine(id) === "NEW" ? newFriendlyIds : legacyFriendlyIds).push(id); + } + + applyRows(await readGroup(newClient, newIds, "id"), newIds, "id"); + applyRows( + await readGroup(newClient, newFriendlyIds, "friendlyId"), + newFriendlyIds, + "friendlyId" + ); + + applyRows(await readGroup(newClient, legacyIds, "id"), legacyIds, "id"); + const legacyIdMisses = legacyIds.filter((id) => !result.has(id)); + applyRows(await readGroup(legacyReplica, legacyIdMisses, "id"), legacyIdMisses, "id"); + + applyRows( + await readGroup(newClient, legacyFriendlyIds, "friendlyId"), + legacyFriendlyIds, + "friendlyId" + ); + const legacyFriendlyMisses = legacyFriendlyIds.filter((id) => !result.has(id)); + applyRows( + await readGroup(legacyReplica, legacyFriendlyMisses, "friendlyId"), + legacyFriendlyMisses, + "friendlyId" + ); + + return result; } static async enqueue(options: CancelDevSessionRunsServiceOptions, runAt?: Date) { @@ -156,3 +209,27 @@ export class CancelDevSessionRunsService extends BaseService { }); } } + +// One `findFirst` for a lone id (byte-identical to the old per-run read), or one +// `findMany` for 2+ ids in the same bucket — the N+1 fix. +async function readGroup( + client: PrismaReplicaClient, + ids: string[], + field: "id" | "friendlyId" +): Promise { + if (ids.length === 0) { + return []; + } + + if (ids.length === 1) { + const row = + field === "id" + ? await client.taskRun.findFirst({ where: { id: ids[0] }, select: RUN_SELECT }) + : await client.taskRun.findFirst({ where: { friendlyId: ids[0] }, select: RUN_SELECT }); + return row ? [row] : []; + } + + return field === "id" + ? client.taskRun.findMany({ where: { id: { in: ids } }, select: RUN_SELECT }) + : client.taskRun.findMany({ where: { friendlyId: { in: ids } }, select: RUN_SELECT }); +} 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/cancelDevSessionRunsGroupedReads.test.ts b/apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts new file mode 100644 index 0000000000..617e68b47c --- /dev/null +++ b/apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts @@ -0,0 +1,183 @@ +// Proves the dev-session-cancel read is grouped into a constant number of queries +// instead of one findFirst per run id (N+1). Real Postgres via testcontainers — the +// DB is never mocked. A call-counting proxy sits over the REAL Prisma client: it +// tallies findFirst/findMany per delegate and forwards every call unchanged, so the +// real query still runs against the real database. This is instrumentation, not a +// mock. +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect, vi } from "vitest"; +import type { PrismaReplicaClient } from "~/db.server"; +import { CancelDevSessionRunsService } from "~/v3/services/cancelDevSessionRuns.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +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: `test-${suffix}`, + pkApiKey: `test-${suffix}`, + shortcode: `test-${suffix}`, + }, + }); + return { organization, project, runtimeEnvironment }; +} + +async function seedRun( + prisma: PrismaClient, + ids: { id: string; friendlyId: string }, + env: { runtimeEnvironmentId: string; projectId: string; organizationId: string } +) { + return prisma.taskRun.create({ + data: { + id: ids.id, + friendlyId: ids.friendlyId, + taskIdentifier: "my-task", + payload: JSON.stringify({ foo: "bar" }), + payloadType: "application/json", + traceId: "1234", + spanId: "1234", + queue: "test", + runtimeEnvironmentId: env.runtimeEnvironmentId, + projectId: env.projectId, + organizationId: env.organizationId, + environmentType: "DEVELOPMENT", + // V1 so the (best-effort, error-swallowed) cancel does not require the V2 engine; + // the unit under test is the READ resolution, not the cancel side effect. + engine: "V1", + status: "EXECUTING", + }, + }); +} + +// Call-counting instrumentation over the REAL Prisma client. Every findFirst/findMany +// invocation is tallied per delegate and then forwarded to the real underlying method +// (`target[key].apply(target, args)`), so the real query still executes against the +// real container. Every other property passes through untouched. +function countingClient(client: PrismaClient) { + const counts = { findFirst: 0, findMany: 0 }; + const realTaskRun = client.taskRun; + const taskRunProxy = new Proxy(realTaskRun, { + get(target, prop, receiver) { + if (prop === "findFirst" || prop === "findMany") { + const key = prop as "findFirst" | "findMany"; + return (...args: unknown[]) => { + counts[key]++; + return (target[key] as (...a: unknown[]) => unknown).apply(target, args); + }; + } + return Reflect.get(target, prop, receiver); + }, + }); + const clientProxy = new Proxy(client, { + get(target, prop, receiver) { + if (prop === "taskRun") { + return taskRunProxy; + } + return Reflect.get(target, prop, receiver); + }, + }); + return { handle: clientProxy as unknown as PrismaReplicaClient, counts }; +} + +describe("CancelDevSessionRunsService grouped reads", () => { + postgresTest( + "reads N runs in a constant number of grouped queries instead of one findFirst per run id", + async ({ prisma }) => { + const { project, organization, runtimeEnvironment } = await seedOrgProjectEnv( + prisma, + "grouped" + ); + + const seeded = await Promise.all( + ["a", "b", "c"].map((suffix) => + seedRun( + prisma, + { id: `run-internal-${suffix}`, friendlyId: `run_run-internal-${suffix}` }, + { + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + } + ) + ) + ); + + const counting = countingClient(prisma); + + const service = new CancelDevSessionRunsService({ + prisma, + replica: prisma, + readThroughDeps: { + splitEnabled: false, + newClient: counting.handle, + }, + }); + + await service.call({ + runIds: seeded.map((r) => r.id), + cancelledAt: new Date(), + reason: "test", + }); + + // The N+1 fix: one grouped findMany for the whole id batch, zero per-run findFirst calls. + expect(counting.counts.findMany).toBe(1); + expect(counting.counts.findFirst).toBe(0); + } + ); + + postgresTest( + "a single run id still resolves correctly (grouped read collapses to the same single read)", + async ({ prisma }) => { + const { project, organization, runtimeEnvironment } = await seedOrgProjectEnv( + prisma, + "single" + ); + + const run = await seedRun( + prisma, + { id: "run-internal-solo", friendlyId: "run_run-internal-solo" }, + { + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + } + ); + + const counting = countingClient(prisma); + + const service = new CancelDevSessionRunsService({ + prisma, + replica: prisma, + readThroughDeps: { + splitEnabled: false, + newClient: counting.handle, + }, + }); + + await service.call({ + runIds: [run.id], + cancelledAt: new Date(), + reason: "test", + }); + + expect(counting.counts.findFirst).toBe(1); + expect(counting.counts.findMany).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..6a0e24b28b --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts @@ -0,0 +1,232 @@ +// RED->GREEN guard: WaitpointPresenter#connectedRunIdsOn must BOUND the fetch of a waitpoint's +// connected-run ids, not just the number displayed. A "displayed count <= 5" assertion would +// false-green on today's buggy code (the existing take:5 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 itself is bounded. +// +// Exercises the dedicated-schema (Prisma `waitpointRunConnection`) branch of #connectedRunIdsOn: +// waitpoint + 8 connected runs seeded on the NEW dedicated run-ops client (RunOpsPrismaClient, +// prisma17), which uses the delegate directly (no raw SQL fallback). +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`]), + }; +}); + +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_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: 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 }; +} + +const CONNECTED_RUN_COUNT = 8; // > CONNECTED_RUNS_DISPLAY_LIMIT (5), proves the fetch isn't bounded + +describe("WaitpointPresenter#connectedRunIdsOn bounds the connected-run-id FETCH", () => { + heteroRunOpsPostgresTest( + "reading a waitpoint with 8 connected runs never IN-lists more than the display 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. On today's buggy code, the Prisma + // `waitpointRunConnection.findMany` branch of #connectedRunIdsOn returns all 8 connected + // run ids with no take/LIMIT, so this IN-list is 8. After the fix it must be + // <= CONNECTED_RUNS_DISPLAY_LIMIT. + expect(calls.length).toBeGreaterThan(0); + for (const call of calls) { + expect(call.where?.id?.in?.length ?? 0).toBeLessThanOrEqual(CONNECTED_RUNS_DISPLAY_LIMIT); + } + } + ); +}); diff --git a/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts b/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts new file mode 100644 index 0000000000..2511a25af1 --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts @@ -0,0 +1,207 @@ +// 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`]), + }; +}); + +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.findRunsByIds.test.ts b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts new file mode 100644 index 0000000000..09dbab106d --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts @@ -0,0 +1,170 @@ +// 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"); + + // 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..948b904748 --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts @@ -0,0 +1,228 @@ +// 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 } { + 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 RunOpsPrismaClient; + return { client, counts }; +} + +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 join query + one target query for the WHOLE batch, never one pair per parent. + expect(counting.counts.waitpointRunConnection.findMany).toBe(1); + expect(counting.counts.taskRun.findMany).toBe(1); + expect(counting.counts.taskRun.findFirst).toBe(0); + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 5d8ea6c8df..c863874220 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -99,13 +99,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; @@ -175,153 +177,192 @@ function stripDedicatedRelations( return { stripped: args, requested }; } -// --- per-model dedicated-schema relation hydrators --- - -// 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); -}; - -// Display connections for a run: WaitpointRunConnection → Waitpoint rows. -const hydrateConnectedWaitpoints: DedicatedRelationHydrator = async ( - client, - parent, - projection -) => { - const join = client.waitpointRunConnection; - if (!join) { - return []; +// --- per-model dedicated-schema relation hydrators (batched across the WHOLE parent array) --- + +// 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 rows = (await client.waitpoint.findMany({ - where: { id: { in: links.map((l) => l.waitpointId) } }, + const targetIds = [...new Set(links.map((l) => l[joinTargetField]))]; + const rows = (await targetDelegate.findMany({ + where: { id: { in: targetIds } }, })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); -}; + 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)); + } + } + 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({ - where: { id: { in: links.map((l) => l.waitpointId) } }, + where: { completedByTaskRunId: { in: parentIds } }, })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); + for (const row of rows) { + const runId = row.completedByTaskRunId as string | undefined; + if (runId && byParent.has(runId)) { + byParent.set(runId, applyProjection(row, 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({ where: { id: { in: runIds } } }) : [] + ) 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 }[]; - if (links.length === 0) { - return []; - } - const rows = (await client.taskRun.findMany({ - where: { id: { in: links.map((l) => l.taskRunId) } }, - })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); -}; +const hydrateConnectedRuns: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateJoinRelation( + client.waitpointRunConnection, + client.taskRun, + parents.map((p) => p.id as string), + "waitpointId", + "taskRunId", + projection + ); // 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({ + where: { id: { in: [...new Set(targetIds)] } }, + })) 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 +1485,65 @@ 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(); for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - TASK_RUN_DEDICATED - ); + byId.set(row.id as string, row); } - return rows; + return byId; } // --- run-ops persistence --- @@ -1563,14 +1654,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[]; } @@ -1949,14 +2038,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 +2106,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 +2352,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 +2370,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..fb1fbd3a6a --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts @@ -0,0 +1,388 @@ +// 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: { id: true, 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"); + + // 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.ts b/internal-packages/run-store/src/runOpsStore.ts index 27cab57d1a..3341bad9b6 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -376,6 +376,55 @@ 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(); + for (const row of rows) { + byId.set(row.id as string, row); + } + return byId; + } + // --------------------------------------------------------------------------- // TaskRun-core: update-family — route by run id in params // --------------------------------------------------------------------------- @@ -1123,23 +1172,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 +1346,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 From e7fb6bdd527da5ecbff9664b22d016944f5cfc84 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 15:50:18 +0100 Subject: [PATCH 02/21] fix(run-store): drop force-injected id from findRunsByIds results when unrequested findRunsByIds forces `id` into the select so it can key the result map, but did not remove it afterwards, so returned rows carried an `id` the declared payload type did not include. Delete the injected id from each value when the caller's select did not request it; the map stays keyed by id. --- .../run-store/src/PostgresRunStore.findRunsByIds.test.ts | 5 +++++ internal-packages/run-store/src/PostgresRunStore.ts | 9 ++++++++- .../run-store/src/runOpsStore.groupedReads.test.ts | 6 +++++- internal-packages/run-store/src/runOpsStore.ts | 9 ++++++++- 4 files changed, 26 insertions(+), 3 deletions(-) diff --git a/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts index 09dbab106d..10e9aa8086 100644 --- a/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts @@ -162,6 +162,11 @@ describe("PostgresRunStore.findRunsByIds", () => { 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.ts b/internal-packages/run-store/src/PostgresRunStore.ts index c863874220..974034ff27 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -1540,8 +1540,15 @@ export class PostgresRunStore implements RunStore { 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) { - byId.set(row.id as string, row); + const key = row.id as string; + if (stripInjectedId) { + delete row.id; + } + byId.set(key, row); } return byId; } diff --git a/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts index fb1fbd3a6a..4b2230bec0 100644 --- a/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts +++ b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts @@ -178,7 +178,7 @@ describe("RoutingRunStore.findRunsByIds — grouped, residency-partitioned read newCounting.counts.findFirst = 0; const result = await router.findRunsByIds(allIds, { - select: { id: true, friendlyId: true }, + select: { friendlyId: true }, }); // All N rows returned, keyed by id. @@ -189,6 +189,10 @@ describe("RoutingRunStore.findRunsByIds — grouped, residency-partitioned read 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); diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 3341bad9b6..1884742356 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -419,8 +419,15 @@ export class RoutingRunStore implements RunStore { 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) { - byId.set(row.id as string, row); + const key = row.id as string; + if (stripInjectedId) { + delete row.id; + } + byId.set(key, row); } return byId; } From fc9588fdf68a3c9e7cf3d4e5eaf9abef90a9fe58 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 17:49:36 +0100 Subject: [PATCH 03/21] chore(webapp): drop cancelDevSessionRuns changes from this PR --- .../services/cancelDevSessionRuns.server.ts | 211 ++++++------------ .../cancelDevSessionRunsGroupedReads.test.ts | 183 --------------- 2 files changed, 67 insertions(+), 327 deletions(-) delete mode 100644 apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts diff --git a/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts b/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts index 3ebc5df76a..3575a75052 100644 --- a/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts +++ b/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts @@ -1,17 +1,10 @@ import { type RunStore } from "@internal/run-store"; -import { ownerEngine } from "@trigger.dev/core/v3/isomorphic"; import { z } from "zod"; -import { - runOpsLegacyReplica as defaultLegacyReplica, - runOpsNewReplica as defaultNewClient, - type PrismaClientOrTransaction, - type PrismaReplicaClient, -} from "~/db.server"; +import { type PrismaClientOrTransaction } from "~/db.server"; import { findLatestSession } from "~/models/runtimeEnvironment.server"; import { logger } from "~/services/logger.server"; import { commonWorker } from "../commonWorker.server"; -import { type ReadThroughDeps } from "../runOpsMigration/readThrough.server"; -import { isSplitEnabled } from "../runOpsMigration/splitMode.server"; +import { type ReadThroughDeps, readThroughRun } from "../runOpsMigration/readThrough.server"; import { BaseService } from "./baseService.server"; import { type CancelableTaskRun, CancelTaskRunService } from "./cancelTaskRun.server"; @@ -24,19 +17,9 @@ export const CancelDevSessionRunsServiceOptions = z.object({ export type CancelDevSessionRunsServiceOptions = z.infer; -const RUN_SELECT = { - id: true, - engine: true, - status: true, - friendlyId: true, - taskEventStore: true, - createdAt: true, - completedAt: true, -} as const; - export class CancelDevSessionRunsService extends BaseService { // Injectable read-through deps for the run-ops TaskRun read. Undefined in production: - // the grouped read then uses its ~/db.server singleton handles and the boot split flag, + // readThroughRun then uses its ~/db.server singleton handles and the boot split flag, // so single-DB is unchanged. Tests inject the hetero new/legacy handles + splitEnabled. readonly #readThroughDeps?: ReadThroughDeps; @@ -86,116 +69,80 @@ export class CancelDevSessionRunsService extends BaseService { const cancelTaskRunService = new CancelTaskRunService(); - // Read every run up front, grouped by id shape and (when the split is on) by - // residency, instead of the old per-run readThroughRun probe. Cancellation stays a - // per-run mutation loop below, unchanged. - const runsById = await this.#readRunsGrouped(options.runIds); + // readThroughRun resolves residency from the run id alone; an env scope is only + // available when a cancelled session was resolved. + const environmentId = cancelledSession?.environmentId ?? ""; for (const runId of options.runIds) { - logger.debug("Cancelling in progress run", { runId }); - - const taskRun = runsById.get(runId); - - if (!taskRun) { - continue; - } - - try { - await cancelTaskRunService.call(taskRun, { - reason: options.reason, - cancelAttempts: true, - cancelledAt: options.cancelledAt, - }); - } catch (e) { - logger.error("Failed to cancel in progress run", { - runId, - error: e, - }); - } + await this.#cancelInProgressRun( + runId, + cancelTaskRunService, + options.cancelledAt, + options.reason, + environmentId + ); } } - // Grouped replacement for readThroughRun-per-id: a constant number of queries - // regardless of N. Same fan-out order as before (NEW-classified ids read only the - // new store; LEGACY-classified ids probe new first, then the legacy replica for - // misses; passthrough reads the single store once), just batched via findMany. - async #readRunsGrouped(runIds: string[]): Promise> { - const result = new Map(); - - if (runIds.length === 0) { - return result; - } - - const deps = this.#readThroughDeps; - const newClient = deps?.newClient ?? defaultNewClient; - const legacyReplica = deps?.legacyReplica ?? defaultLegacyReplica; - const splitEnabled = deps?.splitEnabled ?? (await isSplitEnabled()); - - const friendlyIds: string[] = []; - const internalIds: string[] = []; - - for (const runId of runIds) { - (runId.startsWith("run_") ? friendlyIds : internalIds).push(runId); - } - - const applyRows = ( - rows: CancelableTaskRun[], - requestedIds: string[], - key: "id" | "friendlyId" - ) => { - const byKey = new Map(rows.map((row) => [row[key], row])); - for (const id of requestedIds) { - const row = byKey.get(id); - if (row) { - result.set(id, row); - } - } - }; + async #cancelInProgressRun( + runId: string, + service: CancelTaskRunService, + cancelledAt: Date, + reason: string, + environmentId: string + ) { + logger.debug("Cancelling in progress run", { runId }); + + // Read-through: new store first, legacy read replica for an old + // in-retention run; single plain read in single-DB passthrough. + const where = runId.startsWith("run_") ? { friendlyId: runId } : { id: runId }; + + const result = await readThroughRun({ + runId, + environmentId, + readNew: (client) => + client.taskRun.findFirst({ + where, + select: { + id: true, + engine: true, + status: true, + friendlyId: true, + taskEventStore: true, + createdAt: true, + completedAt: true, + }, + }), + readLegacy: (replica) => + replica.taskRun.findFirst({ + where, + select: { + id: true, + engine: true, + status: true, + friendlyId: true, + taskEventStore: true, + createdAt: true, + completedAt: true, + }, + }), + deps: this.#readThroughDeps, + }); - if (!splitEnabled) { - applyRows(await readGroup(newClient, internalIds, "id"), internalIds, "id"); - applyRows(await readGroup(newClient, friendlyIds, "friendlyId"), friendlyIds, "friendlyId"); - return result; + if (result.source === "not-found" || result.source === "past-retention") { + return; } - // classifyResidency is total (never throws): NEW ids read only the new store, - // LEGACY ids fan out new-then-legacy-replica. - const newIds: string[] = []; - const newFriendlyIds: string[] = []; - const legacyIds: string[] = []; - const legacyFriendlyIds: string[] = []; + const taskRun = result.value; - for (const id of internalIds) { - (ownerEngine(id) === "NEW" ? newIds : legacyIds).push(id); + try { + await service.call(taskRun, { reason, cancelAttempts: true, cancelledAt }); + } catch (e) { + logger.error("Failed to cancel in progress run", { + runId, + error: e, + }); } - for (const id of friendlyIds) { - (ownerEngine(id) === "NEW" ? newFriendlyIds : legacyFriendlyIds).push(id); - } - - applyRows(await readGroup(newClient, newIds, "id"), newIds, "id"); - applyRows( - await readGroup(newClient, newFriendlyIds, "friendlyId"), - newFriendlyIds, - "friendlyId" - ); - - applyRows(await readGroup(newClient, legacyIds, "id"), legacyIds, "id"); - const legacyIdMisses = legacyIds.filter((id) => !result.has(id)); - applyRows(await readGroup(legacyReplica, legacyIdMisses, "id"), legacyIdMisses, "id"); - - applyRows( - await readGroup(newClient, legacyFriendlyIds, "friendlyId"), - legacyFriendlyIds, - "friendlyId" - ); - const legacyFriendlyMisses = legacyFriendlyIds.filter((id) => !result.has(id)); - applyRows( - await readGroup(legacyReplica, legacyFriendlyMisses, "friendlyId"), - legacyFriendlyMisses, - "friendlyId" - ); - - return result; } static async enqueue(options: CancelDevSessionRunsServiceOptions, runAt?: Date) { @@ -209,27 +156,3 @@ export class CancelDevSessionRunsService extends BaseService { }); } } - -// One `findFirst` for a lone id (byte-identical to the old per-run read), or one -// `findMany` for 2+ ids in the same bucket — the N+1 fix. -async function readGroup( - client: PrismaReplicaClient, - ids: string[], - field: "id" | "friendlyId" -): Promise { - if (ids.length === 0) { - return []; - } - - if (ids.length === 1) { - const row = - field === "id" - ? await client.taskRun.findFirst({ where: { id: ids[0] }, select: RUN_SELECT }) - : await client.taskRun.findFirst({ where: { friendlyId: ids[0] }, select: RUN_SELECT }); - return row ? [row] : []; - } - - return field === "id" - ? client.taskRun.findMany({ where: { id: { in: ids } }, select: RUN_SELECT }) - : client.taskRun.findMany({ where: { friendlyId: { in: ids } }, select: RUN_SELECT }); -} diff --git a/apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts b/apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts deleted file mode 100644 index 617e68b47c..0000000000 --- a/apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts +++ /dev/null @@ -1,183 +0,0 @@ -// Proves the dev-session-cancel read is grouped into a constant number of queries -// instead of one findFirst per run id (N+1). Real Postgres via testcontainers — the -// DB is never mocked. A call-counting proxy sits over the REAL Prisma client: it -// tallies findFirst/findMany per delegate and forwards every call unchanged, so the -// real query still runs against the real database. This is instrumentation, not a -// mock. -import { postgresTest } from "@internal/testcontainers"; -import type { PrismaClient } from "@trigger.dev/database"; -import { describe, expect, vi } from "vitest"; -import type { PrismaReplicaClient } from "~/db.server"; -import { CancelDevSessionRunsService } from "~/v3/services/cancelDevSessionRuns.server"; - -vi.setConfig({ testTimeout: 60_000 }); - -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: `test-${suffix}`, - pkApiKey: `test-${suffix}`, - shortcode: `test-${suffix}`, - }, - }); - return { organization, project, runtimeEnvironment }; -} - -async function seedRun( - prisma: PrismaClient, - ids: { id: string; friendlyId: string }, - env: { runtimeEnvironmentId: string; projectId: string; organizationId: string } -) { - return prisma.taskRun.create({ - data: { - id: ids.id, - friendlyId: ids.friendlyId, - taskIdentifier: "my-task", - payload: JSON.stringify({ foo: "bar" }), - payloadType: "application/json", - traceId: "1234", - spanId: "1234", - queue: "test", - runtimeEnvironmentId: env.runtimeEnvironmentId, - projectId: env.projectId, - organizationId: env.organizationId, - environmentType: "DEVELOPMENT", - // V1 so the (best-effort, error-swallowed) cancel does not require the V2 engine; - // the unit under test is the READ resolution, not the cancel side effect. - engine: "V1", - status: "EXECUTING", - }, - }); -} - -// Call-counting instrumentation over the REAL Prisma client. Every findFirst/findMany -// invocation is tallied per delegate and then forwarded to the real underlying method -// (`target[key].apply(target, args)`), so the real query still executes against the -// real container. Every other property passes through untouched. -function countingClient(client: PrismaClient) { - const counts = { findFirst: 0, findMany: 0 }; - const realTaskRun = client.taskRun; - const taskRunProxy = new Proxy(realTaskRun, { - get(target, prop, receiver) { - if (prop === "findFirst" || prop === "findMany") { - const key = prop as "findFirst" | "findMany"; - return (...args: unknown[]) => { - counts[key]++; - return (target[key] as (...a: unknown[]) => unknown).apply(target, args); - }; - } - return Reflect.get(target, prop, receiver); - }, - }); - const clientProxy = new Proxy(client, { - get(target, prop, receiver) { - if (prop === "taskRun") { - return taskRunProxy; - } - return Reflect.get(target, prop, receiver); - }, - }); - return { handle: clientProxy as unknown as PrismaReplicaClient, counts }; -} - -describe("CancelDevSessionRunsService grouped reads", () => { - postgresTest( - "reads N runs in a constant number of grouped queries instead of one findFirst per run id", - async ({ prisma }) => { - const { project, organization, runtimeEnvironment } = await seedOrgProjectEnv( - prisma, - "grouped" - ); - - const seeded = await Promise.all( - ["a", "b", "c"].map((suffix) => - seedRun( - prisma, - { id: `run-internal-${suffix}`, friendlyId: `run_run-internal-${suffix}` }, - { - runtimeEnvironmentId: runtimeEnvironment.id, - projectId: project.id, - organizationId: organization.id, - } - ) - ) - ); - - const counting = countingClient(prisma); - - const service = new CancelDevSessionRunsService({ - prisma, - replica: prisma, - readThroughDeps: { - splitEnabled: false, - newClient: counting.handle, - }, - }); - - await service.call({ - runIds: seeded.map((r) => r.id), - cancelledAt: new Date(), - reason: "test", - }); - - // The N+1 fix: one grouped findMany for the whole id batch, zero per-run findFirst calls. - expect(counting.counts.findMany).toBe(1); - expect(counting.counts.findFirst).toBe(0); - } - ); - - postgresTest( - "a single run id still resolves correctly (grouped read collapses to the same single read)", - async ({ prisma }) => { - const { project, organization, runtimeEnvironment } = await seedOrgProjectEnv( - prisma, - "single" - ); - - const run = await seedRun( - prisma, - { id: "run-internal-solo", friendlyId: "run_run-internal-solo" }, - { - runtimeEnvironmentId: runtimeEnvironment.id, - projectId: project.id, - organizationId: organization.id, - } - ); - - const counting = countingClient(prisma); - - const service = new CancelDevSessionRunsService({ - prisma, - replica: prisma, - readThroughDeps: { - splitEnabled: false, - newClient: counting.handle, - }, - }); - - await service.call({ - runIds: [run.id], - cancelledAt: new Date(), - reason: "test", - }); - - expect(counting.counts.findFirst).toBe(1); - expect(counting.counts.findMany).toBe(0); - } - ); -}); From a32f16992f086dd2b66bf98d24d2af7e969a5856 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 18:09:42 +0100 Subject: [PATCH 04/21] fix(webapp): filter deleted runs out of a waitpoint's connected-runs list The waitpoint detail view capped the connected-run lookup before checking the runs still existed, so a run deleted after connecting could take a slot ahead of a live one and under-count (or empty) the list. Join to the run table so the cap only ever lands on connections whose run still exists. --- .../v3/WaitpointPresenter.server.ts | 41 ++-- ...intPresenter.danglingConnectedRuns.test.ts | 198 ++++++++++++++++++ 2 files changed, 221 insertions(+), 18 deletions(-) create mode 100644 apps/webapp/test/waitpointPresenter.danglingConnectedRuns.test.ts diff --git a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts index c7c4e67605..5e9cd315c4 100644 --- a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts @@ -115,27 +115,32 @@ export class WaitpointPresenter extends BasePresenter { } // 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. + // `WaitpointRunConnection` model (scalar `taskRunId`, no FK -- a row can dangle after its run is + // deleted), the control-plane full schema the implicit `_WaitpointRunConnections` M2M + // (A = TaskRun.id, B = Waitpoint.id). Both branches existence-filter AT THE QUERY via a JOIN to + // TaskRun, so a dangling connection row can never occupy a LIMIT slot ahead of a real one. + // CONNECTED_RUNS_DISPLAY_LIMIT is a compile-time constant int, not an interpolated value. async #connectedRunIdsOn(client: PrismaReplicaClient, waitpointId: string): Promise { - const joinDelegate = ( - client as unknown as { - waitpointRunConnection?: { - findMany: (args: unknown) => Promise<{ taskRunId: string }[]>; - }; - } - ).waitpointRunConnection; - if (joinDelegate && typeof joinDelegate.findMany === "function") { - const links = await joinDelegate.findMany({ - where: { waitpointId }, - select: { taskRunId: true }, - take: CONNECTED_RUNS_DISPLAY_LIMIT, - }); - return links.map((link) => link.taskRunId); + const isDedicated = Boolean( + (client as unknown as { waitpointRunConnection?: unknown }).waitpointRunConnection + ); + + if (isDedicated) { + const rows = await client.$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_DISPLAY_LIMIT} + `; + return rows.map((row) => row.taskRunId); } - // CONNECTED_RUNS_DISPLAY_LIMIT is a compile-time constant int, not an interpolated value. + const rows = await client.$queryRaw<{ A: string }[]>` - SELECT "A" FROM "_WaitpointRunConnections" WHERE "B" = ${waitpointId} + SELECT c."A" AS "A" + FROM "_WaitpointRunConnections" c + JOIN "TaskRun" t ON t."id" = c."A" + WHERE c."B" = ${waitpointId} LIMIT ${CONNECTED_RUNS_DISPLAY_LIMIT} `; return rows.map((row) => row.A); diff --git a/apps/webapp/test/waitpointPresenter.danglingConnectedRuns.test.ts b/apps/webapp/test/waitpointPresenter.danglingConnectedRuns.test.ts new file mode 100644 index 0000000000..c7d0de5dbe --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.danglingConnectedRuns.test.ts @@ -0,0 +1,198 @@ +// 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`]), + }; +}); + +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)) + ); + } + ); +}); From 119c1f639c69e38397a906cc0bd96ad60d9befad Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 18:09:42 +0100 Subject: [PATCH 05/21] fix(run-store): give each parent its own hydrated relation object Batched relation hydration returned the same row object to every parent linking the same target, so two parents connected to one run shared a single instance and an in-place edit to one leaked into the other. Return a shallow clone on the bare-projection path so each parent gets a distinct object. --- ...tgresRunStore.sharedHydratedTarget.test.ts | 139 ++++++++++++++++++ .../run-store/src/PostgresRunStore.ts | 12 +- 2 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 internal-packages/run-store/src/PostgresRunStore.sharedHydratedTarget.test.ts 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 974034ff27..e63541193f 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -121,13 +121,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) { From 3e8faa500701faf84aef257d8b9b33d5bac25455 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 20:53:13 +0100 Subject: [PATCH 06/21] fix(run-store): bound waitpoint connected-run reads to the display limit The connected-runs relation read fetched every connection id for a waitpoint with no existence check or limit, so a heavily-connected waitpoint could pull an unbounded id list into the grouped run lookup. Existence-join to the run table and cap at the shared connected-runs limit on both schema branches, and slice the cross-store union to the same bound. --- ...tgresRunStore.connectedRunsBounded.test.ts | 286 ++++++++++++++++++ .../run-store/src/PostgresRunStore.ts | 31 +- .../run-store/src/runOpsStore.ts | 6 +- 3 files changed, 313 insertions(+), 10 deletions(-) create mode 100644 internal-packages/run-store/src/PostgresRunStore.connectedRunsBounded.test.ts 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.ts b/internal-packages/run-store/src/PostgresRunStore.ts index e63541193f..e09a6a1d0d 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; @@ -1816,23 +1820,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 diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 1884742356..fa068b4c2f 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} @@ -936,7 +937,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( @@ -948,7 +950,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( From e205b800e95a3d20c76ee8bf40934d2e5c6de07f Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 10:03:30 +0100 Subject: [PATCH 07/21] perf(run-store): narrow dedicated hydrator reads to requested columns The dedicated-schema relation hydrators fetched full target rows and then stripped unrequested fields in memory, pulling wide columns even when the caller only selected a couple. Push the caller's select into the target findMany so only the requested columns are read. No change to results. --- ...resRunStore.hydratorSelectPushdown.test.ts | 220 ++++++++++++++++++ .../run-store/src/PostgresRunStore.ts | 44 +++- 2 files changed, 254 insertions(+), 10 deletions(-) create mode 100644 internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts 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..ff1218bbfe --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts @@ -0,0 +1,220 @@ +// 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.ts b/internal-packages/run-store/src/PostgresRunStore.ts index e09a6a1d0d..415340dcd4 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -193,6 +193,24 @@ function stripDedicatedRelations( // --- per-model dedicated-schema relation hydrators (batched across the WHOLE parent array) --- +// 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 }; +} + // 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 @@ -217,9 +235,9 @@ async function batchHydrateJoinRelation( return byParent; } const targetIds = [...new Set(links.map((l) => l[joinTargetField]))]; - const rows = (await targetDelegate.findMany({ - where: { id: { in: targetIds } }, - })) as Record[]; + 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]); @@ -242,9 +260,11 @@ const hydrateAssociatedWaitpoint: DedicatedRelationHydrator = async ( if (parentIds.length === 0) { return byParent; } - const rows = (await client.waitpoint.findMany({ - where: { completedByTaskRunId: { in: parentIds } }, - })) as Record[]; + 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)) { @@ -293,7 +313,11 @@ const hydrateBlockingTaskRuns: DedicatedRelationHydrator = async (client, parent if (nestedTaskRun) { const runIds = [...new Set(edges.map((e) => e.taskRunId as string))]; const runs = ( - runIds.length > 0 ? await client.taskRun.findMany({ where: { id: { in: runIds } } }) : [] + 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])); } @@ -365,9 +389,9 @@ async function batchHydrateEdgeTarget( if (targetIds.length === 0) { return byParent; } - const rows = (await targetDelegate.findMany({ - where: { id: { in: [...new Set(targetIds)] } }, - })) as Record[]; + 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; From b9bf9743c7cd69c65f8868b98f67d7a73a8f0cb4 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 17:08:50 +0100 Subject: [PATCH 08/21] fix(webapp): schedule run-ops mint-kind flips to avoid duplicate runs Flipping an organization's run-ops mint kind (which database new runs are minted on) took effect independently per process as each cached value expired. For a window after a flip, two concurrent triggers using the same idempotency key could mint on different databases, where the per-database unique constraint can't dedupe them, producing a duplicate run. The flip is now a deterministic wall-clock cutover: the admin feature-flag routes stamp the previously-effective kind and a flip timestamp onto the organization, and the mint read resolves the old kind until flippedAt + grace (default 90s, chosen to outlast the caches a flip drains through), after which every process crosses to the new kind together. During the window all processes agree on one database, so a concurrent same-key collision lands on a single database and the existing unique-constraint retry returns the one idempotent run. A same-target re-save carries the in-flight stamp forward, so it can't slide the cutover. --- apps/webapp/app/env.server.ts | 4 + ...i.v1.orgs.$organizationId.feature-flags.ts | 18 ++- ...i.v2.orgs.$organizationId.feature-flags.ts | 13 ++ apps/webapp/app/v3/featureFlags.ts | 9 ++ .../v3/runOpsMigration/mintFlipGrace.test.ts | 153 ++++++++++++++++++ .../app/v3/runOpsMigration/mintFlipGrace.ts | 77 +++++++++ .../runOpsMintKind.flipLatency.test.ts | 16 +- .../runOpsMigration/runOpsMintKind.server.ts | 61 ++++++- 8 files changed, 330 insertions(+), 21 deletions(-) create mode 100644 apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts create mode 100644 apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts 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/routes/admin.api.v1.orgs.$organizationId.feature-flags.ts b/apps/webapp/app/routes/admin.api.v1.orgs.$organizationId.feature-flags.ts index e5fd7f7963..e247493799 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,9 +1,12 @@ 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 { stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; import { validatePartialFeatureFlags } from "~/v3/featureFlags"; const ParamsSchema = z.object({ @@ -82,10 +85,15 @@ export async function action({ request, params }: ActionFunctionArgs) { ? validatePartialFeatureFlags(organization.featureFlags as Record) : { success: false as const }; - const mergedFlags = { - ...(existingFlags.success ? existingFlags.data : {}), - ...validationResult.data, - }; + const mergedFlags = stampMintKindFlip( + existingFlags.success ? existingFlags.data : {}, + { + ...(existingFlags.success ? existingFlags.data : {}), + ...validationResult.data, + }, + Date.now(), + env.RUN_OPS_MINT_FLIP_GRACE_MS + ); // Update the organization's feature flags const updatedOrganization = await prisma.organization.update({ @@ -93,7 +101,7 @@ export async function action({ request, params }: ActionFunctionArgs) { id: organizationId, }, data: { - featureFlags: mergedFlags, + featureFlags: mergedFlags as Prisma.InputJsonValue, }, select: { id: true, 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..cf529024c5 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 { stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; import { flags as getGlobalFlags } from "~/v3/featureFlags.server"; import { FEATURE_FLAG, @@ -119,6 +121,17 @@ export async function action({ request, params }: ActionFunctionArgs) { ); } featureFlags = validationResult.data; + + const existingOrg = await prisma.organization.findFirst({ + where: { id: organizationId }, + select: { featureFlags: true }, + }); + featureFlags = stampMintKindFlip( + existingOrg?.featureFlags as Record | null | undefined, + featureFlags, + Date.now(), + env.RUN_OPS_MINT_FLIP_GRACE_MS + ); } try { 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/mintFlipGrace.test.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts new file mode 100644 index 0000000000..f3ebcce1aa --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts @@ -0,0 +1,153 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { effectiveMintKind, 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("defaults outgoing kind to 'cuid' when runOpsMintKind is absent", () => { + const existing = { runOpsMintKind: "runOpsId" }; + const outgoing: Record = {}; + const result = stampMintKindFlip(existing, outgoing, T, GRACE_MS); + + expect(result.runOpsMintKind).toBe("cuid"); + expect(result.runOpsMintKindPrev).toBe("runOpsId"); + expect(result.runOpsMintKindFlippedAt).toBe(new Date(T).toISOString()); + }); + + 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"); + }); +}); diff --git a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts new file mode 100644 index 0000000000..54d7ec236d --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts @@ -0,0 +1,77 @@ +import type { RunIdMintKind } from "./runOpsMintKind.server"; + +export type { RunIdMintKind }; + +export type MintFlagResolution = { + kind: RunIdMintKind; + prev?: RunIdMintKind; + flippedAtMs?: number; +}; + +const DEFAULT_MINT_KIND: RunIdMintKind = "cuid"; + +// Deterministic wall-clock cutover: during [flippedAt, flippedAt + graceMs) every +// process — stale (hasn't re-read the flag) or fresh (has the new stamp) — resolves to +// the OLD kind. At/after the cutover, every process resolves to the NEW kind. +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; +} + +function resolveEffectiveFromFlags( + flags: Record | null | undefined, + nowMs: number, + graceMs: number +): RunIdMintKind { + 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 effectiveMintKind({ kind, prev, flippedAtMs }, 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 { + const storedKind = readMintKind(existingFlags ?? {}, "runOpsMintKind") ?? DEFAULT_MINT_KIND; + const outgoingKind = readMintKind(outgoingFlags, "runOpsMintKind") ?? DEFAULT_MINT_KIND; + outgoingFlags.runOpsMintKind = outgoingKind; + + 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..51c942385b 100644 --- a/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts @@ -3,8 +3,10 @@ import { env } from "~/env.server"; import { logger } from "~/services/logger.server"; import { BoundedTtlCache } from "~/services/realtime/boundedTtlCache"; import { singleton } from "~/utils/singleton"; -import { FEATURE_FLAG } from "~/v3/featureFlags"; +import { FEATURE_FLAG, FeatureFlagCatalog } from "~/v3/featureFlags"; import { makeFlag } from "~/v3/featureFlags.server"; +import { DEFAULT_CP_CACHE_TTL_MS } from "./controlPlaneCache.server"; +import { effectiveMintKind, type MintFlagResolution } from "./mintFlipGrace"; import { isSplitEnabled } from "./splitMode.server"; export type RunIdMintKind = "cuid" | "runOpsId"; @@ -38,12 +40,38 @@ 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, + } + ); +} + +function readValidatedFlag( + overrides: Record, + key: T +): string | undefined { + const parsed = FeatureFlagCatalog[FEATURE_FLAG[key]].safeParse(overrides[FEATURE_FLAG[key]]); + return parsed.success ? parsed.data : undefined; +} + export async function resolveRunIdMintKind(environment: { organizationId: string; id: string; @@ -54,10 +82,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 +104,26 @@ export async function resolveRunIdMintKind(environment: { }) )?.featureFlags; + const overridesRecord = (overrides as Record) ?? {}; + const kind = await flagFn({ key: FEATURE_FLAG.runOpsMintKind, defaultValue: "cuid", - overrides: (overrides as Record) ?? {}, + overrides: overridesRecord, }); - mintCache.set(orgId, kind); - return kind; + + // Read the grace-linger stamp DIRECTLY from the already-in-memory org overrides — no + // extra DB read, and no forcing these optional fields through flagFn's defaultValue path. + const prev = readValidatedFlag(overridesRecord, "runOpsMintKindPrev") as + | RunIdMintKind + | undefined; + const flippedAtRaw = readValidatedFlag(overridesRecord, "runOpsMintKindFlippedAt"); + const parsedFlippedAt = flippedAtRaw !== undefined ? Date.parse(flippedAtRaw) : NaN; + const flippedAtMs = Number.isNaN(parsedFlippedAt) ? undefined : parsedFlippedAt; + + const resolution: MintFlagResolution = { kind, prev, flippedAtMs }; + mintCache.set(orgId, resolution); + return effectiveMintKind(resolution, Date.now(), env.RUN_OPS_MINT_FLIP_GRACE_MS); }, }); } From 54c79b6e1da5d18fdcd7673436abd67a547370fe Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 17:19:48 +0100 Subject: [PATCH 09/21] test(webapp): assert malformed mint-flip stamp is inert when resolved --- apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts index f3ebcce1aa..4fa7645dd4 100644 --- a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts @@ -149,5 +149,11 @@ describe("stampMintKindFlip", () => { 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"); }); }); From 98996c9c0524fc5f754b8591583e6b8736edf1a4 Mon Sep 17 00:00:00 2001 From: Daniel Sutton Date: Thu, 9 Jul 2026 18:10:15 +0100 Subject: [PATCH 10/21] fix(webapp): stamp the mint-kind flip time from the control-plane DB clock The grace-window cutover was timestamped with the webapp process's own clock, so clock skew between processes could shift the boundary. Read the time from the control-plane database instead, so the flip is anchored to one authoritative clock rather than whichever process handled the flag change. --- .../admin.api.v1.orgs.$organizationId.feature-flags.ts | 6 +++++- .../admin.api.v2.orgs.$organizationId.feature-flags.ts | 4 +++- 2 files changed, 8 insertions(+), 2 deletions(-) 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 e247493799..791047cf6e 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 @@ -85,13 +85,17 @@ export async function action({ request, params }: ActionFunctionArgs) { ? validatePartialFeatureFlags(organization.featureFlags as Record) : { success: false as const }; + // 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 happened to handle this request. + const [{ now: controlPlaneNow }] = await prisma.$queryRaw<{ now: Date }[]>`SELECT now() AS now`; + const mergedFlags = stampMintKindFlip( existingFlags.success ? existingFlags.data : {}, { ...(existingFlags.success ? existingFlags.data : {}), ...validationResult.data, }, - Date.now(), + controlPlaneNow.getTime(), env.RUN_OPS_MINT_FLIP_GRACE_MS ); 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 cf529024c5..708631ad24 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 @@ -126,10 +126,12 @@ export async function action({ request, params }: ActionFunctionArgs) { where: { id: organizationId }, select: { featureFlags: true }, }); + // Anchor the flip stamp to the control-plane DB clock (see the v1 route) rather than this process's. + const [{ now: controlPlaneNow }] = await prisma.$queryRaw<{ now: Date }[]>`SELECT now() AS now`; featureFlags = stampMintKindFlip( existingOrg?.featureFlags as Record | null | undefined, featureFlags, - Date.now(), + controlPlaneNow.getTime(), env.RUN_OPS_MINT_FLIP_GRACE_MS ); } From 2e839a9026437f6b1bc4f20937ed7bd79510f99c Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 13:05:03 +0100 Subject: [PATCH 11/21] docs(webapp): document the mint-kind flip grace clock-domain assumption The grace cutover compares the reader's wall clock against a DB-clock flippedAt stamp, so it assumes clock-synced hosts (skew well under the grace window). Records that invariant and the accepted residual on effectiveMintKind. No behavior change. --- apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts index 54d7ec236d..d868132969 100644 --- a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts @@ -10,9 +10,10 @@ export type MintFlagResolution = { const DEFAULT_MINT_KIND: RunIdMintKind = "cuid"; -// Deterministic wall-clock cutover: during [flippedAt, flippedAt + graceMs) every -// process — stale (hasn't re-read the flag) or fresh (has the new stamp) — resolves to -// the OLD kind. At/after the cutover, every process resolves to the NEW kind. +// 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, From 32e46e935069af15d0fbb69946210c02f98d4194 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 15:53:02 +0100 Subject: [PATCH 12/21] fix(webapp): extend mint-kind flip grace to global flips Global runOpsMintKind flips now stamp the previous kind and flip time like per-org flips, so a global flip is a deterministic cutover and cannot route two concurrent triggers that share an idempotency key to different databases. The resolver reads the grace stamp from the same source as the kind. --- .../app/routes/admin.api.v1.feature-flags.ts | 49 ++++++++++++- .../v3/runOpsMigration/mintFlipGrace.test.ts | 71 ++++++++++++++++++- .../app/v3/runOpsMigration/mintFlipGrace.ts | 33 +++++++-- .../runOpsMigration/runOpsMintKind.server.ts | 50 +++++++------ 4 files changed, 167 insertions(+), 36 deletions(-) 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..7187f72fdb 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,11 @@ 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 +26,50 @@ 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/v3/runOpsMigration/mintFlipGrace.test.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts index 4fa7645dd4..77bfe976ba 100644 --- a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts @@ -1,5 +1,10 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { effectiveMintKind, stampMintKindFlip, type MintFlagResolution } from "./mintFlipGrace"; +import { + effectiveMintKind, + resolveMintFlag, + 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 @@ -157,3 +162,67 @@ describe("stampMintKindFlip", () => { ).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, + }); + }); +}); diff --git a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts index d868132969..1945ca050e 100644 --- a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts @@ -30,19 +30,40 @@ function readMintKind(flags: Record, key: string): RunIdMintKin return value === "cuid" || value === "runOpsId" ? value : undefined; } -function resolveEffectiveFromFlags( - flags: Record | null | undefined, - nowMs: number, - graceMs: number -): RunIdMintKind { +// 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); +} - return effectiveMintKind({ kind, prev, flippedAtMs }, nowMs, graceMs); +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 diff --git a/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts b/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts index 51c942385b..90b96217e1 100644 --- a/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/runOpsMintKind.server.ts @@ -3,10 +3,9 @@ import { env } from "~/env.server"; import { logger } from "~/services/logger.server"; import { BoundedTtlCache } from "~/services/realtime/boundedTtlCache"; import { singleton } from "~/utils/singleton"; -import { FEATURE_FLAG, FeatureFlagCatalog } from "~/v3/featureFlags"; -import { makeFlag } from "~/v3/featureFlags.server"; +import { FEATURE_FLAG } from "~/v3/featureFlags"; import { DEFAULT_CP_CACHE_TTL_MS } from "./controlPlaneCache.server"; -import { effectiveMintKind, type MintFlagResolution } from "./mintFlipGrace"; +import { effectiveMintKind, resolveMintFlag, type MintFlagResolution } from "./mintFlipGrace"; import { isSplitEnabled } from "./splitMode.server"; export type RunIdMintKind = "cuid" | "runOpsId"; @@ -36,7 +35,6 @@ 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", () => @@ -64,14 +62,6 @@ if (env.RUN_OPS_MINT_FLIP_GRACE_MS <= env.RUN_OPS_MINT_FLAG_CACHE_TTL_MS + contr ); } -function readValidatedFlag( - overrides: Record, - key: T -): string | undefined { - const parsed = FeatureFlagCatalog[FEATURE_FLAG[key]].safeParse(overrides[FEATURE_FLAG[key]]); - return parsed.success ? parsed.data : undefined; -} - export async function resolveRunIdMintKind(environment: { organizationId: string; id: string; @@ -106,22 +96,30 @@ export async function resolveRunIdMintKind(environment: { const overridesRecord = (overrides as Record) ?? {}; - const kind = await flagFn({ - key: FEATURE_FLAG.runOpsMintKind, - defaultValue: "cuid", - overrides: overridesRecord, + // 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 }, }); + const globalFlags: Record = {}; + for (const row of globalRows) { + globalFlags[row.key] = row.value; + } - // Read the grace-linger stamp DIRECTLY from the already-in-memory org overrides — no - // extra DB read, and no forcing these optional fields through flagFn's defaultValue path. - const prev = readValidatedFlag(overridesRecord, "runOpsMintKindPrev") as - | RunIdMintKind - | undefined; - const flippedAtRaw = readValidatedFlag(overridesRecord, "runOpsMintKindFlippedAt"); - const parsedFlippedAt = flippedAtRaw !== undefined ? Date.parse(flippedAtRaw) : NaN; - const flippedAtMs = Number.isNaN(parsedFlippedAt) ? undefined : parsedFlippedAt; - - const resolution: MintFlagResolution = { kind, prev, flippedAtMs }; + // 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); }, From 8891394572a6d13a63aff718210eb4c9f60e01af Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 15:53:03 +0100 Subject: [PATCH 13/21] fix(run-store): fall back to the other store when a routed run read misses RoutingRunStore.findRun routed a classifiable id to its single owning store by id shape and returned null on a miss, so a run whose physical residency diverges from its id shape would 404 even though it exists. Read the owning store first, then fall back to the other store only on a miss, mirroring the unrouted fan-out. The found-in-owning-store path is unchanged (single read). --- ...OpsStore.residencyMismatchFallback.test.ts | 278 ++++++++++++++++++ .../run-store/src/runOpsStore.ts | 32 +- 2 files changed, 306 insertions(+), 4 deletions(-) create mode 100644 internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts 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..74122d292a --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts @@ -0,0 +1,278 @@ +// 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< + string, + unknown + > | 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 fa068b4c2f..ec0decce15 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -239,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 @@ -250,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, From 16df23f552c00a5b7addc449a8d6af7a0a2f28b4 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 16:17:17 +0100 Subject: [PATCH 14/21] chore(webapp,run-store): apply oxfmt to new files --- .../app/routes/admin.api.v1.feature-flags.ts | 11 +++-- ...resRunStore.hydratorSelectPushdown.test.ts | 8 +++- ...OpsStore.residencyMismatchFallback.test.ts | 44 +++++++++++++------ 3 files changed, 46 insertions(+), 17 deletions(-) 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 7187f72fdb..b67a9dde02 100644 --- a/apps/webapp/app/routes/admin.api.v1.feature-flags.ts +++ b/apps/webapp/app/routes/admin.api.v1.feature-flags.ts @@ -4,7 +4,11 @@ import { prisma } from "~/db.server"; import { env } from "~/env.server"; import { requireAdminApiRequest } from "~/services/personalAccessToken.server"; import { makeSetMultipleFlags } from "~/v3/featureFlags.server"; -import { FEATURE_FLAG, type FeatureFlagCatalog, 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) { @@ -57,8 +61,9 @@ export async function action({ request }: ActionFunctionArgs) { } // 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`; + const [{ now: controlPlaneNow }] = await prisma.$queryRaw< + { now: Date }[] + >`SELECT now() AS now`; flagsToWrite = stampMintKindFlip( existingGlobal, diff --git a/internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts b/internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts index ff1218bbfe..0db044bd94 100644 --- a/internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.hydratorSelectPushdown.test.ts @@ -79,7 +79,13 @@ function withCapturedTaskRunFindManyArgs(real: RunOpsPrismaClient) { const extended = (real as unknown as { $extends: (config: unknown) => unknown }).$extends({ query: { taskRun: { - findMany({ args, query }: { args: Record; query: (args: unknown) => Promise }) { + findMany({ + args, + query, + }: { + args: Record; + query: (args: unknown) => Promise; + }) { capturedArgs.push(args); return query(args); }, diff --git a/internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts b/internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts index 74122d292a..8897c29566 100644 --- a/internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts +++ b/internal-packages/run-store/src/runOpsStore.residencyMismatchFallback.test.ts @@ -40,13 +40,20 @@ function makeDedicatedStore(prisma17: RunOpsPrismaClient) { } function makeLegacyStore(prisma14: PrismaClient) { - return new PostgresRunStore({ prisma: prisma14, readOnlyPrisma: prisma14, schemaVariant: "legacy" }); + 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 { +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. @@ -139,7 +146,13 @@ function buildCreateRunInput(params: { // 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 } + params: { + runId: string; + friendlyId: string; + environmentId: string; + organizationId: string; + projectId: string; + } ) { await prisma17.taskRun.create({ data: { @@ -258,10 +271,10 @@ describe("RoutingRunStore.findRun — on-miss fan-out for a classifiable id (res newCounts.findRun = 0; legacyCounts.findRun = 0; - const hitLegacy = (await router.findRun({ id: legacyId }, { select: { id: true } })) as Record< - string, - unknown - > | null; + 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); @@ -269,10 +282,15 @@ describe("RoutingRunStore.findRun — on-miss fan-out for a classifiable id (res ); // ── 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(); - }); + 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(); + } + ); }); From bfb9ddf10896673018b130e7ac48bdfddd81cffb Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 16:21:33 +0100 Subject: [PATCH 15/21] fix(webapp): stop unrelated org flag saves from pinning the run-ops mint kind stampMintKindFlip injected the default mint kind into every org feature-flag save, so an unrelated flag change wrote an explicit per-org runOpsMintKind override, pinning the org to that kind and making a later global flip silently skip it. Only stamp when the save actually sets runOpsMintKind. --- .../app/v3/runOpsMigration/mintFlipGrace.test.ts | 13 ++++++++----- apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts | 8 ++++++-- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts index 77bfe976ba..1dd694f614 100644 --- a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts @@ -134,14 +134,17 @@ describe("stampMintKindFlip", () => { expect(result.runOpsMintKindFlippedAt).toBe(new Date(now).toISOString()); }); - it("defaults outgoing kind to 'cuid' when runOpsMintKind is absent", () => { + it("leaves runOpsMintKind untouched when the save omits it (unrelated flag change: no inject, no spurious flip)", () => { const existing = { runOpsMintKind: "runOpsId" }; - const outgoing: Record = {}; + const outgoing: Record = { someOtherFlag: true }; const result = stampMintKindFlip(existing, outgoing, T, GRACE_MS); - expect(result.runOpsMintKind).toBe("cuid"); - expect(result.runOpsMintKindPrev).toBe("runOpsId"); - expect(result.runOpsMintKindFlippedAt).toBe(new Date(T).toISOString()); + // 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", () => { diff --git a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts index 1945ca050e..9c099fc76a 100644 --- a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts @@ -75,9 +75,13 @@ export function stampMintKindFlip( 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; - const outgoingKind = readMintKind(outgoingFlags, "runOpsMintKind") ?? DEFAULT_MINT_KIND; - outgoingFlags.runOpsMintKind = outgoingKind; if (outgoingKind !== storedKind) { // Genuine target change: serve the currently-effective kind through the new grace window. From 4660a0a4dd125903c748aecaa2c99227baabe761 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 17:18:01 +0100 Subject: [PATCH 16/21] fix(webapp): schema-qualify waitpoint connected-run queries The raw queries reading a waitpoint's connected runs relied on search_path, so they broke on deployments using a non-public database schema. Qualify the tables with the configured schema so they resolve correctly. --- .../v3/WaitpointPresenter.server.ts | 44 ++-- ...ointPresenter.connectedRunsBounded.test.ts | 1 + .../waitpointPresenter.controlPlane.test.ts | 1 + ...intPresenter.danglingConnectedRuns.test.ts | 1 + ...dedicatedConnectedRuns.readthrough.test.ts | 1 + .../waitpointPresenter.readthrough.test.ts | 1 + ...waitpointPresenter.schemaQualified.test.ts | 236 ++++++++++++++++++ ...tpointPresenter.splitConnectedRuns.test.ts | 1 + 8 files changed, 270 insertions(+), 16 deletions(-) create mode 100644 apps/webapp/test/waitpointPresenter.schemaQualified.test.ts diff --git a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts index 5e9cd315c4..b0946fe50e 100644 --- a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts @@ -1,5 +1,9 @@ import { isWaitpointOutputTimeout, prettyPrintPacket } from "@trigger.dev/core/v3"; -import { type PrismaClientOrTransaction, type PrismaReplicaClient } from "~/db.server"; +import { + DATABASE_SCHEMA, + type PrismaClientOrTransaction, + type PrismaReplicaClient, +} from "~/db.server"; import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server"; import { generateHttpCallbackUrl } from "~/services/httpCallback.server"; import { logger } from "~/services/logger.server"; @@ -119,30 +123,38 @@ export class WaitpointPresenter extends BasePresenter { // deleted), the control-plane full schema the implicit `_WaitpointRunConnections` M2M // (A = TaskRun.id, B = Waitpoint.id). Both branches existence-filter AT THE QUERY via a JOIN to // TaskRun, so a dangling connection row can never occupy a LIMIT slot ahead of a real one. - // CONNECTED_RUNS_DISPLAY_LIMIT is a compile-time constant int, not an interpolated value. + // Tables are schema-qualified with DATABASE_SCHEMA (trusted boot constant) so a non-`public` + // schema= deployment resolves the right tables instead of leaning on search_path. $queryRawUnsafe, + // not a `sqlDatabaseSchema` Prisma.Sql fragment: `client` may be the dedicated run-ops client (a + // different Prisma runtime) which would bind a foreign runtime's Sql fragment as a param instead + // of inlining it. waitpointId and the constant limit stay bound params ($1/$2). async #connectedRunIdsOn(client: PrismaReplicaClient, waitpointId: string): Promise { const isDedicated = Boolean( (client as unknown as { waitpointRunConnection?: unknown }).waitpointRunConnection ); if (isDedicated) { - const rows = await client.$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_DISPLAY_LIMIT} - `; + const rows = await client.$queryRawUnsafe<{ taskRunId: string }[]>( + `SELECT c."taskRunId" AS "taskRunId" + FROM ${DATABASE_SCHEMA}."WaitpointRunConnection" c + JOIN ${DATABASE_SCHEMA}."TaskRun" t ON t."id" = c."taskRunId" + WHERE c."waitpointId" = $1 + LIMIT $2`, + waitpointId, + CONNECTED_RUNS_DISPLAY_LIMIT + ); return rows.map((row) => row.taskRunId); } - const rows = await client.$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_DISPLAY_LIMIT} - `; + const rows = await client.$queryRawUnsafe<{ A: string }[]>( + `SELECT c."A" AS "A" + FROM ${DATABASE_SCHEMA}."_WaitpointRunConnections" c + JOIN ${DATABASE_SCHEMA}."TaskRun" t ON t."id" = c."A" + WHERE c."B" = $1 + LIMIT $2`, + waitpointId, + CONNECTED_RUNS_DISPLAY_LIMIT + ); return rows.map((row) => row.A); } diff --git a/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts index 6a0e24b28b..7859a816f4 100644 --- a/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts +++ b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts @@ -35,6 +35,7 @@ vi.mock("~/db.server", async () => { runOpsLegacyPrisma: replicaProxy, runOpsLegacyReplica: replicaProxy, sqlDatabaseSchema: Prisma.sql([`public`]), + DATABASE_SCHEMA: "public", }; }); 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 index c7d0de5dbe..0c2565c545 100644 --- a/apps/webapp/test/waitpointPresenter.danglingConnectedRuns.test.ts +++ b/apps/webapp/test/waitpointPresenter.danglingConnectedRuns.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.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..5f327d21fe 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", }; }); diff --git a/apps/webapp/test/waitpointPresenter.schemaQualified.test.ts b/apps/webapp/test/waitpointPresenter.schemaQualified.test.ts new file mode 100644 index 0000000000..0a9a5bc953 --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.schemaQualified.test.ts @@ -0,0 +1,236 @@ +// RED->GREEN guard: #connectedRunIdsOn's raw SQL must schema-qualify every table reference with +// sqlDatabaseSchema (the repo-wide convention every other raw-SQL call site follows — +// task.server.ts, TestPresenter, DeploymentListPresenter). Unqualified names ("WaitpointRunConnection", +// "TaskRun", "_WaitpointRunConnections") resolve against search_path, so a non-`public` schema= +// deployment would hit the wrong schema or miss the tables entirely. +// +// This intercepts the REAL $queryRawUnsafe on BOTH physical clients (pure instrumentation over a +// real testcontainer client — the DB still executes the query and returns real rows) and captures +// the SQL text, then asserts every table reference is schema-qualified. Uses the split topology so +// BOTH raw-SQL branches run: dedicated `WaitpointRunConnection` (prisma17) and legacy implicit M2M +// `_WaitpointRunConnections` (prisma14). $queryRawUnsafe (not a Prisma.Sql fragment) is exactly what +// the fix uses: the two clients are different Prisma runtimes, so a foreign runtime's Sql fragment +// would be bound as a param rather than inlined. +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, + // The real schema constant is derived from DATABASE_URL's `schema=` search param. `public` keeps + // the query runnable against the testcontainer while still exercising qualification. + 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 { 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", + }, + }); +} + +// Capture the SQL text of every `$queryRawUnsafe` call, delegating unchanged to the real client +// (the DB still runs the query -- pure instrumentation, never a mock). Everything else passes +// straight through to the real client. +function capturingQueryRaw(real: T): { client: T; sqls: string[] } { + const sqls: string[] = []; + const client = new Proxy(real, { + get(target, prop, receiver) { + if (prop === "$queryRawUnsafe") { + return (query: string, ...params: unknown[]) => { + sqls.push(query); + return (target as any).$queryRawUnsafe(query, ...params); + }; + } + return Reflect.get(target, prop, receiver); + }, + }) as T; + return { client, sqls }; +} + +describe("WaitpointPresenter#connectedRunIdsOn schema-qualifies every raw-SQL table reference", () => { + heteroRunOpsPostgresTest( + "both the dedicated WaitpointRunConnection and legacy _WaitpointRunConnections branches qualify tables with sqlDatabaseSchema", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "schemaqual"); + + // Waitpoint resident on LEGACY (drives the implicit-M2M `_WaitpointRunConnections` branch). + const waitpoint = await prisma14.waitpoint.create({ + data: { + friendlyId: "waitpoint_schemaqual", + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: "idem-waitpoint_schemaqual", + userProvidedIdempotencyKey: false, + outputType: "application/json", + outputIsError: false, + completedAt: new Date(), + tags: [], + projectId: ctx.projectId, + environmentId: ctx.environmentId, + }, + }); + + // A connected run resident + joined on NEW (dedicated `WaitpointRunConnection` branch). + const newRun = await seedRun(prisma17, ctx, "run_schemaqual_new"); + await prisma17.waitpointRunConnection.create({ + data: { taskRunId: newRun.id, waitpointId: waitpoint.id }, + }); + + // A connected run resident + joined on LEGACY via the implicit M2M. + const legacyRun = await seedRun(prisma14, ctx, "run_schemaqual_legacy"); + await prisma14.waitpoint.update({ + where: { id: waitpoint.id }, + data: { connectedRuns: { connect: [{ id: legacyRun.id }] } }, + }); + + legacyReplicaHolder.client = prisma14; + newClientHolder.client = prisma17; + + const legacy = capturingQueryRaw(prisma14 as unknown as object); + const dedicated = capturingQueryRaw(prisma17 as unknown as object); + + const presenter = new WaitpointPresenter(undefined, undefined, { + splitEnabled: true, + newClient: dedicated.client as unknown as PrismaClient, + legacyReplica: legacy.client as unknown as PrismaClient, + }); + + const result = await presenter.call({ + friendlyId: waitpoint.friendlyId, + environmentId: ctx.environmentId, + projectId: ctx.projectId, + }); + + // The read still works end-to-end (schema-qualified names resolve against the testcontainer's + // public schema), and both stores contributed a connected run. + const returnedIds = (result?.connectedRuns ?? []).map((r) => r.friendlyId).sort(); + expect(returnedIds).toEqual(["run_schemaqual_legacy", "run_schemaqual_new"]); + + const allSqls = [...dedicated.sqls, ...legacy.sqls]; + // Sanity: both raw-SQL branches actually ran. + expect(dedicated.sqls.length).toBeGreaterThan(0); + expect(legacy.sqls.length).toBeGreaterThan(0); + + // Every table reference must be schema-qualified. The buggy (unqualified) code emits e.g. + // `FROM "WaitpointRunConnection"` / `JOIN "TaskRun"`, so these fail RED. + for (const sql of allSqls) { + // No unqualified table references (a `"` or `.` must precede any of these table names). + expect(sql).not.toMatch(/(? { runOpsLegacyPrisma: replicaProxy, runOpsLegacyReplica: replicaProxy, sqlDatabaseSchema: Prisma.sql([`public`]), + DATABASE_SCHEMA: "public", }; }); From ce00de1665d6d0df9edb4f59de630e99ca5f1cdd Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 17:18:01 +0100 Subject: [PATCH 17/21] fix(webapp): correct mint-flag grace baseline and lock down derived fields An org's first per-org mint-flag override computed its cutover grace window against the wrong baseline, which could skip the grace window and briefly let two concurrent triggers sharing an idempotency key resolve to different databases. Seed the baseline from the effective global flag, strip the derived stamp fields from the request body, and apply the read-and-stamp atomically. --- ...i.v1.orgs.$organizationId.feature-flags.ts | 108 ++++++++++-------- ...i.v2.orgs.$organizationId.feature-flags.ts | 94 ++++++++++----- .../v3/runOpsMigration/mintFlipGrace.test.ts | 91 +++++++++++++++ .../app/v3/runOpsMigration/mintFlipGrace.ts | 15 +++ 4 files changed, 232 insertions(+), 76 deletions(-) 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 791047cf6e..5430aacde2 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 @@ -6,8 +6,12 @@ import { env } from "~/env.server"; import { prisma } from "~/db.server"; import { requireAdminApiRequest } from "~/services/personalAccessToken.server"; import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; -import { stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; +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(), @@ -51,20 +55,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(); @@ -80,40 +70,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 }; - - // 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 happened to handle this request. - const [{ now: controlPlaneNow }] = await prisma.$queryRaw<{ now: Date }[]>`SELECT now() AS now`; - - const mergedFlags = stampMintKindFlip( - existingFlags.success ? existingFlags.data : {}, - { - ...(existingFlags.success ? existingFlags.data : {}), - ...validationResult.data, - }, - controlPlaneNow.getTime(), - env.RUN_OPS_MINT_FLIP_GRACE_MS - ); + // 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 as Prisma.InputJsonValue, - }, - 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 708631ad24..443b837f93 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 @@ -6,7 +6,10 @@ import { env } from "~/env.server"; import { prisma } from "~/db.server"; import { requireUser } from "~/services/session.server"; import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; -import { stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; +import { + selectMintBaselineSource, + stampMintKindFlip, +} from "~/v3/runOpsMigration/mintFlipGrace"; import { flags as getGlobalFlags } from "~/v3/featureFlags.server"; import { FEATURE_FLAG, @@ -105,47 +108,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; - const existingOrg = await prisma.organization.findFirst({ - where: { id: organizationId }, - select: { featureFlags: true }, - }); - // Anchor the flip stamp to the control-plane DB clock (see the v1 route) rather than this process's. - const [{ now: controlPlaneNow }] = await prisma.$queryRaw<{ now: Date }[]>`SELECT now() AS now`; - featureFlags = stampMintKindFlip( - existingOrg?.featureFlags as Record | null | undefined, - featureFlags, + controlPlaneResolver.invalidateOrganization(organizationId); + return json({ success: true }); + } + + 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 ); - } - try { - await prisma.organization.update({ + 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/runOpsMigration/mintFlipGrace.test.ts b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts index 1dd694f614..667068842f 100644 --- a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.test.ts @@ -1,7 +1,9 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { effectiveMintKind, + readMintResolution, resolveMintFlag, + selectMintBaselineSource, stampMintKindFlip, type MintFlagResolution, } from "./mintFlipGrace"; @@ -229,3 +231,92 @@ describe("resolveMintFlag", () => { }); }); }); + +// #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 index 9c099fc76a..7e2366edcd 100644 --- a/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts +++ b/apps/webapp/app/v3/runOpsMigration/mintFlipGrace.ts @@ -58,6 +58,21 @@ export function resolveMintFlag( 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, From c60dd3b5ce5e70eee520ebddea8325eb81ac04c6 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 17:18:01 +0100 Subject: [PATCH 18/21] fix(run-store): bound connectedRuns relation hydrator on the dedicated schema The dedicated-schema hydrator for a waitpoint's connectedRuns relation fetched an unbounded list of connections and full run rows. Bound it per parent with a window function so a waitpoint with many connections cannot pull a large result set for a display-only field. --- ...nStore.connectedRunsIncludeBounded.test.ts | 127 ++++++++++++++++++ .../PostgresRunStore.groupedRelations.test.ts | 19 ++- .../run-store/src/PostgresRunStore.ts | 51 +++++-- 3 files changed, 183 insertions(+), 14 deletions(-) create mode 100644 internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts 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..1b70e62636 --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts @@ -0,0 +1,127 @@ +// 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. + expect(connectedRuns.length).toBeLessThanOrEqual(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.groupedRelations.test.ts b/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts index 948b904748..6323ca4858 100644 --- a/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts @@ -19,11 +19,12 @@ type CallCounts = { findMany: number; findFirst: number }; function countingClient( real: RunOpsPrismaClient, delegateNames: string[] -): { client: RunOpsPrismaClient; counts: Record } { +): { 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, @@ -42,10 +43,17 @@ function countingClient( 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 }; + return { client, counts, queryRaw }; } function seedEnvironmentDedicated(suffix: string) { @@ -219,8 +227,11 @@ describe("PostgresRunStore dedicated relation hydrators — grouped batch reads" expect(connected[0].id).toBe(runIds[i]); } - // GROUPED: one join query + one target query for the WHOLE batch, never one pair per parent. - expect(counting.counts.waitpointRunConnection.findMany).toBe(1); + // 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.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 415340dcd4..fa8df51f28 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -336,16 +336,47 @@ const hydrateBlockingTaskRuns: DedicatedRelationHydrator = async (client, parent return byParent; }; -// Display connections for a waitpoint: WaitpointRunConnection → TaskRun rows. -const hydrateConnectedRuns: DedicatedRelationHydrator = async (client, parents, projection) => - batchHydrateJoinRelation( - client.waitpointRunConnection, - client.taskRun, - parents.map((p) => p.id as string), - "waitpointId", - "taskRunId", - projection - ); +// 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 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)); + } + } + return byParent; +}; // Snapshots that completed a waitpoint: CompletedWaitpoint join → TaskRunExecutionSnapshot rows. const hydrateCompletedExecutionSnapshots: DedicatedRelationHydrator = async ( From bb008c9ed38ad71aea0344a982f5e5cf5149627d Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 17:28:41 +0100 Subject: [PATCH 19/21] chore(webapp,run-store): apply formatter --- .../admin.api.v1.orgs.$organizationId.feature-flags.ts | 5 +---- .../admin.api.v2.orgs.$organizationId.feature-flags.ts | 5 +---- .../src/PostgresRunStore.connectedRunsIncludeBounded.test.ts | 5 ++++- 3 files changed, 6 insertions(+), 9 deletions(-) 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 5430aacde2..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 @@ -6,10 +6,7 @@ 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 { selectMintBaselineSource, stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; import { validatePartialFeatureFlags } from "~/v3/featureFlags"; import { flags as getGlobalFlags } from "~/v3/featureFlags.server"; 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 443b837f93..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 @@ -6,10 +6,7 @@ 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 { selectMintBaselineSource, stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace"; import { flags as getGlobalFlags } from "~/v3/featureFlags.server"; import { FEATURE_FLAG, diff --git a/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts b/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts index 1b70e62636..6a5da36064 100644 --- a/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts @@ -89,7 +89,10 @@ describe("PostgresRunStore.findWaitpoint — connectedRuns relation is bounded", } // More REAL connected runs than the limit. - const realRunIds = Array.from({ length: CONNECTED_RUNS_LIMIT + 3 }, (_, i) => `run_cr_real_${i}`); + 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({ From 9c10e14fd935ed847592fa9950b89908cd8340ea Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 17:40:50 +0100 Subject: [PATCH 20/21] test(run-store): assert exact connectedRuns cap in bound test --- .../src/PostgresRunStore.connectedRunsIncludeBounded.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts b/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts index 6a5da36064..a2ce43103c 100644 --- a/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.connectedRunsIncludeBounded.test.ts @@ -118,8 +118,8 @@ describe("PostgresRunStore.findWaitpoint — connectedRuns relation is bounded", expect(waitpoint).not.toBeNull(); const connectedRuns = waitpoint!.connectedRuns; - // Bounded. - expect(connectedRuns.length).toBeLessThanOrEqual(CONNECTED_RUNS_LIMIT); + // 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); From 7db8e466c32da455cb00e96a44c51fcb92496502 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 10 Jul 2026 18:39:23 +0100 Subject: [PATCH 21/21] refactor(webapp): read waitpoint connected runs via ORM instead of raw SQL Replaces the raw connected-runs query, which JOINed the waitpoint connection table onto the large TaskRun table, with two indexed ORM reads joined in memory: read the connection rows, then resolve the runs by id. The id lookup can only plan as a primary-key lookup, so the query planner never scans TaskRun, and Prisma qualifies the schema per client so the previous manual schema handling is gone. --- .../v3/WaitpointPresenter.server.ts | 112 ++++----- ...ointPresenter.connectedRunsBounded.test.ts | 35 +-- .../waitpointPresenter.readthrough.test.ts | 42 +++- ...waitpointPresenter.schemaQualified.test.ts | 236 ------------------ 4 files changed, 107 insertions(+), 318 deletions(-) delete mode 100644 apps/webapp/test/waitpointPresenter.schemaQualified.test.ts diff --git a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts index b0946fe50e..29ef066572 100644 --- a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts @@ -1,9 +1,5 @@ import { isWaitpointOutputTimeout, prettyPrintPacket } from "@trigger.dev/core/v3"; -import { - DATABASE_SCHEMA, - type PrismaClientOrTransaction, - type PrismaReplicaClient, -} from "~/db.server"; +import { type PrismaClientOrTransaction, type PrismaReplicaClient } from "~/db.server"; import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server"; import { generateHttpCallbackUrl } from "~/services/httpCallback.server"; import { logger } from "~/services/logger.server"; @@ -15,10 +11,12 @@ import { waitpointStatusToApiStatus } from "./WaitpointListPresenter.server"; export type WaitpointDetail = NonNullable>>; -// Single-sourced bound for connected run friendlyIds: applied at the FETCH in #connectedRunIdsOn, -// not just at display time. +// 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, @@ -84,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[] = @@ -99,17 +96,8 @@ 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; - } - const runs = await client.taskRun.findMany({ - where: { id: { in: runIds } }, - select: { friendlyId: true }, - take: CONNECTED_RUNS_DISPLAY_LIMIT, - }); - for (const run of runs) { - friendlyIds.add(run.friendlyId); + for (const friendlyId of await this.#connectedRunFriendlyIdsOn(client, waitpointId)) { + friendlyIds.add(friendlyId); } if (friendlyIds.size >= CONNECTED_RUNS_DISPLAY_LIMIT) { break; @@ -118,44 +106,56 @@ export class WaitpointPresenter extends BasePresenter { 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 (scalar `taskRunId`, no FK -- a row can dangle after its run is - // deleted), the control-plane full schema the implicit `_WaitpointRunConnections` M2M - // (A = TaskRun.id, B = Waitpoint.id). Both branches existence-filter AT THE QUERY via a JOIN to - // TaskRun, so a dangling connection row can never occupy a LIMIT slot ahead of a real one. - // Tables are schema-qualified with DATABASE_SCHEMA (trusted boot constant) so a non-`public` - // schema= deployment resolves the right tables instead of leaning on search_path. $queryRawUnsafe, - // not a `sqlDatabaseSchema` Prisma.Sql fragment: `client` may be the dedicated run-ops client (a - // different Prisma runtime) which would bind a foreign runtime's Sql fragment as a param instead - // of inlining it. waitpointId and the constant limit stay bound params ($1/$2). - async #connectedRunIdsOn(client: PrismaReplicaClient, waitpointId: string): Promise { - const isDedicated = Boolean( - (client as unknown as { waitpointRunConnection?: unknown }).waitpointRunConnection - ); - - if (isDedicated) { - const rows = await client.$queryRawUnsafe<{ taskRunId: string }[]>( - `SELECT c."taskRunId" AS "taskRunId" - FROM ${DATABASE_SCHEMA}."WaitpointRunConnection" c - JOIN ${DATABASE_SCHEMA}."TaskRun" t ON t."id" = c."taskRunId" - WHERE c."waitpointId" = $1 - LIMIT $2`, - waitpointId, - CONNECTED_RUNS_DISPLAY_LIMIT - ); - return rows.map((row) => row.taskRunId); + // 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 (dedicated) { + const connections = await dedicated.findMany({ + where: { waitpointId }, + select: { taskRunId: true }, + take: CONNECTED_RUNS_CONNECTION_SCAN_LIMIT, + }); + 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.$queryRawUnsafe<{ A: string }[]>( - `SELECT c."A" AS "A" - FROM ${DATABASE_SCHEMA}."_WaitpointRunConnections" c - JOIN ${DATABASE_SCHEMA}."TaskRun" t ON t."id" = c."A" - WHERE c."B" = $1 - LIMIT $2`, - waitpointId, - CONNECTED_RUNS_DISPLAY_LIMIT - ); - 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/test/waitpointPresenter.connectedRunsBounded.test.ts b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts index 7859a816f4..a8263b80e3 100644 --- a/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts +++ b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts @@ -1,12 +1,11 @@ -// RED->GREEN guard: WaitpointPresenter#connectedRunIdsOn must BOUND the fetch of a waitpoint's +// 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 on today's buggy code (the existing take:5 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 itself is bounded. +// 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 of #connectedRunIdsOn: -// waitpoint + 8 connected runs seeded on the NEW dedicated run-ops client (RunOpsPrismaClient, -// prisma17), which uses the delegate directly (no raw SQL fallback). +// 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 })); @@ -66,7 +65,7 @@ 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, + CONNECTED_RUNS_CONNECTION_SCAN_LIMIT, WaitpointPresenter, } from "~/presenters/v3/WaitpointPresenter.server"; @@ -184,11 +183,13 @@ function capturingTaskRunFindMany(real: RunOpsPrismaClient): { return { client, calls }; } -const CONNECTED_RUN_COUNT = 8; // > CONNECTED_RUNS_DISPLAY_LIMIT (5), proves the fetch isn't bounded +// 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#connectedRunIdsOn bounds the connected-run-id FETCH", () => { +describe("WaitpointPresenter bounds the connected-run-id FETCH", () => { heteroRunOpsPostgresTest( - "reading a waitpoint with 8 connected runs never IN-lists more than the display limit", + "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"); @@ -220,13 +221,15 @@ describe("WaitpointPresenter#connectedRunIdsOn bounds the connected-run-id FETCH }); // The guard: assert the FETCH (the IN-list built from the connected-run-id gather) is - // bounded, not just the eventual displayed count. On today's buggy code, the Prisma - // `waitpointRunConnection.findMany` branch of #connectedRunIdsOn returns all 8 connected - // run ids with no take/LIMIT, so this IN-list is 8. After the fix it must be - // <= CONNECTED_RUNS_DISPLAY_LIMIT. + // 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_DISPLAY_LIMIT); + expect(call.where?.id?.in?.length ?? 0).toBeLessThanOrEqual( + CONNECTED_RUNS_CONNECTION_SCAN_LIMIT + ); } } ); diff --git a/apps/webapp/test/waitpointPresenter.readthrough.test.ts b/apps/webapp/test/waitpointPresenter.readthrough.test.ts index 5f327d21fe..72b794b30d 100644 --- a/apps/webapp/test/waitpointPresenter.readthrough.test.ts +++ b/apps/webapp/test/waitpointPresenter.readthrough.test.ts @@ -68,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) { @@ -82,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); }; @@ -96,6 +110,8 @@ function recording(client: PrismaClient, opts: { forbidden?: boolean } = {}) { }, }) as unknown as PrismaReplicaClient, calls, + hydrateCalls, + connectedRunsCalls, }; } @@ -255,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); } ); @@ -286,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.schemaQualified.test.ts b/apps/webapp/test/waitpointPresenter.schemaQualified.test.ts deleted file mode 100644 index 0a9a5bc953..0000000000 --- a/apps/webapp/test/waitpointPresenter.schemaQualified.test.ts +++ /dev/null @@ -1,236 +0,0 @@ -// RED->GREEN guard: #connectedRunIdsOn's raw SQL must schema-qualify every table reference with -// sqlDatabaseSchema (the repo-wide convention every other raw-SQL call site follows — -// task.server.ts, TestPresenter, DeploymentListPresenter). Unqualified names ("WaitpointRunConnection", -// "TaskRun", "_WaitpointRunConnections") resolve against search_path, so a non-`public` schema= -// deployment would hit the wrong schema or miss the tables entirely. -// -// This intercepts the REAL $queryRawUnsafe on BOTH physical clients (pure instrumentation over a -// real testcontainer client — the DB still executes the query and returns real rows) and captures -// the SQL text, then asserts every table reference is schema-qualified. Uses the split topology so -// BOTH raw-SQL branches run: dedicated `WaitpointRunConnection` (prisma17) and legacy implicit M2M -// `_WaitpointRunConnections` (prisma14). $queryRawUnsafe (not a Prisma.Sql fragment) is exactly what -// the fix uses: the two clients are different Prisma runtimes, so a foreign runtime's Sql fragment -// would be bound as a param rather than inlined. -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, - // The real schema constant is derived from DATABASE_URL's `schema=` search param. `public` keeps - // the query runnable against the testcontainer while still exercising qualification. - 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 { 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", - }, - }); -} - -// Capture the SQL text of every `$queryRawUnsafe` call, delegating unchanged to the real client -// (the DB still runs the query -- pure instrumentation, never a mock). Everything else passes -// straight through to the real client. -function capturingQueryRaw(real: T): { client: T; sqls: string[] } { - const sqls: string[] = []; - const client = new Proxy(real, { - get(target, prop, receiver) { - if (prop === "$queryRawUnsafe") { - return (query: string, ...params: unknown[]) => { - sqls.push(query); - return (target as any).$queryRawUnsafe(query, ...params); - }; - } - return Reflect.get(target, prop, receiver); - }, - }) as T; - return { client, sqls }; -} - -describe("WaitpointPresenter#connectedRunIdsOn schema-qualifies every raw-SQL table reference", () => { - heteroRunOpsPostgresTest( - "both the dedicated WaitpointRunConnection and legacy _WaitpointRunConnections branches qualify tables with sqlDatabaseSchema", - async ({ prisma14, prisma17 }) => { - const ctx = await seedParents(prisma14, "schemaqual"); - - // Waitpoint resident on LEGACY (drives the implicit-M2M `_WaitpointRunConnections` branch). - const waitpoint = await prisma14.waitpoint.create({ - data: { - friendlyId: "waitpoint_schemaqual", - type: "MANUAL", - status: "COMPLETED", - idempotencyKey: "idem-waitpoint_schemaqual", - userProvidedIdempotencyKey: false, - outputType: "application/json", - outputIsError: false, - completedAt: new Date(), - tags: [], - projectId: ctx.projectId, - environmentId: ctx.environmentId, - }, - }); - - // A connected run resident + joined on NEW (dedicated `WaitpointRunConnection` branch). - const newRun = await seedRun(prisma17, ctx, "run_schemaqual_new"); - await prisma17.waitpointRunConnection.create({ - data: { taskRunId: newRun.id, waitpointId: waitpoint.id }, - }); - - // A connected run resident + joined on LEGACY via the implicit M2M. - const legacyRun = await seedRun(prisma14, ctx, "run_schemaqual_legacy"); - await prisma14.waitpoint.update({ - where: { id: waitpoint.id }, - data: { connectedRuns: { connect: [{ id: legacyRun.id }] } }, - }); - - legacyReplicaHolder.client = prisma14; - newClientHolder.client = prisma17; - - const legacy = capturingQueryRaw(prisma14 as unknown as object); - const dedicated = capturingQueryRaw(prisma17 as unknown as object); - - const presenter = new WaitpointPresenter(undefined, undefined, { - splitEnabled: true, - newClient: dedicated.client as unknown as PrismaClient, - legacyReplica: legacy.client as unknown as PrismaClient, - }); - - const result = await presenter.call({ - friendlyId: waitpoint.friendlyId, - environmentId: ctx.environmentId, - projectId: ctx.projectId, - }); - - // The read still works end-to-end (schema-qualified names resolve against the testcontainer's - // public schema), and both stores contributed a connected run. - const returnedIds = (result?.connectedRuns ?? []).map((r) => r.friendlyId).sort(); - expect(returnedIds).toEqual(["run_schemaqual_legacy", "run_schemaqual_new"]); - - const allSqls = [...dedicated.sqls, ...legacy.sqls]; - // Sanity: both raw-SQL branches actually ran. - expect(dedicated.sqls.length).toBeGreaterThan(0); - expect(legacy.sqls.length).toBeGreaterThan(0); - - // Every table reference must be schema-qualified. The buggy (unqualified) code emits e.g. - // `FROM "WaitpointRunConnection"` / `JOIN "TaskRun"`, so these fail RED. - for (const sql of allSqls) { - // No unqualified table references (a `"` or `.` must precede any of these table names). - expect(sql).not.toMatch(/(?