refactor(run-engine): route TaskRun reads through the run store

Relocate the direct TaskRun reads in the engine and its systems to the
RunStore read methods, preserving the exact client (writer, replica, or
transaction) at each site. Behavior-preserving; the engine test suite is
unchanged.

Author: d-cs

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<Array<{ id: string; orgId: string }>> {
-    const runs = await this.readOnlyPrisma.taskRun.findMany({
+    const runs = await this.runStore.findRuns({
       where: {
         id: { in: runIds },
         completedAt: {

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<RetryOutcome> {
   // 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;

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 });

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<DebounceResult> {
-    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", {

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}`);

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");

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,