fix(ratelimit): preserve bucket policy across runtime changes

src/ratelimit/token-bucket.ts

@@ -44,6 +44,8 @@ interface BucketState {
   tokens: number;
   updatedAtMs: number;
   lastSeenAtMs: number;
+  /** Policy that governed elapsed time since updatedAtMs. */
+  policy: RateLimitPolicy;
 }
 
 const SURFACE_ORDER: readonly RateLimitSurface[] = [
@@ -91,6 +93,10 @@ function bucketKey(surface: RateLimitSurface, principal: RateLimitPrincipal): st
   return `${surface}\0${principal.fingerprint}`;
 }
 
+function samePolicy(left: RateLimitPolicy, right: RateLimitPolicy): boolean {
+  return left.requestsPerMinute === right.requestsPerMinute && left.burst === right.burst;
+}
+
 /**
  * Process-local synchronous token buckets.
  *
@@ -141,19 +147,29 @@ export class TokenBucketLimiter {
     if (!state) {
       if (this.buckets.size >= this.maxBuckets) this.evictOneStale(now);
       if (this.buckets.size < this.maxBuckets) {
-        state = { tokens: policy.burst, updatedAtMs: now, lastSeenAtMs: now };
+        state = {
+          tokens: policy.burst,
+          updatedAtMs: now,
+          lastSeenAtMs: now,
+          policy,
+        };
         this.buckets.set(key, state);
       } else {
         source = "overflow";
         state = this.overflowBuckets.get(surface);
         if (!state) {
-          state = { tokens: policy.burst, updatedAtMs: now, lastSeenAtMs: now };
+          state = {
+            tokens: policy.burst,
+            updatedAtMs: now,
+            lastSeenAtMs: now,
+            policy,
+          };
           this.overflowBuckets.set(surface, state);
         }
       }
     }
 
-    this.refill(state, policy, now);
+    this.applyPolicy(state, policy, now);
     state.lastSeenAtMs = now;
     const allowed = state.tokens >= normalizedCost;
     if (allowed) state.tokens -= normalizedCost;
@@ -205,6 +221,18 @@ export class TokenBucketLimiter {
     return this.lastNowMs;
   }
 
+  /**
+   * Refill elapsed time under the policy that actually governed that interval. Only after the
+   * refill is committed do we install the new policy. Increasing burst never mints tokens;
+   * decreasing burst clamps existing balance to the new capacity.
+   */
+  private applyPolicy(state: BucketState, nextPolicy: RateLimitPolicy, now: number): void {
+    this.refill(state, state.policy, now);
+    if (samePolicy(state.policy, nextPolicy)) return;
+    state.tokens = Math.min(nextPolicy.burst, Math.max(0, state.tokens));
+    state.policy = nextPolicy;
+  }
+
   private refill(state: BucketState, policy: RateLimitPolicy, now: number): void {
     const elapsedMs = Math.max(0, now - state.updatedAtMs);
     const refill = elapsedMs * policy.requestsPerMinute / 60_000;

tests/ratelimit-policy-transition.test.ts

@@ -0,0 +1,106 @@
+import { describe, expect, test } from "bun:test";
+import { PrincipalFingerprinter, TokenBucketLimiter } from "../src/ratelimit";
+
+const SECRET = new Uint8Array(32).fill(0x33);
+
+function testPrincipals() {
+  const fingerprinter = new PrincipalFingerprinter(SECRET);
+  return {
+    a: fingerprinter.admissionKey("policy-transition-a"),
+    b: fingerprinter.admissionKey("policy-transition-b"),
+    c: fingerprinter.admissionKey("policy-transition-c"),
+  };
+}
+
+describe("token-bucket policy transitions", () => {
+  test("low-to-high changes do not retroactively refill elapsed time at the new rate", () => {
+    let now = 0;
+    const limiter = new TokenBucketLimiter({ now: () => now });
+    const { a } = testPrincipals();
+    const low = { requestsPerMinute: 1, burst: 1 };
+    const high = { requestsPerMinute: 600, burst: 10 };
+
+    expect(limiter.consume("responses-http", a, low).allowed).toBe(true);
+    now = 10_000;
+
+    const transitioned = limiter.consume("responses-http", a, high);
+    expect(transitioned).toMatchObject({
+      allowed: false,
+      limit: 10,
+      remaining: 0,
+      reason: "rate_limited",
+    });
+  });
+
+  test("high-to-low changes preserve refill earned under the prior policy before clamping", () => {
+    let now = 0;
+    const limiter = new TokenBucketLimiter({ now: () => now });
+    const { a } = testPrincipals();
+    const high = { requestsPerMinute: 600, burst: 10 };
+    const low = { requestsPerMinute: 1, burst: 1 };
+
+    for (let i = 0; i < high.burst; i += 1) {
+      expect(limiter.consume("responses-http", a, high).allowed).toBe(true);
+    }
+    expect(limiter.consume("responses-http", a, high).allowed).toBe(false);
+
+    now = 1_000;
+    expect(limiter.consume("responses-http", a, low)).toMatchObject({
+      allowed: true,
+      limit: 1,
+      remaining: 0,
+    });
+  });
+
+  test("raising burst capacity does not mint free tokens", () => {
+    const limiter = new TokenBucketLimiter({ now: () => 0 });
+    const { a } = testPrincipals();
+
+    expect(limiter.consume("images", a, { requestsPerMinute: 60, burst: 1 }).allowed).toBe(true);
+    expect(limiter.consume("images", a, { requestsPerMinute: 60, burst: 10 })).toMatchObject({
+      allowed: false,
+      limit: 10,
+      remaining: 0,
+    });
+  });
+
+  test("lowering burst capacity clamps the existing balance before admission", () => {
+    const limiter = new TokenBucketLimiter({ now: () => 0 });
+    const { a } = testPrincipals();
+
+    expect(limiter.consume("search", a, { requestsPerMinute: 60, burst: 10 })).toMatchObject({
+      allowed: true,
+      remaining: 9,
+    });
+    expect(limiter.consume("search", a, { requestsPerMinute: 60, burst: 2 })).toMatchObject({
+      allowed: true,
+      limit: 2,
+      remaining: 1,
+    });
+  });
+
+  test("shared overflow buckets use the same explicit policy transition", () => {
+    let now = 0;
+    const limiter = new TokenBucketLimiter({
+      maxBuckets: 1,
+      staleAfterMs: 60_000,
+      now: () => now,
+    });
+    const { a, b, c } = testPrincipals();
+    const low = { requestsPerMinute: 1, burst: 1 };
+    const high = { requestsPerMinute: 600, burst: 10 };
+
+    expect(limiter.consume("chat-completions", a, low).source).toBe("principal");
+    expect(limiter.consume("chat-completions", b, low)).toMatchObject({
+      allowed: true,
+      source: "overflow",
+    });
+
+    now = 10_000;
+    expect(limiter.consume("chat-completions", c, high)).toMatchObject({
+      allowed: false,
+      source: "overflow",
+      limit: 10,
+    });
+  });
+});