fix: extract key when BYOK header is given with delegated auth (#25688)

Previously we were only extracting the API when _not_ delegating auth;
this is incorrect.

We need to extract the key _always_ when BYOK is intended.

---------

Signed-off-by: Danny Kopping <danny@coder.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Danny Kopping
2026-05-26 19:46:26 +02:00
committed by GitHub
co-authored by Claude Opus 4.7
parent d80b484487
commit 5d8ca2e5ce
2 changed files with 44 additions and 27 deletions
+21 -5
View File
@@ -196,6 +196,10 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) {
expectAbsent []string
}{
{
// Delegated + centralized: identity comes from the
// api key ID on the context, in lieu of a session
// token. No header credentials are sent and SessionKey
// is empty downstream.
name: "valid centralized",
applyMocks: func(t *testing.T, client *mock.MockDRPCClient, pool *mock.MockPooler, mockH *mockHandler) {
client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).DoAndReturn(
@@ -208,7 +212,12 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) {
Username: "u",
}, nil
})
pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mockH, nil)
pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
func(_ context.Context, req aibridged.Request, _ aibridged.ClientFunc, _ aibridged.MCPProxyBuilder) (http.Handler, error) {
assert.Empty(t, req.SessionKey,
"delegated centralized request carries no session token")
return mockH, nil
})
},
expectStatus: http.StatusOK,
expectHandled: true,
@@ -222,18 +231,25 @@ func TestServeHTTP_DelegatedAPIKey(t *testing.T) {
name: "valid BYOK preserves user credentials",
reqHeaders: map[string]string{
// Marks BYOK; this header must be stripped before
// forwarding upstream.
agplaibridge.HeaderCoderToken: "should-not-be-present",
// forwarding upstream. Its value is what gets
// surfaced downstream as the SessionKey because
// ExtractAuthToken prefers HeaderCoderToken.
agplaibridge.HeaderCoderToken: "coder-token-byok",
// The user's own LLM credential; must be preserved.
"Authorization": "Bearer sk-ant-oat01-user-token",
},
applyMocks: func(_ *testing.T, client *mock.MockDRPCClient, pool *mock.MockPooler, mockH *mockHandler) {
applyMocks: func(t *testing.T, client *mock.MockDRPCClient, pool *mock.MockPooler, mockH *mockHandler) {
client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).Return(&proto.IsAuthorizedResponse{
OwnerId: uuid.NewString(),
ApiKeyId: testKeyID,
Username: "u",
}, nil)
pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mockH, nil)
pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
func(_ context.Context, req aibridged.Request, _ aibridged.ClientFunc, _ aibridged.MCPProxyBuilder) (http.Handler, error) {
assert.Equal(t, "coder-token-byok", req.SessionKey,
"BYOK delegated request must still surface the extracted Coder token as SessionKey")
return mockH, nil
})
},
expectStatus: http.StatusOK,
expectHandled: true,
+23 -22
View File
@@ -64,30 +64,31 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
// the user's own LLM credentials in Authorization/X-Api-Key when BYOK
// is in effect.
var (
authReq *proto.IsAuthorizedRequest
sessionKey string
delegated bool
authReq *proto.IsAuthorizedRequest
)
if delegatedID, ok := agplaibridge.DelegatedAPIKeyIDFromContext(ctx); ok {
authReq = &proto.IsAuthorizedRequest{KeyId: delegatedID}
delegated = true
// SessionKey is consumed only by the injected MCP path, which is
// not available to delegated callers (they have no secret).
} else {
key := strings.TrimSpace(agplaibridge.ExtractAuthToken(r.Header))
if key == "" {
// Some clients (e.g. Claude) send a HEAD request
// without credentials to check connectivity.
if r.Method == http.MethodHead {
logger.Info(ctx, "unauthenticated HEAD request")
} else {
logger.Warn(ctx, "no auth key provided")
}
http.Error(rw, ErrNoAuthKey.Error(), http.StatusBadRequest)
return
delegatedID, delegated := agplaibridge.DelegatedAPIKeyIDFromContext(ctx)
key := strings.TrimSpace(agplaibridge.ExtractAuthToken(r.Header))
// When a BYOK header is present, a key is ALWAYS required.
// Delegated auth only requires a key when using BYOK.
if key == "" && !delegated {
// Some clients (e.g. Claude) send a HEAD request
// without credentials to check connectivity.
if r.Method == http.MethodHead {
logger.Info(ctx, "unauthenticated HEAD request")
} else {
logger.Warn(ctx, "no auth key provided")
}
http.Error(rw, ErrNoAuthKey.Error(), http.StatusBadRequest)
return
}
if delegated {
authReq = &proto.IsAuthorizedRequest{KeyId: delegatedID}
} else {
authReq = &proto.IsAuthorizedRequest{Key: key}
sessionKey = key
}
// Strip every header that may carry the Coder token so it is never
@@ -151,7 +152,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
}
handler, err := s.GetRequestHandler(ctx, Request{
SessionKey: sessionKey,
SessionKey: key,
APIKeyID: resp.ApiKeyId,
InitiatorID: id,
})