@limrun/api 0.13.1 → 0.13.2

package/src/instance-client.ts

@@ -4,6 +4,16 @@ import { exec } from 'node:child_process';
 import { startTcpTunnel } from './tunnel';
 import type { Tunnel } from './tunnel';
 
+/**
+ * Connection state of the instance client
+ */
+export type ConnectionState = 'connecting' | 'connected' | 'disconnected' | 'reconnecting';
+
+/**
+ * Callback function for connection state changes
+ */
+export type ConnectionStateCallback = (state: ConnectionState) => void;
+
 /**
  * A client for interacting with a Limbar instance
  */
@@ -29,6 +39,17 @@ export type InstanceClient = {
    * rejects with an Error on failure.
    */
   sendAsset: (url: string) => Promise<void>;
+
+  /**
+   * Get current connection state
+   */
+  getConnectionState: () => ConnectionState;
+
+  /**
+   * Register callback for connection state changes
+   * @returns A function to unregister the callback
+   */
+  onConnectionStateChange: (callback: ConnectionStateCallback) => () => void;
 };
 
 /**
@@ -62,6 +83,21 @@ export type InstanceClientOptions = {
    * @default 'info'
    */
   logLevel?: LogLevel;
+  /**
+   * Maximum number of reconnection attempts
+   * @default 6
+   */
+  maxReconnectAttempts?: number;
+  /**
+   * Initial reconnection delay in milliseconds
+   * @default 1000
+   */
+  reconnectDelay?: number;
+  /**
+   * Maximum reconnection delay in milliseconds
+   * @default 30000
+   */
+  maxReconnectDelay?: number;
 };
 
 type ScreenshotRequest = {
@@ -111,7 +147,15 @@ type ServerMessage =
 export async function createInstanceClient(options: InstanceClientOptions): Promise<InstanceClient> {
   const serverAddress = `${options.endpointUrl}?token=${options.token}`;
   const logLevel = options.logLevel ?? 'info';
+  const maxReconnectAttempts = options.maxReconnectAttempts ?? 6;
+  const reconnectDelay = options.reconnectDelay ?? 1000;
+  const maxReconnectDelay = options.maxReconnectDelay ?? 30000;
+
   let ws: WebSocket | undefined = undefined;
+  let connectionState: ConnectionState = 'connecting';
+  let reconnectAttempts = 0;
+  let reconnectTimeout: NodeJS.Timeout | undefined;
+  let intentionalDisconnect = false;
 
   const screenshotRequests: Map<
     string,
@@ -129,6 +173,8 @@ export async function createInstanceClient(options: InstanceClientOptions): Prom
     }
   > = new Map();
 
+  const stateChangeCallbacks: Set<ConnectionStateCallback> = new Set();
+
   // Logger functions
   const logger = {
     debug: (...args: any[]) => {
@@ -145,104 +191,213 @@ export async function createInstanceClient(options: InstanceClientOptions): Prom
     },
   };
 
+  const updateConnectionState = (newState: ConnectionState): void => {
+    if (connectionState !== newState) {
+      connectionState = newState;
+      logger.debug(`Connection state changed to: ${newState}`);
+      stateChangeCallbacks.forEach((callback) => {
+        try {
+          callback(newState);
+        } catch (err) {
+          logger.error('Error in connection state callback:', err);
+        }
+      });
+    }
+  };
+
+  const failPendingRequests = (reason: string): void => {
+    screenshotRequests.forEach((request) => request.rejecter(new Error(reason)));
+    screenshotRequests.clear();
+    assetRequests.forEach((request) => request.rejecter(new Error(reason)));
+    assetRequests.clear();
+  };
+
+  const cleanup = (): void => {
+    if (reconnectTimeout) {
+      clearTimeout(reconnectTimeout);
+      reconnectTimeout = undefined;
+    }
+    if (pingInterval) {
+      clearInterval(pingInterval);
+      pingInterval = undefined;
+    }
+    if (ws) {
+      ws.removeAllListeners();
+      if (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING) {
+        ws.close();
+      }
+      ws = undefined;
+    }
+  };
+
+  let pingInterval: NodeJS.Timeout | undefined;
+
   return new Promise<InstanceClient>((resolveConnection, rejectConnection) => {
-    logger.debug(`Attempting to connect to WebSocket server at ${serverAddress}...`);
-    ws = new WebSocket(serverAddress);
-    ws.on('message', (data: Data) => {
-      let message: ServerMessage;
-      try {
-        message = JSON.parse(data.toString());
-      } catch (e) {
-        logger.error({ data, error: e }, 'Failed to parse JSON message');
+    let hasResolved = false;
+
+    // Reconnection logic with exponential backoff
+    const scheduleReconnect = (): void => {
+      if (intentionalDisconnect) {
+        logger.debug('Skipping reconnection (intentional disconnect)');
         return;
       }
 
-      switch (message.type) {
-        case 'screenshot': {
-          if (!('dataUri' in message) || typeof message.dataUri !== 'string' || !('id' in message)) {
-            logger.warn('Received invalid screenshot message:', message);
-            break;
-          }
+      if (reconnectAttempts >= maxReconnectAttempts) {
+        logger.error(`Max reconnection attempts (${maxReconnectAttempts}) reached. Giving up.`);
+        updateConnectionState('disconnected');
+        return;
+      }
 
-          const screenshotMessage = message as ScreenshotResponse;
-          const request = screenshotRequests.get(screenshotMessage.id);
+      const currentDelay = Math.min(reconnectDelay * Math.pow(2, reconnectAttempts), maxReconnectDelay);
 
-          if (!request) {
-            logger.warn(
-              `Received screenshot data for unknown or already handled session: ${screenshotMessage.id}`,
-            );
-            break;
-          }
+      reconnectAttempts++;
+      logger.debug(`Scheduling reconnection attempt ${reconnectAttempts} in ${currentDelay}ms...`);
+      updateConnectionState('reconnecting');
 
-          logger.debug(`Received screenshot data URI for session ${screenshotMessage.id}.`);
-          request.resolver({ dataUri: screenshotMessage.dataUri });
-          screenshotRequests.delete(screenshotMessage.id);
-          break;
-        }
-        case 'screenshotError': {
-          if (!('message' in message) || !('id' in message)) {
-            logger.warn('Received invalid screenshot error message:', message);
-            break;
-          }
+      reconnectTimeout = setTimeout(() => {
+        logger.debug(`Attempting to reconnect (attempt ${reconnectAttempts})...`);
+        setupWebSocket();
+      }, currentDelay);
+    };
 
-          const errorMessage = message as ScreenshotErrorResponse;
-          const request = screenshotRequests.get(errorMessage.id);
+    const setupWebSocket = (): void => {
+      cleanup();
+      updateConnectionState('connecting');
 
-          if (!request) {
-            logger.warn(
-              `Received screenshot error for unknown or already handled session: ${errorMessage.id}`,
-            );
-            break;
-          }
+      ws = new WebSocket(serverAddress);
 
-          logger.error(
-            `Server reported an error capturing screenshot for session ${errorMessage.id}:`,
-            errorMessage.message,
-          );
-          request.rejecter(new Error(errorMessage.message));
-          screenshotRequests.delete(errorMessage.id);
-          break;
+      ws.on('message', (data: Data) => {
+        let message: ServerMessage;
+        try {
+          message = JSON.parse(data.toString());
+        } catch (e) {
+          logger.error({ data, error: e }, 'Failed to parse JSON message');
+          return;
         }
-        case 'assetResult': {
-          logger.debug('Received assetResult:', message);
-          const request = assetRequests.get(message.url as string);
-          if (!request) {
-            logger.warn(`Received assetResult for unknown or already handled url: ${message.url}`);
+
+        switch (message.type) {
+          case 'screenshot': {
+            if (!('dataUri' in message) || typeof message.dataUri !== 'string' || !('id' in message)) {
+              logger.warn('Received invalid screenshot message:', message);
+              break;
+            }
+
+            const screenshotMessage = message as ScreenshotResponse;
+            const request = screenshotRequests.get(screenshotMessage.id);
+
+            if (!request) {
+              logger.warn(
+                `Received screenshot data for unknown or already handled session: ${screenshotMessage.id}`,
+              );
+              break;
+            }
+
+            logger.debug(`Received screenshot data URI for session ${screenshotMessage.id}.`);
+            request.resolver({ dataUri: screenshotMessage.dataUri });
+            screenshotRequests.delete(screenshotMessage.id);
+            break;
+          }
+          case 'screenshotError': {
+            if (!('message' in message) || !('id' in message)) {
+              logger.warn('Received invalid screenshot error message:', message);
+              break;
+            }
+
+            const errorMessage = message as ScreenshotErrorResponse;
+            const request = screenshotRequests.get(errorMessage.id);
+
+            if (!request) {
+              logger.warn(
+                `Received screenshot error for unknown or already handled session: ${errorMessage.id}`,
+              );
+              break;
+            }
+
+            logger.error(
+              `Server reported an error capturing screenshot for session ${errorMessage.id}:`,
+              errorMessage.message,
+            );
+            request.rejecter(new Error(errorMessage.message));
+            screenshotRequests.delete(errorMessage.id);
             break;
           }
-          if (message.result === 'success') {
-            logger.debug('Asset result is success');
-            request.resolver();
+          case 'assetResult': {
+            logger.debug('Received assetResult:', message);
+            const request = assetRequests.get(message.url as string);
+            if (!request) {
+              logger.warn(`Received assetResult for unknown or already handled url: ${message.url}`);
+              break;
+            }
+            if (message.result === 'success') {
+              logger.debug('Asset result is success');
+              request.resolver();
+              assetRequests.delete(message.url as string);
+              break;
+            }
+            const errorMessage =
+              typeof message.message === 'string' && message.message ?
+                message.message
+              : `Asset processing failed: ${JSON.stringify(message)}`;
+            logger.debug('Asset result is failure', errorMessage);
+            request.rejecter(new Error(errorMessage));
             assetRequests.delete(message.url as string);
             break;
           }
-          const errorMessage =
-            typeof message.message === 'string' && message.message ?
-              message.message
-            : `Asset processing failed: ${JSON.stringify(message)}`;
-          logger.debug('Asset result is failure', errorMessage);
-          request.rejecter(new Error(errorMessage));
-          assetRequests.delete(message.url as string);
-          break;
+          default:
+            logger.warn(`Received unexpected message type: ${message.type}`);
+            break;
         }
-        default:
-          logger.warn(`Received unexpected message type: ${message.type}`);
-          break;
-      }
-    });
+      });
 
-    ws.on('error', (err: Error) => {
-      logger.error('WebSocket error:', err.message);
-      if (ws && (ws.readyState === WebSocket.CONNECTING || ws.readyState === WebSocket.OPEN)) {
-        rejectConnection(err);
-      }
-      screenshotRequests.forEach((request) => request.rejecter(err));
-    });
+      ws.on('error', (err: Error) => {
+        logger.error('WebSocket error:', err.message);
+        if (!hasResolved && (ws?.readyState === WebSocket.CONNECTING || ws?.readyState === WebSocket.OPEN)) {
+          rejectConnection(err);
+        }
+      });
+
+      ws.on('close', () => {
+        if (pingInterval) {
+          clearInterval(pingInterval);
+          pingInterval = undefined;
+        }
+
+        const shouldReconnect = !intentionalDisconnect && connectionState !== 'disconnected';
+        updateConnectionState('disconnected');
+
+        logger.debug('Disconnected from server.');
 
-    ws.on('close', () => {
-      logger.debug('Disconnected from server.');
-      screenshotRequests.forEach((request) => request.rejecter('Disconnected from server'));
-    });
+        failPendingRequests('Connection closed');
+
+        if (shouldReconnect) {
+          scheduleReconnect();
+        }
+      });
+
+      ws.on('open', () => {
+        logger.debug(`Connected to ${serverAddress}`);
+        reconnectAttempts = 0;
+        updateConnectionState('connected');
+
+        pingInterval = setInterval(() => {
+          if (ws && ws.readyState === WebSocket.OPEN) {
+            (ws as any).ping();
+          }
+        }, 30_000);
+
+        if (!hasResolved) {
+          hasResolved = true;
+          resolveConnection({
+            screenshot,
+            disconnect,
+            startAdbTunnel,
+            sendAsset,
+            getConnectionState,
+            onConnectionStateChange,
+          });
+        }
+      });
+    };
 
     const screenshot = async (): Promise<ScreenshotData> => {
       if (!ws || ws.readyState !== WebSocket.OPEN) {
@@ -287,11 +442,22 @@ export async function createInstanceClient(options: InstanceClientOptions): Prom
     };
 
     const disconnect = (): void => {
-      if (ws) {
-        logger.debug('Closing WebSocket connection.');
-        ws.close();
-      }
-      screenshotRequests.forEach((request) => request.rejecter('Websocket connection closed'));
+      intentionalDisconnect = true;
+      cleanup();
+      updateConnectionState('disconnected');
+      failPendingRequests('Intentional disconnect');
+      logger.debug('Intentionally disconnected from WebSocket.');
+    };
+
+    const getConnectionState = (): ConnectionState => {
+      return connectionState;
+    };
+
+    const onConnectionStateChange = (callback: ConnectionStateCallback): (() => void) => {
+      stateChangeCallbacks.add(callback);
+      return () => {
+        stateChangeCallbacks.delete(callback);
+      };
     };
 
     /**
@@ -299,20 +465,28 @@ export async function createInstanceClient(options: InstanceClientOptions): Prom
      * client to it.
      */
     const startAdbTunnel = async (): Promise<Tunnel> => {
-      const { address, close } = await startTcpTunnel(options.adbUrl, options.token, '127.0.0.1', 0);
+      const tunnel = await startTcpTunnel(options.adbUrl, options.token, '127.0.0.1', 0, {
+        maxReconnectAttempts,
+        reconnectDelay,
+        maxReconnectDelay,
+        logLevel,
+      });
       try {
         await new Promise<void>((resolve, reject) => {
-          exec(`${options.adbPath ?? 'adb'} connect ${address.address}:${address.port}`, (err) => {
-            if (err) return reject(err);
-            resolve();
-          });
+          exec(
+            `${options.adbPath ?? 'adb'} connect ${tunnel.address.address}:${tunnel.address.port}`,
+            (err) => {
+              if (err) return reject(err);
+              resolve();
+            },
+          );
         });
-        logger.debug(`ADB connected on ${address.address}`);
+        logger.debug(`ADB connected on ${tunnel.address.address}`);
       } catch (err) {
-        close();
+        tunnel.close();
         throw err;
       }
-      return { address, close };
+      return tunnel;
     };
 
     const sendAsset = async (url: string): Promise<void> => {
@@ -332,14 +506,8 @@ export async function createInstanceClient(options: InstanceClientOptions): Prom
         assetRequests.set(url, { resolver: resolve, rejecter: reject });
       });
     };
-    ws.on('open', () => {
-      logger.debug(`Connected to ${serverAddress}`);
-      resolveConnection({
-        screenshot,
-        disconnect,
-        startAdbTunnel,
-        sendAsset,
-      });
-    });
+
+    // Start the initial connection
+    setupWebSocket();
   });
 }