From 49d0a4a08399d3bd18c34d8f0ba5c553fd854262 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Thu, 27 Aug 2026 13:10:34 +0100 Subject: [PATCH] fix: enable Copilot HTTP transport fallback (#28494) ## Description Copilot provider model metadata returned by `/models` can advertise WebSocket support (for example, `ws:/responses`), so clients attempt WebSocket inference even though AI Gateway supports HTTP transport only. The Copilot CLI and Copilot in VS Code use different retry mechanisms, and only the CLI fallback worked previously. See the investigation details below for more information. Authenticate every configured Copilot provider request with the Coder token validated during CONNECT while preserving Copilot provider credentials and preventing Coder credentials from being forwarded upstream. Reject unsupported WebSocket upgrades on bridged inference routes with `501 Not Implemented` so clients can fall back to HTTP. ## Changes - Add the CONNECT-authenticated Coder token to every configured Copilot provider request. - Preserve Copilot provider credentials and strip Coder credentials from provider credential headers. - Pass through Copilot `/_ping` connectivity checks and `/auto` model-selection requests. - Return `501 Not Implemented` for WebSocket upgrades on bridged inference routes.
Investigation details These logs show the behavior before this change using `GPT-5.3-Codex`, whose model metadata advertises `ws:/responses`. **Copilot CLI** 1. The CLI attempts a WebSocket upgrade. AI Gateway treats the empty `GET /responses` body as an inference request and returns 500. ``` 2026-08-26 11:08:28.883 [info] api: 2026-08-26 11:08:28.883 [warn] coderd.ai-gateway.pool: failed to create interceptor request_id=dcfc945c-fc10-4b8f-abc8-ae31aa21ad44 aibridgeproxy_id=6f651c82-68b2-45ee-b5e6-dc75445e2ef6 path=/copilot/responses ... 2026-08-26 11:08:28.883 [info] api: error= unmarshal request body: 2026-08-26 11:08:28.883 [info] api: github.com/coder/coder/v2/aibridge/provider.(*Copilot).CreateInterceptor 2026-08-26 11:08:28.883 [info] api: /home/coder/coder/aibridge/provider/copilot.go:179 2026-08-26 11:08:28.883 [info] api: - empty request body: 2026-08-26 11:08:28.883 [info] api: github.com/coder/coder/v2/aibridge/intercept/responses.NewRequestPayload 2026-08-26 11:08:28.884 [info] api: /home/coder/coder/aibridge/intercept/responses/reqpayload.go:48 2026-08-26 11:08:28.884 [info] api: 2026-08-26 11:08:28.883 [warn] coderd: GET user_agent="copilot/1.0.80 (client/github/cli linux v24.18.1) term/unknown" host=127.0.0.1:3000 received_host=127.0.0.1:3000 path=/api/v2/ai-gateway/copilot/responses proto=HTTP/1.1 remote_addr=127.0.0.1 start="2026-08-26 11:08:28.880115732 +0000 UTC m=+402.868920240" response_body="failed to create \"/copilot/responses\" interceptor\n" took=3.644201ms status_code=500 latency_ms=3 params_*=copilot/responses request_id=dcfc945c-fc10-4b8f-abc8-ae31aa21ad44 2026-08-26 11:08:28.884 [info] api: 2026-08-26 11:08:28.884 [erro] coderd.aibridgeproxyd: received error response from aibridged connect_id=17dda9cb-ae6a-416b-a812-b3b831755a2b request_id=6f651c82-68b2-45ee-b5e6-dc75445e2ef6 provider=copilot status=500 response_body="failed to create \"/copilot/responses\" interceptor\n" ``` 2. The CLI immediately retries inference with `POST /responses`, which succeeds. ``` 2026-08-26 11:08:28.886 [info] api: 2026-08-26 11:08:28.886 [debu] coderd.aibridgeproxyd: request CONNECT authenticated connect_id=7809a618-2c24-4f5f-a2f2-97af934b8c3e host=api.business.githubcopilot.com:443 provider=copilot 2026-08-26 11:08:28.890 [info] api: 2026-08-26 11:08:28.890 [info] coderd.aibridgeproxyd: routing MITM request to AI Gateway connect_id=7809a618-2c24-4f5f-a2f2-97af934b8c3e request_id=eee98705-ffc2-4a00-8201-e316f997e19d host=api.business.githubcopilot.com method=POST path=/responses provider=copilot gateway_target_url=http://127.0.0.1:3000/api/v2/ai-gateway/copilot/responses 2026-08-26 11:08:31.429 [info] api: 2026-08-26 11:08:31.428 [debu] coderd.aibridgeproxyd: received response from aibridged connect_id=7809a618-2c24-4f5f-a2f2-97af934b8c3e request_id=eee98705-ffc2-4a00-8201-e316f997e19d provider=copilot status=200 ``` **Copilot in VS Code** 1. VS Code attempts a WebSocket upgrade. AI Gateway treats the empty `GET /responses` body as an inference request and returns 500. ``` 2026-08-26 11:09:20.736 [info] api: 2026-08-26 11:09:20.736 [warn] coderd.ai-gateway.pool: failed to create interceptor request_id=4aab1c00-fb45-4533-b8c4-bcb150fd12f3 aibridgeproxy_id=d6867370-cefa-49bb-85ca-98dd4258e74c path=/copilot/responses ... 2026-08-26 11:09:20.736 [info] api: error= unmarshal request body: 2026-08-26 11:09:20.736 [info] api: github.com/coder/coder/v2/aibridge/provider.(*Copilot).CreateInterceptor 2026-08-26 11:09:20.736 [info] api: /home/coder/coder/aibridge/provider/copilot.go:179 2026-08-26 11:09:20.736 [info] api: - empty request body: 2026-08-26 11:09:20.736 [info] api: github.com/coder/coder/v2/aibridge/intercept/responses.NewRequestPayload 2026-08-26 11:09:20.736 [info] api: /home/coder/coder/aibridge/intercept/responses/reqpayload.go:48 2026-08-26 11:09:20.736 [info] api: 2026-08-26 11:09:20.736 [warn] coderd: GET user_agent=node host=127.0.0.1:3000 received_host=127.0.0.1:3000 path=/api/v2/ai-gateway/copilot/responses proto=HTTP/1.1 remote_addr=127.0.0.1 start="2026-08-26 11:09:20.732174884 +0000 UTC m=+454.720979381" response_body="failed to create \"/copilot/responses\" interceptor\n" took=4.225651ms status_code=500 latency_ms=4 params_*=copilot/responses request_id=4aab1c00-fb45-4533-b8c4-bcb150fd12f3 2026-08-26 11:09:20.737 [info] api: 2026-08-26 11:09:20.736 [erro] coderd.aibridgeproxyd: received error response from aibridged connect_id=072a6497-7072-42d8-b52d-a3afe17d5a75 request_id=d6867370-cefa-49bb-85ca-98dd4258e74c provider=copilot status=500 response_body="failed to create \"/copilot/responses\" interceptor\n" ``` 2. VS Code checks connectivity with `GET /_ping` before retrying inference over HTTP. AI Gateway rejects the check because the proxy did not forward the Coder token from the authenticated CONNECT session. ``` 2026-08-26 11:09:21.813 [info] api: 2026-08-26 11:09:21.812 [info] coderd.aibridgeproxyd: routing MITM request to AI Gateway connect_id=2dd5a340-7ad9-43a6-b99d-fd34b838e6a2 request_id=8fda1e3e-4ca8-4908-a2ef-23859743dc38 host=api.business.githubcopilot.com method=GET path=/_ping provider=copilot gateway_target_url=http://127.0.0.1:3000/api/v2/ai-gateway/copilot/_ping 2026-08-26 11:09:21.813 [info] api: 2026-08-26 11:09:21.813 [warn] coderd.ai-gateway: no auth key provided method=GET path=/copilot/_ping aibridgeproxy_id=8fda1e3e-4ca8-4908-a2ef-23859743dc38 request_id=057e3288-afa9-4f8d-9698-e4a3b206b199 aibridgeproxy_id=8fda1e3e-4ca8-4908-a2ef-23859743dc38 2026-08-26 11:09:21.814 [info] api: 2026-08-26 11:09:21.814 [warn] coderd.aibridgeproxyd: received error response from aibridged connect_id=2dd5a340-7ad9-43a6-b99d-fd34b838e6a2 request_id=8fda1e3e-4ca8-4908-a2ef-23859743dc38 provider=copilot status=400 response_body="no authentication key provided\n" ```
Closes https://linear.app/codercom/issue/AIGOV-629/ai-gateway-lacks-support-for-new-copilot-endpoints-behind-mitm > [!NOTE] > Generated by Coder Agents on behalf of @ssncferreira. (cherry picked from commit 849543d8de4cfb629fc6ca826673c3454b43ac4e) --- aibridge/bridge.go | 20 ++ aibridge/bridge_internal_test.go | 30 +++ aibridge/bridge_test.go | 63 +++++- aibridge/provider/copilot.go | 2 + enterprise/aibridgeproxyd/aibridgeproxyd.go | 71 ++++--- .../aibridgeproxyd/aibridgeproxyd_test.go | 184 +++++++++++++++--- enterprise/aibridgeproxyd/reload.go | 11 +- .../aibridgeproxyd/reload_internal_test.go | 24 +-- 8 files changed, 335 insertions(+), 70 deletions(-) diff --git a/aibridge/bridge.go b/aibridge/bridge.go index 5f3a5fbce92..dff41ba117f 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" @@ -248,6 +249,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 @@ -381,6 +394,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. func writeRequestBodyTooLarge(w http.ResponseWriter) { 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 d8e9103a7cb..c82ca422534 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" codertestutil "github.com/coder/coder/v2/testutil" @@ -186,11 +188,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", @@ -243,6 +246,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 { @@ -263,7 +283,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) @@ -273,6 +293,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..fb7506b1d80 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,8 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. req.URL = parsedGatewayTargetURL req.Host = parsedGatewayTargetURL.Host - 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()) @@ -1019,24 +1028,34 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. return req, nil } -// 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) != "" { +// 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) { + // 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) + + 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. + 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) } } diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 580b020ce15..99fc40bfce1 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,146 @@ 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 + authorization string + apiKey string + coderToken string + expectCoderToken *string + expectAuthorization *string + expectAPIKey *string + }{ + { + name: "NoProviderCredential", + host: aibridgeproxyd.HostCopilot, + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "StripCoderBearer", + host: aibridgeproxyd.HostCopilot, + authorization: "Bearer " + coderToken, + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "StripCoderAPIKey", + host: aibridgeproxyd.HostCopilot, + apiKey: coderToken, + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "PreserveProviderBearer", + host: aibridgeproxyd.HostCopilot, + authorization: "Bearer copilot-token", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: stringPtr("Bearer copilot-token"), + expectAPIKey: nil, + }, + { + name: "ReplaceClientCoderToken", + host: aibridgeproxyd.HostCopilot, + coderToken: "other-coder-token", + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "CustomCopilotProvider", + host: "copilot.example.com", + providerType: aibridge.ProviderCopilot, + expectCoderToken: stringPtr(coderToken), + expectAuthorization: nil, + expectAPIKey: nil, + }, + { + name: "NonCopilotProvider", + host: aibridgeproxyd.HostCopilot, + providerType: aibridge.ProviderOpenAI, + expectCoderToken: nil, + expectAuthorization: nil, + expectAPIKey: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var receivedCoderToken, receivedAuthorization, receivedAPIKey string + aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + 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(), http.MethodGet, "https://"+tt.host, 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) + 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 04b1f5438b0..9dc54c5fff6 100644 --- a/enterprise/aibridgeproxyd/reload.go +++ b/enterprise/aibridgeproxyd/reload.go @@ -119,7 +119,7 @@ func (s *Server) mitmHostsCondition() goproxy.ReqConditionFunc { // defense-in-depth measure even though the refresh function should // mark duplicates as errors. 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 { @@ -129,15 +129,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 5ccba37ec7b..537392fc447 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) { @@ -148,7 +148,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) { @@ -162,7 +162,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) }) }