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

## What

Rate-limit the API by **environment** rather than per API key.

Previously the limiter keyed its bucket on the hash of the full
`Authorization` header — one bucket per key. With additional environment
API keys (`tr_*_sk_*`), an environment can mint many keys and each got
its own full bucket, so more keys = higher effective rate limit. This
collapses all of an environment's keys onto a single shared
per-environment bucket, so the ceiling is exactly the configured limit
regardless of key mix.

## How

- `authorizationRateLimitMiddleware` now lets the override return `{
config?, identifier? }`. `identifier`, when present, is the rate limit
bucket key; otherwise it falls back to the hashed `Authorization` header
(unchanged legacy behavior, still used by `engineRateLimiter` and any
unauthenticated fallthrough).
- `apiRateLimiter`'s override resolves the environment id and uses it as
the identifier:
- **Additional keys** (`isAdditionalApiKey`) resolve via a new
`resolveAdditionalApiKeyRateLimitScope()` — a **scope-agnostic** keyHash
→ (environmentId, org limiter config) lookup. It is deliberately
permissive (restricted keys resolve too) because it's used **only for
bucketing, never as an auth decision** — request auth still goes through
the RBAC bearer controller, which enforces scopes. Revoked/expired keys
are excluded so they can't hold a bucket warm.
- **Root/legacy keys** reuse the environment already resolved by
`authenticateAuthorizationHeader` and key on `environment.id` too.
- The identifier is always the stable environment id, never the secret
key (which can rotate and would split the bucket).
- The whole override result is cached per key by the existing SWR cache,
so **no extra per-request lookup and no separate Redis mapping** is
added.

## Behavior notes

- Root + additional keys of the same environment now share one bucket
(ceiling = configured limit, not a multiple of it). Restricted
additional keys are included — they were the biggest gap, since they
authenticate via the RBAC controller and previously fell back to per-key
buckets.
- **Public JWTs** keep their existing fixed-window, per-token bucketing.
- One-time bucket reset on deploy (bucket keys change); harmless.

## Tests

- New: two tokens resolving to the same identifier share one bucket.
- New: with no identifier, bucketing stays per-key (legacy behavior
preserved).
- Updated existing override tests to the new `{ config }` return shape.

