feat: forward client request headers to upstream providers in bridge routes (#214)

* feat: forward client request headers to upstream providers in bridge routes

* chore: address comments

* test: update openai tests to table driven format

Author: ssncferreira

intercept/chatcompletions/base.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/coder/aibridge/config"
 	aibcontext "github.com/coder/aibridge/context"
+	"github.com/coder/aibridge/intercept"
 	"github.com/coder/aibridge/intercept/apidump"
 	"github.com/coder/aibridge/mcp"
 	"github.com/coder/aibridge/recorder"
@@ -29,6 +30,10 @@ type interceptionBase struct {
 	req *ChatCompletionNewParamsWrapper
 	cfg config.OpenAI
 
+	// clientHeaders are the original HTTP headers from the client request.
+	clientHeaders  http.Header
+	authHeaderName string
+
 	logger slog.Logger
 	tracer trace.Tracer
 
@@ -41,10 +46,21 @@ func (i *interceptionBase) newCompletionsService() openai.ChatCompletionService
 
 	// Add extra headers if configured.
 	// Some providers require additional headers that are not added by the SDK.
+	// TODO(ssncferreira): remove as part of https://github.com/coder/aibridge/issues/192
 	for key, value := range i.cfg.ExtraHeaders {
 		opts = append(opts, option.WithHeader(key, value))
 	}
 
+	// Forward client headers to upstream. This middleware runs after the SDK
+	// has built the request, and replaces the outgoing headers with the sanitized
+	// client headers plus provider auth.
+	if i.clientHeaders != nil {
+		opts = append(opts, option.WithMiddleware(func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) {
+			req.Header = intercept.BuildUpstreamHeaders(req.Header, i.clientHeaders, i.authHeaderName)
+			return next(req)
+		}))
+	}
+
 	// Add API dump middleware if configured
 	if mw := apidump.NewBridgeMiddleware(i.cfg.APIDumpDir, config.ProviderOpenAI, i.Model(), i.id, i.logger, quartz.NewReal()); mw != nil {
 		opts = append(opts, option.WithMiddleware(mw))

intercept/chatcompletions/blocking.go

@@ -28,12 +28,21 @@ type BlockingInterception struct {
 	interceptionBase
 }
 
-func NewBlockingInterceptor(id uuid.UUID, req *ChatCompletionNewParamsWrapper, cfg config.OpenAI, tracer trace.Tracer) *BlockingInterception {
+func NewBlockingInterceptor(
+	id uuid.UUID,
+	req *ChatCompletionNewParamsWrapper,
+	cfg config.OpenAI,
+	clientHeaders http.Header,
+	authHeaderName string,
+	tracer trace.Tracer,
+) *BlockingInterception {
 	return &BlockingInterception{interceptionBase: interceptionBase{
-		id:     id,
-		req:    req,
-		cfg:    cfg,
-		tracer: tracer,
+		id:             id,
+		req:            req,
+		cfg:            cfg,
+		clientHeaders:  clientHeaders,
+		authHeaderName: authHeaderName,
+		tracer:         tracer,
 	}}
 }
 
@@ -78,6 +87,9 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req
 
 		var opts []option.RequestOption
 		opts = append(opts, option.WithRequestTimeout(time.Second*600))
+
+		// TODO(ssncferreira): inject actor headers directly in the client-header
+		//   middleware instead of using SDK options.
 		if actor := aibcontext.ActorFromContext(r.Context()); actor != nil && i.cfg.SendActorHeaders {
 			opts = append(opts, intercept.ActorHeadersAsOpenAIOpts(actor)...)
 		}

intercept/chatcompletions/streaming.go

@@ -33,12 +33,21 @@ type StreamingInterception struct {
 	interceptionBase
 }
 
-func NewStreamingInterceptor(id uuid.UUID, req *ChatCompletionNewParamsWrapper, cfg config.OpenAI, tracer trace.Tracer) *StreamingInterception {
+func NewStreamingInterceptor(
+	id uuid.UUID,
+	req *ChatCompletionNewParamsWrapper,
+	cfg config.OpenAI,
+	clientHeaders http.Header,
+	authHeaderName string,
+	tracer trace.Tracer,
+) *StreamingInterception {
 	return &StreamingInterception{interceptionBase: interceptionBase{
-		id:     id,
-		req:    req,
-		cfg:    cfg,
-		tracer: tracer,
+		id:             id,
+		req:            req,
+		cfg:            cfg,
+		clientHeaders:  clientHeaders,
+		authHeaderName: authHeaderName,
+		tracer:         tracer,
 	}}
 }
 
@@ -115,6 +124,9 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re
 	for {
 		// TODO add outer loop span (https://github.com/coder/aibridge/issues/67)
 		var opts []option.RequestOption
+
+		// TODO(ssncferreira): inject actor headers directly in the client-header
+		//   middleware instead of using SDK options.
 		if actor := aibcontext.ActorFromContext(r.Context()); actor != nil && i.cfg.SendActorHeaders {
 			opts = append(opts, intercept.ActorHeadersAsOpenAIOpts(actor)...)
 		}

intercept/chatcompletions/streaming_test.go

@@ -81,16 +81,16 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) {
 				Stream: true,
 			}
 
+			// Create test request
+			w := httptest.NewRecorder()
+			httpReq := httptest.NewRequest(http.MethodPost, "/chat/completions", nil)
+
 			tracer := otel.Tracer("test")
-			interceptor := NewStreamingInterceptor(uuid.New(), req, cfg, tracer)
+			interceptor := NewStreamingInterceptor(uuid.New(), req, cfg, httpReq.Header, "Authorization", tracer)
 
 			logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
 			interceptor.Setup(logger, &testutil.MockRecorder{}, nil)
 
-			// Create test request
-			w := httptest.NewRecorder()
-			httpReq := httptest.NewRequest(http.MethodPost, "/chat/completions", nil)
-
 			// Process the request
 			err := interceptor.ProcessRequest(w, httpReq)
 

intercept/client_headers.go

@@ -0,0 +1,72 @@
+package intercept
+
+import "net/http"
+
+// hopByHopHeaders are connection-level headers specific to the connection
+// between client and AI Bridge, not meant for the upstream.
+// See https://www.rfc-editor.org/rfc/rfc2616#section-13.5.1
+var hopByHopHeaders = []string{
+	"Connection",
+	"Keep-Alive",
+	"Proxy-Authenticate",
+	"Proxy-Authorization",
+	"Te",
+	"Trailer",
+	"Transfer-Encoding",
+	"Upgrade",
+}
+
+// nonForwardedHeaders are transport-level headers managed by aibridge or
+// Go's HTTP transport that must not be forwarded to the upstream provider.
+var nonForwardedHeaders = []string{
+	"Host",
+	"Accept-Encoding",
+	"Content-Length",
+}
+
+// authHeaders are headers that carry authentication credentials from the
+// client. The upstream request is built by the SDK, which sets the correct
+// provider credentials via option.WithAPIKey. Client auth headers are
+// stripped here and the provider credentials are re-injected by
+// BuildUpstreamHeaders from the SDK-built request.
+var authHeaders = []string{
+	"Authorization",
+	"X-Api-Key",
+}
+
+// PrepareClientHeaders returns a copy of the client headers with hop-by-hop,
+// transport, and auth headers removed.
+func PrepareClientHeaders(clientHeaders http.Header) http.Header {
+	prepared := clientHeaders.Clone()
+	for _, h := range hopByHopHeaders {
+		prepared.Del(h)
+	}
+	for _, h := range nonForwardedHeaders {
+		prepared.Del(h)
+	}
+	for _, h := range authHeaders {
+		prepared.Del(h)
+	}
+	return prepared
+}
+
+// BuildUpstreamHeaders produces the header set for an upstream SDK request.
+// It starts from the prepared client headers, then preserves specific
+// headers from the SDK-built request that must not be overwritten.
+func BuildUpstreamHeaders(sdkHeader http.Header, clientHeaders http.Header, authHeaderName string) http.Header {
+	headers := PrepareClientHeaders(clientHeaders)
+
+	// Preserve the auth header set by the SDK from the provider configuration.
+	if v := sdkHeader.Get(authHeaderName); v != "" {
+		headers.Set(authHeaderName, v)
+	}
+
+	// Preserve actor headers injected by aibridge as per-request SDK options.
+	for name, values := range sdkHeader {
+		if IsActorHeader(name) {
+			headers[name] = values
+		}
+	}
+
+	return headers
+}