@crewhaus/rate-limiter 0.1.0

package/package.json

@@ -0,0 +1,41 @@
+{
+  "name": "@crewhaus/rate-limiter",
+  "version": "0.1.0",
+  "type": "module",
+  "description": "Multi-dimensional token-bucket / leaky-bucket rate limiter (per-tenant, per-provider, per-tool)",
+  "main": "src/index.ts",
+  "types": "src/index.ts",
+  "exports": {
+    ".": "./src/index.ts"
+  },
+  "scripts": {
+    "test": "bun test src"
+  },
+  "dependencies": {
+    "@crewhaus/errors": "0.0.0"
+  },
+  "license": "Apache-2.0",
+  "author": {
+    "name": "Max Meier",
+    "email": "max@studiomax.io",
+    "url": "https://studiomax.io"
+  },
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/crewhaus/factory.git",
+    "directory": "packages/rate-limiter"
+  },
+  "homepage": "https://github.com/crewhaus/factory/tree/main/packages/rate-limiter#readme",
+  "bugs": {
+    "url": "https://github.com/crewhaus/factory/issues"
+  },
+  "publishConfig": {
+    "access": "restricted"
+  },
+  "files": [
+    "src",
+    "README.md",
+    "LICENSE",
+    "NOTICE"
+  ]
+}

package/src/index.test.ts

