Update server.js

server.js

@@ -1,4 +1,5 @@
 import express from 'express';
+// Import the 'fal' object directly for configuration within the retry loop
 import { fal } from '@fal-ai/client';
 
 // --- Key Management Setup ---
@@ -7,132 +8,367 @@ const FAL_KEY_STRING = process.env.FAL_KEY;
 // Read the custom API Key for proxy authentication
 const API_KEY = process.env.API_KEY;
 
-// --- (Initial checks for FAL_KEY_STRING, API_KEY, parsing falKeys remain the same) ---
+// --- Initial Environment Variable Checks ---
 if (!FAL_KEY_STRING) {
     console.error("ERROR: FAL_KEY environment variable is not set.");
     console.error("Ensure FAL_KEY contains a comma-separated list of your Fal AI keys.");
-    process.exit(1);
+    process.exit(1); // Exit if no Fal keys are provided
 }
+
+// Parse the comma-separated keys from FAL_KEY_STRING
 const falKeys = FAL_KEY_STRING.split(',')
-    .map(key => key.trim())
-    .filter(key => key.length > 0);
+    .map(key => key.trim()) // Remove leading/trailing whitespace
+    .filter(key => key.length > 0); // Remove any empty strings resulting from extra commas
+
 if (falKeys.length === 0) {
     console.error("ERROR: No valid Fal keys found in the FAL_KEY environment variable after parsing.");
     console.error("Ensure FAL_KEY is a comma-separated list, e.g., 'key1,key2,key3'.");
-    process.exit(1);
+    process.exit(1); // Exit if parsing results in zero valid keys
 }
+
 if (!API_KEY) {
     console.error("ERROR: API_KEY environment variable is not set.");
-    process.exit(1);
+    process.exit(1); // Exit if the proxy auth key is missing
 }
-// --- (End initial checks) ---
+// --- End Initial Checks ---
 
 
 let currentKeyIndex = 0;
+// Keep track of keys that failed persistently during runtime
 const invalidKeys = new Set();
+
 console.log(`Loaded ${falKeys.length} Fal AI Key(s) from the FAL_KEY environment variable.`);
 
-// --- (getNextValidKey function remains the same) ---
+/**
+ * Gets the next available valid Fal AI key in a round-robin fashion.
+ * Skips keys that have been marked as invalid.
+ * @returns {object|null} An object containing the key and its original index { key, index }, or null if no valid keys remain.
+ */
 function getNextValidKey() {
+    // Check if all keys have been marked as invalid
     if (invalidKeys.size >= falKeys.length) {
         console.error("All Fal AI keys are marked as invalid.");
-        return null;
+        return null; // No valid keys left
     }
+
     const initialIndex = currentKeyIndex;
-    let attempts = 0;
+    let attempts = 0; // Prevent infinite loops in edge cases
     while (attempts < falKeys.length) {
         const keyIndex = currentKeyIndex % falKeys.length;
         const key = falKeys[keyIndex];
+
+        // Move to the next index for the *next* call, regardless of validity
         currentKeyIndex = (keyIndex + 1) % falKeys.length;
+
+        // Check if the current key is NOT in the invalid set
         if (!invalidKeys.has(key)) {
+            // Found a valid key
             console.log(`Using Fal Key index: ${keyIndex} (from FAL_KEY list)`);
-            return { key, index: keyIndex };
+            return { key, index: keyIndex }; // Return the key and its original index
         } else {
             console.log(`Skipping invalid Fal Key index: ${keyIndex}`);
         }
+
         attempts++;
+        // Safety check: If we've looped back to the start after trying, break.
         if (currentKeyIndex === initialIndex && attempts > 0) {
             console.warn("Looped through all keys, potentially all are invalid.");
             break;
         }
     }
+
+    // If we exit the loop, it means no valid key was found
     console.error("Could not find a valid Fal AI key after checking all potentially available keys.");
     return null;
 }
 
