From 9c832cd4cd27789d333c34bd2ec0164a2292cfbb Mon Sep 17 00:00:00 2001 From: anivar Date: Fri, 31 Oct 2025 19:05:27 +0530 Subject: [PATCH] feat(api-graphql): add WebSocket health monitoring and manual reconnection - Add getConnectionHealth() method for current connection state - Add getPersistentConnectionHealth() for cross-session health tracking - Add isConnected() for quick connection status check - Add reconnect() for manual WebSocket reconnection - Implement persistent keep-alive tracking using AsyncStorage/localStorage - Use 65-second threshold for health monitoring - Track keep-alive timestamps for connection health assessment Fixes #9749, #4459, #5403, #7057 --- .../Providers/AWSWebSocketProvider/index.ts | 129 ++++++++++++++++++ packages/api-graphql/src/types/index.ts | 19 +++ 2 files changed, 148 insertions(+) diff --git a/packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts b/packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts index f5738ca4475..230c046a2da 100644 --- a/packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts +++ b/packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts @@ -3,6 +3,7 @@ import { Observable, Subscription, SubscriptionLike } from 'rxjs'; import { GraphQLError } from 'graphql'; import { ConsoleLogger, Hub, HubPayload } from '@aws-amplify/core'; +import type { KeyValueStorageInterface } from '@aws-amplify/core'; import { CustomUserAgentDetails, DocumentType, @@ -18,6 +19,7 @@ import { ConnectionState, PubSubContentObserver, } from '../../types/PubSub'; +import type { WebSocketHealthState } from '../../types'; import { AMPLIFY_SYMBOL, CONNECTION_INIT_TIMEOUT, @@ -49,6 +51,35 @@ import { } from './appsyncUrl'; import { awsRealTimeHeaderBasedAuth } from './authHeaders'; +// Storage key for persistent keep-alive tracking +const KEEP_ALIVE_STORAGE_KEY = 'AWS_AMPLIFY_LAST_KEEP_ALIVE'; + +// Platform-safe storage implementation +let platformStorage: Pick< + KeyValueStorageInterface, + 'setItem' | 'getItem' +> | null = null; + +try { + // Try to import AsyncStorage for React Native (optional dependency) + const AsyncStorage = + // eslint-disable-next-line import/no-extraneous-dependencies + require('@react-native-async-storage/async-storage').default; + platformStorage = AsyncStorage; +} catch (e) { + // Fallback for web/other platforms - use localStorage if available + if (typeof localStorage !== 'undefined') { + platformStorage = { + setItem: (key: string, value: string) => { + localStorage.setItem(key, value); + + return Promise.resolve(); + }, + getItem: (key: string) => Promise.resolve(localStorage.getItem(key)), + }; + } +} + const dispatchApiEvent = (payload: HubPayload) => { Hub.dispatch('api', payload, 'PubSub', AMPLIFY_SYMBOL); }; @@ -106,6 +137,7 @@ export abstract class AWSWebSocketProvider { /** * Mark the socket closed and release all active listeners */ + close() { // Mark the socket closed both in status and the connection monitor this.socketStatus = SOCKET_STATUS.CLOSED; @@ -681,6 +713,15 @@ export abstract class AWSWebSocketProvider { if (type === MESSAGE_TYPES.GQL_CONNECTION_KEEP_ALIVE) { this.maintainKeepAlive(); + // Persist keep-alive timestamp for cross-session tracking + if (platformStorage) { + platformStorage + .setItem(KEEP_ALIVE_STORAGE_KEY, `${Date.now()}`) + .catch(error => { + this.logger.warn('Failed to persist keep-alive timestamp:', error); + }); + } + return; } @@ -1025,4 +1066,92 @@ export abstract class AWSWebSocketProvider { } } }; + + // WebSocket Health & Control API + + /** + * Get current WebSocket health state + */ + getConnectionHealth(): WebSocketHealthState { + const timeSinceLastKeepAlive = Date.now() - this.keepAliveTimestamp; + + const isHealthy = + this.connectionState === ConnectionState.Connected && + timeSinceLastKeepAlive < DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT; + + return { + isHealthy, + connectionState: this.connectionState || ConnectionState.Disconnected, + lastKeepAliveTime: this.keepAliveTimestamp, + timeSinceLastKeepAlive, + }; + } + + /** + * Get persistent WebSocket health state (survives app restarts) + */ + async getPersistentConnectionHealth(): Promise { + let persistentKeepAliveTime = 0; + + // Try to get persistent keep-alive timestamp + if (platformStorage) { + try { + const persistentKeepAlive = await platformStorage.getItem( + KEEP_ALIVE_STORAGE_KEY, + ); + if (persistentKeepAlive) { + persistentKeepAliveTime = Number(persistentKeepAlive) || 0; + } + } catch (error) { + this.logger.warn( + 'Failed to retrieve persistent keep-alive timestamp:', + error, + ); + } + } + + // Use the more recent timestamp (in-memory vs persistent) + const lastKeepAliveTime = Math.max( + this.keepAliveTimestamp, + persistentKeepAliveTime, + ); + + const timeSinceLastKeepAlive = + lastKeepAliveTime > 0 ? Date.now() - lastKeepAliveTime : Infinity; // If no keep-alive has been received, treat as unhealthy + + // Health check includes persistent data + const isHealthy = + this.connectionState === ConnectionState.Connected && + timeSinceLastKeepAlive < DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT; + + return { + isHealthy, + connectionState: this.connectionState || ConnectionState.Disconnected, + lastKeepAliveTime: lastKeepAliveTime > 0 ? lastKeepAliveTime : undefined, + timeSinceLastKeepAlive: + lastKeepAliveTime > 0 ? timeSinceLastKeepAlive : undefined, + }; + } + + /** + * Check if WebSocket is currently connected + */ + isConnected(): boolean { + return this.awsRealTimeSocket?.readyState === WebSocket.OPEN; + } + + /** + * Manually reconnect WebSocket + */ + async reconnect(): Promise { + this.logger.info('Manual WebSocket reconnection requested'); + + // Close existing connection if any + if (this.isConnected()) { + await this.close(); + } + + // Trigger reconnection through the reconnection monitor + this.reconnectionMonitor.record(ReconnectEvent.START_RECONNECT); + } } diff --git a/packages/api-graphql/src/types/index.ts b/packages/api-graphql/src/types/index.ts index 2fa4ce08800..ded69f67c2b 100644 --- a/packages/api-graphql/src/types/index.ts +++ b/packages/api-graphql/src/types/index.ts @@ -522,3 +522,22 @@ export interface AuthModeParams extends Record { export type GenerateServerClientParams = { config: ResourcesConfig; } & CommonPublicClientOptions; + +// WebSocket health and control types +export interface WebSocketHealthState { + isHealthy: boolean; + connectionState: import('./PubSub').ConnectionState; + lastKeepAliveTime?: number; + timeSinceLastKeepAlive?: number; +} + +export interface WebSocketControl { + reconnect(): Promise; + disconnect(): void; + isConnected(): boolean; + getConnectionHealth(): WebSocketHealthState; + getPersistentConnectionHealth(): Promise; + onConnectionStateChange( + callback: (state: import('./PubSub').ConnectionState) => void, + ): () => void; +}