feat: add request/response dump support to aibridgeproxyd (#24837)

Closes https://github.com/coder/coder/issues/24335

enterprise/aibridgeproxyd/aibridgeproxyd.go

@@ -37,6 +37,13 @@ const (
 	HostCopilot   = "api.individual.githubcopilot.com"
 )
 
+// RoundTripDumper captures an HTTP request/response pair to disk.
+type RoundTripDumper interface {
+	DumpRequest(*http.Request) error
+	DumpResponse(*http.Response) error
+	DumpError(error) error
+}
+
 const (
 	// ProxyAuthRealm is the realm used in Proxy-Authenticate challenges.
 	// The realm helps clients identify which credentials to use.
@@ -125,6 +132,9 @@ type Server struct {
 	caCert []byte
 	// allowedPrivateRanges are CIDR ranges exempt from the blocked IP denylist.
 	allowedPrivateRanges []net.IPNet
+	// newDumper creates a RoundTripDumper for a given provider and request
+	// ID. Nil when dumping is disabled.
+	newDumper func(provider, requestID string) RoundTripDumper
 	// Metrics is the Prometheus metrics for the proxy. If nil, metrics are disabled.
 	metrics *Metrics
 }
@@ -147,6 +157,9 @@ type requestContext struct {
 	// Set in handleRequest for MITM'd requests.
 	// Sent to aibridged via custom header for cross-service correlation.
 	RequestID uuid.UUID
+	// Dumper captures request/response pairs to disk when API dump is
+	// enabled. Nil when dumping is disabled.
+	Dumper RoundTripDumper
 }
 
 // Options configures the AI Bridge Proxy server.
@@ -193,6 +206,11 @@ type Options struct {
 	// access to specific internal networks while keeping all other private
 	// ranges blocked. If empty, all private ranges are blocked.
 	AllowedPrivateCIDRs []string
+	// NewDumper, when non-nil, is called for each MITM request to create
+	// a RoundTripDumper that writes .req.txt and .resp.txt files. The
+	// caller is responsible for constructing the dumper with the correct
+	// base path.
+	NewDumper func(provider, requestID string) RoundTripDumper
 	// Metrics is the prometheus metrics instance for recording proxy metrics.
 	// If nil, metrics will not be recorded.
 	Metrics *Metrics
@@ -307,6 +325,7 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error)
 		aibridgeProviderFromHost: aibridgeProviderFromHost,
 		caCert:                   certPEM,
 		allowedPrivateRanges:     allowedPrivateRanges,
+		newDumper:                opts.NewDumper,
 		metrics:                  opts.Metrics,
 	}
 
@@ -452,6 +471,7 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Server, error)
 		slog.F("domain_allowlist", mitmHosts),
 		slog.F("upstream_proxy", opts.UpstreamProxy),
 		slog.F("allowed_private_cidrs", opts.AllowedPrivateCIDRs),
+		slog.F("api_dump_enabled", opts.NewDumper != nil),
 	)
 
 	go func() {
@@ -967,6 +987,15 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.
 		slog.F("aibridged_url", aiBridgeParsedURL.String()),
 	)
 
+	// Dump the outgoing request when API dumping is enabled.
+	if s.newDumper != nil {
+		d := s.newDumper(reqCtx.Provider, reqCtx.RequestID.String())
+		reqCtx.Dumper = d
+		if err := d.DumpRequest(req); err != nil {
+			logger.Warn(s.ctx, "failed to dump request", slog.Error(err))
+		}
+	}
+
 	// Record MITM request handling.
 	if s.metrics != nil {
 		s.metrics.MITMRequestsTotal.WithLabelValues(reqCtx.Provider).Inc()
@@ -1039,6 +1068,13 @@ func (s *Server) handleResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *htt
 		s.metrics.MITMResponsesTotal.WithLabelValues(strconv.Itoa(resp.StatusCode), provider).Inc()
 	}
 
+	// Dump the response to disk when a dumper was created for this request.
+	if reqCtx != nil && reqCtx.Dumper != nil {
+		if err := reqCtx.Dumper.DumpResponse(resp); err != nil {
+			logger.Warn(s.ctx, "failed to dump response", slog.Error(err))
+		}
+	}
+
 	return resp
 }
 

enterprise/aibridgeproxyd/aibridgeproxyd_test.go

@@ -26,6 +26,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/xerrors"
 
