feat: add rate limit management service

Introduce the RateLimitSystem class to handle overriding and resetting queue rate limits. This service updates the database state and synchronizes the changes with the Redis cache to ensure runtime rate limiting policies are enforced correctly.

Include unit tests to verify database interactions and side-effect calls to the queue service.

Author: aquiveal

apps/webapp/app/v3/services/rateLimitSystem.server.ts

@@ -0,0 +1,41 @@
+import { PrismaClient, Prisma } from "@trigger.dev/database";
+import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import { removeQueueRateLimits, updateQueueRateLimits } from "../runQueue.server";
+
+export class RateLimitSystem {
+  constructor(
+    private prisma: PrismaClient
+  ) {}
+
+  async overrideQueueRateLimit(
+    environment: AuthenticatedEnvironment,
+    queueName: string,
+    rateLimits: Array<{ limit: number; window: number }>
+  ) {
+    const queue = await this.prisma.taskQueue.updateMany({
+      where: {
+        runtimeEnvironmentId: environment.id,
+        name: queueName,
+      },
+      data: {
+        rateLimit: rateLimits,
+      },
+    });
+
+    await updateQueueRateLimits(environment, queueName, rateLimits);
+  }
+
+  async resetQueueRateLimit(environment: AuthenticatedEnvironment, queueName: string) {
+    await this.prisma.taskQueue.updateMany({
+      where: {
+        runtimeEnvironmentId: environment.id,
+        name: queueName,
+      },
+      data: {
+        rateLimit: Prisma.DbNull,
+      },
+    });
+
+    await removeQueueRateLimits(environment, queueName);
+  }
+}

apps/webapp/test/rateLimitSystem.server.test.ts

@@ -0,0 +1,82 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { RateLimitSystem } from "../app/v3/services/rateLimitSystem.server";
+import { PrismaClient, Prisma } from "@trigger.dev/database";
+import { Redis } from "ioredis";
+import { AuthenticatedEnvironment } from "../app/services/apiAuth.server";
+import * as runQueueServer from "../app/v3/runQueue.server";
+
+vi.mock("../app/v3/runQueue.server", () => ({
+  updateQueueRateLimits: vi.fn(),
+  removeQueueRateLimits: vi.fn(),
+}));
+
+describe("RateLimitSystem", () => {
+  let prismaMock: any;
+  let redisMock: any;
+  let rateLimitSystem: RateLimitSystem;
+  let mockEnvironment: AuthenticatedEnvironment;
+
+  beforeEach(() => {
+    prismaMock = {
+      taskQueue: {
+        updateMany: vi.fn().mockResolvedValue({ count: 1 }),
+      },
+    };
+
+    rateLimitSystem = new RateLimitSystem(prismaMock as unknown as PrismaClient);
+
+    mockEnvironment = {
+      id: "env-123",
+    } as AuthenticatedEnvironment;
+
+    vi.clearAllMocks();
+  });
+
+  describe("overrideQueueRateLimit", () => {
+    it("should update the rateLimit field in the database and call the Redis sync method", async () => {
+      const queueName = "test-queue";
+      const rateLimits = [{ limit: 10, window: 60 }];
+
+      await rateLimitSystem.overrideQueueRateLimit(mockEnvironment, queueName, rateLimits);
+
+      expect(prismaMock.taskQueue.updateMany).toHaveBeenCalledWith({
+        where: {
+          runtimeEnvironmentId: mockEnvironment.id,
+          name: queueName,
+        },
+        data: {
+          rateLimit: rateLimits,
+        },
+      });
+
+      expect(runQueueServer.updateQueueRateLimits).toHaveBeenCalledWith(
+        mockEnvironment,
+        queueName,
+        rateLimits
+      );
+    });
+  });
+
+  describe("resetQueueRateLimit", () => {
+    it("should clear the rateLimit field in the database and call the Redis sync method", async () => {
+      const queueName = "test-queue";
+
+      await rateLimitSystem.resetQueueRateLimit(mockEnvironment, queueName);
+
+      expect(prismaMock.taskQueue.updateMany).toHaveBeenCalledWith({
+        where: {
+          runtimeEnvironmentId: mockEnvironment.id,
+          name: queueName,
+        },
+        data: {
+          rateLimit: Prisma.DbNull,
+        },
+      });
+
+      expect(runQueueServer.removeQueueRateLimits).toHaveBeenCalledWith(
+        mockEnvironment,
+        queueName
+      );
+    });
+  });
+});