Skip to content

chore: migrate single rate limits to multi rate limits #3328

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chronark merged 9 commits into from
Jun 18, 2025
Merged

chore: migrate single rate limits to multi rate limits #3328

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
732e3cd
chore: remove ratelimit v1 entirely from dashboard
chronark Jun 13, 2025
b58fcfb
chore: migrate to multi-ratelimits
chronark Jun 16, 2025
8b2df8e
refactor: bind checkbox to form and add correct default value
ogzhanolguncu Jun 16, 2025
0f5d549
Merge branch 'main' into multi-ratelimit
chronark Jun 17, 2025
0fa7682
fix: pass autoApply through to the UI and to trpc
chronark Jun 17, 2025
b0d0113
[autofix.ci] apply automated fixes
autofix-ci[bot] Jun 17, 2025
85b45ab
[autofix.ci] apply automated fixes (attempt 2/3)
autofix-ci[bot] Jun 17, 2025
37b3919
[autofix.ci] apply automated fixes (attempt 3/3)
autofix-ci[bot] Jun 17, 2025
60ba33b
Merge branch 'main' into multi-ratelimit
chronark Jun 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion apps/api/src/pkg/cache/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ export function initCache(c: Context<HonoEnv>, metrics: Metrics): C<CacheNamespa
cloudflareApiKey: c.env.CLOUDFLARE_API_KEY,
zoneId: c.env.CLOUDFLARE_ZONE_ID,
domain: "cache.unkey.dev",
cacheBuster: "v8",
cacheBuster: "v9",
})
: undefined;

Expand Down
9 changes: 5 additions & 4 deletions apps/api/src/pkg/cache/namespaces.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ type CachedIdentity = Pick<Identity, "id" | "externalId" | "meta">;

