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

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

@@ -12,9 +12,14 @@ const DEFAULT_PER_SECOND_LIMIT = 1;
 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);
@@ -49,7 +54,6 @@ class ApiKeyManager {
                     result.set(keyHash, info.remaining);
                 }
             }
-            console.log(`[ApiKeyManager] Loaded state for ${result.size} keys from ${STATE_FILE}`);
             return result;
         }
         catch (err) {
@@ -90,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
      */
@@ -142,7 +230,6 @@ class ApiKeyManager {
         state.inUse = false;
         // Save state after update
         this.saveState();
-        console.debug(`[ApiKeyManager] Key ${apiKey.slice(0, 8)}... quota: ${state.remainingQuota}/${state.monthLimit}`);
     }
     /**
      * Mark a key as invalid (e.g., 401 Unauthorized)
@@ -180,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.70",
+  "version": "2.1.0",
   "description": "Brave Search MCP Server: web results, images, videos, rich results, AI summaries, and more.",
   "keywords": [
     "api",