diff --git a/aibridge/bridge.go b/aibridge/bridge.go index e7b9101a2d7e2..f736dd6c422eb 100644 --- a/aibridge/bridge.go +++ b/aibridge/bridge.go @@ -16,6 +16,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" @@ -227,6 +228,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 + } + interceptor, err := p.CreateInterceptor(w, r.WithContext(ctx), tracer) if err != nil { span.SetStatus(codes.Error, fmt.Sprintf("failed to create interceptor: %v", err)) @@ -346,6 +359,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 new file mode 100644 index 0000000000000..0b07ed98c3a2f --- /dev/null +++ b/aibridge/bridge_internal_test.go @@ -0,0 +1,39 @@ +package aibridge + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +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)) + }) + } +} diff --git a/aibridge/bridge_test.go b/aibridge/bridge_test.go index 9ac7ea9ec3ddb..c8057b02ff305 100644 --- a/aibridge/bridge_test.go +++ b/aibridge/bridge_test.go @@ -12,10 +12,12 @@ 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/config" + "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/internal/testutil" "github.com/coder/coder/v2/aibridge/provider" ) @@ -123,11 +125,12 @@ func TestPassthroughRoutesForProviders(t *testing.T) { upstreamRespBody := "upstream response" tests := []struct { - name string - baseURLPath string - requestPath string - provider func(string) provider.Provider - expectPath string + name string + baseURLPath string + requestMethod string + requestPath string + provider func(string) provider.Provider + expectPath string }{ { name: "openAI_no_base_path", @@ -180,6 +183,23 @@ func TestPassthroughRoutesForProviders(t *testing.T) { }, expectPath: "/v1/models", }, + { + name: "copilot_ping", + requestPath: "/copilot/_ping", + provider: func(baseURL string) provider.Provider { + return aibridge.NewCopilotProvider(config.Copilot{BaseURL: baseURL}) + }, + expectPath: "/_ping", + }, + { + name: "copilot_auto", + requestMethod: http.MethodPost, + requestPath: "/copilot/auto", + provider: func(baseURL string) provider.Provider { + return aibridge.NewCopilotProvider(config.Copilot{BaseURL: baseURL}) + }, + expectPath: "/auto", + }, } for _, tc := range tests { @@ -200,7 +220,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) @@ -210,6 +230,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 fd317aadabac9..7859342e81e68 100644 --- a/aibridge/provider/copilot.go +++ b/aibridge/provider/copilot.go @@ -97,6 +97,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 1079cb6bde71e..a48a190deb55e 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. @@ -651,13 +659,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") @@ -681,7 +689,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") @@ -932,14 +940,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 } @@ -988,7 +996,8 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. req.URL = aiBridgeParsedURL req.Host = aiBridgeParsedURL.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()) @@ -1015,24 +1024,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 e84f7bd747f71..e433bd6692ae1 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} } } @@ -1547,13 +1547,13 @@ func TestProxy_MITM_BYOKInjection(t *testing.T) { srv := newTestProxy(t, withCoderAccessURL(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) @@ -1577,6 +1577,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, + withCoderAccessURL(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 04b1f5438b0ec..9dc54c5fff647 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 5ccba37ec7bd0..537392fc44751 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) }) }