Update server.js

server.js

@@ -1,101 +1,49 @@
 import express from 'express';
 import { fal } from '@fal-ai/client';
 
-// --- Multi-Key Configuration ---
-// *** 使用 FAL_KEY 环境变量读取逗号分隔的密钥 ***
-const rawFalKeys = process.env.FAL_KEY; // Expect comma-separated keys: key1,key2,key3 in FAL_KEY
-const API_KEY = process.env.API_KEY; // Custom API Key for proxy auth remains the same
-
-if (!rawFalKeys) {
-    // *** 更新错误信息以引用 FAL_KEY ***
-    console.error("Error: FAL_KEY environment variable is not set (should be comma-separated).");
-    process.exit(1);
-}
-
-if (!API_KEY) {
-    console.error("Error: API_KEY environment variable is not set.");
-    process.exit(1);
-}
-
-// Parse and prepare the keys
-let falKeys = rawFalKeys.split(',')
-    .map(key => key.trim())
-    .filter(key => key.length > 0)
-    .map(key => ({
-        key: key,
-        failed: false,          // Track if the key is currently considered failed
-        failedTimestamp: 0      // Timestamp when the key was marked as failed
-    }));
-
-if (falKeys.length === 0) {
-    // *** 更新错误信息以引用 FAL_KEY ***
-    console.error("Error: No valid keys found in FAL_KEY after processing the environment variable.");
-    process.exit(1);
-}
-
-let currentKeyIndex = 0;
-const failedKeyCooldown = 60 * 1000; // Cooldown period in milliseconds (e.g., 60 seconds) before retrying a failed key
-
-// *** 更新日志信息以引用 FAL_KEY ***
-console.log(`Loaded ${falKeys.length} FAL API Key(s) from FAL_KEY environment variable.`);
-console.log(`Failed key cooldown period: ${failedKeyCooldown / 1000} seconds.`);
-
-// NOTE: We will configure fal client per request now, so initial global config is removed.
+// --- Express App Setup ---
+const app = express();
+app.use(express.json({ limit: '50mb' }));
+app.use(express.urlencoded({ extended: true, limit: '50mb' }));
 
-// --- Key Management Functions ---
+const PORT = process.env.PORT || 3000;
 
-/**
- * Selects the next available FAL key using round-robin and skipping recently failed keys.
- * @returns {object | null} Key info object { key, failed, failedTimestamp } or null if all keys are failed.
- */
-function getNextKey() {
-    const totalKeys = falKeys.length;
-    if (totalKeys === 0) return null;
-
-    let attempts = 0;
-    while (attempts < totalKeys) {
-        const keyIndex = currentKeyIndex % totalKeys;
-        const keyInfo = falKeys[keyIndex];
-        // Increment index for the *next* call, ensuring round-robin
-        currentKeyIndex = (currentKeyIndex + 1) % totalKeys;
-
-        // Check if key is marked as failed and if cooldown has passed
-        if (keyInfo.failed) {
-            const now = Date.now();
-            if (now - keyInfo.failedTimestamp < failedKeyCooldown) {
-                // console.log(`Key index ${keyIndex} is in cooldown. Skipping.`);
-                attempts++;
-                continue; // Skip this key, it's still in cooldown
-            } else {
-                console.log(`Cooldown finished for key index ${keyIndex}. Resetting failure status.`);
-                keyInfo.failed = false; // Cooldown expired, reset status
-                keyInfo.failedTimestamp = 0;
-            }
-        }
-        // console.log(`Selected key index: ${keyIndex}`);
-        return keyInfo; // Return the valid key info object
-    }
+// === 全局定义限制 === (Remains the same)
+const PROMPT_LIMIT = 4800;
+const SYSTEM_PROMPT_LIMIT = 4800;
+// === 限制定义结束 ===
 
-    console.warn("All FAL keys are currently marked as failed and in cooldown.");
-    return null; // All keys are currently failed and within cooldown
-}
+// 定义 fal-ai/any-llm 支持的模型列表 (Remains the same)
+const FAL_SUPPORTED_MODELS = [
+    "anthropic/claude-3.7-sonnet",
+    "anthropic/claude-3.5-sonnet",
+    "anthropic/claude-3-5-haiku",
+    "anthropic/claude-3-haiku",
+    "google/gemini-pro-1.5",
+    "google/gemini-flash-1.5",
+    "google/gemini-flash-1.5-8b",
+    "google/gemini-2.0-flash-001",
+    "meta-llama/llama-3.2-1b-instruct",
+    "meta-llama/llama-3.2-3b-instruct",
+    "meta-llama/llama-3.1-8b-instruct",
+    "meta-llama/llama-3.1-70b-instruct",
+    "openai/gpt-4o-mini",
+    "openai/gpt-4o",
+    "deepseek/deepseek-r1",
+    "meta-llama/llama-4-maverick",
+    "meta-llama/llama-4-scout"
+];
 
