From 914a0f78300d2934da47c228e265d1892c908f4b Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Thu, 16 Apr 2026 14:08:23 +0200 Subject: [PATCH] chore: follow-ups from #23948 (#24377) A couple follow-ups from #23948 --------- Signed-off-by: Danny Kopping --- cli/server.go | 21 +- cli/server_aibridge_internal_test.go | 27 ++ enterprise/aibridgeproxyd/aibridgeproxyd.go | 36 +- .../aibridgeproxyd/aibridgeproxyd_test.go | 310 ++++++++++-------- 4 files changed, 221 insertions(+), 173 deletions(-) diff --git a/cli/server.go b/cli/server.go index 6bf924bcce..7054bd42cb 100644 --- a/cli/server.go +++ b/cli/server.go @@ -2961,9 +2961,10 @@ func ReadAIBridgeProvidersFromEnv(logger slog.Logger, environ []string) ([]coder provider.Type = v.Value case "NAME": provider.Name = v.Value - case "KEY": // Alias for a single key. - provider.Key = v.Value - case "KEYS": + case "KEY", "KEYS": + if provider.Key != "" { + return nil, xerrors.Errorf("provider %d: KEY and KEYS are mutually exclusive, use one or the other", providerNum) + } provider.Key = v.Value case "BASE_URL": provider.BaseURL = v.Value @@ -2971,13 +2972,15 @@ func ReadAIBridgeProvidersFromEnv(logger slog.Logger, environ []string) ([]coder provider.BedrockBaseURL = v.Value case "BEDROCK_REGION": provider.BedrockRegion = v.Value - case "BEDROCK_ACCESS_KEY": // Alias for a single key. + case "BEDROCK_ACCESS_KEY", "BEDROCK_ACCESS_KEYS": + if provider.BedrockAccessKey != "" { + 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 - case "BEDROCK_ACCESS_KEYS": - provider.BedrockAccessKey = v.Value - case "BEDROCK_ACCESS_KEY_SECRET": // Alias for a single key secret. - provider.BedrockAccessKeySecret = v.Value - case "BEDROCK_ACCESS_KEY_SECRETS": + case "BEDROCK_ACCESS_KEY_SECRET", "BEDROCK_ACCESS_KEY_SECRETS": + if provider.BedrockAccessKeySecret != "" { + 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 case "BEDROCK_MODEL": provider.BedrockModel = v.Value diff --git a/cli/server_aibridge_internal_test.go b/cli/server_aibridge_internal_test.go index 35e9e6400e..bfa3d2d765 100644 --- a/cli/server_aibridge_internal_test.go +++ b/cli/server_aibridge_internal_test.go @@ -194,6 +194,33 @@ func TestReadAIBridgeProvidersFromEnv(t *testing.T) { }, }, }, + { + name: "ConflictKeyAndKeys", + env: []string{ + "CODER_AIBRIDGE_PROVIDER_0_TYPE=openai", + "CODER_AIBRIDGE_PROVIDER_0_KEY=sk-single", + "CODER_AIBRIDGE_PROVIDER_0_KEYS=sk-multi", + }, + errContains: "KEY and KEYS are mutually exclusive", + }, + { + name: "ConflictBedrockAccessKeyAndKeys", + env: []string{ + "CODER_AIBRIDGE_PROVIDER_0_TYPE=anthropic", + "CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEY=AKID1", + "CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEYS=AKID2", + }, + errContains: "BEDROCK_ACCESS_KEY and BEDROCK_ACCESS_KEYS are mutually exclusive", + }, + { + name: "ConflictBedrockSecretAndSecrets", + env: []string{ + "CODER_AIBRIDGE_PROVIDER_0_TYPE=anthropic", + "CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEY_SECRET=s1", + "CODER_AIBRIDGE_PROVIDER_0_BEDROCK_ACCESS_KEY_SECRETS=s2", + }, + errContains: "BEDROCK_ACCESS_KEY_SECRET and BEDROCK_ACCESS_KEY_SECRETS are mutually exclusive", + }, } for _, tt := range tests { diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 720ebaff6b..85e9d4ad48 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -27,7 +27,6 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog/v3" - "github.com/coder/aibridge" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" ) @@ -175,8 +174,9 @@ type Options struct { // Only requests to these domains will be MITM'd and forwarded to aibridged. // Requests to other domains will be tunneled directly without decryption. DomainAllowlist []string - // AIBridgeProviderFromHost maps a hostname to a known aibridge provider name. - // If nil, the default provider mapping is used. + // AIBridgeProviderFromHost maps a hostname to a known aibridge provider + // name. Must be non-nil; the caller derives it from the configured + // provider list. AIBridgeProviderFromHost func(host string) string // UpstreamProxy is the URL of an upstream HTTP proxy to chain tunneled // (non-allowlisted) requests through. If empty, tunneled requests connect @@ -251,11 +251,10 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) return nil, xerrors.New("domain allowlist is empty, at least one domain is required") } - // Use custom provider mapper if provided, otherwise use default. - aibridgeProviderFromHost := opts.AIBridgeProviderFromHost - if aibridgeProviderFromHost == nil { - aibridgeProviderFromHost = defaultAIBridgeProvider + if opts.AIBridgeProviderFromHost == nil { + return nil, xerrors.New("AIBridgeProviderFromHost is required") } + aibridgeProviderFromHost := opts.AIBridgeProviderFromHost // Validate that all allowlisted domains have correct aibridge provider mappings. for _, domain := range opts.DomainAllowlist { @@ -764,29 +763,6 @@ func newProxyAuthRequiredResponse(req *http.Request) *http.Response { } } -// defaultAIBridgeProvider maps the request host to the aibridge provider name. -// - Known AI providers return their provider name, used to route to the -// corresponding aibridge endpoint. -// - Unknown hosts return empty string and are passed through directly. -func defaultAIBridgeProvider(host string) string { - switch strings.ToLower(host) { - case HostAnthropic: - return aibridge.ProviderAnthropic - case HostOpenAI: - return aibridge.ProviderOpenAI - case HostCopilot: - return aibridge.ProviderCopilot - case agplaibridge.HostCopilotBusiness: - return agplaibridge.ProviderCopilotBusiness - case agplaibridge.HostCopilotEnterprise: - return agplaibridge.ProviderCopilotEnterprise - case agplaibridge.HostChatGPT: - return agplaibridge.ProviderChatGPT - default: - return "" - } -} - // tunneledMiddleware is a CONNECT middleware that handles tunneled (non-allowlisted) // connections. These connections are not MITM'd and are tunneled directly to their // destination. This middleware records metrics for tunneled CONNECT sessions. diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 58594d6163..0f966f3ff0 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -31,6 +31,7 @@ import ( "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/aibridge" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/enterprise/aibridgeproxyd" "github.com/coder/coder/v2/testutil" @@ -189,6 +190,27 @@ func withAIBridgeProviderFromHost(fn func(string) string) testProxyOption { } } +// testProviderFromHost maps well-known AI provider hostnames to +// provider names for test use. Unknown hosts return "". +func testProviderFromHost(host string) string { + switch strings.ToLower(host) { + case aibridgeproxyd.HostAnthropic: + return aibridge.ProviderAnthropic + case aibridgeproxyd.HostOpenAI: + return aibridge.ProviderOpenAI + case aibridgeproxyd.HostCopilot: + return aibridge.ProviderCopilot + case agplaibridge.HostCopilotBusiness: + return agplaibridge.ProviderCopilotBusiness + case agplaibridge.HostCopilotEnterprise: + return agplaibridge.ProviderCopilotEnterprise + case agplaibridge.HostChatGPT: + return agplaibridge.ProviderChatGPT + default: + return "" + } +} + func withUpstreamProxy(upstreamProxy string) testProxyOption { return func(cfg *testProxyConfig) { cfg.upstreamProxy = upstreamProxy @@ -474,13 +496,14 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - TLSCertFile: "/nonexistent/cert.pem", - TLSKeyFile: "/nonexistent/key.pem", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + TLSCertFile: "/nonexistent/cert.pem", + TLSKeyFile: "/nonexistent/key.pem", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.Error(t, err) require.Contains(t, err.Error(), "load listener TLS certificate") @@ -543,11 +566,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.NoError(t, err) require.Equal(t, "localhost", srv.CoderAccessURL().Hostname()) @@ -561,11 +585,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "https://localhost", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "https://localhost", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.NoError(t, err) require.Equal(t, "localhost", srv.CoderAccessURL().Hostname()) @@ -579,11 +604,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.NoError(t, err) require.Equal(t, "localhost", srv.CoderAccessURL().Hostname()) @@ -626,11 +652,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: ":0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: "/nonexistent/cert.pem", - MITMKeyFile: "/nonexistent/key.pem", - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: ":0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: "/nonexistent/cert.pem", + MITMKeyFile: "/nonexistent/key.pem", + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.Error(t, err) require.Contains(t, err.Error(), "failed to load MITM certificate") @@ -710,11 +737,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{"unknown.example.com"}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{"unknown.example.com"}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.Error(t, err) require.Contains(t, err.Error(), `domain "unknown.example.com" is in allowlist but has no provider mapping`) @@ -727,12 +755,13 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "://invalid-url", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "://invalid-url", }) require.Error(t, err) require.Contains(t, err.Error(), "invalid upstream proxy URL") @@ -745,13 +774,14 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "https://proxy.example.com:8080", - UpstreamProxyCA: "/nonexistent/ca.pem", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "https://proxy.example.com:8080", + UpstreamProxyCA: "/nonexistent/ca.pem", }) require.Error(t, err) require.Contains(t, err.Error(), "failed to read upstream proxy CA certificate") @@ -764,12 +794,13 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "http://:@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "http://:@proxy.example.com:8080", }) require.Error(t, err) require.Contains(t, err.Error(), "invalid credentials: both username and password are empty") @@ -782,12 +813,13 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) _, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AllowedPrivateCIDRs: []string{"not-a-cidr"}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + AllowedPrivateCIDRs: []string{"not-a-cidr"}, }) require.Error(t, err) require.Contains(t, err.Error(), "invalid allowed private CIDR") @@ -800,11 +832,12 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.NoError(t, err) require.NotNil(t, srv) @@ -818,13 +851,14 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - TLSCertFile: listenerCertFile, - TLSKeyFile: listenerKeyFile, - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + TLSCertFile: listenerCertFile, + TLSKeyFile: listenerKeyFile, + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.NoError(t, err) require.NotNil(t, srv) @@ -837,12 +871,13 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "http://proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "http://proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -856,13 +891,14 @@ func TestNew(t *testing.T) { // Use the shared MITM certificate as the upstream proxy CA (it's a valid PEM cert) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "https://proxy.example.com:8080", - UpstreamProxyCA: mitmCertFile, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "https://proxy.example.com:8080", + UpstreamProxyCA: mitmCertFile, }) require.NoError(t, err) require.NotNil(t, srv) @@ -875,12 +911,13 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "http://proxyuser:proxypass@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "http://proxyuser:proxypass@proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -893,12 +930,13 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "http://proxyuser:@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "http://proxyuser:@proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -912,12 +950,13 @@ func TestNew(t *testing.T) { // Username only (no colon) should also succeed (password is optional) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "http://proxyuser@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "http://proxyuser@proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -930,12 +969,13 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - UpstreamProxy: "http://:proxypass@proxy.example.com:8080", + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + UpstreamProxy: "http://:proxypass@proxy.example.com:8080", }) require.NoError(t, err) require.NotNil(t, srv) @@ -952,12 +992,13 @@ func TestNew(t *testing.T) { metrics := aibridgeproxyd.NewMetrics(reg) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - Metrics: metrics, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + Metrics: metrics, }) require.NoError(t, err) require.NotNil(t, srv) @@ -970,12 +1011,13 @@ func TestNew(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - AllowedPrivateCIDRs: []string{"127.0.0.1/32"}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + AllowedPrivateCIDRs: []string{"127.0.0.1/32"}, }) require.NoError(t, err) require.NotNil(t, srv) @@ -992,11 +1034,12 @@ func TestClose(t *testing.T) { logger := slogtest.Make(t, nil) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, }) require.NoError(t, err) @@ -1019,12 +1062,13 @@ func TestClose(t *testing.T) { metrics := aibridgeproxyd.NewMetrics(reg) srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - MITMCertFile: mitmCertFile, - MITMKeyFile: mitmKeyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, - Metrics: metrics, + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + MITMCertFile: mitmCertFile, + MITMKeyFile: mitmKeyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + AIBridgeProviderFromHost: testProviderFromHost, + Metrics: metrics, }) require.NoError(t, err) @@ -1406,8 +1450,7 @@ func TestProxy_MITM(t *testing.T) { withCoderAccessURL(aibridgedServer.URL), withAllowedPorts(allowedPorts...), withDomainAllowlist(domainAllowlist...), - // Use default provider mapping to test real AI provider routing. - withAIBridgeProviderFromHost(nil), + withAIBridgeProviderFromHost(testProviderFromHost), withMetrics(metrics), ) @@ -1546,7 +1589,7 @@ func TestProxy_MITM_BYOKInjection(t *testing.T) { srv := newTestProxy(t, withCoderAccessURL(aibridgedServer.URL), withDomainAllowlist(aibridgeproxyd.HostCopilot), - withAIBridgeProviderFromHost(nil), + withAIBridgeProviderFromHost(testProviderFromHost), ) certPool := getProxyCertPool(t) @@ -2014,8 +2057,7 @@ func TestUpstreamProxy(t *testing.T) { withDomainAllowlist(domainAllowlist...), withUpstreamProxy(upstreamProxyURLStr), withAllowedPorts("80", "443", parsedTargetURL.Port()), - // Use default provider mapping to test real AI provider routing. - withAIBridgeProviderFromHost(nil), + withAIBridgeProviderFromHost(testProviderFromHost), } if upstreamProxyCAFile != "" { proxyOpts = append(proxyOpts, withUpstreamProxyCA(upstreamProxyCAFile))