From 5a0b836096adbf94a2e068b2fdf7dc9a30807842 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 12 Mar 2026 20:03:33 +0000 Subject: [PATCH 1/4] feat: aibrige BYOK --- config/config.go | 4 ++ intercept/messages/base.go | 8 ++- provider/anthropic.go | 24 ++++++++ provider/anthropic_test.go | 111 +++++++++++++++++++++++++++++++++++++ 4 files changed, 146 insertions(+), 1 deletion(-) diff --git a/config/config.go b/config/config.go index f4107e42..af16502b 100644 --- a/config/config.go +++ b/config/config.go @@ -15,6 +15,10 @@ type Anthropic struct { CircuitBreaker *CircuitBreaker SendActorHeaders bool ExtraHeaders map[string]string + // BYOKBearerToken is set in BYOK mode when the user authenticates + // with an OAuth token (e.g. Claude Max/Pro subscription). When set, + // the SDK uses Authorization: Bearer instead of X-Api-Key. + BYOKBearerToken string } type AWSBedrock struct { diff --git a/intercept/messages/base.go b/intercept/messages/base.go index 09372ec7..b9a05a28 100644 --- a/intercept/messages/base.go +++ b/intercept/messages/base.go @@ -205,7 +205,13 @@ func (i *interceptionBase) isSmallFastModel() bool { } func (i *interceptionBase) newMessagesService(ctx context.Context, opts ...option.RequestOption) (anthropic.MessageService, error) { - opts = append(opts, option.WithAPIKey(i.cfg.Key)) + // BYOK with OAuth token (Claude Max/Pro) uses Authorization: Bearer. + // Otherwise use X-Api-Key (centralized or BYOK with personal API key). + if i.cfg.BYOKBearerToken != "" { + opts = append(opts, option.WithAuthToken(i.cfg.BYOKBearerToken)) + } else { + opts = append(opts, option.WithAPIKey(i.cfg.Key)) + } opts = append(opts, option.WithBaseURL(i.cfg.BaseURL)) // Add extra headers if configured. diff --git a/provider/anthropic.go b/provider/anthropic.go index 4a79cd42..173b9b4e 100644 --- a/provider/anthropic.go +++ b/provider/anthropic.go @@ -110,6 +110,24 @@ func (p *Anthropic) CreateInterceptor(w http.ResponseWriter, r *http.Request, tr cfg := p.cfg cfg.ExtraHeaders = extractAnthropicHeaders(r) + // In centralized mode, http.go strips Authorization and X-Api-Key + // (they carried the Coder token), so neither header is present + // here and cfg keeps the centralized key. + // + // In BYOK mode, http.go only strips the BYOK header and leaves + // the user's LLM credentials intact: + // - Authorization: Bearer → subscription (Claude + // Max/Pro). Set BYOKBearerToken so the SDK uses + // WithAuthToken(), and clear the centralized key. + // - X-Api-Key: → personal API key. Overwrite the + // centralized key with the user's key. + if bearer := r.Header.Get("Authorization"); bearer != "" { + cfg.BYOKBearerToken = strings.TrimPrefix(bearer, "Bearer ") + cfg.Key = "" + } else if apiKey := r.Header.Get("X-Api-Key"); apiKey != "" { + cfg.Key = apiKey + } + var interceptor intercept.Interceptor if req.Stream { interceptor = messages.NewStreamingInterceptor(id, &req, payload, cfg, p.bedrockCfg, r.Header, p.AuthHeader(), tracer) @@ -137,6 +155,12 @@ func (p *Anthropic) InjectAuthHeader(headers *http.Header) { headers = &http.Header{} } + // BYOK: if the request already carries user-supplied credentials, + // do not overwrite them with the centralized key. + if headers.Get("X-Api-Key") != "" || headers.Get("Authorization") != "" { + return + } + headers.Set(p.AuthHeader(), p.cfg.Key) } diff --git a/provider/anthropic_test.go b/provider/anthropic_test.go index d34fd029..e815f28b 100644 --- a/provider/anthropic_test.go +++ b/provider/anthropic_test.go @@ -110,6 +110,71 @@ func TestAnthropic_CreateInterceptor(t *testing.T) { assert.Empty(t, receivedHeaders.Get("Authorization"), "client Authorization header must not reach upstream") }) + byokTests := []struct { + name string + setHeaders map[string]string + wantXApiKey string + wantAuthorization string + }{ + { + name: "Messages_BYOK_BearerToken", + setHeaders: map[string]string{"Authorization": "Bearer user-oauth-token"}, + wantAuthorization: "Bearer user-oauth-token", + }, + { + name: "Messages_BYOK_APIKey", + setHeaders: map[string]string{"X-Api-Key": "user-api-key"}, + wantXApiKey: "user-api-key", + }, + { + name: "Messages_Centralized_UsesCentralizedKey", + setHeaders: map[string]string{}, + wantXApiKey: "test-key", + }, + } + + for _, tc := range byokTests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var receivedHeaders http.Header + + mockUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"id":"msg-123","type":"message","role":"assistant","content":[{"type":"text","text":"Hello!"}],"model":"claude-opus-4-5","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":5}}`)) + })) + t.Cleanup(mockUpstream.Close) + + provider := NewAnthropic(config.Anthropic{ + BaseURL: mockUpstream.URL, + Key: "test-key", + }, nil) + + body := `{"model": "claude-opus-4-5", "max_tokens": 1024, "messages": [{"role": "user", "content": "hello"}], "stream": false}` + req := httptest.NewRequest(http.MethodPost, routeMessages, bytes.NewBufferString(body)) + for k, v := range tc.setHeaders { + req.Header.Set(k, v) + } + w := httptest.NewRecorder() + + interceptor, err := provider.CreateInterceptor(w, req, testTracer) + require.NoError(t, err) + require.NotNil(t, interceptor) + + logger := slog.Make() + interceptor.Setup(logger, &testutil.MockRecorder{}, nil) + + processReq := httptest.NewRequest(http.MethodPost, routeMessages, nil) + err = interceptor.ProcessRequest(w, processReq) + require.NoError(t, err) + + assert.Equal(t, tc.wantXApiKey, receivedHeaders.Get("X-Api-Key")) + assert.Equal(t, tc.wantAuthorization, receivedHeaders.Get("Authorization")) + }) + } + t.Run("UnknownRoute", func(t *testing.T) { t.Parallel() @@ -124,6 +189,52 @@ func TestAnthropic_CreateInterceptor(t *testing.T) { }) } +func TestAnthropic_InjectAuthHeader_BYOK(t *testing.T) { + t.Parallel() + + provider := NewAnthropic(config.Anthropic{Key: "centralized-key"}, nil) + + tests := []struct { + name string + presetHeaders map[string]string + wantXApiKey string + wantAuthorization string + }{ + { + name: "no pre-existing auth headers injects centralized key", + presetHeaders: map[string]string{}, + wantXApiKey: "centralized-key", + }, + { + name: "pre-existing X-Api-Key is not overwritten", + presetHeaders: map[string]string{"X-Api-Key": "user-api-key"}, + wantXApiKey: "user-api-key", + }, + { + name: "pre-existing Authorization prevents centralized key injection", + presetHeaders: map[string]string{"Authorization": "Bearer user-oauth-token"}, + wantXApiKey: "", + wantAuthorization: "Bearer user-oauth-token", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + headers := http.Header{} + for k, v := range tc.presetHeaders { + headers.Set(k, v) + } + + provider.InjectAuthHeader(&headers) + + assert.Equal(t, tc.wantXApiKey, headers.Get("X-Api-Key")) + assert.Equal(t, tc.wantAuthorization, headers.Get("Authorization")) + }) + } +} + func TestExtractAnthropicHeaders(t *testing.T) { t.Parallel() From 597cbf5a15b13dc14f2939925b97b428a9a8d8f0 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Mon, 16 Mar 2026 20:20:35 +0000 Subject: [PATCH 2/4] fix: reconcile with BuildUpstreamHeaders --- provider/anthropic.go | 6 ++++-- provider/anthropic_test.go | 12 ++++++------ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/provider/anthropic.go b/provider/anthropic.go index 173b9b4e..f35bf64a 100644 --- a/provider/anthropic.go +++ b/provider/anthropic.go @@ -121,18 +121,20 @@ func (p *Anthropic) CreateInterceptor(w http.ResponseWriter, r *http.Request, tr // WithAuthToken(), and clear the centralized key. // - X-Api-Key: → personal API key. Overwrite the // centralized key with the user's key. + authHeaderName := p.AuthHeader() if bearer := r.Header.Get("Authorization"); bearer != "" { cfg.BYOKBearerToken = strings.TrimPrefix(bearer, "Bearer ") cfg.Key = "" + authHeaderName = "Authorization" } else if apiKey := r.Header.Get("X-Api-Key"); apiKey != "" { cfg.Key = apiKey } var interceptor intercept.Interceptor if req.Stream { - interceptor = messages.NewStreamingInterceptor(id, &req, payload, cfg, p.bedrockCfg, r.Header, p.AuthHeader(), tracer) + interceptor = messages.NewStreamingInterceptor(id, &req, payload, cfg, p.bedrockCfg, r.Header, authHeaderName, tracer) } else { - interceptor = messages.NewBlockingInterceptor(id, &req, payload, cfg, p.bedrockCfg, r.Header, p.AuthHeader(), tracer) + interceptor = messages.NewBlockingInterceptor(id, &req, payload, cfg, p.bedrockCfg, r.Header, authHeaderName, tracer) } span.SetAttributes(interceptor.TraceAttributes(r)...) return interceptor, nil diff --git a/provider/anthropic_test.go b/provider/anthropic_test.go index e815f28b..484589cc 100644 --- a/provider/anthropic_test.go +++ b/provider/anthropic_test.go @@ -86,9 +86,8 @@ func TestAnthropic_CreateInterceptor(t *testing.T) { body := `{"model": "claude-opus-4-5", "max_tokens": 1024, "messages": [{"role": "user", "content": "hello"}], "stream": false}` req := httptest.NewRequest(http.MethodPost, routeMessages, bytes.NewBufferString(body)) req.Header.Set("Anthropic-Beta", betaHeader) - // Simulate a client sending its own auth credential, which must be replaced - // by aibridge with the configured provider key. - req.Header.Set("Authorization", "Bearer fake-client-bearer") + // Simulate BYOK: the client sends its own bearer token for upstream auth. + req.Header.Set("Authorization", "Bearer user-oauth-token") w := httptest.NewRecorder() interceptor, err := provider.CreateInterceptor(w, req, testTracer) @@ -105,9 +104,10 @@ func TestAnthropic_CreateInterceptor(t *testing.T) { // Verify the full Anthropic-Beta header (all betas) was forwarded unchanged. assert.Equal(t, betaHeader, receivedHeaders.Get("Anthropic-Beta"), "Anthropic-Beta header must be forwarded unchanged to upstream") - // Verify aibridge's configured key was used and the client's auth credential was not forwarded. - assert.Equal(t, "test-key", receivedHeaders.Get("X-Api-Key"), "upstream must receive configured provider key") - assert.Empty(t, receivedHeaders.Get("Authorization"), "client Authorization header must not reach upstream") + // The client sent Authorization: Bearer, so BYOK bearer mode is active. + // The SDK uses Authorization (not X-Api-Key) for bearer auth. + assert.Empty(t, receivedHeaders.Get("X-Api-Key"), "X-Api-Key must not be set in BYOK bearer mode") + assert.Equal(t, "Bearer user-oauth-token", receivedHeaders.Get("Authorization"), "upstream must receive the client's bearer token") }) byokTests := []struct { From c7f7b60a234036b4075580ee37af7c56cb41aff4 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 18 Mar 2026 16:04:29 +0000 Subject: [PATCH 3/4] chore: add logging --- intercept/messages/base.go | 7 +++++++ passthrough.go | 12 ++++++++++++ utils/mask.go | 10 ++++++++++ 3 files changed, 29 insertions(+) create mode 100644 utils/mask.go diff --git a/intercept/messages/base.go b/intercept/messages/base.go index b9a05a28..f51aef66 100644 --- a/intercept/messages/base.go +++ b/intercept/messages/base.go @@ -19,6 +19,7 @@ import ( aibconfig "github.com/coder/aibridge/config" aibcontext "github.com/coder/aibridge/context" "github.com/coder/aibridge/intercept" + "github.com/coder/aibridge/utils" "github.com/coder/aibridge/intercept/apidump" "github.com/coder/aibridge/mcp" "github.com/coder/aibridge/recorder" @@ -208,8 +209,14 @@ func (i *interceptionBase) newMessagesService(ctx context.Context, opts ...optio // BYOK with OAuth token (Claude Max/Pro) uses Authorization: Bearer. // Otherwise use X-Api-Key (centralized or BYOK with personal API key). if i.cfg.BYOKBearerToken != "" { + i.logger.Debug(ctx, "using byok oauth bearer auth", + slog.F("bearer_hint", utils.MaskSecret(i.cfg.BYOKBearerToken)), + ) opts = append(opts, option.WithAuthToken(i.cfg.BYOKBearerToken)) } else { + i.logger.Debug(ctx, "using api key auth", + slog.F("api_key_hint", utils.MaskSecret(i.cfg.Key)), + ) opts = append(opts, option.WithAPIKey(i.cfg.Key)) } opts = append(opts, option.WithBaseURL(i.cfg.BaseURL)) diff --git a/passthrough.go b/passthrough.go index c6b59edd..28ab3b4b 100644 --- a/passthrough.go +++ b/passthrough.go @@ -9,6 +9,7 @@ import ( "cdr.dev/slog/v3" "github.com/coder/aibridge/intercept/apidump" + "github.com/coder/aibridge/utils" "github.com/coder/aibridge/metrics" "github.com/coder/aibridge/provider" "github.com/coder/aibridge/tracing" @@ -96,6 +97,16 @@ func newPassthroughRouter(provider provider.Provider, logger slog.Logger, m *met // Inject provider auth. provider.InjectAuthHeader(&req.Header) + + if authz := req.Header.Get("Authorization"); authz != "" { + logger.Debug(ctx, "passthrough using oauth bearer auth", + slog.F("bearer_hint", utils.MaskSecret(authz)), + ) + } else { + logger.Debug(ctx, "passthrough using api key auth", + slog.F("api_key_hint", utils.MaskSecret(req.Header.Get("X-Api-Key"))), + ) + } }, ErrorHandler: func(rw http.ResponseWriter, req *http.Request, e error) { logger.Warn(req.Context(), "reverse proxy error", slog.Error(e), slog.F("path", req.URL.Path)) @@ -117,3 +128,4 @@ func newPassthroughRouter(provider provider.Provider, logger slog.Logger, m *met proxy.ServeHTTP(w, r) } } + diff --git a/utils/mask.go b/utils/mask.go new file mode 100644 index 00000000..72de7eec --- /dev/null +++ b/utils/mask.go @@ -0,0 +1,10 @@ +package utils + +// MaskSecret returns the first 4 and last 4 characters of s +// separated by "...", or the full string if 8 characters or fewer. +func MaskSecret(s string) string { + if len(s) <= 8 { + return s + } + return s[:4] + "..." + s[len(s)-4:] +} From 5bdd5a0f2cf8d68a8c58eae42f38693d4d0d1061 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 18 Mar 2026 16:10:11 +0000 Subject: [PATCH 4/4] ci: make fmt --- intercept/messages/base.go | 2 +- passthrough.go | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/intercept/messages/base.go b/intercept/messages/base.go index f51aef66..9d134bc0 100644 --- a/intercept/messages/base.go +++ b/intercept/messages/base.go @@ -19,11 +19,11 @@ import ( aibconfig "github.com/coder/aibridge/config" aibcontext "github.com/coder/aibridge/context" "github.com/coder/aibridge/intercept" - "github.com/coder/aibridge/utils" "github.com/coder/aibridge/intercept/apidump" "github.com/coder/aibridge/mcp" "github.com/coder/aibridge/recorder" "github.com/coder/aibridge/tracing" + "github.com/coder/aibridge/utils" "github.com/coder/quartz" "github.com/tidwall/sjson" diff --git a/passthrough.go b/passthrough.go index 28ab3b4b..dc5af371 100644 --- a/passthrough.go +++ b/passthrough.go @@ -9,10 +9,10 @@ import ( "cdr.dev/slog/v3" "github.com/coder/aibridge/intercept/apidump" - "github.com/coder/aibridge/utils" "github.com/coder/aibridge/metrics" "github.com/coder/aibridge/provider" "github.com/coder/aibridge/tracing" + "github.com/coder/aibridge/utils" "github.com/coder/quartz" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -128,4 +128,3 @@ func newPassthroughRouter(provider provider.Provider, logger slog.Logger, m *met proxy.ServeHTTP(w, r) } } -