@alwaysai/device-agent 0.1.0 → 0.1.2

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

@@ -17,10 +17,12 @@ import {
   LiveUpdatesToggleMessage,
   SignedUrlsRequestMessage,
   ClientMessage,
-  getDeviceTopic
+  getDeviceTopic,
+  AppInstallCloudResponseMessage
 } from '@alwaysai/device-agent-schemas';
 import { getDeviceUuid } from '../util/get-device-id';
 import { logger } from '../util/logger';
+import { cloudModeReady } from '../util/cloud-mode-ready';
 import { AgentConfigFile } from '../infrastructure/agent-config';
 import {
   startApp,
@@ -29,17 +31,19 @@ import {
   updateModelsWithPresignedUrls,
   installApp,
   uninstallApp,
-  updateAppCfg
+  updateAppCfg,
+  setEnv
 } from '../application-control';
-import { AppConfigUpdate, ShadowHandler, ShadowTopics } from './shadow-handler';
+import { ShadowHandler, ShadowTopics, ShadowUpdate } from './shadow-handler';
 import { Publisher } from './publisher';
 import { LiveUpdatesHandler } from './live-updates-handler';
 import { bootstrapProvision } from './bootstrap-provision';
-import { AppConfig } from '@alwaysai/app-configuration-schemas';
 import { CmdStatusManager } from './cmd-status';
 import { PassthroughHandler, runChannel } from './passthrough-handler';
 import { ALWAYSAI_ANALYTICS_PASSTHROUGH } from '../environment';
 import { getAppInstallStatusMessage } from './messages';
+import { ModelsInstallCloudResponseMessage } from '@alwaysai/device-agent-schemas/lib/schemas/client/application-action-schema';
+import sleep from '../util/sleep';
 
 export class DeviceAgentCloudConnection {
   private shadowHandler: ShadowHandler;
@@ -52,9 +56,11 @@ export class DeviceAgentCloudConnection {
   private host = getIoTCoreEndpointUrl();
   private readonly toDeviceTopic = getDeviceTopic(this.clientId);
   // FIXME: Add support for multiple simultaneous project updates
-  private newAppCfgQueue: AppConfig[] = [];
+  private appCfgUpdateQueue: ShadowUpdate[] = [];
 
-  private async handleAppStateControl(payload: AppStateControlPacket) {
+  private handleAppStateControl = async (
+    payload: AppStateControlPacket
+  ): Promise<boolean> => {
     const { baseCommand, projectId } = payload;
     switch (baseCommand) {
       case keyMirrors.appStateControl.start:
@@ -67,29 +73,27 @@ export class DeviceAgentCloudConnection {
         await restartApp({ projectId });
         break;
     }
-  }
+    return true;
+  };
 
-  private async handleAppVersionControl(payload: AppVersionControlPacket) {
+  private handleAppVersionControl = async (
+    payload: AppVersionControlPacket
+  ): Promise<boolean> => {
     switch (payload.baseCommand) {
       case keyMirrors.appVersionControl.install: {
         const { projectId, appReleaseHash } = payload;
-        await this.cmdStatusMgr.start(projectId);
-        await this.liveUpdatesHandler.enableAppInstallStatus({
-          projectId,
-          appReleaseHash
-        });
 
         const signedUrlsRequest = { projectId, appReleaseHash };
         await this.publishCloudRequest({
           messageType: keyMirrors.agentMessageType.signed_urls_request,
           signedUrlsRequest
         });
-        break;
+        return false;
       }
       case keyMirrors.appVersionControl.uninstall: {
         const { projectId } = payload;
         await this.atomicApplicationUninstall(projectId);
-        break;
+        return true;
       }
       default:
         logger.warn(
@@ -99,14 +103,13 @@ export class DeviceAgentCloudConnection {
             2
           )}`
         );
+        return true;
     }
-  }
-
-  private handleDeviceCommand = async (packet: any) => {
-    // TODO
   };
 
-  private async handleAgentCommand(message: LiveUpdatesToggleMessage) {
+  private handleAgentCommand = async (
+    message: LiveUpdatesToggleMessage
+  ): Promise<void> => {
     switch (message.messageType) {
       case keyMirrors.clientMessageType.live_state_updates:
         await this.liveUpdatesHandler.handleToggles(message.liveUpdatesToggles);
@@ -116,7 +119,63 @@ export class DeviceAgentCloudConnection {
           `Invalid agent action message type from message '${message}'`
         );
     }
-  }
+  };
+
+  private handleAppInstallCloudResponse = async (
+    payload: AppInstallCloudResponseMessage
+  ): Promise<boolean> => {
+    const {
+      projectId,
+      appReleaseHash,
+      appInstallPayload,
+      modelsInstallPayload
+    } = payload.appInstallCloudResponse;
+    const signedUrlsPayload = {
+      appInstallPayload,
+      modelsInstallPayload
+    };
+    await this.atomicApplicationUpdate(
+      installApp,
+      [{ projectId, appReleaseHash, signedUrlsPayload }],
+      projectId
+    );
+    return true;
+  };
+
+  private handleModelsInstallCloudResponse = async (
+    payload: ModelsInstallCloudResponseMessage
+  ): Promise<boolean> => {
+    const update = this.appCfgUpdateQueue.shift();
+    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.modelsInstallCloudResponse.projectId;
+    if (appCfgUpdate) {
+      await this.atomicApplicationUpdate(
+        updateModelsWithPresignedUrls,
+        [
+          {
+            projectId,
+            modelInstallPayloads: payload.modelsInstallCloudResponse.newModels,
+            newAppCfg: appCfgUpdate.newAppCfg
+          }
+        ],
+        projectId
+      );
+    }
+
+    if (envVarUpdate) {
+      await this.atomicApplicationUpdate(
+        setEnv,
+        [{ projectId, envVars: envVarUpdate.envVars }],
+        projectId
+      );
+    }
+    return true;
+  };
 
   private async publishCloudRequest(payload: SignedUrlsRequestMessage) {
     this.publisher.publishToCloud(payload);
@@ -128,45 +187,12 @@ export class DeviceAgentCloudConnection {
   }
 
   private async atomicApplicationUninstall(projectId: string) {
-    const appReleaseHash = await AgentConfigFile().getAppVersion({
-      projectId
-    });
-    await this.cmdStatusMgr.start(projectId);
-    await this.liveUpdatesHandler.enableAppInstallStatus({
-      projectId,
-      appReleaseHash
-    });
     try {
       await uninstallApp({ projectId });
       this.shadowHandler.clearAppConfig(projectId);
-
-      await this.cmdStatusMgr.stop(projectId);
-      await this.liveUpdatesHandler.disableAppInstallStatus({
-        projectId
-      });
-      // Send final status message
-      this.publisher.publishToClient(
-        await getAppInstallStatusMessage(
-          keyMirrors.appInstallStatus.success,
-          '',
-          appReleaseHash
-        )
-      );
     } catch (e) {
       logger.error(`Failed to uninstall ${projectId}: ${e.message}`);
-      const message: string = e.message;
-      await this.cmdStatusMgr.stop(projectId);
-      await this.liveUpdatesHandler.disableAppInstallStatus({
-        projectId
-      });
-      // Send final status message
-      this.publisher.publishToClient(
-        await getAppInstallStatusMessage(
-          keyMirrors.appInstallStatus.failure,
-          message,
-          appReleaseHash
-        )
-      );
+      throw e;
     }
   }
 
@@ -174,37 +200,14 @@ export class DeviceAgentCloudConnection {
   private async atomicApplicationUpdate<T extends any[], R extends any>(
     func: (...args: T) => R,
     args: T,
-    projectId: string,
-    appReleaseHash: string
+    projectId: string
   ) {
-    await this.cmdStatusMgr.start(projectId);
-    await this.liveUpdatesHandler.enableAppInstallStatus({
-      projectId,
-      appReleaseHash
-    });
-
-    // Install the app and models
     try {
       const out: R = await func(...args);
-      await this.cmdStatusMgr.stop(projectId);
-      await this.liveUpdatesHandler.disableAppInstallStatus({
-        projectId
-      });
-      // Send final status message
-      this.publisher.publishToClient(
-        await getAppInstallStatusMessage(
-          keyMirrors.appInstallStatus.success,
-          '',
-          appReleaseHash
-        )
-      );
-
-      // update app config shadow for project
-      await this.shadowHandler.publishAppConfig(projectId);
+      await this.shadowHandler.updateProjectShadow(projectId);
       return out;
     } catch (e) {
-      logger.error(`Failed to install ${projectId}: ${e.message}`);
-      const message: string = e.message;
+      logger.error(`Failed to update ${projectId}:\n${e.message}\n${e.stack}`);
 
       // uninstall the failed app to put system back in good state
       // TODO: Replace this with rollback
@@ -212,7 +215,38 @@ export class DeviceAgentCloudConnection {
         await uninstallApp({ projectId });
       } finally {
         this.shadowHandler.clearAppConfig(projectId);
+      }
+      throw e;
+    }
+  }
 
+  // eslint-disable-next-line
+  private async atomicCmd<T extends any[]>(props: {
+    func: (...args: T) => Promise<boolean>;
+    args: T;
+    projectId: string;
+    appReleaseHash?: string;
+  }) {
+    const { func, args, projectId } = props;
+    let appReleaseHash = props.appReleaseHash;
+    if (appReleaseHash === undefined) {
+      try {
+        appReleaseHash = await AgentConfigFile().getAppVersion({
+          projectId
+        });
+      } catch (e) {
+        logger.warn(`Unable to get app release hash for ${projectId}!`);
+        appReleaseHash = '';
+      }
+    }
+    try {
+      await this.cmdStatusMgr.start(projectId);
+      await this.liveUpdatesHandler.enableAppInstallStatus({
+        projectId,
+        appReleaseHash
+      });
+      const completed = await func(...args);
+      if (completed) {
         await this.cmdStatusMgr.stop(projectId);
         await this.liveUpdatesHandler.disableAppInstallStatus({
           projectId
@@ -220,30 +254,48 @@ export class DeviceAgentCloudConnection {
         // Send final status message
         this.publisher.publishToClient(
           await getAppInstallStatusMessage(
-            keyMirrors.appInstallStatus.failure,
-            message,
+            keyMirrors.appInstallStatus.success,
+            '',
             appReleaseHash
           )
         );
       }
-    }
-  }
+      return completed;
+    } catch (e) {
+      logger.error(
+        `Failed to execute cmd for ${projectId}:\n${e.message}\n${e.stack}`
+      );
+      const message: string = e.message;
 
-  private async handleAppConfigUpdates(appCfgUpdates: AppConfigUpdate[]) {
-    for (const appConfigUpdate of appCfgUpdates) {
-      const { projectId, newAppCfg, updatedModels } = appConfigUpdate;
-      const appReleaseHash = await AgentConfigFile().getAppVersion({
+      // uninstall the failed app to put system back in good state
+      await this.cmdStatusMgr.stop(projectId);
+      await this.liveUpdatesHandler.disableAppInstallStatus({
         projectId
       });
-      await this.cmdStatusMgr.start(projectId);
-      await this.liveUpdatesHandler.enableAppInstallStatus({
-        projectId,
-        appReleaseHash
-      });
+      // Send final status message
+      this.publisher.publishToClient(
+        await getAppInstallStatusMessage(
+          keyMirrors.appInstallStatus.failure,
+          message,
+          appReleaseHash
+        )
+      );
+    }
+  }
 
-      if (updatedModels && Object.keys(updatedModels).length) {
-        // Publish request for model urls
-        this.newAppCfgQueue.push(newAppCfg);
+  private handleAppConfigUpdates = async (
+    updates: ShadowUpdate[]
+  ): Promise<boolean> => {
+    for (const update of updates) {
+      const { projectId, appCfgUpdate, envVarUpdate } = update;
+
+      if (
+        appCfgUpdate &&
+        appCfgUpdate.updatedModels &&
+        Object.keys(appCfgUpdate.updatedModels).length
+      ) {
+        // When there are model updates request signed URLs and wait to apply config changes
+        const { updatedModels } = appCfgUpdate;
 
         logger.debug(
           `Requesting presigned urls from cloud for model versions: ${JSON.stringify(
@@ -257,28 +309,34 @@ export class DeviceAgentCloudConnection {
             models: updatedModels
           }
         });
-      } else {
-        // FIXME: do we need to send this up to the cloud?
-        // should it be something other than appReleaseHash?
-        const appReleaseHash = await AgentConfigFile().getAppVersion({
-          projectId
-        });
 
+        this.appCfgUpdateQueue.push(update);
+        return false;
+      }
+
+      if (appCfgUpdate) {
         await this.atomicApplicationUpdate(
           updateAppCfg,
           [
             {
               projectId,
-              appReleaseHash,
-              newAppCfg
+              newAppCfg: appCfgUpdate.newAppCfg
             }
           ],
-          projectId,
-          appReleaseHash
+          projectId
+        );
+      }
+
+      if (envVarUpdate) {
+        await this.atomicApplicationUpdate(
+          setEnv,
+          [{ projectId, envVars: envVarUpdate.envVars }],
+          projectId
         );
       }
     }
-  }
+    return true;
+  };
 
   /*=================================================================
                             Public interface
@@ -323,6 +381,10 @@ export class DeviceAgentCloudConnection {
     return this.cmdStatusMgr.isCmdInProgress(projectId);
   }
 
+  public async updateProjectShadow(projectId: string) {
+    await this.shadowHandler.updateProjectShadow(projectId);
+  }
+
   public async handleClientMessage({
     topic,
     message
@@ -339,70 +401,61 @@ export class DeviceAgentCloudConnection {
           2
         )}`
       );
+      // TODO: Send generic error response
       return;
     }
     const payload = message.payload;
     switch (payload.messageType) {
       case keyMirrors.clientMessageType.app_state_control: {
-        await this.handleAppStateControl(payload.appStateControl);
+        const projectId = payload.appStateControl.projectId;
+        await this.atomicCmd({
+          func: this.handleAppStateControl,
+          args: [payload.appStateControl],
+          projectId
+        });
         break;
       }
       case keyMirrors.clientMessageType.app_version_control: {
-        await this.handleAppVersionControl(payload.appVersionControl);
+        const projectId = payload.appVersionControl.projectId;
+        const appReleaseHash =
+          payload.appVersionControl.baseCommand ===
+          keyMirrors.appVersionControl.install
+            ? payload.appVersionControl.appReleaseHash
+            : undefined;
+        await this.atomicCmd({
+          func: this.handleAppVersionControl,
+          args: [payload.appVersionControl],
+          projectId,
+          appReleaseHash
+        });
         break;
       }
       case keyMirrors.clientMessageType.live_state_updates: {
+        // TODO: Send response?
         await this.handleAgentCommand(payload);
         break;
       }
       case keyMirrors.clientMessageType.app_install_cloud_response: {
-        const {
-          projectId,
-          appReleaseHash,
-          appInstallPayload,
-          modelsInstallPayload
-        } = payload.appInstallCloudResponse;
-        const signedUrlsPayload = {
-          appInstallPayload,
-          modelsInstallPayload
-        };
-        await this.atomicApplicationUpdate(
-          installApp,
-          [{ projectId, appReleaseHash, signedUrlsPayload }],
+        const { projectId, appReleaseHash } = payload.appInstallCloudResponse;
+        await this.atomicCmd({
+          func: this.handleAppInstallCloudResponse,
+          args: [payload],
           projectId,
           appReleaseHash
-        );
+        });
 
         break;
       }
       case keyMirrors.clientMessageType.models_install_cloud_response: {
-        const { projectId, newModels } = payload.modelsInstallCloudResponse;
-        const appReleaseHash = await AgentConfigFile().getAppVersion({
+        // This message doesn't have appReleaseHash in it's payload, but
+        // atomicCmd should be able to read it from the installed app
+        const { projectId } = payload.modelsInstallCloudResponse;
+        await this.atomicCmd({
+          func: this.handleModelsInstallCloudResponse,
+          args: [payload],
           projectId
         });
 
-        const newAppCfg = this.newAppCfgQueue.shift();
-        if (newAppCfg === undefined) {
-          logger.error(
-            'Unknown error while updating models via application config! No config present for model update.'
-          );
-          return;
-        }
-
-        await this.atomicApplicationUpdate(
-          updateModelsWithPresignedUrls,
-          [
-            {
-              projectId,
-              modelInstallPayloads: newModels,
-              newAppCfg,
-              appReleaseHash
-            }
-          ],
-          projectId,
-          appReleaseHash
-        );
-
         break;
       }
       default:
@@ -423,12 +476,25 @@ export class DeviceAgentCloudConnection {
     switch (topic) {
       case this.shadowHandler.shadowTopics.projects.getAccepted:
       case this.shadowHandler.shadowTopics.projects.updateDelta: {
-        const appConfigUpdates = await this.shadowHandler.handleShadowTopic({
+        const shadowUpdates = await this.shadowHandler.handleShadowTopic({
           topic,
           payload: message.state,
           clientToken: message.clientToken
         });
-        await this.handleAppConfigUpdates(appConfigUpdates);
+        if (shadowUpdates.length) {
+          // FIXME: Take project ID of first shadow update. Most likely there will only be one update
+          // so this should be sufficient for now.
+          const projectId = shadowUpdates[0].projectId;
+          try {
+            await this.atomicCmd({
+              func: this.handleAppConfigUpdates,
+              args: [shadowUpdates],
+              projectId
+            });
+          } catch (e) {
+            logger.error(`Error handling shadow message: ${e.message}`);
+          }
+        }
         break;
       }
       case this.shadowHandler.shadowTopics.projects.getRejected:
@@ -483,14 +549,17 @@ export class DeviceAgentCloudConnection {
     });
   }
 
-  public stop() {
+  public async stop() {
+    // FIXME: This method is currently only used by the CLI, and shadow messages
+    // can be lost since we aren't waiting for responses so sleep for a short
+    // time to receive them
+    await sleep(1000);
     this.device.end();
   }
 }
 
 export async function runDeviceAgentCloudInterface() {
-  // FIXME: Check for KeyPath as well
-  if (existsSync(DEVICE_CERTIFICATE_FILE_PATH)) {
+  if (cloudModeReady()) {
     const deviceAgent = new DeviceAgentCloudConnection();
     await deviceAgent.setupHandlers();
     if (ALWAYSAI_ANALYTICS_PASSTHROUGH === true) {

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

@@ -0,0 +1,68 @@
+import { LiveUpdatesHandler } from './live-updates-handler';
+import { Publisher } from './publisher';
+
+const testTrueToggles = {
+  deviceStats: true,
+  appState: true
+};
+
+const testFalseToggles = {
+  deviceStats: false,
+  appState: false
+};
+
+const mockClient = jest.fn();
+const clientId = 'test-client';
+
+jest.spyOn(global, 'setTimeout');
+
+// 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(() => {
+    publisher = new Publisher(mockClient, clientId);
+    liveUpdatesHandler = new LiveUpdatesHandler(publisher);
+    jest.clearAllMocks();
+  });
+
+  test('ignore subsequent enables', async () => {
+    void liveUpdatesHandler.handleToggles(testTrueToggles);
+    // called twice, once for device stats, once for app state
+    expect(mockStartPublishingLiveUpdates).toBeCalledTimes(2);
+    // restartLiveUpdatesTimeout is always called once when handleToggles is called
+    expect(jest.mocked(setTimeout)).toBeCalledTimes(1);
+
+    // Second call -> should not call startPublishingLiveUpdates should not be called
+    jest.clearAllMocks();
+    void liveUpdatesHandler.handleToggles(testTrueToggles);
+    expect(mockStartPublishingLiveUpdates).toBeCalledTimes(0);
+    expect(jest.mocked(setTimeout)).toHaveBeenCalledTimes(1);
+  });
+
+  test('test disable live updates', async () => {
+    // Test calling handleToggles one time, enabling it
+    void liveUpdatesHandler.handleToggles(testTrueToggles);
+    expect(mockStartPublishingLiveUpdates).toBeCalledTimes(2);
+    expect(jest.mocked(setTimeout)).toHaveBeenCalledTimes(1);
+    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);
+    expect(mockStartPublishingLiveUpdates).toBeCalledTimes(0);
+    expect(jest.mocked(setTimeout)).toHaveBeenCalledTimes(1);
+    expect(liveUpdatesHandler.getDeviceStatsLiveUpdates()).toBe(false);
+    expect(liveUpdatesHandler.getAppStateLiveUpdates()).toBe(false);
+    expect(liveUpdatesHandler.getAppLogsLiveUpdates()).toBe(false);
+  });
+});