feat: add ApiKeyPool with weighted round-robin and quota cooldown

Author: jddunn

src/core/providers/ApiKeyPool.ts

@@ -0,0 +1,117 @@
+/**
+ * @module core/providers/ApiKeyPool
+ *
+ * Weighted round-robin API key pool with quota-exhaustion cooldown.
+ *
+ * Accepts a single key or comma-separated keys. The first key gets
+ * higher rotation weight (configurable). Exhausted keys are temporarily
+ * removed from rotation and re-enter after a cooldown period.
+ *
+ * This is a core AgentOS primitive -- every provider that accepts an
+ * API key can use it for automatic multi-key rotation and failover.
+ *
+ * @example
+ * ```ts
+ * // Single key (backward compatible, zero overhead):
+ * const pool = new ApiKeyPool('sk_abc');
+ * pool.next(); // 'sk_abc'
+ *
+ * // Multiple keys with round-robin + first-key priority:
+ * const pool = new ApiKeyPool('sk_primary,sk_backup,sk_overflow');
+ * pool.next(); // weighted rotation, primary selected ~2x more
+ *
+ * // Mark exhausted on quota error:
+ * pool.markExhausted(key); // skipped for 15min cooldown
+ * pool.next(); // returns next available key
+ * ```
+ */
+
+/** Configuration for an ApiKeyPool instance. */
+export interface ApiKeyPoolConfig {
+  /** Cooldown before retrying an exhausted key. Default: 15 minutes. */
+  cooldownMs?: number;
+  /** Weight multiplier for the first key. Default: 2. Set to 1 for equal weighting. */
+  primaryWeight?: number;
+}
+
+interface KeyState {
+  key: string;
+  exhaustedUntil: number;
+}
+
+const DEFAULT_COOLDOWN_MS = 15 * 60_000;
+const DEFAULT_PRIMARY_WEIGHT = 2;
+
+export class ApiKeyPool {
+  private readonly keys: KeyState[];
+  private readonly weightedSlots: number[];
+  private slotIndex = 0;
+  private readonly cooldownMs: number;
+
+  constructor(keys: string | string[], config?: ApiKeyPoolConfig) {
+    this.cooldownMs = config?.cooldownMs ?? DEFAULT_COOLDOWN_MS;
+    const primaryWeight = config?.primaryWeight ?? DEFAULT_PRIMARY_WEIGHT;
+
+    const keyList = Array.isArray(keys)
+      ? keys
+      : keys.split(',').map((k) => k.trim()).filter((k) => k.length > 0);
+
+    this.keys = keyList.map((key) => ({ key, exhaustedUntil: 0 }));
+
+    // Build weighted slot array: first key appears `primaryWeight` times, rest once each.
+    this.weightedSlots = [];
+    for (let i = 0; i < this.keys.length; i++) {
+      const times = i === 0 ? primaryWeight : 1;
+      for (let t = 0; t < times; t++) {
+        this.weightedSlots.push(i);
+      }
+    }
+  }
+
+  /** Number of keys in the pool (including temporarily exhausted ones). */
+  get size(): number {
+    return this.keys.length;
+  }
+
+  /** Whether any key exists at all. */
+  get hasKeys(): boolean {
+    return this.keys.length > 0;
+  }
+
+  /**
+   * Get the next available key via weighted round-robin.
+   * Skips keys currently in cooldown. Returns empty string if pool is empty.
+   */
+  next(): string {
+    if (this.keys.length === 0) return '';
+    if (this.keys.length === 1) return this.keys[0].key;
+
+    const now = Date.now();
+    const totalSlots = this.weightedSlots.length;
+
+    for (let i = 0; i < totalSlots; i++) {
+      const slotIdx = (this.slotIndex + i) % totalSlots;
+      const keyIdx = this.weightedSlots[slotIdx];
+      const state = this.keys[keyIdx];
+      if (state.exhaustedUntil <= now) {
+        this.slotIndex = (slotIdx + 1) % totalSlots;
+        return state.key;
+      }
+    }
+
+    // All keys exhausted -- return the one whose cooldown expires soonest.
+    const sorted = [...this.keys].sort((a, b) => a.exhaustedUntil - b.exhaustedUntil);
+    return sorted[0].key;
+  }
+
+  /**
+   * Mark a key as quota-exhausted. It will be skipped during
+   * rotation until the cooldown expires.
+   */
+  markExhausted(key: string): void {
+    const state = this.keys.find((k) => k.key === key);
+    if (state) {
+      state.exhaustedUntil = Date.now() + this.cooldownMs;
+    }
+  }
+}