-// --- (isKeyRelatedError function remains the same) ---
+/**
+ * Checks if an error object likely indicates an issue with the Fal AI API key.
+ * This is heuristic-based and may need refinement based on observed Fal errors.
+ * @param {Error|object|null} error - The error object caught.
+ * @returns {boolean} True if the error seems key-related, false otherwise.
+ */
 function isKeyRelatedError(error) {
-    if (!error) return false;
+    if (!error) return false; // Handle null/undefined errors
+
     const message = error.message?.toLowerCase() || '';
+    // Check common HTTP status properties
     const status = error.status || error.statusCode;
+
+    // Check for specific HTTP status codes (401: Unauthorized, 403: Forbidden, 429: Too Many Requests/Quota)
     if (status === 401 || status === 403 || status === 429) {
         console.warn(`Detected potential key-related error (HTTP Status: ${status}).`);
         return true;
     }
+
+    // Check for common error message patterns (case-insensitive)
     const keyErrorPatterns = [
         'invalid api key', 'authentication failed', 'permission denied',
         'quota exceeded', 'forbidden', 'unauthorized', 'rate limit',
         'credentials', 'api key missing', 'invalid credential',
-        'exhausted balance', 'user is locked' // Add specific messages if observed
+        'exhausted balance', 'user is locked' // Added based on observed logs
     ];
     if (keyErrorPatterns.some(pattern => message.includes(pattern))) {
         console.warn(`Detected potential key-related error (message contains relevant pattern: "${message}")`);
         return true;
     }
-    // Also check the body.detail if status is 403, as seen in the logs
+
+    // Check the body.detail if status is 403 (as seen in Fal errors)
     if (status === 403 && error.body?.detail) {
-        const detailMessage = error.body.detail.toLowerCase();
+        const detailMessage = String(error.body.detail).toLowerCase();
          if (keyErrorPatterns.some(pattern => detailMessage.includes(pattern))) {
             console.warn(`Detected potential key-related error (body.detail contains relevant pattern: "${detailMessage}")`);
             return true;
          }
     }
+
+    // Add more specific checks based on observed Fal AI errors if needed
+    // e.g., if (error.code === 'FAL_AUTH_FAILURE') return true;
+
     return false;
 }
 // --- End Key Management Setup ---
 
 const app = express();
+// Increase payload size limits if needed
 app.use(express.json({ limit: '50mb' }));
 app.use(express.urlencoded({ extended: true, limit: '50mb' }));
+
 const PORT = process.env.PORT || 3000;
 
-// --- (apiKeyAuth middleware remains the same) ---
+// API Key Authentication Middleware
 const apiKeyAuth = (req, res, next) => {
     const authHeader = req.headers['authorization'];
+
     if (!authHeader) {
         console.warn('Unauthorized: No Authorization header provided');
         return res.status(401).json({ error: 'Unauthorized: No API Key provided' });
     }
+
+    // Expecting "Bearer YOUR_API_KEY"
     const authParts = authHeader.split(' ');
     if (authParts.length !== 2 || authParts[0].toLowerCase() !== 'bearer') {
         console.warn('Unauthorized: Invalid Authorization header format. Expected "Bearer <key>".');
         return res.status(401).json({ error: 'Unauthorized: Invalid Authorization header format' });
     }
+
     const providedKey = authParts[1];
     if (providedKey !== API_KEY) {
         console.warn('Unauthorized: Invalid API Key provided.');
         return res.status(401).json({ error: 'Unauthorized: Invalid API Key' });
     }
+
+    // Key is valid, proceed to the next middleware or route handler
     next();
 };
+
+// Apply API Key Authentication to relevant endpoints
 app.use(['/v1/models', '/v1/chat/completions'], apiKeyAuth);
 