@@ -155,6 +156,7 @@ type testProxyConfig struct {
 	upstreamProxy            string
 	upstreamProxyCA          string
 	allowedPrivateCIDRs      []string
+	newDumper                func(string, string) aibridgeproxyd.RoundTripDumper
 	metrics                  *aibridgeproxyd.Metrics
 }
 
@@ -229,6 +231,12 @@ func withAllowedPrivateCIDRs(cidrs ...string) testProxyOption {
 	}
 }
 
+func withNewDumper(fn func(string, string) aibridgeproxyd.RoundTripDumper) testProxyOption {
+	return func(cfg *testProxyConfig) {
+		cfg.newDumper = fn
+	}
+}
+
 func withMetrics(metrics *aibridgeproxyd.Metrics) testProxyOption {
 	return func(cfg *testProxyConfig) {
 		cfg.metrics = metrics
@@ -279,6 +287,7 @@ func newTestProxy(t *testing.T, opts ...testProxyOption) *aibridgeproxyd.Server
 		UpstreamProxy:            cfg.upstreamProxy,
 		UpstreamProxyCA:          cfg.upstreamProxyCA,
 		AllowedPrivateCIDRs:      cfg.allowedPrivateCIDRs,
+		NewDumper:                cfg.newDumper,
 		Metrics:                  cfg.metrics,
 	}
 	if cfg.certStore != nil {
@@ -2353,3 +2362,133 @@ func TestProxy_PrivateIPBlocking(t *testing.T) {
 		})
 	}
 }
+
+// TestProxy_APIDump verifies that when NewDumper is configured, the proxy
+// calls DumpRequest and DumpResponse for MITM'd requests.
+func TestProxy_APIDump(t *testing.T) {
+	t.Parallel()
+
+	aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(`{"ok":true}`))
+	}))
+	t.Cleanup(aibridgedServer.Close)
+
+	var (
+		dumpedProvider  string
+		dumpedRequestID string
+		reqDumped       bool
+		respDumped      bool
+	)
+
+	srv := newTestProxy(t,
+		withCoderAccessURL(aibridgedServer.URL),
+		withAllowedPorts("443"),
+		withDomainAllowlist(aibridgeproxyd.HostAnthropic),
+		withAIBridgeProviderFromHost(testProviderFromHost),
+		withNewDumper(func(provider, requestID string) aibridgeproxyd.RoundTripDumper {
+			dumpedProvider = provider
+			dumpedRequestID = requestID
+			return &mockDumper{
+				onRequest:  func() { reqDumped = true },
+				onResponse: func() { respDumped = true },
+			}
+		}),
+	)
+
+	certPool := getProxyCertPool(t)
+	client := newProxyClient(t, srv, makeProxyAuthHeader("coder-token"), certPool, false)
+
+	req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "https://api.anthropic.com/v1/messages", strings.NewReader(`{}`))
+	require.NoError(t, err)
+	req.Header.Set("Authorization", "Bearer user-llm-token")
+
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	_, err = io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	assert.Equal(t, "anthropic", dumpedProvider)
+	assert.NotEmpty(t, dumpedRequestID)
+	_, err = uuid.Parse(dumpedRequestID)
+	require.NoError(t, err, "request ID passed to NewDumper must be a valid UUID")
+	assert.True(t, reqDumped, "DumpRequest should have been called")
+	assert.True(t, respDumped, "DumpResponse should have been called")
+}
+
+// TestProxy_APIDump_ErrorsDoNotAffectProxy verifies that dump failures
+// do not break the proxied request/response flow.
+func TestProxy_APIDump_ErrorsDoNotAffectProxy(t *testing.T) {
+	t.Parallel()
+
+	aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(`{"ok":true}`))
+	}))
+	t.Cleanup(aibridgedServer.Close)
+
+	srv := newTestProxy(t,
+		withCoderAccessURL(aibridgedServer.URL),
+		withAllowedPorts("443"),
+		withDomainAllowlist(aibridgeproxyd.HostAnthropic),
+		withAIBridgeProviderFromHost(testProviderFromHost),
+		withNewDumper(func(_, _ string) aibridgeproxyd.RoundTripDumper {
+			return &failingDumper{}
+		}),
+	)
+
+	certPool := getProxyCertPool(t)
+	client := newProxyClient(t, srv, makeProxyAuthHeader("coder-token"), certPool, false)
+
+	req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "https://api.anthropic.com/v1/messages", strings.NewReader(`{}`))
+	require.NoError(t, err)
+	req.Header.Set("Authorization", "Bearer user-token")
+
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+
+	// The proxy must return the upstream response despite dump errors.
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	require.JSONEq(t, `{"ok":true}`, string(body))
+}
+
+type mockDumper struct {
+	onRequest  func()
+	onResponse func()
+	onError    func()
+}
+
+func (m *mockDumper) DumpRequest(_ *http.Request) error {
+	if m.onRequest != nil {
+		m.onRequest()
+	}
+	return nil
+}
+
+func (m *mockDumper) DumpResponse(_ *http.Response) error {
+	if m.onResponse != nil {
+		m.onResponse()
+	}
+	return nil
+}
+
+func (m *mockDumper) DumpError(_ error) error {
+	if m.onError != nil {
+		m.onError()
+	}
+	return nil
+}
+
+// failingDumper always returns errors, used to verify dump failures
+// do not affect proxy behavior.
+type failingDumper struct{}
+
+func (*failingDumper) DumpRequest(*http.Request) error   { return xerrors.New("dump request failed") }
+func (*failingDumper) DumpResponse(*http.Response) error { return xerrors.New("dump response failed") }
+func (*failingDumper) DumpError(error) error             { return xerrors.New("dump error failed") }

