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

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

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+API rate limits now apply per environment, so creating extra API keys no longer increases how many requests an environment can make.

apps/webapp/app/models/runtimeEnvironment.server.ts

@@ -301,6 +301,89 @@ export async function findEnvironmentByApiKeyWithResolution(
   return resolveEnvironmentByApiKey(apiKey, branchName, tx, additionalApiKeyLookupEnabled);
 }
 
+export type PrivateApiKeyRateLimitScope = {
+  environmentId: string;
+  apiRateLimiterConfig: unknown;
+};
+
+export async function resolvePrivateApiKeyRateLimitScope(
+  apiKey: string,
+  tx: PrismaClientOrTransaction = $replica
+): Promise<PrivateApiKeyRateLimitScope | null> {
+  const now = new Date();
+
+  if (isAdditionalApiKey(apiKey)) {
+    const match = await tx.apiKey.findFirst({
+      where: {
+        keyHash: hashApiKey(apiKey),
+        revokedAt: null,
+        OR: [{ expiresAt: null }, { expiresAt: { gt: now } }],
+      },
+      select: {
+        runtimeEnvironment: {
+          select: {
+            id: true,
+            project: { select: { deletedAt: true } },
+            organization: { select: { apiRateLimiterConfig: true } },
+          },
+        },
+      },
+    });
+
+    if (!match?.runtimeEnvironment || match.runtimeEnvironment.project.deletedAt) {
+      return null;
+    }
+
+    return {
+      environmentId: match.runtimeEnvironment.id,
+      apiRateLimiterConfig: match.runtimeEnvironment.organization.apiRateLimiterConfig,
+    };
+  }
+
+  const environment = await tx.runtimeEnvironment.findFirst({
+    where: { apiKey },
+    select: {
+      id: true,
+      project: { select: { deletedAt: true } },
+      organization: { select: { apiRateLimiterConfig: true } },
+    },
+  });
+
+  if (environment) {
+    if (environment.project.deletedAt) {
+      return null;
+    }
+
+    return {
+      environmentId: environment.id,
+      apiRateLimiterConfig: environment.organization.apiRateLimiterConfig,
+    };
+  }
+
+  const revokedApiKey = await tx.revokedApiKey.findFirst({
+    where: { apiKey, expiresAt: { gt: now } },
+    select: {
+      runtimeEnvironment: {
+        select: {
+          id: true,
+          project: { select: { deletedAt: true } },
+          organization: { select: { apiRateLimiterConfig: true } },
+        },
+      },
+    },
+  });
+
+  const revokedEnvironment = revokedApiKey?.runtimeEnvironment;
+  if (!revokedEnvironment || revokedEnvironment.project.deletedAt) {
+    return null;
+  }
+
+  return {
+    environmentId: revokedEnvironment.id,
+    apiRateLimiterConfig: revokedEnvironment.organization.apiRateLimiterConfig,
+  };
+}
+
 /**
  * @deprecated We don't use public API keys (`pk_*` tokens) anymore — public
  * access goes through public JWTs (see `isPublicJWT` / `validatePublicJwtKey`).

apps/webapp/app/presenters/v3/LimitsPresenter.server.ts

@@ -1,6 +1,5 @@
 import { Ratelimit } from "@upstash/ratelimit";
 import type { RuntimeEnvironmentType } from "@trigger.dev/database";
-import { createHash } from "node:crypto";
 import { env } from "~/env.server";
 import { getCurrentPlan } from "~/services/platform.v3.server";
 import {
@@ -90,13 +89,11 @@ export class LimitsPresenter extends BasePresenter {
     projectId,
     environmentId,
     environmentType,
-    environmentApiKey,
   }: {
     organizationId: string;
     projectId: string;
     environmentId: string;
     environmentType: RuntimeEnvironmentType;
-    environmentApiKey: string;
   }): Promise<LimitsResult> {
     // Get organization with all limit-related fields
     const organization = await this._replica.organization.findFirstOrThrow({
@@ -168,10 +165,10 @@ export class LimitsPresenter extends BasePresenter {
       where: { organizationId },
     });
 
-    // Get current rate limit tokens for this environment's API key
+    // Get current rate limit tokens for this environment's API bucket
     const apiRateLimitTokens = await getRateLimitRemainingTokens(
       "api",
-      environmentApiKey,
+      environmentId,
       apiRateLimitConfig
     );
     // Batch rate limiter uses environment ID directly (not hashed) with a different key prefix
@@ -454,20 +451,14 @@ function resolveBatchConcurrencyConfig(batchConcurrencyConfig?: unknown): {
 
 /**
  * Query the current remaining tokens for a rate limiter using the Upstash getRemaining method.
- * This uses the same configuration and hashing logic as the rate limit middleware.
+ * The API limiter uses the environment ID as the bucket identifier for private API keys.
  */
 async function getRateLimitRemainingTokens(
   keyPrefix: string,
-  apiKey: string,
+  identifier: string,
   config: RateLimiterConfig
 ): Promise<number | null> {
   try {
-    // Hash the authorization header the same way the rate limiter does
-    const authorizationValue = `Bearer ${apiKey}`;
-    const hash = createHash("sha256");
-    hash.update(authorizationValue);
-    const hashedKey = hash.digest("hex");
-
     // Create a Ratelimit instance with the same configuration
     const limiter = createLimiterFromConfig(config);
     const ratelimit = new Ratelimit({
@@ -478,9 +469,9 @@ async function getRateLimitRemainingTokens(
       prefix: `ratelimit:${keyPrefix}`,
     });
 
-    // Use the getRemaining method to get the current remaining tokens
+    // Use the same identifier as the API rate-limit middleware.
     // getRemaining returns a Promise<number>
-    const remaining = await ratelimit.getRemaining(hashedKey);
+    const remaining = await ratelimit.getRemaining(identifier);
     return remaining;
   } catch (error) {
     logger.warn("Failed to get rate limit remaining tokens", {

apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx

@@ -82,9 +82,9 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
     presenter.call({
       organizationId: project.organizationId,
       projectId: project.id,
-      environmentId: environment.id,
+      // API traffic for a branch is bucketed on the parent environment id.
+      environmentId: environment.parentEnvironmentId ?? environment.id,
       environmentType: environment.type,
-      environmentApiKey: environment.apiKey,
     })
   );
 

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

@@ -1,11 +1,14 @@
 import { tryCatch } from "@trigger.dev/core/v3";
+import { trail } from "agentcrumbs"; // @crumbs
 import { env } from "~/env.server";
+import { resolvePrivateApiKeyRateLimitScope } from "~/models/runtimeEnvironment.server";
 import { batchStreamGrants } from "~/runEngine/concerns/batchStreamGrantsInstance.server";
 import { authenticateAuthorizationHeader } from "./apiAuth.server";
 import { authorizationRateLimitMiddleware } from "./authorizationRateLimitMiddleware.server";
 import type { Duration } from "./rateLimiter.server";
 
 const BATCH_STREAM_ITEMS_PATH = /^\/api\/v3\/batches\/([^/]+)\/items$/;
+const crumb = trail("webapp"); // @crumbs
 
 export const apiRateLimiter = authorizationRateLimitMiddleware({
   redis: {
@@ -29,6 +32,27 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
     maxItems: 1000,
   },
   limiterConfigOverride: async (authorizationValue) => {
+    const rawApiKey = authorizationValue.replace(/^Bearer /, "");
+
+    if (rawApiKey.startsWith("tr_")) {
+      const scope = await resolvePrivateApiKeyRateLimitScope(rawApiKey);
+
+      if (!scope) {
+        return;
+      }
+
+      // #region @crumbs
+      crumb("resolved private API key rate limit scope", {
+        environmentId: scope.environmentId,
+      });
+      // #endregion @crumbs
+
+      return {
+        config: scope.apiRateLimiterConfig,
+        identifier: scope.environmentId,
+      };
+    }
+
     const authenticatedEnv = await authenticateAuthorizationHeader(authorizationValue, {
       allowPublicKey: true,
       allowJWT: true,
@@ -40,13 +64,19 @@ 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;
     }
+
+    return {
+      config: authenticatedEnv.environment.organization.apiRateLimiterConfig,
+      // Public keys are browser-distributed, so keep them on per-key buckets.
+      identifier: authenticatedEnv.type === "PRIVATE" ? authenticatedEnv.environment.id : undefined,
+    };
   },
   pathMatchers: [/^\/api/],
   // Allow /api/v1/tasks/:id/callback/:secret

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

@@ -52,7 +52,14 @@ export const RateLimiterConfig = z.discriminatedUnion("type", [
 
 export type RateLimiterConfig = z.infer<typeof RateLimiterConfig>;
 
-type LimitConfigOverrideFunction = (authorizationValue: string) => Promise<unknown>;
+type RateLimitOverride = {
+  config?: unknown;
+  identifier?: string;
+};
+
+type LimitConfigOverrideFunction = (
+  authorizationValue: string
+) => Promise<RateLimitOverride | undefined>;
 
 type Options = {
   redis: RedisWithClusterOptions;
@@ -80,16 +87,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 +123,24 @@ async function resolveLimitConfig(
         });
       }
 
-      return defaultLimiter;
+      return { config: defaultLimiter } satisfies ResolvedRateLimit;
     }
 
-    const parsedOverride = RateLimiterConfig.safeParse(override);
+    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 +151,22 @@ async function resolveLimitConfig(
       });
     }
 
-    return parsedOverride.data;
+    return { config: parsedOverride.data, identifier } satisfies ResolvedRateLimit;
   });
 
-  return cacheResult.val ?? defaultLimiter;
+  // Defensive read: the cache is keyed on a shared Redis namespace, so during a
+  // deploy an entry could have been written by a server running a different
+  // code version (a different stored shape). Re-validate here so a stale/foreign
+  // entry can never reach createLimiterFromConfig with an undefined config and
+  // throw. The cache key is also versioned (see RedisCacheStore keyPrefix), so
+  // this is belt-and-suspenders.
+  const cached = cacheResult.val;
+  const parsedConfig = RateLimiterConfig.safeParse(cached?.config);
+
+  return {
+    config: parsedConfig.success ? parsedConfig.data : defaultLimiter,
+    identifier: typeof cached?.identifier === "string" ? cached.identifier : undefined,
+  };
 }
 
 /**
@@ -169,14 +200,17 @@ export function authorizationRateLimitMiddleware({
   const memory = createLRUMemoryStore(limiterCache?.maxItems ?? 1000);
   const redisCacheStore = new RedisCacheStore({
     connection: {
-      keyPrefix: `cache:${keyPrefix}:rate-limit-cache:`,
+      // Versioned namespace: the cached value shape is part of this key. Bump
+      // the version whenever ResolvedRateLimit changes so a rolling deploy never
+      // reads entries written in a previous shape (and vice versa).
+      keyPrefix: `cache:${keyPrefix}:rate-limit-cache:v2:`,
       ...redis,
     },
   });
 
   // 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 +303,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 +312,8 @@ export function authorizationRateLimitMiddleware({
       limiterConfigOverride
     );
 
+    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