diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd.go b/enterprise/aibridgeproxyd/aibridgeproxyd.go index 45f27effb7..d6afd9dedc 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd.go @@ -14,6 +14,7 @@ import ( "net/url" "os" "slices" + "strconv" "strings" "sync" "time" @@ -70,6 +71,8 @@ type Server struct { // caCert is the PEM-encoded CA certificate loaded during initialization. // This is served to clients who need to trust the proxy. caCert []byte + // Metrics is the Prometheus metrics for the proxy. If nil, metrics are disabled. + metrics *Metrics } // requestContext holds metadata propagated through the proxy request/response chain. @@ -126,6 +129,9 @@ type Options struct { // proxies with certificates not trusted by the system. If empty, the system // certificate pool is used. UpstreamProxyCA string + // Metrics is the prometheus metrics instance for recording proxy metrics. + // If nil, metrics will not be recorded. + Metrics *Metrics } func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) { @@ -254,6 +260,7 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) coderAccessURL: coderAccessURL, aibridgeProviderFromHost: aibridgeProviderFromHost, caCert: certPEM, + metrics: opts.Metrics, } // Reject CONNECT requests to non-standard ports. @@ -269,6 +276,11 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error) srv.authMiddleware, ) + // Tunnel CONNECT requests for non-allowlisted domains directly to their destination. + // goproxy calls handlers in registration order: this must come after the MITM handler + // so it only handles requests that weren't matched by the allowlist. + proxy.OnRequest().HandleConnectFunc(srv.tunneledMiddleware) + // Handle decrypted requests: route to aibridged for known AI providers, or tunnel to original destination. proxy.OnRequest().DoFunc(srv.handleRequest) // Handle responses from aibridged. @@ -320,6 +332,12 @@ func (s *Server) Close() error { return nil } s.logger.Info(s.ctx, "closing aibridgeproxyd server") + + // Unregister metrics to clean up Prometheus registry. + if s.metrics != nil { + s.metrics.Unregister() + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() return s.httpServer.Shutdown(ctx) @@ -492,6 +510,11 @@ func (s *Server) authMiddleware(host string, ctx *goproxy.ProxyCtx) (*goproxy.Co logger.Debug(s.ctx, "request CONNECT authenticated") + // Record successful MITM CONNECT session establishment. + if s.metrics != nil { + s.metrics.ConnectSessionsTotal.WithLabelValues(RequestTypeMITM).Inc() + } + return goproxy.MitmConnect, host } @@ -565,6 +588,20 @@ func defaultAIBridgeProvider(host string) string { } } +// tunneledMiddleware is a CONNECT middleware that handles tunneled (non-allowlisted) +// connections. These connections are not MITM'd and are tunneled directly to their +// destination. This middleware records metrics for tunneled CONNECT sessions. +func (s *Server) tunneledMiddleware(host string, _ *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) { + // Record tunneled CONNECT session establishment. + if s.metrics != nil { + s.metrics.ConnectSessionsTotal.WithLabelValues(RequestTypeTunneled).Inc() + } + + // Return OkConnect to allow the tunnel to be established. + // goproxy will create a tunnel between the client and the destination. + return goproxy.OkConnect, host +} + // handleRequest intercepts HTTP requests after MITM decryption. // - Requests to known AI providers are rewritten to aibridged, with the Coder token // (from ctx.UserData, set during CONNECT) set in the X-Coder-Token header. @@ -580,6 +617,7 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. slog.F("method", req.Method), slog.F("path", originalPath), ) + resp := goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusProxyAuthRequired, "Proxy authentication required") resp.Header.Set("Proxy-Authenticate", `Basic realm="Coder AI Bridge Proxy"`) return req, resp @@ -657,6 +695,12 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http. slog.F("aibridged_url", aiBridgeParsedURL.String()), ) + // Record MITM request handling. + if s.metrics != nil { + s.metrics.MITMRequestsTotal.WithLabelValues(reqCtx.Provider).Inc() + s.metrics.InflightMITMRequests.WithLabelValues(reqCtx.Provider).Inc() + } + return req, nil } @@ -678,13 +722,30 @@ func (s *Server) handleResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *htt provider = reqCtx.Provider } - s.logger.Debug(s.ctx, "received response from aibridged", + logger := s.logger.With( slog.F("connect_id", connectSessionID.String()), slog.F("request_id", requestID.String()), - slog.F("status", resp.StatusCode), slog.F("provider", provider), + slog.F("status", resp.StatusCode), ) + switch { + case resp.StatusCode >= http.StatusInternalServerError: + logger.Error(s.ctx, "received error response from aibridged") + case resp.StatusCode >= http.StatusBadRequest: + logger.Warn(s.ctx, "received error response from aibridged") + default: + logger.Debug(s.ctx, "received response from aibridged") + } + + if s.metrics != nil && provider != "" { + // Decrement inflight requests gauge now that the request is complete. + s.metrics.InflightMITMRequests.WithLabelValues(provider).Dec() + + // Record response by status code. + s.metrics.MITMResponsesTotal.WithLabelValues(strconv.Itoa(resp.StatusCode), provider).Inc() + } + return resp } diff --git a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go index 6081d480f3..b5eba5b304 100644 --- a/enterprise/aibridgeproxyd/aibridgeproxyd_test.go +++ b/enterprise/aibridgeproxyd/aibridgeproxyd_test.go @@ -25,6 +25,7 @@ import ( "time" "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/require" "golang.org/x/xerrors" @@ -110,6 +111,7 @@ type testProxyConfig struct { aibridgeProviderFromHost func(string) string upstreamProxy string upstreamProxyCA string + metrics *aibridgeproxyd.Metrics } type testProxyOption func(*testProxyConfig) @@ -156,6 +158,12 @@ func withUpstreamProxyCA(upstreamProxyCA string) testProxyOption { } } +func withMetrics(metrics *aibridgeproxyd.Metrics) testProxyOption { + return func(cfg *testProxyConfig) { + cfg.metrics = metrics + } +} + // newTestProxy creates a new AI Bridge Proxy server for testing. // It uses the shared test CA and registers cleanup automatically. // It waits for the proxy server to be ready before returning. @@ -187,6 +195,7 @@ func newTestProxy(t *testing.T, opts ...testProxyOption) *aibridgeproxyd.Server AIBridgeProviderFromHost: cfg.aibridgeProviderFromHost, UpstreamProxy: cfg.upstreamProxy, UpstreamProxyCA: cfg.upstreamProxyCA, + Metrics: cfg.metrics, } if cfg.certStore != nil { aibridgeOpts.CertStore = cfg.certStore @@ -623,29 +632,89 @@ func TestNew(t *testing.T) { require.NoError(t, err) require.NotNil(t, srv) }) + + t.Run("SuccessWithMetrics", func(t *testing.T) { + t.Parallel() + + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + // Create metrics instance to verify it can be passed and stored. + reg := prometheus.NewRegistry() + metrics := aibridgeproxyd.NewMetrics(reg) + + srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + Metrics: metrics, + }) + require.NoError(t, err) + require.NotNil(t, srv) + }) } func TestClose(t *testing.T) { t.Parallel() - certFile, keyFile := getSharedTestCA(t) - logger := slogtest.Make(t, nil) + t.Run("Success", func(t *testing.T) { + t.Parallel() - srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ - ListenAddr: "127.0.0.1:0", - CoderAccessURL: "http://localhost:3000", - CertFile: certFile, - KeyFile: keyFile, - DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + }) + require.NoError(t, err) + + err = srv.Close() + require.NoError(t, err) + + // Calling Close again should not error. + err = srv.Close() + require.NoError(t, err) }) - require.NoError(t, err) - err = srv.Close() - require.NoError(t, err) + t.Run("WithMetrics", func(t *testing.T) { + t.Parallel() - // Calling Close again should not error - err = srv.Close() - require.NoError(t, err) + certFile, keyFile := getSharedTestCA(t) + logger := slogtest.Make(t, nil) + + // Create metrics instance to verify Close() properly unregisters them. + reg := prometheus.NewRegistry() + metrics := aibridgeproxyd.NewMetrics(reg) + + srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{ + ListenAddr: "127.0.0.1:0", + CoderAccessURL: "http://localhost:3000", + CertFile: certFile, + KeyFile: keyFile, + DomainAllowlist: []string{aibridgeproxyd.HostAnthropic, aibridgeproxyd.HostOpenAI}, + Metrics: metrics, + }) + require.NoError(t, err) + + err = srv.Close() + require.NoError(t, err) + + // Verify metrics were unregistered by attempting to register new metrics + // with the same registry. This should succeed if the old metrics were + // properly unregistered. + newMetrics := aibridgeproxyd.NewMetrics(reg) + require.NotNil(t, newMetrics, "should be able to create new metrics after Close() unregisters old ones") + + // Calling Close again should not error. + err = srv.Close() + require.NoError(t, err) + }) } func TestProxy_CertCaching(t *testing.T) { @@ -913,6 +982,7 @@ func TestProxy_MITM(t *testing.T) { buildTargetURL func(tunneledURL *url.URL) (string, error) tunneled bool expectedPath string + provider string }{ { name: "MitmdAnthropic", @@ -922,6 +992,7 @@ func TestProxy_MITM(t *testing.T) { return "https://api.anthropic.com/v1/messages", nil }, expectedPath: "/api/v2/aibridge/anthropic/v1/messages", + provider: "anthropic", }, { name: "MitmdAnthropicNonDefaultPort", @@ -931,6 +1002,7 @@ func TestProxy_MITM(t *testing.T) { return "https://api.anthropic.com:8443/v1/messages", nil }, expectedPath: "/api/v2/aibridge/anthropic/v1/messages", + provider: "anthropic", }, { name: "MitmdOpenAI", @@ -940,6 +1012,7 @@ func TestProxy_MITM(t *testing.T) { return "https://api.openai.com/v1/chat/completions", nil }, expectedPath: "/api/v2/aibridge/openai/v1/chat/completions", + provider: "openai", }, { name: "MitmdOpenAINonDefaultPort", @@ -949,6 +1022,7 @@ func TestProxy_MITM(t *testing.T) { return "https://api.openai.com:8443/v1/chat/completions", nil }, expectedPath: "/api/v2/aibridge/openai/v1/chat/completions", + provider: "openai", }, { name: "TunneledUnknownHost", @@ -965,6 +1039,10 @@ func TestProxy_MITM(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() + // Create metrics for verification. + reg := prometheus.NewRegistry() + metrics := aibridgeproxyd.NewMetrics(reg) + // Track what aibridged receives. var receivedPath, receivedCoderToken, receivedRequestID string @@ -1003,6 +1081,7 @@ func TestProxy_MITM(t *testing.T) { withDomainAllowlist(domainAllowlist...), // Use default provider mapping to test real AI provider routing. withAIBridgeProviderFromHost(nil), + withMetrics(metrics), ) // Build the target URL: @@ -1036,12 +1115,25 @@ func TestProxy_MITM(t *testing.T) { require.NoError(t, err) require.Equal(t, http.StatusOK, resp.StatusCode) + // Gather metrics for verification. + gatheredMetrics, err := reg.Gather() + require.NoError(t, err) + if tt.tunneled { // Verify request went to target server, not aibridged. require.Equal(t, "hello from tunneled", string(body)) require.Empty(t, receivedPath, "aibridged should not receive tunneled requests") require.Empty(t, receivedCoderToken, "tunneled requests are not authenticated by the proxy") require.Empty(t, receivedRequestID, "tunneled requests should not have request ID header") + + // Verify metrics for tunneled requests. + require.True(t, testutil.PromCounterHasValue(t, gatheredMetrics, 1, "connect_sessions_total", aibridgeproxyd.RequestTypeTunneled)) + + // Verify MITM-specific metrics were not set. + require.False(t, testutil.PromCounterGathered(t, gatheredMetrics, "connect_sessions_total", aibridgeproxyd.RequestTypeMITM)) + require.False(t, testutil.PromCounterGathered(t, gatheredMetrics, "mitm_requests_total", tt.provider)) + require.False(t, testutil.PromGaugeGathered(t, gatheredMetrics, "inflight_mitm_requests", tt.provider)) + require.False(t, testutil.PromCounterGathered(t, gatheredMetrics, "mitm_responses_total", "200", tt.provider)) } else { // Verify the request was routed to aibridged correctly. require.Equal(t, "hello from aibridged", string(body)) @@ -1050,6 +1142,15 @@ func TestProxy_MITM(t *testing.T) { require.NotEmpty(t, receivedRequestID, "MITM'd requests must include request ID header") _, err := uuid.Parse(receivedRequestID) require.NoError(t, err, "request ID must be a valid UUID") + + // Verify metrics for MITM requests. + require.True(t, testutil.PromCounterHasValue(t, gatheredMetrics, 1, "connect_sessions_total", aibridgeproxyd.RequestTypeMITM)) + require.True(t, testutil.PromCounterHasValue(t, gatheredMetrics, 1, "mitm_requests_total", tt.provider)) + require.True(t, testutil.PromGaugeHasValue(t, gatheredMetrics, 0, "inflight_mitm_requests", tt.provider)) + require.True(t, testutil.PromCounterHasValue(t, gatheredMetrics, 1, "mitm_responses_total", "200", tt.provider)) + + // Verify tunneled counter was not set. + require.False(t, testutil.PromCounterGathered(t, gatheredMetrics, "connect_sessions_total", aibridgeproxyd.RequestTypeTunneled)) } }) } diff --git a/enterprise/aibridgeproxyd/metrics.go b/enterprise/aibridgeproxyd/metrics.go new file mode 100644 index 0000000000..55a1fa4177 --- /dev/null +++ b/enterprise/aibridgeproxyd/metrics.go @@ -0,0 +1,70 @@ +package aibridgeproxyd + +import ( + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" +) + +const ( + RequestTypeMITM = "mitm" + RequestTypeTunneled = "tunneled" +) + +// Metrics holds all prometheus metrics for aibridgeproxyd. +type Metrics struct { + registerer prometheus.Registerer + + // ConnectSessionsTotal counts CONNECT sessions established. + // Labels: type (mitm/tunneled) + ConnectSessionsTotal *prometheus.CounterVec + + // MITMRequestsTotal counts MITM requests handled by the proxy. + // Labels: provider + MITMRequestsTotal *prometheus.CounterVec + + // InflightMITMRequests tracks the number of MITM requests currently being processed. + // Labels: provider + InflightMITMRequests *prometheus.GaugeVec + + // MITMResponsesTotal counts MITM responses by HTTP status code. + // Labels: code (HTTP status code), provider + // Cardinality is bounded: ~100 used status codes x few providers. + MITMResponsesTotal *prometheus.CounterVec +} + +// NewMetrics creates and registers all metrics for aibridgeproxyd. +func NewMetrics(reg prometheus.Registerer) *Metrics { + factory := promauto.With(reg) + + return &Metrics{ + registerer: reg, + + ConnectSessionsTotal: factory.NewCounterVec(prometheus.CounterOpts{ + Name: "connect_sessions_total", + Help: "Total number of CONNECT sessions established.", + }, []string{"type"}), + + MITMRequestsTotal: factory.NewCounterVec(prometheus.CounterOpts{ + Name: "mitm_requests_total", + Help: "Total number of MITM requests handled by the proxy.", + }, []string{"provider"}), + + InflightMITMRequests: factory.NewGaugeVec(prometheus.GaugeOpts{ + Name: "inflight_mitm_requests", + Help: "Number of MITM requests currently being processed.", + }, []string{"provider"}), + + MITMResponsesTotal: factory.NewCounterVec(prometheus.CounterOpts{ + Name: "mitm_responses_total", + Help: "Total number of MITM responses by HTTP status code class.", + }, []string{"code", "provider"}), + } +} + +// Unregister removes all metrics from the registerer. +func (m *Metrics) Unregister() { + m.registerer.Unregister(m.ConnectSessionsTotal) + m.registerer.Unregister(m.MITMRequestsTotal) + m.registerer.Unregister(m.InflightMITMRequests) + m.registerer.Unregister(m.MITMResponsesTotal) +} diff --git a/enterprise/cli/aibridgeproxyd.go b/enterprise/cli/aibridgeproxyd.go index 94fd66516b..589c7efe68 100644 --- a/enterprise/cli/aibridgeproxyd.go +++ b/enterprise/cli/aibridgeproxyd.go @@ -5,6 +5,7 @@ package cli import ( "context" + "github.com/prometheus/client_golang/prometheus" "golang.org/x/xerrors" "github.com/coder/coder/v2/enterprise/aibridgeproxyd" @@ -17,6 +18,9 @@ func newAIBridgeProxyDaemon(coderAPI *coderd.API) (*aibridgeproxyd.Server, error logger := coderAPI.Logger.Named("aibridgeproxyd") + reg := prometheus.WrapRegistererWithPrefix("coder_aibridgeproxyd_", coderAPI.PrometheusRegistry) + metrics := aibridgeproxyd.NewMetrics(reg) + srv, err := aibridgeproxyd.New(ctx, logger, aibridgeproxyd.Options{ ListenAddr: coderAPI.DeploymentValues.AI.BridgeProxyConfig.ListenAddr.String(), CoderAccessURL: coderAPI.AccessURL.String(), @@ -25,6 +29,7 @@ func newAIBridgeProxyDaemon(coderAPI *coderd.API) (*aibridgeproxyd.Server, error DomainAllowlist: coderAPI.DeploymentValues.AI.BridgeProxyConfig.DomainAllowlist.Value(), UpstreamProxy: coderAPI.DeploymentValues.AI.BridgeProxyConfig.UpstreamProxy.String(), UpstreamProxyCA: coderAPI.DeploymentValues.AI.BridgeProxyConfig.UpstreamProxyCA.String(), + Metrics: metrics, }) if err != nil { return nil, xerrors.Errorf("failed to start in-memory aibridgeproxy daemon: %w", err)