@splashcodex/api-key-manager 5.2.0 → 5.3.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,1071 +1,1155 @@
1
- /**
2
- * Universal ApiKeyManager v5.0 — Ecosystem 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
6
- * NEW in v5.0: Provider Presets (GeminiManager, OpenAIManager, MultiManager),
7
- * Built-in Persistence (FileStorage, MemoryStorage)
8
- * Gemini-Specific: finishReason handling, Safety blocks, RECITATION detection
9
- */
10
-
11
- import { EventEmitter } from 'events';
12
-
13
- // ─── Re-exports: Persistence ─────────────────────────────────────────────────
14
- // Persistence adapters can be imported from root or via subpath
15
- export { FileStorage } from './persistence/file';
16
- export type { FileStorageOptions } from './persistence/file';
17
- export { MemoryStorage } from './persistence/memory';
18
-
19
- // ─── Presets ─────────────────────────────────────────────────────────────────
20
- // Presets are available via subpath imports to avoid circular dependencies:
21
- // import { GeminiManager } from '@splashcodex/api-key-manager/presets/gemini';
22
- // import { OpenAIManager } from '@splashcodex/api-key-manager/presets/openai';
23
- // import { MultiManager } from '@splashcodex/api-key-manager/presets/multi';
24
- // Or import all at once:
25
- // import { GeminiManager, OpenAIManager, MultiManager } from '@splashcodex/api-key-manager/presets';
26
-
27
-
28
-
29
- // ─── Interfaces & Types ──────────────────────────────────────────────────────
30
-
31
- export interface KeyState {
32
- key: string;
33
- failCount: number; // Consecutive failures
34
- failedAt: number | null; // Timestamp of last failure
35
- isQuotaError: boolean; // Was last error a 429?
36
- circuitState: 'CLOSED' | 'OPEN' | 'HALF_OPEN' | 'DEAD';
37
- lastUsed: number;
38
- successCount: number;
39
- totalRequests: number;
40
- halfOpenTestTime: number | null;
41
- customCooldown: number | null; // From Retry-After header
42
- // v2.0 Stats
43
- weight: number; // 0.0 - 1.0 (Default 1.0)
44
- averageLatency: number; // Rolling average latency in ms
45
- totalLatency: number; // Sum of all latency checks (for calculating average)
46
- latencySamples: number; // Number of samples
47
- // v3.0 Fields
48
- provider: string; // Provider tag (e.g. 'openai', 'gemini')
49
- }
50
-
51
- export type ErrorType =
52
- | 'QUOTA' // 429 - Rotate key, respect cooldown
53
- | 'TRANSIENT' // 500/503/504 - Retry with backoff
54
- | 'AUTH' // 403 - Key is dead, remove from pool
55
- | 'BAD_REQUEST' // 400 - Do not retry, fix request
56
- | 'SAFETY' // finishReason: SAFETY - Not a key issue
57
- | 'RECITATION' // finishReason: RECITATION - Not a key issue
58
- | 'TIMEOUT' // Request timed out
59
- | 'UNKNOWN'; // Catch-all
60
-
61
- export interface ErrorClassification {
62
- type: ErrorType;
63
- retryable: boolean;
64
- cooldownMs: number;
65
- markKeyFailed: boolean;
66
- markKeyDead: boolean;
67
- }
68
-
69
- export interface ApiKeyManagerStats {
70
- total: number;
71
- healthy: number;
72
- cooling: number;
73
- dead: number;
74
- }
75
-
76
- export interface ExecuteOptions {
77
- timeoutMs?: number; // Timeout per attempt in ms
78
- maxRetries?: number; // Max retry attempts (default: 0 = no retry)
79
- finishReason?: string; // For Gemini finishReason handling
80
- provider?: string; // Filter keys by provider (e.g. 'openai')
81
- }
82
-
83
- export interface ApiKeyManagerOptions {
84
- storage?: any;
85
- strategy?: LoadBalancingStrategy;
86
- fallbackFn?: () => any;
87
- concurrency?: number; // Max concurrent execute() calls
88
- semanticCache?: {
89
- threshold?: number; // Similarity threshold (0.0 - 1.0, default 0.95)
90
- ttlMs?: number; // Cache TTL
91
- getEmbedding: (text: string) => Promise<number[]>;
92
- };
93
- }
94
-
95
- export interface CacheEntry {
96
- vector: number[];
97
- prompt: string;
98
- response: any;
99
- timestamp: number;
100
- }
101
-
102
- // ─── Event Types ─────────────────────────────────────────────────────────────
103
-
104
- export interface ApiKeyManagerEventMap {
105
- keyDead: (key: string) => void;
106
- circuitOpen: (key: string) => void;
107
- circuitHalfOpen: (key: string) => void;
108
- keyRecovered: (key: string) => void;
109
- fallback: (reason: string) => void;
110
- allKeysExhausted: () => void;
111
- retry: (key: string, attempt: number, delayMs: number) => void;
112
- healthCheckFailed: (key: string, error: any) => void;
113
- healthCheckPassed: (key: string) => void;
114
- executeSuccess: (key: string, durationMs: number) => void;
115
- executeFailed: (key: string, error: any) => void;
116
- bulkheadRejected: () => void;
117
- }
118
-
119
- // ─── Config ──────────────────────────────────────────────────────────────────
120
-
121
- const CONFIG = {
122
- MAX_CONSECUTIVE_FAILURES: 5,
123
- COOLDOWN_TRANSIENT: 60 * 1000, // 1 minute
124
- COOLDOWN_QUOTA: 5 * 60 * 1000, // 5 minutes (default if no Retry-After)
125
- COOLDOWN_QUOTA_DAILY: 60 * 60 * 1000, // 1 hour for RPD exhaustion
126
- HALF_OPEN_TEST_DELAY: 60 * 1000, // 1 minute after open
127
- MAX_BACKOFF: 64 * 1000, // 64 seconds max
128
- BASE_BACKOFF: 1000, // 1 second base
129
- DEAD_KEY_TTL: 60 * 60 * 1000, // 1 hour — DEAD keys get retested after this
130
- };
131
-
132
- // Error classification patterns
133
- const ERROR_PATTERNS = {
134
- isQuotaError: /429|quota|exhausted|resource.?exhausted|too.?many.?requests|rate.?limit/i,
135
- isAuthError: /403|permission.?denied|invalid.?api.?key|unauthorized|unauthenticated/i,
136
- isSafetyBlock: /safety|blocked|recitation|harmful/i,
137
- isTransient: /500|502|503|504|internal|unavailable|deadline|timeout|overloaded/i,
138
- isBadRequest: /400|invalid.?argument|failed.?precondition|malformed|not.?found|404/i,
139
- };
140
-
141
- // ─── Custom Errors ───────────────────────────────────────────────────────────
142
-
143
- export class TimeoutError extends Error {
144
- constructor(ms: number) {
145
- super(`Request timed out after ${ms}ms`);
146
- this.name = 'TimeoutError';
147
- }
148
- }
149
-
150
- export class BulkheadRejectionError extends Error {
151
- constructor() {
152
- super('Bulkhead capacity exceededtoo many concurrent requests');
153
- this.name = 'BulkheadRejectionError';
154
- }
155
- }
156
-
157
- export class AllKeysExhaustedError extends Error {
158
- constructor() {
159
- super('All API keys exhausted — no healthy keys available');
160
- this.name = 'AllKeysExhaustedError';
161
- }
162
- }
163
-
164
- // ─── Strategies ──────────────────────────────────────────────────────────────
165
-
166
- /**
167
- * Strategy Interface for selecting the next key
168
- */
169
- export interface LoadBalancingStrategy {
170
- next(candidates: KeyState[]): KeyState | null;
171
- }
172
-
173
- /**
174
- * Standard Strategy: Least Failed > Least Recently Used
175
- */
176
- export class StandardStrategy implements LoadBalancingStrategy {
177
- next(candidates: KeyState[]): KeyState | null {
178
- candidates.sort((a, b) => {
179
- if (a.failCount !== b.failCount) return a.failCount - b.failCount;
180
- return a.lastUsed - b.lastUsed;
181
- });
182
- return candidates[0] || null;
183
- }
184
- }
185
-
186
- /**
187
- * Weighted Strategy: Probabilistic selection based on weight
188
- * Higher weight = Higher chance of selection
189
- */
190
- export class WeightedStrategy implements LoadBalancingStrategy {
191
- next(candidates: KeyState[]): KeyState | null {
192
- if (candidates.length === 0) return null;
193
-
194
- const totalWeight = candidates.reduce((sum, k) => sum + k.weight, 0);
195
- let random = Math.random() * totalWeight;
196
-
197
- for (const key of candidates) {
198
- random -= key.weight;
199
- if (random <= 0) return key;
200
- }
201
-
202
- return candidates[0]; // Fallback
203
- }
204
- }
205
-
206
- /**
207
- * Latency Strategy: Pick lowest average latency with LRU tie-break
208
- */
209
- export class LatencyStrategy implements LoadBalancingStrategy {
210
- next(candidates: KeyState[]): KeyState | null {
211
- if (candidates.length === 0) return null;
212
- candidates.sort((a, b) => {
213
- if (a.averageLatency !== b.averageLatency) return a.averageLatency - b.averageLatency;
214
- return a.lastUsed - b.lastUsed; // LRU tie-break
215
- });
216
- return candidates[0];
217
- }
218
- }
219
-
220
- // ─── Semantic Engine ─────────────────────────────────────────────────────────
221
-
222
- /**
223
- * High-performance Vanilla Semantic Cache
224
- * Implements Cosine Similarity math from scratch.
225
- */
226
- export class SemanticCache {
227
- private entries: CacheEntry[] = [];
228
- private threshold: number;
229
- private ttlMs: number;
230
-
231
- constructor(threshold: number = 0.95, ttlMs: number = 24 * 60 * 60 * 1000) {
232
- this.threshold = threshold;
233
- this.ttlMs = ttlMs;
234
- }
235
-
236
- public set(prompt: string, vector: number[], response: any) {
237
- // Expire old entry for same prompt if exists
238
- this.entries = this.entries.filter(e => e.prompt !== prompt);
239
- this.entries.push({
240
- prompt,
241
- vector,
242
- response,
243
- timestamp: Date.now()
244
- });
245
- // Optional: Cap size to prevent memory leaks
246
- if (this.entries.length > 500) this.entries.shift();
247
- }
248
-
249
- public get(vector: number[]): any | null {
250
- const now = Date.now();
251
- let bestMatch: CacheEntry | null = null;
252
- let highestSimilarity = -1;
253
-
254
- for (let i = this.entries.length - 1; i >= 0; i--) {
255
- const entry = this.entries[i];
256
-
257
- // Check TTL
258
- if (now - entry.timestamp > this.ttlMs) {
259
- this.entries.splice(i, 1);
260
- continue;
261
- }
262
-
263
- const similarity = this.calculateCosineSimilarity(vector, entry.vector);
264
- if (similarity >= this.threshold && similarity > highestSimilarity) {
265
- highestSimilarity = similarity;
266
- bestMatch = entry;
267
- }
268
- }
269
-
270
- return bestMatch ? bestMatch.response : null;
271
- }
272
-
273
- /**
274
- * Vanilla Cosine Similarity: (A·B) / (||A|| * ||B||)
275
- */
276
- private calculateCosineSimilarity(vecA: number[], vecB: number[]): number {
277
- if (vecA.length !== vecB.length) return 0;
278
- let dotProduct = 0;
279
- let normA = 0;
280
- let normB = 0;
281
-
282
- for (let i = 0; i < vecA.length; i++) {
283
- dotProduct += vecA[i] * vecB[i];
284
- normA += vecA[i] * vecA[i];
285
- normB += vecB[i] * vecB[i];
286
- }
287
-
288
- const denominator = Math.sqrt(normA) * Math.sqrt(normB);
289
- return denominator === 0 ? 0 : dotProduct / denominator;
290
- }
291
- }
292
-
293
- // ─── Main Class ──────────────────────────────────────────────────────────────
294
-
295
- export class ApiKeyManager extends EventEmitter {
296
- private keys: KeyState[] = [];
297
- private storageKey = 'api_rotation_state_v2';
298
- private storage: any;
299
- private strategy: LoadBalancingStrategy;
300
- private fallbackFn?: () => any;
301
-
302
- // Bulkhead state
303
- private maxConcurrency: number;
304
- private activeCalls: number = 0;
305
-
306
- // Health check state
307
- private healthCheckFn?: (key: string) => Promise<boolean>;
308
- private healthCheckInterval?: ReturnType<typeof setInterval>;
309
-
310
- // Debounced persistence — avoids writeFileSync on every single call
311
- private _saveTimer?: ReturnType<typeof setTimeout>;
312
- private _saveDirty: boolean = false;
313
-
314
- // Semantic Cache v4
315
- private semanticCache?: SemanticCache;
316
- private getEmbeddingFn?: (text: string) => Promise<number[]>;
317
- private _isResolvingEmbedding: boolean = false; // Recursion guard
318
-
319
- /**
320
- * Constructor supports both legacy positional args and new options object.
321
- *
322
- * @example Legacy (v1/v2 — still works):
323
- * new ApiKeyManager(['key1', 'key2'], storage, strategy)
324
- *
325
- * @example New (v3):
326
- * new ApiKeyManager(keys, { storage, strategy, fallbackFn, concurrency })
327
- */
328
- constructor(
329
- initialKeys: string[] | { key: string; weight?: number; provider?: string }[],
330
- storageOrOptions?: any | ApiKeyManagerOptions,
331
- strategy?: LoadBalancingStrategy
332
- ) {
333
- super();
334
-
335
- // Detect if second arg is options object or legacy storage
336
- let options: ApiKeyManagerOptions = {};
337
- if (storageOrOptions && typeof storageOrOptions === 'object' && ('storage' in storageOrOptions || 'strategy' in storageOrOptions || 'fallbackFn' in storageOrOptions || 'concurrency' in storageOrOptions || 'semanticCache' in storageOrOptions)) {
338
- // New v3 options object
339
- options = storageOrOptions as ApiKeyManagerOptions;
340
- } else {
341
- // Legacy positional args
342
- options = {
343
- storage: storageOrOptions,
344
- strategy: strategy,
345
- };
346
- }
347
-
348
- this.storage = options.storage || {
349
- getItem: () => null,
350
- setItem: () => { },
351
- };
352
- this.strategy = options.strategy || new StandardStrategy();
353
- this.fallbackFn = options.fallbackFn;
354
- this.maxConcurrency = options.concurrency || Infinity;
355
-
356
- // Init Semantic Cache if provided
357
- if (options.semanticCache) {
358
- this.semanticCache = new SemanticCache(
359
- options.semanticCache.threshold,
360
- options.semanticCache.ttlMs
361
- );
362
- this.getEmbeddingFn = options.semanticCache.getEmbedding;
363
- }
364
-
365
- // Normalize input to objects
366
- let inputKeys: { key: string; weight?: number; provider?: string }[] = [];
367
- if (initialKeys.length > 0 && typeof initialKeys[0] === 'string') {
368
- inputKeys = (initialKeys as string[]).flatMap(k => k.split(',').map(s => ({ key: s.trim(), weight: 1.0, provider: 'default' })));
369
- } else {
370
- inputKeys = initialKeys as { key: string; weight?: number; provider?: string }[];
371
- }
372
-
373
- // Deduplicate
374
- const uniqueMap = new Map<string, { weight: number; provider: string }>();
375
- inputKeys.forEach(k => {
376
- if (k.key.length > 0) uniqueMap.set(k.key, { weight: k.weight ?? 1.0, provider: k.provider ?? 'default' });
377
- });
378
-
379
- if (uniqueMap.size < inputKeys.length) {
380
- console.warn(`[ApiKeyManager] Removed ${inputKeys.length - uniqueMap.size} duplicate/empty keys.`);
381
- }
382
-
383
- this.keys = Array.from(uniqueMap.entries()).map(([key, meta]) => ({
384
- key,
385
- failCount: 0,
386
- failedAt: null,
387
- isQuotaError: false,
388
- circuitState: 'CLOSED',
389
- lastUsed: 0,
390
- successCount: 0,
391
- totalRequests: 0,
392
- halfOpenTestTime: null,
393
- customCooldown: null,
394
- weight: meta.weight,
395
- averageLatency: 0,
396
- totalLatency: 0,
397
- latencySamples: 0,
398
- provider: meta.provider,
399
- }));
400
-
401
- this.loadState();
402
- }
403
-
404
- // ─── Error Classification ────────────────────────────────────────────────
405
-
406
- /**
407
- * CLASSIFIES an error to determine handling strategy
408
- */
409
- public classifyError(error: any, finishReason?: string): ErrorClassification {
410
- const status = error?.status || error?.response?.status;
411
- const message = error?.message || error?.error?.message || String(error);
412
-
413
- // 1. Check finishReason first
414
- if (finishReason === 'SAFETY') return { type: 'SAFETY', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
415
- if (finishReason === 'RECITATION') return { type: 'RECITATION', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
416
-
417
- // 2. Check timeout
418
- if (error instanceof TimeoutError || error?.name === 'TimeoutError') {
419
- return { type: 'TIMEOUT', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
420
- }
421
-
422
- // 3. Check HTTP status codes
423
- if (status === 403 || ERROR_PATTERNS.isAuthError.test(message)) {
424
- return { type: 'AUTH', retryable: false, cooldownMs: Infinity, markKeyFailed: true, markKeyDead: true };
425
- }
426
- if (status === 429 || ERROR_PATTERNS.isQuotaError.test(message)) {
427
- const retryAfter = this.parseRetryAfter(error);
428
- return {
429
- type: 'QUOTA',
430
- retryable: true,
431
- cooldownMs: retryAfter || CONFIG.COOLDOWN_QUOTA,
432
- markKeyFailed: true,
433
- markKeyDead: false
434
- };
435
- }
436
- if (status === 400 || ERROR_PATTERNS.isBadRequest.test(message)) {
437
- return { type: 'BAD_REQUEST', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
438
- }
439
- if (ERROR_PATTERNS.isTransient.test(message) || [500, 502, 503, 504].includes(status)) {
440
- return { type: 'TRANSIENT', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
441
- }
442
-
443
- return { type: 'UNKNOWN', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
444
- }
445
-
446
- private parseRetryAfter(error: any): number | null {
447
- const retryAfter = error?.response?.headers?.['retry-after'] ||
448
- error?.headers?.['retry-after'] ||
449
- error?.retryAfter;
450
-
451
- if (!retryAfter) return null;
452
-
453
- const seconds = parseInt(retryAfter, 10);
454
- if (!isNaN(seconds)) return seconds * 1000;
455
-
456
- const date = Date.parse(retryAfter);
457
- if (!isNaN(date)) return Math.max(0, date - Date.now());
458
-
459
- return null;
460
- }
461
-
462
- // ─── Cooldown ────────────────────────────────────────────────────────────
463
-
464
- private isOnCooldown(k: KeyState): boolean {
465
- if (k.circuitState === 'DEAD') return true;
466
- const now = Date.now();
467
-
468
- if (k.circuitState === 'OPEN') {
469
- if (k.halfOpenTestTime && now >= k.halfOpenTestTime) {
470
- k.circuitState = 'HALF_OPEN';
471
- this.emit('circuitHalfOpen', k.key);
472
- return false;
473
- }
474
- return true;
475
- }
476
-
477
- if (k.failedAt) {
478
- if (k.customCooldown && now - k.failedAt < k.customCooldown) return true;
479
- const cooldown = k.isQuotaError ? CONFIG.COOLDOWN_QUOTA : CONFIG.COOLDOWN_TRANSIENT;
480
- if (now - k.failedAt < cooldown) return true;
481
- }
482
-
483
- return false;
484
- }
485
-
486
- // ─── Key Selection ───────────────────────────────────────────────────────
487
-
488
- public getKey(): string | null {
489
- // 1. Filter out dead and cooling down keys
490
- const candidates = this.keys.filter(k => k.circuitState !== 'DEAD' && !this.isOnCooldown(k));
491
-
492
- if (candidates.length === 0) {
493
- // FALLBACK: Return oldest failed key (excluding DEAD)
494
- const nonDead = this.keys.filter(k => k.circuitState !== 'DEAD');
495
- if (nonDead.length === 0) {
496
- this.emit('allKeysExhausted');
497
- return null;
498
- }
499
- return nonDead.sort((a, b) => (a.failedAt || 0) - (b.failedAt || 0))[0]?.key || null;
500
- }
501
-
502
- // 2. Delegate to Strategy
503
- const selected = this.strategy.next(candidates);
504
-
505
- if (selected) {
506
- selected.lastUsed = Date.now();
507
- this.saveState();
508
- return selected.key;
509
- }
510
- return null;
511
- }
512
-
513
- /**
514
- * Get a key filtered by provider tag
515
- */
516
- public getKeyByProvider(provider: string): string | null {
517
- const candidates = this.keys.filter(k =>
518
- k.provider === provider && k.circuitState !== 'DEAD' && !this.isOnCooldown(k)
519
- );
520
-
521
- if (candidates.length === 0) return null;
522
-
523
- const selected = this.strategy.next(candidates);
524
- if (selected) {
525
- selected.lastUsed = Date.now();
526
- this.saveState();
527
- return selected.key;
528
- }
529
- return null;
530
- }
531
-
532
- public getKeyCount(): number {
533
- return this.keys.filter(k => k.circuitState !== 'DEAD').length;
534
- }
535
-
536
- // ─── Mark Success / Failed ───────────────────────────────────────────────
537
-
538
- /**
539
- * Mark success AND update latency stats
540
- * @param durationMs Duration of the request in milliseconds
541
- */
542
- public markSuccess(key: string, durationMs?: number) {
543
- const k = this.keys.find(x => x.key === key);
544
- if (!k) return;
545
-
546
- const wasRecovering = k.circuitState !== 'CLOSED' && k.circuitState !== 'DEAD';
547
- if (wasRecovering) {
548
- console.log(`[Key Recovered] ...${key.slice(-4)}`);
549
- this.emit('keyRecovered', key);
550
- }
551
-
552
- k.circuitState = 'CLOSED';
553
- k.failCount = 0;
554
- k.failedAt = null;
555
- k.isQuotaError = false;
556
- k.customCooldown = null;
557
- k.successCount++;
558
- k.totalRequests++;
559
-
560
- if (durationMs !== undefined) {
561
- k.totalLatency += durationMs;
562
- k.latencySamples++;
563
- k.averageLatency = k.totalLatency / k.latencySamples;
564
- }
565
-
566
- this.saveState();
567
- }
568
-
569
- public markFailed(key: string, classification: ErrorClassification) {
570
- const k = this.keys.find(x => x.key === key);
571
- if (!k || k.circuitState === 'DEAD') return;
572
- if (!classification.markKeyFailed) return;
573
-
574
- k.failedAt = Date.now();
575
- k.failCount++;
576
- k.totalRequests++;
577
- k.isQuotaError = classification.type === 'QUOTA';
578
- k.customCooldown = classification.cooldownMs || null;
579
-
580
- if (classification.markKeyDead) {
581
- k.circuitState = 'DEAD';
582
- console.error(`[Key DEAD] ...${key.slice(-4)} - Permanently removed`);
583
- this.emit('keyDead', key);
584
- } else {
585
- // State Transitions
586
- if (k.circuitState === 'HALF_OPEN') {
587
- k.circuitState = 'OPEN';
588
- k.halfOpenTestTime = Date.now() + CONFIG.HALF_OPEN_TEST_DELAY;
589
- this.emit('circuitOpen', key);
590
- } else if (k.failCount >= CONFIG.MAX_CONSECUTIVE_FAILURES || classification.type === 'QUOTA') {
591
- k.circuitState = 'OPEN';
592
- k.halfOpenTestTime = Date.now() + (classification.cooldownMs || CONFIG.HALF_OPEN_TEST_DELAY);
593
- this.emit('circuitOpen', key);
594
- }
595
- }
596
- this.saveState();
597
- }
598
-
599
- public markFailedLegacy(key: string, isQuota: boolean = false) {
600
- this.markFailed(key, {
601
- type: isQuota ? 'QUOTA' : 'TRANSIENT',
602
- retryable: true,
603
- cooldownMs: isQuota ? CONFIG.COOLDOWN_QUOTA : CONFIG.COOLDOWN_TRANSIENT,
604
- markKeyFailed: true,
605
- markKeyDead: false,
606
- });
607
- }
608
-
609
- // ─── Backoff ─────────────────────────────────────────────────────────────
610
-
611
- public calculateBackoff(attempt: number): number {
612
- const exponential = CONFIG.BASE_BACKOFF * Math.pow(2, attempt);
613
- const capped = Math.min(exponential, CONFIG.MAX_BACKOFF);
614
- const jitter = Math.random() * 1000;
615
- return capped + jitter;
616
- }
617
-
618
- // ─── Stats ───────────────────────────────────────────────────────────────
619
-
620
- public getStats(): ApiKeyManagerStats {
621
- const total = this.keys.length;
622
- const dead = this.keys.filter(k => k.circuitState === 'DEAD').length;
623
- const cooling = this.keys.filter(k => k.circuitState === 'OPEN' || k.circuitState === 'HALF_OPEN').length;
624
- const healthy = total - dead - cooling;
625
- return { total, healthy, cooling, dead };
626
- }
627
-
628
- public _getKeys(): KeyState[] { return this.keys; }
629
-
630
- // ─── execute() Wrapper ───────────────────────────────────────────────────
631
-
632
- /**
633
- * Wraps the entire API call lifecycle into a single method.
634
- *
635
- * @example
636
- * const result = await manager.execute(
637
- * (key) => fetch(`https://api.example.com?key=${key}`),
638
- * { maxRetries: 3, timeoutMs: 5000 }
639
- * );
640
- */
641
- public async execute<T>(
642
- fn: (key: string, signal?: AbortSignal) => Promise<T>,
643
- options?: ExecuteOptions & { prompt?: string }
644
- ): Promise<T> {
645
- const maxRetries = options?.maxRetries ?? 0;
646
- const timeoutMs = options?.timeoutMs;
647
- const finishReason = options?.finishReason;
648
- const prompt = options?.prompt;
649
- const provider = options?.provider;
650
-
651
- // 1. Semantic Cache Check (Mastermind Edition)
652
- // Guard: skip cache if we're already resolving an embedding to prevent
653
- // infinite recursion when getEmbeddingFn calls execute() internally.
654
- let currentPromptVector: number[] | null = null;
655
- if (this.semanticCache && this.getEmbeddingFn && prompt && !this._isResolvingEmbedding) {
656
- try {
657
- this._isResolvingEmbedding = true;
658
- currentPromptVector = await this.getEmbeddingFn(prompt);
659
- const cachedResponse = this.semanticCache.get(currentPromptVector);
660
- if (cachedResponse !== null) {
661
- console.log(`[Semantic Cache HIT] for prompt: "${prompt.slice(0, 30)}..."`);
662
- this.emit('executeSuccess', 'CACHE_HIT', 0);
663
- return cachedResponse as T;
664
- }
665
- } catch (e) {
666
- console.warn('[Semantic Cache Check Failed] Proceeding to live API', e);
667
- } finally {
668
- this._isResolvingEmbedding = false;
669
- }
670
- }
671
-
672
- // 2. Bulkhead check
673
- if (this.activeCalls >= this.maxConcurrency) {
674
- this.emit('bulkheadRejected');
675
- throw new BulkheadRejectionError();
676
- }
677
-
678
- this.activeCalls++;
679
- try {
680
- const result = await this._executeWithRetry(fn, maxRetries, timeoutMs, finishReason, provider);
681
-
682
- // 3. Store in Semantic Cache on success
683
- if (this.semanticCache && prompt && currentPromptVector) {
684
- this.semanticCache.set(prompt, currentPromptVector, result);
685
- }
686
-
687
- return result;
688
- } finally {
689
- this.activeCalls--;
690
- }
691
- }
692
-
693
- /**
694
- * Executes a streaming function (AsyncGenerator) with retry logic and semantic caching.
695
- *
696
- * @example
697
- * const stream = await manager.executeStream(async (key) => {
698
- * return await gemini.generateContentStream({ prompt: "..." });
699
- * }, { prompt: "..." });
700
- *
701
- * for await (const chunk of stream) {
702
- * console.log(chunk.text());
703
- * }
704
- */
705
- public async *executeStream<T>(
706
- fn: (key: string, signal?: AbortSignal) => AsyncGenerator<T, any, unknown>,
707
- options?: ExecuteOptions & { prompt?: string }
708
- ): AsyncGenerator<T, any, unknown> {
709
- const maxRetries = options?.maxRetries ?? 0;
710
- const timeoutMs = options?.timeoutMs;
711
- const finishReason = options?.finishReason;
712
- const prompt = options?.prompt;
713
- const provider = options?.provider;
714
-
715
- // 1. Semantic Cache Check
716
- let currentPromptVector: number[] | null = null;
717
- if (this.semanticCache && this.getEmbeddingFn && prompt && !this._isResolvingEmbedding) {
718
- try {
719
- this._isResolvingEmbedding = true;
720
- currentPromptVector = await this.getEmbeddingFn(prompt);
721
- const cachedResponse = this.semanticCache.get(currentPromptVector);
722
- if (cachedResponse !== null) {
723
- console.log(`[Semantic Cache HIT] Streaming cached response for prompt: "${prompt.slice(0, 30)}..."`);
724
- this.emit('executeSuccess', 'CACHE_HIT_STREAM', 0);
725
- // Replay full response as a single chunk (or iterate if it's an array)
726
- if (Array.isArray(cachedResponse)) {
727
- for (const chunk of cachedResponse) yield chunk as T;
728
- } else {
729
- yield cachedResponse as T;
730
- }
731
- return;
732
- }
733
- } catch (e) {
734
- console.warn('[Semantic Cache Check Failed] Proceeding to live stream', e);
735
- } finally {
736
- this._isResolvingEmbedding = false;
737
- }
738
- }
739
-
740
- // 2. Bulkhead check
741
- if (this.activeCalls >= this.maxConcurrency) {
742
- this.emit('bulkheadRejected');
743
- throw new BulkheadRejectionError();
744
- }
745
-
746
- this.activeCalls++;
747
- const accumulatedChunks: T[] = [];
748
- let lastError: any;
749
-
750
- try {
751
- for (let attempt = 0; attempt <= maxRetries; attempt++) {
752
- const key = provider ? this.getKeyByProvider(provider) : this.getKey();
753
-
754
- if (!key) {
755
- if (this.fallbackFn) {
756
- this.emit('fallback', 'all keys exhausted (stream)');
757
- yield this.fallbackFn();
758
- return;
759
- }
760
- throw new AllKeysExhaustedError();
761
- }
762
-
763
- let iterator: AsyncGenerator<T, any, unknown>;
764
- try {
765
- // Start the generator
766
- iterator = fn(key);
767
-
768
- // PRIME THE ITERATOR:
769
- // Try to get the first chunk. This forces the generator body to
770
- // execute until the first 'yield' or until it throws.
771
- const firstResult = await iterator.next();
772
-
773
- // If we got here, the INITIAL connection/logic succeeded!
774
- if (firstResult.done) return;
775
-
776
- // Yield first chunk
777
- yield firstResult.value;
778
- if (this.semanticCache && prompt) accumulatedChunks.push(firstResult.value);
779
-
780
- // Yield rest of chunks
781
- for await (const chunk of iterator) {
782
- yield chunk;
783
- if (this.semanticCache && prompt) accumulatedChunks.push(chunk);
784
- }
785
-
786
- // Success! Store in cache
787
- if (this.semanticCache && prompt && currentPromptVector && accumulatedChunks.length > 0) {
788
- this.semanticCache.set(prompt, currentPromptVector, accumulatedChunks);
789
- }
790
-
791
- return; // Full success, exit retry loop
792
-
793
- } catch (error: any) {
794
- lastError = error;
795
- const classification = this.classifyError(error, finishReason);
796
-
797
- this.markFailed(key, classification);
798
- this.emit('executeFailed', key, error);
799
-
800
- // Note: If we already yielded the FIRST chunk, we CANNOT retry the connection
801
- // because the user has already received data. Mid-stream failures propagate.
802
- if (accumulatedChunks.length > 0) {
803
- throw error;
804
- }
805
-
806
- if (!classification.retryable || attempt >= maxRetries) {
807
- if (this.fallbackFn && attempt >= maxRetries) {
808
- this.emit('fallback', 'max retries exceeded (stream)');
809
- yield this.fallbackFn();
810
- return;
811
- }
812
- throw error;
813
- }
814
-
815
- const delay = this.calculateBackoff(attempt);
816
- this.emit('retry', key, attempt + 1, delay);
817
- await this._sleep(delay);
818
- }
819
- }
820
- } finally {
821
- this.activeCalls--;
822
- }
823
-
824
- throw lastError;
825
- }
826
-
827
- private async _executeWithRetry<T>(
828
- fn: (key: string, signal?: AbortSignal) => Promise<T>,
829
- maxRetries: number,
830
- timeoutMs?: number,
831
- finishReason?: string,
832
- provider?: string
833
- ): Promise<T> {
834
- let lastError: any;
835
-
836
- for (let attempt = 0; attempt <= maxRetries; attempt++) {
837
- const key = provider ? this.getKeyByProvider(provider) : this.getKey();
838
-
839
- if (!key) {
840
- // All keys exhausted — try fallback
841
- if (this.fallbackFn) {
842
- this.emit('fallback', 'all keys exhausted');
843
- return this.fallbackFn();
844
- }
845
- throw new AllKeysExhaustedError();
846
- }
847
-
848
- try {
849
- const start = Date.now();
850
- let result: T;
851
-
852
- if (timeoutMs) {
853
- result = await this._executeWithTimeout(fn, key, timeoutMs);
854
- } else {
855
- result = await fn(key);
856
- }
857
-
858
- const duration = Date.now() - start;
859
- this.markSuccess(key, duration);
860
- this.emit('executeSuccess', key, duration);
861
- return result;
862
-
863
- } catch (error: any) {
864
- lastError = error;
865
- const classification = this.classifyError(error, finishReason);
866
-
867
- this.markFailed(key, classification);
868
- this.emit('executeFailed', key, error);
869
-
870
- if (!classification.retryable || attempt >= maxRetries) {
871
- // Non-retryable or out of retries
872
- if (this.fallbackFn && attempt >= maxRetries) {
873
- this.emit('fallback', 'max retries exceeded');
874
- return this.fallbackFn();
875
- }
876
- throw error;
877
- }
878
-
879
- // Retry with backoff
880
- const delay = this.calculateBackoff(attempt);
881
- this.emit('retry', key, attempt + 1, delay);
882
- await this._sleep(delay);
883
- }
884
- }
885
-
886
- // Should not reach here, but safety net
887
- throw lastError;
888
- }
889
-
890
- private async _executeWithTimeout<T>(
891
- fn: (key: string, signal?: AbortSignal) => Promise<T>,
892
- key: string,
893
- timeoutMs: number
894
- ): Promise<T> {
895
- const controller = new AbortController();
896
- const timer = setTimeout(() => controller.abort(), timeoutMs);
897
-
898
- try {
899
- const result = await Promise.race([
900
- fn(key, controller.signal),
901
- new Promise<never>((_, reject) => {
902
- controller.signal.addEventListener('abort', () => {
903
- reject(new TimeoutError(timeoutMs));
904
- });
905
- })
906
- ]);
907
- return result;
908
- } finally {
909
- clearTimeout(timer);
910
- }
911
- }
912
-
913
- private _sleep(ms: number): Promise<void> {
914
- return new Promise(resolve => {
915
- const timer = setTimeout(resolve, ms);
916
- if (timer.unref) timer.unref();
917
- });
918
- }
919
-
920
- // ─── Health Checks ───────────────────────────────────────────────────────
921
-
922
- /**
923
- * Set a health check function that tests if a key is operational
924
- */
925
- public setHealthCheck(fn: (key: string) => Promise<boolean>) {
926
- this.healthCheckFn = fn;
927
- }
928
-
929
- /**
930
- * Start periodic health checks
931
- * @param intervalMs How often to run health checks (default: 60s)
932
- */
933
- public startHealthChecks(intervalMs: number = 60_000) {
934
- this.stopHealthChecks(); // Clear any existing interval
935
- this.healthCheckInterval = setInterval(() => this._runHealthChecks(), intervalMs);
936
- if (this.healthCheckInterval.unref) this.healthCheckInterval.unref();
937
- }
938
-
939
- /**
940
- * Stop periodic health checks
941
- */
942
- public stopHealthChecks() {
943
- if (this.healthCheckInterval) {
944
- clearInterval(this.healthCheckInterval);
945
- this.healthCheckInterval = undefined;
946
- }
947
- // Flush any pending state to disk on shutdown
948
- this._flushState();
949
- }
950
-
951
- private async _runHealthChecks() {
952
- if (!this.healthCheckFn) return;
953
-
954
- // Check non-DEAD keys that are in OPEN or HALF_OPEN state
955
- const keysToCheck = this.keys.filter(k =>
956
- k.circuitState === 'OPEN' || k.circuitState === 'HALF_OPEN'
957
- );
958
-
959
- for (const k of keysToCheck) {
960
- try {
961
- const healthy = await this.healthCheckFn(k.key);
962
- if (healthy) {
963
- this.markSuccess(k.key);
964
- this.emit('healthCheckPassed', k.key);
965
- } else {
966
- this.emit('healthCheckFailed', k.key, new Error('Health check returned false'));
967
- }
968
- } catch (error) {
969
- this.emit('healthCheckFailed', k.key, error);
970
- }
971
- }
972
- }
973
-
974
- // ─── Persistence ─────────────────────────────────────────────────────────
975
-
976
- /**
977
- * Debounced save — marks state as dirty and flushes after 500ms of inactivity.
978
- * Under heavy load (multiple getKey/markSuccess/markFailed calls), this coalesces
979
- * dozens of writeFileSync calls into one.
980
- */
981
- private saveState() {
982
- if (!this.storage) return;
983
- this._saveDirty = true;
984
-
985
- if (!this._saveTimer) {
986
- this._saveTimer = setTimeout(() => {
987
- this._flushState();
988
- this._saveTimer = undefined;
989
- }, 500);
990
- if (this._saveTimer.unref) this._saveTimer.unref();
991
- }
992
- }
993
-
994
- /**
995
- * Immediately flush state to storage. Called by the debounce timer
996
- * and by stopHealthChecks() to ensure clean shutdown.
997
- */
998
- public flushState() {
999
- this._flushState();
1000
- }
1001
-
1002
- private _flushState() {
1003
- if (!this._saveDirty || !this.storage) return;
1004
- this._saveDirty = false;
1005
-
1006
- if (this._saveTimer) {
1007
- clearTimeout(this._saveTimer);
1008
- this._saveTimer = undefined;
1009
- }
1010
-
1011
- const state = this.keys.reduce((acc, k) => ({
1012
- ...acc,
1013
- [k.key]: {
1014
- failCount: k.failCount,
1015
- failedAt: k.failedAt,
1016
- isQuotaError: k.isQuotaError,
1017
- circuitState: k.circuitState,
1018
- lastUsed: k.lastUsed,
1019
- successCount: k.successCount,
1020
- totalRequests: k.totalRequests,
1021
- customCooldown: k.customCooldown,
1022
- weight: k.weight,
1023
- averageLatency: k.averageLatency,
1024
- totalLatency: k.totalLatency,
1025
- latencySamples: k.latencySamples,
1026
- provider: k.provider
1027
- }
1028
- }), {});
1029
- this.storage.setItem(this.storageKey, JSON.stringify(state));
1030
- }
1031
-
1032
- private loadState() {
1033
- if (!this.storage) return;
1034
- try {
1035
- const raw = this.storage.getItem(this.storageKey);
1036
- if (!raw) return;
1037
- const data = JSON.parse(raw);
1038
- const now = Date.now();
1039
-
1040
- this.keys.forEach(k => {
1041
- if (!data[k.key]) return;
1042
- Object.assign(k, data[k.key]);
1043
-
1044
- // Resurrect DEAD keys that have exceeded the TTL.
1045
- // This allows keys that were marked dead (e.g., temporary 403)
1046
- // to be retested after a cooldown period instead of staying dead forever.
1047
- if (k.circuitState === 'DEAD' && k.failedAt) {
1048
- const deadDuration = now - k.failedAt;
1049
- if (deadDuration >= CONFIG.DEAD_KEY_TTL) {
1050
- k.circuitState = 'HALF_OPEN';
1051
- k.halfOpenTestTime = null; // Allow immediate test
1052
- k.failCount = 0;
1053
- }
1054
- }
1055
-
1056
- // Clear stale cooldowns from previous sessions.
1057
- // If a key was cooling down and the process restarted after the
1058
- // cooldown would have expired, reset it to CLOSED.
1059
- if (k.circuitState === 'OPEN' && k.failedAt) {
1060
- const cooldown = k.customCooldown || (k.isQuotaError ? CONFIG.COOLDOWN_QUOTA : CONFIG.COOLDOWN_TRANSIENT);
1061
- if (now - k.failedAt >= cooldown) {
1062
- k.circuitState = 'HALF_OPEN';
1063
- k.halfOpenTestTime = null;
1064
- }
1065
- }
1066
- });
1067
- } catch (e) {
1068
- console.error("Failed to load key state");
1069
- }
1070
- }
1071
- }
1
+ /**
2
+ * Universal ApiKeyManager v5.0 — Ecosystem 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
6
+ * NEW in v5.0: Provider Presets (GeminiManager, OpenAIManager, MultiManager),
7
+ * Built-in Persistence (FileStorage, MemoryStorage)
8
+ * Gemini-Specific: finishReason handling, Safety blocks, RECITATION detection
9
+ * Infrastructure: cockatiel (Bulkhead queueing + ExponentialBackoff w/ decorrelated jitter)
10
+ */
11
+
12
+ import { EventEmitter } from 'events';
13
+ import {
14
+ bulkhead as cockatielBulkhead,
15
+ BulkheadRejectedError as CockatielBulkheadRejectedError,
16
+ ExponentialBackoff,
17
+ decorrelatedJitterGenerator,
18
+ type BulkheadPolicy,
19
+ type IBackoff,
20
+ } from 'cockatiel';
21
+ import { z } from 'zod';
22
+
23
+ // ─── Re-exports: Persistence ─────────────────────────────────────────────────
24
+ // Persistence adapters can be imported from root or via subpath
25
+ export { FileStorage } from './persistence/file';
26
+ export type { FileStorageOptions } from './persistence/file';
27
+ export { MemoryStorage } from './persistence/memory';
28
+
29
+ // ─── Presets ─────────────────────────────────────────────────────────────────
30
+ // Presets are available via subpath imports to avoid circular dependencies:
31
+ // import { GeminiManager } from '@splashcodex/api-key-manager/presets/gemini';
32
+ // import { OpenAIManager } from '@splashcodex/api-key-manager/presets/openai';
33
+ // import { MultiManager } from '@splashcodex/api-key-manager/presets/multi';
34
+ // Or import all at once:
35
+ // import { GeminiManager, OpenAIManager, MultiManager } from '@splashcodex/api-key-manager/presets';
36
+
37
+
38
+
39
+ // ─── Interfaces & Types ──────────────────────────────────────────────────────
40
+
41
+ export interface KeyState {
42
+ key: string;
43
+ failCount: number; // Consecutive failures
44
+ failedAt: number | null; // Timestamp of last failure
45
+ isQuotaError: boolean; // Was last error a 429?
46
+ circuitState: 'CLOSED' | 'OPEN' | 'HALF_OPEN' | 'DEAD';
47
+ lastUsed: number;
48
+ successCount: number;
49
+ totalRequests: number;
50
+ halfOpenTestTime: number | null;
51
+ customCooldown: number | null; // From Retry-After header
52
+ // v2.0 Stats
53
+ weight: number; // 0.0 - 1.0 (Default 1.0)
54
+ averageLatency: number; // Rolling average latency in ms
55
+ totalLatency: number; // Sum of all latency checks (for calculating average)
56
+ latencySamples: number; // Number of samples
57
+ // v3.0 Fields
58
+ provider: string; // Provider tag (e.g. 'openai', 'gemini')
59
+ }
60
+
61
+ export type ErrorType =
62
+ | 'QUOTA' // 429 - Rotate key, respect cooldown
63
+ | 'TRANSIENT' // 500/503/504 - Retry with backoff
64
+ | 'AUTH' // 403 - Key is dead, remove from pool
65
+ | 'BAD_REQUEST' // 400 - Do not retry, fix request
66
+ | 'SAFETY' // finishReason: SAFETY - Not a key issue
67
+ | 'RECITATION' // finishReason: RECITATION - Not a key issue
68
+ | 'TIMEOUT' // Request timed out
69
+ | 'UNKNOWN'; // Catch-all
70
+
71
+ export interface ErrorClassification {
72
+ type: ErrorType;
73
+ retryable: boolean;
74
+ cooldownMs: number;
75
+ markKeyFailed: boolean;
76
+ markKeyDead: boolean;
77
+ }
78
+
79
+ export interface ApiKeyManagerStats {
80
+ total: number;
81
+ healthy: number;
82
+ cooling: number;
83
+ dead: number;
84
+ }
85
+
86
+ export interface ExecuteOptions {
87
+ timeoutMs?: number; // Timeout per attempt in ms
88
+ maxRetries?: number; // Max retry attempts (default: 0 = no retry)
89
+ finishReason?: string; // For Gemini finishReason handling
90
+ provider?: string; // Filter keys by provider (e.g. 'openai')
91
+ }
92
+
93
+ export const ApiKeyManagerOptionsSchema = z.object({
94
+ storage: z.any().optional(),
95
+ strategy: z.any().optional(), // Expected: LoadBalancingStrategy
96
+ fallbackFn: z.custom<() => any>((val) => typeof val === 'function', 'Must be a function').optional(),
97
+ /** Max concurrent execute() calls. When limit is reached, excess requests queue (up to concurrencyQueueSize) then reject. */
98
+ concurrency: z.number().int().positive().optional(),
99
+ /**
100
+ * Maximum number of requests to hold in the bulkhead queue when all concurrency slots are busy.
101
+ * - Default: `0` — excess requests are rejected immediately (preserves v3 behavior).
102
+ * - Set to a positive number to allow requests to wait for a free slot.
103
+ *
104
+ * @example
105
+ * // Queue up to 10 waiting requests before rejecting
106
+ * new ApiKeyManager(keys, { concurrency: 5, concurrencyQueueSize: 10 })
107
+ */
108
+ concurrencyQueueSize: z.number().int().nonnegative().optional(),
109
+ semanticCache: z.object({
110
+ threshold: z.number().min(0).max(1).optional(), // Similarity threshold (0.0 - 1.0, default 0.95)
111
+ ttlMs: z.number().int().positive().optional(), // Cache TTL
112
+ getEmbedding: z.custom<(text: string) => Promise<number[]>>((val) => typeof val === 'function', 'Must be a function'),
113
+ }).optional(),
114
+ }).passthrough();
115
+
116
+ export type ApiKeyManagerOptions = z.infer<typeof ApiKeyManagerOptionsSchema>;
117
+
118
+ export interface CacheEntry {
119
+ vector: number[];
120
+ prompt: string;
121
+ response: any;
122
+ timestamp: number;
123
+ }
124
+
125
+ // ─── Event Types ─────────────────────────────────────────────────────────────
126
+
127
+ export interface ApiKeyManagerEventMap {
128
+ keyDead: (key: string) => void;
129
+ circuitOpen: (key: string) => void;
130
+ circuitHalfOpen: (key: string) => void;
131
+ keyRecovered: (key: string) => void;
132
+ fallback: (reason: string) => void;
133
+ allKeysExhausted: () => void;
134
+ retry: (key: string, attempt: number, delayMs: number) => void;
135
+ healthCheckFailed: (key: string, error: any) => void;
136
+ healthCheckPassed: (key: string) => void;
137
+ executeSuccess: (key: string, durationMs: number) => void;
138
+ executeFailed: (key: string, error: any) => void;
139
+ bulkheadRejected: () => void;
140
+ }
141
+
142
+ // ─── Config ──────────────────────────────────────────────────────────────────
143
+
144
+ const CONFIG = {
145
+ MAX_CONSECUTIVE_FAILURES: 5,
146
+ COOLDOWN_TRANSIENT: 60 * 1000, // 1 minute
147
+ COOLDOWN_QUOTA: 5 * 60 * 1000, // 5 minutes (default if no Retry-After)
148
+ COOLDOWN_QUOTA_DAILY: 60 * 60 * 1000, // 1 hour for RPD exhaustion
149
+ HALF_OPEN_TEST_DELAY: 60 * 1000, // 1 minute after open
150
+ MAX_BACKOFF: 64 * 1000, // 64 seconds max
151
+ BASE_BACKOFF: 1000, // 1 second base
152
+ DEAD_KEY_TTL: 60 * 60 * 1000, // 1 hourDEAD keys get retested after this
153
+ };
154
+
155
+ // Error classification patterns
156
+ const ERROR_PATTERNS = {
157
+ isQuotaError: /429|quota|exhausted|resource.?exhausted|too.?many.?requests|rate.?limit/i,
158
+ isAuthError: /403|permission.?denied|invalid.?api.?key|unauthorized|unauthenticated/i,
159
+ isSafetyBlock: /safety|blocked|recitation|harmful/i,
160
+ isTransient: /500|502|503|504|internal|unavailable|deadline|timeout|overloaded/i,
161
+ isBadRequest: /400|invalid.?argument|failed.?precondition|malformed|not.?found|404/i,
162
+ };
163
+
164
+ // ─── Custom Errors ───────────────────────────────────────────────────────────
165
+
166
+ export class TimeoutError extends Error {
167
+ constructor(ms: number) {
168
+ super(`Request timed out after ${ms}ms`);
169
+ this.name = 'TimeoutError';
170
+ }
171
+ }
172
+
173
+ export class BulkheadRejectionError extends Error {
174
+ constructor() {
175
+ super('Bulkhead capacity exceeded — too many concurrent requests');
176
+ this.name = 'BulkheadRejectionError';
177
+ }
178
+ }
179
+
180
+ export class AllKeysExhaustedError extends Error {
181
+ constructor() {
182
+ super('All API keys exhausted — no healthy keys available');
183
+ this.name = 'AllKeysExhaustedError';
184
+ }
185
+ }
186
+
187
+ // ─── Strategies ──────────────────────────────────────────────────────────────
188
+
189
+ /**
190
+ * Strategy Interface for selecting the next key
191
+ */
192
+ export interface LoadBalancingStrategy {
193
+ next(candidates: KeyState[]): KeyState | null;
194
+ }
195
+
196
+ /**
197
+ * Standard Strategy: Least Failed > Least Recently Used
198
+ */
199
+ export class StandardStrategy implements LoadBalancingStrategy {
200
+ next(candidates: KeyState[]): KeyState | null {
201
+ candidates.sort((a, b) => {
202
+ if (a.failCount !== b.failCount) return a.failCount - b.failCount;
203
+ return a.lastUsed - b.lastUsed;
204
+ });
205
+ return candidates[0] || null;
206
+ }
207
+ }
208
+
209
+ /**
210
+ * Weighted Strategy: Probabilistic selection based on weight
211
+ * Higher weight = Higher chance of selection
212
+ */
213
+ export class WeightedStrategy implements LoadBalancingStrategy {
214
+ next(candidates: KeyState[]): KeyState | null {
215
+ if (candidates.length === 0) return null;
216
+
217
+ const totalWeight = candidates.reduce((sum, k) => sum + k.weight, 0);
218
+ let random = Math.random() * totalWeight;
219
+
220
+ for (const key of candidates) {
221
+ random -= key.weight;
222
+ if (random <= 0) return key;
223
+ }
224
+
225
+ return candidates[0]; // Fallback
226
+ }
227
+ }
228
+
229
+ /**
230
+ * Latency Strategy: Pick lowest average latency with LRU tie-break
231
+ */
232
+ export class LatencyStrategy implements LoadBalancingStrategy {
233
+ next(candidates: KeyState[]): KeyState | null {
234
+ if (candidates.length === 0) return null;
235
+ candidates.sort((a, b) => {
236
+ if (a.averageLatency !== b.averageLatency) return a.averageLatency - b.averageLatency;
237
+ return a.lastUsed - b.lastUsed; // LRU tie-break
238
+ });
239
+ return candidates[0];
240
+ }
241
+ }
242
+
243
+ // ─── Semantic Engine ─────────────────────────────────────────────────────────
244
+
245
+ /**
246
+ * High-performance Vanilla Semantic Cache
247
+ * Implements Cosine Similarity math from scratch.
248
+ */
249
+ export class SemanticCache {
250
+ private entries: CacheEntry[] = [];
251
+ private threshold: number;
252
+ private ttlMs: number;
253
+
254
+ constructor(threshold: number = 0.95, ttlMs: number = 24 * 60 * 60 * 1000) {
255
+ this.threshold = threshold;
256
+ this.ttlMs = ttlMs;
257
+ }
258
+
259
+ public set(prompt: string, vector: number[], response: any) {
260
+ // Expire old entry for same prompt if exists
261
+ this.entries = this.entries.filter(e => e.prompt !== prompt);
262
+ this.entries.push({
263
+ prompt,
264
+ vector,
265
+ response,
266
+ timestamp: Date.now()
267
+ });
268
+ // Optional: Cap size to prevent memory leaks
269
+ if (this.entries.length > 500) this.entries.shift();
270
+ }
271
+
272
+ public get(vector: number[]): any | null {
273
+ const now = Date.now();
274
+ let bestMatch: CacheEntry | null = null;
275
+ let highestSimilarity = -1;
276
+
277
+ for (let i = this.entries.length - 1; i >= 0; i--) {
278
+ const entry = this.entries[i];
279
+
280
+ // Check TTL
281
+ if (now - entry.timestamp > this.ttlMs) {
282
+ this.entries.splice(i, 1);
283
+ continue;
284
+ }
285
+
286
+ const similarity = this.calculateCosineSimilarity(vector, entry.vector);
287
+ if (similarity >= this.threshold && similarity > highestSimilarity) {
288
+ highestSimilarity = similarity;
289
+ bestMatch = entry;
290
+ }
291
+ }
292
+
293
+ return bestMatch ? bestMatch.response : null;
294
+ }
295
+
296
+ /**
297
+ * Vanilla Cosine Similarity: (A·B) / (||A|| * ||B||)
298
+ */
299
+ private calculateCosineSimilarity(vecA: number[], vecB: number[]): number {
300
+ if (vecA.length !== vecB.length) return 0;
301
+ let dotProduct = 0;
302
+ let normA = 0;
303
+ let normB = 0;
304
+
305
+ for (let i = 0; i < vecA.length; i++) {
306
+ dotProduct += vecA[i] * vecB[i];
307
+ normA += vecA[i] * vecA[i];
308
+ normB += vecB[i] * vecB[i];
309
+ }
310
+
311
+ const denominator = Math.sqrt(normA) * Math.sqrt(normB);
312
+ return denominator === 0 ? 0 : dotProduct / denominator;
313
+ }
314
+ }
315
+
316
+ // ─── Main Class ──────────────────────────────────────────────────────────────
317
+
318
+ export class ApiKeyManager extends EventEmitter {
319
+ private keys: KeyState[] = [];
320
+ private storageKey = 'api_rotation_state_v2';
321
+ private storage: any;
322
+ private strategy: LoadBalancingStrategy;
323
+ private fallbackFn?: () => any;
324
+
325
+ // Bulkhead state — managed by cockatiel BulkheadPolicy
326
+ // Provides FIFO queueing (requests wait for a slot) instead of immediate rejection
327
+ private bulkheadPolicy: BulkheadPolicy | null = null;
328
+
329
+ // Health check state
330
+ private healthCheckFn?: (key: string) => Promise<boolean>;
331
+ private healthCheckInterval?: ReturnType<typeof setInterval>;
332
+
333
+ // Debounced persistence — avoids writeFileSync on every single call
334
+ private _saveTimer?: ReturnType<typeof setTimeout>;
335
+ private _saveDirty: boolean = false;
336
+
337
+ // Semantic Cache v4
338
+ private semanticCache?: SemanticCache;
339
+ private getEmbeddingFn?: (text: string) => Promise<number[]>;
340
+ private _isResolvingEmbedding: boolean = false; // Recursion guard
341
+
342
+ /**
343
+ * Constructor supports both legacy positional args and new options object.
344
+ *
345
+ * @example Legacy (v1/v2 — still works):
346
+ * new ApiKeyManager(['key1', 'key2'], storage, strategy)
347
+ *
348
+ * @example New (v3):
349
+ * new ApiKeyManager(keys, { storage, strategy, fallbackFn, concurrency })
350
+ */
351
+ constructor(
352
+ initialKeys: string[] | { key: string; weight?: number; provider?: string }[],
353
+ storageOrOptions?: any | ApiKeyManagerOptions,
354
+ strategy?: LoadBalancingStrategy
355
+ ) {
356
+ super();
357
+
358
+ // Detect if second arg is options object or legacy storage
359
+ let options: ApiKeyManagerOptions = {};
360
+ if (storageOrOptions && typeof storageOrOptions === 'object' && ('storage' in storageOrOptions || 'strategy' in storageOrOptions || 'fallbackFn' in storageOrOptions || 'concurrency' in storageOrOptions || 'concurrencyQueueSize' in storageOrOptions || 'semanticCache' in storageOrOptions)) {
361
+ // New v3 options object
362
+ options = storageOrOptions as ApiKeyManagerOptions;
363
+ } else {
364
+ // Legacy positional args
365
+ options = {
366
+ storage: storageOrOptions,
367
+ strategy: strategy,
368
+ };
369
+ }
370
+
371
+ // Validate options with Zod (will throw meaningful errors on invalid configs)
372
+ options = ApiKeyManagerOptionsSchema.parse(options);
373
+
374
+ this.storage = options.storage || {
375
+ getItem: () => null,
376
+ setItem: () => { },
377
+ };
378
+ this.strategy = options.strategy || new StandardStrategy();
379
+ this.fallbackFn = options.fallbackFn;
380
+
381
+ // Build cockatiel bulkhead.
382
+ // queueSize defaults to 0 (immediate rejection — preserves existing API contract).
383
+ // Set concurrencyQueueSize > 0 to opt-in to queuing instead of rejection.
384
+ const maxConcurrency = options.concurrency ?? Infinity;
385
+ const queueSize = options.concurrencyQueueSize ?? 0;
386
+ if (maxConcurrency !== Infinity) {
387
+ this.bulkheadPolicy = cockatielBulkhead(maxConcurrency, queueSize);
388
+ this.bulkheadPolicy.onReject(() => {
389
+ this.emit('bulkheadRejected');
390
+ });
391
+ }
392
+
393
+ // Init Semantic Cache if provided
394
+ if (options.semanticCache) {
395
+ this.semanticCache = new SemanticCache(
396
+ options.semanticCache.threshold,
397
+ options.semanticCache.ttlMs
398
+ );
399
+ this.getEmbeddingFn = options.semanticCache.getEmbedding;
400
+ }
401
+
402
+ // Normalize input to objects
403
+ let inputKeys: { key: string; weight?: number; provider?: string }[] = [];
404
+ if (initialKeys.length > 0 && typeof initialKeys[0] === 'string') {
405
+ inputKeys = (initialKeys as string[]).flatMap(k => k.split(',').map(s => ({ key: s.trim(), weight: 1.0, provider: 'default' })));
406
+ } else {
407
+ inputKeys = initialKeys as { key: string; weight?: number; provider?: string }[];
408
+ }
409
+
410
+ // Deduplicate
411
+ const uniqueMap = new Map<string, { weight: number; provider: string }>();
412
+ inputKeys.forEach(k => {
413
+ if (k.key.length > 0) uniqueMap.set(k.key, { weight: k.weight ?? 1.0, provider: k.provider ?? 'default' });
414
+ });
415
+
416
+ if (uniqueMap.size < inputKeys.length) {
417
+ console.warn(`[ApiKeyManager] Removed ${inputKeys.length - uniqueMap.size} duplicate/empty keys.`);
418
+ }
419
+
420
+ this.keys = Array.from(uniqueMap.entries()).map(([key, meta]) => ({
421
+ key,
422
+ failCount: 0,
423
+ failedAt: null,
424
+ isQuotaError: false,
425
+ circuitState: 'CLOSED',
426
+ lastUsed: 0,
427
+ successCount: 0,
428
+ totalRequests: 0,
429
+ halfOpenTestTime: null,
430
+ customCooldown: null,
431
+ weight: meta.weight,
432
+ averageLatency: 0,
433
+ totalLatency: 0,
434
+ latencySamples: 0,
435
+ provider: meta.provider,
436
+ }));
437
+
438
+ this.loadState();
439
+ }
440
+
441
+ // ─── Error Classification ────────────────────────────────────────────────
442
+
443
+ /**
444
+ * CLASSIFIES an error to determine handling strategy
445
+ */
446
+ public classifyError(error: any, finishReason?: string): ErrorClassification {
447
+ const status = error?.status || error?.response?.status;
448
+ const message = error?.message || error?.error?.message || String(error);
449
+
450
+ // 1. Check finishReason first
451
+ if (finishReason === 'SAFETY') return { type: 'SAFETY', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
452
+ if (finishReason === 'RECITATION') return { type: 'RECITATION', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
453
+
454
+ // 2. Check timeout
455
+ if (error instanceof TimeoutError || error?.name === 'TimeoutError') {
456
+ return { type: 'TIMEOUT', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
457
+ }
458
+
459
+ // 3. Check HTTP status codes
460
+ if (status === 403 || ERROR_PATTERNS.isAuthError.test(message)) {
461
+ return { type: 'AUTH', retryable: false, cooldownMs: Infinity, markKeyFailed: true, markKeyDead: true };
462
+ }
463
+ if (status === 429 || ERROR_PATTERNS.isQuotaError.test(message)) {
464
+ const retryAfter = this.parseRetryAfter(error);
465
+ return {
466
+ type: 'QUOTA',
467
+ retryable: true,
468
+ cooldownMs: retryAfter || CONFIG.COOLDOWN_QUOTA,
469
+ markKeyFailed: true,
470
+ markKeyDead: false
471
+ };
472
+ }
473
+ if (status === 400 || ERROR_PATTERNS.isBadRequest.test(message)) {
474
+ return { type: 'BAD_REQUEST', retryable: false, cooldownMs: 0, markKeyFailed: false, markKeyDead: false };
475
+ }
476
+ if (ERROR_PATTERNS.isTransient.test(message) || [500, 502, 503, 504].includes(status)) {
477
+ return { type: 'TRANSIENT', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
478
+ }
479
+
480
+ return { type: 'UNKNOWN', retryable: true, cooldownMs: CONFIG.COOLDOWN_TRANSIENT, markKeyFailed: true, markKeyDead: false };
481
+ }
482
+
483
+ private parseRetryAfter(error: any): number | null {
484
+ const retryAfter = error?.response?.headers?.['retry-after'] ||
485
+ error?.headers?.['retry-after'] ||
486
+ error?.retryAfter;
487
+
488
+ if (!retryAfter) return null;
489
+
490
+ const seconds = parseInt(retryAfter, 10);
491
+ if (!isNaN(seconds)) return seconds * 1000;
492
+
493
+ const date = Date.parse(retryAfter);
494
+ if (!isNaN(date)) return Math.max(0, date - Date.now());
495
+
496
+ return null;
497
+ }
498
+
499
+ // ─── Cooldown ────────────────────────────────────────────────────────────
500
+
501
+ private isOnCooldown(k: KeyState): boolean {
502
+ if (k.circuitState === 'DEAD') return true;
503
+ const now = Date.now();
504
+
505
+ if (k.circuitState === 'OPEN') {
506
+ if (k.halfOpenTestTime && now >= k.halfOpenTestTime) {
507
+ k.circuitState = 'HALF_OPEN';
508
+ this.emit('circuitHalfOpen', k.key);
509
+ return false;
510
+ }
511
+ return true;
512
+ }
513
+
514
+ if (k.failedAt) {
515
+ if (k.customCooldown && now - k.failedAt < k.customCooldown) return true;
516
+ const cooldown = k.isQuotaError ? CONFIG.COOLDOWN_QUOTA : CONFIG.COOLDOWN_TRANSIENT;
517
+ if (now - k.failedAt < cooldown) return true;
518
+ }
519
+
520
+ return false;
521
+ }
522
+
523
+ // ─── Key Selection ───────────────────────────────────────────────────────
524
+
525
+ public getKey(): string | null {
526
+ // 1. Filter out dead and cooling down keys
527
+ const candidates = this.keys.filter(k => k.circuitState !== 'DEAD' && !this.isOnCooldown(k));
528
+
529
+ if (candidates.length === 0) {
530
+ // FALLBACK: Return oldest failed key (excluding DEAD)
531
+ const nonDead = this.keys.filter(k => k.circuitState !== 'DEAD');
532
+ if (nonDead.length === 0) {
533
+ this.emit('allKeysExhausted');
534
+ return null;
535
+ }
536
+ return nonDead.sort((a, b) => (a.failedAt || 0) - (b.failedAt || 0))[0]?.key || null;
537
+ }
538
+
539
+ // 2. Delegate to Strategy
540
+ const selected = this.strategy.next(candidates);
541
+
542
+ if (selected) {
543
+ selected.lastUsed = Date.now();
544
+ this.saveState();
545
+ return selected.key;
546
+ }
547
+ return null;
548
+ }
549
+
550
+ /**
551
+ * Get a key filtered by provider tag
552
+ */
553
+ public getKeyByProvider(provider: string): string | null {
554
+ const candidates = this.keys.filter(k =>
555
+ k.provider === provider && k.circuitState !== 'DEAD' && !this.isOnCooldown(k)
556
+ );
557
+
558
+ if (candidates.length === 0) return null;
559
+
560
+ const selected = this.strategy.next(candidates);
561
+ if (selected) {
562
+ selected.lastUsed = Date.now();
563
+ this.saveState();
564
+ return selected.key;
565
+ }
566
+ return null;
567
+ }
568
+
569
+ public getKeyCount(): number {
570
+ return this.keys.filter(k => k.circuitState !== 'DEAD').length;
571
+ }
572
+
573
+ // ─── Mark Success / Failed ───────────────────────────────────────────────
574
+
575
+ /**
576
+ * Mark success AND update latency stats
577
+ * @param durationMs Duration of the request in milliseconds
578
+ */
579
+ public markSuccess(key: string, durationMs?: number) {
580
+ const k = this.keys.find(x => x.key === key);
581
+ if (!k) return;
582
+
583
+ const wasRecovering = k.circuitState !== 'CLOSED' && k.circuitState !== 'DEAD';
584
+ if (wasRecovering) {
585
+ console.log(`[Key Recovered] ...${key.slice(-4)}`);
586
+ this.emit('keyRecovered', key);
587
+ }
588
+
589
+ k.circuitState = 'CLOSED';
590
+ k.failCount = 0;
591
+ k.failedAt = null;
592
+ k.isQuotaError = false;
593
+ k.customCooldown = null;
594
+ k.successCount++;
595
+ k.totalRequests++;
596
+
597
+ if (durationMs !== undefined) {
598
+ k.totalLatency += durationMs;
599
+ k.latencySamples++;
600
+ k.averageLatency = k.totalLatency / k.latencySamples;
601
+ }
602
+
603
+ this.saveState();
604
+ }
605
+
606
+ public markFailed(key: string, classification: ErrorClassification) {
607
+ const k = this.keys.find(x => x.key === key);
608
+ if (!k || k.circuitState === 'DEAD') return;
609
+ if (!classification.markKeyFailed) return;
610
+
611
+ k.failedAt = Date.now();
612
+ k.failCount++;
613
+ k.totalRequests++;
614
+ k.isQuotaError = classification.type === 'QUOTA';
615
+ k.customCooldown = classification.cooldownMs || null;
616
+
617
+ if (classification.markKeyDead) {
618
+ k.circuitState = 'DEAD';
619
+ console.error(`[Key DEAD] ...${key.slice(-4)} - Permanently removed`);
620
+ this.emit('keyDead', key);
621
+ } else {
622
+ // State Transitions
623
+ if (k.circuitState === 'HALF_OPEN') {
624
+ k.circuitState = 'OPEN';
625
+ k.halfOpenTestTime = Date.now() + CONFIG.HALF_OPEN_TEST_DELAY;
626
+ this.emit('circuitOpen', key);
627
+ } else if (k.failCount >= CONFIG.MAX_CONSECUTIVE_FAILURES || classification.type === 'QUOTA') {
628
+ k.circuitState = 'OPEN';
629
+ k.halfOpenTestTime = Date.now() + (classification.cooldownMs || CONFIG.HALF_OPEN_TEST_DELAY);
630
+ this.emit('circuitOpen', key);
631
+ }
632
+ }
633
+ this.saveState();
634
+ }
635
+
636
+ public markFailedLegacy(key: string, isQuota: boolean = false) {
637
+ this.markFailed(key, {
638
+ type: isQuota ? 'QUOTA' : 'TRANSIENT',
639
+ retryable: true,
640
+ cooldownMs: isQuota ? CONFIG.COOLDOWN_QUOTA : CONFIG.COOLDOWN_TRANSIENT,
641
+ markKeyFailed: true,
642
+ markKeyDead: false,
643
+ });
644
+ }
645
+
646
+ // ─── Backoff ─────────────────────────────────────────────────────────────
647
+
648
+ /**
649
+ * Calculate exponential backoff with decorrelated jitter using cockatiel.
650
+ * Decorrelated jitter avoids thundering-herd by randomizing retry intervals
651
+ * independent of the previous delay, which is statistically superior to
652
+ * simple `random * exponential` jitter.
653
+ *
654
+ * @param attempt - Zero-indexed attempt number
655
+ */
656
+ private readonly _backoffFactory = new ExponentialBackoff({
657
+ initialDelay: CONFIG.BASE_BACKOFF,
658
+ maxDelay: CONFIG.MAX_BACKOFF,
659
+ generator: decorrelatedJitterGenerator,
660
+ });
661
+
662
+ public calculateBackoff(attempt: number): number {
663
+ // ExponentialBackoff is a linked-list: _backoffFactory.next() returns an
664
+ // IBackoff node with { duration, next() }. Walk `attempt` steps to get
665
+ // the correctly scaled delay. Uses decorrelated jitter (superior to random*exp).
666
+ // We cast via `any` to bypass the interface vs concrete class arg-count conflict.
667
+ // eslint-disable-next-line @typescript-eslint/no-explicit-any
668
+ let node: any = (this._backoffFactory as any).next();
669
+ for (let i = 0; i < attempt; i++) {
670
+ node = node.next();
671
+ }
672
+ return (node as IBackoff<unknown>).duration;
673
+ }
674
+
675
+ // ─── Stats ───────────────────────────────────────────────────────────────
676
+
677
+ public getStats(): ApiKeyManagerStats {
678
+ const total = this.keys.length;
679
+ const dead = this.keys.filter(k => k.circuitState === 'DEAD').length;
680
+ const cooling = this.keys.filter(k => k.circuitState === 'OPEN' || k.circuitState === 'HALF_OPEN').length;
681
+ const healthy = total - dead - cooling;
682
+ return { total, healthy, cooling, dead };
683
+ }
684
+
685
+ public _getKeys(): KeyState[] { return this.keys; }
686
+
687
+ // ─── execute() Wrapper ───────────────────────────────────────────────────
688
+
689
+ /**
690
+ * Wraps the entire API call lifecycle into a single method.
691
+ *
692
+ * @example
693
+ * const result = await manager.execute(
694
+ * (key) => fetch(`https://api.example.com?key=${key}`),
695
+ * { maxRetries: 3, timeoutMs: 5000 }
696
+ * );
697
+ */
698
+ public async execute<T>(
699
+ fn: (key: string, signal?: AbortSignal) => Promise<T>,
700
+ options?: ExecuteOptions & { prompt?: string }
701
+ ): Promise<T> {
702
+ const maxRetries = options?.maxRetries ?? 0;
703
+ const timeoutMs = options?.timeoutMs;
704
+ const finishReason = options?.finishReason;
705
+ const prompt = options?.prompt;
706
+ const provider = options?.provider;
707
+
708
+ // 1. Semantic Cache Check (Mastermind Edition)
709
+ // Guard: skip cache if we're already resolving an embedding to prevent
710
+ // infinite recursion when getEmbeddingFn calls execute() internally.
711
+ let currentPromptVector: number[] | null = null;
712
+ if (this.semanticCache && this.getEmbeddingFn && prompt && !this._isResolvingEmbedding) {
713
+ try {
714
+ this._isResolvingEmbedding = true;
715
+ currentPromptVector = await this.getEmbeddingFn(prompt);
716
+ const cachedResponse = this.semanticCache.get(currentPromptVector);
717
+ if (cachedResponse !== null) {
718
+ console.log(`[Semantic Cache HIT] for prompt: "${prompt.slice(0, 30)}..."`);
719
+ this.emit('executeSuccess', 'CACHE_HIT', 0);
720
+ return cachedResponse as T;
721
+ }
722
+ } catch (e) {
723
+ console.warn('[Semantic Cache Check Failed] Proceeding to live API', e);
724
+ } finally {
725
+ this._isResolvingEmbedding = false;
726
+ }
727
+ }
728
+
729
+ // 2. Bulkhead check — if policy exists, cockatiel queues excess requests
730
+ // instead of immediately rejecting them (queue drains as slots free up)
731
+ if (this.bulkheadPolicy) {
732
+ try {
733
+ const result = await this.bulkheadPolicy.execute(async () => {
734
+ return await this._executeWithSemanticAndRetry<T>(fn, maxRetries, timeoutMs, finishReason, provider, prompt, currentPromptVector);
735
+ });
736
+ return result;
737
+ } catch (err) {
738
+ if (err instanceof CockatielBulkheadRejectedError) {
739
+ throw new BulkheadRejectionError();
740
+ }
741
+ throw err;
742
+ }
743
+ }
744
+
745
+ // No concurrency limit configured — run directly
746
+ return this._executeWithSemanticAndRetry<T>(fn, maxRetries, timeoutMs, finishReason, provider, prompt, currentPromptVector);
747
+ }
748
+
749
+ /**
750
+ * Executes a streaming function (AsyncGenerator) with retry logic and semantic caching.
751
+ *
752
+ * @example
753
+ * const stream = await manager.executeStream(async (key) => {
754
+ * return await gemini.generateContentStream({ prompt: "..." });
755
+ * }, { prompt: "..." });
756
+ *
757
+ * for await (const chunk of stream) {
758
+ * console.log(chunk.text());
759
+ * }
760
+ */
761
+ public async *executeStream<T>(
762
+ fn: (key: string, signal?: AbortSignal) => AsyncGenerator<T, any, unknown>,
763
+ options?: ExecuteOptions & { prompt?: string }
764
+ ): AsyncGenerator<T, any, unknown> {
765
+ const maxRetries = options?.maxRetries ?? 0;
766
+ const timeoutMs = options?.timeoutMs;
767
+ const finishReason = options?.finishReason;
768
+ const prompt = options?.prompt;
769
+ const provider = options?.provider;
770
+
771
+ // 1. Semantic Cache Check
772
+ let currentPromptVector: number[] | null = null;
773
+ if (this.semanticCache && this.getEmbeddingFn && prompt && !this._isResolvingEmbedding) {
774
+ try {
775
+ this._isResolvingEmbedding = true;
776
+ currentPromptVector = await this.getEmbeddingFn(prompt);
777
+ const cachedResponse = this.semanticCache.get(currentPromptVector);
778
+ if (cachedResponse !== null) {
779
+ console.log(`[Semantic Cache HIT] Streaming cached response for prompt: "${prompt.slice(0, 30)}..."`);
780
+ this.emit('executeSuccess', 'CACHE_HIT_STREAM', 0);
781
+ // Replay full response as a single chunk (or iterate if it's an array)
782
+ if (Array.isArray(cachedResponse)) {
783
+ for (const chunk of cachedResponse) yield chunk as T;
784
+ } else {
785
+ yield cachedResponse as T;
786
+ }
787
+ return;
788
+ }
789
+ } catch (e) {
790
+ console.warn('[Semantic Cache Check Failed] Proceeding to live stream', e);
791
+ } finally {
792
+ this._isResolvingEmbedding = false;
793
+ }
794
+ }
795
+
796
+ // 2. Bulkhead check — gate streaming with bulkhead slot management.
797
+ // Async generators can't be wrapped in bulkheadPolicy.execute() (which expects
798
+ // a Promise), so we check availability and track manually.
799
+ const useBulkhead = this.bulkheadPolicy !== null;
800
+ if (useBulkhead) {
801
+ const slots = (this.bulkheadPolicy as any);
802
+ const hasSlot = (slots.executionSlots > 0) || (slots.queueSlots > 0);
803
+ if (!hasSlot) {
804
+ this.emit('bulkheadRejected');
805
+ throw new BulkheadRejectionError();
806
+ }
807
+ }
808
+ const accumulatedChunks: T[] = [];
809
+ let lastError: any;
810
+
811
+ try {
812
+ for (let attempt = 0; attempt <= maxRetries; attempt++) {
813
+ const key = provider ? this.getKeyByProvider(provider) : this.getKey();
814
+
815
+ if (!key) {
816
+ if (this.fallbackFn) {
817
+ this.emit('fallback', 'all keys exhausted (stream)');
818
+ yield this.fallbackFn();
819
+ return;
820
+ }
821
+ throw new AllKeysExhaustedError();
822
+ }
823
+
824
+ let iterator: AsyncGenerator<T, any, unknown>;
825
+ try {
826
+ // Start the generator
827
+ iterator = fn(key);
828
+
829
+ // PRIME THE ITERATOR:
830
+ // Try to get the first chunk. This forces the generator body to
831
+ // execute until the first 'yield' or until it throws.
832
+ const firstResult = await iterator.next();
833
+
834
+ // If we got here, the INITIAL connection/logic succeeded!
835
+ if (firstResult.done) return;
836
+
837
+ // Yield first chunk
838
+ yield firstResult.value;
839
+ if (this.semanticCache && prompt) accumulatedChunks.push(firstResult.value);
840
+
841
+ // Yield rest of chunks
842
+ for await (const chunk of iterator) {
843
+ yield chunk;
844
+ if (this.semanticCache && prompt) accumulatedChunks.push(chunk);
845
+ }
846
+
847
+ // Success! Store in cache
848
+ if (this.semanticCache && prompt && currentPromptVector && accumulatedChunks.length > 0) {
849
+ this.semanticCache.set(prompt, currentPromptVector, accumulatedChunks);
850
+ }
851
+
852
+ return; // Full success, exit retry loop
853
+
854
+ } catch (error: any) {
855
+ lastError = error;
856
+ const classification = this.classifyError(error, finishReason);
857
+
858
+ this.markFailed(key, classification);
859
+ this.emit('executeFailed', key, error);
860
+
861
+ // Note: If we already yielded the FIRST chunk, we CANNOT retry the connection
862
+ // because the user has already received data. Mid-stream failures propagate.
863
+ if (accumulatedChunks.length > 0) {
864
+ throw error;
865
+ }
866
+
867
+ if (!classification.retryable || attempt >= maxRetries) {
868
+ if (this.fallbackFn && attempt >= maxRetries) {
869
+ this.emit('fallback', 'max retries exceeded (stream)');
870
+ yield this.fallbackFn();
871
+ return;
872
+ }
873
+ throw error;
874
+ }
875
+
876
+ const delay = this.calculateBackoff(attempt);
877
+ this.emit('retry', key, attempt + 1, delay);
878
+ await this._sleep(delay);
879
+ }
880
+ }
881
+ } finally {
882
+ // Slot is freed automatically when the generator completes/throws
883
+ }
884
+
885
+ throw lastError;
886
+ }
887
+
888
+ /**
889
+ * Helper: stores result in semantic cache then returns it.
890
+ * Called by the bulkhead execute() callback and the no-bulkhead path.
891
+ */
892
+ private async _executeWithSemanticAndRetry<T>(
893
+ fn: (key: string, signal?: AbortSignal) => Promise<T>,
894
+ maxRetries: number,
895
+ timeoutMs?: number,
896
+ finishReason?: string,
897
+ provider?: string,
898
+ prompt?: string,
899
+ currentPromptVector?: number[] | null
900
+ ): Promise<T> {
901
+ const result = await this._executeWithRetry(fn, maxRetries, timeoutMs, finishReason, provider);
902
+
903
+ // Store in Semantic Cache on success
904
+ if (this.semanticCache && prompt && currentPromptVector) {
905
+ this.semanticCache.set(prompt, currentPromptVector, result);
906
+ }
907
+
908
+ return result;
909
+ }
910
+
911
+ private async _executeWithRetry<T>(
912
+ fn: (key: string, signal?: AbortSignal) => Promise<T>,
913
+ maxRetries: number,
914
+ timeoutMs?: number,
915
+ finishReason?: string,
916
+ provider?: string
917
+ ): Promise<T> {
918
+ let lastError: any;
919
+
920
+ for (let attempt = 0; attempt <= maxRetries; attempt++) {
921
+ const key = provider ? this.getKeyByProvider(provider) : this.getKey();
922
+
923
+ if (!key) {
924
+ // All keys exhausted — try fallback
925
+ if (this.fallbackFn) {
926
+ this.emit('fallback', 'all keys exhausted');
927
+ return this.fallbackFn();
928
+ }
929
+ throw new AllKeysExhaustedError();
930
+ }
931
+
932
+ try {
933
+ const start = Date.now();
934
+ let result: T;
935
+
936
+ if (timeoutMs) {
937
+ result = await this._executeWithTimeout(fn, key, timeoutMs);
938
+ } else {
939
+ result = await fn(key);
940
+ }
941
+
942
+ const duration = Date.now() - start;
943
+ this.markSuccess(key, duration);
944
+ this.emit('executeSuccess', key, duration);
945
+ return result;
946
+
947
+ } catch (error: any) {
948
+ lastError = error;
949
+ const classification = this.classifyError(error, finishReason);
950
+
951
+ this.markFailed(key, classification);
952
+ this.emit('executeFailed', key, error);
953
+
954
+ if (!classification.retryable || attempt >= maxRetries) {
955
+ // Non-retryable or out of retries
956
+ if (this.fallbackFn && attempt >= maxRetries) {
957
+ this.emit('fallback', 'max retries exceeded');
958
+ return this.fallbackFn();
959
+ }
960
+ throw error;
961
+ }
962
+
963
+ // Retry with backoff
964
+ const delay = this.calculateBackoff(attempt);
965
+ this.emit('retry', key, attempt + 1, delay);
966
+ await this._sleep(delay);
967
+ }
968
+ }
969
+
970
+ // Should not reach here, but safety net
971
+ throw lastError;
972
+ }
973
+
974
+ private async _executeWithTimeout<T>(
975
+ fn: (key: string, signal?: AbortSignal) => Promise<T>,
976
+ key: string,
977
+ timeoutMs: number
978
+ ): Promise<T> {
979
+ const controller = new AbortController();
980
+ const timer = setTimeout(() => controller.abort(), timeoutMs);
981
+
982
+ try {
983
+ const result = await Promise.race([
984
+ fn(key, controller.signal),
985
+ new Promise<never>((_, reject) => {
986
+ controller.signal.addEventListener('abort', () => {
987
+ reject(new TimeoutError(timeoutMs));
988
+ });
989
+ })
990
+ ]);
991
+ return result;
992
+ } finally {
993
+ clearTimeout(timer);
994
+ }
995
+ }
996
+
997
+ private _sleep(ms: number): Promise<void> {
998
+ return new Promise(resolve => {
999
+ const timer = setTimeout(resolve, ms);
1000
+ if (timer.unref) timer.unref();
1001
+ });
1002
+ }
1003
+
1004
+ // ─── Health Checks ───────────────────────────────────────────────────────
1005
+
1006
+ /**
1007
+ * Set a health check function that tests if a key is operational
1008
+ */
1009
+ public setHealthCheck(fn: (key: string) => Promise<boolean>) {
1010
+ this.healthCheckFn = fn;
1011
+ }
1012
+
1013
+ /**
1014
+ * Start periodic health checks
1015
+ * @param intervalMs How often to run health checks (default: 60s)
1016
+ */
1017
+ public startHealthChecks(intervalMs: number = 60_000) {
1018
+ this.stopHealthChecks(); // Clear any existing interval
1019
+ this.healthCheckInterval = setInterval(() => this._runHealthChecks(), intervalMs);
1020
+ if (this.healthCheckInterval.unref) this.healthCheckInterval.unref();
1021
+ }
1022
+
1023
+ /**
1024
+ * Stop periodic health checks
1025
+ */
1026
+ public stopHealthChecks() {
1027
+ if (this.healthCheckInterval) {
1028
+ clearInterval(this.healthCheckInterval);
1029
+ this.healthCheckInterval = undefined;
1030
+ }
1031
+ // Flush any pending state to disk on shutdown
1032
+ this._flushState();
1033
+ }
1034
+
1035
+ private async _runHealthChecks() {
1036
+ if (!this.healthCheckFn) return;
1037
+
1038
+ // Check non-DEAD keys that are in OPEN or HALF_OPEN state
1039
+ const keysToCheck = this.keys.filter(k =>
1040
+ k.circuitState === 'OPEN' || k.circuitState === 'HALF_OPEN'
1041
+ );
1042
+
1043
+ for (const k of keysToCheck) {
1044
+ try {
1045
+ const healthy = await this.healthCheckFn(k.key);
1046
+ if (healthy) {
1047
+ this.markSuccess(k.key);
1048
+ this.emit('healthCheckPassed', k.key);
1049
+ } else {
1050
+ this.emit('healthCheckFailed', k.key, new Error('Health check returned false'));
1051
+ }
1052
+ } catch (error) {
1053
+ this.emit('healthCheckFailed', k.key, error);
1054
+ }
1055
+ }
1056
+ }
1057
+
1058
+ // ─── Persistence ─────────────────────────────────────────────────────────
1059
+
1060
+ /**
1061
+ * Debounced save marks state as dirty and flushes after 500ms of inactivity.
1062
+ * Under heavy load (multiple getKey/markSuccess/markFailed calls), this coalesces
1063
+ * dozens of writeFileSync calls into one.
1064
+ */
1065
+ private saveState() {
1066
+ if (!this.storage) return;
1067
+ this._saveDirty = true;
1068
+
1069
+ if (!this._saveTimer) {
1070
+ this._saveTimer = setTimeout(() => {
1071
+ this._flushState();
1072
+ this._saveTimer = undefined;
1073
+ }, 500);
1074
+ if (this._saveTimer.unref) this._saveTimer.unref();
1075
+ }
1076
+ }
1077
+
1078
+ /**
1079
+ * Immediately flush state to storage. Called by the debounce timer
1080
+ * and by stopHealthChecks() to ensure clean shutdown.
1081
+ */
1082
+ public flushState() {
1083
+ this._flushState();
1084
+ }
1085
+
1086
+ private _flushState() {
1087
+ if (!this._saveDirty || !this.storage) return;
1088
+ this._saveDirty = false;
1089
+
1090
+ if (this._saveTimer) {
1091
+ clearTimeout(this._saveTimer);
1092
+ this._saveTimer = undefined;
1093
+ }
1094
+
1095
+ const state = this.keys.reduce((acc, k) => ({
1096
+ ...acc,
1097
+ [k.key]: {
1098
+ failCount: k.failCount,
1099
+ failedAt: k.failedAt,
1100
+ isQuotaError: k.isQuotaError,
1101
+ circuitState: k.circuitState,
1102
+ lastUsed: k.lastUsed,
1103
+ successCount: k.successCount,
1104
+ totalRequests: k.totalRequests,
1105
+ customCooldown: k.customCooldown,
1106
+ weight: k.weight,
1107
+ averageLatency: k.averageLatency,
1108
+ totalLatency: k.totalLatency,
1109
+ latencySamples: k.latencySamples,
1110
+ provider: k.provider
1111
+ }
1112
+ }), {});
1113
+ this.storage.setItem(this.storageKey, JSON.stringify(state));
1114
+ }
1115
+
1116
+ private loadState() {
1117
+ if (!this.storage) return;
1118
+ try {
1119
+ const raw = this.storage.getItem(this.storageKey);
1120
+ if (!raw) return;
1121
+ const data = JSON.parse(raw);
1122
+ const now = Date.now();
1123
+
1124
+ this.keys.forEach(k => {
1125
+ if (!data[k.key]) return;
1126
+ Object.assign(k, data[k.key]);
1127
+
1128
+ // Resurrect DEAD keys that have exceeded the TTL.
1129
+ // This allows keys that were marked dead (e.g., temporary 403)
1130
+ // to be retested after a cooldown period instead of staying dead forever.
1131
+ if (k.circuitState === 'DEAD' && k.failedAt) {
1132
+ const deadDuration = now - k.failedAt;
1133
+ if (deadDuration >= CONFIG.DEAD_KEY_TTL) {
1134
+ k.circuitState = 'HALF_OPEN';
1135
+ k.halfOpenTestTime = null; // Allow immediate test
1136
+ k.failCount = 0;
1137
+ }
1138
+ }
1139
+
1140
+ // Clear stale cooldowns from previous sessions.
1141
+ // If a key was cooling down and the process restarted after the
1142
+ // cooldown would have expired, reset it to CLOSED.
1143
+ if (k.circuitState === 'OPEN' && k.failedAt) {
1144
+ const cooldown = k.customCooldown || (k.isQuotaError ? CONFIG.COOLDOWN_QUOTA : CONFIG.COOLDOWN_TRANSIENT);
1145
+ if (now - k.failedAt >= cooldown) {
1146
+ k.circuitState = 'HALF_OPEN';
1147
+ k.halfOpenTestTime = null;
1148
+ }
1149
+ }
1150
+ });
1151
+ } catch (e) {
1152
+ console.error("Failed to load key state");
1153
+ }
1154
+ }
1155
+ }