Skip to content
Open
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
4 changes: 4 additions & 0 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
15 changes: 14 additions & 1 deletion intercept/messages/base.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import (
"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"

Expand Down Expand Up @@ -205,7 +206,19 @@ 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 != "" {
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))

// Add extra headers if configured.
Expand Down
11 changes: 11 additions & 0 deletions passthrough.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
"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"
Expand Down Expand Up @@ -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))
Expand Down
30 changes: 28 additions & 2 deletions provider/anthropic.go
Original file line number Diff line number Diff line change
Expand Up @@ -110,11 +110,31 @@ 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.
Comment on lines +113 to +115
Copy link
Contributor

Choose a reason for hiding this comment

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

I'm confused by this comment, is this a reference to the coder repo? If yes, I would suggest not using http.go and making it clear that the upstream caller needs to strip Authorization and X-API-Key headers.

IIUC, at this point, the headers are as follows:

  1. Centralized: no Authorization header or X-API-Key
  2. BYOK (oauth): Authorization header includes user's oauth token, no X-API-key
  3. BYOK Api key: no Authorization header, X-API-Key is user's API key

//
// In BYOK mode, http.go only strips the BYOK header and leaves
// the user's LLM credentials intact:
// - Authorization: Bearer <oauth-token> → subscription (Claude
// Max/Pro). Set BYOKBearerToken so the SDK uses
// WithAuthToken(), and clear the centralized key.
// - X-Api-Key: <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
}
Comment on lines +129 to +131
Copy link
Contributor

Choose a reason for hiding this comment

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

From this point on, we no longer know whether this interception is using a centralized (global) key or a BYOK (user's personal) API key, right? This could be useful to store and to show in the logs (the same for BYOK oauth token). For example, if Anthropic returns a 401, we wouldn't know if the failing key is the global key (affecting everyone) or a single user's personal key.

Additionally, this is probably out of scope for this PR, but it might make sense to store this information in the interception so we can later surface it in the UI, wdyt?


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
Expand All @@ -137,6 +157,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)
}

Expand Down
123 changes: 117 additions & 6 deletions provider/anthropic_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -105,11 +104,77 @@ 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 {
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()

Expand All @@ -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()

Expand Down
10 changes: 10 additions & 0 deletions utils/mask.go
Original file line number Diff line number Diff line change
@@ -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 {
Copy link
Contributor

Choose a reason for hiding this comment

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

Is this a good idea? I think logging the auth mode ("centralized", "byok_bearer", "byok_apikey") rather than a hint of the secret might be cleaner 👀
Additionally, if we need to correlate a failure to a specific user, I believe we already log this in some cases.

if len(s) <= 8 {
return s
}
return s[:4] + "..." + s[len(s)-4:]
}
Loading