-/**
- * Marks a specific key as failed.
- * @param {object} keyInfo - The key info object to mark as failed.
- */
-function markKeyFailed(keyInfo) {
-    if (keyInfo && !keyInfo.failed) { // Only mark if not already marked
-        keyInfo.failed = true;
-        keyInfo.failedTimestamp = Date.now();
-        const keyIndex = falKeys.findIndex(k => k.key === keyInfo.key);
-        console.warn(`Marking key index ${keyIndex} (ending ...${keyInfo.key.slice(-4)}) as failed.`);
+// Helper function to get owner from model ID (Remains the same)
+const getOwner = (modelId) => {
+    if (modelId && modelId.includes('/')) {
+        return modelId.split('/')[0];
     }
-}
+    return 'fal-ai';
+};
 
 /**
  * Determines if an error likely indicates an API key issue (auth, quota, etc.).
- * This needs refinement based on actual errors from fal.ai.
  * @param {Error} error - The error object caught from the fal client.
  * @returns {boolean} - True if the error suggests a key failure, false otherwise.
  */
@@ -115,19 +63,10 @@ function isKeyRelatedError(error) {
         errorMessage.includes('quota exceeded')) {
         return true;
     }
-    // Add more specific error messages or codes from fal.ai if known
-    // console.log("Error does not appear to be key-related:", error); // Debugging
     return false;
 }
 
