@alwaysai/device-agent 1.4.0 → 1.5.0

package/src/cloud-connection/device-agent-cloud-connection.ts

@@ -15,7 +15,8 @@ import {
   buildToClientStatusResponseMessage,
   StatusResponsePayload,
   keyMirrors,
-  validateToDeviceAgentMessage
+  validateToDeviceAgentMessage,
+  validateSecureTunnelShadowUpdate
 } from '@alwaysai/device-agent-schemas';
 import { existsSync } from 'fs';
 import {
@@ -46,12 +47,14 @@ import { logger } from '../util/logger';
 import sleep from '../util/sleep';
 import { bootstrapProvision } from './bootstrap-provision';
 import { LiveUpdatesHandler } from './live-updates-handler';
-import { PassthroughHandler, runChannel } from './passthrough-handler';
+import { PassthroughHandler } from './passthrough-handler';
 import { Publisher } from './publisher';
 import { ShadowHandler, ShadowUpdate } from './shadow-handler';
 import { TransactionManager } from './transaction-manager';
 import { exec } from 'child_process';
 import { promisify } from 'util';
+import { pruneModels } from '../application-control/models';
+import { getDeviceAgentVersion } from '../util/check-for-updates';
 
 const exec_promise = promisify(exec);
 
@@ -69,8 +72,6 @@ export class DeviceAgentCloudConnection {
   private readonly secureTunnelNotifyTopic = `$aws/things/${this.clientId}/tunnels/notify`;
   private readonly secureTunnelHandler =
     SecureTunnelHandlerSingleton.getInstance();
-  // FIXME: Add support for multiple simultaneous project updates
-  private appCfgUpdateQueue: ShadowUpdate[] = [];
 
   private handleAppStateControl = async (
     payload: AppStateControlPayload
@@ -152,16 +153,18 @@ export class DeviceAgentCloudConnection {
   };
 
   private handleModelsInstallCloudResponsePayload = async (
-    payload: ModelsInstallResponsePayload
+    payload: ModelsInstallResponsePayload,
+    txId: string
   ): Promise<boolean> => {
-    const update = this.appCfgUpdateQueue.shift();
+    const projectId = payload.modelsInstallResponse.projectId;
+
+    const update = this.txnMgr.getAppCfgUpdateFromTxID(txId);
     if (update === undefined) {
       throw new Error(
         'Unknown error while updating models via application config! No config present for model update.'
       );
     }
     const { appCfgUpdate, envVarUpdate } = update;
-    const projectId = payload.modelsInstallResponse.projectId;
     if (appCfgUpdate) {
       await this.atomicApplicationUpdate(
         async () =>
@@ -185,6 +188,7 @@ export class DeviceAgentCloudConnection {
         true
       );
     }
+
     return true;
   };
 
@@ -305,19 +309,22 @@ export class DeviceAgentCloudConnection {
       );
       this.publisher.publishToCloud(message);
 
-      this.appCfgUpdateQueue.push(update);
+      this.txnMgr.setAppCfgUpdateToTx(txId, update);
+
       return false;
     }
 
     if (appCfgUpdate) {
-      await this.atomicApplicationUpdate(
-        async () =>
-          await updateAppCfg({
-            projectId,
-            newAppCfg: appCfgUpdate.newAppCfg
-          }),
-        projectId
-      );
+      await this.atomicApplicationUpdate(async () => {
+        await pruneModels({
+          projectId,
+          appCfg: appCfgUpdate.newAppCfg
+        });
+        await updateAppCfg({
+          projectId,
+          newAppCfg: appCfgUpdate.newAppCfg
+        });
+      }, projectId);
     }
 
     if (envVarUpdate) {
@@ -369,6 +376,30 @@ export class DeviceAgentCloudConnection {
     }
   }
 
+  public async handleSecureTunnelMessage(payload: any): Promise<void> {
+    logger.info(`Received secure tunnel update: ${JSON.stringify(payload)}`);
+    const state = payload.state;
+    if (!state) {
+      logger.debug(`No state found in message: ${JSON.stringify(payload)}`);
+      return;
+    }
+    const valid = validateSecureTunnelShadowUpdate(state);
+    if (!valid) {
+      logger.error(
+        `Error validating message: ${JSON.stringify(
+          { payload, errors: validateSecureTunnelShadowUpdate.errors },
+          null,
+          2
+        )}`
+      );
+      return;
+    }
+    const secureTunnelUpdate =
+      await this.secureTunnelHandler.syncShadowToDeviceState(payload);
+    await this.shadowHandler.updateSecureTunnelShadow(secureTunnelUpdate);
+    return;
+  }
+
   /*=================================================================
                             Public interface
   =================================================================*/
@@ -526,7 +557,8 @@ export class DeviceAgentCloudConnection {
           );
         }
         await this.txnMgr.runTransactionStep({
-          func: () => this.handleModelsInstallCloudResponsePayload(payload),
+          func: () =>
+            this.handleModelsInstallCloudResponsePayload(payload, txId),
           projectId,
           txId,
           start: false,
@@ -607,10 +639,7 @@ export class DeviceAgentCloudConnection {
     }
   }
 
-  public async handleMessage(
-    topic: string,
-    message: ToDeviceAgentMessage | any
-  ) {
+  public async handleMessage(topic: string, message: any) {
     logger.debug(
       `Received message: ${JSON.stringify({ topic, message }, null, 2)}`
     );
@@ -629,14 +658,7 @@ export class DeviceAgentCloudConnection {
     } else if (
       topic === this.shadowHandler.shadowTopics.secureTunnel.updateDelta
     ) {
-      logger.info(`Received secure tunnel update: ${message}`);
-      const reported = await this.secureTunnelHandler.syncShadowToDeviceState(
-        message
-      );
-      this.publisher.publish(
-        this.shadowHandler.shadowTopics.secureTunnel.update,
-        JSON.stringify({ state: { reported } })
-      );
+      await this.handleSecureTunnelMessage(message);
     } else if (
       topic === this.shadowHandler.shadowTopics.secureTunnel.deleteAccepted
     ) {
@@ -749,6 +771,9 @@ export class DeviceAgentCloudConnection {
 }
 
 export async function runDeviceAgentCloudInterface() {
+  logger.info(
+    `Starting alwaysAI Device Agent v${await getDeviceAgentVersion()}`
+  );
   if (cloudModeReady()) {
     const deviceAgent = new DeviceAgentCloudConnection();
     await deviceAgent.setupHandlers();
@@ -756,7 +781,6 @@ export async function runDeviceAgentCloudInterface() {
       const publisher = deviceAgent.publisher;
       const passthroughHandler = new PassthroughHandler(publisher);
       await passthroughHandler.setup();
-      await runChannel(passthroughHandler);
     }
   } else if (existsSync(BOOTSTRAP_PRIVATE_KEY_FILE_PATH())) {
     await bootstrapProvision();

package/src/cloud-connection/live-updates-handler.test.ts

@@ -1,5 +1,12 @@
 import { LiveUpdatesHandler } from './live-updates-handler';
 import { Publisher } from './publisher';
+import sleep from '../util/sleep';
+import { getDeviceStatsPayload } from './messages';
+
+jest.mock('../util/sleep');
+jest.mock('./messages');
+
+jest.mocked(getDeviceStatsPayload).mockResolvedValue({});
 
 global.setTimeout = jest.fn() as unknown as typeof setTimeout;
 
@@ -18,64 +25,163 @@ const testFalseToggles = {
   appState: false
 };
 
-const mockClient = jest.fn();
+const mockClient = {
+  publish: jest.fn()
+};
 const clientId = 'test-client';
 const emptyTxId = '';
 
-// NOTE: this was the way I found to mock private class functions
-const mockStartPublishingLiveUpdates = jest.spyOn(
-  LiveUpdatesHandler.prototype as any,
-  'startPublishingLiveUpdates'
-);
-mockStartPublishingLiveUpdates.mockResolvedValue(null);
-
 describe('Test Live Updates Handler', () => {
   let liveUpdatesHandler: LiveUpdatesHandler;
   let publisher: Publisher;
 
   beforeEach(() => {
+    mockClient.publish = jest.fn();
+    jest.mocked(sleep).mockImplementation(async () => {
+      return;
+    });
     publisher = new Publisher(mockClient, clientId);
     liveUpdatesHandler = new LiveUpdatesHandler(publisher, clientId);
     jest.clearAllMocks();
   });
 
+  test('enable device stats', async () => {
+    mockClient.publish = jest.fn().mockImplementation(async () => {
+      liveUpdatesHandler.disableDeviceStatsLiveUpdates();
+    });
+
+    const enable = { deviceStats: true };
+    const promises = await liveUpdatesHandler.handleToggles(enable, emptyTxId);
+    await Promise.all(promises);
+
+    expect(jest.mocked(setTimeout)).toBeCalledTimes(1);
+    expect(mockClient.publish).toBeCalledTimes(1);
+    expect(jest.mocked(sleep)).toBeCalledTimes(1);
+  });
+
+  test('enable app state', async () => {
+    mockClient.publish = jest.fn().mockImplementation(async () => {
+      liveUpdatesHandler.disableAppStateLiveUpdates();
+    });
+
+    const enable = { appState: true };
+    const promises = await liveUpdatesHandler.handleToggles(enable, emptyTxId);
+    await Promise.all(promises);
+
+    expect(jest.mocked(setTimeout)).toBeCalledTimes(1);
+    expect(mockClient.publish).toBeCalledTimes(1);
+    expect(jest.mocked(sleep)).toBeCalledTimes(1);
+  });
+
   test('ignore subsequent enables', async () => {
-    void liveUpdatesHandler.handleToggles(testTrueToggles, emptyTxId);
-    // called twice, once for device stats, once for app state
-    expect(mockStartPublishingLiveUpdates).toBeCalledTimes(2);
-    // restartLiveUpdatesTimeout is always called once when handleToggles is called
+    // Block on the first sleep call
+    let doneWaiting;
+    const waitPromise = new Promise(function (resolve) {
+      doneWaiting = resolve;
+    });
+
+    let startSleeping;
+    const sleepPromise = new Promise(function (resolve) {
+      startSleeping = resolve;
+    });
+    jest.mocked(sleep).mockImplementation(async () => {
+      startSleeping();
+      await waitPromise;
+    });
+
+    const enable = { deviceStats: true };
+    await liveUpdatesHandler.handleToggles(enable, emptyTxId);
+    await sleepPromise;
+
     expect(jest.mocked(setTimeout)).toBeCalledTimes(1);
+    expect(mockClient.publish).toBeCalledTimes(1);
+    expect(jest.mocked(sleep)).toBeCalledTimes(1);
 
-    // Second call -> should not call startPublishingLiveUpdates should not be called
+    // Second call
     jest.clearAllMocks();
-    void liveUpdatesHandler.handleToggles(testTrueToggles, emptyTxId);
-    expect(mockStartPublishingLiveUpdates).toBeCalledTimes(0);
+
+    await liveUpdatesHandler.handleToggles(enable, emptyTxId);
+
     expect(jest.mocked(setTimeout)).toHaveBeenCalledTimes(1);
+    expect(mockClient.publish).toBeCalledTimes(0);
+    expect(jest.mocked(sleep)).toBeCalledTimes(0);
+
+    // Disable and clean up
+    liveUpdatesHandler.disableDeviceStatsLiveUpdates();
+    doneWaiting();
   });
 
   test('test disable live updates', async () => {
+    let doneWaiting;
+    const waitPromise = new Promise(function (resolve) {
+      doneWaiting = resolve;
+    });
+
+    let startSleeping;
+    const sleepPromise = new Promise(function (resolve) {
+      startSleeping = resolve;
+    });
+
+    // Wait for two sleep calls, one for Device Stats and one for App State
+    jest
+      .mocked(sleep)
+      .mockImplementationOnce(async () => {
+        await waitPromise;
+      })
+      .mockImplementationOnce(async () => {
+        startSleeping();
+        await waitPromise;
+      });
+
     // Test calling handleToggles one time, enabling it
-    void liveUpdatesHandler.handleToggles(testTrueToggles, emptyTxId);
-    expect(mockStartPublishingLiveUpdates).toBeCalledTimes(2);
+    await liveUpdatesHandler.handleToggles(testTrueToggles, emptyTxId);
+    await sleepPromise;
+
     expect(jest.mocked(setTimeout)).toHaveBeenCalledTimes(1);
+    expect(jest.mocked(sleep)).toBeCalledTimes(2);
     expect(liveUpdatesHandler.getDeviceStatsLiveUpdates()).toBe(true);
     expect(liveUpdatesHandler.getAppStateLiveUpdates()).toBe(true);
     expect(liveUpdatesHandler.getAppLogsLiveUpdates()).toBe(false);
 
     // Check to see that attributes are properly set to false when disabled
     jest.clearAllMocks();
-    void liveUpdatesHandler.handleToggles(testFalseToggles, emptyTxId);
-    expect(mockStartPublishingLiveUpdates).toBeCalledTimes(0);
+    await liveUpdatesHandler.handleToggles(testFalseToggles, emptyTxId);
+
     expect(jest.mocked(setTimeout)).toHaveBeenCalledTimes(1);
+    expect(jest.mocked(sleep)).toBeCalledTimes(0);
     expect(liveUpdatesHandler.getDeviceStatsLiveUpdates()).toBe(false);
     expect(liveUpdatesHandler.getAppStateLiveUpdates()).toBe(false);
     expect(liveUpdatesHandler.getAppLogsLiveUpdates()).toBe(false);
+
+    doneWaiting();
   });
 
   test('timeout turns off live updates', async () => {
     jest.useFakeTimers({ legacyFakeTimers: true });
+    let doneWaiting;
+    const waitPromise = new Promise(function (resolve) {
+      doneWaiting = resolve;
+    });
+
+    let startSleeping;
+    const sleepPromise = new Promise(function (resolve) {
+      startSleeping = resolve;
+    });
+
+    // Wait for two sleep calls, one for Device Stats and one for App State
+    jest
+      .mocked(sleep)
+      .mockImplementationOnce(async () => {
+        await waitPromise;
+      })
+      .mockImplementationOnce(async () => {
+        startSleeping();
+        await waitPromise;
+      });
+
+    await liveUpdatesHandler.handleToggles(testTrueToggles, emptyTxId);
+    await sleepPromise;
 
-    void liveUpdatesHandler.handleToggles(testTrueToggles, emptyTxId);
     expect(liveUpdatesHandler.getDeviceStatsLiveUpdates()).toBe(true);
     expect(liveUpdatesHandler.getAppStateLiveUpdates()).toBe(true);
     expect(liveUpdatesHandler.getAppLogsLiveUpdates()).toBe(false);
@@ -85,5 +191,40 @@ describe('Test Live Updates Handler', () => {
     expect(liveUpdatesHandler.getDeviceStatsLiveUpdates()).toBe(false);
     expect(liveUpdatesHandler.getAppStateLiveUpdates()).toBe(false);
     expect(liveUpdatesHandler.getAppLogsLiveUpdates()).toBe(false);
+
+    doneWaiting();
+  });
+
+  test("failure doesn't kill publish loop", async () => {
+    // Block on the first sleep call
+    let doneWaiting;
+    const waitPromise = new Promise(function (resolve) {
+      doneWaiting = resolve;
+    });
+
+    let startSleeping;
+    const sleepPromise = new Promise(function (resolve) {
+      startSleeping = resolve;
+    });
+    jest
+      .mocked(sleep)
+      .mockImplementationOnce(async () => {
+        throw new Error('Test error!');
+      })
+      .mockImplementationOnce(async () => {
+        startSleeping();
+        await waitPromise;
+      });
+
+    const enable = { deviceStats: true };
+    await liveUpdatesHandler.handleToggles(enable, emptyTxId);
+    await sleepPromise;
+
+    expect(jest.mocked(setTimeout)).toBeCalledTimes(1);
+    expect(mockClient.publish).toBeCalledTimes(2);
+    expect(jest.mocked(sleep)).toBeCalledTimes(2);
+
+    liveUpdatesHandler.disableDeviceStatsLiveUpdates();
+    doneWaiting();
   });
 });

package/src/cloud-connection/live-updates-handler.ts

@@ -16,11 +16,10 @@ import sleep from '../util/sleep';
 import { Publisher } from './publisher';
 import { getAppStatePayload, getDeviceStatsPayload } from './messages';
 import { ALWAYSAI_LIVE_UPDATES_TIMEOUT_MS } from '../environment';
-import { read } from 'fs';
 
 const LIVE_UPDATES_TIMEOUT = ALWAYSAI_LIVE_UPDATES_TIMEOUT_MS
   ? parseInt(ALWAYSAI_LIVE_UPDATES_TIMEOUT_MS)
-  : 60000;
+  : 80000;
 
 export class LiveUpdatesHandler {
   private publisher: Publisher;
@@ -135,22 +134,11 @@ export class LiveUpdatesHandler {
     return -1;
   }
 
-  private setLiveUpdates(toggles: LiveStateUpdatesTogglePayload) {
-    if (toggles.deviceStats !== undefined) {
-      this.liveUpdatesAlive.device_stats = toggles.deviceStats;
-    }
-    if (toggles.appState !== undefined) {
-      this.liveUpdatesAlive.app_state = toggles.appState;
-    }
-  }
-
   private restartLiveUpdatesTimeout() {
     clearTimeout(this.liveUpdatesTimeout);
     this.liveUpdatesTimeout = setTimeout(() => {
-      this.setLiveUpdates({
-        deviceStats: false,
-        appState: false
-      });
+      this.disableAppStateLiveUpdates();
+      this.disableDeviceStatsLiveUpdates();
       this.appLogStreams.clear();
     }, LIVE_UPDATES_TIMEOUT);
   }
@@ -161,17 +149,17 @@ export class LiveUpdatesHandler {
     txId: string
   ) {
     logger.info(`Turned on live updates for ${messageType}`);
-    try {
-      while (this.continuePublishing(messageType, txId)) {
+    while (this.continuePublishing(messageType, txId)) {
+      try {
         const message = await getMessage();
         this.publisher.publishToClient(message, logger.silly);
 
         await sleep(this.getLiveUpdatesInterval(messageType));
+      } catch (e) {
+        logger.error(
+          `Error publishing live updates for ${messageType}: ${e.message}`
+        );
       }
-    } catch (e) {
-      logger.error(
-        `Error publishing live updates for ${messageType}: ${e.message}`
-      );
     }
     logger.info(`Turned off live updates for ${messageType}`);
   }
@@ -188,14 +176,24 @@ export class LiveUpdatesHandler {
     );
   }
 
-  public getDeviceStatsLiveUpdates() {
+  public getDeviceStatsLiveUpdates(): boolean {
     return this.liveUpdatesAlive.device_stats;
   }
 
-  public getAppStateLiveUpdates() {
+  public disableDeviceStatsLiveUpdates() {
+    logger.info('Disabled live updates for device_stats');
+    this.liveUpdatesAlive.device_stats = false;
+  }
+
+  public getAppStateLiveUpdates(): boolean {
     return this.liveUpdatesAlive.app_state;
   }
 
+  public disableAppStateLiveUpdates() {
+    logger.info('Disabled live updates for app_state');
+    this.liveUpdatesAlive.app_state = false;
+  }
+
   public getAppLogsLiveUpdates() {
     return this.liveUpdatesAlive.app_logs;
   }
@@ -239,12 +237,14 @@ export class LiveUpdatesHandler {
     const { deviceStats, appState, appLogs } = toggles;
     this.restartLiveUpdatesTimeout();
 
+    const promises: Promise<void>[] = [];
+
     if (deviceStats !== undefined) {
       const currentDeviceStats = this.getDeviceStatsLiveUpdates();
       this.liveUpdatesAlive.device_stats = deviceStats;
       if (deviceStats && currentDeviceStats !== true) {
         // Don't wait for this call to finish since it loops until disabled
-        void this.startPublishingLiveUpdates(
+        const deviceStatsPromise = this.startPublishingLiveUpdates(
           keyMirrors.toClientMessageType.device_stats,
           async () => {
             const payload = await getDeviceStatsPayload();
@@ -252,6 +252,7 @@ export class LiveUpdatesHandler {
           },
           txId
         );
+        promises.push(deviceStatsPromise);
       }
     }
 
@@ -260,7 +261,7 @@ export class LiveUpdatesHandler {
       this.liveUpdatesAlive.app_state = appState;
       if (appState && currentAppState !== true) {
         // Don't wait for this call to finish since it loops until disabled
-        void this.startPublishingLiveUpdates(
+        const appStatePromise = this.startPublishingLiveUpdates(
           keyMirrors.toClientMessageType.app_state,
           async () => {
             const payload = await getAppStatePayload();
@@ -268,6 +269,7 @@ export class LiveUpdatesHandler {
           },
           txId
         );
+        promises.push(appStatePromise);
       }
     }
 
@@ -275,10 +277,12 @@ export class LiveUpdatesHandler {
       const currentAppLogs = this.getAppLogsLiveUpdates();
       if (appLogs.toggle && currentAppLogs !== true) {
         // Don't wait for this call to finish since it loops until disabled
-        void this.startAppLogStream(appLogs.projectId, txId);
+        const appLogPromise = this.startAppLogStream(appLogs.projectId, txId);
+        promises.push(appLogPromise);
       } else {
         this.appLogStreams.delete(appLogs.projectId);
       }
     }
+    return promises;
   }
 }