feat: add metrics to aibridgeproxy (#21709)

## Description Adds Prometheus metrics to the AI Bridge Proxy for observability into proxy traffic and performance. ## Changes * Add Metrics struct with the following metrics: * `connect_sessions_total`: counts CONNECT sessions by type (mitm/tunneled) * `mitm_requests_total`: counts MITM requests by provider * `inflight_mitm_requests`: gauge tracking in-flight requests by provider * `mitm_request_duration_seconds`: histogram of request latencies by provider * `mitm_responses_total`: counts responses by status code class (2XX/3XX/4XX/5XX) and provider * Register metrics with `coder_aibridgeproxyd_` prefix in CLI * Unregister metrics on server close to prevent registry leaks * Add `tunneledMiddleware` to track non-allowlisted CONNECT sessions * Add tests for metric recording in both MITM and tunneled paths Closes: https://github.com/coder/internal/issues/1185
2026-09-24 15:04:27 +08:00 · 2026-01-29 15:11:36 +00:00
parent d09300eadf
commit 9f6ce7542a
4 changed files with 253 additions and 16 deletions
@@ -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
 }

@@ -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))
 			}
 		})
 	}
@@ -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)
+}
@@ -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)