feat(webapp): share rate limit bucket across additional API keys per environment

Author: carderne

.server-changes/additional-api-key-rate-limit-bucket.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Additional environment API keys now share a single rate limit bucket per environment instead of each key getting its own, so minting more keys no longer multiplies an environment's effective API rate limit. Root and legacy keys are unaffected. The bucket is keyed on the stable environment id resolved from the already-authenticated request, so no extra lookup is added.

apps/webapp/app/services/apiRateLimit.server.ts

@@ -1,4 +1,5 @@
 import { tryCatch } from "@trigger.dev/core/v3";
+import { isAdditionalApiKey } from "@trigger.dev/core/v3/apiKeys";
 import { env } from "~/env.server";
 import { batchStreamGrants } from "~/runEngine/concerns/batchStreamGrantsInstance.server";
 import { authenticateAuthorizationHeader } from "./apiAuth.server";
@@ -40,13 +41,32 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
 
     if (authenticatedEnv.type === "PUBLIC_JWT") {
       return {
-        type: "fixedWindow",
-        window: env.API_RATE_LIMIT_JWT_WINDOW,
-        tokens: env.API_RATE_LIMIT_JWT_TOKENS,
+        config: {
+          type: "fixedWindow",
+          window: env.API_RATE_LIMIT_JWT_WINDOW,
+          tokens: env.API_RATE_LIMIT_JWT_TOKENS,
+        },
       };
-    } else {
-      return authenticatedEnv.environment.organization.apiRateLimiterConfig;
     }
+
+    // Additional API keys (`tr_*_sk_*`) share their environment's rate limit
+    // bucket rather than each getting their own. Keying on the stable
+    // environment id (not the secret key, which can rotate) keeps a single
+    // bucket per environment no matter how many additional keys exist.
+    // Root/legacy keys fall through to the default per-key (hashed header)
+    // bucketing: they already map 1:1 to an environment.
+    //
+    // This reuses the already-authenticated environment above; the whole
+    // override result is cached per key by the middleware's SWR cache, so no
+    // separate lookup or Redis mapping is needed for the identifier.
+    const apiKey = authenticatedEnv.apiKey;
+    const identifier =
+      apiKey && isAdditionalApiKey(apiKey) ? authenticatedEnv.environment.id : undefined;
+
+    return {
+      config: authenticatedEnv.environment.organization.apiRateLimiterConfig,
+      identifier,
+    };
   },
   pathMatchers: [/^\/api/],
   // Allow /api/v1/tasks/:id/callback/:secret

apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts

