Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions intercept/chatcompletions/base.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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

Expand All @@ -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))
Expand Down
22 changes: 17 additions & 5 deletions intercept/chatcompletions/blocking.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}}
}

Expand Down Expand Up @@ -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.
Comment on lines +91 to +92
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will handle this in a follow-up in order to not increase the scope of this PR.

if actor := aibcontext.ActorFromContext(r.Context()); actor != nil && i.cfg.SendActorHeaders {
opts = append(opts, intercept.ActorHeadersAsOpenAIOpts(actor)...)
}
Expand Down
22 changes: 17 additions & 5 deletions intercept/chatcompletions/streaming.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}}
}

Expand Down Expand Up @@ -121,6 +130,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)...)
}
Expand Down
10 changes: 5 additions & 5 deletions intercept/chatcompletions/streaming_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
72 changes: 72 additions & 0 deletions intercept/client_headers.go
Original file line number Diff line number Diff line change
@@ -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{
Comment thread
ssncferreira marked this conversation as resolved.
"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
}
Loading
Loading