@splashcodex/api-key-manager 1.0.0 → 3.0.0

package/src/index.ts

@@ -1,9 +1,15 @@
 /**
- * Universal ApiKeyManager v2.0
- * Implements: Rotation, Circuit Breaker, Persistence, Exponential Backoff
+ * 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
@@ -15,6 +21,13 @@ export interface KeyState {
     totalRequests: number;
     halfOpenTestTime: number | null;
     customCooldown: number | null; // From Retry-After header
+    // v2.0 Stats
+    weight: number;              // 0.0 - 1.0 (Default 1.0)
+    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 =
@@ -24,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 {
@@ -34,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
@@ -53,43 +106,156 @@ 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
+ */
+export interface LoadBalancingStrategy {
+    next(candidates: KeyState[]): KeyState | null;
+}
+
+/**
+ * Standard Strategy: Least Failed > Least Recently Used
+ */
+export class StandardStrategy implements LoadBalancingStrategy {
+    next(candidates: KeyState[]): KeyState | null {
+        candidates.sort((a, b) => {
+            if (a.failCount !== b.failCount) return a.failCount - b.failCount;
+            return a.lastUsed - b.lastUsed;
+        });
+        return candidates[0] || null;
+    }
 }
 
-export class ApiKeyManager {
+/**
+ * Weighted Strategy: Probabilistic selection based on weight
+ * Higher weight = Higher chance of selection
+ */
+export class WeightedStrategy implements LoadBalancingStrategy {
+    next(candidates: KeyState[]): KeyState | null {
+        if (candidates.length === 0) return null;
+
+        const totalWeight = candidates.reduce((sum, k) => sum + k.weight, 0);
+        let random = Math.random() * totalWeight;
+
+        for (const key of candidates) {
+            random -= key.weight;
+            if (random <= 0) return key;
+        }
+
+        return candidates[0]; // Fallback
+    }
+}
+
+/**
+ * Latency Strategy: Pick lowest average latency
+ */
+export class LatencyStrategy implements LoadBalancingStrategy {
+    next(candidates: KeyState[]): KeyState | null {
+        if (candidates.length === 0) return null;
+        candidates.sort((a, b) => a.averageLatency - b.averageLatency);
+        return candidates[0];
+    }
+}
+
+// ─── Main Class ──────────────────────────────────────────────────────────────
+
+export class ApiKeyManager extends EventEmitter {
     private keys: KeyState[] = [];
     private storageKey = 'api_rotation_state_v2';
-    // Simplified Storage interface for Node.js environment
     private storage: any;
+    private strategy: LoadBalancingStrategy;
+    private fallbackFn?: () => any;
+
+    // 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,
+            };
+        }
 
-    constructor(initialKeys: string[], storage?: any) {
-        // Simple in-memory storage mock if none provided (for testing/Node)
-        this.storage = storage || {
+        this.storage = options.storage || {
             getItem: () => null,
             setItem: () => { },
         };
+        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; 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, provider: 'default' })));
+        } else {
+            inputKeys = initialKeys as { key: string; weight?: number; provider?: string }[];
+        }
 
