diff --git a/coderd/aibridged/aibridged_test.go b/coderd/aibridged/aibridged_test.go index caa162888c..8b29b2653a 100644 --- a/coderd/aibridged/aibridged_test.go +++ b/coderd/aibridged/aibridged_test.go @@ -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, diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index 5d75fa623f..640716f37a 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -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, })