fix(aws-amplify): sync Cognito auth config and merge libraryOptions on reconfigure

.changeset/fix-init-singleton-auth-reconfigure.md

@@ -0,0 +1,5 @@
+---
+'aws-amplify': patch
+---
+
+fix(aws-amplify): merge libraryOptions when Auth is overridden and refresh default Cognito auth config on reconfigure.

packages/aws-amplify/__tests__/initSingleton.integration.test.ts

@@ -0,0 +1,123 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Integration-style tests using the real @aws-amplify/core Amplify singleton
+ * (initSingleton.test.ts mocks core).
+ */
+import { Amplify as CoreAmplify, ResourcesConfig } from '@aws-amplify/core';
+
+import { cognitoUserPoolsTokenProvider } from '../src/auth/cognito';
+import { Amplify as DefaultAmplify } from '../src';
+
+const poolAConfig: ResourcesConfig = {
+	Auth: {
+		Cognito: {
+			userPoolClientId: 'client-a',
+			userPoolId: 'pool-a',
+		},
+	},
+};
+
+const poolBConfig: ResourcesConfig = {
+	Auth: {
+		Cognito: {
+			userPoolClientId: 'client-b',
+			userPoolId: 'pool-b',
+		},
+	},
+};
+
+const storageLibraryOptions = {
+	Storage: {
+		S3: {
+			defaultAccessLevel: 'private' as const,
+			isObjectLockEnabled: true,
+		},
+	},
+};
+
+describe('DefaultAmplify.configure integration', () => {
+	let setAuthConfigSpy: jest.SpyInstance;
+	let setKeyValueStorageSpy: jest.SpyInstance;
+
+	beforeEach(() => {
+		CoreAmplify.libraryOptions = {};
+		CoreAmplify.resourcesConfig = {};
+		setAuthConfigSpy = jest.spyOn(
+			cognitoUserPoolsTokenProvider,
+			'setAuthConfig',
+		);
+		setKeyValueStorageSpy = jest.spyOn(
+			cognitoUserPoolsTokenProvider,
+			'setKeyValueStorage',
+		);
+	});
+
+	afterEach(() => {
+		setAuthConfigSpy.mockRestore();
+		setKeyValueStorageSpy.mockRestore();
+		CoreAmplify.libraryOptions = {};
+		CoreAmplify.resourcesConfig = {};
+	});
+
+	it('keeps Storage and refreshes Cognito auth config on partial reconfigure', () => {
+		DefaultAmplify.configure(poolAConfig, storageLibraryOptions);
+
+		expect(CoreAmplify.libraryOptions.Storage).toEqual(
+			storageLibraryOptions.Storage,
+		);
+		expect(CoreAmplify.libraryOptions.Auth?.tokenProvider).toBe(
+			cognitoUserPoolsTokenProvider,
+		);
+		expect(setAuthConfigSpy).toHaveBeenCalledWith(poolAConfig.Auth);
+
+		setAuthConfigSpy.mockClear();
+
+		DefaultAmplify.configure(poolBConfig, { ssr: false });
+
+		expect(CoreAmplify.libraryOptions.Storage).toEqual(
+			storageLibraryOptions.Storage,
+		);
+		expect(CoreAmplify.getConfig().Auth?.Cognito?.userPoolClientId).toBe(
+			'client-b',
+		);
+		expect(setAuthConfigSpy).toHaveBeenCalledWith(poolBConfig.Auth);
+	});
+
+	it('merges prior libraryOptions when libraryOptions.Auth overrides default provider', () => {
+		DefaultAmplify.configure(poolAConfig, storageLibraryOptions);
+
+		setAuthConfigSpy.mockClear();
+
+		DefaultAmplify.configure(poolBConfig, {
+			Auth: {
+				tokenProvider: cognitoUserPoolsTokenProvider,
+				credentialsProvider:
+					CoreAmplify.libraryOptions.Auth!.credentialsProvider!,
+			},
+		});
+
+		expect(CoreAmplify.libraryOptions.Storage).toEqual(
+			storageLibraryOptions.Storage,
+		);
+		expect(setAuthConfigSpy).toHaveBeenCalledWith(poolBConfig.Auth);
+		expect(CoreAmplify.getConfig().Auth?.Cognito?.userPoolClientId).toBe(
+			'client-b',
+		);
+	});
+
+	it('syncs default Cognito auth config when only resource config is passed', () => {
+		DefaultAmplify.configure(poolAConfig, storageLibraryOptions);
+
+		setAuthConfigSpy.mockClear();
+
+		DefaultAmplify.configure(poolBConfig);
+
+		expect(CoreAmplify.libraryOptions.Storage).toEqual(
+			storageLibraryOptions.Storage,
+		);
+		expect(setAuthConfigSpy).toHaveBeenCalledWith(poolBConfig.Auth);
+		expect(CoreAmplify.getConfig().Auth?.Cognito?.userPoolId).toBe('pool-b');
+	});
+});

packages/aws-amplify/__tests__/initSingleton.test.ts

