feat: schema for multiratelimits (#1775)

* feat: schema for multiratelimits

* refactor(agent.ts): add multiRatelimit function to handle ratelimiting
feat(agent-sdk): add service_connect.ts and service_pb.ts for ClusterService
feat(agent-sdk): add service_connect.ts for RatelimitService with multiRatelimit
feat(agent-sdk): add PushPull function to RatelimitService for syncing ratelimits

refactor(service_pb.ts): optimize import statements for better readability
refactor(service_pb.ts): reformat fields initialization for consistency
feat(service_pb.ts): add support for current field in RatelimitResponse
feat(service_pb.ts): introduce RatelimitMultiRequest and RatelimitMultiResponse
feat(service_pb.ts): implement PushPullRequest and PushPullResponse messages

feat(agent-sdk): add generated proto files for vault v1 object and service connect

feat(agent-sdk): add generated TypeScript classes for vault service messages

chore(metrics): add new enum value 'multiRatelimit' to 'metric.agent.latency' operation in metricSchema

refactor(openapi.d.ts): simplify XOR and OneOf type definitions
refactor(openapi.d.ts): reformat PermissionQuery type definition for clarity
refactor(openapi.d.ts): remove unnecessary empty object in parameters field

refactor(openapi.d.ts): update API key creation schema to include new fields and remove deprecated fields
feat(openapi.d.ts): add support for per-key ratelimiting in API key creation schema

refactor(openapi.d.ts): simplify code enum options for error codes in responses
docs(openapi.d.ts): update documentation for key creation in openapi definition

* wip

* fix(ratelimit/middleware.go): remove unnecessary empty line
feat(keys/service.ts): add MissingRatelimitError class for missing ratelimits

test(v1_keys_verifyKey.test.ts): skip tests with ratelimit override
test(v1_keys_verifyKey.test.ts): skip tests with ratelimit
fix(v1_keys_verifyKey.ts): import MissingRatelimitError in service
fix(deployment/docker-compose.yaml): update exposed ports range
chore(internal/db/package.json): remove dotenv from migrate script
fix(internal/db/src/schema/llm-gateway.ts): remove uniqueIndex on subdomain
fix(internal/db/src/schema/secrets.ts): remove index on workspaceId
chore(package.json): add migrate script for docker compose
chore(tools/local/package.json): remove drizzle-orm dependency
feat(tools/local/src/cmd/api.ts): add Logging configuration
fix(tools/local/src/db.ts): use mysqlDrizzle instead of drizzle in connectDatabase
chore(turbo.json): disable cache for migrate script

apps/agent/services/ratelimit/middleware.go

@@ -28,7 +28,6 @@ func (mw *tracingMiddleware) Ratelimit(ctx context.Context, req *ratelimitv1.Rat
 }
 
 func (mw *tracingMiddleware) MultiRatelimit(ctx context.Context, req *ratelimitv1.RatelimitMultiRequest) (res *ratelimitv1.RatelimitMultiResponse, err error) {
-
 	ctx, span := tracing.Start(ctx, tracing.NewSpanName("svc.ratelimit", "MultiRatelimit"))
 	defer span.End()
 

apps/api/src/pkg/agent.ts

@@ -45,5 +45,19 @@ export function connectAgent(
       });
       return res;
     },
+    multiRatelimit: async (...args: Parameters<Ratelimit["multiRatelimit"]>) => {
+      const [req] = args;
+      const start = performance.now();
+      const res = await ratelimit.multiRatelimit(req);
+      metrics.emit({
+        metric: "metric.agent.latency",
+        op: "multiRatelimit",
+        latency: performance.now() - start,
+      });
+      return res;
+    },
+    pushPull: async (..._args: Parameters<Ratelimit["pushPull"]>) => {
+      throw new Error("Not indented to be used");
+    },
   };
 }

apps/api/src/pkg/cache/namespaces.ts

@@ -1,8 +1,10 @@
 import type {
   Api,
   EncryptedKey,
+  Identity,
   Key,
   KeyAuth,
+  Ratelimit,
   RatelimitNamespace,
   RatelimitOverride,
 } from "@unkey/db";
@@ -28,6 +30,8 @@ export type CacheNamespaces = {
     api: Api;
     permissions: string[];
     roles: string[];
+    ratelimits: { [name: string]: Ratelimit };
+    identity: Identity | null;
   } | null;
   apiById: (Api & { keyAuth: KeyAuth | null }) | null;
   keysByOwnerId: {

apps/api/src/pkg/keys/service.ts

@@ -1,5 +1,5 @@
 import type { Cache } from "@/pkg/cache";
-import type { Api, Database, Key } from "@/pkg/db";
+import type { Api, Database, Key, Ratelimit } from "@/pkg/db";
 import type { Metrics } from "@/pkg/metrics";
 import type { RateLimiter } from "@/pkg/ratelimit";
 import type { UsageLimiter } from "@/pkg/usagelimit";
@@ -23,6 +23,19 @@ export class DisabledWorkspaceError extends BaseError<{ workspaceId: string }> {
   }
 }
 
