@alwaysai/device-agent 0.2.0 → 1.0.0

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

@@ -11,14 +11,17 @@ import {
 } from '../util/directories';
 import {
   keyMirrors,
-  validateClientMessage,
-  AppStateControlPacket,
-  AppVersionControlPacket,
-  LiveUpdatesToggleMessage,
-  SignedUrlsRequestMessage,
-  ClientMessage,
-  getDeviceTopic,
-  AppInstallCloudResponseMessage
+  validateToClientMessage,
+  SignedUrlsRequestPayload,
+  getToDeviceTopic,
+  AppInstallResponsePayload,
+  validateToDeviceAgentMessage,
+  ToDeviceAgentMessage,
+  ToCloudMessage,
+  AppStateControlPayload,
+  AppVersionControlInstallPayload,
+  AppVersionControlUninstallPayload,
+  ToClientMessage
 } from '@alwaysai/device-agent-schemas';
 import { getDeviceUuid } from '../util/get-device-id';
 import { logger } from '../util/logger';
@@ -42,26 +45,34 @@ import { bootstrapProvision } from './bootstrap-provision';
 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 { getStatusResponsePayload } from './messages';
+import { ModelsInstallResponsePayload } from '@alwaysai/device-agent-schemas';
 import sleep from '../util/sleep';
