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 00000000000..dad804e40ba --- /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. diff --git a/apps/webapp/app/models/runtimeEnvironment.server.ts b/apps/webapp/app/models/runtimeEnvironment.server.ts index be05adaa8a7..9135872417c 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/presenters/v3/ApiBatchResultsPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts index 0b610215ef9..b3dd39637da 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/presenters/v3/ApiRetrieveRunPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts index fec8dabdb0e..68e3643f9e9 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 c11a04a1581..7e0540674e8 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 3594aa71cea..2e587e8c4a7 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 1ff68e9b96f..c4c3ac88c48 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 e0e88e4dd02..ab777d0b8e9 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 0586ab8eced..bff1bda0177 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 c63f9e39a2a..36ef46d4b4e 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 98ee75cda39..49d8f303560 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 a9381ab60d2..0ebf5054bb1 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; diff --git a/apps/webapp/app/routes/@.runs.$runParam.ts b/apps/webapp/app/routes/@.runs.$runParam.ts index a709191271e..cd1e1eade18 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 c4d17ba875d..af041353ada 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 0924bf3fc91..091312a13b8 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 39c30894416..cd88ef38281 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 3f22929aca9..7ec10835c78 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 c38206473cb..061199f33e9 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 04ae398194f..f1aa4d58967 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 130f6ff163a..4b238869d3a 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 7c5718aeae3..cc1a6d4f9fc 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 44f8c7ef69f..ec8c171fc20 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 1f8a42af08c..eb9e5d974e4 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 0cf92a53b70..01428301432 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 da4bab693ba..afc481a571a 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 a3f35013b78..4c057046479 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 bab59fd0637..d9f6ca9a6d0 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 60505460bd6..9254a74e834 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 199244b1da8..8d7f6b84345 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 63a89d7e0aa..d5d4ab0f2f6 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 fe267d1f9fa..2a6cb34c913 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 46118c1d894..f2268989bdd 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 d6470794a73..81784f9bc3a 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 11074840a38..7cb813a6dec 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 cdee9567b79..c71ad48d121 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 a404e6a76ae..78fe332b8af 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 f4d34907042..be4fdba7fe8 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 0fab90e1457..da77d2cc692 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 614b668f910..06233f88c70 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 66135347253..3a0dfca568e 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 8d0af728df8..cec6c3c4e98 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 c9480299cc0..1ecc7819c23 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 60233d6d38f..24e7a73374f 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 7cda5ac7824..7662a88b4d2 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 c5e467533a3..38e17531f6f 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 fa6ee29f3db..ca92615bb83 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 d7acf18e517..7b37b1bcc00 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 03bfdaccc65..0719a8e6a19 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 b472d7ae8f4..5e0c2b21d6b 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 279e2ffa517..ee5d1c964f4 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 }); diff --git a/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts b/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts index 2bdf95eb9a6..02d0ec957f2 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 a8a7cbf0f3b..031411844b4 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 78455f9b686..89a938da8bf 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 2cc057f10f2..2af44d747bd 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 012c28c08fc..3f29f3faa47 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 e8509d73de4..98ce4dc35ff 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 1ad5174d1c6..b227f382c7b 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 55b969e7e55..a523111b5b2 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 8dbb5007c20..35333f9639b 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 7fc824f3d39..50e041ee64b 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 88e792b4a40..d32652a0b3b 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 4be392535c3..c59be0f3f57 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 f4b3c92ea66..c2f58662491 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 e6deff5dbee..91c877c8133 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 b3db81368b9..dac12768a75 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 082974af388..e2285a4fecc 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/deliverAlert.server.ts b/apps/webapp/app/v3/services/alerts/deliverAlert.server.ts index bc8f9a3a5f2..49f464d6dc8 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/alerts/performTaskRunAlerts.server.ts b/apps/webapp/app/v3/services/alerts/performTaskRunAlerts.server.ts index 9c055346232..31912c39fd0 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 33036871599..c001932baad 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 babdb02ca6a..76d550c7008 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 f779d81641f..c1562275e58 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/cancelTaskAttemptDependencies.server.ts b/apps/webapp/app/v3/services/cancelTaskAttemptDependencies.server.ts index f3ad291ac9b..82b22d5935d 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/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index c4076648819..22a9047c3fe 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 cd55b9ec0f9..bff4b8d65b1 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 63a8b6bb9aa..59c37947178 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 8be2b9557cc..dbc4c576b75 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 0b6149dfae6..79cb4fb0976 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 fb519b43151..a77727c9242 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 0409b6ed956..12ccddbf2e6 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 ab51df5de60..b770ceef177 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 b4ab5235761..6ed83c10807 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 2363d241c0c..dcf2488f273 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 8359cc4a4aa..c472bff53ee 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/presenters/ApiBatchResultsPresenter.test.ts b/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts new file mode 100644 index 00000000000..d0888ba6a18 --- /dev/null +++ b/apps/webapp/test/presenters/ApiBatchResultsPresenter.test.ts @@ -0,0 +1,260 @@ +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"; + +// 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); + +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/realtime/runReaderProjection.test.ts b/apps/webapp/test/realtime/runReaderProjection.test.ts index 07aebf92589..ad6616f5464 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 () => { diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 8d1f4c9c1f8..a6a20b5b9fd 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 6099d5b649b..a64dfb796e1 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 9933a715162..a3d44507a46 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 5b9d851d0f2..bf4b3e68bb4 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 cff29a75a4f..a77e60d05e7 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 26ea7866a67..8791dc1bd12 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 59d72c4c461..741ad8a14f6 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 1aa1738f3b0..977c94a8e83 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 ebd1cbdd80b..faffa2c59e5 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 8b8d4f82fcf..29eba297be5 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}`); diff --git a/internal-packages/run-store/src/NoopRunStore.ts b/internal-packages/run-store/src/NoopRunStore.ts index 3b4fb0a36fe..e27080c9af6 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 f2fb2969e6c..47876b70c8d 100644 --- a/internal-packages/run-store/src/PostgresRunStore.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.test.ts @@ -1528,3 +1528,247 @@ 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"); + }); + + 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 925a39425b6..fcc53c00266 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -617,4 +617,144 @@ 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>; + findRun( + where: Prisma.TaskRunWhereInput, + client?: PrismaClientOrTransaction + ): Promise; + async findRun( + where: Prisma.TaskRunWhereInput, + argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | PrismaClientOrTransaction, + client?: PrismaClientOrTransaction + ): Promise { + const { args, prisma } = this.#resolveReadArgs(argsOrClient, client); + + 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>; + findRunOrThrow( + where: Prisma.TaskRunWhereInput, + client?: PrismaClientOrTransaction + ): Promise; + async findRunOrThrow( + where: Prisma.TaskRunWhereInput, + argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | PrismaClientOrTransaction, + client?: PrismaClientOrTransaction + ): Promise { + const { args, prisma } = this.#resolveReadArgs(argsOrClient, client); + + 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[]>; + 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; + 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); + } + + /** + * 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 e680f254633..4c2d9d554aa 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -319,4 +319,62 @@ 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>; + findRun(where: Prisma.TaskRunWhereInput, client?: PrismaClientOrTransaction): Promise; + + findRunOrThrow( + where: Prisma.TaskRunWhereInput, + args: { select: S }, + client?: PrismaClientOrTransaction + ): Promise>; + findRunOrThrow( + where: Prisma.TaskRunWhereInput, + args: { include: I }, + client?: PrismaClientOrTransaction + ): Promise>; + findRunOrThrow(where: Prisma.TaskRunWhereInput, 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[]>; + findRuns( + args: { + where: Prisma.TaskRunWhereInput; + orderBy?: Prisma.TaskRunOrderByWithRelationInput | Prisma.TaskRunOrderByWithRelationInput[]; + take?: number; + skip?: number; + cursor?: Prisma.TaskRunWhereUniqueInput; + }, + client?: PrismaClientOrTransaction + ): Promise; } diff --git a/scripts/recover-stuck-runs.ts b/scripts/recover-stuck-runs.ts index 15deeb899c9..28bb4e85e46 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;