use new resolver for api keys

Author: carderne

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

@@ -301,57 +301,85 @@ export async function findEnvironmentByApiKeyWithResolution(
   return resolveEnvironmentByApiKey(apiKey, branchName, tx, additionalApiKeyLookupEnabled);
 }
 
-export type AdditionalApiKeyRateLimitScope = {
+export type PrivateApiKeyRateLimitScope = {
   environmentId: string;
-  // Organization rate limiter override (bucket size), if configured.
   apiRateLimiterConfig: unknown;
 };
 
-/**
- * Resolve ONLY the environment id (and its organization's rate limiter config)
- * for an additional API key, for RATE-LIMIT BUCKETING.
- *
- * Deliberately scope-agnostic: unlike `findEnvironmentByApiKey`, a
- * scope-restricted additional key still resolves here, so every additional key
- * for an environment shares that environment's rate limit bucket. This is NOT
- * an authentication or authorization decision and must never be used as one —
- * request auth still goes through the RBAC bearer controller, which enforces
- * scopes. Revoked and expired keys are excluded so they cannot keep a bucket
- * warm.
- */
-export async function resolveAdditionalApiKeyRateLimitScope(
+export async function resolvePrivateApiKeyRateLimitScope(
   apiKey: string,
   tx: PrismaClientOrTransaction = $replica
-): Promise<AdditionalApiKeyRateLimitScope | null> {
-  if (!isAdditionalApiKey(apiKey)) {
-    return null;
-  }
-
+): Promise<PrivateApiKeyRateLimitScope | null> {
   const now = new Date();
 
-  const match = await tx.apiKey.findFirst({
-    where: {
-      keyHash: hashApiKey(apiKey),
-      revokedAt: null,
-      OR: [{ expiresAt: null }, { expiresAt: { gt: now } }],
+  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,
+            organization: { select: { apiRateLimiterConfig: true } },
+          },
+        },
+      },
+    });
+
+    if (!match?.runtimeEnvironment) {
+      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 } },
         },
       },
     },
   });
 
-  if (!match?.runtimeEnvironment) {
+  const revokedEnvironment = revokedApiKey?.runtimeEnvironment;
+  if (!revokedEnvironment || revokedEnvironment.project.deletedAt) {
     return null;
   }
 
   return {
-    environmentId: match.runtimeEnvironment.id,
-    apiRateLimiterConfig: match.runtimeEnvironment.organization.apiRateLimiterConfig,
+    environmentId: revokedEnvironment.id,
+    apiRateLimiterConfig: revokedEnvironment.organization.apiRateLimiterConfig,
   };
 }
 

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

@@ -1,13 +1,14 @@
 import { tryCatch } from "@trigger.dev/core/v3";
-import { isAdditionalApiKey } from "@trigger.dev/core/v3/apiKeys";
+import { trail } from "agentcrumbs"; // @crumbs
 import { env } from "~/env.server";
-import { resolveAdditionalApiKeyRateLimitScope } from "~/models/runtimeEnvironment.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: {
@@ -33,26 +34,19 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
   limiterConfigOverride: async (authorizationValue) => {
     const rawApiKey = authorizationValue.replace(/^Bearer /, "");
 
-    // 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.
-    //
-    // Resolve scope-agnostically for bucketing: a restricted additional key
-    // authenticates at the route level via the RBAC controller (and fails
-    // closed in the legacy header auth below), but for rate limiting it must
-    // still land on its environment's shared bucket — otherwise minting many
-    // restricted keys would multiply the effective limit. This is NOT an auth
-    // decision. The whole override result is cached per key by the
-    // middleware's SWR cache, so no separate lookup or Redis mapping is needed.
-    if (isAdditionalApiKey(rawApiKey)) {
-      const scope = await resolveAdditionalApiKeyRateLimitScope(rawApiKey);
-
-      // Unknown/revoked/expired key: fall back to the default per-key bucket.
+    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,
@@ -78,11 +72,6 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
       };
     }
 
-    // Root/legacy keys also bucket per environment, so an environment's ceiling
-    // is exactly its configured limit regardless of key mix (root + additional
-    // keys share one bucket). The environment is already resolved above, so this
-    // adds no lookup. JWTs intentionally stay on per-token bucketing (handled
-    // above).
     return {
       config: authenticatedEnv.environment.organization.apiRateLimiterConfig,
       identifier: authenticatedEnv.environment.id,

apps/webapp/test/findEnvironmentByApiKey.test.ts

@@ -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", () => {