+// === Global Limits Definition ===
+const PROMPT_LIMIT = 4800; // Max length for the main 'prompt' field
+const SYSTEM_PROMPT_LIMIT = 4800; // Max length for the 'system_prompt' field
+// === End Limits Definition ===
+
+// Define the list of models supported by fal-ai/any-llm (Update as needed)
+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"
+    // Add or remove models here
+];
+
+// Helper function to extract the owner/organization from a model ID string
+const getOwner = (modelId) => {
+    if (modelId && typeof modelId === 'string' && modelId.includes('/')) {
+        return modelId.split('/')[0];
+    }
+    // Default owner if format is unexpected or missing
+    return 'fal-ai';
+}
 
-// --- (Global Limits, FAL_SUPPORTED_MODELS, getOwner remain the same) ---
-const PROMPT_LIMIT = 4800;
-const SYSTEM_PROMPT_LIMIT = 4800;
-const FAL_SUPPORTED_MODELS = [ /* ... model list ... */ ];
-const getOwner = (modelId) => { /* ... */ };
+// GET /v1/models endpoint - Returns the list of supported models
+app.get('/v1/models', (req, res) => {
+    console.log("Received request for GET /v1/models");
+    try {
+        const modelsData = FAL_SUPPORTED_MODELS.map(modelId => ({
+            id: modelId,
+            object: "model",
+            created: Math.floor(Date.now() / 1000), // Use current timestamp
+            owned_by: getOwner(modelId)
+        }));
+        res.json({ object: "list", data: modelsData });
+        console.log("Successfully returned model list.");
+    } catch (error) {
+        console.error("Error processing GET /v1/models:", error);
+        res.status(500).json({ error: "Failed to retrieve model list." });
+    }
+});
 
-// --- (GET /v1/models endpoint remains the same) ---
-app.get('/v1/models', (req, res) => { /* ... */ });
 