@@ -52,7 +52,24 @@ export const RateLimiterConfig = z.discriminatedUnion("type", [
 
 export type RateLimiterConfig = z.infer<typeof RateLimiterConfig>;
 
-type LimitConfigOverrideFunction = (authorizationValue: string) => Promise<unknown>;
+/**
+ * Result of an override lookup for a given Authorization header.
+ *
+ * - `config`: the rate limiter configuration to apply (bucket size). When
+ *   absent, the default limiter is used.
+ * - `identifier`: the value to key the rate limit bucket on. When absent, the
+ *   hashed Authorization header is used (the legacy per-key behavior). Supply a
+ *   stable value (e.g. an environment id) so multiple credentials that should
+ *   share a bucket collapse onto one. Never a secret: it lands in Redis keys.
+ */
+type RateLimitOverride = {
+  config?: unknown;
+  identifier?: string;
+};
+
+type LimitConfigOverrideFunction = (
+  authorizationValue: string
+) => Promise<RateLimitOverride | undefined>;
 
 type Options = {
   redis: RedisWithClusterOptions;
@@ -80,16 +97,22 @@ type Options = {
   };
 };
 
-async function resolveLimitConfig(
+type ResolvedRateLimit = {
+  config: RateLimiterConfig;
+  // Bucket key to use, or undefined to fall back to the hashed Authorization header.
+  identifier?: string;
+};
+
+async function resolveRateLimit(
   authorizationValue: string,
   hashedAuthorizationValue: string,
   defaultLimiter: RateLimiterConfig,
-  cache: UnkeyCache<{ limiter: RateLimiterConfig }>,
+  cache: UnkeyCache<{ limiter: ResolvedRateLimit }>,
   logsEnabled: boolean,
   limiterConfigOverride?: LimitConfigOverrideFunction
-): Promise<RateLimiterConfig> {
+): Promise<ResolvedRateLimit> {
   if (!limiterConfigOverride) {
-    return defaultLimiter;
+    return { config: defaultLimiter };
   }
 
   if (logsEnabled) {
@@ -110,18 +133,26 @@ async function resolveLimitConfig(
         });
       }
 
-      return defaultLimiter;
+      return { config: defaultLimiter } satisfies ResolvedRateLimit;
     }
 
-    const parsedOverride = RateLimiterConfig.safeParse(override);
+    // The identifier (if any) is trusted through even when the config falls back
+    // to the default: bucketing and bucket size are independent concerns.
+    const identifier = override.identifier;
+
+    if (!override.config) {
+      return { config: defaultLimiter, identifier } satisfies ResolvedRateLimit;
+    }
+
+    const parsedOverride = RateLimiterConfig.safeParse(override.config);
 
     if (!parsedOverride.success) {
       logger.error("Error parsing rate limiter override", {
         override,
         errors: parsedOverride.error.errors,
       });
 
-      return defaultLimiter;
+      return { config: defaultLimiter, identifier } satisfies ResolvedRateLimit;
     }
 
     if (logsEnabled && parsedOverride.data) {
@@ -132,10 +163,10 @@ async function resolveLimitConfig(
       });
     }
 
-    return parsedOverride.data;
+    return { config: parsedOverride.data, identifier } satisfies ResolvedRateLimit;
   });
 