-        // 1. Sanitize & Deduplicate Keys "Even Better"
-        const uniqueKeys = new Set<string>();
-        initialKeys.forEach(rawKey => {
-            // Handle comma-separated keys if they appear in a single string
-            const parts = rawKey.split(',').map(s => s.trim()).filter(s => s.length > 0);
-            parts.forEach(p => uniqueKeys.add(p));
+        // Deduplicate
+        const uniqueMap = new Map<string, { weight: number; provider: string }>();
+        inputKeys.forEach(k => {
+            if (k.key.length > 0) uniqueMap.set(k.key, { weight: k.weight ?? 1.0, provider: k.provider ?? 'default' });
         });
 
-        // 2. Warn about duplicates/empty
-        const sanitizedCount = uniqueKeys.size;
-        const rawLength = initialKeys.length; // Might be misleading if split happened, but roughly
-        if (sanitizedCount < rawLength) {
-            console.warn(`[ApiKeyManager] Optimized key pool: Removed ${rawLength - sanitizedCount} duplicate/empty keys.`);
+        if (uniqueMap.size < inputKeys.length) {
+            console.warn(`[ApiKeyManager] Removed ${inputKeys.length - uniqueMap.size} duplicate/empty keys.`);
         }
 
-        this.keys = Array.from(uniqueKeys).map(k => ({
-            key: k,
+        this.keys = Array.from(uniqueMap.entries()).map(([key, meta]) => ({
+            key,
             failCount: 0,
             failedAt: null,
             isQuotaError: false,
@@ -99,11 +265,18 @@ export class ApiKeyManager {
             totalRequests: 0,
             halfOpenTestTime: null,
             customCooldown: null,
+            weight: meta.weight,
+            averageLatency: 0,
+            totalLatency: 0,
+            latencySamples: 0,
+            provider: meta.provider,
         }));
 
         this.loadState();
     }
 
+    // ─── Error Classification ────────────────────────────────────────────────
+
     /**
      * CLASSIFIES an error to determine handling strategy
      */
@@ -111,15 +284,16 @@ export class ApiKeyManager {
         const status = error?.status || error?.response?.status;
         const message = error?.message || error?.error?.message || String(error);
 
-        // 1. Check finishReason first (for 200 responses with content issues)
-        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 };
+        // 1. Check finishReason first
+        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 timeout
+        if (error instanceof TimeoutError || error?.name === 'TimeoutError') {
+            return { type: 'TIMEOUT', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
         }
 
-        // 2. Check HTTP status codes
+        // 3. Check HTTP status codes
         if (status === 403 || ERROR_PATTERNS.isAuthError.test(message)) {
             return { type: 'AUTH', retryable: false, cooldownMs: Infinity, markKeyFailed: true, markKeyDead: true };
         }
@@ -143,9 +317,6 @@ export class ApiKeyManager {
         return { type: 'UNKNOWN', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
     }
 
-    /**
-     * Parses Retry-After header from error response
-     */
     private parseRetryAfter(error: any): number | null {
         const retryAfter = error?.response?.headers?.['retry-after'] ||
             error?.headers?.['retry-after'] ||
@@ -153,43 +324,32 @@ export class ApiKeyManager {
 
         if (!retryAfter) return null;
 
-        // If it's a number (seconds)
         const seconds = parseInt(retryAfter, 10);
         if (!isNaN(seconds)) return seconds * 1000;
 
-        // If it's a date string
         const date = Date.parse(retryAfter);
         if (!isNaN(date)) return Math.max(0, date - Date.now());
 
         return null;
     }
 
-    /**
-     * HEALTH CHECK
-     * Determines if a key is usable based on Circuit Breaker logic
-     */
+    // ─── Cooldown ────────────────────────────────────────────────────────────
+
     private isOnCooldown(k: KeyState): boolean {
-        // Dead keys are NEVER usable
         if (k.circuitState === 'DEAD') return true;
-
         const now = Date.now();
 
         if (k.circuitState === 'OPEN') {
-            // Check if ready for HALF_OPEN test
             if (k.halfOpenTestTime && now >= k.halfOpenTestTime) {
                 k.circuitState = 'HALF_OPEN';
+                this.emit('circuitHalfOpen', k.key);
                 return false;
             }
             return true;
         }
 
-        // Additional safeguard for custom cooldowns
-        if (k.failedAt && k.customCooldown) {
-            if (now - k.failedAt < k.customCooldown) return true;
-        }
-
-        // Standard cooldown check
         if (k.failedAt) {
+            if (k.customCooldown && now - k.failedAt < k.customCooldown) return true;
             const cooldown = k.isQuotaError ? CONFIG.COOLDOWN_QUOTA : CONFIG.COOLDOWN_TRANSIENT;
             if (now - k.failedAt < cooldown) return true;
         }
@@ -197,53 +357,70 @@ export class ApiKeyManager {
         return false;
     }
 
-    /**
-     * CORE ROTATION LOGIC
-     * Returns the best available key
-     */
+    // ─── 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)
-        );
+        const candidates = this.keys.filter(k => k.circuitState !== 'DEAD' && !this.isOnCooldown(k));
 
         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; // All keys are dead!
-
+            if (nonDead.length === 0) {
+                this.emit('allKeysExhausted');
+                return null;
+            }
             return nonDead.sort((a, b) => (a.failedAt || 0) - (b.failedAt || 0))[0]?.key || null;
         }
 
-        // 2. Sort candidates: 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;
-        });
-
-        const selected = candidates[0];
-        selected.lastUsed = Date.now();
-        this.saveState();
+        // 2. Delegate to Strategy
+        const selected = this.strategy.next(candidates);
 
-        return selected.key;
+        if (selected) {
+            selected.lastUsed = Date.now();
+            this.saveState();
+            return selected.key;
+        }
+        return null;
     }
 
     /**
-     * Get count of healthy (non-DEAD) keys
+     * 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 ───────────────────────────────────────────────
+
     /**
-     * FEEDBACK LOOP: Success
+     * Mark success AND update latency stats
+     * @param durationMs Duration of the request in milliseconds
      */
-    public markSuccess(key: string) {
+    public markSuccess(key: string, durationMs?: number) {
         const k = this.keys.find(x => x.key === key);
         if (!k) return;
 
-        if (k.circuitState !== 'CLOSED' && k.circuitState !== 'DEAD') {
+        const wasRecovering = k.circuitState !== 'CLOSED' && k.circuitState !== 'DEAD';
+        if (wasRecovering) {
             console.log(`[Key Recovered] ...${key.slice(-4)}`);
+            this.emit('keyRecovered', key);
         }
 
         k.circuitState = 'CLOSED';
@@ -254,21 +431,18 @@ export class ApiKeyManager {
         k.successCount++;
         k.totalRequests++;
 
+        if (durationMs !== undefined) {
+            k.totalLatency += durationMs;
+            k.latencySamples++;
+            k.averageLatency = k.totalLatency / k.latencySamples;
+        }
+
         this.saveState();
     }
 
-    /**
-     * FEEDBACK LOOP: Failure
-     * Enhanced with error classification
-     */
     public markFailed(key: string, classification: ErrorClassification) {
         const k = this.keys.find(x => x.key === key);
-        if (!k) return;
-
-        // Don't modify DEAD keys
-        if (k.circuitState === 'DEAD') return;
-
-        // If this error shouldn't mark the key as failed, skip
+        if (!k || k.circuitState === 'DEAD') return;
         if (!classification.markKeyFailed) return;
 
         k.failedAt = Date.now();
@@ -277,29 +451,25 @@ export class ApiKeyManager {
         k.isQuotaError = classification.type === 'QUOTA';
         k.customCooldown = classification.cooldownMs || null;
 
-        // Permanent death for auth errors
         if (classification.markKeyDead) {
             k.circuitState = 'DEAD';
-            console.error(`[Key DEAD] ...${key.slice(-4)} - Permanently removed from rotation`);
-            this.saveState();
-            return;
-        }
-
-        // State Transitions
-        if (k.circuitState === 'HALF_OPEN') {
-            k.circuitState = 'OPEN';
-            k.halfOpenTestTime = Date.now() + CONFIG.HALF_OPEN_TEST_DELAY;
-        } 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);
+            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();
     }
 
-    /**
-     * Legacy markFailed for backward compatibility
-     */
     public markFailedLegacy(key: string, isQuota: boolean = false) {
         this.markFailed(key, {
             type: isQuota ? 'QUOTA' : 'TRANSIENT',
@@ -310,9 +480,8 @@ export class ApiKeyManager {
         });
     }
 
-    /**
-     * Calculate backoff delay with jitter
-     */
+    // ─── Backoff ─────────────────────────────────────────────────────────────
+
     public calculateBackoff(attempt: number): number {
         const exponential = CONFIG.BASE_BACKOFF * Math.pow(2, attempt);
         const capped = Math.min(exponential, CONFIG.MAX_BACKOFF);
@@ -320,13 +489,8 @@ export class ApiKeyManager {
         return capped + jitter;
     }
 
+    // ─── Stats ───────────────────────────────────────────────────────────────
 
-
-    // ... inside class ...
-
-    /**
-     * Get health statistics
-     */
     public getStats(): ApiKeyManagerStats {
         const total = this.keys.length;
         const dead = this.keys.filter(k => k.circuitState === 'DEAD').length;
@@ -335,11 +499,183 @@ export class ApiKeyManager {
         return { total, healthy, cooling, dead };
     }
 
-    // Helper for testing
-    public _getKeys(): KeyState[] {
-        return this.keys;
+    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) => ({
@@ -353,6 +689,11 @@ export class ApiKeyManager {
                 successCount: k.successCount,
                 totalRequests: k.totalRequests,
                 customCooldown: k.customCooldown,
+                weight: k.weight,
+                averageLatency: k.averageLatency,
+                totalLatency: k.totalLatency,
+                latencySamples: k.latencySamples,
+                provider: k.provider
             }
         }), {});
         this.storage.setItem(this.storageKey, JSON.stringify(state));