@aws-amplify/pubsub 4.2.5 → 4.2.6-custom-pk.86

package/src/Providers/{AWSAppSyncRealTimeProvider.ts → AWSAppSyncRealTimeProvider/index.ts}

@@ -15,7 +15,7 @@ import { GraphQLError } from 'graphql';
 import * as url from 'url';
 import { v4 as uuid } from 'uuid';
 import { Buffer } from 'buffer';
-import { ProvidertOptions } from '../types';
+import { ProviderOptions } from '../../types';
 import {
 	Logger,
 	Credentials,
@@ -25,135 +25,77 @@ import {
 	USER_AGENT_HEADER,
 	jitteredExponentialRetry,
 	NonRetryableError,
+	ICredentials,
 } from '@aws-amplify/core';
 import Cache from '@aws-amplify/cache';
-import Auth from '@aws-amplify/auth';
-import { AbstractPubSubProvider } from './PubSubProvider';
-import { CONTROL_MSG } from '../index';
+import Auth, { GRAPHQL_AUTH_MODE } from '@aws-amplify/auth';
+import { AbstractPubSubProvider } from '../PubSubProvider';
+import { CONTROL_MSG } from '../../index';
+import {
+	AMPLIFY_SYMBOL,
+	AWS_APPSYNC_REALTIME_HEADERS,
+	CONNECTION_INIT_TIMEOUT,
+	DEFAULT_KEEP_ALIVE_TIMEOUT,
+	MAX_DELAY_MS,
+	MESSAGE_TYPES,
+	NON_RETRYABLE_CODES,
+	SOCKET_STATUS,
+	START_ACK_TIMEOUT,
+	SUBSCRIPTION_STATUS,
+} from './constants';
 
 const logger = new Logger('AWSAppSyncRealTimeProvider');
 
-const AMPLIFY_SYMBOL = (typeof Symbol !== 'undefined' &&
-typeof Symbol.for === 'function'
-	? Symbol.for('amplify_default')
-	: '@@amplify_default') as Symbol;
-
 const dispatchApiEvent = (event: string, data: any, message: string) => {
 	Hub.dispatch('api', { event, data, message }, 'PubSub', AMPLIFY_SYMBOL);
 };
 
-const MAX_DELAY_MS = 5000;
-
-const NON_RETRYABLE_CODES = [400, 401, 403];
-
-type ObserverQuery = {
+export type ObserverQuery = {
 	observer: ZenObservable.SubscriptionObserver<any>;
 	query: string;
 	variables: object;
 	subscriptionState: SUBSCRIPTION_STATUS;
 	subscriptionReadyCallback?: Function;
 	subscriptionFailedCallback?: Function;
-	startAckTimeoutId?;
+	startAckTimeoutId?: ReturnType<typeof setTimeout>;
 };
 
-enum MESSAGE_TYPES {
-	/**
-	 * Client -> Server message.
-	 * This message type is the first message after handshake and this will initialize AWS AppSync RealTime communication
-	 */
-	GQL_CONNECTION_INIT = 'connection_init',
-	/**
-	 * Server -> Client message
-	 * This message type is in case there is an issue with AWS AppSync RealTime when establishing connection
-	 */
-	GQL_CONNECTION_ERROR = 'connection_error',
-	/**
-	 * Server -> Client message.
-	 * This message type is for the ack response from AWS AppSync RealTime for GQL_CONNECTION_INIT message
-	 */
-	GQL_CONNECTION_ACK = 'connection_ack',
-	/**
-	 * Client -> Server message.
-	 * This message type is for register subscriptions with AWS AppSync RealTime
-	 */
-	GQL_START = 'start',
-	/**
-	 * Server -> Client message.
-	 * This message type is for the ack response from AWS AppSync RealTime for GQL_START message
-	 */
-	GQL_START_ACK = 'start_ack',
-	/**
-	 * Server -> Client message.
-	 * This message type is for subscription message from AWS AppSync RealTime
-	 */
-	GQL_DATA = 'data',
-	/**
-	 * Server -> Client message.
-	 * This message type helps the client to know is still receiving messages from AWS AppSync RealTime
-	 */
-	GQL_CONNECTION_KEEP_ALIVE = 'ka',
-	/**
-	 * Client -> Server message.
-	 * This message type is for unregister subscriptions with AWS AppSync RealTime
-	 */
-	GQL_STOP = 'stop',
-	/**
-	 * Server -> Client message.
-	 * This message type is for the ack response from AWS AppSync RealTime for GQL_STOP message
-	 */
-	GQL_COMPLETE = 'complete',
-	/**
-	 * Server -> Client message.
-	 * This message type is for sending error messages from AWS AppSync RealTime to the client
-	 */
-	GQL_ERROR = 'error', // Server -> Client
-}
+const standardDomainPattern =
+	/^https:\/\/\w{26}\.appsync\-api\.\w{2}(?:(?:\-\w{2,})+)\-\d\.amazonaws.com\/graphql$/i;
 
-enum SUBSCRIPTION_STATUS {
-	PENDING,
-	CONNECTED,
-	FAILED,
-}
+const customDomainPath = '/realtime';
 
-enum SOCKET_STATUS {
-	CLOSED,
-	READY,
-	CONNECTING,
+type GraphqlAuthModes = keyof typeof GRAPHQL_AUTH_MODE;
+
+export interface AWSAppSyncRealTimeProviderOptions extends ProviderOptions {
+	appSyncGraphqlEndpoint?: string;
+	authenticationType?: GraphqlAuthModes;
+	query?: string;
+	variables?: object;
+	apiKey?: string;
+	region?: string;
+	graphql_headers?: () => {} | (() => Promise<{}>);
+	additionalHeaders?: { [key: string]: string };
 }
 
-const AWS_APPSYNC_REALTIME_HEADERS = {
-	accept: 'application/json, text/javascript',
-	'content-encoding': 'amz-1.0',
-	'content-type': 'application/json; charset=UTF-8',
-};
-
-/**
- * Time in milleseconds to wait for GQL_CONNECTION_INIT message
- */
-const CONNECTION_INIT_TIMEOUT = 15000;
-
-/**
- * Time in milleseconds to wait for GQL_START_ACK message
- */
-const START_ACK_TIMEOUT = 15000;
-
-/**
- * Default Time in milleseconds to wait for GQL_CONNECTION_KEEP_ALIVE message
- */
-const DEFAULT_KEEP_ALIVE_TIMEOUT = 5 * 60 * 1000;
-
-const standardDomainPattern = /^https:\/\/\w{26}\.appsync\-api\.\w{2}(?:(?:\-\w{2,})+)\-\d\.amazonaws.com\/graphql$/i;
-
-const customDomainPath = '/realtime';
+type AWSAppSyncRealTimeAuthInput =
+	Partial<AWSAppSyncRealTimeProviderOptions> & {
+		canonicalUri: string;
+		payload: string;
+	};
 
 export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
-	private awsRealTimeSocket: WebSocket;
+	private awsRealTimeSocket?: WebSocket;
 	private socketStatus: SOCKET_STATUS = SOCKET_STATUS.CLOSED;
-	private keepAliveTimeoutId;
+	private keepAliveTimeoutId?: ReturnType<typeof setTimeout>;
 	private keepAliveTimeout = DEFAULT_KEEP_ALIVE_TIMEOUT;
 	private subscriptionObserverMap: Map<string, ObserverQuery> = new Map();
 	private promiseArray: Array<{ res: Function; rej: Function }> = [];
 
+	getNewWebSocket(url, protocol) {
+		return new WebSocket(url, protocol);
+	}
+
 	getProviderName() {
 		return 'AWSAppSyncRealTimeProvider';
 	}
@@ -173,12 +115,12 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 
 	subscribe(
 		_topics: string[] | string,
-		options?: ProvidertOptions
+		options?: AWSAppSyncRealTimeProviderOptions
 	): Observable<any> {
-		const { appSyncGraphqlEndpoint } = options;
+		const appSyncGraphqlEndpoint = options?.appSyncGraphqlEndpoint;
 
 		return new Observable(observer => {
-			if (!appSyncGraphqlEndpoint) {
+			if (!options || !appSyncGraphqlEndpoint) {
 				observer.error({
 					errors: [
 						{
@@ -195,7 +137,7 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 					options,
 					observer,
 					subscriptionId,
-				}).catch(err => {
+				}).catch<any>(err => {
 					observer.error({
 						errors: [
 							{
@@ -241,10 +183,15 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		return !this.options
 			.aws_appsync_dangerously_connect_to_http_endpoint_for_testing;
 	}
+
 	private async _startSubscriptionWithAWSAppSyncRealTime({
 		options,
 		observer,
 		subscriptionId,
+	}: {
+		options: AWSAppSyncRealTimeProviderOptions;
+		observer: ZenObservable.SubscriptionObserver<any>;
+		subscriptionId: string;
 	}) {
 		const {
 			appSyncGraphqlEndpoint,
@@ -265,10 +212,10 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		// Having a subscription id map will make it simple to forward messages received
 		this.subscriptionObserverMap.set(subscriptionId, {
 			observer,
-			query,
-			variables,
+			query: query ?? '',
+			variables: variables ?? {},
 			subscriptionState,
-			startAckTimeoutId: null,
+			startAckTimeoutId: undefined,
 		});
 
 		// Preparing payload for subscription message
@@ -314,7 +261,7 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 			});
 		} catch (err) {
 			logger.debug({ err });
-			const { message = '' } = err;
+			const message = err['message'] ?? '';
 			observer.error({
 				errors: [
 					{
@@ -323,7 +270,6 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 				],
 			});
 			observer.complete();
-
 			const { subscriptionFailedCallback } =
 				this.subscriptionObserverMap.get(subscriptionId) || {};
 
@@ -338,17 +284,15 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		// E.g.unsubscribe gets invoked prior to finishing WebSocket handshake or START_ACK.
 		// Both subscriptionFailedCallback and subscriptionReadyCallback are used to synchronized this.
 
-		const {
-			subscriptionFailedCallback,
-			subscriptionReadyCallback,
-		} = this.subscriptionObserverMap.get(subscriptionId);
+		const { subscriptionFailedCallback, subscriptionReadyCallback } =
+			this.subscriptionObserverMap.get(subscriptionId) ?? {};
 
 		// This must be done before sending the message in order to be listening immediately
 		this.subscriptionObserverMap.set(subscriptionId, {
 			observer,
 			subscriptionState,
-			variables,
-			query,
+			query: query ?? '',
+			variables: variables ?? {},
 			subscriptionReadyCallback,
 			subscriptionFailedCallback,
 			startAckTimeoutId: setTimeout(() => {
@@ -361,32 +305,30 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 	}
 
 	// Waiting that subscription has been connected before trying to unsubscribe
-	private async _waitForSubscriptionToBeConnected(subscriptionId) {
-		const { subscriptionState } = this.subscriptionObserverMap.get(
-			subscriptionId
-		);
-		// This in case unsubscribe is invoked before sending start subscription message
-		if (subscriptionState === SUBSCRIPTION_STATUS.PENDING) {
-			return new Promise((res, rej) => {
-				const {
-					observer,
-					subscriptionState,
-					variables,
-					query,
-				} = this.subscriptionObserverMap.get(subscriptionId);
-				this.subscriptionObserverMap.set(subscriptionId, {
-					observer,
-					subscriptionState,
-					variables,
-					query,
-					subscriptionReadyCallback: res,
-					subscriptionFailedCallback: rej,
+	private async _waitForSubscriptionToBeConnected(subscriptionId: string) {
+		const subscriptionObserver =
+			this.subscriptionObserverMap.get(subscriptionId);
+		if (subscriptionObserver) {
+			const { subscriptionState } = subscriptionObserver;
+			// This in case unsubscribe is invoked before sending start subscription message
+			if (subscriptionState === SUBSCRIPTION_STATUS.PENDING) {
+				return new Promise((res, rej) => {
+					const { observer, subscriptionState, variables, query } =
+						subscriptionObserver;
+					this.subscriptionObserverMap.set(subscriptionId, {
+						observer,
+						subscriptionState,
+						variables,
+						query,
+						subscriptionReadyCallback: res,
+						subscriptionFailedCallback: rej,
+					});
 				});
-			});
+			}
 		}
 	}
 
-	private _sendUnsubscriptionMessage(subscriptionId) {
+	private _sendUnsubscriptionMessage(subscriptionId: string) {
 		try {
 			if (
 				this.awsRealTimeSocket &&
@@ -407,7 +349,7 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		}
 	}
 
-	private _removeSubscriptionObserver(subscriptionId) {
+	private _removeSubscriptionObserver(subscriptionId: string) {
 		this.subscriptionObserverMap.delete(subscriptionId);
 
 		// Verifying 1000ms after removing subscription in case there are new subscription unmount/mount
@@ -429,13 +371,13 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 			setTimeout(this._closeSocketIfRequired.bind(this), 1000);
 		} else {
 			logger.debug('closing WebSocket...');
-			clearTimeout(this.keepAliveTimeoutId);
+			if (this.keepAliveTimeoutId) clearTimeout(this.keepAliveTimeoutId);
 			const tempSocket = this.awsRealTimeSocket;
 			// Cleaning callbacks to avoid race condition, socket still exists
-			tempSocket.onclose = undefined;
-			tempSocket.onerror = undefined;
+			tempSocket.onclose = null;
+			tempSocket.onerror = null;
 			tempSocket.close(1000);
-			this.awsRealTimeSocket = null;
+			this.awsRealTimeSocket = undefined;
 			this.socketStatus = SOCKET_STATUS.CLOSED;
 		}
 	}
@@ -472,29 +414,31 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 			if (typeof subscriptionReadyCallback === 'function') {
 				subscriptionReadyCallback();
 			}
-			clearTimeout(startAckTimeoutId);
+			if (startAckTimeoutId) clearTimeout(startAckTimeoutId);
 			dispatchApiEvent(
 				CONTROL_MSG.SUBSCRIPTION_ACK,
 				{ query, variables },
 				'Connection established for subscription'
 			);
 			const subscriptionState = SUBSCRIPTION_STATUS.CONNECTED;
-			this.subscriptionObserverMap.set(id, {
-				observer,
-				query,
-				variables,
-				startAckTimeoutId: null,
-				subscriptionState,
-				subscriptionReadyCallback,
-				subscriptionFailedCallback,
-			});
+			if (observer) {
+				this.subscriptionObserverMap.set(id, {
+					observer,
+					query,
+					variables,
+					startAckTimeoutId: undefined,
+					subscriptionState,
+					subscriptionReadyCallback,
+					subscriptionFailedCallback,
+				});
+			}
 
 			// TODO: emit event on hub but it requires to store the id first
 			return;
 		}
 
 		if (type === MESSAGE_TYPES.GQL_CONNECTION_KEEP_ALIVE) {
-			clearTimeout(this.keepAliveTimeoutId);
+			if (this.keepAliveTimeoutId) clearTimeout(this.keepAliveTimeoutId);
 			this.keepAliveTimeoutId = setTimeout(
 				this._errorDisconnect.bind(this, CONTROL_MSG.TIMEOUT_DISCONNECT),
 				this.keepAliveTimeout
@@ -504,30 +448,32 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 
 		if (type === MESSAGE_TYPES.GQL_ERROR) {
 			const subscriptionState = SUBSCRIPTION_STATUS.FAILED;
-			this.subscriptionObserverMap.set(id, {
-				observer,
-				query,
-				variables,
-				startAckTimeoutId,
-				subscriptionReadyCallback,
-				subscriptionFailedCallback,
-				subscriptionState,
-			});
+			if (observer) {
+				this.subscriptionObserverMap.set(id, {
+					observer,
+					query,
+					variables,
+					startAckTimeoutId,
+					subscriptionReadyCallback,
+					subscriptionFailedCallback,
+					subscriptionState,
+				});
 
-			observer.error({
-				errors: [
-					{
-						...new GraphQLError(
-							`${CONTROL_MSG.CONNECTION_FAILED}: ${JSON.stringify(payload)}`
-						),
-					},
-				],
-			});
-			clearTimeout(startAckTimeoutId);
+				observer.error({
+					errors: [
+						{
+							...new GraphQLError(
+								`${CONTROL_MSG.CONNECTION_FAILED}: ${JSON.stringify(payload)}`
+							),
+						},
+					],
+				});
+				if (startAckTimeoutId) clearTimeout(startAckTimeoutId);
 
-			observer.complete();
-			if (typeof subscriptionFailedCallback === 'function') {
-				subscriptionFailedCallback();
+				observer.complete();
+				if (typeof subscriptionFailedCallback === 'function') {
+					subscriptionFailedCallback();
+				}
 			}
 		}
 	}
@@ -549,39 +495,42 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		this.socketStatus = SOCKET_STATUS.CLOSED;
 	}
 
-	private _timeoutStartSubscriptionAck(subscriptionId) {
-		const { observer, query, variables } =
-			this.subscriptionObserverMap.get(subscriptionId) || {};
-		if (!observer) {
-			return;
-		}
-		this.subscriptionObserverMap.set(subscriptionId, {
-			observer,
-			query,
-			variables,
-			subscriptionState: SUBSCRIPTION_STATUS.FAILED,
-		});
-
-		if (observer && !observer.closed) {
-			observer.error({
-				errors: [
-					{
-						...new GraphQLError(
-							`Subscription timeout ${JSON.stringify({
-								query,
-								variables,
-							})}`
-						),
-					},
-				],
+	private _timeoutStartSubscriptionAck(subscriptionId: string) {
+		const subscriptionObserver =
+			this.subscriptionObserverMap.get(subscriptionId);
+		if (subscriptionObserver) {
+			const { observer, query, variables } = subscriptionObserver;
+			if (!observer) {
+				return;
+			}
+			this.subscriptionObserverMap.set(subscriptionId, {
+				observer,
+				query,
+				variables,
+				subscriptionState: SUBSCRIPTION_STATUS.FAILED,
 			});
-			// Cleanup will be automatically executed
-			observer.complete();
+
+			if (observer && !observer.closed) {
+				observer.error({
+					errors: [
+						{
+							...new GraphQLError(
+								`Subscription timeout ${JSON.stringify({
+									query,
+									variables,
+								})}`
+							),
+						},
+					],
+				});
+				// Cleanup will be automatically executed
+				observer.complete();
+			}
+			logger.debug(
+				'timeoutStartSubscription',
+				JSON.stringify({ query, variables })
+			);
 		}
-		logger.debug(
-			'timeoutStartSubscription',
-			JSON.stringify({ query, variables })
-		);
 	}
 
 	private _initializeWebSocketConnection({
@@ -590,7 +539,7 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		apiKey,
 		region,
 		additionalHeaders,
-	}) {
+	}: AWSAppSyncRealTimeProviderOptions) {
 		if (this.socketStatus === SOCKET_STATUS.READY) {
 			return;
 		}
@@ -617,17 +566,18 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 
 					const payloadQs = Buffer.from(payloadString).toString('base64');
 
-					let discoverableEndpoint = appSyncGraphqlEndpoint;
+					let discoverableEndpoint = appSyncGraphqlEndpoint ?? '';
 
 					if (this.isCustomDomain(discoverableEndpoint)) {
-						discoverableEndpoint = discoverableEndpoint.concat(
-							customDomainPath
-						);
+						discoverableEndpoint =
+							discoverableEndpoint.concat(customDomainPath);
 					} else {
-						discoverableEndpoint = discoverableEndpoint.replace('appsync-api', 'appsync-realtime-api').replace('gogi-beta', 'grt-beta');
+						discoverableEndpoint = discoverableEndpoint
+							.replace('appsync-api', 'appsync-realtime-api')
+							.replace('gogi-beta', 'grt-beta');
 					}
 
-				    // Creating websocket url with required query strings
+					// Creating websocket url with required query strings
 					const protocol = this.isSSLEnabled ? 'wss://' : 'ws://';
 					discoverableEndpoint = discoverableEndpoint
 						.replace('https://', protocol)
@@ -635,7 +585,7 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 
 					const awsRealTimeUrl = `${discoverableEndpoint}?header=${headerQs}&payload=${payloadQs}`;
 
-					await this._initializeRetryableHandshake({ awsRealTimeUrl });
+					await this._initializeRetryableHandshake(awsRealTimeUrl);
 
 					this.promiseArray.forEach(({ res }) => {
 						logger.debug('Notifying connection successful');
@@ -652,30 +602,30 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 					) {
 						this.awsRealTimeSocket.close(3001);
 					}
-					this.awsRealTimeSocket = null;
+					this.awsRealTimeSocket = undefined;
 					this.socketStatus = SOCKET_STATUS.CLOSED;
 				}
 			}
 		});
 	}
 
-	private async _initializeRetryableHandshake({ awsRealTimeUrl }) {
+	private async _initializeRetryableHandshake(awsRealTimeUrl: string) {
 		logger.debug(`Initializaling retryable Handshake`);
 		await jitteredExponentialRetry(
 			this._initializeHandshake.bind(this),
-			[{ awsRealTimeUrl }],
+			[awsRealTimeUrl],
 			MAX_DELAY_MS
 		);
 	}
 
-	private async _initializeHandshake({ awsRealTimeUrl }) {
+	private async _initializeHandshake(awsRealTimeUrl: string) {
 		logger.debug(`Initializing handshake ${awsRealTimeUrl}`);
 		// Because connecting the socket is async, is waiting until connection is open
 		// Step 1: connect websocket
 		try {
 			await (() => {
 				return new Promise<void>((res, rej) => {
-					const newSocket = new WebSocket(awsRealTimeUrl, 'graphql-ws');
+					const newSocket = this.getNewWebSocket(awsRealTimeUrl, 'graphql-ws');
 					newSocket.onerror = () => {
 						logger.debug(`WebSocket connection error`);
 					};
@@ -692,61 +642,66 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 			// Step 2: wait for ack from AWS AppSyncReaTime after sending init
 			await (() => {
 				return new Promise((res, rej) => {
-					let ackOk = false;
-					this.awsRealTimeSocket.onerror = error => {
-						logger.debug(`WebSocket error ${JSON.stringify(error)}`);
-					};
-					this.awsRealTimeSocket.onclose = event => {
-						logger.debug(`WebSocket closed ${event.reason}`);
-						rej(new Error(JSON.stringify(event)));
-					};
-
-					this.awsRealTimeSocket.onmessage = (message: MessageEvent) => {
-						logger.debug(
-							`subscription message from AWS AppSyncRealTime: ${message.data} `
-						);
-						const data = JSON.parse(message.data);
-						const {
-							type,
-							payload: {
-								connectionTimeoutMs = DEFAULT_KEEP_ALIVE_TIMEOUT,
-							} = {},
-						} = data;
-						if (type === MESSAGE_TYPES.GQL_CONNECTION_ACK) {
-							ackOk = true;
-							this.keepAliveTimeout = connectionTimeoutMs;
-							this.awsRealTimeSocket.onmessage = this._handleIncomingSubscriptionMessage.bind(
-								this
+					if (this.awsRealTimeSocket) {
+						let ackOk = false;
+						this.awsRealTimeSocket.onerror = error => {
+							logger.debug(`WebSocket error ${JSON.stringify(error)}`);
+						};
+						this.awsRealTimeSocket.onclose = event => {
+							logger.debug(`WebSocket closed ${event.reason}`);
+							rej(new Error(JSON.stringify(event)));
+						};
+
+						this.awsRealTimeSocket.onmessage = (message: MessageEvent) => {
+							logger.debug(
+								`subscription message from AWS AppSyncRealTime: ${message.data} `
 							);
-							this.awsRealTimeSocket.onerror = err => {
-								logger.debug(err);
-								this._errorDisconnect(CONTROL_MSG.CONNECTION_CLOSED);
-							};
-							this.awsRealTimeSocket.onclose = event => {
-								logger.debug(`WebSocket closed ${event.reason}`);
-								this._errorDisconnect(CONTROL_MSG.CONNECTION_CLOSED);
-							};
-							res('Cool, connected to AWS AppSyncRealTime');
-							return;
-						}
-
-						if (type === MESSAGE_TYPES.GQL_CONNECTION_ERROR) {
+							const data = JSON.parse(message.data);
 							const {
+								type,
 								payload: {
-									errors: [{ errorType = '', errorCode = 0 } = {}] = [],
+									connectionTimeoutMs = DEFAULT_KEEP_ALIVE_TIMEOUT,
 								} = {},
 							} = data;
+							if (type === MESSAGE_TYPES.GQL_CONNECTION_ACK) {
+								ackOk = true;
+								if (this.awsRealTimeSocket) {
+									this.keepAliveTimeout = connectionTimeoutMs;
+									this.awsRealTimeSocket.onmessage =
+										this._handleIncomingSubscriptionMessage.bind(this);
+									this.awsRealTimeSocket.onerror = err => {
+										logger.debug(err);
+										this._errorDisconnect(CONTROL_MSG.CONNECTION_CLOSED);
+									};
+									this.awsRealTimeSocket.onclose = event => {
+										logger.debug(`WebSocket closed ${event.reason}`);
+										this._errorDisconnect(CONTROL_MSG.CONNECTION_CLOSED);
+									};
+								}
+								res('Cool, connected to AWS AppSyncRealTime');
+								return;
+							}
+
+							if (type === MESSAGE_TYPES.GQL_CONNECTION_ERROR) {
+								const {
+									payload: {
+										errors: [{ errorType = '', errorCode = 0 } = {}] = [],
+									} = {},
+								} = data;
+
+								rej({ errorType, errorCode });
+							}
+						};
+
+						const gqlInit = {
+							type: MESSAGE_TYPES.GQL_CONNECTION_INIT,
+						};
+						this.awsRealTimeSocket.send(JSON.stringify(gqlInit));
+
+						setTimeout(checkAckOk.bind(this, ackOk), CONNECTION_INIT_TIMEOUT);
+					}
 
-							rej({ errorType, errorCode });
-						}
-					};
-
-					const gqlInit = {
-						type: MESSAGE_TYPES.GQL_CONNECTION_INIT,
-					};
-					this.awsRealTimeSocket.send(JSON.stringify(gqlInit));
-
-					function checkAckOk() {
+					function checkAckOk(ackOk: boolean) {
 						if (!ackOk) {
 							rej(
 								new Error(
@@ -755,12 +710,13 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 							);
 						}
 					}
-
-					setTimeout(checkAckOk.bind(this), CONNECTION_INIT_TIMEOUT);
 				});
 			})();
 		} catch (err) {
-			const { errorType, errorCode } = err;
+			const { errorType, errorCode } = err as {
+				errorType: string;
+				errorCode: number;
+			};
 
 			if (NON_RETRYABLE_CODES.includes(errorCode)) {
 				throw new NonRetryableError(errorType);
@@ -780,8 +736,10 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		apiKey,
 		region,
 		additionalHeaders,
-	}): Promise<any> {
-		const headerHandler = {
+	}: AWSAppSyncRealTimeProviderOptions): Promise<any> {
+		const headerHandler: {
+			[key in GraphqlAuthModes]: (AWSAppSyncRealTimeAuthInput) => {};
+		} = {
 			API_KEY: this._awsRealTimeApiKeyHeader.bind(this),
 			AWS_IAM: this._awsRealTimeIAMHeader.bind(this),
 			OPENID_CONNECT: this._awsRealTimeOPENIDHeader.bind(this),
@@ -789,29 +747,31 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 			AWS_LAMBDA: this._customAuthHeader,
 		};
 
-		const handler = headerHandler[authenticationType];
-
-		if (typeof handler !== 'function') {
+		if (!authenticationType || !headerHandler[authenticationType]) {
 			logger.debug(`Authentication type ${authenticationType} not supported`);
 			return '';
-		}
+		} else {
+			const handler = headerHandler[authenticationType];
 
-		const { host } = url.parse(appSyncGraphqlEndpoint);
+			const { host } = url.parse(appSyncGraphqlEndpoint ?? '');
 
-		const result = await handler({
-			payload,
-			canonicalUri,
-			appSyncGraphqlEndpoint,
-			apiKey,
-			region,
-			host,
-			additionalHeaders,
-		});
+			logger.debug(`Authenticating with ${authenticationType}`);
 
-		return result;
+			const result = await handler({
+				payload,
+				canonicalUri,
+				appSyncGraphqlEndpoint,
+				apiKey,
+				region,
+				host,
+				additionalHeaders,
+			});
+
+			return result;
+		}
 	}
 
-	private async _awsRealTimeCUPHeader({ host }) {
+	private async _awsRealTimeCUPHeader({ host }: AWSAppSyncRealTimeAuthInput) {
 		const session = await Auth.currentSession();
 		return {
 			Authorization: session.getAccessToken().getJwtToken(),
@@ -819,7 +779,9 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		};
 	}
 
-	private async _awsRealTimeOPENIDHeader({ host }) {
+	private async _awsRealTimeOPENIDHeader({
+		host,
+	}: AWSAppSyncRealTimeAuthInput) {
 		let token;
 		// backwards compatibility
 		const federatedInfo = await Cache.getItem('federatedInfo');
@@ -840,7 +802,10 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		};
 	}
 
-	private async _awsRealTimeApiKeyHeader({ apiKey, host }) {
+	private async _awsRealTimeApiKeyHeader({
+		apiKey,
+		host,
+	}: AWSAppSyncRealTimeAuthInput) {
 		const dt = new Date();
 		const dtStr = dt.toISOString().replace(/[:\-]|\.\d{3}/g, '');
 
@@ -856,7 +821,7 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		canonicalUri,
 		appSyncGraphqlEndpoint,
 		region,
-	}) {
+	}: AWSAppSyncRealTimeAuthInput) {
 		const endpointInfo = {
 			region,
 			service: 'appsync',
@@ -866,11 +831,16 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		if (!credentialsOK) {
 			throw new Error('No credentials');
 		}
-		const creds = await Credentials.get().then(credentials => ({
-			secret_key: credentials.secretAccessKey,
-			access_key: credentials.accessKeyId,
-			session_token: credentials.sessionToken,
-		}));
+		const creds = await Credentials.get().then((credentials: any) => {
+			const { secretAccessKey, accessKeyId, sessionToken } =
+				credentials as ICredentials;
+
+			return {
+				secret_key: secretAccessKey,
+				access_key: accessKeyId,
+				session_token: sessionToken,
+			};
+		});
 
 		const request = {
 			url: `${appSyncGraphqlEndpoint}${canonicalUri}`,
@@ -883,8 +853,11 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 		return signed_params.headers;
 	}
 
-	private _customAuthHeader({ host, additionalHeaders }) {
-		if (!additionalHeaders.Authorization) {
+	private _customAuthHeader({
+		host,
+		additionalHeaders,
+	}: AWSAppSyncRealTimeAuthInput) {
+		if (!additionalHeaders || !additionalHeaders['Authorization']) {
 			throw new Error('No auth token specified');
 		}
 
@@ -899,14 +872,14 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 	 */
 	_ensureCredentials() {
 		return Credentials.get()
-			.then(credentials => {
+			.then((credentials: any) => {
 				if (!credentials) return false;
 				const cred = Credentials.shear(credentials);
 				logger.debug('set credentials for AWSAppSyncRealTimeProvider', cred);
 
 				return true;
 			})
-			.catch(err => {
+			.catch((err: any) => {
 				logger.warn('ensure credentials error', err);
 				return false;
 			});