refactor(webapp): route run metadata, idempotency-key, and reschedule writes through RunStore

Author: d-cs

apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts

@@ -10,6 +10,7 @@ import { getMollifierBuffer } from "~/v3/mollifier/mollifierBuffer.server";
 import { findRunByIdWithMollifierFallback } from "~/v3/mollifier/readFallback.server";
 import { claimOrAwait } from "~/v3/mollifier/idempotencyClaim.server";
 import { makeResolveMollifierFlag } from "~/v3/mollifier/mollifierGate.server";
+import { runStore } from "~/v3/runStore.server";
 import type { TraceEventConcern, TriggerTaskRequest } from "../types";
 
 // In-memory per-org mollifier-enabled check, shared with `evaluateGate`
@@ -190,10 +191,10 @@ export class IdempotencyKeyConcern {
         });
 
         // Update the existing run to remove the idempotency key
-        await this.prisma.taskRun.updateMany({
-          where: { id: existingRun.id, idempotencyKey },
-          data: { idempotencyKey: null, idempotencyKeyExpiresAt: null },
-        });
+        await runStore.clearIdempotencyKey(
+          { byId: { runId: existingRun.id, idempotencyKey } },
+          this.prisma
+        );
 
         return { isCached: false, idempotencyKey, idempotencyKeyExpiresAt };
       }
@@ -207,10 +208,10 @@ export class IdempotencyKeyConcern {
         });
 
         // Update the existing run to remove the idempotency key
-        await this.prisma.taskRun.updateMany({
-          where: { id: existingRun.id, idempotencyKey },
-          data: { idempotencyKey: null, idempotencyKeyExpiresAt: null },
-        });
+        await runStore.clearIdempotencyKey(
+          { byId: { runId: existingRun.id, idempotencyKey } },
+          this.prisma
+        );
 
         return { isCached: false, idempotencyKey, idempotencyKeyExpiresAt };
       }

apps/webapp/app/services/metadata/updateMetadata.server.ts

@@ -13,6 +13,8 @@ import { Effect, Schedule, Duration, Fiber } from "effect";
 import { type RuntimeFiber } from "effect/Fiber";
 import { setTimeout } from "timers/promises";
 import { Logger, LogLevel } from "@trigger.dev/core/logger";
+import type { RunStore } from "@internal/run-store";
+import { runStore as defaultRunStore } from "~/v3/runStore.server";
 
 const RUN_UPDATABLE_WINDOW_MS = 60 * 60 * 1000; // 1 hour
 
