Client-Side Token Bucket for SDKs

// bucket.ts — the pacing primitive an SDK schedules against.
export interface BucketOptions {
  ratePerSec: number;          // 0.9 x published rate
  capacity: number;            // 0.9 x published burst
  maxWaitMs?: number;          // reject rather than queue forever
}

const now = () =>
  typeof performance !== "undefined" ? performance.now() : Number(process.hrtime.bigint() / 1_000_000n);

export class TokenBucket {
  private tokens: number;
  private last = now();
  private baseRate: number;
  private rate: number;
  private pausedUntil = 0;
  private waiters = 0;

  constructor(private opts: BucketOptions) {
    this.tokens = opts.capacity;
    this.baseRate = opts.ratePerSec;
    this.rate = opts.ratePerSec;
  }

  private refill() {
    const t = now();
    // Monotonic source + clamp: a clock correction can never mint tokens.
    const elapsed = Math.max(0, t - this.last);
    this.tokens = Math.min(this.opts.capacity, this.tokens + (elapsed / 1000) * this.rate);
    this.last = t;
  }

  /** Resolves when a token has been taken; rejects if the wait exceeds maxWaitMs. */
  async acquire(cost = 1): Promise<void> {
    const deadline = now() + (this.opts.maxWaitMs ?? 30_000);
    this.waiters += 1;
    try {
      for (;;) {
        const t = now();
        if (t < this.pausedUntil) {                    // the server told us to wait
          await sleep(Math.min(this.pausedUntil - t, deadline - t));
        }
        this.refill();
        if (this.tokens >= cost) { this.tokens -= cost; return; }
        const needMs = ((cost - this.tokens) / this.rate) * 1000;
        if (now() + needMs > deadline) throw new Error("rate_limit_wait_exceeded");
        await sleep(Math.min(needMs, 250));            // re-check periodically
      }
    } finally { this.waiters -= 1; }
  }

  /** The server publishes the truth; trust it over the local estimate. */
  syncFromHeaders(headers: { get(name: string): string | null }) {
    const remaining = Number(headers.get("X-RateLimit-Remaining") ?? NaN);
    if (Number.isFinite(remaining)) {
      // Never increase our own optimism: take the smaller of the two views.
      this.tokens = Math.min(this.tokens, remaining);
    }
    const reset = Number(headers.get("X-RateLimit-Reset") ?? NaN);
    if (Number.isFinite(reset) && remaining === 0) this.pausedUntil = now() + reset * 1000;
  }

  onRejected(retryAfterSeconds?: number) {
    this.rate = Math.max(this.baseRate * 0.1, this.rate * 0.5);   // multiplicative decrease
    this.tokens = 0;
    if (retryAfterSeconds && retryAfterSeconds > 0) {
      this.pausedUntil = now() + retryAfterSeconds * 1000;
    }
  }

  onCleanWindow() { this.rate = Math.min(this.baseRate, this.rate * 1.1); }

  get state() { return { tokens: Math.floor(this.tokens), rate: this.rate, waiters: this.waiters }; }
}

const sleep = (ms: number) => new Promise((r) => setTimeout(r, Math.max(0, ms)));

// bucket.test.ts — deterministic, with an injectable clock.
import { describe, it, expect, vi, beforeEach } from "vitest";
import { TokenBucket } from "./bucket";

beforeEach(() => vi.useFakeTimers());

it("releases the burst immediately, then paces", async () => {
  const b = new TokenBucket({ ratePerSec: 1.5, capacity: 18 });
  const sent: number[] = [];
  const run = async () => { for (let i = 0; i < 22; i += 1) { await b.acquire(); sent.push(Date.now()); } };
  const p = run();
  await vi.advanceTimersByTimeAsync(0);
  expect(sent.length).toBe(18);                       // the whole burst, at once
  await vi.advanceTimersByTimeAsync(3000);            // 4.5 tokens refilled
  expect(sent.length).toBe(22);
  await p;
});

it("never mints tokens when the wall clock jumps", async () => {
  const b = new TokenBucket({ ratePerSec: 1, capacity: 5 });
  for (let i = 0; i < 5; i += 1) await b.acquire();
  vi.setSystemTime(Date.now() - 60_000);              // clock steps backwards
  expect(b.state.tokens).toBe(0);                     // monotonic source is unaffected
});

it("adopts the server's remaining count when it is lower", () => {
  const b = new TokenBucket({ ratePerSec: 10, capacity: 20 });
  b.syncFromHeaders(new Headers({ "X-RateLimit-Remaining": "3" }));
  expect(b.state.tokens).toBe(3);
});

it("halves the rate on rejection and recovers on a clean window", () => {
  const b = new TokenBucket({ ratePerSec: 10, capacity: 20 });
  b.onRejected(2);
  expect(b.state.rate).toBe(5);
  b.onCleanWindow();
  expect(b.state.rate).toBeCloseTo(5.5, 5);
});