-  return cacheResult.val ?? defaultLimiter;
+  return cacheResult.val ?? { config: defaultLimiter };
 }
 
 /**
@@ -176,7 +207,7 @@ export function authorizationRateLimitMiddleware({
 
   // This cache holds the rate limit configuration for each org, so we don't have to fetch it every request
   const cache = createCache({
-    limiter: new Namespace<RateLimiterConfig>(ctx, {
+    limiter: new Namespace<ResolvedRateLimit>(ctx, {
       stores: [memory, redisCacheStore],
       fresh: limiterCache?.fresh ?? 30_000,
       stale: limiterCache?.stale ?? 60_000,
@@ -269,7 +300,7 @@ export function authorizationRateLimitMiddleware({
     hash.update(authorizationValue);
     const hashedAuthorizationValue = hash.digest("hex");
 
-    const limiterConfig = await resolveLimitConfig(
+    const { config: limiterConfig, identifier } = await resolveRateLimit(
       authorizationValue,
       hashedAuthorizationValue,
       defaultLimiter,
@@ -278,6 +309,11 @@ export function authorizationRateLimitMiddleware({
       limiterConfigOverride
     );
 
+    // Bucket key: an override-supplied identifier (e.g. environment id, so all
+    // additional API keys for an environment share one bucket) or, by default,
+    // the hashed Authorization header (legacy per-key behavior).
+    const rateLimitIdentifier = identifier ?? hashedAuthorizationValue;
+
     const limiter = createLimiterFromConfig(limiterConfig);
 
     const rateLimiter = new RateLimiter({
@@ -288,7 +324,7 @@ export function authorizationRateLimitMiddleware({
       logFailure: log.rejections,
     });
 
-    const { success, limit, reset, remaining } = await rateLimiter.limit(hashedAuthorizationValue);
+    const { success, limit, reset, remaining } = await rateLimiter.limit(rateLimitIdentifier);
 
     const $remaining = Math.max(0, remaining); // remaining can be negative if the user has exceeded the limit, so clamp it to 0
 

apps/webapp/test/authorizationRateLimitMiddleware.test.ts

@@ -150,10 +150,12 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware",
         limiterConfigOverride: async (authorizationValue) => {
           if (authorizationValue === "Bearer premium-token") {
             return {
-              type: "tokenBucket",
-              refillRate: 10,
-              interval: "1m",
-              maxTokens: 100,
+              config: {
+                type: "tokenBucket",
+                refillRate: 10,
+                interval: "1m",
+                maxTokens: 100,
+              },
             };
           }
           return undefined;
@@ -184,6 +186,80 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware",
     }
   );
 
+  redisTest(
+    "should share a bucket across tokens that resolve to the same identifier",
+    async ({ redisOptions }) => {
+      const rateLimitMiddleware = authorizationRateLimitMiddleware({
+        redis: { ...redisOptions, tlsDisabled: true },
+        keyPrefix: "test-identifier",
+        defaultLimiter: {
+          type: "tokenBucket",
+          refillRate: 1,
+          interval: "1m",
+          maxTokens: 1,
+        },
+        pathMatchers: [/^\/api/],
+        // Both tokens map to the same environment identifier, so they should
+        // consume from a single shared bucket rather than one bucket each.
+        limiterConfigOverride: async () => ({ identifier: "env_shared" }),
+      });
+
+      app.use(rateLimitMiddleware);
+      app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" }));
+
+      // First token uses the single token in the shared bucket.
+      const first = await request(app)
+        .get("/api/test")
+        .set("Authorization", "Bearer tr_prod_sk_aaaaaaaaaaaaaaaaaaaaaaaa");
+      expect(first.status).toBe(200);
+
+      // A different token that resolves to the same identifier is limited,
+      // because the bucket is shared rather than per-key.
+      const second = await request(app)
+        .get("/api/test")
+        .set("Authorization", "Bearer tr_prod_sk_bbbbbbbbbbbbbbbbbbbbbbbb");
+      expect(second.status).toBe(429);
+    }
+  );
+
+  redisTest(
+    "should key per token when no identifier is supplied",
+    async ({ redisOptions }) => {
+      const rateLimitMiddleware = authorizationRateLimitMiddleware({
+        redis: { ...redisOptions, tlsDisabled: true },
+        keyPrefix: "test-no-identifier",
+        defaultLimiter: {
+          type: "tokenBucket",
+          refillRate: 1,
+          interval: "1m",
+          maxTokens: 1,
+        },
+        pathMatchers: [/^\/api/],
+        // Override supplies a config but no identifier: bucketing stays per-key
+        // (hashed Authorization header), the legacy behavior.
+        limiterConfigOverride: async () => ({
+          config: { type: "tokenBucket", refillRate: 1, interval: "1m", maxTokens: 1 },
+        }),
+      });
+
+      app.use(rateLimitMiddleware);
+      app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" }));
+
+      const first = await request(app).get("/api/test").set("Authorization", "Bearer token-a");
+      expect(first.status).toBe(200);
+
+      // Same token is limited...
+      const firstAgain = await request(app)
+        .get("/api/test")
+        .set("Authorization", "Bearer token-a");
+      expect(firstAgain.status).toBe(429);
+
+      // ...but a different token gets its own bucket.
+      const second = await request(app).get("/api/test").set("Authorization", "Bearer token-b");
+      expect(second.status).toBe(200);
+    }
+  );
+
   describe("Advanced Cases", () => {
     // 1. Test different rate limit configurations
     redisTest("should enforce fixed window rate limiting", async ({ redisOptions }) => {
@@ -375,10 +451,12 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware",
           configOverrideCalls++;
           if (authorizationValue === "Bearer premium-token") {
             return {
-              type: "tokenBucket",
-              refillRate: 10,
-              interval: "1m",
-              maxTokens: 100,
+              config: {
+                type: "tokenBucket",
+                refillRate: 10,
+                interval: "1m",
+                maxTokens: 100,
+              },
             };
           }
           return undefined;