+import { createAppBackup, rollbackApp } from '../application-control/backup';
+import { TransactionQueue } from './transaction-queue';
+import {
+  buildSignedUrlsRequestMessage,
+  buildStatusResponseMessage
+} from './message-builder';
+import { generateTxId } from '@alwaysai/device-agent-schemas';
 
 export class DeviceAgentCloudConnection {
   private shadowHandler: ShadowHandler;
   public publisher: Publisher;
   private cmdStatusMgr: CmdStatusManager;
   private liveUpdatesHandler: LiveUpdatesHandler;
+  private txnQueue: TransactionQueue;
   private device = awsIot.device;
 
   private clientId = getDeviceUuid();
   private host = getIoTCoreEndpointUrl();
-  private readonly toDeviceTopic = getDeviceTopic(this.clientId);
+  private readonly toDeviceTopic = getToDeviceTopic(this.clientId);
   private readonly secureTunnelNotifyTopic = `$aws/things/${this.clientId}/tunnels/notify`;
   // FIXME: Add support for multiple simultaneous project updates
   private appCfgUpdateQueue: ShadowUpdate[] = [];
 
   private handleAppStateControl = async (
-    payload: AppStateControlPacket
+    payload: AppStateControlPayload
   ): Promise<boolean> => {
     const { baseCommand, projectId } = payload;
     switch (baseCommand) {
@@ -79,17 +90,26 @@ export class DeviceAgentCloudConnection {
   };
 
   private handleAppVersionControl = async (
-    payload: AppVersionControlPacket
+    payload:
+      | AppVersionControlInstallPayload
+      | AppVersionControlUninstallPayload,
+    txId: string
   ): Promise<boolean> => {
     switch (payload.baseCommand) {
       case keyMirrors.appVersionControl.install: {
         const { projectId, appReleaseHash } = payload;
 
-        const signedUrlsRequest = { projectId, appReleaseHash };
-        await this.publishCloudRequest({
-          messageType: keyMirrors.agentMessageType.signed_urls_request,
-          signedUrlsRequest
-        });
+        const signedUrlsRequestPayload: SignedUrlsRequestPayload = {
+          signedUrlsRequest: {
+            projectId,
+            appReleaseHash
+          }
+        };
+        const message = await buildSignedUrlsRequestMessage(
+          signedUrlsRequestPayload,
+          txId
+        );
+        await this.publishCloudRequest(message);
         return false;
       }
       case keyMirrors.appVersionControl.uninstall: {
@@ -109,29 +129,15 @@ export class DeviceAgentCloudConnection {
     }
   };
 
-  private handleAgentCommand = async (
-    message: LiveUpdatesToggleMessage
-  ): Promise<void> => {
-    switch (message.messageType) {
-      case keyMirrors.clientMessageType.live_state_updates:
-        await this.liveUpdatesHandler.handleToggles(message.liveUpdatesToggles);
-        break;
-      default:
-        logger.error(
-          `Invalid agent action message type from message '${message}'`
-        );
-    }
-  };
-
-  private handleAppInstallCloudResponse = async (
-    payload: AppInstallCloudResponseMessage
+  private handleAppInstallCloudResponsePayload = async (
+    payload: AppInstallResponsePayload
   ): Promise<boolean> => {
     const {
       projectId,
       appReleaseHash,
       appInstallPayload,
       modelsInstallPayload
-    } = payload.appInstallCloudResponse;
+    } = payload.appInstallResponse;
     const signedUrlsPayload = {
       appInstallPayload,
       modelsInstallPayload
@@ -144,8 +150,8 @@ export class DeviceAgentCloudConnection {
     return true;
   };
 
-  private handleModelsInstallCloudResponse = async (
-    payload: ModelsInstallCloudResponseMessage
+  private handleModelsInstallCloudResponsePayload = async (
+    payload: ModelsInstallResponsePayload
   ): Promise<boolean> => {
     const update = this.appCfgUpdateQueue.shift();
     if (update === undefined) {
@@ -154,14 +160,14 @@ export class DeviceAgentCloudConnection {
       );
     }
     const { appCfgUpdate, envVarUpdate } = update;
-    const projectId = payload.modelsInstallCloudResponse.projectId;
+    const projectId = payload.modelsInstallResponse.projectId;
     if (appCfgUpdate) {
       await this.atomicApplicationUpdate(
         updateModelsWithPresignedUrls,
         [
           {
             projectId,
-            modelInstallPayloads: payload.modelsInstallCloudResponse.newModels,
+            modelInstallPayloads: payload.modelsInstallResponse.newModels,
             newAppCfg: appCfgUpdate.newAppCfg
           }
         ],
@@ -179,8 +185,8 @@ export class DeviceAgentCloudConnection {
     return true;
   };
 
-  private async publishCloudRequest(payload: SignedUrlsRequestMessage) {
-    this.publisher.publishToCloud(payload);
+  private async publishCloudRequest(message: ToCloudMessage) {
+    this.publisher.publishToCloud(message);
   }
 
   private subscribe(topic: string) {
@@ -204,21 +210,48 @@ export class DeviceAgentCloudConnection {
     args: T,
     projectId: string
   ) {
+    // First try to create a backup, so that there is one available if something goes wrong in the next try:catch.
+    if (await AgentConfigFile().isAppPresent({ projectId })) {
+      try {
+        await createAppBackup({ projectId });
+      } catch (e) {
+        logger.error(
+          `Could not create a backup for the project: ${projectId}:\n${e.message}\n${e.stack}`
+        );
+      }
+    }
+
     try {
       const out: R = await func(...args);
       await this.shadowHandler.updateProjectShadow(projectId);
       return out;
-    } catch (e) {
-      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
+    } catch (errorAppUpdate) {
+      logger.error(
+        `Failed to update ${projectId}:\n${JSON.stringify(errorAppUpdate)}}`
+      );
+      // If something goes wrong, first try to rollback
       try {
-        await uninstallApp({ projectId });
-      } finally {
-        this.shadowHandler.clearAppConfig(projectId);
+        await rollbackApp({ projectId });
+        logger.error(
+          `Application update failed, rolled back to previous version: ${errorAppUpdate}`
+        );
+      } catch (errorRollbackApp) {
+        // and if that fails, uninstall the app as a last resort.
+        try {
+          await uninstallApp({ projectId });
+        } finally {
+          this.shadowHandler.clearAppConfig(projectId);
+        }
+        logger.error(
+          `Application update failed, rolled back to previous version: ${errorAppUpdate}`
+        );
+        throw new Error(
+          `Application update and rollback failed, uninstalled the application: ${errorAppUpdate}`
+        );
       }
-      throw e;
+      throw new Error(
+        `Application update failed, rolled the application back: ${errorAppUpdate}`
+      );
     }
   }
 
@@ -227,40 +260,30 @@ export class DeviceAgentCloudConnection {
     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 = '';
-      }
-    }
+    txId: string;
+  }): Promise<boolean> {
+    const { func, args, projectId, txId } = props;
     try {
       await this.cmdStatusMgr.start(projectId);
-      await this.liveUpdatesHandler.enableAppInstallStatus({
-        projectId,
-        appReleaseHash
+      await this.liveUpdatesHandler.enableTransactionStatus({
+        txId
       });
       const completed = await func(...args);
       if (completed) {
         await this.cmdStatusMgr.stop(projectId);
-        await this.liveUpdatesHandler.disableAppInstallStatus({
-          projectId
+        await this.liveUpdatesHandler.disableTransactionStatus({
+          txId
         });
+        const sucessStatusResponsePayload = await getStatusResponsePayload(
+          keyMirrors.statusResponse.success,
+          ''
+        );
         // Send final status message
-        this.publisher.publishToClient(
-          await getAppInstallStatusMessage(
-            keyMirrors.appInstallStatus.success,
-            '',
-            appReleaseHash
-          )
+        const message = await buildStatusResponseMessage(
+          sucessStatusResponsePayload,
+          txId
         );
+        this.publisher.publishToClient(message);
       }
       return completed;
     } catch (e) {
@@ -271,22 +294,26 @@ export class DeviceAgentCloudConnection {
 
       // uninstall the failed app to put system back in good state
       await this.cmdStatusMgr.stop(projectId);
-      await this.liveUpdatesHandler.disableAppInstallStatus({
-        projectId
+      await this.liveUpdatesHandler.disableTransactionStatus({
+        txId
       });
+      const failureStatusResponsePayload = await getStatusResponsePayload(
+        keyMirrors.statusResponse.failure,
+        message
+      );
       // Send final status message
-      this.publisher.publishToClient(
-        await getAppInstallStatusMessage(
-          keyMirrors.appInstallStatus.failure,
-          message,
-          appReleaseHash
-        )
+      const failureStatusResponseMessage = await buildStatusResponseMessage(
+        failureStatusResponsePayload,
+        txId
       );
+      this.publisher.publishToClient(failureStatusResponseMessage);
+      return true;
     }
   }
 
   private handleAppConfigUpdates = async (
-    updates: ShadowUpdate[]
+    updates: ShadowUpdate[],
+    txId: string
   ): Promise<boolean> => {
     for (const update of updates) {
       const { projectId, appCfgUpdate, envVarUpdate } = update;
@@ -304,13 +331,17 @@ export class DeviceAgentCloudConnection {
             updatedModels
           )}`
         );
-        this.publisher.publishToCloud({
-          messageType: keyMirrors.agentMessageType.signed_urls_request,
+        const modelsOnlyUrlsRequestPayload: SignedUrlsRequestPayload = {
           modelsOnlyUrlsRequest: {
             projectId,
             models: updatedModels
           }
-        });
+        };
+        const message = await buildSignedUrlsRequestMessage(
+          modelsOnlyUrlsRequestPayload,
+          txId
+        );
+        this.publisher.publishToCloud(message);
 
         this.appCfgUpdateQueue.push(update);
         return false;
@@ -357,7 +388,11 @@ export class DeviceAgentCloudConnection {
     this.publisher = new Publisher(this.device, this.clientId);
     this.shadowHandler = new ShadowHandler(this.clientId, this.publisher);
     this.cmdStatusMgr = new CmdStatusManager();
-    this.liveUpdatesHandler = new LiveUpdatesHandler(this.publisher);
+    this.liveUpdatesHandler = new LiveUpdatesHandler(
+      this.publisher,
+      this.clientId
+    );
+    this.txnQueue = new TransactionQueue();
 
     this.subscribe(this.toDeviceTopic);
     this.subscribe(this.secureTunnelNotifyTopic);
@@ -388,18 +423,18 @@ export class DeviceAgentCloudConnection {
     await this.shadowHandler.updateProjectShadow(projectId);
   }
 
-  public async handleClientMessage({
+  public async handleDeviceAgentMessage({
     topic,
     message
   }: {
     topic: string;
-    message: ClientMessage;
+    message: ToDeviceAgentMessage;
   }) {
-    const valid = validateClientMessage(message);
+    const valid = validateToDeviceAgentMessage(message);
     if (!valid) {
       logger.error(
         `Error validating message: ${JSON.stringify(
-          { topic, message, errors: validateClientMessage.errors },
+          { topic, message, errors: validateToDeviceAgentMessage.errors },
           null,
           2
         )}`
@@ -407,64 +442,109 @@ export class DeviceAgentCloudConnection {
       // TODO: Send generic error response
       return;
     }
-    const payload = message.payload;
-    switch (payload.messageType) {
-      case keyMirrors.clientMessageType.app_state_control: {
-        const projectId = payload.appStateControl.projectId;
-        await this.atomicCmd({
-          func: this.handleAppStateControl,
-          args: [payload.appStateControl],
-          projectId
-        });
+    const txId = message.txId;
+    switch (message.messageType) {
+      case keyMirrors.toDeviceAgentMessageType.app_state_control: {
+        // txId sent from cloud, just need to continue it
+        const payload = message.payload;
+        const projectId = payload.projectId;
+
+        try {
+          this.txnQueue.addTxIdToQueue(txId, projectId);
+          const completed = await this.atomicCmd({
+            func: this.handleAppStateControl,
+            args: [message.payload],
+            projectId,
+            txId
+          });
+          if (completed) {
+            this.txnQueue.completeTxn(txId);
+          }
+        } catch (e) {
+          logger.error(
+            `Error processing application state control request: ${e}!`
+          );
+        }
+
         break;
       }
-      case keyMirrors.clientMessageType.app_version_control: {
-        const projectId = payload.appVersionControl.projectId;
+      case keyMirrors.toDeviceAgentMessageType.app_version_control: {
+        // txId sent from cloud, just need to continue it
+        const payload = message.payload;
+        const projectId = payload.projectId;
         const appReleaseHash =
-          payload.appVersionControl.baseCommand ===
-          keyMirrors.appVersionControl.install
-            ? payload.appVersionControl.appReleaseHash
+          payload.baseCommand === keyMirrors.appVersionControl.install
+            ? payload.appReleaseHash
             : undefined;
-        await this.atomicCmd({
-          func: this.handleAppVersionControl,
-          args: [payload.appVersionControl],
-          projectId,
-          appReleaseHash
-        });
+        try {
+          this.txnQueue.addTxIdToQueue(txId, projectId);
+          const completed = await this.atomicCmd({
+            func: this.handleAppVersionControl,
+            args: [payload, txId],
+            projectId,
+            txId
+          });
+          if (completed) {
+            this.txnQueue.completeTxn(txId);
+          }
+        } catch (e) {
+          logger.error(`Error processing application install request: ${e}!`);
+        }
+
         break;
       }
-      case keyMirrors.clientMessageType.live_state_updates: {
+      case keyMirrors.toDeviceAgentMessageType.live_state_updates: {
+        const payload = message.payload;
         // TODO: Send response?
-        await this.handleAgentCommand(payload);
+        await this.liveUpdatesHandler.handleToggles(payload, txId);
         break;
       }
-      case keyMirrors.clientMessageType.app_install_cloud_response: {
-        const { projectId, appReleaseHash } = payload.appInstallCloudResponse;
-        await this.atomicCmd({
-          func: this.handleAppInstallCloudResponse,
+      case keyMirrors.toDeviceAgentMessageType.app_install_response: {
+        const payload = message.payload;
+        const { projectId, appReleaseHash } = payload.appInstallResponse;
+        if (txId !== this.txnQueue.getCurrentTxId()) {
+          throw new Error(
+            `App install response received a message for a transaction ID ${txId} that is not currently underway (${this.txnQueue.getCurrentTxId()})!`
+          );
+        }
+        const completed = await this.atomicCmd({
+          func: this.handleAppInstallCloudResponsePayload,
           args: [payload],
           projectId,
-          appReleaseHash
+          txId
         });
+        if (completed) {
+          this.txnQueue.completeTxn(txId);
+        }
 
         break;
       }
-      case keyMirrors.clientMessageType.models_install_cloud_response: {
+      case keyMirrors.toDeviceAgentMessageType.models_install_response: {
         // 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,
+        const payload = message.payload;
+        const { projectId } = payload.modelsInstallResponse;
+        if (txId !== this.txnQueue.getCurrentTxId()) {
+          throw new Error(
+            `Model install response received a message for a transaction ID ${txId} that is not currently underway (${this.txnQueue.getCurrentTxId()})!`
+          );
+        }
+        const completed = await this.atomicCmd({
+          func: this.handleModelsInstallCloudResponsePayload,
           args: [payload],
-          projectId
+          projectId,
+          txId
         });
+        if (completed) {
+          this.txnQueue.completeTxn(txId);
+        }
 
         break;
       }
       default:
         logger.error(
           `Invalid client message: '${JSON.stringify(
-            { topic, message },
+            { topic, message, txId },
             null,
             2
           )}'`
@@ -472,10 +552,14 @@ export class DeviceAgentCloudConnection {
     }
   }
 
-  public async handleMessage(topic: string, message: ClientMessage | any) {
+  public async handleMessage(
+    topic: string,
+    message: ToDeviceAgentMessage | any
+  ) {
     logger.debug(
       `Received message: ${JSON.stringify({ topic, message }, null, 2)}`
     );
+    const txId = message.txId || generateTxId(); // txId should be passed in for DeviceAgentMessage messages but shadow updates won't have this
     switch (topic) {
       case this.shadowHandler.shadowTopics.projects.getAccepted:
       case this.shadowHandler.shadowTopics.projects.updateDelta: {
@@ -489,11 +573,20 @@ export class DeviceAgentCloudConnection {
           // so this should be sufficient for now.
           const projectId = shadowUpdates[0].projectId;
           try {
-            await this.atomicCmd({
-              func: this.handleAppConfigUpdates,
-              args: [shadowUpdates],
-              projectId
-            });
+            try {
+              this.txnQueue.addTxIdToQueue(txId, projectId);
+              const completed = await this.atomicCmd({
+                func: this.handleAppConfigUpdates,
+                args: [shadowUpdates, txId],
+                projectId,
+                txId
+              });
+              if (completed) {
+                this.txnQueue.completeTxn(txId);
+              }
+            } catch (e) {
+              logger.error(`Error processing model update request: ${e}!`);
+            }
           } catch (e) {
             logger.error(`Error handling shadow message: ${e.message}`);
           }
@@ -506,7 +599,7 @@ export class DeviceAgentCloudConnection {
         // Not handling these for now
         break;
       case this.toDeviceTopic:
-        await this.handleClientMessage({
+        await this.handleDeviceAgentMessage({
           topic,
           message
         });

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

@@ -13,6 +13,7 @@ const testFalseToggles = {
 
 const mockClient = jest.fn();
 const clientId = 'test-client';
+const emptyTxId = '';
 
 jest.spyOn(global, 'setTimeout');
 
@@ -29,12 +30,12 @@ describe('Test Live Updates Handler', () => {
 
   beforeEach(() => {
     publisher = new Publisher(mockClient, clientId);
-    liveUpdatesHandler = new LiveUpdatesHandler(publisher);
+    liveUpdatesHandler = new LiveUpdatesHandler(publisher, clientId);
     jest.clearAllMocks();
   });
 
   test('ignore subsequent enables', async () => {
-    void liveUpdatesHandler.handleToggles(testTrueToggles);
+    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
@@ -42,14 +43,14 @@ describe('Test Live Updates Handler', () => {
 
     // Second call -> should not call startPublishingLiveUpdates should not be called
     jest.clearAllMocks();
-    void liveUpdatesHandler.handleToggles(testTrueToggles);
+    void liveUpdatesHandler.handleToggles(testTrueToggles, emptyTxId);
     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);
+    void liveUpdatesHandler.handleToggles(testTrueToggles, emptyTxId);
     expect(mockStartPublishingLiveUpdates).toBeCalledTimes(2);
     expect(jest.mocked(setTimeout)).toHaveBeenCalledTimes(1);
     expect(liveUpdatesHandler.getDeviceStatsLiveUpdates()).toBe(true);
@@ -58,7 +59,7 @@ describe('Test Live Updates Handler', () => {
 
     // Check to see that attributes are properly set to false when disabled
     jest.clearAllMocks();
-    void liveUpdatesHandler.handleToggles(testFalseToggles);
+    void liveUpdatesHandler.handleToggles(testFalseToggles, emptyTxId);
     expect(mockStartPublishingLiveUpdates).toBeCalledTimes(0);
     expect(jest.mocked(setTimeout)).toHaveBeenCalledTimes(1);
     expect(liveUpdatesHandler.getDeviceStatsLiveUpdates()).toBe(false);