fix(webapp): authorize batch parent runs before attaching waitpoints

Mono-RevId: c69d1cf0ad5e5d77b4c00aa3c5c59f6542dea4ce

Author: carderne

.server-changes/batch-parent-environment.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Reject batch-and-wait requests when the parent run belongs to another environment

apps/webapp/app/runEngine/services/batchTrigger.server.ts

@@ -7,7 +7,7 @@ import {
   parsePacket,
   TaskRunErrorCodes,
 } from "@trigger.dev/core/v3";
-import { BatchId, RunId } from "@trigger.dev/core/v3/isomorphic";
+import { BatchId } from "@trigger.dev/core/v3/isomorphic";
 import { type BatchTaskRun, Prisma } from "@trigger.dev/database";
 import { Evt } from "evt";
 import { z } from "zod";
@@ -28,6 +28,7 @@ import type { RunEngine } from "../../v3/runEngine.server";
 import { ServiceValidationError, WithRunEngine } from "../../v3/services/baseService.server";
 import { TriggerTaskService } from "../../v3/services/triggerTask.server";
 import { startActiveSpan } from "../../v3/tracer.server";
+import { resolveBatchParentRun } from "./resolveBatchParentRun.server";
 import { TriggerFailedTaskService } from "./triggerFailedTask.server";
 
 const PROCESSING_BATCH_SIZE = 50;
