triggerdotdev--trigger.dev/apps/webapp/app/runEngine/concerns/batchLimits.server.ts

import { z } from "zod";
import { env } from "~/env.server";
import {
  createLimiterFromConfig,
  RateLimiterConfig,
} from "~/services/authorizationRateLimitMiddleware.server";
import type { Duration } from "~/services/rateLimiter.server";
import { createRedisRateLimitClient, RateLimiter } from "~/services/rateLimiter.server";
import { singleton } from "~/utils/singleton";

const BatchLimitsConfig = z.object({
  processingConcurrency: z.number().int().default(env.BATCH_CONCURRENCY_LIMIT_DEFAULT),
});

/**
 * Batch limits configuration for a plan type
 */
export type BatchLimitsConfig = z.infer<typeof BatchLimitsConfig>;

const batchLimitsRedisClient = singleton("batchLimitsRedisClient", createBatchLimitsRedisClient);

function createBatchLimitsRedisClient() {
  const redisClient = createRedisRateLimitClient({
    port: env.RATE_LIMIT_REDIS_PORT,
    host: env.RATE_LIMIT_REDIS_HOST,
    username: env.RATE_LIMIT_REDIS_USERNAME,
    password: env.RATE_LIMIT_REDIS_PASSWORD,
    tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
    clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
  });

  return redisClient;
}

// Just the org fields this module reads. Compatible with both the full
// Prisma `Organization` payload and the slim `AuthenticatedEnvironment`
// `["organization"]` shape (when passed `batchRateLimitConfig` /
// `batchQueueConcurrencyConfig` as `unknown`).
type OrganizationForBatchLimits = {
  batchRateLimitConfig?: unknown;
  batchQueueConcurrencyConfig?: unknown;
};

function createOrganizationRateLimiter(organization: OrganizationForBatchLimits): RateLimiter {
  const limiterConfig = resolveBatchRateLimitConfig(organization.batchRateLimitConfig);

  const limiter = createLimiterFromConfig(limiterConfig);

  return new RateLimiter({
    redisClient: batchLimitsRedisClient,
    keyPrefix: "ratelimit:batch",
    limiter,
    logSuccess: false,
    logFailure: true,
  });
}

function resolveBatchRateLimitConfig(batchRateLimitConfig?: unknown): RateLimiterConfig {
  const defaultRateLimiterConfig: RateLimiterConfig = {
    type: "tokenBucket",
    refillRate: env.BATCH_RATE_LIMIT_REFILL_RATE,
    interval: env.BATCH_RATE_LIMIT_REFILL_INTERVAL as Duration,
    maxTokens: env.BATCH_RATE_LIMIT_MAX,
  };

  if (!batchRateLimitConfig) {
    return defaultRateLimiterConfig;
  }

  const parsedBatchRateLimitConfig = RateLimiterConfig.safeParse(batchRateLimitConfig);

  if (!parsedBatchRateLimitConfig.success) {
    return defaultRateLimiterConfig;
  }

  return parsedBatchRateLimitConfig.data;
}

/**
 * Get the rate limiter and limits for an organization.
 * Internally looks up the plan type, but doesn't expose it to callers.
 */
export async function getBatchLimits(
  organization: OrganizationForBatchLimits
): Promise<{ rateLimiter: RateLimiter; config: BatchLimitsConfig }> {
  const rateLimiter = createOrganizationRateLimiter(organization);
  const config = resolveBatchLimitsConfig(organization.batchQueueConcurrencyConfig);
  return { rateLimiter, config };
}

function resolveBatchLimitsConfig(batchLimitsConfig?: unknown): BatchLimitsConfig {
  const defaultLimitsConfig: BatchLimitsConfig = {
    processingConcurrency: env.BATCH_CONCURRENCY_LIMIT_DEFAULT,
  };

  if (!batchLimitsConfig) {
    return defaultLimitsConfig;
  }

  const parsedBatchLimitsConfig = BatchLimitsConfig.safeParse(batchLimitsConfig);

  if (!parsedBatchLimitsConfig.success) {
    return defaultLimitsConfig;
  }

  return parsedBatchLimitsConfig.data;
}

/**
 * Error thrown when batch rate limit is exceeded.
 * Contains information for constructing a proper 429 response.
 */
export class BatchRateLimitExceededError extends Error {
  constructor(
    public readonly limit: number,
    public readonly remaining: number,
    public readonly resetAt: Date,
    public readonly itemCount: number
  ) {
    super(`Batch rate limit exceeded. Limit resets at ${resetAt.toISOString()}`);
    this.name = "BatchRateLimitExceededError";
  }
}