triggerdotdev--trigger.dev/apps/webapp/app/runEngine/services/createBatch.server.ts

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";
import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { logger } from "~/services/logger.server";
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
import { mintBatchFriendlyId } from "~/v3/runOpsMigration/mintBatchFriendlyId.server";
import { ServiceValidationError, WithRunEngine } from "../../v3/services/baseService.server";
import { BatchRateLimitExceededError, getBatchLimits } from "../concerns/batchLimits.server";
import { DefaultQueueManager } from "../concerns/queues.server";
import { DefaultTriggerTaskValidator } from "../validators/triggerTaskValidator";

export type CreateBatchServiceOptions = {
  triggerVersion?: string;
  traceContext?: Record<string, string | undefined | Record<string, string | undefined>>;
  spanParentAsLink?: boolean;
  oneTimeUseToken?: string;
  realtimeStreamsVersion?: "v1" | "v2";
  triggerSource?: string;
};

/**
 * Create Batch Service (Phase 1 of 2-phase batch API).
 *
 * This service handles Phase 1 of the streaming batch API:
 * 1. Validates entitlement and queue limits
 * 2. Creates BatchTaskRun in Postgres with status=PENDING, expectedCount set
 * 3. For batchTriggerAndWait: blocks the parent run immediately
 * 4. Initializes batch metadata in Redis
 * 5. Returns batch ID - items are streamed separately via Phase 2
 *
 * The batch is NOT sealed until Phase 2 completes.
 */
export class CreateBatchService extends WithRunEngine {
  public onBatchTaskRunCreated: Evt<BatchTaskRun> = new Evt();
  private readonly queueConcern: DefaultQueueManager;
  private readonly validator: DefaultTriggerTaskValidator;

  constructor(protected readonly _prisma: PrismaClientOrTransaction = prisma) {
    super({ prisma: _prisma });

    this.queueConcern = new DefaultQueueManager(this._prisma, this._engine, this._replica);
    this.validator = new DefaultTriggerTaskValidator();
  }

  /**
   * Create a batch for 2-phase processing.
   * Items will be streamed separately via the StreamBatchItemsService.
   */
  public async call(
    environment: AuthenticatedEnvironment,
    body: CreateBatchRequestBody,
    options: CreateBatchServiceOptions = {}
  ): Promise<CreateBatchResponse> {
    try {
      return await this.traceWithEnv<CreateBatchResponse>(
        "createBatch()",
        environment,
        async (span) => {
          const { id, friendlyId } = await mintBatchFriendlyId({
            environment: {
              organizationId: environment.organizationId,
              id: environment.id,
              orgFeatureFlags: environment.organization.featureFlags,
            },
            parentRunFriendlyId: body.parentRunId,
          });

          span.setAttribute("batchId", friendlyId);
          span.setAttribute("runCount", body.runCount);

          const entitlementValidation = await this.validator.validateEntitlement({
            environment,
          });

          if (!entitlementValidation.ok) {
            throw entitlementValidation.error;
          }

          const planType = entitlementValidation.plan?.type;

          const { config, rateLimiter } = await getBatchLimits(environment.organization);

          // Rate-limit before creating the batch, to stop bursts exceeding the limit.
          const rateResult = await rateLimiter.limit(environment.id, body.runCount);

          if (!rateResult.success) {
            throw new BatchRateLimitExceededError(
              rateResult.limit,
              rateResult.remaining,
              new Date(rateResult.reset),
              body.runCount
            );
          }

          // 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.

          // BatchTaskRun.runtimeEnvironmentId no longer has an FK into RuntimeEnvironment;
          // validate env existence app-side (passthrough when split is off).
          await controlPlaneResolver.assertEnvExists(environment.id);

          // Created PENDING; sealed (status -> PROCESSING) once items are streamed.
          const batch = await this._engine.runStore.createBatchTaskRun({
            id,
            friendlyId,
            runtimeEnvironmentId: environment.id,
            status: "PENDING",
            runCount: body.runCount,
            expectedCount: body.runCount,
            runIds: [],
            batchVersion: "runengine:v2", // 2-phase streaming batch API
            oneTimeUseToken: options.oneTimeUseToken,
            idempotencyKey: body.idempotencyKey,
            sealed: false,
          });

          this.onBatchTaskRunCreated.post(batch);

          // Block parent run if this is a batchTriggerAndWait
          if (body.parentRunId && body.resumeParentOnCompletion) {
            await this._engine.blockRunWithCreatedBatch({
              runId: RunId.fromFriendlyId(body.parentRunId),
              batchId: batch.id,
              environmentId: environment.id,
              projectId: environment.projectId,
              organizationId: environment.organizationId,
            });
          }

          // Initialize batch metadata in Redis (without items)
          const initOptions: InitializeBatchOptions = {
            batchId: id,
            friendlyId,
            environmentId: environment.id,
            environmentType: environment.type,
            organizationId: environment.organizationId,
            projectId: environment.projectId,
            runCount: body.runCount,
            parentRunId: body.parentRunId,
            resumeParentOnCompletion: body.resumeParentOnCompletion,
            triggerVersion: options.triggerVersion,
            traceContext: options.traceContext as Record<string, unknown> | undefined,
            spanParentAsLink: options.spanParentAsLink,
            realtimeStreamsVersion: options.realtimeStreamsVersion,
            idempotencyKey: body.idempotencyKey,
            processingConcurrency: config.processingConcurrency,
            planType,
            triggerSource: options.triggerSource,
          };

          await this._engine.initializeBatch(initOptions);

          logger.info("Batch created", {
            batchId: friendlyId,
            runCount: body.runCount,
            envId: environment.id,
            projectId: environment.projectId,
            parentRunId: body.parentRunId,
            resumeParentOnCompletion: body.resumeParentOnCompletion,
            processingConcurrency: config.processingConcurrency,
          });

          return {
            id: friendlyId,
            runCount: body.runCount,
            isCached: false,
            idempotencyKey: body.idempotencyKey,
          };
        }
      );
    } catch (error) {
      // Handle Prisma unique constraint violations
      if (error instanceof Prisma.PrismaClientKnownRequestError) {
        logger.debug("CreateBatchService: Prisma error", {
          code: error.code,
          message: error.message,
          meta: error.meta,
        });

        if (error.code === "P2002") {
          const target = error.meta?.target;

          if (
            Array.isArray(target) &&
            target.length > 0 &&
            typeof target[0] === "string" &&
            target[0].includes("oneTimeUseToken")
          ) {
            throw new ServiceValidationError(
              "Cannot create batch with a one-time use token as it has already been used."
            );
          } else {
            throw new ServiceValidationError(
              "Cannot create batch as it has already been created with the same idempotency key."
            );
          }
        }
      }

      throw error;
    }
  }
}