@@ -246,16 +246,29 @@ describe('initSingleton (DefaultAmplify)', () => {
 		});
 
 		describe('when ResourcesConfig.Auth is defined', () => {
-			it('should just configure with the provided config and options when libraryOptions.Auth is defined', () => {
+			it('should merge with existing libraryOptions when libraryOptions.Auth is defined', () => {
+				const customTokenProvider = { getTokens: jest.fn() };
+				const storageLibraryOptions = {
+					S3: { defaultAccessLevel: 'private' as const },
+				};
+				AmplifySingleton.libraryOptions = {
+					Storage: storageLibraryOptions,
+				};
 				const libraryOptions = {
-					Auth: { tokenProvider: { getTokens: jest.fn() } },
+					Auth: { tokenProvider: customTokenProvider },
 				};
 				Amplify.configure(mockResourceConfig, libraryOptions);
 
 				expect(mockAmplifySingletonConfigure).toHaveBeenCalledWith(
 					mockResourceConfig,
-					libraryOptions,
+					{
+						Storage: storageLibraryOptions,
+						Auth: libraryOptions.Auth,
+					},
 				);
+				expect(
+					mockCognitoUserPoolsTokenProviderSetAuthConfig,
+				).not.toHaveBeenCalled();
 			});
 
 			describe('when the singleton libraryOptions have not yet been configured with Auth', () => {
@@ -324,38 +337,40 @@ describe('initSingleton (DefaultAmplify)', () => {
 
 				it('should preserve current auth providers (default or otherwise) and configure provider with a new CookieStorage instance', () => {
 					const libraryOptions = { ssr: true };
+					const authLibraryOptions = AmplifySingleton.libraryOptions.Auth;
 					Amplify.configure(mockResourceConfig, libraryOptions);
 
 					expect(
 						mockCognitoUserPoolsTokenProviderSetAuthConfig,
-					).not.toHaveBeenCalled();
+					).toHaveBeenCalledWith(mockResourceConfig.Auth);
 					expect(MockCookieStorage).toHaveBeenCalledWith({ sameSite: 'lax' });
 					expect(
 						mockCognitoUserPoolsTokenProviderSetKeyValueStorage,
 					).toHaveBeenCalledWith(mockCookieStorageInstance);
 					expect(mockAmplifySingletonConfigure).toHaveBeenCalledWith(
 						mockResourceConfig,
 						{
-							Auth: AmplifySingleton.libraryOptions.Auth,
+							Auth: authLibraryOptions,
 							...libraryOptions,
 						},
 					);
 				});
 
 				it('should preserve current auth providers (default or otherwise) and configure provider with defaultStorage', () => {
 					const libraryOptions = { ssr: false };
+					const authLibraryOptions = AmplifySingleton.libraryOptions.Auth;
 					Amplify.configure(mockResourceConfig, libraryOptions);
 
 					expect(
 						mockCognitoUserPoolsTokenProviderSetAuthConfig,
-					).not.toHaveBeenCalled();
+					).toHaveBeenCalledWith(mockResourceConfig.Auth);
 					expect(
 						mockCognitoUserPoolsTokenProviderSetKeyValueStorage,
 					).toHaveBeenCalledWith(defaultStorage);
 					expect(mockAmplifySingletonConfigure).toHaveBeenCalledWith(
 						mockResourceConfig,
 						{
-							Auth: AmplifySingleton.libraryOptions.Auth,
+							Auth: authLibraryOptions,
 							...libraryOptions,
 						},
 					);
@@ -365,30 +380,90 @@ describe('initSingleton (DefaultAmplify)', () => {
 					const libraryOptions = {
 						Storage: { S3: { isObjectLockEnabled: true } },
 					};
+					const authLibraryOptions = AmplifySingleton.libraryOptions.Auth;
 					Amplify.configure(mockResourceConfig, libraryOptions);
 
 					expect(
 						mockCognitoUserPoolsTokenProviderSetAuthConfig,
-					).not.toHaveBeenCalled();
+					).toHaveBeenCalledWith(mockResourceConfig.Auth);
 					expect(
 						mockCognitoUserPoolsTokenProviderSetKeyValueStorage,
 					).not.toHaveBeenCalled();
 					expect(mockAmplifySingletonConfigure).toHaveBeenCalledWith(
 						mockResourceConfig,
 						{
-							Auth: AmplifySingleton.libraryOptions.Auth,
+							Auth: authLibraryOptions,
 							...libraryOptions,
 						},
 					);
 				});
 
-				it('should just configure without touching libraryOptions', () => {
+				it('should preserve non-Auth library options when reconfiguring with partial libraryOptions', () => {
+					const storageLibraryOptions = {
+						S3: { defaultAccessLevel: 'private' as const },
+					};
+					AmplifySingleton.libraryOptions = {
+						Auth: {
+							tokenProvider: cognitoUserPoolsTokenProvider,
+							credentialsProvider: cognitoCredentialsProvider,
+						},
+						Storage: storageLibraryOptions,
+					};
+					const authLibraryOptions = AmplifySingleton.libraryOptions.Auth;
+
+					Amplify.configure(mockResourceConfig, { ssr: true });
+
+					expect(
+						mockCognitoUserPoolsTokenProviderSetAuthConfig,
+					).toHaveBeenCalledWith(mockResourceConfig.Auth);
+					expect(mockAmplifySingletonConfigure).toHaveBeenCalledWith(
+						mockResourceConfig,
+						{
+							Auth: authLibraryOptions,
+							Storage: storageLibraryOptions,
+							ssr: true,
+						},
+					);
+				});
+
+				it('should sync default Cognito auth config when reconfiguring with resource config only', () => {
 					Amplify.configure(mockResourceConfig);
 
+					expect(
+						mockCognitoUserPoolsTokenProviderSetAuthConfig,
+					).toHaveBeenCalledWith(mockResourceConfig.Auth);
 					expect(mockAmplifySingletonConfigure).toHaveBeenCalledWith(
 						mockResourceConfig,
 					);
 				});
+
+				it('should sync default Cognito auth config when libraryOptions.Auth overrides with default provider', () => {
+					const updatedResourceConfig: ResourcesConfig = {
+						Auth: {
+							Cognito: {
+								userPoolClientId: 'newClientId',
+								userPoolId: 'newPoolId',
+							},
+						},
+					};
+					AmplifySingleton.libraryOptions = {
+						Auth: {
+							tokenProvider: cognitoUserPoolsTokenProvider,
+							credentialsProvider: cognitoCredentialsProvider,
+						},
+					};
+
+					Amplify.configure(updatedResourceConfig, {
+						Auth: {
+							tokenProvider: cognitoUserPoolsTokenProvider,
+							credentialsProvider: cognitoCredentialsProvider,
+						},
+					});
+
+					expect(
+						mockCognitoUserPoolsTokenProviderSetAuthConfig,
+					).toHaveBeenCalledWith(updatedResourceConfig.Auth);
+				});
 			});
 
 			it('should invoke AmplifySingleton.configure with other provided library options', () => {

packages/aws-amplify/src/initSingleton.ts

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 import {
 	Amplify,
+	AuthConfig,
 	CookieStorage,
 	LibraryOptions,
 	ResourcesConfig,
@@ -20,6 +21,15 @@ import {
 	cognitoUserPoolsTokenProvider,
 } from './auth/cognito';
 
+const usesDefaultCognitoTokenProvider = (
+	authLibraryOptions?: LibraryOptions['Auth'],
+): boolean =>
+	authLibraryOptions?.tokenProvider === cognitoUserPoolsTokenProvider;
+
+const syncDefaultCognitoAuthConfig = (authConfig: AuthConfig): void => {
+	cognitoUserPoolsTokenProvider.setAuthConfig(authConfig);
+};
+
 export const DefaultAmplify = {
 	/**
 	 * Configures Amplify with the {@link resourceConfig} and {@link libraryOptions}.
@@ -56,10 +66,25 @@ export const DefaultAmplify = {
 			return;
 		}
 
-		// If Auth options are provided, always just configure as is.
-		// Otherwise, we can assume no Auth libraryOptions were provided from here on.
+		// If Auth options are provided, merge with existing libraryOptions so other categories are preserved.
 		if (libraryOptions?.Auth) {
-			Amplify.configure(resolvedResourceConfig, libraryOptions);
+			const mergedLibraryOptions: LibraryOptions = {
+				...Amplify.libraryOptions,
+				...libraryOptions,
+			};
+
+			if (usesDefaultCognitoTokenProvider(mergedLibraryOptions.Auth)) {
+				syncDefaultCognitoAuthConfig(resolvedResourceConfig.Auth);
+
+				if (libraryOptions.ssr !== undefined) {
+					cognitoUserPoolsTokenProvider.setKeyValueStorage(
+						// TODO: allow configure with a public interface
+						resolvedKeyValueStorage,
+					);
+				}
+			}
+
+			Amplify.configure(resolvedResourceConfig, mergedLibraryOptions);
 
 			return;
 		}
@@ -97,16 +122,25 @@ export const DefaultAmplify = {
 				authLibraryOptions.credentialsProvider = resolvedCredentialsProvider;
 			}
 
+			if (usesDefaultCognitoTokenProvider(authLibraryOptions)) {
+				syncDefaultCognitoAuthConfig(resolvedResourceConfig.Auth);
+			}
+
 			Amplify.configure(resolvedResourceConfig, {
-				Auth: authLibraryOptions,
+				...Amplify.libraryOptions,
 				...libraryOptions,
+				Auth: authLibraryOptions,
 			});
 
 			return;
 		}
 
 		// Finally, if there were no libraryOptions given at all, we should simply not touch the currently
 		// configured libraryOptions.
+		if (usesDefaultCognitoTokenProvider(Amplify.libraryOptions.Auth)) {
+			syncDefaultCognitoAuthConfig(resolvedResourceConfig.Auth);
+		}
+
 		Amplify.configure(resolvedResourceConfig);
 	},
 	/**