From 473dbb0624c80bdcacb2e91886e835cf6bdc316b Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 26 Aug 2026 12:54:38 +0000 Subject: [PATCH 1/5] fix: enable Copilot HTTP transport fallback --- aibridge/bridge.go | 20 ++ aibridge/bridge_internal_test.go | 30 +++ aibridge/bridge_test.go | 63 ++++- aibridge/provider/copilot.go | 2 + enterprise/aibridgeproxyd/aibridgeproxyd.go | 51 +++- .../aibridgeproxyd/aibridgeproxyd_test.go | 247 ++++++++++++++++-- enterprise/aibridgeproxyd/reload.go | 11 +- .../aibridgeproxyd/reload_internal_test.go | 24 +- 8 files changed, 392 insertions(+), 56 deletions(-) diff --git a/aibridge/bridge.go b/aibridge/bridge.go index 879f0beb488..2331cd85eeb 100644 --- a/aibridge/bridge.go +++ b/aibridge/bridge.go @@ -18,6 +18,7 @@ import ( "github.com/sony/gobreaker/v2" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/trace" + "golang.org/x/net/http/httpguts" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -249,6 +250,18 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC client := GuessClient(r) sessionID := GuessSessionID(client, r) + if isWebSocketUpgrade(r) { + route := strings.TrimPrefix(r.URL.Path, fmt.Sprintf("/%s", p.Name())) + logger.Debug(ctx, "rejecting unsupported WebSocket upgrade", + slog.F("provider", p.Name()), + slog.F("route", route), + slog.F("client", string(client)), + slog.F("client_session_id", sessionID), + ) + http.Error(w, "WebSocket transport is not supported, use HTTP", http.StatusNotImplemented) + return + } + // Read and validate Agent Firewall correlation headers. The // values are captured here and recorded below; the headers // themselves are stripped from the upstream request by @@ -382,6 +395,13 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC } } +// isWebSocketUpgrade reports whether r is a WebSocket opening handshake. +func isWebSocketUpgrade(r *http.Request) bool { + return r.Method == http.MethodGet && + httpguts.HeaderValuesContainsToken(r.Header.Values("Connection"), "upgrade") && + httpguts.HeaderValuesContainsToken(r.Header.Values("Upgrade"), "websocket") +} + // writeRequestBodyTooLarge writes a human-readable 413 response indicating that // the request body exceeded maxRequestBodyBytes. // diff --git a/aibridge/bridge_internal_test.go b/aibridge/bridge_internal_test.go index 561f758de12..e92d554ec59 100644 --- a/aibridge/bridge_internal_test.go +++ b/aibridge/bridge_internal_test.go @@ -10,6 +10,36 @@ import ( agplaibridge "github.com/coder/coder/v2/coderd/aibridge" ) +func TestIsWebSocketUpgrade(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + method string + connection string + upgrade string + want bool + }{ + {name: "websocket upgrade", method: http.MethodGet, connection: "keep-alive, Upgrade", upgrade: "WebSocket", want: true}, + {name: "non-GET request", method: http.MethodPost, connection: "Upgrade", upgrade: "websocket", want: false}, + {name: "missing connection upgrade", method: http.MethodGet, connection: "keep-alive", upgrade: "websocket", want: false}, + {name: "different upgrade protocol", method: http.MethodGet, connection: "Upgrade", upgrade: "h2c", want: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req, err := http.NewRequestWithContext(t.Context(), tc.method, "/", nil) + require.NoError(t, err) + req.Header.Set("Connection", tc.connection) + req.Header.Set("Upgrade", tc.upgrade) + + assert.Equal(t, tc.want, isWebSocketUpgrade(req)) + }) + } +} + func TestExtractAgentFirewallHeaders(t *testing.T) { t.Parallel() diff --git a/aibridge/bridge_test.go b/aibridge/bridge_test.go index f40a06eecec..b1339d89cdc 100644 --- a/aibridge/bridge_test.go +++ b/aibridge/bridge_test.go @@ -13,11 +13,13 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/trace" "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/aibridge" "github.com/coder/coder/v2/aibridge/aibridgetest" "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/internal/testutil" "github.com/coder/coder/v2/aibridge/provider" "github.com/coder/coder/v2/coderd/httpapi" @@ -187,11 +189,12 @@ func TestPassthroughRoutesForProviders(t *testing.T) { upstreamRespBody := "upstream response" tests := []struct { - name string - baseURLPath string - requestPath string - provider func(*testing.T, string) provider.Provider - expectPath string + name string + baseURLPath string + requestMethod string + requestPath string + provider func(*testing.T, string) provider.Provider + expectPath string }{ { name: "openAI_no_base_path", @@ -244,6 +247,23 @@ func TestPassthroughRoutesForProviders(t *testing.T) { }, expectPath: "/v1/models", }, + { + name: "copilot_ping", + requestPath: "/copilot/_ping", + provider: func(_ *testing.T, baseURL string) provider.Provider { + return aibridge.NewCopilotProvider(config.Copilot{BaseURL: baseURL}) + }, + expectPath: "/_ping", + }, + { + name: "copilot_auto", + requestMethod: http.MethodPost, + requestPath: "/copilot/auto", + provider: func(_ *testing.T, baseURL string) provider.Provider { + return aibridge.NewCopilotProvider(config.Copilot{BaseURL: baseURL}) + }, + expectPath: "/auto", + }, } for _, tc := range tests { @@ -264,7 +284,7 @@ func TestPassthroughRoutesForProviders(t *testing.T) { bridge, err := aibridge.NewRequestBridge(t.Context(), []provider.Provider{prov}, &rec, nil, logger, nil, bridgeTestTracer) require.NoError(t, err) - req := httptest.NewRequest("", tc.requestPath, nil) + req := httptest.NewRequest(tc.requestMethod, tc.requestPath, nil) resp := httptest.NewRecorder() bridge.ServeHTTP(resp, req) @@ -274,6 +294,37 @@ func TestPassthroughRoutesForProviders(t *testing.T) { } } +func TestWebSocketUpgradeRejected(t *testing.T) { + t.Parallel() + + interceptorCalled := false + prov := &testutil.MockProvider{ + NameStr: "test", + Bridged: []string{"/responses"}, + InterceptorFunc: func(http.ResponseWriter, *http.Request, trace.Tracer) (intercept.Interceptor, error) { + interceptorCalled = true + return nil, nil //nolint:nilnil // The interceptor must not be reached. + }, + } + bridge, err := aibridge.NewRequestBridge( + t.Context(), + []provider.Provider{prov}, + nil, nil, slogtest.Make(t, nil), nil, bridgeTestTracer, + ) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodGet, "/test/responses", nil) + req.Header.Set("Connection", "keep-alive, Upgrade") + req.Header.Set("Upgrade", "WebSocket") + resp := httptest.NewRecorder() + + bridge.ServeHTTP(resp, req) + + assert.Equal(t, http.StatusNotImplemented, resp.Code) + assert.Contains(t, resp.Body.String(), "WebSocket transport is not supported, use HTTP") + assert.False(t, interceptorCalled) +} + func TestRequestBodySizeLimit(t *testing.T) { t.Parallel() diff --git a/aibridge/provider/copilot.go b/aibridge/provider/copilot.go index 79dc37ae4be..10d634c7bdf 100644 --- a/aibridge/provider/copilot.go +++ b/aibridge/provider/copilot.go @@ -91,6 +91,8 @@ func (*Copilot) BridgedRoutes() []string { func (*Copilot) PassthroughRoutes() []string { return []string{ + "/_ping", + "/auto", "/models", "/models/", "/agents/", diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 0941c2da4b5..5bc87b5feaf 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -28,6 +28,7 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog/v3" + aibridgeconfig "github.com/coder/coder/v2/aibridge/config" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" ) @@ -132,7 +133,7 @@ type Server struct { // refreshProviders fetches the live provider snapshot on Reload. // Nil disables hot-reload. refreshProviders RefreshProvidersFunc - // providerRouter holds the live (mitmHosts, nameByHost) pair. + // providerRouter holds the live routing snapshot. providerRouter atomic.Pointer[providerRouter] // allowedPorts is the port allowlist for CONNECT requests. Fixed at // construction; not reloadable. @@ -149,19 +150,26 @@ type Server struct { metrics *Metrics } +type routedProvider struct { + name string + providerType string +} + // providerRouter keeps CONNECT matching and provider lookup in sync. type providerRouter struct { - mitmHosts []string // host:port set the goproxy condition matches against. - nameByHost map[string]string // lowercase hostname -> provider name. + mitmHosts []string // host:port set the goproxy condition matches against. + providerByHost map[string]routedProvider // lowercase hostname -> provider. } // emptyProviderRouter is used before the first Reload (or when the // operator deconfigures every provider) so handlers can safely call // loadProviderRouter without a nil check. -var emptyProviderRouter = &providerRouter{nameByHost: map[string]string{}} +var emptyProviderRouter = &providerRouter{ + providerByHost: map[string]routedProvider{}, +} -func (r *providerRouter) providerFromHost(host string) string { - return r.nameByHost[strings.ToLower(host)] +func (r *providerRouter) providerFromHost(host string) routedProvider { + return r.providerByHost[strings.ToLower(host)] } // requestContext holds metadata propagated through the proxy request/response chain. @@ -655,13 +663,13 @@ func (s *Server) authMiddleware(host string, ctx *goproxy.ProxyCtx) (*goproxy.Co provider := s.loadProviderRouter().providerFromHost(ctx.Req.URL.Hostname()) // A concurrent Reload can swap the router between CONNECT matching // and provider lookup, so treat a missing mapping as a runtime miss. - if provider == "" { + if provider.name == "" { logger.Warn(s.ctx, "rejecting CONNECT request with no provider mapping") return goproxy.RejectConnect, host } logger = logger.With( - slog.F("provider", provider), + slog.F("provider", provider.name), ) proxyAuth := ctx.Req.Header.Get("Proxy-Authorization") @@ -685,7 +693,7 @@ func (s *Server) authMiddleware(host string, ctx *goproxy.ProxyCtx) (*goproxy.Co ctx.UserData = &requestContext{ ConnectSessionID: connectSessionID, CoderToken: coderToken, - Provider: provider, + Provider: provider.name, } logger.Debug(s.ctx, "request CONNECT authenticated") @@ -936,14 +944,14 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. } } liveProvider := s.loadProviderRouter().providerFromHost(host) - if liveProvider == "" || liveProvider != reqCtx.Provider { + if liveProvider.name == "" || liveProvider.name != reqCtx.Provider { s.logger.Warn(s.ctx, "provider mapping changed or removed since CONNECT, passing through", slog.F("connect_id", reqCtx.ConnectSessionID.String()), slog.F("host", req.Host), slog.F("method", req.Method), slog.F("path", originalPath), slog.F("connect_provider", reqCtx.Provider), - slog.F("live_provider", liveProvider), + slog.F("live_provider", liveProvider.name), ) return req, nil } @@ -992,7 +1000,13 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. req.URL = parsedGatewayTargetURL req.Host = parsedGatewayTargetURL.Host - injectBYOKHeaderIfNeeded(req.Header, reqCtx.CoderToken) + // Copilot is always BYOK, so every request needs the Coder token to + // authenticate with AI Gateway. Other providers only need it for BYOK requests. + if liveProvider.providerType == aibridgeconfig.ProviderCopilot { + setCopilotAuth(req.Header, reqCtx.CoderToken) + } else { + injectBYOKHeaderIfNeeded(req.Header, reqCtx.CoderToken) + } // Set request ID header to correlate requests between aibridgeproxyd and aibridged. req.Header.Set(agplaibridge.HeaderCoderRequestID, reqCtx.RequestID.String()) @@ -1019,6 +1033,19 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. return req, nil } +// setCopilotAuth sets the Coder auth header and removes the Coder token from +// provider auth headers. +func setCopilotAuth(headers http.Header, coderToken string) { + headers.Set(agplaibridge.HeaderCoderToken, coderToken) + + if extractCoderTokenFromBearerAuth(headers.Get("Authorization")) == coderToken { + headers.Del("Authorization") + } + if strings.TrimSpace(headers.Get("X-Api-Key")) == coderToken { + headers.Del("X-Api-Key") + } +} + // injectBYOKHeaderIfNeeded sets HeaderCoderToken when the // Authorization header carries a bearer token that differs from the // Coder token, indicating the client is using its own LLM diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 580b020ce15..946184d3753 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -193,26 +193,21 @@ func withProviders(providers ...aibridgeproxyd.ReloadedProvider) testProxyOption } // withProviderHosts is a convenience that builds enabled -// ReloadedProvider entries from each host, looking up the well-known -// provider name via testProviderFromHost and falling back to -// "test-provider" for hosts without a well-known mapping. Equivalent -// to passing each entry individually to withProviders. +// ReloadedProvider entries from each host, looking up well-known providers +// via testProviderFromHost. Unknown hosts use a generic name and OpenAI type. func withProviderHosts(hosts ...string) testProxyOption { return func(cfg *testProxyConfig) { providers := make([]aibridgeproxyd.ReloadedProvider, 0, len(hosts)) for _, h := range hosts { - name := testProviderFromHost(h) - if name == "" { - name = "test-provider" - } + provider := testProviderFromHost(h) host, _, splitErr := net.SplitHostPort(h) if splitErr != nil { host = h } providers = append(providers, aibridgeproxyd.ReloadedProvider{ ProviderOutcome: aibridged.ProviderOutcome{ - Name: name, - Type: "openai", + Name: provider.name, + Type: provider.providerType, Status: aibridged.ProviderStatusEnabled, }, Host: strings.ToLower(host), @@ -222,24 +217,29 @@ func withProviderHosts(hosts ...string) testProxyOption { } } -// testProviderFromHost maps well-known AI provider hostnames to -// provider names for test use. Unknown hosts return "". -func testProviderFromHost(host string) string { +type testProvider struct { + name string + providerType string +} + +// testProviderFromHost maps well-known AI provider hostnames to providers for +// test use. Unknown hosts use a generic name and OpenAI type. +func testProviderFromHost(host string) testProvider { switch strings.ToLower(host) { case aibridgeproxyd.HostAnthropic: - return aibridge.ProviderAnthropic + return testProvider{name: aibridge.ProviderAnthropic, providerType: aibridge.ProviderAnthropic} case aibridgeproxyd.HostOpenAI: - return aibridge.ProviderOpenAI + return testProvider{name: aibridge.ProviderOpenAI, providerType: aibridge.ProviderOpenAI} case aibridgeproxyd.HostCopilot: - return aibridge.ProviderCopilot + return testProvider{name: aibridge.ProviderCopilot, providerType: aibridge.ProviderCopilot} case agplaibridge.HostCopilotBusiness: - return agplaibridge.ProviderCopilotBusiness + return testProvider{name: agplaibridge.ProviderCopilotBusiness, providerType: aibridge.ProviderCopilot} case agplaibridge.HostCopilotEnterprise: - return agplaibridge.ProviderCopilotEnterprise + return testProvider{name: agplaibridge.ProviderCopilotEnterprise, providerType: aibridge.ProviderCopilot} case agplaibridge.HostChatGPT: - return agplaibridge.ProviderChatGPT + return testProvider{name: agplaibridge.ProviderChatGPT, providerType: aibridge.ProviderOpenAI} default: - return "" + return testProvider{name: "test-provider", providerType: aibridge.ProviderOpenAI} } } @@ -1613,13 +1613,13 @@ func TestProxy_MITM_BYOKInjection(t *testing.T) { srv := newTestProxy(t, withGatewayURL(aibridgedServer.URL), - withProviderHosts(aibridgeproxyd.HostCopilot), + withProviderHosts(aibridgeproxyd.HostOpenAI), ) certPool := getProxyCertPool(t) client := newProxyClient(t, srv, makeProxyAuthHeader(coderToken), certPool, false) - req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "https://"+aibridgeproxyd.HostCopilot+"/chat/completions", strings.NewReader(`{}`)) + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "https://"+aibridgeproxyd.HostOpenAI+"/chat/completions", strings.NewReader(`{}`)) require.NoError(t, err) req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", tt.authzHeader) @@ -1643,6 +1643,209 @@ func TestProxy_MITM_BYOKInjection(t *testing.T) { } } +func TestProxy_MITM_CopilotAuth(t *testing.T) { + t.Parallel() + + const coderToken = "coder-token" + stringPtr := func(value string) *string { return &value } + tests := []struct { + name string + host string + providerType string + method string + path string + authorization string + apiKey string + coderToken string + expectCoderToken *string + expectAuthorization *string + expectAPIKey *string + }{ + { + name: "ping passthrough route", + host: aibridgeproxyd.HostCopilot, + method: http.MethodGet, + path: "/_ping", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "models passthrough route", + host: aibridgeproxyd.HostCopilot, + method: http.MethodGet, + path: "/models", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "auto passthrough route", + host: aibridgeproxyd.HostCopilot, + method: http.MethodPost, + path: "/auto", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "responses bridged route", + host: aibridgeproxyd.HostCopilot, + method: http.MethodPost, + path: "/responses", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "messages bridged route", + host: aibridgeproxyd.HostCopilot, + method: http.MethodPost, + path: "/v1/messages", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "unknown route", + host: aibridgeproxyd.HostCopilot, + method: http.MethodGet, + path: "/unknown", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "Copilot strips Coder bearer token", + host: aibridgeproxyd.HostCopilot, + method: http.MethodPost, + path: "/responses", + authorization: "Bearer " + coderToken, + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "Copilot strips Coder API key", + host: aibridgeproxyd.HostCopilot, + method: http.MethodPost, + path: "/responses", + apiKey: coderToken, + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "Copilot preserves provider bearer token", + host: aibridgeproxyd.HostCopilot, + method: http.MethodPost, + path: "/responses", + authorization: "Bearer copilot-token", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: stringPtr("Bearer copilot-token"), + expectAPIKey: nil, + }, + { + name: "Copilot replaces client Coder token", + host: aibridgeproxyd.HostCopilot, + method: http.MethodGet, + path: "/models", + coderToken: "other-coder-token", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "Custom Copilot provider", + host: "copilot.example.com", + providerType: aibridge.ProviderCopilot, + method: http.MethodGet, + path: "/_ping", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "Non-Copilot provider on Copilot host", + host: aibridgeproxyd.HostCopilot, + providerType: aibridge.ProviderOpenAI, + method: http.MethodGet, + path: "/_ping", + expectCoderToken: nil, + expectAuthorization: nil, + expectAPIKey: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var receivedPath, receivedCoderToken, receivedAuthorization, receivedAPIKey string + aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedPath = r.URL.Path + receivedCoderToken = r.Header.Get(agplaibridge.HeaderCoderToken) + receivedAuthorization = r.Header.Get("Authorization") + receivedAPIKey = r.Header.Get("X-Api-Key") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(aibridgedServer.Close) + + provider := testProviderFromHost(tt.host) + if tt.providerType != "" { + provider.providerType = tt.providerType + } + srv := newTestProxy(t, + withGatewayURL(aibridgedServer.URL), + withProviders(aibridgeproxyd.ReloadedProvider{ + ProviderOutcome: aibridged.ProviderOutcome{ + Name: provider.name, + Type: provider.providerType, + Status: aibridged.ProviderStatusEnabled, + }, + Host: tt.host, + }), + ) + + certPool := getProxyCertPool(t) + client := newProxyClient(t, srv, makeProxyAuthHeader(coderToken), certPool, false) + + req, err := http.NewRequestWithContext(t.Context(), tt.method, "https://"+tt.host+tt.path, nil) + require.NoError(t, err) + if tt.authorization != "" { + req.Header.Set("Authorization", tt.authorization) + } + if tt.apiKey != "" { + req.Header.Set("X-Api-Key", tt.apiKey) + } + if tt.coderToken != "" { + req.Header.Set(agplaibridge.HeaderCoderToken, tt.coderToken) + } + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "/"+provider.name+tt.path, receivedPath) + if tt.expectAuthorization == nil { + require.Empty(t, receivedAuthorization) + } else { + require.Equal(t, *tt.expectAuthorization, receivedAuthorization) + } + if tt.expectAPIKey == nil { + require.Empty(t, receivedAPIKey) + } else { + require.Equal(t, *tt.expectAPIKey, receivedAPIKey) + } + if tt.expectCoderToken == nil { + require.Empty(t, receivedCoderToken) + } else { + require.Equal(t, *tt.expectCoderToken, receivedCoderToken) + } + }) + } +} + // TestListenerTLS verifies that the proxy works correctly when its listener is wrapped in TLS. // It tests both tunneled and MITM'd requests through an HTTPS proxy listener. func TestListenerTLS(t *testing.T) { diff --git a/enterprise/aibridgeproxyd/reload.go b/enterprise/aibridgeproxyd/reload.go index 5498fcdf64c..54aa69f8e17 100644 --- a/enterprise/aibridgeproxyd/reload.go +++ b/enterprise/aibridgeproxyd/reload.go @@ -123,7 +123,7 @@ func (s *Server) mitmHostsCondition() goproxy.ReqConditionFunc { // defense-in-depth even though the refresh function should mark // duplicates as proxy-excluded. func buildProviderRouter(reload ProviderReload, allowedPorts []string) (*providerRouter, error) { - nameByHost := make(map[string]string, len(reload.Providers)) + providerByHost := make(map[string]routedProvider, len(reload.Providers)) domains := make([]string, 0, len(reload.Providers)) for _, p := range reload.Providers { if p.Status != aibridged.ProviderStatusEnabled { @@ -133,15 +133,18 @@ func buildProviderRouter(reload ProviderReload, allowedPorts []string) (*provide if host == "" { continue } - if _, exists := nameByHost[host]; exists { + if _, exists := providerByHost[host]; exists { continue } - nameByHost[host] = p.Name + providerByHost[host] = routedProvider{name: p.Name, providerType: p.Type} domains = append(domains, host) } mitmHosts, err := convertDomainsToHosts(domains, allowedPorts) if err != nil { return nil, err } - return &providerRouter{mitmHosts: mitmHosts, nameByHost: nameByHost}, nil + return &providerRouter{ + mitmHosts: mitmHosts, + providerByHost: providerByHost, + }, nil } diff --git a/enterprise/aibridgeproxyd/reload_internal_test.go b/enterprise/aibridgeproxyd/reload_internal_test.go index a209cba038d..af83ce6761a 100644 --- a/enterprise/aibridgeproxyd/reload_internal_test.go +++ b/enterprise/aibridgeproxyd/reload_internal_test.go @@ -40,7 +40,7 @@ func TestServerReloadSwapsProviderRouter(t *testing.T) { srv.providerRouter.Store(emptyProviderRouter) require.NoError(t, srv.Reload(ctx)) - assert.Equal(t, "old", srv.loadProviderRouter().providerFromHost("old.example.com")) + assert.Equal(t, routedProvider{name: "old", providerType: "openai"}, srv.loadProviderRouter().providerFromHost("old.example.com")) assert.Empty(t, srv.loadProviderRouter().providerFromHost("new.example.com")) reload = ProviderReload{Providers: []ReloadedProvider{enabledProvider("new", "new.example.com")}} @@ -48,7 +48,7 @@ func TestServerReloadSwapsProviderRouter(t *testing.T) { router := srv.loadProviderRouter() assert.Empty(t, router.providerFromHost("old.example.com")) - assert.Equal(t, "new", router.providerFromHost("new.example.com")) + assert.Equal(t, routedProvider{name: "new", providerType: "openai"}, router.providerFromHost("new.example.com")) assert.Equal(t, []string{"new.example.com:443"}, router.mitmHosts) } @@ -74,14 +74,14 @@ func TestServerReloadPreservesProviderRouterOnRefreshError(t *testing.T) { require.NoError(t, srv.Reload(ctx)) before := srv.loadProviderRouter() - assert.Equal(t, "old", before.providerFromHost("old.example.com")) + assert.Equal(t, routedProvider{name: "old", providerType: "openai"}, before.providerFromHost("old.example.com")) failRefresh = true require.ErrorIs(t, srv.Reload(ctx), refreshErr) after := srv.loadProviderRouter() assert.Same(t, before, after) - assert.Equal(t, "old", after.providerFromHost("old.example.com")) + assert.Equal(t, routedProvider{name: "old", providerType: "openai"}, after.providerFromHost("old.example.com")) assert.Equal(t, []string{"old.example.com:443"}, after.mitmHosts) } @@ -95,7 +95,7 @@ func TestBuildProviderRouter(t *testing.T) { reload := ProviderReload{Providers: []ReloadedProvider{ enabledProvider("openai", "api.openai.com"), - enabledProvider("anthropic", "api.anthropic.com"), + {ProviderOutcome: aibridged.ProviderOutcome{Name: "anthropic", Type: "anthropic", Status: aibridged.ProviderStatusEnabled}, Host: "api.anthropic.com"}, enabledProvider("custom", "custom-llm.example.com"), // Host is populated on the non-enabled rows so the Status // guard, not the empty-host guard, is what excludes them. @@ -106,9 +106,9 @@ func TestBuildProviderRouter(t *testing.T) { router, err := buildProviderRouter(reload, []string{"443"}) require.NoError(t, err) - assert.Equal(t, "openai", router.providerFromHost("api.openai.com")) - assert.Equal(t, "anthropic", router.providerFromHost("api.anthropic.com")) - assert.Equal(t, "custom", router.providerFromHost("custom-llm.example.com")) + assert.Equal(t, routedProvider{name: "openai", providerType: "openai"}, router.providerFromHost("api.openai.com")) + assert.Equal(t, routedProvider{name: "anthropic", providerType: "anthropic"}, router.providerFromHost("api.anthropic.com")) + assert.Equal(t, routedProvider{name: "custom", providerType: "openai"}, router.providerFromHost("custom-llm.example.com")) assert.Empty(t, router.providerFromHost("unknown.com")) assert.Empty(t, router.providerFromHost("disabled.example.com"), "disabled provider must not be routable even with a populated Host") @@ -130,8 +130,8 @@ func TestBuildProviderRouter(t *testing.T) { router, err := buildProviderRouter(reload, []string{"443"}) require.NoError(t, err) - assert.Equal(t, "provider", router.providerFromHost("API.Example.COM")) - assert.Equal(t, "provider", router.providerFromHost("api.example.com")) + assert.Equal(t, routedProvider{name: "provider", providerType: "openai"}, router.providerFromHost("API.Example.COM")) + assert.Equal(t, routedProvider{name: "provider", providerType: "openai"}, router.providerFromHost("api.example.com")) }) t.Run("DefensiveDeduplicatesSameHost", func(t *testing.T) { @@ -147,7 +147,7 @@ func TestBuildProviderRouter(t *testing.T) { router, err := buildProviderRouter(reload, []string{"443"}) require.NoError(t, err) - assert.Equal(t, "first", router.providerFromHost("api.example.com")) + assert.Equal(t, routedProvider{name: "first", providerType: "openai"}, router.providerFromHost("api.example.com")) }) t.Run("SkipsRowsWithEmptyHost", func(t *testing.T) { @@ -161,7 +161,7 @@ func TestBuildProviderRouter(t *testing.T) { router, err := buildProviderRouter(reload, []string{"443"}) require.NoError(t, err) - assert.Equal(t, "good", router.providerFromHost("api.good.example.com")) + assert.Equal(t, routedProvider{name: "good", providerType: "openai"}, router.providerFromHost("api.good.example.com")) assert.Equal(t, []string{"api.good.example.com:443"}, router.mitmHosts) }) } From 49f88a89d12d2ff1d5343048bd238b978272bf7f Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 27 Aug 2026 11:25:54 +0000 Subject: [PATCH 2/5] chore: address comments --- enterprise/aibridgeproxyd/aibridgeproxyd.go | 29 +++++--- .../aibridgeproxyd/aibridgeproxyd_test.go | 69 +------------------ 2 files changed, 23 insertions(+), 75 deletions(-) diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 5bc87b5feaf..03af51e3226 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -1000,13 +1000,8 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. req.URL = parsedGatewayTargetURL req.Host = parsedGatewayTargetURL.Host - // Copilot is always BYOK, so every request needs the Coder token to - // authenticate with AI Gateway. Other providers only need it for BYOK requests. - if liveProvider.providerType == aibridgeconfig.ProviderCopilot { - setCopilotAuth(req.Header, reqCtx.CoderToken) - } else { - injectBYOKHeaderIfNeeded(req.Header, reqCtx.CoderToken) - } + // Prepare Coder authentication for centralized and BYOK requests. + prepareAIGatewayAuth(req.Header, reqCtx.CoderToken, liveProvider.providerType) // Set request ID header to correlate requests between aibridgeproxyd and aibridged. req.Header.Set(agplaibridge.HeaderCoderRequestID, reqCtx.RequestID.String()) @@ -1033,8 +1028,24 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. return req, nil } -// setCopilotAuth sets the Coder auth header and removes the Coder token from -// provider auth headers. +// prepareAIGatewayAuth prepares the Coder authentication headers for AI +// Gateway. Copilot is always BYOK, while other providers may use centralized +// or BYOK authentication. +func prepareAIGatewayAuth(headers http.Header, coderToken, providerType string) { + if providerType == aibridgeconfig.ProviderCopilot { + setCopilotAuth(headers, coderToken) + return + } + + // For other providers, only add the Coder token when a separate provider + // credential indicates BYOK. + injectBYOKHeaderIfNeeded(headers, coderToken) +} + +// setCopilotAuth adds the Coder token required by AI Gateway to every Copilot +// request. Unlike other providers, Copilot is always BYOK, even when a route +// does not include a provider credential (e.g., /_ping). It also prevents the +// Coder token from being forwarded to Copilot as a provider credential. func setCopilotAuth(headers http.Header, coderToken string) { headers.Set(agplaibridge.HeaderCoderToken, coderToken) diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 946184d3753..950c3f77cc2 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -1652,8 +1652,6 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { name string host string providerType string - method string - path string authorization string apiKey string coderToken string @@ -1662,55 +1660,8 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { expectAPIKey *string }{ { - name: "ping passthrough route", + name: "Copilot request without provider credential", host: aibridgeproxyd.HostCopilot, - method: http.MethodGet, - path: "/_ping", - expectCoderToken: stringPtr(coderToken), - expectAuthorization: nil, - expectAPIKey: nil, - }, - { - name: "models passthrough route", - host: aibridgeproxyd.HostCopilot, - method: http.MethodGet, - path: "/models", - expectCoderToken: stringPtr(coderToken), - expectAuthorization: nil, - expectAPIKey: nil, - }, - { - name: "auto passthrough route", - host: aibridgeproxyd.HostCopilot, - method: http.MethodPost, - path: "/auto", - expectCoderToken: stringPtr(coderToken), - expectAuthorization: nil, - expectAPIKey: nil, - }, - { - name: "responses bridged route", - host: aibridgeproxyd.HostCopilot, - method: http.MethodPost, - path: "/responses", - expectCoderToken: stringPtr(coderToken), - expectAuthorization: nil, - expectAPIKey: nil, - }, - { - name: "messages bridged route", - host: aibridgeproxyd.HostCopilot, - method: http.MethodPost, - path: "/v1/messages", - expectCoderToken: stringPtr(coderToken), - expectAuthorization: nil, - expectAPIKey: nil, - }, - { - name: "unknown route", - host: aibridgeproxyd.HostCopilot, - method: http.MethodGet, - path: "/unknown", expectCoderToken: stringPtr(coderToken), expectAuthorization: nil, expectAPIKey: nil, @@ -1718,8 +1669,6 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { { name: "Copilot strips Coder bearer token", host: aibridgeproxyd.HostCopilot, - method: http.MethodPost, - path: "/responses", authorization: "Bearer " + coderToken, expectCoderToken: stringPtr(coderToken), expectAuthorization: nil, @@ -1728,8 +1677,6 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { { name: "Copilot strips Coder API key", host: aibridgeproxyd.HostCopilot, - method: http.MethodPost, - path: "/responses", apiKey: coderToken, expectCoderToken: stringPtr(coderToken), expectAuthorization: nil, @@ -1738,8 +1685,6 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { { name: "Copilot preserves provider bearer token", host: aibridgeproxyd.HostCopilot, - method: http.MethodPost, - path: "/responses", authorization: "Bearer copilot-token", expectCoderToken: stringPtr(coderToken), expectAuthorization: stringPtr("Bearer copilot-token"), @@ -1748,8 +1693,6 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { { name: "Copilot replaces client Coder token", host: aibridgeproxyd.HostCopilot, - method: http.MethodGet, - path: "/models", coderToken: "other-coder-token", expectCoderToken: stringPtr(coderToken), expectAuthorization: nil, @@ -1759,8 +1702,6 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { name: "Custom Copilot provider", host: "copilot.example.com", providerType: aibridge.ProviderCopilot, - method: http.MethodGet, - path: "/_ping", expectCoderToken: stringPtr(coderToken), expectAuthorization: nil, expectAPIKey: nil, @@ -1769,8 +1710,6 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { name: "Non-Copilot provider on Copilot host", host: aibridgeproxyd.HostCopilot, providerType: aibridge.ProviderOpenAI, - method: http.MethodGet, - path: "/_ping", expectCoderToken: nil, expectAuthorization: nil, expectAPIKey: nil, @@ -1781,9 +1720,8 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - var receivedPath, receivedCoderToken, receivedAuthorization, receivedAPIKey string + var receivedCoderToken, receivedAuthorization, receivedAPIKey string aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - receivedPath = r.URL.Path receivedCoderToken = r.Header.Get(agplaibridge.HeaderCoderToken) receivedAuthorization = r.Header.Get("Authorization") receivedAPIKey = r.Header.Get("X-Api-Key") @@ -1810,7 +1748,7 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { certPool := getProxyCertPool(t) client := newProxyClient(t, srv, makeProxyAuthHeader(coderToken), certPool, false) - req, err := http.NewRequestWithContext(t.Context(), tt.method, "https://"+tt.host+tt.path, nil) + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://"+tt.host, nil) require.NoError(t, err) if tt.authorization != "" { req.Header.Set("Authorization", tt.authorization) @@ -1826,7 +1764,6 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { defer resp.Body.Close() require.Equal(t, http.StatusOK, resp.StatusCode) - require.Equal(t, "/"+provider.name+tt.path, receivedPath) if tt.expectAuthorization == nil { require.Empty(t, receivedAuthorization) } else { From 4259311faccb4ae84348543232ec5ac47334cb47 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 27 Aug 2026 11:31:49 +0000 Subject: [PATCH 3/5] test(enterprise/aibridgeproxyd): improve Copilot auth test names --- enterprise/aibridgeproxyd/aibridgeproxyd_test.go | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 950c3f77cc2..99fc40bfce1 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -1660,14 +1660,14 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { expectAPIKey *string }{ { - name: "Copilot request without provider credential", + name: "NoProviderCredential", host: aibridgeproxyd.HostCopilot, expectCoderToken: stringPtr(coderToken), expectAuthorization: nil, expectAPIKey: nil, }, { - name: "Copilot strips Coder bearer token", + name: "StripCoderBearer", host: aibridgeproxyd.HostCopilot, authorization: "Bearer " + coderToken, expectCoderToken: stringPtr(coderToken), @@ -1675,7 +1675,7 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { expectAPIKey: nil, }, { - name: "Copilot strips Coder API key", + name: "StripCoderAPIKey", host: aibridgeproxyd.HostCopilot, apiKey: coderToken, expectCoderToken: stringPtr(coderToken), @@ -1683,7 +1683,7 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { expectAPIKey: nil, }, { - name: "Copilot preserves provider bearer token", + name: "PreserveProviderBearer", host: aibridgeproxyd.HostCopilot, authorization: "Bearer copilot-token", expectCoderToken: stringPtr(coderToken), @@ -1691,7 +1691,7 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { expectAPIKey: nil, }, { - name: "Copilot replaces client Coder token", + name: "ReplaceClientCoderToken", host: aibridgeproxyd.HostCopilot, coderToken: "other-coder-token", expectCoderToken: stringPtr(coderToken), @@ -1699,7 +1699,7 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { expectAPIKey: nil, }, { - name: "Custom Copilot provider", + name: "CustomCopilotProvider", host: "copilot.example.com", providerType: aibridge.ProviderCopilot, expectCoderToken: stringPtr(coderToken), @@ -1707,7 +1707,7 @@ func TestProxy_MITM_CopilotAuth(t *testing.T) { expectAPIKey: nil, }, { - name: "Non-Copilot provider on Copilot host", + name: "NonCopilotProvider", host: aibridgeproxyd.HostCopilot, providerType: aibridge.ProviderOpenAI, expectCoderToken: nil, From db477ee0b95eea4ead04b57f2edb7539f5b0be7f Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 27 Aug 2026 11:48:45 +0000 Subject: [PATCH 4/5] refactor(enterprise/aibridgeproxyd): consolidate auth preparation --- enterprise/aibridgeproxyd/aibridgeproxyd.go | 44 ++++++--------------- 1 file changed, 11 insertions(+), 33 deletions(-) diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 03af51e3226..b073991ac2b 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -1033,48 +1033,26 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. // or BYOK authentication. func prepareAIGatewayAuth(headers http.Header, coderToken, providerType string) { if providerType == aibridgeconfig.ProviderCopilot { - setCopilotAuth(headers, coderToken) + headers.Set(agplaibridge.HeaderCoderToken, coderToken) + + if extractCoderTokenFromBearerAuth(headers.Get("Authorization")) == coderToken { + headers.Del("Authorization") + } + if strings.TrimSpace(headers.Get("X-Api-Key")) == coderToken { + headers.Del("X-Api-Key") + } return } // For other providers, only add the Coder token when a separate provider // credential indicates BYOK. - injectBYOKHeaderIfNeeded(headers, coderToken) -} - -// setCopilotAuth adds the Coder token required by AI Gateway to every Copilot -// request. Unlike other providers, Copilot is always BYOK, even when a route -// does not include a provider credential (e.g., /_ping). It also prevents the -// Coder token from being forwarded to Copilot as a provider credential. -func setCopilotAuth(headers http.Header, coderToken string) { - headers.Set(agplaibridge.HeaderCoderToken, coderToken) - - if extractCoderTokenFromBearerAuth(headers.Get("Authorization")) == coderToken { - headers.Del("Authorization") - } - if strings.TrimSpace(headers.Get("X-Api-Key")) == coderToken { - headers.Del("X-Api-Key") - } -} - -// injectBYOKHeaderIfNeeded sets HeaderCoderToken when the -// Authorization header carries a bearer token that differs from the -// Coder token, indicating the client is using its own LLM -// credentials. Clients that can set custom headers -// do this themselves; this handles clients that cannot. -// -// In centralized mode, Authorization carries the Coder token -// itself, so aibridged discovers it via ExtractAuthToken -// without any extra header. -func injectBYOKHeaderIfNeeded(header http.Header, coderToken string) { - // Don’t overwrite the header if it’s already set. - if header.Get(agplaibridge.HeaderCoderToken) != "" { + if headers.Get(agplaibridge.HeaderCoderToken) != "" { return } - bearer := extractCoderTokenFromBearerAuth(header.Get("Authorization")) + bearer := extractCoderTokenFromBearerAuth(headers.Get("Authorization")) if bearer != "" && bearer != coderToken { - header.Set(agplaibridge.HeaderCoderToken, coderToken) + headers.Set(agplaibridge.HeaderCoderToken, coderToken) } } From 9654aa3219a956ffda6a584f477ed872bb4f7894 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 27 Aug 2026 11:53:10 +0000 Subject: [PATCH 5/5] docs(enterprise/aibridgeproxyd): explain Copilot auth handling --- enterprise/aibridgeproxyd/aibridgeproxyd.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index b073991ac2b..fb7506b1d80 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -1032,6 +1032,9 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. // Gateway. Copilot is always BYOK, while other providers may use centralized // or BYOK authentication. func prepareAIGatewayAuth(headers http.Header, coderToken, providerType string) { + // Copilot is always BYOK, even when a route does not include a provider + // credential (e.g., /_ping). Prevent the Coder token from being forwarded + // to Copilot as a provider credential. if providerType == aibridgeconfig.ProviderCopilot { headers.Set(agplaibridge.HeaderCoderToken, coderToken)