src/core/providers/__tests__/ApiKeyPool.spec.ts

@@ -0,0 +1,107 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { ApiKeyPool } from '../ApiKeyPool.js';
+
+describe('ApiKeyPool', () => {
+  describe('construction', () => {
+    it('parses comma-separated string into keys', () => {
+      const pool = new ApiKeyPool('sk_a,sk_b,sk_c');
+      expect(pool.size).toBe(3);
+      expect(pool.hasKeys).toBe(true);
+    });
+
+    it('accepts a single key string', () => {
+      const pool = new ApiKeyPool('sk_a');
+      expect(pool.size).toBe(1);
+    });
+
+    it('accepts an array of keys', () => {
+      const pool = new ApiKeyPool(['sk_a', 'sk_b']);
+      expect(pool.size).toBe(2);
+    });
+
+    it('trims whitespace from keys', () => {
+      const pool = new ApiKeyPool(' sk_a , sk_b ');
+      expect(pool.next()).toBe('sk_a');
+    });
+
+    it('filters empty segments', () => {
+      const pool = new ApiKeyPool('sk_a,,sk_b,');
+      expect(pool.size).toBe(2);
+    });
+
+    it('handles empty string', () => {
+      const pool = new ApiKeyPool('');
+      expect(pool.size).toBe(0);
+      expect(pool.hasKeys).toBe(false);
+    });
+  });
+
+  describe('round-robin', () => {
+    it('single key always returns the same key', () => {
+      const pool = new ApiKeyPool('sk_only');
+      expect(pool.next()).toBe('sk_only');
+      expect(pool.next()).toBe('sk_only');
+      expect(pool.next()).toBe('sk_only');
+    });
+
+    it('rotates through keys in order', () => {
+      const pool = new ApiKeyPool('sk_a,sk_b,sk_c', { primaryWeight: 1 });
+      expect(pool.next()).toBe('sk_a');
+      expect(pool.next()).toBe('sk_b');
+      expect(pool.next()).toBe('sk_c');
+      expect(pool.next()).toBe('sk_a');
+    });
+
+    it('gives first key higher weight with default primaryWeight=2', () => {
+      const pool = new ApiKeyPool('sk_a,sk_b');
+      const counts: Record<string, number> = { sk_a: 0, sk_b: 0 };
+      for (let i = 0; i < 90; i++) counts[pool.next()]++;
+      expect(counts.sk_a).toBeGreaterThan(counts.sk_b);
+    });
+  });
+
+  describe('exhaustion', () => {
+    beforeEach(() => { vi.useFakeTimers(); });
+    afterEach(() => { vi.useRealTimers(); });
+
+    it('skips exhausted key', () => {
+      const pool = new ApiKeyPool('sk_a,sk_b', { primaryWeight: 1 });
+      pool.next(); // sk_a
+      pool.markExhausted('sk_a');
+      expect(pool.next()).toBe('sk_b');
+      expect(pool.next()).toBe('sk_b');
+    });
+
+    it('re-includes key after cooldown expires', () => {
+      const pool = new ApiKeyPool('sk_a,sk_b', { primaryWeight: 1, cooldownMs: 1000 });
+      pool.next(); // sk_a
+      pool.markExhausted('sk_a');
+      expect(pool.next()).toBe('sk_b');
+
+      vi.advanceTimersByTime(1001);
+      const next3 = [pool.next(), pool.next()];
+      expect(next3).toContain('sk_a');
+    });
+
+    it('returns least-exhausted key when all are exhausted', () => {
+      const pool = new ApiKeyPool('sk_a,sk_b', { primaryWeight: 1, cooldownMs: 10_000 });
+      pool.markExhausted('sk_a');
+      vi.advanceTimersByTime(5000);
+      pool.markExhausted('sk_b');
+      expect(pool.next()).toBe('sk_a');
+    });
+  });
+
+  describe('edge cases', () => {
+    it('next() returns empty string for empty pool', () => {
+      const pool = new ApiKeyPool('');
+      expect(pool.next()).toBe('');
+    });
+
+    it('markExhausted on unknown key is a no-op', () => {
+      const pool = new ApiKeyPool('sk_a');
+      pool.markExhausted('sk_unknown');
+      expect(pool.next()).toBe('sk_a');
+    });
+  });
+});