chore: import upstream snapshot with attribution

2026-07-13 13:32:57 +08:00
commit cd420f9332
4811 changed files with 884702 additions and 0 deletions
@@ -0,0 +1,123 @@
+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";
+  }
+}