feat: allow-developer-to-set-a-custom-refill-time-when-using-the

apps/api/src/pkg/key_migration/handler.ts

@@ -96,6 +96,7 @@ export async function migrateKey(
         expires: message.expires ? new Date(message.expires) : null,
         refillInterval: message.refill?.interval,
         refillAmount: message.refill?.amount,
+        refillDay: message.refill?.refillDay,
         enabled: message.enabled,
         remaining: message.remaining,
         ratelimitAsync: message.ratelimit?.async,

apps/api/src/pkg/key_migration/message.ts

@@ -14,7 +14,7 @@ export type MessageBody = {
   permissions?: string[];
   expires?: number;
   remaining?: number;
-  refill?: { interval: "daily" | "monthly"; amount: number };
+  refill?: { interval: "daily" | "monthly"; amount: number; refillDay?: number };
   ratelimit?: { async: boolean; limit: number; duration: number };
   enabled: boolean;
   environment?: string;

apps/api/src/routes/schema.ts

@@ -66,6 +66,11 @@ export const keySchema = z
           description: "Resets `remaining` to this value every interval.",
           example: 100,
         }),
+        refillDay: z.number().min(1).max(31).default(1).nullable().openapi({
+          description:
+            "The day verifications will refill each month, when interval is set to 'monthly'. Value is not zero-indexed making 1 the first day of the month. If left blank it will default to the last day of the month.",
+          example: 15,
+        }),
         lastRefillAt: z.number().int().optional().openapi({
           description: "The unix timestamp in miliseconds when the key was last refilled.",
           example: 100,
@@ -76,10 +81,12 @@ export const keySchema = z
         description:
           "Unkey allows you to refill remaining verifications on a key on a regular interval.",
         example: {
-          interval: "daily",
+          interval: "monthly",
           amount: 10,
+          refillDay: 10,
         },
       }),
+
     ratelimit: z
       .object({
         async: z.boolean().openapi({

apps/api/src/routes/v1_apis_listKeys.ts

@@ -317,6 +317,7 @@ export const registerV1ApisListKeys = (app: App) =>
             ? {
                 interval: k.refillInterval,
                 amount: k.refillAmount,
+                refillDay: k.refillInterval === "monthly" && k.refillDay ? k.refillDay : null,
                 lastRefillAt: k.lastRefillAt?.getTime(),
               }
             : undefined,

apps/api/src/routes/v1_keys_createKey.error.test.ts

@@ -11,7 +11,6 @@ import type { V1KeysCreateKeyRequest, V1KeysCreateKeyResponse } from "./v1_keys_
 test("when the api does not exist", async (t) => {
   const h = await IntegrationHarness.init(t);
   const apiId = newId("api");
-
   const root = await h.createRootKey([`api.${apiId}.create_key`]);
   /* The code snippet is making a POST request to the "/v1/keys.createKey" endpoint with the specified headers. It is using the `h.post` method from the `Harness` instance to send the request. The generic types `<V1KeysCreateKeyRequest, V1KeysCreateKeyResponse>` specify the request payload and response types respectively. */
 
@@ -119,3 +118,35 @@ test("when key recovery is not enabled", async (t) => {
     },
   });
 });
+
+test("reject invalid refill config when daily interval has non-null refillDay", 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<V1KeysCreateKeyRequest, V1KeysCreateKeyResponse>({
+    url: "/v1/keys.createKey",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${root.key}`,
+    },
+    body: {
+      byteLength: 16,
+      apiId: h.resources.userApi.id,
+      remaining: 10,
+      refill: {
+        amount: 100,
+        refillDay: 4,
+        interval: "daily",
+      },
+    },
+  });
+  expect(res.status).toEqual(400);
+  expect(res.body).toMatchObject({
+    error: {
+      code: "BAD_REQUEST",
+      docs: "https://unkey.dev/docs/api-reference/errors/code/BAD_REQUEST",
+      message: "when interval is set to 'daily', 'refillDay' must be null.",
+    },
+  });
+});

apps/api/src/routes/v1_keys_createKey.happy.test.ts

@@ -467,4 +467,37 @@ describe("with externalId", () => {
       expect(key!.identity!.id).toEqual(identity.id);
     });
   });
+  describe("Should default last day of month if none provided", () => {
+    test("should provide default value", async (t) => {
+      const h = await IntegrationHarness.init(t);
+      const date = new Date();
+      const lastDate = new Date(date.getFullYear(), date.getMonth() + 1, 0).getDate();
+      const root = await h.createRootKey([`api.${h.resources.userApi.id}.create_key`]);
+
+      const res = await h.post<V1KeysCreateKeyRequest, V1KeysCreateKeyResponse>({
+        url: "/v1/keys.createKey",
+        headers: {
+          "Content-Type": "application/json",
+          Authorization: `Bearer ${root.key}`,
+        },
+        body: {
+          apiId: h.resources.userApi.id,
+          remaining: 10,
+          refill: {
+            interval: "monthly",
+            amount: 20,
+            refillDay: undefined,
+          },
+        },
+      });
+
+      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),
+      });
+      expect(key).toBeDefined();
+      expect(key!.refillDay).toEqual(lastDate);
+    });
+  });
 });

apps/api/src/routes/v1_keys_createKey.ts

@@ -118,16 +118,27 @@ When validating a key, we will return this back to you, so you can clearly ident
                   description:
                     "The number of verifications to refill for each occurrence is determined individually for each key.",
                 }),
+                refillDay: z
+                  .number()
+                  .min(1)
+                  .max(31)
+                  .optional()
+                  .openapi({
+                    description: `The day of the month, when we will refill the remaining verifications. To refill on the 15th of each month, set 'refillDay': 15.
+                    If the day does not exist, for example you specified the 30th and it's february, we will refill them on the last day of the month instead.`,
+                  }),
               })
               .optional()
               .openapi({
                 description:
                   "Unkey enables you to refill verifications for each key at regular intervals.",
                 example: {
-                  interval: "daily",
+                  interval: "monthly",
                   amount: 100,
+                  refillDay: 15,
                 },
               }),
+
             ratelimit: z
               .object({
                 async: z
@@ -309,6 +320,12 @@ export const registerV1KeysCreateKey = (app: App) =>
         message: "remaining must be set if you are using refill.",
       });
     }
+    if (req.refill?.refillDay && req.refill.interval === "daily") {
+      throw new UnkeyApiError({
+        code: "BAD_REQUEST",
+        message: "when interval is set to 'daily', 'refillDay' must be null.",
+      });
+    }
     /**
      * Set up an api for production
      */
@@ -325,7 +342,15 @@ export const registerV1KeysCreateKey = (app: App) =>
         ? upsertIdentity(db.primary, authorizedWorkspaceId, externalId)
         : Promise.resolve(null),
     ]);