Base: `feat/multi-keys-surface`. Closes TRI-12888.
This commit is contained in:
Chris Arderne
2026-08-06 16:05:27 +01:00
committed by GitHub
parent 9409ddf9bc
commit 088f68b373
9 changed files with 324 additions and 63 deletions
@@ -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.
-10
View File
@@ -1,6 +1,5 @@
import type { PrismaClient, RuntimeEnvironment } from "@trigger.dev/database";
import type { HostRbacController } from "@trigger.dev/rbac";
import { trail } from "agentcrumbs"; // @crumbs
import { customAlphabet } from "nanoid";
import { MAX_API_KEY_TASK_IDENTIFIERS } from "~/consts";
import { prisma } from "~/db.server";
@@ -11,8 +10,6 @@ import { rbac } from "~/services/rbac.server";
import { generateAdditionalApiKey, generateRootApiKey } from "~/utils/apiKeys";
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
const crumb = trail("webapp"); // @crumbs
const apiKeyId = customAlphabet(
"1234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ",
12
@@ -220,12 +217,6 @@ export async function createEnvironmentApiKey(
})();
telemetryRecorder.recordOperation("create", "success");
crumb("environment API key created", {
apiKeyId: apiKey.id,
environmentId,
presetId: apiKey.presetId,
}); // @crumbs
return { apiKey, plaintext: generated.apiKey };
}
@@ -267,7 +258,6 @@ export async function revokeEnvironmentApiKey(
}
telemetryRecorder.recordOperation("revoke", "success");
crumb("environment API key revoked", { apiKeyId, environmentId }); // @crumbs
}
export function createApiKeyForEnv(envType: RuntimeEnvironment["type"]) {
@@ -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`).
@@ -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,21 @@ export class LimitsPresenter extends BasePresenter {
where: { organizationId },
});
// Get current rate limit tokens for this environment's API key
const runtimeEnv = await this._replica.runtimeEnvironment.findFirst({
where: { id: environmentId },
select: {
id: true,
parentEnvironmentId: true,
maximumConcurrencyLimit: true,
concurrencyLimitBurstFactor: true,
},
});
const apiRateLimitEnvironmentId = runtimeEnv?.parentEnvironmentId ?? environmentId;
// Get current rate limit tokens for this environment's API bucket
const apiRateLimitTokens = await getRateLimitRemainingTokens(
"api",
environmentApiKey,
apiRateLimitEnvironmentId,
apiRateLimitConfig
);
// Batch rate limiter uses environment ID directly (not hashed) with a different key prefix
@@ -181,15 +189,6 @@ export class LimitsPresenter extends BasePresenter {
);
// Get current queue size for this environment
// We need the runtime environment fields for the engine query
const runtimeEnv = await this._replica.runtimeEnvironment.findFirst({
where: { id: environmentId },
select: {
id: true,
maximumConcurrencyLimit: true,
concurrencyLimitBurstFactor: true,
},
});
let currentQueueSize = 0;
if (runtimeEnv) {
@@ -454,20 +453,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 +471,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", {
@@ -78,7 +78,6 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
projectId: project.id,
environmentId: environment.id,
environmentType: environment.type,
environmentApiKey: environment.apiKey,
})
);
@@ -1,5 +1,6 @@
import { tryCatch } from "@trigger.dev/core/v3";
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";
@@ -29,6 +30,21 @@ 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;
}
return {
config: scope.apiRateLimiterConfig,
identifier: scope.environmentId,
};
}
const authenticatedEnv = await authenticateAuthorizationHeader(authorizationValue, {
allowPublicKey: true,
allowJWT: true,
@@ -40,13 +56,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
@@ -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,10 +123,16 @@ 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", {
@@ -121,7 +140,7 @@ async function resolveLimitConfig(
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
@@ -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,75 @@ 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 +446,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;
@@ -1,7 +1,10 @@
import { postgresTest } from "@internal/testcontainers";
import { type PrismaClient } from "@trigger.dev/database";
import { describe, expect, it, vi } from "vitest";
import { findEnvironmentByApiKey } from "~/models/runtimeEnvironment.server";
import {
findEnvironmentByApiKey,
resolvePrivateApiKeyRateLimitScope,
} from "~/models/runtimeEnvironment.server";
import { generateAdditionalApiKey, hashApiKey } from "~/utils/apiKeys";
import { createTestOrgProjectWithMember, uniqueId } from "./fixtures/environmentVariablesFixtures";
@@ -143,6 +146,36 @@ describe("findEnvironmentByApiKey — PREVIEW (regression guard)", () => {
expect(resolved?.apiKey).toBe(previewParent.apiKey);
}
);
postgresTest(
"rate limit scope resolves root and additional keys to the preview parent",
async ({ prisma }) => {
const { organization, project, user } = await createTestOrgProjectWithMember(prisma);
const previewParent = await createEnv(prisma, project.id, organization.id, {
type: "PREVIEW",
isBranchableEnvironment: true,
});
const additional = generateAdditionalApiKey("PREVIEW").apiKey;
await prisma.apiKey.create({
data: {
name: "Preview integration",
keyHash: hashApiKey(additional),
lastFour: additional.slice(-4),
runtimeEnvironmentId: previewParent.id,
createdByUserId: user.id,
presetId: null,
scopes: ["admin"],
},
});
const rootScope = await resolvePrivateApiKeyRateLimitScope(previewParent.apiKey, prisma);
const additionalScope = await resolvePrivateApiKeyRateLimitScope(additional, prisma);
expect(rootScope?.environmentId).toBe(previewParent.id);
expect(additionalScope?.environmentId).toBe(previewParent.id);
}
);
});
describe("findEnvironmentByApiKey — non-branchable", () => {
@@ -366,4 +399,30 @@ describe("findEnvironmentByApiKey — additional and disabled keys", () => {
).resolves.toMatchObject({ id: environment.id });
}
);
postgresTest("does not resolve additional keys for deleted projects", async ({ prisma }) => {
const { organization, project, user } = await createTestOrgProjectWithMember(prisma);
const environment = await createEnv(prisma, project.id, organization.id, {
type: "PRODUCTION",
});
const additional = generateAdditionalApiKey("PRODUCTION").apiKey;
await prisma.apiKey.create({
data: {
name: "Deleted project key",
keyHash: hashApiKey(additional),
lastFour: additional.slice(-4),
runtimeEnvironmentId: environment.id,
createdByUserId: user.id,
presetId: null,
scopes: ["admin"],
},
});
await prisma.project.update({
where: { id: project.id },
data: { deletedAt: new Date() },
});
await expect(resolvePrivateApiKeyRateLimitScope(additional, prisma)).resolves.toBeNull();
});
});