opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: support multiple keys per AI Bridge provider (#24683)

Browse Source 
## Description

Adds support for configuring multiple API keys per AI Bridge provider. This PR introduces the configuration parsing and validation only; wiring the key pools into the aibridge providers will happen in upstream PRs.

## Changes

Providers now accept a comma-separated list of keys via the `KEYS` env var (or a single key via the existing `KEY` var). The two are mutually exclusive. Bedrock follows the same pattern with `BEDROCK_ACCESS_KEYS` / `BEDROCK_ACCESS_KEY_SECRETS`, with an additional validation that the two slices have matching lengths.

Key validation at startup checks for empty values, duplicates, and a maximum of 5 keys per provider. 

Related to: https://github.com/coder/internal/issues/1445

> [!NOTE]
> Initially generated by Coder Agents, modified and reviewed by @ssncferreira
This commit is contained in:
Susana Ferreira
2026-04-30 09:19:32 +01:00
committed by GitHub
parent 123c8dfc02
commit 101a4082dd
5 changed files with 239 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/server_aibridge_internal_test.go
+112 -16
View File
@@ -40,7 +40,7 @@ func TestReadAIBridgeProvidersFromEnv(t *testing.T) {
{
Type: aibridge.ProviderAnthropic,
Name: "anthropic-zdr",
Key: "sk-ant-xxx",
Keys: []string{"sk-ant-xxx"},
BaseURL: "https://api.anthropic.com/",
DumpDir: "/tmp/aibridge-dump",
},
@@ -99,14 +99,14 @@ func TestReadAIBridgeProvidersFromEnv(t *testing.T) {
},
expected: []codersdk.AIBridgeProviderConfig{
{
Type: aibridge.ProviderAnthropic,
Name: "anthropic-bedrock",
BedrockRegion: "us-west-2",
BedrockAccessKey: "AKID",
BedrockAccessKeySecret: "secret",
BedrockModel: "anthropic.claude-3-sonnet",
BedrockSmallFastModel: "anthropic.claude-3-haiku",
BedrockBaseURL: "https://bedrock.us-west-2.amazonaws.com",
Type: aibridge.ProviderAnthropic,
Name: "anthropic-bedrock",
BedrockRegion: "us-west-2",
BedrockAccessKeys: []string{"AKID"},
BedrockAccessKeySecrets: []string{"secret"},
BedrockModel: "anthropic.claude-3-sonnet",
BedrockSmallFastModel: "anthropic.claude-3-haiku",
BedrockBaseURL: "https://bedrock.us-west-2.amazonaws.com",
},
},
},
@@ -173,7 +173,7 @@ func TestReadAIBridgeProvidersFromEnv(t *testing.T) {
"SOME_OTHER_VAR=hello",
},
expected: []codersdk.AIBridgeProviderConfig{
{Type: aibridge.ProviderOpenAI, Name: aibridge.ProviderOpenAI, Key: "sk-xxx"},
{Type: aibridge.ProviderOpenAI, Name: aibridge.ProviderOpenAI, Keys: []string{"sk-xxx"}},
},
},
{
@@ -188,11 +188,11 @@ func TestReadAIBridgeProvidersFromEnv(t *testing.T) {
},
expected: []codersdk.AIBridgeProviderConfig{
{
Type: aibridge.ProviderAnthropic,
Name: aibridge.ProviderAnthropic,
Key: "sk-ant-xxx",
BedrockAccessKey: "AKID",
BedrockAccessKeySecret: "secret",
Type: aibridge.ProviderAnthropic,
Name: aibridge.ProviderAnthropic,
Keys: []string{"sk-ant-xxx"},
BedrockAccessKeys: []string{"AKID"},
BedrockAccessKeySecrets: []string{"secret"},
},
},
},
@@ -223,6 +223,102 @@ func TestReadAIBridgeProvidersFromEnv(t *testing.T) {
},
errContains: "BEDROCK_ACCESS_KEY_SECRET and BEDROCK_ACCESS_KEY_SECRETS are mutually exclusive",
},
{
name: "CopilotRejectsKey",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=copilot",
"CODER_AIBRIDGE_PROVIDER_0_KEY=sk-xxx",
},
errContains: "KEY/KEYS are not supported for TYPE",
},
{
name: "CopilotRejectsKeys",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=copilot",
"CODER_AIBRIDGE_PROVIDER_0_KEYS=sk-a,sk-b",
},
errContains: "KEY/KEYS are not supported for TYPE",
},
{
name: "MultipleKeysCommaSeparated",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=openai",
"CODER_AIBRIDGE_PROVIDER_0_KEYS=sk-a,sk-b,sk-c",
},
expected: []codersdk.AIBridgeProviderConfig{
{Type: aibridge.ProviderOpenAI, Name: aibridge.ProviderOpenAI, Keys: []string{"sk-a", "sk-b", "sk-c"}},
},
},
{
name: "KeysWhitespaceTrimmed",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=openai",
"CODER_AIBRIDGE_PROVIDER_0_KEYS= sk-a , sk-b ",
},
expected: []codersdk.AIBridgeProviderConfig{
{Type: aibridge.ProviderOpenAI, Name: aibridge.ProviderOpenAI, Keys: []string{"sk-a", "sk-b"}},
},
},
{
name: "KeysEmptyAfterTrim",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=openai",
"CODER_AIBRIDGE_PROVIDER_0_KEYS=sk-a,,sk-b",
},
errContains: "key at index 1 is empty",
},
{
name: "KeysDuplicate",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=openai",
"CODER_AIBRIDGE_PROVIDER_0_KEYS=sk-a,sk-b,sk-a",
},
errContains: "duplicate key at index 2",
},
{
name: "KeysTooMany",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=openai",
"CODER_AIBRIDGE_PROVIDER_0_KEYS=sk-1,sk-2,sk-3,sk-4,sk-5,sk-6",
},
errContains: "too many keys (6), maximum is 5",
},
{
name: "BedrockMultipleKeys",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=anthropic",
"CODER_AIBRIDGE_PROVIDER_0_BEDROCK_REGION=us-west-2",
"CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEYS=AKID1,AKID2",
"CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEY_SECRETS=secret1,secret2",
},
expected: []codersdk.AIBridgeProviderConfig{
{
Type: aibridge.ProviderAnthropic,
Name: aibridge.ProviderAnthropic,
BedrockRegion: "us-west-2",
BedrockAccessKeys: []string{"AKID1", "AKID2"},
BedrockAccessKeySecrets: []string{"secret1", "secret2"},
},
},
},
{
name: "BedrockKeyCountMismatch",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=anthropic",
"CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEYS=AKID1,AKID2",
"CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEY_SECRET=secret1",
},
errContains: "BEDROCK_ACCESS_KEYS count (2) must match BEDROCK_ACCESS_KEY_SECRETS count (1)",
},
{
name: "BedrockKeysTooMany",
env: []string{
"CODER_AIBRIDGE_PROVIDER_0_TYPE=anthropic",
"CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEYS=AKID1,AKID2,AKID3,AKID4,AKID5,AKID6",
"CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEY_SECRETS=s1,s2,s3,s4,s5,s6",
},
errContains: "too many keys (6), maximum is 5",
},
}
for _, tt := range tests {
@@ -256,7 +352,7 @@ func TestReadAIBridgeProvidersFromEnv(t *testing.T) {
expected = append(expected, codersdk.AIBridgeProviderConfig{
Type: aibridge.ProviderOpenAI,
Name: fmt.Sprintf("p%d", i),
Key: fmt.Sprintf("sk-%d", i),
Keys: []string{fmt.Sprintf("sk-%d", i)},
})
}
providers, err := ReadAIBridgeProvidersFromEnv(slogtest.Make(t, nil), env)