feat(orchestrator): per-client token-bucket rate limiter for LLM routes

orchestrator/.env.example

@@ -19,6 +19,10 @@ MONTHLY_BUDGET=200
 CIRCUIT_BREAKER_THRESHOLD=3
 CIRCUIT_BREAKER_COOLDOWN_MS=60000
 
+# Rate Limiting (per-client token bucket on the LLM-backed routes)
+RATE_LIMIT_CAPACITY=60
+RATE_LIMIT_REFILL_PER_SEC=1
+
 # Redis (for Hybrid RAG keyword search)
 REDIS_URL=redis://localhost:6379
 

orchestrator/README.md

@@ -97,6 +97,10 @@ AI_ML_SERVICE_URL=http://localhost:8000
 CIRCUIT_BREAKER_THRESHOLD=3
 CIRCUIT_BREAKER_COOLDOWN_MS=60000
 
+# Rate limiting (per-client token bucket on LLM-backed routes)
+RATE_LIMIT_CAPACITY=60
+RATE_LIMIT_REFILL_PER_SEC=1
+
 # Cost budgets (USD)
 DAILY_BUDGET=10
 MONTHLY_BUDGET=200
@@ -180,6 +184,21 @@ Full system health report.
 }
 ```
 
+### `GET /api/rate-limit`
+
+Current token-bucket rate-limiter configuration and the number of active client buckets.
+
+**Response:**
+```json
+{
+  "activeKeys": 3,
+  "capacity": 60,
+  "refillPerSec": 1
+}
+```
+
+The expensive LLM-backed routes (`/api/supervisor/process`, `/api/agent/run`, `/api/batch/process`) are guarded by a per-client token bucket. Each response carries `X-RateLimit-Limit` and `X-RateLimit-Remaining` headers; when the bucket is empty the route returns **`429 Too Many Requests`** with a `Retry-After` header and `{ "error": "Rate limit exceeded", "retryAfterMs": <ms> }`. Tune with `RATE_LIMIT_CAPACITY` / `RATE_LIMIT_REFILL_PER_SEC`.
+
 ### `POST /api/supervisor/process`
 
 Route a request through the full supervisor pipeline (classify, budget check, decompose, dispatch, aggregate).

orchestrator/__tests__/rate-limiter.test.js

@@ -0,0 +1,202 @@
+/**
+ * Unit tests for the token-bucket RateLimiter and its Express middleware.
+ * Run: cd orchestrator && npx jest rate-limiter --no-cache
+ *
+ * A controllable `now()` clock is injected so refill behavior is deterministic
+ * (no real timers / sleeps).
+ */
+
+const { RateLimiter, rateLimitMiddleware } = require("../core/rate-limiter");
+
+// Minimal Express res/next stubs so the middleware can be exercised in isolation.
+function makeRes() {
+  return {
+    statusCode: 200,
+    headers: {},
+    body: undefined,
+    set(key, value) {
+      this.headers[key] = String(value);
+      return this;
+    },
+    status(code) {
+      this.statusCode = code;
+      return this;
+    },
+    json(payload) {
+      this.body = payload;
+      return this;
+    },
+  };
+}
+
+// ============ RateLimiter ============
+describe("RateLimiter", () => {
+  let clock;
+  const now = () => clock;
+
+  beforeEach(() => {
+    clock = 1_000;
+  });
+
+  test("allows requests up to capacity then blocks", () => {
+    const rl = new RateLimiter({ capacity: 3, refillPerSec: 1, now });
+    expect(rl.tryConsume("a").allowed).toBe(true);
+    expect(rl.tryConsume("a").allowed).toBe(true);
+    expect(rl.tryConsume("a").allowed).toBe(true);
+    expect(rl.tryConsume("a").allowed).toBe(false);
+  });
+
+  test("reports remaining tokens and the limit", () => {
+    const rl = new RateLimiter({ capacity: 5, refillPerSec: 1, now });
+    const first = rl.tryConsume("a");
+    expect(first.limit).toBe(5);
+    expect(first.remaining).toBe(4);
+    expect(rl.tryConsume("a").remaining).toBe(3);
+  });
+
+  test("a blocked request reports a positive retryAfterMs", () => {
+    const rl = new RateLimiter({ capacity: 1, refillPerSec: 1, now });
+    rl.tryConsume("a"); // drains the bucket
+    const blocked = rl.tryConsume("a");
+    expect(blocked.allowed).toBe(false);
+    expect(blocked.remaining).toBe(0);
+    expect(blocked.retryAfterMs).toBe(1000); // 1 token / 1 per sec
+  });
+
+  test("refills tokens over time at refillPerSec", () => {
+    const rl = new RateLimiter({ capacity: 5, refillPerSec: 1, now });
+    for (let i = 0; i < 5; i++) rl.tryConsume("a"); // drain
+    expect(rl.tryConsume("a").allowed).toBe(false);
+
+    clock += 2_000; // 2 seconds → +2 tokens
+    const after = rl.tryConsume("a");
+    expect(after.allowed).toBe(true);
+    expect(after.remaining).toBe(1);
+  });
+
+  test("refill never exceeds capacity", () => {
+    const rl = new RateLimiter({ capacity: 3, refillPerSec: 10, now });
+    rl.tryConsume("a"); // tokens: 2
+    clock += 60_000; // would add 600 tokens, but caps at capacity
+    expect(rl.getState("a").tokens).toBe(3);
+  });
+
+  test("keeps a separate bucket per key", () => {
+    const rl = new RateLimiter({ capacity: 1, refillPerSec: 1, now });
+    expect(rl.tryConsume("a").allowed).toBe(true);
+    expect(rl.tryConsume("a").allowed).toBe(false);
+    // A different key is unaffected.
+    expect(rl.tryConsume("b").allowed).toBe(true);
+  });
+
+  test("supports a cost greater than one", () => {
+    const rl = new RateLimiter({ capacity: 10, refillPerSec: 1, now });
+    const r = rl.tryConsume("a", 4);
+    expect(r.allowed).toBe(true);
+    expect(r.remaining).toBe(6);
+    // Not enough tokens for a cost of 7 → blocked, bucket unchanged.
+    const blocked = rl.tryConsume("a", 7);
+    expect(blocked.allowed).toBe(false);
+    expect(rl.getState("a").tokens).toBe(6);
+  });
+
+  test("getState returns a full bucket for an unknown key", () => {
+    const rl = new RateLimiter({ capacity: 4, refillPerSec: 1, now });
+    expect(rl.getState("never-seen")).toEqual({ tokens: 4, capacity: 4 });
+  });
+
+  test("reset restores a single key; resetAll clears every bucket", () => {
+    const rl = new RateLimiter({ capacity: 1, refillPerSec: 1, now });
+    rl.tryConsume("a");
+    rl.tryConsume("b");
+    rl.reset("a");
+    expect(rl.tryConsume("a").allowed).toBe(true); // refilled
+    expect(rl.tryConsume("b").allowed).toBe(false); // still drained
+
+    rl.resetAll();
+    expect(rl.getStats().activeKeys).toBe(0);
+  });
+
+  test("getStats reports active keys and configuration", () => {
+    const rl = new RateLimiter({ capacity: 7, refillPerSec: 2, now });
+    rl.tryConsume("a");
+    rl.tryConsume("b");
+    expect(rl.getStats()).toEqual({
+      activeKeys: 2,
+      capacity: 7,
+      refillPerSec: 2,
+    });
+  });
+
+  test("rejects invalid configuration", () => {
+    expect(() => new RateLimiter({ capacity: 0 })).toThrow();
+    expect(() => new RateLimiter({ capacity: 5, refillPerSec: 0 })).toThrow();
+  });
+});
+
+// ============ rateLimitMiddleware ============
+describe("rateLimitMiddleware", () => {
+  let clock;
+  const now = () => clock;
+
+  beforeEach(() => {
+    clock = 1_000;
+  });
+
+  test("allows a request, sets rate-limit headers, and calls next", () => {
+    const rl = new RateLimiter({ capacity: 2, refillPerSec: 1, now });
+    const mw = rateLimitMiddleware(rl);
+    const req = { ip: "1.2.3.4" };
+    const res = makeRes();
+    let nextCalled = false;
+
+    mw(req, res, () => {
+      nextCalled = true;
+    });
+
+    expect(nextCalled).toBe(true);
+    expect(res.headers["X-RateLimit-Limit"]).toBe("2");
+    expect(res.headers["X-RateLimit-Remaining"]).toBe("1");
+  });
+
+  test("blocks with 429, Retry-After, and does not call next", () => {
+    const rl = new RateLimiter({ capacity: 1, refillPerSec: 1, now });
+    const mw = rateLimitMiddleware(rl);
+    const req = { ip: "1.2.3.4" };
+    let nextCount = 0;
+    const next = () => {
+      nextCount += 1;
+    };
+
+    mw(req, makeRes(), next); // consumes the only token
+    const res = makeRes();
+    mw(req, res, next); // blocked
+
+    expect(nextCount).toBe(1);
+    expect(res.statusCode).toBe(429);
+    expect(res.headers["Retry-After"]).toBe("1");
+    expect(res.body.error).toMatch(/rate limit/i);
+    expect(res.body.retryAfterMs).toBe(1000);
+  });
+
+  test("derives the bucket key from a custom keyFn", () => {
+    const rl = new RateLimiter({ capacity: 1, refillPerSec: 1, now });
+    const mw = rateLimitMiddleware(rl, {
+      keyFn: (req) => req.body.userId,
+    });
+
+    let blocked = false;
+    mw({ body: { userId: "u1" } }, makeRes(), () => {});
+    mw({ body: { userId: "u1" } }, makeRes(), () => {
+      blocked = true;
+    });
+    expect(blocked).toBe(false); // second u1 call was blocked → next not called
+
+    // A different user has an independent bucket.
+    let allowed = false;
+    mw({ body: { userId: "u2" } }, makeRes(), () => {
+      allowed = true;
+    });
+    expect(allowed).toBe(true);
+  });
+});

orchestrator/core/rate-limiter.js

@@ -0,0 +1,126 @@
+/**
+ * Token-bucket rate limiter.
+ *
+ * Sits alongside the other reliability primitives (circuit breaker, cost
+ * tracker, dead-letter queue) and protects the orchestrator's expensive
+ * LLM-backed routes from runaway usage / abuse. Each key (by default the client
+ * IP) gets its own bucket that holds up to `capacity` tokens and refills at
+ * `refillPerSec`. A request consumes `cost` tokens; if the bucket can't cover
+ * the cost the request is rejected with the time until enough tokens refill.
+ *
+ * The clock is injectable (`now`) so behavior is fully deterministic in tests.
+ */
+class RateLimiter {
+  constructor({ capacity = 60, refillPerSec = 1, now } = {}) {
+    if (!Number.isFinite(capacity) || capacity <= 0) {
+      throw new Error("RateLimiter: capacity must be a positive number");
+    }
+    if (!Number.isFinite(refillPerSec) || refillPerSec <= 0) {
+      throw new Error("RateLimiter: refillPerSec must be a positive number");
+    }
+    this.capacity = capacity;
+    this.refillPerSec = refillPerSec;
+    this.now = now || (() => Date.now());
+    this.buckets = new Map(); // key -> { tokens, lastRefill }
+  }
+
+  // Lazily creates a bucket (full) and brings its token count up to `now`.
+  _refill(key) {
+    const ts = this.now();
+    let bucket = this.buckets.get(key);
+    if (!bucket) {
+      bucket = { tokens: this.capacity, lastRefill: ts };
+      this.buckets.set(key, bucket);
+      return bucket;
+    }
+    const elapsedSec = (ts - bucket.lastRefill) / 1000;
+    if (elapsedSec > 0) {
+      bucket.tokens = Math.min(
+        this.capacity,
+        bucket.tokens + elapsedSec * this.refillPerSec,
+      );
+      bucket.lastRefill = ts;
+    }
+    return bucket;
+  }
+
+  /**
+   * Attempt to consume `cost` tokens for `key`.
+   * @returns {{allowed: boolean, remaining: number, limit: number, retryAfterMs: number}}
+   */
+  tryConsume(key, cost = 1) {
+    const bucket = this._refill(key);
+    if (bucket.tokens >= cost) {
+      bucket.tokens -= cost;
+      return {
+        allowed: true,
+        remaining: bucket.tokens,
+        limit: this.capacity,
+        retryAfterMs: 0,
+      };
+    }
+    const deficit = cost - bucket.tokens;
+    return {
+      allowed: false,
+      remaining: bucket.tokens,
+      limit: this.capacity,
+      retryAfterMs: Math.ceil((deficit / this.refillPerSec) * 1000),
+    };
+  }
+
+  /** Current token count for a key (refilled to now). Unknown keys read full. */
+  getState(key) {
+    if (!this.buckets.has(key)) {
+      return { tokens: this.capacity, capacity: this.capacity };
+    }
+    const bucket = this._refill(key);
+    return { tokens: bucket.tokens, capacity: this.capacity };
+  }
+
+  /** Forget a single key's bucket (it will read as full again). */
+  reset(key) {
+    this.buckets.delete(key);
+  }
+
+  /** Forget every bucket. */
+  resetAll() {
+    this.buckets.clear();
+  }
+
+  getStats() {
+    return {
+      activeKeys: this.buckets.size,
+      capacity: this.capacity,
+      refillPerSec: this.refillPerSec,
+    };
+  }
+}
+
+/**
+ * Express middleware factory backed by a {@link RateLimiter}.
+ *
+ * @param {RateLimiter} limiter
+ * @param {object}   [options]
+ * @param {(req)=>string} [options.keyFn]  derive the bucket key (default: req.ip)
+ * @param {number}        [options.cost=1] tokens this route consumes
+ */
+function rateLimitMiddleware(limiter, { keyFn, cost = 1 } = {}) {
+  return (req, res, next) => {
+    const key = (keyFn ? keyFn(req) : req.ip) || "global";
+    const result = limiter.tryConsume(key, cost);
+
+    res.set("X-RateLimit-Limit", result.limit);
+    res.set("X-RateLimit-Remaining", Math.max(0, Math.floor(result.remaining)));
+
+    if (!result.allowed) {
+      res.set("Retry-After", Math.ceil(result.retryAfterMs / 1000));
+      return res.status(429).json({
+        error: "Rate limit exceeded",
+        retryAfterMs: result.retryAfterMs,
+      });
+    }
+    return next();
+  };
+}
+
+module.exports = { RateLimiter, rateLimitMiddleware };