mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
feat: support multiple keys per AI Bridge provider (#24683)
## 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:
+81
-7
@@ -2975,10 +2975,14 @@ func ReadAIBridgeProvidersFromEnv(logger slog.Logger, environ []string) ([]coder
|
||||
case "NAME":
|
||||
provider.Name = v.Value
|
||||
case "KEY", "KEYS":
|
||||
if provider.Key != "" {
|
||||
if len(provider.Keys) > 0 {
|
||||
return nil, xerrors.Errorf("provider %d: KEY and KEYS are mutually exclusive, use one or the other", providerNum)
|
||||
}
|
||||
provider.Key = v.Value
|
||||
if key == "KEYS" {
|
||||
provider.Keys = strings.Split(v.Value, ",")
|
||||
} else {
|
||||
provider.Keys = []string{v.Value}
|
||||
}
|
||||
case "BASE_URL":
|
||||
provider.BaseURL = v.Value
|
||||
case "DUMP_DIR":
|
||||
@@ -2988,15 +2992,23 @@ func ReadAIBridgeProvidersFromEnv(logger slog.Logger, environ []string) ([]coder
|
||||
case "BEDROCK_REGION":
|
||||
provider.BedrockRegion = v.Value
|
||||
case "BEDROCK_ACCESS_KEY", "BEDROCK_ACCESS_KEYS":
|
||||
if provider.BedrockAccessKey != "" {
|
||||
if len(provider.BedrockAccessKeys) > 0 {
|
||||
return nil, xerrors.Errorf("provider %d: BEDROCK_ACCESS_KEY and BEDROCK_ACCESS_KEYS are mutually exclusive, use one or the other", providerNum)
|
||||
}
|
||||
provider.BedrockAccessKey = v.Value
|
||||
if key == "BEDROCK_ACCESS_KEYS" {
|
||||
provider.BedrockAccessKeys = strings.Split(v.Value, ",")
|
||||
} else {
|
||||
provider.BedrockAccessKeys = []string{v.Value}
|
||||
}
|
||||
case "BEDROCK_ACCESS_KEY_SECRET", "BEDROCK_ACCESS_KEY_SECRETS":
|
||||
if provider.BedrockAccessKeySecret != "" {
|
||||
if len(provider.BedrockAccessKeySecrets) > 0 {
|
||||
return nil, xerrors.Errorf("provider %d: BEDROCK_ACCESS_KEY_SECRET and BEDROCK_ACCESS_KEY_SECRETS are mutually exclusive, use one or the other", providerNum)
|
||||
}
|
||||
provider.BedrockAccessKeySecret = v.Value
|
||||
if key == "BEDROCK_ACCESS_KEY_SECRETS" {
|
||||
provider.BedrockAccessKeySecrets = strings.Split(v.Value, ",")
|
||||
} else {
|
||||
provider.BedrockAccessKeySecrets = []string{v.Value}
|
||||
}
|
||||
case "BEDROCK_MODEL":
|
||||
provider.BedrockModel = v.Value
|
||||
case "BEDROCK_SMALL_FAST_MODEL":
|
||||
@@ -3029,6 +3041,19 @@ func ReadAIBridgeProvidersFromEnv(logger slog.Logger, environ []string) ([]coder
|
||||
i, p.Type, aibridge.ProviderAnthropic)
|
||||
}
|
||||
|
||||
if p.Type == aibridge.ProviderCopilot && len(p.Keys) > 0 {
|
||||
return nil, xerrors.Errorf("provider %d (%s): KEY/KEYS are not supported for TYPE %q",
|
||||
i, p.Type, aibridge.ProviderCopilot)
|
||||
}
|
||||
|
||||
if err := validateProviderCredentialList(i, p.Type, p.Keys); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := validateBedrockCredentials(i, p.Type, p.BedrockAccessKeys, p.BedrockAccessKeySecrets); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if p.Name == "" {
|
||||
p.Name = p.Type
|
||||
}
|
||||
@@ -3043,10 +3068,59 @@ func ReadAIBridgeProvidersFromEnv(logger slog.Logger, environ []string) ([]coder
|
||||
|
||||
func hasBedrockFields(p codersdk.AIBridgeProviderConfig) bool {
|
||||
return p.BedrockBaseURL != "" || p.BedrockRegion != "" ||
|
||||
p.BedrockAccessKey != "" || p.BedrockAccessKeySecret != "" ||
|
||||
len(p.BedrockAccessKeys) > 0 || len(p.BedrockAccessKeySecrets) > 0 ||
|
||||
p.BedrockModel != "" || p.BedrockSmallFastModel != ""
|
||||
}
|
||||
|
||||
// maxKeysPerProvider is the maximum number of keys allowed per
|
||||
// provider. This bounds the failover pool size and keeps the
|
||||
// configuration manageable.
|
||||
const maxKeysPerProvider = 5
|
||||
|
||||
// validateProviderCredentialList checks that a list of credentials
|
||||
// belonging to a provider is well-formed: no empty values, no
|
||||
// duplicates, and within the maximum count. Trims whitespace in
|
||||
// place.
|
||||
func validateProviderCredentialList(providerIndex int, providerType string, keys []string) error {
|
||||
if len(keys) > maxKeysPerProvider {
|
||||
return xerrors.Errorf("provider %d (%s): too many keys (%d), maximum is %d",
|
||||
providerIndex, providerType, len(keys), maxKeysPerProvider)
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{}, len(keys))
|
||||
for i, key := range keys {
|
||||
trimmed := strings.TrimSpace(key)
|
||||
if trimmed == "" {
|
||||
return xerrors.Errorf("provider %d (%s): key at index %d is empty",
|
||||
providerIndex, providerType, i)
|
||||
}
|
||||
keys[i] = trimmed
|
||||
if _, exists := seen[trimmed]; exists {
|
||||
return xerrors.Errorf("provider %d (%s): duplicate key at index %d",
|
||||
providerIndex, providerType, i)
|
||||
}
|
||||
seen[trimmed] = struct{}{}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateBedrockCredentials checks that Bedrock access keys and
|
||||
// secrets are paired correctly (same count) and that each list is
|
||||
// well-formed.
|
||||
func validateBedrockCredentials(providerIndex int, providerType string, accessKeys, secrets []string) error {
|
||||
if len(accessKeys) != len(secrets) {
|
||||
return xerrors.Errorf("provider %d (%s): BEDROCK_ACCESS_KEYS count (%d) must match BEDROCK_ACCESS_KEY_SECRETS count (%d)",
|
||||
providerIndex, providerType, len(accessKeys), len(secrets))
|
||||
}
|
||||
|
||||
if err := validateProviderCredentialList(providerIndex, providerType, accessKeys); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return validateProviderCredentialList(providerIndex, providerType, secrets)
|
||||
}
|
||||
|
||||
var reInvalidPortAfterHost = regexp.MustCompile(`invalid port ".+" after host`)
|
||||
|
||||
// If the user provides a postgres URL with a password that contains special
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user