perf(run-store,webapp): scope idempotency dedup to one table for non-v2 orgs, add cross-table tests

The idempotency-key dedup is a non-id predicate, so RunStore read BOTH run
tables in parallel on every idempotency-keyed trigger, including orgs not cut
over to v2 (whose runs only live in TaskRun, so the task_run_v2 query is always
empty; while native realtime is off that is every org). Add an optional
`tables: "legacy" | "both"` scope to findRun and pass "legacy" from the
idempotency concern when the org is not on v2, keeping the trigger hot path
single-table.

Backfills cross-table tests the audit flagged as missing: findRun legacy-scope
skips task_run_v2, and clearIdempotencyKey fans out across both tables
(byPredicate hits v2; a mixed byFriendlyIds array clears both).

Author: d-cs

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

@@ -302,6 +302,20 @@ export class IdempotencyKeyConcern {
       return { isCached: false, idempotencyKey, idempotencyKeyExpiresAt };
     }
 
+    // Resolve the org's v2-cutover state ONCE: it gates both the cross-table
+    // idempotency lookup below and the pre-gate claim further down. While the
+    // org is not on v2 (the default, and every org while native realtime is
+    // off) the run can only be in the legacy table, so scope the dedup read to
+    // "legacy" and skip the empty task_run_v2 query on the trigger hot path.
+    const orgFeatureFlags =
+      (request.environment.organization?.featureFlags as
+        | Record<string, unknown>
+        | null
+        | undefined) ?? null;
+    const orgUsesV2 = shouldUseV2RunTable(orgFeatureFlags, {
+      nativeRealtimeEnabled: env.REALTIME_BACKEND_NATIVE_ENABLED === "1",
+    });
+
     const existingRun = idempotencyKey
       ? await runStore.findRun(
           {
@@ -313,6 +327,7 @@ export class IdempotencyKeyConcern {
             include: {
               associatedWaitpoint: true,
             },
+            tables: orgUsesV2 ? "both" : "legacy",
           },
           this.prisma
         )
@@ -408,11 +423,6 @@ export class IdempotencyKeyConcern {
     // TaskRun, ksuid -> task_run_v2); the per-table idempotency unique
     // constraints can't see each other, so neither INSERT raises P2002 and two
     // runs share one key. The Redis claim is the only backstop in that window.
-    const orgFeatureFlags =
-      (request.environment.organization?.featureFlags as
-        | Record<string, unknown>
-        | null
-        | undefined) ?? null;
     // v2-cutover orgs: an idempotency-keyed trigger can straddle a `runTableV2`
     // flag flip into different physical tables (cuid -> TaskRun, ksuid ->
     // task_run_v2), and the per-table idempotency-key unique constraints can't
@@ -430,9 +440,7 @@ export class IdempotencyKeyConcern {
     // a pathological client. shouldUseV2RunTable is checked first so a v2 org
     // skips the mollifier-flag resolve entirely.
     const claimEligible =
-      shouldUseV2RunTable(orgFeatureFlags, {
-        nativeRealtimeEnabled: env.REALTIME_BACKEND_NATIVE_ENABLED === "1",
-      }) ||
+      orgUsesV2 ||
       (!request.body.options?.resumeParentOnCompletion &&
         !request.body.options?.debounce &&
         !request.options?.oneTimeUseToken &&

internal-packages/run-store/src/PostgresRunStore.test.ts

@@ -2048,6 +2048,127 @@ describe("PostgresRunStore — table routing by id format", () => {
     }
   );
 
+  postgresTest(
+    "findRun tables:'legacy' skips the task_run_v2 query (idempotency hot-path scope)",
+    async ({ prisma }) => {
+      const { organization, project, environment } = await seedEnvironment(prisma);
+      const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
+
+      // A v2 (ksuid) run carrying an idempotency key — it lives only in
+      // task_run_v2.
+      const ksuid = RunId.generateKsuid();
+      await seedRoutedRun(prisma, {
+        id: ksuid.id,
+        friendlyId: ksuid.friendlyId,
+        organizationId: organization.id,
+        projectId: project.id,
+        runtimeEnvironmentId: environment.id,
+        idempotencyKey: "idem-scope",
+        taskIdentifier: "my-task",
+      });
+
+      const where = {
+        runtimeEnvironmentId: environment.id,
+        idempotencyKey: "idem-scope",
+        taskIdentifier: "my-task",
+      };
+
+      // Default (both tables) and an explicit "both" find the v2 run.
+      expect((await store.findRun(where))?.id).toBe(ksuid.id);
+      expect(
+        (await store.findRun(where, { select: { id: true }, tables: "both" }))?.id
+      ).toBe(ksuid.id);
+
+      // "legacy" scope skips task_run_v2 entirely, so the v2 run is NOT found.
+      // This is the hot-path optimisation for an org not cut over to v2: its
+      // runs only live in TaskRun, so the second (v2) query is always empty and
+      // can be skipped. (If a caller mis-scopes a genuinely-v2 org to legacy it
+      // would miss the run — hence it is gated on shouldUseV2RunTable upstream.)
+      expect(await store.findRun(where, { select: { id: true }, tables: "legacy" })).toBeNull();
+    }
+  );
+
+  postgresTest(
+    "clearIdempotencyKey fans out across both tables (byPredicate hits v2; byFriendlyIds partitions a mixed array)",
+    async ({ prisma }) => {
+      const { organization, project, environment } = await seedEnvironment(prisma);
+      const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
+
+      // byPredicate carries no id, so it must reach task_run_v2 to clear a v2 run.
+      const v2Pred = RunId.generateKsuid();
+      await seedRoutedRun(prisma, {
+        id: v2Pred.id,
+        friendlyId: v2Pred.friendlyId,
+        organizationId: organization.id,
+        projectId: project.id,
+        runtimeEnvironmentId: environment.id,
+        idempotencyKey: "kp-v2",
+        taskIdentifier: "my-task",
+      });
+
+      const predResult = await store.clearIdempotencyKey({
+        byPredicate: {
+          idempotencyKey: "kp-v2",
+          taskIdentifier: "my-task",
+          runtimeEnvironmentId: environment.id,
+        },
+      });
+      expect(predResult.count).toBe(1);
+      expect(
+        (
+          await prisma.taskRunV2.findFirst({
+            where: { id: v2Pred.id },
+            select: { idempotencyKey: true },
+          })
+        )?.idempotencyKey
+      ).toBeNull();
+
+      // byFriendlyIds with a MIXED (ksuid + cuid) array must clear rows in BOTH
+      // physical tables — the partition + sum is the cross-table behaviour.
+      const v2F = RunId.generateKsuid();
+      const legacyF = RunId.generate();
+      await seedRoutedRun(prisma, {
+        id: v2F.id,
+        friendlyId: v2F.friendlyId,
+        organizationId: organization.id,
+        projectId: project.id,
+        runtimeEnvironmentId: environment.id,
+        idempotencyKey: "kf-v2",
+        taskIdentifier: "my-task",
+      });
+      await seedRoutedRun(prisma, {
+        id: legacyF.id,
+        friendlyId: legacyF.friendlyId,
+        organizationId: organization.id,
+        projectId: project.id,
+        runtimeEnvironmentId: environment.id,
+        idempotencyKey: "kf-legacy",
+        taskIdentifier: "my-task",
+      });
+
+      const friendlyResult = await store.clearIdempotencyKey({
+        byFriendlyIds: [v2F.friendlyId, legacyF.friendlyId],
+      });
+      expect(friendlyResult.count).toBe(2);
+      expect(
+        (
+          await prisma.taskRunV2.findFirst({
+            where: { id: v2F.id },
+            select: { idempotencyKey: true },
+          })
+        )?.idempotencyKey
+      ).toBeNull();
+      expect(
+        (
+          await prisma.taskRun.findFirst({
+            where: { id: legacyF.id },
+            select: { idempotencyKey: true },
+          })
+        )?.idempotencyKey
+      ).toBeNull();
+    }
+  );
+
   postgresTest(
     "expireRunsBatch with a mixed array updates both tables and returns the combined count",
     async ({ prisma }) => {

internal-packages/run-store/src/PostgresRunStore.ts

@@ -13,6 +13,7 @@ import type {
   CreateFailedRunInput,
   CreateRunInput,
   ExpireSnapshotInput,
+  FindRunTableScope,
   LockRunData,
   ReadClient,
   RescheduleSnapshotInput,
@@ -114,8 +115,16 @@ export class PostgresRunStore implements RunStore {
   async #findFirstAcrossTables(
     prisma: ReadClient,
     where: Prisma.TaskRunWhereInput,
-    args: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }
+    args: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude },
+    tables: FindRunTableScope = "both"
   ): Promise<unknown> {
+    // Legacy-only scope: the caller knows the run can't be in task_run_v2 (e.g.
+    // idempotency dedup for an org not cut over to v2), so skip the second,
+    // empty v2 query and keep this a single-table read on the hot path.
+    if (tables === "legacy") {
+      return prisma.taskRun.findFirst({ where, ...args });
+    }
+
     const v2Model = prisma.taskRunV2 as unknown as typeof prisma.taskRun;
 
     const [legacyRun, v2Run] = await Promise.all([
@@ -755,12 +764,12 @@ export class PostgresRunStore implements RunStore {
 
   findRun<S extends Prisma.TaskRunSelect>(
     where: Prisma.TaskRunWhereInput,
-    args: { select: S },
+    args: { select: S; tables?: FindRunTableScope },
     client?: ReadClient
   ): Promise<Prisma.TaskRunGetPayload<{ select: S }> | null>;
   findRun<I extends Prisma.TaskRunInclude>(
     where: Prisma.TaskRunWhereInput,
-    args: { include: I },
+    args: { include: I; tables?: FindRunTableScope },
     client?: ReadClient
   ): Promise<Prisma.TaskRunGetPayload<{ include: I }> | null>;
   findRun(
@@ -769,10 +778,12 @@ export class PostgresRunStore implements RunStore {
   ): Promise<TaskRun | null>;
   async findRun(
     where: Prisma.TaskRunWhereInput,
-    argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | ReadClient,
+    argsOrClient?:
+      | { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude; tables?: FindRunTableScope }
+      | ReadClient,
     client?: ReadClient
   ): Promise<unknown> {
-    const { args, prisma } = this.#resolveReadArgs(argsOrClient, client);
+    const { args, prisma, tables } = this.#resolveReadArgs(argsOrClient, client);
 
     const routingKey = this.#routingKeyOf(where);
     if (routingKey !== undefined) {
@@ -781,8 +792,8 @@ export class PostgresRunStore implements RunStore {
     }
 
     // Non-id predicate (e.g. idempotency-key dedup): the match can be in
-    // either table, so read both.
-    return this.#findFirstAcrossTables(prisma, where, args);
+    // either table, so read both (unless the caller scopes to legacy-only).
+    return this.#findFirstAcrossTables(prisma, where, args, tables);
   }
 
   findRunOrThrow<S extends Prisma.TaskRunSelect>(
@@ -1298,30 +1309,40 @@ export class PostgresRunStore implements RunStore {
    */
   #resolveReadArgs(
     argsOrClient:
-      | { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }
+      | { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude; tables?: FindRunTableScope }
       | ReadClient
       | undefined,
     client: ReadClient | undefined
   ): {
     args: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude };
     prisma: ReadClient;
+    tables: FindRunTableScope;
   } {
     const isProjection =
       typeof argsOrClient === "object" &&
       argsOrClient !== null &&
-      ("select" in argsOrClient || "include" in argsOrClient);
+      ("select" in argsOrClient || "include" in argsOrClient || "tables" in argsOrClient);
 
     if (isProjection) {
+      // Split the table-scope hint out of the args that get spread into Prisma
+      // (which would reject an unknown `tables` field).
+      const { tables, ...prismaArgs } = argsOrClient as {
+        select?: Prisma.TaskRunSelect;
+        include?: Prisma.TaskRunInclude;
+        tables?: FindRunTableScope;
+      };
       return {
-        args: argsOrClient as { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude },
+        args: prismaArgs,
         prisma: client ?? this.readOnlyPrisma,
+        tables: tables ?? "both",
       };
     }
 
     // No projection: the second positional arg, when present, is the client.
     return {
       args: {},
       prisma: (argsOrClient as ReadClient | undefined) ?? this.readOnlyPrisma,
+      tables: "both",
     };
   }
 }

internal-packages/run-store/src/types.ts

@@ -232,6 +232,18 @@ export type ClearIdempotencyKeyInput =
 
 export type TaskRunWithWaitpoint = TaskRun & { associatedWaitpoint: Waitpoint | null };
 
+/**
+ * Which physical run tables a non-id `findRun` predicate should read.
+ *
+ * Defaults to `"both"` (the safe cross-table behaviour). A caller that KNOWS
+ * the run can only be in the legacy table — e.g. the idempotency-key dedup for
+ * an org that is not cut over to `task_run_v2` — can pass `"legacy"` to skip the
+ * second (empty) `task_run_v2` query and keep the trigger hot path single-table.
+ * Only meaningful for non-id predicates; id/friendlyId reads already route to
+ * exactly one table by id format.
+ */
+export type FindRunTableScope = "both" | "legacy";
+
 export interface RunStore {
   // Create
   createRun(params: CreateRunInput, tx?: PrismaClientOrTransaction): Promise<TaskRunWithWaitpoint>;
@@ -332,12 +344,12 @@ export interface RunStore {
   // Read
   findRun<S extends Prisma.TaskRunSelect>(
     where: Prisma.TaskRunWhereInput,
-    args: { select: S },
+    args: { select: S; tables?: FindRunTableScope },
     client?: ReadClient
   ): Promise<Prisma.TaskRunGetPayload<{ select: S }> | null>;
   findRun<I extends Prisma.TaskRunInclude>(
     where: Prisma.TaskRunWhereInput,
-    args: { include: I },
+    args: { include: I; tables?: FindRunTableScope },
     client?: ReadClient
   ): Promise<Prisma.TaskRunGetPayload<{ include: I }> | null>;
   findRun(where: Prisma.TaskRunWhereInput, client?: ReadClient): Promise<TaskRun | null>;