From 5aa219a06fa2937455bc1288aff6b49eedf351a8 Mon Sep 17 00:00:00 2001 From: Flo Date: Tue, 28 Oct 2025 13:04:04 +0100 Subject: [PATCH 1/3] feat: add identity-based credits support to API v1 - Update cache to support credit-based invalidation - Enhance usage limiter to work with both key and credit IDs - Update all v1 key endpoints to support new credits table - Add comprehensive tests for credit operations (create, update, verify) - Implement double-read strategy (credits table takes precedence over legacy remaining) - Support refill configuration in credits table --- apps/api/src/pkg/cache/index.ts | 2 +- apps/api/src/pkg/cache/namespaces.ts | 8 +- apps/api/src/pkg/keys/service.ts | 31 +- apps/api/src/pkg/usagelimit/client.ts | 32 +- apps/api/src/pkg/usagelimit/durable_object.ts | 130 ++++++- apps/api/src/pkg/usagelimit/interface.ts | 13 +- .../src/routes/v1_apis_listKeys.happy.test.ts | 157 ++++++++ apps/api/src/routes/v1_apis_listKeys.ts | 26 +- .../routes/v1_keys_createKey.happy.test.ts | 160 +++++++- apps/api/src/routes/v1_keys_createKey.ts | 37 +- .../routes/v1_keys_deleteKey.happy.test.ts | 65 ++++ apps/api/src/routes/v1_keys_deleteKey.ts | 12 +- .../src/routes/v1_keys_getKey.happy.test.ts | 121 ++++++ apps/api/src/routes/v1_keys_getKey.ts | 28 +- .../src/routes/v1_keys_getVerifications.ts | 2 + .../routes/v1_keys_updateKey.happy.test.ts | 288 ++++++++++++++- apps/api/src/routes/v1_keys_updateKey.ts | 103 +++++- .../v1_keys_updateRemaining.happy.test.ts | 251 ++++++++++++- .../api/src/routes/v1_keys_updateRemaining.ts | 120 ++++-- apps/api/src/routes/v1_keys_verifyKey.test.ts | 349 ++++++++++++++++++ .../src/routes/v1_keys_whoami.happy.test.ts | 134 +++++++ apps/api/src/routes/v1_keys_whoami.ts | 100 ++++- .../api/src/routes/v1_migrations_createKey.ts | 27 +- 23 files changed, 2090 insertions(+), 106 deletions(-) diff --git a/apps/api/src/pkg/cache/index.ts b/apps/api/src/pkg/cache/index.ts index 1260505a2b..6df3e80e8d 100644 --- a/apps/api/src/pkg/cache/index.ts +++ b/apps/api/src/pkg/cache/index.ts @@ -35,7 +35,7 @@ export function initCache(c: Context, metrics: Metrics): C }; + ratelimits: { + [name: string]: Pick; + }; identity: CachedIdentity | null; } | null; apiById: (Api & { keyAuth: KeyAuth | null }) | null; @@ -59,6 +64,7 @@ export type CacheNamespaces = { roles: string[]; identity: CachedIdentity | null; ratelimits: Ratelimit[]; + credits: Credits | null; } >; total: number; diff --git a/apps/api/src/pkg/keys/service.ts b/apps/api/src/pkg/keys/service.ts index f0499a91e1..d6938ac7fa 100644 --- a/apps/api/src/pkg/keys/service.ts +++ b/apps/api/src/pkg/keys/service.ts @@ -211,6 +211,7 @@ export class KeyService { where: (table, { and, eq, isNull }) => and(eq(table.hash, hash), isNull(table.deletedAtM)), with: { encrypted: true, + credits: true, workspace: { columns: { id: true, @@ -278,7 +279,7 @@ export class KeyService { } /** - * Createa a unique set of all permissions, whether they're attached directly or connected + * Create a unique set of all permissions, whether they're attached directly or connected * through a role. */ const permissions = new Set([ @@ -290,7 +291,7 @@ export class KeyService { /** * Merge ratelimits from the identity and the key - * Key limits take pecedence + * Key limits take precedence */ const ratelimits: { [name: string]: Pick; @@ -299,11 +300,13 @@ export class KeyService { for (const rl of dbRes.identity?.ratelimits ?? []) { ratelimits[rl.name] = rl; } + for (const rl of dbRes.ratelimits ?? []) { ratelimits[rl.name] = rl; } return { + credits: dbRes.credits, workspace: dbRes.workspace, forWorkspace: dbRes.forWorkspace, key: dbRes, @@ -320,6 +323,7 @@ export class KeyService { if (cached) { return cached; } + const hash = await sha256(key); this.hashCache.set(key, hash); return hash; @@ -616,11 +620,30 @@ export class KeyService { } let remaining: number | undefined = undefined; - if (data.key.remaining !== null) { + if (data.key.remaining !== null || data.credits != null) { const t0 = performance.now(); const cost = req.remaining?.cost ?? DEFAULT_REMAINING_COST; + + // V1 Migration: Priority order for credits + // 1. Key credits (credits table with keyId) + // 2. key.remaining (legacy system) + const keyCredits = data.credits; + const legacyRemaining = data.key.remaining; + + let creditId: string | undefined = undefined; + let keyId: string | undefined = undefined; + + if (keyCredits != null) { + // Use key-based credits (credits table) + creditId = keyCredits.id; + } else if (legacyRemaining !== null) { + // Use legacy key.remaining + keyId = data.key.id; + } + const limited = await this.usageLimiter.limit({ - keyId: data.key.id, + creditId, + keyId, cost, }); diff --git a/apps/api/src/pkg/usagelimit/client.ts b/apps/api/src/pkg/usagelimit/client.ts index fcff086593..fd1a996632 100644 --- a/apps/api/src/pkg/usagelimit/client.ts +++ b/apps/api/src/pkg/usagelimit/client.ts @@ -36,20 +36,31 @@ export class DurableUsageLimiter implements UsageLimiter { public async limit(req: LimitRequest): Promise { const start = performance.now(); + // Use creditId if present (new system), otherwise fall back to keyId (legacy) + const identifier = req.creditId ?? req.keyId; + if (!identifier) { + this.logger.error("usagelimit called without keyId or creditId"); + return { valid: false }; + } + try { const url = `https://${this.domain}/limit`; - const res = await this.getStub(req.keyId) + const res = await this.getStub(identifier) .fetch(url, { method: "POST", - headers: { "Content-Type": "application/json", "Unkey-Request-Id": this.requestId }, + headers: { + "Content-Type": "application/json", + "Unkey-Request-Id": this.requestId, + }, body: JSON.stringify(req), }) .catch(async (e) => { this.logger.warn("calling the usagelimit DO failed, retrying ...", { keyId: req.keyId, + creditId: req.creditId, error: (e as Error).message, }); - return await this.getStub(req.keyId).fetch(url, { + return await this.getStub(identifier).fetch(url, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(req), @@ -57,19 +68,30 @@ export class DurableUsageLimiter implements UsageLimiter { }); return limitResponseSchema.parse(await res.json()); } catch (e) { - this.logger.error("usagelimit failed", { keyId: req.keyId, error: (e as Error).message }); + this.logger.error("usagelimit failed", { + keyId: req.keyId, + creditId: req.creditId, + error: (e as Error).message, + }); return { valid: false }; } finally { this.metrics.emit({ metric: "metric.usagelimit", latency: performance.now() - start, keyId: req.keyId, + creditId: req.creditId, }); } } public async revalidate(req: RevalidateRequest): Promise { - const obj = this.namespace.get(this.namespace.idFromName(req.keyId)); + const identifier = req.creditId ?? req.keyId; + if (!identifier) { + this.logger.error("revalidate called without keyId or creditId"); + return; + } + + const obj = this.namespace.get(this.namespace.idFromName(identifier)); const url = `https://${this.domain}/revalidate`; await obj.fetch(url, { method: "POST", diff --git a/apps/api/src/pkg/usagelimit/durable_object.ts b/apps/api/src/pkg/usagelimit/durable_object.ts index 7606dd8345..805e5baff3 100644 --- a/apps/api/src/pkg/usagelimit/durable_object.ts +++ b/apps/api/src/pkg/usagelimit/durable_object.ts @@ -1,4 +1,14 @@ -import { type Database, type Key, and, createConnection, eq, gt, schema, sql } from "@/pkg/db"; +import { + type Credits, + type Database, + type Key, + and, + createConnection, + eq, + gt, + schema, + sql, +} from "@/pkg/db"; import { ConsoleLogger } from "@unkey/worker-logging"; import type { Env } from "../env"; import { limitRequestSchema, revalidateRequestSchema } from "./interface"; @@ -8,6 +18,7 @@ export class DurableObjectUsagelimiter implements DurableObject { private readonly db: Database; private lastRevalidate = 0; private key: Key | undefined = undefined; + private credit: Credits | undefined = undefined; private readonly env: Env; constructor(state: DurableObjectState, env: Env) { @@ -37,30 +48,119 @@ export class DurableObjectUsagelimiter implements DurableObject { environment: this.env.ENVIRONMENT, }, }); + const url = new URL(request.url); switch (url.pathname) { case "/revalidate": { const req = revalidateRequestSchema.parse(await request.json()); + const creditId = req.creditId; + const keyId = req.keyId; - this.key = await this.db.query.keys.findFirst({ - where: (table, { and, eq, isNull }) => - and(eq(table.id, req.keyId), isNull(table.deletedAtM)), - }); + if (creditId) { + // New credits system + this.credit = await this.db.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + } else if (keyId) { + // Legacy key-based system + this.key = await this.db.query.keys.findFirst({ + where: (table, { and, eq, isNull }) => + and(eq(table.id, keyId), isNull(table.deletedAtM)), + }); + } this.lastRevalidate = Date.now(); return Response.json({}); } case "/limit": { const req = limitRequestSchema.parse(await request.json()); + + // Determine if we're using new credits system or legacy key-based system + const creditId = req.creditId; + const keyId = req.keyId; + + if (creditId) { + // New credits table system + if (!this.credit) { + this.credit = await this.db.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + this.lastRevalidate = Date.now(); + } + + if (!this.credit) { + logger.error("credit not found", { creditId }); + return Response.json({ + valid: false, + }); + } + + const currentRemaining = this.credit.remaining; + + if (currentRemaining < req.cost && req.cost !== 0) { + return Response.json({ + valid: false, + remaining: Math.max(0, currentRemaining), + }); + } + + this.credit.remaining = Math.max(0, currentRemaining - req.cost); + + this.state.waitUntil( + this.db + .update(schema.credits) + .set({ remaining: sql`${schema.credits.remaining}-${req.cost}` }) + .where( + and( + eq(schema.credits.id, this.credit.id), + gt(schema.credits.remaining, 0), // prevent negative remaining + ), + ) + .execute(), + ); + + // revalidate every minute + if (Date.now() - this.lastRevalidate > 60_000) { + logger.info("revalidating in the background", { + creditId: this.credit.id, + }); + + this.state.waitUntil( + this.db.query.credits + .findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }) + .execute() + .then((credit) => { + this.credit = credit; + this.lastRevalidate = Date.now(); + }), + ); + } + + return Response.json({ + valid: true, + remaining: this.credit.remaining, + }); + } + + if (!keyId) { + logger.error("Neither creditId nor keyId provided"); + return Response.json({ + valid: false, + }); + } + + // Legacy key-based system if (!this.key) { this.key = await this.db.query.keys.findFirst({ where: (table, { and, eq, isNull }) => - and(eq(table.id, req.keyId), isNull(table.deletedAtM)), + and(eq(table.id, keyId), isNull(table.deletedAtM)), }); this.lastRevalidate = Date.now(); } if (!this.key) { - logger.error("key not found", { keyId: req.keyId }); + logger.error("key not found", { keyId }); return Response.json({ valid: false, }); @@ -75,14 +175,16 @@ export class DurableObjectUsagelimiter implements DurableObject { }); } - if (this.key.remaining <= 0 && req.cost !== 0) { + const currentRemaining = this.key.remaining; + + if (currentRemaining < req.cost && req.cost !== 0) { return Response.json({ valid: false, - remaining: 0, + remaining: Math.max(0, currentRemaining), }); } - this.key.remaining = Math.max(0, this.key.remaining - req.cost); + this.key.remaining = Math.max(0, currentRemaining - req.cost); this.state.waitUntil( this.db @@ -96,14 +198,18 @@ export class DurableObjectUsagelimiter implements DurableObject { ) .execute(), ); + // revalidate every minute if (Date.now() - this.lastRevalidate > 60_000) { - logger.info("revalidating in the background", { keyId: this.key.id }); + logger.info("revalidating in the background", { + keyId: this.key.id, + }); + this.state.waitUntil( this.db.query.keys .findFirst({ where: (table, { and, eq, isNull }) => - and(eq(table.id, req.keyId), isNull(table.deletedAtM)), + and(eq(table.id, keyId), isNull(table.deletedAtM)), }) .execute() .then((key) => { diff --git a/apps/api/src/pkg/usagelimit/interface.ts b/apps/api/src/pkg/usagelimit/interface.ts index 5493e0efdb..dbee99757d 100644 --- a/apps/api/src/pkg/usagelimit/interface.ts +++ b/apps/api/src/pkg/usagelimit/interface.ts @@ -1,7 +1,15 @@ import { z } from "zod"; export const limitRequestSchema = z.object({ - keyId: z.string(), + /** + * For legacy key-based credits stored in keys.remaining + */ + keyId: z.string().optional(), + /** + * For new credits system stored in credits table. + * When present, this takes precedence over keyId. + */ + creditId: z.string().optional(), cost: z.number(), }); export type LimitRequest = z.infer; @@ -13,7 +21,8 @@ export const limitResponseSchema = z.object({ export type LimitResponse = z.infer; export const revalidateRequestSchema = z.object({ - keyId: z.string(), + keyId: z.string().optional(), + creditId: z.string().optional(), }); export type RevalidateRequest = z.infer; diff --git a/apps/api/src/routes/v1_apis_listKeys.happy.test.ts b/apps/api/src/routes/v1_apis_listKeys.happy.test.ts index 36a1f9522f..df76a8d06d 100644 --- a/apps/api/src/routes/v1_apis_listKeys.happy.test.ts +++ b/apps/api/src/routes/v1_apis_listKeys.happy.test.ts @@ -404,3 +404,160 @@ test("retrieves a key in plain text", async (t) => { expect(listKeysRes.status).toBe(200); expect(listKeysRes.body.keys.at(0)?.plaintext).toEqual(key); }); + +test("returns credits from new credits table", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + + // Create key + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + }); + + // Create credits in new table + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 1000, + refillAmount: 500, + refillDay: 15, // monthly refill + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + updatedAt: null, + }); + + await revalidateKeyCount(h.db.primary, h.resources.userKeyAuth.id); + + const root = await h.createRootKey([ + `api.${h.resources.userApi.id}.read_api`, + `api.${h.resources.userApi.id}.read_key`, + ]); + + const res = await h.get({ + url: `/v1/apis.listKeys?apiId=${h.resources.userApi.id}`, + headers: { + Authorization: `Bearer ${root.key}`, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + const foundKey = res.body.keys.find((k) => k.id === keyId); + expect(foundKey).toBeDefined(); + expect(foundKey!.remaining).toBe(1000); + expect(foundKey!.refill).toBeDefined(); + expect(foundKey!.refill!.interval).toBe("monthly"); + expect(foundKey!.refill!.amount).toBe(500); + expect(foundKey!.refill!.refillDay).toBe(15); + expect(foundKey!.refill!.lastRefillAt).toBeDefined(); +}); + +test("returns credits from legacy remaining field when no credits table entry", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + + // Create key with legacy credits fields + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + remaining: 750, + refillAmount: 100, + refillDay: null, // daily refill + }); + + await revalidateKeyCount(h.db.primary, h.resources.userKeyAuth.id); + + const root = await h.createRootKey([ + `api.${h.resources.userApi.id}.read_api`, + `api.${h.resources.userApi.id}.read_key`, + ]); + + const res = await h.get({ + url: `/v1/apis.listKeys?apiId=${h.resources.userApi.id}`, + headers: { + Authorization: `Bearer ${root.key}`, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + const foundKey = res.body.keys.find((k) => k.id === keyId); + expect(foundKey).toBeDefined(); + expect(foundKey!.remaining).toBe(750); + expect(foundKey!.refill).toBeDefined(); + expect(foundKey!.refill!.interval).toBe("daily"); + expect(foundKey!.refill!.amount).toBe(100); +}); + +test("prefers new credits table over legacy remaining field", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + + // Create key with legacy credits fields + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + remaining: 999, // This should be ignored + refillAmount: 99, // This should be ignored + }); + + // Create credits in new table (should take precedence) + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 2000, + refillAmount: 1000, + refillDay: null, // daily refill + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + updatedAt: null, + }); + + await revalidateKeyCount(h.db.primary, h.resources.userKeyAuth.id); + + const root = await h.createRootKey([ + `api.${h.resources.userApi.id}.read_api`, + `api.${h.resources.userApi.id}.read_key`, + ]); + + const res = await h.get({ + url: `/v1/apis.listKeys?apiId=${h.resources.userApi.id}`, + headers: { + Authorization: `Bearer ${root.key}`, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + const foundKey = res.body.keys.find((k) => k.id === keyId); + expect(foundKey).toBeDefined(); + // Should use values from credits table, not legacy fields + expect(foundKey!.remaining).toBe(2000); + expect(foundKey!.refill).toBeDefined(); + expect(foundKey!.refill!.amount).toBe(1000); +}); diff --git a/apps/api/src/routes/v1_apis_listKeys.ts b/apps/api/src/routes/v1_apis_listKeys.ts index ddc71a6ab0..50c50980f9 100644 --- a/apps/api/src/routes/v1_apis_listKeys.ts +++ b/apps/api/src/routes/v1_apis_listKeys.ts @@ -189,6 +189,7 @@ export const registerV1ApisListKeys = (app: App) => with: { identity: true, encrypted: true, + credits: true, roles: { with: { role: { @@ -216,7 +217,10 @@ export const registerV1ApisListKeys = (app: App) => }); if (!keySpace) { - throw new UnkeyApiError({ code: "NOT_FOUND", message: "keyspace not found" }); + throw new UnkeyApiError({ + code: "NOT_FOUND", + message: "keyspace not found", + }); } /** @@ -230,6 +234,7 @@ export const registerV1ApisListKeys = (app: App) => ...k.permissions.map((p) => p.permission.slug), ...k.roles.flatMap((r) => r.role.permissions.map((p) => p.permission.slug)), ]); + return { ...k, identity: k.identity @@ -242,6 +247,7 @@ export const registerV1ApisListKeys = (app: App) => permissions: Array.from(permissions.values()), roles: k.roles.map((r) => r.role.name), ratelimits: k.ratelimits, + credits: k.credits ?? null, }; }), total: keySpace.sizeApprox, @@ -306,6 +312,11 @@ export const registerV1ApisListKeys = (app: App) => return c.json({ keys: data.keys.map((k) => { const ratelimit = k.ratelimits.find((rl) => rl.name === "default"); + const remaining = k.credits?.remaining ?? k.remaining ?? undefined; + const refillAmount = k.credits?.refillAmount ?? k.refillAmount ?? undefined; + const refillDay = k.credits?.refillDay ?? k.refillDay ?? null; + const lastRefillAt = k.credits?.refilledAt ?? k.lastRefillAt?.getTime(); + return { id: k.id, start: k.start, @@ -326,14 +337,13 @@ export const registerV1ApisListKeys = (app: App) => refillInterval: ratelimit.duration, } : undefined, - - remaining: k.remaining ?? undefined, - refill: k.refillAmount + remaining, + refill: refillAmount ? { - interval: k.refillDay ? ("monthly" as const) : ("daily" as const), - amount: k.refillAmount, - refillDay: k.refillDay, - lastRefillAt: k.lastRefillAt?.getTime(), + interval: refillDay ? ("monthly" as const) : ("daily" as const), + amount: refillAmount, + refillDay, + lastRefillAt: lastRefillAt, } : undefined, environment: k.environment ?? undefined, diff --git a/apps/api/src/routes/v1_keys_createKey.happy.test.ts b/apps/api/src/routes/v1_keys_createKey.happy.test.ts index 28bea35fac..e453339d08 100644 --- a/apps/api/src/routes/v1_keys_createKey.happy.test.ts +++ b/apps/api/src/routes/v1_keys_createKey.happy.test.ts @@ -8,6 +8,7 @@ import { IntegrationHarness } from "src/pkg/testutil/integration-harness"; import { KeyV1 } from "@unkey/keys"; import type { V1KeysCreateKeyRequest, V1KeysCreateKeyResponse } from "./v1_keys_createKey"; +import type { V1KeysVerifyKeyRequest, V1KeysVerifyKeyResponse } from "./v1_keys_verifyKey"; test("creates key", async (t) => { const h = await IntegrationHarness.init(t); @@ -648,9 +649,166 @@ describe("with externalId", () => { const key = await h.db.primary.query.keys.findFirst({ where: (table, { eq }) => eq(table.id, res.body.keyId), + with: { + credits: true, + }, }); expect(key).toBeDefined(); - expect(key!.refillDay).toEqual(1); + // Since key has remaining, it uses the new credits table + expect(key!.credits).toBeDefined(); + expect(key!.credits!.refillDay).toEqual(1); }); }); }); + +describe("creates key with new credits table", () => { + test("creates credits table entry when remaining is specified", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey([`api.${h.resources.userApi.id}.create_key`]); + + const res = await h.post({ + url: "/v1/keys.createKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + apiId: h.resources.userApi.id, + remaining: 100, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + const key = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, res.body.keyId), + with: { + credits: true, + }, + }); + + expect(key).toBeDefined(); + expect(key!.remaining).toBeNull(); + + // Verify credits table entry was created + expect(key!.credits).toBeDefined(); + expect(key!.credits!.remaining).toEqual(100); + expect(key!.credits!.keyId).toEqual(res.body.keyId); + expect(key!.credits!.workspaceId).toEqual(h.resources.userWorkspace.id); + expect(key!.credits!.identityId).toBeNull(); + }); + + test("creates credits with refill configuration", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey([`api.${h.resources.userApi.id}.create_key`]); + + const res = await h.post({ + url: "/v1/keys.createKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + apiId: h.resources.userApi.id, + remaining: 50, + refill: { + interval: "monthly", + amount: 100, + refillDay: 15, + }, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + const key = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, res.body.keyId), + with: { + credits: true, + }, + }); + + expect(key).toBeDefined(); + expect(key!.credits).toBeDefined(); + expect(key!.credits!.remaining).toEqual(50); + expect(key!.credits!.refillAmount).toEqual(100); + expect(key!.credits!.refillDay).toEqual(15); + expect(key!.credits!.refilledAt).toBeDefined(); + }); + + test("does not create credits table entry when remaining is not specified", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey([`api.${h.resources.userApi.id}.create_key`]); + + const res = await h.post({ + url: "/v1/keys.createKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + apiId: h.resources.userApi.id, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + const key = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, res.body.keyId), + with: { + credits: true, + }, + }); + + expect(key).toBeDefined(); + expect(key!.remaining).toBeNull(); + expect(key!.credits).toBeNull(); + }); + + test("creates key with remaining and can verify it", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey([`api.${h.resources.userApi.id}.create_key`]); + + const res = await h.post({ + url: "/v1/keys.createKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + apiId: h.resources.userApi.id, + remaining: 5, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + // Verify we can use the key + const verifyRes = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key: res.body.key, + apiId: h.resources.userApi.id, + }, + }); + + expect(verifyRes.status, `expected 200, received: ${JSON.stringify(verifyRes, null, 2)}`).toBe( + 200, + ); + expect(verifyRes.body.valid).toBe(true); + expect(verifyRes.body.remaining).toEqual(4); + + // Verify the credits table was updated + const key = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, res.body.keyId), + with: { + credits: true, + }, + }); + + expect(key!.credits!.remaining).toEqual(4); + }); +}); diff --git a/apps/api/src/routes/v1_keys_createKey.ts b/apps/api/src/routes/v1_keys_createKey.ts index 3a8372a674..2f04422837 100644 --- a/apps/api/src/routes/v1_keys_createKey.ts +++ b/apps/api/src/routes/v1_keys_createKey.ts @@ -318,12 +318,14 @@ export const registerV1KeysCreateKey = (app: App) => message: "remaining must be greater than 0.", }); } + if ((req.remaining === null || req.remaining === undefined) && req.refill?.interval) { throw new UnkeyApiError({ code: "BAD_REQUEST", message: "remaining must be set if you are using refill.", }); } + if (req.refill?.refillDay && req.refill.interval === "daily") { throw new UnkeyApiError({ code: "BAD_REQUEST", @@ -336,7 +338,6 @@ export const registerV1KeysCreateKey = (app: App) => const authorizedWorkspaceId = auth.authorizedWorkspaceId; const rootKeyId = auth.key.id; - const externalId = req.externalId ?? req.ownerId; const [permissionIds, roleIds, identity] = await Promise.all([ @@ -377,14 +378,30 @@ export const registerV1KeysCreateKey = (app: App) => expires: req.expires ? new Date(req.expires) : null, createdAtM: Date.now(), updatedAtM: null, - remaining: req.remaining, - refillDay: req.refill?.interval === "daily" ? null : (req?.refill?.refillDay ?? 1), - refillAmount: req.refill?.amount, - lastRefillAt: req.refill?.interval ? new Date() : null, enabled: req.enabled, environment: req.environment ?? null, identityId: identity?.id, }); + + if (req.remaining) { + await db.primary.insert(schema.credits).values({ + id: newId("credit"), + remaining: req.remaining, + workspaceId: authorizedWorkspaceId, + createdAt: Date.now(), + keyId: kId, + identityId: null, + refillAmount: req.refill?.amount, + refillDay: req.refill + ? req.refill.interval === "daily" + ? null + : (req.refill.refillDay ?? 1) + : null, + refilledAt: req.refill?.interval ? Date.now() : null, + updatedAt: null, + }); + } + if (req.ratelimit) { await db.primary.insert(schema.ratelimits).values({ id: newId("ratelimit"), @@ -582,7 +599,10 @@ async function getPermissionIds( }); } if (!val.valid) { - throw new UnkeyApiError({ code: "INSUFFICIENT_PERMISSIONS", message: val.message }); + throw new UnkeyApiError({ + code: "INSUFFICIENT_PERMISSIONS", + message: val.message, + }); } const missingPermissionSlugs = permissionsSlugs.filter( @@ -634,7 +654,10 @@ async function getRoleIds( }); } if (!val.valid) { - throw new UnkeyApiError({ code: "INSUFFICIENT_PERMISSIONS", message: val.message }); + throw new UnkeyApiError({ + code: "INSUFFICIENT_PERMISSIONS", + message: val.message, + }); } const missingRoles = roleNames.filter((name) => !roles.some((role) => role.name === name)); diff --git a/apps/api/src/routes/v1_keys_deleteKey.happy.test.ts b/apps/api/src/routes/v1_keys_deleteKey.happy.test.ts index a7ed444001..e6e3fe793e 100644 --- a/apps/api/src/routes/v1_keys_deleteKey.happy.test.ts +++ b/apps/api/src/routes/v1_keys_deleteKey.happy.test.ts @@ -77,3 +77,68 @@ test("hard deletes key", async (t) => { }); expect(found).toBeUndefined(); }); + +test("permanent delete removes credits from credits table", async (t) => { + const h = await IntegrationHarness.init(t); + const keyId = newId("test"); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + + // Create key + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + }); + + // Create credits in new table + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 1000, + refillAmount: null, + refillDay: null, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + updatedAt: null, + }); + + // Verify credits exist + const creditsBefore = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.keyId, keyId), + }); + expect(creditsBefore).toBeDefined(); + expect(creditsBefore!.id).toBe(creditId); + + const root = await h.createRootKey([`api.${h.resources.userApi.id}.delete_key`]); + const res = await h.post({ + url: "/v1/keys.deleteKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId, + permanent: true, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + // Verify key is deleted + const foundKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(foundKey).toBeUndefined(); + + // Verify credits are also deleted + const creditsAfter = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.keyId, keyId), + }); + expect(creditsAfter).toBeUndefined(); +}); diff --git a/apps/api/src/routes/v1_keys_deleteKey.ts b/apps/api/src/routes/v1_keys_deleteKey.ts index 6ac781a3c7..f361b501a0 100644 --- a/apps/api/src/routes/v1_keys_deleteKey.ts +++ b/apps/api/src/routes/v1_keys_deleteKey.ts @@ -70,6 +70,7 @@ export const registerV1KeysDeleteKey = (app: App) => with: { identity: true, encrypted: true, + credits: true, permissions: { with: { permission: true, @@ -93,6 +94,7 @@ export const registerV1KeysDeleteKey = (app: App) => } return { key: dbRes, + credits: dbRes.credits, api: dbRes.keyAuth.api, permissions: dbRes.permissions.map((p) => p.permission.name), roles: dbRes.roles.map((r) => r.role.name), @@ -114,7 +116,10 @@ export const registerV1KeysDeleteKey = (app: App) => }); } if (!data.val) { - throw new UnkeyApiError({ code: "NOT_FOUND", message: `key ${keyId} not found` }); + throw new UnkeyApiError({ + code: "NOT_FOUND", + message: `key ${keyId} not found`, + }); } const { key, api } = data.val; @@ -136,6 +141,11 @@ export const registerV1KeysDeleteKey = (app: App) => await db.primary.transaction(async (tx) => { if (permanent) { await tx.delete(schema.keys).where(eq(schema.keys.id, key.id)); + await tx.delete(schema.credits).where(eq(schema.credits.keyId, key.id)); + await tx.delete(schema.ratelimits).where(eq(schema.ratelimits.keyId, key.id)); + await tx.delete(schema.encryptedKeys).where(eq(schema.encryptedKeys.keyId, key.id)); + await tx.delete(schema.keysPermissions).where(eq(schema.keysPermissions.keyId, key.id)); + await tx.delete(schema.keysRoles).where(eq(schema.keysRoles.keyId, key.id)); } else { await tx .update(schema.keys) diff --git a/apps/api/src/routes/v1_keys_getKey.happy.test.ts b/apps/api/src/routes/v1_keys_getKey.happy.test.ts index c42ac3e9b8..edbfeb8e3c 100644 --- a/apps/api/src/routes/v1_keys_getKey.happy.test.ts +++ b/apps/api/src/routes/v1_keys_getKey.happy.test.ts @@ -66,3 +66,124 @@ test("returns identity", async (t) => { expect(res.body.identity!.id).toEqual(identity.id); expect(res.body.identity!.externalId).toEqual(identity.externalId); }); + +test("returns credits from new credits table", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey(["api.*.read_key"]); + + const keyId = newId("test"); + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + createdAtM: Date.now(), + }); + + // Create credits in new table with monthly refill + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 5000, + refillAmount: 2000, + refillDay: 10, // monthly refill + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + updatedAt: null, + }); + + const res = await h.get({ + url: `/v1/keys.getKey?keyId=${keyId}`, + headers: { + Authorization: `Bearer ${root.key}`, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.remaining).toBe(5000); + expect(res.body.refill).toBeDefined(); + expect(res.body.refill!.interval).toBe("monthly"); + expect(res.body.refill!.amount).toBe(2000); + expect(res.body.refill!.refillDay).toBe(10); +}); + +test("returns credits from legacy remaining field when no credits table entry", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey(["api.*.read_key"]); + + const keyId = newId("test"); + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + remaining: 800, + refillAmount: 100, + refillDay: null, // daily refill + createdAtM: Date.now(), + }); + + const res = await h.get({ + url: `/v1/keys.getKey?keyId=${keyId}`, + headers: { + Authorization: `Bearer ${root.key}`, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.remaining).toBe(800); + expect(res.body.refill).toBeDefined(); + expect(res.body.refill!.interval).toBe("daily"); + expect(res.body.refill!.amount).toBe(100); +}); + +test("prefers new credits table over legacy remaining field", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey(["api.*.read_key"]); + + const keyId = newId("test"); + // Create key with legacy credits (should be ignored) + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + remaining: 123, // This should be ignored + refillAmount: 50, // This should be ignored + createdAtM: Date.now(), + }); + + // Create credits in new table (should take precedence) + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 7000, + refillAmount: 3000, + refillDay: null, // daily refill + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + updatedAt: null, + }); + + const res = await h.get({ + url: `/v1/keys.getKey?keyId=${keyId}`, + headers: { + Authorization: `Bearer ${root.key}`, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + // Should use values from credits table, not legacy fields + expect(res.body.remaining).toBe(7000); + expect(res.body.refill).toBeDefined(); + expect(res.body.refill!.amount).toBe(3000); +}); diff --git a/apps/api/src/routes/v1_keys_getKey.ts b/apps/api/src/routes/v1_keys_getKey.ts index 08fea4f3d3..68517a31f2 100644 --- a/apps/api/src/routes/v1_keys_getKey.ts +++ b/apps/api/src/routes/v1_keys_getKey.ts @@ -58,6 +58,7 @@ export const registerV1KeysGetKey = (app: App) => encrypted: true, permissions: { with: { permission: true } }, roles: { with: { role: true } }, + credits: true, keyAuth: { with: { api: true, @@ -73,6 +74,7 @@ export const registerV1KeysGetKey = (app: App) => return { key: dbRes, api: dbRes.keyAuth.api, + credits: dbRes.credits, permissions: dbRes.permissions.map((p) => p.permission.name), roles: dbRes.roles.map((p) => p.role.name), ratelimits: dbRes.ratelimits, @@ -146,6 +148,21 @@ export const registerV1KeysGetKey = (app: App) => const ratelimit = key.ratelimits.find((rl) => rl.name === "default"); + // Prefer new credits table over legacy fields + const refillAmount = data.credits?.refillAmount ?? key.refillAmount; + const refillDay = data.credits?.refillDay ?? key.refillDay; + const lastRefillAt = data.credits?.refilledAt ?? key.lastRefillAt?.getTime(); + + let refill = undefined; + if (refillAmount) { + refill = { + interval: refillDay ? ("monthly" as const) : ("daily" as const), + amount: refillAmount, + refillDay: refillDay ?? null, + lastRefillAt: lastRefillAt ?? undefined, + }; + } + return c.json({ id: key.id, start: key.start, @@ -157,15 +174,8 @@ export const registerV1KeysGetKey = (app: App) => createdAt: key.createdAtM, updatedAt: key.updatedAtM ?? undefined, expires: key.expires?.getTime() ?? undefined, - remaining: key.remaining ?? undefined, - refill: key.refillAmount - ? { - interval: key.refillDay ? ("monthly" as const) : ("daily" as const), - amount: key.refillAmount, - refillDay: key.refillDay, - lastRefillAt: key.lastRefillAt?.getTime(), - } - : undefined, + remaining: data.credits?.remaining ?? key.remaining ?? undefined, + refill: refill, ratelimit: ratelimit ? { async: false, diff --git a/apps/api/src/routes/v1_keys_getVerifications.ts b/apps/api/src/routes/v1_keys_getVerifications.ts index 4141c86793..aead96eaa0 100644 --- a/apps/api/src/routes/v1_keys_getVerifications.ts +++ b/apps/api/src/routes/v1_keys_getVerifications.ts @@ -102,6 +102,7 @@ export const registerV1KeysGetVerifications = (app: App) => with: { identity: true, encrypted: true, + credits: true, permissions: { with: { permission: true } }, roles: { with: { role: true } }, keyAuth: { @@ -119,6 +120,7 @@ export const registerV1KeysGetVerifications = (app: App) => return { key: dbRes, api: dbRes.keyAuth.api, + credits: dbRes.credits, permissions: dbRes.permissions.map((p) => p.permission.name), roles: dbRes.roles.map((p) => p.role.name), identity: dbRes.identity diff --git a/apps/api/src/routes/v1_keys_updateKey.happy.test.ts b/apps/api/src/routes/v1_keys_updateKey.happy.test.ts index 35d29b8048..210120345e 100644 --- a/apps/api/src/routes/v1_keys_updateKey.happy.test.ts +++ b/apps/api/src/routes/v1_keys_updateKey.happy.test.ts @@ -90,13 +90,78 @@ test("update all", async (t) => { where: (table, { eq }) => eq(table.id, key.id), with: { ratelimits: true, + credits: true, }, }); expect(found).toBeDefined(); expect(found?.name).toEqual("newName"); expect(found?.ownerId).toEqual("newOwnerId"); expect(found?.meta).toEqual(JSON.stringify({ new: "meta" })); - expect(found?.remaining).toEqual(0); + // Since key didn't have legacy remaining, the new credits table is used + expect(found?.remaining).toBeNull(); + expect(found?.credits).toBeDefined(); + expect(found?.credits?.remaining).toEqual(0); + expect(found!.ratelimits.length).toBe(1); + expect(found!.ratelimits[0].name).toBe("default"); + expect(found!.ratelimits[0].limit).toBe(10); + expect(found!.ratelimits[0].duration).toBe(1000); +}); + +test("update all with legacy remaining", async (t) => { + const h = await IntegrationHarness.init(t); + + const key = { + id: newId("test"), + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + remaining: 100, + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + const root = await h.createRootKey([`api.${h.resources.userApi.id}.update_key`]); + + const res = await h.post({ + url: "/v1/keys.updateKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + name: "newName", + ownerId: "newOwnerId", + expires: null, + meta: { new: "meta" }, + ratelimit: { + type: "fast", + limit: 10, + refillRate: 5, + refillInterval: 1000, + }, + remaining: 50, + enabled: true, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + const found = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, key.id), + with: { + ratelimits: true, + credits: true, + }, + }); + expect(found).toBeDefined(); + expect(found?.name).toEqual("newName"); + expect(found?.ownerId).toEqual("newOwnerId"); + expect(found?.meta).toEqual(JSON.stringify({ new: "meta" })); + // Since key has legacy remaining, it should stay in legacy system + expect(found?.remaining).toEqual(50); + expect(found?.credits).toBeNull(); expect(found!.ratelimits.length).toBe(1); expect(found!.ratelimits[0].name).toBe("default"); expect(found!.ratelimits[0].limit).toBe(10); @@ -1029,6 +1094,7 @@ test("update ratelimit should not disable it", async (t) => { expect(verify.body.ratelimit!.limit).toBe(5); expect(verify.body.ratelimit!.remaining).toBe(4); }); + describe("When refillDay is omitted.", () => { test("should provide default value", async (t) => { const h = await IntegrationHarness.init(t); @@ -1147,3 +1213,223 @@ describe("update name", () => { expect(verify.body.ratelimit!.remaining).toBe(9); }); }); + +describe("update remaining with new credits table", () => { + test("updates credits table when key has new credits", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = { + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 100, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + const root = await h.createRootKey([`api.${h.resources.userApi.id}.update_key`]); + + const res = await h.post({ + url: "/v1/keys.updateKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + remaining: 200, + enabled: true, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + // Verify the credits table was updated + const updatedCredit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + expect(updatedCredit?.remaining).toEqual(200); + + // Verify the keys table was NOT updated (should remain null) + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.remaining).toBeNull(); + }); + + test("creates new credits table entry when key has no credits", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = { + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + + const root = await h.createRootKey([`api.${h.resources.userApi.id}.update_key`]); + + const res = await h.post({ + url: "/v1/keys.updateKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + remaining: 150, + enabled: true, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + // Verify a new credits table entry was created + const credit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.keyId, keyId), + }); + expect(credit).toBeDefined(); + expect(credit?.remaining).toEqual(150); + expect(credit?.keyId).toEqual(keyId); + expect(credit?.workspaceId).toEqual(h.resources.userWorkspace.id); + + // Verify the keys table was NOT updated (should remain null) + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.remaining).toBeNull(); + }); + + test("updates legacy key.remaining when it exists", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = { + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + remaining: 100, + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + + const root = await h.createRootKey([`api.${h.resources.userApi.id}.update_key`]); + + const res = await h.post({ + url: "/v1/keys.updateKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + remaining: 50, + enabled: true, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + // Verify the keys table was updated + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.remaining).toEqual(50); + + // Verify NO credits table entry was created (we should stay in legacy system) + const credit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.keyId, keyId), + }); + expect(credit).toBeUndefined(); + }); + + test("updates refill in credits table when key has new credits", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = { + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 100, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + const root = await h.createRootKey([`api.${h.resources.userApi.id}.update_key`]); + + const res = await h.post({ + url: "/v1/keys.updateKey", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + refill: { + interval: "monthly", + amount: 500, + refillDay: 15, + }, + enabled: true, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + + // Verify the credits table was updated + const updatedCredit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + expect(updatedCredit?.refillAmount).toEqual(500); + expect(updatedCredit?.refillDay).toEqual(15); + + // Verify the keys table was NOT updated + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.refillAmount).toBeNull(); + expect(updatedKey?.refillDay).toBeNull(); + }); +}); diff --git a/apps/api/src/routes/v1_keys_updateKey.ts b/apps/api/src/routes/v1_keys_updateKey.ts index ee148d6da7..d002bca574 100644 --- a/apps/api/src/routes/v1_keys_updateKey.ts +++ b/apps/api/src/routes/v1_keys_updateKey.ts @@ -4,7 +4,7 @@ import { createRoute, z } from "@hono/zod-openapi"; import { insertUnkeyAuditLog } from "@/pkg/audit"; import { rootKeyAuth } from "@/pkg/auth/root_key"; import { UnkeyApiError, openApiErrorResponses } from "@/pkg/errors"; -import { type Key, and, schema } from "@unkey/db"; +import { type Credits, type Key, and, schema } from "@unkey/db"; import { eq } from "@unkey/db"; import { newId } from "@unkey/id"; import { buildUnkeyQuery } from "@unkey/rbac"; @@ -289,6 +289,7 @@ export const registerV1KeysUpdate = (app: App) => const key = await db.primary.query.keys.findFirst({ where: (table, { eq, and, isNull }) => and(eq(table.id, req.keyId), isNull(table.deletedAtM)), with: { + credits: true, keyAuth: { with: { api: true, @@ -347,12 +348,19 @@ export const registerV1KeysUpdate = (app: App) => const rootKeyId = auth.key.id; const changes: Partial = {}; + const creditChanges: Partial = {}; + + const hasOldCredits = key.remaining !== null; + const hasNewCredits = key.credits !== null; + if (typeof req.name !== "undefined") { changes.name = req.name; } + if (typeof req.meta !== "undefined") { changes.meta = req.meta === null ? null : JSON.stringify(req.meta); } + if (typeof req.externalId !== "undefined") { if (req.externalId === null) { changes.identityId = null; @@ -372,12 +380,46 @@ export const registerV1KeysUpdate = (app: App) => changes.ownerId = req.ownerId; } } + if (typeof req.expires !== "undefined") { changes.expires = req.expires === null ? null : new Date(req.expires); } + + let creditDeleted = false; + if (typeof req.remaining !== "undefined") { - changes.remaining = req.remaining; + // Key has new credit system. + if (hasNewCredits) { + if (req.remaining === null) { + await db.primary.delete(schema.credits).where(eq(schema.credits.id, key.credits.id)); + creditDeleted = true; + } else { + creditChanges.remaining = req.remaining as number; + } + } else if (hasOldCredits) { + // Key has old credit system, so we work based off that + changes.remaining = req.remaining; + } else { + // Key doesn't have old system so we use new system from the getgo + await db.primary.insert(schema.credits).values({ + id: newId("credit"), + keyId: key.id, + workspaceId: auth.authorizedWorkspaceId, + createdAt: Date.now(), + refilledAt: Date.now(), + remaining: req.remaining as number, + identityId: null, + refillAmount: req.refill?.amount ?? null, + refillDay: req.refill + ? req.refill.interval === "monthly" + ? (req.refill.refillDay ?? 1) + : null + : null, + updatedAt: null, + }); + } } + if (typeof req.ratelimit !== "undefined") { if (req.ratelimit === null) { await db.primary @@ -414,21 +456,48 @@ export const registerV1KeysUpdate = (app: App) => } if (typeof req.refill !== "undefined") { - if (req.refill === null) { - changes.refillAmount = null; - changes.refillDay = null; - changes.lastRefillAt = null; - } else { - changes.refillAmount = req.refill.amount; - changes.refillDay = req.refill.refillDay ?? 1; + if (hasNewCredits || !hasOldCredits) { + if (req.refill === null) { + creditChanges.refillAmount = null; + creditChanges.refillDay = null; + creditChanges.refilledAt = null; + } else { + creditChanges.refillAmount = req.refill.amount; + creditChanges.refillDay = + req.refill.interval === "monthly" ? (req.refill.refillDay ?? 1) : null; + } + } + + if (hasOldCredits) { + if (req.refill === null) { + changes.refillAmount = null; + changes.refillDay = null; + changes.lastRefillAt = null; + } else { + changes.refillAmount = req.refill.amount; + changes.refillDay = + req.refill.interval === "monthly" ? (req.refill.refillDay ?? 1) : null; + } } } + if (typeof req.enabled !== "undefined") { changes.enabled = req.enabled; } - if (Object.keys(changes).length > 0) { + + if (Object.keys(changes).length) { await db.primary.update(schema.keys).set(changes).where(eq(schema.keys.id, key.id)); } + + // Since we don't know if we already have a row for this key we just update on keyId + // Skip if we deleted the credits row earlier + if (Object.keys(creditChanges).length && !creditDeleted) { + await db.primary + .update(schema.credits) + .set(creditChanges) + .where(eq(schema.credits.keyId, key.id)); + } + await insertUnkeyAuditLog(c, undefined, { workspaceId: authorizedWorkspaceId, event: "key.update", @@ -458,7 +527,19 @@ export const registerV1KeysUpdate = (app: App) => userAgent: c.get("userAgent"), }, }); - c.executionCtx.waitUntil(usageLimiter.revalidate({ keyId: key.id })); + + // Revalidate usage limiter with appropriate ID based on which credit system is in use + const keyAfterUpdate = await db.readonly.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, key.id), + with: { + credits: true, + }, + }); + c.executionCtx.waitUntil( + usageLimiter.revalidate( + keyAfterUpdate?.credits ? { creditId: keyAfterUpdate.credits.id } : { keyId: key.id }, + ), + ); c.executionCtx.waitUntil(cache.keyByHash.remove(key.hash)); c.executionCtx.waitUntil(cache.keyById.remove(key.id)); diff --git a/apps/api/src/routes/v1_keys_updateRemaining.happy.test.ts b/apps/api/src/routes/v1_keys_updateRemaining.happy.test.ts index 649fb98557..33055cca33 100644 --- a/apps/api/src/routes/v1_keys_updateRemaining.happy.test.ts +++ b/apps/api/src/routes/v1_keys_updateRemaining.happy.test.ts @@ -14,8 +14,9 @@ import type { test("increment", async (t) => { const h = await IntegrationHarness.init(t); + const keyId = newId("test"); const key = { - id: newId("test"), + id: keyId, keyAuthId: h.resources.userKeyAuth.id, workspaceId: h.resources.userWorkspace.id, start: "test", @@ -42,13 +43,24 @@ test("increment", async (t) => { expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); expect(res.body.remaining).toEqual(110); + + // Verify legacy keys.remaining was updated + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + with: { + credits: true, + }, + }); + expect(updatedKey?.remaining).toEqual(110); + expect(updatedKey?.credits).toBeNull(); }); test("decrement", async (t) => { const h = await IntegrationHarness.init(t); + const keyId = newId("test"); const key = { - id: newId("test"), + id: keyId, keyAuthId: h.resources.userKeyAuth.id, workspaceId: h.resources.userWorkspace.id, start: "test", @@ -75,13 +87,24 @@ test("decrement", async (t) => { expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); expect(res.body.remaining).toEqual(90); + + // Verify legacy keys.remaining was updated + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + with: { + credits: true, + }, + }); + expect(updatedKey?.remaining).toEqual(90); + expect(updatedKey?.credits).toBeNull(); }); test("set", async (t) => { const h = await IntegrationHarness.init(t); + const keyId = newId("test"); const key = { - id: newId("test"), + id: keyId, keyAuthId: h.resources.userKeyAuth.id, workspaceId: h.resources.userWorkspace.id, start: "test", @@ -108,6 +131,16 @@ test("set", async (t) => { expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); expect(res.body.remaining).toEqual(10); + + // Verify legacy keys.remaining was updated + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + with: { + credits: true, + }, + }); + expect(updatedKey?.remaining).toEqual(10); + expect(updatedKey?.credits).toBeNull(); }); test("invalid operation", async (t) => { @@ -142,3 +175,215 @@ test("invalid operation", async (t) => { expect(res.status).toEqual(400); }); + +test("increment with new credits table", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = { + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 100, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + const root = await h.createRootKey(["api.*.update_key"]); + const res = await h.post({ + url: "/v1/keys.updateRemaining", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + op: "increment", + value: 10, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.remaining).toEqual(110); + + // Verify the credits table was updated + const updatedCredit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + expect(updatedCredit?.remaining).toEqual(110); + + // Verify the keys table was NOT updated (should remain null) + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.remaining).toBeNull(); +}); + +test("decrement with new credits table", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = { + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 100, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + const root = await h.createRootKey(["api.*.update_key"]); + const res = await h.post({ + url: "/v1/keys.updateRemaining", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + op: "decrement", + value: 10, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.remaining).toEqual(90); + + // Verify the credits table was updated + const updatedCredit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + expect(updatedCredit?.remaining).toEqual(90); + + // Verify the keys table was NOT updated (should remain null) + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.remaining).toBeNull(); +}); + +test("set operation creates new credits table entry when key has no credits", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = { + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + + const root = await h.createRootKey(["api.*.update_key"]); + const res = await h.post({ + url: "/v1/keys.updateRemaining", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + op: "set", + value: 50, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.remaining).toEqual(50); + + // Verify a new credits table entry was created + const credit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.keyId, keyId), + }); + expect(credit).toBeDefined(); + expect(credit?.remaining).toEqual(50); + expect(credit?.keyId).toEqual(keyId); + expect(credit?.workspaceId).toEqual(h.resources.userWorkspace.id); + + // Verify the keys table was NOT updated (should remain null) + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.remaining).toBeNull(); +}); + +test("set operation updates legacy key.remaining when it exists", async (t) => { + const h = await IntegrationHarness.init(t); + + const keyId = newId("test"); + const key = { + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + name: "test", + hash: await sha256(new KeyV1({ byteLength: 16 }).toString()), + remaining: 100, + createdAtM: Date.now(), + }; + await h.db.primary.insert(schema.keys).values(key); + + const root = await h.createRootKey(["api.*.update_key"]); + const res = await h.post({ + url: "/v1/keys.updateRemaining", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + keyId: key.id, + op: "set", + value: 50, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.remaining).toEqual(50); + + // Verify the keys table was updated + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.remaining).toEqual(50); + + // Verify NO credits table entry was created (we should stay in legacy system) + const credit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.keyId, keyId), + }); + expect(credit).toBeUndefined(); +}); diff --git a/apps/api/src/routes/v1_keys_updateRemaining.ts b/apps/api/src/routes/v1_keys_updateRemaining.ts index 05de238b97..3866c6d517 100644 --- a/apps/api/src/routes/v1_keys_updateRemaining.ts +++ b/apps/api/src/routes/v1_keys_updateRemaining.ts @@ -5,6 +5,7 @@ import { insertUnkeyAuditLog } from "@/pkg/audit"; import { rootKeyAuth } from "@/pkg/auth/root_key"; import { eq, schema, sql } from "@/pkg/db"; import { UnkeyApiError, openApiErrorResponses } from "@/pkg/errors"; +import { newId } from "@unkey/id"; import { buildUnkeyQuery } from "@unkey/rbac"; const route = createRoute({ @@ -74,6 +75,7 @@ export const registerV1KeysUpdateRemaining = (app: App) => const key = await db.readonly.query.keys.findFirst({ where: (table, { eq, isNull, and }) => and(eq(table.id, req.keyId), isNull(table.deletedAtM)), with: { + credits: true, keyAuth: { with: { api: true, @@ -88,12 +90,14 @@ export const registerV1KeysUpdateRemaining = (app: App) => message: `key ${req.keyId} not found`, }); } + const auth = await rootKeyAuth( c, buildUnkeyQuery(({ or }) => or("*", "api.*.update_key", `api.${key.keyAuth.api.id}.update_key`), ), ); + if (key.workspaceId !== auth.authorizedWorkspaceId) { throw new UnkeyApiError({ code: "NOT_FOUND", @@ -101,72 +105,122 @@ export const registerV1KeysUpdateRemaining = (app: App) => }); } + const hasNewCredits = key.credits !== null; + const hasOldCredits = key.remaining !== null; + const authorizedWorkspaceId = auth.authorizedWorkspaceId; const rootKeyId = auth.key.id; switch (req.op) { case "increment": { - if (key.remaining === null) { + if (key.remaining === null && key.credits === null) { throw new UnkeyApiError({ code: "BAD_REQUEST", message: "cannot increment a key with unlimited remaining requests, please 'set' a value instead.", }); } + if (req.value === null) { throw new UnkeyApiError({ code: "BAD_REQUEST", message: "cannot increment a key by null.", }); } - await db.primary - .update(schema.keys) - .set({ - remaining: sql`remaining_requests + ${req.value}`, - }) - .where(eq(schema.keys.id, req.keyId)); + + if (hasNewCredits) { + await db.primary + .update(schema.credits) + .set({ + remaining: sql`remaining + ${req.value}`, + }) + .where(eq(schema.credits.id, key.credits.id)); + } else { + await db.primary + .update(schema.keys) + .set({ + remaining: sql`remaining_requests + ${req.value}`, + }) + .where(eq(schema.keys.id, req.keyId)); + } + break; } case "decrement": { - if (key.remaining === null) { + if (key.remaining === null && key.credits === null) { throw new UnkeyApiError({ code: "BAD_REQUEST", message: "cannot decrement a key with unlimited remaining requests, please 'set' a value instead.", }); } + if (req.value === null) { throw new UnkeyApiError({ code: "BAD_REQUEST", message: "cannot decrement a key by null.", }); } - await db.primary - .update(schema.keys) - .set({ - remaining: sql`remaining_requests - ${req.value}`, - }) - .where(eq(schema.keys.id, req.keyId)); + + if (hasNewCredits) { + await db.primary + .update(schema.credits) + .set({ + remaining: sql`remaining - ${req.value}`, + }) + .where(eq(schema.credits.id, key.credits.id)); + } else { + await db.primary + .update(schema.keys) + .set({ + remaining: sql`remaining_requests - ${req.value}`, + }) + .where(eq(schema.keys.id, req.keyId)); + } + break; } case "set": { - await db.primary - .update(schema.keys) - .set({ + if (hasOldCredits) { + await db.primary + .update(schema.keys) + .set({ + remaining: req.value, + }) + .where(eq(schema.keys.id, req.keyId)); + } else if (hasNewCredits) { + if (req.value === null) { + await db.primary.delete(schema.credits).where(eq(schema.credits.id, key.credits.id)); + } else { + await db.primary + .update(schema.credits) + .set({ + remaining: req.value, + }) + .where(eq(schema.credits.id, key.credits.id)); + } + } else if (req.value !== null) { + await db.primary.insert(schema.credits).values({ + id: newId("credit"), + keyId: req.keyId, remaining: req.value, - }) - .where(eq(schema.keys.id, req.keyId)); + workspaceId: auth.authorizedWorkspaceId, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + } break; } } - await Promise.all([ - usageLimiter.revalidate({ keyId: key.id }), - cache.keyByHash.remove(key.hash), - cache.keyById.remove(key.id), - ]); - const keyAfterUpdate = await db.readonly.query.keys.findFirst({ where: (table, { eq }) => eq(table.id, req.keyId), + with: { + credits: true, + }, }); if (!keyAfterUpdate) { throw new UnkeyApiError({ @@ -175,6 +229,18 @@ export const registerV1KeysUpdateRemaining = (app: App) => }); } + await Promise.all([ + usageLimiter.revalidate( + keyAfterUpdate.credits + ? { creditId: keyAfterUpdate.credits.id } + : { keyId: keyAfterUpdate.id }, + ), + cache.keyByHash.remove(key.hash), + cache.keyById.remove(key.id), + ]); + + const actualRemaining = keyAfterUpdate.credits?.remaining ?? keyAfterUpdate.remaining; + await insertUnkeyAuditLog(c, undefined, { actor: { type: "key", @@ -182,7 +248,7 @@ export const registerV1KeysUpdateRemaining = (app: App) => }, event: "key.update", workspaceId: authorizedWorkspaceId, - description: `Changed remaining to ${keyAfterUpdate.remaining}`, + description: `Changed remaining to ${actualRemaining}`, resources: [ { type: "keyAuth", @@ -200,6 +266,6 @@ export const registerV1KeysUpdateRemaining = (app: App) => }); return c.json({ - remaining: keyAfterUpdate.remaining, + remaining: actualRemaining, }); }); diff --git a/apps/api/src/routes/v1_keys_verifyKey.test.ts b/apps/api/src/routes/v1_keys_verifyKey.test.ts index ce2df35ddc..a816874b9f 100644 --- a/apps/api/src/routes/v1_keys_verifyKey.test.ts +++ b/apps/api/src/routes/v1_keys_verifyKey.test.ts @@ -491,6 +491,45 @@ describe("with default ratelimit", () => { }); describe("with remaining", () => { + test("works with legacy keys.remaining", async (t) => { + const h = await IntegrationHarness.init(t); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + const keyId = newId("test"); + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + remaining: 10, + }); + + const res = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + apiId: h.resources.userApi.id, + }, + }); + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.valid).toBe(true); + expect(res.body.remaining).toEqual(9); + + // Verify the legacy keys.remaining was decremented + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + with: { + credits: true, + }, + }); + expect(updatedKey?.remaining).toEqual(9); + expect(updatedKey?.credits).toBeNull(); + }); + test("custom cost works", async (t) => { const h = await IntegrationHarness.init(t); const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); @@ -548,6 +587,316 @@ describe("with remaining", () => { expect(res.body.valid).toBe(true); expect(res.body.remaining).toEqual(0); }); + + test("rejects when cost exceeds remaining with legacy keys", async (t) => { + const h = await IntegrationHarness.init(t); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + const keyId = newId("test"); + + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + remaining: 50, + }); + + // Try to use 100 credits when only 50 are remaining + const res = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + apiId: h.resources.userApi.id, + remaining: { cost: 100 }, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.valid).toBe(false); + expect(res.body.code).toEqual("USAGE_EXCEEDED"); + + // Verify the legacy keys.remaining was not decremented + const updatedKey = await h.db.primary.query.keys.findFirst({ + where: (table, { eq }) => eq(table.id, keyId), + }); + expect(updatedKey?.remaining).toEqual(50); + }); +}); + +describe("with remaining in new credits table", () => { + test("usage limiting works correctly with new credits table", async (t) => { + const h = await IntegrationHarness.init(t); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + const keyId = newId("test"); + + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + }); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 100, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + const res = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + apiId: h.resources.userApi.id, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.valid).toBe(true); + expect(res.body.remaining).toEqual(99); + + // Verify the credits table was decremented + const updatedCredit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + expect(updatedCredit?.remaining).toEqual(99); + }); + + test("custom cost works with new credits table", async (t) => { + const h = await IntegrationHarness.init(t); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + const keyId = newId("test"); + + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + }); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 50, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + const res = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + apiId: h.resources.userApi.id, + remaining: { cost: 5 }, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.valid).toBe(true); + expect(res.body.remaining).toEqual(45); + + // Verify the credits table was decremented by the custom cost + const updatedCredit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + expect(updatedCredit?.remaining).toEqual(45); + }); + + test("cost=0 works with new credits table when remaining=0", async (t) => { + const h = await IntegrationHarness.init(t); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + const keyId = newId("test"); + + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + }); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 0, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + const res = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + apiId: h.resources.userApi.id, + remaining: { cost: 0 }, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.valid).toBe(true); + expect(res.body.remaining).toEqual(0); + + // Verify the credits table was not decremented + const updatedCredit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + expect(updatedCredit?.remaining).toEqual(0); + }); + + test("key becomes invalid when credits run out", async (t) => { + const h = await IntegrationHarness.init(t); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + const keyId = newId("test"); + + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + }); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 1, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + // First request should succeed and use the last credit + const firstRes = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + apiId: h.resources.userApi.id, + }, + }); + + expect(firstRes.status, `expected 200, received: ${JSON.stringify(firstRes, null, 2)}`).toBe( + 200, + ); + expect(firstRes.body.valid).toBe(true); + expect(firstRes.body.remaining).toEqual(0); + + // Second request should fail with no credits remaining + const secondRes = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + apiId: h.resources.userApi.id, + }, + }); + + expect(secondRes.status, `expected 200, received: ${JSON.stringify(secondRes, null, 2)}`).toBe( + 200, + ); + expect(secondRes.body.valid).toBe(false); + expect(secondRes.body.code).toEqual("USAGE_EXCEEDED"); + }); + + test("rejects when cost exceeds remaining credits", async (t) => { + const h = await IntegrationHarness.init(t); + const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); + const keyId = newId("test"); + + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + hash: await sha256(key), + start: key.slice(0, 8), + workspaceId: h.resources.userWorkspace.id, + createdAtM: Date.now(), + }); + + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId: keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 99, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }); + + // Try to use 101 credits when only 99 are remaining + const res = await h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + apiId: h.resources.userApi.id, + remaining: { cost: 101 }, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.valid).toBe(false); + expect(res.body.code).toEqual("USAGE_EXCEEDED"); + + // Verify the credits were not decremented + const updatedCredit = await h.db.primary.query.credits.findFirst({ + where: (table, { eq }) => eq(table.id, creditId), + }); + expect(updatedCredit?.remaining).toEqual(99); + }); }); describe("with ratelimit", () => { diff --git a/apps/api/src/routes/v1_keys_whoami.happy.test.ts b/apps/api/src/routes/v1_keys_whoami.happy.test.ts index 0ae1fb2a3e..70a092ef8b 100644 --- a/apps/api/src/routes/v1_keys_whoami.happy.test.ts +++ b/apps/api/src/routes/v1_keys_whoami.happy.test.ts @@ -84,3 +84,137 @@ test("returns identity", async (t) => { expect(res.body.identity!.id).toEqual(identity.id); expect(res.body.identity!.externalId).toEqual(identity.externalId); }); + +test("returns credits from new credits table", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey(["api.*.read_key"]); + + const key = new KeyV1({ byteLength: 16 }).toString(); + const hash = await sha256(key); + const keyId = newId("test"); + + // Create key + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + hash: hash, + enabled: true, + createdAtM: Date.now(), + }); + + // Create credits in new table + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 2500, + refillAmount: null, + refillDay: null, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + updatedAt: null, + }); + + const res = await h.post({ + url: "/v1/keys.whoami", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + key: key, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.remaining).toBe(2500); +}); + +test("returns credits from legacy remaining field when no credits table entry", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey(["api.*.read_key"]); + + const key = new KeyV1({ byteLength: 16 }).toString(); + const hash = await sha256(key); + + // Create key with legacy credits field + await h.db.primary.insert(schema.keys).values({ + id: newId("test"), + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + hash: hash, + enabled: true, + remaining: 1500, + createdAtM: Date.now(), + }); + + const res = await h.post({ + url: "/v1/keys.whoami", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + key: key, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + expect(res.body.remaining).toBe(1500); +}); + +test("prefers new credits table over legacy remaining field", async (t) => { + const h = await IntegrationHarness.init(t); + const root = await h.createRootKey(["api.*.read_key"]); + + const key = new KeyV1({ byteLength: 16 }).toString(); + const hash = await sha256(key); + const keyId = newId("test"); + + // Create key with legacy credits field (should be ignored) + await h.db.primary.insert(schema.keys).values({ + id: keyId, + keyAuthId: h.resources.userKeyAuth.id, + workspaceId: h.resources.userWorkspace.id, + start: "test", + hash: hash, + enabled: true, + remaining: 999, // This should be ignored + createdAtM: Date.now(), + }); + + // Create credits in new table (should take precedence) + const creditId = newId("credit"); + await h.db.primary.insert(schema.credits).values({ + id: creditId, + keyId, + workspaceId: h.resources.userWorkspace.id, + remaining: 3000, + refillAmount: null, + refillDay: null, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + updatedAt: null, + }); + + const res = await h.post({ + url: "/v1/keys.whoami", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + key: key, + }, + }); + + expect(res.status, `expected 200, received: ${JSON.stringify(res, null, 2)}`).toBe(200); + // Should use value from credits table, not legacy field + expect(res.body.remaining).toBe(3000); +}); diff --git a/apps/api/src/routes/v1_keys_whoami.ts b/apps/api/src/routes/v1_keys_whoami.ts index dff60eecb3..974fa1c67c 100644 --- a/apps/api/src/routes/v1_keys_whoami.ts +++ b/apps/api/src/routes/v1_keys_whoami.ts @@ -3,6 +3,7 @@ import { createRoute, z } from "@hono/zod-openapi"; import { rootKeyAuth } from "@/pkg/auth/root_key"; import { UnkeyApiError, openApiErrorResponses } from "@/pkg/errors"; +import type { Ratelimit } from "@unkey/db"; import { sha256 } from "@unkey/hash"; import { buildUnkeyQuery } from "@unkey/rbac"; @@ -106,29 +107,105 @@ export const registerV1KeysWhoAmI = (app: App) => const { cache, db } = c.get("services"); const hash = await sha256(secret); const { val: data, err } = await cache.keyByHash.swr(hash, async () => { - const dbRes = await db.readonly.query.keys.findFirst({ - where: (table, { eq, and, isNull }) => and(eq(table.hash, hash), isNull(table.deletedAtM)), + const query = db.readonly.query.keys.findFirst({ + where: (table, { and, eq, isNull }) => and(eq(table.hash, hash), isNull(table.deletedAtM)), with: { + encrypted: true, + credits: true, + workspace: { + columns: { + id: true, + enabled: true, + }, + }, + forWorkspace: { + columns: { + id: true, + enabled: true, + }, + }, + roles: { + with: { + role: { + with: { + permissions: { + with: { + permission: true, + }, + }, + }, + }, + }, + }, + permissions: { + with: { + permission: true, + }, + }, keyAuth: { with: { api: true, }, }, - identity: true, + ratelimits: true, + identity: { + with: { + ratelimits: true, + }, + }, }, }); - if (!dbRes) { + const dbRes = await query; + + if (!dbRes?.keyAuth?.api) { + return null; + } + if (dbRes.keyAuth.deletedAtM) { + return null; + } + if (dbRes.keyAuth.api.deletedAtM) { return null; } + /** + * Create a unique set of all permissions, whether they're attached directly or connected + * through a role. + */ + const permissions = new Set([ + ...dbRes.permissions.filter((p) => p.permission).map((p) => p.permission.slug), + ...dbRes.roles.flatMap((r) => + r.role.permissions.filter((p) => p.permission).map((p) => p.permission.slug), + ), + ]); + + /** + * Merge ratelimits from the identity and the key + * Key limits take precedence + */ + const ratelimits: { + [name: string]: Pick; + } = {}; + + for (const rl of dbRes.identity?.ratelimits ?? []) { + ratelimits[rl.name] = rl; + } + + for (const rl of dbRes.ratelimits ?? []) { + ratelimits[rl.name] = rl; + } + return { - key: { - ...dbRes, - }, - api: dbRes.keyAuth.api, + credits: dbRes.credits, + workspace: dbRes.workspace, + forWorkspace: dbRes.forWorkspace, + key: dbRes, identity: dbRes.identity, - } as any; // this was necessary so that we don't need to return the workspace and other types defined in keyByHash + api: dbRes.keyAuth.api, + permissions: Array.from(permissions.values()), + roles: dbRes.roles.map((r) => r.role.name), + ratelimits, + }; }); if (err) { @@ -137,12 +214,14 @@ export const registerV1KeysWhoAmI = (app: App) => message: `unable to load key: ${err.message}`, }); } + if (!data) { throw new UnkeyApiError({ code: "NOT_FOUND", message: "Key not found", }); } + const { api, key } = data; const auth = await rootKeyAuth( c, @@ -155,6 +234,7 @@ export const registerV1KeysWhoAmI = (app: App) => message: "Key not found", }); } + let meta = key.meta ? JSON.parse(key.meta) : undefined; if (!meta || Object.keys(meta).length === 0) { meta = undefined; @@ -163,7 +243,7 @@ export const registerV1KeysWhoAmI = (app: App) => return c.json({ id: key.id, name: key.name ?? undefined, - remaining: key.remaining ?? undefined, + remaining: data?.credits?.remaining ?? key.remaining ?? undefined, identity: data.identity ? { id: data.identity.id, diff --git a/apps/api/src/routes/v1_migrations_createKey.ts b/apps/api/src/routes/v1_migrations_createKey.ts index cc21f50d59..be1b1b271a 100644 --- a/apps/api/src/routes/v1_migrations_createKey.ts +++ b/apps/api/src/routes/v1_migrations_createKey.ts @@ -8,6 +8,7 @@ import { retry } from "@/pkg/util/retry"; import { DatabaseError } from "@planetscale/database"; import { DrizzleQueryError, + type InsertCredits, type InsertEncryptedKey, type InsertIdentity, type InsertKey, @@ -324,6 +325,7 @@ export const registerV1MigrationsCreateKeys = (app: App) => const createIdentities: Array = []; const keys: Array = []; const encryptedKeys: Array = []; + const credits: Array = []; const roleConnections: Array = []; const permissionConnections: Array = []; @@ -474,9 +476,6 @@ export const registerV1MigrationsCreateKeys = (app: App) => workspaceId: authorizedWorkspaceId, forWorkspaceId: null, expires: key.expires ? new Date(key.expires) : null, - remaining: key.remaining ?? null, - refillDay: key.refill?.interval === "daily" ? null : (key?.refill?.refillDay ?? 1), - refillAmount: key.refill?.amount ?? null, deletedAtM: null, enabled: key.enabled ?? true, environment: key.environment ?? null, @@ -485,6 +484,21 @@ export const registerV1MigrationsCreateKeys = (app: App) => lastRefillAt: null, }); + if (key.remaining) { + credits.push({ + id: newId("credit"), + remaining: key.remaining, + workspaceId: authorizedWorkspaceId, + createdAt: Date.now(), + keyId: key.keyId, + identityId: null, + refillAmount: key.refill?.amount ?? null, + refillDay: key.refill?.interval === "daily" ? null : (key?.refill?.refillDay ?? 1), + refilledAt: key.refill?.interval ? Date.now() : null, + updatedAt: null, + }); + } + for (const role of key.roles ?? []) { const roleId = roles[role]; roleConnections.push({ @@ -495,6 +509,7 @@ export const registerV1MigrationsCreateKeys = (app: App) => workspaceId: authorizedWorkspaceId, }); } + for (const permission of key.permissions ?? []) { const permissionId = permissions[permission]; permissionConnections.push({ @@ -571,9 +586,15 @@ export const registerV1MigrationsCreateKeys = (app: App) => if (encryptedKeys.length > 0) { await tx.insert(schema.encryptedKeys).values(encryptedKeys); } + + if (credits.length > 0) { + await tx.insert(schema.credits).values(credits); + } + if (roleConnections.length > 0) { await tx.insert(schema.keysRoles).values(roleConnections); } + if (permissionConnections.length > 0) { await tx.insert(schema.keysPermissions).values(permissionConnections); } From 86602970dd98b351908b9cf379ba203f7c6915cf Mon Sep 17 00:00:00 2001 From: Flo Date: Mon, 3 Nov 2025 12:48:50 +0100 Subject: [PATCH 2/3] fix: coderabbit comments --- apps/api/src/pkg/usagelimit/durable_object.ts | 59 ++++++++++--------- .../src/routes/v1_apis_listKeys.happy.test.ts | 2 + apps/api/src/routes/v1_keys_updateKey.ts | 24 +++++--- .../api/src/routes/v1_keys_updateRemaining.ts | 16 ++--- 4 files changed, 57 insertions(+), 44 deletions(-) diff --git a/apps/api/src/pkg/usagelimit/durable_object.ts b/apps/api/src/pkg/usagelimit/durable_object.ts index 805e5baff3..9dcb3dd1ac 100644 --- a/apps/api/src/pkg/usagelimit/durable_object.ts +++ b/apps/api/src/pkg/usagelimit/durable_object.ts @@ -5,7 +5,6 @@ import { and, createConnection, eq, - gt, schema, sql, } from "@/pkg/db"; @@ -103,20 +102,22 @@ export class DurableObjectUsagelimiter implements DurableObject { }); } - this.credit.remaining = Math.max(0, currentRemaining - req.cost); + if (req.cost !== 0) { + this.credit.remaining = Math.max(0, currentRemaining - req.cost); - this.state.waitUntil( - this.db - .update(schema.credits) - .set({ remaining: sql`${schema.credits.remaining}-${req.cost}` }) - .where( - and( - eq(schema.credits.id, this.credit.id), - gt(schema.credits.remaining, 0), // prevent negative remaining - ), - ) - .execute(), - ); + this.state.waitUntil( + this.db + .update(schema.credits) + .set({ remaining: sql`${schema.credits.remaining}-${req.cost}` }) + .where( + and( + eq(schema.credits.id, this.credit.id), + sql`${schema.credits.remaining} >= ${req.cost}`, // ensure sufficient credits + ), + ) + .execute(), + ); + } // revalidate every minute if (Date.now() - this.lastRevalidate > 60_000) { @@ -184,20 +185,22 @@ export class DurableObjectUsagelimiter implements DurableObject { }); } - this.key.remaining = Math.max(0, currentRemaining - req.cost); - - this.state.waitUntil( - this.db - .update(schema.keys) - .set({ remaining: sql`${schema.keys.remaining}-${req.cost}` }) - .where( - and( - eq(schema.keys.id, this.key.id), - gt(schema.keys.remaining, 0), // prevent negative remaining - ), - ) - .execute(), - ); + if (req.cost !== 0) { + this.key.remaining = Math.max(0, currentRemaining - req.cost); + + this.state.waitUntil( + this.db + .update(schema.keys) + .set({ remaining: sql`${schema.keys.remaining}-${req.cost}` }) + .where( + and( + eq(schema.keys.id, this.key.id), + sql`${schema.keys.remaining} >= ${req.cost}`, // ensure sufficient credits + ), + ) + .execute(), + ); + } // revalidate every minute if (Date.now() - this.lastRevalidate > 60_000) { diff --git a/apps/api/src/routes/v1_apis_listKeys.happy.test.ts b/apps/api/src/routes/v1_apis_listKeys.happy.test.ts index df76a8d06d..ee79b0664b 100644 --- a/apps/api/src/routes/v1_apis_listKeys.happy.test.ts +++ b/apps/api/src/routes/v1_apis_listKeys.happy.test.ts @@ -560,4 +560,6 @@ test("prefers new credits table over legacy remaining field", async (t) => { expect(foundKey!.remaining).toBe(2000); expect(foundKey!.refill).toBeDefined(); expect(foundKey!.refill!.amount).toBe(1000); + expect(foundKey!.refill!.interval).toBe("daily"); + expect(foundKey!.refill!.lastRefillAt).toBeDefined(); }); diff --git a/apps/api/src/routes/v1_keys_updateKey.ts b/apps/api/src/routes/v1_keys_updateKey.ts index d002bca574..277f5dd17d 100644 --- a/apps/api/src/routes/v1_keys_updateKey.ts +++ b/apps/api/src/routes/v1_keys_updateKey.ts @@ -386,6 +386,7 @@ export const registerV1KeysUpdate = (app: App) => } let creditDeleted = false; + let creditCreatedId: string | undefined = undefined; if (typeof req.remaining !== "undefined") { // Key has new credit system. @@ -401,8 +402,10 @@ export const registerV1KeysUpdate = (app: App) => changes.remaining = req.remaining; } else { // Key doesn't have old system so we use new system from the getgo + const newCreditId = newId("credit"); + creditCreatedId = newCreditId; await db.primary.insert(schema.credits).values({ - id: newId("credit"), + id: newCreditId, keyId: key.id, workspaceId: auth.authorizedWorkspaceId, createdAt: Date.now(), @@ -529,15 +532,20 @@ export const registerV1KeysUpdate = (app: App) => }); // Revalidate usage limiter with appropriate ID based on which credit system is in use - const keyAfterUpdate = await db.readonly.query.keys.findFirst({ - where: (table, { eq }) => eq(table.id, key.id), - with: { - credits: true, - }, - }); + // Determine credit state from in-memory tracking to avoid stale reads from db.readonly + let creditIdForRevalidation: string | undefined; + if (creditCreatedId) { + // A new credit was created + creditIdForRevalidation = creditCreatedId; + } else if (hasNewCredits && !creditDeleted) { + // Credit existed and wasn't deleted + creditIdForRevalidation = key.credits.id; + } + // else: credit was deleted or never existed, use keyId + c.executionCtx.waitUntil( usageLimiter.revalidate( - keyAfterUpdate?.credits ? { creditId: keyAfterUpdate.credits.id } : { keyId: key.id }, + creditIdForRevalidation ? { creditId: creditIdForRevalidation } : { keyId: key.id }, ), ); c.executionCtx.waitUntil(cache.keyByHash.remove(key.hash)); diff --git a/apps/api/src/routes/v1_keys_updateRemaining.ts b/apps/api/src/routes/v1_keys_updateRemaining.ts index 3866c6d517..0088e79415 100644 --- a/apps/api/src/routes/v1_keys_updateRemaining.ts +++ b/apps/api/src/routes/v1_keys_updateRemaining.ts @@ -180,14 +180,7 @@ export const registerV1KeysUpdateRemaining = (app: App) => break; } case "set": { - if (hasOldCredits) { - await db.primary - .update(schema.keys) - .set({ - remaining: req.value, - }) - .where(eq(schema.keys.id, req.keyId)); - } else if (hasNewCredits) { + if (hasNewCredits) { if (req.value === null) { await db.primary.delete(schema.credits).where(eq(schema.credits.id, key.credits.id)); } else { @@ -198,6 +191,13 @@ export const registerV1KeysUpdateRemaining = (app: App) => }) .where(eq(schema.credits.id, key.credits.id)); } + } else if (hasOldCredits) { + await db.primary + .update(schema.keys) + .set({ + remaining: req.value, + }) + .where(eq(schema.keys.id, req.keyId)); } else if (req.value !== null) { await db.primary.insert(schema.credits).values({ id: newId("credit"), From ab009c0df87dd382da60dc45f01406ef321b5e65 Mon Sep 17 00:00:00 2001 From: Flo Date: Mon, 3 Nov 2025 12:59:39 +0100 Subject: [PATCH 3/3] fix: coderabbit comments --- apps/api/src/routes/v1_keys_updateKey.ts | 48 ++++++++++++------- .../api/src/routes/v1_keys_updateRemaining.ts | 33 ++++++++----- .../api/src/routes/v1_migrations_createKey.ts | 6 ++- 3 files changed, 58 insertions(+), 29 deletions(-) diff --git a/apps/api/src/routes/v1_keys_updateKey.ts b/apps/api/src/routes/v1_keys_updateKey.ts index 277f5dd17d..c7ac5d939c 100644 --- a/apps/api/src/routes/v1_keys_updateKey.ts +++ b/apps/api/src/routes/v1_keys_updateKey.ts @@ -404,22 +404,38 @@ export const registerV1KeysUpdate = (app: App) => // Key doesn't have old system so we use new system from the getgo const newCreditId = newId("credit"); creditCreatedId = newCreditId; - await db.primary.insert(schema.credits).values({ - id: newCreditId, - keyId: key.id, - workspaceId: auth.authorizedWorkspaceId, - createdAt: Date.now(), - refilledAt: Date.now(), - remaining: req.remaining as number, - identityId: null, - refillAmount: req.refill?.amount ?? null, - refillDay: req.refill - ? req.refill.interval === "monthly" - ? (req.refill.refillDay ?? 1) - : null - : null, - updatedAt: null, - }); + // Use upsert to prevent race condition with concurrent updates + await db.primary + .insert(schema.credits) + .values({ + id: newCreditId, + keyId: key.id, + workspaceId: auth.authorizedWorkspaceId, + createdAt: Date.now(), + refilledAt: Date.now(), + remaining: req.remaining as number, + identityId: null, + refillAmount: req.refill?.amount ?? null, + refillDay: req.refill + ? req.refill.interval === "monthly" + ? (req.refill.refillDay ?? 1) + : null + : null, + updatedAt: null, + }) + .onDuplicateKeyUpdate({ + set: { + remaining: req.remaining as number, + refillAmount: req.refill?.amount ?? null, + refillDay: req.refill + ? req.refill.interval === "monthly" + ? (req.refill.refillDay ?? 1) + : null + : null, + refilledAt: req.refill?.interval ? Date.now() : null, + updatedAt: Date.now(), + }, + }); } } diff --git a/apps/api/src/routes/v1_keys_updateRemaining.ts b/apps/api/src/routes/v1_keys_updateRemaining.ts index 0088e79415..7de5a4cf5b 100644 --- a/apps/api/src/routes/v1_keys_updateRemaining.ts +++ b/apps/api/src/routes/v1_keys_updateRemaining.ts @@ -199,18 +199,27 @@ export const registerV1KeysUpdateRemaining = (app: App) => }) .where(eq(schema.keys.id, req.keyId)); } else if (req.value !== null) { - await db.primary.insert(schema.credits).values({ - id: newId("credit"), - keyId: req.keyId, - remaining: req.value, - workspaceId: auth.authorizedWorkspaceId, - createdAt: Date.now(), - refilledAt: Date.now(), - identityId: null, - refillAmount: null, - refillDay: null, - updatedAt: null, - }); + // Use upsert to prevent race condition with concurrent set operations + await db.primary + .insert(schema.credits) + .values({ + id: newId("credit"), + keyId: req.keyId, + remaining: req.value, + workspaceId: auth.authorizedWorkspaceId, + createdAt: Date.now(), + refilledAt: Date.now(), + identityId: null, + refillAmount: null, + refillDay: null, + updatedAt: null, + }) + .onDuplicateKeyUpdate({ + set: { + remaining: req.value, + updatedAt: Date.now(), + }, + }); } break; } diff --git a/apps/api/src/routes/v1_migrations_createKey.ts b/apps/api/src/routes/v1_migrations_createKey.ts index be1b1b271a..4b942edca2 100644 --- a/apps/api/src/routes/v1_migrations_createKey.ts +++ b/apps/api/src/routes/v1_migrations_createKey.ts @@ -493,7 +493,11 @@ export const registerV1MigrationsCreateKeys = (app: App) => keyId: key.keyId, identityId: null, refillAmount: key.refill?.amount ?? null, - refillDay: key.refill?.interval === "daily" ? null : (key?.refill?.refillDay ?? 1), + refillDay: key.refill + ? key.refill.interval === "daily" + ? null + : (key.refill.refillDay ?? 1) + : null, refilledAt: key.refill?.interval ? Date.now() : null, updatedAt: null, });