@@ -24,6 +26,7 @@ type BufferedRunMetadataChangeOperation = {
 
 export type UpdateMetadataServiceOptions = {
   prisma: PrismaClientOrTransaction;
+  runStore?: RunStore;
   flushIntervalMs?: number;
   flushEnabled?: boolean;
   flushLoggingEnabled?: boolean;
@@ -49,6 +52,7 @@ export class UpdateMetadataService {
   private _bufferedOperations: Map<string, BufferedRunMetadataChangeOperation[]> = new Map();
   private _flushFiber: RuntimeFiber<void> | null = null;
   private readonly _prisma: PrismaClientOrTransaction;
+  private readonly _runStore: RunStore;
   private readonly flushIntervalMs: number;
   private readonly flushEnabled: boolean;
   private readonly flushLoggingEnabled: boolean;
@@ -57,6 +61,7 @@ export class UpdateMetadataService {
 
   constructor(private readonly options: UpdateMetadataServiceOptions) {
     this._prisma = options.prisma;
+    this._runStore = options.runStore ?? defaultRunStore;
     this.flushIntervalMs = options.flushIntervalMs ?? 5000;
     this.flushEnabled = options.flushEnabled ?? true;
     this.flushLoggingEnabled = options.flushLoggingEnabled ?? true;
@@ -260,17 +265,16 @@ export class UpdateMetadataService {
       const writeTime = new Date();
       const result = yield* _(
         Effect.tryPromise(() =>
-          this._prisma.taskRun.updateMany({
-            where: {
-              id: runId,
-              metadataVersion: run.metadataVersion,
-            },
-            data: {
-              metadata: newMetadataPacket.data,
+          this._runStore.updateMetadata(
+            runId,
+            {
+              metadata: newMetadataPacket.data!,
               metadataVersion: { increment: 1 },
               updatedAt: writeTime,
             },
-          })
+            { expectedMetadataVersion: run.metadataVersion },
+            this._prisma
+          )
         )
       );
 
@@ -469,20 +473,19 @@ export class UpdateMetadataService {
       // Update with optimistic locking; updatedAt stamped explicitly so the caller can
       // publish the exact committed watermark without a follow-up read.
       const writeTime = new Date();
-      const result = await this._prisma.taskRun.updateMany({
-        where: {
-          id: runId,
-          metadataVersion: run.metadataVersion,
-        },
-        data: {
-          metadata: newMetadataPacket.data,
+      const result = await this._runStore.updateMetadata(
+        runId,
+        {
+          metadata: newMetadataPacket.data!,
           metadataType: newMetadataPacket.dataType,
           metadataVersion: {
             increment: 1,
           },
           updatedAt: writeTime,
         },
-      });
+        { expectedMetadataVersion: run.metadataVersion },
+        this._prisma
+      );
 
       if (result.count === 0) {
         if (this.flushLoggingEnabled) {
@@ -564,19 +567,19 @@ export class UpdateMetadataService {
       // Update the metadata without version check; updatedAt stamped explicitly so the
       // caller can publish the exact committed watermark.
       const writeTime = new Date();
-      await this._prisma.taskRun.update({
-        where: {
-          id: runId,
-        },
-        data: {
-          metadata: metadataPacket?.data,
+      await this._runStore.updateMetadata(
+        runId,
+        {
+          metadata: metadataPacket?.data!,
           metadataType: metadataPacket?.dataType,
           metadataVersion: {
             increment: 1,
           },
           updatedAt: writeTime,
         },
-      });
+        {},
+        this._prisma
+      );
       updatedAtMs = writeTime.getTime();
     }
 

apps/webapp/app/v3/services/batchTriggerV3.server.ts

@@ -408,10 +408,10 @@ export class BatchTriggerV3Service extends BaseService {
 
     // Expire the cached runs that are no longer valid
     if (expiredRunIds.size) {
-      await this._prisma.taskRun.updateMany({
-        where: { friendlyId: { in: Array.from(expiredRunIds) } },
-        data: { idempotencyKey: null },
-      });
+      await this.runStore.clearIdempotencyKey(
+        { byFriendlyIds: Array.from(expiredRunIds) },
+        this._prisma
+      );
     }
 
     return runs;

apps/webapp/app/v3/services/rescheduleTaskRun.server.ts

@@ -17,15 +17,14 @@ export class RescheduleTaskRunService extends BaseService {
       throw new ServiceValidationError(`Invalid delay: ${body.delay}`);
     }
 
-    const updatedRun = await this._prisma.taskRun.update({
-      where: {
-        id: taskRun.id,
-      },
-      data: {
+    const updatedRun = await this.runStore.rescheduleRun(
+      taskRun.id,
+      {
         delayUntil: delay,
         queueTimestamp: delay,
       },
-    });
+      this._prisma
+    );
 
     if (updatedRun.engine === "V1") {
       await EnqueueDelayedRunService.reschedule(taskRun.id, delay);

apps/webapp/app/v3/services/resetIdempotencyKey.server.ts

@@ -9,17 +9,16 @@ export class ResetIdempotencyKeyService extends BaseService {
     taskIdentifier: string,
     authenticatedEnv: AuthenticatedEnvironment
   ): Promise<{ id: string }> {
-    const { count: pgCount } = await this._prisma.taskRun.updateMany({
-      where: {
-        idempotencyKey,
-        taskIdentifier,
-        runtimeEnvironmentId: authenticatedEnv.id,
-      },
-      data: {
-        idempotencyKey: null,
-        idempotencyKeyExpiresAt: null,
+    const { count: pgCount } = await this.runStore.clearIdempotencyKey(
+      {
+        byPredicate: {
+          idempotencyKey,
+          taskIdentifier,
+          runtimeEnvironmentId: authenticatedEnv.id,
+        },
       },
-    });
+      this._prisma
+    );
 
     // Buffer-side reset: the key may belong to a buffered run that
     // hasn't materialised yet. The PG updateMany above can't see it.
@@ -75,17 +74,16 @@ export class ResetIdempotencyKeyService extends BaseService {
       // lookup against the writer when there's nothing to find;
       // otherwise the exact write the customer asked for (i.e., not
       // duplicative — without it the reset is silently lost).
-      const { count: handoffPgCount } = await this._prisma.taskRun.updateMany({
-        where: {
-          idempotencyKey,
-          taskIdentifier,
-          runtimeEnvironmentId: authenticatedEnv.id,
-        },
-        data: {
-          idempotencyKey: null,
-          idempotencyKeyExpiresAt: null,
+      const { count: handoffPgCount } = await this.runStore.clearIdempotencyKey(
+        {
+          byPredicate: {
+            idempotencyKey,
+            taskIdentifier,
+            runtimeEnvironmentId: authenticatedEnv.id,
+          },
         },
-      });
+        this._prisma
+      );
       if (handoffPgCount > 0) {
         logger.info(
           `Reset idempotency key via handoff re-check: ${idempotencyKey} for task: ${taskIdentifier} in env: ${authenticatedEnv.id}, affected ${handoffPgCount} run(s)`