@splashcodex/api-key-manager 2.0.0 → 3.0.0

package/src/index.ts

@@ -1,9 +1,15 @@
 /**
- * Universal ApiKeyManager v2.0
- * Implements: Rotation, Circuit Breaker, Persistence, Exponential Backoff, Strategies
+ * Universal ApiKeyManager v3.0
+ * Implements: Rotation, Circuit Breaker, Persistence, Exponential Backoff, Strategies,
+ *             Event Emitter, Fallback, execute(), Timeout, Auto-Retry, Provider Tags,
+ *             Health Checks, Bulkhead/Concurrency
  * Gemini-Specific: finishReason handling, Safety blocks, RECITATION detection
  */
 
+import { EventEmitter } from 'events';
+
+// ─── Interfaces & Types ──────────────────────────────────────────────────────
+
 export interface KeyState {
     key: string;
     failCount: number;           // Consecutive failures
@@ -20,6 +26,8 @@ export interface KeyState {
     averageLatency: number;      // Rolling average latency in ms
     totalLatency: number;        // Sum of all latency checks (for calculating average)
     latencySamples: number;      // Number of samples
+    // v3.0 Fields
+    provider: string;            // Provider tag (e.g. 'openai', 'gemini')
 }
 
 export type ErrorType =
@@ -29,6 +37,7 @@ export type ErrorType =
     | 'BAD_REQUEST' // 400 - Do not retry, fix request
     | 'SAFETY'      // finishReason: SAFETY - Not a key issue
     | 'RECITATION'  // finishReason: RECITATION - Not a key issue
+    | 'TIMEOUT'     // Request timed out
     | 'UNKNOWN';    // Catch-all
 
 export interface ErrorClassification {
@@ -39,6 +48,45 @@ export interface ErrorClassification {
     markKeyDead: boolean;
 }
 
+export interface ApiKeyManagerStats {
+    total: number;
+    healthy: number;
+    cooling: number;
+    dead: number;
+}
+
+export interface ExecuteOptions {
+    timeoutMs?: number;      // Timeout per attempt in ms
+    maxRetries?: number;     // Max retry attempts (default: 0 = no retry)
+    finishReason?: string;   // For Gemini finishReason handling
+}
+
+export interface ApiKeyManagerOptions {
+    storage?: any;
+    strategy?: LoadBalancingStrategy;
+    fallbackFn?: () => any;
+    concurrency?: number;    // Max concurrent execute() calls
+}
+
+// ─── Event Types ─────────────────────────────────────────────────────────────
+
+export interface ApiKeyManagerEventMap {
+    keyDead: (key: string) => void;
+    circuitOpen: (key: string) => void;
+    circuitHalfOpen: (key: string) => void;
+    keyRecovered: (key: string) => void;
+    fallback: (reason: string) => void;
+    allKeysExhausted: () => void;
+    retry: (key: string, attempt: number, delayMs: number) => void;
+    healthCheckFailed: (key: string, error: any) => void;
+    healthCheckPassed: (key: string) => void;
+    executeSuccess: (key: string, durationMs: number) => void;
+    executeFailed: (key: string, error: any) => void;
+    bulkheadRejected: () => void;
+}
+
+// ─── Config ──────────────────────────────────────────────────────────────────
+
 const CONFIG = {
     MAX_CONSECUTIVE_FAILURES: 5,
     COOLDOWN_TRANSIENT: 60 * 1000,        // 1 minute
@@ -58,13 +106,31 @@ const ERROR_PATTERNS = {
     isBadRequest: /400|invalid.?argument|failed.?precondition|malformed|not.?found|404/i,
 };
 
-export interface ApiKeyManagerStats {
-    total: number;
-    healthy: number;
-    cooling: number;
-    dead: number;
+// ─── Custom Errors ───────────────────────────────────────────────────────────
+
+export class TimeoutError extends Error {
+    constructor(ms: number) {
+        super(`Request timed out after ${ms}ms`);
+        this.name = 'TimeoutError';
+    }
+}
+
+export class BulkheadRejectionError extends Error {
+    constructor() {
+        super('Bulkhead capacity exceeded — too many concurrent requests');
+        this.name = 'BulkheadRejectionError';
+    }
 }
 
+export class AllKeysExhaustedError extends Error {
+    constructor() {
+        super('All API keys exhausted — no healthy keys available');
+        this.name = 'AllKeysExhaustedError';
+    }
+}
+
+// ─── Strategies ──────────────────────────────────────────────────────────────
+
 /**
  * Strategy Interface for selecting the next key
  */
@@ -77,7 +143,6 @@ export interface LoadBalancingStrategy {
  */
 export class StandardStrategy implements LoadBalancingStrategy {
     next(candidates: KeyState[]): KeyState | null {
-        // Sort: Pristine > Fewest Failures > Least Recently Used
         candidates.sort((a, b) => {
             if (a.failCount !== b.failCount) return a.failCount - b.failCount;
             return a.lastUsed - b.lastUsed;
@@ -112,49 +177,84 @@ export class WeightedStrategy implements LoadBalancingStrategy {
 export class LatencyStrategy implements LoadBalancingStrategy {
     next(candidates: KeyState[]): KeyState | null {
         if (candidates.length === 0) return null;
-        // Sort by averageLatency (lowest first)
-        // If latency is 0 (untried), treat as high priority or neutral?
-        // Let's treat 0 as "unknown, give it a shot" -> insert at top or mixed?
-        // Simple: Sort ASC. 0 comes first.
         candidates.sort((a, b) => a.averageLatency - b.averageLatency);
         return candidates[0];
     }
 }
 
-export class ApiKeyManager {
+// ─── Main Class ──────────────────────────────────────────────────────────────
+
+export class ApiKeyManager extends EventEmitter {
     private keys: KeyState[] = [];
     private storageKey = 'api_rotation_state_v2';
     private storage: any;
     private strategy: LoadBalancingStrategy;
+    private fallbackFn?: () => any;
 
-    constructor(initialKeys: string[] | { key: string; weight?: number }[], storage?: any, strategy?: LoadBalancingStrategy) {
-        // Simple in-memory storage mock if none provided
-        this.storage = storage || {
+    // Bulkhead state
+    private maxConcurrency: number;
+    private activeCalls: number = 0;
+
+    // Health check state
+    private healthCheckFn?: (key: string) => Promise<boolean>;
+    private healthCheckInterval?: ReturnType<typeof setInterval>;
+
+    /**
+     * Constructor supports both legacy positional args and new options object.
+     *
+     * @example Legacy (v1/v2 — still works):
+     *   new ApiKeyManager(['key1', 'key2'], storage, strategy)
+     *
+     * @example New (v3):
+     *   new ApiKeyManager(keys, { storage, strategy, fallbackFn, concurrency })
+     */
+    constructor(
+        initialKeys: string[] | { key: string; weight?: number; provider?: string }[],
+        storageOrOptions?: any | ApiKeyManagerOptions,
+        strategy?: LoadBalancingStrategy
+    ) {
+        super();
+
+        // Detect if second arg is options object or legacy storage
+        let options: ApiKeyManagerOptions = {};
+        if (storageOrOptions && typeof storageOrOptions === 'object' && ('storage' in storageOrOptions || 'strategy' in storageOrOptions || 'fallbackFn' in storageOrOptions || 'concurrency' in storageOrOptions)) {
+            // New v3 options object
+            options = storageOrOptions as ApiKeyManagerOptions;
+        } else {
+            // Legacy positional args
+            options = {
+                storage: storageOrOptions,
+                strategy: strategy,
+            };
+        }
+
+        this.storage = options.storage || {
             getItem: () => null,
             setItem: () => { },
         };
-
-        this.strategy = strategy || new StandardStrategy();
+        this.strategy = options.strategy || new StandardStrategy();
+        this.fallbackFn = options.fallbackFn;
+        this.maxConcurrency = options.concurrency || Infinity;
 
         // Normalize input to objects
-        let inputKeys: { key: string; weight?: number }[] = [];
+        let inputKeys: { key: string; weight?: number; provider?: string }[] = [];
         if (initialKeys.length > 0 && typeof initialKeys[0] === 'string') {
-            inputKeys = (initialKeys as string[]).flatMap(k => k.split(',').map(s => ({ key: s.trim(), weight: 1.0 })));
+            inputKeys = (initialKeys as string[]).flatMap(k => k.split(',').map(s => ({ key: s.trim(), weight: 1.0, provider: 'default' })));
         } else {
-            inputKeys = initialKeys as { key: string; weight?: number }[];
+            inputKeys = initialKeys as { key: string; weight?: number; provider?: string }[];
         }
 
         // Deduplicate
-        const uniqueMap = new Map<string, number>();
+        const uniqueMap = new Map<string, { weight: number; provider: string }>();
         inputKeys.forEach(k => {
-            if (k.key.length > 0) uniqueMap.set(k.key, k.weight ?? 1.0);
+            if (k.key.length > 0) uniqueMap.set(k.key, { weight: k.weight ?? 1.0, provider: k.provider ?? 'default' });
         });
 
         if (uniqueMap.size < inputKeys.length) {
             console.warn(`[ApiKeyManager] Removed ${inputKeys.length - uniqueMap.size} duplicate/empty keys.`);
         }
 
-        this.keys = Array.from(uniqueMap.entries()).map(([key, weight]) => ({
+        this.keys = Array.from(uniqueMap.entries()).map(([key, meta]) => ({
             key,
             failCount: 0,
             failedAt: null,
@@ -165,15 +265,18 @@ export class ApiKeyManager {
             totalRequests: 0,
             halfOpenTestTime: null,
             customCooldown: null,
-            weight: weight,
+            weight: meta.weight,
             averageLatency: 0,
             totalLatency: 0,
-            latencySamples: 0
+            latencySamples: 0,
+            provider: meta.provider,
         }));
 
         this.loadState();
     }
 
+    // ─── Error Classification ────────────────────────────────────────────────
+
     /**
      * CLASSIFIES an error to determine handling strategy
      */
@@ -185,7 +288,12 @@ export class ApiKeyManager {
         if (finishReason === 'SAFETY') return { type: 'SAFETY', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
         if (finishReason === 'RECITATION') return { type: 'RECITATION', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
 
-        // 2. Check HTTP status codes
+        // 2. Check timeout
+        if (error instanceof TimeoutError || error?.name === 'TimeoutError') {
+            return { type: 'TIMEOUT', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
+        }
+
+        // 3. Check HTTP status codes
         if (status === 403 || ERROR_PATTERNS.isAuthError.test(message)) {
             return { type: 'AUTH', retryable: false, cooldownMs: Infinity, markKeyFailed: true, markKeyDead: true };
         }
@@ -225,6 +333,8 @@ export class ApiKeyManager {
         return null;
     }
 
+    // ─── Cooldown ────────────────────────────────────────────────────────────
+
     private isOnCooldown(k: KeyState): boolean {
         if (k.circuitState === 'DEAD') return true;
         const now = Date.now();
@@ -232,6 +342,7 @@ export class ApiKeyManager {
         if (k.circuitState === 'OPEN') {
             if (k.halfOpenTestTime && now >= k.halfOpenTestTime) {
                 k.circuitState = 'HALF_OPEN';
+                this.emit('circuitHalfOpen', k.key);
                 return false;
             }
             return true;
@@ -246,6 +357,8 @@ export class ApiKeyManager {
         return false;
     }
 
+    // ─── Key Selection ───────────────────────────────────────────────────────
+
     public getKey(): string | null {
         // 1. Filter out dead and cooling down keys
         const candidates = this.keys.filter(k => k.circuitState !== 'DEAD' && !this.isOnCooldown(k));
@@ -253,7 +366,10 @@ export class ApiKeyManager {
         if (candidates.length === 0) {
             // FALLBACK: Return oldest failed key (excluding DEAD)
             const nonDead = this.keys.filter(k => k.circuitState !== 'DEAD');
-            if (nonDead.length === 0) return null;
+            if (nonDead.length === 0) {
+                this.emit('allKeysExhausted');
+                return null;
+            }
             return nonDead.sort((a, b) => (a.failedAt || 0) - (b.failedAt || 0))[0]?.key || null;
         }
 
@@ -268,10 +384,31 @@ export class ApiKeyManager {
         return null;
     }
 
+    /**
+     * Get a key filtered by provider tag
+     */
+    public getKeyByProvider(provider: string): string | null {
+        const candidates = this.keys.filter(k =>
+            k.provider === provider && k.circuitState !== 'DEAD' && !this.isOnCooldown(k)
+        );
+
+        if (candidates.length === 0) return null;
+
+        const selected = this.strategy.next(candidates);
+        if (selected) {
+            selected.lastUsed = Date.now();
+            this.saveState();
+            return selected.key;
+        }
+        return null;
+    }
+
     public getKeyCount(): number {
         return this.keys.filter(k => k.circuitState !== 'DEAD').length;
     }
 
+    // ─── Mark Success / Failed ───────────────────────────────────────────────
+
     /**
      * Mark success AND update latency stats
      * @param durationMs Duration of the request in milliseconds
@@ -280,7 +417,11 @@ export class ApiKeyManager {
         const k = this.keys.find(x => x.key === key);
         if (!k) return;
 
-        if (k.circuitState !== 'CLOSED' && k.circuitState !== 'DEAD') console.log(`[Key Recovered] ...${key.slice(-4)}`);
+        const wasRecovering = k.circuitState !== 'CLOSED' && k.circuitState !== 'DEAD';
+        if (wasRecovering) {
+            console.log(`[Key Recovered] ...${key.slice(-4)}`);
+            this.emit('keyRecovered', key);
+        }
 
         k.circuitState = 'CLOSED';
         k.failCount = 0;
@@ -313,14 +454,17 @@ export class ApiKeyManager {
         if (classification.markKeyDead) {
             k.circuitState = 'DEAD';
             console.error(`[Key DEAD] ...${key.slice(-4)} - Permanently removed`);
+            this.emit('keyDead', key);
         } else {
             // State Transitions
             if (k.circuitState === 'HALF_OPEN') {
                 k.circuitState = 'OPEN';
                 k.halfOpenTestTime = Date.now() + CONFIG.HALF_OPEN_TEST_DELAY;
+                this.emit('circuitOpen', key);
             } else if (k.failCount >= CONFIG.MAX_CONSECUTIVE_FAILURES || classification.type === 'QUOTA') {
                 k.circuitState = 'OPEN';
                 k.halfOpenTestTime = Date.now() + (classification.cooldownMs || CONFIG.HALF_OPEN_TEST_DELAY);
+                this.emit('circuitOpen', key);
             }
         }
         this.saveState();
@@ -336,6 +480,8 @@ export class ApiKeyManager {
         });
     }
 
+    // ─── Backoff ─────────────────────────────────────────────────────────────
+
     public calculateBackoff(attempt: number): number {
         const exponential = CONFIG.BASE_BACKOFF * Math.pow(2, attempt);
         const capped = Math.min(exponential, CONFIG.MAX_BACKOFF);
@@ -343,6 +489,8 @@ export class ApiKeyManager {
         return capped + jitter;
     }
 
+    // ─── Stats ───────────────────────────────────────────────────────────────
+
     public getStats(): ApiKeyManagerStats {
         const total = this.keys.length;
         const dead = this.keys.filter(k => k.circuitState === 'DEAD').length;
@@ -353,6 +501,181 @@ export class ApiKeyManager {
 
     public _getKeys(): KeyState[] { return this.keys; }
 
+    // ─── execute() Wrapper ───────────────────────────────────────────────────
+
+    /**
+     * Wraps the entire API call lifecycle into a single method.
+     *
+     * @example
+     *   const result = await manager.execute(
+     *     (key) => fetch(`https://api.example.com?key=${key}`),
+     *     { maxRetries: 3, timeoutMs: 5000 }
+     *   );
+     */
+    public async execute<T>(
+        fn: (key: string, signal?: AbortSignal) => Promise<T>,
+        options?: ExecuteOptions
+    ): Promise<T> {
+        const maxRetries = options?.maxRetries ?? 0;
+        const timeoutMs = options?.timeoutMs;
+        const finishReason = options?.finishReason;
+
+        // Bulkhead check
+        if (this.activeCalls >= this.maxConcurrency) {
+            this.emit('bulkheadRejected');
+            throw new BulkheadRejectionError();
+        }
+
+        this.activeCalls++;
+        try {
+            return await this._executeWithRetry(fn, maxRetries, timeoutMs, finishReason);
+        } finally {
+            this.activeCalls--;
+        }
+    }
+
+    private async _executeWithRetry<T>(
+        fn: (key: string, signal?: AbortSignal) => Promise<T>,
+        maxRetries: number,
+        timeoutMs?: number,
+        finishReason?: string
+    ): Promise<T> {
+        let lastError: any;
+
+        for (let attempt = 0; attempt <= maxRetries; attempt++) {
+            const key = this.getKey();
+
+            if (!key) {
+                // All keys exhausted — try fallback
+                if (this.fallbackFn) {
+                    this.emit('fallback', 'all keys exhausted');
+                    return this.fallbackFn();
+                }
+                throw new AllKeysExhaustedError();
+            }
+
+            try {
+                const start = Date.now();
+                let result: T;
+
+                if (timeoutMs) {
+                    result = await this._executeWithTimeout(fn, key, timeoutMs);
+                } else {
+                    result = await fn(key);
+                }
+
+                const duration = Date.now() - start;
+                this.markSuccess(key, duration);
+                this.emit('executeSuccess', key, duration);
+                return result;
+
+            } catch (error: any) {
+                lastError = error;
+                const classification = this.classifyError(error, finishReason);
+
+                this.markFailed(key, classification);
+                this.emit('executeFailed', key, error);
+
+                if (!classification.retryable || attempt >= maxRetries) {
+                    // Non-retryable or out of retries
+                    if (this.fallbackFn && attempt >= maxRetries) {
+                        this.emit('fallback', 'max retries exceeded');
+                        return this.fallbackFn();
+                    }
+                    throw error;
+                }
+
+                // Retry with backoff
+                const delay = this.calculateBackoff(attempt);
+                this.emit('retry', key, attempt + 1, delay);
+                await this._sleep(delay);
+            }
+        }
+
+        // Should not reach here, but safety net
+        throw lastError;
+    }
+
+    private async _executeWithTimeout<T>(
+        fn: (key: string, signal?: AbortSignal) => Promise<T>,
+        key: string,
+        timeoutMs: number
+    ): Promise<T> {
+        const controller = new AbortController();
+        const timer = setTimeout(() => controller.abort(), timeoutMs);
+
+        try {
+            const result = await Promise.race([
+                fn(key, controller.signal),
+                new Promise<never>((_, reject) => {
+                    controller.signal.addEventListener('abort', () => {
+                        reject(new TimeoutError(timeoutMs));
+                    });
+                })
+            ]);
+            return result;
+        } finally {
+            clearTimeout(timer);
+        }
+    }
+
+    private _sleep(ms: number): Promise<void> {
+        return new Promise(resolve => setTimeout(resolve, ms));
+    }
+
+    // ─── Health Checks ───────────────────────────────────────────────────────
+
+    /**
+     * Set a health check function that tests if a key is operational
+     */
+    public setHealthCheck(fn: (key: string) => Promise<boolean>) {
+        this.healthCheckFn = fn;
+    }
+
+    /**
+     * Start periodic health checks
+     * @param intervalMs How often to run health checks (default: 60s)
+     */
+    public startHealthChecks(intervalMs: number = 60_000) {
+        this.stopHealthChecks(); // Clear any existing interval
+        this.healthCheckInterval = setInterval(() => this._runHealthChecks(), intervalMs);
+    }
+
+    /**
+     * Stop periodic health checks
+     */
+    public stopHealthChecks() {
+        if (this.healthCheckInterval) {
+            clearInterval(this.healthCheckInterval);
+            this.healthCheckInterval = undefined;
+        }
+    }
+
+    private async _runHealthChecks() {
+        if (!this.healthCheckFn) return;
+
+        // Check non-DEAD keys that are in OPEN or HALF_OPEN state
+        const keysToCheck = this.keys.filter(k =>
+            k.circuitState === 'OPEN' || k.circuitState === 'HALF_OPEN'
+        );
+
+        for (const k of keysToCheck) {
+            try {
+                const healthy = await this.healthCheckFn(k.key);
+                if (healthy) {
+                    this.markSuccess(k.key);
+                    this.emit('healthCheckPassed', k.key);
+                } else {
+                    this.emit('healthCheckFailed', k.key, new Error('Health check returned false'));
+                }
+            } catch (error) {
+                this.emit('healthCheckFailed', k.key, error);
+            }
+        }
+    }
+
+    // ─── Persistence ─────────────────────────────────────────────────────────
+
     private saveState() {
         if (!this.storage) return;
         const state = this.keys.reduce((acc, k) => ({
@@ -369,7 +692,8 @@ export class ApiKeyManager {
                 weight: k.weight,
                 averageLatency: k.averageLatency,
                 totalLatency: k.totalLatency,
-                latencySamples: k.latencySamples
+                latencySamples: k.latencySamples,
+                provider: k.provider
             }
         }), {});
         this.storage.setItem(this.storageKey, JSON.stringify(state));