export type CacheNamespaces = {
keyById: {
key: Key & { encrypted: EncryptedKey | null };
key: Key & { encrypted: EncryptedKey | null; ratelimits: Ratelimit[] };
api: Api;
permissions: string[];
roles: string[];
Expand All @@ -30,16 +30,16 @@ export type CacheNamespaces = {
id: string;
enabled: boolean;
} | null;
key: Key & { encrypted: EncryptedKey | null };
key: Key & { encrypted: EncryptedKey | null; ratelimits: Ratelimit[] };
api: Api;
permissions: string[];
roles: string[];
ratelimits: { [name: string]: Pick<Ratelimit, "name" | "limit" | "duration"> };
ratelimits: { [name: string]: Pick<Ratelimit, "name" | "limit" | "duration" | "autoApply"> };
identity: CachedIdentity | null;
Comment on lines 31 to 38
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Jun 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick (assertive)

Duplicate ratelimit representations increase cache size

Both key.ratelimits: Ratelimit[] and a flattened ratelimits map live in the same entry.
Unless consumers need both shapes, keeping one reduces memory and invalidation surface.

Evaluate dropping the redundant field or deriving it on demand.

🤖 Prompt for AI Agents 
In apps/api/src/pkg/cache/namespaces.ts between lines 31 and 38, the cache entry
contains duplicate ratelimit data in both key.ratelimits (an array) and a
separate ratelimits map, which unnecessarily increases cache size. To fix this,
remove one of these representations—either drop the key.ratelimits array or the
ratelimits map—and adjust the code to derive the removed structure on demand if
needed by consumers, thereby reducing memory usage and simplifying cache
invalidation.

} | null;
apiById: (Api & { keyAuth: KeyAuth | null }) | null;
keysByOwnerId: {
key: Key & { encrypted: EncryptedKey | null };
key: Key & { encrypted: EncryptedKey | null; ratelimits: Ratelimit[] };
api: Api;
}[];
verificationsByKeyId: {
Expand All @@ -58,6 +58,7 @@ export type CacheNamespaces = {
permissions: string[];
roles: string[];
identity: CachedIdentity | null;
ratelimits: Ratelimit[];
}
>;
total: number;
Expand Down
13 changes: 10 additions & 3 deletions apps/api/src/pkg/key_migration/handler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,11 +83,18 @@ export async function migrateKey(
refillDay: message.refill?.refillDay,
enabled: message.enabled,
remaining: message.remaining,
ratelimitAsync: message.ratelimit?.async,
ratelimitLimit: message.ratelimit?.limit,
ratelimitDuration: message.ratelimit?.duration,
environment: message.environment,
});
if (message.ratelimit) {
await tx.insert(schema.ratelimits).values({
id: newId("ratelimit"),
workspaceId: message.workspaceId,
keyId: keyId,
limit: message.ratelimit.limit,
duration: message.ratelimit.duration,
name: "default",
});
}

if (message.encrypted) {
await tx.insert(schema.encryptedKeys).values({
Expand Down
42 changes: 18 additions & 24 deletions apps/api/src/pkg/keys/service.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -287,19 +287,10 @@ export class KeyService {
* Merge ratelimits from the identity and the key
* Key limits take pecedence
*/
const ratelimits: { [name: string]: Pick<Ratelimit, "name" | "limit" | "duration"> } = {};

if (
dbRes.ratelimitAsync !== null &&
dbRes.ratelimitDuration !== null &&
dbRes.ratelimitLimit !== null
) {
ratelimits.default = {
name: "default",
limit: dbRes.ratelimitLimit,
duration: dbRes.ratelimitDuration,
};
}
const ratelimits: {
[name: string]: Pick<Ratelimit, "name" | "limit" | "duration" | "autoApply">;
} = {};

for (const rl of dbRes.identity?.ratelimits ?? []) {
ratelimits[rl.name] = rl;
}
Expand Down Expand Up @@ -553,16 +544,19 @@ export class KeyService {
*/

const ratelimits: {
[name: string | "default"]: Required<RatelimitRequest>;
[name: string]: Required<RatelimitRequest>;
} = {};
if (data.ratelimits && "default" in data.ratelimits && typeof req.ratelimits === "undefined") {
ratelimits.default = {
identity: data.key.id,
name: data.ratelimits.default.name,
cost: req.ratelimit?.cost ?? DEFAULT_RATELIMIT_COST,
limit: data.ratelimits.default.limit,
duration: data.ratelimits.default.duration,
};

for (const rl of Object.values(data.ratelimits)) {
if (rl.autoApply) {
ratelimits[rl.name] = {
identity: data.identity?.id ?? data.key.id,
name: rl.name,
cost: DEFAULT_RATELIMIT_COST,
limit: rl.limit,
duration: rl.duration,
};
}
}

for (const r of req.ratelimits ?? []) {
Expand Down Expand Up @@ -681,7 +675,7 @@ export class KeyService {
private async ratelimit(
c: Context,
key: Key,
ratelimits: { [name: string | "default"]: Required<RatelimitRequest> },
ratelimits: { [name: string]: Required<RatelimitRequest> },
): Promise<[boolean, VerifyKeyResult["ratelimit"]]> {
if (Object.keys(ratelimits).length === 0) {
return [true, undefined];
Expand All @@ -695,7 +689,7 @@ export class KeyService {
c,
Object.values(ratelimits).map((r) => ({
name: r.name,
async: !!key.ratelimitAsync,
async: false,
workspaceId: key.workspaceId,
identifier: r.identity,
cost: r.cost,
Expand Down
44 changes: 24 additions & 20 deletions apps/api/src/routes/legacy_apis_listKeys.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,9 @@ export const registerLegacyApisListKeys = (app: App) =>
limit: limit,
orderBy: schema.keys.id,
offset: offset ? offset : undefined,
with: {
ratelimits: true,
},
});
metrics.emit({
metric: "metric.db.read",
Expand All @@ -130,29 +133,30 @@ export const registerLegacyApisListKeys = (app: App) =>
.where(and(eq(schema.keys.keyAuthId, api.keyAuthId), isNull(schema.keys.deletedAtM)));

return c.json({
keys: keys.map((k) => ({
id: k.id,
start: k.start,
apiId: api.id,
workspaceId: k.workspaceId,
name: k.name ?? undefined,
ownerId: k.ownerId ?? undefined,
meta: k.meta ? JSON.parse(k.meta) : undefined,
createdAt: k.createdAtM ?? undefined,
expires: k.expires?.getTime() ?? undefined,
ratelimit:
k.ratelimitAsync !== null && k.ratelimitLimit !== null && k.ratelimitDuration !== null
keys: keys.map((k) => {
const ratelimit = k.ratelimits.find((rl) => rl.name === "default");
return {
id: k.id,
start: k.start,
apiId: api.id,
workspaceId: k.workspaceId,
name: k.name ?? undefined,
ownerId: k.ownerId ?? undefined,
meta: k.meta ? JSON.parse(k.meta) : undefined,
createdAt: k.createdAtM ?? undefined,
expires: k.expires?.getTime() ?? undefined,
ratelimit: ratelimit
? {
async: k.ratelimitAsync,
type: k.ratelimitAsync ? "fast" : ("consistent" as unknown),
limit: k.ratelimitLimit,
duration: k.ratelimitDuration,
refillRate: k.ratelimitLimit,
refillInterval: k.ratelimitDuration,
async: false,
limit: ratelimit.limit,
duration: ratelimit.duration,
refillRate: ratelimit.limit,
refillInterval: ratelimit.duration,
}
: undefined,
remaining: k.remaining ?? undefined,
})),
remaining: k.remaining ?? undefined,
};
}),
// @ts-ignore, mysql sucks
total: Number.parseInt(total.at(0)?.count ?? "0"),
cursor: keys.at(-1)?.id ?? undefined,
Expand Down
14 changes: 11 additions & 3 deletions apps/api/src/routes/legacy_keys_createKey.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,12 +221,20 @@ export const registerLegacyKeysCreate = (app: App) =>
forWorkspaceId: null,
expires: req.expires ? new Date(req.expires) : null,
createdAtM: Date.now(),
ratelimitLimit: req.ratelimit?.limit,
ratelimitDuration: req.ratelimit?.refillRate,
ratelimitAsync: req.ratelimit?.type === "fast",
remaining: req.remaining,
deletedAtM: null,
});
if (req.ratelimit) {
await tx.insert(schema.ratelimits).values({
id: newId("ratelimit"),
workspaceId: authorizedWorkspaceId,
name: "default",
limit: req.ratelimit.limit,
duration: req.ratelimit.refillRate,
keyId: keyId,
});
}

await insertUnkeyAuditLog(c, tx, {
workspaceId: authorizedWorkspaceId,
actor: { type: "key", id: rootKeyId },
Expand Down
85 changes: 45 additions & 40 deletions apps/api/src/routes/v1_apis_listKeys.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,7 @@ export const registerV1ApisListKeys = (app: App) =>
permission: true,
},
},
ratelimits: true,
},
limit: limit,
orderBy: schema.keys.id,
Expand Down Expand Up @@ -233,6 +234,7 @@ export const registerV1ApisListKeys = (app: App) =>
: null,
permissions: Array.from(permissions.values()),
roles: k.roles.map((r) => r.role.name),
ratelimits: k.ratelimits,
};
}),
total: keySpace.sizeApprox,
Expand Down Expand Up @@ -293,50 +295,53 @@ export const registerV1ApisListKeys = (app: App) =>
}),
);
}

return c.json({
keys: data.keys.map((k) => ({
id: k.id,
start: k.start,
apiId: api.id,
workspaceId: k.workspaceId,
name: k.name ?? undefined,
ownerId: k.ownerId ?? undefined,
meta: k.meta ? JSON.parse(k.meta) : undefined,
createdAt: k.createdAtM ?? undefined,
updatedAt: k.updatedAtM ?? undefined,
expires: k.expires?.getTime() ?? undefined,
ratelimit:
k.ratelimitAsync !== null && k.ratelimitLimit !== null && k.ratelimitDuration !== null
keys: data.keys.map((k) => {
const ratelimit = k.ratelimits.find((rl) => rl.name === "default");
return {
id: k.id,
start: k.start,
apiId: api.id,
workspaceId: k.workspaceId,
name: k.name ?? undefined,
ownerId: k.ownerId ?? undefined,
meta: k.meta ? JSON.parse(k.meta) : undefined,
createdAt: k.createdAtM ?? undefined,
updatedAt: k.updatedAtM ?? undefined,
expires: k.expires?.getTime() ?? undefined,
ratelimit: ratelimit
? {
async: false,
limit: ratelimit.limit,
duration: ratelimit.duration,
refillRate: ratelimit.limit,
refillInterval: ratelimit.duration,
}
: undefined,
Comment on lines +314 to +321
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Jun 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick (assertive)

Avoid legacy alias duplication

refillRate and refillInterval mirror limit/duration.
Carrying both inflates the payload and may mislead API consumers that both sets are independently authoritative.

Consider deprecating the aliases in the response or gating them behind a version flag.

🤖 Prompt for AI Agents 
In apps/api/src/routes/v1_apis_listKeys.ts around lines 314 to 321, the response
includes both `refillRate`/`refillInterval` and `limit`/`duration` which are
duplicates. Remove `refillRate` and `refillInterval` from the response object to
avoid redundancy and potential confusion for API consumers. If backward
compatibility is needed, gate these aliases behind a version flag instead of
always including them.


remaining: k.remaining ?? undefined,
refill: k.refillAmount
? {
async: k.ratelimitAsync,
type: k.ratelimitAsync ? "fast" : ("consistent" as unknown),
limit: k.ratelimitLimit,
duration: k.ratelimitDuration,
refillRate: k.ratelimitLimit,
refillInterval: k.ratelimitDuration,
interval: k.refillDay ? ("monthly" as const) : ("daily" as const),
amount: k.refillAmount,
refillDay: k.refillDay,
lastRefillAt: k.lastRefillAt?.getTime(),
}
: undefined,
remaining: k.remaining ?? undefined,
refill: k.refillAmount
? {
interval: k.refillDay ? ("monthly" as const) : ("daily" as const),
amount: k.refillAmount,
refillDay: k.refillDay,
lastRefillAt: k.lastRefillAt?.getTime(),
}
: undefined,
environment: k.environment ?? undefined,
plaintext: plaintext[k.id] ?? undefined,
roles: k.roles,
permissions: k.permissions,
identity: k.identity
? {
id: k.identity.id,
externalId: k.identity.externalId,
meta: k.identity.meta ?? undefined,
}
: undefined,
})),
environment: k.environment ?? undefined,
plaintext: plaintext[k.id] ?? undefined,
roles: k.roles,
permissions: k.permissions,
identity: k.identity
? {
id: k.identity.id,
externalId: k.identity.externalId,
meta: k.identity.meta ?? undefined,
}
: undefined,
};
}),
total: data.total,
cursor: data.keys.at(-1)?.id ?? undefined,
});
Expand Down
16 changes: 13 additions & 3 deletions apps/api/src/routes/v1_keys_createKey.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -373,9 +373,6 @@ export const registerV1KeysCreateKey = (app: App) =>
expires: req.expires ? new Date(req.expires) : null,
createdAtM: Date.now(),
updatedAtM: null,
ratelimitAsync: req.ratelimit?.async ?? req.ratelimit?.type === "fast",
ratelimitLimit: req.ratelimit?.limit ?? req.ratelimit?.refillRate,
ratelimitDuration: req.ratelimit?.duration ?? req.ratelimit?.refillInterval,
remaining: req.remaining,
refillDay: req.refill?.interval === "daily" ? null : (req?.refill?.refillDay ?? 1),
refillAmount: req.refill?.amount,
Expand All @@ -384,6 +381,19 @@ export const registerV1KeysCreateKey = (app: App) =>
environment: req.environment ?? null,
identityId: identity?.id,
});
if (req.ratelimit) {
await db.primary.insert(schema.ratelimits).values({
id: newId("ratelimit"),
keyId: kId,
limit: req.ratelimit.limit ?? req.ratelimit.refillRate!,
duration: req.ratelimit.duration ?? req.ratelimit.refillInterval!,
workspaceId: authorizedWorkspaceId,
name: "default",
autoApply: true,
identityId: null,
createdAt: Date.now(),
});
}
Comment on lines +384 to +396
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Jun 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Insert + ratelimit must be atomic & validated

  1. The key row is inserted before the ratelimit row and the two calls are not wrapped in a transaction.
    • If the second insert fails (e.g. DB constraint, network hiccup) the system will end up with a key that has no ratelimit, violating the invariant introduced by this PR.
  2. limit / duration silently fall back to refillRate / refillInterval, but nothing guarantees that either pair is supplied – undefined will happily be inserted and blow up at runtime.
- await db.primary.insert(schema.keys).values({ … });
- if (req.ratelimit) {
-   await db.primary.insert(schema.ratelimits).values({ … });
- }
+ await db.primary.transaction(async (tx) => {
+   await tx.insert(schema.keys).values({ … });
+   if (req.ratelimit) {
+     if (
+       req.ratelimit.limit === undefined &&
+       req.ratelimit.refillRate === undefined
+     ) {
+       throw new UnkeyApiError({
+         code: "BAD_REQUEST",
+         message: "`ratelimit.limit` is required",
+       });
+     }
+     if (
+       req.ratelimit.duration === undefined &&
+       req.ratelimit.refillInterval === undefined
+     ) {
+       throw new UnkeyApiError({
+         code: "BAD_REQUEST",
+         message: "`ratelimit.duration` is required",
+       });
+     }
+     await tx.insert(schema.ratelimits).values({ … });
+   }
+ });

This keeps the DB consistent and surfaces invalid payloads early.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if (req.ratelimit) {
await db.primary.insert(schema.ratelimits).values({
id: newId("ratelimit"),
keyId: kId,
limit: req.ratelimit.limit ?? req.ratelimit.refillRate!,
duration: req.ratelimit.duration ?? req.ratelimit.refillInterval!,
workspaceId: authorizedWorkspaceId,
name: "default",
autoApply: true,
identityId: null,
createdAt: Date.now(),
});
}
// Replace the separate key insert + ratelimit insert with a single transaction:
await db.primary.transaction(async (tx) => {
// Insert the key row
await tx.insert(schema.keys).values({
// … (all the original key fields, e.g. id: kId, workspaceId, etc.)
});
// If a ratelimit was requested, validate and insert it
if (req.ratelimit) {
if (
req.ratelimit.limit === undefined &&
req.ratelimit.refillRate === undefined
) {
throw new UnkeyApiError({
code: "BAD_REQUEST",
message: "`ratelimit.limit` is required",
});
}
if (
req.ratelimit.duration === undefined &&
req.ratelimit.refillInterval === undefined
) {
throw new UnkeyApiError({
code: "BAD_REQUEST",
message: "`ratelimit.duration` is required",
});
}
await tx.insert(schema.ratelimits).values({
id: newId("ratelimit"),
keyId: kId,
limit: req.ratelimit.limit ?? req.ratelimit.refillRate!,
duration: req.ratelimit.duration ?? req.ratelimit.refillInterval!,
workspaceId: authorizedWorkspaceId,
name: "default",
autoApply: true,
identityId: null,
createdAt: Date.now(),
});
}
});
🤖 Prompt for AI Agents 
In apps/api/src/routes/v1_keys_createKey.ts around lines 384 to 396, the
insertion of the ratelimit row is not wrapped in a transaction with the key
insertion, risking inconsistent state if the second insert fails. Also, the
limit and duration fields fall back silently without validation, potentially
inserting undefined values. To fix this, wrap both inserts in a single
transaction to ensure atomicity, and add explicit validation to confirm that
either limit/duration or refillRate/refillInterval are provided before
inserting, throwing an error if neither is present.


if (req.recoverable && api.keyAuth?.storeEncryptedKeys) {
const perm = rbac.evaluatePermissions(
Expand Down
2 changes: 2 additions & 0 deletions apps/api/src/routes/v1_keys_deleteKey.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ export const registerV1KeysDeleteKey = (app: App) =>
api: true,
},
},
ratelimits: true,
},
Comment on lines +84 to 85
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Jun 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick (assertive)

Loading ratelimits eagerly is OK but might be overkill

For a delete operation we never touch ratelimit rows. Fetching them adds an extra join with no functional gain and bloats the cache entry.
Optional: drop ratelimits: true from the with clause.

🤖 Prompt for AI Agents 
In apps/api/src/routes/v1_keys_deleteKey.ts at lines 84 to 85, the code eagerly
loads ratelimits with the delete operation, which is unnecessary since ratelimit
rows are not used here. To fix this, remove the `ratelimits: true` entry from
the `with` clause to avoid the extra join and reduce cache bloat.

});
if (!dbRes) {
Expand All @@ -98,6 +99,7 @@ export const registerV1KeysDeleteKey = (app: App) =>
meta: dbRes.identity.meta,
}
: null,
ratelimits: dbRes.ratelimits,
};
});

Expand Down
Loading
Loading