+    const date = new Date();
+    let lastDayOfMonth = new Date(date.getFullYear(), date.getMonth() + 1, 0).getDate();
     const newKey = await retry(5, async (attempt) => {
+      if (req.refill?.refillDay && req?.refill?.refillDay <= lastDayOfMonth) {
+        if (req.refill.interval === "monthly") {
+          lastDayOfMonth = req.refill.refillDay;
+        }
+      }
+
       if (attempt > 1) {
         logger.warn("retrying key creation", {
           attempt,
@@ -357,6 +382,7 @@ export const registerV1KeysCreateKey = (app: App) =>
         ratelimitDuration: req.ratelimit?.duration ?? req.ratelimit?.refillInterval,
         remaining: req.remaining,
         refillInterval: req.refill?.interval,
+        refillDay: req.refill?.interval === "monthly" ? lastDayOfMonth : null,
         refillAmount: req.refill?.amount,
         lastRefillAt: req.refill?.interval ? new Date() : null,
         deletedAt: null,

apps/api/src/routes/v1_keys_getKey.ts

@@ -155,6 +155,7 @@ export const registerV1KeysGetKey = (app: App) =>
           ? {
               interval: key.refillInterval,
               amount: key.refillAmount,
+              refillDay: key.refillInterval === "monthly" ? key.refillDay : null,
               lastRefillAt: key.lastRefillAt?.getTime(),
             }
           : undefined,

apps/api/src/routes/v1_keys_updateKey.error.test.ts

@@ -3,6 +3,7 @@ import { expect, test } from "vitest";
 import { newId } from "@unkey/id";
 import { IntegrationHarness } from "src/pkg/testutil/integration-harness";
 
+import type { ErrorResponse } from "@/pkg/errors";
 import type { V1KeysUpdateKeyRequest, V1KeysUpdateKeyResponse } from "./v1_keys_updateKey";
 
 test("when the key does not exist", async (t) => {
@@ -31,3 +32,34 @@ test("when the key does not exist", async (t) => {
     },
   });
 });
+test("reject invalid refill config", async (t) => {
+  const h = await IntegrationHarness.init(t);
+  const keyId = newId("test");
+  const root = await h.createRootKey([`api.${h.resources.userApi.id}.update_key`]);
+  /* The code snippet is making a POST request to the "/v1/keys.createKey" endpoint with the specified headers. It is using the `h.post` method from the `Harness` instance to send the request. The generic types `<V1KeysCreateKeyRequest, V1KeysCreateKeyResponse>` specify the request payload and response types respectively. */
+
+  const res = await h.post<V1KeysUpdateKeyRequest, ErrorResponse>({
+    url: "/v1/keys.updateKey",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${root.key}`,
+    },
+    body: {
+      keyId,
+      remaining: 10,
+      refill: {
+        amount: 100,
+        refillDay: 4,
+        interval: "daily",
+      },
+    },
+  });
+  expect(res.status).toEqual(400);
+  expect(res.body).toMatchObject({
+    error: {
+      code: "BAD_REQUEST",
+      docs: "https://unkey.dev/docs/api-reference/errors/code/BAD_REQUEST",
+      message: "when interval is set to 'daily', 'refillDay' must be null.",
+    },
+  });
+});

apps/api/src/routes/v1_keys_updateKey.happy.test.ts

@@ -1013,3 +1013,49 @@ 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);
+
+    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()),
+
+      createdAt: new Date(),
+    };
+    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<V1KeysUpdateKeyRequest, V1KeysUpdateKeyResponse>({
+      url: "/v1/keys.updateKey",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${root.key}`,
+      },
+      body: {
+        keyId: key.id,
+        remaining: 10,
+        refill: {
+          interval: "monthly",
+          amount: 130,
+          refillDay: undefined,
+        },
+        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),
+    });
+    expect(found).toBeDefined();
+    expect(found?.remaining).toEqual(10);
+    expect(found?.refillAmount).toEqual(130);
+    expect(found?.refillInterval).toEqual("monthly");
+    expect(found?.refillDay).toEqual(1);
+  });
+});

apps/api/src/routes/v1_keys_updateKey.ts

@@ -151,6 +151,10 @@ This field will become required in a future version.`,
                     description:
                       "The amount of verifications to refill for each occurrence is determined individually for each key.",
                   }),
+                  refillDay: z.number().min(1).max(31).optional().openapi({
+                    description:
+                      "The day verifications will refill each month, when interval is set to 'monthly'",
+                  }),
                 })
                 .nullable()
                 .optional()
@@ -276,9 +280,7 @@ export const registerV1KeysUpdate = (app: App) =>
   app.openapi(route, async (c) => {
     const req = c.req.valid("json");
     const { cache, db, usageLimiter, analytics, rbac } = c.get("services");
-
     const auth = await rootKeyAuth(c);
-
     const key = await db.primary.query.keys.findFirst({
       where: (table, { eq }) => eq(table.id, req.keyId),
       with: {
@@ -335,7 +337,12 @@ export const registerV1KeysUpdate = (app: App) =>
         message: "Cannot set refill on a key with unlimited requests",
       });
     }
-
+    if (req.refill?.refillDay && req.refill.interval === "daily") {
+      throw new UnkeyApiError({
+        code: "BAD_REQUEST",
+        message: "when interval is set to 'daily', 'refillDay' must be null.",
+      });
+    }
     const authorizedWorkspaceId = auth.authorizedWorkspaceId;
     const rootKeyId = auth.key.id;
 
@@ -346,7 +353,11 @@ export const registerV1KeysUpdate = (app: App) =>
         : externalId === null
           ? null
           : (await upsertIdentity(db.primary, authorizedWorkspaceId, externalId)).id;
-
+    const date = new Date();
+    let lastDayOfMonth = new Date(date.getFullYear(), date.getMonth() + 1, 0).getDate();
+    if (req.refill?.refillDay && req?.refill?.refillDay <= lastDayOfMonth) {
+      lastDayOfMonth = req.refill.refillDay;
+    }
     await db.primary
       .update(schema.keys)
       .set({
@@ -377,6 +388,7 @@ export const registerV1KeysUpdate = (app: App) =>
             : req.ratelimit?.duration ?? req.ratelimit?.refillInterval ?? null,
         refillInterval: req.refill === null ? null : req.refill?.interval,
         refillAmount: req.refill === null ? null : req.refill?.amount,
+        refillDay: req.refill?.interval === "monthly" ? lastDayOfMonth : null,
         lastRefillAt: req.refill == null || req.refill?.amount == null ? null : new Date(),
         enabled: req.enabled,
       })

apps/api/src/routes/v1_migrations_createKey.error.test.ts

@@ -112,3 +112,39 @@ test("reject invalid ratelimit config", async (t) => {
   expect(res.status).toEqual(400);
   expect(res.body.error.code).toEqual("BAD_REQUEST");
 });
+test("reject invalid refill config when daily interval has non-null refillDay", async (t) => {
+  const h = await IntegrationHarness.init(t);
+  const { key } = await h.createRootKey(["*"]);
+
+  const res = await h.post<V1MigrationsCreateKeysRequest, ErrorResponse>({
+    url: "/v1/migrations.createKeys",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${key}`,
+    },
+    body: [
+      {
+        start: "x",
+        hash: {
+          value: "x",
+          variant: "sha256_base64",
+        },
+        apiId: h.resources.userApi.id,
+        remaining: 10,
+        refill: {
+          amount: 100,
+          refillDay: 4,
+          interval: "daily",
+        },
+      },
+    ],
+  });
+  expect(res.status).toEqual(400);
+  expect(res.body).toMatchObject({
+    error: {
+      code: "BAD_REQUEST",
+      docs: "https://unkey.dev/docs/api-reference/errors/code/BAD_REQUEST",
+      message: "when interval is set to 'daily', 'refillDay' must be null.",
+    },
+  });
+});