enterprise/cli/aibridgeproxyd.go

@@ -5,12 +5,14 @@ package cli
 import (
 	"context"
 	"net/url"
+	"path/filepath"
 	"strings"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/aibridge"
+	"github.com/coder/coder/v2/aibridge/intercept/apidump"
 	"github.com/coder/coder/v2/enterprise/aibridgeproxyd"
 	"github.com/coder/coder/v2/enterprise/coderd"
 )
@@ -26,6 +28,13 @@ func newAIBridgeProxyDaemon(coderAPI *coderd.API, providers []aibridge.Provider)
 	reg := prometheus.WrapRegistererWithPrefix("coder_aibridgeproxyd_", coderAPI.PrometheusRegistry)
 	metrics := aibridgeproxyd.NewMetrics(reg)
 
+	var newDumper func(provider, requestID string) aibridgeproxyd.RoundTripDumper
+	if dumpDir := coderAPI.DeploymentValues.AI.BridgeProxyConfig.APIDumpDir.String(); dumpDir != "" {
+		newDumper = func(provider, requestID string) aibridgeproxyd.RoundTripDumper {
+			return apidump.NewDumper(filepath.Join(dumpDir, provider, requestID), logger)
+		}
+	}
+
 	srv, err := aibridgeproxyd.New(ctx, logger, aibridgeproxyd.Options{
 		ListenAddr:               coderAPI.DeploymentValues.AI.BridgeProxyConfig.ListenAddr.String(),
 		TLSCertFile:              coderAPI.DeploymentValues.AI.BridgeProxyConfig.TLSCertFile.String(),
@@ -38,6 +47,7 @@ func newAIBridgeProxyDaemon(coderAPI *coderd.API, providers []aibridge.Provider)
 		UpstreamProxy:            coderAPI.DeploymentValues.AI.BridgeProxyConfig.UpstreamProxy.String(),
 		UpstreamProxyCA:          coderAPI.DeploymentValues.AI.BridgeProxyConfig.UpstreamProxyCA.String(),
 		AllowedPrivateCIDRs:      coderAPI.DeploymentValues.AI.BridgeProxyConfig.AllowedPrivateCIDRs.Value(),
+		NewDumper:                newDumper,
 		Metrics:                  metrics,
 	})
 	if err != nil {

enterprise/cli/testdata/coder_server_--help.golden

@@ -175,6 +175,12 @@ AI BRIDGE OPTIONS:
           exporting these records to external SIEM or observability systems.
 
 AI BRIDGE PROXY OPTIONS: 
+      --aibridge-proxy-dump-dir string, $CODER_AIBRIDGE_PROXY_DUMP_DIR
+          Directory for dumping MITM request/response pairs to disk for
+          debugging. When set, each proxied request produces .req.txt and
+          .resp.txt files organized by provider. Sensitive headers are redacted.
+          Leave empty to disable.
+
       --aibridge-proxy-allowed-private-cidrs string-array, $CODER_AIBRIDGE_PROXY_ALLOWED_PRIVATE_CIDRS
           Comma-separated list of CIDR ranges that are permitted even though
           they fall within blocked private/reserved IP ranges. By default all