@dyingc/brave-search-mcp-server 2.0.69 → 2.1.0

package/README.md

@@ -129,6 +129,48 @@ The server supports the following environment variables:
 - `BRAVE_MCP_RETRY_BASE_DELAY`: Base delay for retry in milliseconds (default: 1000, range: 100-60000)
 - `BRAVE_MCP_RETRY_MAX_DELAY`: Maximum delay for retry in milliseconds (default: 30000, range: 1000-300000)
 
+### Multiple API Keys
+
+When using multiple API keys, the server automatically load balances requests across keys based on remaining monthly quota. Keys with more available quota are prioritized.
+
+**Environment Variables for Multiple Keys:**
+
+1. `BRAVE_API_KEYS`: Comma-separated list of API keys
+   ```bash
+   export BRAVE_API_KEYS="key1,key2,key3"
+   ```
+
+2. Numbered variables (legacy):
+   ```bash
+   export BRAVE_API_KEY_1="key1"
+   export BRAVE_API_KEY_2="key2"
+   export BRAVE_API_KEY_3="key3"
+   ```
+
+**State Persistence:**
+
+The server persists API key quota state to `~/.brave-mcp/state.json` to maintain consistent load balancing across server restarts:
+
+- State is saved after each API request
+- On startup, state is loaded if the timestamp is within the current month
+- State automatically resets at the beginning of each month (matches API quota cycle)
+- If the state file is missing or corrupted, the server starts with default quota (2000 requests per key)
+
+**State File Format:**
+```json
+{
+  "version": 1,
+  "keys": {
+    "BSKa1234": {
+      "remaining": 1500,
+      "timestamp": "2025-02-03T10:00:00Z"
+    }
+  }
+}
+```
+
+**Note:** For security, only the first 8 characters of each API key are stored in the state file.
+
 ### Command Line Options
 
 ```bash

package/dist/config.js

@@ -37,6 +37,18 @@ export const configSchema = z.object({
         .default(false)
         .describe('Whether the server should be stateless')
         .optional(),
+    apiKeySelectionStrategy: z
+        .enum(['greedy-max', 'greedy-min', 'soft-max'])
+        .default('soft-max')
+        .describe('API key selection strategy')
+        .optional(),
+    apiKeyTemperature: z
+        .number()
+        .min(0)
+        .max(5.0)
+        .default(1.0)
+        .describe('Temperature for soft-max strategy (0-5.0, 0 = hard limit)')
+        .optional(),
 });
 const state = {
     transport: 'stdio',
@@ -51,6 +63,8 @@ const state = {
     retryMaxAttempts: 5,
     retryBaseDelay: 1000,
     retryMaxDelay: 30000,
+    apiKeySelectionStrategy: 'soft-max',
+    apiKeyTemperature: 1.0,
 };
 export function isToolPermittedByUser(toolName) {
     return state.enabledTools.length > 0
@@ -101,6 +115,8 @@ export function getOptions() {
         .option('--retry-max-attempts <number>', 'Maximum retry attempts for rate-limited requests', process.env.BRAVE_MCP_MAX_RETRIES ?? '5')
         .option('--retry-base-delay <number>', 'Base delay for retry in milliseconds', process.env.BRAVE_MCP_RETRY_BASE_DELAY ?? '1000')
         .option('--retry-max-delay <number>', 'Maximum delay for retry in milliseconds', process.env.BRAVE_MCP_RETRY_MAX_DELAY ?? '30000')
+        .option('--api-key-selection-strategy <strategy>', 'API key selection strategy (greedy-max, greedy-min, soft-max)', process.env.BRAVE_API_KEY_SELECTION_STRATEGY ?? 'soft-max')
+        .option('--api-key-temperature <number>', 'Temperature for soft-max strategy (0-5.0, default: 1.0, 0 = hard limit)', process.env.BRAVE_API_KEY_TEMPERATURE ?? '1.0')
         .allowUnknownOption()
         .parse(process.argv);
     const options = program.opts();
@@ -150,6 +166,20 @@ export function getOptions() {
         return false;
     }
     options.retryMaxDelay = retryMaxDelay;
+    // Validate API key selection strategy
+    const validStrategies = ['greedy-max', 'greedy-min', 'soft-max'];
+    if (!validStrategies.includes(options.apiKeySelectionStrategy)) {
+        console.error(`Invalid --api-key-selection-strategy value: '${options.apiKeySelectionStrategy}'. Must be one of: ${validStrategies.join(', ')}`);
+        return false;
+    }
+    // Validate temperature (range 0-5.0)
+    const temperature = parseFloat(options.apiKeyTemperature);
+    if (isNaN(temperature) || temperature < 0 || temperature > 5.0) {
+        console.error(`Invalid --api-key-temperature value: '${options.apiKeyTemperature}'. Must be a number between 0 and 5.0`);
+        return false;
+    }
+    options.apiKeyTemperature = temperature;
+    options.apiKeySelectionStrategy = options.apiKeySelectionStrategy;
     if (options.transport === 'http') {
         if (options.port < 1 || options.port > 65535) {
             console.error(`Invalid --port value: '${options.port}'. Must be a valid port number between 1 and 65535.`);
@@ -174,6 +204,8 @@ export function getOptions() {
     state.retryMaxAttempts = options.retryMaxAttempts;
     state.retryBaseDelay = options.retryBaseDelay;
     state.retryMaxDelay = options.retryMaxDelay;
+    state.apiKeySelectionStrategy = options.apiKeySelectionStrategy;
+    state.apiKeyTemperature = options.apiKeyTemperature;
     state.ready = true;
     return options;
 }

package/dist/index.js

@@ -8,8 +8,11 @@ async function main() {
         console.error('Invalid configuration');
         process.exit(1);
     }
-    // Initialize API Key Manager with parsed keys
-    initApiKeyManager(options.braveApiKeys);
+    // Initialize API Key Manager with parsed keys and selection strategy
+    initApiKeyManager(options.braveApiKeys, {
+        selectionStrategy: options.apiKeySelectionStrategy,
+        temperature: options.apiKeyTemperature,
+    });
     // default to stdio server unless http is explicitly requested
     if (options.transport === 'http') {
         httpServer.start();

package/dist/utils/apiKeyManager.js

@@ -2,17 +2,30 @@
  * API Key Manager for multi-key support with quota tracking
  * Designed for free tier keys (2000 requests/month per key)
  */
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import * as os from 'node:os';
 // Free tier constants
 const FREE_TIER_MONTHLY_QUOTA = 2000;
 const DEFAULT_PER_SECOND_LIMIT = 1;
+// State file path
+const STATE_FILE = path.join(os.homedir(), '.brave-mcp', 'state.json');
 class ApiKeyManager {
     keys;
-    constructor(apiKeys, initialQuota = FREE_TIER_MONTHLY_QUOTA) {
+    selectionStrategy;
+    temperature;
+    constructor(apiKeys, options = {}) {
         this.keys = new Map();
+        const savedStates = this.loadState();
+        const { initialQuota = FREE_TIER_MONTHLY_QUOTA, selectionStrategy = 'soft-max', temperature = 1.0, } = options;
+        this.selectionStrategy = selectionStrategy;
+        this.temperature = temperature;
         for (const key of apiKeys) {
+            const keyHash = key.slice(0, 8);
+            const savedRemaining = savedStates.get(keyHash);
             this.keys.set(key, {
                 key,
-                remainingQuota: initialQuota,
+                remainingQuota: savedRemaining ?? initialQuota,
                 monthLimit: FREE_TIER_MONTHLY_QUOTA,
                 remainingSecond: DEFAULT_PER_SECOND_LIMIT,
                 secondLimit: DEFAULT_PER_SECOND_LIMIT,
@@ -21,6 +34,59 @@ class ApiKeyManager {
             });
         }
     }
+    /**
+     * Load saved state from file, filtering for current month only
+     * Returns Map of keyHash -> remaining quota
+     */
+    loadState() {
+        try {
+            if (!fs.existsSync(STATE_FILE)) {
+                return new Map();
+            }
+            const data = JSON.parse(fs.readFileSync(STATE_FILE, 'utf-8'));
+            const now = new Date();
+            const result = new Map();
+            for (const [keyHash, info] of Object.entries(data.keys)) {
+                const timestamp = new Date(info.timestamp);
+                // Only use state from current month
+                if (timestamp.getMonth() === now.getMonth() &&
+                    timestamp.getFullYear() === now.getFullYear()) {
+                    result.set(keyHash, info.remaining);
+                }
+            }
+            return result;
+        }
+        catch (err) {
+            console.warn(`[ApiKeyManager] Failed to load state: ${err.message}`);
+            return new Map();
+        }
+    }
+    /**
+     * Save current state to file
+     */
+    saveState() {
+        try {
+            const dir = path.dirname(STATE_FILE);
+            if (!fs.existsSync(dir)) {
+                fs.mkdirSync(dir, { recursive: true });
+            }
+            const data = {
+                version: 1,
+                keys: {},
+            };
+            for (const [key, state] of this.keys.entries()) {
+                const keyHash = key.slice(0, 8);
+                data.keys[keyHash] = {
+                    remaining: state.remainingQuota,
+                    timestamp: new Date().toISOString(),
+                };
+            }
+            fs.writeFileSync(STATE_FILE, JSON.stringify(data, null, 2));
+        }
+        catch (err) {
+            console.warn(`[ApiKeyManager] Failed to save state: ${err.message}`);
+        }
+    }
     /**
      * Select the best available key based on remaining quota
      * Returns null if no valid keys available
@@ -28,20 +94,104 @@ class ApiKeyManager {
     selectBestKey() {
         // Filter valid keys with remaining quota and not in use
         const availableKeys = Array.from(this.keys.values()).filter((k) => k.isValid && k.remainingQuota > 0 && !k.inUse);
+        // No available keys
         if (availableKeys.length === 0) {
             return null;
         }
-        // Sort by remaining quota (descending)
-        availableKeys.sort((a, b) => b.remainingQuota - a.remainingQuota);
-        // Get keys with max quota
-        const maxQuota = availableKeys[0].remainingQuota;
-        const topKeys = availableKeys.filter((k) => k.remainingQuota === maxQuota);
-        // Random selection if tied
-        const selected = topKeys.length > 1 ? topKeys[Math.floor(Math.random() * topKeys.length)] : topKeys[0];
-        // Mark as in-use
+        // Single available key - return directly
+        if (availableKeys.length === 1) {
+            availableKeys[0].inUse = true;
+            return availableKeys[0].key;
+        }
+        let selected;
+        switch (this.selectionStrategy) {
+            case 'greedy-min':
+                selected = this.selectGreedyMin(availableKeys);
+                break;
+            case 'soft-max':
+                selected = this.selectKeyWithSoftMax(availableKeys, this.temperature);
+                break;
+            case 'greedy-max':
+            default:
+                selected = this.selectGreedyMax(availableKeys);
+                break;
+        }
+        // selected should never be null here since we checked availableKeys.length >= 2
+        if (!selected) {
+            return null;
+        }
         selected.inUse = true;
         return selected.key;
     }
+    /**
+     * Select key with maximum remaining quota (greedy-max)
+     */
+    selectGreedyMax(keys) {
+        keys.sort((a, b) => b.remainingQuota - a.remainingQuota);
+        const maxQuota = keys[0].remainingQuota;
+        const topKeys = keys.filter((k) => k.remainingQuota === maxQuota);
+        return topKeys.length > 1 ? topKeys[Math.floor(Math.random() * topKeys.length)] : topKeys[0];
+    }
+    /**
+     * Select key with minimum remaining quota (greedy-min)
+     */
+    selectGreedyMin(keys) {
+        keys.sort((a, b) => a.remainingQuota - b.remainingQuota);
+        const minQuota = keys[0].remainingQuota;
+        const bottomKeys = keys.filter((k) => k.remainingQuota === minQuota);
+        return bottomKeys.length > 1
+            ? bottomKeys[Math.floor(Math.random() * bottomKeys.length)]
+            : bottomKeys[0];
+    }
+    /**
+     * Calculate soft max probabilities for key selection
+     * Uses total normalization + standard soft max formula
+     */
+    calculateSoftMaxProbabilities(keys, temperature) {
+        // Calculate total quota for normalization
+        const totalQuota = keys.reduce((sum, k) => sum + k.remainingQuota, 0);
+        // Normalize by total (sum of normalized values = 1), then invert (lower quota = higher score)
+        // normalized_quota = quota / total_quota ∈ [0, 1]
+        // score = -normalized_quota ∈ [-1, 0], lower quota gives score closer to 0
+        const scores = keys.map((k) => -k.remainingQuota / totalQuota);
+        // Standard soft max: scale by temperature
+        const scaledScores = scores.map((s) => s / temperature);
+        // Numerical stability: subtract max value
+        const maxScore = Math.max(...scaledScores);
+        const expScores = scaledScores.map((s) => Math.exp(s - maxScore));
+        const sumExp = expScores.reduce((sum, exp) => sum + exp, 0);
+        const probabilities = expScores.map((exp) => exp / sumExp);
+        return keys.map((key, index) => ({
+            key,
+            probability: probabilities[index],
+        }));
+    }
+    /**
+     * Select key using soft max probability distribution
+     * temperature = 0: hard limit (equivalent to greedy-min)
+     */
+    selectKeyWithSoftMax(keys, temperature) {
+        // Empty array
+        if (keys.length === 0) {
+            return null;
+        }
+        // temperature = 0: hard limit, equivalent to greedy-min
+        if (temperature === 0) {
+            return this.selectGreedyMin(keys);
+        }
+        const keyProbabilities = this.calculateSoftMaxProbabilities(keys, temperature);
+        // Weighted random sampling (roulette wheel selection)
+        const random = Math.random();
+        let cumulativeProbability = 0;
+        for (const { key, probability } of keyProbabilities) {
+            cumulativeProbability += probability;
+            if (random <= cumulativeProbability) {
+                return key;
+            }
+        }
+        // Fallback to last key
+        return keyProbabilities[keyProbabilities.length - 1].key;
+    }
     /**
      * Release a key after request completion
      */
@@ -78,7 +228,8 @@ class ApiKeyManager {
         state.monthLimit = limit.perMonth;
         state.secondLimit = limit.perSecond;
         state.inUse = false;
-        console.debug(`[ApiKeyManager] Key ${apiKey.slice(0, 8)}... quota: ${state.remainingQuota}/${state.monthLimit}`);
+        // Save state after update
+        this.saveState();
     }
     /**
      * Mark a key as invalid (e.g., 401 Unauthorized)
@@ -116,8 +267,8 @@ class ApiKeyManager {
 }
 // Singleton instance
 let apiKeyManagerInstance = null;
-export function initApiKeyManager(apiKeys) {
-    apiKeyManagerInstance = new ApiKeyManager(apiKeys);
+export function initApiKeyManager(apiKeys, options) {
+    apiKeyManagerInstance = new ApiKeyManager(apiKeys, options);
 }
 export function getApiKeyManager() {
     return apiKeyManagerInstance;

package/package.json

@@ -2,7 +2,7 @@
   "name": "@dyingc/brave-search-mcp-server",
   "mcpName": "io.github.brave/brave-search-mcp-server",
   "private": false,
-  "version": "2.0.69",
+  "version": "2.1.0",
   "description": "Brave Search MCP Server: web results, images, videos, rich results, AI summaries, and more.",
   "keywords": [
     "api",