-// --- Express App Setup ---
-const app = express();
-app.use(express.json({ limit: '50mb' }));
-app.use(express.urlencoded({ extended: true, limit: '50mb' }));
-
-const PORT = process.env.PORT || 3000;
-
-// API Key 鉴权中间件 (Remains the same, checks custom API_KEY)
+// API Key 鉴权中间件 (Modified to extract FAL key directly)
 const apiKeyAuth = (req, res, next) => {
     const authHeader = req.headers['authorization'];
 
@@ -143,50 +82,19 @@ const apiKeyAuth = (req, res, next) => {
     }
 
     const providedKey = authParts[1];
-    if (providedKey !== API_KEY) {
-        console.warn('Unauthorized: Invalid API Key');
-        return res.status(401).json({ error: 'Unauthorized: Invalid API Key' });
+    if (!providedKey || providedKey.trim() === '') {
+        console.warn('Unauthorized: Empty API Key');
+        return res.status(401).json({ error: 'Unauthorized: Empty API Key' });
     }
 
+    // Store the FAL key in the request object for later use
+    req.falKey = providedKey;
+    console.log(`Received request with FAL key ending in ...${providedKey.slice(-4)}`);
     next();
 };
 
 app.use(['/v1/models', '/v1/chat/completions'], apiKeyAuth);
 
-// === 全局定义限制 === (Remains the same)
-const PROMPT_LIMIT = 4800;
-const SYSTEM_PROMPT_LIMIT = 4800;
-// === 限制定义结束 ===
-
-// 定义 fal-ai/any-llm 支持的模型列表 (Remains the same)
-const FAL_SUPPORTED_MODELS = [
-    "anthropic/claude-3.7-sonnet",
-    "anthropic/claude-3.5-sonnet",
-    "anthropic/claude-3-5-haiku",
-    "anthropic/claude-3-haiku",
-    "google/gemini-pro-1.5",
-    "google/gemini-flash-1.5",
-    "google/gemini-flash-1.5-8b",
-    "google/gemini-2.0-flash-001",
-    "meta-llama/llama-3.2-1b-instruct",
-    "meta-llama/llama-3.2-3b-instruct",
-    "meta-llama/llama-3.1-8b-instruct",
-    "meta-llama/llama-3.1-70b-instruct",
-    "openai/gpt-4o-mini",
-    "openai/gpt-4o",
-    "deepseek/deepseek-r1",
-    "meta-llama/llama-4-maverick",
-    "meta-llama/llama-4-scout"
-];
-
-// Helper function to get owner from model ID (Remains the same)
-const getOwner = (modelId) => {
-    if (modelId && modelId.includes('/')) {
-        return modelId.split('/')[0];
-    }
-    return 'fal-ai';
-};
-
 // GET /v1/models endpoint (Remains the same)
 app.get('/v1/models', (req, res) => {
     console.log("Received request for GET /v1/models");
@@ -323,74 +231,53 @@ function convertMessagesToFalPrompt(messages) {
 }
 // === convertMessagesToFalPrompt 函数结束 ===
 
-
 /**
- * Wraps the fal.ai API call with retry logic using available keys.
+ * Makes a call to the fal.ai API using the provided key.
  * @param {'stream' | 'subscribe'} operation - The fal operation to perform.
  * @param {string} functionId - The fal function ID (e.g., "fal-ai/any-llm").
  * @param {object} params - The parameters for the fal function call (input, logs, etc.).
- * @returns {Promise<any>} - The result from the successful fal call (stream or subscription result).
- * @throws {Error} - Throws an error if all keys fail or a non-key-related error occurs.
+ * @param {string} falKey - The FAL API key to use for this request.
+ * @returns {Promise<any>} - The result from the fal call (stream or subscription result).
+ * @throws {Error} - Throws an error if the call fails.
  */
-async function tryFalCallWithFailover(operation, functionId, params) {
-    const maxRetries = falKeys.length; // Try each key at most once per request cycle
-    let lastError = null;
-
-    for (let i = 0; i < maxRetries; i++) {
-        const keyInfo = getNextKey();
-        if (!keyInfo) {
-            throw new Error(lastError ? `All FAL keys failed. Last error: ${lastError.message}` : "All FAL keys are currently unavailable (failed or in cooldown).");
-        }
-
-        const currentFalKey = keyInfo.key;
-        console.log(`Attempt ${i + 1}/${maxRetries}: Using key ending in ...${currentFalKey.slice(-4)}`);
-
-        try {
-            // --- Configure fal client with the selected key for this attempt ---
-            // WARNING: This global config change might have concurrency issues in high-load scenarios
-            // if the fal client library doesn't isolate requests properly.
-            fal.config({ credentials: currentFalKey });
-
-            if (operation === 'stream') {
-                const streamResult = await fal.stream(functionId, params);
-                console.log(`Successfully initiated stream with key ending in ...${currentFalKey.slice(-4)}`);
-                return streamResult;
-            } else { // 'subscribe' (non-stream)
-                const result = await fal.subscribe(functionId, params);
-                console.log(`Successfully completed subscribe request with key ending in ...${currentFalKey.slice(-4)}`);
-
-                if (result && result.error) {
-                     console.warn(`Fal-ai returned an application error (non-stream) with key ...${currentFalKey.slice(-4)}: ${JSON.stringify(result.error)}`);
-                }
-                return result;
-            }
-        } catch (error) {
-            console.error(`Error using key ending in ...${currentFalKey.slice(-4)}:`, error.message || error);
-            lastError = error;
-
-            if (isKeyRelatedError(error)) {
-                markKeyFailed(keyInfo);
-                console.log(`Key marked as failed. Trying next key if available...`);
-            } else {
-                console.error("Non-key related error occurred. Aborting retries.");
-                throw error;
+async function callFalApi(operation, functionId, params, falKey) {
+    try {
+        // Configure fal client with the provided key
+        fal.config({ credentials: falKey });
+        
+        if (operation === 'stream') {
+            const streamResult = await fal.stream(functionId, params);
+            console.log(`Successfully initiated stream with key ending in ...${falKey.slice(-4)}`);
+            return streamResult;
+        } else { // 'subscribe' (non-stream)
+            const result = await fal.subscribe(functionId, params);
+            console.log(`Successfully completed subscribe request with key ending in ...${falKey.slice(-4)}`);
+
+            if (result && result.error) {
+                console.warn(`Fal-ai returned an application error (non-stream) with key ...${falKey.slice(-4)}: ${JSON.stringify(result.error)}`);
             }
+            return result;
         }
+    } catch (error) {
+        console.error(`Error using key ending in ...${falKey.slice(-4)}:`, error.message || error);
+        
+        if (isKeyRelatedError(error)) {
+            console.error(`Key-related error detected with key ending in ...${falKey.slice(-4)}`);
+        }
+        
+        throw error;
     }
-
-    console.error("All FAL keys failed after attempting each one.");
-    throw new Error(lastError ? `All FAL keys failed. Last error: ${lastError.message}` : "All FAL API keys failed.");
 }
 
-
-// POST /v1/chat/completions endpoint (Modified to use tryFalCallWithFailover)
+// POST /v1/chat/completions endpoint (Modified to use client-provided key)
 app.post('/v1/chat/completions', async (req, res) => {
     const { model, messages, stream = false, reasoning = false, ...restOpenAIParams } = req.body;
+    const falKey = req.falKey; // Get the FAL key from the request object
 
     console.log(`Received chat completion request for model: ${model}, stream: ${stream}`);
 
     if (!FAL_SUPPORTED_MODELS.includes(model)) {
-         console.warn(`Warning: Requested model '${model}' is not in the explicitly supported list.`);
+        console.warn(`Warning: Requested model '${model}' is not in the explicitly supported list.`);
     }
     if (!model || !messages || !Array.isArray(messages) || messages.length === 0) {
         console.error("Invalid request parameters:", { model, messages: Array.isArray(messages) ? messages.length : typeof messages });
@@ -420,9 +307,9 @@ app.post('/v1/chat/completions', async (req, res) => {
             let falStream;
 
             try {
-                 falStream = await tryFalCallWithFailover('stream', "fal-ai/any-llm", { input: falInput });
+                falStream = await callFalApi('stream', "fal-ai/any-llm", { input: falInput }, falKey);
 
-                 for await (const event of falStream) {
+                for await (const event of falStream) {
                     const currentOutput = (event && typeof event.output === 'string') ? event.output : '';
                     const isPartial = (event && typeof event.partial === 'boolean') ? event.partial : true;
                     const errorInfo = (event && event.error) ? event.error : null;
@@ -448,84 +335,84 @@ app.post('/v1/chat/completions', async (req, res) => {
                         const openAIChunk = { id: `chatcmpl-${Date.now()}`, object: "chat.completion.chunk", created: Math.floor(Date.now() / 1000), model: model, choices: [{ index: 0, delta: { content: deltaContent }, finish_reason: isPartial === false ? "stop" : null }] };
                         res.write(`data: ${JSON.stringify(openAIChunk)}\n\n`);
                     }
-                 }
-                 res.write(`data: [DONE]\n\n`);
-                 res.end();
-                 console.log("Stream finished successfully.");
+                }
+                res.write(`data: [DONE]\n\n`);
+                res.end();
+                console.log("Stream finished successfully.");
 
             } catch (streamError) {
                 console.error('Error during stream processing:', streamError);
                 if (!res.writableEnded) {
-                     try {
-                         const errorDetails = (streamError instanceof Error) ? streamError.message : JSON.stringify(streamError);
-                         const finalErrorChunk = { error: { message: "Stream failed", type: "proxy_error", details: errorDetails } };
-                         res.write(`data: ${JSON.stringify(finalErrorChunk)}\n\n`);
-                         res.write(`data: [DONE]\n\n`);
-                         res.end();
-                     } catch (finalError) {
-                         console.error('Error sending final stream error message to client:', finalError);
-                         if (!res.writableEnded) { res.end(); }
-                     }
-                 }
+                    try {
+                        const errorDetails = (streamError instanceof Error) ? streamError.message : JSON.stringify(streamError);
+                        const finalErrorChunk = { error: { message: "Stream failed", type: "proxy_error", details: errorDetails } };
+                        res.write(`data: ${JSON.stringify(finalErrorChunk)}\n\n`);
+                        res.write(`data: [DONE]\n\n`);
+                        res.end();
+                    } catch (finalError) {
+                        console.error('Error sending final stream error message to client:', finalError);
+                        if (!res.writableEnded) { res.end(); }
+                    }
+                }
             }
 
         } else { // Non-stream
-             console.log("Executing non-stream request with failover...");
-             const result = await tryFalCallWithFailover('subscribe', "fal-ai/any-llm", { input: falInput, logs: true });
-
-             console.log("Received non-stream result from fal-ai via failover wrapper.");
-
-             if (result && result.error) {
-                 console.error("Fal-ai returned an application error in non-stream mode (after successful API call):", result.error);
-                 return res.status(500).json({
-                      object: "error",
-                      message: `Fal-ai application error: ${JSON.stringify(result.error)}`,
-                      type: "fal_ai_error",
-                      param: null,
-                      code: result.error.code || null
-                 });
-             }
-
-             const openAIResponse = {
-                 id: `chatcmpl-${result?.requestId || Date.now()}`,
-                 object: "chat.completion",
-                 created: Math.floor(Date.now() / 1000),
-                 model: model,
-                 choices: [{
-                     index: 0,
-                     message: {
-                         role: "assistant",
-                         content: result?.output || ""
-                     },
-                     finish_reason: "stop"
-                 }],
-                 usage: {
-                     prompt_tokens: null,
-                     completion_tokens: null,
-                     total_tokens: null
-                 },
-                 system_fingerprint: null,
-                 ...(result?.reasoning && { fal_reasoning: result.reasoning }),
-             };
-             res.json(openAIResponse);
-             console.log("Returned non-stream response successfully.");
+            console.log("Executing non-stream request...");
+            const result = await callFalApi('subscribe', "fal-ai/any-llm", { input: falInput, logs: true }, falKey);
+
+            console.log("Received non-stream result from fal-ai.");
+
+            if (result && result.error) {
+                console.error("Fal-ai returned an application error in non-stream mode (after successful API call):", result.error);
+                return res.status(500).json({
+                    object: "error",
+                    message: `Fal-ai application error: ${JSON.stringify(result.error)}`,
+                    type: "fal_ai_error",
+                    param: null,
+                    code: result.error.code || null
+                });
+            }
+
+            const openAIResponse = {
+                id: `chatcmpl-${result?.requestId || Date.now()}`,
+                object: "chat.completion",
+                created: Math.floor(Date.now() / 1000),
+                model: model,
+                choices: [{
+                    index: 0,
+                    message: {
+                        role: "assistant",
+                        content: result?.output || ""
+                    },
+                    finish_reason: "stop"
+                }],
+                usage: {
+                    prompt_tokens: null,
+                    completion_tokens: null,
+                    total_tokens: null
+                },
+                system_fingerprint: null,
+                ...(result?.reasoning && { fal_reasoning: result.reasoning }),
+            };
+            res.json(openAIResponse);
+            console.log("Returned non-stream response successfully.");
         }
 
     } catch (error) {
         console.error('Unhandled error in /v1/chat/completions:', error);
         if (!res.headersSent) {
             const errorMessage = (error instanceof Error) ? error.message : JSON.stringify(error);
-            const errorType = error.message?.includes("All FAL keys failed") ? "api_key_error" : "proxy_internal_error";
+            const errorType = isKeyRelatedError(error) ? "api_key_error" : "proxy_internal_error";
             res.status(500).json({
-                 error: {
-                     message: `Internal Server Error in Proxy: ${errorMessage}`,
-                     type: errorType,
-                     details: error.stack // Optional: include stack in dev/debug mode
-                 }
+                error: {
+                    message: `Internal Server Error in Proxy: ${errorMessage}`,
+                    type: errorType,
+                    details: error.stack // Optional: include stack in dev/debug mode
+                }
             });
         } else if (!res.writableEnded) {
-             console.error("Headers already sent, attempting to end response after error.");
-             res.end();
+            console.error("Headers already sent, attempting to end response after error.");
+            res.end();
         }
     }
 });
@@ -533,11 +420,9 @@ app.post('/v1/chat/completions', async (req, res) => {
 // --- Server Start ---
 app.listen(PORT, () => {
     console.log(`===========================================================`);
-    console.log(` Fal OpenAI Proxy Server (Multi-Key Failover)`);
+    console.log(` Fal OpenAI Proxy Server (Direct Key Mode)`);
     console.log(` Listening on port: ${PORT}`);
-    // *** 更新日志信息以引用 FAL_KEY ***
-    console.log(` Loaded ${falKeys.length} FAL API Key(s) from FAL_KEY.`);
-    console.log(` API Key Auth Enabled: ${API_KEY ? 'Yes' : 'No'}`);
+    console.log(` Using client-provided FAL API keys directly`);
     console.log(` Limits: System Prompt=${SYSTEM_PROMPT_LIMIT}, Prompt=${PROMPT_LIMIT}`);
     console.log(` Chat Completions: POST http://localhost:${PORT}/v1/chat/completions`);
     console.log(` Models Endpoint: GET http://localhost:${PORT}/v1/models`);
@@ -546,5 +431,5 @@ app.listen(PORT, () => {
 
 // Root path response
 app.get('/', (req, res) => {
-    res.send('Fal OpenAI Proxy (Multi-Key Failover) is running.');
+    res.send('Fal OpenAI Proxy (Direct Key Mode) is running.');
 });