feat: add API key load balancing for providers (#676)

Support multiple API keys per provider with random selection for load
balancing. When AI_MODELS_CONFIG has multiple apiKeyEnv values for
a provider, requests will randomly select one available key.

- Update schema to accept apiKeyEnv as string or string array
- Add random key selection in resolveApiKey()
- Update validation to check at least one key exists
- Add tests for array format support

lib/ai-providers.ts

@@ -34,8 +34,9 @@ export interface ClientOverrides {
     vertexApiKey?: string | null // Express Mode API key
     // Custom headers (e.g., for EdgeOne cookie auth)
     headers?: Record<string, string>
-    // Custom env var names for server models (allows multiple API keys per provider)
-    apiKeyEnv?: string
+    // Custom env var name(s) for server models
+    // Can be a single string or array of strings for load balancing
+    apiKeyEnv?: string | string[]
     baseUrlEnv?: string
 }
 
@@ -99,10 +100,12 @@ export function resolveBaseURL(
 /**
  * Resolve API key from custom env var name or default env var.
  * Supports multiple API keys per provider via ai-models.json apiKeyEnv config.
+ * When multiple keys are configured, randomly selects one for load balancing.
  *
  * Priority:
  * 1. User-provided API key (overrides.apiKey)
- * 2. Custom env var from ai-models.json (overrides.apiKeyEnv)
+ * 2. Custom env var(s) from ai-models.json (overrides.apiKeyEnv)
+ *    - If array, randomly picks one with a valid value
  * 3. Default provider env var (defaultEnvVar)
  */
 function resolveApiKey(
@@ -110,7 +113,30 @@ function resolveApiKey(
     defaultEnvVar: string,
 ): string | undefined {
     if (overrides?.apiKey) return overrides.apiKey
-    if (overrides?.apiKeyEnv) return process.env[overrides.apiKeyEnv]
+
+    if (overrides?.apiKeyEnv) {
+        // Handle array of env var names - randomly select one
+        if (Array.isArray(overrides.apiKeyEnv)) {
+            // Filter to only env vars that have values
+            const validEnvVars = overrides.apiKeyEnv.filter(
+                (envVar) => process.env[envVar],
+            )
+            if (validEnvVars.length > 0) {
+                // Randomly select one
+                const selectedEnvVar =
+                    validEnvVars[
+                        Math.floor(Math.random() * validEnvVars.length)
+                    ]
+                console.log(
+                    `[API Key Routing] Selected ${selectedEnvVar} from ${validEnvVars.length} available keys`,
+                )
+                return process.env[selectedEnvVar]
+            }
+        } else {
+            return process.env[overrides.apiKeyEnv]
+        }
+    }
+
     return process.env[defaultEnvVar]
 }
 
@@ -516,12 +542,24 @@ function detectProvider(): ProviderName | null {
 /**
  * Validate that required API keys are present for the selected provider
  * @param provider - The provider to validate
- * @param customApiKeyEnv - Optional custom env var name (from ai-models.json apiKeyEnv)
+ * @param customApiKeyEnv - Optional custom env var name(s) (from ai-models.json apiKeyEnv)
  */
 function validateProviderCredentials(
     provider: ProviderName,
-    customApiKeyEnv?: string,
+    customApiKeyEnv?: string | string[],
 ): void {
+    // Handle array of env var names - at least one must be set
+    if (Array.isArray(customApiKeyEnv)) {
+        const hasAnyKey = customApiKeyEnv.some((envVar) => process.env[envVar])
+        if (!hasAnyKey) {
+            throw new Error(
+                `At least one of [${customApiKeyEnv.join(", ")}] environment variables is required for ${provider} provider. ` +
+                    `Please set at least one in your .env.local file.`,
+            )
+        }
+        return
+    }
+
     // Use custom env var name if provided, otherwise use default
     const requiredVar = customApiKeyEnv || PROVIDER_ENV_VARS[provider]
     if (requiredVar && !process.env[requiredVar]) {

lib/server-model-config.ts

@@ -14,9 +14,12 @@ export const ServerProviderSchema = z.object({
     name: z.string().min(1),
     provider: ProviderNameSchema,
     models: z.array(z.string().min(1)),
-    // Optional: custom environment variable name for API key
-    // e.g., "OPENAI_API_KEY_TEAM_A" instead of default "OPENAI_API_KEY"
-    apiKeyEnv: z.string().min(1).optional(),
+    // Optional: custom environment variable name(s) for API key
+    // Can be a single string or array of strings for load balancing
+    // e.g., "OPENAI_API_KEY_TEAM_A" or ["OPENAI_KEY_1", "OPENAI_KEY_2"]
+    apiKeyEnv: z
+        .union([z.string().min(1), z.array(z.string().min(1)).min(1)])
+        .optional(),
     // Optional: custom environment variable name for base URL
     baseUrlEnv: z.string().min(1).optional(),
     // Optional: mark the first model in this provider as the default
@@ -36,8 +39,9 @@ export interface FlattenedServerModel {
     provider: ProviderName
     providerLabel: string
     isDefault: boolean
-    // Custom env var names for credentials (optional)
-    apiKeyEnv?: string
+    // Custom env var name(s) for API key (optional)
+    // Can be a single string or array of strings for load balancing
+    apiKeyEnv?: string | string[]
     baseUrlEnv?: string
 }
 

lib/types/model-config.ts

@@ -73,8 +73,9 @@ export interface FlattenedModel {
     source?: "user" | "server"
     // Whether this model is the server default (matches AI_MODEL env var)
     isDefault?: boolean
-    // Custom env var names for server models (allows multiple API keys per provider)
-    apiKeyEnv?: string
+    // Custom env var name(s) for server models
+    // Can be a single string or array of strings for load balancing
+    apiKeyEnv?: string | string[]
     baseUrlEnv?: string
 }