diff --git a/cli/server.go b/cli/server.go index 35eec130b4..4c8c41f415 100644 --- a/cli/server.go +++ b/cli/server.go @@ -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 diff --git a/cli/server_aibridge_internal_test.go b/cli/server_aibridge_internal_test.go index 49789f1a83..6ab6bca585 100644 --- a/cli/server_aibridge_internal_test.go +++ b/cli/server_aibridge_internal_test.go @@ -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) diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 8dca6ab102..be2f738ceb 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -4110,20 +4110,27 @@ type AIBridgeProviderConfig struct { // Name is the unique instance identifier used for routing. // Defaults to Type if not provided. Name string `json:"name"` - // Key is the API key for authenticating with the upstream provider. - Key string `json:"-"` + // Keys holds one or more API keys for authenticating with the + // upstream provider. When multiple keys are configured, they + // form a key pool for automatic failover. + Keys []string `json:"-"` // BaseURL is the base URL of the upstream provider API. BaseURL string `json:"base_url"` // DumpDir is the directory path for dumping API requests and responses. DumpDir string `json:"dump_dir,omitempty"` // Bedrock fields (only applicable when Type == "anthropic"). - BedrockBaseURL string `json:"-"` - BedrockRegion string `json:"bedrock_region,omitempty"` - BedrockAccessKey string `json:"-"` - BedrockAccessKeySecret string `json:"-"` - BedrockModel string `json:"bedrock_model,omitempty"` - BedrockSmallFastModel string `json:"bedrock_small_fast_model,omitempty"` + BedrockBaseURL string `json:"-"` + BedrockRegion string `json:"bedrock_region,omitempty"` + // BedrockAccessKeys and BedrockAccessKeySecrets hold one or + // more AWS credential pairs for authenticating with Bedrock. + // When multiple pairs are configured, they form a key pool + // for automatic failover. The two slices must have the same + // length. + BedrockAccessKeys []string `json:"-"` + BedrockAccessKeySecrets []string `json:"-"` + BedrockModel string `json:"bedrock_model,omitempty"` + BedrockSmallFastModel string `json:"bedrock_small_fast_model,omitempty"` } type AIBridgeProxyConfig struct { diff --git a/enterprise/cli/aibridged.go b/enterprise/cli/aibridged.go index a51e0e01a7..8c02db7d55 100644 --- a/enterprise/cli/aibridged.go +++ b/enterprise/cli/aibridged.go @@ -110,12 +110,18 @@ func buildProviders(cfg codersdk.AIBridgeConfig) ([]aibridge.Provider, error) { if name == "" { name = p.Type } + // Currently, only the first key is used, if any. + // TODO(ssncferreira): pass a keypool.Pool instead. + var key string + if len(p.Keys) > 0 { + key = p.Keys[0] + } switch p.Type { case aibridge.ProviderOpenAI: providers = append(providers, aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{ Name: name, BaseURL: p.BaseURL, - Key: p.Key, + Key: key, APIDumpDir: p.DumpDir, CircuitBreaker: cbConfig, SendActorHeaders: cfg.SendActorHeaders.Value(), @@ -124,7 +130,7 @@ func buildProviders(cfg codersdk.AIBridgeConfig) ([]aibridge.Provider, error) { providers = append(providers, aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{ Name: name, BaseURL: p.BaseURL, - Key: p.Key, + Key: key, APIDumpDir: p.DumpDir, CircuitBreaker: cbConfig, SendActorHeaders: cfg.SendActorHeaders.Value(), @@ -148,14 +154,23 @@ func buildProviders(cfg codersdk.AIBridgeConfig) ([]aibridge.Provider, error) { // AIBridgeProviderConfig into an aibridge AWSBedrockConfig. // Returns nil if no Bedrock fields are set. func bedrockConfigFromProvider(p codersdk.AIBridgeProviderConfig) *aibridge.AWSBedrockConfig { - if p.BedrockRegion == "" && p.BedrockBaseURL == "" && p.BedrockAccessKey == "" && p.BedrockAccessKeySecret == "" { + if p.BedrockRegion == "" && p.BedrockBaseURL == "" && len(p.BedrockAccessKeys) == 0 && len(p.BedrockAccessKeySecrets) == 0 { return nil } + // Currently, only the first key pair is used, if any. + // TODO(ssncferreira): pass a keypool.Pool instead. + var accessKey, accessKeySecret string + if len(p.BedrockAccessKeys) > 0 { + accessKey = p.BedrockAccessKeys[0] + } + if len(p.BedrockAccessKeySecrets) > 0 { + accessKeySecret = p.BedrockAccessKeySecrets[0] + } return &aibridge.AWSBedrockConfig{ BaseURL: p.BedrockBaseURL, Region: p.BedrockRegion, - AccessKey: p.BedrockAccessKey, - AccessKeySecret: p.BedrockAccessKeySecret, + AccessKey: accessKey, + AccessKeySecret: accessKeySecret, Model: p.BedrockModel, SmallFastModel: p.BedrockSmallFastModel, } diff --git a/enterprise/cli/aibridged_internal_test.go b/enterprise/cli/aibridged_internal_test.go index e923c106af..0504fdb705 100644 --- a/enterprise/cli/aibridged_internal_test.go +++ b/enterprise/cli/aibridged_internal_test.go @@ -46,13 +46,13 @@ func TestBuildProviders(t *testing.T) { { Type: aibridge.ProviderAnthropic, Name: "anthropic-zdr", - Key: "sk-zdr", + Keys: []string{"sk-zdr"}, DumpDir: "/tmp/anthropic-dump", }, { Type: aibridge.ProviderOpenAI, Name: "openai-azure", - Key: "sk-azure", + Keys: []string{"sk-azure"}, BaseURL: "https://azure.openai.com", DumpDir: "/tmp/openai-dump", }, @@ -72,7 +72,7 @@ func TestBuildProviders(t *testing.T) { t.Parallel() cfg := codersdk.AIBridgeConfig{ Providers: []codersdk.AIBridgeProviderConfig{ - {Type: aibridge.ProviderOpenAI, Name: aibridge.ProviderOpenAI, Key: "sk-indexed"}, + {Type: aibridge.ProviderOpenAI, Name: aibridge.ProviderOpenAI, Keys: []string{"sk-indexed"}}, }, } cfg.LegacyOpenAI.Key = serpent.String("sk-legacy") @@ -86,7 +86,7 @@ func TestBuildProviders(t *testing.T) { t.Parallel() cfg := codersdk.AIBridgeConfig{ Providers: []codersdk.AIBridgeProviderConfig{ - {Type: aibridge.ProviderAnthropic, Name: aibridge.ProviderAnthropic, Key: "sk-indexed"}, + {Type: aibridge.ProviderAnthropic, Name: aibridge.ProviderAnthropic, Keys: []string{"sk-indexed"}}, }, } cfg.LegacyAnthropic.Key = serpent.String("sk-legacy") @@ -100,7 +100,7 @@ func TestBuildProviders(t *testing.T) { t.Parallel() cfg := codersdk.AIBridgeConfig{ Providers: []codersdk.AIBridgeProviderConfig{ - {Type: aibridge.ProviderAnthropic, Name: "anthropic-zdr", Key: "sk-zdr"}, + {Type: aibridge.ProviderAnthropic, Name: "anthropic-zdr", Keys: []string{"sk-zdr"}}, }, } cfg.LegacyOpenAI.Key = serpent.String("sk-openai") @@ -220,9 +220,9 @@ func TestDomainsFromProviders(t *testing.T) { providers, err := buildProviders(codersdk.AIBridgeConfig{ Providers: []codersdk.AIBridgeProviderConfig{ - {Type: aibridge.ProviderOpenAI, Name: "openai", Key: "k"}, - {Type: aibridge.ProviderAnthropic, Name: "anthropic", Key: "k"}, - {Type: aibridge.ProviderOpenAI, Name: "custom", Key: "k", BaseURL: "https://custom-llm.example.com:8443/api"}, + {Type: aibridge.ProviderOpenAI, Name: "openai", Keys: []string{"k"}}, + {Type: aibridge.ProviderAnthropic, Name: "anthropic", Keys: []string{"k"}}, + {Type: aibridge.ProviderOpenAI, Name: "custom", Keys: []string{"k"}, BaseURL: "https://custom-llm.example.com:8443/api"}, }, }) require.NoError(t, err) @@ -244,8 +244,8 @@ func TestDomainsFromProviders(t *testing.T) { providers, err := buildProviders(codersdk.AIBridgeConfig{ Providers: []codersdk.AIBridgeProviderConfig{ - {Type: aibridge.ProviderOpenAI, Name: "first", Key: "k", BaseURL: "https://api.example.com/v1"}, - {Type: aibridge.ProviderOpenAI, Name: "second", Key: "k", BaseURL: "https://api.example.com/v2"}, + {Type: aibridge.ProviderOpenAI, Name: "first", Keys: []string{"k"}, BaseURL: "https://api.example.com/v1"}, + {Type: aibridge.ProviderOpenAI, Name: "second", Keys: []string{"k"}, BaseURL: "https://api.example.com/v2"}, }, }) require.NoError(t, err) @@ -269,7 +269,7 @@ func TestDomainsFromProviders(t *testing.T) { providers, err := buildProviders(codersdk.AIBridgeConfig{ Providers: []codersdk.AIBridgeProviderConfig{ - {Type: aibridge.ProviderOpenAI, Name: "provider", Key: "k", BaseURL: "https://API.Example.COM/v1"}, + {Type: aibridge.ProviderOpenAI, Name: "provider", Keys: []string{"k"}, BaseURL: "https://API.Example.COM/v1"}, }, }) require.NoError(t, err)