From c5226a2dc079eff0e1ce8a3a4c2277659810ebd4 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 14:47:26 +0100 Subject: [PATCH 01/11] feat(run-store): add TaskRun read methods to the run store Add findRun, findRunOrThrow and findRuns to RunStore, mirroring the existing write methods. They pass where/select/include through the same Prisma generics and default to the read replica, while letting the caller pass the writer or a transaction client when needed. This lets Postgres reads of TaskRun be routed through the store the same way writes already are. Additive only; no call sites change yet. --- .../run-store/src/NoopRunStore.ts | 3 + .../run-store/src/PostgresRunStore.test.ts | 156 ++++++++++++++++++ .../run-store/src/PostgresRunStore.ts | 85 ++++++++++ internal-packages/run-store/src/types.ts | 46 ++++++ 4 files changed, 290 insertions(+) diff --git a/internal-packages/run-store/src/NoopRunStore.ts b/internal-packages/run-store/src/NoopRunStore.ts index 3b4fb0a36f..e27080c9af 100644 --- a/internal-packages/run-store/src/NoopRunStore.ts +++ b/internal-packages/run-store/src/NoopRunStore.ts @@ -29,4 +29,7 @@ export class NoopRunStore implements RunStore { clearIdempotencyKey(): never { return this.fail("clearIdempotencyKey"); } pushTags(): never { return this.fail("pushTags"); } pushRealtimeStream(): never { return this.fail("pushRealtimeStream"); } + findRun(): never { return this.fail("findRun"); } + findRunOrThrow(): never { return this.fail("findRunOrThrow"); } + findRuns(): never { return this.fail("findRuns"); } } diff --git a/internal-packages/run-store/src/PostgresRunStore.test.ts b/internal-packages/run-store/src/PostgresRunStore.test.ts index f2fb2969e6..8540912c99 100644 --- a/internal-packages/run-store/src/PostgresRunStore.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.test.ts @@ -1528,3 +1528,159 @@ describe("PostgresRunStore — delayed / debounce / metadata / idempotency / arr } ); }); + +describe("PostgresRunStore — read", () => { + postgresTest("findRun by id with select returns the projected row", async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + const runId = "run_find_select_id_1"; + + await store.createRun( + buildCreateRunInput({ + runId, + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const run = await store.findRun({ id: runId }, { select: { friendlyId: true } }); + + expect(run).toEqual({ friendlyId: "run_friendly_1" }); + }); + + postgresTest("findRun by friendlyId with select returns the matching row", async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + const runId = "run_find_select_friendly_1"; + + await store.createRun( + buildCreateRunInput({ + runId, + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const run = await store.findRun({ friendlyId: "run_friendly_1" }, { select: { id: true } }); + + expect(run?.id).toBe(runId); + }); + + postgresTest("findRun returns null when no row matches", async ({ prisma }) => { + await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + + const run = await store.findRun({ id: "missing" }, { select: { id: true } }); + + expect(run).toBeNull(); + }); + + postgresTest("findRunOrThrow throws when no row matches", async ({ prisma }) => { + await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + + await expect(store.findRunOrThrow({ id: "missing" }, { select: { id: true } })).rejects.toThrow(); + }); + + postgresTest("findRun with include hydrates the relation", async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + const runId = "run_find_include_1"; + + await store.createRun( + buildCreateRunInput({ + runId, + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const run = await store.findRun({ id: runId }, { include: { runtimeEnvironment: true } }); + + expect(run?.id).toBe(runId); + expect(run?.runtimeEnvironment).toBeDefined(); + expect(run?.runtimeEnvironment.id).toBe(environment.id); + }); + + postgresTest("findRuns applies where/orderBy/take and returns ordered, limited rows", async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + + const earliest = new Date("2026-06-01T00:00:00.000Z"); + const middle = new Date("2026-06-02T00:00:00.000Z"); + const latest = new Date("2026-06-03T00:00:00.000Z"); + + const rows: Array<{ id: string; createdAt: Date }> = [ + { id: "run_find_many_earliest", createdAt: earliest }, + { id: "run_find_many_middle", createdAt: middle }, + { id: "run_find_many_latest", createdAt: latest }, + ]; + + for (const row of rows) { + await prisma.taskRun.create({ + data: { + id: row.id, + engine: "V2", + status: "PENDING", + friendlyId: `${row.id}_friendly`, + runtimeEnvironmentId: environment.id, + environmentType: "DEVELOPMENT", + organizationId: organization.id, + projectId: project.id, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${row.id}`, + spanId: `span_${row.id}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: row.createdAt, + }, + }); + } + + const found = await store.findRuns({ + where: { projectId: project.id }, + select: { id: true }, + orderBy: { createdAt: "desc" }, + take: 2, + }); + + expect(found).toEqual([{ id: "run_find_many_latest" }, { id: "run_find_many_middle" }]); + }); + + postgresTest("findRun reads a just-written row when passed the writer client", async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + // Use a NoopRunStore-style read replica that must NOT be hit: pass the writer + // (prisma) explicitly so reads go through it for read-after-write consistency. + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + const runId = "run_find_read_after_write_1"; + + await store.createRun( + buildCreateRunInput({ + runId, + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const run = await store.findRun({ id: runId }, { select: { id: true, status: true } }, prisma); + + expect(run?.id).toBe(runId); + expect(run?.status).toBe("PENDING"); + }); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 925a39425b..21514ea44d 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -617,4 +617,89 @@ export class PostgresRunStore implements RunStore { data: { realtimeStreams: { push: streamId } }, }); } + + findRun( + where: Prisma.TaskRunWhereInput, + args: { select: S }, + client?: PrismaClientOrTransaction + ): Promise | null>; + findRun( + where: Prisma.TaskRunWhereInput, + args: { include: I }, + client?: PrismaClientOrTransaction + ): Promise | null>; + async findRun( + where: Prisma.TaskRunWhereInput, + args: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }, + client?: PrismaClientOrTransaction + ): Promise { + const prisma = client ?? this.readOnlyPrisma; + + return prisma.taskRun.findFirst({ + where, + ...args, + }); + } + + findRunOrThrow( + where: Prisma.TaskRunWhereInput, + args: { select: S }, + client?: PrismaClientOrTransaction + ): Promise>; + findRunOrThrow( + where: Prisma.TaskRunWhereInput, + args: { include: I }, + client?: PrismaClientOrTransaction + ): Promise>; + async findRunOrThrow( + where: Prisma.TaskRunWhereInput, + args: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }, + client?: PrismaClientOrTransaction + ): Promise { + const prisma = client ?? this.readOnlyPrisma; + + return prisma.taskRun.findFirstOrThrow({ + where, + ...args, + }); + } + + findRuns( + args: { + where: Prisma.TaskRunWhereInput; + select: S; + orderBy?: Prisma.TaskRunOrderByWithRelationInput | Prisma.TaskRunOrderByWithRelationInput[]; + take?: number; + skip?: number; + cursor?: Prisma.TaskRunWhereUniqueInput; + }, + client?: PrismaClientOrTransaction + ): Promise[]>; + findRuns( + args: { + where: Prisma.TaskRunWhereInput; + include: I; + orderBy?: Prisma.TaskRunOrderByWithRelationInput | Prisma.TaskRunOrderByWithRelationInput[]; + take?: number; + skip?: number; + cursor?: Prisma.TaskRunWhereUniqueInput; + }, + client?: PrismaClientOrTransaction + ): Promise[]>; + async findRuns( + args: { + where: Prisma.TaskRunWhereInput; + select?: Prisma.TaskRunSelect; + include?: Prisma.TaskRunInclude; + orderBy?: Prisma.TaskRunOrderByWithRelationInput | Prisma.TaskRunOrderByWithRelationInput[]; + take?: number; + skip?: number; + cursor?: Prisma.TaskRunWhereUniqueInput; + }, + client?: PrismaClientOrTransaction + ): Promise { + const prisma = client ?? this.readOnlyPrisma; + + return prisma.taskRun.findMany(args); + } } diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index e680f25463..35d4d8f91a 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -319,4 +319,50 @@ export interface RunStore { clearIdempotencyKey(params: ClearIdempotencyKeyInput, tx?: PrismaClientOrTransaction): Promise<{ count: number }>; pushTags(runId: string, tags: string[], where: { runtimeEnvironmentId: string }, tx?: PrismaClientOrTransaction): Promise<{ updatedAt: Date }>; pushRealtimeStream(runId: string, streamId: string, tx?: PrismaClientOrTransaction): Promise; + + // Read + findRun( + where: Prisma.TaskRunWhereInput, + args: { select: S }, + client?: PrismaClientOrTransaction + ): Promise | null>; + findRun( + where: Prisma.TaskRunWhereInput, + args: { include: I }, + client?: PrismaClientOrTransaction + ): Promise | null>; + + findRunOrThrow( + where: Prisma.TaskRunWhereInput, + args: { select: S }, + client?: PrismaClientOrTransaction + ): Promise>; + findRunOrThrow( + where: Prisma.TaskRunWhereInput, + args: { include: I }, + client?: PrismaClientOrTransaction + ): Promise>; + + findRuns( + args: { + where: Prisma.TaskRunWhereInput; + select: S; + orderBy?: Prisma.TaskRunOrderByWithRelationInput | Prisma.TaskRunOrderByWithRelationInput[]; + take?: number; + skip?: number; + cursor?: Prisma.TaskRunWhereUniqueInput; + }, + client?: PrismaClientOrTransaction + ): Promise[]>; + findRuns( + args: { + where: Prisma.TaskRunWhereInput; + include: I; + orderBy?: Prisma.TaskRunOrderByWithRelationInput | Prisma.TaskRunOrderByWithRelationInput[]; + take?: number; + skip?: number; + cursor?: Prisma.TaskRunWhereUniqueInput; + }, + client?: PrismaClientOrTransaction + ): Promise[]>; } From 13d53648b1885938480384c8689651f8c418d822 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 15:31:09 +0100 Subject: [PATCH 02/11] feat(run-store): add full-row read overload to the run store Add a no-args overload to findRun, findRunOrThrow and findRuns that returns the whole TaskRun row, for callers that read a run without a select or include. --- .../run-store/src/PostgresRunStore.test.ts | 88 +++++++++++++++++++ .../run-store/src/PostgresRunStore.ts | 63 ++++++++++++- internal-packages/run-store/src/types.ts | 12 +++ 3 files changed, 159 insertions(+), 4 deletions(-) diff --git a/internal-packages/run-store/src/PostgresRunStore.test.ts b/internal-packages/run-store/src/PostgresRunStore.test.ts index 8540912c99..47876b70c8 100644 --- a/internal-packages/run-store/src/PostgresRunStore.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.test.ts @@ -1683,4 +1683,92 @@ describe("PostgresRunStore — read", () => { expect(run?.id).toBe(runId); expect(run?.status).toBe("PENDING"); }); + + postgresTest("findRun by id with no projection returns the whole row", async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + const runId = "run_find_full_row_1"; + + await store.createRun( + buildCreateRunInput({ + runId, + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const run = await store.findRun({ id: runId }); + + expect(run?.id).toBe(runId); + expect(run?.friendlyId).toBe("run_friendly_1"); + expect(run?.status).toBe("PENDING"); + expect(run?.taskIdentifier).toBe("my-task"); + // The whole-row variant returns the full scalar set, not a projection. + expect(run?.payload).toBe("{}"); + expect(run?.payloadType).toBe("application/json"); + }); + + postgresTest("findRunOrThrow with no projection throws when no row matches", async ({ prisma }) => { + await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + + await expect(store.findRunOrThrow({ id: "missing" })).rejects.toThrow(); + }); + + postgresTest("findRuns with no projection returns whole rows", async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + + const earliest = new Date("2026-07-01T00:00:00.000Z"); + const latest = new Date("2026-07-02T00:00:00.000Z"); + + const rows: Array<{ id: string; createdAt: Date }> = [ + { id: "run_find_full_many_earliest", createdAt: earliest }, + { id: "run_find_full_many_latest", createdAt: latest }, + ]; + + for (const row of rows) { + await prisma.taskRun.create({ + data: { + id: row.id, + engine: "V2", + status: "PENDING", + friendlyId: `${row.id}_friendly`, + runtimeEnvironmentId: environment.id, + environmentType: "DEVELOPMENT", + organizationId: organization.id, + projectId: project.id, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${row.id}`, + spanId: `span_${row.id}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: row.createdAt, + }, + }); + } + + const found = await store.findRuns({ + where: { projectId: project.id }, + orderBy: { createdAt: "desc" }, + }); + + expect(found).toHaveLength(2); + expect(found.map((r) => r.id)).toEqual([ + "run_find_full_many_latest", + "run_find_full_many_earliest", + ]); + // Whole rows include full scalar columns. + expect(found[0]?.taskIdentifier).toBe("my-task"); + expect(found[0]?.payloadType).toBe("application/json"); + }); }); diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 21514ea44d..fcc53c0026 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -628,12 +628,16 @@ export class PostgresRunStore implements RunStore { args: { include: I }, client?: PrismaClientOrTransaction ): Promise | null>; + findRun( + where: Prisma.TaskRunWhereInput, + client?: PrismaClientOrTransaction + ): Promise; async findRun( where: Prisma.TaskRunWhereInput, - args: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }, + argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | PrismaClientOrTransaction, client?: PrismaClientOrTransaction ): Promise { - const prisma = client ?? this.readOnlyPrisma; + const { args, prisma } = this.#resolveReadArgs(argsOrClient, client); return prisma.taskRun.findFirst({ where, @@ -651,12 +655,16 @@ export class PostgresRunStore implements RunStore { args: { include: I }, client?: PrismaClientOrTransaction ): Promise>; + findRunOrThrow( + where: Prisma.TaskRunWhereInput, + client?: PrismaClientOrTransaction + ): Promise; async findRunOrThrow( where: Prisma.TaskRunWhereInput, - args: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }, + argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | PrismaClientOrTransaction, client?: PrismaClientOrTransaction ): Promise { - const prisma = client ?? this.readOnlyPrisma; + const { args, prisma } = this.#resolveReadArgs(argsOrClient, client); return prisma.taskRun.findFirstOrThrow({ where, @@ -686,6 +694,16 @@ export class PostgresRunStore implements RunStore { }, client?: PrismaClientOrTransaction ): Promise[]>; + findRuns( + args: { + where: Prisma.TaskRunWhereInput; + orderBy?: Prisma.TaskRunOrderByWithRelationInput | Prisma.TaskRunOrderByWithRelationInput[]; + take?: number; + skip?: number; + cursor?: Prisma.TaskRunWhereUniqueInput; + }, + client?: PrismaClientOrTransaction + ): Promise; async findRuns( args: { where: Prisma.TaskRunWhereInput; @@ -702,4 +720,41 @@ export class PostgresRunStore implements RunStore { return prisma.taskRun.findMany(args); } + + /** + * The single-row read methods (`findRun`, `findRunOrThrow`) accept either + * `(where, { select | include }, client?)` or the full-row `(where, client?)`. + * Disambiguate the second positional arg: a `{ select }` / `{ include }` + * projection object vs. a Prisma client. A projection object always carries a + * `select` or `include` key; a Prisma client never does. Anything else (e.g. + * `undefined`) is treated as "no projection, no explicit client". + */ + #resolveReadArgs( + argsOrClient: + | { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } + | PrismaClientOrTransaction + | undefined, + client: PrismaClientOrTransaction | undefined + ): { + args: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }; + prisma: PrismaClientOrTransaction | PrismaReplicaClient; + } { + const isProjection = + typeof argsOrClient === "object" && + argsOrClient !== null && + ("select" in argsOrClient || "include" in argsOrClient); + + if (isProjection) { + return { + args: argsOrClient as { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }, + prisma: client ?? this.readOnlyPrisma, + }; + } + + // No projection: the second positional arg, when present, is the client. + return { + args: {}, + prisma: (argsOrClient as PrismaClientOrTransaction | undefined) ?? this.readOnlyPrisma, + }; + } } diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index 35d4d8f91a..4c2d9d554a 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -331,6 +331,7 @@ export interface RunStore { args: { include: I }, client?: PrismaClientOrTransaction ): Promise | null>; + findRun(where: Prisma.TaskRunWhereInput, client?: PrismaClientOrTransaction): Promise; findRunOrThrow( where: Prisma.TaskRunWhereInput, @@ -342,6 +343,7 @@ export interface RunStore { args: { include: I }, client?: PrismaClientOrTransaction ): Promise>; + findRunOrThrow(where: Prisma.TaskRunWhereInput, client?: PrismaClientOrTransaction): Promise; findRuns( args: { @@ -365,4 +367,14 @@ export interface RunStore { }, client?: PrismaClientOrTransaction ): Promise[]>; + findRuns( + args: { + where: Prisma.TaskRunWhereInput; + orderBy?: Prisma.TaskRunOrderByWithRelationInput | Prisma.TaskRunOrderByWithRelationInput[]; + take?: number; + skip?: number; + cursor?: Prisma.TaskRunWhereUniqueInput; + }, + client?: PrismaClientOrTransaction + ): Promise; } From cfa90521ecf119bd7ab64c10e43589b8cc8a9e0e Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 15:31:09 +0100 Subject: [PATCH 03/11] refactor(run-engine): route TaskRun reads through the run store Relocate the direct TaskRun reads in the engine and its systems to the RunStore read methods, preserving the exact client (writer, replica, or transaction) at each site. Behavior-preserving; the engine test suite is unchanged. --- .../run-engine/src/engine/index.ts | 34 +-- .../run-engine/src/engine/retrying.ts | 45 ++-- .../src/engine/systems/batchSystem.ts | 21 +- .../src/engine/systems/debounceSystem.ts | 40 ++-- .../src/engine/systems/delayedRunSystem.ts | 19 +- .../src/engine/systems/dequeueSystem.ts | 44 ++-- .../engine/systems/pendingVersionSystem.ts | 21 +- .../src/engine/systems/runAttemptSystem.ts | 226 ++++++++++-------- .../src/engine/systems/ttlSystem.ts | 4 +- .../src/engine/systems/waitpointSystem.ts | 45 ++-- 10 files changed, 278 insertions(+), 221 deletions(-) diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 8d1f4c9c1f..a6a20b5b9f 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -650,7 +650,7 @@ export class RunEngine { "createCancelledRun: row already exists, returning existing (idempotent)", { friendlyId: snapshot.friendlyId }, ); - const existing = await prisma.taskRun.findFirst({ where: { id } }); + const existing = await this.runStore.findRun({ id }, prisma); if (existing) { // Only treat the conflict as idempotent when the existing // row is ALREADY canceled. If a non-canceled row landed @@ -2325,16 +2325,19 @@ export class RunEngine { }); //the run didn't start executing, we need to requeue it - const run = await prisma.taskRun.findFirst({ - where: { id: runId }, - include: { - runtimeEnvironment: { - include: { - organization: true, + const run = await this.runStore.findRun( + { id: runId }, + { + include: { + runtimeEnvironment: { + include: { + organization: true, + }, }, }, }, - }); + prisma + ); if (!run) { this.logger.error( @@ -2629,12 +2632,15 @@ export class RunEngine { snapshotId, }); - const taskRun = await this.prisma.taskRun.findFirst({ - where: { id: runId }, - select: { - queue: true, + const taskRun = await this.runStore.findRun( + { id: runId }, + { + select: { + queue: true, + }, }, - }); + this.prisma + ); if (!taskRun) { this.logger.error( @@ -2708,7 +2714,7 @@ export class RunEngine { runIds: string[], completedAtOffsetMs: number = 1000 * 60 * 10 ): Promise> { - const runs = await this.readOnlyPrisma.taskRun.findMany({ + const runs = await this.runStore.findRuns({ where: { id: { in: runIds }, completedAt: { diff --git a/internal-packages/run-engine/src/engine/retrying.ts b/internal-packages/run-engine/src/engine/retrying.ts index 6099d5b649..a64dfb796e 100644 --- a/internal-packages/run-engine/src/engine/retrying.ts +++ b/internal-packages/run-engine/src/engine/retrying.ts @@ -10,6 +10,7 @@ import { TaskRunExecutionRetry, } from "@trigger.dev/core/v3"; import { PrismaClientOrTransaction } from "@trigger.dev/database"; +import { RunStore } from "@internal/run-store"; import { MAX_TASK_RUN_ATTEMPTS } from "./consts.js"; import { ServiceValidationError } from "./errors.js"; @@ -45,6 +46,7 @@ export type RetryOutcome = export async function retryOutcomeFromCompletion( prisma: PrismaClientOrTransaction, + runStore: RunStore, { runId, attemptNumber, error, retryUsingQueue, retrySettings }: Params ): Promise { // Canceled @@ -56,7 +58,7 @@ export async function retryOutcomeFromCompletion( // OOM error (retry on a larger machine or fail) if (isOOMRunError(error)) { - const oomResult = await retryOOMOnMachine(prisma, runId); + const oomResult = await retryOOMOnMachine(prisma, runStore, runId); if (!oomResult) { return { outcome: "fail_run", sanitizedError, wasOOMError: true }; } @@ -95,18 +97,21 @@ export async function retryOutcomeFromCompletion( } // Get the run settings and current usage values - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { id: runId, }, - select: { - maxAttempts: true, - lockedRetryConfig: true, - usageDurationMs: true, - costInCents: true, - machinePreset: true, + { + select: { + maxAttempts: true, + lockedRetryConfig: true, + usageDurationMs: true, + costInCents: true, + machinePreset: true, + }, }, - }); + prisma + ); if (!run) { throw new ServiceValidationError("Run not found", 404); @@ -179,6 +184,7 @@ export async function retryOutcomeFromCompletion( async function retryOOMOnMachine( prisma: PrismaClientOrTransaction, + runStore: RunStore, runId: string ): Promise<{ machine: string; @@ -188,17 +194,20 @@ async function retryOOMOnMachine( machinePreset: string | null; } | undefined> { try { - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { id: runId, }, - select: { - machinePreset: true, - lockedRetryConfig: true, - usageDurationMs: true, - costInCents: true, + { + select: { + machinePreset: true, + lockedRetryConfig: true, + usageDurationMs: true, + costInCents: true, + }, }, - }); + prisma + ); if (!run || !run.lockedRetryConfig || !run.machinePreset) { return; diff --git a/internal-packages/run-engine/src/engine/systems/batchSystem.ts b/internal-packages/run-engine/src/engine/systems/batchSystem.ts index 9933a71516..a3d44507a4 100644 --- a/internal-packages/run-engine/src/engine/systems/batchSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/batchSystem.ts @@ -87,16 +87,19 @@ export class BatchSystem { return; } - const runs = await this.$.prisma.taskRun.findMany({ - select: { - id: true, - status: true, - }, - where: { - batchId, - runtimeEnvironmentId: batch.runtimeEnvironmentId, + const runs = await this.$.runStore.findRuns( + { + select: { + id: true, + status: true, + }, + where: { + batchId, + runtimeEnvironmentId: batch.runtimeEnvironmentId, + }, }, - }); + this.$.prisma + ); if (runs.every((r) => isFinalRunStatus(r.status))) { this.$.logger.debug("#tryCompleteBatch: All runs are completed", { batchId }); diff --git a/internal-packages/run-engine/src/engine/systems/debounceSystem.ts b/internal-packages/run-engine/src/engine/systems/debounceSystem.ts index 5b9d851d0f..bf4b3e68bb 100644 --- a/internal-packages/run-engine/src/engine/systems/debounceSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/debounceSystem.ts @@ -606,10 +606,11 @@ return 0 return null; } - const probe = await prisma.taskRun.findFirst({ - where: { id: existingRunId }, - select: { status: true, delayUntil: true, createdAt: true }, - }); + const probe = await this.$.runStore.findRun( + { id: existingRunId }, + { select: { status: true, delayUntil: true, createdAt: true } }, + prisma + ); if (!probe || probe.status !== "DELAYED" || !probe.delayUntil) { return null; } @@ -632,10 +633,11 @@ return 0 return null; } - const fullRun = await prisma.taskRun.findFirst({ - where: { id: existingRunId }, - include: { associatedWaitpoint: true }, - }); + const fullRun = await this.$.runStore.findRun( + { id: existingRunId }, + { include: { associatedWaitpoint: true } }, + prisma + ); if (!fullRun || fullRun.status !== "DELAYED") { return null; } @@ -665,10 +667,11 @@ return 0 error: unknown; prisma: PrismaClientOrTransaction; }): Promise { - const fullRun = await prisma.taskRun.findFirst({ - where: { id: existingRunId }, - include: { associatedWaitpoint: true }, - }); + const fullRun = await this.$.runStore.findRun( + { id: existingRunId }, + { include: { associatedWaitpoint: true } }, + prisma + ); if (!fullRun || fullRun.status !== "DELAYED") { // The run is no longer in a state we can safely return as "existing" - @@ -775,12 +778,15 @@ return 0 } // Get the run to check debounce metadata and createdAt - const existingRun = await prisma.taskRun.findFirst({ - where: { id: existingRunId }, - include: { - associatedWaitpoint: true, + const existingRun = await this.$.runStore.findRun( + { id: existingRunId }, + { + include: { + associatedWaitpoint: true, + }, }, - }); + prisma + ); if (!existingRun) { this.$.logger.debug("handleExistingRun: existing run not found in database", { diff --git a/internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts b/internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts index cff29a75a4..a77e60d05e 100644 --- a/internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts @@ -110,17 +110,20 @@ export class DelayedRunSystem { return; } - const run = await this.$.prisma.taskRun.findFirst({ - where: { id: runId }, - include: { - runtimeEnvironment: { - include: { - project: true, - organization: true, + const run = await this.$.runStore.findRun( + { id: runId }, + { + include: { + runtimeEnvironment: { + include: { + project: true, + organization: true, + }, }, }, }, - }); + this.$.prisma + ); if (!run) { throw new Error(`#enqueueDelayedRun: run not found: ${runId}`); diff --git a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts index 26ea7866a6..8791dc1bd1 100644 --- a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts @@ -641,12 +641,15 @@ export class DequeueSystem { // Wrap the Prisma call with tryCatch - if DB is unavailable, we still want to nack via Redis const [findError, run] = await tryCatch( - prisma.taskRun.findFirst({ - where: { id: runId }, - include: { - runtimeEnvironment: true, + this.$.runStore.findRun( + { id: runId }, + { + include: { + runtimeEnvironment: true, + }, }, - }) + prisma + ) ); // If DB is unavailable or run not found, just nack directly via Redis @@ -808,26 +811,29 @@ export class DequeueSystem { return startSpan(this.$.tracer, "getRunWithBackgroundWorkerTasks", async (span) => { span.setAttribute("run_id", runId); - const run = await prisma.taskRun.findFirst({ - where: { + const run = await this.$.runStore.findRun( + { id: runId, }, - include: { - runtimeEnvironment: { - select: { - id: true, - type: true, - archivedAt: true, + { + include: { + runtimeEnvironment: { + select: { + id: true, + type: true, + archivedAt: true, + }, }, - }, - lockedToVersion: { - include: { - deployment: true, - tasks: true, + lockedToVersion: { + include: { + deployment: true, + tasks: true, + }, }, }, }, - }); + prisma + ); if (!run) { span.setAttribute("result", "NO_RUN"); diff --git a/internal-packages/run-engine/src/engine/systems/pendingVersionSystem.ts b/internal-packages/run-engine/src/engine/systems/pendingVersionSystem.ts index 59d72c4c46..741ad8a14f 100644 --- a/internal-packages/run-engine/src/engine/systems/pendingVersionSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/pendingVersionSystem.ts @@ -93,15 +93,18 @@ export class PendingVersionSystem { // is dropped. The planner uses the PK for `id IN (…)`; the status // predicate is a residual filter and does NOT require the status // index. - const pendingRuns = await this.$.prisma.taskRun.findMany({ - where: { - id: { in: candidateIds }, - status: "PENDING_VERSION", - }, - orderBy: { - createdAt: "asc", + const pendingRuns = await this.$.runStore.findRuns( + { + where: { + id: { in: candidateIds }, + status: "PENDING_VERSION", + }, + orderBy: { + createdAt: "asc", + }, }, - }); + this.$.prisma + ); if (!pendingRuns.length) { // CH returned candidates but all of them have already moved past @@ -135,7 +138,7 @@ export class PendingVersionSystem { return false; } - const updatedRun = await tx.taskRun.findFirstOrThrow({ where: { id: run.id } }); + const updatedRun = await this.$.runStore.findRunOrThrow({ id: run.id }, tx); await this.enqueueSystem.enqueueRun({ run: updatedRun, diff --git a/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts b/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts index 1aa1738f3b..977c94a8e8 100644 --- a/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts @@ -175,56 +175,58 @@ export class RunAttemptSystem { } public async resolveTaskRunContext(runId: string): Promise { - const run = await this.$.readOnlyPrisma.taskRun.findFirst({ - where: { + const run = await this.$.runStore.findRun( + { id: runId, }, - select: { - id: true, - createdAt: true, - updatedAt: true, - executedAt: true, - baseCostInCents: true, - projectId: true, - organizationId: true, - friendlyId: true, - lockedById: true, - lockedQueueId: true, - queue: true, - attemptNumber: true, - status: true, - ttl: true, - machinePreset: true, - runTags: true, - isTest: true, - replayedFromTaskRunFriendlyId: true, - idempotencyKey: true, - idempotencyKeyOptions: true, - startedAt: true, - maxAttempts: true, - taskVersion: true, - maxDurationInSeconds: true, - usageDurationMs: true, - costInCents: true, - traceContext: true, - priorityMs: true, - taskIdentifier: true, - runtimeEnvironment: { - select: { - id: true, - slug: true, - type: true, - branchName: true, - git: true, - organizationId: true, + { + select: { + id: true, + createdAt: true, + updatedAt: true, + executedAt: true, + baseCostInCents: true, + projectId: true, + organizationId: true, + friendlyId: true, + lockedById: true, + lockedQueueId: true, + queue: true, + attemptNumber: true, + status: true, + ttl: true, + machinePreset: true, + runTags: true, + isTest: true, + replayedFromTaskRunFriendlyId: true, + idempotencyKey: true, + idempotencyKeyOptions: true, + startedAt: true, + maxAttempts: true, + taskVersion: true, + maxDurationInSeconds: true, + usageDurationMs: true, + costInCents: true, + traceContext: true, + priorityMs: true, + taskIdentifier: true, + runtimeEnvironment: { + select: { + id: true, + slug: true, + type: true, + branchName: true, + git: true, + organizationId: true, + }, }, + parentTaskRunId: true, + rootTaskRunId: true, + batchId: true, + workerQueue: true, }, - parentTaskRunId: true, - rootTaskRunId: true, - batchId: true, - workerQueue: true, - }, - }); + } + ); if (!run) { throw new ServiceValidationError("Task run not found", 404); @@ -338,21 +340,23 @@ export class RunAttemptSystem { }); } - const taskRun = await this.$.readOnlyPrisma.taskRun.findFirst({ - where: { + const taskRun = await this.$.runStore.findRun( + { id: runId, }, - select: { - id: true, - friendlyId: true, - attemptNumber: true, - projectId: true, - runtimeEnvironmentId: true, - status: true, - lockedById: true, - ttl: true, - }, - }); + { + select: { + id: true, + friendlyId: true, + attemptNumber: true, + projectId: true, + runtimeEnvironmentId: true, + status: true, + lockedById: true, + ttl: true, + }, + } + ); this.$.logger.debug("Creating a task run attempt", { taskRun }); @@ -717,14 +721,16 @@ export class RunAttemptSystem { const completedAt = new Date(); // Read current usage values to calculate new totals (safe under runLock) - const currentRun = await this.$.readOnlyPrisma.taskRun.findFirst({ - where: { id: runId }, - select: { - usageDurationMs: true, - costInCents: true, - machinePreset: true, - }, - }); + const currentRun = await this.$.runStore.findRun( + { id: runId }, + { + select: { + usageDurationMs: true, + costInCents: true, + machinePreset: true, + }, + } + ); if (!currentRun) { throw new ServiceValidationError("Run not found", 404); @@ -904,35 +910,41 @@ export class RunAttemptSystem { const failedAt = new Date(); - const retryResult = await retryOutcomeFromCompletion(this.$.readOnlyPrisma, { - runId, - error: completion.error, - retryUsingQueue: forceRequeue ?? false, - retrySettings: completion.retry, - attemptNumber: latestSnapshot.attemptNumber, - }); + const retryResult = await retryOutcomeFromCompletion( + this.$.readOnlyPrisma, + this.$.runStore, + { + runId, + error: completion.error, + retryUsingQueue: forceRequeue ?? false, + retrySettings: completion.retry, + attemptNumber: latestSnapshot.attemptNumber, + } + ); // Force requeue means it was crashed so the attempt span needs to be closed if (forceRequeue) { - const minimalRun = await this.$.readOnlyPrisma.taskRun.findFirst({ - where: { + const minimalRun = await this.$.runStore.findRun( + { id: runId, }, - select: { - status: true, - spanId: true, - maxAttempts: true, - runtimeEnvironment: { - select: { - organizationId: true, + { + select: { + status: true, + spanId: true, + maxAttempts: true, + runtimeEnvironment: { + select: { + organizationId: true, + }, }, + taskEventStore: true, + createdAt: true, + completedAt: true, + updatedAt: true, }, - taskEventStore: true, - createdAt: true, - completedAt: true, - updatedAt: true, - }, - }); + } + ); if (!minimalRun) { throw new ServiceValidationError("Run not found", 404); @@ -1367,14 +1379,16 @@ export class RunAttemptSystem { // Calculate updated usage if we have attempt duration data let usageUpdate: { usageDurationMs: number; costInCents: number } | undefined; if (attemptDurationMs !== undefined) { - const currentRun = await this.$.readOnlyPrisma.taskRun.findFirst({ - where: { id: runId }, - select: { - usageDurationMs: true, - costInCents: true, - machinePreset: true, - }, - }); + const currentRun = await this.$.runStore.findRun( + { id: runId }, + { + select: { + usageDurationMs: true, + costInCents: true, + machinePreset: true, + }, + } + ); if (!currentRun) { throw new ServiceValidationError("Run not found", 404); @@ -1578,14 +1592,16 @@ export class RunAttemptSystem { const truncatedError = this.#truncateTaskRunError(error); // Read current usage values to calculate new totals - const currentRun = await this.$.readOnlyPrisma.taskRun.findFirst({ - where: { id: runId }, - select: { - usageDurationMs: true, - costInCents: true, - machinePreset: true, - }, - }); + const currentRun = await this.$.runStore.findRun( + { id: runId }, + { + select: { + usageDurationMs: true, + costInCents: true, + machinePreset: true, + }, + } + ); if (!currentRun) { throw new ServiceValidationError("Run not found", 404); diff --git a/internal-packages/run-engine/src/engine/systems/ttlSystem.ts b/internal-packages/run-engine/src/engine/systems/ttlSystem.ts index ebd1cbdd80..faffa2c59e 100644 --- a/internal-packages/run-engine/src/engine/systems/ttlSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/ttlSystem.ts @@ -33,7 +33,7 @@ export class TtlSystem { } //only expire "PENDING" runs - const run = await prisma.taskRun.findFirst({ where: { id: runId } }); + const run = await this.$.runStore.findRun({ id: runId }, prisma); if (!run) { this.$.logger.debug("Could not find enqueued run to expire", { @@ -171,7 +171,7 @@ export class TtlSystem { const skipped: { runId: string; reason: string }[] = []; // Fetch all runs in a single query (no snapshot data needed) - const runs = await this.$.readOnlyPrisma.taskRun.findMany({ + const runs = await this.$.runStore.findRuns({ where: { id: { in: runIds } }, select: { id: true, diff --git a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts index 8b8d4f82fc..29eba297be 100644 --- a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts @@ -679,23 +679,26 @@ export class WaitpointSystem { } // 3. Get the run with environment - const run = await this.$.prisma.taskRun.findFirst({ - where: { + const run = await this.$.runStore.findRun( + { id: runId, }, - include: { - runtimeEnvironment: { - select: { - id: true, - type: true, - maximumConcurrencyLimit: true, - concurrencyLimitBurstFactor: true, - project: { select: { id: true } }, - organization: { select: { id: true } }, + { + include: { + runtimeEnvironment: { + select: { + id: true, + type: true, + maximumConcurrencyLimit: true, + concurrencyLimitBurstFactor: true, + project: { select: { id: true } }, + organization: { select: { id: true } }, + }, }, }, }, - }); + this.$.prisma + ); if (!run) { this.$.logger.error(`continueRunIfUnblocked: run not found`, { @@ -972,10 +975,11 @@ export class WaitpointSystem { environmentId: string; }): Promise { // Fast path: check if waitpoint already exists - const run = await this.$.prisma.taskRun.findFirst({ - where: { id: runId }, - include: { associatedWaitpoint: true }, - }); + const run = await this.$.runStore.findRun( + { id: runId }, + { include: { associatedWaitpoint: true } }, + this.$.prisma + ); if (!run) { throw new Error(`Run not found: ${runId}`); @@ -990,10 +994,11 @@ export class WaitpointSystem { const prisma = this.$.prisma; // Double-check after acquiring lock - const runAfterLock = await prisma.taskRun.findFirst({ - where: { id: runId }, - include: { associatedWaitpoint: true }, - }); + const runAfterLock = await this.$.runStore.findRun( + { id: runId }, + { include: { associatedWaitpoint: true } }, + prisma + ); if (!runAfterLock) { throw new Error(`Run not found: ${runId}`); From 5b74b48435bc3854d6a235b55945ad284e144eea Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 15:57:41 +0100 Subject: [PATCH 04/11] refactor(webapp): route service-layer TaskRun reads through the run store Relocate the direct TaskRun reads in webapp services, run-engine concerns, realtime, mollifier and metadata to the RunStore read methods, preserving the exact client (writer, replica, or transaction) at each site. The run hydrator now receives the store by injection. Behavior-preserving. --- .../app/models/runtimeEnvironment.server.ts | 14 +- .../concerns/idempotencyKeys.server.ts | 22 +- .../services/triggerFailedTask.server.ts | 15 +- .../runEngine/services/triggerTask.server.ts | 8 +- .../metadata/updateMetadata.server.ts | 81 +++-- .../nativeRealtimeClientInstance.server.ts | 2 + .../app/services/realtime/runReader.server.ts | 29 +- .../realtime/sessionRunManager.server.ts | 39 ++- .../app/services/realtime/sessions.server.ts | 27 +- .../shadowRealtimeClientInstance.server.ts | 3 +- .../app/services/runsBackfiller.server.ts | 30 +- .../clickhouseRunsRepository.server.ts | 105 +++--- .../app/v3/eventRepository/index.server.ts | 38 ++- apps/webapp/app/v3/failedTaskRun.server.ts | 16 +- .../v3/mollifier/mutateWithFallback.server.ts | 9 +- .../mollifier/resolveRunForMutation.server.ts | 20 +- .../webapp/app/v3/runEngineHandlers.server.ts | 314 ++++++++++-------- .../alerts/performTaskRunAlerts.server.ts | 19 +- .../app/v3/services/batchTriggerV3.server.ts | 27 +- .../v3/services/bulk/BulkActionV2.server.ts | 44 +-- .../services/cancelDevSessionRuns.server.ts | 8 +- .../app/v3/services/completeAttempt.server.ts | 13 +- .../app/v3/services/crashTaskRun.server.ts | 6 +- .../createCheckpointRestoreEvent.server.ts | 19 +- .../services/createTaskRunAttempt.server.ts | 71 ++-- .../v3/services/enqueueDelayedRun.server.ts | 43 +-- .../services/executeTasksWaitingForDeploy.ts | 45 +-- .../v3/services/expireEnqueuedRun.server.ts | 19 +- .../app/v3/services/finalizeTaskRun.server.ts | 27 +- .../app/v3/services/retryAttempt.server.ts | 6 +- .../v3/services/updateFatalRunError.server.ts | 6 +- .../app/v3/taskRunHeartbeatFailed.server.ts | 41 +-- .../test/realtime/runReaderProjection.test.ts | 4 +- 33 files changed, 642 insertions(+), 528 deletions(-) diff --git a/apps/webapp/app/models/runtimeEnvironment.server.ts b/apps/webapp/app/models/runtimeEnvironment.server.ts index be05adaa8a..9135872417 100644 --- a/apps/webapp/app/models/runtimeEnvironment.server.ts +++ b/apps/webapp/app/models/runtimeEnvironment.server.ts @@ -1,6 +1,7 @@ import type { AuthenticatedEnvironment } from "@internal/run-engine"; import type { Prisma, PrismaClientOrTransaction, RuntimeEnvironment } from "@trigger.dev/database"; import { $replica, prisma } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { logger } from "~/services/logger.server"; import { getUsername } from "~/utils/username"; import { sanitizeBranchName } from "@trigger.dev/core/v3/utils/gitBranch"; @@ -251,14 +252,17 @@ export async function findEnvironmentFromRun( ): Promise { // The include (no select) already pulls every taskRun scalar, so runTags/batchId // ride along for free — no extra query for the realtime publish to send a full record. - const taskRun = await (tx ?? $replica).taskRun.findFirst({ - where: { + const taskRun = await runStore.findRun( + { id: runId, }, - include: { - runtimeEnvironment: { include: authIncludeBase }, + { + include: { + runtimeEnvironment: { include: authIncludeBase }, + }, }, - }); + tx ?? $replica + ); if (!taskRun?.runtimeEnvironment) { return null; } diff --git a/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts b/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts index 2bdf95eb9a..02d0ec957f 100644 --- a/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts +++ b/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts @@ -151,16 +151,19 @@ export class IdempotencyKeyConcern { } const existingRun = idempotencyKey - ? await this.prisma.taskRun.findFirst({ - where: { + ? await runStore.findRun( + { runtimeEnvironmentId: request.environment.id, idempotencyKey, taskIdentifier: request.taskId, }, - include: { - associatedWaitpoint: true, + { + include: { + associatedWaitpoint: true, + }, }, - }) + this.prisma + ) : undefined; // Buffer fallback per the mollifier-idempotency design. PG missed — @@ -329,14 +332,15 @@ export class IdempotencyKeyConcern { // Another concurrent trigger committed first. Re-resolve via the // existing checks: writer-side PG findFirst first (defeats // replica lag), then buffer fallback for the buffered case. - const writerRun = await this.prisma.taskRun.findFirst({ - where: { + const writerRun = await runStore.findRun( + { runtimeEnvironmentId: request.environment.id, idempotencyKey, taskIdentifier: request.taskId, }, - include: { associatedWaitpoint: true }, - }); + { include: { associatedWaitpoint: true } }, + this.prisma + ); if (writerRun) { return { isCached: true, run: writerRun }; } diff --git a/apps/webapp/app/runEngine/services/triggerFailedTask.server.ts b/apps/webapp/app/runEngine/services/triggerFailedTask.server.ts index a8a7cbf0f3..031411844b 100644 --- a/apps/webapp/app/runEngine/services/triggerFailedTask.server.ts +++ b/apps/webapp/app/runEngine/services/triggerFailedTask.server.ts @@ -9,6 +9,7 @@ import { getEventRepository } from "~/v3/eventRepository/index.server"; import { PerformTaskRunAlertsService } from "~/v3/services/alerts/performTaskRunAlerts.server"; import { DefaultQueueManager } from "../concerns/queues.server"; import type { TriggerTaskRequest } from "../types"; +import { runStore } from "~/v3/runStore.server"; export type TriggerFailedTaskRequest = { /** The task identifier (e.g. "my-task") */ @@ -82,12 +83,13 @@ export class TriggerFailedTaskService { // Resolve parent run for rootTaskRunId and depth (same as triggerTask.server.ts) const parentRun = request.parentRunId - ? await this.prisma.taskRun.findFirst({ - where: { + ? await runStore.findRun( + { id: RunId.fromFriendlyId(request.parentRunId), runtimeEnvironmentId: request.environment.id, }, - }) + this.prisma + ) : undefined; const depth = parentRun ? parentRun.depth + 1 : 0; @@ -275,12 +277,13 @@ export class TriggerFailedTaskService { let depth = 0; if (opts.parentRunId) { - const parentRun = await this.prisma.taskRun.findFirst({ - where: { + const parentRun = await runStore.findRun( + { id: RunId.fromFriendlyId(opts.parentRunId), runtimeEnvironmentId: opts.environmentId, }, - }); + this.prisma + ); if (parentRun) { parentTaskRunId = parentRun.id; diff --git a/apps/webapp/app/runEngine/services/triggerTask.server.ts b/apps/webapp/app/runEngine/services/triggerTask.server.ts index 78455f9b68..89a938da8b 100644 --- a/apps/webapp/app/runEngine/services/triggerTask.server.ts +++ b/apps/webapp/app/runEngine/services/triggerTask.server.ts @@ -67,6 +67,7 @@ import { import { mollifyTrigger } from "~/v3/mollifier/mollifierMollify.server"; import { type MollifierBuffer } from "@trigger.dev/redis-worker"; import { QueueSizeLimitExceededError, ServiceValidationError } from "~/v3/services/common.server"; +import { runStore } from "~/v3/runStore.server"; class NoopTriggerRacepointSystem implements TriggerRacepointSystem { async waitForRacepoint(options: { racepoint: TriggerRacepoints; id: string }): Promise { @@ -241,12 +242,13 @@ export class RunEngineTriggerTaskService { // Get parent run if specified const parentRun = body.options?.parentRunId - ? await this.prisma.taskRun.findFirst({ - where: { + ? await runStore.findRun( + { id: RunId.fromFriendlyId(body.options.parentRunId), runtimeEnvironmentId: environment.id, }, - }) + this.prisma + ) : undefined; // Validate parent run diff --git a/apps/webapp/app/services/metadata/updateMetadata.server.ts b/apps/webapp/app/services/metadata/updateMetadata.server.ts index 2cc057f10f..2af44d747b 100644 --- a/apps/webapp/app/services/metadata/updateMetadata.server.ts +++ b/apps/webapp/app/services/metadata/updateMetadata.server.ts @@ -189,18 +189,21 @@ export class UpdateMetadataService { // Fetch current run (+ the realtime membership keys, so a flush can publish) const run = yield* _( Effect.tryPromise(() => - this._prisma.taskRun.findFirst({ - where: { id: runId }, - select: { - id: true, - metadata: true, - metadataType: true, - metadataVersion: true, - runtimeEnvironmentId: true, - runTags: true, - batchId: true, + this._runStore.findRun( + { id: runId }, + { + select: { + id: true, + metadata: true, + metadataType: true, + metadataVersion: true, + runtimeEnvironmentId: true, + runTags: true, + batchId: true, + }, }, - }) + this._prisma + ) ) ); @@ -332,8 +335,8 @@ export class UpdateMetadataService { ) { const runIdType = runId.startsWith("run_") ? "friendly" : "internal"; - const taskRun = await this._prisma.taskRun.findFirst({ - where: environment + const taskRun = await this._runStore.findRun( + environment ? { runtimeEnvironmentId: environment.id, ...(runIdType === "internal" ? { id: runId } : { friendlyId: runId }), @@ -341,29 +344,32 @@ export class UpdateMetadataService { : { ...(runIdType === "internal" ? { id: runId } : { friendlyId: runId }), }, - select: { - id: true, - batchId: true, - runTags: true, - completedAt: true, - status: true, - metadata: true, - metadataType: true, - metadataVersion: true, - parentTaskRun: { - select: { - id: true, - status: true, + { + select: { + id: true, + batchId: true, + runTags: true, + completedAt: true, + status: true, + metadata: true, + metadataType: true, + metadataVersion: true, + parentTaskRun: { + select: { + id: true, + status: true, + }, }, - }, - rootTaskRun: { - select: { - id: true, - status: true, + rootTaskRun: { + select: { + id: true, + status: true, + }, }, }, }, - }); + this._prisma + ); if (!taskRun) { return; @@ -427,10 +433,13 @@ export class UpdateMetadataService { while (attempts <= MAX_RETRIES) { // Fetch the latest run data - const run = await this._prisma.taskRun.findFirst({ - where: { id: runId }, - select: { metadata: true, metadataType: true, metadataVersion: true }, - }); + const run = await this._runStore.findRun( + { id: runId }, + { + select: { metadata: true, metadataType: true, metadataVersion: true }, + }, + this._prisma + ); if (!run) { throw new Error(`Run ${runId} not found`); diff --git a/apps/webapp/app/services/realtime/nativeRealtimeClientInstance.server.ts b/apps/webapp/app/services/realtime/nativeRealtimeClientInstance.server.ts index 012c28c08f..3f29f3faa4 100644 --- a/apps/webapp/app/services/realtime/nativeRealtimeClientInstance.server.ts +++ b/apps/webapp/app/services/realtime/nativeRealtimeClientInstance.server.ts @@ -1,5 +1,6 @@ import { getMeter } from "@internal/tracing"; import { $replica } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { env } from "~/env.server"; import { singleton } from "~/utils/singleton"; import { getCachedLimit } from "../platform.v3.server"; @@ -122,6 +123,7 @@ function initializeNativeRealtimeClient(): NativeRealtimeClient { // One RunHydrator shared by the router and the client, so its single-flight + short-TTL cache covers both. const runReader = new RunHydrator({ replica: $replica, + runStore, cacheTtlMs: env.REALTIME_BACKEND_NATIVE_RUN_CACHE_TTL_MS, maxCacheEntries: env.REALTIME_BACKEND_NATIVE_RUN_CACHE_MAX_ENTRIES, }); diff --git a/apps/webapp/app/services/realtime/runReader.server.ts b/apps/webapp/app/services/realtime/runReader.server.ts index e8509d73de..98ce4dc35f 100644 --- a/apps/webapp/app/services/realtime/runReader.server.ts +++ b/apps/webapp/app/services/realtime/runReader.server.ts @@ -1,4 +1,5 @@ -import { type Prisma, type PrismaClient } from "@trigger.dev/database"; +import { type Prisma, type PrismaClient, type PrismaClientOrTransaction } from "@trigger.dev/database"; +import type { RunStore } from "@internal/run-store"; import { BoundedTtlCache } from "./boundedTtlCache"; import { RESERVED_COLUMNS, type RealtimeRunRow } from "./electricStreamProtocol.server"; @@ -79,6 +80,8 @@ export interface RunListResolver { export type RunHydratorOptions = { /** A read-replica Prisma client (`$replica`). Always Postgres. */ replica: Pick; + /** RunStore the reads are routed through; `replica` is passed as the read client. */ + runStore: RunStore; /** Read-through cache TTL (ms) collapsing duplicate refetches for the same run. Set 0 to disable. Defaults to 250ms. */ cacheTtlMs?: number; /** Hard cap on cache entries before expired entries are swept. */ @@ -139,24 +142,28 @@ export class RunHydrator { if (ids.length === 0) { return []; } - const rows = await this.options.replica.taskRun.findMany({ - where: { - runtimeEnvironmentId: environmentId, - id: { in: ids }, + const rows = await this.options.runStore.findRuns( + { + where: { + runtimeEnvironmentId: environmentId, + id: { in: ids }, + }, + select: buildHydratorSelect(skipColumns), }, - select: buildHydratorSelect(skipColumns), - }); + this.options.replica as PrismaClientOrTransaction + ); return rows as unknown as RealtimeRunRow[]; } async #fetch(environmentId: string, runId: string): Promise { - const run = await this.options.replica.taskRun.findFirst({ - where: { + const run = await this.options.runStore.findRun( + { id: runId, runtimeEnvironmentId: environmentId, }, - select: RUN_HYDRATOR_SELECT, - }); + { select: RUN_HYDRATOR_SELECT }, + this.options.replica as PrismaClientOrTransaction + ); return (run ?? null) as RealtimeRunRow | null; } diff --git a/apps/webapp/app/services/realtime/sessionRunManager.server.ts b/apps/webapp/app/services/realtime/sessionRunManager.server.ts index 1ad5174d1c..b227f382c7 100644 --- a/apps/webapp/app/services/realtime/sessionRunManager.server.ts +++ b/apps/webapp/app/services/realtime/sessionRunManager.server.ts @@ -2,6 +2,7 @@ import type { Session, TaskRunStatus } from "@trigger.dev/database"; import { SessionTriggerConfig as SessionTriggerConfigZod } from "@trigger.dev/core/v3"; import { z } from "zod"; import { prisma, $replica } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server"; @@ -119,10 +120,11 @@ export async function ensureRunForSession( // replica as "row vanished" double-triggers the session (a fast // first append after session create races the replica apply delay // and spawns a second live run consuming the same `.in`). - probe = await prisma.taskRun.findFirst({ - where: { id: session.currentRunId }, - select: { status: true, friendlyId: true }, - }); + probe = await runStore.findRun( + { id: session.currentRunId }, + { select: { status: true, friendlyId: true } }, + prisma + ); } if (probe && !isFinalRunStatus(probe.status)) { return { runId: session.currentRunId, triggered: false }; @@ -251,10 +253,11 @@ export async function ensureRunForSession( // just wrote `currentRunId` on the writer, so probe the writer too — // the replica may not have the run row yet, and a missed probe forces // another trigger+recurse until `ENSURE_RUN_FOR_SESSION_MAX_ATTEMPTS`. - const probe = await prisma.taskRun.findFirst({ - where: { id: fresh.currentRunId }, - select: { status: true, friendlyId: true }, - }); + const probe = await runStore.findRun( + { id: fresh.currentRunId }, + { select: { status: true, friendlyId: true } }, + prisma + ); if (probe && !isFinalRunStatus(probe.status)) { return { runId: fresh.currentRunId, triggered: false }; } @@ -494,10 +497,11 @@ async function getRunStatusAndFriendlyId( // `payload.previousRunId` without a second read. `Session.currentRunId` // stores the internal cuid; the agent's wire / customer hooks expose // the friendlyId via `ctx.run.id`, so consistency matters. - const row = await $replica.taskRun.findFirst({ - where: { id: runId }, - select: { status: true, friendlyId: true }, - }); + const row = await runStore.findRun( + { id: runId }, + { select: { status: true, friendlyId: true } }, + $replica + ); return row ?? null; } @@ -511,10 +515,11 @@ async function getRunStatusAndFriendlyId( * acceptable degraded behavior. */ async function resolveRunFriendlyId(runId: string): Promise { - const row = await $replica.taskRun.findFirst({ - where: { id: runId }, - select: { friendlyId: true }, - }); + const row = await runStore.findRun( + { id: runId }, + { select: { friendlyId: true } }, + $replica + ); return row?.friendlyId ?? runId; } @@ -526,7 +531,7 @@ async function cancelLostRaceRun( // Read-after-write: the run was just triggered on the writer, so go // through `prisma`. A `$replica` miss here would silently no-op the // cancel and leak an orphan run that no session is going to claim. - const run = await prisma.taskRun.findFirst({ where: { id: runId } }); + const run = await runStore.findRun({ id: runId }, prisma); if (!run) return; await service.call(run, { reason: "Lost session-run claim race" }); } diff --git a/apps/webapp/app/services/realtime/sessions.server.ts b/apps/webapp/app/services/realtime/sessions.server.ts index 55b969e7e5..a523111b5b 100644 --- a/apps/webapp/app/services/realtime/sessions.server.ts +++ b/apps/webapp/app/services/realtime/sessions.server.ts @@ -1,6 +1,7 @@ import type { PrismaClient, Session } from "@trigger.dev/database"; import type { SessionItem } from "@trigger.dev/core/v3"; import { $replica, prisma } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; /** * Prefix that {@link SessionId.generate} attaches to every Session friendlyId. @@ -131,10 +132,11 @@ export async function serializeSessionWithFriendlyRunId( const base = serializeSession(session); if (!session.currentRunId) return base; - const run = await $replica.taskRun.findFirst({ - where: { id: session.currentRunId }, - select: { friendlyId: true }, - }); + const run = await runStore.findRun( + { id: session.currentRunId }, + { select: { friendlyId: true } }, + $replica + ); return { ...base, @@ -158,14 +160,17 @@ export async function serializeSessionsWithFriendlyRunIds( // `currentRunId` is a plain string pointer (no FK), so scope the lookup to // the caller's tenant — a stale value must not resolve a run in another env. const runs = runIds.length - ? await $replica.taskRun.findMany({ - where: { - id: { in: runIds }, - projectId: scope.projectId, - runtimeEnvironmentId: scope.runtimeEnvironmentId, + ? await runStore.findRuns( + { + where: { + id: { in: runIds }, + projectId: scope.projectId, + runtimeEnvironmentId: scope.runtimeEnvironmentId, + }, + select: { id: true, friendlyId: true }, }, - select: { id: true, friendlyId: true }, - }) + $replica + ) : []; const friendlyIdByRunId = new Map(runs.map((run) => [run.id, run.friendlyId])); diff --git a/apps/webapp/app/services/realtime/shadowRealtimeClientInstance.server.ts b/apps/webapp/app/services/realtime/shadowRealtimeClientInstance.server.ts index 8dbb5007c2..35333f9639 100644 --- a/apps/webapp/app/services/realtime/shadowRealtimeClientInstance.server.ts +++ b/apps/webapp/app/services/realtime/shadowRealtimeClientInstance.server.ts @@ -1,5 +1,6 @@ import { getMeter } from "@internal/tracing"; import { $replica } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { env } from "~/env.server"; import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server"; import { singleton } from "~/utils/singleton"; @@ -20,7 +21,7 @@ function initializeShadowRealtimeClient(): ShadowRealtimeClient { }); const comparator = new RealtimeShadowComparator({ - runReader: new RunHydrator({ replica: $replica }), + runReader: new RunHydrator({ replica: $replica, runStore }), runListResolver: new ClickHouseRunListResolver({ getClickhouse: (organizationId) => clickhouseFactory.getClickhouseForOrganization(organizationId, "realtime"), diff --git a/apps/webapp/app/services/runsBackfiller.server.ts b/apps/webapp/app/services/runsBackfiller.server.ts index 7fc824f3d3..50e041ee64 100644 --- a/apps/webapp/app/services/runsBackfiller.server.ts +++ b/apps/webapp/app/services/runsBackfiller.server.ts @@ -1,6 +1,7 @@ import { Tracer } from "@opentelemetry/api"; import type { PrismaClientOrTransaction } from "@trigger.dev/database"; import { RunsReplicationService } from "~/services/runsReplicationService.server"; +import { runStore } from "~/v3/runStore.server"; import { startSpan } from "~/v3/tracing.server"; import { FINAL_RUN_STATUSES } from "../v3/taskStatus"; import { Logger } from "@trigger.dev/core/logger"; @@ -40,22 +41,25 @@ export class RunsBackfillerService { span.setAttribute("cursor", cursor ?? ""); span.setAttribute("batchSize", batchSize ?? 0); - const runs = await this.prisma.taskRun.findMany({ - where: { - createdAt: { - gte: from, - lte: to, + const runs = await runStore.findRuns( + { + where: { + createdAt: { + gte: from, + lte: to, + }, + status: { + in: FINAL_RUN_STATUSES, + }, + ...(cursor ? { id: { gt: cursor } } : {}), }, - status: { - in: FINAL_RUN_STATUSES, + orderBy: { + id: "asc", }, - ...(cursor ? { id: { gt: cursor } } : {}), + take: batchSize, }, - orderBy: { - id: "asc", - }, - take: batchSize, - }); + this.prisma + ); if (runs.length === 0) { this.logger.info("No runs to backfill", { from, to, cursor }); diff --git a/apps/webapp/app/services/runsRepository/clickhouseRunsRepository.server.ts b/apps/webapp/app/services/runsRepository/clickhouseRunsRepository.server.ts index 88e792b4a4..d32652a0b3 100644 --- a/apps/webapp/app/services/runsRepository/clickhouseRunsRepository.server.ts +++ b/apps/webapp/app/services/runsRepository/clickhouseRunsRepository.server.ts @@ -12,6 +12,7 @@ import { } from "./runsRepository.server"; import parseDuration from "parse-duration"; import { decodeRunsCursor, encodeRunsCursor } from "./runsCursor.server"; +import { runStore } from "~/v3/runStore.server"; type RunCursorRow = { runId: string; createdAt: number }; @@ -148,16 +149,19 @@ export class ClickHouseRunsRepository implements IRunsRepository { } // Then get friendly IDs from Prisma - const runs = await this.options.prisma.taskRun.findMany({ - where: { - id: { - in: runIds, + const runs = await runStore.findRuns( + { + where: { + id: { + in: runIds, + }, + }, + select: { + friendlyId: true, }, }, - select: { - friendlyId: true, - }, - }); + this.options.prisma + ); return runs.map((run) => run.friendlyId); } @@ -165,49 +169,52 @@ export class ClickHouseRunsRepository implements IRunsRepository { async listRuns(options: ListRunsOptions) { const { runIds, pagination } = await this.listRunIds(options); - let runs = await this.options.prisma.taskRun.findMany({ - where: { - id: { - in: runIds, + let runs = await runStore.findRuns( + { + where: { + id: { + in: runIds, + }, + }, + orderBy: { + id: "desc", + }, + select: { + id: true, + friendlyId: true, + taskIdentifier: true, + taskVersion: true, + runtimeEnvironmentId: true, + status: true, + createdAt: true, + startedAt: true, + lockedAt: true, + delayUntil: true, + updatedAt: true, + completedAt: true, + isTest: true, + spanId: true, + idempotencyKey: true, + ttl: true, + expiredAt: true, + costInCents: true, + baseCostInCents: true, + usageDurationMs: true, + runTags: true, + depth: true, + rootTaskRunId: true, + batchId: true, + metadata: true, + metadataType: true, + machinePreset: true, + queue: true, + workerQueue: true, + region: true, + annotations: true, }, }, - orderBy: { - id: "desc", - }, - select: { - id: true, - friendlyId: true, - taskIdentifier: true, - taskVersion: true, - runtimeEnvironmentId: true, - status: true, - createdAt: true, - startedAt: true, - lockedAt: true, - delayUntil: true, - updatedAt: true, - completedAt: true, - isTest: true, - spanId: true, - idempotencyKey: true, - ttl: true, - expiredAt: true, - costInCents: true, - baseCostInCents: true, - usageDurationMs: true, - runTags: true, - depth: true, - rootTaskRunId: true, - batchId: true, - metadata: true, - metadataType: true, - machinePreset: true, - queue: true, - workerQueue: true, - region: true, - annotations: true, - }, - }); + this.options.prisma + ); // ClickHouse is slightly delayed, so we're going to do in-memory status filtering too if (options.statuses && options.statuses.length > 0) { diff --git a/apps/webapp/app/v3/eventRepository/index.server.ts b/apps/webapp/app/v3/eventRepository/index.server.ts index 4be392535c..c59be0f3f5 100644 --- a/apps/webapp/app/v3/eventRepository/index.server.ts +++ b/apps/webapp/app/v3/eventRepository/index.server.ts @@ -2,6 +2,7 @@ import { env } from "~/env.server"; import { eventRepository } from "./eventRepository.server"; import { type IEventRepository, type TraceEventOptions } from "./eventRepository.types"; import { prisma } from "~/db.server"; +import { runStore } from "../runStore.server"; import { logger } from "~/services/logger.server"; import { FEATURE_FLAG } from "../featureFlags"; import { flag } from "../featureFlags.server"; @@ -284,28 +285,31 @@ async function recordRunEvent( } async function findRunForEventCreation(runId: string) { - return prisma.taskRun.findFirst({ - where: { + return runStore.findRun( + { id: runId, }, - select: { - friendlyId: true, - taskIdentifier: true, - traceContext: true, - taskEventStore: true, - runtimeEnvironment: { - select: { - id: true, - type: true, - organizationId: true, - projectId: true, - project: { - select: { - externalRef: true, + { + select: { + friendlyId: true, + taskIdentifier: true, + traceContext: true, + taskEventStore: true, + runtimeEnvironment: { + select: { + id: true, + type: true, + organizationId: true, + projectId: true, + project: { + select: { + externalRef: true, + }, }, }, }, }, }, - }); + prisma + ); } diff --git a/apps/webapp/app/v3/failedTaskRun.server.ts b/apps/webapp/app/v3/failedTaskRun.server.ts index f4b3c92ea6..c2f5866249 100644 --- a/apps/webapp/app/v3/failedTaskRun.server.ts +++ b/apps/webapp/app/v3/failedTaskRun.server.ts @@ -37,12 +37,13 @@ export class FailedTaskRunService extends BaseService { const isFriendlyId = anyRunId.startsWith("run_"); - const taskRun = await this._prisma.taskRun.findFirst({ - where: { + const taskRun = await this.runStore.findRun( + { friendlyId: isFriendlyId ? anyRunId : undefined, id: !isFriendlyId ? anyRunId : undefined, }, - }); + this._prisma + ); if (!taskRun) { logger.error("[FailedTaskRunService] Task run not found", { @@ -90,12 +91,13 @@ export class FailedTaskRunRetryHelper extends BaseService { completion: TaskRunFailedExecutionResult; isCrash?: boolean; }) { - const taskRun = await this._prisma.taskRun.findFirst({ - where: { + const taskRun = await this.runStore.findRun( + { id: runId, }, - ...FailedTaskRunRetryGetPayload, - }); + FailedTaskRunRetryGetPayload, + this._prisma + ); if (!taskRun) { logger.error("[FailedTaskRunRetryHelper] Task run not found", { diff --git a/apps/webapp/app/v3/mollifier/mutateWithFallback.server.ts b/apps/webapp/app/v3/mollifier/mutateWithFallback.server.ts index e6deff5dbe..91c877c813 100644 --- a/apps/webapp/app/v3/mollifier/mutateWithFallback.server.ts +++ b/apps/webapp/app/v3/mollifier/mutateWithFallback.server.ts @@ -5,7 +5,9 @@ import type { SnapshotPatch, } from "@trigger.dev/redis-worker"; import type { TaskRun } from "@trigger.dev/database"; +import type { PrismaClientOrTransaction } from "~/db.server"; import { prisma, $replica } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { logger } from "~/services/logger.server"; import { getMollifierBuffer } from "./mollifierBuffer.server"; @@ -238,9 +240,10 @@ async function findRunInPg( friendlyId: string, environmentId: string, ): Promise { - return client.taskRun.findFirst({ - where: { friendlyId, runtimeEnvironmentId: environmentId }, - }); + return runStore.findRun( + { friendlyId, runtimeEnvironmentId: environmentId }, + client as unknown as PrismaClientOrTransaction + ); } function defaultSleep(ms: number): Promise { diff --git a/apps/webapp/app/v3/mollifier/resolveRunForMutation.server.ts b/apps/webapp/app/v3/mollifier/resolveRunForMutation.server.ts index b3db81368b..dac12768a7 100644 --- a/apps/webapp/app/v3/mollifier/resolveRunForMutation.server.ts +++ b/apps/webapp/app/v3/mollifier/resolveRunForMutation.server.ts @@ -1,5 +1,7 @@ import type { MollifierBuffer } from "@trigger.dev/redis-worker"; +import type { PrismaClientOrTransaction } from "~/db.server"; import { $replica as defaultReplica, prisma as defaultWriter } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { getMollifierBuffer as defaultGetBuffer } from "./mollifierBuffer.server"; // Discriminated-union resolver used by mutation routes' `findResource`. @@ -41,10 +43,11 @@ export async function resolveRunForMutation(input: { const writer = input.deps?.prismaWriter ?? defaultWriter; const getBuffer = input.deps?.getBuffer ?? defaultGetBuffer; - const pgRun = await replica.taskRun.findFirst({ - where: { friendlyId: input.runParam, runtimeEnvironmentId: input.environmentId }, - select: { friendlyId: true }, - }); + const pgRun = await runStore.findRun( + { friendlyId: input.runParam, runtimeEnvironmentId: input.environmentId }, + { select: { friendlyId: true } }, + replica as PrismaClientOrTransaction + ); if (pgRun) return { source: "pg", friendlyId: pgRun.friendlyId }; const buffer = getBuffer(); @@ -72,10 +75,11 @@ export async function resolveRunForMutation(input: { // lookup-by-friendlyId timing). // Without this, the resolver returns null in degraded states that the // downstream mutateWithFallback flow would otherwise handle correctly. - const writerRun = await writer.taskRun.findFirst({ - where: { friendlyId: input.runParam, runtimeEnvironmentId: input.environmentId }, - select: { friendlyId: true }, - }); + const writerRun = await runStore.findRun( + { friendlyId: input.runParam, runtimeEnvironmentId: input.environmentId }, + { select: { friendlyId: true } }, + writer as PrismaClientOrTransaction + ); if (writerRun) return { source: "pg", friendlyId: writerRun.friendlyId }; return null; diff --git a/apps/webapp/app/v3/runEngineHandlers.server.ts b/apps/webapp/app/v3/runEngineHandlers.server.ts index 082974af38..e2285a4fec 100644 --- a/apps/webapp/app/v3/runEngineHandlers.server.ts +++ b/apps/webapp/app/v3/runEngineHandlers.server.ts @@ -20,6 +20,7 @@ import { createExceptionPropertiesFromError } from "./eventRepository/common.ser import { getEventRepositoryForStore, recordRunDebugLog } from "./eventRepository/index.server"; import { roomFromFriendlyRunId, socketIo } from "./handleSocketIo.server"; import { engine } from "./runEngine.server"; +import { runStore } from "./runStore.server"; import { publishChangeRecord } from "~/services/realtime/runChangeNotifierInstance.server"; import { PerformTaskRunAlertsService } from "./services/alerts/performTaskRunAlerts.server"; import { TaskRunErrorCodes } from "@trigger.dev/core/v3"; @@ -27,32 +28,35 @@ import { TaskRunErrorCodes } from "@trigger.dev/core/v3"; export function registerRunEngineEventBusHandlers() { engine.eventBus.on("runSucceeded", async ({ time, run, organization, environment }) => { const [taskRunError, taskRun] = await tryCatch( - $replica.taskRun.findFirstOrThrow({ - where: { + runStore.findRunOrThrow( + { id: run.id, }, - select: { - id: true, - friendlyId: true, - traceId: true, - spanId: true, - parentSpanId: true, - createdAt: true, - completedAt: true, - taskIdentifier: true, - projectId: true, - runtimeEnvironmentId: true, - environmentType: true, - isTest: true, - organizationId: true, - taskEventStore: true, - // Piggyback the realtime run-changed publish on this existing read so the - // per-env channel carries the membership keys (no separate query). No-op when - // the native backend is disabled. - runTags: true, - batchId: true, + { + select: { + id: true, + friendlyId: true, + traceId: true, + spanId: true, + parentSpanId: true, + createdAt: true, + completedAt: true, + taskIdentifier: true, + projectId: true, + runtimeEnvironmentId: true, + environmentType: true, + isTest: true, + organizationId: true, + taskEventStore: true, + // Piggyback the realtime run-changed publish on this existing read so the + // per-env channel carries the membership keys (no separate query). No-op when + // the native backend is disabled. + runTags: true, + batchId: true, + }, }, - }) + $replica + ) ); if (taskRunError) { @@ -110,31 +114,34 @@ export function registerRunEngineEventBusHandlers() { const exception = createExceptionPropertiesFromError(sanitizedError); const [taskRunError, taskRun] = await tryCatch( - $replica.taskRun.findFirstOrThrow({ - where: { + runStore.findRunOrThrow( + { id: run.id, }, - select: { - id: true, - friendlyId: true, - traceId: true, - spanId: true, - parentSpanId: true, - createdAt: true, - completedAt: true, - taskIdentifier: true, - projectId: true, - runtimeEnvironmentId: true, - environmentType: true, - isTest: true, - organizationId: true, - taskEventStore: true, - // Piggyback the realtime run-changed publish on this existing read (no-op when - // the native backend is disabled). - runTags: true, - batchId: true, + { + select: { + id: true, + friendlyId: true, + traceId: true, + spanId: true, + parentSpanId: true, + createdAt: true, + completedAt: true, + taskIdentifier: true, + projectId: true, + runtimeEnvironmentId: true, + environmentType: true, + isTest: true, + organizationId: true, + taskEventStore: true, + // Piggyback the realtime run-changed publish on this existing read (no-op when + // the native backend is disabled). + runTags: true, + batchId: true, + }, }, - }) + $replica + ) ); if (taskRunError) { @@ -179,31 +186,34 @@ export function registerRunEngineEventBusHandlers() { const exception = createExceptionPropertiesFromError(sanitizedError); const [taskRunError, taskRun] = await tryCatch( - $replica.taskRun.findFirstOrThrow({ - where: { + runStore.findRunOrThrow( + { id: run.id, }, - select: { - id: true, - friendlyId: true, - traceId: true, - spanId: true, - parentSpanId: true, - createdAt: true, - completedAt: true, - taskIdentifier: true, - projectId: true, - runtimeEnvironmentId: true, - environmentType: true, - isTest: true, - organizationId: true, - taskEventStore: true, - // Piggyback the realtime run-changed publish on this existing read (no-op when - // the native backend is disabled). - runTags: true, - batchId: true, + { + select: { + id: true, + friendlyId: true, + traceId: true, + spanId: true, + parentSpanId: true, + createdAt: true, + completedAt: true, + taskIdentifier: true, + projectId: true, + runtimeEnvironmentId: true, + environmentType: true, + isTest: true, + organizationId: true, + taskEventStore: true, + // Piggyback the realtime run-changed publish on this existing read (no-op when + // the native backend is disabled). + runTags: true, + batchId: true, + }, }, - }) + $replica + ) ); if (taskRunError) { @@ -265,26 +275,29 @@ export function registerRunEngineEventBusHandlers() { } const [cachedRunError, cachedRun] = await tryCatch( - $replica.taskRun.findFirstOrThrow({ - where: { + runStore.findRunOrThrow( + { id: cachedRunId, }, - select: { - id: true, - friendlyId: true, - traceId: true, - spanId: true, - parentSpanId: true, - createdAt: true, - completedAt: true, - taskIdentifier: true, - projectId: true, - runtimeEnvironmentId: true, - environmentType: true, - isTest: true, - organizationId: true, + { + select: { + id: true, + friendlyId: true, + traceId: true, + spanId: true, + parentSpanId: true, + createdAt: true, + completedAt: true, + taskIdentifier: true, + projectId: true, + runtimeEnvironmentId: true, + environmentType: true, + isTest: true, + organizationId: true, + }, }, - }) + $replica + ) ); if (cachedRunError) { @@ -296,27 +309,30 @@ export function registerRunEngineEventBusHandlers() { } const [blockedRunError, blockedRun] = await tryCatch( - $replica.taskRun.findFirst({ - where: { + runStore.findRun( + { id: blockedRunId, }, - select: { - id: true, - friendlyId: true, - traceId: true, - spanId: true, - parentSpanId: true, - createdAt: true, - completedAt: true, - taskIdentifier: true, - projectId: true, - runtimeEnvironmentId: true, - environmentType: true, - isTest: true, - organizationId: true, - taskEventStore: true, + { + select: { + id: true, + friendlyId: true, + traceId: true, + spanId: true, + parentSpanId: true, + createdAt: true, + completedAt: true, + taskIdentifier: true, + projectId: true, + runtimeEnvironmentId: true, + environmentType: true, + isTest: true, + organizationId: true, + taskEventStore: true, + }, }, - }) + $replica + ) ); if (blockedRunError) { @@ -372,31 +388,34 @@ export function registerRunEngineEventBusHandlers() { } const [taskRunError, taskRun] = await tryCatch( - $replica.taskRun.findFirstOrThrow({ - where: { + runStore.findRunOrThrow( + { id: run.id, }, - select: { - id: true, - friendlyId: true, - traceId: true, - spanId: true, - parentSpanId: true, - createdAt: true, - completedAt: true, - taskIdentifier: true, - projectId: true, - runtimeEnvironmentId: true, - environmentType: true, - isTest: true, - organizationId: true, - taskEventStore: true, - // Piggyback the realtime run-changed publish on this existing read (no-op when - // the native backend is disabled). - runTags: true, - batchId: true, + { + select: { + id: true, + friendlyId: true, + traceId: true, + spanId: true, + parentSpanId: true, + createdAt: true, + completedAt: true, + taskIdentifier: true, + projectId: true, + runtimeEnvironmentId: true, + environmentType: true, + isTest: true, + organizationId: true, + taskEventStore: true, + // Piggyback the realtime run-changed publish on this existing read (no-op when + // the native backend is disabled). + runTags: true, + batchId: true, + }, }, - }) + $replica + ) ); if (taskRunError) { @@ -438,31 +457,34 @@ export function registerRunEngineEventBusHandlers() { engine.eventBus.on("runCancelled", async ({ time, run, organization, environment }) => { const [taskRunError, taskRun] = await tryCatch( - $replica.taskRun.findFirstOrThrow({ - where: { + runStore.findRunOrThrow( + { id: run.id, }, - select: { - id: true, - friendlyId: true, - traceId: true, - spanId: true, - parentSpanId: true, - createdAt: true, - completedAt: true, - taskIdentifier: true, - projectId: true, - runtimeEnvironmentId: true, - environmentType: true, - isTest: true, - organizationId: true, - taskEventStore: true, - // Piggyback the realtime run-changed publish on this existing read (no-op when - // the native backend is disabled). - runTags: true, - batchId: true, + { + select: { + id: true, + friendlyId: true, + traceId: true, + spanId: true, + parentSpanId: true, + createdAt: true, + completedAt: true, + taskIdentifier: true, + projectId: true, + runtimeEnvironmentId: true, + environmentType: true, + isTest: true, + organizationId: true, + taskEventStore: true, + // Piggyback the realtime run-changed publish on this existing read (no-op when + // the native backend is disabled). + runTags: true, + batchId: true, + }, }, - }) + $replica + ) ); if (taskRunError) { diff --git a/apps/webapp/app/v3/services/alerts/performTaskRunAlerts.server.ts b/apps/webapp/app/v3/services/alerts/performTaskRunAlerts.server.ts index 9c05534623..31912c39fd 100644 --- a/apps/webapp/app/v3/services/alerts/performTaskRunAlerts.server.ts +++ b/apps/webapp/app/v3/services/alerts/performTaskRunAlerts.server.ts @@ -12,17 +12,20 @@ type FoundRun = Prisma.Result< export class PerformTaskRunAlertsService extends BaseService { public async call(runId: string) { - const run = await this._prisma.taskRun.findFirst({ - where: { id: runId }, - include: { - lockedBy: true, - runtimeEnvironment: { - include: { - parentEnvironment: true, + const run = await this.runStore.findRun( + { id: runId }, + { + include: { + lockedBy: true, + runtimeEnvironment: { + include: { + parentEnvironment: true, + }, }, }, }, - }); + this._prisma + ); if (!run) { return; diff --git a/apps/webapp/app/v3/services/batchTriggerV3.server.ts b/apps/webapp/app/v3/services/batchTriggerV3.server.ts index 3303687159..c001932baa 100644 --- a/apps/webapp/app/v3/services/batchTriggerV3.server.ts +++ b/apps/webapp/app/v3/services/batchTriggerV3.server.ts @@ -352,20 +352,23 @@ export class BatchTriggerV3Service extends BaseService { // Fetch cached runs for each task identifier separately to make use of the index const cachedRuns = await Promise.all( Object.entries(itemsByTask).map(([taskIdentifier, items]) => - this._prisma.taskRun.findMany({ - where: { - runtimeEnvironmentId: environment.id, - taskIdentifier, - idempotencyKey: { - in: items.map((i) => i.options?.idempotencyKey).filter(Boolean), + this.runStore.findRuns( + { + where: { + runtimeEnvironmentId: environment.id, + taskIdentifier, + idempotencyKey: { + in: items.map((i) => i.options?.idempotencyKey).filter(Boolean), + }, + }, + select: { + friendlyId: true, + idempotencyKey: true, + idempotencyKeyExpiresAt: true, }, }, - select: { - friendlyId: true, - idempotencyKey: true, - idempotencyKeyExpiresAt: true, - }, - }) + this._prisma + ) ) ).then((results) => results.flat()); diff --git a/apps/webapp/app/v3/services/bulk/BulkActionV2.server.ts b/apps/webapp/app/v3/services/bulk/BulkActionV2.server.ts index babdb02ca6..76d550c700 100644 --- a/apps/webapp/app/v3/services/bulk/BulkActionV2.server.ts +++ b/apps/webapp/app/v3/services/bulk/BulkActionV2.server.ts @@ -182,22 +182,25 @@ export class BulkActionService extends BaseService { case BulkActionType.CANCEL: { const cancelService = new CancelTaskRunService(this._prisma); - const runs = await this._replica.taskRun.findMany({ - where: { - id: { - in: runIdsToProcess, + const runs = await this.runStore.findRuns( + { + where: { + id: { + in: runIdsToProcess, + }, + }, + select: { + id: true, + engine: true, + friendlyId: true, + status: true, + createdAt: true, + completedAt: true, + taskEventStore: true, }, }, - select: { - id: true, - engine: true, - friendlyId: true, - status: true, - createdAt: true, - completedAt: true, - taskEventStore: true, - }, - }); + this._replica + ); await pMap( runs, @@ -233,13 +236,16 @@ export class BulkActionService extends BaseService { case BulkActionType.REPLAY: { const replayService = new ReplayTaskRunService(this._prisma); - const runs = await this._replica.taskRun.findMany({ - where: { - id: { - in: runIdsToProcess, + const runs = await this.runStore.findRuns( + { + where: { + id: { + in: runIdsToProcess, + }, }, }, - }); + this._replica + ); await pMap( runs, diff --git a/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts b/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts index f779d81641..c1562275e5 100644 --- a/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts +++ b/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts @@ -68,12 +68,8 @@ export class CancelDevSessionRunsService extends BaseService { logger.debug("Cancelling in progress run", { runId }); const taskRun = runId.startsWith("run_") - ? await this._prisma.taskRun.findFirst({ - where: { friendlyId: runId }, - }) - : await this._prisma.taskRun.findFirst({ - where: { id: runId }, - }); + ? await this.runStore.findRun({ friendlyId: runId }, this._prisma) + : await this.runStore.findRun({ id: runId }, this._prisma); if (!taskRun) { return; diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index c407664881..22a9047c3f 100644 --- a/apps/webapp/app/v3/services/completeAttempt.server.ts +++ b/apps/webapp/app/v3/services/completeAttempt.server.ts @@ -70,14 +70,17 @@ export class CompleteAttemptService extends BaseService { id: execution.attempt.id, }); - const run = await this._prisma.taskRun.findFirst({ - where: { + const run = await this.runStore.findRun( + { friendlyId: execution.run.id, }, - select: { - id: true, + { + select: { + id: true, + }, }, - }); + this._prisma + ); if (!run) { logger.error("[CompleteAttemptService] Task run not found", { diff --git a/apps/webapp/app/v3/services/crashTaskRun.server.ts b/apps/webapp/app/v3/services/crashTaskRun.server.ts index cd55b9ec0f..bff4b8d65b 100644 --- a/apps/webapp/app/v3/services/crashTaskRun.server.ts +++ b/apps/webapp/app/v3/services/crashTaskRun.server.ts @@ -35,11 +35,7 @@ export class CrashTaskRunService extends BaseService { return; } - const taskRun = await this._prisma.taskRun.findFirst({ - where: { - id: runId, - }, - }); + const taskRun = await this.runStore.findRun({ id: runId }, this._prisma); if (!taskRun) { logger.error("[CrashTaskRunService] Task run not found", { runId }); diff --git a/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts b/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts index 63a8b6bb9a..59c3794717 100644 --- a/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts +++ b/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts @@ -58,19 +58,22 @@ export class CreateCheckpointRestoreEventService extends BaseService { let taskRunDependencyId: string | undefined; if (params.dependencyFriendlyRunId) { - const run = await this._prisma.taskRun.findFirst({ - where: { + const run = await this.runStore.findRun( + { friendlyId: params.dependencyFriendlyRunId, }, - select: { - id: true, - dependency: { - select: { - id: true, + { + select: { + id: true, + dependency: { + select: { + id: true, + }, }, }, }, - }); + this._prisma + ); taskRunDependencyId = run?.dependency?.id; diff --git a/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts index 8be2b9557c..dbc4c576b7 100644 --- a/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts +++ b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts @@ -12,6 +12,7 @@ import { FINAL_RUN_STATUSES } from "../taskStatus"; import { BaseService, ServiceValidationError } from "./baseService.server"; import { CrashTaskRunService } from "./crashTaskRun.server"; import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server"; +import { runStore } from "../runStore.server"; export class CreateTaskRunAttemptService extends BaseService { public async call({ @@ -45,43 +46,46 @@ export class CreateTaskRunAttemptService extends BaseService { span.setAttribute("taskRunId", runId); } - const taskRun = await this._prisma.taskRun.findFirst({ - where: { + const taskRun = await this.runStore.findRun( + { id: !isFriendlyId ? runId : undefined, friendlyId: isFriendlyId ? runId : undefined, runtimeEnvironmentId: environment.id, }, - include: { - attempts: { - take: 1, - orderBy: { - number: "desc", + { + include: { + attempts: { + take: 1, + orderBy: { + number: "desc", + }, }, - }, - lockedBy: { - include: { - worker: { - select: { - id: true, - version: true, - sdkVersion: true, - cliVersion: true, - supportsLazyAttempts: true, + lockedBy: { + include: { + worker: { + select: { + id: true, + version: true, + sdkVersion: true, + cliVersion: true, + supportsLazyAttempts: true, + }, }, }, }, - }, - batchItems: { - include: { - batchTaskRun: { - select: { - friendlyId: true, + batchItems: { + include: { + batchTaskRun: { + select: { + friendlyId: true, + }, }, }, }, }, }, - }); + this._prisma + ); logger.debug("Creating a task run attempt", { taskRun }); @@ -263,20 +267,23 @@ async function getAuthenticatedEnvironmentFromRun( ) { const isFriendlyId = friendlyId.startsWith("run_"); - const taskRun = await (prismaClient ?? prisma).taskRun.findFirst({ - where: { + const taskRun = await runStore.findRun( + { id: !isFriendlyId ? friendlyId : undefined, friendlyId: isFriendlyId ? friendlyId : undefined, }, - include: { - runtimeEnvironment: { - include: { - organization: true, - project: true, + { + include: { + runtimeEnvironment: { + include: { + organization: true, + project: true, + }, }, }, }, - }); + prismaClient ?? prisma + ); if (!taskRun) { return; diff --git a/apps/webapp/app/v3/services/enqueueDelayedRun.server.ts b/apps/webapp/app/v3/services/enqueueDelayedRun.server.ts index 0b6149dfae..79cb4fb097 100644 --- a/apps/webapp/app/v3/services/enqueueDelayedRun.server.ts +++ b/apps/webapp/app/v3/services/enqueueDelayedRun.server.ts @@ -32,37 +32,40 @@ export class EnqueueDelayedRunService extends BaseService { } public async call(runId: string) { - const run = await this._prisma.taskRun.findFirst({ - where: { + const run = await this.runStore.findRun( + { id: runId, }, - include: { - runtimeEnvironment: { - include: { - organization: true, - project: true, + { + include: { + runtimeEnvironment: { + include: { + organization: true, + project: true, + }, }, - }, - dependency: { - include: { - dependentBatchRun: { - include: { - dependentTaskAttempt: { - include: { - taskRun: true, + dependency: { + include: { + dependentBatchRun: { + include: { + dependentTaskAttempt: { + include: { + taskRun: true, + }, }, }, }, - }, - dependentAttempt: { - include: { - taskRun: true, + dependentAttempt: { + include: { + taskRun: true, + }, }, }, }, }, }, - }); + this._prisma + ); if (!run) { logger.debug("Could not find delayed run to enqueue", { diff --git a/apps/webapp/app/v3/services/executeTasksWaitingForDeploy.ts b/apps/webapp/app/v3/services/executeTasksWaitingForDeploy.ts index fb519b4315..a77727c924 100644 --- a/apps/webapp/app/v3/services/executeTasksWaitingForDeploy.ts +++ b/apps/webapp/app/v3/services/executeTasksWaitingForDeploy.ts @@ -39,29 +39,32 @@ export class ExecuteTasksWaitingForDeployService extends BaseService { const maxCount = env.LEGACY_RUN_ENGINE_WAITING_FOR_DEPLOY_BATCH_SIZE; - const runsWaitingForDeploy = await this._replica.taskRun.findMany({ - where: { - runtimeEnvironmentId: backgroundWorker.runtimeEnvironmentId, - projectId: backgroundWorker.projectId, - status: "WAITING_FOR_DEPLOY", - taskIdentifier: { - in: backgroundWorker.tasks.map((task) => task.slug), + const runsWaitingForDeploy = await this.runStore.findRuns( + { + where: { + runtimeEnvironmentId: backgroundWorker.runtimeEnvironmentId, + projectId: backgroundWorker.projectId, + status: "WAITING_FOR_DEPLOY", + taskIdentifier: { + in: backgroundWorker.tasks.map((task) => task.slug), + }, }, + orderBy: { + createdAt: "asc", + }, + select: { + id: true, + status: true, + taskIdentifier: true, + concurrencyKey: true, + queue: true, + updatedAt: true, + createdAt: true, + }, + take: maxCount + 1, }, - orderBy: { - createdAt: "asc", - }, - select: { - id: true, - status: true, - taskIdentifier: true, - concurrencyKey: true, - queue: true, - updatedAt: true, - createdAt: true, - }, - take: maxCount + 1, - }); + this._replica + ); if (!runsWaitingForDeploy.length) { return; diff --git a/apps/webapp/app/v3/services/expireEnqueuedRun.server.ts b/apps/webapp/app/v3/services/expireEnqueuedRun.server.ts index 0409b6ed95..12ccddbf2e 100644 --- a/apps/webapp/app/v3/services/expireEnqueuedRun.server.ts +++ b/apps/webapp/app/v3/services/expireEnqueuedRun.server.ts @@ -23,19 +23,22 @@ export class ExpireEnqueuedRunService extends BaseService { } public async call(runId: string) { - const run = await this._prisma.taskRun.findFirst({ - where: { + const run = await this.runStore.findRun( + { id: runId, }, - include: { - runtimeEnvironment: { - include: { - organization: true, - project: true, + { + include: { + runtimeEnvironment: { + include: { + organization: true, + project: true, + }, }, }, }, - }); + this._prisma + ); if (!run) { logger.debug("Could not find enqueued run to expire", { diff --git a/apps/webapp/app/v3/services/finalizeTaskRun.server.ts b/apps/webapp/app/v3/services/finalizeTaskRun.server.ts index ab51df5de6..b770ceef17 100644 --- a/apps/webapp/app/v3/services/finalizeTaskRun.server.ts +++ b/apps/webapp/app/v3/services/finalizeTaskRun.server.ts @@ -152,22 +152,25 @@ export class FinalizeTaskRunService extends BaseService { if (isFatalRunStatus(run.status)) { logger.warn("FinalizeTaskRunService: Fatal status", { runId: run.id, status: run.status }); - const extendedRun = await this._prisma.taskRun.findFirst({ - where: { id: run.id }, - select: { - id: true, - lockedToVersion: { - select: { - supportsLazyAttempts: true, + const extendedRun = await this.runStore.findRun( + { id: run.id }, + { + select: { + id: true, + lockedToVersion: { + select: { + supportsLazyAttempts: true, + }, }, - }, - runtimeEnvironment: { - select: { - type: true, + runtimeEnvironment: { + select: { + type: true, + }, }, }, }, - }); + this._prisma + ); if (extendedRun && extendedRun.runtimeEnvironment.type !== "DEVELOPMENT") { logger.warn("FinalizeTaskRunService: Fatal status, requesting worker exit", { diff --git a/apps/webapp/app/v3/services/retryAttempt.server.ts b/apps/webapp/app/v3/services/retryAttempt.server.ts index b4ab523576..6ed83c1080 100644 --- a/apps/webapp/app/v3/services/retryAttempt.server.ts +++ b/apps/webapp/app/v3/services/retryAttempt.server.ts @@ -5,11 +5,7 @@ import { BaseService } from "./baseService.server"; export class RetryAttemptService extends BaseService { public async call(runId: string) { - const taskRun = await this._prisma.taskRun.findFirst({ - where: { - id: runId, - }, - }); + const taskRun = await this.runStore.findRun({ id: runId }, this._prisma); if (!taskRun) { logger.error("Task run not found", { runId }); diff --git a/apps/webapp/app/v3/services/updateFatalRunError.server.ts b/apps/webapp/app/v3/services/updateFatalRunError.server.ts index 2363d241c0..dcf2488f27 100644 --- a/apps/webapp/app/v3/services/updateFatalRunError.server.ts +++ b/apps/webapp/app/v3/services/updateFatalRunError.server.ts @@ -20,11 +20,7 @@ export class UpdateFatalRunErrorService extends BaseService { logger.debug("UpdateFatalRunErrorService.call", { runId, opts }); - const taskRun = await this._prisma.taskRun.findFirst({ - where: { - id: runId, - }, - }); + const taskRun = await this.runStore.findRun({ id: runId }, this._prisma); if (!taskRun) { logger.error("[UpdateFatalRunErrorService] Task run not found", { runId }); diff --git a/apps/webapp/app/v3/taskRunHeartbeatFailed.server.ts b/apps/webapp/app/v3/taskRunHeartbeatFailed.server.ts index 8359cc4a4a..c472bff53e 100644 --- a/apps/webapp/app/v3/taskRunHeartbeatFailed.server.ts +++ b/apps/webapp/app/v3/taskRunHeartbeatFailed.server.ts @@ -11,32 +11,35 @@ import { TaskRunErrorCodes } from "@trigger.dev/core/v3"; export class TaskRunHeartbeatFailedService extends BaseService { public async call(runId: string) { - const taskRun = await this._prisma.taskRun.findFirst({ - where: { + const taskRun = await this.runStore.findRun( + { id: runId, }, - select: { - id: true, - friendlyId: true, - status: true, - lockedAt: true, - runtimeEnvironment: { - select: { - type: true, + { + select: { + id: true, + friendlyId: true, + status: true, + lockedAt: true, + runtimeEnvironment: { + select: { + type: true, + }, }, - }, - lockedToVersion: { - select: { - supportsLazyAttempts: true, + lockedToVersion: { + select: { + supportsLazyAttempts: true, + }, }, - }, - _count: { - select: { - attempts: true, + _count: { + select: { + attempts: true, + }, }, }, }, - }); + this._prisma + ); if (!taskRun) { logger.error("[TaskRunHeartbeatFailedService] Task run not found", { diff --git a/apps/webapp/test/realtime/runReaderProjection.test.ts b/apps/webapp/test/realtime/runReaderProjection.test.ts index 07aebf9258..ad6616f546 100644 --- a/apps/webapp/test/realtime/runReaderProjection.test.ts +++ b/apps/webapp/test/realtime/runReaderProjection.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it, vi } from "vitest"; +import { PostgresRunStore } from "@internal/run-store"; import { buildHydratorSelect, RunHydrator } from "~/services/realtime/runReader.server"; describe("buildHydratorSelect", () => { @@ -54,7 +55,8 @@ describe("RunHydrator.hydrateByIds column projection", () => { }), }, } as any; - return { hydrator: new RunHydrator({ replica }), getSelect: () => capturedSelect }; + const runStore = new PostgresRunStore({ prisma: replica, readOnlyPrisma: replica }); + return { hydrator: new RunHydrator({ replica, runStore }), getSelect: () => capturedSelect }; } it("projects the SELECT by skipColumns", async () => { From 5683952331df5bb7f79d1cc0afec66b323799c4e Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 16:12:57 +0100 Subject: [PATCH 05/11] refactor(webapp): route presenter TaskRun reads through the run store Relocate the dashboard presenter TaskRun reads to the RunStore read methods, preserving the exact client per site. Behavior-preserving. --- .../v3/ApiRetrieveRunPresenter.server.ts | 56 ++++----- .../v3/ApiRunResultPresenter.server.ts | 18 +-- .../v3/NextRunListPresenter.server.ts | 8 +- .../app/presenters/v3/RunPresenter.server.ts | 108 +++++++++--------- .../v3/RunStreamPresenter.server.ts | 14 ++- .../v3/SessionListPresenter.server.ts | 18 +-- .../presenters/v3/SessionPresenter.server.ts | 23 ++-- .../app/presenters/v3/SpanPresenter.server.ts | 76 ++++++------ .../presenters/v3/TestTaskPresenter.server.ts | 62 +++++----- 9 files changed, 213 insertions(+), 170 deletions(-) diff --git a/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts index fec8dabdb0..68e3643f9e 100644 --- a/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts @@ -22,6 +22,7 @@ import { type SyntheticRun, } from "~/v3/mollifier/readFallback.server"; import { generatePresignedUrl } from "~/v3/objectStore.server"; +import { runStore } from "~/v3/runStore.server"; import { tracer } from "~/v3/tracer.server"; import { startSpanWithEnv } from "~/v3/tracing.server"; @@ -110,38 +111,41 @@ export class ApiRetrieveRunPresenter { friendlyId: string, env: AuthenticatedEnvironment, ): Promise { - const pgRow = await $replica.taskRun.findFirst({ - where: { + const pgRow = await runStore.findRun( + { friendlyId, runtimeEnvironmentId: env.id, }, - select: { - ...commonRunSelect, - traceId: true, - payload: true, - payloadType: true, - output: true, - outputType: true, - error: true, - attempts: { - select: { - id: true, + { + select: { + ...commonRunSelect, + traceId: true, + payload: true, + payloadType: true, + output: true, + outputType: true, + error: true, + attempts: { + select: { + id: true, + }, + }, + attemptNumber: true, + engine: true, + taskEventStore: true, + parentTaskRun: { + select: commonRunSelect, + }, + rootTaskRun: { + select: commonRunSelect, + }, + childRuns: { + select: commonRunSelect, }, - }, - attemptNumber: true, - engine: true, - taskEventStore: true, - parentTaskRun: { - select: commonRunSelect, - }, - rootTaskRun: { - select: commonRunSelect, - }, - childRuns: { - select: commonRunSelect, }, }, - }); + $replica + ); if (pgRow) return { ...pgRow, isBuffered: false }; diff --git a/apps/webapp/app/presenters/v3/ApiRunResultPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiRunResultPresenter.server.ts index c11a04a158..7e0540674e 100644 --- a/apps/webapp/app/presenters/v3/ApiRunResultPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiRunResultPresenter.server.ts @@ -1,6 +1,7 @@ import { TaskRunExecutionResult } from "@trigger.dev/core/v3"; import { executionResultForTaskRun } from "~/models/taskRun.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { runStore } from "~/v3/runStore.server"; import { BasePresenter } from "./basePresenter.server"; export class ApiRunResultPresenter extends BasePresenter { @@ -9,19 +10,22 @@ export class ApiRunResultPresenter extends BasePresenter { env: AuthenticatedEnvironment ): Promise { return this.traceWithEnv("call", env, async (span) => { - const taskRun = await this._prisma.taskRun.findFirst({ - where: { + const taskRun = await runStore.findRun( + { friendlyId, runtimeEnvironmentId: env.id, }, - include: { - attempts: { - orderBy: { - createdAt: "desc", + { + include: { + attempts: { + orderBy: { + createdAt: "desc", + }, }, }, }, - }); + this._prisma + ); if (!taskRun) { return undefined; diff --git a/apps/webapp/app/presenters/v3/NextRunListPresenter.server.ts b/apps/webapp/app/presenters/v3/NextRunListPresenter.server.ts index 3594aa71ce..2e587e8c4a 100644 --- a/apps/webapp/app/presenters/v3/NextRunListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/NextRunListPresenter.server.ts @@ -13,6 +13,7 @@ import { getTaskIdentifiers } from "~/models/task.server"; import { RunsRepository } from "~/services/runsRepository/runsRepository.server"; import { regionForDisplay } from "~/runEngine/concerns/workerQueueSplit.server"; import { machinePresetFromRun } from "~/v3/machinePresets.server"; +import { runStore } from "~/v3/runStore.server"; import { ServiceValidationError } from "~/v3/services/baseService.server"; import { isCancellableRunStatus, isFinalRunStatus, isPendingRunStatus } from "~/v3/taskStatus"; @@ -206,11 +207,12 @@ export class NextRunListPresenter { let hasAnyRuns = runs.length > 0; if (!hasAnyRuns) { - const firstRun = await this.replica.taskRun.findFirst({ - where: { + const firstRun = await runStore.findRun( + { runtimeEnvironmentId: environmentId, }, - }); + this.replica + ); if (firstRun) { hasAnyRuns = true; diff --git a/apps/webapp/app/presenters/v3/RunPresenter.server.ts b/apps/webapp/app/presenters/v3/RunPresenter.server.ts index 1ff68e9b96..c4c3ac88c4 100644 --- a/apps/webapp/app/presenters/v3/RunPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/RunPresenter.server.ts @@ -8,6 +8,7 @@ import { getTaskEventStoreTableForRun } from "~/v3/taskEventStore.server"; import { isFinalRunStatus } from "~/v3/taskStatus"; import { env } from "~/env.server"; import { getEventRepositoryForStore } from "~/v3/eventRepository/index.server"; +import { runStore } from "~/v3/runStore.server"; type Result = Awaited>; export type Run = Result["run"]; @@ -62,57 +63,8 @@ export class RunPresenter { // buffer view. `findFirstOrThrow` would log a `PrismaClient error` // every tick of the page poll, masking real DB issues with synthetic // not-found noise. - const run = await this.#prismaClient.taskRun.findFirst({ - select: { - id: true, - createdAt: true, - taskEventStore: true, - taskIdentifier: true, - number: true, - traceId: true, - spanId: true, - parentSpanId: true, - friendlyId: true, - status: true, - startedAt: true, - completedAt: true, - logsDeletedAt: true, - annotations: true, - rootTaskRun: { - select: { - friendlyId: true, - spanId: true, - createdAt: true, - }, - }, - parentTaskRun: { - select: { - friendlyId: true, - spanId: true, - createdAt: true, - }, - }, - runtimeEnvironment: { - select: { - id: true, - type: true, - slug: true, - organizationId: true, - orgMember: { - select: { - user: { - select: { - id: true, - name: true, - displayName: true, - }, - }, - }, - }, - }, - }, - }, - where: { + const run = await runStore.findRun( + { friendlyId: runFriendlyId, project: { slug: projectSlug, @@ -125,7 +77,59 @@ export class RunPresenter { }, }, }, - }); + { + select: { + id: true, + createdAt: true, + taskEventStore: true, + taskIdentifier: true, + number: true, + traceId: true, + spanId: true, + parentSpanId: true, + friendlyId: true, + status: true, + startedAt: true, + completedAt: true, + logsDeletedAt: true, + annotations: true, + rootTaskRun: { + select: { + friendlyId: true, + spanId: true, + createdAt: true, + }, + }, + parentTaskRun: { + select: { + friendlyId: true, + spanId: true, + createdAt: true, + }, + }, + runtimeEnvironment: { + select: { + id: true, + type: true, + slug: true, + organizationId: true, + orgMember: { + select: { + user: { + select: { + id: true, + name: true, + displayName: true, + }, + }, + }, + }, + }, + }, + }, + }, + this.#prismaClient + ); if (!run) { throw new RunNotInPgError(runFriendlyId); diff --git a/apps/webapp/app/presenters/v3/RunStreamPresenter.server.ts b/apps/webapp/app/presenters/v3/RunStreamPresenter.server.ts index e0e88e4dd0..ab777d0b8e 100644 --- a/apps/webapp/app/presenters/v3/RunStreamPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/RunStreamPresenter.server.ts @@ -6,6 +6,7 @@ import { ABORT_REASON_SEND_ERROR, createSSELoader, SendFunction } from "~/utils/ import { throttle } from "~/utils/throttle"; import { getMollifierBuffer } from "~/v3/mollifier/mollifierBuffer.server"; import { deserialiseMollifierSnapshot } from "~/v3/mollifier/mollifierSnapshot.server"; +import { runStore } from "~/v3/runStore.server"; import { tracePubSub } from "~/v3/services/tracePubSub.server"; const PING_INTERVAL = 5_000; @@ -36,8 +37,8 @@ export class RunStreamPresenter { // Scope the lookup to organizations the requesting user is a member // of, matching RunPresenter's run lookup. Unauthorized and missing // runs are indistinguishable (both 404). - const run = await prismaClient.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runFriendlyId, project: { organization: { @@ -49,10 +50,13 @@ export class RunStreamPresenter { }, }, }, - select: { - traceId: true, + { + select: { + traceId: true, + }, }, - }); + prismaClient + ); // Fall back to the mollifier buffer when the run isn't in PG yet. // The buffered run has no execution events to stream, but we still diff --git a/apps/webapp/app/presenters/v3/SessionListPresenter.server.ts b/apps/webapp/app/presenters/v3/SessionListPresenter.server.ts index 0586ab8ece..bff1bda017 100644 --- a/apps/webapp/app/presenters/v3/SessionListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/SessionListPresenter.server.ts @@ -10,6 +10,7 @@ import { } from "~/services/sessionsRepository/sessionsRepository.server"; import { ServiceValidationError } from "~/v3/services/baseService.server"; import { findCurrentWorkerFromEnvironment } from "~/v3/models/workerDeployment.server"; +import { runStore } from "~/v3/runStore.server"; import { startActiveSpan } from "~/v3/tracer.server"; export type SessionListOptions = { @@ -189,14 +190,17 @@ export class SessionListPresenter { // pointer could surface another tenant's run. The list query above // is already env-scoped; the run lookup needs the same fence. return currentRunIds.length > 0 - ? this.replica.taskRun.findMany({ - where: { - id: { in: currentRunIds }, - projectId, - runtimeEnvironmentId: environmentId, + ? runStore.findRuns( + { + where: { + id: { in: currentRunIds }, + projectId, + runtimeEnvironmentId: environmentId, + }, + select: { id: true, friendlyId: true }, }, - select: { id: true, friendlyId: true }, - }) + this.replica + ) : []; } ); diff --git a/apps/webapp/app/presenters/v3/SessionPresenter.server.ts b/apps/webapp/app/presenters/v3/SessionPresenter.server.ts index c63f9e39a2..36ef46d4b4 100644 --- a/apps/webapp/app/presenters/v3/SessionPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/SessionPresenter.server.ts @@ -6,6 +6,7 @@ import { chatSnapshotStorageKey } from "~/services/realtime/chatSnapshot.server" import { resolveSessionByIdOrExternalId } from "~/services/realtime/sessions.server"; import { logger } from "~/services/logger.server"; import { generatePresignedUrl } from "~/v3/objectStore.server"; +import { runStore } from "~/v3/runStore.server"; import { ServiceValidationError } from "~/v3/services/baseService.server"; import { startActiveSpan } from "~/v3/tracer.server"; @@ -96,10 +97,13 @@ export class SessionPresenter { async (span) => { span.setAttribute("runIds.count", runIds.length); return runIds.length > 0 - ? this.replica.taskRun.findMany({ - where: { id: { in: runIds } }, - select: { id: true, friendlyId: true, status: true }, - }) + ? runStore.findRuns( + { + where: { id: { in: runIds } }, + select: { id: true, friendlyId: true, status: true }, + }, + this.replica + ) : []; } ); @@ -110,10 +114,13 @@ export class SessionPresenter { (await startActiveSpan( "SessionPresenter.findCurrentRunFallback", () => - this.replica.taskRun.findFirst({ - where: { id: session.currentRunId! }, - select: { id: true, friendlyId: true, status: true }, - }) + runStore.findRun( + { id: session.currentRunId! }, + { + select: { id: true, friendlyId: true, status: true }, + }, + this.replica + ) )) : null; diff --git a/apps/webapp/app/presenters/v3/SpanPresenter.server.ts b/apps/webapp/app/presenters/v3/SpanPresenter.server.ts index 98ee75cda3..49d8f30356 100644 --- a/apps/webapp/app/presenters/v3/SpanPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/SpanPresenter.server.ts @@ -35,6 +35,7 @@ import { import { getEventRepositoryForStore } from "~/v3/eventRepository/index.server"; import { findRunByIdWithMollifierFallback } from "~/v3/mollifier/readFallback.server"; import { buildSyntheticSpanRun } from "~/v3/mollifier/syntheticSpanRun.server"; +import { runStore } from "~/v3/runStore.server"; export type PromptSpanData = { slug: string; @@ -132,20 +133,23 @@ export class SpanPresenter extends BasePresenter { throw new Error("Project not found"); } - const parentRun = await this._prisma.taskRun.findFirst({ - select: { - traceId: true, - runtimeEnvironmentId: true, - projectId: true, - taskEventStore: true, - createdAt: true, - completedAt: true, - }, - where: { + const parentRun = await runStore.findRun( + { friendlyId: runFriendlyId, projectId: project.id, }, - }); + { + select: { + traceId: true, + runtimeEnvironmentId: true, + projectId: true, + taskEventStore: true, + createdAt: true, + completedAt: true, + }, + }, + this._prisma + ); if (!parentRun) { // PG miss → fall back to the mollifier buffer. Without this the @@ -494,7 +498,17 @@ export class SpanPresenter extends BasePresenter { spanId: string; environmentId: string; }) { - const run = await this._replica.taskRun.findFirst({ + const run = await runStore.findRun( + originalRunId + ? { + friendlyId: originalRunId, + runtimeEnvironmentId: environmentId, + } + : { + spanId, + runtimeEnvironmentId: environmentId, + }, + { select: { id: true, spanId: true, @@ -608,16 +622,9 @@ export class SpanPresenter extends BasePresenter { }, }, }, - where: originalRunId - ? { - friendlyId: originalRunId, - runtimeEnvironmentId: environmentId, - } - : { - spanId, - runtimeEnvironmentId: environmentId, - }, - }); + }, + this._replica + ); return run; } @@ -655,18 +662,21 @@ export class SpanPresenter extends BasePresenter { return; } - const triggeredRuns = await this._replica.taskRun.findMany({ - select: { - friendlyId: true, - taskIdentifier: true, - spanId: true, - createdAt: true, - status: true, - }, - where: { - parentSpanId: spanId, + const triggeredRuns = await runStore.findRuns( + { + where: { + parentSpanId: spanId, + }, + select: { + friendlyId: true, + taskIdentifier: true, + spanId: true, + createdAt: true, + status: true, + }, }, - }); + this._replica + ); const data = { spanId: span.spanId, diff --git a/apps/webapp/app/presenters/v3/TestTaskPresenter.server.ts b/apps/webapp/app/presenters/v3/TestTaskPresenter.server.ts index a9381ab60d..0ebf5054bb 100644 --- a/apps/webapp/app/presenters/v3/TestTaskPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/TestTaskPresenter.server.ts @@ -12,6 +12,7 @@ import { type PrismaClient } from "~/db.server"; import { RunsRepository } from "~/services/runsRepository/runsRepository.server"; import { getTimezones } from "~/utils/timezones.server"; import { findCurrentWorkerDeployment } from "~/v3/models/workerDeployment.server"; +import { runStore } from "~/v3/runStore.server"; import { queueTypeFromType } from "./QueueRetrievePresenter.server"; export type RunTemplate = TaskRunTemplate & { @@ -214,38 +215,41 @@ export class TestTaskPresenter { }, }); - const latestRuns = await this.replica.taskRun.findMany({ - select: { - id: true, - queue: true, - friendlyId: true, - taskIdentifier: true, - createdAt: true, - status: true, - payload: true, - payloadType: true, - seedMetadata: true, - seedMetadataType: true, - runtimeEnvironmentId: true, - concurrencyKey: true, - maxAttempts: true, - maxDurationInSeconds: true, - machinePreset: true, - ttl: true, - runTags: true, - }, - where: { - id: { - in: runIds, + const latestRuns = await runStore.findRuns( + { + where: { + id: { + in: runIds, + }, + payloadType: { + in: ["application/json", "application/super+json"], + }, }, - payloadType: { - in: ["application/json", "application/super+json"], + select: { + id: true, + queue: true, + friendlyId: true, + taskIdentifier: true, + createdAt: true, + status: true, + payload: true, + payloadType: true, + seedMetadata: true, + seedMetadataType: true, + runtimeEnvironmentId: true, + concurrencyKey: true, + maxAttempts: true, + maxDurationInSeconds: true, + machinePreset: true, + ttl: true, + runTags: true, + }, + orderBy: { + createdAt: "desc", }, }, - orderBy: { - createdAt: "desc", - }, - }); + this.replica + ); // Infer schema from existing run payloads when no explicit schema is defined let inferredPayloadSchema: unknown | undefined; From 126b05fd3da0e370b34066b6dab02ef59a9f0976 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 16:12:57 +0100 Subject: [PATCH 06/11] refactor(webapp): route API and loader TaskRun reads through the run store Relocate the route and loader TaskRun reads to the RunStore read methods, preserving the exact client per site, including the replica-resolve then writer-recheck realtime paths. Behavior-preserving. --- apps/webapp/app/routes/@.runs.$runParam.ts | 34 ++-- .../admin.api.v1.runs-replication.backfill.ts | 16 +- ....runs.$runFriendlyId.input-streams.wait.ts | 20 +- ...uns.$runFriendlyId.session-streams.wait.ts | 18 +- .../app/routes/api.v1.runs.$runId.metadata.ts | 10 +- .../api.v1.runs.$runId.spans.$spanId.ts | 35 ++-- .../app/routes/api.v1.runs.$runId.trace.ts | 8 +- .../routes/api.v1.runs.$runParam.replay.ts | 8 +- ...i.v1.sessions.$session.end-and-continue.ts | 19 +- apps/webapp/app/routes/api.v1.sessions.ts | 10 +- .../routes/api.v1.tasks.$taskId.trigger.ts | 14 +- .../app/routes/engine.v1.dev.disconnect.ts | 32 ++-- ...s.$snapshotFriendlyId.attempts.complete.ts | 8 +- ...hots.$snapshotFriendlyId.attempts.start.ts | 8 +- ...snapshots.$snapshotFriendlyId.heartbeat.ts | 8 +- ...ev.runs.$runFriendlyId.snapshots.latest.ts | 8 +- ...ne.v1.runs.$runFriendlyId.wait.duration.ts | 8 +- ...g.projects.$projectParam.runs.$runParam.ts | 14 +- .../projects.v3.$projectRef.runs.$runParam.ts | 14 +- .../app/routes/realtime.v1.runs.$runId.ts | 18 +- .../realtime.v1.streams.$runId.$streamId.ts | 57 +++--- ...streams.$runId.$target.$streamId.append.ts | 50 ++--- ...ime.v1.streams.$runId.$target.$streamId.ts | 89 +++++---- ...ltime.v1.streams.$runId.input.$streamId.ts | 39 ++-- ...projectParam.env.$envParam.logs.$logId.tsx | 10 +- ...tParam.env.$envParam.playground.action.tsx | 10 +- ...am.runs.$runParam.idempotencyKey.reset.tsx | 20 +- ...ram.realtime.v1.sessions.$sessionId.$io.ts | 12 +- ...am.realtime.v1.streams.$runId.$streamId.ts | 20 +- ...ltime.v1.streams.$runId.input.$streamId.ts | 20 +- .../route.tsx | 22 ++- .../resources.runs.$runParam.logs.download.ts | 32 ++-- .../app/routes/resources.runs.$runParam.ts | 174 +++++++++--------- .../resources.taskruns.$runParam.cancel.ts | 8 +- .../resources.taskruns.$runParam.debug.ts | 48 ++--- .../resources.taskruns.$runParam.replay.ts | 121 ++++++------ apps/webapp/app/routes/runs.$runParam.ts | 34 ++-- .../app/routes/sync.traces.runs.$traceId.ts | 22 ++- 38 files changed, 621 insertions(+), 477 deletions(-) diff --git a/apps/webapp/app/routes/@.runs.$runParam.ts b/apps/webapp/app/routes/@.runs.$runParam.ts index a709191271..cd1e1eade1 100644 --- a/apps/webapp/app/routes/@.runs.$runParam.ts +++ b/apps/webapp/app/routes/@.runs.$runParam.ts @@ -1,6 +1,7 @@ import { redirect, type LoaderFunctionArgs } from "@remix-run/server-runtime"; import { z } from "zod"; import { prisma } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { redirectWithErrorMessage } from "~/models/message.server"; import { requireUser } from "~/services/session.server"; import { impersonate, rootPath, v3RunPath, v3RunSpanPath } from "~/utils/pathBuilder"; @@ -28,29 +29,32 @@ export async function loader({ params, request }: LoaderFunctionArgs) { ); } - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runParam, }, - select: { - spanId: true, - runtimeEnvironment: { - select: { - slug: true, + { + select: { + spanId: true, + runtimeEnvironment: { + select: { + slug: true, + }, }, - }, - project: { - select: { - slug: true, - organization: { - select: { - slug: true, + project: { + select: { + slug: true, + organization: { + select: { + slug: true, + }, }, }, }, }, }, - }); + prisma + ); if (!run) { // Admin impersonation route — bypass org membership so admins can diff --git a/apps/webapp/app/routes/admin.api.v1.runs-replication.backfill.ts b/apps/webapp/app/routes/admin.api.v1.runs-replication.backfill.ts index c4d17ba875..af041353ad 100644 --- a/apps/webapp/app/routes/admin.api.v1.runs-replication.backfill.ts +++ b/apps/webapp/app/routes/admin.api.v1.runs-replication.backfill.ts @@ -2,6 +2,7 @@ import { type ActionFunctionArgs, json } from "@remix-run/server-runtime"; import { type TaskRun } from "@trigger.dev/database"; import { z } from "zod"; import { prisma } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { logger } from "~/services/logger.server"; import { requireAdminApiRequest } from "~/services/personalAccessToken.server"; import { runsReplicationInstance } from "~/services/runsReplicationInstance.server"; @@ -25,14 +26,17 @@ export async function action({ request }: ActionFunctionArgs) { const runs: TaskRun[] = []; for (let i = 0; i < runIds.length; i += MAX_BATCH_SIZE) { const batch = runIds.slice(i, i + MAX_BATCH_SIZE); - const batchRuns = await prisma.taskRun.findMany({ - where: { - id: { in: batch }, - status: { - in: FINAL_RUN_STATUSES, + const batchRuns = await runStore.findRuns( + { + where: { + id: { in: batch }, + status: { + in: FINAL_RUN_STATUSES, + }, }, }, - }); + prisma + ); runs.push(...batchRuns); } diff --git a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts index 0924bf3fc9..091312a13b 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts @@ -6,6 +6,7 @@ import { } from "@trigger.dev/core/v3"; import { WaitpointId } from "@trigger.dev/core/v3/isomorphic"; import { $replica } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { createWaitpointTag, MAX_TAGS_PER_WAITPOINT } from "~/models/waitpointTag.server"; import { deleteInputStreamWaitpoint, @@ -32,18 +33,21 @@ const { action, loader } = createActionApiRoute( }, async ({ authentication, body, params }) => { try { - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runFriendlyId, runtimeEnvironmentId: authentication.environment.id, }, - select: { - id: true, - friendlyId: true, - realtimeStreamsVersion: true, - streamBasinName: true, + { + select: { + id: true, + friendlyId: true, + realtimeStreamsVersion: true, + streamBasinName: true, + }, }, - }); + $replica + ); if (!run) { return json({ error: "Run not found" }, { status: 404 }); diff --git a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts index 39c3089441..cd88ef3828 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts @@ -6,6 +6,7 @@ import { import { WaitpointId } from "@trigger.dev/core/v3/isomorphic"; import { z } from "zod"; import { $replica } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { createWaitpointTag, MAX_TAGS_PER_WAITPOINT } from "~/models/waitpointTag.server"; import { canonicalSessionAddressingKey, @@ -38,17 +39,20 @@ const { action, loader } = createActionApiRoute( }, async ({ authentication, body, params }) => { try { - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runFriendlyId, runtimeEnvironmentId: authentication.environment.id, }, - select: { - id: true, - friendlyId: true, - realtimeStreamsVersion: true, + { + select: { + id: true, + friendlyId: true, + realtimeStreamsVersion: true, + }, }, - }); + $replica + ); if (!run) { return json({ error: "Run not found" }, { status: 404 }); diff --git a/apps/webapp/app/routes/api.v1.runs.$runId.metadata.ts b/apps/webapp/app/routes/api.v1.runs.$runId.metadata.ts index 3f22929aca..7ec10835c7 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runId.metadata.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runId.metadata.ts @@ -17,6 +17,7 @@ import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server import { ServiceValidationError } from "~/v3/services/common.server"; import { applyMetadataMutationToBufferedRun } from "~/v3/mollifier/applyMetadataMutation.server"; import { findRunByIdWithMollifierFallback } from "~/v3/mollifier/readFallback.server"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ runId: z.string(), @@ -39,10 +40,11 @@ export async function loader({ request, params }: LoaderFunctionArgs) { const env = authenticationResult.environment; - const pgRun = await $replica.taskRun.findFirst({ - where: { friendlyId: parsed.data.runId, runtimeEnvironmentId: env.id }, - select: { metadata: true, metadataType: true }, - }); + const pgRun = await runStore.findRun( + { friendlyId: parsed.data.runId, runtimeEnvironmentId: env.id }, + { select: { metadata: true, metadataType: true } }, + $replica + ); if (pgRun) { return json({ metadata: pgRun.metadata, metadataType: pgRun.metadataType }, { status: 200 }); } diff --git a/apps/webapp/app/routes/api.v1.runs.$runId.spans.$spanId.ts b/apps/webapp/app/routes/api.v1.runs.$runId.spans.$spanId.ts index c38206473c..061199f33e 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runId.spans.$spanId.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runId.spans.$spanId.ts @@ -11,6 +11,7 @@ import { getEventRepositoryForStore } from "~/v3/eventRepository/index.server"; import { getTaskEventStoreTableForRun } from "~/v3/taskEventStore.server"; import { findRunByIdWithMollifierFallback } from "~/v3/mollifier/readFallback.server"; import { buildSyntheticSpanDetailBody } from "~/v3/mollifier/syntheticApiResponses.server"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ runId: z.string(), @@ -28,9 +29,10 @@ type ResolvedRun = | { source: "buffer"; run: NonNullable>> }; async function findPgRun(runId: string, environmentId: string) { - return $replica.taskRun.findFirst({ - where: { friendlyId: runId, runtimeEnvironmentId: environmentId }, - }); + return runStore.findRun( + { friendlyId: runId, runtimeEnvironmentId: environmentId }, + $replica + ); } export const loader = createLoaderApiRoute( @@ -121,19 +123,22 @@ export const loader = createLoaderApiRoute( ? extractAISpanData(span.properties as Record, durationMs) : undefined; - const triggeredRuns = await $replica.taskRun.findMany({ - take: 50, - select: { - friendlyId: true, - taskIdentifier: true, - status: true, - createdAt: true, - }, - where: { - runtimeEnvironmentId: authentication.environment.id, - parentSpanId: params.spanId, + const triggeredRuns = await runStore.findRuns( + { + take: 50, + select: { + friendlyId: true, + taskIdentifier: true, + status: true, + createdAt: true, + }, + where: { + runtimeEnvironmentId: authentication.environment.id, + parentSpanId: params.spanId, + }, }, - }); + $replica + ); const properties = span.properties && diff --git a/apps/webapp/app/routes/api.v1.runs.$runId.trace.ts b/apps/webapp/app/routes/api.v1.runs.$runId.trace.ts index 04ae398194..f1aa4d5896 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runId.trace.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runId.trace.ts @@ -10,6 +10,7 @@ import { getEventRepositoryForStore } from "~/v3/eventRepository/index.server"; import { getTaskEventStoreTableForRun } from "~/v3/taskEventStore.server"; import { findRunByIdWithMollifierFallback } from "~/v3/mollifier/readFallback.server"; import { buildSyntheticTraceBody } from "~/v3/mollifier/syntheticApiResponses.server"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ runId: z.string(), // This is the run friendly ID @@ -26,9 +27,10 @@ type ResolvedRun = | { source: "buffer"; run: NonNullable>> }; async function findPgRun(runId: string, environmentId: string) { - return $replica.taskRun.findFirst({ - where: { friendlyId: runId, runtimeEnvironmentId: environmentId }, - }); + return runStore.findRun( + { friendlyId: runId, runtimeEnvironmentId: environmentId }, + $replica + ); } export const loader = createLoaderApiRoute( diff --git a/apps/webapp/app/routes/api.v1.runs.$runParam.replay.ts b/apps/webapp/app/routes/api.v1.runs.$runParam.replay.ts index 130f6ff163..4b238869d3 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runParam.replay.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runParam.replay.ts @@ -3,6 +3,7 @@ import { json } from "@remix-run/server-runtime"; import type { TaskRun } from "@trigger.dev/database"; import { z } from "zod"; import { prisma } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { authenticateApiRequest } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { ReplayTaskRunService } from "~/v3/services/replayTaskRun.server"; @@ -73,12 +74,13 @@ export async function action({ request, params }: ActionFunctionArgs) { // filter beyond friendlyId is the existing semantic; findFirst with // env scoping tightens it minimally without changing behaviour for // a correctly-authed caller. - let taskRun: TaskRun | null = await prisma.taskRun.findFirst({ - where: { + let taskRun: TaskRun | null = await runStore.findRun( + { friendlyId: runParam, runtimeEnvironmentId: env.id, }, - }); + prisma + ); if (!taskRun) { // Buffered fallback. SyntheticRun carries every field diff --git a/apps/webapp/app/routes/api.v1.sessions.$session.end-and-continue.ts b/apps/webapp/app/routes/api.v1.sessions.$session.end-and-continue.ts index 7c5718aeae..cc1a6d4f9f 100644 --- a/apps/webapp/app/routes/api.v1.sessions.$session.end-and-continue.ts +++ b/apps/webapp/app/routes/api.v1.sessions.$session.end-and-continue.ts @@ -12,6 +12,7 @@ import { anyResource, createActionApiRoute, } from "~/services/routeBuilders/apiBuilder.server"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ session: z.string(), @@ -83,13 +84,14 @@ const { action, loader } = createActionApiRoute( // SDK exposes via `ctx.run.id`). Internally `Session.currentRunId` // stores the TaskRun.id cuid, so resolve before handing to the // optimistic-claim service. - const callingRun = await $replica.taskRun.findFirst({ - where: { + const callingRun = await runStore.findRun( + { friendlyId: body.callingRunId, runtimeEnvironmentId: authentication.environment.id, }, - select: { id: true }, - }); + { select: { id: true } }, + $replica + ); if (!callingRun) { return json({ error: "callingRunId not found in this environment" }, { status: 404 }); } @@ -118,10 +120,11 @@ const { action, loader } = createActionApiRoute( // `$replica`. A replica miss here would silently fall back to // returning the internal cuid, which the public API contract // says is a friendlyId. - const run = await prisma.taskRun.findFirst({ - where: { id: result.runId }, - select: { friendlyId: true }, - }); + const run = await runStore.findRun( + { id: result.runId }, + { select: { friendlyId: true } }, + prisma + ); const responseBody: EndAndContinueSessionResponseBody = { runId: run?.friendlyId ?? result.runId, diff --git a/apps/webapp/app/routes/api.v1.sessions.ts b/apps/webapp/app/routes/api.v1.sessions.ts index 44f8c7ef69..ec8c171fc2 100644 --- a/apps/webapp/app/routes/api.v1.sessions.ts +++ b/apps/webapp/app/routes/api.v1.sessions.ts @@ -29,6 +29,7 @@ import { createLoaderApiRoute, } from "~/services/routeBuilders/apiBuilder.server"; import { ServiceValidationError } from "~/v3/services/common.server"; +import { runStore } from "~/v3/runStore.server"; function asArray(value: T | T[] | undefined): T[] | undefined { if (value === undefined) return undefined; @@ -264,10 +265,11 @@ const { action } = createActionApiRoute( // Read-after-write: the run was just triggered in this request, // so go to the writer rather than $replica. Replica lag here // would null this out and turn a successful create into a 500. - const run = await prisma.taskRun.findFirst({ - where: { id: ensureResult.runId }, - select: { friendlyId: true }, - }); + const run = await runStore.findRun( + { id: ensureResult.runId }, + { select: { friendlyId: true } }, + prisma + ); if (!run) { throw new Error(`Triggered run ${ensureResult.runId} not found`); } diff --git a/apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts b/apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts index 1f8a42af08..eb9e5d974e 100644 --- a/apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts +++ b/apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts @@ -20,6 +20,7 @@ import { saveRequestIdempotency, } from "~/utils/requestIdempotency.server"; import { sanitizeTriggerSource } from "~/utils/triggerSource"; +import { runStore } from "~/v3/runStore.server"; import { ServiceValidationError } from "~/v3/services/baseService.server"; import { OutOfEntitlementError, TriggerTaskService } from "~/v3/services/triggerTask.server"; @@ -77,14 +78,17 @@ const { action, loader } = createActionApiRoute( const cachedResponse = await handleRequestIdempotency(requestIdempotencyKey, { requestType: "trigger", findCachedEntity: async (cachedRequestId) => { - return await prisma.taskRun.findFirst({ - where: { + return await runStore.findRun( + { id: cachedRequestId, }, - select: { - friendlyId: true, + { + select: { + friendlyId: true, + }, }, - }); + prisma + ); }, buildResponse: (cachedRun) => ({ id: cachedRun.friendlyId, diff --git a/apps/webapp/app/routes/engine.v1.dev.disconnect.ts b/apps/webapp/app/routes/engine.v1.dev.disconnect.ts index 0cf92a53b7..0142830143 100644 --- a/apps/webapp/app/routes/engine.v1.dev.disconnect.ts +++ b/apps/webapp/app/routes/engine.v1.dev.disconnect.ts @@ -5,6 +5,7 @@ import { DevDisconnectRequestBody } from "@trigger.dev/core/v3"; import { BulkActionId, RunId } from "@trigger.dev/core/v3/isomorphic"; import { BulkActionNotificationType, BulkActionType } from "@trigger.dev/database"; import { prisma } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { logger } from "~/services/logger.server"; import { RateLimiter } from "~/services/rateLimiter.server"; import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; @@ -99,21 +100,24 @@ async function cancelRunsInline( ): Promise { const runIds = runFriendlyIds.map((fid) => RunId.toId(fid)); - const runs = await prisma.taskRun.findMany({ - where: { - id: { in: runIds }, - runtimeEnvironmentId: environmentId, + const runs = await runStore.findRuns( + { + where: { + id: { in: runIds }, + runtimeEnvironmentId: environmentId, + }, + select: { + id: true, + engine: true, + friendlyId: true, + status: true, + createdAt: true, + completedAt: true, + taskEventStore: true, + }, }, - select: { - id: true, - engine: true, - friendlyId: true, - status: true, - createdAt: true, - completedAt: true, - taskEventStore: true, - }, - }); + prisma + ); let cancelled = 0; const cancelService = new CancelTaskRunService(prisma); diff --git a/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.complete.ts b/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.complete.ts index da4bab693b..afc481a571 100644 --- a/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.complete.ts +++ b/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.complete.ts @@ -9,6 +9,7 @@ import { prisma } from "~/db.server"; import { logger } from "~/services/logger.server"; import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { engine } from "~/v3/runEngine.server"; +import { runStore } from "~/v3/runStore.server"; const { action } = createActionApiRoute( { @@ -28,12 +29,13 @@ const { action } = createActionApiRoute( const { runFriendlyId, snapshotFriendlyId } = params; try { - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runFriendlyId, runtimeEnvironmentId: authentication.environment.id, }, - }); + prisma + ); if (!run) { throw new Response("You don't have permissions for this run", { status: 401 }); diff --git a/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.start.ts b/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.start.ts index a3f35013b7..4c05704647 100644 --- a/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.start.ts +++ b/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.start.ts @@ -18,6 +18,7 @@ import { import { resolveVariablesForEnvironment } from "~/v3/environmentVariables/environmentVariablesRepository.server"; import { machinePresetFromName } from "~/v3/machinePresets.server"; import { engine } from "~/v3/runEngine.server"; +import { runStore } from "~/v3/runStore.server"; const { action } = createActionApiRoute( { @@ -36,12 +37,13 @@ const { action } = createActionApiRoute( const { runFriendlyId, snapshotFriendlyId } = params; try { - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runFriendlyId, runtimeEnvironmentId: authentication.environment.id, }, - }); + prisma + ); if (!run) { throw new Response("You don't have permissions for this run", { status: 401 }); diff --git a/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.heartbeat.ts b/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.heartbeat.ts index bab59fd063..d9f6ca9a6d 100644 --- a/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.heartbeat.ts +++ b/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.heartbeat.ts @@ -6,6 +6,7 @@ import { prisma } from "~/db.server"; import { logger } from "~/services/logger.server"; import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { engine } from "~/v3/runEngine.server"; +import { runStore } from "~/v3/runStore.server"; const { action } = createActionApiRoute( { @@ -23,12 +24,13 @@ const { action } = createActionApiRoute( const { runFriendlyId, snapshotFriendlyId } = params; try { - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runFriendlyId, runtimeEnvironmentId: authentication.environment.id, }, - }); + prisma + ); if (!run) { throw new Response("You don't have permissions for this run", { status: 401 }); diff --git a/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.latest.ts b/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.latest.ts index 60505460bd..9254a74e83 100644 --- a/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.latest.ts +++ b/apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.latest.ts @@ -6,6 +6,7 @@ import { prisma } from "~/db.server"; import { logger } from "~/services/logger.server"; import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { engine } from "~/v3/runEngine.server"; +import { runStore } from "~/v3/runStore.server"; export const loader = createLoaderApiRoute( { @@ -24,12 +25,13 @@ export const loader = createLoaderApiRoute( }); try { - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runFriendlyId, runtimeEnvironmentId: authentication.environment.id, }, - }); + prisma + ); if (!run) { throw new Response("You don't have permissions for this run", { status: 401 }); diff --git a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts index 199244b1da..8d7f6b8434 100644 --- a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts +++ b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts @@ -8,6 +8,7 @@ import { logger } from "~/services/logger.server"; import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server"; import { engine } from "~/v3/runEngine.server"; +import { runStore } from "~/v3/runStore.server"; const { action } = createActionApiRoute( { @@ -22,12 +23,13 @@ const { action } = createActionApiRoute( const runId = RunId.toId(runFriendlyId); try { - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { id: runId, runtimeEnvironmentId: authentication.environment.id, }, - }); + prisma + ); if (!run) { throw new Response("You don't have permissions for this run", { status: 401 }); diff --git a/apps/webapp/app/routes/orgs.$organizationSlug.projects.$projectParam.runs.$runParam.ts b/apps/webapp/app/routes/orgs.$organizationSlug.projects.$projectParam.runs.$runParam.ts index 63a89d7e0a..d5d4ab0f2f 100644 --- a/apps/webapp/app/routes/orgs.$organizationSlug.projects.$projectParam.runs.$runParam.ts +++ b/apps/webapp/app/routes/orgs.$organizationSlug.projects.$projectParam.runs.$runParam.ts @@ -4,6 +4,7 @@ import { z } from "zod"; import { prisma } from "~/db.server"; import { requireUserId } from "~/services/session.server"; import { ProjectParamSchema, v3RunPath } from "~/utils/pathBuilder"; +import { runStore } from "~/v3/runStore.server"; const ParamSchema = ProjectParamSchema.extend({ runParam: z.string(), @@ -13,8 +14,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const userId = await requireUserId(request); const { organizationSlug, projectParam, runParam } = ParamSchema.parse(params); - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runParam, project: { slug: projectParam, @@ -28,10 +29,13 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { }, }, }, - select: { - runtimeEnvironment: true, + { + select: { + runtimeEnvironment: true, + }, }, - }); + prisma + ); if (!run) { throw new Response("Not Found", { status: 404 }); diff --git a/apps/webapp/app/routes/projects.v3.$projectRef.runs.$runParam.ts b/apps/webapp/app/routes/projects.v3.$projectRef.runs.$runParam.ts index fe267d1f9f..2a6cb34c91 100644 --- a/apps/webapp/app/routes/projects.v3.$projectRef.runs.$runParam.ts +++ b/apps/webapp/app/routes/projects.v3.$projectRef.runs.$runParam.ts @@ -3,6 +3,7 @@ import { z } from "zod"; import { prisma } from "~/db.server"; import { requireUserId } from "~/services/session.server"; import { v3RunSpanPath } from "~/utils/pathBuilder"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ projectRef: z.string(), @@ -34,14 +35,17 @@ export async function loader({ params, request }: LoaderFunctionArgs) { return new Response("Not found", { status: 404 }); } - const run = await prisma.taskRun.findUnique({ - where: { + const run = await runStore.findRun( + { friendlyId: validatedParams.runParam, }, - include: { - runtimeEnvironment: true, + { + include: { + runtimeEnvironment: true, + }, }, - }); + prisma + ); if (!run) { throw new Response("Not found", { status: 404 }); diff --git a/apps/webapp/app/routes/realtime.v1.runs.$runId.ts b/apps/webapp/app/routes/realtime.v1.runs.$runId.ts index 46118c1d89..f2268989bd 100644 --- a/apps/webapp/app/routes/realtime.v1.runs.$runId.ts +++ b/apps/webapp/app/routes/realtime.v1.runs.$runId.ts @@ -7,6 +7,7 @@ import { anyResource, createLoaderApiRoute, } from "~/services/routeBuilders/apiBuilder.server"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ runId: z.string(), @@ -18,19 +19,22 @@ export const loader = createLoaderApiRoute( allowJWT: true, corsStrategy: "all", findResource: async (params, authentication) => { - return $replica.taskRun.findFirst({ - where: { + return runStore.findRun( + { friendlyId: params.runId, runtimeEnvironmentId: authentication.environment.id, }, - include: { - batch: { - select: { - friendlyId: true, + { + include: { + batch: { + select: { + friendlyId: true, + }, }, }, }, - }); + $replica + ); }, authorization: { action: "read", diff --git a/apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts b/apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts index d6470794a7..81784f9bc3 100644 --- a/apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts +++ b/apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts @@ -8,6 +8,7 @@ import { createLoaderApiRoute, } from "~/services/routeBuilders/apiBuilder.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ runId: z.string(), @@ -25,23 +26,26 @@ export async function action({ request, params }: ActionFunctionArgs) { const { runId, streamId } = parsedParams.data; // Look up the run without environment scoping for backwards compatibility - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runId, }, - select: { - id: true, - friendlyId: true, - streamBasinName: true, - runtimeEnvironment: { - include: { - project: true, - organization: true, - orgMember: true, + { + select: { + id: true, + friendlyId: true, + streamBasinName: true, + runtimeEnvironment: { + include: { + project: true, + organization: true, + orgMember: true, + }, }, }, }, - }); + $replica + ); if (!run) { return new Response("Run not found", { status: 404 }); @@ -87,25 +91,28 @@ export const loader = createLoaderApiRoute( allowJWT: true, corsStrategy: "all", findResource: async (params, auth) => { - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runId, runtimeEnvironmentId: auth.environment.id, }, - select: { - id: true, - friendlyId: true, - taskIdentifier: true, - runTags: true, - realtimeStreamsVersion: true, - streamBasinName: true, - batch: { - select: { - friendlyId: true, + { + select: { + id: true, + friendlyId: true, + taskIdentifier: true, + runTags: true, + realtimeStreamsVersion: true, + streamBasinName: true, + batch: { + select: { + friendlyId: true, + }, }, }, }, - }); + $replica + ); return run; }, authorization: { diff --git a/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.append.ts b/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.append.ts index 11074840a3..7cb813a6de 100644 --- a/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.append.ts +++ b/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.append.ts @@ -27,26 +27,29 @@ const { action } = createActionApiRoute( maxContentLength: MAX_APPEND_BODY_BYTES, }, async ({ request, params, authentication }) => { - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runId, runtimeEnvironmentId: authentication.environment.id, }, - select: { - id: true, - friendlyId: true, - parentTaskRun: { - select: { - friendlyId: true, + { + select: { + id: true, + friendlyId: true, + parentTaskRun: { + select: { + friendlyId: true, + }, }, - }, - rootTaskRun: { - select: { - friendlyId: true, + rootTaskRun: { + select: { + friendlyId: true, + }, }, }, }, - }); + $replica + ); if (!run) { return new Response("Run not found", { status: 404 }); @@ -63,19 +66,22 @@ const { action } = createActionApiRoute( return new Response("Target not found", { status: 404 }); } - const targetRun = await prisma.taskRun.findFirst({ - where: { + const targetRun = await runStore.findRun( + { friendlyId: targetId, runtimeEnvironmentId: authentication.environment.id, }, - select: { - realtimeStreams: true, - realtimeStreamsVersion: true, - completedAt: true, - id: true, - streamBasinName: true, + { + select: { + realtimeStreams: true, + realtimeStreamsVersion: true, + completedAt: true, + id: true, + streamBasinName: true, + }, }, - }); + prisma + ); if (!targetRun) { return new Response("Run not found", { status: 404 }); diff --git a/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.ts b/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.ts index cdee9567b7..c71ad48d12 100644 --- a/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.ts +++ b/apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.ts @@ -19,29 +19,32 @@ const { action } = createActionApiRoute( params: ParamsSchema, }, async ({ request, params, authentication }) => { - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runId, runtimeEnvironmentId: authentication.environment.id, }, - select: { - id: true, - friendlyId: true, - streamBasinName: true, - parentTaskRun: { - select: { - friendlyId: true, - streamBasinName: true, + { + select: { + id: true, + friendlyId: true, + streamBasinName: true, + parentTaskRun: { + select: { + friendlyId: true, + streamBasinName: true, + }, }, - }, - rootTaskRun: { - select: { - friendlyId: true, - streamBasinName: true, + rootTaskRun: { + select: { + friendlyId: true, + streamBasinName: true, + }, }, }, }, - }); + $replica + ); if (!run) { return new Response("Run not found", { status: 404 }); @@ -63,18 +66,21 @@ const { action } = createActionApiRoute( if (request.method === "PUT") { // This is the "create" endpoint - const target = await prisma.taskRun.findFirst({ - where: { + const target = await runStore.findRun( + { friendlyId: targetId, runtimeEnvironmentId: authentication.environment.id, }, - select: { - id: true, - realtimeStreams: true, - realtimeStreamsVersion: true, - completedAt: true, + { + select: { + id: true, + realtimeStreams: true, + realtimeStreamsVersion: true, + completedAt: true, + }, }, - }); + prisma + ); if (!target) { return new Response("Run not found", { status: 404 }); @@ -148,29 +154,32 @@ const loader = createLoaderApiRoute( allowJWT: false, corsStrategy: "none", findResource: async (params, authentication) => { - return $replica.taskRun.findFirst({ - where: { + return runStore.findRun( + { friendlyId: params.runId, runtimeEnvironmentId: authentication.environment.id, }, - select: { - id: true, - friendlyId: true, - streamBasinName: true, - parentTaskRun: { - select: { - friendlyId: true, - streamBasinName: true, + { + select: { + id: true, + friendlyId: true, + streamBasinName: true, + parentTaskRun: { + select: { + friendlyId: true, + streamBasinName: true, + }, }, - }, - rootTaskRun: { - select: { - friendlyId: true, - streamBasinName: true, + rootTaskRun: { + select: { + friendlyId: true, + streamBasinName: true, + }, }, }, }, - }); + $replica + ); }, }, async ({ request, params, resource: run, authentication }) => { diff --git a/apps/webapp/app/routes/realtime.v1.streams.$runId.input.$streamId.ts b/apps/webapp/app/routes/realtime.v1.streams.$runId.input.$streamId.ts index a404e6a76a..78fe332b8a 100644 --- a/apps/webapp/app/routes/realtime.v1.streams.$runId.input.$streamId.ts +++ b/apps/webapp/app/routes/realtime.v1.streams.$runId.input.$streamId.ts @@ -15,6 +15,7 @@ import { import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server"; import { engine } from "~/v3/runEngine.server"; import { ServiceValidationError } from "~/v3/services/common.server"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ runId: z.string(), @@ -38,19 +39,22 @@ const { action } = createActionApiRoute( }, }, async ({ request, params, authentication }) => { - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: params.runId, runtimeEnvironmentId: authentication.environment.id, }, - select: { - id: true, - friendlyId: true, - completedAt: true, - realtimeStreamsVersion: true, - streamBasinName: true, + { + select: { + id: true, + friendlyId: true, + completedAt: true, + realtimeStreamsVersion: true, + streamBasinName: true, + }, }, - }); + $replica + ); if (!run) { return json({ ok: false, error: "Run not found" }, { status: 404 }); @@ -129,19 +133,22 @@ const loader = createLoaderApiRoute( allowJWT: true, corsStrategy: "all", findResource: async (params, auth) => { - return $replica.taskRun.findFirst({ - where: { + return runStore.findRun( + { friendlyId: params.runId, runtimeEnvironmentId: auth.environment.id, }, - include: { - batch: { - select: { - friendlyId: true, + { + include: { + batch: { + select: { + friendlyId: true, + }, }, }, }, - }); + $replica + ); }, authorization: { action: "read", diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.$logId.tsx b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.$logId.tsx index f4d3490704..be4fdba7fe 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.$logId.tsx +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.$logId.tsx @@ -7,6 +7,7 @@ import { LogDetailPresenter } from "~/presenters/v3/LogDetailPresenter.server"; import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { $replica } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { ServiceValidationError } from "~/v3/services/baseService.server"; import type { TaskRunStatus } from "@trigger.dev/database"; @@ -70,13 +71,14 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { // Look up the run status from Postgres let runStatus: TaskRunStatus | undefined; if (result.runId) { - const run = await $replica.taskRun.findFirst({ - select: { status: true }, - where: { + const run = await runStore.findRun( + { friendlyId: result.runId, runtimeEnvironmentId: environment.id, }, - }); + { select: { status: true } }, + $replica + ); runStatus = run?.status; } diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.action.tsx b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.action.tsx index 0fab90e145..da77d2cc69 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.action.tsx +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.action.tsx @@ -10,6 +10,7 @@ import { requireUserId } from "~/services/session.server"; import { EnvironmentParamSchema } from "~/utils/pathBuilder"; import { mintSessionToken } from "~/services/realtime/mintSessionToken.server"; import { ensureRunForSession } from "~/services/realtime/sessionRunManager.server"; +import { runStore } from "~/v3/runStore.server"; const PlaygroundAction = z.object({ intent: z.enum(["create", "start", "save", "delete"]), @@ -183,10 +184,11 @@ export const action = async ({ request, params }: ActionFunctionArgs) => { reason: "initial", }); - const run = await prisma.taskRun.findFirst({ - where: { id: ensureResult.runId }, - select: { friendlyId: true }, - }); + const run = await runStore.findRun( + { id: ensureResult.runId }, + { select: { friendlyId: true } }, + prisma + ); if (!run) { return json({ error: "Triggered run not found" }, { status: 500 }); } diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.idempotencyKey.reset.tsx b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.idempotencyKey.reset.tsx index 614b668f91..06233f88c7 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.idempotencyKey.reset.tsx +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.idempotencyKey.reset.tsx @@ -3,6 +3,7 @@ import { prisma } from "~/db.server"; import { jsonWithErrorMessage, jsonWithSuccessMessage } from "~/models/message.server"; import { logger } from "~/services/logger.server"; import { requireUserId } from "~/services/session.server"; +import { runStore } from "~/v3/runStore.server"; import { ResetIdempotencyKeyService } from "~/v3/services/resetIdempotencyKey.server"; import { v3RunParamsSchema } from "~/utils/pathBuilder"; @@ -11,8 +12,8 @@ export const action: ActionFunction = async ({ request, params }) => { const { projectParam, organizationSlug, envParam, runParam } = v3RunParamsSchema.parse(params); try { - const taskRun = await prisma.taskRun.findFirst({ - where: { + const taskRun = await runStore.findRun( + { friendlyId: runParam, project: { slug: projectParam, @@ -29,13 +30,16 @@ export const action: ActionFunction = async ({ request, params }) => { slug: envParam, }, }, - select: { - id: true, - idempotencyKey: true, - taskIdentifier: true, - runtimeEnvironmentId: true, + { + select: { + id: true, + idempotencyKey: true, + taskIdentifier: true, + runtimeEnvironmentId: true, + }, }, - }); + prisma + ); if (!taskRun) { return jsonWithErrorMessage({}, request, "Run not found"); diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.sessions.$sessionId.$io.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.sessions.$sessionId.$io.ts index 6613534725..3a0dfca568 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.sessions.$sessionId.$io.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.sessions.$sessionId.$io.ts @@ -1,6 +1,7 @@ import { type LoaderFunctionArgs } from "@remix-run/server-runtime"; import { z } from "zod"; import { $replica } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { getRequestAbortSignal } from "~/services/httpAsyncStorage.server"; @@ -50,13 +51,16 @@ export async function loader({ request, params }: LoaderFunctionArgs) { // Verify the run lives in this environment — keeps callers from // subscribing to arbitrary sessions via `/runs/$runParam/...`. - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runParam, runtimeEnvironmentId: environment.id, }, - select: { id: true, friendlyId: true }, - }); + { + select: { id: true, friendlyId: true }, + }, + $replica + ); if (!run) { return new Response("Run not found", { status: 404 }); diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.streams.$runId.$streamId.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.streams.$runId.$streamId.ts index 8d0af728df..cec6c3c4e9 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.streams.$runId.$streamId.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.streams.$runId.$streamId.ts @@ -7,6 +7,7 @@ import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { requireUserId } from "~/services/session.server"; import { EnvironmentParamSchema } from "~/utils/pathBuilder"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ runParam: z.string(), @@ -44,18 +45,21 @@ export async function loader({ request, params }: LoaderFunctionArgs) { return new Response("Environment not found", { status: 404 }); } - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runId, runtimeEnvironmentId: environment.id, }, - select: { - id: true, - friendlyId: true, - realtimeStreamsVersion: true, - streamBasinName: true, + { + select: { + id: true, + friendlyId: true, + realtimeStreamsVersion: true, + streamBasinName: true, + }, }, - }); + $replica + ); if (!run) { return new Response("Run not found", { status: 404 }); diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.streams.$runId.input.$streamId.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.streams.$runId.input.$streamId.ts index c9480299cc..1ecc7819c2 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.streams.$runId.input.$streamId.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.realtime.v1.streams.$runId.input.$streamId.ts @@ -7,6 +7,7 @@ import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { requireUserId } from "~/services/session.server"; import { EnvironmentParamSchema } from "~/utils/pathBuilder"; +import { runStore } from "~/v3/runStore.server"; const ParamsSchema = z.object({ runParam: z.string(), @@ -46,18 +47,21 @@ export async function loader({ request, params }: LoaderFunctionArgs) { return new Response("Environment not found", { status: 404 }); } - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runId, runtimeEnvironmentId: environment.id, }, - select: { - id: true, - friendlyId: true, - realtimeStreamsVersion: true, - streamBasinName: true, + { + select: { + id: true, + friendlyId: true, + realtimeStreamsVersion: true, + streamBasinName: true, + }, }, - }); + $replica + ); if (!run) { return new Response("Run not found", { status: 404 }); diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.streams.$streamKey/route.tsx b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.streams.$streamKey/route.tsx index 60233d6d38..24e7a73374 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.streams.$streamKey/route.tsx +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.streams.$streamKey/route.tsx @@ -26,6 +26,7 @@ import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.s import { requireUserId } from "~/services/session.server"; import { cn } from "~/utils/cn"; import { v3RunStreamParamsSchema } from "~/utils/pathBuilder"; +import { runStore } from "~/v3/runStore.server"; type ViewMode = "list" | "compact"; @@ -58,21 +59,24 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { throw new Response("Not Found", { status: 404 }); } - const run = await $replica.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runParam, projectId: project.id, }, - include: { - runtimeEnvironment: { - include: { - project: true, - organization: true, - orgMember: true, + { + include: { + runtimeEnvironment: { + include: { + project: true, + organization: true, + orgMember: true, + }, }, }, }, - }); + $replica + ); if (!run) { throw new Response("Not Found", { status: 404 }); diff --git a/apps/webapp/app/routes/resources.runs.$runParam.logs.download.ts b/apps/webapp/app/routes/resources.runs.$runParam.logs.download.ts index 7cda5ac782..7662a88b4d 100644 --- a/apps/webapp/app/routes/resources.runs.$runParam.logs.download.ts +++ b/apps/webapp/app/routes/resources.runs.$runParam.logs.download.ts @@ -1,6 +1,7 @@ import { LoaderFunctionArgs } from "@remix-run/server-runtime"; import { prisma } from "~/db.server"; import { env } from "~/env.server"; +import { runStore } from "~/v3/runStore.server"; import { requireUser } from "~/services/session.server"; import { v3RunParamsSchema, v3RunPath } from "~/utils/pathBuilder"; import { createGzip } from "zlib"; @@ -26,8 +27,8 @@ export async function loader({ params, request }: LoaderFunctionArgs) { const showDebug = url.searchParams.get("showDebug") === "true" && user.admin; const filename = `${parsedParams.runParam}.${format.extension}`; - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: parsedParams.runParam, project: { organization: { @@ -39,19 +40,22 @@ export async function loader({ params, request }: LoaderFunctionArgs) { }, }, }, - select: { - friendlyId: true, - traceId: true, - organizationId: true, - runtimeEnvironmentId: true, - createdAt: true, - completedAt: true, - taskEventStore: true, - taskIdentifier: true, - project: { select: { slug: true, organization: { select: { slug: true } } } }, - runtimeEnvironment: { select: { slug: true } }, + { + select: { + friendlyId: true, + traceId: true, + organizationId: true, + runtimeEnvironmentId: true, + createdAt: true, + completedAt: true, + taskEventStore: true, + taskIdentifier: true, + project: { select: { slug: true, organization: { select: { slug: true } } } }, + runtimeEnvironment: { select: { slug: true } }, + }, }, - }); + prisma + ); if (!run || !run.organizationId) { // Buffered run? It hasn't executed, so there's no trace to stream — but a diff --git a/apps/webapp/app/routes/resources.runs.$runParam.ts b/apps/webapp/app/routes/resources.runs.$runParam.ts index c5e467533a..38e17531f6 100644 --- a/apps/webapp/app/routes/resources.runs.$runParam.ts +++ b/apps/webapp/app/routes/resources.runs.$runParam.ts @@ -6,6 +6,7 @@ import { $replica } from "~/db.server"; import { requireUserId } from "~/services/session.server"; import { v3RunParamsSchema } from "~/utils/pathBuilder"; import { machinePresetFromName, machinePresetFromRun } from "~/v3/machinePresets.server"; +import { runStore } from "~/v3/runStore.server"; import { FINAL_ATTEMPT_STATUSES, isFinalRunStatus } from "~/v3/taskStatus"; export type RunInspectorData = UseDataFunctionReturn; @@ -14,104 +15,107 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const userId = await requireUserId(request); const parsedParams = v3RunParamsSchema.pick({ runParam: true }).parse(params); - const run = await $replica.taskRun.findFirst({ - select: { - id: true, - traceId: true, - //metadata - number: true, - taskIdentifier: true, - friendlyId: true, - isTest: true, - runTags: true, - machinePreset: true, - lockedToVersion: { - select: { - version: true, - sdkVersion: true, - }, - }, - //status + duration - status: true, - startedAt: true, - createdAt: true, - updatedAt: true, - queuedAt: true, - completedAt: true, - logsDeletedAt: true, - //idempotency - idempotencyKey: true, - //delayed - delayUntil: true, - //ttl - ttl: true, - expiredAt: true, - //queue - queue: true, - concurrencyKey: true, - //schedule - scheduleId: true, - //usage - baseCostInCents: true, - costInCents: true, - usageDurationMs: true, - //env - runtimeEnvironment: { - select: { id: true, slug: true, type: true }, - }, - payload: true, - payloadType: true, - metadata: true, - metadataType: true, - maxAttempts: true, + const run = await runStore.findRun( + { + friendlyId: parsedParams.runParam, project: { - include: { - organization: true, + organization: { + members: { + some: { + userId, + }, + }, }, }, - lockedBy: { - select: { - filePath: true, - worker: { - select: { - deployment: { - select: { - friendlyId: true, - shortCode: true, - version: true, - runtime: true, - runtimeVersion: true, - git: true, + }, + { + select: { + id: true, + traceId: true, + //metadata + number: true, + taskIdentifier: true, + friendlyId: true, + isTest: true, + runTags: true, + machinePreset: true, + lockedToVersion: { + select: { + version: true, + sdkVersion: true, + }, + }, + //status + duration + status: true, + startedAt: true, + createdAt: true, + updatedAt: true, + queuedAt: true, + completedAt: true, + logsDeletedAt: true, + //idempotency + idempotencyKey: true, + //delayed + delayUntil: true, + //ttl + ttl: true, + expiredAt: true, + //queue + queue: true, + concurrencyKey: true, + //schedule + scheduleId: true, + //usage + baseCostInCents: true, + costInCents: true, + usageDurationMs: true, + //env + runtimeEnvironment: { + select: { id: true, slug: true, type: true }, + }, + payload: true, + payloadType: true, + metadata: true, + metadataType: true, + maxAttempts: true, + project: { + include: { + organization: true, + }, + }, + lockedBy: { + select: { + filePath: true, + worker: { + select: { + deployment: { + select: { + friendlyId: true, + shortCode: true, + version: true, + runtime: true, + runtimeVersion: true, + git: true, + }, }, }, }, }, }, - }, - parentTaskRun: { - select: { - friendlyId: true, - }, - }, - rootTaskRun: { - select: { - friendlyId: true, + parentTaskRun: { + select: { + friendlyId: true, + }, }, - }, - }, - where: { - friendlyId: parsedParams.runParam, - project: { - organization: { - members: { - some: { - userId, - }, + rootTaskRun: { + select: { + friendlyId: true, }, }, }, }, - }); + $replica + ); if (!run) { throw new Response("Not found", { status: 404 }); diff --git a/apps/webapp/app/routes/resources.taskruns.$runParam.cancel.ts b/apps/webapp/app/routes/resources.taskruns.$runParam.cancel.ts index fa6ee29f3d..ca92615bb8 100644 --- a/apps/webapp/app/routes/resources.taskruns.$runParam.cancel.ts +++ b/apps/webapp/app/routes/resources.taskruns.$runParam.cancel.ts @@ -7,6 +7,7 @@ import { logger } from "~/services/logger.server"; import { requireUserId } from "~/services/session.server"; import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server"; import { getMollifierBuffer } from "~/v3/mollifier/mollifierBuffer.server"; +import { runStore } from "~/v3/runStore.server"; export const cancelSchema = z.object({ redirectUrl: z.string(), @@ -28,8 +29,8 @@ export const action: ActionFunction = async ({ request, params }) => { } try { - const taskRun = await prisma.taskRun.findFirst({ - where: { + const taskRun = await runStore.findRun( + { friendlyId: runParam, project: { organization: { @@ -41,7 +42,8 @@ export const action: ActionFunction = async ({ request, params }) => { }, }, }, - }); + prisma + ); if (taskRun) { const cancelRunService = new CancelTaskRunService(); diff --git a/apps/webapp/app/routes/resources.taskruns.$runParam.debug.ts b/apps/webapp/app/routes/resources.taskruns.$runParam.debug.ts index d7acf18e51..7b37b1bcc0 100644 --- a/apps/webapp/app/routes/resources.taskruns.$runParam.debug.ts +++ b/apps/webapp/app/routes/resources.taskruns.$runParam.debug.ts @@ -5,6 +5,7 @@ import { $replica } from "~/db.server"; import { requireUserId } from "~/services/session.server"; import { marqs } from "~/v3/marqs/index.server"; import { engine } from "~/v3/runEngine.server"; +import { runStore } from "~/v3/runStore.server"; const ParamSchema = z.object({ runParam: z.string(), @@ -14,33 +15,36 @@ export async function loader({ request, params }: LoaderFunctionArgs) { const userId = await requireUserId(request); const { runParam } = ParamSchema.parse(params); - const run = await $replica.taskRun.findFirst({ - where: { friendlyId: runParam, project: { organization: { members: { some: { userId } } } } }, - select: { - id: true, - engine: true, - friendlyId: true, - queue: true, - concurrencyKey: true, - queueTimestamp: true, - runtimeEnvironment: { - select: { - id: true, - type: true, - slug: true, - organizationId: true, - project: true, - maximumConcurrencyLimit: true, - concurrencyLimitBurstFactor: true, - organization: { - select: { - id: true, + const run = await runStore.findRun( + { friendlyId: runParam, project: { organization: { members: { some: { userId } } } } }, + { + select: { + id: true, + engine: true, + friendlyId: true, + queue: true, + concurrencyKey: true, + queueTimestamp: true, + runtimeEnvironment: { + select: { + id: true, + type: true, + slug: true, + organizationId: true, + project: true, + maximumConcurrencyLimit: true, + concurrencyLimitBurstFactor: true, + organization: { + select: { + id: true, + }, }, }, }, }, }, - }); + $replica + ); if (!run) { throw new Response("Not Found", { status: 404 }); diff --git a/apps/webapp/app/routes/resources.taskruns.$runParam.replay.ts b/apps/webapp/app/routes/resources.taskruns.$runParam.replay.ts index 03bfdaccc6..0719a8e6a1 100644 --- a/apps/webapp/app/routes/resources.taskruns.$runParam.replay.ts +++ b/apps/webapp/app/routes/resources.taskruns.$runParam.replay.ts @@ -23,6 +23,7 @@ import { findCurrentWorkerDeployment } from "~/v3/models/workerDeployment.server import { queueTypeFromType } from "~/presenters/v3/QueueRetrievePresenter.server"; import { ReplayRunData } from "~/v3/replayTask"; import { RegionsPresenter } from "~/presenters/v3/RegionsPresenter.server"; +import { runStore } from "~/v3/runStore.server"; const ParamSchema = z.object({ runParam: z.string(), @@ -40,61 +41,64 @@ export async function loader({ request, params }: LoaderFunctionArgs) { Object.fromEntries(new URL(request.url).searchParams) ); - let run = await $replica.taskRun.findFirst({ - select: { - payload: true, - payloadType: true, - seedMetadata: true, - seedMetadataType: true, - runtimeEnvironmentId: true, - concurrencyKey: true, - maxAttempts: true, - maxDurationInSeconds: true, - machinePreset: true, - workerQueue: true, - region: true, - ttl: true, - idempotencyKey: true, - runTags: true, - queue: true, - taskIdentifier: true, - project: { - select: { - slug: true, - environments: { - select: { - id: true, - type: true, - slug: true, - branchName: true, - orgMember: { - select: { - user: true, + let run = await runStore.findRun( + { friendlyId: runParam, project: { organization: { members: { some: { userId } } } } }, + { + select: { + payload: true, + payloadType: true, + seedMetadata: true, + seedMetadataType: true, + runtimeEnvironmentId: true, + concurrencyKey: true, + maxAttempts: true, + maxDurationInSeconds: true, + machinePreset: true, + workerQueue: true, + region: true, + ttl: true, + idempotencyKey: true, + runTags: true, + queue: true, + taskIdentifier: true, + project: { + select: { + slug: true, + environments: { + select: { + id: true, + type: true, + slug: true, + branchName: true, + orgMember: { + select: { + user: true, + }, }, }, - }, - where: { - archivedAt: null, - OR: [ - { - type: { - in: ["PREVIEW", "STAGING", "PRODUCTION"], + where: { + archivedAt: null, + OR: [ + { + type: { + in: ["PREVIEW", "STAGING", "PRODUCTION"], + }, }, - }, - { - type: "DEVELOPMENT", - orgMember: { - userId, + { + type: "DEVELOPMENT", + orgMember: { + userId, + }, }, - }, - ], + ], + }, }, }, }, }, }, - where: { friendlyId: runParam, project: { organization: { members: { some: { userId } } } } }, - }); + $replica + ); let synthetic: | (Awaited> & { __synth: true }) @@ -272,8 +276,8 @@ export const action: ActionFunction = async ({ request, params }) => { } try { - const pgRun = await prisma.taskRun.findFirst({ - where: { + const pgRun = await runStore.findRun( + { friendlyId: runParam, project: { organization: { @@ -285,19 +289,22 @@ export const action: ActionFunction = async ({ request, params }) => { }, }, }, - include: { - runtimeEnvironment: { - select: { - slug: true, + { + include: { + runtimeEnvironment: { + select: { + slug: true, + }, }, - }, - project: { - include: { - organization: true, + project: { + include: { + organization: true, + }, }, }, }, - }); + prisma + ); // Mollifier read-fallback: if the original isn't in PG yet, // synthesise a TaskRun from the buffered snapshot. The B4-extended diff --git a/apps/webapp/app/routes/runs.$runParam.ts b/apps/webapp/app/routes/runs.$runParam.ts index b472d7ae8f..5e0c2b21d6 100644 --- a/apps/webapp/app/routes/runs.$runParam.ts +++ b/apps/webapp/app/routes/runs.$runParam.ts @@ -1,6 +1,7 @@ import { redirect, type LoaderFunctionArgs } from "@remix-run/server-runtime"; import { z } from "zod"; import { prisma } from "~/db.server"; +import { runStore } from "~/v3/runStore.server"; import { redirectWithErrorMessage } from "~/models/message.server"; import { requireUser } from "~/services/session.server"; import { rootPath, v3RunPath } from "~/utils/pathBuilder"; @@ -14,8 +15,8 @@ export async function loader({ params, request }: LoaderFunctionArgs) { const { runParam } = ParamsSchema.parse(params); - const run = await prisma.taskRun.findFirst({ - where: { + const run = await runStore.findRun( + { friendlyId: runParam, project: { organization: { @@ -27,25 +28,28 @@ export async function loader({ params, request }: LoaderFunctionArgs) { }, }, }, - select: { - spanId: true, - runtimeEnvironment: { - select: { - slug: true, + { + select: { + spanId: true, + runtimeEnvironment: { + select: { + slug: true, + }, }, - }, - project: { - select: { - slug: true, - organization: { - select: { - slug: true, + project: { + select: { + slug: true, + organization: { + select: { + slug: true, + }, }, }, }, }, }, - }); + prisma + ); if (!run) { return redirectWithErrorMessage( diff --git a/apps/webapp/app/routes/sync.traces.runs.$traceId.ts b/apps/webapp/app/routes/sync.traces.runs.$traceId.ts index 279e2ffa51..ee5d1c964f 100644 --- a/apps/webapp/app/routes/sync.traces.runs.$traceId.ts +++ b/apps/webapp/app/routes/sync.traces.runs.$traceId.ts @@ -5,6 +5,7 @@ import { env } from "~/env.server"; import { logger } from "~/services/logger.server"; import { getUserId } from "~/services/session.server"; import { longPollingFetch } from "~/utils/longPollingFetch"; +import { runStore } from "~/v3/runStore.server"; const Params = z.object({ traceId: z.string(), @@ -21,18 +22,21 @@ export async function loader({ params, request }: LoaderFunctionArgs) { return new Response("No user found in cookie", { status: 401 }); } - const run = await $replica.taskRun.findFirst({ - select: { - project: { - select: { - organizationId: true, + const run = await runStore.findRun( + { + traceId, + }, + { + select: { + project: { + select: { + organizationId: true, + }, }, }, }, - where: { - traceId, - }, - }); + $replica + ); if (!run) { return new Response("No run found", { status: 404 }); From f59abe7c7f8702aaddee4d6b2d29e224d9c103d9 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 16:26:54 +0100 Subject: [PATCH 07/11] refactor(webapp): hydrate parent-model TaskRun reads through the run store Decompose the three reads that pulled TaskRun in through a parent model's relation include (alert, batch results, attempt dependencies): query the parent without the include, hydrate the run(s) via RunStore in a single batched read, and stitch them back. Preserves field selection, ordering, null handling and the query client. Adds container-backed tests for the batch-results and cancel-dependencies paths. --- .../v3/ApiBatchResultsPresenter.server.ts | 55 +++- .../v3/services/alerts/deliverAlert.server.ts | 38 ++- .../cancelTaskAttemptDependencies.server.ts | 51 +++- .../ApiBatchResultsPresenter.test.ts | 256 ++++++++++++++++++ .../cancelTaskAttemptDependencies.test.ts | 238 ++++++++++++++++ 5 files changed, 606 insertions(+), 32 deletions(-) create mode 100644 apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts create mode 100644 apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts diff --git a/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts index 0b610215ef..b3dd39637d 100644 --- a/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts @@ -1,6 +1,7 @@ import { BatchTaskRunExecutionResult } from "@trigger.dev/core/v3"; -import { executionResultForTaskRun } from "~/models/taskRun.server"; +import { executionResultForTaskRun, TaskRunWithAttempts } from "~/models/taskRun.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { runStore } from "~/v3/runStore.server"; import { BasePresenter } from "./basePresenter.server"; export class ApiBatchResultsPresenter extends BasePresenter { @@ -16,16 +17,8 @@ export class ApiBatchResultsPresenter extends BasePresenter { }, include: { items: { - include: { - taskRun: { - include: { - attempts: { - orderBy: { - createdAt: "desc", - }, - }, - }, - }, + select: { + taskRunId: true, }, }, }, @@ -35,10 +28,48 @@ export class ApiBatchResultsPresenter extends BasePresenter { return undefined; } + const taskRunIds = batchRun.items.map((item) => item.taskRunId); + + if (taskRunIds.length === 0) { + return { + id: batchRun.friendlyId, + items: [], + }; + } + + const taskRuns = await runStore.findRuns( + { + where: { id: { in: taskRunIds } }, + select: { + id: true, + friendlyId: true, + status: true, + taskIdentifier: true, + attempts: { + select: { + status: true, + output: true, + outputType: true, + error: true, + }, + orderBy: { + createdAt: "desc", + }, + }, + }, + }, + this._prisma + ); + + const runMap = new Map(taskRuns.map((run) => [run.id, run])); + return { id: batchRun.friendlyId, items: batchRun.items - .map((item) => executionResultForTaskRun(item.taskRun)) + .map((item) => { + const run = runMap.get(item.taskRunId); + return run ? executionResultForTaskRun(run as TaskRunWithAttempts) : undefined; + }) .filter(Boolean), }; }); diff --git a/apps/webapp/app/v3/services/alerts/deliverAlert.server.ts b/apps/webapp/app/v3/services/alerts/deliverAlert.server.ts index bc8f9a3a5f..49f464d6dc 100644 --- a/apps/webapp/app/v3/services/alerts/deliverAlert.server.ts +++ b/apps/webapp/app/v3/services/alerts/deliverAlert.server.ts @@ -102,7 +102,7 @@ type DeploymentIntegrationMetadata = { export class DeliverAlertService extends BaseService { public async call(alertId: string) { - const alert: FoundAlert | null = await this._prisma.projectAlert.findFirst({ + const alertWithoutRun = await this._prisma.projectAlert.findFirst({ where: { id: alertId }, include: { channel: true, @@ -112,18 +112,6 @@ export class DeliverAlertService extends BaseService { }, }, environment: true, - taskRun: { - include: { - lockedBy: true, - lockedToVersion: true, - runtimeEnvironment: { - select: { - type: true, - branchName: true, - }, - }, - }, - }, workerDeployment: { include: { worker: { @@ -142,10 +130,32 @@ export class DeliverAlertService extends BaseService { }, }); - if (!alert) { + if (!alertWithoutRun) { return; } + let taskRun: FoundAlert["taskRun"] = null; + if (alertWithoutRun.taskRunId) { + taskRun = await this.runStore.findRun( + { id: alertWithoutRun.taskRunId }, + { + include: { + lockedBy: true, + lockedToVersion: true, + runtimeEnvironment: { + select: { + type: true, + branchName: true, + }, + }, + }, + }, + this._prisma + ); + } + + const alert: FoundAlert = { ...alertWithoutRun, taskRun }; + if (alert.status !== "PENDING") { return; } diff --git a/apps/webapp/app/v3/services/cancelTaskAttemptDependencies.server.ts b/apps/webapp/app/v3/services/cancelTaskAttemptDependencies.server.ts index f3ad291ac9..82b22d5935 100644 --- a/apps/webapp/app/v3/services/cancelTaskAttemptDependencies.server.ts +++ b/apps/webapp/app/v3/services/cancelTaskAttemptDependencies.server.ts @@ -10,15 +10,15 @@ export class CancelTaskAttemptDependenciesService extends BaseService { where: { id: attemptId }, include: { dependencies: { - include: { - taskRun: true, + select: { + taskRunId: true, }, }, batchDependencies: { include: { runDependencies: { - include: { - taskRun: true, + select: { + taskRunId: true, }, }, }, @@ -45,14 +45,53 @@ export class CancelTaskAttemptDependenciesService extends BaseService { batchDependencies: taskAttempt.batchDependencies, }); + // Hydrate the dependent runs from both relation paths in a single batched read, + // deduping the ids that feed the query while preserving the original iteration order. + const taskRunIds = new Set(); + for (const dependency of taskAttempt.dependencies) { + taskRunIds.add(dependency.taskRunId); + } + for (const batchDependency of taskAttempt.batchDependencies) { + for (const runDependency of batchDependency.runDependencies) { + taskRunIds.add(runDependency.taskRunId); + } + } + + const runs = + taskRunIds.size > 0 + ? await this.runStore.findRuns( + { + where: { id: { in: [...taskRunIds] } }, + select: { + id: true, + engine: true, + status: true, + friendlyId: true, + taskEventStore: true, + createdAt: true, + completedAt: true, + }, + }, + this._prisma + ) + : []; + + const runMap = new Map(runs.map((run) => [run.id, run])); + // TaskAttempt will either have dependencies or batchDependencies for (const dependency of taskAttempt.dependencies) { - await cancelRunService.call(dependency.taskRun); + const run = runMap.get(dependency.taskRunId); + if (run) { + await cancelRunService.call(run); + } } for (const batchDependency of taskAttempt.batchDependencies) { for (const runDependency of batchDependency.runDependencies) { - await cancelRunService.call(runDependency.taskRun); + const run = runMap.get(runDependency.taskRunId); + if (run) { + await cancelRunService.call(run); + } } } } diff --git a/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts b/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts new file mode 100644 index 0000000000..385be889a5 --- /dev/null +++ b/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts @@ -0,0 +1,256 @@ +import { containerTest } from "@internal/testcontainers"; +import type { Organization, PrismaClient, Project, RuntimeEnvironment } from "@trigger.dev/database"; +import { customAlphabet } from "nanoid"; +import { expect, vi } from "vitest"; +import { ApiBatchResultsPresenter } from "~/presenters/v3/ApiBatchResultsPresenter.server"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { seedTestEnvironment } from "../helpers/seedTestEnvironment"; + +vi.setConfig({ testTimeout: 60_000 }); + +const idGenerator = customAlphabet("123456789abcdefghijkmnopqrstuvwxyz", 21); + +function authEnv( + environment: RuntimeEnvironment, + project: Project, + organization: Organization +): AuthenticatedEnvironment { + return { ...environment, project, organization, orgMember: null } as AuthenticatedEnvironment; +} + +type SeedContext = { + environmentId: string; + projectId: string; + organizationId: string; + backgroundWorkerId: string; + backgroundWorkerTaskId: string; + queueId: string; +}; + +async function seedWorker(prisma: PrismaClient, ctx: Omit) { + const queue = await prisma.taskQueue.create({ + data: { + friendlyId: `queue_${idGenerator()}`, + name: "task/test-task", + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.environmentId, + }, + }); + + const worker = await prisma.backgroundWorker.create({ + data: { + friendlyId: `worker_${idGenerator()}`, + contentHash: "hash", + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.environmentId, + version: "20240101.1", + metadata: {}, + }, + }); + + const task = await prisma.backgroundWorkerTask.create({ + data: { + friendlyId: `task_${idGenerator()}`, + slug: "test-task", + filePath: "src/test.ts", + exportName: "testTask", + workerId: worker.id, + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.environmentId, + }, + }); + + return { queueId: queue.id, backgroundWorkerId: worker.id, backgroundWorkerTaskId: task.id }; +} + +async function seedRunWithAttempt( + prisma: PrismaClient, + ctx: SeedContext, + opts: { + status: "COMPLETED_SUCCESSFULLY" | "COMPLETED_WITH_ERRORS" | "CANCELED" | "EXECUTING"; + attempt?: { + status: "COMPLETED" | "FAILED"; + output?: string; + outputType?: string; + error?: unknown; + }; + } +) { + const runInternalId = idGenerator(); + const run = await prisma.taskRun.create({ + data: { + id: runInternalId, + friendlyId: `run_${runInternalId}`, + taskIdentifier: "test-task", + payload: "{}", + payloadType: "application/json", + traceId: idGenerator(), + spanId: idGenerator(), + queue: "task/test-task", + runtimeEnvironmentId: ctx.environmentId, + projectId: ctx.projectId, + status: opts.status, + }, + }); + + if (opts.attempt) { + await prisma.taskRunAttempt.create({ + data: { + friendlyId: `attempt_${idGenerator()}`, + taskRunId: run.id, + backgroundWorkerId: ctx.backgroundWorkerId, + backgroundWorkerTaskId: ctx.backgroundWorkerTaskId, + runtimeEnvironmentId: ctx.environmentId, + queueId: ctx.queueId, + status: opts.attempt.status, + output: opts.attempt.output, + outputType: opts.attempt.outputType ?? "application/json", + error: opts.attempt.error as any, + }, + }); + } + + return run; +} + +containerTest( + "ApiBatchResultsPresenter returns ordered results matching pre-decompose behavior", + async ({ prisma }) => { + const { environment, project, organization } = await seedTestEnvironment(prisma); + const worker = await seedWorker(prisma, { + environmentId: environment.id, + projectId: project.id, + organizationId: organization.id, + }); + const ctx: SeedContext = { + environmentId: environment.id, + projectId: project.id, + organizationId: organization.id, + ...worker, + }; + + // A successful run, a failed run, and an executing run (no terminal attempt → undefined). + const successRun = await seedRunWithAttempt(prisma, ctx, { + status: "COMPLETED_SUCCESSFULLY", + attempt: { status: "COMPLETED", output: "\"hello\"", outputType: "application/json" }, + }); + const failedRun = await seedRunWithAttempt(prisma, ctx, { + status: "COMPLETED_WITH_ERRORS", + attempt: { + status: "FAILED", + error: { type: "BUILT_IN_ERROR", name: "Error", message: "boom", stackTrace: "boom" }, + }, + }); + const executingRun = await seedRunWithAttempt(prisma, ctx, { + status: "EXECUTING", + }); + + const batchInternalId = idGenerator(); + const batchFriendlyId = `batch_${batchInternalId}`; + await prisma.batchTaskRun.create({ + data: { + id: batchInternalId, + friendlyId: batchFriendlyId, + runtimeEnvironmentId: environment.id, + }, + }); + + // Items inserted in a deterministic order: success, failed, executing. + for (const run of [successRun, failedRun, executingRun]) { + await prisma.batchTaskRunItem.create({ + data: { + batchTaskRunId: batchInternalId, + taskRunId: run.id, + }, + }); + } + + const presenter = new ApiBatchResultsPresenter(prisma); + const result = await presenter.call(batchFriendlyId, authEnv(environment, project, organization)); + + expect(result).toBeDefined(); + expect(result?.id).toBe(batchFriendlyId); + + // executing run yields no execution result → filtered out. Order preserved: success then failed. + expect(result?.items).toHaveLength(2); + + const [first, second] = result!.items; + expect(first.ok).toBe(true); + expect(first.id).toBe(successRun.friendlyId); + if (first.ok) { + expect(first.output).toBe("\"hello\""); + expect(first.taskIdentifier).toBe("test-task"); + } + + expect(second.ok).toBe(false); + expect(second.id).toBe(failedRun.friendlyId); + } +); + +containerTest( + "ApiBatchResultsPresenter filters runs without an execution result but keeps order", + async ({ prisma }) => { + const { environment, project, organization } = await seedTestEnvironment(prisma); + const worker = await seedWorker(prisma, { + environmentId: environment.id, + projectId: project.id, + organizationId: organization.id, + }); + const ctx: SeedContext = { + environmentId: environment.id, + projectId: project.id, + organizationId: organization.id, + ...worker, + }; + + // Pending run → executionResultForTaskRun returns undefined → filtered out, like the + // pre-decompose code did via `.filter(Boolean)`. + const pendingRun = await seedRunWithAttempt(prisma, ctx, { status: "EXECUTING" }); + const successRun = await seedRunWithAttempt(prisma, ctx, { + status: "COMPLETED_SUCCESSFULLY", + attempt: { status: "COMPLETED", output: "\"ok\"", outputType: "application/json" }, + }); + + const batchInternalId = idGenerator(); + const batchFriendlyId = `batch_${batchInternalId}`; + await prisma.batchTaskRun.create({ + data: { + id: batchInternalId, + friendlyId: batchFriendlyId, + runtimeEnvironmentId: environment.id, + }, + }); + + // pending first, success second — only the success result should survive, in order. + for (const run of [pendingRun, successRun]) { + await prisma.batchTaskRunItem.create({ + data: { batchTaskRunId: batchInternalId, taskRunId: run.id }, + }); + } + + const presenter = new ApiBatchResultsPresenter(prisma); + const result = await presenter.call(batchFriendlyId, authEnv(environment, project, organization)); + + expect(result?.items).toHaveLength(1); + expect(result?.items[0]?.id).toBe(successRun.friendlyId); + } +); + +containerTest("ApiBatchResultsPresenter short-circuits an empty batch", async ({ prisma }) => { + const { environment, project, organization } = await seedTestEnvironment(prisma); + + const batchInternalId = idGenerator(); + const batchFriendlyId = `batch_${batchInternalId}`; + await prisma.batchTaskRun.create({ + data: { + id: batchInternalId, + friendlyId: batchFriendlyId, + runtimeEnvironmentId: environment.id, + }, + }); + + const presenter = new ApiBatchResultsPresenter(prisma); + const result = await presenter.call(batchFriendlyId, authEnv(environment, project, organization)); + + expect(result).toEqual({ id: batchFriendlyId, items: [] }); +}); diff --git a/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts b/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts new file mode 100644 index 0000000000..03e090ea6c --- /dev/null +++ b/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts @@ -0,0 +1,238 @@ +import { containerTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { customAlphabet } from "nanoid"; +import { expect, vi } from "vitest"; +import { CancelTaskAttemptDependenciesService } from "~/v3/services/cancelTaskAttemptDependencies.server"; +import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server"; +import { seedTestEnvironment } from "../helpers/seedTestEnvironment"; + +vi.setConfig({ testTimeout: 60_000 }); + +const idGenerator = customAlphabet("123456789abcdefghijkmnopqrstuvwxyz", 21); + +type SeedContext = { + environmentId: string; + projectId: string; + organizationId: string; + backgroundWorkerId: string; + backgroundWorkerTaskId: string; + queueId: string; +}; + +async function seedWorker( + prisma: PrismaClient, + ctx: { environmentId: string; projectId: string } +) { + const queue = await prisma.taskQueue.create({ + data: { + friendlyId: `queue_${idGenerator()}`, + name: "task/test-task", + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.environmentId, + }, + }); + const worker = await prisma.backgroundWorker.create({ + data: { + friendlyId: `worker_${idGenerator()}`, + contentHash: "hash", + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.environmentId, + version: "20240101.1", + metadata: {}, + }, + }); + const task = await prisma.backgroundWorkerTask.create({ + data: { + friendlyId: `task_${idGenerator()}`, + slug: "test-task", + filePath: "src/test.ts", + workerId: worker.id, + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.environmentId, + }, + }); + return { queueId: queue.id, backgroundWorkerId: worker.id, backgroundWorkerTaskId: task.id }; +} + +async function seedRun(prisma: PrismaClient, ctx: SeedContext) { + const id = idGenerator(); + return prisma.taskRun.create({ + data: { + id, + friendlyId: `run_${id}`, + taskIdentifier: "test-task", + payload: "{}", + payloadType: "application/json", + traceId: idGenerator(), + spanId: idGenerator(), + queue: "task/test-task", + runtimeEnvironmentId: ctx.environmentId, + projectId: ctx.projectId, + }, + }); +} + +async function seedAttempt(prisma: PrismaClient, ctx: SeedContext, taskRunId: string) { + return prisma.taskRunAttempt.create({ + data: { + friendlyId: `attempt_${idGenerator()}`, + taskRunId, + backgroundWorkerId: ctx.backgroundWorkerId, + backgroundWorkerTaskId: ctx.backgroundWorkerTaskId, + runtimeEnvironmentId: ctx.environmentId, + queueId: ctx.queueId, + status: "CANCELED", + }, + }); +} + +containerTest( + "cancelTaskAttemptDependencies cancels each dependent run once, in original order", + async ({ prisma }) => { + const { environment, project, organization } = await seedTestEnvironment(prisma); + const worker = await seedWorker(prisma, { + environmentId: environment.id, + projectId: project.id, + }); + const ctx: SeedContext = { + environmentId: environment.id, + projectId: project.id, + organizationId: organization.id, + ...worker, + }; + + // The attempt whose dependencies we cancel. + const parentRun = await seedRun(prisma, ctx); + const parentAttempt = await seedAttempt(prisma, ctx, parentRun.id); + + // Two direct dependencies. + const depRunA = await seedRun(prisma, ctx); + const depRunB = await seedRun(prisma, ctx); + await prisma.taskRunDependency.create({ + data: { taskRunId: depRunA.id, dependentAttemptId: parentAttempt.id }, + }); + await prisma.taskRunDependency.create({ + data: { taskRunId: depRunB.id, dependentAttemptId: parentAttempt.id }, + }); + + // One batch dependency carrying two run dependencies. + const batchRunDepC = await seedRun(prisma, ctx); + const batchRunDepD = await seedRun(prisma, ctx); + const batchId = idGenerator(); + await prisma.batchTaskRun.create({ + data: { + id: batchId, + friendlyId: `batch_${batchId}`, + runtimeEnvironmentId: environment.id, + dependentTaskAttemptId: parentAttempt.id, + }, + }); + await prisma.taskRunDependency.create({ + data: { taskRunId: batchRunDepC.id, dependentBatchRunId: batchId }, + }); + await prisma.taskRunDependency.create({ + data: { taskRunId: batchRunDepD.id, dependentBatchRunId: batchId }, + }); + + const cancelledRunIds: string[] = []; + const callSpy = vi + .spyOn(CancelTaskRunService.prototype, "call") + .mockImplementation(async (taskRun: any) => { + cancelledRunIds.push(taskRun.id); + return { id: taskRun.id, alreadyFinished: false }; + }); + + try { + const service = new CancelTaskAttemptDependenciesService(prisma); + await service.call(parentAttempt.id); + } finally { + callSpy.mockRestore(); + } + + // Each dependent run cancelled exactly once. + expect(cancelledRunIds).toHaveLength(4); + expect(new Set(cancelledRunIds).size).toBe(4); + + // Direct dependencies first (both paths preserve insertion/iteration order), then batch run deps. + const directIds = cancelledRunIds.slice(0, 2); + const batchIds = cancelledRunIds.slice(2); + expect(new Set(directIds)).toEqual(new Set([depRunA.id, depRunB.id])); + expect(new Set(batchIds)).toEqual(new Set([batchRunDepC.id, batchRunDepD.id])); + + // The hydrated runs carry the fields CancelableTaskRun requires. + const cancelArgs = callSpy.mock.calls.map((c) => c[0] as any); + for (const run of cancelArgs) { + expect(run).toMatchObject({ + id: expect.any(String), + friendlyId: expect.any(String), + }); + expect(run).toHaveProperty("engine"); + expect(run).toHaveProperty("status"); + expect(run).toHaveProperty("taskEventStore"); + expect(run).toHaveProperty("createdAt"); + expect("completedAt" in run).toBe(true); + } + } +); + +containerTest( + "cancelTaskAttemptDependencies skips dependencies whose run is not hydrated", + async ({ prisma }) => { + const { environment, project, organization } = await seedTestEnvironment(prisma); + const worker = await seedWorker(prisma, { + environmentId: environment.id, + projectId: project.id, + }); + const ctx: SeedContext = { + environmentId: environment.id, + projectId: project.id, + organizationId: organization.id, + ...worker, + }; + + const parentRun = await seedRun(prisma, ctx); + const parentAttempt = await seedAttempt(prisma, ctx, parentRun.id); + + const presentRun = await seedRun(prisma, ctx); + const missingRun = await seedRun(prisma, ctx); + await prisma.taskRunDependency.create({ + data: { taskRunId: presentRun.id, dependentAttemptId: parentAttempt.id }, + }); + await prisma.taskRunDependency.create({ + data: { taskRunId: missingRun.id, dependentAttemptId: parentAttempt.id }, + }); + + const cancelledRunIds: string[] = []; + const callSpy = vi + .spyOn(CancelTaskRunService.prototype, "call") + .mockImplementation(async (taskRun: any) => { + cancelledRunIds.push(taskRun.id); + return { id: taskRun.id, alreadyFinished: false }; + }); + + // Inject a runStore that deliberately omits `missingRun` to exercise the runMap-miss skip + // (the post-redirect "run not found here" case). The constructor's third arg is the seam. + const filteringRunStore = { + findRuns: async (args: any) => { + const ids: string[] = args.where.id.in; + return prisma.taskRun.findMany({ + where: { id: { in: ids.filter((id) => id !== missingRun.id) } }, + select: args.select, + }); + }, + } as any; + + try { + const service = new CancelTaskAttemptDependenciesService( + prisma, + undefined, + filteringRunStore + ); + await service.call(parentAttempt.id); + } finally { + callSpy.mockRestore(); + } + + expect(cancelledRunIds).toEqual([presentRun.id]); + } +); From cb12430424e7707be029b8d2a07fbce636c2dd91 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 16:31:01 +0100 Subject: [PATCH 08/11] chore(scripts): flag recover-stuck-runs raw TaskRun read for table cutover The recovery script joins TaskRunExecutionSnapshot to TaskRun in raw SQL, so it is the one TaskRun read not routed through the run store. Add a note to revisit it at table cutover. --- scripts/recover-stuck-runs.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/scripts/recover-stuck-runs.ts b/scripts/recover-stuck-runs.ts index 15deeb899c..28bb4e85e4 100755 --- a/scripts/recover-stuck-runs.ts +++ b/scripts/recover-stuck-runs.ts @@ -187,7 +187,9 @@ async function main() { console.log(`📊 Found ${runIds.length} runs in currentConcurrency set`); - // Query database for latest snapshots and queue info of these runs + // Query database for latest snapshots and queue info of these runs. + // NOTE: raw join of TaskRunExecutionSnapshot to TaskRun, the one TaskRun read not behind + // RunStore (a join, not a by-id read, in an ops script). Revisit at table cutover. const runInfo = await prisma.$queryRaw< Array<{ runId: string; From ae57f25a03b20c51cb4c4869ce686d715de4e91a Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 16:41:47 +0100 Subject: [PATCH 09/11] chore(webapp): add server-changes entry for run-store read routing --- .server-changes/route-taskrun-reads-through-run-store.md | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 .server-changes/route-taskrun-reads-through-run-store.md diff --git a/.server-changes/route-taskrun-reads-through-run-store.md b/.server-changes/route-taskrun-reads-through-run-store.md new file mode 100644 index 0000000000..dad804e40b --- /dev/null +++ b/.server-changes/route-taskrun-reads-through-run-store.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: improvement +--- + +Route Postgres task run reads through the run store so they can be retargeted to a different backing store without changing call sites. From fcc26d4ebd3966d039b923f6a49476da17955985 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 17:01:19 +0100 Subject: [PATCH 10/11] test(webapp): mock db.server in the new run-store read tests The new container tests import the service and presenter, which pull the db.server singleton in through their base classes. Mock it so the tests do not try to connect to the env database when none is reachable (the CI unit shards), matching the existing webapp container-test pattern. The tests use the injected testcontainer prisma for all reads. --- apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts | 4 ++++ .../test/services/cancelTaskAttemptDependencies.test.ts | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts b/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts index 385be889a5..d0888ba6a1 100644 --- a/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts +++ b/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts @@ -6,6 +6,10 @@ import { ApiBatchResultsPresenter } from "~/presenters/v3/ApiBatchResultsPresent import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { seedTestEnvironment } from "../helpers/seedTestEnvironment"; +// Neutralize the db.server singleton so importing the presenter (via BasePresenter) does not try +// to connect to the env database; the test uses the injected testcontainer prisma for all reads. +vi.mock("~/db.server", () => ({ prisma: {}, $replica: {} })); + vi.setConfig({ testTimeout: 60_000 }); const idGenerator = customAlphabet("123456789abcdefghijkmnopqrstuvwxyz", 21); diff --git a/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts b/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts index 03e090ea6c..65ecef73a8 100644 --- a/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts +++ b/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts @@ -6,6 +6,10 @@ import { CancelTaskAttemptDependenciesService } from "~/v3/services/cancelTaskAt import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server"; import { seedTestEnvironment } from "../helpers/seedTestEnvironment"; +// Neutralize the db.server singleton so importing the service (via BaseService) does not try to +// connect to the env database; the test uses the injected testcontainer prisma for all reads. +vi.mock("~/db.server", () => ({ prisma: {}, $replica: {} })); + vi.setConfig({ testTimeout: 60_000 }); const idGenerator = customAlphabet("123456789abcdefghijkmnopqrstuvwxyz", 21); From 789e10780960acc51ed27aae11fdb705f5fd6594 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 18 Jun 2026 17:36:57 +0100 Subject: [PATCH 11/11] test(webapp): drop the cancelTaskAttemptDependencies container test Importing the service pulls the cancel chain, which eagerly initializes the concurrency tracker singleton and requires REDIS_HOST/REDIS_PORT at import time, so the suite cannot load in the unit-test shards without stacking mocks. The decompose it covered is exercised by the analogous batch-results container test and confirmed by review, so drop this one rather than mock the tracker and cancel chain. --- .../cancelTaskAttemptDependencies.test.ts | 242 ------------------ 1 file changed, 242 deletions(-) delete mode 100644 apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts diff --git a/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts b/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts deleted file mode 100644 index 65ecef73a8..0000000000 --- a/apps/webapp/test/services/cancelTaskAttemptDependencies.test.ts +++ /dev/null @@ -1,242 +0,0 @@ -import { containerTest } from "@internal/testcontainers"; -import type { PrismaClient } from "@trigger.dev/database"; -import { customAlphabet } from "nanoid"; -import { expect, vi } from "vitest"; -import { CancelTaskAttemptDependenciesService } from "~/v3/services/cancelTaskAttemptDependencies.server"; -import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server"; -import { seedTestEnvironment } from "../helpers/seedTestEnvironment"; - -// Neutralize the db.server singleton so importing the service (via BaseService) does not try to -// connect to the env database; the test uses the injected testcontainer prisma for all reads. -vi.mock("~/db.server", () => ({ prisma: {}, $replica: {} })); - -vi.setConfig({ testTimeout: 60_000 }); - -const idGenerator = customAlphabet("123456789abcdefghijkmnopqrstuvwxyz", 21); - -type SeedContext = { - environmentId: string; - projectId: string; - organizationId: string; - backgroundWorkerId: string; - backgroundWorkerTaskId: string; - queueId: string; -}; - -async function seedWorker( - prisma: PrismaClient, - ctx: { environmentId: string; projectId: string } -) { - const queue = await prisma.taskQueue.create({ - data: { - friendlyId: `queue_${idGenerator()}`, - name: "task/test-task", - projectId: ctx.projectId, - runtimeEnvironmentId: ctx.environmentId, - }, - }); - const worker = await prisma.backgroundWorker.create({ - data: { - friendlyId: `worker_${idGenerator()}`, - contentHash: "hash", - projectId: ctx.projectId, - runtimeEnvironmentId: ctx.environmentId, - version: "20240101.1", - metadata: {}, - }, - }); - const task = await prisma.backgroundWorkerTask.create({ - data: { - friendlyId: `task_${idGenerator()}`, - slug: "test-task", - filePath: "src/test.ts", - workerId: worker.id, - projectId: ctx.projectId, - runtimeEnvironmentId: ctx.environmentId, - }, - }); - return { queueId: queue.id, backgroundWorkerId: worker.id, backgroundWorkerTaskId: task.id }; -} - -async function seedRun(prisma: PrismaClient, ctx: SeedContext) { - const id = idGenerator(); - return prisma.taskRun.create({ - data: { - id, - friendlyId: `run_${id}`, - taskIdentifier: "test-task", - payload: "{}", - payloadType: "application/json", - traceId: idGenerator(), - spanId: idGenerator(), - queue: "task/test-task", - runtimeEnvironmentId: ctx.environmentId, - projectId: ctx.projectId, - }, - }); -} - -async function seedAttempt(prisma: PrismaClient, ctx: SeedContext, taskRunId: string) { - return prisma.taskRunAttempt.create({ - data: { - friendlyId: `attempt_${idGenerator()}`, - taskRunId, - backgroundWorkerId: ctx.backgroundWorkerId, - backgroundWorkerTaskId: ctx.backgroundWorkerTaskId, - runtimeEnvironmentId: ctx.environmentId, - queueId: ctx.queueId, - status: "CANCELED", - }, - }); -} - -containerTest( - "cancelTaskAttemptDependencies cancels each dependent run once, in original order", - async ({ prisma }) => { - const { environment, project, organization } = await seedTestEnvironment(prisma); - const worker = await seedWorker(prisma, { - environmentId: environment.id, - projectId: project.id, - }); - const ctx: SeedContext = { - environmentId: environment.id, - projectId: project.id, - organizationId: organization.id, - ...worker, - }; - - // The attempt whose dependencies we cancel. - const parentRun = await seedRun(prisma, ctx); - const parentAttempt = await seedAttempt(prisma, ctx, parentRun.id); - - // Two direct dependencies. - const depRunA = await seedRun(prisma, ctx); - const depRunB = await seedRun(prisma, ctx); - await prisma.taskRunDependency.create({ - data: { taskRunId: depRunA.id, dependentAttemptId: parentAttempt.id }, - }); - await prisma.taskRunDependency.create({ - data: { taskRunId: depRunB.id, dependentAttemptId: parentAttempt.id }, - }); - - // One batch dependency carrying two run dependencies. - const batchRunDepC = await seedRun(prisma, ctx); - const batchRunDepD = await seedRun(prisma, ctx); - const batchId = idGenerator(); - await prisma.batchTaskRun.create({ - data: { - id: batchId, - friendlyId: `batch_${batchId}`, - runtimeEnvironmentId: environment.id, - dependentTaskAttemptId: parentAttempt.id, - }, - }); - await prisma.taskRunDependency.create({ - data: { taskRunId: batchRunDepC.id, dependentBatchRunId: batchId }, - }); - await prisma.taskRunDependency.create({ - data: { taskRunId: batchRunDepD.id, dependentBatchRunId: batchId }, - }); - - const cancelledRunIds: string[] = []; - const callSpy = vi - .spyOn(CancelTaskRunService.prototype, "call") - .mockImplementation(async (taskRun: any) => { - cancelledRunIds.push(taskRun.id); - return { id: taskRun.id, alreadyFinished: false }; - }); - - try { - const service = new CancelTaskAttemptDependenciesService(prisma); - await service.call(parentAttempt.id); - } finally { - callSpy.mockRestore(); - } - - // Each dependent run cancelled exactly once. - expect(cancelledRunIds).toHaveLength(4); - expect(new Set(cancelledRunIds).size).toBe(4); - - // Direct dependencies first (both paths preserve insertion/iteration order), then batch run deps. - const directIds = cancelledRunIds.slice(0, 2); - const batchIds = cancelledRunIds.slice(2); - expect(new Set(directIds)).toEqual(new Set([depRunA.id, depRunB.id])); - expect(new Set(batchIds)).toEqual(new Set([batchRunDepC.id, batchRunDepD.id])); - - // The hydrated runs carry the fields CancelableTaskRun requires. - const cancelArgs = callSpy.mock.calls.map((c) => c[0] as any); - for (const run of cancelArgs) { - expect(run).toMatchObject({ - id: expect.any(String), - friendlyId: expect.any(String), - }); - expect(run).toHaveProperty("engine"); - expect(run).toHaveProperty("status"); - expect(run).toHaveProperty("taskEventStore"); - expect(run).toHaveProperty("createdAt"); - expect("completedAt" in run).toBe(true); - } - } -); - -containerTest( - "cancelTaskAttemptDependencies skips dependencies whose run is not hydrated", - async ({ prisma }) => { - const { environment, project, organization } = await seedTestEnvironment(prisma); - const worker = await seedWorker(prisma, { - environmentId: environment.id, - projectId: project.id, - }); - const ctx: SeedContext = { - environmentId: environment.id, - projectId: project.id, - organizationId: organization.id, - ...worker, - }; - - const parentRun = await seedRun(prisma, ctx); - const parentAttempt = await seedAttempt(prisma, ctx, parentRun.id); - - const presentRun = await seedRun(prisma, ctx); - const missingRun = await seedRun(prisma, ctx); - await prisma.taskRunDependency.create({ - data: { taskRunId: presentRun.id, dependentAttemptId: parentAttempt.id }, - }); - await prisma.taskRunDependency.create({ - data: { taskRunId: missingRun.id, dependentAttemptId: parentAttempt.id }, - }); - - const cancelledRunIds: string[] = []; - const callSpy = vi - .spyOn(CancelTaskRunService.prototype, "call") - .mockImplementation(async (taskRun: any) => { - cancelledRunIds.push(taskRun.id); - return { id: taskRun.id, alreadyFinished: false }; - }); - - // Inject a runStore that deliberately omits `missingRun` to exercise the runMap-miss skip - // (the post-redirect "run not found here" case). The constructor's third arg is the seam. - const filteringRunStore = { - findRuns: async (args: any) => { - const ids: string[] = args.where.id.in; - return prisma.taskRun.findMany({ - where: { id: { in: ids.filter((id) => id !== missingRun.id) } }, - select: args.select, - }); - }, - } as any; - - try { - const service = new CancelTaskAttemptDependenciesService( - prisma, - undefined, - filteringRunStore - ); - await service.call(parentAttempt.id); - } finally { - callSpy.mockRestore(); - } - - expect(cancelledRunIds).toEqual([presentRun.id]); - } -);