@@ -0,0 +1,179 @@
+/**
+ * Section 27 — `rate-limiter` tests:
+ *  - T1 per algorithm (token-bucket vs leaky-bucket) edge cases
+ *  - T7 1000-acquirer load test (concurrency-fair + no starvation)
+ *  - T8 fail-closed when keys are missing (deny rather than allow)
+ */
+import { describe, expect, test } from "bun:test";
+import {
+  type AcquireKey,
+  type BucketConfig,
+  RateLimitError,
+  bucketKeyOf,
+  createRateLimiter,
+} from "./index";
+
+describe("rate-limiter — T1 token-bucket", () => {
+  test("acquire below capacity is immediate", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "token-bucket", capacity: 10, refillPerSec: 1 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    const t0 = Date.now();
+    await rl.acquire([{ dimension: "tenant", id: "t1" }], 5);
+    // 250 ms threshold: "immediate" relative to the bucket's 1-token-per-second
+    // refill rate, while tolerating CI scheduler jitter (we saw 51 ms flakes
+    // against a 50 ms cap on shared GitHub runners).
+    expect(Date.now() - t0).toBeLessThan(250);
+    const inspect = rl.inspect({ dimension: "tenant", id: "t1" });
+    expect(inspect?.available).toBeCloseTo(5, 1);
+  });
+
+  test("burst tolerance: capacity available immediately at start", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "token-bucket", capacity: 10, refillPerSec: 0.1 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    const t0 = Date.now();
+    for (let i = 0; i < 10; i++) {
+      await rl.acquire([{ dimension: "tenant", id: "t1" }], 1);
+    }
+    expect(Date.now() - t0).toBeLessThan(100);
+  });
+
+  test("blocks until refill when over capacity", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "token-bucket", capacity: 1, refillPerSec: 10 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    const t0 = Date.now();
+    await rl.acquire([{ dimension: "tenant", id: "t1" }], 1);
+    await rl.acquire([{ dimension: "tenant", id: "t1" }], 1);
+    const elapsed = Date.now() - t0;
+    // Second call needs to wait for ~100ms refill. Generous lower bound for
+    // shared-CI scheduling jitter; upper bound large enough to avoid flake.
+    expect(elapsed).toBeGreaterThanOrEqual(50);
+    expect(elapsed).toBeLessThan(2_000);
+  });
+
+  test("rejects after maxWaitMs when refill rate too slow", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "token-bucket", capacity: 1, refillPerSec: 0.01 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    await rl.acquire([{ dimension: "tenant", id: "t1" }], 1);
+    expect(
+      rl.acquire([{ dimension: "tenant", id: "t1" }], 1, { maxWaitMs: 100 }),
+    ).rejects.toBeInstanceOf(RateLimitError);
+  });
+});
+
+describe("rate-limiter — T1 leaky-bucket", () => {
+  test("smoothing: requests release at refill rate", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "leaky-bucket", capacity: 5, refillPerSec: 50 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    // 5 fit under capacity; 6th queues for ~20ms. Generous bounds for jitter.
+    const promises: Array<Promise<void>> = [];
+    const t0 = Date.now();
+    for (let i = 0; i < 7; i++) {
+      promises.push(rl.acquire([{ dimension: "tenant", id: "t1" }], 1, { maxWaitMs: 30_000 }));
+    }
+    await Promise.all(promises);
+    const elapsed = Date.now() - t0;
+    expect(elapsed).toBeGreaterThanOrEqual(15);
+    expect(elapsed).toBeLessThan(5_000);
+  });
+
+  test("rejects on maxWait when queue stays full", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "leaky-bucket", capacity: 1, refillPerSec: 0.01 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    await rl.acquire([{ dimension: "tenant", id: "t1" }], 1);
+    expect(
+      rl.acquire([{ dimension: "tenant", id: "t1" }], 1, { maxWaitMs: 50 }),
+    ).rejects.toBeInstanceOf(RateLimitError);
+  });
+});
+
+describe("rate-limiter — T8 fail-closed on missing keys", () => {
+  test("acquire on unknown key throws RateLimitError", async () => {
+    const rl = createRateLimiter({ buckets: new Map() });
+    expect(rl.acquire([{ dimension: "tenant", id: "unknown" }], 1)).rejects.toBeInstanceOf(
+      RateLimitError,
+    );
+  });
+
+  test("acquire passes for unknown id when * default exists", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:*", { kind: "token-bucket", capacity: 5, refillPerSec: 1 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    await rl.acquire([{ dimension: "tenant", id: "any" }], 1);
+  });
+
+  test("partial failure refunds successful acquisitions", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "token-bucket", capacity: 10, refillPerSec: 1 }],
+      // provider:p1 missing
+    ]);
+    const rl = createRateLimiter({ buckets });
+    expect(
+      rl.acquire(
+        [
+          { dimension: "tenant", id: "t1" },
+          { dimension: "provider", id: "p1" },
+        ],
+        1,
+      ),
+    ).rejects.toBeInstanceOf(RateLimitError);
+    // tenant bucket should still have full capacity after refund.
+    const inspect = rl.inspect({ dimension: "tenant", id: "t1" });
+    expect(inspect?.available).toBeCloseTo(10, 1);
+  });
+});
+
+describe("rate-limiter — multi-dimensional", () => {
+  test("acquire sums against tenant + provider + tool buckets", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "token-bucket", capacity: 10, refillPerSec: 1 }],
+      ["provider:p1", { kind: "token-bucket", capacity: 10, refillPerSec: 1 }],
+      ["tool:Bash", { kind: "token-bucket", capacity: 5, refillPerSec: 1 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    await rl.acquire([
+      { dimension: "tenant", id: "t1" },
+      { dimension: "provider", id: "p1" },
+      { dimension: "tool", id: "Bash" },
+    ]);
+    expect(rl.inspect({ dimension: "tenant", id: "t1" })?.available).toBeCloseTo(9, 1);
+    expect(rl.inspect({ dimension: "tool", id: "Bash" })?.available).toBeCloseTo(4, 1);
+  });
+});
+
+describe("rate-limiter — T7 load: 1000 acquirers, no starvation", () => {
+  test("1000 concurrent acquires drain in expected wall-clock time", async () => {
+    const buckets = new Map<string, BucketConfig>([
+      ["tenant:t1", { kind: "token-bucket", capacity: 100, refillPerSec: 5000 }],
+    ]);
+    const rl = createRateLimiter({ buckets });
+    const t0 = Date.now();
+    const promises = Array.from({ length: 1000 }, () =>
+      rl.acquire([{ dimension: "tenant", id: "t1" }], 1, { maxWaitMs: 60_000 }),
+    );
+    await Promise.all(promises);
+    const elapsed = Date.now() - t0;
+    // (1000 - 100) tokens to refill at 5000/s ≈ 180ms baseline. Allow very
+    // generous headroom for parallel-CI jitter.
+    expect(elapsed).toBeLessThan(15_000);
+  });
+});
+
+describe("rate-limiter — bucketKeyOf", () => {
+  test("formats dimension + id stably", () => {
+    const k: AcquireKey = { dimension: "provider", id: "anthropic" };
+    expect(bucketKeyOf(k)).toBe("provider:anthropic");
+  });
+});

