feat: extract, log, and strip aibridgeproxy request ID header in aibridged (#23731)

## Problem

`aibridgeproxyd` sends `X-AI-Bridge-Request-Id` on every MITM request to
`aibridged` for cross-service log correlation, but aibridged never reads
it. The header is silently forwarded to upstream LLM providers.

## Changes

* Renamed the header to `X-Coder-AI-Governance-Request-Id` to match the
existing `X-Coder-AI-Governance-*` convention.
* `aibridged` now extracts the header, logs it and strips it before
forwarding upstream.
* Added `TestServeHTTP_StripInternalHeaders` to verify no `X-Coder-*`
headers leak to upstream
This commit is contained in:
Susana Ferreira
2026-03-30 15:21:30 +01:00
committed by GitHub
parent 7fb93dbf0e
commit 0fb3e5cba5
5 changed files with 79 additions and 7 deletions
+4
View File
@@ -16,6 +16,10 @@ import (
// that use per-user LLM credentials but cannot set custom headers.
const HeaderCoderToken = "X-Coder-AI-Governance-Token" //nolint:gosec // This is a header name, not a credential.
// HeaderCoderRequestID is a header set by aibridgeproxyd on each
// request forwarded to aibridged for cross-service log correlation.
const HeaderCoderRequestID = "X-Coder-AI-Governance-Request-Id"
// IsBYOK reports whether the request is using BYOK mode, determined
// by the presence of the X-Coder-AI-Governance-Token header.
func IsBYOK(header http.Header) bool {
+62
View File
@@ -574,3 +574,65 @@ func TestRouting(t *testing.T) {
})
}
}
// TestServeHTTP_StripInternalHeaders verifies that internal X-Coder-*
// headers are never forwarded to upstream LLM providers.
func TestServeHTTP_StripInternalHeaders(t *testing.T) {
t.Parallel()
cases := []struct {
name string
header string
value string
}{
{
name: "X-Coder-AI-Governance-Token",
header: agplaibridge.HeaderCoderToken,
value: "coder-token",
},
{
name: "X-Coder-AI-Governance-Request-Id",
header: agplaibridge.HeaderCoderRequestID,
value: uuid.NewString(),
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
mockH := &mockHandler{}
srv, client, pool := newTestServer(t)
conn := &mockDRPCConn{}
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil)
pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(mockH, nil)
httpSrv := httptest.NewServer(srv)
t.Cleanup(httpSrv.Close)
ctx := testutil.Context(t, testutil.WaitShort)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, httpSrv.URL+"/anthropic/v1/messages", nil)
require.NoError(t, err)
// Always set a valid auth token so the request reaches
// the upstream handler.
req.Header.Set("Authorization", "Bearer coder-token")
req.Header.Set(tc.header, tc.value)
resp, err := http.DefaultClient.Do(req)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
require.NotNil(t, mockH.headersReceived)
// Assert no X-Coder-* headers were forwarded upstream.
for name := range mockH.headersReceived {
require.NotContains(t, name, "X-Coder-",
"internal header %q must not be forwarded to upstream providers", name)
}
})
}
}
+10
View File
@@ -37,6 +37,16 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
logger := s.logger.With(slog.F("path", r.URL.Path))
// Extract and strip proxy request ID for cross-service log
// correlation. Absent for direct requests not routed through
// aibridgeproxyd.
if proxyReqID := r.Header.Get(agplaibridge.HeaderCoderRequestID); proxyReqID != "" {
// Inject into context so downstream loggers include it.
ctx = slog.With(ctx, slog.F("aibridgeproxy_id", proxyReqID))
logger = logger.With(slog.F("aibridgeproxy_id", proxyReqID))
}
r.Header.Del(agplaibridge.HeaderCoderRequestID)
byok := agplaibridge.IsBYOK(r.Header)
authMode := "centralized"
if byok {
+2 -6
View File
@@ -39,9 +39,6 @@ const (
)
const (
// HeaderAIBridgeRequestID is the header used to correlate requests
// between aibridgeproxyd and aibridged.
HeaderAIBridgeRequestID = "X-AI-Bridge-Request-Id"
// ProxyAuthRealm is the realm used in Proxy-Authenticate challenges.
// The realm helps clients identify which credentials to use.
ProxyAuthRealm = `"Coder AI Bridge Proxy"`
@@ -981,9 +978,8 @@ func (s *Server) handleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.
injectBYOKHeaderIfNeeded(req.Header, reqCtx.CoderToken)
// Set custom header for cross-service log correlation.
// This allows correlating aibridgeproxyd logs with aibridged logs.
req.Header.Set(HeaderAIBridgeRequestID, reqCtx.RequestID.String())
// Set request ID header to correlate requests between aibridgeproxyd and aibridged.
req.Header.Set(agplaibridge.HeaderCoderRequestID, reqCtx.RequestID.String())
logger.Info(s.ctx, "routing MITM request to aibridged",
slog.F("aibridged_url", aiBridgeParsedURL.String()),
@@ -1377,7 +1377,7 @@ func TestProxy_MITM(t *testing.T) {
receivedPath = r.URL.Path
receivedAuthz = r.Header.Get("Authorization")
receivedBYOK = r.Header.Get(agplaibridge.HeaderCoderToken)
receivedRequestID = r.Header.Get(aibridgeproxyd.HeaderAIBridgeRequestID)
receivedRequestID = r.Header.Get(agplaibridge.HeaderCoderRequestID)
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("hello from aibridged"))
}))