-// --- (convertMessagesToFalPrompt function remains the same) ---
-function convertMessagesToFalPrompt(messages) { /* ... */ }
+/**
+ * Converts OpenAI-style messages array to Fal AI's prompt and system_prompt format.
+ * Implements System prompt top-priority, separator, and recency-based history filling.
+ * Includes robustness checks for input validation and internal errors.
+ * @param {Array<object>} messages - Array of message objects ({ role: string, content: string|null }).
+ * @returns {object} An object containing { system_prompt: string, prompt: string }.
+ * @throws {Error} If input is fundamentally invalid (e.g., not an array) or an unexpected internal processing error occurs.
+ */
+function convertMessagesToFalPrompt(messages) {
+    // --- Optional Debug Log: Uncomment to see the exact input causing issues ---
+    // console.log(">>> Entering convertMessagesToFalPrompt with messages:", JSON.stringify(messages, null, 2));
+
+    // --- Input Validation ---
+    if (!Array.isArray(messages)) {
+        console.error("!!! ERROR in convertMessagesToFalPrompt: Input 'messages' is not an array.");
+        // Throw an error here because this is a fundamental type mismatch from the caller.
+        throw new Error("Invalid input: 'messages' must be an array.");
+    }
+    if (messages.length === 0) {
+        console.warn("Warning in convertMessagesToFalPrompt: Input 'messages' array is empty.");
+        // Return empty strings if no messages, this is valid input.
+        return { system_prompt: "", prompt: "" };
+    }
+    // --- End Input Validation ---
+
+    // --- Main Processing Logic ---
+    try { // *** Wrap core logic in try...catch to handle internal errors ***
+        let fixed_system_prompt_content = "";
+        const conversation_message_blocks = [];
+        // console.log(`Original messages count: ${messages.length}`);
+
+        // 1. Separate System messages, format User/Assistant messages
+        for (const message of messages) {
+            // ** Validate individual message structure **
+            if (!message || typeof message !== 'object' || typeof message.role !== 'string') {
+                console.warn(`--> Skipping invalid message object in convertMessagesToFalPrompt: ${JSON.stringify(message)}`);
+                continue; // Skip this malformed message, proceed with others
+            }
+
+            // ** Safely handle content (null/undefined/non-string become empty string) **
+            let content = (message.content === null || message.content === undefined) ? "" : String(message.content);
+
+            switch (message.role) {
+                case 'system':
+                    fixed_system_prompt_content += `System: ${content}\n\n`;
+                    break;
+                case 'user':
+                    conversation_message_blocks.push(`Human: ${content}\n\n`);
+                    break;
+                case 'assistant':
+                    conversation_message_blocks.push(`Assistant: ${content}\n\n`);
+                    break;
+                default:
+                    console.warn(`--> Unsupported role encountered in convertMessagesToFalPrompt: '${message.role}'. Skipping message.`);
+                    continue; // Skip messages with unsupported roles
+            }
+        }
+
+        // 2. Truncate combined system messages if they exceed the limit
+        if (fixed_system_prompt_content.length > SYSTEM_PROMPT_LIMIT) {
+            const originalLength = fixed_system_prompt_content.length;
+            fixed_system_prompt_content = fixed_system_prompt_content.substring(0, SYSTEM_PROMPT_LIMIT);
+            console.warn(`Combined system messages truncated from ${originalLength} to ${SYSTEM_PROMPT_LIMIT} characters.`);
+        }
+        fixed_system_prompt_content = fixed_system_prompt_content.trim();
+
+        // 3. Calculate remaining space in system_prompt for history
+        let space_occupied_by_fixed_system = 0;
+        if (fixed_system_prompt_content.length > 0) {
+            space_occupied_by_fixed_system = fixed_system_prompt_content.length + 4; // Approx for spacing
+        }
+        const remaining_system_limit = Math.max(0, SYSTEM_PROMPT_LIMIT - space_occupied_by_fixed_system);
+
+        // 4. Fill history backwards (recency): Prioritize 'prompt', then 'system_prompt' overflow
+        const prompt_history_blocks = [];
+        const system_prompt_history_blocks = [];
+        let current_prompt_length = 0;
+        let current_system_history_length = 0;
+        let promptFull = (PROMPT_LIMIT <= 0);
+        let systemHistoryFull = (remaining_system_limit <= 0);
+
+        for (let i = conversation_message_blocks.length - 1; i >= 0; i--) {
+            const message_block = conversation_message_blocks[i];
+            const block_length = (typeof message_block === 'string') ? message_block.length : 0; // Ensure it's a string
+
+            if (block_length === 0) continue;
+            if (promptFull && systemHistoryFull) break;
+
+            // Try fitting into the main 'prompt' first
+            if (!promptFull) {
+                if (current_prompt_length + block_length <= PROMPT_LIMIT) {
+                    prompt_history_blocks.unshift(message_block);
+                    current_prompt_length += block_length;
+                    continue;
+                } else {
+                    promptFull = true;
+                }
+            }
+
+            // If prompt is full, try fitting into the 'system_prompt' remaining space
+            if (!systemHistoryFull) {
+                if (current_system_history_length + block_length <= remaining_system_limit) {
+                    system_prompt_history_blocks.unshift(message_block);
+                    current_system_history_length += block_length;
+                    continue;
+                } else {
+                    systemHistoryFull = true;
+                }
+            }
+        }
+
+        // 5. Combine the final prompt and system_prompt parts
+        const system_prompt_history_content = system_prompt_history_blocks.join('').trim();
+        const final_prompt = prompt_history_blocks.join('').trim();
+        const SEPARATOR = "\n\n------- Earlier Conversation History -------\n\n";
+        let final_system_prompt = "";
+        const hasFixedSystem = fixed_system_prompt_content.length > 0;
+        const hasSystemHistory = system_prompt_history_content.length > 0;
+
+        if (hasFixedSystem && hasSystemHistory) {
+            final_system_prompt = fixed_system_prompt_content + SEPARATOR + system_prompt_history_content;
+        } else if (hasFixedSystem) {
+            final_system_prompt = fixed_system_prompt_content;
+        } else if (hasSystemHistory) {
+            final_system_prompt = system_prompt_history_content;
+        }
+
+        // 6. ** Crucially, always return an object **
+        const result = {
+            system_prompt: final_system_prompt,
+            prompt: final_prompt
+        };
+        // console.log("<<< Exiting convertMessagesToFalPrompt successfully."); // Optional success log
+        return result;
+
+    } catch (internalError) { // *** Catch any unexpected errors during processing ***
+        console.error("!!! CRITICAL ERROR inside convertMessagesToFalPrompt processing:", internalError);
+        // Log the input that caused the error for debugging
+        console.error("!!! Failing messages input was:", JSON.stringify(messages, null, 2));
+        // Re-throw the error so the main handler knows something went wrong during setup.
+        // This will be caught by the outer try...catch in the route handler.
+        throw new Error(`Failed to process messages internally: ${internalError.message}`);
+    }
+}
+// === End convertMessagesToFalPrompt function ===
 
 
 /**
- * MODIFIED: Makes a request to the Fal AI API, handling key rotation and retries on key-related errors.
+ * Makes a request to the Fal AI API, handling key rotation and retries on key-related errors.
  * For stream requests, returns the stream AND the key info used.
  * @param {object} falInput - The input object for the Fal AI API call.
  * @param {boolean} [stream=false] - Whether to make a streaming request.
@@ -166,10 +402,10 @@ async function makeFalRequestWithRetry(falInput, stream = false) {
             if (stream) {
                 const falStream = await fal.stream("fal-ai/any-llm", { input: falInput });
                 console.log(`Successfully initiated stream with key index ${keyInfo.index}.`);
-                // **MODIFIED: Return stream AND key info**
+                // Return stream AND key info
                 return { stream: falStream, keyUsed: keyInfo.key, indexUsed: keyInfo.index };
             } else {
-                // Non-stream logic remains the same
+                // Non-stream logic
                 console.log(`Executing non-stream request with key index ${keyInfo.index}...`);
                 const result = await fal.subscribe("fal-ai/any-llm", { input: falInput, logs: true });
                 console.log(`Successfully received non-stream result with key index ${keyInfo.index}.`);
@@ -186,7 +422,7 @@ async function makeFalRequestWithRetry(falInput, stream = false) {
                 return result; // Return only result for non-stream
             }
         } catch (error) {
-            // This catch block now primarily handles errors during *request initiation*
+            // This catch block handles errors during *request initiation*
             console.error(`Error caught during request initiation using Fal Key index ${keyInfo.index}:`, error.message || error);
             if (isKeyRelatedError(error)) {
                 console.warn(`Marking Fal Key index ${keyInfo.index} as invalid due to caught initiation error.`);
@@ -209,41 +445,52 @@ app.post('/v1/chat/completions', async (req, res) => {
     const { model, messages, stream = false, reasoning = false, ...restOpenAIParams } = req.body;
     console.log(`--> POST /v1/chat/completions | Model: ${model} | Stream: ${stream}`);
 
-    // --- (Input validation for model, messages remains the same) ---
+    // --- Input validation for model and messages ---
      if (!model || !messages || !Array.isArray(messages) || messages.length === 0) {
         console.error("Invalid request: Missing 'model' or 'messages' array is empty/invalid.");
         return res.status(400).json({ error: 'Bad Request: `model` and a non-empty `messages` array are required.' });
     }
+    if (!FAL_SUPPORTED_MODELS.includes(model)) {
+        console.warn(`Warning: Requested model '${model}' is not in the explicitly supported list. Proxy will still attempt the request.`);
+    }
+    // --- End Input Validation ---
 
-    let keyUsedForRequest = null; // Variable to store the key used for this request, if successful initiation
+    let keyUsedForRequest = null; // Variable to store the key used if initiation succeeds
     let indexUsedForRequest = null;
 
     try {
+        // --- Prepare Fal AI Input ---
+        // This might throw an error if convertMessagesToFalPrompt fails internally
         const { prompt, system_prompt } = convertMessagesToFalPrompt(messages);
-        const falInput = { /* ... falInput setup ... */ };
-        falInput.model = model;
-        falInput.prompt = prompt;
+
+        const falInput = {
+            model: model,
+            prompt: prompt,
+            reasoning: !!reasoning,
+        };
         if (system_prompt && system_prompt.length > 0) {
              falInput.system_prompt = system_prompt;
         }