@@ -90,6 +91,13 @@ export class RunEngineBatchTriggerService extends WithRunEngine {
         "call()",
         environment,
         async (span) => {
+          const parentRunInternalId = await resolveBatchParentRun({
+            runStore: this._engine.runStore,
+            environmentId: environment.id,
+            parentRunId: body.parentRunId,
+            resumeParentOnCompletion: body.resumeParentOnCompletion,
+          });
+
           const { friendlyId } = await mintBatchFriendlyId({
             environment: {
               organizationId: environment.organizationId,
@@ -112,7 +120,8 @@ export class RunEngineBatchTriggerService extends WithRunEngine {
             payloadPacket,
             environment,
             body,
-            options
+            options,
+            parentRunInternalId
           );
 
           if (!batch) {
@@ -165,7 +174,8 @@ export class RunEngineBatchTriggerService extends WithRunEngine {
     payloadPacket: IOPacket,
     environment: AuthenticatedEnvironment,
     body: BatchTriggerTaskV2RequestBody,
-    options: BatchTriggerTaskServiceOptions = {}
+    options: BatchTriggerTaskServiceOptions = {},
+    parentRunInternalId?: string
   ) {
     // BatchTaskRun.runtimeEnvironmentId no longer has an FK into RuntimeEnvironment;
     // validate env existence app-side (covers both create arms below).
@@ -187,9 +197,9 @@ export class RunEngineBatchTriggerService extends WithRunEngine {
 
       this.onBatchTaskRunCreated.post(batch);
 
-      if (body.parentRunId && body.resumeParentOnCompletion) {
+      if (parentRunInternalId) {
         await this._engine.blockRunWithCreatedBatch({
-          runId: RunId.fromFriendlyId(body.parentRunId),
+          runId: parentRunInternalId,
           batchId: batch.id,
           environmentId: environment.id,
           projectId: environment.projectId,
@@ -278,9 +288,9 @@ export class RunEngineBatchTriggerService extends WithRunEngine {
 
       this.onBatchTaskRunCreated.post(batch);
 
-      if (body.parentRunId && body.resumeParentOnCompletion) {
+      if (parentRunInternalId) {
         await this._engine.blockRunWithCreatedBatch({
-          runId: RunId.fromFriendlyId(body.parentRunId),
+          runId: parentRunInternalId,
           batchId: batch.id,
           environmentId: environment.id,
           projectId: environment.projectId,

apps/webapp/app/runEngine/services/createBatch.server.ts

@@ -1,6 +1,5 @@
 import type { InitializeBatchOptions } from "@internal/run-engine";
 import { type CreateBatchRequestBody, type CreateBatchResponse } from "@trigger.dev/core/v3";
-import { RunId } from "@trigger.dev/core/v3/isomorphic";
 import { type BatchTaskRun, Prisma } from "@trigger.dev/database";
 import { Evt } from "evt";
 import { prisma, type PrismaClientOrTransaction } from "~/db.server";
@@ -14,6 +13,7 @@ import { ServiceValidationError, WithRunEngine } from "../../v3/services/baseSer
 import { BatchRateLimitExceededError, getBatchLimits } from "../concerns/batchLimits.server";
 import { DefaultQueueManager } from "../concerns/queues.server";
 import { DefaultTriggerTaskValidator } from "../validators/triggerTaskValidator";
+import { resolveBatchParentRun } from "./resolveBatchParentRun.server";
 
 export type CreateBatchServiceOptions = {
   triggerVersion?: string;
@@ -101,6 +101,13 @@ export class CreateBatchService extends WithRunEngine {
           // Note: Queue size limits are validated per-queue when batch items are processed,
           // since we don't know which queues items will go to until they're streamed.
 
+          const parentRunInternalId = await resolveBatchParentRun({
+            runStore: this._engine.runStore,
+            environmentId: environment.id,
+            parentRunId: body.parentRunId,
+            resumeParentOnCompletion: body.resumeParentOnCompletion,
+          });
+
           // BatchTaskRun.runtimeEnvironmentId no longer has an FK into RuntimeEnvironment;
           // validate env existence app-side (passthrough when split is off).
           await controlPlaneResolver.assertEnvExists(environment.id);
@@ -125,14 +132,14 @@ export class CreateBatchService extends WithRunEngine {
           await batchStreamGrants.mint(environment.id, friendlyId);
 
           // Block parent run if this is a batchTriggerAndWait
-          if (body.parentRunId && body.resumeParentOnCompletion) {
+          if (parentRunInternalId) {
             await this._engine.scheduleExpireBatch({
               batchId: batch.id,
               availableAt: new Date(Date.now() + env.BATCH_SEAL_TIMEOUT_MS),
             });
 
             await this._engine.blockRunWithCreatedBatch({
-              runId: RunId.fromFriendlyId(body.parentRunId),
+              runId: parentRunInternalId,
               batchId: batch.id,
               environmentId: environment.id,
               projectId: environment.projectId,

apps/webapp/app/runEngine/services/resolveBatchParentRun.server.test.ts

@@ -0,0 +1,83 @@
+import type { RunStore } from "@internal/run-store";
+import { describe, expect, it, vi } from "vitest";
+import type { ServiceValidationError } from "../../v3/services/baseService.server";
+import { resolveBatchParentRun } from "./resolveBatchParentRun.server";
+
+type BatchParentRunStore = Pick<RunStore, "findRun" | "findRunOnPrimary">;
+
+function createRunStore({ replicaRun, primaryRun }: { replicaRun?: object; primaryRun?: object }) {
+  return {
+    findRun: vi.fn().mockResolvedValue(replicaRun ?? null),
+    findRunOnPrimary: vi.fn().mockResolvedValue(primaryRun ?? null),
+  } as unknown as BatchParentRunStore;
+}
+
+describe("resolveBatchParentRun", () => {
+  it("does not resolve an unused parent run", async () => {
+    const runStore = createRunStore({});
+
+    await expect(
+      resolveBatchParentRun({
+        runStore,
+        environmentId: "env_1",
+        parentRunId: "run_parent",
+        resumeParentOnCompletion: false,
+      })
+    ).resolves.toBeUndefined();
+
+    expect(runStore.findRun).not.toHaveBeenCalled();
+  });
+
+  it("returns a parent run from the calling environment", async () => {
+    const runStore = createRunStore({ replicaRun: { id: "parent" } });
+
+    await expect(
+      resolveBatchParentRun({
+        runStore,
+        environmentId: "env_1",
+        parentRunId: "run_parent",
+        resumeParentOnCompletion: true,
+      })
+    ).resolves.toBe("parent");
+
+    expect(runStore.findRun).toHaveBeenCalledWith(
+      { id: "parent", runtimeEnvironmentId: "env_1" },
+      { select: { id: true } }
+    );
+    expect(runStore.findRunOnPrimary).not.toHaveBeenCalled();
+  });
+
+  it("falls back to the owning primary", async () => {
+    const runStore = createRunStore({ primaryRun: { id: "parent" } });
+
+    await expect(
+      resolveBatchParentRun({
+        runStore,
+        environmentId: "env_1",
+        parentRunId: "run_parent",
+        resumeParentOnCompletion: true,
+      })
+    ).resolves.toBe("parent");
+
+    expect(runStore.findRunOnPrimary).toHaveBeenCalledWith(
+      { id: "parent", runtimeEnvironmentId: "env_1" },
+      { select: { id: true } }
+    );
+  });
+
+  it("rejects a parent outside the calling environment", async () => {
+    const runStore = createRunStore({});
+
+    const result = resolveBatchParentRun({
+      runStore,
+      environmentId: "env_1",
+      parentRunId: "run_foreign",
+      resumeParentOnCompletion: true,
+    });
+
+    await expect(result).rejects.toMatchObject<ServiceValidationError>({
+      message: "Parent run not found in the calling environment",
+      status: 404,
+    });
+  });
+});

apps/webapp/app/runEngine/services/resolveBatchParentRun.server.ts

@@ -0,0 +1,34 @@
+import type { RunStore } from "@internal/run-store";
+import { RunId } from "@trigger.dev/core/v3/isomorphic";
+import { ServiceValidationError } from "../../v3/services/baseService.server";
+
+type BatchParentRunStore = Pick<RunStore, "findRun" | "findRunOnPrimary">;
+
+export async function resolveBatchParentRun({
+  runStore,
+  environmentId,
+  parentRunId,
+  resumeParentOnCompletion,
+}: {
+  runStore: BatchParentRunStore;
+  environmentId: string;
+  parentRunId?: string;
+  resumeParentOnCompletion?: boolean;
+}): Promise<string | undefined> {
+  if (!parentRunId || !resumeParentOnCompletion) {
+    return;
+  }
+
+  const runId = RunId.fromFriendlyId(parentRunId);
+  const where = { id: runId, runtimeEnvironmentId: environmentId };
+  const args = { select: { id: true } } as const;
+
+  const parentRun =
+    (await runStore.findRun(where, args)) ?? (await runStore.findRunOnPrimary(where, args));
+
+  if (!parentRun) {
+    throw new ServiceValidationError("Parent run not found in the calling environment", 404);
+  }
+
+  return parentRun.id;
+}