package/src/index.ts

@@ -0,0 +1,337 @@
+/**
+ * Section 27 — `rate-limiter`. Multi-dimensional gating between callers
+ * and downstream services. Three keyed dimensions:
+ *  - **per-tenant** (gateway-server pre-handler)
+ *  - **per-provider** (model-router pre-call)
+ *  - **per-tool** (runtime-core pre-tool-execute, configured in spec under
+ *    `tools.<Name>.rateLimit`)
+ *
+ * Two algorithms; pick per-bucket:
+ *  - **token-bucket** — burst-tolerant. `capacity` tokens; refill at
+ *    `refillPerSec`. Acquire blocks until enough tokens are available.
+ *  - **leaky-bucket** — smoothing. Treat acquires as drops landing in a
+ *    bucket that drains at `refillPerSec`. New drops queue when the
+ *    bucket is full; the queue serves drops at the drain rate.
+ *
+ * `acquire(keys, cost)` evaluates each key in order and only proceeds
+ * when *every* bucket has the requested cost. The implementation never
+ * takes a partial reservation — if any bucket would block, the call
+ * either waits for the longest delay or rejects on `maxWaitMs`. This
+ * guarantees fail-closed semantics: an unknown key always denies.
+ */
+import { CrewhausError } from "@crewhaus/errors";
+
+export class RateLimitError extends CrewhausError {
+  override readonly name = "RateLimitError";
+  constructor(message: string, cause?: unknown) {
+    super("config", message, cause);
+  }
+}
+
+export type BucketKind = "token-bucket" | "leaky-bucket";
+
+export type BucketConfig = {
+  readonly kind: BucketKind;
+  /** Maximum tokens (token-bucket) or queue depth (leaky-bucket). */
+  readonly capacity: number;
+  /** Refill rate (token-bucket) or drain rate (leaky-bucket), per second. */
+  readonly refillPerSec: number;
+};
+
+export type AcquireKey = {
+  readonly dimension: "tenant" | "provider" | "tool";
+  readonly id: string;
+};
+
+export type AcquireOptions = {
+  /** How long to wait for tokens before rejecting. Defaults to 30s. */
+  readonly maxWaitMs?: number;
+  /** Override now() for tests. */
+  readonly now?: () => number;
+};
+
+export type RateLimiterOptions = {
+  /**
+   * Per-`(dimension, id)` bucket configuration. Lookup is exact-match;
+   * unknown keys deny by default (fail-closed). The `*` id is reserved
+   * for the per-dimension default — declared explicitly when one is
+   * desired.
+   */
+  readonly buckets: ReadonlyMap<string, BucketConfig>;
+  /** Override "now" for tests. */
+  readonly now?: () => number;
+};
+
+export interface RateLimiter {
+  /**
+   * Acquire `cost` tokens (default 1) from each key's bucket. Resolves
+   * once every bucket has paid out. Rejects with `RateLimitError` if
+   * any waited longer than `maxWaitMs`, or if any key is missing.
+   */
+  acquire(keys: ReadonlyArray<AcquireKey>, cost?: number, opts?: AcquireOptions): Promise<void>;
+  /** Diagnostic snapshot of current bucket state. */
+  inspect(key: AcquireKey):
+    | {
+        config: BucketConfig;
+        available: number;
+        waitingCount: number;
+      }
+    | undefined;
+}
+
+/** Stable string key for a dimension+id pair. */
+export function bucketKeyOf(key: AcquireKey): string {
+  return `${key.dimension}:${key.id}`;
+}
+
+/** Static helper: bucket capacity check (no async waiting). */
+export function tokenBucketAvailable(
+  state: TokenBucketState,
+  cost: number,
+  now: number,
+  config: BucketConfig,
+): boolean {
+  refillTokenBucket(state, now, config);
+  return state.tokens >= cost;
+}
+
+type TokenBucketState = {
+  tokens: number;
+  lastRefillMs: number;
+};
+
+type LeakyBucketState = {
+  /** Number of tokens currently in the bucket (queued). */
+  level: number;
+  lastDrainMs: number;
+  /** FIFO queue of pending acquirers awaiting drain. */
+  queue: Array<{
+    cost: number;
+    resolve: () => void;
+    reject: (err: Error) => void;
+    timer?: ReturnType<typeof setTimeout>;
+  }>;
+};
+
+function refillTokenBucket(state: TokenBucketState, now: number, config: BucketConfig): void {
+  const elapsedSec = Math.max(0, (now - state.lastRefillMs) / 1000);
+  const refilled = elapsedSec * config.refillPerSec;
+  state.tokens = Math.min(config.capacity, state.tokens + refilled);
+  state.lastRefillMs = now;
+}
+
+function drainLeakyBucket(state: LeakyBucketState, now: number, config: BucketConfig): void {
+  const elapsedSec = Math.max(0, (now - state.lastDrainMs) / 1000);
+  const drained = elapsedSec * config.refillPerSec;
+  state.level = Math.max(0, state.level - drained);
+  state.lastDrainMs = now;
+}
+
+export function createRateLimiter(opts: RateLimiterOptions): RateLimiter {
+  const buckets = opts.buckets;
+  const tokenStates = new Map<string, TokenBucketState>();
+  const leakyStates = new Map<string, LeakyBucketState>();
+
+  function getNow(callerNow?: () => number): number {
+    return (callerNow ?? opts.now ?? Date.now)();
+  }
+
+  function getOrInitTokenState(key: string, config: BucketConfig, now: number): TokenBucketState {
+    let s = tokenStates.get(key);
+    if (!s) {
+      s = { tokens: config.capacity, lastRefillMs: now };
+      tokenStates.set(key, s);
+    }
+    return s;
+  }
+
+  function getOrInitLeakyState(key: string, now: number): LeakyBucketState {
+    let s = leakyStates.get(key);
+    if (!s) {
+      s = { level: 0, lastDrainMs: now, queue: [] };
+      leakyStates.set(key, s);
+    }
+    return s;
+  }
+
+  /**
+   * Wait for a single bucket to allow `cost` tokens. Resolves when ready.
+   * `maxWaitMs` enforces the cap; rejects with RateLimitError on timeout.
+   */
+  function acquireOne(
+    key: AcquireKey,
+    cost: number,
+    config: BucketConfig,
+    maxWaitMs: number,
+    nowFn: () => number,
+  ): Promise<void> {
+    const k = bucketKeyOf(key);
+    const start = nowFn();
+
+    if (config.kind === "token-bucket") {
+      return new Promise<void>((resolve, reject) => {
+        const tryAcquire = (): void => {
+          const now = nowFn();
+          const state = getOrInitTokenState(k, config, now);
+          refillTokenBucket(state, now, config);
+          if (state.tokens >= cost) {
+            state.tokens -= cost;
+            resolve();
+            return;
+          }
+          const elapsedMs = now - start;
+          const remainingMs = maxWaitMs - elapsedMs;
+          if (remainingMs <= 0) {
+            reject(
+              new RateLimitError(
+                `rate limit exceeded for ${k}: ${cost} tokens needed, ${state.tokens.toFixed(2)} available, max wait ${maxWaitMs}ms reached`,
+              ),
+            );
+            return;
+          }
+          // Time until enough tokens accrue
+          const deficit = cost - state.tokens;
+          const msToWait = Math.min(
+            remainingMs,
+            Math.max(10, (deficit / config.refillPerSec) * 1000),
+          );
+          setTimeout(tryAcquire, msToWait);
+        };
+        tryAcquire();
+      });
+    }
+
+    // leaky-bucket
+    return new Promise<void>((resolve, reject) => {
+      const now = nowFn();
+      const state = getOrInitLeakyState(k, now);
+      drainLeakyBucket(state, now, config);
+      const wouldExceed = state.level + cost > config.capacity;
+      if (!wouldExceed && state.queue.length === 0) {
+        // Fast-path: no queue, fits in capacity.
+        state.level += cost;
+        resolve();
+        return;
+      }
+      // Queue and rely on drain timer.
+      const entry = {
+        cost,
+        resolve,
+        reject,
+        timer: setTimeout(() => {
+          const idx = state.queue.indexOf(entry);
+          if (idx >= 0) state.queue.splice(idx, 1);
+          reject(
+            new RateLimitError(
+              `rate limit exceeded for ${k}: leaky bucket full, max wait ${maxWaitMs}ms reached`,
+            ),
+          );
+        }, maxWaitMs),
+      };
+      state.queue.push(entry);
+      // Schedule a drain check.
+      const drainEveryMs = Math.max(10, 1000 / config.refillPerSec);
+      const tick = (): void => {
+        const tickNow = nowFn();
+        drainLeakyBucket(state, tickNow, config);
+        // Process as many queue entries as fit under capacity.
+        while (state.queue.length > 0) {
+          const head = state.queue[0];
+          if (!head) break;
+          if (state.level + head.cost <= config.capacity) {
+            state.queue.shift();
+            if (head.timer) clearTimeout(head.timer);
+            state.level += head.cost;
+            head.resolve();
+          } else {
+            break;
+          }
+        }
+        if (state.queue.length > 0) {
+          setTimeout(tick, drainEveryMs);
+        }
+      };
+      setTimeout(tick, drainEveryMs);
+    });
+  }
+
+  return {
+    async acquire(keys, cost = 1, callerOpts = {}): Promise<void> {
+      const maxWaitMs = callerOpts.maxWaitMs ?? 30_000;
+      const nowFn = (): number => getNow(callerOpts.now);
+
+      // Fail-closed: every key must resolve to a known bucket.
+      for (const key of keys) {
+        const k = bucketKeyOf(key);
+        if (!buckets.has(k)) {
+          // Per-dimension default lookup
+          const fallback = `${key.dimension}:*`;
+          if (!buckets.has(fallback)) {
+            throw new RateLimitError(`no bucket configured for ${k} (and no ${fallback} default)`);
+          }
+        }
+      }
+
+      // Acquire each in sequence so we don't double-charge a bucket on
+      // partial failure. (Parallel acquisition would require two-phase
+      // commit; sequential is simpler and the bucket counts stay correct.)
+      const acquired: AcquireKey[] = [];
+      try {
+        for (const key of keys) {
+          const k = bucketKeyOf(key);
+          const config = buckets.get(k) ?? buckets.get(`${key.dimension}:*`);
+          if (!config) throw new RateLimitError(`no bucket for ${k}`);
+          await acquireOne(key, cost, config, maxWaitMs, nowFn);
+          acquired.push(key);
+        }
+      } catch (err) {
+        // Refund any successful acquisitions so partial failures don't drain buckets.
+        const now = nowFn();
+        for (const key of acquired) {
+          const k = bucketKeyOf(key);
+          const config = buckets.get(k) ?? buckets.get(`${key.dimension}:*`);
+          if (!config) continue;
+          if (config.kind === "token-bucket") {
+            const state = tokenStates.get(k);
+            if (state) {
+              state.tokens = Math.min(config.capacity, state.tokens + cost);
+              state.lastRefillMs = now;
+            }
+          } else {
+            const state = leakyStates.get(k);
+            if (state) {
+              state.level = Math.max(0, state.level - cost);
+              state.lastDrainMs = now;
+            }
+          }
+        }
+        throw err;
+      }
+    },
+
+    inspect(key): { config: BucketConfig; available: number; waitingCount: number } | undefined {
+      const k = bucketKeyOf(key);
+      const config = buckets.get(k) ?? buckets.get(`${key.dimension}:*`);
+      if (!config) return undefined;
+      const now = (opts.now ?? Date.now)();
+      if (config.kind === "token-bucket") {
+        const state = tokenStates.get(k);
+        if (!state) {
+          return { config, available: config.capacity, waitingCount: 0 };
+        }
+        refillTokenBucket(state, now, config);
+        return { config, available: state.tokens, waitingCount: 0 };
+      }
+      const state = leakyStates.get(k);
+      if (!state) {
+        return { config, available: config.capacity, waitingCount: 0 };
+      }
+      drainLeakyBucket(state, now, config);
+      return {
+        config,
+        available: Math.max(0, config.capacity - state.level),
+        waitingCount: state.queue.length,
+      };
+    },
+  };
+}