BID · Console
Baseline · Intelligence · Decision
src/tools/retrieval/rate-limiter.ts 1,675 bytes · typescript
/**
 * Generic token-bucket rate limiter. Used by the dispatcher to enforce
 * each connector's declared rate envelope (Std 4 / Std 5 — approved
 * tools must respect declared constraints).
 *
 * Not connector-specific; any consumer can grab a `RateLimiter` and
 * `acquire()` before each outbound call.
 */

export interface RateLimitConfig {
  readonly requestsPerSecond: number;
  readonly burstSize?: number;
}

export class RateLimiter {
  private tokens: number;
  private lastRefill: number;
  private readonly capacity: number;
  private readonly rate: number;

  constructor(cfg: RateLimitConfig) {
    this.rate = Math.max(0.01, cfg.requestsPerSecond);
    this.capacity = Math.max(1, cfg.burstSize ?? Math.ceil(cfg.requestsPerSecond));
    this.tokens = this.capacity;
    this.lastRefill = Date.now();
  }

  /** Wait until a token is available, then consume it. */
  async acquire(): Promise<void> {
    while (true) {
      this.refill();
      if (this.tokens >= 1) {
        this.tokens -= 1;
        return;
      }
      const needed = 1 - this.tokens;
      const waitMs = Math.ceil((needed / this.rate) * 1000);
      await new Promise(r => setTimeout(r, waitMs));
    }
  }

  /** Try to consume a token without blocking; returns true on success. */
  tryAcquire(): boolean {
    this.refill();
    if (this.tokens >= 1) {
      this.tokens -= 1;
      return true;
    }
    return false;
  }

  private refill(): void {
    const now = Date.now();
    const elapsedSec = (now - this.lastRefill) / 1000;
    if (elapsedSec <= 0) return;
    this.tokens = Math.min(this.capacity, this.tokens + elapsedSec * this.rate);
    this.lastRefill = now;
  }
}