feat(auth): add custom headers support for Cognito Auth requests (#14756)

* chore: reorder jest.mock declarations to match import order

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* feat(auth): add custom headers support for Cognito Auth requests

Add optional `headers` function to `LibraryAuthOptions` that allows
attaching custom headers to all Auth requests to Cognito, consistent
with the existing REST and GraphQL custom headers pattern.

Closes #12308
Closes #13197

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* chore: add changeset for auth custom headers

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix(auth): preserve existing request headers in cache middleware

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: piotrekwitkowski <piotrekwitkowski@users.noreply.github.com>

Author: piotrekwitkowski

.changeset/auth-custom-headers.md

@@ -0,0 +1,6 @@
+---
+'@aws-amplify/auth': minor
+'@aws-amplify/core': patch
+---
+
+Add custom headers support for Cognito Auth requests via `libraryOptions.Auth.headers`

packages/auth/__tests__/foundation/factories/serviceClients/cognitoIdentityProvider/shared/handler/cognitoUserPoolTransferHandler.test.ts

@@ -1,10 +1,12 @@
+import { Amplify } from '@aws-amplify/core';
 import { unauthenticatedHandler } from '@aws-amplify/core/internals/aws-client-utils';
 import { composeTransferHandler } from '@aws-amplify/core/internals/aws-client-utils/composers';
 
 import { cognitoUserPoolTransferHandler } from '../../../../../../../src/foundation/factories/serviceClients/cognitoIdentityProvider/shared/handler';
 
-jest.mock('@aws-amplify/core/internals/aws-client-utils/composers');
+jest.mock('@aws-amplify/core');
 jest.mock('@aws-amplify/core/internals/aws-client-utils');
+jest.mock('@aws-amplify/core/internals/aws-client-utils/composers');
 
 const mockComposeTransferHandler = jest.mocked(composeTransferHandler);
 const mockUnauthenticatedHandler = jest.mocked(unauthenticatedHandler);
@@ -16,7 +18,7 @@ describe('cognitoUserPoolTransferHandler', () => {
 		const _ = cognitoUserPoolTransferHandler;
 	});
 
-	it('adds the disableCacheMiddlewareFactory at module loading', () => {
+	it('adds the disableCacheMiddlewareFactory at module loading', async () => {
 		expect(mockComposeTransferHandler).toHaveBeenCalledTimes(1);
 
 		const [core, middleware] = mockComposeTransferHandler.mock.calls[0];
@@ -33,11 +35,45 @@ describe('cognitoUserPoolTransferHandler', () => {
 			headers: {},
 		};
 
-		disableCacheMiddleware(mockRequest);
+		await disableCacheMiddleware(mockRequest);
 
 		expect(mockNext).toHaveBeenCalledWith(mockRequest);
 		expect(mockRequest.headers).toEqual({
 			'cache-control': 'no-store',
 		});
 	});
+
+	it('attaches custom headers from libraryOptions when configured', async () => {
+		const mockHeaders = jest.fn().mockResolvedValue({
+			'custom-header': 'custom-value',
+		});
+		(Amplify as any).libraryOptions = {
+			Auth: { headers: mockHeaders },
+		};
+
+		const [, middleware] = mockComposeTransferHandler.mock.calls[0];
+		const disableCacheMiddlewareFactory = middleware[0] as any;
+		const middlewareFn = disableCacheMiddlewareFactory()(jest.fn());
+		const mockRequest = { headers: {} as Record<string, string> };
+
+		await middlewareFn(mockRequest);
+
+		expect(mockHeaders).toHaveBeenCalled();
+		expect(mockRequest.headers['custom-header']).toBe('custom-value');
+	});
+
+	it('does not attach custom headers when not configured', async () => {
+		(Amplify as any).libraryOptions = {};
+
+		const [, middleware] = mockComposeTransferHandler.mock.calls[0];
+		const disableCacheMiddlewareFactory = middleware[0] as any;
+		const mockNext = jest.fn();
+		const middlewareFn = disableCacheMiddlewareFactory()(mockNext);
+		const mockRequest = { headers: {} as Record<string, string> };
+
+		await middlewareFn(mockRequest);
+
+		expect(mockRequest.headers['custom-header']).toBeUndefined();
+		expect(mockNext).toHaveBeenCalledWith(mockRequest);
+	});
 });

packages/auth/src/foundation/factories/serviceClients/cognitoIdentityProvider/shared/handler/cognitoUserPoolTransferHandler.ts

@@ -1,6 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+import { Amplify } from '@aws-amplify/core';
 import { composeTransferHandler } from '@aws-amplify/core/internals/aws-client-utils/composers';
 import {
 	HttpRequest,
@@ -18,7 +19,11 @@ const disableCacheMiddlewareFactory: Middleware<
 	Record<string, unknown>
 > = () => (next, _) =>
 	async function disableCacheMiddleware(request) {
-		request.headers['cache-control'] = 'no-store';
+		request.headers = {
+			...request.headers,
+			'cache-control': 'no-store',
+			...(await Amplify.libraryOptions?.Auth?.headers?.()),
+		};
 
 		return next(request);
 	};

packages/core/src/singleton/Auth/types.ts

@@ -1,6 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+import { Headers } from '../../clients';
 import { StrictUnion } from '../../types';
 import { AtLeastOne } from '../types';
 
@@ -56,6 +57,7 @@ export interface AuthSession {
 export interface LibraryAuthOptions {
 	tokenProvider?: TokenProvider;
 	credentialsProvider?: CredentialsAndIdentityIdProvider;
+	headers?(): Promise<Headers>;
 }
 
 export interface Identity {