+export class MissingRatelimitError extends BaseError<{ name: string }> {
+  public readonly retry = false;
+  public readonly name = MissingRatelimitError.name;
+  constructor(name: string) {
+    super({
+      message: `ratelimit "${name}" does not exist`,
+      context: {
+        name,
+      },
+    });
+  }
+}
+
 type NotFoundResponse = {
   valid: false;
   code: "NOT_FOUND";
@@ -67,6 +80,13 @@ type ValidResponse = {
 };
 type VerifyKeyResult = NotFoundResponse | InvalidResponse | ValidResponse;
 
+type RatelimitRequest = {
+  name: string;
+  cost?: number;
+  limit?: number;
+  duration?: number;
+};
+
 export class KeyService {
   private readonly cache: Cache;
   private readonly logger: Logger;
@@ -105,8 +125,14 @@ export class KeyService {
       apiId?: string;
       permissionQuery?: PermissionQuery;
       ratelimit?: { cost?: number };
+      ratelimits?: Array<RatelimitRequest>;
     },
-  ): Promise<Result<VerifyKeyResult, SchemaError | FetchError | DisabledWorkspaceError>> {
+  ): Promise<
+    Result<
+      VerifyKeyResult,
+      SchemaError | FetchError | DisabledWorkspaceError | MissingRatelimitError
+    >
+  > {
     try {
       const res = await this._verifyKey(c, req);
       if (res.err) {
@@ -180,8 +206,14 @@ export class KeyService {
       apiId?: string;
       permissionQuery?: PermissionQuery;
       ratelimit?: { cost?: number };
+      ratelimits?: Array<RatelimitRequest>;
     },
-  ): Promise<Result<VerifyKeyResult, FetchError | SchemaError | DisabledWorkspaceError>> {
+  ): Promise<
+    Result<
+      VerifyKeyResult,
+      FetchError | SchemaError | DisabledWorkspaceError | MissingRatelimitError
+    >
+  > {
     const keyHash = await this.hash(req.key);
     const { val: data, err } = await this.cache.keyByHash.swr(keyHash, async (hash) => {
       const dbStart = performance.now();
@@ -224,6 +256,12 @@ export class KeyService {
               api: true,
             },
           },
+          ratelimits: true,
+          identity: {
+            with: {
+              ratelimits: true,
+            },
+          },
         },
       });
       this.metrics.emit({
@@ -246,13 +284,28 @@ export class KeyService {
         ...dbRes.permissions.map((p) => p.permission.name),
         ...dbRes.roles.flatMap((r) => r.role.permissions.map((p) => p.permission.name)),
       ]);
+
+      /**
+       * Merge ratelimits from the identity and the key
+       * Key limits take pecedence
+       */
+      const ratelimits: { [name: string]: Ratelimit } = {};
+      for (const rl of dbRes.identity?.ratelimits ?? []) {
+        ratelimits[rl.name] = rl;
+      }
+      for (const rl of dbRes.ratelimits ?? []) {
+        ratelimits[rl.name] = rl;
+      }
+
       return {
         workspace: dbRes.workspace,
         forWorkspace: dbRes.forWorkspace,
         key: dbRes,
         api: dbRes.keyAuth.api,
         permissions: Array.from(permissions.values()),
         roles: dbRes.roles.map((r) => r.role.name),
+        identity: dbRes.identity,
+        ratelimits,
       };
     });
 
@@ -377,7 +430,49 @@ export class KeyService {
     /**
      * Ratelimiting
      */
-    const [pass, ratelimit] = await this.ratelimit(c, data.key, { cost: req.ratelimit?.cost });
+
+    const ratelimits: {
+      [name: string | "default"]: Required<RatelimitRequest>;
+    } = {};
+    if (
+      data.key.ratelimitAsync !== null &&
+      data.key.ratelimitDuration !== null &&
+      data.key.ratelimitLimit !== null
+    ) {
+      ratelimits.default = {
+        name: "default",
+        cost: req.ratelimit?.cost ?? 1,
+        limit: data.key.ratelimitLimit,
+        duration: data.key.ratelimitDuration,
+      };
+    }
+    for (const r of req.ratelimits ?? []) {
+      if (typeof r.limit !== "undefined" && typeof r.duration !== "undefined") {
+        ratelimits[r.name] = {
+          name: r.name,
+          cost: r.cost ?? 1,
+          limit: r.limit,
+          duration: r.duration,
+        };
+        continue;
+      }
+
+      const configured = data.ratelimits[r.name];
+      if (configured) {
+        ratelimits[configured.name] = {
+          name: configured.name,
+          cost: r.cost ?? 1,
+          limit: configured.limit,
+          duration: configured.duration,
+        };
+        continue;
+      }
+
+      return Err(new MissingRatelimitError(r.name));
+    }
+    console.warn(JSON.stringify({ req, ratelimits }, null, 2));
+
+    const [pass, ratelimit] = await this.ratelimit(c, data.key, ratelimits);
     if (!pass) {
       return Ok({
         key: data.key,
@@ -433,30 +528,27 @@ export class KeyService {
   private async ratelimit(
     c: Context,
     key: Key,
-    opts?: { cost?: number },
+    ratelimits: { [name: string | "default"]: Required<RatelimitRequest> },
   ): Promise<[boolean, VerifyKeyResult["ratelimit"]]> {
-    if (
-      key.ratelimitAsync === null ||
-      key.ratelimitLimit === null ||
-      key.ratelimitDuration === null
-    ) {
+    if (Object.keys(ratelimits).length === 0) {
       return [true, undefined];
     }
     if (!this.rateLimiter) {
-      this.logger.warn("ratelimiting is not enabled, but a key has ratelimiting enabled");
+      this.logger.error("ratelimiting is not enabled, but a key has ratelimiting enabled");
       return [true, undefined];
     }
 
-    const res = await this.rateLimiter.limit(c, {
-      workspaceId: key.workspaceId,
-      identifier: key.id,
-      limit: key.ratelimitLimit,
-      interval: key.ratelimitDuration,
-      cost: opts?.cost ?? 1,
-      // root keys are sharded per edge colo
-      shard: key.forWorkspaceId ? "edge" : undefined,
-      async: key.ratelimitAsync,
-    });
+    const res = await this.rateLimiter.multiLimit(
+      c,
+      Object.values(ratelimits).map((r) => ({
+        async: !!key.ratelimitAsync,
+        workspaceId: key.workspaceId,
+        identifier: key.id,
+        cost: r.cost,
+        interval: r.duration,
+        limit: r.limit,
+      })),
+    );
 
     if (res.err) {
       this.logger.error("ratelimiting failed", {
@@ -470,9 +562,9 @@ export class KeyService {
     return [
       res.val.pass,
       {
-        remaining: key.ratelimitLimit - res.val.current,
-        limit: key.ratelimitLimit,
-        reset: res.val.reset,
+        remaining: -1,
+        limit: -1,
+        reset: -1,
       },
     ];
   }

apps/api/src/pkg/ratelimit/client.ts

@@ -85,6 +85,23 @@ export class DurableRateLimiter implements RateLimiter {
     return res;
   }
 
+  /**
+   * Do not use
+   */
+  public async multiLimit(
+    c: Context,
+    req: Array<RatelimitRequest>,
+  ): Promise<Result<RatelimitResponse, RatelimitError>> {
+    const res = await Promise.all(req.map((r) => this.limit(c, r)));
+    for (const r of res) {
+      if (!r.val?.current) {
+        return r;
+      }
+    }
+
+    return Ok({ current: -1, pass: true, reset: -1 });
+  }
+
   private async _limit(
     c: Context,
     req: RatelimitRequest,

apps/api/src/pkg/ratelimit/interface.ts

@@ -35,4 +35,8 @@ export type RatelimitResponse = z.infer<typeof ratelimitResponseSchema>;
 
 export interface RateLimiter {
   limit: (c: Context, req: RatelimitRequest) => Promise<Result<RatelimitResponse, RatelimitError>>;
+  multiLimit: (
+    c: Context,
+    req: Array<RatelimitRequest>,
+  ) => Promise<Result<RatelimitResponse, RatelimitError>>;
 }

apps/api/src/pkg/ratelimit/noop.ts

@@ -9,4 +9,10 @@ export class NoopRateLimiter implements RateLimiter {
   ): Promise<Result<RatelimitResponse, RatelimitError>> {
     return Ok({ current: 0, pass: true, reset: 0 });
   }
+  public async multiLimit(
+    _c: Context,
+    _req: Array<RatelimitRequest>,
+  ): Promise<Result<RatelimitResponse, RatelimitError>> {
+    return Ok({ current: 0, pass: true, reset: 0 });
+  }
 }