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
This commit is contained in:
Dayuan Jiang
2026-02-02 15:54:35 +09:00
committed by GitHub
parent 4624ad40a1
commit cd33e131ef
5 changed files with 143 additions and 14 deletions

View File

@@ -45,6 +45,72 @@ describe("ServerModelsConfigSchema", () => {
ServerModelsConfigSchema.parse(invalidConfig as any),
).toThrow()
})
it("accepts apiKeyEnv as single string", () => {
const config: ServerModelsConfig = {
providers: [
{
name: "OpenAI Server",
provider: "openai",
models: ["gpt-4o"],
apiKeyEnv: "OPENAI_API_KEY_TEAM_A",
},
],
}
const parsed = ServerModelsConfigSchema.parse(config)
expect(parsed.providers[0].apiKeyEnv).toBe("OPENAI_API_KEY_TEAM_A")
})
it("accepts apiKeyEnv as array of strings for load balancing", () => {
const config: ServerModelsConfig = {
providers: [
{
name: "OpenAI Server",
provider: "openai",
models: ["gpt-4o"],
apiKeyEnv: ["OPENAI_KEY_1", "OPENAI_KEY_2", "OPENAI_KEY_3"],
},
],
}
const parsed = ServerModelsConfigSchema.parse(config)
expect(parsed.providers[0].apiKeyEnv).toEqual([
"OPENAI_KEY_1",
"OPENAI_KEY_2",
"OPENAI_KEY_3",
])
})
it("rejects empty array for apiKeyEnv", () => {
const config = {
providers: [
{
name: "OpenAI Server",
provider: "openai",
models: ["gpt-4o"],
apiKeyEnv: [],
},
],
}
expect(() => ServerModelsConfigSchema.parse(config)).toThrow()
})
it("rejects empty string in apiKeyEnv array", () => {
const config = {
providers: [
{
name: "OpenAI Server",
provider: "openai",
models: ["gpt-4o"],
apiKeyEnv: ["VALID_KEY", ""],
},
],
}
expect(() => ServerModelsConfigSchema.parse(config)).toThrow()
})
})
describe("loadFlattenedServerModels", () => {
@@ -82,4 +148,24 @@ describe("loadFlattenedServerModels", () => {
expect(defaultModel.provider).toBe("openai")
expect(defaultModel.modelId).toBe("gpt-4o") // First model of default provider
})
it("preserves apiKeyEnv array in flattened models for load balancing", async () => {
const config: ServerModelsConfig = {
providers: [
{
name: "OpenAI LoadBalanced",
provider: "openai",
models: ["gpt-4o"],
apiKeyEnv: ["OPENAI_KEY_1", "OPENAI_KEY_2"],
},
],
}
process.env.AI_MODELS_CONFIG = JSON.stringify(config)
process.env.AI_MODELS_CONFIG_PATH = "" // Clear file path
const models = await loadFlattenedServerModels()
expect(models.length).toBe(1)
expect(models[0].apiKeyEnv).toEqual(["OPENAI_KEY_1", "OPENAI_KEY_2"])
})
})