@splashcodex/api-key-manager 2.0.0 → 4.0.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/src/index.ts CHANGED
@@ -1,9 +1,15 @@
1
1
  /**
2
- * Universal ApiKeyManager v2.0
3
- * Implements: Rotation, Circuit Breaker, Persistence, Exponential Backoff, Strategies
2
+ * Universal ApiKeyManager v4.0 — Mastermind Edition
3
+ * Implements: Rotation, Circuit Breaker, Persistence, Exponential Backoff, Strategies,
4
+ * Event Emitter, Fallback, execute(), Timeout, Auto-Retry, Provider Tags,
5
+ * Health Checks, Bulkhead/Concurrency
4
6
  * Gemini-Specific: finishReason handling, Safety blocks, RECITATION detection
5
7
  */
6
8
 
9
+ import { EventEmitter } from 'events';
10
+
11
+ // ─── Interfaces & Types ──────────────────────────────────────────────────────
12
+
7
13
  export interface KeyState {
8
14
  key: string;
9
15
  failCount: number; // Consecutive failures
@@ -20,6 +26,8 @@ export interface KeyState {
20
26
  averageLatency: number; // Rolling average latency in ms
21
27
  totalLatency: number; // Sum of all latency checks (for calculating average)
22
28
  latencySamples: number; // Number of samples
29
+ // v3.0 Fields
30
+ provider: string; // Provider tag (e.g. 'openai', 'gemini')
23
31
  }
24
32
 
25
33
  export type ErrorType =
@@ -29,6 +37,7 @@ export type ErrorType =
29
37
  | 'BAD_REQUEST' // 400 - Do not retry, fix request
30
38
  | 'SAFETY' // finishReason: SAFETY - Not a key issue
31
39
  | 'RECITATION' // finishReason: RECITATION - Not a key issue
40
+ | 'TIMEOUT' // Request timed out
32
41
  | 'UNKNOWN'; // Catch-all
33
42
 
34
43
  export interface ErrorClassification {
@@ -39,6 +48,58 @@ export interface ErrorClassification {
39
48
  markKeyDead: boolean;
40
49
  }
41
50
 
51
+ export interface ApiKeyManagerStats {
52
+ total: number;
53
+ healthy: number;
54
+ cooling: number;
55
+ dead: number;
56
+ }
57
+
58
+ export interface ExecuteOptions {
59
+ timeoutMs?: number; // Timeout per attempt in ms
60
+ maxRetries?: number; // Max retry attempts (default: 0 = no retry)
61
+ finishReason?: string; // For Gemini finishReason handling
62
+ provider?: string; // Filter keys by provider (e.g. 'openai')
63
+ }
64
+
65
+ export interface ApiKeyManagerOptions {
66
+ storage?: any;
67
+ strategy?: LoadBalancingStrategy;
68
+ fallbackFn?: () => any;
69
+ concurrency?: number; // Max concurrent execute() calls
70
+ semanticCache?: {
71
+ threshold?: number; // Similarity threshold (0.0 - 1.0, default 0.95)
72
+ ttlMs?: number; // Cache TTL
73
+ getEmbedding: (text: string) => Promise<number[]>;
74
+ };
75
+ }
76
+
77
+ export interface CacheEntry {
78
+ vector: number[];
79
+ prompt: string;
80
+ response: any;
81
+ timestamp: number;
82
+ }
83
+
84
+ // ─── Event Types ─────────────────────────────────────────────────────────────
85
+
86
+ export interface ApiKeyManagerEventMap {
87
+ keyDead: (key: string) => void;
88
+ circuitOpen: (key: string) => void;
89
+ circuitHalfOpen: (key: string) => void;
90
+ keyRecovered: (key: string) => void;
91
+ fallback: (reason: string) => void;
92
+ allKeysExhausted: () => void;
93
+ retry: (key: string, attempt: number, delayMs: number) => void;
94
+ healthCheckFailed: (key: string, error: any) => void;
95
+ healthCheckPassed: (key: string) => void;
96
+ executeSuccess: (key: string, durationMs: number) => void;
97
+ executeFailed: (key: string, error: any) => void;
98
+ bulkheadRejected: () => void;
99
+ }
100
+
101
+ // ─── Config ──────────────────────────────────────────────────────────────────
102
+
42
103
  const CONFIG = {
43
104
  MAX_CONSECUTIVE_FAILURES: 5,
44
105
  COOLDOWN_TRANSIENT: 60 * 1000, // 1 minute
@@ -58,13 +119,31 @@ const ERROR_PATTERNS = {
58
119
  isBadRequest: /400|invalid.?argument|failed.?precondition|malformed|not.?found|404/i,
59
120
  };
60
121
 
61
- export interface ApiKeyManagerStats {
62
- total: number;
63
- healthy: number;
64
- cooling: number;
65
- dead: number;
122
+ // ─── Custom Errors ───────────────────────────────────────────────────────────
123
+
124
+ export class TimeoutError extends Error {
125
+ constructor(ms: number) {
126
+ super(`Request timed out after ${ms}ms`);
127
+ this.name = 'TimeoutError';
128
+ }
129
+ }
130
+
131
+ export class BulkheadRejectionError extends Error {
132
+ constructor() {
133
+ super('Bulkhead capacity exceeded — too many concurrent requests');
134
+ this.name = 'BulkheadRejectionError';
135
+ }
66
136
  }
67
137
 
138
+ export class AllKeysExhaustedError extends Error {
139
+ constructor() {
140
+ super('All API keys exhausted — no healthy keys available');
141
+ this.name = 'AllKeysExhaustedError';
142
+ }
143
+ }
144
+
145
+ // ─── Strategies ──────────────────────────────────────────────────────────────
146
+
68
147
  /**
69
148
  * Strategy Interface for selecting the next key
70
149
  */
@@ -77,7 +156,6 @@ export interface LoadBalancingStrategy {
77
156
  */
78
157
  export class StandardStrategy implements LoadBalancingStrategy {
79
158
  next(candidates: KeyState[]): KeyState | null {
80
- // Sort: Pristine > Fewest Failures > Least Recently Used
81
159
  candidates.sort((a, b) => {
82
160
  if (a.failCount !== b.failCount) return a.failCount - b.failCount;
83
161
  return a.lastUsed - b.lastUsed;
@@ -112,49 +190,171 @@ export class WeightedStrategy implements LoadBalancingStrategy {
112
190
  export class LatencyStrategy implements LoadBalancingStrategy {
113
191
  next(candidates: KeyState[]): KeyState | null {
114
192
  if (candidates.length === 0) return null;
115
- // Sort by averageLatency (lowest first)
116
- // If latency is 0 (untried), treat as high priority or neutral?
117
- // Let's treat 0 as "unknown, give it a shot" -> insert at top or mixed?
118
- // Simple: Sort ASC. 0 comes first.
119
193
  candidates.sort((a, b) => a.averageLatency - b.averageLatency);
120
194
  return candidates[0];
121
195
  }
122
196
  }
123
197
 
124
- export class ApiKeyManager {
198
+ // ─── Semantic Engine ─────────────────────────────────────────────────────────
199
+
200
+ /**
201
+ * High-performance Vanilla Semantic Cache
202
+ * Implements Cosine Similarity math from scratch.
203
+ */
204
+ export class SemanticCache {
205
+ private entries: CacheEntry[] = [];
206
+ private threshold: number;
207
+ private ttlMs: number;
208
+
209
+ constructor(threshold: number = 0.95, ttlMs: number = 24 * 60 * 60 * 1000) {
210
+ this.threshold = threshold;
211
+ this.ttlMs = ttlMs;
212
+ }
213
+
214
+ public set(prompt: string, vector: number[], response: any) {
215
+ // Expire old entry for same prompt if exists
216
+ this.entries = this.entries.filter(e => e.prompt !== prompt);
217
+ this.entries.push({
218
+ prompt,
219
+ vector,
220
+ response,
221
+ timestamp: Date.now()
222
+ });
223
+ // Optional: Cap size to prevent memory leaks
224
+ if (this.entries.length > 500) this.entries.shift();
225
+ }
226
+
227
+ public get(vector: number[]): any | null {
228
+ const now = Date.now();
229
+ let bestMatch: CacheEntry | null = null;
230
+ let highestSimilarity = -1;
231
+
232
+ for (let i = this.entries.length - 1; i >= 0; i--) {
233
+ const entry = this.entries[i];
234
+
235
+ // Check TTL
236
+ if (now - entry.timestamp > this.ttlMs) {
237
+ this.entries.splice(i, 1);
238
+ continue;
239
+ }
240
+
241
+ const similarity = this.calculateCosineSimilarity(vector, entry.vector);
242
+ if (similarity >= this.threshold && similarity > highestSimilarity) {
243
+ highestSimilarity = similarity;
244
+ bestMatch = entry;
245
+ }
246
+ }
247
+
248
+ return bestMatch ? bestMatch.response : null;
249
+ }
250
+
251
+ /**
252
+ * Vanilla Cosine Similarity: (A·B) / (||A|| * ||B||)
253
+ */
254
+ private calculateCosineSimilarity(vecA: number[], vecB: number[]): number {
255
+ if (vecA.length !== vecB.length) return 0;
256
+ let dotProduct = 0;
257
+ let normA = 0;
258
+ let normB = 0;
259
+
260
+ for (let i = 0; i < vecA.length; i++) {
261
+ dotProduct += vecA[i] * vecB[i];
262
+ normA += vecA[i] * vecA[i];
263
+ normB += vecB[i] * vecB[i];
264
+ }
265
+
266
+ const denominator = Math.sqrt(normA) * Math.sqrt(normB);
267
+ return denominator === 0 ? 0 : dotProduct / denominator;
268
+ }
269
+ }
270
+
271
+ // ─── Main Class ──────────────────────────────────────────────────────────────
272
+
273
+ export class ApiKeyManager extends EventEmitter {
125
274
  private keys: KeyState[] = [];
126
275
  private storageKey = 'api_rotation_state_v2';
127
276
  private storage: any;
128
277
  private strategy: LoadBalancingStrategy;
278
+ private fallbackFn?: () => any;
279
+
280
+ // Bulkhead state
281
+ private maxConcurrency: number;
282
+ private activeCalls: number = 0;
129
283
 
130
- constructor(initialKeys: string[] | { key: string; weight?: number }[], storage?: any, strategy?: LoadBalancingStrategy) {
131
- // Simple in-memory storage mock if none provided
132
- this.storage = storage || {
284
+ // Health check state
285
+ private healthCheckFn?: (key: string) => Promise<boolean>;
286
+ private healthCheckInterval?: ReturnType<typeof setInterval>;
287
+
288
+ // Semantic Cache v4
289
+ private semanticCache?: SemanticCache;
290
+ private getEmbeddingFn?: (text: string) => Promise<number[]>;
291
+ private _isResolvingEmbedding: boolean = false; // Recursion guard
292
+
293
+ /**
294
+ * Constructor supports both legacy positional args and new options object.
295
+ *
296
+ * @example Legacy (v1/v2 — still works):
297
+ * new ApiKeyManager(['key1', 'key2'], storage, strategy)
298
+ *
299
+ * @example New (v3):
300
+ * new ApiKeyManager(keys, { storage, strategy, fallbackFn, concurrency })
301
+ */
302
+ constructor(
303
+ initialKeys: string[] | { key: string; weight?: number; provider?: string }[],
304
+ storageOrOptions?: any | ApiKeyManagerOptions,
305
+ strategy?: LoadBalancingStrategy
306
+ ) {
307
+ super();
308
+
309
+ // Detect if second arg is options object or legacy storage
310
+ let options: ApiKeyManagerOptions = {};
311
+ if (storageOrOptions && typeof storageOrOptions === 'object' && ('storage' in storageOrOptions || 'strategy' in storageOrOptions || 'fallbackFn' in storageOrOptions || 'concurrency' in storageOrOptions || 'semanticCache' in storageOrOptions)) {
312
+ // New v3 options object
313
+ options = storageOrOptions as ApiKeyManagerOptions;
314
+ } else {
315
+ // Legacy positional args
316
+ options = {
317
+ storage: storageOrOptions,
318
+ strategy: strategy,
319
+ };
320
+ }
321
+
322
+ this.storage = options.storage || {
133
323
  getItem: () => null,
134
324
  setItem: () => { },
135
325
  };
136
-
137
- this.strategy = strategy || new StandardStrategy();
326
+ this.strategy = options.strategy || new StandardStrategy();
327
+ this.fallbackFn = options.fallbackFn;
328
+ this.maxConcurrency = options.concurrency || Infinity;
329
+
330
+ // Init Semantic Cache if provided
331
+ if (options.semanticCache) {
332
+ this.semanticCache = new SemanticCache(
333
+ options.semanticCache.threshold,
334
+ options.semanticCache.ttlMs
335
+ );
336
+ this.getEmbeddingFn = options.semanticCache.getEmbedding;
337
+ }
138
338
 
139
339
  // Normalize input to objects
140
- let inputKeys: { key: string; weight?: number }[] = [];
340
+ let inputKeys: { key: string; weight?: number; provider?: string }[] = [];
141
341
  if (initialKeys.length > 0 && typeof initialKeys[0] === 'string') {
142
- inputKeys = (initialKeys as string[]).flatMap(k => k.split(',').map(s => ({ key: s.trim(), weight: 1.0 })));
342
+ inputKeys = (initialKeys as string[]).flatMap(k => k.split(',').map(s => ({ key: s.trim(), weight: 1.0, provider: 'default' })));
143
343
  } else {
144
- inputKeys = initialKeys as { key: string; weight?: number }[];
344
+ inputKeys = initialKeys as { key: string; weight?: number; provider?: string }[];
145
345
  }
146
346
 
147
347
  // Deduplicate
148
- const uniqueMap = new Map<string, number>();
348
+ const uniqueMap = new Map<string, { weight: number; provider: string }>();
149
349
  inputKeys.forEach(k => {
150
- if (k.key.length > 0) uniqueMap.set(k.key, k.weight ?? 1.0);
350
+ if (k.key.length > 0) uniqueMap.set(k.key, { weight: k.weight ?? 1.0, provider: k.provider ?? 'default' });
151
351
  });
152
352
 
153
353
  if (uniqueMap.size < inputKeys.length) {
154
354
  console.warn(`[ApiKeyManager] Removed ${inputKeys.length - uniqueMap.size} duplicate/empty keys.`);
155
355
  }
156
356
 
157
- this.keys = Array.from(uniqueMap.entries()).map(([key, weight]) => ({
357
+ this.keys = Array.from(uniqueMap.entries()).map(([key, meta]) => ({
158
358
  key,
159
359
  failCount: 0,
160
360
  failedAt: null,
@@ -165,15 +365,18 @@ export class ApiKeyManager {
165
365
  totalRequests: 0,
166
366
  halfOpenTestTime: null,
167
367
  customCooldown: null,
168
- weight: weight,
368
+ weight: meta.weight,
169
369
  averageLatency: 0,
170
370
  totalLatency: 0,
171
- latencySamples: 0
371
+ latencySamples: 0,
372
+ provider: meta.provider,
172
373
  }));
173
374
 
174
375
  this.loadState();
175
376
  }
176
377
 
378
+ // ─── Error Classification ────────────────────────────────────────────────
379
+
177
380
  /**
178
381
  * CLASSIFIES an error to determine handling strategy
179
382
  */
@@ -185,7 +388,12 @@ export class ApiKeyManager {
185
388
  if (finishReason === 'SAFETY') return { type: 'SAFETY', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
186
389
  if (finishReason === 'RECITATION') return { type: 'RECITATION', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
187
390
 
188
- // 2. Check HTTP status codes
391
+ // 2. Check timeout
392
+ if (error instanceof TimeoutError || error?.name === 'TimeoutError') {
393
+ return { type: 'TIMEOUT', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
394
+ }
395
+
396
+ // 3. Check HTTP status codes
189
397
  if (status === 403 || ERROR_PATTERNS.isAuthError.test(message)) {
190
398
  return { type: 'AUTH', retryable: false, cooldownMs: Infinity, markKeyFailed: true, markKeyDead: true };
191
399
  }
@@ -225,6 +433,8 @@ export class ApiKeyManager {
225
433
  return null;
226
434
  }
227
435
 
436
+ // ─── Cooldown ────────────────────────────────────────────────────────────
437
+
228
438
  private isOnCooldown(k: KeyState): boolean {
229
439
  if (k.circuitState === 'DEAD') return true;
230
440
  const now = Date.now();
@@ -232,6 +442,7 @@ export class ApiKeyManager {
232
442
  if (k.circuitState === 'OPEN') {
233
443
  if (k.halfOpenTestTime && now >= k.halfOpenTestTime) {
234
444
  k.circuitState = 'HALF_OPEN';
445
+ this.emit('circuitHalfOpen', k.key);
235
446
  return false;
236
447
  }
237
448
  return true;
@@ -246,6 +457,8 @@ export class ApiKeyManager {
246
457
  return false;
247
458
  }
248
459
 
460
+ // ─── Key Selection ───────────────────────────────────────────────────────
461
+
249
462
  public getKey(): string | null {
250
463
  // 1. Filter out dead and cooling down keys
251
464
  const candidates = this.keys.filter(k => k.circuitState !== 'DEAD' && !this.isOnCooldown(k));
@@ -253,7 +466,10 @@ export class ApiKeyManager {
253
466
  if (candidates.length === 0) {
254
467
  // FALLBACK: Return oldest failed key (excluding DEAD)
255
468
  const nonDead = this.keys.filter(k => k.circuitState !== 'DEAD');
256
- if (nonDead.length === 0) return null;
469
+ if (nonDead.length === 0) {
470
+ this.emit('allKeysExhausted');
471
+ return null;
472
+ }
257
473
  return nonDead.sort((a, b) => (a.failedAt || 0) - (b.failedAt || 0))[0]?.key || null;
258
474
  }
259
475
 
@@ -268,10 +484,31 @@ export class ApiKeyManager {
268
484
  return null;
269
485
  }
270
486
 
487
+ /**
488
+ * Get a key filtered by provider tag
489
+ */
490
+ public getKeyByProvider(provider: string): string | null {
491
+ const candidates = this.keys.filter(k =>
492
+ k.provider === provider && k.circuitState !== 'DEAD' && !this.isOnCooldown(k)
493
+ );
494
+
495
+ if (candidates.length === 0) return null;
496
+
497
+ const selected = this.strategy.next(candidates);
498
+ if (selected) {
499
+ selected.lastUsed = Date.now();
500
+ this.saveState();
501
+ return selected.key;
502
+ }
503
+ return null;
504
+ }
505
+
271
506
  public getKeyCount(): number {
272
507
  return this.keys.filter(k => k.circuitState !== 'DEAD').length;
273
508
  }
274
509
 
510
+ // ─── Mark Success / Failed ───────────────────────────────────────────────
511
+
275
512
  /**
276
513
  * Mark success AND update latency stats
277
514
  * @param durationMs Duration of the request in milliseconds
@@ -280,7 +517,11 @@ export class ApiKeyManager {
280
517
  const k = this.keys.find(x => x.key === key);
281
518
  if (!k) return;
282
519
 
283
- if (k.circuitState !== 'CLOSED' && k.circuitState !== 'DEAD') console.log(`[Key Recovered] ...${key.slice(-4)}`);
520
+ const wasRecovering = k.circuitState !== 'CLOSED' && k.circuitState !== 'DEAD';
521
+ if (wasRecovering) {
522
+ console.log(`[Key Recovered] ...${key.slice(-4)}`);
523
+ this.emit('keyRecovered', key);
524
+ }
284
525
 
285
526
  k.circuitState = 'CLOSED';
286
527
  k.failCount = 0;
@@ -313,14 +554,17 @@ export class ApiKeyManager {
313
554
  if (classification.markKeyDead) {
314
555
  k.circuitState = 'DEAD';
315
556
  console.error(`[Key DEAD] ...${key.slice(-4)} - Permanently removed`);
557
+ this.emit('keyDead', key);
316
558
  } else {
317
559
  // State Transitions
318
560
  if (k.circuitState === 'HALF_OPEN') {
319
561
  k.circuitState = 'OPEN';
320
562
  k.halfOpenTestTime = Date.now() + CONFIG.HALF_OPEN_TEST_DELAY;
563
+ this.emit('circuitOpen', key);
321
564
  } else if (k.failCount >= CONFIG.MAX_CONSECUTIVE_FAILURES || classification.type === 'QUOTA') {
322
565
  k.circuitState = 'OPEN';
323
566
  k.halfOpenTestTime = Date.now() + (classification.cooldownMs || CONFIG.HALF_OPEN_TEST_DELAY);
567
+ this.emit('circuitOpen', key);
324
568
  }
325
569
  }
326
570
  this.saveState();
@@ -336,6 +580,8 @@ export class ApiKeyManager {
336
580
  });
337
581
  }
338
582
 
583
+ // ─── Backoff ─────────────────────────────────────────────────────────────
584
+
339
585
  public calculateBackoff(attempt: number): number {
340
586
  const exponential = CONFIG.BASE_BACKOFF * Math.pow(2, attempt);
341
587
  const capped = Math.min(exponential, CONFIG.MAX_BACKOFF);
@@ -343,6 +589,8 @@ export class ApiKeyManager {
343
589
  return capped + jitter;
344
590
  }
345
591
 
592
+ // ─── Stats ───────────────────────────────────────────────────────────────
593
+
346
594
  public getStats(): ApiKeyManagerStats {
347
595
  const total = this.keys.length;
348
596
  const dead = this.keys.filter(k => k.circuitState === 'DEAD').length;
@@ -353,6 +601,212 @@ export class ApiKeyManager {
353
601
 
354
602
  public _getKeys(): KeyState[] { return this.keys; }
355
603
 
604
+ // ─── execute() Wrapper ───────────────────────────────────────────────────
605
+
606
+ /**
607
+ * Wraps the entire API call lifecycle into a single method.
608
+ *
609
+ * @example
610
+ * const result = await manager.execute(
611
+ * (key) => fetch(`https://api.example.com?key=${key}`),
612
+ * { maxRetries: 3, timeoutMs: 5000 }
613
+ * );
614
+ */
615
+ public async execute<T>(
616
+ fn: (key: string, signal?: AbortSignal) => Promise<T>,
617
+ options?: ExecuteOptions & { prompt?: string }
618
+ ): Promise<T> {
619
+ const maxRetries = options?.maxRetries ?? 0;
620
+ const timeoutMs = options?.timeoutMs;
621
+ const finishReason = options?.finishReason;
622
+ const prompt = options?.prompt;
623
+ const provider = options?.provider;
624
+
625
+ // 1. Semantic Cache Check (Mastermind Edition)
626
+ // Guard: skip cache if we're already resolving an embedding to prevent
627
+ // infinite recursion when getEmbeddingFn calls execute() internally.
628
+ let currentPromptVector: number[] | null = null;
629
+ if (this.semanticCache && this.getEmbeddingFn && prompt && !this._isResolvingEmbedding) {
630
+ try {
631
+ this._isResolvingEmbedding = true;
632
+ currentPromptVector = await this.getEmbeddingFn(prompt);
633
+ const cachedResponse = this.semanticCache.get(currentPromptVector);
634
+ if (cachedResponse !== null) {
635
+ console.log(`[Semantic Cache HIT] for prompt: "${prompt.slice(0, 30)}..."`);
636
+ this.emit('executeSuccess', 'CACHE_HIT', 0);
637
+ return cachedResponse as T;
638
+ }
639
+ } catch (e) {
640
+ console.warn('[Semantic Cache Check Failed] Proceeding to live API', e);
641
+ } finally {
642
+ this._isResolvingEmbedding = false;
643
+ }
644
+ }
645
+
646
+ // 2. Bulkhead check
647
+ if (this.activeCalls >= this.maxConcurrency) {
648
+ this.emit('bulkheadRejected');
649
+ throw new BulkheadRejectionError();
650
+ }
651
+
652
+ this.activeCalls++;
653
+ try {
654
+ const result = await this._executeWithRetry(fn, maxRetries, timeoutMs, finishReason, provider);
655
+
656
+ // 3. Store in Semantic Cache on success
657
+ if (this.semanticCache && prompt && currentPromptVector) {
658
+ this.semanticCache.set(prompt, currentPromptVector, result);
659
+ }
660
+
661
+ return result;
662
+ } finally {
663
+ this.activeCalls--;
664
+ }
665
+ }
666
+
667
+ private async _executeWithRetry<T>(
668
+ fn: (key: string, signal?: AbortSignal) => Promise<T>,
669
+ maxRetries: number,
670
+ timeoutMs?: number,
671
+ finishReason?: string,
672
+ provider?: string
673
+ ): Promise<T> {
674
+ let lastError: any;
675
+
676
+ for (let attempt = 0; attempt <= maxRetries; attempt++) {
677
+ const key = provider ? this.getKeyByProvider(provider) : this.getKey();
678
+
679
+ if (!key) {
680
+ // All keys exhausted — try fallback
681
+ if (this.fallbackFn) {
682
+ this.emit('fallback', 'all keys exhausted');
683
+ return this.fallbackFn();
684
+ }
685
+ throw new AllKeysExhaustedError();
686
+ }
687
+
688
+ try {
689
+ const start = Date.now();
690
+ let result: T;
691
+
692
+ if (timeoutMs) {
693
+ result = await this._executeWithTimeout(fn, key, timeoutMs);
694
+ } else {
695
+ result = await fn(key);
696
+ }
697
+
698
+ const duration = Date.now() - start;
699
+ this.markSuccess(key, duration);
700
+ this.emit('executeSuccess', key, duration);
701
+ return result;
702
+
703
+ } catch (error: any) {
704
+ lastError = error;
705
+ const classification = this.classifyError(error, finishReason);
706
+
707
+ this.markFailed(key, classification);
708
+ this.emit('executeFailed', key, error);
709
+
710
+ if (!classification.retryable || attempt >= maxRetries) {
711
+ // Non-retryable or out of retries
712
+ if (this.fallbackFn && attempt >= maxRetries) {
713
+ this.emit('fallback', 'max retries exceeded');
714
+ return this.fallbackFn();
715
+ }
716
+ throw error;
717
+ }
718
+
719
+ // Retry with backoff
720
+ const delay = this.calculateBackoff(attempt);
721
+ this.emit('retry', key, attempt + 1, delay);
722
+ await this._sleep(delay);
723
+ }
724
+ }
725
+
726
+ // Should not reach here, but safety net
727
+ throw lastError;
728
+ }
729
+
730
+ private async _executeWithTimeout<T>(
731
+ fn: (key: string, signal?: AbortSignal) => Promise<T>,
732
+ key: string,
733
+ timeoutMs: number
734
+ ): Promise<T> {
735
+ const controller = new AbortController();
736
+ const timer = setTimeout(() => controller.abort(), timeoutMs);
737
+
738
+ try {
739
+ const result = await Promise.race([
740
+ fn(key, controller.signal),
741
+ new Promise<never>((_, reject) => {
742
+ controller.signal.addEventListener('abort', () => {
743
+ reject(new TimeoutError(timeoutMs));
744
+ });
745
+ })
746
+ ]);
747
+ return result;
748
+ } finally {
749
+ clearTimeout(timer);
750
+ }
751
+ }
752
+
753
+ private _sleep(ms: number): Promise<void> {
754
+ return new Promise(resolve => setTimeout(resolve, ms));
755
+ }
756
+
757
+ // ─── Health Checks ───────────────────────────────────────────────────────
758
+
759
+ /**
760
+ * Set a health check function that tests if a key is operational
761
+ */
762
+ public setHealthCheck(fn: (key: string) => Promise<boolean>) {
763
+ this.healthCheckFn = fn;
764
+ }
765
+
766
+ /**
767
+ * Start periodic health checks
768
+ * @param intervalMs How often to run health checks (default: 60s)
769
+ */
770
+ public startHealthChecks(intervalMs: number = 60_000) {
771
+ this.stopHealthChecks(); // Clear any existing interval
772
+ this.healthCheckInterval = setInterval(() => this._runHealthChecks(), intervalMs);
773
+ }
774
+
775
+ /**
776
+ * Stop periodic health checks
777
+ */
778
+ public stopHealthChecks() {
779
+ if (this.healthCheckInterval) {
780
+ clearInterval(this.healthCheckInterval);
781
+ this.healthCheckInterval = undefined;
782
+ }
783
+ }
784
+
785
+ private async _runHealthChecks() {
786
+ if (!this.healthCheckFn) return;
787
+
788
+ // Check non-DEAD keys that are in OPEN or HALF_OPEN state
789
+ const keysToCheck = this.keys.filter(k =>
790
+ k.circuitState === 'OPEN' || k.circuitState === 'HALF_OPEN'
791
+ );
792
+
793
+ for (const k of keysToCheck) {
794
+ try {
795
+ const healthy = await this.healthCheckFn(k.key);
796
+ if (healthy) {
797
+ this.markSuccess(k.key);
798
+ this.emit('healthCheckPassed', k.key);
799
+ } else {
800
+ this.emit('healthCheckFailed', k.key, new Error('Health check returned false'));
801
+ }
802
+ } catch (error) {
803
+ this.emit('healthCheckFailed', k.key, error);
804
+ }
805
+ }
806
+ }
807
+
808
+ // ─── Persistence ─────────────────────────────────────────────────────────
809
+
356
810
  private saveState() {
357
811
  if (!this.storage) return;
358
812
  const state = this.keys.reduce((acc, k) => ({
@@ -369,7 +823,8 @@ export class ApiKeyManager {
369
823
  weight: k.weight,
370
824
  averageLatency: k.averageLatency,
371
825
  totalLatency: k.totalLatency,
372
- latencySamples: k.latencySamples
826
+ latencySamples: k.latencySamples,
827
+ provider: k.provider
373
828
  }
374
829
  }), {});
375
830
  this.storage.setItem(this.storageKey, JSON.stringify(state));