-        falInput.reasoning = !!reasoning;
+        // --- End Prepare Input ---
 
         console.log("Attempting Fal request with key rotation/retry logic...");
         console.log(`Prepared Input Lengths - System Prompt: ${system_prompt?.length || 0}, Prompt: ${prompt?.length || 0}`);
 
+        // --- Handle Stream vs Non-Stream ---
         if (stream) {
+            // Set headers for Server-Sent Events (SSE)
             res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
-            /* ... other headers ... */
             res.setHeader('Cache-Control', 'no-cache');
             res.setHeader('Connection', 'keep-alive');
-            res.setHeader('Access-Control-Allow-Origin', '*');
-            res.flushHeaders();
+            res.setHeader('Access-Control-Allow-Origin', '*'); // Adjust CORS for production if needed
+            res.flushHeaders(); // Send headers immediately
 
-            let previousOutput = '';
+            let previousOutput = ''; // Track previous output for delta calculation
             let streamResult; // To hold the object { stream, keyUsed, indexUsed }
 
             try {
-                // **MODIFIED: Get stream and key info**
+                // **Initiate the stream using the retry helper**
+                // This can throw if initiation fails after all retries
                 streamResult = await makeFalRequestWithRetry(falInput, true);
                 const falStream = streamResult.stream;
                 keyUsedForRequest = streamResult.keyUsed; // Store the key used for this stream
@@ -251,19 +498,24 @@ app.post('/v1/chat/completions', async (req, res) => {
 
                 // Process the stream events asynchronously
                 for await (const event of falStream) {
-                    // --- (Stream event processing logic remains the same) ---
+                    // Safely extract data from the event
                     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;
 
+                    // Handle errors reported *within* a stream event payload
                     if (errorInfo) {
-                        // Log error from within the stream, but might continue processing
                         console.error("Error received *within* fal stream event payload:", errorInfo);
-                        const errorChunk = { /* ... error chunk details ... */ };
+                        const errorChunk = {
+                            id: `chatcmpl-${Date.now()}-error`, object: "chat.completion.chunk", created: Math.floor(Date.now() / 1000), model: model,
+                            choices: [{ index: 0, delta: {}, finish_reason: "error", message: { role: 'assistant', content: `Fal Stream Event Error: ${JSON.stringify(errorInfo)}` } }]
+                        };
                         if (!res.writableEnded) { res.write(`data: ${JSON.stringify(errorChunk)}\n\n`); }
                         else { console.warn("Stream ended before writing event error."); }
-                        // Decide whether to break or continue based on error severity if needed
+                        // Potentially break or add logic based on error type here
                     }
+
+                    // Calculate the delta (new content)
                     let deltaContent = '';
                     if (currentOutput.startsWith(previousOutput)) {
                         deltaContent = currentOutput.substring(previousOutput.length);
@@ -273,46 +525,48 @@ app.post('/v1/chat/completions', async (req, res) => {
                          previousOutput = '';
                     }
                     previousOutput = currentOutput;
+
+                    // Send OpenAI-compatible SSE chunk
                     if (deltaContent || !isPartial) {
-                        const openAIChunk = { /* ... chunk details ... */ };
-                        openAIChunk.id = `chatcmpl-${Date.now()}`;
-                        openAIChunk.object = "chat.completion.chunk";
-                        openAIChunk.created = Math.floor(Date.now() / 1000);
-                        openAIChunk.model = model;
-                        openAIChunk.choices = [{ index: 0, delta: { content: deltaContent }, finish_reason: isPartial === false ? "stop" : null }];
+                        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
+                             }]
+                        };
                         if (!res.writableEnded) { res.write(`data: ${JSON.stringify(openAIChunk)}\n\n`); }
                         else { console.warn("Stream ended before writing data chunk."); }
                     }
-                    // --- (End stream event processing) ---
                 } // End for-await loop
 
-                // Send [DONE] marker
+                // Send the final [DONE] marker
                 if (!res.writableEnded) {
                     res.write(`data: [DONE]\n\n`);
-                    res.end();
+                    res.end(); // Close the connection
                     console.log("<-- Stream finished successfully and [DONE] sent.");
                 } else {
-                     console.log("<-- Stream processing finished, but connection was already ended before [DONE].");
+                    console.log("<-- Stream processing finished, but connection was already ended before [DONE].");
                 }
 
             } catch (streamError) {
-                // **MODIFIED CATCH BLOCK for stream processing errors**
-                // This catches errors from makeFalRequestWithRetry (initiation failure)
-                // OR errors thrown during the 'for await...of falStream' loop.
-                console.error('Error during stream request processing:', streamError.message || streamError);
-
-                // **NEW: Check if the error is key-related and invalidate the key if needed**
-                // We only do this if keyUsedForRequest has been set (meaning initiation succeeded)
-                // And if the error occurred *during* the stream processing, not during initiation
-                // (initiation errors are handled inside makeFalRequestWithRetry)
-                // The check `keyUsedForRequest !== null` helps distinguish.
+                // **Catch block for errors during stream processing OR initiation failure**
+                console.error('Error during stream request processing/initiation:', streamError.message || streamError);
+
+                // **Check if the error is key-related AND if initiation succeeded**
+                // This ensures we only invalidate the key if the error happened *during* processing
+                // using a key that successfully initiated the stream.
                 if (keyUsedForRequest && isKeyRelatedError(streamError)) {
                     console.warn(`Marking Fal Key index ${indexUsedForRequest} as invalid due to error during stream processing.`);
                     invalidKeys.add(keyUsedForRequest);
                 }
-                // else: The error was either not key-related, or occurred during initiation (already handled)
+                // else: Error was either non-key-related or happened during initiation (already handled/logged in makeFalRequestWithRetry).
 
-                // --- (Error reporting logic to client remains the same) ---
+                // --- Report error back to the client ---
                  try {
                      if (!res.headersSent) {
                          // Error likely during initiation (caught from makeFalRequestWithRetry)
@@ -323,7 +577,7 @@ app.post('/v1/chat/completions', async (req, res) => {
                          // Error during stream processing after headers sent
                          const errorDetails = (streamError instanceof Error) ? streamError.message : JSON.stringify(streamError);
                          res.write(`data: ${JSON.stringify({ error: { message: "Stream processing error after initiation", type: "proxy_error", details: errorDetails } })}\n\n`);
-                         res.write(`data: [DONE]\n\n`);
+                         res.write(`data: [DONE]\n\n`); // Still send DONE for robust client handling
                          res.end();
                          console.log("<-- Stream error sent within stream, stream ended.");
                      } else {
@@ -333,27 +587,40 @@ app.post('/v1/chat/completions', async (req, res) => {
                     console.error('Error sending stream error message to client:', finalError);
                     if (!res.writableEnded) { res.end(); }
                  }
-                 // --- (End error reporting) ---
+                 // --- End error reporting ---
             }
 
         } else {
-            // --- Non-Stream Logic (remains the same, uses makeFalRequestWithRetry directly) ---
+            // --- Non-Stream Request ---
             try {
+                // Get the result using the retry helper
                 const result = await makeFalRequestWithRetry(falInput, false);
-                const openAIResponse = { /* ... construct response ... */ };
-                openAIResponse.id = `chatcmpl-${result.requestId || Date.now()}`;
-                openAIResponse.object = "chat.completion";
-                openAIResponse.created = Math.floor(Date.now() / 1000);
-                openAIResponse.model = model;
-                openAIResponse.choices = [{ index: 0, message: { role: "assistant", content: result.output || "" }, finish_reason: "stop" }];
-                openAIResponse.usage = { prompt_tokens: null, completion_tokens: null, total_tokens: null };
-                openAIResponse.system_fingerprint = null;
-                if (result.reasoning) { openAIResponse.fal_reasoning = result.reasoning; }
+
+                // Construct OpenAI compatible response
+                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 || "" // Ensure content is string
+                         },
+                        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("<-- Non-stream response sent successfully.");
+
             } catch (error) {
-                 console.error('Error during non-stream request processing:', error.message || error);
+                // Catches errors from makeFalRequestWithRetry (e.g., all keys failed or non-key error)
+                console.error('Error during non-stream request processing:', error.message || error);
                 if (!res.headersSent) {
                     const errorMessage = (error instanceof Error) ? error.message : JSON.stringify(error);
                     const finalMessage = errorMessage.includes("No valid Fal AI keys available") || errorMessage.includes("Request failed after trying")
@@ -366,13 +633,15 @@ app.post('/v1/chat/completions', async (req, res) => {
                     if (!res.writableEnded) { res.end(); }
                 }
             }
-        }
+        } // --- End Stream/Non-Stream Logic ---
 
     } catch (error) {
-        // --- (Outer catch block for setup errors remains the same) ---
+        // **Catch block for errors BEFORE Fal request attempt**
+        // (e.g., errors from convertMessagesToFalPrompt, JSON parsing errors)
         console.error('Unhandled error before initiating Fal request (likely setup or input conversion):', error.message || error);
         if (!res.headersSent) {
             const errorMessage = (error instanceof Error) ? error.message : JSON.stringify(error);
+            // Use 500 Internal Server Error for issues within the proxy itself during setup
             res.status(500).json({ error: 'Internal Server Error in Proxy Setup', details: errorMessage });
             console.log("<-- Proxy setup error response sent (500).");
         } else {
@@ -382,6 +651,26 @@ app.post('/v1/chat/completions', async (req, res) => {
     }
 });
 
-// --- (Server listen and root path handler remain the same) ---
-app.listen(PORT, () => { /* ... startup messages ... */ });
-app.get('/', (req, res) => { /* ... root message ... */ });
+// Start the Express server
+app.listen(PORT, () => {
+    console.log(`=====================================================================`);
+    console.log(` Fal OpenAI Proxy Server (Multi-Key Rotation & Failover)`);
+    console.log(`---------------------------------------------------------------------`);
+    console.log(` Listening on port   : ${PORT}`);
+    console.log(` Reading Fal Keys from : FAL_KEY environment variable (comma-separated)`);
+    console.log(` Loaded Keys Count   : ${falKeys.length}`);
+    console.log(` Invalid Keys Set    : Initialized (size: ${invalidKeys.size})`);
+    console.log(` Proxy API Key Auth  : ${API_KEY ? 'Enabled (using API_KEY env var)' : 'DISABLED'}`);
+    console.log(` Input Limits        : System Prompt=${SYSTEM_PROMPT_LIMIT}, Prompt=${PROMPT_LIMIT}`);
+    console.log(` Concurrency Warning : Global Fal client reconfigured per request attempt!`);
+    console.log(`---------------------------------------------------------------------`);
+    console.log(` Endpoints Available:`);
+    console.log(`   POST http://localhost:${PORT}/v1/chat/completions`);
+    console.log(`   GET  http://localhost:${PORT}/v1/models`);
+    console.log(`=====================================================================`);
+});
+
+// Root path handler for basic health check / info
+app.get('/', (req, res) => {
+    res.send(`Fal OpenAI Proxy (Multi-Key Rotation/Failover from FAL_KEY) is running. Loaded ${falKeys.length} key(s). Currently ${invalidKeys.size} key(s) marked as invalid.`);
+});