buroventures-harald-code-core
Version:
Harald Code Core - Core functionality for AI-powered coding assistant
232 lines • 8.28 kB
JavaScript
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { getErrorMessage } from '../utils/errors.js';
export class ApiKeyRotationManager {
settings;
onSettingsUpdate;
fallbackApiKey;
constructor(config, fallbackApiKey) {
this.settings = {
apiKeys: [],
currentKeyIndex: 0,
autoRotateOnRateLimit: true,
resetRotationDaily: true,
dailyUsageTracking: {},
...config.settings,
};
this.onSettingsUpdate = config.onSettingsUpdate;
this.fallbackApiKey = fallbackApiKey;
// Reset daily tracking if needed
this.resetDailyTrackingIfNeeded();
}
/**
* Get the current active API key
*/
getCurrentApiKey() {
const keys = this.getAvailableKeys();
if (keys.length === 0) {
return this.fallbackApiKey;
}
const currentIndex = this.settings.currentKeyIndex || 0;
if (currentIndex >= keys.length) {
// Reset to first key if index is out of bounds
this.settings.currentKeyIndex = 0;
this.saveSettings();
return keys[0];
}
return keys[currentIndex];
}
/**
* Get all available API keys (from settings + fallback)
*/
getAvailableKeys() {
const keys = this.settings.apiKeys || [];
// Always include the fallback key (from environment) as the first key if it exists
if (this.fallbackApiKey && !keys.includes(this.fallbackApiKey)) {
return [this.fallbackApiKey, ...keys];
}
return keys;
}
/**
* Rotate to the next available API key
*/
async rotateToNextKey() {
const keys = this.getAvailableKeys();
if (keys.length <= 1) {
console.warn('Cannot rotate: only one API key available');
return this.getCurrentApiKey();
}
const currentIndex = this.settings.currentKeyIndex || 0;
const nextIndex = (currentIndex + 1) % keys.length;
this.settings.currentKeyIndex = nextIndex;
await this.saveSettings();
const newKey = keys[nextIndex];
console.log(`Rotated to API key ${nextIndex + 1} of ${keys.length}`);
return newKey;
}
/**
* Handle rate limit error and potentially rotate keys
*/
async handleRateLimit(error) {
if (!this.settings.autoRotateOnRateLimit) {
return this.getCurrentApiKey();
}
console.log('Rate limit detected, attempting to rotate API key...');
// Track the failed request for current key
await this.trackUsage(this.getCurrentApiKey() || '', true);
// Rotate to next key
return await this.rotateToNextKey();
}
/**
* Track API usage for a key
*/
async trackUsage(apiKey, isRateLimit = false) {
if (!this.settings.dailyUsageTracking) {
this.settings.dailyUsageTracking = {};
}
const today = new Date().toISOString().split('T')[0];
const keyHash = this.hashApiKey(apiKey);
if (!this.settings.dailyUsageTracking[keyHash] ||
this.settings.dailyUsageTracking[keyHash].date !== today) {
this.settings.dailyUsageTracking[keyHash] = {
date: today,
requests: 0
};
}
this.settings.dailyUsageTracking[keyHash].requests += 1;
if (isRateLimit) {
console.log(`Rate limit hit for key ending in ...${apiKey.slice(-4)} (${this.settings.dailyUsageTracking[keyHash].requests} requests today)`);
}
await this.saveSettings();
}
/**
* Add a new API key to the rotation
*/
async addApiKey(apiKey) {
if (!apiKey || typeof apiKey !== 'string') {
throw new Error('Invalid API key provided');
}
if (!this.settings.apiKeys) {
this.settings.apiKeys = [];
}
if (this.settings.apiKeys.includes(apiKey)) {
console.log('API key already exists in rotation');
return;
}
this.settings.apiKeys.push(apiKey);
await this.saveSettings();
console.log(`Added new API key to rotation (total: ${this.settings.apiKeys.length})`);
}
/**
* Remove an API key from rotation
*/
async removeApiKey(apiKey) {
if (!this.settings.apiKeys) {
return;
}
const index = this.settings.apiKeys.indexOf(apiKey);
if (index === -1) {
console.log('API key not found in rotation');
return;
}
this.settings.apiKeys.splice(index, 1);
// Adjust current index if needed
if (this.settings.currentKeyIndex && this.settings.currentKeyIndex >= index) {
this.settings.currentKeyIndex = Math.max(0, this.settings.currentKeyIndex - 1);
}
await this.saveSettings();
console.log(`Removed API key from rotation (remaining: ${this.settings.apiKeys.length})`);
}
/**
* Get rotation status information
*/
getRotationStatus() {
const keys = this.getAvailableKeys();
const currentKey = this.getCurrentApiKey() || '';
return {
totalKeys: keys.length,
currentKeyIndex: this.settings.currentKeyIndex || 0,
currentKeyPreview: currentKey ? `...${currentKey.slice(-4)}` : 'none',
autoRotateEnabled: this.settings.autoRotateOnRateLimit || false,
dailyUsage: this.settings.dailyUsageTracking || {}
};
}
/**
* Reset daily usage tracking if it's a new day
*/
resetDailyTrackingIfNeeded() {
if (!this.settings.resetRotationDaily || !this.settings.dailyUsageTracking) {
return;
}
const today = new Date().toISOString().split('T')[0];
const hasOldData = Object.values(this.settings.dailyUsageTracking).some(usage => usage.date !== today);
if (hasOldData) {
console.log('Resetting daily usage tracking for new day');
this.settings.dailyUsageTracking = {};
this.settings.currentKeyIndex = 0; // Reset to first key each day
this.saveSettings();
}
}
/**
* Create a hash of the API key for tracking (privacy-safe)
*/
hashApiKey(apiKey) {
// Simple hash for privacy - just use last 8 characters
return apiKey.slice(-8);
}
/**
* Save settings changes
*/
async saveSettings() {
if (this.onSettingsUpdate) {
try {
await this.onSettingsUpdate(this.settings);
}
catch (error) {
console.error('Failed to save rotation settings:', getErrorMessage(error));
}
}
}
/**
* Update rotation settings
*/
async updateSettings(newSettings) {
this.settings = { ...this.settings, ...newSettings };
await this.saveSettings();
}
/**
* Check if rate limiting error should trigger rotation
*/
static isRateLimitError(error) {
if (!error)
return false;
const errorMessage = error instanceof Error
? error.message.toLowerCase()
: String(error).toLowerCase();
// Check for various rate limit indicators
return (errorMessage.includes('rate limit') ||
errorMessage.includes('quota exceeded') ||
errorMessage.includes('too many requests') ||
errorMessage.includes('429') ||
// Cerebras specific messages
errorMessage.includes('daily limit') ||
errorMessage.includes('usage limit'));
}
/**
* Validate API key format (basic validation)
*/
static validateApiKey(apiKey) {
if (!apiKey || typeof apiKey !== 'string') {
return false;
}
// Basic validation - should be at least 20 characters and start with expected prefix
return apiKey.length >= 20 && (apiKey.startsWith('csk-') || // Cerebras keys
apiKey.startsWith('sk-') || // OpenAI keys
apiKey.startsWith('api-') // Generic API keys
);
}
}
//# sourceMappingURL=apiKeyRotationManager.js.map