moving

packages/agent/tests/credential-pool.test.ts

@@ -0,0 +1,369 @@
+import { describe, it, expect } from 'vitest';
+import {
+  CredentialPool,
+  loadCredentialPool,
+  extractStatusCode,
+  type VaultLike,
+} from '../src/credential-pool.js';
+
+// ── Helpers ──────────────────────────────────────────────────────────────
+
+function createPool(keyCount = 3, nowFn?: () => number): CredentialPool {
+  const pool = new CredentialPool({ provider: 'anthropic' }, nowFn);
+  for (let i = 0; i < keyCount; i++) {
+    pool.addCredential(`key-${i}`, `sk-${i}`);
+  }
+  return pool;
+}
+
+function mockVault(keys: Record<string, string>): VaultLike {
+  return {
+    get(name: string) {
+      return name in keys ? { value: keys[name] } : null;
+    },
+    has(name: string) {
+      return name in keys;
+    },
+  };
+}
+
+// ── Basic Operations ─────────────────────────────────────────────────────
+
+describe('CredentialPool', () => {
+  it('starts empty', () => {
+    const pool = new CredentialPool({ provider: 'test' });
+    expect(pool.size).toBe(0);
+    expect(pool.getKey()).toBeNull();
+  });
+
+  it('returns added keys', () => {
+    const pool = createPool(2);
+    const key = pool.getKey();
+    expect(key).toBe('sk-0');
+  });
+
+  it('prevents duplicate additions', () => {
+    const pool = new CredentialPool({ provider: 'test' });
+    pool.addCredential('k1', 'val1');
+    pool.addCredential('k1', 'val2');
+    expect(pool.size).toBe(1);
+  });
+
+  it('getNameForKey returns the credential name', () => {
+    const pool = createPool(2);
+    expect(pool.getNameForKey('sk-0')).toBe('key-0');
+    expect(pool.getNameForKey('sk-1')).toBe('key-1');
+    expect(pool.getNameForKey('unknown')).toBeNull();
+  });
+});
+
+// ── Round-Robin ──────────────────────────────────────────────────────────
+
+describe('round-robin', () => {
+  it('rotates through keys in order', () => {
+    const pool = createPool(3);
+    expect(pool.getKey()).toBe('sk-0');
+    expect(pool.getKey()).toBe('sk-1');
+    expect(pool.getKey()).toBe('sk-2');
+    expect(pool.getKey()).toBe('sk-0'); // wraps around
+  });
+
+  it('skips keys in cooldown', () => {
+    const pool = createPool(3);
+    pool.reportError('sk-1', 429);
+    expect(pool.getKey()).toBe('sk-0');
+    expect(pool.getKey()).toBe('sk-2'); // sk-1 skipped
+    expect(pool.getKey()).toBe('sk-0');
+  });
+
+  it('skips disabled keys', () => {
+    const pool = createPool(3);
+    pool.reportError('sk-1', 401);
+    expect(pool.getKey()).toBe('sk-0');
+    expect(pool.getKey()).toBe('sk-2');
+    expect(pool.getKey()).toBe('sk-0');
+  });
+});
+
+// ── Cooldown: 429 (Rate Limit) ───────────────────────────────────────────
+
+describe('429 cooldown', () => {
+  it('puts key in 1-hour cooldown on 429', () => {
+    const now = 1000000;
+    const pool = createPool(2, () => now);
+
+    pool.reportError('sk-0', 429, 'Rate limit exceeded');
+    expect(pool.getKey()).toBe('sk-1'); // sk-0 is in cooldown
+
+    const status = pool.getStatus();
+    const k0 = status.entries.find(e => e.name === 'key-0');
+    expect(k0?.status).toBe('cooldown');
+    expect(k0?.lastError).toBe('Rate limit exceeded');
+  });
+
+  it('auto-recovers after cooldown expires', () => {
+    let now = 1000000;
+    const pool = createPool(2, () => now);
+
+    pool.reportError('sk-0', 429);
+    expect(pool.getKey()).toBe('sk-1');
+
+    // Fast-forward past the 1-hour cooldown
+    now += 60 * 60 * 1000 + 1;
+
+    // sk-0 should be available again
+    const key = pool.getKey();
+    // After recovery, round-robin continues — could be sk-0 or sk-1
+    expect(['sk-0', 'sk-1']).toContain(key);
+
+    const status = pool.getStatus();
+    const k0 = status.entries.find(e => e.name === 'key-0');
+    expect(k0?.status).toBe('active');
+  });
+
+  it('does NOT recover before cooldown expires', () => {
+    let now = 1000000;
+    const pool = createPool(1, () => now);
+
+    pool.reportError('sk-0', 429);
+    expect(pool.getKey()).toBeNull(); // only key is in cooldown
+
+    // Advance 30 minutes — not enough
+    now += 30 * 60 * 1000;
+    expect(pool.getKey()).toBeNull();
+  });
+});
+
+// ── Cooldown: 402 (Payment Required) ─────────────────────────────────────
+
+describe('402 cooldown', () => {
+  it('puts key in 24-hour cooldown on 402', () => {
+    const now = 1000000;
+    const pool = createPool(2, () => now);
+
+    pool.reportError('sk-0', 402, 'Insufficient funds');
+
+    const status = pool.getStatus();
+    const k0 = status.entries.find(e => e.name === 'key-0');
+    expect(k0?.status).toBe('cooldown');
+    expect(k0?.cooldownUntil).toBe(now + 24 * 60 * 60 * 1000);
+  });
+
+  it('recovers after 24 hours', () => {
+    let now = 1000000;
+    const pool = createPool(1, () => now);
+
+    pool.reportError('sk-0', 402);
+    expect(pool.getKey()).toBeNull();
+
+    // Fast-forward 24 hours
+    now += 24 * 60 * 60 * 1000 + 1;
+    expect(pool.getKey()).toBe('sk-0');
+  });
+});
+
+// ── Permanent Disable: 401 ───────────────────────────────────────────────
+
+describe('401 permanent disable', () => {
+  it('permanently disables key on 401', () => {
+    const pool = createPool(2);
+    pool.reportError('sk-0', 401, 'Invalid API key');
+
+    const status = pool.getStatus();
+    const k0 = status.entries.find(e => e.name === 'key-0');
+    expect(k0?.status).toBe('disabled');
+    expect(k0?.cooldownUntil).toBeNull();
+  });
+
+  it('never recovers a disabled key', () => {
+    let now = 1000000;
+    const pool = createPool(1, () => now);
+
+    pool.reportError('sk-0', 401);
+    expect(pool.getKey()).toBeNull();
+
+    // Even after 100 hours
+    now += 100 * 60 * 60 * 1000;
+    expect(pool.getKey()).toBeNull();
+  });
+});
+
+// ── Other Errors ─────────────────────────────────────────────────────────
+
+describe('other errors', () => {
+  it('puts key in 5-minute cooldown for 500/503', () => {
+    const now = 1000000;
+    const pool = createPool(2, () => now);
+
+    pool.reportError('sk-0', 500);
+
+    const status = pool.getStatus();
+    const k0 = status.entries.find(e => e.name === 'key-0');
+    expect(k0?.status).toBe('cooldown');
+    expect(k0?.cooldownUntil).toBe(now + 5 * 60 * 1000);
+  });
+});
+
+// ── Success Tracking ─────────────────────────────────────────────────────
+
+describe('success tracking', () => {
+  it('increments success count', () => {
+    const pool = createPool(1);
+    pool.reportSuccess('sk-0');
+    pool.reportSuccess('sk-0');
+
+    const status = pool.getStatus();
+    expect(status.entries[0].successCount).toBe(2);
+  });
+
+  it('tracks errors separately', () => {
+    const pool = createPool(1);
+    pool.reportSuccess('sk-0');
+    pool.reportError('sk-0', 429);
+
+    const status = pool.getStatus();
+    expect(status.entries[0].successCount).toBe(1);
+    expect(status.entries[0].errorCount).toBe(1);
+  });
+});
+
+// ── hasAvailableKeys ─────────────────────────────────────────────────────
+
+describe('hasAvailableKeys', () => {
+  it('returns true when active keys exist', () => {
+    const pool = createPool(2);
+    expect(pool.hasAvailableKeys()).toBe(true);
+  });
+
+  it('returns false when all keys are disabled', () => {
+    const pool = createPool(2);
+    pool.reportError('sk-0', 401);
+    pool.reportError('sk-1', 401);
+    expect(pool.hasAvailableKeys()).toBe(false);
+  });
+
+  it('returns true when a cooldown is about to expire', () => {
+    let now = 1000000;
+    const pool = createPool(1, () => now);
+    pool.reportError('sk-0', 429);
+
+    // Still in cooldown
+    expect(pool.hasAvailableKeys()).toBe(false);
+
+    // Past cooldown
+    now += 60 * 60 * 1000 + 1;
+    expect(pool.hasAvailableKeys()).toBe(true);
+  });
+
+  it('reportError returns whether other keys are available', () => {
+    const pool = createPool(3);
+    expect(pool.reportError('sk-0', 429)).toBe(true);  // sk-1 and sk-2 still active
+    expect(pool.reportError('sk-1', 429)).toBe(true);  // sk-2 still active
+    expect(pool.reportError('sk-2', 429)).toBe(false); // all in cooldown
+  });
+});
+
+// ── Pool Status ──────────────────────────────────────────────────────────
+
+describe('getStatus', () => {
+  it('returns correct counts', () => {
+    const pool = createPool(4);
+    pool.reportError('sk-0', 401);  // disabled
+    pool.reportError('sk-1', 429);  // cooldown
+
+    const status = pool.getStatus();
+    expect(status.provider).toBe('anthropic');
+    expect(status.totalKeys).toBe(4);
+    expect(status.activeKeys).toBe(2);
+    expect(status.cooldownKeys).toBe(1);
+    expect(status.disabledKeys).toBe(1);
+  });
+
+  it('recovers expired cooldowns in status', () => {
+    let now = 1000000;
+    const pool = createPool(1, () => now);
+    pool.reportError('sk-0', 429);
+
+    now += 60 * 60 * 1000 + 1;
+    const status = pool.getStatus();
+    expect(status.activeKeys).toBe(1);
+    expect(status.cooldownKeys).toBe(0);
+  });
+});
+
+// ── Vault Loader ─────────────────────────────────────────────────────────
+
+describe('loadCredentialPool', () => {
+  it('loads single key', () => {
+    const vault = mockVault({ 'anthropic': 'sk-ant-primary' });
+    const pool = loadCredentialPool(vault, 'anthropic');
+    expect(pool.size).toBe(1);
+    expect(pool.getKey()).toBe('sk-ant-primary');
+  });
+
+  it('loads multiple keys following convention', () => {
+    const vault = mockVault({
+      'anthropic': 'sk-ant-1',
+      'anthropic-2': 'sk-ant-2',
+      'anthropic-3': 'sk-ant-3',
+    });
+    const pool = loadCredentialPool(vault, 'anthropic');
+    expect(pool.size).toBe(3);
+    expect(pool.getKey()).toBe('sk-ant-1');
+    expect(pool.getKey()).toBe('sk-ant-2');
+    expect(pool.getKey()).toBe('sk-ant-3');
+  });
+
+  it('stops at first gap in numbering', () => {
+    const vault = mockVault({
+      'openai': 'sk-1',
+      'openai-2': 'sk-2',
+      // openai-3 missing
+      'openai-4': 'sk-4',
+    });
+    const pool = loadCredentialPool(vault, 'openai');
+    expect(pool.size).toBe(2); // only primary + -2
+  });
+
+  it('returns empty pool when no keys exist', () => {
+    const vault = mockVault({});
+    const pool = loadCredentialPool(vault, 'anthropic');
+    expect(pool.size).toBe(0);
+    expect(pool.getKey()).toBeNull();
+  });
+
+  it('respects maxKeys limit', () => {
+    const keys: Record<string, string> = { 'test': 'k0' };
+    for (let i = 2; i <= 20; i++) keys[`test-${i}`] = `k${i}`;
+
+    const vault = mockVault(keys);
+    const pool = loadCredentialPool(vault, 'test', 5);
+    expect(pool.size).toBe(5); // primary + 2,3,4,5
+  });
+});
+
+// ── extractStatusCode ────────────────────────────────────────────────────
+
+describe('extractStatusCode', () => {
+  it('extracts from .status property', () => {
+    expect(extractStatusCode({ status: 429 })).toBe(429);
+  });
+
+  it('extracts from .statusCode property', () => {
+    expect(extractStatusCode({ statusCode: 402 })).toBe(402);
+  });
+
+  it('extracts from error message', () => {
+    expect(extractStatusCode(new Error('Server returned 401 Unauthorized'))).toBe(401);
+  });
+
+  it('extracts rate limit 429 from message', () => {
+    expect(extractStatusCode(new Error('HTTP 429 Too Many Requests'))).toBe(429);
+  });
+
+  it('returns null for unknown errors', () => {
+    expect(extractStatusCode(new Error('Network timeout'))).toBeNull();
+    expect(extractStatusCode('string error')).toBeNull();
+    expect(extractStatusCode(null)).toBeNull();
+  });
+});